From 33daaad7d0c88a38ee1594869957b8266e39f464 Mon Sep 17 00:00:00 2001 From: clr Date: Wed, 17 Sep 2025 10:00:10 -0700 Subject: [PATCH] dynamo: Handle objects in graph that do not support weakref (#163168) We are seeing crashes of the form ``` Traceback (most recent call last): File "/packages/aps_ads_vm/launcher_multiapp-inplace#link-tree/torch/_dynamo/symbolic_convert.py", line 1487, in run while self.step(): File "/packages/aps_ads_vm/launcher_multiapp-inplace#link-tree/torch/_dynamo/symbolic_convert.py", line 1348, in step self.dispatch_table[inst.opcode](self, inst) File "/packages/aps_ads_vm/launcher_multiapp-inplace#link-tree/torch/_dynamo/symbolic_convert.py", line 2437, in LOAD_ATTR self._load_attr(inst) File "/packages/aps_ads_vm/launcher_multiapp-inplace#link-tree/torch/_dynamo/symbolic_convert.py", line 2425, in _load_attr result = BuiltinVariable(getattr).call_function( File "/packages/aps_ads_vm/launcher_multiapp-inplace#link-tree/torch/_dynamo/variables/builtin.py", line 1347, in call_function return handler(tx, args, kwargs) File "/packages/aps_ads_vm/launcher_multiapp-inplace#link-tree/torch/_dynamo/variables/builtin.py", line 967, in tx, [v.realize() for v in args], kwargs File "/packages/aps_ads_vm/launcher_multiapp-inplace#link-tree/torch/_dynamo/variables/builtin.py", line 967, in tx, [v.realize() for v in args], kwargs File "/packages/aps_ads_vm/launcher_multiapp-inplace#link-tree/torch/_dynamo/variables/lazy.py", line 72, in realize self._cache.realize() File "/packages/aps_ads_vm/launcher_multiapp-inplace#link-tree/torch/_dynamo/variables/lazy.py", line 33, in realize self.vt = builder.VariableBuilder(tx, self.source)(self.value) File "/packages/aps_ads_vm/launcher_multiapp-inplace#link-tree/torch/_dynamo/variables/builder.py", line 445, in __call__ vt = self._wrap(value) File "/packages/aps_ads_vm/launcher_multiapp-inplace#link-tree/torch/_dynamo/variables/builder.py", line 1043, in _wrap torch._dynamo.utils.store_user_object_weakref(value) File "/packages/aps_ads_vm/launcher_multiapp-inplace#link-tree/torch/_dynamo/utils.py", line 4694, in store_user_object_weakref user_obj_id_to_weakref[obj_id] = weakref.ref(obj) torch._dynamo.exc.InternalTorchDynamoError: TypeError: cannot create weak reference to 'torch.Event' object ``` This pull request makes us gracefully graph break, vs explicitly crashing. I've added a test which reproduces the issue. There is a side discussion re: how did torch.Event support ever work here, since it appears you cannot take a weakref to a torch.Event Pull Request resolved: https://github.com/pytorch/pytorch/pull/163168 Approved by: https://github.com/Lucaskabela, https://github.com/jansel --- test/dynamo/test_compile.py | 21 +++++++++++++++++++++ torch/_dynamo/graph_break_registry.json | 8 ++++++++ torch/_dynamo/utils.py | 13 ++++++++++++- 3 files changed, 41 insertions(+), 1 deletion(-) diff --git a/test/dynamo/test_compile.py b/test/dynamo/test_compile.py index 1f7290c51dd8..7df0ba2f1d3e 100644 --- a/test/dynamo/test_compile.py +++ b/test/dynamo/test_compile.py @@ -234,6 +234,27 @@ class InPlaceCompilationTests(TestCase): with self.assertRaises(IndexError): fn(torch.randn(10), 99) + def test_list_bad_weakref(self): + import weakref + + a = torch.Event() + with self.assertRaises(TypeError): + weakref.ref(a) + + @torch.compile(backend="eager") + class Mod(torch.nn.Module): + def __init__(self, event): + super().__init__() + self.event = event + + def forward(self, x): + return x * int(self.event.query()) + + e = torch.Event() + m = Mod(e) + a = torch.randn(10) + self.assertEqual(m(a), a) + # The private variants of the below functions are extensively tested # So as long as the signatures match we're good diff --git a/torch/_dynamo/graph_break_registry.json b/torch/_dynamo/graph_break_registry.json index 28fd02294ad3..4a873bbb0ce3 100644 --- a/torch/_dynamo/graph_break_registry.json +++ b/torch/_dynamo/graph_break_registry.json @@ -2718,5 +2718,13 @@ "Explanation": "Dyanmo does not support tracing mutations on a class when its __dict__ is materialized", "Hints": [] } + ], + "GB0272": [ + { + "Gb_type": "Failed to make weakref to User Object", + "Context": "user_objected: {obj}", + "Explanation": "Object does not allow us to make a weakref to it", + "Hints": [] + } ] } diff --git a/torch/_dynamo/utils.py b/torch/_dynamo/utils.py index 9d3b6bcd43cf..063a33833fe8 100644 --- a/torch/_dynamo/utils.py +++ b/torch/_dynamo/utils.py @@ -4696,7 +4696,18 @@ def get_user_object_from_id(obj_id: int) -> Any: def store_user_object_weakref(obj: object) -> None: obj_id = id(obj) - user_obj_id_to_weakref[obj_id] = weakref.ref(obj) + try: + user_obj_id_to_weakref[obj_id] = weakref.ref(obj) + except TypeError as e: + from .exc import unimplemented_v2 + + unimplemented_v2( + gb_type="Failed to make weakref to User Object", + context=f"user_objected: {obj}", + explanation="Object does not allow us to make a weakref to it", + hints=[], + from_exc=e, + ) class CompileTimeInstructionCounter: