mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
9db7bcb3fe
This PR:
* Expands `Hooks` with a new, optional `frame_traced_fn` field. It should be a callable receiving the list of traced code objects
* Maintains a list of `traced_code` objects in the `TracingContext` of an `OutputGraph`
* Whenever an `inline_call()` is encountered, the corresponding code object is added to this set
* `OutputGraph`'s associated `f_code` is added to the list just before the hook is called
I believe use of this hook should enable the source code hashing that vLLM does in a better way than monkey-patching `inline_call()`.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/153622
Approved by: https://github.com/jansel
67 lines
1.5 KiB
Python
67 lines
1.5 KiB
Python
# Owner(s): ["module: dynamo"]
|
|
import importlib
|
|
import os
|
|
import sys
|
|
import types
|
|
from contextlib import contextmanager
|
|
|
|
import torch
|
|
import torch._dynamo
|
|
|
|
|
|
g_tensor_export = torch.ones(10)
|
|
|
|
|
|
tensor_for_import_testing = torch.ones(10, 10)
|
|
|
|
|
|
def inner_func():
|
|
return torch.is_grad_enabled()
|
|
|
|
|
|
def outer_func(func):
|
|
def wrapped(*args):
|
|
a = func(*args)
|
|
torch._dynamo.graph_break()
|
|
return torch.sin(a + 1), inner_func()
|
|
|
|
return wrapped
|
|
|
|
|
|
# Create a dummy python module and function to test skipfiles rules.
|
|
module_code = """
|
|
def add(x):
|
|
return x + 1
|
|
"""
|
|
|
|
|
|
def add(x):
|
|
return x + 1
|
|
|
|
|
|
def break_it(x):
|
|
return x.sum().item()
|
|
|
|
|
|
def create_dummy_module_and_function():
|
|
module = types.ModuleType("dummy_module")
|
|
module.__spec__ = importlib.machinery.ModuleSpec(
|
|
"dummy_module", None, origin=os.path.abspath(__file__)
|
|
)
|
|
exec(module_code, module.__dict__)
|
|
sys.modules["dummy_module"] = module
|
|
# Need to override the original function since its __code__.co_filename is not a regular python file name,
|
|
# and the skipfiles rules use filename when checking SKIP_DIRS.
|
|
module.add = add
|
|
return module, module.add
|
|
|
|
|
|
@contextmanager
|
|
def install_guard_manager_testing_hook(hook_fn):
|
|
old_value = torch._dynamo.guards.guard_manager_testing_hook_fn
|
|
try:
|
|
torch._dynamo.guards.guard_manager_testing_hook_fn = hook_fn
|
|
yield
|
|
finally:
|
|
torch._dynamo.guards.guard_manager_testing_hook_fn = old_value
|