已合并
feat: add ExternalStream class with unsupported-scenario docs #43139
yuht9创建于 26 天前
feat: add ExternalStream class with unsupported-scenario docs #43139
已合并
共 5 个文件变更+159-3
| @@ -1,3 +1,6 @@ | |||
| 1 | +import contextlib | ||
| 2 | +import ctypes | ||
| 3 | + | ||
| 1 | import torch | 4 | import torch |
| 2 | 5 | ||
| 3 | import torch_npu | 6 | import torch_npu |
| @@ -47,5 +50,71 @@ class TestNpuStream(TestCase): | |||
| 47 | self.assertTrue((s.stream_id >> 5) == 4) | 50 | self.assertTrue((s.stream_id >> 5) == 4) |
| 48 | 51 | ||
| 49 | 52 | ||
| 53 | +class TestExternalStream(TestCase): | ||
| 54 | + | ||
| 55 | + | ||
| 56 | + def _get_external_stream(self, device=None): | ||
| 57 | + rt = torch_npu.npu.npurt() | ||
| 58 | + stream = ctypes.c_void_p(0) | ||
| 59 | + with torch_npu.npu.device(device): | ||
| 60 | + try: | ||
| 61 | + ret = rt.npuStreamCreate(ctypes.addressof(stream)) | ||
| 62 | + self.assertEqual(int(ret), 0) | ||
| 63 | + self.assertNotEqual(stream.value, 0) | ||
| 64 | + yield stream.value | ||
| 65 | + finally: | ||
| 66 | + if stream.value: | ||
| 67 | + ret = rt.npuStreamDestroy(stream.value) | ||
| 68 | + self.assertEqual(int(ret), 0) | ||
| 69 | + | ||
| 70 | + def test_external_stream_creation(self): | ||
| 71 | + with self._get_external_stream() as stream_v: | ||
| 72 | + ext_stream = torch_npu.npu.ExternalStream(stream_v) | ||
| 73 | + self.assertEqual(stream_v, ext_stream.npu_stream) | ||
| 74 | + self.assertEqual(ext_stream.device.index, torch_npu.npu.current_device()) | ||
| 75 | + | ||
| 76 | + def test_external_stream_as_current(self): | ||
| 77 | + with self._get_external_stream() as stream_v: | ||
| 78 | + ext_stream = torch_npu.npu.ExternalStream(stream_v) | ||
| 79 | + with torch_npu.npu.stream(ext_stream): | ||
| 80 | + self.assertEqual( | ||
| 81 | + torch_npu.npu.current_stream().npu_stream, | ||
| 82 | + ext_stream.npu_stream) | ||
| 83 | + | ||
| 84 | + def test_external_stream_op(self): | ||
| 85 | + with self._get_external_stream() as stream_v: | ||
| 86 | + ext_stream = torch_npu.npu.ExternalStream(stream_v) | ||
| 87 | + with torch_npu.npu.stream(ext_stream): | ||
| 88 | + x = torch.randn(10, device='npu') | ||
| 89 | + y = x + 1 | ||
| 90 | + torch_npu.npu.synchronize() | ||
| 91 | + self.assertEqual(y.sum().item(), x.sum().item() + x.numel()) | ||
| 92 | + | ||
| 93 | + def test_external_stream_same_ptr(self): | ||
| 94 | + with self._get_external_stream() as stream_v: | ||
| 95 | + ext1 = torch_npu.npu.ExternalStream(stream_v) | ||
| 96 | + ext2 = torch_npu.npu.ExternalStream(stream_v) | ||
| 97 | + self.assertEqual(ext1.npu_stream, ext2.npu_stream) | ||
| 98 | + self.assertEqual(ext1, ext2) | ||
| 99 | + | ||
| 100 | + def test_external_stream_isinstance(self): | ||
| 101 | + with self._get_external_stream() as stream_v: | ||
| 102 | + ext_stream = torch_npu.npu.ExternalStream(stream_v) | ||
| 103 | + self.assertIsInstance(ext_stream, torch_npu.npu.Stream) | ||
| 104 | + self.assertIsInstance(ext_stream, torch_npu.npu.ExternalStream) | ||
| 105 | + | ||
| 106 | + def test_external_stream_synchronize_restriction(self): | ||
| 107 | + with self._get_external_stream() as stream_v: | ||
| 108 | + ext_stream = torch_npu.npu.ExternalStream(stream_v) | ||
| 109 | + with self.assertRaisesRegex(RuntimeError, "NPUStream::synchronize"): | ||
| 110 | + ext_stream.synchronize() | ||
| 111 | + | ||
| 112 | + def test_external_stream_query_restriction(self): | ||
| 113 | + with self._get_external_stream() as stream_v: | ||
| 114 | + ext_stream = torch_npu.npu.ExternalStream(stream_v) | ||
| 115 | + with self.assertRaisesRegex(RuntimeError, "Cannot query"): | ||
| 116 | + ext_stream.query() | ||
| 117 | + | ||
| 118 | + | ||
| 50 | if __name__ == "__main__": | 119 | if __name__ == "__main__": |
| 51 | run_tests() | 120 | run_tests() |
| @@ -719,6 +719,9 @@ | |||
| 719 | "torch_npu.npu.graphs.make_graphed_callables": { | 719 | "torch_npu.npu.graphs.make_graphed_callables": { |
| 720 | "signature": "(callables, sample_args, num_warmup_iters=3, allow_unused_input=False, pool=None)" | 720 | "signature": "(callables, sample_args, num_warmup_iters=3, allow_unused_input=False, pool=None)" |
| 721 | }, | 721 | }, |
| 722 | + "torch_npu.npu.streams.ExternalStream": { | ||
| 723 | + "signature": "(stream_ptr, device=None, **kwargs)" | ||
| 724 | + }, | ||
| 722 | "torch_npu.npu.streams.ExternalEvent": { | 725 | "torch_npu.npu.streams.ExternalEvent": { |
| 723 | "signature": "()" | 726 | "signature": "()" |
| 724 | }, | 727 | }, |
| @@ -863,6 +866,9 @@ | |||
| 863 | "torch_npu.npu.Event.ipc_handle": { | 866 | "torch_npu.npu.Event.ipc_handle": { |
| 864 | "signature": "(self)" | 867 | "signature": "(self)" |
| 865 | }, | 868 | }, |
| 869 | + "torch_npu.npu.ExternalStream": { | ||
| 870 | + "signature": "(stream_ptr, device=None, **kwargs)" | ||
| 871 | + }, | ||
| 866 | "torch_npu.npu.ExternalEvent": { | 872 | "torch_npu.npu.ExternalEvent": { |
| 867 | "signature": "()" | 873 | "signature": "()" |
| 868 | }, | 874 | }, |
| @@ -56,12 +56,26 @@ static PyObject *THNPStream_pynew( | |||
| 56 | return nullptr; | 56 | return nullptr; |
| 57 | } | 57 | } |
| 58 | 58 | ||
| 59 | + if (stream_ptr) { | ||
| 60 | + TORCH_CHECK( | ||
| 61 | + priority == 0, | ||
| 62 | + "Priority was explicitly set for an external stream", | ||
| 63 | + PTA_ERROR(ErrCode::PARAM)); | ||
| 64 | + TORCH_CHECK( | ||
| 65 | + is_sync_launch == 0, | ||
| 66 | + "is_sync_launch was explicitly set for an external stream", | ||
| 67 | + PTA_ERROR(ErrCode::PARAM)); | ||
| 68 | + } | ||
| 69 | + | ||
| 59 | c10_npu::NPUStream stream = | 70 | c10_npu::NPUStream stream = |
| 60 | (stream_id || device_index || device_type) ? | 71 | (stream_id || device_index || device_type) ? |
| 61 | c10_npu::NPUStream::unpack3( | 72 | c10_npu::NPUStream::unpack3( |
| 62 | stream_id, device_index, static_cast<c10::DeviceType>(device_type)) : | 73 | stream_id, device_index, static_cast<c10::DeviceType>(device_type)) : |
| 74 | + (stream_ptr ? | ||
| 75 | + c10_npu::getStreamFromExternal( | ||
| 76 | + reinterpret_cast<aclrtStream>(stream_ptr), current_device) : | ||
| 63 | (is_sync_launch ? c10_npu::getNPUStreamFromSyncLaunchPool() : | 77 | (is_sync_launch ? c10_npu::getNPUStreamFromSyncLaunchPool() : |
| 64 | - c10_npu::getStreamFromPool(priority)); | 78 | + c10_npu::getStreamFromPool(priority))); |
| 65 | 79 | ||
| 66 | THNPStream *self = (THNPStream *)ptr.get(); | 80 | THNPStream *self = (THNPStream *)ptr.get(); |
| 67 | self->stream_id = static_cast<int64_t>(stream.id()); | 81 | self->stream_id = static_cast<int64_t>(stream.id()); |
| @@ -115,6 +115,7 @@ __all__ = [ | |||
| 115 | "is_current_stream_capturing", | 115 | "is_current_stream_capturing", |
| 116 | "make_graphed_callables", | 116 | "make_graphed_callables", |
| 117 | "ExternalEvent", | 117 | "ExternalEvent", |
| 118 | + "ExternalStream", | ||
| 118 | "graph_task_group_begin", | 119 | "graph_task_group_begin", |
| 119 | "graph_task_group_end", | 120 | "graph_task_group_end", |
| 120 | "graph_task_update_begin", | 121 | "graph_task_update_begin", |
| @@ -164,7 +165,7 @@ from .utils import (obfuscation_initialize, obfuscation_calculate, obfuscation_f | |||
| 164 | finalize_dump, set_dump, get_npu_overflow_flag, clear_npu_overflow_flag, | 165 | finalize_dump, set_dump, get_npu_overflow_flag, clear_npu_overflow_flag, |
| 165 | check_uce_in_memory, stress_detect, _get_uce_addr, ipc_collect, set_op_timeout_ms) | 166 | check_uce_in_memory, stress_detect, _get_uce_addr, ipc_collect, set_op_timeout_ms) |
| 166 | from ._recovery import restart_device, stop_device | 167 | from ._recovery import restart_device, stop_device |
| 167 | -from .streams import Stream, Event, SyncLaunchStream, ExternalEvent | 168 | +from .streams import Stream, Event, SyncLaunchStream, ExternalStream, ExternalEvent |
| 168 | from .mstx import mstx | 169 | from .mstx import mstx |
| 169 | from .npu_config import * # noqa: F403 | 170 | from .npu_config import * # noqa: F403 |
| 170 | from .autocast_utils import * # noqa: F403 | 171 | from .autocast_utils import * # noqa: F403 |
| @@ -3,7 +3,7 @@ import ctypes | |||
| 3 | import torch_npu | 3 | import torch_npu |
| 4 | import torch_npu._C | 4 | import torch_npu._C |
| 5 | 5 | ||
| 6 | -__all__ = ["Stream", "Event", "SyncLaunchStream", "ExternalEvent"] | 6 | +__all__ = ["Stream", "Event", "SyncLaunchStream", "ExternalStream", "ExternalEvent"] |
| 7 | 7 | ||
| 8 | 8 | ||
| 9 | class Stream(torch_npu._C._NPUStreamBase): | 9 | class Stream(torch_npu._C._NPUStreamBase): |
| @@ -110,6 +110,72 @@ class Stream(torch_npu._C._NPUStreamBase): | |||
| 110 | .format(self.device, self.npu_stream)) | 110 | .format(self.device, self.npu_stream)) |
| 111 | 111 | ||
| 112 | 112 | ||
| 113 | +class ExternalStream(Stream): | ||
| 114 | + r"""Wrapper around an externally allocated NPU stream. | ||
| 115 | + | ||
| 116 | + This class is used to wrap streams allocated in other libraries in order | ||
| 117 | + to facilitate data exchange and multi-library interactions. | ||
| 118 | + | ||
| 119 | + .. note:: This class doesn't manage the stream life-cycle, it is the user | ||
| 120 | + responsibility to keep the referenced stream alive while this class is | ||
| 121 | + being used. | ||
| 122 | + | ||
| 123 | + .. note:: ``priority`` and ``is_sync_launch`` must keep their defaults | ||
| 124 | + (``0``) when wrapping an external stream; passing non-default values | ||
| 125 | + raises a ``RuntimeError``. | ||
| 126 | + | ||
| 127 | + Args: | ||
| 128 | + stream_ptr(int): Integer representation of the `aclrtStream` value. | ||
| 129 | + allocated externally. | ||
| 130 | + device(torch.device or int, optional): the device where the stream | ||
| 131 | + was originally allocated. If device is specified incorrectly, | ||
| 132 | + subsequent launches using this stream may fail. | ||
| 133 | + | ||
| 134 | + Unsupported scenarios (raise ``RuntimeError``): | ||
| 135 | + | ||
| 136 | + .. list-table:: | ||
| 137 | + :widths: 35 65 | ||
| 138 | + :header-rows: 1 | ||
| 139 | + | ||
| 140 | + * - API | ||
| 141 | + - Description | ||
| 142 | + * - ``synchronize()`` | ||
| 143 | + - Caller must synchronize the external stream externally. | ||
| 144 | + * - ``query()`` | ||
| 145 | + - Caller must track stream completion externally. | ||
| 146 | + * - ``record_event()`` / ``Event.record(ext)`` / | ||
| 147 | + ``ExternalEvent.record(ext)`` | ||
| 148 | + - Event recording on an external stream is not supported. | ||
| 149 | + * - ``wait_event(event)`` / ``wait_stream(stream)`` / | ||
| 150 | + ``Event.wait(ext)`` / ``ExternalEvent.wait(ext)`` | ||
| 151 | + - Event waiting on an external stream is not supported. | ||
| 152 | + * - ``Tensor.record_stream(ext)`` | ||
| 153 | + - Caching allocator stream tracking is not supported. | ||
| 154 | + * - ``NPUGraph.capture_begin()`` / ``capture_end()`` | ||
| 155 | + - Graph capture on an external stream is not supported. | ||
| 156 | + * - ``graph_task_group_begin`` / ``end`` / | ||
| 157 | + ``graph_task_update_begin`` / ``end`` | ||
| 158 | + - Graph task group APIs are not supported. | ||
| 159 | + * - ``super_kernel_scope_begin`` / ``end`` | ||
| 160 | + - Super kernel scope is not supported. | ||
| 161 | + * - ``launch_callback`` | ||
| 162 | + - Host callback launch is not supported. | ||
| 163 | + * - ``subscribe_report`` / ``unsubscribe_report`` | ||
| 164 | + - Task report subscribe/unsubscribe is not supported. | ||
| 165 | + * - ``LaunchRecordEventTask`` / ``LaunchWaitEventTask`` | ||
| 166 | + - Async task queue event ops are not supported. | ||
| 167 | + | ||
| 168 | + The root cause: the NPU backend routes the above APIs through torch_npu's | ||
| 169 | + internal ``AsyncTaskQueue`` (taskqueue), which only tracks streams it | ||
| 170 | + created itself. An external stream is outside the taskqueue's tracking | ||
| 171 | + scope. For full capabilities, use a torch_npu-managed ``Stream`` instead. | ||
| 172 | + """ | ||
| 173 | + | ||
| 174 | + def __new__(cls, stream_ptr, device=None, **kwargs): | ||
| 175 | + with torch_npu.npu.device(device): | ||
| 176 | + return super().__new__(cls, stream_ptr=stream_ptr, **kwargs) | ||
 | |||
| 177 | + | ||
| 178 | + | ||
| 113 | class Event(torch_npu._C._NPUEventBase): | 179 | class Event(torch_npu._C._NPUEventBase): |
| 114 | r"""Wrapper around a NPU event. | 180 | r"""Wrapper around a NPU event. |
| 115 | 181 | ||
🟡 Medium Priority
changed line → affected behavior/contract → failure mode → suggested fix
问题链：
ExternalStream.__new__（streams.py:174-176）接受stream_ptr参数并透传至 C++ 构造函数THNPStream_pynew（Stream.cpp:29），其中stream_ptr被声明为uint64_t stream_ptr = 0。当用户传入
stream_ptr=0（例如未初始化的c_void_p值）时：而
getStreamFromExternal内部虽有空指针校验（NPUStream.cpp:486：TORCH_CHECK(stream != nullptr, ...)），但因三元表达式已将stream_ptr=0过滤，该校验永远不会被触发。失效模式： 用户显式构造
ExternalStream(0)期望得到一个包装外部流的对象或收到明确错误，结果却静默得到一个普通池流，且该流的stream_id类型为NORMAL/HIGH而非EXT，后续调用synchronize()/query()不会触发预期的RuntimeError（测试test_external_stream_synchronize_restriction/test_external_stream_query_restriction专门验证了对真正外部流应抛错）。这在多库交互场景下可能导致难以排查的行为异常。建议：在 ExternalStream.new 开头增加 stream_ptr 非零校验，提前以明确的 ValueError 拒绝零值指针，避免静默回退为普通池流。