Pull Request已成功合入, 合并人@CANN-robot
(感谢 Chang-an-HW 的贡献)变更摘要
本 PR 为 Eager 自定义算子在 RT2 动态图链路下补齐辅流申请能力。核心是在 EagerOpExecutionContext 上新增 request_attached_stream(key) 接口,并在 RT2 侧引入按执行器隔离、按 key 缓存的辅流容器 Rt2AttachedStreamCollection,通过 AttachedStreamProvider 接口和 CreateAttachedStreamProvider 初始化图节点注入到 EagerArgsHandler;同时调整 V1 侧 key 校验与辅流清理时机,并补充 RT2 动态 shape 样例与单元测试。
主要改动
-
新增 Python/绑定层辅流接口: 在
_ge_custom_op_native.pyi中新增EagerOpExecutionContext.request_attached_stream(key),并在context_binding.cc中实现BorrowedEagerOpExecutionContext::RequestAttachedStream,空 key 抛ValueError,链路不支持/建流失败/超配额时返回None。 -
RT2 辅流容器与生命周期管理: 新增
Rt2AttachedStreamCollection(实现AttachedStreamProvider),按 key 惰性建流并复用、单实例辅流上限kMaxAttachedStreamNum = 8、Destroy()先aclrtSynchronizeStream再销毁且幂等;配套新增attached_stream_provider_kernel.cc在 Init 图创建容器并由 Chain deleter 持有,执行器析构时释放,实现 per-executor 隔离。 -
RT2 图构建与执行注入链路: 新增
bg_attached_stream_provider获取当前执行器辅流容器;custom_node_converter.cc在custom_executor_func之后追加该附加输入,custom_op_kernel.cc扩展CustomOpInput::kAttachedStreamProvider并通过EagerArgsHandler::SetAttachedStreamProvider注入;Release()与析构不清理 provider,保证跨轮次复用。 -
V1 侧 key 校验与辅流清理调整:
AttachedStreamCollection::IsValidKey改为仅校验空 key,不再拦截CANN-FMK-前缀(保留为软约束);DavinciModel新增UnbindAndDestroyAttachedStreams并在DavinciModelFinalizer中于DestroyStream之前调用,重复调用无副作用。 -
样例与测试补充: 新增
eager_attached_stream_dynamic_shape动态 shape 样例(kernel、算子、session_run、run.sh日志断言、RtcKernelLoader等),并更新eager_attached_stream_add注释;新增/更新 RT2 辅流容器、CustomNodeKernel、EagerArgsHandler、custom_node_converter、DavinciModel等单元测试。


/approve


Pull Request
描述
2026829评审通过
为 Eager 自定义算子(
EagerExecuteOp)在 RT2 动态图执行链路 补齐按 key 申请框架托管物理辅流的能力。此前该能力只在 V1 静态图链路可用,RT2 下EagerOpExecutionContext::RequestAttachedStream因EagerArgsHandler未提供 provider 而恒返回nullptr。公共 API 与 ABI 零改动,用户算子代码零改动:
RequestAttachedStream签名、AdditionalInputIndex/AdditionalOutputIndex布局、ArgsHandler::GetAttachedStreamProvider()均不变,仅扩展 v2 内部枚举CustomOpInput。核心链路
关键设计
ModelV2Executor:容器由 Init 图节点的输出 Chain 持有。StreamExecutor为每条 aclrtStream 各建一个执行器,故不同执行流即使用相同 key 也拿到不同辅流;换来天然隔离且无需加锁(ExecuteCustomOp已注册kKernelUseMemory,固定在唯一 memory worker 上串行执行)。rtModel_t,不做aclmdlRIBindStream;StreamActive在 RT2 已退化为透传,不存在 V1 的 WAIT_ACTIVE 激活坑,按默认 flag 建流即可。Destroy()(先aclrtSynchronizeStream再aclrtDestroyStream,幂等)。现有卸载路径(UnloadRt2Model→DeleteExecutor、StreamExecutor::Erase)都在UnLoad()后立即销毁执行器,故无需额外 DeInit 节点。Execute每轮都会被调用(V1 只在下沉期一次),故 key 查找用透明比较容器(std::less<>+string_view),命中路径零堆分配,host 微基准约 25.3 ns/次。nullptr,由算子自行降级。Execute返回前用 event 把辅流 join 回ctx->GetStream(),已写入 API 文档约束说明。顺带修复的既存缺陷
RT2 内嵌 V1 静态子图场景下 Eager 算子申请的辅流确定性泄漏:
DavinciModelFinalizer未清理辅流,而DestroyResources()会置has_finalized_ = true,导致~DavinciModel()中含UnbindAndDestroy()的兜底清理块被整体跳过。修复:新增幂等的DavinciModel::UnbindAndDestroyAttachedStreams(),由 Finalizer 在DestroyStream()之前调用。行为变化(需评审确认)
RequestAttachedStream由恒nullptr变为可返回有效句柄——对以nullptr判定"链路不支持"的算子属行为变化。CANN-FMK-前缀的硬拦截(后续 HCCL 等框架组件也要通过本接口建流),改为 API 文档中的命名约定,key 校验统一为"仅判空"。属放宽型变化,原被拒绝的 key 现在会成功。对既有动态图流程的影响
bg::GetAttachedStreamProvider全仓唯一调用点在LoweringCustomNode,而它只注册给kCustomOpKernelLibName+kOnDeviceHbm,这类图不会新增任何节点或输入。SerializeComputeNodeInfo只遍历 IR 计算节点);profiling/dump/trace 只按 IR 输入数遍历,不受附加输入影响。本 PR 不包含
Python 绑定(因 ST fallback 编译问题临时移除)、动态 shape 端到端样例、OM2 路径(其 context 由 codegen SO 内部构造且跨 OM2 C ABI)。
变更类型
关联的Issue
无
如何测试
Execute每轮都被调用(RT2 链路生效)、辅流句柄多轮恒定不变(跨轮次跨 shape 复用)、异 key 隔离、每轮精度全部通过。核对清单
其他信息
无