perf: env_step sub-stage timing panel + optional numba TBB layer (+23% throughput on many-core) (#72)

Author: wlgys8

docs/source/en/user_guide/getting_started/installation.md

@@ -92,6 +92,7 @@ training backend extra on top of them, and each option only selects or appends t
 | `--gpu` | `cuda`<br>`rocm` | Select the torch wheel flavor explicitly, overriding auto-detection |
 | `--skrl-jax` | — | SKRL on JAX backend, Linux only |
 | `--rslrl` | — | RSL-RL on PyTorch backend |
+| `--tbb` | — | Install Intel TBB for the numba parallel kernels; recommended for training on many-core servers (64+ cores), auto-activates once installed |
 | `--docs` | — | Add the toolchain (sphinx) needed to build the documentation locally |
 | `-h`, `--help` | — | Show the help message |
 

docs/source/zh_CN/user_guide/getting_started/installation.md

@@ -114,6 +114,7 @@ sh install.sh --all
 | `--gpu` | `cuda`<br>`rocm` | 指定 torch wheel 来源，覆盖自动探测 |
 | `--skrl-jax` | — | SKRL（JAX）训练后端，仅 Linux |
 | `--rslrl` | — | RSL-RL（PyTorch）训练后端 |
+| `--tbb` | — | 为 numba 并行 kernel 安装 Intel TBB 线程层；多核服务器（64 核以上）训练推荐启用，安装后自动生效 |
 | `--docs` | — | 追加本地构建文档所需的工具链（Sphinx） |
 | `-h`、`--help` | — | 显示帮助 |
 

install.sh

@@ -11,6 +11,8 @@
 #
 # Training backends are enabled with --skrl-torch (default), --skrl-jax or
 # --rslrl, named after the extras that install them; multiple flags combine.
+# --tbb installs Intel TBB so numba parallel kernels use its threading layer
+# (avoids the OpenMP per-region thread-wakeup storm on many-core hosts).
 
 set -e
 
@@ -25,6 +27,7 @@ usage() {
     echo "  --skrl-torch           Training backend: SKRL on PyTorch (default)"
     echo "  --skrl-jax             Training backend: SKRL on JAX (Linux only)"
     echo "  --rslrl                Training backend: RSL-RL on PyTorch"
+    echo "  --tbb                  Install Intel TBB for numba parallel kernels (recommended on many-core hosts)"
     echo "  -h, --help             Show this help"
 }
 
@@ -34,6 +37,7 @@ SKRL_TORCH=""
 SKRL_JAX=""
 RSLRL=""
 DOCS=""
+TBB=""
 
 while [ $# -gt 0 ]; do
     case "$1" in
@@ -57,6 +61,9 @@ while [ $# -gt 0 ]; do
         --rslrl)
             RSLRL=1
             ;;
+        --tbb)
+            TBB=1
+            ;;
         -h|--help)
             usage
             exit 0
@@ -162,6 +169,7 @@ else
     [ -n "$SKRL_JAX" ] && EXTRAS="$EXTRAS --extra skrl-jax"
     [ -n "$RSLRL" ] && EXTRAS="$EXTRAS --extra rslrl"
     [ -n "$DOCS" ] && EXTRAS="$EXTRAS --extra docs"
+    [ -n "$TBB" ] && EXTRAS="$EXTRAS --extra tbb"
     set -x
     uv sync --all-packages --no-default-groups --extra "$GPU"$EXTRAS
 fi

motrix_env_core/pyproject.toml

@@ -21,3 +21,10 @@ dependencies = [
     "omegaconf>=2.3,<2.4",
     "typing-extensions>=4.1",
 ]
+
+[project.optional-dependencies]
+# TBB threading layer for numba parallel kernels: on many-core machines the
+# default OpenMP layer pays a per-region thread-wakeup storm that dominates
+# short kernels. The wheel drops libtbb.so.12 into <venv>/lib, so the runtime
+# preloads it by full path (see motrix_env_core.numba.threading).
+tbb = ["tbb==2021.13.0"]

motrix_env_core/src/motrix_env_core/numba/manager/env.py

@@ -52,6 +52,7 @@
     TerminationTermCfg,
 )
 from motrix_env_core.numba.program import NumbaTaskProgram
+from motrix_env_core.numba.threading import preload_tbb
 from motrix_env_core.sim import (
     ModelQuery,
     PhysicsReadProgram,
@@ -343,6 +344,12 @@ def __init__(self, cfg: EnvCfgType, num_envs: int = 1, backend: str | None = Non
         """Construct a manager environment using the selected simulator backend."""
         if not isinstance(cfg, ManagerBasedEnvCfg):
             raise TypeError(f"{type(cfg).__name__} must inherit ManagerBasedEnvCfg.")
+        # Prefer the TBB threading layer for the parallel step kernels when the
+        # tbb extra is installed: the default OpenMP layer pays a thread-wakeup
+        # storm per kernel call that dominates short kernels on many-core
+        # machines. Must happen before the first parallel kernel executes.
+        if preload_tbb():
+            logger.info("Manager env %r: numba TBB threading layer available", type(self).__name__)
         super().__init__(cfg, num_envs)
         self._rand_seed = 1 if seed is None else seed
         from motrix_env_core.sim.registry import create_sim_backend, default_sim_backend_name

motrix_env_core/src/motrix_env_core/numba/threading.py

@@ -0,0 +1,59 @@
+# Copyright Motphys Technology Co., Ltd. 2025, 2026
+# SPDX-License-Identifier: Apache-2.0
+
+"""Runtime threading-layer helpers for numba parallel kernels.
+
+On many-core machines the default OpenMP layer pays a per-region thread-wakeup
+storm (``OMP_WAIT_POLICY=PASSIVE`` parks every worker between calls, and each
+``prange`` region must wake them all before any work starts). For short kernels
+called every environment step this overhead dominates the kernel itself. The
+TBB layer keeps a resident work-stealing pool and does not have this problem.
+
+The PyPI ``tbb`` wheel installs ``libtbb.so.12`` into ``<venv>/lib`` instead of
+``site-packages``, so numba's plain ``CDLL("libtbb.so.12")`` lookup cannot find
+it without ``LD_LIBRARY_PATH``. Preloading the library by full path with
+``RTLD_GLOBAL`` registers its SONAME in the process, and numba's later lookup by
+the same SONAME reuses the loaded handle.
+"""
+
+from __future__ import annotations
+
+import ctypes
+import glob
+import sys
+from pathlib import Path
+
+_TBB_SONAME = "libtbb.so.12"
+_tbb_preload_result: bool | None = None
+
+
+def preload_tbb() -> bool:
+    """Preload ``libtbb.so.12`` from the active environment, if present.
+
+    Returns ``True`` when the library was loaded (or was already loaded), which
+    lets numba select the TBB threading layer for parallel kernels. Returns
+    ``False`` when the ``tbb`` extra is not installed — callers keep the numba
+    default layer in that case. The probe runs once and is cached: every
+    ``ManagerEnv`` construction calls this, and the glob + CDLL work should
+    not repeat per instance.
+    """
+    global _tbb_preload_result
+    if _tbb_preload_result is None:
+        _tbb_preload_result = _probe_tbb()
+    return _tbb_preload_result
+
+
+def _probe_tbb() -> bool:
+    if not sys.platform.startswith("linux"):
+        return False
+    candidates = sorted(glob.glob(str(Path(sys.prefix) / "lib" / f"{_TBB_SONAME}*")), reverse=True)
+    for path in candidates:
+        try:
+            ctypes.CDLL(path, mode=ctypes.RTLD_GLOBAL)
+        except OSError:
+            continue
+        return True
+    return False
+
+
+__all__ = ["preload_tbb"]

motrix_env_core/src/motrix_env_core/perf.py

@@ -187,6 +187,27 @@ def snapshot(self) -> tuple[PerfNode, ...]:
             raise RuntimeError("Cannot snapshot Perf while a scope is active.")
         return tuple(root.snapshot() for root in self._roots.values())
 
+    def stage_mean_ms(self, root: str) -> dict[str, float]:
+        """Per-call mean milliseconds of every sub-scope under ``root``.
+
+        Keys are dotted paths relative to the root's children (``stage`` or
+        ``stage.sub``), so nested scope names — which may themselves contain
+        underscores — stay unambiguously separable. Returns ``{}`` when the
+        root has not run.
+        """
+        for node in self._roots.values():
+            if node.name == root:
+                means: dict[str, float] = {}
+
+                def emit(subtree: _MutablePerfNode, path: tuple[str, ...]) -> None:
+                    for child in subtree.children.values():
+                        means[".".join((*path, child.name))] = child.total_ns / 1e6 / max(child.count, 1)
+                        emit(child, (*path, child.name))
+
+                emit(node, ())
+                return means
+        return {}
+
     def call(self, name: str, function: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs) -> _R:
         """Call a function inside a named scope."""
         if not self._enabled:

motrix_env_core/tests/test_numba_threading.py

@@ -0,0 +1,10 @@
+# Copyright Motphys Technology Co., Ltd. 2025, 2026
+# SPDX-License-Identifier: Apache-2.0
+
+from motrix_env_core.numba.threading import preload_tbb
+
+
+def test_preload_tbb_returns_bool_without_raising() -> None:
+    # Whether the library exists depends on the tbb extra; the contract is
+    # only that discovery is side-effect-safe and reports success as bool.
+    assert isinstance(preload_tbb(), bool)

motrix_env_core/tests/test_perf.py

@@ -139,3 +139,29 @@ def test_active_perf_scope_contributes_to_activated_profiler() -> None:
 def test_active_perf_scope_is_disabled_without_an_activated_profiler() -> None:
     with active_perf_scope("backend"):
         pass
+
+
+def test_stage_mean_ms_flattens_nested_scopes_to_dotted_paths() -> None:
+    # two step calls; each scope entry/exit consumes one clock tick (8 per
+    # iteration). physics runs 3ms per call, transition 9ms with a 2ms read.
+    base = (0, 2_000_000, 5_000_000, 6_000_000, 10_000_000, 12_000_000, 15_000_000, 26_000_000)
+    ticks = base + tuple(t + 26_000_000 for t in base)
+    perf = Perf(enabled=True, clock=_Clock(*ticks))
+    for _ in range(2):
+        with perf.scope("step"):
+            with perf.scope("physics"):
+                pass
+            with perf.scope("transition"):
+                with perf.scope("read"):
+                    pass
+
+    means = perf.stage_mean_ms("step")
+
+    assert means["physics"] == pytest.approx(3.0)
+    assert means["transition"] == pytest.approx(9.0)
+    assert means["transition.read"] == pytest.approx(2.0)
+    assert "step" not in means  # the root itself is not part of the paths
+
+
+def test_stage_mean_ms_returns_empty_for_unrun_root() -> None:
+    assert Perf(enabled=True).stage_mean_ms("step") == {}

motrix_env_motrixsim/src/motrix_env_motrixsim/sim_data.py

@@ -11,6 +11,7 @@
 import numpy as np
 import numpy.typing as npt
 
+from motrix_env_core.perf import active_perf_scope
 from motrix_env_core.sim import PhysicsReadProgram, SimDataQuery, SimDataQueryCompiler
 
 FloatArray: TypeAlias = npt.NDArray[np.float32]
@@ -370,8 +371,10 @@ def execute(self, env_ids: npt.NDArray[np.int64] | None = None) -> None:
         if env_ids is not None and (env_ids.dtype != np.int64 or env_ids.ndim != 1):
             raise TypeError("Partial simulator read env_ids must be a one-dimensional int64 ndarray.")
         # One native call writes the authoritative arena — the full batch, or
-        # only the rows selected by env_ids.
-        self._read.program.execute(self._source, env_ids=env_ids)
+        # only the rows selected by env_ids. The perf scope nests under the
+        # caller's stage (e.g. transition/reset in the console timing tree).
+        with active_perf_scope("read"):
+            self._read.program.execute(self._source, env_ids=env_ids)
 
 
 def compile_read_program(