feat(sc): telemetry for rollout checkpointing - #3925
Conversation
30ba01f to
a9f5575
Compare
| async with self._data_plane_checkpoint_barrier.mutation( | ||
| "group_removals" | ||
| ): |
There was a problem hiding this comment.
1 action item.
TL;DR — this new mutation("group_removals") wrapper makes the first stale-group eviction raise RuntimeError, killing the train pump on 8 of the 9 shipped single-controller configs.
PR-introduced: the wrapper is new here, the inner section is pre-existing.
DataPlaneCheckpointBarrier sections are deliberately non-reentrant so one task can't hold two cuts and lose track of which is live. _current_task() enforces that via _section_holders.
- This block opens a barrier section.
InOrderSampler.evict— andBaseSampler.evict, inherited byWindowedSamplerandWeightFifoSamplersince_GatedSamplerdoesn't override it — callself._buffer.remove(stale_idxs, remove_in_dp=True).TQReplayBuffer.removeopens its own section.mutation()registers the holder at L317 and only discards it at L325, so the outer holder is still registered when the inner section opens on the same asyncio task._current_task()raises:
RuntimeError: this task already holds a data-plane barrier section; pass the DataPlaneMutationCut you already have instead of opening another
Only ReadyFirstSampler.evict returns 0 without touching the buffer. Affected shipped configs: grpo_math_1B_megatron_single_controller.yaml:26, ppo_math_1B_megatron_single_controller.yaml:34, grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.yaml:27, mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.yaml:12, ppo-qwen2.5-1.5b-gsm8k-2n8g-...-async-single-controller.yaml:62, grpo_qwen3_0_6b_megatron_generation_single_controller.yaml:66 (in_order), plus grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.yaml:24 and grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.yaml:27 (windowed).
Not rare: eviction is steady-state for async off-policy training — this same block increments evicted_stale_prompt_groups and prints evicted N stale prompt group(s). It raises the first time anything is actually stale. CI misses it because FakeBuffer.remove in tests/unit/single_controller/test_sampler_interface.py is a plain async def that never opens a section, so the stub is more permissive than the contract it stands in for.
AI-1
Two shapes, and I don't want to over-prescribe:
- Drop the wrapper —
remove()already opens and tags its own"group_removals"section. - Thread the existing cut down (
evict(cut, ...)→remove(cut, ...)), which is what the error message prescribes and what preserves atomicity across the select-then-remove window, if that was the intent.
The three other group_removals sites in this file (L1509, L1802, L3557) all bind as cut and pass it into the callee, so option 2 matches the established idiom. Could you confirm which atomicity the wrapper was added for? That decides the fix.
| reused_siblings += reused | ||
| redispatched_siblings += group.expected_generations - reused |
There was a problem hiding this comment.
1 action item. CI blocker.
uv run --frozen pyrefly check reports exactly three errors repo-wide, and two are these lines:
ERROR `reused_siblings` is not mutable from the current scope [unknown-name]
--> nemo_rl/algorithms/single_controller.py:1149:13
ERROR `redispatched_siblings` is not mutable from the current scope [unknown-name]
--> nemo_rl/algorithms/single_controller.py:1150:13
pyrefly==0.24.2 is pinned (pyproject.toml:277) and runs as the pyrefly-typecheck pre-commit hook, and this file is in pyrefly.toml project-includes.
This is an upstream false positive — the runtime semantics are correct, and 0.24.2 flags any nonlocal augmented assignment (I reproduced it on a two-line snippet). But it still fails CI, and no other pyrefly-checked file under nemo_rl/ uses nonlocal, which is why nothing else trips it.
AI-2
Either accumulate into a mutable container (a small dataclass or a Counter) so no nonlocal is needed, or add a scoped # pyrefly: ignore[unknown-name] with a comment pointing at the upstream bug.
| stacklevel=2, | ||
| ) | ||
| else: | ||
| vllm_cfg = generation_config["vllm_cfg"] |
There was a problem hiding this comment.
1 action item. CI blocker (third of the three pyrefly errors).
ERROR TypedDict `GenerationConfig` does not have key `vllm_cfg` [typed-dict-key-error]
--> nemo_rl/algorithms/single_controller_utils/setup.py:1008:42
GenerationConfig(TypedDict) genuinely has no vllm_cfg key. The pre-existing sibling 85 lines down at setup.py:1093 already handles this.
AI-3
Match the existing sibling:
| vllm_cfg = generation_config["vllm_cfg"] | |
| vllm_cfg = cast(dict[str, Any], generation_config)["vllm_cfg"] |
| async with barrier.mutation() as cut: | ||
| cut.require_live() | ||
|
|
||
| def test_reports_mutations_blocked_by_checkpoint(self): |
There was a problem hiding this comment.
1 action item.
TL;DR — inserting this test here stole the previous test's asyncio.run(exercise()), so test_same_task_cannot_nest_barrier_sections now builds a coroutine and returns without awaiting it — and that is the exact guard the evict() bug elsewhere in this PR trips.
PR-introduced. test_same_task_cannot_nest_barrier_sections ends at L448 with its "a rejected nested section must not poison later acquisitions" block, and the very next line is this def. The asyncio.run(exercise()) that used to terminate that test now terminates this one.
All four parametrisations (mutation/checkpoint × mutation/checkpoint) now pass vacuously. The only trace is a RuntimeWarning: coroutine 'exercise' was never awaited, which is easy to miss among the pydantic deprecation warnings. I confirmed it is dead by injecting raise AssertionError into its body — it still passed.
AI-4
Restore asyncio.run(exercise()) to the end of test_same_task_cannot_nest_barrier_sections (after L448) and give this new test its own runner.
| self._rollout_completion_durations_s: deque[float] = deque(maxlen=10_000) | ||
| self._rollout_queue_wait_durations_s: deque[float] = deque(maxlen=10_000) |
There was a problem hiding this comment.
1 action item.
TL;DR — 16 tests in test_rollout_pump.py fail at HEAD with AttributeError: 'SingleControllerActor' object has no attribute '_rollout_queue_wait_durations_s'; these two deques are set in __init__, but the tests hand-build the controller and never get them.
PR-introduced. These deques are read from _record_rollout_timing on every dispatched prompt, but test_rollout_pump.py builds its controllers via object.__new__(cls) and populates state manually in _init_pump_ledgers, so the attribute never exists and the pump's TaskGroup tears down.
I reproduced the full sweep at a pristine HEAD:
$ cd tests && uv run --no-sync pytest unit/single_controller/test_rollout_pump.py \
unit/single_controller/test_checkpoint_dispatch_races.py \
unit/experience/test_rollout_reassembler_actor.py -p no:randomly --maxfail=100
20 failed, 54 passed
All three files run in CI via tests/unit/L0_Unit_Tests_Other.sh.
AI-5
Nothing wrong with the code here — the fix is in the test fixture (see my comment on _init_pump_ledgers). Flagging it at the definition so the coupling is visible: any new __init__ state read from the pump hot path has to be mirrored there.
| }, | ||
| step=self._train_steps, |
There was a problem hiding this comment.
1 action item.
TL;DR — _log_telemetry_metrics is wrapped in a broad try/except so telemetry can never kill a run, but step=self._train_steps is evaluated as an argument, i.e. outside that guard — so a telemetry-only failure escalates to a run-killing AttributeError.
PR-introduced. This is the product-side cause of the 3 failures in test_checkpoint_dispatch_races.py:
AttributeError: 'SingleControllerActor' object has no attribute '_train_steps', raised from this line rather than warned about.
The tests are hand-building controllers, so in CI it surfaces as a test gap — but the same shape means any future path that reaches _redispatch_restored_rollouts before _train_steps is set takes down the run instead of degrading to a warning, which defeats the point of the guard in _log_telemetry_metrics.
AI-7
Read the step inside _log_telemetry_metrics, or use getattr(self, "_train_steps", 0) here, so the existing fail-soft intent actually holds on the restore path. This blocks nothing legitimate — on every real path _train_steps is set, so the guard only changes behaviour in the case that is currently a crash.
There was a problem hiding this comment.
The rebase rewrote e847da0 into 0f8451c04a, and this half did not come across: at HEAD _log_telemetry_metrics still takes a required step: int and guards only its body, so all four call sites still evaluate self._train_steps as an argument — :821, :1267, :3989 and :4217.
What did land is the test-side fix, _init_recovery_telemetry, which is the right fix and turns the three failures green.
There is no way to trigger the original problem anyway — _train_steps is set in __init__ before any of those paths run — so leaving it is fine. Flagging only so this thread does not close on a premise that no longer holds. If you would rather close it properly, the cheap version is at the definition: step: Optional[int] = None, then read step if step is not None else self._train_steps inside the existing try.
|
|
||
| def record_recovery_siblings(self, *, reused: int, redispatched: int) -> None: | ||
| del reused, redispatched | ||
|
|
There was a problem hiding this comment.
1 action item. 3 of the 20 L0 failures.
The three recovery tests build controllers via object.__new__(controller_cls) and never set the telemetry state that _redispatch_restored_rollouts now reads, so all three fail with AttributeError: 'SingleControllerActor' object has no attribute '_train_steps'.
AI-8
Add a helper and call it at the three object.__new__(controller_cls) sites:
def _init_recovery_telemetry(controller: Any) -> None:
"""Give a hand-built controller the telemetry state restore logs through.
``_redispatch_restored_rollouts`` reads ``self._train_steps`` to stamp the
redispatch sample. That read is an *argument* to ``_log_telemetry_metrics``,
so it happens outside that method's try/except and raises rather than warns.
"""
controller._train_steps = 0
controller._telemetry_sample_index = 0
controller._telemetry_started_at = 0.0
controller._logger = SimpleNamespace(log_metrics=lambda *a, **k: None)Verified: with this plus the test_rollout_pump.py fixture fix, both files go green.
| group_min_wv: int | ||
| group_max_wv: int | ||
| staging_keys: list[str] | ||
| canonical_output_tokens: int = 0 |
There was a problem hiding this comment.
1 action item. 1 of the 20 L0 failures.
test_rollout_reassembler_actor.py::test_rpc_dataclass_fields_are_classified fails at HEAD because this new field isn't in the hand-maintained field inventory. That test is doing exactly its job — its docstring says a new field "must be a deliberate choice" — so this just needs classifying.
AI-9
| canonical_output_tokens: int = 0 | |
| "staging_keys", | |
| # A count, not a payload: the canonical token ids stay in staging. | |
| "canonical_output_tokens", | |
| "metrics", |
(in the inventory in tests/unit/experience/test_rollout_reassembler_actor.py, not here — anchoring on the field that triggers it.)
| if self._pending_step is not None and self._pending_step != step: | ||
| self._flush_pending_metrics_locked() | ||
| if self._pending_step is None: | ||
| self._pending_step = step |
There was a problem hiding this comment.
1 action item.
TL;DR — training metrics move off W&B's built-in _step onto a custom nemo_rl/step axis, so the UI's default "Step" x-axis becomes a row counter for every NeMo-RL W&B user, not just single-controller.
PR-introduced. Before this PR WandbLogger passed the trainer step straight through (git show 9534f00f8:nemo_rl/utils/logger.py — self.run.log(metrics, step=step, commit=True)). HEAD buffers rows here and flushes via _flush_pending_metrics_locked with no step=, stamping nemo_rl/step into the payload instead. W&B then auto-increments _step once per log() call.
Consequence: any saved view, report, or panel pinned to the default Step axis — and any run-to-run comparison spanning this commit — silently misaligns. Nothing errors; the curves just stop meaning what they used to.
AI-10
Not asking you to revert it — the custom axis is the right call for interleaved telemetry, and test_independent_events_do_not_reuse_wandb_internal_step pins the invariant nicely. Please just add a line to the PR description and to the W&B section of the logging docs saying _step is no longer the trainer step and panels should use nemo_rl/step.
Follow-up (fine for a later PR): the implicit-boundary flush this function implements — a metric for step N+1 arriving while step N is still buffered — has no test. Only step_finished=True and finish() are covered, yet the implicit path is what non-trainer callers (validation loops, log_plot, log_histogram) rely on since they never pass step_finished. Mutation-checked: collapsing this to self._pending_step = step is not caught by the current suite. Separately, log_plot/log_histogram were rewritten to buffer and have zero W&B-backend coverage — every existing plot/histogram test targets TensorBoard or MLflow.
github-actions
Bot
added
the
Documentation
e847da0 to
2df5b12
Compare
2df5b12 to
70940a9
Compare
|
/ok to test 70940a9 |
macandro96
added
the
CI:Lfast
70940a9 to
deabd07
Compare
deabd07 to
a4c10a7
Compare
a4c10a7 to
23d1295
Compare
|
/ok to test 23d1295 |
terrykong
left a comment
terrykong
left a comment
There was a problem hiding this comment.
This PR adds telemetry for rollout checkpointing in the single controller: a telemetry_interval_s sampling pump, throughput/save/restore/outcome metric families, checkpoint-barrier mutation counters in the replay buffer, and a new metrics doc section. It also moves WandbLogger.log_plot and log_histogram onto the buffered-step path so plots and scalars land in the same W&B history row.
Worth calling out: your earlier self-review round holds up. Checked against the code at HEAD rather than the "addressed in ..." replies — the rebase rewrote the commits those replies cite — 10 of the 11 findings are genuinely fixed. The 11th has its own thread reply: the test-side half landed and the production-side half did not, and it is harmless, so it just needs the thread not to close on a stale premise.
Local evidence at a4c10a78:
tests/unit/utils/test_logger.py+tests/unit/single_controller/: 1014 passed.tests/unit/experience/+tests/unit/data_plane/: 648 passed, 15 skipped. All 15 skips are missing dependencies (4xnemo_gym, 11xmooncake), none GPU.- ruff 0.9.9 clean on all three pinned hooks. Pyrefly reported no error in any file this PR touches, but the dependency set was incomplete, so treat that as a partial check rather than a clean bill.
Two housekeeping items. The PR is currently marked as conflicting with main — please rebase and resolve before the next round. And the review is pinned to a4c10a78; the head has since moved to 23d12951, which differs only in .github/workflows/cicd-main.yml and two functional-test scripts (13 lines), so every comment below still points at the code it describes. /ok to test 23d1295 is now running, which is what closes out the two pyrefly findings from the last round and the canonical_output_tokens arms that no local lane executed.
One coverage note: tests/unit/data_plane/test_rollout_reassembler.py is module-level skipped locally for missing nemo_gym, so the new canonical_output_tokens assertion only runs in CI's Nemo_Gym lane.
Nineteen inline comments: two getattr guards that cannot do what they look like they do, an ungated per-tick print, one docs/code mismatch in the recovery metrics table (the most substantive one — a round-trip test on the real restore path is attached, and it asks that the re-admission discard raise instead of dropping prompts, with a startup check that the restored sampler cursor is not behind the restored buffer), a naming ask for the canonical_* throughput keys, a note that the long-stalled recovery test closes #4047 and should say so, an undocumented on-switch, a hand-copied kind list, three test-coverage gaps, and six placement/duplication notes on the telemetry constants (including the mutation-kind list, which should be an enum so a wrong kind fails instead of vanishing from the series, and the hard-coded three-failure cap on the snapshot pump) — one of which is a wall-clock axis that restarts at 0 on every resume, so before- and after-restart series overlap.
Generated by Claude Code
| rollout_checkpoint_cfg = getattr( | ||
| self._master_config, "rollout_checkpointing", None | ||
| ) | ||
| telemetry_interval_s = getattr( | ||
| rollout_checkpoint_cfg, "telemetry_interval_s", None | ||
| ) |
There was a problem hiding this comment.
1 action item.
TL;DR — this getattr guard cannot prevent the crash it looks like it is guarding against; five lines down the same config is read plainly and raises AttributeError anyway.
New in this PR. In the exact case the guard defends against, telemetry_interval_s comes back None, telemetry quietly turns itself off, and then self._master_config.rollout_checkpointing.snapshot_attempt_interval_s blows up regardless. Neither fallback is reachable: rollout_checkpointing is a required field with a default factory, and telemetry_interval_s is a declared field. So the guard buys nothing and only hides which knob went missing.
AI-1
| rollout_checkpoint_cfg = getattr( | |
| self._master_config, "rollout_checkpointing", None | |
| ) | |
| telemetry_interval_s = getattr( | |
| rollout_checkpoint_cfg, "telemetry_interval_s", None | |
| ) | |
| telemetry_interval_s = ( | |
| self._master_config.rollout_checkpointing.telemetry_interval_s | |
| ) |
| self._rollout_checkpoint_load_metrics = getattr( | ||
| actor_args, "rollout_checkpoint_load_metrics", None | ||
| ) |
There was a problem hiding this comment.
1 action item.
TL;DR — same defensive getattr on a declared dataclass field that is always set, so it can only turn a config bug into a silent None.
New in this PR. rollout_checkpoint_load_metrics is a declared field with a = None default and is always populated at construction. The line right above reads its neighbour plainly as actor_args.data_plane_checkpoint_metadata, and driver and actor run the same image, so there is no version-skew case to cover. This is the same thing you raised on yourself as AI-11 and fixed in telemetry_snapshot — this site and the one at line 655 were just missed.
AI-1
| self._rollout_checkpoint_load_metrics = getattr( | |
| actor_args, "rollout_checkpoint_load_metrics", None | |
| ) | |
| self._rollout_checkpoint_load_metrics = ( | |
| actor_args.rollout_checkpoint_load_metrics | |
| ) |
| step=self._train_steps, | ||
| prefix="rollout/throughput", | ||
| ) | ||
| print(f"rollout_throughput_metrics={metrics}", flush=True) |
There was a problem hiding this comment.
1 action item.
TL;DR — this print sends the same dict to stdout that the line above already sends to W&B/TensorBoard, on every telemetry tick, with no way to turn it off.
New in this PR. _log_telemetry_metrics(metrics, ...) five lines above already logs metrics, so this is a second copy of the same payload going to the job log. It is also the only unconditional per-tick print in this file — the rest are event-driven. self._async_cfg.diagnostics already exists for exactly this and gates a print at line 2047.
AI-1
| print(f"rollout_throughput_metrics={metrics}", flush=True) | |
| if self._async_cfg.diagnostics: | |
| print(f"rollout_throughput_metrics={metrics}", flush=True) |
| | `timing/rollout_checkpoint` | `snapshot_rows`, `replay_rows`, `staging_rows`, `replay_groups`, `ledger_groups`, `controller_sidecar_bytes` | Logical volume captured by the snapshot. `controller_sidecar_bytes` excludes the native TQ payload because the current TQ checkpoint API does not report bytes written. NeMo-RL deliberately does not recursively scan the shared checkpoint directory because that scan would perturb the benchmark. | | ||
| | `rollout/checkpoint_outcome` | `completed`, `skipped`, `failed`, `reason_*`, `seconds_since_previous_success`, `seconds_since_last_success` | Result, actionable reason, and effective cadence of every scheduled checkpoint attempt. | | ||
| | `timing/rollout_recovery` | `snapshot_resolution_seconds`, `dataloader_load_seconds`, `tq_load_seconds`, `replay_metadata_load_seconds`, `recovery_prepare_seconds`, `total_load_seconds` | Rollout-state restore latency. `total_load_seconds` is the sum of these non-overlapping restore phases. | | ||
| | `timing/rollout_recovery` | `groups_reused`, `groups_redispatched`, `siblings_reused`, `siblings_redispatched`, `redispatch_schedule_seconds` | Completed groups restored without generation and unfinished sibling work preserved or repeated after restart. | |
There was a problem hiding this comment.
1 action item.
TL;DR — groups_considered and groups_redispatched are equal on every checkpoint this code writes, so one is redundant; the discard that could separate them should raise instead of dropping prompts; and the names make two unrelated things both read as “reused”.
New in this PR. This row documents five recovery keys as if they arrive together.
The two counts cannot differ on a checkpoint this code writes. They would only diverge through the discard at single_controller.py:1149-1152, which releases a RESERVED group when the replay buffer already holds a group stamped for the step the batch is being re-admitted to. That needs a stamp past the sampler cursor, and nothing produces one: a re-admitted batch gets dispatch_index + 1, the borrow only moves groups between steps that were already admitted (promote_ready_group picks among slots that exist), and the repayment spare takes the lender's stamp (:1987). The cursor is saved with every checkpoint (:3710, :4296) and restored at :479.
A round-trip test on the real path — the real rollout pump doing a borrow and repayment, the real _capture_rollout_checkpoint_cut, then _maybe_restore_replay_buffer, _maybe_restore_rollout_recovery and _redispatch_restored_rollouts in a fresh controller; only generation and the tensor converter are faked — gives:
restore with the saved cursor -> 14 of 14 prompts back, discarded 0, considered 4, redispatched 4
restore with the cursor rebuilt from
trainer_version (the fallback at
:477 for a checkpoint that has no
sampler_dispatch_index) -> prompts 12, 13, 14 discarded, considered 4, redispatched 1
The second row is the only way the branch fires: a cursor that lands below stamps already in the buffer. Test file: https://github.com/terrykong/gh-pages-poc/blob/47cd8814384382893c2a56a86f939caf76b48936/terryk/pr-3925-test_restore_after_borrow.py — walk-through with the figures: https://terrykong.github.io/gh-pages-poc/terryk/pr-3925-can-restore-drop-data.html
Separately, the names collide. groups_reused is a whole group that needed no generation; siblings_reused is partial credit inside a group that is being redone. And groups_reused / groups_redispatched read like complements but count disjoint populations — the replay buffer versus the recovery ledger, which the code calls “unfinished … next to canonical”.
AI-1
First, make the discard at :1149-1152 fail instead of dropping. It cannot run on a correct cursor, and on a wrong one it loses prompts with only a print. A re-admitted batch finding groups already stamped for its step means the sampler cursor is behind the buffer, and that should stop the run:
if buffered:
raise RuntimeError(
f"target_step={target_step} already holds {buffered} group(s) at re-admission; "
"the restored sampler cursor is behind the replay buffer"
)The same branch in the live pump at :2132-2140 has the same property — pre-existing, so a follow-up is fine. The root cause — a restored cursor behind the restored buffer — can be rejected at startup instead; that ask is the comment at :637. With both, groups_redispatched has nothing left to count. Drop it, and rename the rest so the population and the outcome are both visible:
| today | proposed |
|---|---|
groups_reused |
groups_complete_restored |
groups_considered |
groups_unfinished_found |
groups_redispatched |
remove — always equals groups_unfinished_found once the discard raises |
siblings_reused |
unchanged — unambiguous once the group-level key is renamed |
siblings_redispatched |
siblings_rerun |
And take the round-trip test into the PR: nothing in CI runs a borrow through a checkpoint, and nothing re-admits a RESERVED batch against a real buffer — every scenario in the matrix builds groups with admitted=True, and the one test that builds RESERVED groups replaces the admission with a stub. Touches the emission dicts at :818 and :1259-1269 plus this table row — not a suggestion block, it spans two files and several rows here.
Context — no action. Two smaller mismatches in the same row. groups_reused comes from _log_rollout_restore_metrics while the other four come from _redispatch_restored_rollouts, and the former returns early when _rollout_checkpoint_load_metrics is None — so a run can show the redispatch keys with no groups_reused at all, which reads as “nothing was reused” rather than “never recorded”. And groups_considered is emitted today but absent from this table.
What the round-trip test runs, and how to reproduce it
Real: _rollout_pump (admission, spare pool, the borrow at :1983, repayment), TQReplayBuffer, RolloutRecoveryLedger, InOrderSampler, DataPlaneCheckpointBarrier, RolloutManager.generate_and_push, _capture_rollout_checkpoint_cut, _maybe_restore_replay_buffer, _maybe_restore_rollout_recovery, _redispatch_restored_rollouts, NoOpDataPlaneClient save/load. Faked: run_rollout (finishes at once, except one prompt that fails on purpose and one that stays in flight) and record_to_train_batch. There is no train pump, so after the restore the test bumps trainer_version once — what training step 0 would do — so the in-order gate opens.
Setup: group size 2, 3 groups per step, in_order with lookahead 2, on_dropped_prompt: replace, pool of 3. Prompt 1 (step 0) fails; step 0 borrows prompt 9's finished group from step 2; spare 3 is dispatched for step 2 and is still in flight when the checkpoint is taken. Checkpoint holds 8 finished groups, 1 ADMITTED, 3 RESERVED, cursor 2. Run from the repo root with PYTHONPATH=. pytest <file> -q -s.
| """ | ||
|
|
||
| snapshot_attempt_interval_s: Annotated[Optional[float], Field(gt=0)] = None | ||
| telemetry_interval_s: Annotated[Optional[float], Field(gt=0)] = None |
There was a problem hiding this comment.
1 action item.
TL;DR — the knob that turns all of this new telemetry on is not documented anywhere, and it ships off by default, so nothing in the new metrics table is reachable out of the box.
New in this PR. grep -rn "telemetry_interval_s" docs/ returns nothing: the PR adds this field, a 45-line table in the metrics doc, and the sampling pump, but never says how to switch it on. All three shipped YAMLs set it to null, so a reader who finds the table has no way to get any of it.
AI-1
Add a short paragraph to docs/guides/single-controller.md: name rollout_checkpointing.telemetry_interval_s, say it is null (off) by default, give a sample value in seconds, and link to the metrics table. Not a suggestion block — that file is not in this diff.
| # Logger this module also uses as `self._logger`. | ||
| log = logging.getLogger(__name__) | ||
|
|
||
| _MAX_CONSECUTIVE_ROLLOUT_CHECKPOINT_FAILURES = 3 |
There was a problem hiding this comment.
1 action item.
TL;DR — the number of failed snapshots a run tolerates before it aborts is an operator decision, and it is a magic 3 at the top of the controller; it belongs on RolloutCheckpointConfig next to the interval it governs.
Pre-existing line, but this PR is what wires it into telemetry: every failure it counts is now a rollout/checkpoint_outcome row (:4014-4018), and when the count hits this cap the run dies (:4033-4037) with RuntimeError("periodic rollout checkpoint failed 3 consecutive times"). Only OSError and TimeoutError are counted — anything else re-raises on the first hit (:4019-4020) — so this is "how many transient storage hiccups in a row before giving up", which is exactly the kind of value that differs between a fast local disk and a shared filesystem under load. Nobody can change it without editing the module, and the config docstring for the pump does not mention that the limit exists.
AI-1
Move it to RolloutCheckpointConfig as max_consecutive_failures: Annotated[int, Field(ge=1)] = 3, read it in _rollout_checkpoint_pump, and say in the class docstring what happens at the limit. Not a suggestion block: two files. If you prefer to keep it a constant, it still should not live here — rollout_checkpoint.py owns the other snapshot constants (see the comment at :185).
| def telemetry_snapshot(self) -> dict[str, int]: | ||
| """Return cumulative canonical-publication and recovery counters.""" | ||
| return { | ||
| "canonical_groups_finalized": self._canonical_groups_finalized, |
There was a problem hiding this comment.
1 action item.
TL;DR — canonical is the code's internal word for "the training-ready TQ rows", and these metric keys are the only place a user meets it; committed is already the public word for the same event, so name the keys after it.
New in this PR. A reader of the W&B chart has no way to learn what canonical_output_tokens_per_second means without opening replay_buffer.py. The event being counted is the one RolloutOutcome.COMMITTED already names — "the group reached the buffer" — and it is recorded at exactly the two places a group is committed: after TQReplayBuffer.commit() in generate_and_push (an estimate: mean_gen_tokens_per_sample × len(completions)) and after the finalizer's canonical commit (an exact count from the staged rows). Two different quantities under one name is worth one sentence in the docs row.
AI-1
Rename the four W&B keys and the docs row; the internal attributes can keep their names:
| today | proposed |
|---|---|
canonical_groups_finalized |
committed_groups |
canonical_output_tokens |
committed_output_tokens |
canonical_groups_per_second |
committed_groups_per_second |
canonical_output_tokens_per_second |
committed_output_tokens_per_second |
Then generation_output_tokens_per_second vs committed_output_tokens_per_second reads as "generated vs kept" with no glossary. Touches this dict, the rate keys in _collect_and_log_rollout_throughput_metrics, and the docs row, where "finalized token … throughput available for training" is already the plain-English version — add that the ordinary path estimates the token count while the token-capture path counts it. Not a suggestion block: three files.
|
|
||
| replay_restore_started = time.monotonic() | ||
| restored_replay_groups = await self._maybe_restore_replay_buffer() | ||
| replay_restore_seconds = time.monotonic() - replay_restore_started |
There was a problem hiding this comment.
1 action item.
TL;DR — after the replay buffer and ledger are restored, nothing checks that the sampler cursor is at or past every stamp they contain; a cursor behind the buffer is the one state that makes re-admission drop prompts, and it can be rejected here, at startup.
New lines in this PR, on the path where the check belongs. The cursor is restored before this point from one of two places: the trainer checkpoint's GRPOSaveState.sampler_dispatch_index (written at :4296) or, when a rollout snapshot is selected, the snapshot's manifest.json via setup.py. When neither has it, __init__ rebuilds it from trainer_version, and that is the case the round-trip test shows dropping every reserved prompt of the re-admitted batch (see the comment on the metrics table). The buffer and ledger are restored a few lines below this one, so after _maybe_restore_rollout_recovery both sides are in memory and the invariant is one comparison: every target_step in the buffer and on every ADMITTED ledger group is <= self._sampler.dispatch_index, because a stamp is only ever handed out by commit_admission advancing that cursor.
AI-1
Right after _maybe_restore_rollout_recovery returns, fail if the restored cursor is behind the restored work:
restored_stamps = [
step for step in self._buffer.target_step_list if step is not None
] + [
group.target_step
for group in self._rollout_manager.recovery_ledger.groups()
if group.target_step is not None
]
if restored_stamps and max(restored_stamps) > self._sampler.dispatch_index:
raise RuntimeError(
f"restored sampler cursor {self._sampler.dispatch_index} is behind the "
f"restored replay state (max target_step {max(restored_stamps)}); "
"the checkpoint's sampler_dispatch_index is missing or stale"
)Not a suggestion block: it goes below these lines, after the two restore calls. With this in place the discard at :1149-1152 cannot be reached from a restore, which is what makes it safe to turn that branch into a hard error too (the other comment).
|
|
||
|
|
||
| @pytest.mark.parametrize("case", LONG_STALLED_CASES, ids=lambda case: case.id) | ||
| def test_long_stalled_partial_group_is_selected_or_deliberately_evicted(case, tmp_path): |
There was a problem hiding this comment.
1 action item.
TL;DR — this test shows a recovered group can be evicted as stale right after restore, and the recovery counters do not show that; one sentence in the metrics doc should.
This commit adds S_LONG_STALLED_PARTIAL — a half-generated group seven steps behind the newest — and pins what happens to it after a restore, per sampler: ready_first keeps it and selects it (:139-142); windowed, weight_fifo and in_order evict it as stale right after restore and count the eviction (:143-146). Good — that is the intended behaviour, recorded.
It also means the doc row for the recovery counters (metrics.md:42) reads as if a recovered group is a trained group, and it is not: on three of the four samplers the straggler is relaunched, its missing sibling regenerated, and the finished group is then evicted by evict before anything trains on it. None of that shows up in timing/rollout_recovery — groups_reused and siblings_reused count what was restored, and the eviction lands in the sampler's own accounting.
AI-1
Add one sentence to the recovery-counters row: "Recovered groups can still be evicted by the sampler's staleness rule right after restore (all gated samplers do this to a long-stalled group); the recovery counters are an upper bound on what trains." Not a suggestion block: a docs file.
There was a problem hiding this comment.
1 action item.
TL;DR — this file closes #4047; say so in the PR description so the issue closes on merge.
test_long_stalled_partial_group_is_selected_or_deliberately_evicted is the test #4047 asked for: a partly generated group several steps behind the newest (S_LONG_STALLED_PARTIAL), run over the sampler list, asserting per sampler whether the recovered straggler trains or is deliberately evicted and counted. The PR description's Issues section is empty, so the issue stays open after merge.
AI-1
Add Closes #4047 under Issues in the PR description. Not a suggestion block: it is the PR body, not a file.
There was a problem hiding this comment.
addressed in the PR description
Signed-off-by: Anish Mahishi <amahishi@nvidia.com> (cherry picked from commit 52aa315)
Signed-off-by: Anish Mahishi <amahishi@nvidia.com> (cherry picked from commit 815f3d9)
Signed-off-by: Anish Mahishi <amahishi@nvidia.com> (cherry picked from commit 9853cf4)
Signed-off-by: Anish Mahishi <amahishi@nvidia.com> (cherry picked from commit 67f8278)
Signed-off-by: Anish Mahishi <amahishi@nvidia.com> (cherry picked from commit eafe0a5)
(cherry picked from commit af43edd) Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
1784afe to
195df04
Compare
|
/ok to test a3dfa6c |
2 participants
What does this PR do?
Summary
Adds observability for Single Controller rollout checkpointing and recovery.
This PR measures whether periodic checkpointing affects rollout throughput, where checkpoint and restore time is spent, how much work is blocked by checkpoint barriers, and how effectively completed rollout work is reused after restart.
Stacked on #3924.
What changed
Rollout throughput
Adds metrics for:
Raw generation throughput uses cumulative counters exposed through the generation interface. vLLM implements these counters today; other inference backends can add support through the same interface. With token capture, committed tokens are counted exactly from valid staged rows; without token capture, they are estimated from the per-sample mean and completion count.
Checkpoint latency and backpressure
Adds phase-level checkpoint timing:
Also records:
controller_sidecar_bytesintentionally excludes the native TQ payload because the current TQ checkpoint API does not report its size. The implementation avoids recursively scanning the shared checkpoint directory because doing so could perturb checkpoint benchmarks.Checkpoint outcomes
Every scheduled checkpoint attempt records:
Checkpoint invariant violations continue to fail immediately. Retryable I/O and timeout failures retain bounded retry behavior controlled by
rollout_checkpointing.max_consecutive_failures.Restore and recovery
Adds restore timing for:
Recovery-efficiency metrics report:
The restore path now rejects a sampler cursor that is older than restored rollout work instead of silently dropping colliding prompts. Coverage includes a real rollout-pump borrow and repayment round trip through checkpoint capture and restore.
Logging behavior
Independent checkpoint, recovery, and throughput events use absolute wall-clock time for visualization so telemetry continues advancing during a long trainer step and across process restart.
nemo_rl/stepremains the trainer-correlated step axis. W&B's internal_stepis a monotonically increasing event-row index.Wall-clock sampling is disabled by default. Set
rollout_checkpointing.telemetry_interval_sto a positive interval in seconds to enable it; this does not change checkpoint cadence.Primary dashboard
The intended high-level dashboard contains:
The remaining controller, buffer, vLLM request, KV-cache, and per-mutation metrics are diagnostic drill-down signals.
Additional fixes
This PR also incorporates review follow-ups that:
Validation
Completed locally:
ruff checkruff format --checkgit diff --checkRecommended unit tests:
Issues
Closes #4047
Before your PR is "Ready for review"
Pre checks: