Skip to content

perf(mujoco): 单卡默认 pool 线程数收敛到有效 CPU 数,消除 2× 超订 (#1328) - #1333

Merged
TATP-233 merged 1 commit into
dev/issue-1042-manager-based-apifrom
perf/issue-1328-mujoco-pool-thread-sizing
Aug 27, 2026
Merged

perf(mujoco): 单卡默认 pool 线程数收敛到有效 CPU 数,消除 2× 超订 (#1328)#1333
TATP-233 merged 1 commit into
dev/issue-1042-manager-based-apifrom
perf/issue-1328-mujoco-pool-thread-sizing

Conversation

@TATP-233

Copy link
Copy Markdown
Collaborator

Fixes #1328

改动

  • src/unilab/base/backend/mujoco/backend.py:新增 _effective_cpu_count()len(os.sched_getaffinity(0)),尊重 taskset/cgroup;macOS 等无此 syscall 的平台回退 os.cpu_count()),单卡(cpu_ids=None)默认 sizing 从 min(num_envs, 2*cpu_count()) 改为 min(num_envs, _effective_cpu_count())。显式 cpu_ids 路径不变(nthread = len(cpu_ids))。
  • tests/base/backend/test_mujoco_cpu_affinity_wiring.py:新增两条 sizing 规则单测(默认 = 有效 CPU 数;被 num_envs 截断)。
  • tests/base/backend/test_mujoco_chunk_size_wiring.py:原有测试改为 monkeypatch _effective_cpu_count(原先 patch cpu_count,在 ≥4 核机器上对本 sizing 已失效,会导致 benchmark 分支被短路)。

chunk tuner 的 cache key 含 nthreadchunk_tuner.py),默认线程数变化后自动重调,无脏 cache。

Validation

  • make test-all @ df9def32286 passed, 28 skipped, 1 xfailed(All checks passed)。
    • 过程中发现并清理了工作区里 src/unilab/envs/{locomotion,manipulation,motion_tracking} 三个纯 __pycache__ 残留目录(非 git 跟踪文件),它们以 namespace package 形式导致 test_removed_legacy_env_packages_stay_removed 假失败;与代码改动无关。
    • 另注:在 origin/main 上验证时发现一个 main 预置的测试污染——scripts/benchmark/env/benchmark_env_step.py import 时全局且不可撤销地 patch unilab.base.backend.create_backend(mujoco_warp 可导入时),致 tests/base/test_motrix_backend_options.py 两个用例在全量套件中失败(单跑通过);base 分支(dev/issue-1042)上不复现。建议单独建 issue 跟进。
  • 真实训练前后对比(本机 9950X3D + RTX 4090,32 逻辑核;uv run train --algo sac --task g1_motion_tracking --sim mujoco algo.num_envs=4096 algo.max_iterations=2000,tfevents 后半程 median):
指标 before (nthread=64) after (nthread=32) 变化
perf/steps_per_sec 86,099 / 86,162(两次) 96,262 +11.7%
perf/iter_ms 47.6 42.6 −10.5%
timing/collector_env_step_ms 47.7 41.9 −12.2%
timing/learner_train_ms 20.3 20.4 持平
整机 CPU%(/proc/stat,稳态 median) 37.8% 45.2% +7.4pt

满足 acceptance:默认 nthread == min(num_envs, 有效 CPU 数);显式 cpu_ids 仍为 len(cpu_ids);稳态 steps/s +11.7%(≥5%)。本机为 collector-bound(learner 仅 ~20ms/iter),收益集中在 env_step 物理相,与 issue 的 Linux 机型诊断一致。

范围外(issue non-goals,未动)

…rsubscription (#1328)

Single-GPU (cpu_ids=None) BatchEnvPool sizing now uses the CPUs actually
usable by this process (len(os.sched_getaffinity(0)), falling back to
os.cpu_count() where unavailable) instead of 2*cpu_count(). Measured on
this machine (sac + g1_motion_tracking + mujoco, num_envs=4096, 2000 iters,
steady-state second-half median): 86,162 -> 96,262 steps/s (+11.7%),
collector env_step 47.6ms -> 41.9ms. Explicit cpu_ids still fixes
nthread = len(cpu_ids).
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