Repository navigation
feat(checkpoint): spans and metrics for partial-rollout checkpoints - #3903
Conversation
ananthsub
added
feature
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
fd8308c to
9775d38
Compare
dc49274 to
f9affd7
Compare
f9affd7 to
f8f2b11
Compare
14f23fa to
98f6598
Compare
98f6598 to
bba4db3
Compare
6c2e46e to
3589813
Compare
3589813 to
8176162
Compare
8176162 to
5163155
Compare
5163155 to
9ae6dbc
Compare
macandro96
left a comment
macandro96
left a comment
There was a problem hiding this comment.
Thanks for this! Inline comments cover the main issues (mostly metrics recording the wrong thing). A few more notes:
Docs: fern/versions/latest/pages/observability/span-groups.mdx doesn't list the new checkpoint group or how to enable it (default,checkpoint), and metrics.mdx omits the gym.checkpoint.* instruments.
Optional follow-ups (non-blocking)
Telemetry:
- A result on the collection span (
published/not_ready/error), so there's one place to see whether a checkpoint was saved. - Finer operation outcomes: use the
ControlError.code(timeout, conflict, ...) instead of onlyok/error. - Separate
in_flight,restoredandtotalrow attributes on the collection span;in_flightcurrently excludes restored rows that the commit includes. - Count attempts separately from actual writes/installs, which would catch repeated side effects like the volume double count.
- A counter for idempotent replays returned by
_recorded().
Hang detection:
- How long the control lock has been held, and by which operation. Most waits under the lock are deadline-bounded, but two in
_reopenaren't:asyncio.to_thread(self._write.stop.stop)andasyncio.wait([cleanup], timeout=None)whendeadline_tsisNone. Both should be short or unreachable today, but this would surface a slow filesystem or a future regression. - Age of the oldest background task (export / write / cleanup / install), since these are shielded and outlive request deadlines. Log once when it crosses a threshold.
Race tests (races here show up as wrong state, not errors, so metrics won't catch them):
- Resume or lease expiry while prepare is waiting for readiness.
- Lease expiry during a commit's cleanup.
- Commit retry after a timeout while the original write is still running.
- Restore failing while install is still running, racing with a resume.
| await _resume_in_order(participants, body, deadline_ts, failures) | ||
| if failures: | ||
| raise CoordinationError("resume", failures) |
There was a problem hiding this comment.
Failed resume is recorded as a success. CoordinationError is raised after _coordinating("resume") has exited, so operation() always records outcome=ok for coordinate.resume, and the span ends without error. Move the if failures: raise inside the with block.
There was a problem hiding this comment.
Fixed. resume now raises the CoordinationError inside _coordinating("resume"), so a resume that left a participant closed records outcome=error for coordinate.resume and its span ends with the error.
The other stages are still all reopened before it raises.
Covered by test_a_resume_that_fails_on_a_participant_records_an_error_and_ends_its_span_with_it: the agent stage fails, and the test checks the error outcome and the span's error status. On the previous head the span status was unset and the outcome was ok.
| record_checkpoint_event( | ||
| "lease_expired", participant_kind=self.participant.kind, phase=self.phase.value |
There was a problem hiding this comment.
lease_expired is recorded repeatedly, and before recovery succeeds. It's recorded before _reopen, and the loop records it again on every retry (e.g. while _reopen raises DeadlineExceededError because cleanup is still deleting restored state). One expiry can count many times. The metric description in gym_metrics.py also says "a participant resumed on its own", which may never happen if reopen keeps failing. Suggest recording the expiry once per lease and tracking reopen attempts/outcome separately.
There was a problem hiding this comment.
Fixed. The lease watcher now counts lease_expired once per checkpoint, the first time it finds the lease expired, and no longer on each retry of the reopen loop.
Each attempt to reopen is now its own operation, lease_reopen, so it gets a gym.checkpoint.lease_reopen span and a point in gym.checkpoint.operation_duration_ms with outcome=error when it runs out of time waiting for the cleanup and ok when it reopens.
The description in gym_metrics.py now says what is counted: the lease ran out before the controller resumed the participant, so the participant starts reopening on its own, and the attempts are the lease_reopen operation.
Covered by test_a_lease_expiry_is_counted_once_and_each_reopen_attempt_with_its_outcome: the commit's cleanup is held, so reopening after the lease expires runs out of time at least twice before the cleanup is released and it succeeds. The test checks one lease_expired, at least two lease_reopen errors, and one lease_reopen success. On the previous head lease_expired was counted three times.
| record_checkpoint_volume( | ||
| operation="commit", kind=self.participant.kind, records=manifest["record_count"], size_bytes=size | ||
| ) |
There was a problem hiding this comment.
Commit volume is double-counted on retry. If the _within(cleanup) below hits the deadline after the write finished, the retry reuses the finished write via _start_write, awaits it again, and records the same records/bytes a second time. Suggest recording volume once, tied to the write's completion (e.g. a flag on _Write). Restore (~L886) looks to have the same pattern if a retry re-reads after a failed install.
There was a problem hiding this comment.
Fixed for both. A commit now counts its records and bytes from a done callback on the write's task, so the count happens once, when the write finishes, however many commit calls await it.
A write that fails or is stopped by resume counts nothing, and a write that finishes after its commit call timed out still counts once.
A restore now counts once it has succeeded, after the install, so a restore whose install failed counts nothing and its retry counts the records once.
Covered by:
test_a_commit_retried_after_its_cleanup_ran_out_of_time_counts_its_records_once: the write finishes, the cleanup times out, and the retry reuses the write. The test checks one record and the records file's size in bytes. On the previous head the record was counted twice;test_a_restore_retried_after_its_install_failed_counts_its_records_once: the first install fails, the retry succeeds, and the test checks one restored record. On the previous head it was counted twice.
| def _records_size(checkpoint_dir: Path, kind: str, instance: str, records_file: str) -> int: | ||
| path = participant_dir(checkpoint_dir, kind=kind, instance=instance) / records_file | ||
| try: | ||
| return path.stat().st_size | ||
| except OSError: | ||
| return 0 |
There was a problem hiding this comment.
Blocking stat() on the event loop while the control lock is held (called at L773 and L882). On Lustre/NFS a stat can stall, which blocks every control call, the lease watcher, and data-plane requests on this loop. Suggest returning the byte count from write_participant_state / read_participant_state, which already run in a thread.
There was a problem hiding this comment.
Fixed as you suggested. write_participant_state now returns the manifest and the size of the records file, counted as each line is written, and read_participant_state returns the manifest, the records kept, and the size of the whole file, counted as it streams.
Both already run in a thread, so _records_size and its stat() are gone and nothing touches the filesystem on the event loop for these numbers.
A retried write that finds the manifest already published returns the size of its own records, which has the same digest and so the same size.
Every caller across the stack now unpacks the new return values; the callers that ignored the return value are unchanged.
Covered by:
test_the_store_returns_the_bytes_it_wrote_and_read, which checks that the first write, the repeated write, and a read that keeps one record all return the records file's size. On the previous head the write returned only the manifest;test_commit_and_restore_measure_their_records_without_a_stat_on_the_event_loop, which fails if a records file is stat'ed on the event loop's thread during a commit and a restore, and checks the bytes metric against the file's size. On the previous head both the commit and the restore stat'ed it on the loop.
| request = PrepareRequest(checkpoint_id=request.checkpoint_id, deadline_ts=request.deadline_ts) | ||
| with self._operation("prepare", request.checkpoint_id) as span: | ||
| result = await self._prepare(request) | ||
| span.set(phase=result["phase"], blockers=len(result["report"]["blockers"])) |
There was a problem hiding this comment.
Blocker count is truncated. report["blockers"] is capped at MAX_REPORTED_BLOCKERS (100); report["blocker_count"] is the full count. Suggest blockers=result["report"]["blocker_count"].
There was a problem hiding this comment.
Fixed: the prepare span now sets blockers from report["blocker_count"].
Covered by test_a_prepare_span_counts_every_blocker_not_only_the_listed_ones, with fifty more blockers than MAX_REPORTED_BLOCKERS. On the previous head the span said 100.
| def response(self) -> JSONResponse: | ||
| from nemo_gym.telemetry.gym_metrics import record_checkpoint_event | ||
|
|
||
| record_checkpoint_event("refused", code=self.code) |
There was a problem hiding this comment.
refused counts every control error. Every ControlError passes through response(), so this also counts deadline exceeded, export failed, invalid state, unauthorized, and bad rollout IDs, not just refusals. That contradicts the metric description ("without failing a call"). Suggest either a broader control_error event keyed by code, or recording refused only for true refusal subclasses. Recording at the exception handler instead of inside response() would also keep response() side-effect free and avoid double counting if it's called twice.
There was a problem hiding this comment.
Replaced. refused is gone, and ControlError.response() no longer records anything.
A new counter, gym.checkpoint.control_errors_total, counts control calls that failed, by nemo.gym.checkpoint.error_code, which is the code in the reply.
It is recorded once per failed call, in the control routes: each route's handler counts a ControlError that escapes it and re-raises it to the app's exception handler, which builds the reply.
It is not recorded in that exception handler itself, because the handler also replies to ControlErrors raised by data-plane routes, for example an environment server refusing an episode, while the admission middlewares of the resources, model and agent servers reply to their refusals without it.
Counting there would have counted some data-plane refusals and not others.
With several workers, the worker that received the call counts it once; the coordinator's own dispatch is not a route.
Data-plane refusals are no longer counted at all; gym.checkpoint.events_total keeps only events that do not fail a call.
Covered by:
test_each_failed_control_call_counts_its_error_code_once: a commit in the wrong phase, an unauthorized prepare, and a retire past its deadline each count their code once, and a successful prepare counts nothing. On the previous head the counter did not exist;test_building_an_error_reply_records_nothing, which callsresponse()twice and checks that no metric was recorded. On the previous head it recordedrefusedtwice;test_a_failed_control_call_is_counted_and_a_refused_data_request_is_not, which replaces the old refusal test.
| ) | ||
|
|
||
|
|
||
| def _reset_for_testing() -> None: |
There was a problem hiding this comment.
Looks like a rebase leftover. This second _reset_for_testing() is identical to the one at L534, and L478-484 repeat the rollout constants already defined at L151-160. The later definitions silently override the earlier ones. Suggest deleting this block.
There was a problem hiding this comment.
Removed. It was a rebase leftover: the second _reset_for_testing() and the repeated rollout constants were identical to the earlier definitions, and the module now matches main there, with one _reset_for_testing() at the end and the rollout constants at the top.
Nothing depended on the override, since both copies were the same; the telemetry and rollout metric tests pass unchanged.
| RECORDS = "nemo.gym.checkpoint.records" | ||
| BYTES = "nemo.gym.checkpoint.bytes" | ||
| PHASE = "nemo.gym.checkpoint.phase" |
There was a problem hiding this comment.
nit: RECORDS, BYTES and PHASE aren't used; OperationSpan.set builds names with an f-string. Suggest removing them (or using them).
There was a problem hiding this comment.
Used them instead: there is now one source for the attribute names.
gym_metrics.py defines CHECKPOINT_ATTRIBUTE_PREFIX (nemo.gym.checkpoint.) and the attribute names its recorders use, and nemo_gym/_checkpoint/telemetry.py builds CHECKPOINT_ID, KIND and INSTANCE from it.
OperationSpan.set prefixes each attribute name with it, so RECORDS, BYTES and PHASE were removed rather than kept as a second spelling of names set already builds.
| """Partial-rollout checkpoint spans: each participant's control operations (prepare, commit, restore, | ||
| resume, retire) with their export, write, read, and install steps, and the controller's coordination | ||
| calls. Enable it with ``default,checkpoint``; the checkpoint metrics are recorded either way.""" |
There was a problem hiding this comment.
nit: this list omits forget, which also gets a span and duration metric (same in the gym_metrics.py docstring). The public docs also need updating; see the review summary.
There was a problem hiding this comment.
Added forget to the checkpoint span group's docstring in span_groups.py and to the operation list in the gym_metrics.py docstring, which now also lists lease_reopen and the coordination calls.
The public pages are updated too: span-groups.mdx lists the checkpoint group and how to enable it, and metrics.mdx lists every gym.checkpoint.* instrument; details in the reply to the review summary.
|
Thanks @macandro96. Docs.
Race tests. The four races are now deterministic tests in #4197.
All four leave the right state. Optional follow-ups. The other telemetry and hang-detection suggestions (a result on the collection span, outcomes by error code, separate row attributes, attempts counted separately from writes, a replay counter, control-lock hold time, and the age of the oldest background task) are tracked for later and are not part of this PR. |
|
🌿 Preview your docs: https://nvidia-preview-pr-3903-a51b8ca748cb.docs.buildwithfern.com/nemo/gym |
Spans and metrics for every coordination call and for each participant's part in it, so a slow or failing checkpoint shows where its time went. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA-NeMo/Gym/actions/runs/39291983064 |
What changed and why
Checkpoints emitted nothing. That left three failure modes invisible:
This PR adds spans and metrics through the existing nemo-lens telemetry, on the same pattern as the sandbox metrics.
checkpointspan group. It is opt-in withdefault,checkpoint, because thedefaultpreset is deliberately coarse. It sits next tomain's per-serverstartupspan group (feat(telemetry): export per-server startup spans and metrics #4354), which is in thedefaultpreset and stays on by default; the two groups are independent.gym.checkpoint.<operation>forprepare,commit,restore,resume,retireandforget, andlease_reopenfor each attempt to reopen after a lease expires. Each carries the checkpoint ID, participant kind, server instance, outcome and record count.wait_ready,export,write,read,install, and the policy model'sgeneration_cutround, which also records each cut's disposition.gym.checkpoint.coordinate.<operation>, with one child per prepare stage. Rollout collection's own spans come with feat(checkpoint): checkpoint evaluation runs from rollout collection #3893, which is no longer below this PR.gym.checkpoint.operation_duration_ms, by operation, participant kind and outcome;gym.checkpoint.records_totalandgym.checkpoint.bytes_total(unitBy), for commit and restore, counted once per write that finishes and per restore whose install succeeds, so a retry never counts the same volume twice;gym.checkpoint.control_errors_total, by error code: every failed control call, counted once in the control routes;gym.checkpoint.events_total, which counts:lease_expired: once per checkpoint whose lease expired, which aborts that checkpoint; each attempt to reopen afterwards is alease_reopenoperation with its own outcome;prepare_not_ready: a prepare reached its deadline with blockers, with the full blocker count;generation_cut: each cut, by disposition.staton the event loop while the control lock is held.span-groups.mdxlists thecheckpointgroup and how to enable it, andmetrics.mdxlists everygym.checkpoint.*instrument with its unit and attributes.keyandtoken, which the span helpers redact.The spans were used to profile commit and restore at 16,000 rollouts. The fixes that profile led to are in the core, policy model, environment, resources, and agent PRs below (cheaper record payloads and ledger imports).
How it works
Where this PR sits in the overall flow
The highlighted part is what this PR adds.
flowchart LR C["Controller<br/>NeMo RL, or rollout collection"] CO["Coordination<br/>prepare, commit, restore, resume,<br/>retire, forget"] K["Control plane on every server<br/>phases, lease, storage,<br/>retire: refuse, stop, free;<br/>forget"] subgraph G["One participant per Gym server"] E["Environment server<br/>episode steps"] M["Policy model<br/>held responses, generation cuts"] A["Agent<br/>sessions parked at boundaries"] R["Resources server<br/>session state"] end W["Inference worker<br/>stages cut prefixes"] D[("Checkpoint directory<br/>records, then manifest")] L[("Capture ledger<br/>retire and delete from #3938, #3939")] C --> CO --> K K --> E & M & A & R M --> W M --> L G --> D classDef this fill:#fde68a,stroke:#b45309,stroke-width:2px,color:#1f2937 class CO,K thisOne checkpoint, a crash, and the restore, end to end:
This PR
The spans of one checkpoint. Control calls carry the trace context, so the controller's spans and every participant's spans form one trace across Gym's processes.
Where this sits in the stack
This is one PR in a stack of draft PRs for partial-rollout checkpointing. Each PR is one commit, and its base is the branch of the PR before it, so each diff shows only that PR's change.
The stack builds on the token-capture cleanup PRs #3938 (capture ledger
retireanddelete) and #3939 (complete-recordretireanddelete), which have merged. #3882 to #3889 have merged too, so #3903's base ismain. The checkpoint stack uses those operations to free the ledgers of retired attempts and to clear a dead execution's capture files before a restore. The capture-ledger cache bound (#4268) and the striped lock files (#3937) are independent of the stack.ananthsub/partial-ckpt-core): feat(checkpoint): add the participant control plane, episode steps, and coordination (merged)ananthsub/partial-ckpt-policy-model): feat(checkpoint): make policy model servers checkpoint participants with generation cuts (merged)ananthsub/partial-ckpt-model-worker-cuts): feat(token-capture): add worker staging helpers for generation cuts (merged)ananthsub/partial-ckpt-environment): feat(checkpoint): continue environment server episodes from their boundaries (merged)ananthsub/partial-ckpt-resources): feat(checkpoint): add the resources server participant with asynchronous session hooks (merged)ananthsub/partial-ckpt-verifier-declarations): feat(checkpoint): declare twenty training verifiers stateless with replayable verification (merged)ananthsub/partial-ckpt-agent): feat(checkpoint): add the agent session participant and Simple Agent continuation (merged)ananthsub/partial-ckpt-e2e): test(checkpoint): add a process-level end-to-end suite driven by coordination (merged)ananthsub/partial-ckpt-telemetry): feat(checkpoint): spans and metrics for partial-rollout checkpoints (this PR)ananthsub/partial-ckpt-multi-worker): feat(checkpoint): resources, agent, and environment servers with several workersananthsub/partial-ckpt-fault-tests): test(checkpoint): check checkpoint invariants over random sequences of operations#3893 (
ananthsub/partial-ckpt-rollout-collection, checkpoint evaluation runs from rollout collection) is independent of this chain and targetsmain. Its two spans,gym.checkpoint.collection.checkpointandgym.checkpoint.collection.restore, are added to it once #3903 merges, or in a small follow-up if #3893 merges first.Server ports that build on the stack: #3894 (Workplace Assistant), #3895 (Gymnasium), #3896 (Blackjack), #3897 (indirect prompt injection), #3898 (proof refinement). Session routing for several workers is #3905, against main.
Issue
No tracking issue exists. Observability was a planned follow-up for the stack.
Validation
Run on this branch, on the current
main:RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/unit_tests/telemetry tests/unit_tests/test_checkpoint_*.py: passed.tests/unit_tests/telemetry/test_checkpoint_telemetry.pydrives a real resources participant through its control routes and asserts against a real in-memory span exporter and metric reader:outcome=error, including a resume that fails on a participant;staton the event loop;checkpointnext tomain'sstartup; thedefaultpreset is unchanged, withstartupin it andcheckpointout of it.tests/unit_tests, eight processes, withouttest_opensandbox_compose.py, which needs the optionalopensandboxpackage): passed except a sandbox retry test, which fails the same way onmainin this development environment.pre-commit run --files <files changed by this PR>: all hooks passed, and no hook modified a file.Signed-off-byline.Rollout evidence
The 16,000-rollout scale test was profiled with the console exporter. The per-participant breakdown is in the next PR. Real-model runs are in #3889, #3893 and #3909.
Compatibility
checkpoint, opt-in. Existing presets are unchanged, somain'sstartupgroup stays on by default.gym.checkpoint.*, listed inmetrics.mdx.gym.checkpoint.control_errors_totalwith thenemo.gym.checkpoint.error_codeattribute; there is norefusedevent, and refusals of data-plane requests are not counted.write_participant_stateandread_participant_statein the privatenemo_gym/_checkpoint/store.pyalso return the byte count; the on-disk format is unchanged.