Repository navigation
Conversation
pthombre
added
feature
675c3f4 to
4d55d7a
Compare
|
/ok to test 4d55d7a |
|
🌿 Preview your docs: https://nvidia-preview-pr-5264-4d55d7ab04e7.docs.buildwithfern.com/nemo/gym |
svcnvidia-nemo-ci
left a comment
svcnvidia-nemo-ci
left a comment
There was a problem hiding this comment.
HOLD — one BLOCKER and four RISK findings.
I reviewed all 34 changed files and traced the affected restore, grant, budget, export and release paths. A queued replacement activation can lose its saved continuation; the inline findings cover that and three other checkpoint reliability defects.
RISK — Provide representative real smoke-rollout evidence
The PR reports fake-based tests and treats real rollouts as unnecessary protocol validation. However, it also changes existing episode behavior: single-agent-turn budget charging affects replacements using existing agents, independently of a sandbox harness opting in. The trusted quality bar requires representative real model rollouts and inspection of agent and verifier behavior for behavior-changing environment or agent work.
Could you add that evidence after correcting the defects? Use a real model with an existing checkpointable agent/resources pairing, record the configuration and commands, and inspect an uninterrupted rollout plus a rollout checkpointed during activation and continued after restore. Show that the restored agent retains its progress, verification still produces the expected task result, and cleanup completes. Also exercise the remaining-budget timeout case: the replacement must receive only the unspent budget, excluding checkpoint parking. Preserve the uninterrupted pairing's reward, observations and cleanup behavior. No verifier or reward formula changed here, so this request is for smoke-rollout evidence rather than a benchmark reward-distribution comparison.
I couldn't execute tests or rollouts in this read-only review. All acceptance checks below are proposed; suggested code needs verification before adoption.
Note: formal approval is not available for this request, so this review is posted as a comment only.
Review generated by codex. This was a full review. Request a full review with /review reset=true.
| @asynccontextmanager | ||
| async def queued(self, session: _Session) -> AsyncIterator[None]: | ||
| previous = session.state | ||
| session.state = "at_boundary" |
There was a problem hiding this comment.
BLOCKER — Preserve a restored continuation while its activation queues.
Restore a session with boundary={"step": 4}, enter its activation, and wait for a semaphore inside activation.queued() before processing the continuation. The activation constructor moves the saved boundary into activation.continuation and clears session.continuation. This new queue context then makes the session checkpoint-ready, but export finds neither a boundary callback nor a continuation and writes boundary=None. A second restore silently loses step 4 and can replay already completed work. Queuing must preserve the latest resumable boundary, not merely readiness.
A small correction is to retain the restored boundary as the activation's initial snapshot. In Activation.__init__, immediately after moving the continuation:
if self.continuation is not None:
saved_boundary = self.continuation
session.boundary = lambda: saved_boundarySource tracing shows that activation.boundary() subsequently replaces this callback, while activation teardown clears it. Fresh queued activations without a continuation should still export no loop boundary; a fresh or restored activation that records a newer boundary must export that newer boundary instead.
Proposed acceptance: install a record with {"step": 4}, start its replacement activation, block inside queued() on a zero-capacity semaphore, then prepare/export. Assert the exported boundary is exactly {"step": 4} and a further restore delivers it as the next activation's continuation. Also check fresh/no-boundary export remains None and a newly recorded {"step": 5} supersedes the restored boundary. Retain the existing queue test's assertions that prepare completes while queued and a slot granted during the checkpoint cannot start work before resume.
There was a problem hiding this comment.
Reproduced on 4d55d7a: restoring {"step": 4}, entering the activation and blocking in queued() leaves activation.continuation == {"step": 4}, but prepare/export writes boundary=None, because the constructor (agent.py#L181) cleared session.continuation and _export_boundary (#L785-L789) only reads session.boundary or session.continuation. Your fix works: session.boundary = lambda c=self.continuation: c keeps {"step": 4} while queued, a recorded {"step": 5} supersedes it, teardown clears it, and a further restore delivers {"step": 4}. The same hole exists today through awaiting_model().
One severity note: no in-tree agent calls queued() yet (Simple Agent calls commit_boundary() first), so I'd grade this RISK rather than BLOCKER. Since this PR introduces queued(), the fix and the restore, queue, checkpoint, restore test should land here.
| the agent re-binds its restored session to that sandbox and to this attempt's tool grants. | ||
| Without a borrowed sandbox nothing changed hands, and nothing is seeded. | ||
| """ | ||
| if not handles.borrowed_sandbox: |
There was a problem hiding this comment.
RISK — Refresh MCP grants even without a borrowed sandbox.
An exported agent may declare x-ng-checkpoint-rebinds-tools: 1 while using resources MCP tools without borrowing a resources sandbox. _seed_agent now allows that episode to continue, but this guard skips both re-seeds after restore. The agent therefore keeps the previous attempt's token. Resources admission rejects that token because the restored session belongs to the replacement attempt. The new MCP capability test explicitly admits this no-sandbox combination but doesn't continue it after restore.
Could you make reseed_after_restore() cover attempt-bound MCP grants as well as borrowed sandboxes? For the current protocol, replacing the guard with the following covers both reasons to re-seed:
if not handles.borrowed_sandbox and "mcp" not in self.config.resources_tool_transports:
returnUpdate its docstring/log message accordingly. This relies on the existing seed contract: resources re-seeding must preserve restored task state, and the agent rebind must take the replacement grants. Keep no-sandbox direct-HTTP continuations unseeded, borrowed-sandbox continuations re-seeded, and agents without the rebind header marked as restarts. Do not re-seed completed verify/return stages.
Proposed acceptance: for both environment variants, checkpoint a no-sandbox MCP session whose agent declares rebinding, restore at invoke_agent, and assert the replacement receives a newly minted token and a tool call succeeds against the preserved resources state. The old token must remain refused. Cover close_agent as well, and retain the existing no-sandbox/direct-HTTP no-reseed and non-rebinding MCP restart assertions.
There was a problem hiding this comment.
Confirmed for both environment variants: with ["mcp"], an agent that declares re-binding and no borrowed sandbox, a restore at invoke_agent makes no /seed_session or /v1/agent_sessions call (app.py#L369), and resources admission then refuses the old attempt's token (resources.py#L206-L214).
Two ways to close it. Your guard ("mcp" in resources_tool_transports) re-mints the token in both variants, but it relies on a repeat /seed_session of a restored session keeping task state without re-running setup. No docstring in this PR states that rule (only the sandbox-owner base in #5266 is meant to do it), so please document it on seed_session and sandbox_access. The alternative is to honour the rebinds header only when a sandbox is lent (seed.sandbox_access is not None); that turns the MCP-rebind cases of test_an_episode_that_grants_mcp_access_restarts_unless_its_agent_rebinds_grants into restarts, so they would need a sandbox. No in-tree agent declares the header today, so this is latent until a harness does.
| self.checkpoint_charge_elapsed(request, handles.elapsed_s) | ||
|
|
||
| def bookmark(next_stage: str) -> dict[str, Any]: | ||
| elapsed = self.checkpoint_elapsed(request) |
There was a problem hiding this comment.
RISK — Capture elapsed time at the checkpoint cut, not only before the replay step.
With a 30-second episode budget, enter invoke_agent near time zero, let the agent work for 20 seconds, then checkpoint while that replay step is still running. The saved environment bookmark still has the elapsed value recorded before invocation. Environment export writes that unchanged bookmark, even though the agent's checkpoint can preserve its later progress. On restore, line 180 charges approximately zero, so the continued agent gets another 30 seconds instead of approximately 10. Long activations—the main partial-rollout use case—therefore still receive a renewed budget.
Could you arrange for the persisted elapsed value to be sampled at the actual checkpoint cut for an in-flight replay step? bookmark(), EnvironmentParticipant.export_records() and the EpisodeSteps elapsed/restore contract need to preserve the original protocol boundary while carrying the current consumed budget. Merely moving this call into the existing boundary lambda won't fix it: checkpoint_boundary() evaluates that lambda when the boundary is recorded, not at export. The exact owner/runtime representation needs a decision, so I don't have a safe local replacement block to suggest.
Preserve zero elapsed for older bookmarks, unbounded episodes, cumulative charging over repeated restores, and exclusion of checkpoint parking. Do not advance the saved stage or replay a completed state-changing verification just to refresh its time.
Proposed acceptance: use a controlled clock or held activation with a 30-second budget; spend 20 seconds inside invoke_agent, prepare/export before the invocation completes, and assert the saved budget charge is approximately 20 seconds. Restore and hold the activation for another 11 seconds: it must time out. Add a second checkpoint/restore to verify charges accumulate rather than reset, and a parked interval to verify it contributes no elapsed time. The existing budget tests cover prefilled elapsed values and completed seed work, not this in-flight case.
There was a problem hiding this comment.
Confirmed: with a 3 s budget and the agent held about 1.2 s inside invoke_agent, the exported bookmark carries elapsed_s of about 0.0003 while steps.elapsed(key) reads 1.2 at the same cut, so the replacement gets nearly the full budget back. The bookmark dict is built when the boundary is recorded (app.py#L239) and steps.exported() returns it as stored (steps.py#L345).
The deadline is already suspended at the cut for replay steps, so sampling at export is accurate. I'd keep the bookmark layout out of core: an optional EpisodeRecord.elapsed_s (a default keeps old checkpoints loading), filled in EnvironmentParticipant.export_records from steps.elapsed(key) and charged in begin(), keeps repeated restores cumulative and parked time excluded, and removes checkpoint_elapsed and checkpoint_charge_elapsed from every environment author. If you'd rather defer this to the OpenCode PR, please drop the budget wording from the description, including "no behavior change until a server opts in", since the charge applies to every restored single_agent_turn episode with a timeout.
| with self._operation("release", request.checkpoint_id) as span: | ||
| span.set(kept=len(request.kept)) | ||
| async with self._lock: | ||
| if self.phase != CheckpointPhase.IDLE or self.checkpoint_id is not None: |
There was a problem hiding this comment.
RISK — Fence obsolete release requests after a newer checkpoint resumes.
The idle check prevents release during a checkpoint, but not after it. A concrete sequence is: commit/resume C1; send a still-valid release(checkpoint_id="C1", kept=["C1"]) whose delivery is delayed; commit/resume C2; then admit the delayed C1 release before its deadline. Both phase conditions pass, and the hook receives {C1}, authorizing deletion of C2's external snapshots although the controller retains C2. This violates the safety invariant in this method's docstring. Resume clears the active identity, so no comparison here detects checkpoint progression.
Could you fence releases against subsequent checkpoint progression, retaining the necessary identity or generation across resume? A release authorized before a subsequent checkpoint must never garbage-collect that checkpoint's external state. The representation needs a protocol decision, so I don't have a compatibility-safe replacement block to suggest. Calling _admit() unchanged isn't sufficient: normal releases intentionally name an already-resumed checkpoint.
Preserve current-identity release after resume, identical retries, best-effort participant fan-out and cancellation-aware all-worker forwarding. Final empty retention must remain possible; prepared, committed-but-not-resumed and restored phases must still refuse release.
Proposed acceptance: use a fake external-state hook holding C1 and C2 snapshots. Delay C1's release until C2 has committed/resumed, keeping the old request's deadline valid; assert it is refused without invoking deletion and C2 remains restorable. Release with the current authorization and {C2}; assert C1 is deleted and C2 retained. Repeat that release, exercise final empty retention, and retain the existing phase-refusal and all-worker forwarding checks.
There was a problem hiding this comment.
Reproduced at the control-plane level: commit and resume c1, build release(c1, kept=[c1]) with a long deadline, commit and resume c2, then deliver it. It is admitted (control.py#L820-L832) and a hook that deletes everything outside kept removes c2's state, because resume forgets the active id (#L1179) and the request's checkpoint_id is never compared.
It needs delivery after a whole later checkpoint cycle but still within the request's deadline, so I'd grade it NOTE unless a controller can fire releases without awaiting them. Refusing a release whose checkpoint_id is not the most recently resumed id (the participant already keeps _resumed_ids in order) closes it; please keep an empty kept allowed for the end-of-run call. Either fence it, or narrow this docstring's "must never" wording and the "safe to retry" line at coordination.py#L7.
Every protocol change the environment-server checkpointing of sandboxed harnesses needs between the four server types, reviewed together and tested with fakes: - Control plane: `release(kept)` admitted only while idle, fanned out to every participant best effort (failures returned per participant) and forwarded to every worker; export hooks receive the `checkpoint_id` they write (keyword-only); `CheckpointSettings.state_dir` for servers that record external state such as snapshots; optional `release_external_state(kept)` hooks. In-tree overrides (genrm_compare, gymnasium, workplace_assistant, example_session_state_mgmt, simple_agent, tool_simulation_agent) take the new keyword. - Per-instance checkpoint mode: `checkpoint_mode_in_effect()` on agents, resources and environment servers, read by participant setup, status/discovery and seed replies; defaults to the ClassVar. - Park signal for agent sessions: `Activation.park_requested`, `Activation.wait_park_requested()` (set by prepare, at activation entry while closed, and by retire/close; cleared by resume; never for an episode that restarts) and `Activation.queued()`, so an activation waiting for the agent's own concurrency limit counts as parked. - Agent base: a replacement attempt's seed of a restored, unclaimed session calls `_rebind_restored_agent_session_state` instead of answering 409 (the default still refuses); `_seed_restarts(body)` marks a session the agent cannot capture as a restart (header on the seed reply, never exported or waited for); `rebinds_tool_accesses` is reported with `x-ng-checkpoint-rebinds-tools`. - single_agent_turn (and the legacy flat-row server, which inherits it): the bookmark records `borrowed_sandbox` and `elapsed_s`; a restore at `invoke_agent` or `close_agent` re-seeds both sessions in a replay step (`reseed_after_restore`), so the agent re-binds to the forked sandbox; a replacement attempt only has the rest of the episode budget; resources MCP grants force a restart only when the agent seed reply does not say it re-binds tool grants. Provenance: the control-plane part of #4220 and #4222, generalized from LegacyRun to agent sessions. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Pranav Thombre <pthombre@nvidia.com>
| resources_verify_step_type: StepType = "wait" | ||
| # Whether the agent borrows a sandbox the resources server lent in its seed reply. | ||
| # A restore rebuilds that sandbox under a new id, so the replacement attempt seeds both sessions again. | ||
| borrowed_sandbox: bool = False |
| borrowed_sandbox: bool = False | ||
| # Seconds of the episode's time budget spent when the boundary was recorded, not counting parked time. | ||
| # A replacement attempt that continues the boundary has only the rest. | ||
| elapsed_s: float = 0.0 |
There was a problem hiding this comment.
elapsed_s naming a bit general
| # replaces it. An agent that keeps the token it was seeded with would see every tool call fail | ||
| # after a restore, so the episode restarts from its input after a crash instead. | ||
| # An agent that takes the new token its replacement attempt's seed carries can continue. | ||
| await self.checkpoint_restart(request) |
There was a problem hiding this comment.
two conditions both lead to "await self.checkpoint_restart(request)", combine for clearity?
| # They are restarts, which never hold up a checkpoint and start over from their input after a crash. | ||
| checkpoint_mode: ClassVar[CheckpointMode] = "restart_only" | ||
|
|
||
| def checkpoint_mode_in_effect(self) -> CheckpointMode: |
There was a problem hiding this comment.
what case needs to override?
| @@ -421,6 +476,14 @@ async def _seed_agent( | |||
| handles.agent_cookies = _cookies(agent_create_http_response) | |||
| if seed_restarts(agent_create_http_response.headers): | |||
There was a problem hiding this comment.
If a checkpoint opens while this seed is still in its replay step, mark_restart waits until that checkpoint resumes, but the agent session is already open. Does this checkpoint then export the session at the seed boundary? A restore would fresh-seed into that restored session, and the default rebind answers 409. On main this restart was marked before the session existed. Would it be safer to record the restart before the session can be exported?
| if self._live(session): | ||
| session.state = previous | ||
| raise | ||
| if session.park_requested: |
There was a problem hiding this comment.
park_requested and wait_park_requested both leave restart sessions running, but queued() parks on the raw flag, which close_admission sets for every session. If a restart episode is granted its slot during a checkpoint, does it sit in park() until resume while its deadline keeps running? Should this skip park() when restarts_session is true?
yaoyu-33
left a comment
yaoyu-33
left a comment
There was a problem hiding this comment.
Verdict: SHIP WITH CARE. No blocker on this head. The RISKs below should be fixed or accepted deliberately before merge.
This PR adds the protocol that the rest of the stack needs for harnesses that borrow a sandbox: the release(kept) control operation, checkpoint_id in the export hooks, the park signal and Activation.queued(), agent re-bind and _seed_restarts, and the re-seed, remaining-budget and MCP-grant changes in single_agent_turn. Only this PR is reviewed here; the later stack PRs are not.
Merge conflict. The branch conflicts with main (mergeStateStatus: DIRTY). Could you rebase onto main and re-run the checks on the new head?
Earlier review. I reproduced its four inline findings on this head with probe tests and replied on each thread. The queued() boundary loss is real, but I'd grade it RISK rather than BLOCKER: nothing in-tree calls queued() yet, and Simple Agent calls commit_boundary() first. The MCP re-seed gap and the stale release are latent until a harness opts in. The elapsed-budget gap is not latent: it affects every restored single_agent_turn episode that has a timeout.
New inline comments.
- RISK: a slot taken inside
queued()leaks when the session is retired while it parks. - RISK: two test gaps that let a regressed re-seed or a no-op
ResourcesParticipant.releasepass every test. - NOTEs: restart-mark timing, multi-worker park signal, startup validation of hooks, the
release_external_statecontract, and a private import.
PROCESS
- The description lists suites but no exact commands. The 16-node run covers the whole stack, and no configuration or logs are attached to this PR.
- "No behavior change until a server opts in" is not accurate: a replacement attempt of any restored
single_agent_turnepisode with a timeout now gets only the remaining budget (and the in-flight activation is not charged yet, see the thread onapp.py). AGENTS.md "Quality bar" asks for smoke rollouts on behavior-changing environment work. One real-model run through Simple Agent andsingle_agent_turnwith an episode timeout and one kill/restore mid-activation, with the commands and config recorded, would be enough here, and it would also show the elapsed gap. - The
stop_rolloutdocstring path edit inbase_responses_api_agent.pyis unrelated to this PR. - Docs:
span-groups.mdx:93andtelemetry.py:10list the control operations withoutrelease, which now has a span. The docs for the rest of the new surface (the rebinds header, the hooks,state_dir) are in #5269; please make sure they land before a harness author depends on them. state_dirhas no reader in this PR. When its consumer lands, require an absolute path and keep it per server and per run.
Checked, no finding
- Simple Agent and tool_simulation_agent behave the same before and after: ten seed cases (fresh, duplicate, other task or attempt, restored and unclaimed or claimed) give identical results on the base and on this head.
- Every in-tree override of
export_session_statesandexport_agent_sessionstakescheckpoint_id. Nohttpx.AsyncClient,ray.get, blocking call or undecoded subprocess output is added, and no verifier or reward code changes. Old bookmarks load through the new defaults. - Passing on this head: the server_runtime and model checkpoint suites,
single_agent_turnand its legacy variant,simple_agent,tool_simulation_agent, and the touched resources-server suites. - Not verified: GPU tests, real model rollouts and real sandboxes.
| raise | ||
| if session.park_requested: | ||
| # A slot granted while a checkpoint is open: the checkpoint may have counted this session as parked. | ||
| await self.park(session) |
There was a problem hiding this comment.
RISK: a slot taken inside queued() leaks when the session is retired or closed while it parks. The park at L622 runs after the yield, so a retire or close during it raises out of __aexit__ before a caller's try/finally: release() is entered. A probe with Semaphore(1), await sem.acquire() inside async with act.queued():, a checkpoint open and a retire while parked leaves the semaphore at 0, so every occurrence removes one slot from the agent's concurrency limit for good.
Could queued() own the slot (for example an activation.slot(sem) helper that releases on any exception), or could its docstring require the flag-and-release shape (acquired = True in the body, release in an outer finally), with a test for it?
|
|
||
|
|
||
| @pytest.mark.parametrize("server_type", [SingleAgentTurnEnvironmentServer, SingleAgentTurnLegacyEnvironmentServer]) | ||
| async def test_a_restore_at_invoke_agent_re_seeds_both_sessions_onto_the_forked_sandbox(server_type: type) -> None: |
There was a problem hiding this comment.
RISK (tests): nothing proves borrowed_sandbox is recorded from a real seed. Every re-seed test writes borrowed_sandbox=True into the restored handles by hand (_restored_handles(borrowed_sandbox=True)). Replacing handles.borrowed_sandbox = seed.sandbox_access is not None with pass keeps all 138 tests green in both variants, yet production would then never re-seed.
Could you add a fresh-run test whose resources seed returns sandbox_access, asserts that the exported bookmark carries borrowed_sandbox: true, restores it and asserts the re-seed? A check that a checkpoint does not wait for the re-seed belongs next to it: it is a replay step, and turning it into wait also passes today.
| # The checkpoint just written is not published until resume: a list computed before it would drop it. | ||
| await controller.release(ReleaseRequest(**control("c1", kept=["c1"]))) | ||
| await controller.resume(CheckpointRequest(**control("c2"))) | ||
| reply = await controller.release(ReleaseRequest(**control("c2", kept=["c1", "c2"]))) |
There was a problem hiding this comment.
test_checkpoint_harness_protocol.py:117-133
RISK (tests): the real ResourcesParticipant.release and the checkpoint_id it forwards are never exercised. These tests release through the fake _Releasing participant, and every export_session_states fake accepts and ignores the keyword. Replacing the body of ResourcesParticipant.release with pass leaves 293 tests passing across the control, coordination, resources, worker and protocol suites, and passing a wrong checkpoint_id to the hook passes too.
Could you add one test with a real ResourcesParticipant in exported mode (release reaches release_external_state(kept), stateless and restart_only skip it, and the hook sees the committed checkpoint's id), and one over the multi-worker export body?
| handles.agent_cookies = _cookies(agent_create_http_response) | ||
| if seed_restarts(agent_create_http_response.headers): | ||
| await self.checkpoint_restart(request) | ||
| elif "mcp" in self.config.resources_tool_transports and not seed_rebinds_tools( |
There was a problem hiding this comment.
NOTE: the restart mark now comes after the agent's reply, so a checkpoint can record the episode in between. The removed code marked the restart before the agent received the token ("no checkpoint ever records the episode past its seed"). Now a checkpoint that lands while the agent-seed POST is in flight exports the environment at seed together with the agent session that holds the old attempt's token. A probe shows restarts=[] and exported=[('rollout-a2', 'seed')] here, against restarts=['rollout-a2'] and nothing exported on the base.
After a crash the replacement's re-seed carries a new token, which an exported agent that does not re-bind answers with 409, so the episode fails instead of restarting (that half is from reading the code; I did not run it). It needs an out-of-tree exported MCP agent that neither restarts nor re-binds. Defaulting _seed_restarts on the agent to "has an MCP grant and does not set rebinds_tool_accesses" would let the agent register the session as a restart itself.
| # replaces it. An agent that keeps the token it was seeded with would see every tool call fail | ||
| # after a restore, so the episode restarts from its input after a crash instead. | ||
| # An agent that takes the new token its replacement attempt's seed carries can continue. | ||
| await self.checkpoint_restart(request) |
There was a problem hiding this comment.
NOTE (docs): the README is now wrong. It says an episode that grants mcp access restarts from its input after a crash. After this change that holds only unless the agent declares x-ng-checkpoint-rebinds-tools. Could you reword it, and mention the re-seed of a borrowed sandbox?
| assert len(agent.rebound) == 1 | ||
|
|
||
|
|
||
| async def test_a_re_bind_is_refused_for_another_attempt_and_by_an_agent_that_does_not_re_bind(tmp_path: Path) -> None: |
There was a problem hiding this comment.
test_checkpoint_harness_protocol.py:475-485
NOTE (tests): the "agent that does not re-bind" case does not test the default gate. _BorrowingAgent(rebinds=False) overrides the hook and delegates to super(), so _rebinds_restored_sessions() is true for it. Making that method always return True passes every test here and in Simple Agent's suite, yet it turns Simple Agent's identical re-seed of a restored session from 200 into 409, which the description says is unchanged.
Could you add a case with an agent that does not override the hook (Simple Agent, or a stub): an identical seed answers 200 and a different seed answers 409?
|
|
||
| async def export_session_states(self, session_ids: list[str]) -> dict[str, JsonValue]: | ||
| """Return each session's state for a checkpoint, leaving out sessions this server no longer holds.""" | ||
| async def export_session_states(self, session_ids: list[str], *, checkpoint_id: str) -> dict[str, JsonValue]: |
There was a problem hiding this comment.
NOTE: an out-of-tree override with the old signature starts normally and fails at the first commit. setup_resources_checkpoint (L581-L589) checks that the hooks are overridden, not that they accept checkpoint_id. A server with export_session_states(self, session_ids) serves traffic until its first commit, after prepare has closed admission, and then fails with TypeError: unexpected keyword argument 'checkpoint_id', possibly hours into a run.
The break is documented in the description and every in-tree override is updated, so this is only about failing early. An inspect.signature check at startup (accepts checkpoint_id or **kwargs, ValueError naming the class) for export_session_states and for export_agent_sessions (base_responses_api_agent.py:670) would do it, or a one-release shim that retries without the keyword.
| # such as a new MCP session token, in ``_rebind_restored_agent_session_state``. | ||
| # The seed reply says so, and an environment server then continues an episode with tool grants after a restore | ||
| # instead of starting it over. | ||
| rebinds_tool_accesses: ClassVar[bool] = False |
There was a problem hiding this comment.
base_responses_api_agent.py:614
NOTE: rebinds_tool_accesses = True without overriding _rebind_restored_agent_session_state is accepted silently. The seed reply then tells the environment server not to restart MCP episodes, and after a restore the default hook answers 409, so those episodes fail instead of restarting. checkpoint_mode is already validated at startup; a ValueError in setup_agent_checkpoint when rebinds_tool_accesses is set but _rebinds_restored_sessions() is false would catch it there.
| a terminal step among them, are refused. | ||
| """ | ||
|
|
||
| async def release_external_state(self, kept: frozenset[str]) -> None: |
There was a problem hiding this comment.
base_resources_server.py:689-695
NOTE (docs): release_external_state needs a fuller contract. Implementers will depend on these behaviors, and the docstring states none of them:
- The hook runs under the control-plane lock, so a slow one delays the next prepare until the release deadline (a prepare sent during a 2 s hook completed at 2 s).
- At the deadline it is cancelled mid-way (409
checkpoint_deadline_exceeded), so it must be idempotent and cancellation-safe. - Admission stays open, so it runs beside live seeds and tool calls.
- With several workers, each worker runs it concurrently against the same external store and
state_dir. - Retention must be by reference: a later checkpoint can re-export a snapshot created under an older checkpoint id.
Could you add these here and to the agent hook's docstring?
| seed_restarts, | ||
| seed_verify_step_type, | ||
| ) | ||
| from nemo_gym.server_runtime._checkpoint.steps import seed_rebinds_tools |
There was a problem hiding this comment.
NOTE: import seed_rebinds_tools from the public API. AGENTS.md has code outside nemo_gym/ import from nemo_gym.server_runtime and add missing names to _PUBLIC_API. This is the only server import of a _checkpoint submodule, and the sibling seed_restarts is exported ("seed_restarts": "_checkpoint.steps", __init__.py:149). Adding "seed_rebinds_tools": "_checkpoint.steps" there and importing it with the other names fixes it.
4d55d7a to
02f62ff
Compare
github-actions
Bot
added
the
sla:review-overdue
…t instead of waiting for it A rollout checkpoint lands right after a training step dispatches its rollouts, while hundreds of first seeds are still running: the resources server creating each sandbox and the agent setting its harness up in it. Prepare waited for every one, 70-80 s at 16 nodes, while every rollout already running sat parked. Such a seed has done no work of its episode worth saving. When admission closes, the agent and resources participants now record the rollouts whose first seed is still in flight and report their episodes as restarts of that checkpoint: prepare does not wait for them, and the commit leaves them out, including the sessions they register before resume. After a crash those rollouts start over from their input; the live run is unchanged, and after resume the next checkpoint captures them like any other. A seed of a session that exists, such as one re-binding a restored session to its rebuilt sandbox, still blocks prepare, so restored progress stays in the checkpoint. At 16 nodes on AgentENV (runs 8639483 and 8641901 against 8634011), steady-state checkpoints park rollouts for 7-36 s instead of 79-107 s, and leave 200-370 fewer sandboxes to snapshot. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Pranav Thombre <pthombre@nvidia.com>
What changed and why
The partial-rollout checkpoint stack on
maincheckpoints environment-server episodes whose agent and resources state live in Gym processes. It cannot yet checkpoint an episode whose agent harness runs in a sandbox the resources server lends (for example OpenCode on SWE-bench):This PR is every protocol change between the four server types that the rest of the stack needs. The pieces are small and only useful together, so they are reviewed together and tested with fakes. Nothing here takes a snapshot or pauses a process.
release(kept): a new control operation, admitted only while the participant is idle with no checkpoint open. The controller calls it after saving a checkpoint, with every checkpoint it may still restore, so participants can free external state that only dropped checkpoints need. Coordination fans it out best effort and reports failures per participant. Multi-worker servers forward it to every worker.checkpoint_idthey write (keyword-only), so a server can name external snapshots after the checkpoint and find them again atrelease.CheckpointSettings.state_dir: a directory stable across restarts, for servers that index external state.release_external_state(kept)on resources servers and agents (no-op by default).exportedandstatelessmodes;restart_onlyalready did this). The controller leaves those rollouts out of the commit, so prepare no longer waits for sandbox creation, repository setup and harness setup. In large runs a checkpoint often lands just after a batch is dispatched, with hundreds of first seeds in flight, and waiting for them was most of the time rollouts stayed parked. The set is fixed when admission closes and cleared at resume: the seed keeps running, and its session is in the next checkpoint. A seed of a session that exists, such as the re-seed of a restored session, still blocks prepare, so restored progress stays in the checkpoint.checkpoint_mode_in_effect()on agents, resources and environment servers, defaulting to thecheckpoint_modeClassVar. A sandbox owner (feat(resources): sandbox owner base for resources servers #5266) reportsrestart_onlywhen its provider cannot snapshot.Activation.park_requestedandwait_park_requested()tell agent code running outside Gym (a harness in a sandbox) to park.Activation.queued()wraps the agent's own concurrency limit, so an activation still waiting for a slot counts as parked instead of holding up prepare. An episode that will restart is never asked to park._rebind_restored_agent_session_stateinstead of answering 409. The default still refuses, so Simple Agent is unchanged._seed_restarts(body): an agent can mark a session it cannot capture as a restart; that session is never exported or waited for.rebinds_tool_accesses: an agent that re-renders tool grants at re-bind says so in its seed reply (x-ng-checkpoint-rebinds-tools).single_agent_turnand its legacy variant)borrowed_sandboxandelapsed_s(budget spent, excluding time parked for checkpoints). Both have defaults, so existing bookmarks load.invoke_agentorclose_agentwith a borrowed sandbox re-seeds the resources session and then the agent session inside a replay step, so the agent re-binds to the forked sandbox.This generalizes the control-plane part of #4220 and #4222 (legacy stack) to agent sessions.
How it works
Where this PR sits in the overall flow
The highlighted parts are what this PR changes.
flowchart LR C["Controller<br/>NeMo RL, or rollout collection"] CO["Coordination<br/>prepare, commit, resume, restore,<br/>retire, forget, release"] E["Environment server<br/>bookmarks, re-seed after a restore"] A["Harness agent, e.g. OpenCode<br/>run until parked, re-bind, resume"] S["Process supervisor in the task sandbox<br/>pause, continue, stop"] R["Resources server, e.g. SWE-bench<br/>owns and lends the task sandbox"] OB["Sandbox owner base<br/>export, restore, capture_once"] K["Sandbox checkpointer<br/>snapshot index, release, gc"] P["Sandbox provider, e.g. OpenSandbox<br/>snapshot, fork, delete, list"] M["Policy model<br/>held calls, capture ledger"] D[("Checkpoint directory<br/>and state_dir")] C --> CO --> E & A & R & M E -- "seed, re-seed" --> R E -- "seed, re-seed, activation" --> A R -. "sandbox_access" .-> A A --> S R --- OB --> K --> P A -- "model calls" --> M CO --> D classDef this fill:khaki,stroke:darkgoldenrod,stroke-width:2px,color:black class CO,E,A thisThis PR
A crash after a checkpoint, then the restore. The new parts are the re-seed in the replay step and the agent's re-bind.
sequenceDiagram participant C as Controller participant E as Environment server participant R as Resources server participant A as Agent Note over C,A: checkpoint C1 committed while the episode was at invoke_agent<br/>bookmark records borrowed_sandbox and elapsed_s Note over C,A: crash, every Gym process dies C->>E: restore C1 (the bookmark becomes unclaimed) C->>R: restore C1 (the sandbox is forked, S1 becomes S2) C->>A: restore C1 (session installed, not yet connected) C->>E: run, attempt N+1 E->>E: continue at invoke_agent with the remaining budget rect rgb(253, 230, 138) Note over E: replay step, a crash here just re-seeds again E->>R: seed (same session, attempt N+1) R-->>E: sandbox_access(S2), no setup run E->>A: seed (same session, attempt N+1, S2, fresh grants) A->>A: re-bind the restored session A-->>E: reply declares that it re-binds tools end E->>A: activation (the harness continues in S2) C->>E: release(kept=[C1]) after the next saveHow a restored episode continues, by the boundary it was saved at:
seedinvoke_agentclose_agentverify,returnWhere this sits in the stack
This is one PR in a stack for checkpointing environment-server episodes whose agent harness runs in a sandbox the resources server lends (first pair: OpenCode and SWE-bench on OpenSandbox; #5410 to #5413 add AgentENV, where OpenCode resumes in place after a restore). Each PR is based on the one before it, so each diff shows only that PR's change. The stack is based on
mainand runs parallel to thelegacy_agentpath of #4220 to #4936, which it does not change or depend on.Issue
No tracking issue exists. This implements the environment-server plan for checkpointing sandboxed harnesses, on top of the merged partial-rollout checkpoint stack (#3882 to #3909, #4997).
Validation
checkpoint_idin the export hooks, the park signal, re-bind and per-session restarts, and for first seeds in flight becoming restarts (agent, resources, multi-worker) while re-seeds of existing sessions still block.main: the stack's suites pass, and two 16-node AgentENV RL training runs of the rebased stack, a checkpoint run and a restore from its checkpoint, commit every checkpoint with no sandbox loss; the restore resumes in place every held session whose image can run the model relay.Compatibility
export_session_statesandexport_agent_sessionstake a keyword-onlycheckpoint_id. Out-of-tree overrides must add it; in-tree overrides are updated here.release, with the same auth as the others; a controller that never calls it sees no change.checkpoint.state_dir, bookmark fields with defaults), so checkpoints written before this PR restore.x-ng-checkpoint-rebinds-tools. Episodes with resources MCP grants still restart unless their agent declares it; no in-tree agent does until the OpenCode PR.🤖 Generated with Claude Code