Skip to content

feat(sandbox): add Sandbox pause and resume - #2636

Merged
ananthsub merged 11 commits into
mainfrom
hemild/feat-sandbox-pause-resume
Sep 21, 2026
Merged

ananthsub merged 11 commits into
mainfrom
hemild/feat-sandbox-pause-resume

Conversation

@hemildesai

@hemildesai hemildesai commented Aug 20, 2026 •

Copy link
Copy Markdown
Contributor

What does this PR do?

  • Adds public async and sync pause() / resume() methods behind an optional sandbox-provider capability.
  • Implements snapshot-backed OpenSandbox pause/resume, including paused status mapping, per-sandbox PTY detachment, and SDK handle refresh after endpoints are rebuilt.
  • Supports descriptor handoff across processes: reconnect accepts a paused sandbox without probing its suspended exec daemon, while resume() keeps health checks enabled and returns an exec-ready sandbox.
  • Adds lifecycle timeout configuration and focused API, provider, and PTY regression coverage.
  • Adds a snapshot management helper to the OpenSandbox provider (python -m nemo_gym.sandbox.providers.opensandbox.snapshots with list / delete / cleanup [--kill-paused] [--dry-run]) to inventory and reclaim snapshots and paused sandboxes left behind by interrupted runs; validated live against an OpenSandbox deployment.

Test plan

  • .venv/bin/pytest -q tests/unit_tests --ignore=tests/unit_tests/test_enroot_provider.py — 2,516 passed, 4 skipped, 15 subtests passed.
  • .venv/bin/pytest -q tests/unit_tests/test_enroot_provider.py -k 'not test_create_start_timeout_cleans_up' — 60 passed, 1 deselected. The deselected macOS /proc-dependent test fails identically on untouched main.
  • .venv/bin/pre-commit run --all-files — all hooks passed.
  • Live OpenSandbox E2E against the supplied deployment in two separate OS processes with default health checks: process A created a sandbox, wrote proof state, paused, serialized its descriptor, and exited; process B connected while paused, resumed, and its first exec immediately read the exact preserved proof before stopping the sandbox.

Checklist

  • I have read the contributing guidelines.
  • The change is focused; unrelated "drive-by" edits are tracked as separate issues/PRs.
  • Tests added or updated and pass locally, or N/A for docs-only / non-code changes (so CI unit/server checks pass when applicable).
  • Pre-commit checks pass locally (pre-commit run --all-files) (so CI lint/format/copyright pass).
  • All commits have DCO sign-off (git commit -s) (so the DCO check passes).

Generated with Claude Code

@copy-pr-bot

copy-pr-bot Bot commented Aug 20, 2026

Copy link
Copy Markdown

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.

@hemildesai hemildesai changed the title feat(sandbox): add OpenSandbox pause and resume feat(sandbox): add Sandbox pause and resume Aug 21, 2026
hemildesai added a commit that referenced this pull request Aug 29, 2026
Wire sandbox pause/resume (#2636) into the SWE-bench eval flow so an
in-flight run can be paused, its launcher killed, and relaunched without
losing agent progress:

- nemo_gym/sandbox/manifest.py: append-only jsonl manifest of per-rollout
  sandbox lifecycle transitions (created/paused/resumed/resume_failed/done);
  the latest record per rollout is its current state.
- swebench resources server: a pause_resume config section. With
  manifest_fpath set, seed_session/verify append created/done transitions;
  with resume=true, seed_session resumes the recorded paused sandbox (fresh
  PTY, conda re-activated, anti-cheating skipped so its git reset --hard
  cannot wipe the restored work) and reports resumed=true. Any resume
  failure appends resume_failed and falls back to a fresh sandbox rather
  than failing the rollout.
- opencode sandboxed agent: on a resumed sandbox, reinstall only if the
  binary is missing, rediscover the previous per-run XDG data dir, and
  continue the restored opencode session via `opencode run --continue`
  (or start fresh when there is no session to continue); the new exec gets
  a full sandbox_timeout budget so paused wall-time costs nothing.
- scripts/sandbox_pause_resume.py: idempotent client-side drivers to pause
  every active sandbox in the manifest (warning when a paused sandbox's ttl
  expires within the pause window) and to delete leftover snapshots and
  kill leftover sandboxes after the run.

Already-graded rollouts are skipped on relaunch by the existing
resume_from_cache support (dedup by task and rollout index against the
materialized inputs), so no new skip logic is needed.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Hemil Desai <hemild@nvidia.com>
@hemildesai
hemildesai force-pushed the hemild/feat-sandbox-pause-resume branch 2 times, most recently from 34bc723 to d15160e Compare September 4, 2026 19:32
@hemildesai
hemildesai marked this pull request as ready for review September 8, 2026 19:07
@github-actions github-actions Bot added the sla:triage-overdue Review assignment is over the one-business-day SLA label Sep 9, 2026
@ananthsub
ananthsub self-requested a review September 10, 2026 00:05
@github-actions github-actions Bot added sla:review-overdue Review response is over the one-business-day SLA and removed sla:triage-overdue Review assignment is over the one-business-day SLA labels Sep 10, 2026
# Servers disagree on state casing, so match paused sandboxes client-side.
for item in await list_all("sandboxes", []):
status = item.get("status")
state = status.get("state") if isinstance(status, dict) else None

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

--snapshot-id X --kill-paused --reap appears scoped to one snapshot, but sandbox_id remains None, so this condition selects every paused sandbox in the deployment. Could we reject --kill-paused with --snapshot-id unless a sandbox scope is available, or otherwise require an explicit global-delete option?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed, that combination was a footgun. --kill-paused is now rejected alongside --snapshot-id with a parser error pointing at --sandbox-id for scoping, and the --kill-paused help text says explicitly that without --sandbox-id it means every paused sandbox. I kept the unscoped form since that is the cleanup use case for interrupted runs, but it is gated behind --reap and the audit output lists each sandbox it would kill first.

Test case added to test_cli_rejects_bad_arguments. cd18548

)

# Pause invalidates its WebSockets. Detach local PTY clients
# without ending the server sessions being suspended.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

“Server sessions being suspended” implies that the remote PTY process survives pause. OpenSandbox only commits the root filesystem, releases the runtime Pod, and recreates a new runtime; process and PTY state are not checkpointed. Setting _owned = False merely suppresses the explicit PTY delete request. Could we correct this comment and test/document that callers must establish a new PTY session after resume?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You are right, the comment overstated what pause preserves. Reworded it to say pause commits the root filesystem and replaces the runtime, so processes and PTY sessions do not survive; we drop the local clients without a delete request because the new runtime never had those sessions. Documented the same on SupportsSandboxPauseResume, AsyncSandbox.pause(), and the provider resume() docstring (open a new PTY after resume), and fixed the assertion message in test_provider_pause_detaches_only_the_target_sandbox_pty_sessions to match. cd18548

@github-actions github-actions Bot removed the sla:review-overdue Review response is over the one-business-day SLA label Sep 11, 2026
@hemildesai

Copy link
Copy Markdown
Contributor Author

/ok to test cd18548

@hemildesai

Copy link
Copy Markdown
Contributor Author

/claude review

@claude

claude Bot commented Sep 11, 2026

Copy link
Copy Markdown
Contributor

SHIP — LGTM, no reliability concerns.

Reviewed the pause/resume lifecycle addition and the snapshots.py cleanup utility. The correctness-sensitive paths hold up:

  • pause() uses the non-retrying _await_sdk_call (correct: pause is mutating and its outcome must not be retried), then polls status() under a single asyncio.timeout, terminating on ERROR/STOPPED. PTY sessions for the target sandbox are detached with _owned=False before close(), so no spurious delete is sent to a runtime that never had them. lifecycle_timeout.expired() correctly distinguishes the outer lifecycle timeout from an inner per-request timeout and re-raises the latter unwrapped.
  • resume() bounds the request + endpoint rebuild + readiness check under one deadline, and on timeout keeps the old handle.raw (state unknown → reconnect-and-check contract), replacing it only after a successful resume. Verified by test_resume_failure_keeps_the_paused_handle / test_resume_timeout_bounds_the_full_lifecycle.
  • connect() now always passes skip_health_check=True to the SDK and gates check_ready() on status() != PAUSED — necessary so a paused sandbox (no live exec daemon) connects without hanging on a dead endpoint. The failure path closes the handle under a bounded wait_for and re-raises the original error.
  • SupportsSandboxPauseResume is a @runtime_checkable Protocol, exported consistently across both __init__.py files and __all__; the facade raises NotImplementedError for providers without it.
  • snapshots.py uses aiohttp (not httpx), bounds delete concurrency with an asyncio.Semaphore, collects all pages before mutating, and sweeps to catch pagination stragglers with a bounded retry cap. aiohttp/PyYAML are already core deps. Secrets are kept out of stdout/stderr (asserted in tests).

Async conventions (aiohttp, asyncio.timeout, semaphore, no ray.get() in async), config placement (pause_resume_timeout_s defaulted in the exemplar YAML with a validator), and test coverage all look right.

bxyu-nvidia
bxyu-nvidia previously approved these changes Sep 13, 2026

@ananthsub ananthsub left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Follow-up re-review found additional lifecycle ownership issues, so I am holding approval until these are addressed.

Comment thread nemo_gym/sandbox/providers/base.py Outdated
class SupportsSandboxPauseResume(Protocol):
"""Optional provider capability for preserving and restoring a sandbox.

Pause preserves the filesystem, not running processes or PTY sessions;

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These are Kubernetes rootfs-snapshot semantics, not a provider-neutral pause contract. The same OpenSandbox provider can target its Docker backend, where container.pause() / container.unpause() freezes and resumes the existing processes. In that case the provider drops the local PTY without deleting or retaining access to the preserved server session, while this guidance tells callers to open another PTY. Please keep the protocol contract backend-neutral and document/handle the runtime-specific PTY and process behavior in the OpenSandbox provider.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fair, that was Kubernetes-specific wording in a neutral contract. SupportsSandboxPauseResume and AsyncSandbox.pause() now only say that PTY sessions are detached and that process survival and re-attachability are provider/backend specific. The OpenSandbox provider documents both backends in pause(): the Kubernetes backend replaces the runtime, so callers open a new PTY after resume; the Docker backend freezes the container, so the detached sessions stay re-attachable by id via attach_pty(). The behavior itself was already right for both (we drop the local client and never send a delete), so this is docs plus the test comment. 701c738

hemildesai and others added 8 commits September 21, 2026 11:40
Signed-off-by: Hemil Desai <hemild@nvidia.com>
Refresh the provider-owned handle in place and remove no-op lifecycle retry
wrappers. Keep tests focused on observable pause/resume behavior.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: Hemil Desai <hemild@nvidia.com>
Apply one lifecycle deadline across the pause request and status polling. Preserve shorter per-request timeouts and add regression coverage.

Signed-off-by: Hemil Desai <hemild@nvidia.com>
Allow paused descriptors to reconnect without probing execd, while preserving readiness checks for running and resumed sandboxes.

Signed-off-by: Hemil Desai <hemild@nvidia.com>
Limit local SDK cleanup after reconnect failures and preserve the original readiness error when cleanup stalls or fails.

Signed-off-by: Hemil Desai <hemild@nvidia.com>
Adds nemo_gym/sandbox/providers/opensandbox/snapshots.py with
list / delete / cleanup subcommands to inventory and reclaim snapshots
(and, via --kill-paused, paused sandboxes) left behind by interrupted
pause/resume runs. Uses the public management API only; idempotent and
safe to re-run; --dry-run previews.

Signed-off-by: Hemil Desai <hemild@nvidia.com>
Rewrite the OpenSandbox snapshot helper in the shape of the run-scoped
sandbox cleanup script so operators learn one interface:

- Talk to the management API directly over aiohttp instead of the SDK
  manager, so the script runs from any Python with aiohttp and PyYAML,
  including as a job submitted by file path.
- Take the connection from --connection-config YAML or --domain/--api-key
  with the same validation and error messages as cleanup_sandboxes.py.
- Audit by default and delete only with --reap, replacing the
  list/delete/cleanup subcommands and --dry-run. --sandbox-id, --state,
  --snapshot-id and --kill-paused keep the previous filtering surface.
- Scope --kill-paused to --sandbox-id when one is given.
- Refuse redirects, treat 404 as already gone, bound delete concurrency,
  and delete snapshots before the paused sandboxes that hold them.

Add unit tests mirroring test_opensandbox_cleanup.py; the script is at
100% line coverage.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Hemil Desai <hemild@nvidia.com>
- Bound the whole resume lifecycle with pause_resume_timeout_s. The SDK
  applies resume_timeout only to the readiness check, so the resume
  request and endpoint rebuild could run past the configured deadline.
  Wrap Sandbox.resume() in asyncio.timeout() like pause() and document
  that a timeout leaves server-side state unknown: reconnect and check
  status() before retrying.
- Correct the pause() PTY comment and docstrings: OpenSandbox commits the
  root filesystem and replaces the runtime, so processes and PTY sessions
  do not survive pause. Callers open a new PTY after resume.
- snapshots.py: reject --kill-paused with --snapshot-id. Snapshot ids
  carry no sandbox scope, so the combination would select every paused
  sandbox in the deployment.
- snapshots.py: numbered pages shift while another actor deletes, so an
  item can be skipped and the run still succeed. Re-list and sweep up to
  REAP_SWEEPS times until a fresh listing is empty or a sweep stops
  progressing, matching cleanup_sandboxes.py.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Hemil Desai <hemild@nvidia.com>
@claude

claude Bot commented Sep 21, 2026

Copy link
Copy Markdown
Contributor

SHIP — sandbox pause/resume + snapshot cleanup. No blocking correctness or async concerns.

Reviewed the full diff against base and the head-version provider.py. Notable things I checked and found sound:

  • connect() rework (provider.py:1049): now always calls the SDK with skip_health_check=True, then does its own status() + check_ready() unless the sandbox is PAUSED. The precedence not skip_health_check and status != PAUSED is correct, and both failure branches (sandbox is None → close transport; else best-effort sandbox.close() under a bounded timeout) preserve the original exception. Tests updated consistently (test_connect_*).
  • pause() (provider.py:1095): correctly not retried (outcome-unknown mutating call), PTY sessions detached with _owned = False so no server-side DELETE fires (matches the K8s-vs-Docker backend note), then polls status() to PAUSED with ERROR/STOPPED treated as terminal. The lifecycle_timeout.expired() guard properly distinguishes the outer deadline from an inner per-request TimeoutError.
  • resume() (provider.py:1143): installs the new handle before the best-effort close of the old one, so a cancellation during close can't drop a successful resume; owned transport is closed on any BaseException. Good.
  • _sandbox_id is stamped on sessions before they enter _pty_sessions (provider.py create/attach paths), so pause()'s per-sandbox filter is reliable.
  • snapshots.py: standalone maintenance CLI (documented as such — runs from any Python with aiohttp/PyYAML), so the "route async HTTP through the global aiohttp client" rule for server paths doesn't bind it; it uses a bounded aiohttp.ClientSession + semaphore, collect-then-delete pagination, bounded re-list sweeps against page-shift races, and redacts the API key from output (verified by tests).
  • pause_resume_timeout_s default lives on the @dataclass OpenSandboxOperationConfig, which already carries Python defaults for every field, and is mirrored in the exemplar opensandbox.yaml — consistent with the existing pattern, and validated (> 0).

NOTE (author's call, not blocking): connect()/resume() call .with_transport_if_missing() / .close_transport_if_owned() on the real SDK ConnectionConfig, and test_resume_timeout_closes_an_sdk_owned_transport reaches into SDK internals (transport, _owns_transport). These couple the runtime paths to SDK behavior only pin-guarded by opensandbox>=0.1.15; if a future SDK renames those, the failure surfaces as an AttributeError at connect/resume time rather than at import. Worth a tighter upper pin or a capability check if the SDK's transport API isn't yet stable. (Could not execute the SDK-backed tests here — the opensandbox package isn't installed in this review environment; the pure-Python snapshots and facade tests are self-contained and read as correct.)

…freeze

- snapshots.py: cleanup_snapshots() now rejects blank selectors itself
  and strips the rest, so a direct caller passing sandbox_id="" can no
  longer fall through to an unfiltered listing and delete every snapshot.
  Explicit-id mode is keyed on `snapshot_ids is not None`: an empty list
  names nothing rather than meaning "list everything".
- pause(): detach the local PTY clients before the pause request, while
  execd can still answer the WebSocket close handshake, and bound each
  detach by close_timeout_s outside the lifecycle deadline. On the Docker
  backend the freeze stops execd, so a close issued afterwards could wait
  on a handshake that never comes and charge it to pause_resume_timeout_s,
  reporting a timeout for a sandbox that was already paused. Server
  sessions are still never deleted: if the pause request fails the
  sandbox keeps running and they remain re-attachable by id.

Tests: blank and unstripped selectors, empty explicit id list, detach
ordering, and a hung detach that no longer blocks a completed pause.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Hemil Desai <hemild@nvidia.com>
@hemildesai

Copy link
Copy Markdown
Contributor Author

/claude review

@hemildesai

Copy link
Copy Markdown
Contributor Author

/ok to test 5ee9daa

- connect(): reuse _await_sdk_call with close_timeout_s for the failed
  connect cleanup instead of a hand-rolled wait_for with min() arithmetic,
  matching how resume() closes its old handle.
- Trim docstrings and comments across the PR to the "why" a reader needs;
  drop the duplicated kill-paused rationale from the CLI.
- Tests: reset all FakeSandbox class state in the fixture instead of in
  one test; shorten test comments.

No behavior change beyond the connect cleanup bound now being
close_timeout_s rather than min(connect timeout, close timeout).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Hemil Desai <hemild@nvidia.com>
@hemildesai

Copy link
Copy Markdown
Contributor Author

/ok to test a024393

@claude

claude Bot commented Sep 21, 2026

Copy link
Copy Markdown
Contributor

Review incomplete: the PR head advanced during review (analysis was pinned to 5ee9daafcb400340709489fba2f1de9c9d0aa1c9, but the final revision check found head at a0243938e3906e5f8b34e88dea2936db615656a7). The captured diff and changed-file list no longer match the current head, so I'm withholding findings rather than posting comments against stale code. Re-run the review against the new head to get a fresh result.

@ananthsub
ananthsub added this pull request to the merge queue Sep 21, 2026
@github-actions

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA-NeMo/Gym/actions/runs/35667634095

Merged via the queue into main with commit fac62a4 Sep 21, 2026
41 checks passed
@ananthsub
ananthsub deleted the hemild/feat-sandbox-pause-resume branch September 21, 2026 23:40
ananthsub pushed a commit to ananthsub/Gym that referenced this pull request Sep 23, 2026
…ession (NVIDIA-NeMo#3641)

## What changed and why

Follow-up to NVIDIA-NeMo#2636 (sandbox pause/resume). On cell-3, re-attaching to a
PTY session opened before a pause took about **95s** to fail, and the
error was misleading:

```
SandboxPtyError: PTY session already has an attached client
```

On the Kubernetes backend, pause commits the rootfs and replaces the
runtime, so the old session no longer exists. But execd answers a
WebSocket attach to a **missing** session with the same `1008` close it
uses for a session **held by another client**. I checked this live: a
random session id, a random id with `takeover=1`, and a genuinely held
session all close `1008` with an empty reason within 0.3s. So Gym
treated it as a stale holder and waited out every retry: the session's
own takeover retries (2+5+10s), the provider's retries, and a final
sandbox-death check.

execd's HTTP API does tell the two apart, which the same live probe
showed:

| Request | Answer |
|---|---|
| `GET /pty/<missing>` | `404 {"code":"CONTEXT_NOT_FOUND","message":"pty
session … not found"}` |
| `GET /pty/<live>` | `200 {"session_id":…,"running":true,…}` |

`OpenSandboxProvider.attach_pty()` now makes that `GET` before dialing.
Only execd's own 404, with the `CONTEXT_NOT_FOUND` code, counts as
missing. A proxy 404 while a route is still registering, or a failed
check, lets the attach proceed exactly as before. A missing session now
fails immediately with `PTY session <id> not found`, with or without
`takeover`. Before, without `takeover`, a missing session was also
misreported as "already has an attached client".

Cost: one extra small GET per attach. Attaches are reconnect operations,
not per-command.

No issue filed. This was found while validating NVIDIA-NeMo#2636 end to end.

## Validation

**Unit tests** (`test_opensandbox_pty.py`,
`test_opensandbox_provider.py`, `test_sandbox.py`): 246 passed. New
cases cover:
- A missing session fails with "not found", makes no WebSocket dial, and
never consults sandbox status. Covered for both `takeover` values.
- A proxy 404 without execd's code, and a failed check, both keep the
existing takeover retries.

Two existing tests now count the existence-check client. `pre-commit
run` on the changed files passes.

**Live on cell-3** (tenant namespace, Kubernetes backend), running the
pause/resume e2e across two OS processes:

| Step | Result |
|---|---|
| pause, then reconnect from a new process | paused |
| resume | running |
| file written before pause | preserved |
| non-PTY `exec`, foreground path | rc 3; stdout and stderr correct |
| non-PTY `exec`, background path (what `opensandbox.yaml` enables) | rc
3; combined stream in stdout, as documented |
| attach to the pre-pause PTY session | **fails in 0.3s** with `not
found` (was about 95s with "already has an attached client") |
| fresh PTY after resume | works |

A separate probe confirmed the normal case is unaffected. Detaching from
a live session and re-attaching by id succeeds with and without
`takeover`, in about 0.6s, with shell state intact. All sandboxes were
deleted afterwards.

Rollout evidence: N/A. This is a sandbox provider change with no
environment, agent, or benchmark behavior change.

## Noticed, not changed here

The foreground `exec` path appends the SDK's error to stderr, so a
command that exits non-zero gets a synthetic `CommandExecError: <code>`
line after its own stderr. This behavior dates from the original sandbox
API, and it's out of scope here.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Signed-off-by: Hemil Desai <hemild@nvidia.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

This branch was previously deployed

1 inactive deployment
public — a0243938 Deployed Sep 21, 2026 by copy-pr-bot[bot] via release / finalize / notify #3398
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants