feat: disaggregated sandbox backends on OpenSandbox (ns_tools + math_formal_lean)

[hemildesai]

What this adds

Two opt-in sandbox backends that move code execution off the resources server's own host and onto OpenSandbox pods, plus the provider method they both need.

ns_tools — sandbox_type: sandbox_pool

A fixed set of long-lived, shared pods, each running the NeMo-Skills sandbox HTTP protocol, with sessions multiplexed across them by sticky routing. Sharing is what makes large batches feasible: many concurrent sessions ride K pods instead of one pod per session.

• Sticky session routing. A new session pins to the least-loaded healthy pod and stays there, so stateful ipython state survives across tool calls.
• Prewarmed-pool claims with fallback. When pool_ref is set, a slot is filled by claiming a prewarmed pod from a server-side Pool via extensions.poolRef, which drops warmup to allocation time. A full or unavailable pool degrades to a direct create by default (pool_fallback: true), or fails the slot when that is turned off. Fallback pods go through the normal prepare step, so pool configs should still carry their setup/service settings.
• Health and healing. A pod is evicted after three consecutive failed probes and healed in the same slot, with concurrent heals spaced by a create-rate limit so a mass heal cannot storm the create path. An idle sweep drops stale session pins.
• Per-run attribution labels on every pod, so an epilogue reaper can delete exactly one run's sandboxes.
• Transport. Rides a shared aiohttp session, per this repo's guidance that httpx/httpcore connection pooling collapses at high concurrency, while keeping httpx exception types so the NeMo-Skills client contract is unchanged. Infra failures normalize into that contract and degrade rewards rather than crashing the server.

math_formal_lean — sandbox_backend: gym_sandbox

Lean4 compilation on OpenSandbox pods via provider exec, reproducing the NS server's lake env ... lean invocation and its process_status/stdout/stderr contract exactly.

• pool_size: 0 — a fresh pod per verify, created under a bounded semaphore and destroyed in finally.
• pool_size: N — a warm pool built at server startup and reused across verifies. A cold pod's first import Mathlib lazy-pulls several GB of olean files one page fault at a time (measured at ~900s), while a warmed pod compiles in ~4s, so pool pods bulk-prefetch that tree once at prepare and then serve verifies back-to-back.
• Failed pods (including TTL expiry) are killed and replaced in place; the verify retries once on a second pod before degrading. Infra failures never raise into verify().
• A third value, ns_http_proxy, speaks the NS protocol through a full base_url plus headers. It is useful as a parity oracle against the default path.

nemo_gym/sandbox — OpenSandbox provider endpoint()

Implements the existing SupportsSandboxEndpoint protocol for the OpenSandbox provider (Docker already had it). It resolves the SDK's server-proxy route for a declared port, absolutizes a scheme-less URL using the configured domain, and carries the auth header the proxy requires. The key is attached only in server-proxy mode — a direct endpoint terminates at the sandbox itself, which runs untrusted code and must never be handed the credential (the same scoping rule as #2462). This is the only change to nemo_gym/sandbox/ — everything else the backends rely on was already on main.

Default behavior is unchanged

Both backends are opt-in and off by default. With no environment variables set, the resolved configs are byte-identical to the current ones, and neither backend module is imported or constructed:

• ns_tools defaults to sandbox_type: local
• math_formal_lean defaults to sandbox_backend: ns_http

The opensandbox SDK is declared in each server's requirements.txt; the root pyproject.toml already pinned it.

Config surface

Everything deployment-specific is env-fed with an empty default, so nothing about any particular deployment is baked into the repo:

Setting	Env var
Service domain / API key	OPENSANDBOX_BASE_URL, OPENSANDBOX_API_KEY
Sandbox image	NS_SANDBOX_IMAGE
Server-side pool name / fallback	NS_SANDBOX_POOL_REF, NS_SANDBOX_POOL_FALLBACK
Pool size / TTL	NS_SANDBOX_POOL_SIZE, NS_SANDBOX_TTL_S
Backend selection	NS_TOOLS_SANDBOX_TYPE, MATH_FORMAL_LEAN_BACKEND
Lean concurrency / pool size	LEAN_SANDBOX_MAX_CONCURRENT, LEAN_SANDBOX_POOL_SIZE

Values that must not arrive as strings go through oc.decode or explicit int()/float() coercion at the constructor, since Hydra hands env-fed values over as strings and the create API rejects non-string env maps. Selecting either backend with an empty domain, API key, or image is a hard startup error, never a silent no-op.

Sizing and prewarming notes

• Pool size is a capacity decision, not a concurrency limit: each pod's NS server multiplexes many sessions, so K is chosen from pod memory/CPU headroom rather than from request count.
• request_timeout for creates is raised well above the SDK default, because the first create against a new image tag blocks on image conversion and the default fails the whole cold warmup wave.
• For Lean, prefer pool_ref or a warm pool_size in any latency-sensitive setting: the first compile on a genuinely cold pod is dominated by lazy olean pulls, far beyond any verify's admission window.
• Warmup is non-blocking. start() kicks pod creation from server lifespan startup and returns immediately, so server boot is never gated on the pool being full.

Test coverage

Suite	Count
resources_servers/ns_tools/tests/test_sandbox_pool.py	21
resources_servers/math_formal_lean/tests/test_sandbox_backends.py	14
tests/unit_tests/test_opensandbox_endpoint.py	9

All pass, along with the pre-existing suites for the touched servers (41 total in ns_tools, 56 in math_formal_lean) and tests/unit_tests -k sandbox (301 passed). No network is required: routing, eviction, and pool-claim logic are driven directly, and transports are faked. The sandbox_pool tests that need nemo_skills skip gracefully when it is not installed.

🤖 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.

[hemildesai]

/claude review

[claude]

SHIP WITH CARE — solid, heavily-tested disaggregated-sandbox backend. Scoring/verify contract preserved; async cancellation handling is careful and well-covered. One correctness RISK worth a look before high-concurrency training use.

Correctness / scoring (highest scrutiny) — clean

• GymSandboxLean4Client._map_result faithfully reproduces the NS process_status contract: rc 0 → completed, rc 124/137/-9 → timeout (partial stdout preserved), else → failed; truncation matches max_output_characters. determine_proof_status consumes it identically to the HTTP backend, so rewards are unchanged when the backend is swapped.
• Infra failures degrade to error/timeout inside the client and never raise into verify() — matches the "graceful error handling" convention.
• Both new backends are strictly opt-in (sandbox_backend/sandbox_type default to the existing local path) and hard-fail at startup on empty creds/image rather than silently no-op'ing.

Async correctness — clean

• No ray.get() in async paths; no missing await. await_cleanup + asyncio.shield correctly finishes owned teardown before propagating caller cancellation, and the BaseException widening in provider.close()/_create_once closes the right resources on cancellation. Cancellation paths are directly tested (test_*_cancelled_*).
• httpx.AsyncClient is used only in the pre-existing local Lean4SandboxClient (unchanged, single-host localhost); the new pool/transport code uses aiohttp. See note below on the global-client convention.

RISK (inline)

• gym_sandbox.py 502 retry may double-execute stateful ipython sessions — the "pod never received the request" assumption isn't a guarantee gateways make. Details inline.

NOTE (non-blocking)

• CLAUDE.md says all async HTTP should go through the global aiohttp client (nemo_gym.server_utils.request()); sandbox_pool.py and gym_sandbox.py instead build their own aiohttp.ClientSessions. Given these are long-lived pods (K connections, not the 16k-fan-out case the convention targets) and need per-request pod routing/headers that request() doesn't model, a dedicated session is defensible — but it's a deviation from the documented pattern; confirm it's deliberate.
• opensandbox==0.1.15 / tenacity are correctly declared only in the two per-server requirements.txt and imported lazily behind the opt-in branch, keeping the core library importable without them. Good.

[hemildesai]

/claude review

[hemildesai]

/ok to test 6bf85fe81b

[copy-pr-bot]

/ok to test 6bf85fe81b

@hemildesai, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

[hemildesai]

/claude review

[claude]

SHIP WITH CARE — solid, opt-in, and unusually well-tested. The default paths are unchanged and the new provider-backed backends are gated behind config, so blast radius for existing users is zero unless they set sandbox_backend=gym_sandbox / sandbox_type=sandbox_pool.

What I verified holds up:

• Scorer contract preserved. GymSandboxLean4Client._map_result maps rc→{completed,timeout,failed,error} and preserves partial stdout + the Execution timed out after N seconds suffix, so determine_proof_status/format_error_feedback and the \bsorry\b scan behave identically to the NS HTTP backend. Infra failures degrade to an error dict and never raise into verify(). No silent score corruption.
• Async hygiene. New HTTP is aiohttp with a long-lived, bounded-connector session (not httpx.AsyncClient) — consistent with the convention. httpx is imported for exception types only (documented). No ray.get(); no missing awaits spotted.
• Cancellation/cleanup. await_cleanup + _stop_task/_close_task shielding is a clean pattern and is exercised directly (cancelled admission, cancelled close, dead-pod replace, pool heal rotation). The BaseException widening in OpenSandboxProvider.close/_create_once correctly ensures a cancel between create and verify can't strand a remote sandbox.
• Security. The management API key is injected into headers only in use_server_proxy mode and popped for direct endpoints (endpoint()), so untrusted in-sandbox code never sees it. Covered by test_direct_mode_endpoint_does_not_inject_the_key and friends.
• Deps declared (opensandbox, tenacity) in both per-server requirements.txt and root pyproject.toml. Config defaults live in YAML; new TypedDict/config fields default to empty and hard-fail at startup when the backend is selected without creds/image.

NOTE (please confirm, couldn't verify here): GymSandbox._send_request does request.pop("session_id", None), mutating the caller's dict. This is safe only if the nemo_skills LocalSandbox retry layer rebuilds the request rather than re-calling _send_request with the same object — a replayed call would lose sticky-session routing. nemo_skills isn't installed in this review env so I couldn't check the base directly; the presence of test_502_does_not_replay_stateful_code suggests you designed against replay, so this is likely fine — just worth a one-line confirmation.

No BLOCKER or RISK-level findings. Operability caveat is the usual one: the pool/heal/routing machinery is unit-tested but not exercised against a live sandbox service, so watch the first real rollout at concurrency.

[hemildesai]

Confirmed against NeMo-Skills da85a881: each execution and restoration path builds a fresh request dict immediately before a single _send_request() call, and transport failures do not retry that object. The base Sandbox._send_request() also pops session_id. So the pinned contract cannot replay a mutated request or lose sticky routing; no code change is needed.

[hemildesai]

/claude review

[hemildesai]

/ok to test ace4c1b

[claude]

SHIP WITH CARE — two opt-in OpenSandbox backends (Lean per-verify/pooled compilation; ns_tools disaggregated pod pool) plus core AsyncSandbox/provider cancellation-safety hardening. Default paths (math_formal_lean → ns_http, ns_tools → local) are untouched, so blast radius is confined to users who set the new env vars. Config follows the YAML-as-source-of-truth convention (TypedDict-style knobs live in the exemplar YAML, empty creds are a hard startup error not a silent no-op). Cancellation/teardown correctness is unusually well covered — await_cleanup shielding, cancelled-lease replacement, cancelled-close-finishes-cleanup, and BaseException propagation in the provider are all exercised. httpx is confined to exception-type imports; the actual transport correctly rides nemo_gym.sandbox + aiohttp per the async HTTP rule.

One RISK posted inline (gym_sandbox._send_request mutates the caller's request dict via pop, which can drop the session pin on a normalized-timeout retry and silently corrupt stateful-tool rewards). Non-blocking pending confirmation of the parent NS client's retry behavior — the non-mutating rewrite is free insurance either way.

Notes (author's call, not blocking):

• sandbox_pool.py config coercion is inconsistent: pool_fallback gets explicit string→bool handling (correct — bool("false") is truthy) but size/ttl_s/port rely on int()/float() of env strings (fine) while warmup_fill_concurrency, health_interval_s, etc. are typed as numbers with no coercion. Since these come from YAML ${oc.decode} they're fine today, but a future plain ${oc.env} wiring would pass a str through to arithmetic. Low likelihood given current configs.
• Two requirements.txt add tenacity>=9.1.4 but I found no tenacity import in either server's new code — appears to be a transitive dep of the opensandbox SDK pulled in explicitly. If it's not directly used, it's harmless but unnecessary.

Verifier/scorer logic (verify(), _map_result status mapping, reward aggregation) is unchanged in intent — _map_result faithfully preserves the NS timeout/partial-stdout/truncation contract, and the tests assert that byte-for-byte. No async-hang or broken-public-API concerns in the core sandbox changes.

[init-nikhil]

is this still needed given the this is being set on line 663 now?

[init-nikhil]

nice handling.

[init-nikhil]

Could we rename this to await_uninterruptibly, since it applies to any task and defers cancellation until the task finishes?

[init-nikhil]

Could we move this lifecycle wiring into SimpleResourcesServer as opt-in behavior, so resource servers needing sandbox warmup don’t have to implement it themselves?

[init-nikhil]

Could you post the link to your run? (for the run numbers in https://terryk.gitlab-master-pages.nvidia.com/nemo-html/hemild/nemo-sandbox/rlvr-disagg-sandbox-design.html)