Skip to content

fix(sandboxed-gym): apply independent upstream fixes from the RL fork - #1734

Open
SandyChapman wants to merge 3 commits into
mainfrom
aalgo-584-sandboxed-gym-upstream-fixes/schapman
Open

fix(sandboxed-gym): apply independent upstream fixes from the RL fork#1734
SandyChapman wants to merge 3 commits into
mainfrom
aalgo-584-sandboxed-gym-upstream-fixes/schapman

Conversation

@SandyChapman

@SandyChapman SandyChapman commented Sep 2, 2026

Copy link
Copy Markdown
Contributor

Summary

packages/sandboxed_gym was vendored from soluwalana/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 — ports soluwalana/RL@6e268a2. Drops max_restarts=-1, max_task_retries=-1 from SandboxedGymActor: the host handle lives only in the actor process, so a restarted actor provisions a second sandbox and leaks the first until its ttl_s. The sibling SandboxEpisodeBrokerActor already documented this reasoning, so the two files contradicted each other. Adds install_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 of soluwalana/RL@d90a51e. Logs the resolved health/rollout URLs and names the polled URL in the readiness TimeoutError.
  • make Gym install into the venv it activated — ports soluwalana/RL@4968fab + @874f947. Sets uv_pip_set_python, absent from this repo entirely.

Design calls worth a reviewer's attention

  • install_termination_cleanup lives in the Ray-free orchestrator module, called from the actor, rather than on the actor as upstream has it. This package keeps the actor a thin wrapper, and ray is an optional extra, so logic placed there cannot be tested from a plain checkout.
  • It restores the default signal action before running cleanup, deviating from upstream, which restores 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. The cleanup is wrapped in try/finally so a failed destroy still re-raises. An independent review flagged the upstream ordering as a major defect, separately from this change.
  • Not applied to the Evaluator's in-process path. plugins/nemo_evaluator/jobs/gym_sandbox.py embeds SandboxedGymOrchestrator directly rather than going through the Ray actor, and a library has no business replacing an embedding host's signal handlers. Its try/finally covers the ordinary path; if signal coverage is wanted there it belongs in Evaluator's task runner. The sandboxed-gym serve CLI is likewise untouched — a reasonable follow-up, but out of scope here.
  • uv_pip_set_python is conditional today and load-bearing shortly. docker/rl/Dockerfile.nmp-rl-base:102 is the only image here setting an absolute UV_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

  • The result-ordering fix. Both sides fixed the same asyncio.as_completed mispairing independently and incompatibly (upstream tags [_rowidx, result]; we copy _ng_task_index onto the result). Reconciling them is a wire-format decision for the shared wheel, not a drift item.
  • The environment-wheel feature (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.
  • The requires_opensandbox skipif 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

  • Code change (feature, bug fix, or refactor)
  • Code change with documentation updates
  • Documentation only
  • Contributor tooling or automation
  • CI, build, or test infrastructure

Quality Gates

  • Tests added or updated for changed behavior
  • Existing tests cover changed behavior — justification:
  • Tests not applicable — justification:
  • Documentation updated for user-visible behavior
  • Documentation not applicable — justification: no user-visible surface changes. All four are internal defaults and diagnostics; the package README describes architecture, which is unchanged.

Verification

  • Pull request title follows the repository's Conventional Commit format
  • Every commit includes an appropriate Signed-off-by: trailer
  • uv run pre-commit run -a passes, or any blocked checks are identified below
  • Targeted tests pass, or tests are marked not applicable above
  • No secrets, API keys, or credentials are included

Targeted validation:

  • uv run --frozen pytest packages/sandboxed_gym/tests/ -q209 passed, 2 skipped (2 skips are the pre-existing requires_opensandbox guards). Baseline on main is 196.
  • Each commit is independently green — 205 → 207 → 209 in a scratch worktree, so the history bisects.
  • 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 unresolved import ray (optional extra, excluded from the repo-wide run).
  • uv run pre-commit run -aexit 0, no failures, run inside flox activate so uv is the pinned 0.9.14. Outside Flox the uv-lock hook fails on the local uv being 0.9.30; that is a toolchain-version check only — this branch touches no pyproject.toml or uv.lock, and Check for uv.lock drift passes either way.
  • Every new test was mutation-checked: the production change was reverted and the corresponding test confirmed to go red, then restored.
  • uv_pip_set_python verified honoured, not assumed: 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; both images pin nemo-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

    • Sandbox hosts now clean up reliably when processes exit or receive termination signals.
    • Failed sandbox jobs no longer automatically restart after actor crashes.
    • Runtime setup now consistently configures the Python executable for package installation.
  • Diagnostics

    • Host startup logs resolved health and rollout URLs.
    • Readiness timeout errors now identify the health URL and underlying connection issue.
  • Tests

    • Added coverage for cleanup behavior, actor failure handling, runtime defaults, and readiness diagnostics.

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>
@github-actions github-actions Bot added the fix label Sep 2, 2026
@github-actions

github-actions Bot commented Sep 2, 2026

Copy link
Copy Markdown
Contributor
Suite Lines Covered Line Rate Branch Rate
Unit Tests 38392/48975 78.4% 62.5%
Integration Tests 23082/46213 50.0% 22.6%

@SandyChapman
SandyChapman marked this pull request as ready for review September 3, 2026 16:42
@SandyChapman
SandyChapman requested review from a team as code owners September 3, 2026 16:42
@coderabbitai

coderabbitai Bot commented Sep 3, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: b1b0950a-3172-4720-a9dd-da1ebfcdc3f3

📥 Commits

Reviewing files that changed from the base of the PR and between 3b68137 and 87534fa.

📒 Files selected for processing (7)
  • packages/sandboxed_gym/src/sandboxed_gym/host/opensandbox.py
  • packages/sandboxed_gym/src/sandboxed_gym/orchestrator.py
  • packages/sandboxed_gym/src/sandboxed_gym/ray/gym_actor.py
  • packages/sandboxed_gym/tests/test_opensandbox_host_provider.py
  • packages/sandboxed_gym/tests/test_ray_actors.py
  • packages/sandboxed_gym/tests/test_sandboxed_gym_host.py
  • packages/sandboxed_gym/tests/test_termination_cleanup.py

Included review availability: Your plan provides up to 12 included reviews per hour; 9 remain after this review.


📝 Walkthrough

Walkthrough

Changes

Sandbox lifecycle and host operation

Layer / File(s) Summary
Runtime defaults and host diagnostics
packages/sandboxed_gym/src/sandboxed_gym/orchestrator.py, packages/sandboxed_gym/src/sandboxed_gym/host/opensandbox.py, packages/sandboxed_gym/tests/test_sandboxed_gym_host.py, packages/sandboxed_gym/tests/test_opensandbox_host_provider.py
Runtime defaults set uv_pip_set_python=True. Host logs include resolved URLs. Readiness timeouts include the health URL and polling error.
Termination cleanup handling
packages/sandboxed_gym/src/sandboxed_gym/orchestrator.py, packages/sandboxed_gym/tests/test_termination_cleanup.py
Termination cleanup handles signals, restores default actions, runs shutdown, re-raises signals, and supports atexit fallback.
Actor startup and failure policy
packages/sandboxed_gym/src/sandboxed_gym/ray/gym_actor.py, packages/sandboxed_gym/tests/test_ray_actors.py
SandboxedGymActor registers shutdown cleanup after startup. Ray actor restarts and task retries are removed.

Suggested reviewers: a2bondar

Merge Risk: ⚪ Minimal · up to 87534

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)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 36 functions across 7 files. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title accurately summarizes the pull request as a set of upstream fixes applied to sandboxed-gym. It is concise and related to the main changes.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
  • Fix all pre-merge checks with AI
✨ Finishing Touches 💡 1
📝 Generate docstrings 💡
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch aalgo-584-sandboxed-gym-upstream-fixes/schapman

Comment @coderabbitai help to get the list of available commands.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant