Skip to content

perf(env): confine DP collector host compute to the per-rank CPU block - #1324

Merged
TATP-233 merged 1 commit into
dev/issue-1304-motion-numba-body-statefrom
perf/dp-collector-numba-cpu-block
Aug 26, 2026
Merged

perf(env): confine DP collector host compute to the per-rank CPU block#1324
TATP-233 merged 1 commit into
dev/issue-1304-motion-numba-body-statefrom
perf/dp-collector-numba-cpu-block

Conversation

@TATP-233

Copy link
Copy Markdown
Collaborator

问题

多 rank off-policy DP 下,每个 rank 的 collector 通过 EnvCfg.cpu_ids 把 MuJoCo BatchEnvPool 的物理 worker 逐核绑定到本 rank 的 CPU 核区(#959/#966),但 collector 进程内的 host 侧计算没有跟随:

  • Numba 并行 kernel(base/backend/body_state.py、motion-tracking kernels.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 之前完成:

  1. os.sched_setaffinity(0, ids) 绑定调用线程,之后创建的线程(含 Numba 惰性启动的线程池)继承受限 mask;
  2. 通过 /proc/self/task 逐线程绑定已存在的线程(覆盖 import 期创建的 OpenBLAS 池),单线程退出竞态按 OSError 跳过;
  3. 未显式设置 NUMBA_NUM_THREADSnumba.set_num_threads(len(cpu_ids)),把 Numba 池规模从宿主核数收敛到核区长度(与 motion kernel runtime 的既有策略一致);
  4. CPU id 不在 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.pyNpEnv.__init__ 冷路径接入。
  • src/unilab/base/base.pyEnvCfg.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.py 10 项全过,含真实子进程 placement 验证。

说明

  • Base 分支为 dev/issue-1304-motion-numba-body-state(栈式):本 PR 修复的 Numba kernel 只存在于该分支链,尚未进 main。
  • 按 maintainer 要求不等待远程 CI,请人工审核并手动合并。

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.
@TATP-233
TATP-233 merged commit c2b110d into dev/issue-1304-motion-numba-body-state Aug 26, 2026
@TATP-233
TATP-233 deleted the perf/dp-collector-numba-cpu-block branch August 26, 2026 19:23
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant