CANN SHMEM Python 扩展接口全面测试与验证指南从单卡 smoke 到跨机 ROCE handle_wait【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem导读本指南围绕 CANN SHMEM 仓库中的 examples/python_extension 示例目录展开系统讲解其 Python 扩展接口shmem.core高层封装的功能覆盖、测试运行方式与验证口径。读者将掌握如何在多卡、单卡与 Host 对称堆场景下跑通全部用例如何对比putmem_on_stream的 Python/C 两条调用路径的性能差异如何通过标准多机torchrun在 ROCE 环境下验证handle_wait异步完成语义以及如何理解各用例背后的aclshmemx_*底层接口与对称堆、多实例、引擎配置、集合同步、profiling 等能力。1. 目录能力全景Python 扩展测试覆盖了哪些接口examples/python_extension下的测试不仅回归原有的 Python 接口还覆盖了新增的 UID 属性、多实例multi-instance、mem_type对称堆、引擎配置、集合同步、profiling 与handle_wait接口。整体结构如下examples/python_extension/ ├── run.sh # 一键运行全部用例安装 wheel torchrun 串行执行 ├── perf/ # putmem_on_stream 的 Python/C 性能对比 │ ├── run.sh # 性能测试入口默认 2 PE │ ├── benchmark_putmem_on_stream.py # 计时、轮次调度、结果聚合与门限判定 │ ├── putmem_on_stream_cpp.cpp # C 参考路径辅助动态库 │ └── build_cpp_ref.sh # 编译 C 参考库 libputmem_on_stream_cpp_ref.so └── test/ ├── init_test.py # 初始化冒烟 ├── qp_num_test.py # QP 数量检查 ├── tls_test.py # config store TLS 开关检查 ├── unique_id_test.py # UID 生成/属性检查 └── core/ # 核心功能用例torchrun 多 PE 运行 ├── test_init_final.py ├── test_memory.py # malloc/calloc/align/free、MemType、peer buffer ├── test_rma.py ├── test_direct.py ├── test_sync_config_prof.py # 引擎配置 集合同步 profiling ├── test_multi_instance.py # 多实例隔离 └── test_handle_wait.py # Handle stream RMA handle_wait1.1 Python 绑定的底层实现位置高层shmem.core封装最终通过 pybind11 绑定进入 C 层。绑定模块位于 src/host/python_wrapper/pyshmem.cpp其中直接引用aclshmemx_*系列原生接口并对入参做 Python 侧校验。例如instance_id在进入原生init之前会被校验为[0, 254]的整数对应MAX_INSTANCE_COUNT(255)非整数、布尔值或越界值会抛出value_error参见 pyshmem.cpp。因此测试用例中的“非法参数必须抛AclshmemInvalid/ValueError”断言验证的正是这一层参数防护逻辑。2. 前置条件与环境准备在运行任何用例前需要确认已安装与当前 CANN、Python 版本及 CPU 架构匹配的 SHMEM wheel 包cann-shmem。每个本地进程可独占一张 NPU运行前已加载 CANN 与 SHMEM 环境变量。多 PE 用例通过torchrun启动所有 PE 必须按相同顺序执行集合操作如barrier、sync否则会出现超时或挂起。run.sh的pre_check阶段会自动完成 wheel 管理若已安装旧版cann-shmem则先卸载再从项目根目录dist/下查找最新的cann_shmem-*.whl并以pip install --force-reinstall --no-deps安装见 run.sh。注意--no-deps的含义torch、torch-npu等运行时依赖是环境前提重装 wheel 时不会用不匹配的版本替换它们。若dist/下找不到 wheel脚本会提示先执行bash scripts/build.sh -python_extension构建。3. 运行全部用例3.1 默认双卡全量回归默认使用两张 NPUbash examples/python_extension/run.sh执行顺序对应 run.sh 中run_py_test与run_core_py_testtest/init_test.py初始化冒烟test/qp_num_test.pyQP 数量检查test/tls_test.pyconfig store TLS 配置检查test/unique_id_test.pyUID 获取与属性检查test/core/test_init_final.py、test_memory.py、test_rma.py、test_direct.pytest/core/test_sync_config_prof.py以SHMEM_CYCLE_PROF_PE0运行使 profiling 数据只在 PE 0 采集test/core/test_multi_instance.pytest/core/test_handle_wait.py。3.2 单卡 smoke用于绑定、单 PE 集合路径和多实例隔离的快速检查NPROC_PER_NODE1 bash examples/python_extension/run.shNPROC_PER_NODE会被透传给每个torchrun的--nproc-per-node默认值为 2见 run.sh。3.3 打开 Host 对称堆测试若当前 CANN 运行时支持 Host 对称堆可同时打开HOST_SIDE分配测试SHMEM_TEST_HOST_HEAP1 NPROC_PER_NODE2 bash examples/python_extension/run.sh该开关驱动 test_memory.py 中的条件分支仅在SHMEM_TEST_HOST_HEAP1时额外以core.MemType.HOST_SIDE分配并释放一块 4096 字节缓冲区并校验地址非零与mem_type正确。因此不要在运行时未提供 Host 对称堆的环境中盲目打开该变量否则用例会失败。3.4 单用例超时控制run.sh默认给每个torchrun用例设置 10 分钟超时对应TEST_TIMEOUT${SHMEM_TEST_TIMEOUT:-10m}配合timeout --signalTERM可用SHMEM_TEST_TIMEOUT调整SHMEM_TEST_TIMEOUT20m NPROC_PER_NODE2 bash examples/python_extension/run.sh超时通常意味着PE 之间集合调用顺序不一致、某个进程提前异常退出或运行环境的通信链路未就绪。3.5 handle_wait 的默认单 PE smoke 策略脚本默认单独用一个 PE 跑 MTEhandle_waitsmokeHANDLE_WAIT_NPROC_PER_NODE${HANDLE_WAIT_NPROC_PER_NODE:-1}避免在未启用 ROCE 时把该接口用于多 PE。如需覆盖HANDLE_WAIT_NPROC_PER_NODE2 bash examples/python_extension/run.sh但多 PEhandle_wait应优先按第 5 节在真实 ROCE 环境验证。4. putmem_on_streamPython/C 性能对比4.1 对比设计perf目录的核心思路是在同一 SHMEM 实例、同一组对称缓冲区、相同数据长度、相同 PE ring 路由、同一个显式 ACL stream 上对比两条路径C 参考路径辅助动态库libputmem_on_stream_cpp_ref.so在 C 循环中直接调用aclshmemx_putmem_on_stream。其实现见 putmem_on_stream_cpp.cpp函数签名通过extern C导出循环内依次下发aclshmemx_putmem_on_stream随后aclrtSynchronizeStream同步并返回微秒级耗时。Python 路径Python 循环调用高层shmem.core.put经 pybind11 进入同一个aclshmemx_putmem_on_stream。见 benchmark_putmem_on_stream.py 中的run_python以time.perf_counter_ns()计时循环iterations次后acl.rt.synchronize_stream(stream)同步。两条路径均在计时区间末尾同步相同 stream。测试先预热warmup次 Python put warmup次 C put中间以dist.barrier()对齐随后交替执行 C/Python 路径以降低系统性的先跑/后跑偏差每轮内再以dist.barrier()同步所有 PE见benchmark_putmem_on_stream.py第 218-265 行的交替调度逻辑。4.2 结果口径与正确性校验以各轮中所有 PE 的最大耗时作为整组完成时间模拟整组同步完成语义再取多轮结果的中位数测试结束后抽样校验接收缓冲区首尾字节actual_first/actual_last必须等于previous_pe 1确保 RMA 结果正确bit-exact默认门限为Python overhead 不超过 5%即python_us_per_op / cpp_us_per_op - 1 5%rank 0 负责聚合各 PE 数据输出 JSON含 world_size、bytes、iterations、warmup、rounds、stream、route、两个us_per_op、overhead_percent、passed、逐 rank 明细并写入输出文件见aggregate_results。4.3 运行 2/4/8 PE 对比NPROC_PER_NODE2 bash examples/python_extension/perf/run.sh NPROC_PER_NODE4 bash examples/python_extension/perf/run.sh NPROC_PER_NODE8 bash examples/python_extension/perf/run.sh默认配置为8 MiB8388608 字节、100 次迭代、10 次预热、7 轮测量均可通过环境变量覆盖PERF_BYTES8388608 \ PERF_ITERATIONS100 \ PERF_WARMUP10 \ PERF_ROUNDS7 \ PERF_THRESHOLD_PERCENT5 \ NPROC_PER_NODE4 \ bash examples/python_extension/perf/run.sh各变量的映射关系与 perf/run.sh 一一对应PERF_BYTES→--bytesPERF_ITERATIONS→--iterationsPERF_WARMUP→--warmupPERF_ROUNDS→--roundsPERF_THRESHOLD_PERCENT→--threshold-percent另有PERF_OUTPUT_DIR控制输出目录默认perf/results。脚本会先执行build_cpp_ref.sh编译 C 参考库再启动torchrunJSON 结果写入examples/python_extension/perf/results/putmem_on_stream_${NPROC_PER_NODE}pe.json。4.4 C 参考库的构建细节build_cpp_ref.sh 要求先设置ASCEND_HOME_PATH即 source CANN 的set_env.sh通过shmem-config --include与shmem-config --lib定位已安装的 SHMEM 头文件与库目录然后以-stdc17 -O3 -fPIC -shared编译链接-lshmem与-lascendcl。编译器按$CXX→g→c→bisheng的顺序选择。4.5 验收证据的硬性约束JSON 结果默认写入examples/python_extension/perf/results/。只有实际具备对应数量 NPU 的环境才能作为相应 2/4/8 卡验收证据不得使用较少卡数的结果替代——例如 4 卡环境下跑出putmem_on_stream_8pe.json是不合规的。5. 跨机 ROCE handle_wait 验证5.1 用例语义与完成顺序test/core/test_handle_wait.py与 C 侧 examples/rdma_handlewait_test/use_handlewait 使用相同的完成顺序stream RMA - handle_wait - stream synchronize - 数据可见性检查即先在指定 ACL stream 上下发 RMAcore.put(recv_buffer, send_buffer, remote_penext_pe, streamstream)随后在同一 stream上等待handle_waitcore.handle_wait(core.Handle(core.ACLSHMEM_TEAM_WORLD), stream)再torch.npu.synchronize()最后用acl.rt.memcpyD2H读取远端写入的字节做可见性校验期望值严格等于previous_pe 1见 test_handle_wait.py。这套语义验证的是handle_wait 返回后同一 stream 上先前的 RMA 结果对 Host 可见。5.2 MTE 与 ROCE 两种初始化路径用例通过环境变量SHMEM_TEST_ENGINE选择引擎默认MTE不区分大小写MTE走core.init(rank, nranks, mem_size, uid, initializer_methoduid)高层初始化ROCE走低层属性构造——ash.InitAttr()ash.aclshmemx_set_attr_uniqueid_args(pe, world_size, G_ASH_SIZE, native_uid, attr)再设置attr.option_attr.data_op_engine_type ash.OpEngineType.ROCE最后ash.aclshmemx_init_attr(ash.InitMode.UNIQUEID, attr)其他取值抛出ValueError(SHMEM_TEST_ENGINE must be MTE or ROCE.)。UID 通过dist.broadcast_object_list从 rank 0 广播给所有 PE_broadcast_unique_id。5.3 两机 ROCE 启动方式在两台已配置 ROCE 的机器上以标准多机torchrun参数启动同一脚本。例如每节点一张 NPUSHMEM_TEST_ENGINEROCE torchrun \ --nnodes2 \ --nproc-per-node1 \ --node-rank${NODE_RANK} \ --master-addr${MASTER_ADDR} \ --master-port${MASTER_PORT} \ examples/python_extension/test/core/test_handle_wait.pyNODE_RANK在两台机器上分别为 0 和 1MASTER_ADDR必须是两台机器都可达的 rank 0 地址。CANN/驱动、ROCE 网卡和 SHMEM 路由需提前配置完成。此外用例在执行前会调用ash.set_conf_store_tls(False, )关闭 config store TLS测试代码统一关闭 TLS 是为了在多进程场景下规避 TLS 初始化的串行/互斥成本与复杂度这也与仓库中其他多 PE 用例的做法一致。6. 用例与覆盖接口对照| 用例 | 覆盖内容 | | - | - | |test/core/test_memory.py|aclshmemx_malloc/calloc/align/free、MemTypeDEVICE_SIDE/HOST_SIDE、peer bufferget_peer_buffer | |test/core/test_sync_config_prof.py| 4 个 config 接口set_mte_config/set_sdma_config/set_rdma_config/set_udma_config、4 个 Host 集合同步接口barrier/barrier_all/sync/sync_all、2 个 stream barrierbarrier_on_stream/barrier_all_on_stream、get_prof/show_prof| |test/core/test_multi_instance.py| UID 属性构造链路、instance_ctx_get/set即set_instance/current_instance/multi_instance上下文管理器、heap 隔离、指定实例 finalize | |test/core/test_handle_wait.py|Handle、stream RMA、handle_wait、完成后数据可见性 | | 原有test/core用例 | 初始化、direct、RMA/Signal 及兼容性回归 |6.1 test_memory.py对称堆的完整生命周期用例依次验证详见 test_memory.pycore.buffer(g_malloc_size, mem_typecore.MemType.DEVICE_SIDE)校验 addr 非零、length 正确、mem_type正确、ownedTrue、instance_id core.current_instance()、初始release_calledFalsecore.Buffer(addr, length)直接构造做整数边界防护保留INTPTR_MAX/SIZE_T_MAX边界值拒绝 0 地址、0 长度、布尔值、地址溢出、长度溢出等非法输入抛AclshmemInvalid非 owning 的Buffer直接构造不允许free并发 free 防护两个线程同时free同一 buffer断言恰好一个进入原生 free、另一个被拒绝且release_called状态被更新对应 C 侧释放路径的 GIL 释放与原子占用语义core.calloc(count, size, mem_type)与core.align(alignment, size, mem_type)校验 calloc 总长度、align 地址按 256 对齐SHMEM_TEST_HOST_HEAP1时增加MemType.HOST_SIDE分配/释放core.get_peer_buffer(buf, next_pe)返回 peer PE 的同一对称地址视图非 owning、禁止 free且free状态会向 peer buffer 传播重复 free 被拒绝。6.2 test_sync_config_prof.py配置、集合与 profiling引擎 workspace 最小值校验RDMA/UDMA workspace 小于 128 字节时在进入原生代码前即被拒绝4 个 config 接口更新运行时状态供后续 device 侧操作使用所有 PE 必须以相同顺序执行集合调用Team 句柄校验越界 team、从未创建的 team、被销毁的 team 在使用barrier/sync/barrier_on_stream/handle_wait时均应被拒绝core.Handle构造与team_idsetter 对非法 team 抛ValueErrorprofilingSHMEM_CYCLE_PROF_PE0下仅配置的 PE 能取到core.ProfDatape_id、ccount/cycles矩阵维度为ACLSHMEM_CYCLE_PROF_MAX_BLOCK × ACLSHMEM_CYCLE_PROF_FRAME_CNT即 64×1024并支持并发快照与show_prof打印。6.3 test_multi_instance.py多实例隔离核心验证点详见 test_multi_instance.py以instance_id0初始化后current_instance()0core.multi_instance(id)上下文管理器可切换活动实例退出后自动恢复生命周期互斥用mock.patch打桩原生入口断言multi_instance持有锁期间set_mte_config、init构造 UID 属性、finalize等生命周期调用不会进入原生代码直到上下文释放防止实例切换过程中被并发生命周期调用破坏两个实例的 heap 地址不同instance_0_buffer.addr ! instance_1_buffer.addrheap 完全隔离跨实例的get_peer_buffer、put、get、put_signal、signal_op、signal_wait、free一律抛AclshmemInvalidset_instance/current_instance显式切换指定实例 finalizecore.finalize(instance_id1)后自动回落到实例 0。7. 常见失败模式排查结合 run.sh 与各用例的断言设计可归纳出几类典型失败信号用例超时默认 10 分钟优先检查所有 PE 是否按相同顺序执行集合调用barrier/sync/barrier_all检查是否有进程提前异常退出检查通信链路如多机场景下的MASTER_ADDR可达性、ROCE 路由是否就绪。性能门限未通过overhead_percent PERF_THRESHOLD_PERCENT。可增大PERF_BYTES大数据块摊薄 Python 侧 per-call 开销或调整PERF_THRESHOLD_PERCENT但注意结果 JSON 会记录门限值验收时需说明调整理由。HOST_SIDE 分配失败确认 CANN 运行时确实支持 Host 对称堆再设置SHMEM_TEST_HOST_HEAP1。跨机 handle_wait 失败确认已在两台机器上配置 ROCE 且使用SHMEM_TEST_ENGINEROCE走属性构造初始化路径单机 MTE smokeHANDLE_WAIT_NPROC_PER_NODE1不能替代跨机 ROCE 验证。8. 延伸阅读examples/python_extension/README.md本指南的原始出处含运行命令总览examples/python_extension/test/core全部核心用例源码可作为 Python 接口的使用范式examples/python_extension/perf性能对比基准实现examples/rdma_handlewait_test/use_handlewaitC 侧handle_wait对照用例src/host/python_wrapper/pyshmem.cpppybind11 绑定与入参校验实现src/python/shmem/coreshmem.core高层封装源码examples/python_extension/torch_testallgather/kv_shuffle的 PyTorch 集成示例。【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考