diff --git a/nemo_gym/config_types.py b/nemo_gym/config_types.py index 329fa777fc..f8e345e35b 100644 --- a/nemo_gym/config_types.py +++ b/nemo_gym/config_types.py @@ -126,7 +126,12 @@ class AgentServerRef(BaseModel): name: str -ServerRef = Union[ModelServerRef, ResourcesServerRef, AgentServerRef] +class SandboxServerRef(BaseModel): + type: Literal["sandbox_servers"] + name: str + + +ServerRef = Union[ModelServerRef, ResourcesServerRef, AgentServerRef, SandboxServerRef] ServerRefTypeAdapter = TypeAdapter(ServerRef) @@ -595,6 +600,7 @@ class BaseServerTypeConfig(BaseModel): Literal["responses_api_models"], Literal["resources_servers"], Literal["responses_api_agents"], + Literal["sandbox_servers"], ] ] @@ -623,10 +629,19 @@ class ResponsesAPIAgentServerTypeConfig(BaseServerTypeConfig): responses_api_agents: Dict[str, BaseRunServerTypeConfig] = Field(min_length=1, max_length=1) +class SandboxServerTypeConfig(BaseServerTypeConfig): + SERVER_TYPE: ClassVar[Literal["sandbox_servers"]] = "sandbox_servers" + + model_config = ConfigDict(extra="allow") + + sandbox_servers: Dict[str, BaseRunServerTypeConfig] = Field(min_length=1, max_length=1) + + ServerTypeConfig = Union[ ResponsesAPIModelServerTypeConfig, ResourcesServerTypeConfig, ResponsesAPIAgentServerTypeConfig, + SandboxServerTypeConfig, ] @@ -709,7 +724,7 @@ def is_almost_server(server_type_config_dict: Any) -> bool: return False # Check for server type. - server_type_keys = ["responses_api_models", "resources_servers", "responses_api_agents"] + server_type_keys = ["responses_api_models", "resources_servers", "responses_api_agents", "sandbox_servers"] has_server_type = any(key in server_type_config_dict for key in server_type_keys) if not has_server_type: diff --git a/nemo_gym/sandbox/__init__.py b/nemo_gym/sandbox/__init__.py index e85f5728db..f59e647c09 100644 --- a/nemo_gym/sandbox/__init__.py +++ b/nemo_gym/sandbox/__init__.py @@ -17,6 +17,7 @@ from nemo_gym.sandbox.api import AsyncSandbox, Sandbox from nemo_gym.sandbox.config import resolve_provider_config, resolve_provider_metadata from nemo_gym.sandbox.providers import ( + ConnectableProvider, ExecResult, SandboxCreateError, SandboxCreateVerificationError, @@ -31,12 +32,17 @@ list_providers, register_provider, ) +from nemo_gym.sandbox.ref import SCOPE_OPERATE, SCOPE_OWNER, SandboxRef from nemo_gym.sandbox.utils import rewrite_image __all__ = [ "Sandbox", "AsyncSandbox", + "ConnectableProvider", + "SandboxRef", + "SCOPE_OWNER", + "SCOPE_OPERATE", "ExecResult", "SandboxCreateError", "SandboxCreateVerificationError", diff --git a/nemo_gym/sandbox/api.py b/nemo_gym/sandbox/api.py index 329baa8ad9..38faab4981 100644 --- a/nemo_gym/sandbox/api.py +++ b/nemo_gym/sandbox/api.py @@ -24,6 +24,7 @@ from typing import Any, TypeVar from nemo_gym.sandbox.providers import ( + ConnectableProvider, SandboxExecResult, SandboxHandle, SandboxProvider, @@ -131,6 +132,41 @@ async def stop(self) -> None: await self._provider.aclose() self._closed = True + async def serialize(self, *, scope: str | None = None) -> dict[str, Any]: + """Return a JSON descriptor another process can rebuild this box from. + + Requires a provider that supports the connect capability (the remote + provider, or an external-control-plane provider such as OpenSandbox). For + the remote provider, ``scope`` mints a co-lease (``scope="operate"``). + """ + provider = self._provider + if not isinstance(provider, ConnectableProvider): + name = getattr(provider, "name", type(provider).__name__) + raise RuntimeError(f"provider {name!r} does not support serialize()/connect()") + descriptor = await provider.serialize_handle(self._require_handle(), scope=scope) + # Carry the working directory so a reattached sandbox lands in the same + # place, even for providers whose descriptor does not include it (the + # remote provider's SandboxRef already has it; e.g. OpenSandbox does not). + if isinstance(descriptor, dict) and descriptor.get("workdir") is None and self._spec is not None: + descriptor = {**descriptor, "workdir": self._spec.workdir} + return descriptor + + @classmethod + async def connect(cls, descriptor: Mapping[str, Any] | Any, *, provider: SandboxProvider) -> "AsyncSandbox": + """Rebuild a sandbox in this process from a descriptor produced by + :meth:`serialize`, using ``provider`` (which must support connect).""" + if not isinstance(provider, ConnectableProvider): + name = getattr(provider, "name", type(provider).__name__) + raise RuntimeError(f"provider {name!r} does not support serialize()/connect()") + if not isinstance(descriptor, Mapping) and hasattr(descriptor, "to_dict"): + descriptor = descriptor.to_dict() + handle = await provider.connect(descriptor) + workdir = descriptor.get("workdir") if isinstance(descriptor, Mapping) else None + sandbox = cls(provider, SandboxSpec(workdir=workdir)) + sandbox._handle = handle + sandbox._stopped = False + return sandbox + async def __aenter__(self) -> "AsyncSandbox": return self diff --git a/nemo_gym/sandbox/providers/__init__.py b/nemo_gym/sandbox/providers/__init__.py index d1b2eed496..bdc9205f9c 100644 --- a/nemo_gym/sandbox/providers/__init__.py +++ b/nemo_gym/sandbox/providers/__init__.py @@ -15,6 +15,7 @@ """Sandbox provider registry.""" from nemo_gym.sandbox.providers.base import ( + ConnectableProvider, ExecResult, SandboxCreateError, SandboxCreateVerificationError, @@ -34,6 +35,7 @@ __all__ = [ + "ConnectableProvider", "ExecResult", "SandboxCreateError", "SandboxCreateVerificationError", diff --git a/nemo_gym/sandbox/providers/base.py b/nemo_gym/sandbox/providers/base.py index b8bca468bb..52ef17995c 100644 --- a/nemo_gym/sandbox/providers/base.py +++ b/nemo_gym/sandbox/providers/base.py @@ -18,7 +18,7 @@ from dataclasses import dataclass, field from enum import Enum from pathlib import Path -from typing import Any, Protocol +from typing import Any, Protocol, runtime_checkable class SandboxStatus(str, Enum): @@ -168,3 +168,25 @@ async def close(self, handle: SandboxHandle) -> None: async def aclose(self) -> None: """Close provider-scoped resources such as SDK clients.""" ... + + +@runtime_checkable +class ConnectableProvider(Protocol): + """Optional capability: rebuild a handle in another process from a descriptor. + + Providers whose sandboxes are reachable by id (external control plane, e.g. + OpenSandbox and Fargate, and the sandbox server's remote provider) implement + this. A provider that does not implement it can only be shared by fronting it + with a sandbox server. Membership is checked with ``isinstance`` because the + protocol is ``runtime_checkable``. + """ + + async def serialize_handle(self, handle: SandboxHandle, *, scope: str | None = None) -> dict[str, Any]: + """Return a JSON-serializable descriptor that ``connect`` can rebuild a + handle from. ``scope`` is honored by providers that mint leases (the + remote provider) and ignored by the rest.""" + ... + + async def connect(self, descriptor: Mapping[str, Any]) -> SandboxHandle: + """Rebuild a live handle in this process from a descriptor.""" + ... diff --git a/nemo_gym/sandbox/providers/opensandbox/provider.py b/nemo_gym/sandbox/providers/opensandbox/provider.py index 7d9925f7a2..c744a2ca7e 100644 --- a/nemo_gym/sandbox/providers/opensandbox/provider.py +++ b/nemo_gym/sandbox/providers/opensandbox/provider.py @@ -545,6 +545,32 @@ async def aclose(self) -> None: """Close provider-owned resources.""" return None + async def serialize_handle(self, handle: SandboxHandle, *, scope: str | None = None) -> dict[str, Any]: + """Return a descriptor for reattaching to this sandbox by id. + + OpenSandbox sandboxes are reachable by id from any process that has the + connection config, so the id alone is enough to reconnect and no sandbox + server is needed to share one. ``scope`` is ignored: OpenSandbox has no + lease concept of its own. + """ + return {"sandbox_id": handle.sandbox_id} + + async def connect(self, descriptor: Mapping[str, Any]) -> SandboxHandle: + """Rebuild a live handle from an OpenSandbox sandbox id via the SDK.""" + Sandbox, _, _, _, _ = _require_opensandbox_sdk() + sandbox_id = str(descriptor["sandbox_id"]) + timeout_s = self._create.connect_attempt_timeout_s + sandbox = await asyncio.wait_for( + Sandbox.connect( + sandbox_id, + connection_config=self._connection_config(request_timeout_s=timeout_s), + connect_timeout=timedelta(seconds=timeout_s), + skip_health_check=True, + ), + timeout=timeout_s, + ) + return SandboxHandle(sandbox_id=str(sandbox.id), provider_name=self.name, raw=sandbox) + async def _await_sdk_call( self, awaitable: Any, diff --git a/nemo_gym/sandbox/providers/remote/__init__.py b/nemo_gym/sandbox/providers/remote/__init__.py new file mode 100644 index 0000000000..f157a5ba3a --- /dev/null +++ b/nemo_gym/sandbox/providers/remote/__init__.py @@ -0,0 +1,20 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Remote sandbox provider: forwards operations to a sandbox server over HTTP.""" + +from nemo_gym.sandbox.providers.remote.provider import RemoteSandboxProvider + + +__all__ = ["RemoteSandboxProvider"] diff --git a/nemo_gym/sandbox/providers/remote/provider.py b/nemo_gym/sandbox/providers/remote/provider.py new file mode 100644 index 0000000000..87e66a853b --- /dev/null +++ b/nemo_gym/sandbox/providers/remote/provider.py @@ -0,0 +1,229 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Sandbox provider that forwards operations to a sandbox server over HTTP. + +This is the client half of the sandbox server. It implements the ordinary +``SandboxProvider`` protocol, so ``AsyncSandbox`` drives a server-owned sandbox +with the same code that drives an in-process one; the only difference is that +``handle.raw`` is a serializable ``SandboxRef`` instead of provider-owned state. +It also implements ``ConnectableProvider``, so a handle can be serialized to a +ref and rebuilt in another process. + +The sandbox server owns the real provider (docker, opensandbox, ...). This class +never imports one; it only speaks the server's small HTTP surface. +""" + +from __future__ import annotations + +import base64 +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +from nemo_gym.sandbox.providers.base import ( + SandboxCreateError, + SandboxExecResult, + SandboxHandle, + SandboxSpec, + SandboxStatus, +) +from nemo_gym.sandbox.providers.remote.schemas import ( + CreateSandboxRequest, + DownloadResponse, + ExecRequest, + ExecResponse, + LeaseRequest, + StatusResponse, + UploadRequest, +) +from nemo_gym.sandbox.ref import SCOPE_OWNER, SandboxRef +from nemo_gym.sandbox.transport import SandboxHttpTransport + + +class RemoteSandboxProvider: + """Forwards sandbox operations to a sandbox server. + + Depends only on an injected :class:`SandboxHttpTransport`, so it stays in the + sandbox library without importing the server framework. The server layer + builds it with a Gym-wired transport (``nemo_gym.sandbox_client``). + """ + + name = "remote" + + def __init__(self, *, server_url: str, transport: SandboxHttpTransport, api_key: str | None = None) -> None: + if not server_url: + raise ValueError("RemoteSandboxProvider requires a server_url") + self._server_url = server_url.rstrip("/") + self._transport = transport + self._api_key = api_key + + def _headers(self, lease_token: str | None = None) -> dict[str, str]: + headers: dict[str, str] = {} + if self._api_key: + headers["Authorization"] = f"Bearer {self._api_key}" + if lease_token: + headers["X-NeMo-Gym-Sandbox-Lease"] = lease_token + return headers + + def _ref(self, handle_or_descriptor: Any) -> SandboxRef: + raw = getattr(handle_or_descriptor, "raw", handle_or_descriptor) + if isinstance(raw, SandboxRef): + return raw + if isinstance(raw, Mapping): + return SandboxRef.from_dict(dict(raw)) + raise TypeError(f"RemoteSandboxProvider expected a SandboxRef, got {type(raw).__name__}") + + async def create(self, spec: SandboxSpec) -> SandboxHandle: + """Ask the server to create a sandbox; return an owner-scoped handle. + + Files are not sent here: ``AsyncSandbox.start`` uploads them afterward + through :meth:`upload_file`, matching the in-process providers. + """ + body = CreateSandboxRequest( + image=spec.image, + ttl_s=spec.ttl_s, + ready_timeout_s=spec.ready_timeout_s, + workdir=spec.workdir, + env=dict(spec.env), + metadata=dict(spec.metadata), + resources={ + "cpu": spec.resources.cpu, + "memory_mib": spec.resources.memory_mib, + "disk_gib": spec.resources.disk_gib, + "gpu": spec.resources.gpu, + "gpu_type": spec.resources.gpu_type, + }, + entrypoint=spec.entrypoint, + provider_options=dict(spec.provider_options), + ) + resp = await self._transport.request( + "POST", + f"{self._server_url}/sandboxes", + json=body.model_dump(mode="json"), + headers=self._headers(), + ) + await self._transport.raise_for_status(resp) + ref = SandboxRef.from_dict(await resp.json()) + if not ref.can_close: + raise SandboxCreateError("sandbox server did not return an owner-scoped ref for create()") + return SandboxHandle(sandbox_id=ref.sandbox_id, provider_name=self.name, raw=ref) + + async def exec( + self, + handle: SandboxHandle, + command: str, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_s: int | float | None = None, + user: str | int | None = None, + ) -> SandboxExecResult: + ref = self._ref(handle) + body = ExecRequest(command=command, cwd=cwd, env=env, timeout_s=timeout_s, user=user) + resp = await self._transport.request( + "POST", + f"{self._server_url}/sandboxes/{ref.sandbox_id}/exec", + json=body.model_dump(mode="json"), + headers=self._headers(ref.lease_token), + ) + await self._transport.raise_for_status(resp) + result = ExecResponse.model_validate(await resp.json()) + return SandboxExecResult( + stdout=result.stdout, + stderr=result.stderr, + return_code=result.return_code, + error_type=result.error_type, + ) + + async def upload_file(self, handle: SandboxHandle, source_path: Path, target_path: str) -> None: + ref = self._ref(handle) + body = UploadRequest( + target_path=target_path, + contents_b64=base64.b64encode(Path(source_path).read_bytes()).decode("ascii"), + ) + resp = await self._transport.request( + "POST", + f"{self._server_url}/sandboxes/{ref.sandbox_id}/upload", + json=body.model_dump(mode="json"), + headers=self._headers(ref.lease_token), + ) + await self._transport.raise_for_status(resp) + + async def download_file(self, handle: SandboxHandle, source_path: str, target_path: Path) -> None: + ref = self._ref(handle) + resp = await self._transport.request( + "GET", + f"{self._server_url}/sandboxes/{ref.sandbox_id}/download", + params={"remote_path": source_path}, + headers=self._headers(ref.lease_token), + ) + await self._transport.raise_for_status(resp) + data = DownloadResponse.model_validate(await resp.json()) + target = Path(target_path) + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(base64.b64decode(data.contents_b64)) + + async def status(self, handle: SandboxHandle) -> SandboxStatus: + ref = self._ref(handle) + resp = await self._transport.request( + "GET", + f"{self._server_url}/sandboxes/{ref.sandbox_id}/status", + headers=self._headers(ref.lease_token), + ) + if resp.status != 200: + return SandboxStatus.UNKNOWN + data = StatusResponse.model_validate(await resp.json()) + try: + return SandboxStatus(data.status) + except ValueError: + return SandboxStatus.UNKNOWN + + async def close(self, handle: SandboxHandle) -> None: + """End the lifecycle for an owner ref, or release the lease for a co-lessee.""" + ref = self._ref(handle) + path = f"{self._server_url}/sandboxes/{ref.sandbox_id}" + if ref.scope != SCOPE_OWNER: + path = f"{path}/leases/release" + resp = await self._transport.request("DELETE", path, headers=self._headers(ref.lease_token)) + await self._transport.raise_for_status(resp) + + async def serialize_handle(self, handle: SandboxHandle, *, scope: str | None = None) -> dict[str, Any]: + """Return a ``SandboxRef`` descriptor for this handle. + + With no ``scope`` (or the handle's own scope) the current ref is returned. + A different scope mints a new co-lease on the server and returns that ref, + which is how a verifier is handed operate rights on the owner's box. + """ + ref = self._ref(handle) + if scope is None or scope == ref.scope: + return ref.to_dict() + body = LeaseRequest(scope=scope) + resp = await self._transport.request( + "POST", + f"{self._server_url}/sandboxes/{ref.sandbox_id}/leases", + json=body.model_dump(mode="json"), + headers=self._headers(ref.lease_token), + ) + await self._transport.raise_for_status(resp) + return await resp.json() + + async def connect(self, descriptor: Mapping[str, Any]) -> SandboxHandle: + """Rebuild a handle from a ``SandboxRef`` descriptor. No network call: + the ref already carries the id and lease, and calls are lazy.""" + ref = self._ref(descriptor) + return SandboxHandle(sandbox_id=ref.sandbox_id, provider_name=self.name, raw=ref) + + async def aclose(self) -> None: + return None diff --git a/nemo_gym/sandbox/providers/remote/schemas.py b/nemo_gym/sandbox/providers/remote/schemas.py new file mode 100644 index 0000000000..c8aef837a8 --- /dev/null +++ b/nemo_gym/sandbox/providers/remote/schemas.py @@ -0,0 +1,100 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Wire schema for the sandbox server HTTP protocol. + +One source of truth for the request/response bodies, imported by both the +server (which validates incoming requests and serializes responses) and the +remote provider client (which builds requests and parses responses). +""" + +from __future__ import annotations + +from typing import Any, Optional + +from pydantic import BaseModel, Field + +from nemo_gym.sandbox.ref import SCOPE_OPERATE + + +class CreateSandboxRequest(BaseModel): + """POST /sandboxes — mirrors the createable fields of ``SandboxSpec``.""" + + image: Optional[str] = None + ttl_s: Optional[float] = None + ready_timeout_s: Optional[float] = None + workdir: Optional[str] = None + env: dict[str, str] = Field(default_factory=dict) + metadata: dict[str, str] = Field(default_factory=dict) + resources: dict[str, Any] = Field(default_factory=dict) + entrypoint: Optional[list[str]] = None + provider_options: dict[str, Any] = Field(default_factory=dict) + + +class ExecRequest(BaseModel): + """POST /sandboxes/{id}/exec.""" + + command: str + cwd: Optional[str] = None + env: Optional[dict[str, str]] = None + timeout_s: Optional[float] = None + user: Optional[Any] = None + + +class ExecResponse(BaseModel): + """Result of exec (mirrors ``SandboxExecResult``).""" + + stdout: Optional[str] = None + stderr: Optional[str] = None + return_code: int + error_type: Optional[str] = None + + +class UploadRequest(BaseModel): + """POST /sandboxes/{id}/upload — one base64-encoded file.""" + + target_path: str + contents_b64: str + + +class DownloadResponse(BaseModel): + """GET /sandboxes/{id}/download — one base64-encoded file.""" + + contents_b64: str + + +class StatusResponse(BaseModel): + """GET /sandboxes/{id}/status.""" + + status: str + + +class LeaseRequest(BaseModel): + """POST /sandboxes/{id}/leases — mint a co-lease on an existing sandbox.""" + + scope: str = SCOPE_OPERATE + ttl_s: Optional[float] = None + + +class ReleaseResponse(BaseModel): + """DELETE /sandboxes/{id}/leases/release.""" + + released: bool + remaining_leases: int + + +class DeleteResponse(BaseModel): + """DELETE /sandboxes/{id}.""" + + deleted: bool diff --git a/nemo_gym/sandbox/ref.py b/nemo_gym/sandbox/ref.py new file mode 100644 index 0000000000..69002aed43 --- /dev/null +++ b/nemo_gym/sandbox/ref.py @@ -0,0 +1,80 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Serializable reference to a sandbox owned by a sandbox server. + +A ``SandboxRef`` names one sandbox by a stable id plus a signed lease token, so +a sandbox created by one Gym server (the owner) can be operated by another (a +verifier) over HTTP. It is the descriptor the remote provider produces when a +handle is serialized. Unlike ``SandboxHandle.raw`` (live, process-bound state), +a ``SandboxRef`` is plain JSON and travels in request bodies. + +Scope semantics: +- ``owner`` may run commands, transfer files, and end the sandbox lifecycle. +- ``operate`` may run commands and transfer files, but not end the lifecycle. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + + +SCOPE_OWNER = "owner" +SCOPE_OPERATE = "operate" + + +@dataclass(frozen=True) +class SandboxRef: + """A serializable capability to operate one sandbox on a sandbox server. + + Never carries provider ``raw`` state; ``lease_token`` is the server-signed + grant that authorizes operations and encodes ``{sandbox_id, rollout_id, + scope}``. + """ + + server_url: str + sandbox_id: str + lease_token: str + provider_name: str = "" + scope: str = SCOPE_OPERATE + workdir: str | None = None + extra: dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + return { + "server_url": self.server_url, + "sandbox_id": self.sandbox_id, + "lease_token": self.lease_token, + "provider_name": self.provider_name, + "scope": self.scope, + "workdir": self.workdir, + "extra": dict(self.extra), + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "SandboxRef": + return cls( + server_url=str(data["server_url"]), + sandbox_id=str(data["sandbox_id"]), + lease_token=str(data["lease_token"]), + provider_name=str(data.get("provider_name") or ""), + scope=str(data.get("scope") or SCOPE_OPERATE), + workdir=data.get("workdir"), + extra=dict(data.get("extra") or {}), + ) + + @property + def can_close(self) -> bool: + return self.scope == SCOPE_OWNER diff --git a/nemo_gym/sandbox/transport.py b/nemo_gym/sandbox/transport.py new file mode 100644 index 0000000000..fff9a8c43a --- /dev/null +++ b/nemo_gym/sandbox/transport.py @@ -0,0 +1,44 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""HTTP transport interface for the remote sandbox provider. + +The remote provider needs to make async HTTP calls, but ``nemo_gym.sandbox`` is +a library that must not depend on the server framework (``nemo_gym.server_utils``). +So the provider depends only on this small transport interface, and the concrete +aiohttp-backed transport is injected from the server layer (see +``nemo_gym.sandbox_client``). +""" + +from __future__ import annotations + +from typing import Any, Protocol, runtime_checkable + + +@runtime_checkable +class SandboxHttpTransport(Protocol): + """Minimal async HTTP surface the remote sandbox provider needs. + + Implemented in the server layer over Gym's global aiohttp client, so the + sandbox library never imports it directly. + """ + + async def request(self, method: str, url: str, **kwargs: Any) -> Any: + """Perform an HTTP request and return a response exposing ``status`` and + an awaitable ``json()`` (aiohttp ``ClientResponse`` shape).""" + ... + + async def raise_for_status(self, response: Any) -> None: + """Raise if the response carries an error status.""" + ... diff --git a/nemo_gym/sandbox_client.py b/nemo_gym/sandbox_client.py new file mode 100644 index 0000000000..dad2f440d7 --- /dev/null +++ b/nemo_gym/sandbox_client.py @@ -0,0 +1,63 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Client-side helpers for talking to a sandbox server. + +This is the one bridge between the ``nemo_gym.sandbox`` library and the server +framework (``nemo_gym.server_utils``): it wires Gym's global aiohttp client into +the remote provider's transport. The sandbox library depends only on the +transport interface, and ``server_utils`` does not import the sandbox library, +so the dependency runs one way (server framework -> sandbox library). + +A Gym server that operates a sandbox server (an agent that creates a box, a +resources server that reattaches to one or spins up its own) uses these helpers +rather than constructing the remote provider directly. +""" + +from __future__ import annotations + +from typing import Any + +from nemo_gym.sandbox import AsyncSandbox, SandboxRef +from nemo_gym.sandbox.providers.remote import RemoteSandboxProvider +from nemo_gym.server_utils import raise_for_status, request + + +class GymSandboxHttpTransport: + """SandboxHttpTransport backed by Gym's global aiohttp client.""" + + async def request(self, method: str, url: str, **kwargs: Any) -> Any: + return await request(method, url, **kwargs) + + async def raise_for_status(self, response: Any) -> None: + await raise_for_status(response) + + +_TRANSPORT = GymSandboxHttpTransport() + + +def gym_sandbox_transport() -> GymSandboxHttpTransport: + """The shared Gym-backed sandbox HTTP transport.""" + return _TRANSPORT + + +def make_remote_provider(server_url: str, *, api_key: str | None = None) -> RemoteSandboxProvider: + """Build a remote sandbox provider wired to Gym's HTTP client.""" + return RemoteSandboxProvider(server_url=server_url, transport=_TRANSPORT, api_key=api_key) + + +async def connect_sandbox(ref: SandboxRef, *, api_key: str | None = None) -> AsyncSandbox: + """Reattach to a server-owned sandbox by reference, wired to Gym's HTTP client.""" + provider = make_remote_provider(ref.server_url, api_key=api_key) + return await AsyncSandbox.connect(ref, provider=provider) diff --git a/sandbox_servers/sandbox_server/app.py b/sandbox_servers/sandbox_server/app.py new file mode 100644 index 0000000000..d52db7aee1 --- /dev/null +++ b/sandbox_servers/sandbox_server/app.py @@ -0,0 +1,291 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Sandbox server: owns physical sandboxes and lends them across Gym servers. + +It fronts exactly one sandbox provider (docker, opensandbox, ...) and exposes +its lifecycle and operations over HTTP. The value it adds over an in-process +``nemo_gym.sandbox.AsyncSandbox`` is that a sandbox created by one server can be +operated by another by reference: the client holds a small signed ``SandboxRef``, +not provider-owned in-process state. + +Ownership model: +- create mints an OWNER lease; only an owner lease may DELETE (destroy) the box. +- a co-lease (POST /sandboxes/{id}/leases, scope=operate) may exec/upload/ + download but not destroy; this is what a verifier reattaches with. +- leases are signed capabilities bound to the sandbox's rollout id, so a ref + leaked from one rollout can't touch another rollout's box. +- ttl_s reaps orphaned boxes regardless of leases (crash safety). +""" + +from __future__ import annotations + +import asyncio +import base64 +import contextlib +import tempfile +import time +from contextlib import asynccontextmanager +from pathlib import Path +from typing import Any, Optional +from uuid import uuid4 + +from fastapi import FastAPI, HTTPException, Request +from itsdangerous import BadSignature, URLSafeTimedSerializer +from pydantic import ConfigDict + +from nemo_gym.config_types import BaseRunServerInstanceConfig +from nemo_gym.global_config import get_global_config_dict +from nemo_gym.sandbox import ( + SCOPE_OPERATE, + SCOPE_OWNER, + AsyncSandbox, + SandboxRef, + SandboxResources, + SandboxSpec, + create_provider, + resolve_provider_config, +) +from nemo_gym.sandbox.providers.remote.schemas import ( + CreateSandboxRequest, + DeleteResponse, + DownloadResponse, + ExecRequest, + ExecResponse, + LeaseRequest, + ReleaseResponse, + StatusResponse, + UploadRequest, +) +from nemo_gym.server_utils import SimpleServer, is_nemo_gym_fastapi_entrypoint + + +LEASE_HEADER = "X-NeMo-Gym-Sandbox-Lease" +_LEASE_SALT = "nemo-gym-sandbox-lease" +ROLLOUT_METADATA_KEY = "ng_rollout_id" + + +class SandboxServerConfig(BaseRunServerInstanceConfig): + model_config = ConfigDict(extra="allow") + + # The single provider this server fronts: an inline single-key mapping + # ({docker: {}}) or the name of a composed provider block. Resolved via + # nemo_gym.sandbox.resolve_provider_config against the merged global config. + sandbox_provider: Any + # Pool cap: max concurrently-live sandboxes. Admission control in one place + # instead of every consumer inventing its own limit. + max_concurrent: int = 64 + # Orphan backstop applied when a create request omits ttl_s. + default_ttl_s: Optional[float] = 1800.0 + # Lease-signing secret. Left unset -> a per-process random secret (fine for + # a single-process local server; set it for multi-worker/persistent runs). + lease_secret: Optional[str] = None + + +class _Entry: + """One live sandbox and its bookkeeping.""" + + def __init__(self, sandbox: AsyncSandbox, rollout_id: str, workdir: Optional[str], expires_at: float) -> None: + self.sandbox = sandbox + self.rollout_id = rollout_id + self.workdir = workdir + self.expires_at = expires_at + self.leases = 1 # the owner lease + + +class SandboxServer(SimpleServer): + config: SandboxServerConfig + + def model_post_init(self, context: Any) -> None: + self._entries: dict[str, _Entry] = {} + self._lock = asyncio.Lock() + self._sem = asyncio.Semaphore(self.config.max_concurrent) + self._signer = URLSafeTimedSerializer(self.config.lease_secret or uuid4().hex, salt=_LEASE_SALT) + self._provider_config = resolve_provider_config(self.config.sandbox_provider, get_global_config_dict()) + return super().model_post_init(context) + + # -- lease helpers ----------------------------------------------------- + + def _mint(self, sandbox_id: str, rollout_id: str, scope: str) -> str: + return self._signer.dumps({"sid": sandbox_id, "rid": rollout_id, "scope": scope}) + + def _check(self, request: Request, sandbox_id: str, *, require_owner: bool = False) -> _Entry: + token = request.headers.get(LEASE_HEADER) + if not token: + raise HTTPException(status_code=401, detail="missing sandbox lease token") + try: + payload = self._signer.loads(token, max_age=24 * 3600) + except BadSignature as e: + raise HTTPException(status_code=401, detail="invalid sandbox lease token") from e + if payload.get("sid") != sandbox_id: + raise HTTPException(status_code=403, detail="lease token does not match this sandbox") + entry = self._entries.get(sandbox_id) + if entry is None: + raise HTTPException(status_code=404, detail=f"sandbox {sandbox_id!r} not found") + if payload.get("rid") != entry.rollout_id: + raise HTTPException(status_code=403, detail="lease token rollout id does not match the sandbox") + if require_owner and payload.get("scope") != SCOPE_OWNER: + raise HTTPException(status_code=403, detail="this operation requires an owner lease") + return entry + + def _server_url(self) -> str: + host = self.config.host or "127.0.0.1" + if host in ("0.0.0.0", "::", ""): + host = "127.0.0.1" + return f"http://{host}:{self.config.port}" + + def _ref(self, sandbox_id: str, rollout_id: str, scope: str, workdir: Optional[str]) -> dict: + return SandboxRef( + server_url=self._server_url(), + sandbox_id=sandbox_id, + lease_token=self._mint(sandbox_id, rollout_id, scope), + provider_name=next(iter(self._provider_config)), + scope=scope, + workdir=workdir, + ).to_dict() + + # -- routes ------------------------------------------------------------ + + def setup_webserver(self) -> FastAPI: + @asynccontextmanager + async def lifespan(app: FastAPI): + reaper = asyncio.ensure_future(self._reap_loop()) + try: + yield + finally: + reaper.cancel() + with contextlib.suppress(asyncio.CancelledError): + await reaper + await self._shutdown_all() + + app = FastAPI(lifespan=lifespan) + self.setup_liveness(app) + app.post("/sandboxes")(self.create_sandbox) + app.post("/sandboxes/{sandbox_id}/exec")(self.exec_sandbox) + app.post("/sandboxes/{sandbox_id}/upload")(self.upload_sandbox) + app.get("/sandboxes/{sandbox_id}/download")(self.download_sandbox) + app.get("/sandboxes/{sandbox_id}/status")(self.status_sandbox) + app.post("/sandboxes/{sandbox_id}/leases")(self.grant_lease) + app.delete("/sandboxes/{sandbox_id}/leases/release")(self.release_lease) + app.delete("/sandboxes/{sandbox_id}")(self.delete_sandbox) + return app + + async def create_sandbox(self, body: CreateSandboxRequest) -> dict: + rollout_id = str(body.metadata.get(ROLLOUT_METADATA_KEY) or "") + ttl_s = body.ttl_s if body.ttl_s is not None else self.config.default_ttl_s + spec = SandboxSpec( + image=body.image, + ttl_s=ttl_s, + ready_timeout_s=body.ready_timeout_s, + workdir=body.workdir, + env=dict(body.env), + metadata=dict(body.metadata), + resources=SandboxResources.from_mapping(body.resources), + entrypoint=body.entrypoint, + provider_options=dict(body.provider_options), + ) + await self._sem.acquire() + try: + provider = create_provider(self._provider_config) + sandbox = await AsyncSandbox(provider, spec).start() + except Exception: + self._sem.release() + raise + sandbox_id = sandbox._handle.sandbox_id # provider-neutral id + expires_at = time.monotonic() + ttl_s if ttl_s else float("inf") + async with self._lock: + self._entries[sandbox_id] = _Entry(sandbox, rollout_id, body.workdir, expires_at) + return self._ref(sandbox_id, rollout_id, SCOPE_OWNER, body.workdir) + + async def exec_sandbox(self, sandbox_id: str, body: ExecRequest, request: Request) -> ExecResponse: + entry = self._check(request, sandbox_id) + result = await entry.sandbox.exec( + body.command, cwd=body.cwd, env=body.env, timeout_s=body.timeout_s, user=body.user + ) + return ExecResponse( + stdout=result.stdout, + stderr=result.stderr, + return_code=result.return_code, + error_type=result.error_type, + ) + + async def upload_sandbox(self, sandbox_id: str, body: UploadRequest, request: Request) -> dict: + entry = self._check(request, sandbox_id) + with tempfile.TemporaryDirectory(prefix="nemo-gym-sbxsrv-up-") as tmp: + src = Path(tmp) / "payload" + src.write_bytes(base64.b64decode(body.contents_b64)) + await entry.sandbox.upload(src, body.target_path) + return {} + + async def download_sandbox(self, sandbox_id: str, remote_path: str, request: Request) -> DownloadResponse: + entry = self._check(request, sandbox_id) + with tempfile.TemporaryDirectory(prefix="nemo-gym-sbxsrv-dl-") as tmp: + dst = Path(tmp) / "payload" + await entry.sandbox.download(remote_path, dst) + return DownloadResponse(contents_b64=base64.b64encode(dst.read_bytes()).decode("ascii")) + + async def status_sandbox(self, sandbox_id: str, request: Request) -> StatusResponse: + entry = self._check(request, sandbox_id) + return StatusResponse(status=(await entry.sandbox.status()).value) + + async def grant_lease(self, sandbox_id: str, body: LeaseRequest, request: Request) -> dict: + entry = self._check(request, sandbox_id) + scope = SCOPE_OWNER if body.scope == SCOPE_OWNER else SCOPE_OPERATE + async with self._lock: + entry.leases += 1 + return self._ref(sandbox_id, entry.rollout_id, scope, entry.workdir) + + async def release_lease(self, sandbox_id: str, request: Request) -> ReleaseResponse: + entry = self._check(request, sandbox_id) + async with self._lock: + entry.leases = max(0, entry.leases - 1) + return ReleaseResponse(released=True, remaining_leases=entry.leases) + + async def delete_sandbox(self, sandbox_id: str, request: Request) -> DeleteResponse: + self._check(request, sandbox_id, require_owner=True) + await self._destroy(sandbox_id) + return DeleteResponse(deleted=True) + + # -- lifecycle --------------------------------------------------------- + + async def _destroy(self, sandbox_id: str) -> None: + async with self._lock: + entry = self._entries.pop(sandbox_id, None) + if entry is None: + return + try: + await entry.sandbox.stop() + finally: + self._sem.release() + + async def _reap_loop(self) -> None: + while True: + await asyncio.sleep(10.0) + now = time.monotonic() + expired = [sid for sid, e in list(self._entries.items()) if e.expires_at <= now] + for sid in expired: + with contextlib.suppress(Exception): + await self._destroy(sid) + + async def _shutdown_all(self) -> None: + for sid in list(self._entries): + with contextlib.suppress(Exception): + await self._destroy(sid) + + +if __name__ == "__main__": + SandboxServer.run_webserver() +elif is_nemo_gym_fastapi_entrypoint(__file__): + app = SandboxServer.run_webserver() # noqa: F401 diff --git a/sandbox_servers/sandbox_server/configs/sandbox_server.yaml b/sandbox_servers/sandbox_server/configs/sandbox_server.yaml new file mode 100644 index 0000000000..a87d9aee75 --- /dev/null +++ b/sandbox_servers/sandbox_server/configs/sandbox_server.yaml @@ -0,0 +1,18 @@ +# A sandbox server fronting the local Docker provider. Compose it into a run +# with +config_paths and reference it by name from a server that operates it. +# +# One server instance fronts exactly ONE provider. To offer another backend +# (e.g. opensandbox on a cluster), compose a second sandbox_servers block with +# its own provider; operators pick one by name. +pool_sandbox_server: + sandbox_servers: + sandbox_server: + entrypoint: app.py + sandbox_provider: + docker: + create: + # A container that needs to reach a service on the host uses host + # networking; drop this for isolated boxes. + network: host + max_concurrent: 16 + default_ttl_s: 1800 diff --git a/sandbox_servers/sandbox_server/requirements.txt b/sandbox_servers/sandbox_server/requirements.txt new file mode 100644 index 0000000000..00ed83213e --- /dev/null +++ b/sandbox_servers/sandbox_server/requirements.txt @@ -0,0 +1 @@ +-e nemo-gym[dev] @ ../../ diff --git a/tests/unit_tests/test_sandbox.py b/tests/unit_tests/test_sandbox.py index ade373adf6..16b946007b 100644 --- a/tests/unit_tests/test_sandbox.py +++ b/tests/unit_tests/test_sandbox.py @@ -1216,3 +1216,59 @@ async def exec( assert exc_info.value.messages[0]["extra"]["submission"] == "final answer" finally: env.cleanup() + + +@requires_tenacity +def test_opensandbox_implements_connectable_provider(monkeypatch) -> None: + asyncio.run(_assert_opensandbox_implements_connectable_provider(monkeypatch)) + + +async def _assert_opensandbox_implements_connectable_provider(monkeypatch) -> None: + from nemo_gym.sandbox import ConnectableProvider + + opensandbox_provider_module, OpenSandboxProvider, *_unused = _require_opensandbox_provider() + + class FakeConnectionConfig: + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + + class FakeSDKSandbox: + connect_calls: list[dict[str, Any]] = [] + + def __init__(self, sandbox_id: str) -> None: + self.id = sandbox_id + + @classmethod + async def connect(cls, sandbox_id: str, **kwargs: Any) -> "FakeSDKSandbox": + cls.connect_calls.append({"sandbox_id": sandbox_id, **kwargs}) + return cls(sandbox_id) + + monkeypatch.setattr( + opensandbox_provider_module, + "_require_opensandbox_sdk", + lambda: (FakeSDKSandbox, FakeConnectionConfig, object, object, object), + ) + + provider = OpenSandboxProvider( + connection={"domain": "sandbox.example", "protocol": "https"}, + create={"connect_attempt_timeout_s": 1}, + probe={"command": None}, + ) + + # The provider satisfies the optional capability protocol. + assert isinstance(provider, ConnectableProvider) + + # serialize_handle returns a descriptor of just the id. + descriptor = await provider.serialize_handle( + SandboxHandle(sandbox_id="sdk-sandbox-9", provider_name="opensandbox", raw=object()) + ) + assert descriptor == {"sandbox_id": "sdk-sandbox-9"} + + # connect rebuilds a live handle by reconnecting to that id via the SDK. + handle = await provider.connect(descriptor) + assert handle.sandbox_id == "sdk-sandbox-9" + assert isinstance(handle.raw, FakeSDKSandbox) + connect_call = FakeSDKSandbox.connect_calls[0] + assert connect_call["sandbox_id"] == "sdk-sandbox-9" + assert connect_call["skip_health_check"] is True + assert connect_call["connection_config"].kwargs["domain"] == "sandbox.example" diff --git a/tests/unit_tests/test_sandbox_connect.py b/tests/unit_tests/test_sandbox_connect.py new file mode 100644 index 0000000000..33676e5cf1 --- /dev/null +++ b/tests/unit_tests/test_sandbox_connect.py @@ -0,0 +1,169 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the ConnectableProvider capability and the serialize/connect facade. + +Hermetic: a fake in-process provider stands in for a reattachable backend, so +these exercise the facade and protocol without any sandbox server. +""" + +from __future__ import annotations + +import asyncio +from pathlib import Path +from typing import Any +from uuid import uuid4 + +import pytest + +from nemo_gym.sandbox import ( + AsyncSandbox, + ConnectableProvider, + SandboxExecResult, + SandboxHandle, + SandboxSpec, + SandboxStatus, +) + + +_STORE: dict[str, dict[str, Any]] = {} + + +class FakeConnectableProvider: + """In-memory provider that supports connect; boxes live in a global store + keyed by id, so a second instance can reconnect by id.""" + + name = "fake_connectable" + + async def create(self, spec: SandboxSpec) -> SandboxHandle: + sid = f"fake-{uuid4().hex[:8]}" + _STORE[sid] = {"files": {}, "closed": False} + return SandboxHandle(sandbox_id=sid, provider_name=self.name, raw=sid) + + async def exec(self, handle, command, *, cwd=None, env=None, timeout_s=None, user=None) -> SandboxExecResult: + box = _STORE.get(handle.sandbox_id) + if box is None or box["closed"]: + return SandboxExecResult(stdout=None, stderr="no such sandbox", return_code=1) + if command.startswith("cat "): + data = box["files"].get(command[4:].strip()) + if data is None: + return SandboxExecResult(stdout=None, stderr="no such file", return_code=1) + return SandboxExecResult(stdout=data.decode(), stderr=None, return_code=0) + if command == "pwd": + return SandboxExecResult(stdout=cwd or "", stderr=None, return_code=0) + return SandboxExecResult(stdout=f"ran: {command}", stderr=None, return_code=0) + + async def upload_file(self, handle, source_path, target_path) -> None: + _STORE[handle.sandbox_id]["files"][target_path] = Path(source_path).read_bytes() + + async def download_file(self, handle, source_path, target_path) -> None: + Path(target_path).parent.mkdir(parents=True, exist_ok=True) + Path(target_path).write_bytes(_STORE[handle.sandbox_id]["files"][source_path]) + + async def status(self, handle) -> SandboxStatus: + box = _STORE.get(handle.sandbox_id) + if box is None: + return SandboxStatus.UNKNOWN + return SandboxStatus.STOPPED if box["closed"] else SandboxStatus.RUNNING + + async def close(self, handle) -> None: + box = _STORE.get(handle.sandbox_id) + if box is not None: + box["closed"] = True + + async def aclose(self) -> None: + return None + + async def serialize_handle(self, handle, *, scope=None) -> dict[str, Any]: + return {"sandbox_id": handle.sandbox_id} + + async def connect(self, descriptor) -> SandboxHandle: + sid = str(descriptor["sandbox_id"]) + if sid not in _STORE: + raise RuntimeError(f"no such sandbox {sid!r}") + return SandboxHandle(sandbox_id=sid, provider_name=self.name, raw=sid) + + +class _OpsOnlyProvider: + """A provider without the connect capability (negative case).""" + + name = "ops_only" + + async def create(self, spec): + return SandboxHandle(sandbox_id="x", provider_name=self.name, raw="x") + + async def exec(self, *a, **k): + return SandboxExecResult(stdout="", stderr=None, return_code=0) + + async def upload_file(self, *a, **k): + return None + + async def download_file(self, *a, **k): + return None + + async def status(self, handle): + return SandboxStatus.RUNNING + + async def close(self, handle): + return None + + async def aclose(self): + return None + + +def test_connectable_provider_isinstance() -> None: + assert isinstance(FakeConnectableProvider(), ConnectableProvider) + assert not isinstance(_OpsOnlyProvider(), ConnectableProvider) + + +def test_serialize_then_connect_round_trip(tmp_path: Path) -> None: + async def _run() -> None: + sandbox = await AsyncSandbox(FakeConnectableProvider(), SandboxSpec(workdir="/w")).start() + payload = tmp_path / "f.txt" + payload.write_text("hello-connect") + await sandbox.upload(payload, "/w/f.txt") + + descriptor = await sandbox.serialize() + assert "sandbox_id" in descriptor + # The facade carries workdir even though the provider descriptor omits it. + assert descriptor["workdir"] == "/w" + + # A separate provider instance rebuilds a working handle from the descriptor. + reattached = await AsyncSandbox.connect(descriptor, provider=FakeConnectableProvider()) + result = await reattached.exec("cat /w/f.txt") + assert result.return_code == 0 + assert result.stdout == "hello-connect" + + # The reattached sandbox defaults exec to the original working directory. + assert (await reattached.exec("pwd")).stdout == "/w" + + asyncio.run(_run()) + + +def test_serialize_requires_connect_capability() -> None: + async def _run() -> None: + sandbox = await AsyncSandbox(_OpsOnlyProvider(), SandboxSpec()).start() + with pytest.raises(RuntimeError, match="serialize"): + await sandbox.serialize() + + asyncio.run(_run()) + + +def test_connect_requires_connect_capability() -> None: + async def _run() -> None: + with pytest.raises(RuntimeError, match="serialize"): + await AsyncSandbox.connect({"sandbox_id": "x"}, provider=_OpsOnlyProvider()) + + asyncio.run(_run()) diff --git a/tests/unit_tests/test_sandbox_server.py b/tests/unit_tests/test_sandbox_server.py new file mode 100644 index 0000000000..220ae994b1 --- /dev/null +++ b/tests/unit_tests/test_sandbox_server.py @@ -0,0 +1,244 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Standalone tests for the sandbox server. + +Hermetic: a fake in-process provider backs the server, so no docker, cloud, or +external OpenSandbox server is needed. Covers SandboxRef serialization and a +full in-process HTTP round-trip against the real server app (driven via httpx +ASGI transport) through RemoteSandboxProvider. +""" + +from __future__ import annotations + +import asyncio +import importlib.util +import sys +from pathlib import Path +from typing import Any +from unittest.mock import MagicMock +from uuid import uuid4 + +import httpx +import pytest + +from nemo_gym.sandbox import ( + SCOPE_OPERATE, + SCOPE_OWNER, + AsyncSandbox, + SandboxExecResult, + SandboxHandle, + SandboxRef, + SandboxSpec, + SandboxStatus, + register_provider, +) +from nemo_gym.sandbox.providers.remote import RemoteSandboxProvider +from nemo_gym.server_utils import ServerClient + + +# --- load the sandbox server app (lives outside the nemo_gym package) -------- + +_APP_PATH = Path(__file__).resolve().parents[2] / "sandbox_servers" / "sandbox_server" / "app.py" +_spec = importlib.util.spec_from_file_location("sandbox_server_app", _APP_PATH) +sandbox_server_app = importlib.util.module_from_spec(_spec) +# Register before executing so pydantic can resolve the config's annotations +# (the module uses ``from __future__ import annotations``, so forward refs are +# looked up via ``sys.modules[cls.__module__]``). +sys.modules["sandbox_server_app"] = sandbox_server_app +_spec.loader.exec_module(sandbox_server_app) +SandboxServer = sandbox_server_app.SandboxServer +SandboxServerConfig = sandbox_server_app.SandboxServerConfig + + +# --- a fake provider backing the server (in-memory) -------------------------- + +_FAKE_STORE: dict[str, dict[str, Any]] = {} + + +class FakeProvider: + """In-memory provider used to back the sandbox server in tests.""" + + name = "fake" + + def __init__(self, **kwargs: Any) -> None: + pass + + async def create(self, spec: SandboxSpec) -> SandboxHandle: + sid = f"fake-{uuid4().hex[:8]}" + _FAKE_STORE[sid] = {"files": {}, "closed": False, "workdir": spec.workdir} + return SandboxHandle(sandbox_id=sid, provider_name=self.name, raw=sid) + + async def exec(self, handle, command, *, cwd=None, env=None, timeout_s=None, user=None) -> SandboxExecResult: + box = _FAKE_STORE.get(handle.sandbox_id) + if box is None or box["closed"]: + return SandboxExecResult(stdout=None, stderr="no such sandbox", return_code=1) + if command.startswith("cat "): + data = box["files"].get(command[4:].strip()) + if data is None: + return SandboxExecResult(stdout=None, stderr="no such file", return_code=1) + return SandboxExecResult(stdout=data.decode(), stderr=None, return_code=0) + if command.startswith("echo "): + return SandboxExecResult(stdout=command[len("echo ") :], stderr=None, return_code=0) + return SandboxExecResult(stdout=f"ran: {command}", stderr=None, return_code=0) + + async def upload_file(self, handle, source_path, target_path) -> None: + _FAKE_STORE[handle.sandbox_id]["files"][target_path] = Path(source_path).read_bytes() + + async def download_file(self, handle, source_path, target_path) -> None: + Path(target_path).parent.mkdir(parents=True, exist_ok=True) + Path(target_path).write_bytes(_FAKE_STORE[handle.sandbox_id]["files"][source_path]) + + async def status(self, handle) -> SandboxStatus: + box = _FAKE_STORE.get(handle.sandbox_id) + if box is None: + return SandboxStatus.UNKNOWN + return SandboxStatus.STOPPED if box["closed"] else SandboxStatus.RUNNING + + async def close(self, handle) -> None: + box = _FAKE_STORE.get(handle.sandbox_id) + if box is not None: + box["closed"] = True + + async def aclose(self) -> None: + return None + + +# --- httpx ASGI transport adapter (aiohttp-shaped responses) ----------------- + + +class _Resp: + def __init__(self, response: httpx.Response) -> None: + self._response = response + self.status = response.status_code + + async def json(self) -> Any: + return self._response.json() + + +class ASGITransportAdapter: + """A SandboxHttpTransport that drives an ASGI app in-process, exposing the + aiohttp ``ClientResponse`` shape the remote provider expects.""" + + def __init__(self, app: Any) -> None: + self._client = httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://testserver") + + async def request(self, method, url, *, json=None, params=None, headers=None) -> _Resp: + response = await self._client.request(method, url, json=json, params=params, headers=headers) + return _Resp(response) + + async def raise_for_status(self, response: _Resp) -> None: + response._response.raise_for_status() + + async def aclose(self) -> None: + await self._client.aclose() + + +def _build_server(monkeypatch: pytest.MonkeyPatch) -> SandboxServer: + register_provider("fake", FakeProvider, override=True) + # The server resolves its provider against the global config; an inline + # single-key mapping does not need it, so stub it to stay hermetic. + monkeypatch.setattr(sandbox_server_app, "get_global_config_dict", lambda: {}) + config = SandboxServerConfig( + host="", + port=0, + entrypoint="", + name="sbx", + sandbox_provider={"fake": {}}, + max_concurrent=4, + default_ttl_s=None, + ) + return SandboxServer(config=config, server_client=MagicMock(spec=ServerClient)) + + +# --- tests ------------------------------------------------------------------- + + +def test_sandbox_ref_roundtrip_and_scope() -> None: + ref = SandboxRef( + server_url="http://host:8080", + sandbox_id="abc", + lease_token="tok", + provider_name="fake", + scope=SCOPE_OWNER, + workdir="/work", + ) + restored = SandboxRef.from_dict(ref.to_dict()) + assert restored == ref + assert restored.can_close is True + operate = SandboxRef.from_dict({"server_url": "u", "sandbox_id": "s", "lease_token": "t"}) + assert operate.scope == SCOPE_OPERATE + assert operate.can_close is False + + +def test_server_end_to_end_over_http(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + async def _run() -> None: + server = _build_server(monkeypatch) + app = server.setup_webserver() + transport = ASGITransportAdapter(app) + try: + provider = RemoteSandboxProvider(server_url="http://testserver", transport=transport) + sandbox = await AsyncSandbox(provider, SandboxSpec(image="fake-image", workdir="/w")).start() + + # exec + result = await sandbox.exec("echo hi") + assert result.return_code == 0 + assert result.stdout == "hi" + + # upload + download round-trip through the server + payload = tmp_path / "in.txt" + payload.write_text("through-the-server") + await sandbox.upload(payload, "/w/in.txt") + out = tmp_path / "out.txt" + await sandbox.download("/w/in.txt", out) + assert out.read_text() == "through-the-server" + + # status + assert await sandbox.status() == SandboxStatus.RUNNING + + # owner ref serializes to itself; an operate co-lease is a fresh ref + owner_ref = SandboxRef.from_dict(await sandbox.serialize()) + assert owner_ref.scope == SCOPE_OWNER + operate_ref = SandboxRef.from_dict(await sandbox.serialize(scope=SCOPE_OPERATE)) + assert operate_ref.scope == SCOPE_OPERATE + assert operate_ref.sandbox_id == owner_ref.sandbox_id + assert operate_ref.lease_token != owner_ref.lease_token + + # a second server reattaches with the co-lease and operates the box + co_provider = RemoteSandboxProvider(server_url="http://testserver", transport=transport) + co_sandbox = await AsyncSandbox.connect(operate_ref, provider=co_provider) + co_result = await co_sandbox.exec("echo from-colease") + assert co_result.stdout == "from-colease" + + # a bad lease token is rejected + bad_ref = SandboxRef( + server_url="http://testserver", + sandbox_id=owner_ref.sandbox_id, + lease_token="not-a-valid-token", + scope=SCOPE_OPERATE, + ) + bad_provider = RemoteSandboxProvider(server_url="http://testserver", transport=transport) + bad_sandbox = await AsyncSandbox.connect(bad_ref, provider=bad_provider) + with pytest.raises(httpx.HTTPStatusError): + await bad_sandbox.exec("echo nope") + + # releasing the co-lease does not destroy the box; owner close does + await co_sandbox.stop() + assert await sandbox.status() == SandboxStatus.RUNNING + await sandbox.stop() + finally: + await transport.aclose() + + asyncio.run(_run())