fix(sandboxed-gym): frame rollout responses and honour an offline environment

Two fixes landed on NeMo-RL's fork of this code while the branch that deletes that
fork sat unmerged. Neither exists here, so completing that consolidation would drop
both -- the exact drift consolidating is meant to end.

Response framing. `/rollouts/run` emits a whitespace heartbeat while a batch runs,
because the body can take minutes and the status line has to leave first. The body
went out under HTTP/1.0 with neither a length nor chunking, delimited only by the
connection closing -- and the OpenSandbox proxy does not wait for that close. It
returned 200 with the lone " " heartbeat as the finished response, and the caller
failed with `JSONDecodeError: Expecting value: line 1 column 1 (char 0)`. Chunked
framing is the delimiter and it is 1.1-only, so the handler moves to 1.1 and takes
only that half: `parse_request` declines connection reuse, which this server gains
nothing from -- a rollout runs for minutes, so a handshake per request rounds to
zero -- and which would cost a real hazard, because an unread request body (the 413
path declines to read one on purpose) is what the server would otherwise parse as
the next request line. Every response is now explicitly framed through one of
`_send_empty`, `_send_body`, `_send_json` or the chunked path. A 1.0 caller cannot
parse chunks, so it gets the batch buffered behind a Content-Length instead: a late
answer is a loud failure where a truncated one is a wrong reward.

Offline environments. `nmp/rl`'s grpo_config already sets `environment_offline` on a
manifest whose wheelhouse is a complete closure, and this package had no such field.
`NemoGymSandboxedConfig` is `extra="forbid"`, so that key does not pass through
unnoticed -- it fails validation, and a sandboxed run with an offline environment
could not start once NeMo-RL consumes this package. The field now exists on both the
caller dialect and the serve config, reaches the host as `NMP_ENVIRONMENT_OFFLINE`,
and sets `UV_OFFLINE` alongside the wheelhouse's `UV_FIND_LINKS`. Index fallback is
otherwise retained, which is a liability only when an index is configured but
unreachable: uv then fails to resolve `uv venv --seed` rather than falling back to
`--find-links`. Not derivable from the package format -- a wheels-v1 package can
ship wheels and still need an index for its agent -- so the caller decides, and an
operator's own `UV_OFFLINE` wins.

Signed-off-by: Sandy Chapman <schapman@nvidia.com>

Author: SandyChapman

packages/sandboxed_gym/src/sandboxed_gym/host/models.py

@@ -223,6 +223,10 @@ class NemoGymSandboxedConfig(BaseModel):
     sandboxed: bool = False
     host_provider: str = "opensandbox"
     environment_path: str | None = None
+    # Whether the environment package's wheelhouse is a complete closure. A wheels-v1 package
+    # can ship wheels and still need an index for its agent, so this is not derivable from the
+    # format -- the caller who built it decides.
+    environment_offline: bool = False
     sandbox: SandboxConfig | None = None
     job_id: str = DEFAULT_JOB_ID
     episode_broker: dict[str, Any] = Field(default_factory=dict)

packages/sandboxed_gym/src/sandboxed_gym/orchestrator.py

@@ -34,6 +34,7 @@
 )
 from sandboxed_gym.host.provider import SandboxedGymHostProvider, get_host_provider
 from sandboxed_gym.runtime.gym_host_runtime import (
+    ENVIRONMENT_OFFLINE_ENV_KEY,
     ENVIRONMENT_PACKAGE_REQUIRED_ENV_KEY,
     GYM_GLOBAL_CONFIG_ENV_KEY,
 )
@@ -276,6 +277,8 @@ def build_gym_host_spec(
         # The mount path is ``/job/environment`` with or without a FileSet. This flag is how the
         # host distinguishes a required package from an image-bundled tree at the same path.
         bootstrap_extra[ENVIRONMENT_PACKAGE_REQUIRED_ENV_KEY] = "true"
+    if cfg.environment_offline:
+        bootstrap_extra[ENVIRONMENT_OFFLINE_ENV_KEY] = "true"
     bootstrap_env = build_bootstrap_env(
         cfg.job_id,
         cfg.environment_path or sandbox.env_mount_path,

packages/sandboxed_gym/src/sandboxed_gym/runtime/gym_host_runtime.py

@@ -40,6 +40,7 @@
 #: for image-bundled Gym too; this flag is how a missing ``nemo-environment.yaml`` becomes a
 #: FileSet error instead of a silent fallback to the image-shipped environment.
 ENVIRONMENT_PACKAGE_REQUIRED_ENV_KEY = "NMP_ENVIRONMENT_PACKAGE_REQUIRED"
+ENVIRONMENT_OFFLINE_ENV_KEY = "NMP_ENVIRONMENT_OFFLINE"
 UV_CACHE_DIR_KEY = "uv_cache_dir"
 UV_VENV_DIR_KEY = "uv_venv_dir"
 # Writable /job/work subdirectory where wheels are installed for the running Gym host.
@@ -58,6 +59,7 @@
 MODEL_CALLS_RESULT_KEY = "_nmp_model_calls"
 # uv setting that points Gym's per-server dependency resolver at the staged wheelhouse.
 UV_FIND_LINKS_ENV_KEY = "UV_FIND_LINKS"
+UV_OFFLINE_ENV_KEY = "UV_OFFLINE"
 NEMO_GYM_EXTRA_ROOTS_ENV_KEY = "NEMO_GYM_EXTRA_ROOTS"
 #: Which agent, resources server, and model to run. Gym has no schema for this key, so
 #: the host pops it and rewrites ``config_paths`` before Gym parses the dict.
@@ -220,6 +222,11 @@ def _environment_package_required() -> bool:
     return os.environ.get(ENVIRONMENT_PACKAGE_REQUIRED_ENV_KEY, "").strip().lower() in {"1", "true", "yes"}
 
 
+def _environment_offline() -> bool:
+    """Whether to resolve from the wheelhouse and uv's cache alone."""
+    return os.environ.get(ENVIRONMENT_OFFLINE_ENV_KEY, "").strip().lower() in {"1", "true", "yes"}
+
+
 def _load_runtime_environment_package(
     environment_path: str,
     *,
@@ -287,6 +294,18 @@ def _install_wheels_v1_dependencies(package: EnvironmentPackage | None, work_pat
     # for image-owned Gym and its core dependencies.
     os.environ[UV_FIND_LINKS_ENV_KEY] = wheels_dir
 
+    if _environment_offline():
+        # That fallback is a liability when the index is configured but unreachable: uv fails
+        # to resolve `uv venv --seed` rather than falling back to --find-links.
+        if UV_OFFLINE_ENV_KEY in os.environ:
+            print(
+                f"gym-host: offline requested, but {UV_OFFLINE_ENV_KEY}="
+                f"{os.environ[UV_OFFLINE_ENV_KEY]} is already set",
+                flush=True,
+            )
+        else:
+            os.environ[UV_OFFLINE_ENV_KEY] = "1"
+
     # Gym starts agent and resource servers as child Python processes. Prepend the wheel target so
     # those processes can import the environment's vendored dependencies.
     existing_pythonpath = os.environ.get("PYTHONPATH")
@@ -606,33 +625,46 @@ def run_rollouts_sync(
 
 
 class Handler(BaseHTTPRequestHandler):
+    # A rollout body's length is not known when the status line has to leave, and HTTP/1.0
+    # cannot delimit such a body except by closing the connection -- so a hop is free to treat
+    # the first bytes it sees as the whole response, which is how a lone " " heartbeat reached
+    # a caller as a finished body and failed to parse. Chunked is the delimiter, and it is
+    # 1.1-only.
+    protocol_version = "HTTP/1.1"
+    # Set per request in do_POST; read by the heartbeat in _await_results.
+    _chunked: bool = False
     max_request_bytes: int = 268_435_456
     max_response_bytes: int = 268_435_456
     heartbeat_interval_s: float = _HEARTBEAT_INTERVAL_S
     rollout_deadline_s: float = _DEFAULT_ROLLOUT_DEADLINE_S
 
+    def parse_request(self) -> bool:
+        """Accept the request, then decline to reuse its connection.
+
+        Reuse gains this server nothing -- a rollout runs for minutes, so a handshake per
+        request rounds to zero -- and would cost: an unread request body is what the server
+        would parse as the next request line, and the 413 path declines to read one on
+        purpose. Set here so it also covers request lines that never reach a handler, and
+        early enough for _announce_close to advertise it.
+        """
+        parsed = super().parse_request()
+        self.close_connection = True
+        return parsed
+
     def do_GET(self) -> None:
         if not self.path.startswith("/health"):
-            self.send_response(404)
-            self.end_headers()
+            self._send_empty(404)
             return
         # Must match do_POST: a host that passes /health and then 503s every rollout is
         # invisible to wait_ready.
         if not _READY or _HEAD_SERVER_CONFIG is None or _ROLLOUT_HELPER is None:
-            body = json.dumps({"status": "starting"}).encode("utf-8")
-            self.send_response(503)
+            self._send_json(503, {"status": "starting"})
         else:
-            body = json.dumps({"status": "ready"}).encode("utf-8")
-            self.send_response(200)
-        self.send_header("Content-Type", "application/json")
-        self.send_header("Content-Length", str(len(body)))
-        self.end_headers()
-        self.wfile.write(body)
+            self._send_json(200, {"status": "ready"})
 
     def do_POST(self) -> None:
         if not self.path.startswith("/rollouts/run"):
-            self.send_response(404)
-            self.end_headers()
+            self._send_empty(404)
             return
         if not _READY or _HEAD_SERVER_CONFIG is None or _ROLLOUT_HELPER is None:
             self._send_json(503, _runtime_error("bootstrap_failed", "Gym host not ready"))
@@ -690,14 +722,25 @@ def do_POST(self) -> None:
         # Committed to 200 before the work is done, so the first byte leaves immediately and no hop
         # can mistake a long batch for a dead one. Everything judgeable from the request alone was
         # rejected with a real status above; failures from here travel in the body as
-        # {"error": ...}, which the caller already treats as fatal. No Content-Length: the body is
-        # delimited by the close that `Connection: close` promises, which is what allows the
-        # heartbeats below to precede a payload of unknown length.
+        # {"error": ...}, which the caller already treats as fatal.
+        self._chunked = self.request_version >= "HTTP/1.1"
+        if not self._chunked:
+            # A 1.0 caller cannot parse chunks, and 1.0's own answer for a body of unknown
+            # length is the close-delimited one that lost the response in the first place. So
+            # buffer instead: the proxy's first-byte cap now covers the whole rollout, but a
+            # late answer is a loud failure where a truncated one is a wrong reward.
+            self._send_body(200, self._await_results(future, started))
+            return
+
         self.send_response(200)
         self.send_header("Content-Type", "application/json")
-        self.send_header("Connection", "close")
+        self.send_header("Transfer-Encoding", "chunked")
+        self._announce_close()
         self.end_headers()
-        self.wfile.write(self._await_results(future, started))
+        self._write_chunk(self._await_results(future, started))
+        # Terminator. Only this ends the body.
+        self.wfile.write(b"0\r\n\r\n")
+        self.wfile.flush()
 
     def _await_results(self, future: concurrent.futures.Future[list[dict]], started: float) -> bytes:
         """Wait for ``future``, heartbeating while it runs, and return the body to send.
@@ -718,9 +761,12 @@ def _await_results(self, future: concurrent.futures.Future[list[dict]], started:
             # its own, which is not this loop's tick.
             done, _ = concurrent.futures.wait([future], timeout=min(self.heartbeat_interval_s, remaining))
             if not done:
+                if not self._chunked:
+                    # Nothing is on the wire yet on the buffered path, so there is nothing to
+                    # heartbeat into; that caller waits for the whole body.
+                    continue
                 try:
-                    self.wfile.write(b" ")
-                    self.wfile.flush()
+                    self._write_chunk(b" ")
                 except OSError as exc:
                     # The caller is gone. Nothing will read this batch, so stop paying for it:
                     # cancelling the future propagates to the collector task on the shared loop.
@@ -765,14 +811,40 @@ def _await_results(self, future: concurrent.futures.Future[list[dict]], started:
     def _error_body(self, code: str, message: str) -> bytes:
         return json.dumps(_runtime_error(code, message)).encode("utf-8")
 
-    def _send_json(self, status: int, payload: dict[str, Any]) -> None:
-        body = json.dumps(payload).encode("utf-8")
+    def _write_chunk(self, data: bytes) -> None:
+        """Emit one HTTP chunk. Empty writes are dropped: a zero-length chunk ends the body."""
+        if not data:
+            return
+        self.wfile.write(b"%X\r\n%s\r\n" % (len(data), data))
+        self.wfile.flush()
+
+    def _announce_close(self) -> None:
+        """Tell the peer the connection is ending rather than let it find out.
+
+        A pooling proxy told nothing here returns the socket to its pool and fails on the
+        next request it sends down it.
+        """
+        if self.close_connection:
+            self.send_header("Connection", "close")
+
+    def _send_empty(self, status: int) -> None:
+        """Frame a bodiless response: zero bytes still has to say it is zero bytes."""
+        self.send_response(status)
+        self._announce_close()
+        self.send_header("Content-Length", "0")
+        self.end_headers()
+
+    def _send_body(self, status: int, body: bytes) -> None:
         self.send_response(status)
+        self._announce_close()
         self.send_header("Content-Type", "application/json")
         self.send_header("Content-Length", str(len(body)))
         self.end_headers()
         self.wfile.write(body)
 
+    def _send_json(self, status: int, payload: dict[str, Any]) -> None:
+        self._send_body(status, json.dumps(payload).encode("utf-8"))
+
     def log_message(self, format: str, *args: Any) -> None:
         return
 

packages/sandboxed_gym/src/sandboxed_gym/serve_config.py

@@ -30,6 +30,9 @@ class SandboxedGymServeConfig(BaseModel):
     job_id: str = DEFAULT_JOB_ID
     host_provider: str = "opensandbox"
     environment_path: str | None = None
+    # Whether that package's wheelhouse is a complete closure. A wheels-v1 package can ship
+    # wheels and still need an index for its agent, so this is not derivable from the format.
+    environment_offline: bool = False
     sandbox: SandboxConfig
     episode_broker: EpisodeBrokerConfig | dict[str, Any] = Field(default_factory=dict)
     gym_global_config: dict[str, Any] = Field(default_factory=dict)