Star843
1301
代码
代码
Issues1417
Pull Requests2256
流水线
项目讨论
Wiki
分析
Star843
1301
已合并
feat: add ExternalStream class with unsupported-scenario docs #43143
yuht9创建于 22 天前
feat: add ExternalStream class with unsupported-scenario docs #43143
已合并
yuht9创建于 22 天前
从yuht9/pytorchv2.7.1合入到Ascend/pytorchv2.7.1
Y
yuht9
22 天前

【合入来源】

如有社区issue，请关联issue链接
请勿携带内部流程信息（需求链接、问题单、内部issue等）

【修改方案】

ExternalStream 是 torch_npu.npu.Stream 的子类，用于包装外部库创建的 aclrtStream 句柄，使 torch_npu 能在该流上 launch 算子，实现多库数据交换。

类定义

class ExternalStream(Stream):
    def __new__(cls, stream_ptr, device=None, **kwargs):
        with torch_npu.npu.device(device):
            return super().__new__(cls, stream_ptr=stream_ptr, **kwargs)

构造参数

参数 类型 必填 说明
stream_ptr int 是 外部 aclrtStream 的整数值
device torch.device / int 否 流所在设备，默认当前设备
priority int 否 必须为 0（外部流不支持优先级）
is_sync_launch int 否 必须为 0（外部流不支持 sync launch）

【资料变更】

不涉及

【接口变更】

不涉及

【功能验证】

DT验证通过

【CheckList】

PR提交人对以下CheckList自检项进行全量自检，自检通过或不涉及，均修改 [ ] 为 [x]

likedislike
Pull Request已成功合入, 合并人@ascend-robot
（感谢 yuht9 的贡献）
Yyuht9
22 天前 创建了 pull request，commit 66b9fb32
Yyuht9
22 天前 关联了issue：[Feature]: add ExternalStream class
atomgit-bot
atomgit-bot
22 天前 评论：

变更摘要

此 PR 新增了 ExternalStream 类，作为 torch_npu.npu.Stream 的子类，用于包装外部库创建的 aclrtStream 句柄，使 torch_npu 能够在该外部流上执行算子以实现多库数据交换。同时扩展了 C++ 层的流构造逻辑以支持通过 stream_ptr 创建外部流，并添加了对应的单元测试覆盖创建、切换、运算、同步限制等场景。

主要改动

  • 新增 ExternalStream 类：在 torch_npu/npu/streams.py 中定义 ExternalStream(Stream) 子类，通过 __new__ 接收 stream_ptr 整数值和可选的 device 参数，并在设备上下文中调用父类构造；类文档详细列出了 synchronize()、query()、record_event() 等不支持的使用场景及其根因说明。

  • C++ 层支持 stream_ptr 构造外部流：在 torch_npu/csrc/npu/Stream.cpp 的 THNPStream_pynew 函数中增加对 stream_ptr 参数的判断分支，当传入 stream_ptr 时校验 priority 和 is_sync_launch 必须为默认值，并调用 c10_npu::getStreamFromExternal 创建流对象。

  • 模块导出更新：在 torch_npu/npu/__init__.py 中将 ExternalStream 加入 __all__ 列表并从 streams 模块导入，确保用户可通过 torch_npu.npu.ExternalStream 访问。

  • 新增测试用例：在 test/npu/test_stream.py 中新增 TestExternalStream 测试类，覆盖外部流创建、设为当前流、执行张量运算、相同指针返回同一对象、isinstance 检查以及 synchronize() 和 query() 的限制行为验证。

likedislike
ascend-robot
ascend-robot成员
22 天前 评论：
atomgit-bot
atomgit-bot
22 天前 评论：

代码审查

All four changed files are reviewed. Let me now submit the closing summary.

审查总结

按优先级统计

优先级 数量
P0 0
P1 0
P2 1
P3 0

各文件审查结果

  • test/npu/test_stream.py — 新增 TestExternalStream 测试类，7 个测试用例覆盖创建、设为当前流、执行算子、相同指针、isinstance、synchronize/query 限制。_get_external_stream 上下文管理器正确创建和销毁原始流。无问题。

  • torch_npu/csrc/npu/Stream.cpp — 新增 stream_ptr 参数解析和两条 TORCH_CHECK 校验，修改三元表达式链以支持外部流路径。存在 1 个 P2 问题：三元链中 stream_id || device_index || device_type 优先于 stream_ptr，当两者同时传入时外部流创建被绕过。

  • torch_npu/npu/__init__.py — 在 __all__ 和 import 行中新增 ExternalStream。无问题。

  • torch_npu/npu/streams.py — 新增 ExternalStream(Stream) 类，通过 __new__ 将 device 参数转为设备上下文并透传 stream_ptr 到 C++ 层。设计合理，参数约束由 C++ 层兜底。无问题。

整体风险评估

低风险。唯一的 P2 问题需要调用者同时传入 stream_ptr 和 stream_id 才会触发，正常通过 ExternalStream Python API 使用不会命中该路径。建议修复三元表达式顺序以消除隐患。

类型 数量
🔴 阻塞 0
🟡 建议 2

💬 仅评论

likedislike
此处折叠了43条消息 查看更多
H
hbhu_bin成员
22 天前 评论：

/lgtm

likedislike
ascend-robotascend-robot成员
22 天前 添加了label：lgtm
ascend-robotascend-robot成员
22 天前 合入了pull request
ascend-robot
ascend-robot成员
22 天前 评论：

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.

likedislike
ascend-robot
ascend-robot成员
22 天前 评论：
流水线 pytorch_gitcode_PR_multiVersion#13320 [ commitID：cd784bd9 ] 已完成
likedislike