From 29cb8d73244c7d7c9262d4db618d01230a164a7a Mon Sep 17 00:00:00 2001 From: Piotr Sarna Date: Tue, 20 Oct 2026 12:08:43 +0200 Subject: [PATCH] sandbox: add pty capability to e2b sandbox provider So that it is usable with sandbox/shell.py script Signed-off-by: Piotr Sarna --- nemo_gym/sandbox/providers/e2b/provider.py | 10 ++ nemo_gym/sandbox/providers/e2b/pty.py | 186 ++++++++++++++++++++ nemo_gym/sandbox/shell.py | 5 + tests/unit_tests/test_e2b_pty.py | 188 +++++++++++++++++++++ 4 files changed, 389 insertions(+) create mode 100644 nemo_gym/sandbox/providers/e2b/pty.py create mode 100644 tests/unit_tests/test_e2b_pty.py diff --git a/nemo_gym/sandbox/providers/e2b/provider.py b/nemo_gym/sandbox/providers/e2b/provider.py index 5999e80394f..24988335268 100644 --- a/nemo_gym/sandbox/providers/e2b/provider.py +++ b/nemo_gym/sandbox/providers/e2b/provider.py @@ -49,10 +49,13 @@ SandboxCreateError, SandboxExecResult, SandboxHandle, + SandboxPtySession, + SandboxPtySpec, SandboxSpec, SandboxStatus, ) from nemo_gym.sandbox.providers.e2b._sdk import require_e2b_sdk +from nemo_gym.sandbox.providers.e2b.pty import E2BPtySession, validate_pty_spec LOGGER = logging.getLogger(__name__) @@ -472,6 +475,13 @@ async def status(self, handle: SandboxHandle) -> SandboxStatus: return SandboxStatus.UNKNOWN return SandboxStatus.RUNNING if running else SandboxStatus.STOPPED + async def create_pty(self, handle: SandboxHandle, spec: SandboxPtySpec) -> SandboxPtySession: + """Open an interactive terminal through the SDK's ``sandbox.pty``.""" + validate_pty_spec(spec) + e2b = _require_e2b_sdk() + sandbox = self._sandbox(handle) + return await E2BPtySession(sandbox, spec, pty_size=e2b.PtySize).start() + async def close(self, handle: SandboxHandle) -> None: e2b = _require_e2b_sdk() sandbox = handle.raw diff --git a/nemo_gym/sandbox/providers/e2b/pty.py b/nemo_gym/sandbox/providers/e2b/pty.py new file mode 100644 index 00000000000..941b1e2354f --- /dev/null +++ b/nemo_gym/sandbox/providers/e2b/pty.py @@ -0,0 +1,186 @@ +# 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. + +"""Interactive sessions over the E2B SDK's ``sandbox.pty``.""" + +import asyncio +import contextlib +import shlex +import uuid +from typing import Any + +from nemo_gym.sandbox.providers.base import SandboxPtyError, SandboxPtySpec + + +_EOF = object() +_KILL_SIGNALS = {"SIGTERM", "SIGKILL"} + + +class E2BPtySession: + """One client-owned E2B PTY process. + + E2B always starts ``/bin/bash -i -l`` for a PTY, so the image needs bash and + a ``SandboxPtySpec.command`` is run by typing ``exec bash -c `` into + that shell (it replaces the shell, so its exit code is the session's). + """ + + mode = "pty" + + def __init__(self, sandbox: Any, spec: SandboxPtySpec, *, pty_size: Any) -> None: + self.session_id = uuid.uuid4().hex + self._sandbox = sandbox + self._spec = spec + self._pty_size = pty_size + self._rows = spec.rows + self._cols = spec.cols + self._output: asyncio.Queue[bytes | object] = asyncio.Queue() + self._exit: asyncio.Future[int] = asyncio.get_running_loop().create_future() + self._closed = False + self._handle: Any | None = None + self._waiter: asyncio.Task[None] | None = None + + @property + def closed(self) -> bool: + return self._closed + + async def start(self) -> "E2BPtySession": + spec = self._spec + envs = {str(k): str(v) for k, v in (spec.env or {}).items()} + kwargs: dict[str, Any] = {} + if spec.user is not None: + kwargs["user"] = str(spec.user) + # timeout=0 lifts the SDK's default 60 s bound on the running stream. + self._handle = await self._sandbox.pty.create( + self._pty_size(rows=self._rows, cols=self._cols), + self._on_data, + cwd=spec.cwd, + envs=envs or None, + timeout=0, + **kwargs, + ) + self._waiter = asyncio.create_task(self._wait_for_exit()) + if spec.command is not None: + # One quoted unit: a bare `exec a; b` would replace the shell with + # `a` and never run `b`, losing the command's own exit status. + await self.write(f"exec bash -c {shlex.quote(spec.command)}\n".encode()) + return self + + def _on_data(self, data: bytes) -> None: + self._output.put_nowait(bytes(data)) + + async def _wait_for_exit(self) -> None: + try: + try: + result = await self._handle.wait() + code = result.exit_code + except Exception as exc: # noqa: BLE001 - CommandExitException carries the exit code + code = getattr(exc, "exit_code", None) + if code is None: + if not self._exit.done(): + self._exit.set_exception(SandboxPtyError(f"E2B PTY stream failed: {exc}")) + return + if not self._exit.done(): + self._exit.set_result(code) + finally: + self._output.put_nowait(_EOF) + + async def read(self, *, timeout_s: float | None = None) -> bytes: + if timeout_s is None: + value = await self._output.get() + else: + async with asyncio.timeout(timeout_s): + value = await self._output.get() + if value is _EOF: + # Keep the sentinel so later reads and other iterators also see EOF. + self._output.put_nowait(_EOF) + if self._exit.done() and self._exit.exception() is not None: + raise self._exit.exception() # type: ignore[misc] + return b"" + return value # type: ignore[return-value] + + async def read_stderr(self, *, timeout_s: float | None = None) -> bytes: + # PTY mode merges stderr into the terminal stream: block until the + # process exits, then report EOF. + with contextlib.suppress(SandboxPtyError): + await self.wait_exit(timeout_s=timeout_s) + return b"" + + def __aiter__(self) -> "E2BPtySession": + return self + + async def __anext__(self) -> bytes: + chunk = await self.read() + if not chunk: + raise StopAsyncIteration + return chunk + + async def write(self, data: bytes) -> None: + if self._closed: + raise SandboxPtyError("E2B PTY session is closed") + await self._sandbox.pty.send_stdin(self._handle.pid, data) + + async def resize(self, rows: int, cols: int) -> None: + if self._closed: + raise SandboxPtyError("E2B PTY session is closed") + if (rows, cols) == (self._rows, self._cols): + return + await self._sandbox.pty.resize(self._handle.pid, self._pty_size(rows=rows, cols=cols)) + self._rows, self._cols = rows, cols + + async def send_signal(self, signal: str) -> None: + name = signal.upper() + if name == "SIGINT": + await self.write(b"\x03") + elif name in _KILL_SIGNALS: + # E2B's PTY kill is SIGKILL; there is no gentler signal to send. + await self._sandbox.pty.kill(self._handle.pid) + else: + raise NotImplementedError(f"E2B PTY cannot send {signal}") + + async def wait_exit(self, *, timeout_s: float | None = None) -> int: + if timeout_s is None: + return await asyncio.shield(self._exit) + async with asyncio.timeout(timeout_s): + return await asyncio.shield(self._exit) + + async def run_detached(self, command: str, *, poll_interval_s: float = 15.0) -> tuple[bytes, int | None]: + del command, poll_interval_s + raise NotImplementedError("E2B PTY does not support detached PTY execution") + + async def close(self) -> None: + if self._closed: + return + self._closed = True + if self._handle is not None and not self._exit.done(): + with contextlib.suppress(Exception): + await self._sandbox.pty.kill(self._handle.pid) + if self._waiter is not None: + with contextlib.suppress(asyncio.CancelledError, Exception): + await asyncio.wait_for(self._waiter, timeout=10) + if not self._exit.done(): + self._exit.cancel() + + async def __aenter__(self) -> "E2BPtySession": + return self + + async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + await self.close() + + +def validate_pty_spec(spec: SandboxPtySpec) -> None: + """Reject spec fields the E2B SDK cannot honor, rather than dropping them.""" + if not spec.pty: + raise ValueError("The e2b provider only supports PTY mode; pipe mode (pty=False) is not available") + if spec.command is not None and not spec.command.strip(): + raise ValueError("SandboxPtySpec.command must not be blank") diff --git a/nemo_gym/sandbox/shell.py b/nemo_gym/sandbox/shell.py index 97dc18be345..8b49adf8570 100644 --- a/nemo_gym/sandbox/shell.py +++ b/nemo_gym/sandbox/shell.py @@ -196,6 +196,11 @@ def resize() -> None: def main(argv: list[str] | None = None) -> int: args = _parse_args(argv) + # Providers whose HTTP goes through Gym's global aiohttp client (e2b, + # opensandbox) lazily load the global config from sys.argv with Hydra, where + # our "--config" is an ambiguous prefix of Hydra's own options. The flags are + # already consumed above, so hide them. + sys.argv[1:] = [] if not sys.stdin.isatty() or not sys.stdout.isatty(): print("sandbox shell requires an interactive stdin and stdout", file=sys.stderr) return 2 diff --git a/tests/unit_tests/test_e2b_pty.py b/tests/unit_tests/test_e2b_pty.py new file mode 100644 index 00000000000..dd7d68e1922 --- /dev/null +++ b/tests/unit_tests/test_e2b_pty.py @@ -0,0 +1,188 @@ +# 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 + +"""Unit tests for the e2b PTY session (SDK faked; no network).""" + +import asyncio +from dataclasses import dataclass + +import pytest + +from nemo_gym.sandbox.providers.base import SandboxHandle, SandboxPtyError, SandboxPtySpec, SupportsSandboxPty +from nemo_gym.sandbox.providers.e2b import provider as e2b_provider +from nemo_gym.sandbox.providers.e2b.provider import E2BProvider + + +@dataclass +class FakePtySize: + rows: int + cols: int + + +class FakeExit(Exception): + def __init__(self, exit_code: int) -> None: + super().__init__(f"exit {exit_code}") + self.exit_code = exit_code + + +class FakeResult: + def __init__(self, exit_code: int = 0) -> None: + self.exit_code = exit_code + + +class FakePtyHandle: + pid = 4242 + + def __init__(self) -> None: + self.exited = asyncio.Event() + self.exit_code = 0 + + async def wait(self) -> FakeResult: + await self.exited.wait() + if self.exit_code: + raise FakeExit(self.exit_code) + return FakeResult(0) + + +class FakePty: + def __init__(self) -> None: + self.handle = FakePtyHandle() + self.on_data = None + self.create_kwargs: dict = {} + self.stdin: list[bytes] = [] + self.sizes: list[FakePtySize] = [] + self.killed = 0 + + async def create(self, size, on_data, **kwargs): + self.create_kwargs = {"size": size, **kwargs} + self.on_data = on_data + return self.handle + + async def send_stdin(self, pid: int, data: bytes) -> None: + assert pid == self.handle.pid + self.stdin.append(data) + + async def resize(self, pid: int, size) -> None: + self.sizes.append(size) + + async def kill(self, pid: int) -> bool: + self.killed += 1 + self.handle.exited.set() + return True + + +class FakeSandbox: + sandbox_id = "sbx-1" + + def __init__(self) -> None: + self.pty = FakePty() + + +@pytest.fixture +def provider_and_handle(monkeypatch: pytest.MonkeyPatch): + sdk = type("SDK", (), {"PtySize": FakePtySize}) + monkeypatch.setattr(e2b_provider, "_require_e2b_sdk", lambda: sdk) + sandbox = FakeSandbox() + return E2BProvider(), SandboxHandle(sandbox_id="sbx-1", provider_name="e2b", raw=sandbox), sandbox + + +def test_provider_advertises_pty_support() -> None: + assert isinstance(E2BProvider(), SupportsSandboxPty) + + +async def test_session_streams_output_and_input(provider_and_handle) -> None: + provider, handle, sandbox = provider_and_handle + session = await provider.create_pty(handle, SandboxPtySpec(rows=30, cols=100, cwd="/work", env={"A": "1"})) + + assert session.mode == "pty" + kwargs = sandbox.pty.create_kwargs + assert (kwargs["size"].rows, kwargs["size"].cols) == (30, 100) + assert kwargs["cwd"] == "/work" + assert kwargs["envs"] == {"A": "1"} + # The SDK's default 60 s bound on the running stream would kill long sessions. + assert kwargs["timeout"] == 0 + + sandbox.pty.on_data(b"hello\r\n") + assert await session.read(timeout_s=1) == b"hello\r\n" + + await session.write(b"ls\n") + assert sandbox.pty.stdin == [b"ls\n"] + + await session.resize(50, 132) + await session.resize(50, 132) # unchanged size is not re-sent + assert [(s.rows, s.cols) for s in sandbox.pty.sizes] == [(50, 132)] + + sandbox.pty.handle.exit_code = 7 + sandbox.pty.handle.exited.set() + assert await session.wait_exit(timeout_s=1) == 7 + assert await session.read(timeout_s=1) == b"" + assert [chunk async for chunk in session] == [] + await session.close() + + +async def test_clean_exit_reports_zero(provider_and_handle) -> None: + provider, handle, sandbox = provider_and_handle + session = await provider.create_pty(handle, SandboxPtySpec()) + sandbox.pty.handle.exited.set() + assert await session.wait_exit(timeout_s=1) == 0 + await session.close() + + +async def test_command_runs_as_one_quoted_unit(provider_and_handle) -> None: + provider, handle, sandbox = provider_and_handle + session = await provider.create_pty(handle, SandboxPtySpec(command="echo it's; exit 3")) + # `exec a; b` alone would replace the shell with `a` and never run `b`. + assert sandbox.pty.stdin == [b"exec bash -c 'echo it'\"'\"'s; exit 3'\n"] + await session.close() + + +async def test_signals(provider_and_handle) -> None: + provider, handle, sandbox = provider_and_handle + session = await provider.create_pty(handle, SandboxPtySpec()) + + await session.send_signal("SIGINT") + assert sandbox.pty.stdin == [b"\x03"] + with pytest.raises(NotImplementedError): + await session.send_signal("SIGHUP") + await session.send_signal("SIGKILL") + assert sandbox.pty.killed == 1 + await session.close() + + +async def test_close_kills_a_live_session_and_blocks_further_use(provider_and_handle) -> None: + provider, handle, sandbox = provider_and_handle + session = await provider.create_pty(handle, SandboxPtySpec()) + await session.close() + await session.close() # idempotent + + assert session.closed + assert sandbox.pty.killed == 1 + with pytest.raises(SandboxPtyError): + await session.write(b"x") + with pytest.raises(SandboxPtyError): + await session.resize(1, 1) + + +async def test_pipe_mode_is_rejected(provider_and_handle) -> None: + provider, handle, _ = provider_and_handle + with pytest.raises(ValueError, match="pipe mode"): + await provider.create_pty(handle, SandboxPtySpec(pty=False)) + + +async def test_detached_runs_are_unsupported(provider_and_handle) -> None: + provider, handle, _ = provider_and_handle + session = await provider.create_pty(handle, SandboxPtySpec()) + with pytest.raises(NotImplementedError): + await session.run_detached("true") + await session.close()