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..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 @@ -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,11 @@ 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 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 428f77398..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 @@ -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,10 @@ SAC / FlashSAC 每个 vectorized env tick 记录热路径互斥阶段，终端 `Action+Backend` 明细属于安静诊断项：它们仍是 canonical TensorBoard / W&B 字段， 但只有达到 collector cycle 的 1% 时才出现在终端。`Backend Step`、 `Update State` 与 `Reset Done` 是嵌套诊断而非额外 cycle 切片，同样按该显著性规则显示。 +对于事件型异步后端（mjwarp），设备物理等待发生在 step 输出的第一次 host 读取 +——update_state 内 termination 阶段的 queue drain（若 update_state 没有 host 读取， +则为 autoreset 的 done-mask 门控）——并归因到 `Backend Step`，因此该行反映真实设备 +工作而非仅 kernel 发射开销。 `Learner Action Wait` 特意不叫 “Inference Wait”：它不是纯 inference latency。如果 collector 在 learner update 期间先完成 `Env Step + Replay Write` 并提交下一 request，这一项会包含 diff --git a/src/unilab/base/torch_env.py b/src/unilab/base/torch_env.py index d90ac69ed..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,8 +325,17 @@ 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. 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 + + phase = time.perf_counter() + phase_cpu = _cpu_time() if did_reset: self._reset_done_envs() reset_done_ms = (time.perf_counter() - phase) * 1000.0 @@ -645,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/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/base/test_torch_env.py b/tests/base/test_torch_env.py index 1e8659191..42b469203 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,66 @@ 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_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 ( 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")