OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

Add DISABLE_JUSTKNOBS to torch/_utils_internal.py and use it for dynamo _maybe_set_eval_frame (#162298)

Browse Source 
If JustKnobs is disabled (as it always is in OSS), we can easily avoid an extra layer of Python function call.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/162298
Approved by: https://github.com/ezyang
This commit is contained in:
Scott Wolchok 2025-09-11 16:51:51 -07:00 committed by PyTorch MergeBot
parent 090e6838a0
commit 6b608dfe81
2 changed files with 18 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
torch/_dynamo/eval_frame.py
View File

@ -67,7 +67,7 @@ from torch._dispatch.python import enable_python_dispatcher
from torch._dynamo.types import ConvertFrameReturn, FrameAction, FrameExecStrategy
from torch._export.utils import _compiling_state_context
from torch._subclasses.fake_tensor import unset_fake_temporarily
from torch._utils_internal import justknobs_check, log_export_usage
from torch._utils_internal import DISABLE_JUSTKNOBS, justknobs_check, log_export_usage
from torch.export.dynamic_shapes import (
_combine_args,
_DimHint,
@ -145,16 +145,20 @@ cached_backends: dict[int, CompilerFn] = {}
unset = Unset.token
def _maybe_set_eval_frame(callback: DynamoCallback) -> DynamoCallback:
# A wrapper on set_eval_frame that is guarded by a Justknob.
# Users can disable torchDynamo by setting the JK to False.
if not justknobs_check("pytorch/compiler:enable_compiler_set_eval_frame"):
torch._dynamo.utils.warn_once(
"Dynamo disabled by Justknob: enable_compiler_set_eval_frame, skipping set_eval_frame"
)
return callback
else:
return set_eval_frame(callback)
if DISABLE_JUSTKNOBS:
_maybe_set_eval_frame = set_eval_frame
else:
def _maybe_set_eval_frame(callback: DynamoCallback) -> DynamoCallback:
# A wrapper on set_eval_frame that is guarded by a Justknob.
# Users can disable torchDynamo by setting the JK to False.
if not justknobs_check("pytorch/compiler:enable_compiler_set_eval_frame"):
torch._dynamo.utils.warn_once(
"Dynamo disabled by Justknob: enable_compiler_set_eval_frame, skipping set_eval_frame"
)
return callback
else:
return set_eval_frame(callback)
@dataclass

3
torch/_utils_internal.py
View File

@ -176,6 +176,9 @@ def log_torch_jit_trace_exportability(
return
DISABLE_JUSTKNOBS = True
def justknobs_check(name: str, default: bool = True) -> bool:
"""
This function can be used to killswitch functionality in FB prod,