[pthombre]

What changed and why

A resources server that owns one sandbox per session and lends it to the agent, as SWE-bench does, needs the same checkpoint wiring every time: a shared provider, a SandboxSessionCheckpointer (#4220), the export/restore/retire/park/release hooks, the working directory, /sandbox_access, and a 409 instead of grading an empty patch. This PR writes that once, as SandboxOwnerResourcesServer (nemo_gym/sandbox/owner.py), and moves the SWE-bench server and the e2e sandbox_notes_server onto it. The base also gets right what each such server would otherwise get wrong on its own.

The owner base. A subclass starts the session's sandbox at seed and ends it when it is done with it:

• start_session_sandbox(session_id, spec, *, setup=None, workdir=None) starts the sandbox on the shared provider, runs setup against it, and only then tracks it. A setup that raises stops the sandbox and leaves nothing for a checkpoint to export. The working directory defaults to pwd after the setup.
• session_sandbox(session_id) returns the sandbox, including one a restore rebuilt: a replacement attempt re-seeding a restored session reuses it, and its setup does not run again.
• require_session_sandbox(session_id) answers 409 when the session has no sandbox, so a resent /verify is refused instead of graded as reward 0.
• end_session_sandbox(session_id) stops the sandbox and forgets the session; its snapshots stay until a release drops the checkpoints that name them.
• shared_provider() is for sandboxes that are not session state, such as SWE-bench's grading sandbox.
• The checkpoint hooks are implemented once: a commit snapshots every live session's sandbox, a restore forks them all or none, retire stops the sandbox and keeps its snapshots, a commit with stop parks the exported sandboxes (sandbox_checkpoint.on_stop), and release_external_state deletes the snapshots of checkpoints the controller no longer keeps.

Per-session data. Some servers learn at seed what verify needs and cannot recompute once the agent has worked; the SWE family keeps base_commit and pristine_untracked in process memory, which a restore loses. A setup may return a small JSON object (SessionData). The base exports it next to the snapshot, restores it all or nothing with the sandbox, and drops it when the session ends; session_data(session_id) reads it. SWE-bench returns none (verify diffs the working tree in the sandbox), so its exported state keeps its shape.

Checkpoint only what can be captured. With checkpointing on, the server decides its mode at startup through a new SimpleResourcesServer.checkpoint_mode_in_effect() hook, whose default is the declared checkpoint_mode:

• An instance that lends no sandbox (lends_sandboxes() false) runs restart_only. SWE-bench's golden-patch instance only grades the dataset's own patch, with 8 workers; training configs load it next to the agent-facing instance, so it must neither need a state_dir nor trip the worker guard.
• A provider that cannot snapshot (docker, apptainer, enroot) gives restart_only and one warning: swebench: sandbox provider 'apptainer' cannot snapshot; its rollouts restart after a crash instead of continuing. Before, every commit with a session in flight failed. Building the provider at startup also turns a bad provider config into a startup error.
• Otherwise exported. Release and park do not reach the workers of a multi-worker server yet, so an exported owner refuses to start with num_workers > 1, and requires checkpoint.state_dir.

Leaving without a verdict. /close_session with an empty body ends the cookie's session, the convention swebench_pro already serves. A checkpointed agent run that fails or is cancelled closes its session this way instead of stopping a sandbox the owner may still export (#4274). A typed ResourcesCloseSessionRequest gets the base default.

SWE-bench on the base. SWE-bench keeps its sandbox spec, the multilingual patch, the anti-cheat scrub (now inside setup, so a failed seed no longer leaks a tracked sandbox), patch extraction, and grading. lends_sandboxes() returns not is_verifying_golden_patch. SandboxCheckpointConfig moves to nemo_gym/sandbox/checkpoint.py.

The second user. The e2e sandbox_notes_server (#4221) runs on the base, so the process-level sandbox scenarios (restore, drift, lease expiry, release, a failed snapshot, the harness runs) exercise it. Its setup returns a seed token and its verify rewards nothing without it, so every crash scenario proves the session data survives a process-level restore.

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"]
  CO["Coordination<br/>prepare, commit, restore, resume,<br/>retire, release"]
  subgraph G["One participant per Gym server"]
    A["Agent<br/>runs the harness in the lent sandbox"]
    R["SandboxOwnerResourcesServer<br/>SWE-bench, notes server:<br/>seed, snapshot, fork, release,<br/>/sandbox_access, /close_session"]
  end
  P["Shared provider<br/>OpenSandbox: snapshots;<br/>others: restart_only"]
  S[("Sandbox per session<br/>+ session data")]
  I[("Snapshot index<br/>checkpoint.state_dir")]
  C --> CO --> G
  A -- "seed, /sandbox_access, verify, close" --> R
  R --> P --> S
  R --> I
  classDef this fill:#fde68a,stroke:#b45309,stroke-width:2px,color:#1f2937
  class R this

Loading

Which mode a server runs in

flowchart TD
  S["server starts"] --> CK{checkpointing on?}
  CK -- no --> OFF["unchanged: no provider, state_dir<br/>or worker checks at startup"]
  CK -- yes --> L{"lends_sandboxes()?"}
  L -- "no (golden-patch instance)" --> RO["restart_only"]
  L -- yes --> P{"provider can snapshot?"}
  P -- no --> W["restart_only + one warning"]
  P -- yes --> N{"num_workers > 1?"}
  N -- yes --> X["refuse to start"]
  N -- no --> E["exported<br/>(checkpoint.state_dir required)"]

Loading

A session through a checkpoint, a crash, and verify

sequenceDiagram
  participant A as Agent
  participant R as SandboxOwnerResourcesServer
  participant P as Provider
  A->>R: /seed_session
  R->>P: start sandbox S1, run setup (patch, anti-cheat)
  Note over R: tracked only now, with setup's data
  R-->>A: sandbox_handle S1
  Note over R: commit: snapshot S1, index it, export {sandbox, workdir, data}
  Note over R: crash, fresh process, restore
  R->>P: fork the snapshot into S2, stop S1
  Note over R: S2 installed with the workdir and data, all sessions or none
  A->>R: /seed_session (replacement attempt)
  R-->>A: S2 (setup does not run again)
  A->>R: /sandbox_access
  R-->>A: S2 and its workdir
  A->>R: /verify
  Note over R: require S2 (409 without it), extract, end the session
  Note over R: release(kept) later deletes the dropped checkpoints' snapshots

Loading

Where this sits in the stack

Stack #4578, on ananthsub/partial-ckpt-fault-tests:

1. feat(sandbox): sandbox and snapshot lifecycle for partial-rollout checkpoints #4220: the sandbox session checkpointer, release, the snapshot index, the OpenSandbox lease
2. feat(checkpoint): process-level evidence for resources-owned sandbox checkpoints, their release, and lease expiry #4221: process-level evidence for resources-owned sandboxes
3. feat(checkpoint): ask a long-running agent harness to park when a checkpoint opens #4222: the park request for long-running agent harnesses
4. feat(sandbox): sandbox owner base for resources servers, with SWE-bench on it #4576 (this PR): the sandbox owner base, with SWE-bench and the notes server on it
5. feat(checkpoint): harness driver for sandboxed agents that pauses them at checkpoints, with OpenCode on it #4274: the harness driver, with OpenCode on it
6. fix(checkpoint): leave the calls a checkpoint abandoned out of the capture ledger #4451: calls a checkpoint abandoned stay out of the capture ledger

Issue

Part of #3284 (sandboxed blackbox agents).

Validation

Run in the r3Run nemo-rl+Gym container on Slurm against this head, on ananthsub/partial-ckpt-fault-tests 8947bb2ca (job 1289959 suites, job 1289960 e2e):

• tests/unit_tests/test_sandbox_owner.py (new, 15) with a fake snapshot provider and a fake provider without snapshots:
  • export and restore of the sandbox, its working directory and the setup's data; a state without data restores as {}; a restore with one malformed state installs nothing and forks nothing;
  • a re-seed of a restored session reuses the fork and does not run the setup again; end_session_sandbox drops the data;
  • a setup that raises, or returns something other than a JSON object, stops the sandbox and tracks nothing, so the next commit exports nothing;
  • retire and the end of the episode keep the snapshots, and a release deletes them; a /verify without a sandbox answers 409;
  • the empty-body /close_session stops the sandbox, succeeds again when already closed, and the closed session is not exported; a typed body gets the base default, a malformed one a 422;
  • a provider without snapshots runs restart_only: the seed reply carries the restart header, a commit with the session in flight succeeds and writes no record, no state_dir is needed, and the warning is logged once;
  • an exported owner with 2 or 8 workers refuses to start, a restart_only one with 2 starts; an instance that lends no sandbox starts restart_only with 8 workers and builds no provider, checkpointer or state_dir; with checkpointing off nothing is checked at startup.
• resources_servers/swebench/tests 39 passed in a venv with the swebench harness, including a golden-patch instance that starts restart_only with checkpointing on and 8 workers. The existing checkpoint tests changed only where they patch the provider resolution.
• tests/e2e/checkpoint (NEMO_GYM_CHECKPOINT_E2E=1): 61 passed, 12 skipped (real-model and scale cases), including the sandbox scenarios on the notes server, now built on the base, whose verify rewards nothing without the seed token its setup returned.
• tests/unit_tests/test_checkpoint_*.py 783 passed; sandbox 227 passed; Simple Agent 73; the session-state and single-turn examples, rollout collection and ViBench 532.
• ruff check and ruff format --check with CI's ruff 0.9.9 clean on the changed files. Every commit carries Signed-off-by.

Rollout evidence

On the cell-5 OpenSandbox cell, with the OpenCode agent (#4274) and a nano-3.5 SWE SFT policy on 8 SWE-bench Verified rows, on this series' tip dc39546d6 (jobs 1289962, 1289967):

• Checkpoint with stop: SWE-bench snapshotted the 6 sandboxes of the runs still in flight (5 OpenCode runs interrupted, 1 run past OpenCode), the commit took 21.5 s, Gym was killed, and the restore forked them in 10.4 s; every replacement re-seeded into its fork and verified.
• Checkpoint and continue, then a crash: commit 21.3 s, resume, Gym killed with 3 rollouts running, restore 11.5 s; all 3 continued in their forks. A release keeping the checkpoint deleted nothing; the final release deleted all 6 snapshots (404 on the cell) and emptied the index. The lease probe saw the cell delete a sandbox 42 s after its renewal stopped. The agent log has no "Failed to stop sandbox": only the owner stops the agent's sandbox.
• OpenSandbox with its snapshot API hidden, as a provider that cannot snapshot: SWE-bench logged resources: sandbox provider 'opensandbox' cannot snapshot; its rollouts restart after a crash instead of continuing once at startup, a commit with every rollout in flight succeeded with no record, and all 8 verified.
• No sandbox of either run was left on the cell.

Compatibility

• Unchanged with checkpointing off.
• With checkpointing on: a lending server on a provider without snapshots now runs restart_only instead of failing commits; an instance that lends no sandbox runs restart_only; an exported owner refuses num_workers > 1 and requires checkpoint.state_dir.
• New core API: nemo_gym.sandbox.owner (SandboxOwnerResourcesServer, SandboxOwnerConfig, SessionData) and SimpleResourcesServer.checkpoint_mode_in_effect(). SandboxCheckpointConfig moved from resources_servers/swebench/app.py to nemo_gym/sandbox/checkpoint.py.
• SWE-bench's exported session state keeps its shape ({sandbox, workdir}).

🤖 Generated with Claude Code

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.