已合并
fix: add PyObject_ClearWeakRefs in NPU Stream/Event tp_dealloc #41596
21xiu创建于 7月14日
fix: add PyObject_ClearWeakRefs in NPU Stream/Event tp_dealloc #41596
已合并
共 3 个文件变更+35-2
| @@ -1,6 +1,8 @@ | |||
| 1 | # Owner(s): ["module: dynamo"] | 1 | # Owner(s): ["module: dynamo"] |
| 2 | import functools | 2 | import functools |
| 3 | +import gc | ||
| 3 | import unittest | 4 | import unittest |
| 5 | +import weakref | ||
| 4 | import torch | 6 | import torch |
| 5 | import torch._dynamo.test_case | 7 | import torch._dynamo.test_case |
| 6 | import torch_npu | 8 | import torch_npu |
| @@ -25,6 +27,34 @@ class StreamintoDynamoTests(torch._dynamo.test_case.TestCase): | |||
| 25 | output1 = model_1(inp) | 27 | output1 = model_1(inp) |
| 26 | torch.allclose(output, output1) | 28 | torch.allclose(output, output1) |
| 27 | 29 | ||
| 30 | + def _assert_weakref_callback_fires(self, factory): | ||
| 31 | + """Backend Stream/Event tp_dealloc overrides must call | ||
| 32 | + PyObject_ClearWeakRefs so that weakrefs to destroyed instances | ||
| 33 | + are properly cleared. Without it, the weakref's wr_object is | ||
| 34 | + left dangling and any later access (e.g. the dynamo external- | ||
| 35 | + object registry being cleared at interpreter finalization) | ||
| 36 | + hits a use-after-free. | ||
| 37 | + | ||
| 38 | + The callback firing is the only reliable Python-level signal — | ||
| 39 | + `weakref.ref(s)() is None` can return True even with the bug | ||
| 40 | + present, because other CPython paths may null wr_object | ||
| 41 | + without invoking callbacks (Objects/weakrefobject.c). | ||
| 42 | + """ | ||
| 43 | + called = [] | ||
| 44 | + obj = factory() | ||
| 45 | + # ``_weakref_keepalive`` must outlive ``del obj`` so the callback | ||
| 46 | + # has a chance to fire; the binding is load-bearing. | ||
| 47 | + _weakref_keepalive = weakref.ref(obj, lambda _ref: called.append(True)) | ||
| 48 | + del obj | ||
| 49 | + gc.collect() | ||
| 50 | + self.assertEqual(called, [True]) | ||
| 51 | + del _weakref_keepalive | ||
| 52 | + | ||
| 53 | + | ||
| 54 | + def test_npu_stream_event_weakref_callback(self): | ||
| 55 | + self._assert_weakref_callback_fires(torch.npu.Stream) | ||
| 56 | + self._assert_weakref_callback_fires(torch.npu.Event) | ||
| 57 | + | ||
| 28 | 58 | ||
| 29 | if __name__ == "__main__": | 59 | if __name__ == "__main__": |
| 30 | from torch._dynamo.test_case import run_tests | 60 | from torch._dynamo.test_case import run_tests |
| @@ -118,6 +118,7 @@ static PyObject* THNPEvent_from_ipc_handle( | |||
| 118 | static void THNPEvent_dealloc(THNPEvent *self) | 118 | static void THNPEvent_dealloc(THNPEvent *self) |
| 119 | { | 119 | { |
| 120 | self->npu_event.~NPUEvent(); | 120 | self->npu_event.~NPUEvent(); |
| 121 | + PyObject_ClearWeakRefs((PyObject*)self); | ||
| 121 | Py_TYPE(self)->tp_free((PyObject*)self); | 122 | Py_TYPE(self)->tp_free((PyObject*)self); |
| 122 | } | 123 | } |
| 123 | 124 | ||
| @@ -262,7 +263,7 @@ PyTypeObject THNPEventType = { | |||
| 262 | nullptr, /* tp_traverse */ | 263 | nullptr, /* tp_traverse */ |
| 263 | nullptr, /* tp_clear */ | 264 | nullptr, /* tp_clear */ |
| 264 | nullptr, /* tp_richcompare */ | 265 | nullptr, /* tp_richcompare */ |
| 265 | - 0, /* tp_weaklistoffset */ | 266 | + 0, /* tp_weaklistoffset (inherited from THPStreamType via tp_base) */ |
| 266 | nullptr, /* tp_iter */ | 267 | nullptr, /* tp_iter */ |
| 267 | nullptr, /* tp_iternext */ | 268 | nullptr, /* tp_iternext */ |
| 268 | THNPEvent_methods, /* tp_methods */ | 269 | THNPEvent_methods, /* tp_methods */ |
| @@ -76,6 +76,8 @@ static PyObject *THNPStream_pynew( | |||
| 76 | static void THNPStream_dealloc(THNPStream *self) | 76 | static void THNPStream_dealloc(THNPStream *self) |
| 77 | { | 77 | { |
| 78 | self->npu_stream.~NPUStream(); | 78 | self->npu_stream.~NPUStream(); |
| 79 | + PyObject_ClearWeakRefs((PyObject*)self); | ||
| 80 | + Py_CLEAR(self->context); | ||
| 79 | Py_TYPE(self)->tp_free((PyObject*)self); | 81 | Py_TYPE(self)->tp_free((PyObject*)self); |
| 80 | } | 82 | } |
| 81 | 83 | ||
| @@ -192,7 +194,7 @@ PyTypeObject THNPStreamType = { | |||
| 192 | nullptr, /* tp_traverse */ | 194 | nullptr, /* tp_traverse */ |
| 193 | nullptr, /* tp_clear */ | 195 | nullptr, /* tp_clear */ |
| 194 | nullptr, /* tp_richcompare */ | 196 | nullptr, /* tp_richcompare */ |
| 195 | - 0, /* tp_weaklistoffset */ | 197 | + 0, /* tp_weaklistoffset (inherited from THPStreamType via tp_base) */ |
| 196 | nullptr, /* tp_iter */ | 198 | nullptr, /* tp_iter */ |
| 197 | nullptr, /* tp_iternext */ | 199 | nullptr, /* tp_iternext */ |
| 198 | THNPStream_methods, /* tp_methods */ | 200 | THNPStream_methods, /* tp_methods */ |