perf(env): confine DP collector host compute to the per-rank CPU block - #1324
Merged
TATP-233 merged 1 commit intoAug 26, 2026
Conversation
Multi-rank off-policy collectors pin their MuJoCo BatchEnvPool workers to a per-rank CPU block via EnvCfg.cpu_ids, but the collector's host-side compute did not follow: Numba's parallel kernels sized their pool from the host CPU count and drifted across rank boundaries, and the OpenBLAS pool spawned at import kept the host-wide mask. NpEnv.__init__ now applies apply_env_cpu_runtime(cfg.cpu_ids) on the cold path: the process is confined to the block (existing threads pinned individually via /proc/self/task, later threads — including Numba's lazily-launched pool — inherit the mask) and Numba's pool is sized to len(cpu_ids) unless NUMBA_NUM_THREADS is set explicitly. cpu_ids=None keeps the single-rank path bit-identical. Backend-agnostic: any env declaring cpu_ids (e.g. motrix once it grows affinity support) gets the same confinement.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
问题
多 rank off-policy DP 下,每个 rank 的 collector 通过
EnvCfg.cpu_ids把 MuJoCo BatchEnvPool 的物理 worker 逐核绑定到本 rank 的 CPU 核区(#959/#966),但 collector 进程内的 host 侧计算没有跟随:base/backend/body_state.py、motion-trackingkernels.py)的线程池按宿主 CPU 总数建池、位置交给 OS 调度,会漂移到其他 rank 的核区,与相邻 rank 已绑定的物理 worker 互相抢占。import numpy时 OpenBLAS 就建好的线程池(早于 env 构造)同样保留全机 mask。方案
通用 env 层修改,不依赖具体后端:
EnvCfg.cpu_ids是唯一事实来源,NpEnv.__init__冷路径首先调用新的apply_env_cpu_runtime(cfg.cpu_ids)(src/unilab/base/cpu_runtime.py),在 managers 加载、backend materialize、首个 Numba 并行 kernel 之前完成:os.sched_setaffinity(0, ids)绑定调用线程,之后创建的线程(含 Numba 惰性启动的线程池)继承受限 mask;/proc/self/task逐线程绑定已存在的线程(覆盖 import 期创建的 OpenBLAS 池),单线程退出竞态按 OSError 跳过;NUMBA_NUM_THREADS时numba.set_num_threads(len(cpu_ids)),把 Numba 池规模从宿主核数收敛到核区长度(与 motion kernel runtime 的既有策略一致);sched_getaffinity内时 fail-closed 抛ValueError(先于 backend pool 创建失败);无sched_setaffinity的平台告警降级为只限制 Numba 线程数。cpu_ids=None(单 rank 默认)完全 no-op,单卡路径 bit-identical。motrix 后端目前不消费cpu_ids;一旦其接入该字段(#962),env 层 confinement 自动生效,无需本 PR 之外的改动。改动
src/unilab/base/cpu_runtime.py(新增):apply_env_cpu_runtime+_confine_existing_threads。src/unilab/base/np_env.py:NpEnv.__init__冷路径接入。src/unilab/base/base.py:EnvCfg.cpu_ids注释更新为扩展后的语义。tests/base/test_cpu_runtime.py(新增):mock 单测覆盖 no-op/生效/NUMBA_NUM_THREADS 优先/不可用 CPU fail-closed/无 affinity 平台降级/逐线程绑定跳过失败项/NpEnv 接线;另有一个 Linux 子进程实测,断言主线程、import 期 OpenBLAS 池与 Numba 池全部落在核区内且get_num_threads()==len(block)。docs/sphinx/.../zh_CN/2-user_guide/2-algorithms/3-sac.md:补充核区对 collector 进程与 Numba 池的约束说明(en 页未记载该字段,不做强制镜像)。Validation
make test-all通过(ruff format/check、mypy、pyright、pytest -m "not slow" --cov、benchmark smoke 33/34 + 34/35,唯一 skip 为平台可选 mlx)。tests/base/test_cpu_runtime.py10 项全过,含真实子进程 placement 验证。说明
dev/issue-1304-motion-numba-body-state(栈式):本 PR 修复的 Numba kernel 只存在于该分支链,尚未进 main。