unilab.base.cpu_runtime¶
Process-local CPU confinement for envs that own an explicit CPU block.
Multi-rank off-policy data-parallel runs partition host CPUs so each rank’s
collector owns one contiguous block (training.dp_collector_cpu_ids, routed
into EnvCfg.cpu_ids). The MuJoCo BatchEnvPool already pins its physics
workers to that block, but the collector’s host-side compute did not follow:
Numba parallel kernels (unisim.backend.body_state and the
motion-tracking kernels) size their pool from the host CPU count and leave
placement to the OS, so they drift across rank boundaries and compete with
sibling ranks’ pinned physics workers.
apply_env_cpu_runtime is the generic env-level counterpart, applied once on
the env-construction cold path — before managers, backend materialization, or
the first Numba parallel call exist:
os.sched_setaffinitypins the calling thread to the block, and every already-running thread (e.g. BLAS pools spawned atimport numpy) is pinned individually via/proc/self/task; threads spawned later (including Numba’s lazily-launched pool) inherit the mask.numba.set_num_threads(len(cpu_ids))sizes Numba’s pool to the block instead of the host, unless the operator pinnedNUMBA_NUM_THREADS(mirroring the motion-kernel runtime policy).
The function is backend-agnostic: EnvCfg.cpu_ids is the single source of
truth, so any backend whose env declares a block gets the same confinement.
Functions
|
Confine this process's host-side compute to the env-owned CPU block. |
- unilab.base.cpu_runtime.apply_env_cpu_runtime(cpu_ids)[source]¶
Confine this process’s host-side compute to the env-owned CPU block.
Cold path only: call from env construction, before the backend pool, managers, or any Numba parallel kernel exist.
None(the single-rank default) is a no-op so the default path stays bit-identical.Structural validation (non-empty, unique, non-negative ints) is owned by
EnvCfg.validate; this function fails closed on CPU ids that are not available to the process.