From 10c20a8dd7434ee5a4236c0b905297e4ca1748e9 Mon Sep 17 00:00:00 2001 From: zilch Date: Mon, 19 Oct 2026 01:23:34 +0800 Subject: [PATCH] fix(fastsac): clarify first-update readiness and throughput status --- motrix_rl/src/motrix_rl/console.py | 12 ++-- motrix_rl/tests/test_console.py | 19 ++++++ .../fastsac/async_impl/train.py | 32 ++++++---- .../fastsac/async_impl/worker.py | 6 +- .../tests/test_fastsac_live_metrics.py | 61 ++++++++++++++++++- 5 files changed, 112 insertions(+), 18 deletions(-) diff --git a/motrix_rl/src/motrix_rl/console.py b/motrix_rl/src/motrix_rl/console.py index 4bbf69b2..a1d3a850 100644 --- a/motrix_rl/src/motrix_rl/console.py +++ b/motrix_rl/src/motrix_rl/console.py @@ -57,7 +57,7 @@ class TrainingPanelStats: learn_ms: float learn_percent: float warming: bool = False - # Rolling-window iteration rate; None = same "warming up" semantics as + # Rolling-window iteration rate; None = same unavailable-rate semantics as # steps_per_second. Falls back to the cumulative average when absent. iterations_per_second: float | None = None training_metrics: Mapping[str, Any] | None = None @@ -194,7 +194,7 @@ def _iter_rate_text(stats: TrainingPanelStats) -> str: if stats.iterations_per_second is not None: return f"{stats.iterations_per_second:,.0f} iter/s" if stats.steps_per_second is None: - return "warming up" + return "warming up" if stats.warming else "measuring…" return f"{stats.iteration / max(stats.elapsed_seconds, 1e-9):,.0f} iter/s" @@ -214,7 +214,11 @@ def si(n: float) -> str: pct = 100.0 * stats.iteration / max(stats.total_iterations, 1) eta = _eta_seconds(stats) eta_text = f" - eta ~{hms(eta)}" if eta is not None else "" - sps_text = f"{stats.steps_per_second:.0f} env-steps/s" if stats.steps_per_second is not None else "warming up" + sps_text = ( + f"{stats.steps_per_second:.0f} env-steps/s" + if stats.steps_per_second is not None + else ("warming up" if stats.warming else "measuring…") + ) header = ( f" {title} - iter {stats.iteration}/{stats.total_iterations} ({pct:.1f}%) - " f"{sps_text} - {hms(stats.elapsed_seconds)}{eta_text}" @@ -773,7 +777,7 @@ def memory_style(memory: MemoryUsage | None) -> str: Text( f"{stats.steps_per_second:,.0f} env-steps/s" if stats.steps_per_second is not None - else "warming up", + else ("warming up" if stats.warming else "measuring…"), style="bold cyan", ), Text(_iter_rate_text(stats), style="white"), diff --git a/motrix_rl/tests/test_console.py b/motrix_rl/tests/test_console.py index c64c0f07..5b9c8953 100644 --- a/motrix_rl/tests/test_console.py +++ b/motrix_rl/tests/test_console.py @@ -37,6 +37,25 @@ def test_format_value_uses_scientific_notation_for_large_floats() -> None: assert _format_value(-1_234_567.0, signed=True) == "-1.235e+06" +@pytest.mark.parametrize("rich", [False, True]) +@pytest.mark.parametrize( + ("rate", "warming", "expected"), + [(None, False, "measuring…"), (None, True, "warming up"), (0.0, False, "0 env-steps/s")], +) +def test_throughput_distinguishes_missing_measurement_from_replay_warmup(rich, rate, warming, expected): + stats = _panel_stats(steps_per_second=rate, iterations_per_second=rate, warming=warming) + if rich: + console = Console(width=120, force_terminal=False) + with console.capture() as capture: + console.print(render_training_panel(stats)) + text = capture.get() + else: + text = format_training_panel(stats) + assert expected in text + if not warming: + assert "warming up" not in text + + def test_format_training_panel_groups_timing_metrics() -> None: stats = TrainingPanelStats( iteration=1, diff --git a/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/async_impl/train.py b/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/async_impl/train.py index 8ca0474d..b6a12047 100644 --- a/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/async_impl/train.py +++ b/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/async_impl/train.py @@ -474,10 +474,16 @@ def _drain_worker_stats() -> None: pass def _panel_data_ready() -> bool: - """The training panel's quiescence gate (see _render_panel).""" + """Wait for all learners' first updates, including lazy compilation.""" return ( all(v is not None for v in per_collector_stats.values()) - and any(v is not None for v in learner_payloads.values()) + # Initial payloads do not mean lazy compilation/first updates + # completed. Short runs may legitimately finish in replay warmup. + and all( + v is not None + and (v["training_started"] or control.collector_steps >= num_iterations * num_collectors) + for v in learner_payloads.values() + ) and control.collector_steps > 0 ) @@ -487,7 +493,8 @@ def _handoff_gate_line() -> str: f"c{i} {'✓' if per_collector_stats[i] is not None else '…'}" for i in range(num_collectors) ) payload_marks = " ".join( - f"l{r} {'✓' if learner_payloads[r] is not None else '…'}" for r in range(num_learners) + f"l{r} " + ("✓" if p is not None and p["training_started"] else "initializing first update…") + for r, p in learner_payloads.items() ) return f"panel data — stats: {stats_marks} | payloads: {payload_marks} | steps: {control.collector_steps}" @@ -547,6 +554,10 @@ def _handoff_gate_line() -> str: if abort_reason is None: while not _panel_data_ready(): _drain_worker_stats() + # A run shorter than replay warmup may end without an + # update payload; let normal shutdown collect its result. + if all(not p.is_alive() for p in p_learners): + break try: process_name, traceback_text = error_queue.get_nowait() except Empty: @@ -643,15 +654,12 @@ def _render_panel(force: bool = False) -> None: panel_rate_samples.pop(0) t0, s0 = panel_rate_samples[0] t1, s1 = panel_rate_samples[-1] - # Short-window init: as soon as steps flow, compute over the - # available span (floor 1s against div-by-tiny) so the panel - # shows a real — if rough — rate immediately; None (warming up) - # only means "no steps flowed in the window" (pre-start or a - # genuine throughput stall). - if s1 > s0: - span = max(t1 - t0, 1.0) - sps = (s1 - s0) * num_envs / span - iters = (s1 - s0) / span + # An empty initial window has no rate yet. Once a time span + # exists, zero progress is a measured zero, not replay warmup. + span = t1 - t0 + if span > 0: + sps = (s1 - s0) * num_envs / max(span, 1.0) + iters = (s1 - s0) / max(span, 1.0) else: sps = None iters = None diff --git a/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/async_impl/worker.py b/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/async_impl/worker.py index e818ed62..cd5cfab4 100644 --- a/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/async_impl/worker.py +++ b/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/async_impl/worker.py @@ -689,6 +689,7 @@ def run_learner_process( # global progress basis (identical on every DDP rank; the single # learner's own position when num_collectors == 1). trained = learner.maybe_train(step) + first_update = trained and not train_started if trained: train_started = True elapsed_learn_s = time.perf_counter() - t_l @@ -711,7 +712,9 @@ def run_learner_process( else: # Data arrived, but replay/batch readiness still gated training. learner_gate_wait_samples_ms.append(idle_ms) - if step >= next_log: + if step >= next_log or first_update: + # Publish first-update completion immediately: the startup + # handoff must not wait an entire user logging interval. # One compact payload per log window to the parent (panel + # TensorBoard live there; see train.py). All ranks send. now = time.time() @@ -745,6 +748,7 @@ def run_learner_process( "step": step, "sps": sps, "updates": learner.update_idx, + "training_started": train_started, "metrics": metrics_log, "warming": warming, "learn_ms": timing_mean(learner_train_samples_ms) if learner_train_samples_ms else 0.0, diff --git a/motrix_rl_builtin/tests/test_fastsac_live_metrics.py b/motrix_rl_builtin/tests/test_fastsac_live_metrics.py index ae7bc5eb..1de7b789 100644 --- a/motrix_rl_builtin/tests/test_fastsac_live_metrics.py +++ b/motrix_rl_builtin/tests/test_fastsac_live_metrics.py @@ -400,6 +400,7 @@ class _Stage(NamedTuple): payload: bool = True episodes: int = 2 mean_return: float = 1.5 + updates: int | None = None def _snapshot(collector_id: int, stage: _Stage) -> dict: @@ -423,6 +424,7 @@ def _payload(step: int, checkpoint_path: str | None) -> dict: "step": step, "sps": 100.0, "updates": 2 * step, + "training_started": True, "metrics": None if warming else {"qf_loss": 0.25, "actor_loss": -1.5}, "warming": warming, "learn_ms": 1.0, @@ -471,7 +473,14 @@ def learner(self, process) -> int: for collector in collectors: collector["stats_queue"].put(_snapshot(collector["collector_id"], stage)) if stage.payload: - args["panel_queue"].put(_payload(control.collector_steps // NUM_COLLECTORS, checkpoint_path)) + payload = _payload(control.collector_steps // NUM_COLLECTORS, checkpoint_path) + if stage.updates is not None: + payload["updates"] = stage.updates + payload["training_started"] = stage.updates > 0 + args["panel_queue"].put(payload) + if payload["updates"] == 0: + self._wait(lambda: args["panel_queue"].empty() or control.stop) + continue logged += 1 if self.live: self._wait(lambda: self.sink.count("scalars") >= logged or control.stop) @@ -551,6 +560,56 @@ def _train_async(tmp_path, monkeypatch, workers: _Workers, *, num_iterations: in return async_train.Trainer(context=context), writer +def test_async_panel_waits_for_first_learner_update(tmp_path, monkeypatch): + from motrix_rl_builtin.fastsac.async_impl import train as async_train + + sink = _RecordingSink(tmp_path / "metrics.jsonl") + workers = _Workers(sink, [_Stage((10, 10), updates=0), _Stage((20, 20))]) + trainer, _ = _train_async(tmp_path, monkeypatch, workers, num_iterations=30) + panels = [] + monkeypatch.setattr(async_train, "emit_training_panel", lambda live, stats, **kwargs: panels.append(stats)) + + trainer.train() + + assert panels + assert all(panel.iteration >= 20 for panel in panels) + + +def test_async_run_can_finish_without_any_training_update(tmp_path, monkeypatch): + sink = _RecordingSink(tmp_path / "metrics.jsonl") + workers = _Workers(sink, [_Stage((10, 10), updates=0)]) + trainer, _ = _train_async(tmp_path, monkeypatch, workers, num_iterations=10) + + trainer.train() + + assert checkpoints.final_checkpoint_path("pt", tmp_path).exists() + + +def test_async_panel_reports_zero_throughput_after_collection_stalls(tmp_path, monkeypatch): + from motrix_rl_builtin.fastsac.async_impl import train as async_train + + panels = [] + + class StalledWorkers(_Workers): + def learner(self, process): + result = super().learner(process) + # Hold fixed collection progress across two measured windows. + self._wait(lambda: len(panels) >= 3 or process.arguments["control"].stop) + return result + + sink = _RecordingSink(tmp_path / "metrics.jsonl") + workers = StalledWorkers(sink, [_Stage((20, 20))]) + trainer, _ = _train_async(tmp_path, monkeypatch, workers, num_iterations=30) + monkeypatch.setattr(async_train, "emit_training_panel", lambda live, stats, **kwargs: panels.append(stats)) + + trainer.train() + + assert len(panels) >= 3 + assert panels[-1].steps_per_second == 0.0 + assert panels[-1].iterations_per_second == 0.0 + assert not panels[-1].warming + + def test_async_trainer_emits_tensorboard_tags_events_and_collector_env_steps(tmp_path, monkeypatch): sink = _RecordingSink(tmp_path / "metrics.jsonl") # Uneven collectors in the first stage: 21 batches -> step 10, but 21 * 3 env steps.