fix(sandboxed-gym): apply independent upstream fixes from the RL fork - #1734
Open
SandyChapman wants to merge 3 commits into
Open
fix(sandboxed-gym): apply independent upstream fixes from the RL fork#1734SandyChapman wants to merge 3 commits into
SandyChapman wants to merge 3 commits into
Conversation
Ports soluwalana/RL@6e268a2 ("duplicate sandboxes + orphan sandbox when training pod is killed"), which this package's vendored copy predates. Both halves of that fix are absent here. The actor carried max_restarts=-1, max_task_retries=-1. Its host handle lives only in the actor process, so a restarted actor cannot name the sandbox its predecessor created: __init__ provisions a second one and the first survives to its ttl_s, silently doubling the pods a job holds. A crash should fail the job instead. The sibling broker actor already documents exactly this reasoning for why it must not be restartable, so the two files contradicted each other. Ray also tears an actor's worker down without running user teardown, so a cancelled, evicted or preempted job leaked its sandbox until ttl_s. The new install_termination_cleanup registers atexit plus SIGTERM/SIGINT handlers that destroy the host and re-raise, keeping the signal in the exit status so a cancelled job does not read as a clean stop. It lives in the Ray-free orchestrator module and is called from the actor, rather than sitting on the actor as upstream has it, for two reasons: this package deliberately keeps the actor a thin wrapper, and ray is an optional extra, so logic placed there cannot be tested from a plain checkout. Deviates from upstream in restoring the default signal action *before* running the cleanup rather than after. Upstream leaves its handler installed for the whole of shutdown -- the slowest part, since it waits on the sandbox being destroyed -- so a second SIGTERM re-enters and stacks another destroy on the in-flight one. Restoring first means a second signal terminates instead, which is both safer and the conventional response to a repeated signal. The cleanup is wrapped in try/finally so a failed destroy still re-raises. Signals are installed on a best-effort basis: only the main thread may install handlers and Ray does not promise to run an actor method there, so a ValueError degrades to atexit rather than failing spinup. Signed-off-by: Sandy Chapman <schapman@nvidia.com>
Ports the diagnostics half of soluwalana/RL@d90a51e ("add logging around health checks and rollouts"). wait_ready reported only the host id on timeout, and the resolved health and rollout URLs were never logged anywhere. Those URLs come back from the SDK's route resolution onto a handle the caller may never print, so they cannot be reconstructed from outside the process. A protocol or port mismatch therefore presents identically to a slow sandbox -- a bare timeout against an address that appears in no log -- and sends the reader to the wrong problem. Worth more here than it looks: this is the OpenSandbox path, which has never been exercised end to end, so its first failures will be diagnosed from logs alone. The tests need no OpenSandbox SDK. The module keeps its SDK types under TYPE_CHECKING and builds the driver lazily, so the provider imports and constructs from a plain checkout; the two existing provider tests are guarded by a skipif that is stricter than the module actually requires. Signed-off-by: Sandy Chapman <schapman@nvidia.com>
Ports soluwalana/RL@4968fab and @874f947 ("make gym install deps in the venv created"), which set uv_pip_set_python. The key appeared nowhere in this repo. Gym's `uv pip install` names no target, so it resolves one from the environment. An absolute UV_PYTHON outranks the venv Gym has just activated, and the install lands in the read-only interpreter tree; every Gym server then dies with "Permission denied ... site-packages", which surfaces to the operator only as "Process `policy_model` finished unexpectedly!". Setting the key makes Gym pass `--python <venv>/bin/python` explicitly. Conditional today, load-bearing shortly. docker/rl/Dockerfile.nmp-rl-base is the only image here setting an absolute UV_PYTHON (/opt/cpython/bin/python3.13, deliberately, so RL's checked-in .python-version cannot pin an unpatched interpreter). The gym-host and gym-tasks images do not, which is why the live Docker validation of #1400 could not have caught this -- and why it becomes live as soon as the library is shared with the RL actor. Verified honoured rather than assumed: uv_pip_set_python has been an upstream NeMo-Gym global-config key since NVIDIA-NeMo/Gym@146b1a5b5 (2025-12-19), read in nemo_gym/cli/setup_command.py with a default of False, and both images pin nemo-gym==0.5.0. A setdefault, so a caller with its own answer keeps it. Signed-off-by: Sandy Chapman <schapman@nvidia.com>
Contributor
|
SandyChapman
marked this pull request as ready for review
September 3, 2026 16:42
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (7)
Included review availability: Your plan provides up to 12 included reviews per hour; 9 remain after this review. 📝 WalkthroughWalkthroughChangesSandbox lifecycle and host operation
Suggested reviewers: Merge Risk: ⚪ Minimal · up to The sandbox lifecycle, runtime defaults, and readiness diagnostics changes are mergeable with no actionable risk identified. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
15 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
packages/sandboxed_gymwas vendored fromsoluwalana/RL@67821a9(2026-08-03) and has taken no upstream changes since; fourteen commits have landed there in the meantime. This applies the four that are independent of each other and of the rollout transport path — an actor that leaks sandboxes on restart and on termination, a readiness timeout that cannot be diagnosed, and a Gym config default without which every Gym server dies during startup on the training image.The remaining upstream work (heartbeat, chunking, proxy-error classification) rewrites the rollout path and is deliberately left for a second PR; see AALGO-584.
Related Issue
Part of AALGO-584 (drift analysis and full item list attached there). Parent: AALGO-583.
Changes
Three commits, each naming the upstream commit it ports:
stop the Gym actor leaking sandboxes— portssoluwalana/RL@6e268a2. Dropsmax_restarts=-1, max_task_retries=-1fromSandboxedGymActor: the host handle lives only in the actor process, so a restarted actor provisions a second sandbox and leaks the first until itsttl_s. The siblingSandboxEpisodeBrokerActoralready documented this reasoning, so the two files contradicted each other. Addsinstall_termination_cleanup(atexit + SIGTERM/SIGINT) so a cancelled, evicted or preempted job destroys its host instead of leaking it.name the health URL a job host failed to reach— ports the diagnostics half ofsoluwalana/RL@d90a51e. Logs the resolved health/rollout URLs and names the polled URL in the readinessTimeoutError.make Gym install into the venv it activated— portssoluwalana/RL@4968fab+@874f947. Setsuv_pip_set_python, absent from this repo entirely.Design calls worth a reviewer's attention
install_termination_cleanuplives in the Ray-freeorchestratormodule, called from the actor, rather than on the actor as upstream has it. This package keeps the actor a thin wrapper, andrayis an optional extra, so logic placed there cannot be tested from a plain checkout.shutdown()— the slowest part, since it waits on the sandbox being destroyed — so a second SIGTERM re-enters and stacks another destroy on the in-flight one. Restoring first means a second signal terminates instead. The cleanup is wrapped intry/finallyso a failed destroy still re-raises. An independent review flagged the upstream ordering as amajordefect, separately from this change.plugins/nemo_evaluator/jobs/gym_sandbox.pyembedsSandboxedGymOrchestratordirectly rather than going through the Ray actor, and a library has no business replacing an embedding host's signal handlers. Itstry/finallycovers the ordinary path; if signal coverage is wanted there it belongs in Evaluator's task runner. Thesandboxed-gym serveCLI is likewise untouched — a reasonable follow-up, but out of scope here.uv_pip_set_pythonis conditional today and load-bearing shortly.docker/rl/Dockerfile.nmp-rl-base:102is the only image here setting an absoluteUV_PYTHON; the gym-host and gym-tasks images do not. That is why the live Docker validation on feat(evaluator): run Gym evaluations inside a sandboxed Gym host #1400 could not have caught this, and why it goes live the moment the library is shared with the RL actor.Deliberately not changed
asyncio.as_completedmispairing independently and incompatibly (upstream tags[_rowidx, result]; we copy_ng_task_indexonto the result). Reconciling them is a wire-format decision for the shared wheel, not a drift item.gym_env_package.py), which reintroduces the NeMo-RL dependency this package exists to avoid and overlaps feat(evaluator): Install wheels-v1 Gym environments from FileSets #1522.requires_opensandboxskipif on two existing provider tests is stricter than the module needs — the provider imports and constructs without the SDK. Left alone as unrelated.Type of Change
Quality Gates
Verification
Signed-off-by:traileruv run pre-commit run -apasses, or any blocked checks are identified belowTargeted validation:
uv run --frozen pytest packages/sandboxed_gym/tests/ -q→ 209 passed, 2 skipped (2 skips are the pre-existingrequires_opensandboxguards). Baseline onmainis 196.uv run ruff check packages/sandboxed_gym/→ All checks passed.ruff format --check→ 51 files already formatted.uv run --frozen ty check→ exit 0. The one diagnostic on changed files is the pre-existing unresolvedimport ray(optional extra, excluded from the repo-wide run).uv run pre-commit run -a→ exit 0, no failures, run insideflox activatesouvis the pinned 0.9.14. Outside Flox theuv-lockhook fails on the local uv being 0.9.30; that is a toolchain-version check only — this branch touches nopyproject.tomloruv.lock, andCheck for uv.lock driftpasses either way.uv_pip_set_pythonverified honoured, not assumed: an upstream NeMo-Gym global-config key sinceNVIDIA-NeMo/Gym@146b1a5b5(2025-12-19), read innemo_gym/cli/setup_command.pywith a default ofFalse; both images pinnemo-gym==0.5.0.Not verified
Nothing here has been exercised against a real OpenSandbox deployment — the same gap #1400 shipped with. The termination and readiness paths in particular are the ones a live cluster would exercise differently, and the readiness-diagnostics change exists precisely because that path has never run end to end.
Summary by CodeRabbit
Reliability
Diagnostics
Tests