From a352acc07cd6db897c7a458c2be6bb07ded3e260 Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Wed, 21 Oct 2026 23:35:47 +0800 Subject: [PATCH 1/4] fix(envs): attribute the autoreset done-gate host read to step core timing For an event-backed asynchronous backend step (mjwarp), step_tensor only launches device work; the physics wait surfaces at the first host read of the step outputs, which is the autoreset done-mask gate. That window was measured inside reset_done_ms, so the collector Env Step breakdown showed real device work as Reset Done while Backend Step collapsed to launch overhead and disappeared under the terminal materiality filter. Measure the gate separately and add it to step_core_ms/step_core_cpu_ms. Issue #2495. --- src/unilab/base/torch_env.py | 10 +++++++++- tests/base/test_torch_env.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/src/unilab/base/torch_env.py b/src/unilab/base/torch_env.py index d90ac69ed..19916915e 100644 --- a/src/unilab/base/torch_env.py +++ b/src/unilab/base/torch_env.py @@ -317,8 +317,16 @@ def step(self, actions: torch.Tensor) -> TorchEnvState: phase = time.perf_counter() phase_cpu = _cpu_time() # Preserve the semantic autoreset gate: an empty done mask must not - # consume reset RNG or publish a terminal carrier. + # consume reset RNG or publish a terminal carrier. The done-mask host + # read is the first synchronization point after an event-backed + # asynchronous backend step, so its wait belongs to the backend step + # phase rather than to reset bookkeeping. did_reset = self._autoreset and bool((self._state.terminated | self._state.truncated).any()) + step_core_ms += (time.perf_counter() - phase) * 1000.0 + step_core_cpu_ms += (_cpu_time() - phase_cpu) * 1000.0 + + phase = time.perf_counter() + phase_cpu = _cpu_time() if did_reset: self._reset_done_envs() reset_done_ms = (time.perf_counter() - phase) * 1000.0 diff --git a/tests/base/test_torch_env.py b/tests/base/test_torch_env.py index 1e8659191..5bc73785e 100644 --- a/tests/base/test_torch_env.py +++ b/tests/base/test_torch_env.py @@ -3,6 +3,8 @@ from __future__ import annotations import dataclasses +import itertools +import types from dataclasses import dataclass from typing import Any, cast from unittest.mock import MagicMock @@ -19,6 +21,7 @@ TensorRuntimeDiagnostic, ) +import unilab.base.torch_env as torch_env_module from unilab.base.base import EnvCfg, EnvPlayCapabilities from unilab.base.scene import SceneCfg from unilab.base.torch_env import TorchEnv, TorchEnvState @@ -220,6 +223,31 @@ def test_step_without_done_preserves_eager_autoreset_gate() -> None: assert state.info["timing"]["reset_done_count"] == 0.0 +def test_step_done_gate_host_read_is_attributed_to_step_core(monkeypatch) -> None: + wall = itertools.count(start=0.0, step=0.001) + cpu = itertools.count(start=0.0, step=0.001) + monkeypatch.setattr( + torch_env_module, + "time", + types.SimpleNamespace(perf_counter=lambda: next(wall)), + ) + monkeypatch.setattr(torch_env_module, "_cpu_time", lambda: next(cpu)) + + env = _StubTorchEnv(terminate=False) + env.init_state() + state = env.step(_actions(env)) + + timing = state.info["timing"] + # The fake clocks advance 1 ms per call, so each measured window spans + # exactly 1 ms. The backend call window and the done-gate host read both + # belong to step_core; reset bookkeeping keeps only its own empty window. + assert timing["step_core_ms"] == pytest.approx(2.0) + assert timing["step_core_cpu_ms"] == pytest.approx(2.0) + assert timing["reset_done_ms"] == pytest.approx(1.0) + assert timing["reset_done_cpu_ms"] == pytest.approx(1.0) + assert timing["env_step_other_cpu_ms"] >= 0.0 + + def test_reset_done_count_remains_device_resident_before_timing_publication() -> None: env = _StubTorchEnv() env.init_state() From 24b8869610d5d54c3adf44bb3feaa90193fb0d2c Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Wed, 21 Oct 2026 23:35:47 +0800 Subject: [PATCH 2/4] feat(logging): show CUDA MPS status in the off-policy System column Render the rank-local cuda_process_sharing runtime-manifest evidence as a CUDA MPS row: on for a validated request, ambient for a transparent environment-level daemon, off otherwise. This closes the terminal-side observation gap from #2479 where environment-level MPS silently changed benchmark behavior without any run-level signal. Issue #2495. --- src/unilab/rl/logging/offpolicy.py | 14 +++++++++ tests/rl/logging/test_offpolicy_logger.py | 35 +++++++++++++++++++++++ 2 files changed, 49 insertions(+) diff --git a/src/unilab/rl/logging/offpolicy.py b/src/unilab/rl/logging/offpolicy.py index f4368f2d5..037364dce 100644 --- a/src/unilab/rl/logging/offpolicy.py +++ b/src/unilab/rl/logging/offpolicy.py @@ -872,6 +872,17 @@ def _build_reward_table(self, snapshot: _TerminalSnapshot | None = None) -> Tabl source_iteration=max(source_iterations) if source_iterations else None, ) + def _cuda_mps_status_label(self) -> str | None: + """Render rank-local CUDA MPS state from runtime-manifest evidence.""" + evidence = self._runtime_manifest.get("cuda_process_sharing") + if not isinstance(evidence, dict): + return None + if evidence.get("effective") == "mps": + return "[green]on[/]" + if evidence.get("ambient_server"): + return "[yellow]ambient[/]" + return "[dim]off[/]" + def _build_timing_table(self, snapshot: _TerminalSnapshot | None = None) -> Table: snapshot = snapshot or self._terminal_snapshot terminal_width = self._console.width @@ -1006,6 +1017,9 @@ def _show_collector_row(key: str) -> bool: system_items.append(("Envs", f"{self.num_envs:,}")) if batch_size_per_rank > 0: system_items.append(("Batch/Rank", f"{batch_size_per_rank:,}")) + cuda_mps_status = self._cuda_mps_status_label() + if cuda_mps_status is not None: + system_items.append(("CUDA MPS", cuda_mps_status)) column_items: tuple[list[tuple[str, str]], ...] if terminal_width >= 120: column_items = (learner_items, collector_items, system_items) diff --git a/tests/rl/logging/test_offpolicy_logger.py b/tests/rl/logging/test_offpolicy_logger.py index 152a37eef..be3ab07a8 100644 --- a/tests/rl/logging/test_offpolicy_logger.py +++ b/tests/rl/logging/test_offpolicy_logger.py @@ -324,6 +324,41 @@ def test_offpolicy_terminal_render_is_width_aware() -> None: assert "Wait Ms" not in output +def test_offpolicy_logger_system_column_shows_cuda_mps_status() -> None: + logger = OffPolicyLogger(log_backend="none") + logger._console = Console(width=140, record=True, force_terminal=False) + + def system_cells() -> tuple[list[str], list[str]]: + table = logger._build_timing_table() + labels = [cell for cell in table.columns[4].cells if cell] + values = [cell for cell in table.columns[5].cells if cell] + return labels, values + + labels, _ = system_cells() + assert "CUDA MPS" not in labels + + logger.update_runtime_manifest( + {"cuda_process_sharing": {"effective": "mps", "ambient_server": False}} + ) + labels, values = system_cells() + assert labels[-1] == "CUDA MPS" + assert values[-1] == "[green]on[/]" + + logger.update_runtime_manifest( + {"cuda_process_sharing": {"effective": None, "ambient_server": True}} + ) + labels, values = system_cells() + assert labels[-1] == "CUDA MPS" + assert values[-1] == "[yellow]ambient[/]" + + logger.update_runtime_manifest( + {"cuda_process_sharing": {"effective": None, "ambient_server": False}} + ) + labels, values = system_cells() + assert labels[-1] == "CUDA MPS" + assert values[-1] == "[dim]off[/]" + + def test_offpolicy_logger_rejects_retired_collector_timing_names() -> None: logger = OffPolicyLogger(log_backend="none") From bfd228d9cab60f24ff64e641359fe2a2531138fa Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Wed, 21 Oct 2026 23:35:47 +0800 Subject: [PATCH 3/4] docs(logging): document Backend Step sync attribution and the CUDA MPS row Issue #2495. --- .../source/en/2-user_guide/1-training/3-logging.md | 11 ++++++++--- .../source/zh_CN/2-user_guide/1-training/3-logging.md | 7 ++++++- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/docs/sphinx/source/en/2-user_guide/1-training/3-logging.md b/docs/sphinx/source/en/2-user_guide/1-training/3-logging.md index e6532f765..78fb3c5d2 100644 --- a/docs/sphinx/source/en/2-user_guide/1-training/3-logging.md +++ b/docs/sphinx/source/en/2-user_guide/1-training/3-logging.md @@ -82,8 +82,10 @@ The terminal has explicit width tiers: uses `Iter Wall` as its denominator. - `Collector (own clock)` is measured in the collector subprocess and runs in parallel with the learner. - - `System` contains the buffer size, timeout rate, environment count, and - per-rank batch size. + - `System` contains the buffer size, timeout rate, environment count, + per-rank batch size, and the rank-local CUDA MPS status (`on` for a + validated request, `ambient` for a transparent environment-level daemon, + `off` otherwise). Metric rendering separates semantic groups: losses and policy state use stable short labels, training diagnostics are indented, and replay-ingress gauges render as @@ -265,7 +267,10 @@ The terminal keeps the core cycle rows above visible. `Inference Request`, canonical TensorBoard / W&B fields, but appear in the terminal only when they reach at least 1% of the collector cycle. `Backend Step`, `Update State`, and `Reset Done` follow the same materiality rule because they are nested diagnostics -rather than additional cycle slices. +rather than additional cycle slices. For an event-backed asynchronous backend +(mjwarp), the device physics wait surfaces at the first host read of the step +outputs — the autoreset done-mask gate — and is attributed to `Backend Step`, so +the row reports real device work instead of only the kernel-launch overhead. `Learner Action Wait` is deliberately not named “Inference Wait”: it is not pure inference latency. If the collector finishes `Env Step + Replay Write` and submits its diff --git a/docs/sphinx/source/zh_CN/2-user_guide/1-training/3-logging.md b/docs/sphinx/source/zh_CN/2-user_guide/1-training/3-logging.md index 428f77398..632302fe3 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/1-training/3-logging.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/1-training/3-logging.md @@ -76,7 +76,9 @@ key 的契约如下： - 120 列及以上时，timing 使用完整三列： - `Learner (Iter Wall)`：learner 主线程的一圈；所有行使用 `Iter Wall` 作分母。 - `Collector (own clock)`：collector 子进程自己的采集时钟，与 learner 并行。 - - `System`：buffer 大小、timeout rate、env 数和每 rank batch 大小。 + - `System`：buffer 大小、timeout rate、env 数、每 rank batch 大小，以及 rank + 本地 CUDA MPS 状态（验证通过的请求显示 `on`，环境级透明 daemon 显示 + `ambient`，否则 `off`）。 Metric 渲染按语义分组：loss 与 policy 状态使用稳定短标签，training 诊断缩进， replay-ingress gauge 显示为 `Ingress Backpressure`、`Ingress Depth` 等标签；它们 @@ -238,6 +240,9 @@ SAC / FlashSAC 每个 vectorized env tick 记录热路径互斥阶段，终端 `Action+Backend` 明细属于安静诊断项：它们仍是 canonical TensorBoard / W&B 字段， 但只有达到 collector cycle 的 1% 时才出现在终端。`Backend Step`、 `Update State` 与 `Reset Done` 是嵌套诊断而非额外 cycle 切片，同样按该显著性规则显示。 +对于事件型异步后端（mjwarp），设备物理等待发生在 step 输出的第一次 host 读取 +——autoreset 的 done-mask 门控——并归因到 `Backend Step`，因此该行反映真实设备 +工作而非仅 kernel 发射开销。 `Learner Action Wait` 特意不叫 “Inference Wait”：它不是纯 inference latency。如果 collector 在 learner update 期间先完成 `Env Step + Replay Write` 并提交下一 request，这一项会包含 From 286d809104c58dbdf3be8affcad398c1877472fb Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Thu, 22 Oct 2026 11:26:00 +0800 Subject: [PATCH 4/4] fix(envs): re-attribute the update-state queue drain to step core timing The first host read of an event-backed asynchronous backend step (mjwarp) lands inside ManagerBasedRlEnv.update_state: the termination phase window, already measured as update_state_queue_drain_ms, is where the host drains the queued device work. Move that window from update_state_ms into step_core_ms via the new TorchEnv._collect_step_core_drain_ms protocol so the collector Env Step breakdown shows the physics wait as Backend Step (2048-env probe: step_core 0.05 -> 5.52 ms, update_state 8.5 -> 3.8 ms). Only receipt-publishing backends re-attribute; synchronous producers drain nothing and keep their existing attribution. Issue #2495. --- .../en/2-user_guide/1-training/3-logging.md | 5 +-- .../2-user_guide/1-training/3-logging.md | 3 +- src/unilab/base/torch_env.py | 29 ++++++++++++--- src/unilab/envs/manager_based_rl_env.py | 14 ++++++++ tests/base/test_torch_env.py | 35 +++++++++++++++++++ tests/envs/test_manager_based_rl_env.py | 3 ++ 6 files changed, 82 insertions(+), 7 deletions(-) diff --git a/docs/sphinx/source/en/2-user_guide/1-training/3-logging.md b/docs/sphinx/source/en/2-user_guide/1-training/3-logging.md index 78fb3c5d2..d1672e1df 100644 --- a/docs/sphinx/source/en/2-user_guide/1-training/3-logging.md +++ b/docs/sphinx/source/en/2-user_guide/1-training/3-logging.md @@ -269,8 +269,9 @@ reach at least 1% of the collector cycle. `Backend Step`, `Update State`, and `Reset Done` follow the same materiality rule because they are nested diagnostics rather than additional cycle slices. For an event-backed asynchronous backend (mjwarp), the device physics wait surfaces at the first host read of the step -outputs — the autoreset done-mask gate — and is attributed to `Backend Step`, so -the row reports real device work instead of only the kernel-launch overhead. +outputs — the update-state termination drain, or the autoreset done-mask gate +when update_state performs no host read — and is attributed to `Backend Step`, +so the row reports real device work instead of only the kernel-launch overhead. `Learner Action Wait` is deliberately not named “Inference Wait”: it is not pure inference latency. If the collector finishes `Env Step + Replay Write` and submits its diff --git a/docs/sphinx/source/zh_CN/2-user_guide/1-training/3-logging.md b/docs/sphinx/source/zh_CN/2-user_guide/1-training/3-logging.md index 632302fe3..a37fd0b55 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/1-training/3-logging.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/1-training/3-logging.md @@ -241,7 +241,8 @@ SAC / FlashSAC 每个 vectorized env tick 记录热路径互斥阶段，终端 但只有达到 collector cycle 的 1% 时才出现在终端。`Backend Step`、 `Update State` 与 `Reset Done` 是嵌套诊断而非额外 cycle 切片，同样按该显著性规则显示。 对于事件型异步后端（mjwarp），设备物理等待发生在 step 输出的第一次 host 读取 -——autoreset 的 done-mask 门控——并归因到 `Backend Step`，因此该行反映真实设备 +——update_state 内 termination 阶段的 queue drain（若 update_state 没有 host 读取， +则为 autoreset 的 done-mask 门控）——并归因到 `Backend Step`，因此该行反映真实设备 工作而非仅 kernel 发射开销。 `Learner Action Wait` 特意不叫 “Inference Wait”：它不是纯 inference latency。如果 collector diff --git a/src/unilab/base/torch_env.py b/src/unilab/base/torch_env.py index 19916915e..35f939a96 100644 --- a/src/unilab/base/torch_env.py +++ b/src/unilab/base/torch_env.py @@ -308,6 +308,14 @@ def step(self, actions: torch.Tensor) -> TorchEnvState: self._validate_state(self._state) update_state_ms = (time.perf_counter() - phase) * 1000.0 update_state_cpu_ms = (_cpu_time() - phase_cpu) * 1000.0 + if self.produces_completion_receipts: + # An event-backed asynchronous step drains its device work at the + # first update-state host read; re-attribute that measured window + # to the backend step phase so update_state reports manager work. + drain_ms = self._collect_step_core_drain_ms(self._state) + if drain_ms > 0.0: + step_core_ms += drain_ms + update_state_ms = max(update_state_ms - drain_ms, 0.0) self._state.info["steps"] += 1 self.step_counter += 1 @@ -317,10 +325,11 @@ def step(self, actions: torch.Tensor) -> TorchEnvState: phase = time.perf_counter() phase_cpu = _cpu_time() # Preserve the semantic autoreset gate: an empty done mask must not - # consume reset RNG or publish a terminal carrier. The done-mask host - # read is the first synchronization point after an event-backed - # asynchronous backend step, so its wait belongs to the backend step - # phase rather than to reset bookkeeping. + # consume reset RNG or publish a terminal carrier. When update_state + # made no host read of the step outputs, this gate is the first + # synchronization point after an event-backed asynchronous backend + # step, so its wait belongs to the backend step phase rather than to + # reset bookkeeping. did_reset = self._autoreset and bool((self._state.terminated | self._state.truncated).any()) step_core_ms += (time.perf_counter() - phase) * 1000.0 step_core_cpu_ms += (_cpu_time() - phase_cpu) * 1000.0 @@ -653,6 +662,18 @@ def _collect_reset_backend_timing_ms(self) -> dict[str, float]: """Return backend-owned sub-timings from the most recent tensor reset.""" return {} + def _collect_step_core_drain_ms(self, state: TorchEnvState) -> float: + """Return the update-state window that drained an async backend step. + + Only consulted when the backend publishes event-backed completion + receipts (``produces_completion_receipts``); synchronous backends drain + nothing, so the default is zero. An env whose update-state sub-phase + hosts the first read of asynchronous step outputs overrides this to + report that measured window. + """ + del state + return 0.0 + def _clear_step_final_observation(self) -> None: assert self._state is not None self._state.final_observation = None diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index e4a0ff419..9bd1ae60b 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -1469,6 +1469,20 @@ def _collect_reset_backend_timing_ms(self) -> dict[str, float]: timing.update(self._last_reset_manager_timing_ms) return timing + def _collect_step_core_drain_ms(self, state: TorchEnvState) -> float: + # Termination is the first update-state reduction after the backend + # step; for an event-backed asynchronous step its window (measured as + # update_state_queue_drain_ms) drains the queued device work, which is + # backend step time rather than Manager work. + timing = state.info.get("timing") + if not isinstance(timing, dict): + return 0.0 + value = timing.get("update_state_queue_drain_ms") + if isinstance(value, bool) or not isinstance(value, (int, float)): + return 0.0 + value_float = float(value) + return value_float if math.isfinite(value_float) else 0.0 + def reset( self, env_ids: torch.Tensor | None = None, diff --git a/tests/base/test_torch_env.py b/tests/base/test_torch_env.py index 5bc73785e..42b469203 100644 --- a/tests/base/test_torch_env.py +++ b/tests/base/test_torch_env.py @@ -248,6 +248,41 @@ def test_step_done_gate_host_read_is_attributed_to_step_core(monkeypatch) -> Non assert timing["env_step_other_cpu_ms"] >= 0.0 +def test_step_reattributes_update_state_drain_to_step_core(monkeypatch) -> None: + wall = itertools.count(start=0.0, step=0.001) + cpu = itertools.count(start=0.0, step=0.001) + monkeypatch.setattr( + torch_env_module, + "time", + types.SimpleNamespace(perf_counter=lambda: next(wall)), + ) + monkeypatch.setattr(torch_env_module, "_cpu_time", lambda: next(cpu)) + + class _DrainEnv(_StubTorchEnv): + def _collect_step_core_drain_ms(self, state: TorchEnvState) -> float: + return 0.6 + + async_backend = _backend() + async_backend.backend_type = "mjwarp" + env = _DrainEnv(backend=async_backend, terminate=False) + env.init_state() + timing = env.step(_actions(env)).info["timing"] + + # Fake clock: every window spans 1 ms. Backend launch (1 ms) + the drained + # update-state window (0.6 ms) + the done-gate host read (1 ms) all belong + # to step_core; update_state keeps only the remaining manager work. + assert timing["step_core_ms"] == pytest.approx(2.6) + assert timing["update_state_ms"] == pytest.approx(0.4) + assert timing["reset_done_ms"] == pytest.approx(1.0) + + # Synchronous producers publish no completion receipt and drain nothing. + sync_env = _DrainEnv(terminate=False) + sync_env.init_state() + sync_timing = sync_env.step(_actions(sync_env)).info["timing"] + assert sync_timing["step_core_ms"] == pytest.approx(2.0) + assert sync_timing["update_state_ms"] == pytest.approx(1.0) + + def test_reset_done_count_remains_device_resident_before_timing_publication() -> None: env = _StubTorchEnv() env.init_state() diff --git a/tests/envs/test_manager_based_rl_env.py b/tests/envs/test_manager_based_rl_env.py index 2519f3142..25382bb2c 100644 --- a/tests/envs/test_manager_based_rl_env.py +++ b/tests/envs/test_manager_based_rl_env.py @@ -2335,6 +2335,9 @@ def test_update_state_publishes_drain_vs_host_termination_attribution() -> None: assert timing["update_state_state_replace_ms"] >= 0.0 assert timing["update_state_child_sum_ms"] >= 0.0 assert timing["update_state_command_preflight_ms"] >= 0.0 + assert env._collect_step_core_drain_ms(state) == pytest.approx( + timing["update_state_queue_drain_ms"] + ) observation_timing = { key: state.info["update_state_timing"][key] for key in (