diff --git a/docs/source/guides/runtime-providers.md b/docs/source/guides/runtime-providers.md index 257d719f2..e719eb0bd 100644 --- a/docs/source/guides/runtime-providers.md +++ b/docs/source/guides/runtime-providers.md @@ -14,6 +14,7 @@ is a one-line change. | `UVProvider` | Local process via `uv` (no container) | core | ✅ | | `DaytonaProvider` | Daytona cloud sandboxes | `pip install openenv[daytona]` | ✅ | | `ACASandboxProvider` | Azure Container Apps Sandboxes | `pip install openenv[aca]` | ✅ | +| `HFSandboxProvider` | Hugging Face Sandboxes | `pip install openenv[hf-sandbox]` | ✅ | | `ModalProvider` | Modal sandboxes | `pip install openenv[modal]` | ✅ | | `KubernetesProvider` | Kubernetes cluster | core | 🚧 planned | @@ -159,6 +160,26 @@ from openenv.core.containers.runtime import DockerSwarmProvider provider = DockerSwarmProvider() ``` +### HFSandboxProvider + +Runs a registry image in a Hugging Face Sandbox. Pooled mode is the default and +is appropriate for trusted same-user rollout fan-out. Dedicated mode allocates +one VM per sandbox and is suited to untrusted or GPU workloads. + +```python +from openenv.core.containers.runtime.hf_sandbox_provider import HFSandboxProvider + +provider = HFSandboxProvider( + image="hf.co/spaces/openenv/coding_env", + mode="dedicated", +) +base_url = provider.start_container() +provider.wait_for_ready(base_url) +``` + +Plain `env_vars` are visible as Job environment metadata. Use the dedicated +mode's `secrets=` argument for encrypted runtime secrets. + ### KubernetesProvider 🚧 Not yet implemented. The class exists as a placeholder for the planned diff --git a/docs/source/reference/core.md b/docs/source/reference/core.md index bf9de12cf..1b0c747dd 100644 --- a/docs/source/reference/core.md +++ b/docs/source/reference/core.md @@ -240,4 +240,6 @@ For a high-level explanation of how MCP-backed environments move through `step() [[autodoc]] openenv.core.containers.runtime.aca_provider.ACASandboxProvider +[[autodoc]] openenv.core.containers.runtime.hf_sandbox_provider.HFSandboxProvider + [[autodoc]] openenv.core.containers.runtime.modal_provider.ModalProvider diff --git a/pyproject.toml b/pyproject.toml index 5d7277ae8..8877c3b82 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,9 @@ modal = [ inspect = [ "inspect-ai>=0.3.0", ] +hf-sandbox = [ + "huggingface_hub>=1.22.0", +] [project.scripts] openenv = "openenv.cli.__main__:main" diff --git a/src/openenv/core/containers/runtime/hf_sandbox_provider.py b/src/openenv/core/containers/runtime/hf_sandbox_provider.py index 7010abe87..9c45abd3c 100644 --- a/src/openenv/core/containers/runtime/hf_sandbox_provider.py +++ b/src/openenv/core/containers/runtime/hf_sandbox_provider.py @@ -14,7 +14,7 @@ import threading import time from contextlib import suppress -from typing import Any +from typing import Any, Literal import requests import uvicorn @@ -27,12 +27,18 @@ _DEFAULT_PORT = 8000 -_SERVER_COMMAND = ( +_POOLED_SERVER_COMMAND = ( 'export SBX_PROXY_DIR="${SBX_PROXY_DIR:-$HOME/.sbx/proxy}"; ' 'mkdir -p "$SBX_PROXY_DIR"; ' - 'nohup server >"$HOME/openenv-server.log" 2>&1 &' + 'exec server >"$HOME/openenv-server.log" 2>&1' ) +_DEDICATED_SERVER_COMMAND = ( + 'unset SBX_PROXY_DIR; exec server >"$HOME/openenv-server.log" 2>&1' +) +# Compatibility alias for callers that imported the old private constant. +_SERVER_COMMAND = _POOLED_SERVER_COMMAND _MAX_WS_MESSAGE_SIZE = 100 * 1024 * 1024 +_MAX_HTTP_RESPONSE_SIZE = 100 * 1024 * 1024 _HOP_BY_HOP_HEADERS = { "connection", "content-encoding", @@ -50,6 +56,11 @@ _POOL_LOCK = threading.Lock() +class _NoRedirectWebSocketConnect(ws_connect): + def process_redirect(self, exc: Exception) -> Exception | str: + return exc + + def _find_available_port() -> int: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: sock.bind(("127.0.0.1", 0)) @@ -81,6 +92,22 @@ def _get_sandbox_pool_cls() -> Any: return SandboxPool +def _get_sandbox_cls() -> Any: + try: + from huggingface_hub import Sandbox + except ImportError as exc: + raise RuntimeError( + "Dedicated HFSandboxProvider execution requires huggingface_hub>=1.22.0. " + "Install it with `pip install 'openenv[hf-sandbox]'`." + ) from exc + if not hasattr(Sandbox, "create"): + raise RuntimeError( + "Dedicated HFSandboxProvider execution requires huggingface_hub>=1.22.0. " + "Install it with `pip install 'openenv[hf-sandbox]'`." + ) + return Sandbox + + def _get_pool(image: str, flavor: str) -> Any: key = (image, flavor) with _POOL_LOCK: @@ -133,23 +160,40 @@ async def proxy_http(path: str, request: Request) -> Response: data=body, headers=headers, timeout=60.0, - allow_redirects=True, + allow_redirects=False, + stream=True, ) except requests.RequestException: return Response( content=b"upstream HF job unreachable", status_code=502, ) - response_headers = { - key: value - for key, value in upstream.headers.items() - if key.lower() not in _HOP_BY_HOP_HEADERS - } - return Response( - content=upstream.content, - status_code=upstream.status_code, - headers=response_headers, - ) + try: + if 300 <= upstream.status_code < 400: + return Response( + content=b"upstream HF job redirects are not allowed", + status_code=502, + ) + content = upstream.raw.read( + _MAX_HTTP_RESPONSE_SIZE + 1, decode_content=True + ) + if len(content) > _MAX_HTTP_RESPONSE_SIZE: + return Response( + content=b"upstream HF job response exceeded proxy limit", + status_code=502, + ) + response_headers = { + key: value + for key, value in upstream.headers.items() + if key.lower() not in _HOP_BY_HOP_HEADERS + } + return Response( + content=content, + status_code=upstream.status_code, + headers=response_headers, + ) + finally: + upstream.close() @app.websocket("/{path:path}") async def proxy_websocket(path: str, websocket: WebSocket) -> None: @@ -202,7 +246,7 @@ async def _connect_upstream_websocket( last_error: Exception | None = None for _ in range(5): try: - return await ws_connect( + return await _NoRedirectWebSocketConnect( target, additional_headers=headers, max_size=_MAX_WS_MESSAGE_SIZE, @@ -243,14 +287,59 @@ def __init__( *, image: str, env_vars: dict[str, str] | None = None, + secrets: dict[str, str] | None = None, flavor: str = "cpu-basic", + mode: Literal["pooled", "dedicated"] = "pooled", ): + if mode not in {"pooled", "dedicated"}: + raise ValueError("HFSandboxProvider mode must be 'pooled' or 'dedicated'") + if mode == "pooled" and secrets: + raise ValueError("Encrypted secrets require dedicated HF sandbox mode") self.image = image - self.env_vars = env_vars + self._env_vars = dict(env_vars) if env_vars is not None else None + self._secrets = dict(secrets) if secrets is not None else None self.flavor = flavor + self._mode = mode self._sandbox: Any = None + self._server_process: Any = None self._proxy: _LocalAuthProxy | None = None + @property + def mode(self) -> Literal["pooled", "dedicated"]: + return self._mode + + @property + def env_vars(self) -> dict[str, str] | None: + return dict(self._env_vars) if self._env_vars is not None else None + + @property + def secrets(self) -> dict[str, str] | None: + return dict(self._secrets) if self._secrets is not None else None + + @property + def isolation_mode(self) -> Literal["pooled", "dedicated", "unknown"]: + if self._sandbox is None: + return self._mode + try: + host_id = self._sandbox.host_id + except (AttributeError, RuntimeError): + return "unknown" + return "dedicated" if host_id is None else "pooled" + + def _create_sandbox(self, *, image: str, env_vars: dict[str, str] | None) -> Any: + if self.mode == "dedicated": + Sandbox = _get_sandbox_cls() + create_kwargs: dict[str, Any] = { + "image": image, + "flavor": self.flavor, + "env": env_vars, + } + if self._secrets is not None: + create_kwargs["secrets"] = dict(self._secrets) + return Sandbox.create(**create_kwargs) + pool = _get_pool(image, self.flavor) + return pool.create(env=env_vars) + def start_container( self, image: str | None = None, @@ -272,9 +361,10 @@ def start_container( ) effective_image = self.image if image is None else image - effective_env = self.env_vars if env_vars is None else env_vars - pool = _get_pool(effective_image, self.flavor) - self._sandbox = pool.create(env=effective_env) + effective_env = self._env_vars if env_vars is None else env_vars + self._sandbox = self._create_sandbox( + image=effective_image, env_vars=effective_env + ) try: if not hasattr(self._sandbox, "proxy_url_for") or not hasattr( self._sandbox, "proxy_headers" @@ -283,9 +373,14 @@ def start_container( "HFSandboxProvider requires a huggingface_hub version with " "Sandbox.proxy_url_for and Sandbox.proxy_headers support." ) - self._sandbox.run( - _SERVER_COMMAND, + self._server_process = self._sandbox.run( + ( + _DEDICATED_SERVER_COMMAND + if self.mode == "dedicated" + else _POOLED_SERVER_COMMAND + ), shell=True, + background=True, ) self._proxy = _LocalAuthProxy( target_url=self._sandbox.proxy_url_for(_DEFAULT_PORT, "/"), @@ -293,25 +388,38 @@ def start_container( ) return self._proxy.start() except Exception: - self.stop_container() + with suppress(Exception): + self.stop_container() raise def stop_container(self) -> None: + proxy_error: Exception | None = None if self._proxy is not None: - self._proxy.stop() - self._proxy = None + try: + self._proxy.stop() + except Exception as exc: + proxy_error = exc + finally: + self._proxy = None if self._sandbox is not None: sandbox = self._sandbox + sandbox.kill() + if getattr(sandbox, "_killed", None) is not True: + raise RuntimeError( + "Hugging Face did not confirm termination of the sandbox; " + "the provider retained its handle so cleanup can be retried." + ) self._sandbox = None - with suppress(Exception): - sandbox.kill() + self._server_process = None + if proxy_error is not None: + raise proxy_error def wait_for_ready(self, base_url: str, timeout_s: float = 120.0) -> None: deadline = time.time() + timeout_s health_url = f"{base_url}/health" while time.time() < deadline: try: - response = requests.get(health_url, timeout=5.0) + response = requests.get(health_url, timeout=5.0, allow_redirects=False) if response.status_code == 200: return except requests.exceptions.RequestException: diff --git a/tests/test_core/test_hf_sandbox_provider.py b/tests/test_core/test_hf_sandbox_provider.py new file mode 100644 index 000000000..aa5b69eb7 --- /dev/null +++ b/tests/test_core/test_hf_sandbox_provider.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Unit tests for Hugging Face sandbox isolation modes.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest +from openenv.core.containers.runtime import hf_sandbox_provider as provider_module +from openenv.core.containers.runtime.hf_sandbox_provider import HFSandboxProvider +from starlette.testclient import TestClient + + +def _proxy() -> MagicMock: + proxy = MagicMock() + proxy.start.return_value = "http://127.0.0.1:12345" + return proxy + + +def test_hf_sandbox_provider_defaults_to_pooled_mode() -> None: + sandbox = MagicMock() + sandbox.proxy_url_for.return_value = "https://sandbox.example/proxy" + sandbox.proxy_headers = {"Authorization": "Bearer token"} + pool = MagicMock() + pool.create.return_value = sandbox + proxy = _proxy() + + with patch.object(provider_module, "_get_pool", return_value=pool) as get_pool: + with patch.object( + provider_module, "_LocalAuthProxy", return_value=proxy + ) as proxy_cls: + instance = HFSandboxProvider( + image="hf.co/spaces/openenv/coding_env", env_vars={"MODE": "test"} + ) + base_url = instance.start_container() + + assert instance.mode == "pooled" + assert base_url == "http://127.0.0.1:12345" + get_pool.assert_called_once_with("hf.co/spaces/openenv/coding_env", "cpu-basic") + pool.create.assert_called_once_with(env={"MODE": "test"}) + command = sandbox.run.call_args.args[0] + assert "SBX_PROXY_DIR" in command + assert sandbox.run.call_args.kwargs == {"shell": True, "background": True} + proxy_cls.assert_called_once_with( + target_url="https://sandbox.example/proxy", + headers={"Authorization": "Bearer token"}, + ) + + +def test_hf_sandbox_provider_dedicated_mode_uses_sandbox_create() -> None: + sandbox = MagicMock() + sandbox.proxy_url_for.return_value = "https://sandbox.example/proxy" + sandbox.proxy_headers = {} + sandbox_cls = MagicMock() + sandbox_cls.create.return_value = sandbox + proxy = _proxy() + + with patch.object( + provider_module, "_get_sandbox_cls", return_value=sandbox_cls + ) as get_sandbox_cls: + with patch.object(provider_module, "_get_pool") as get_pool: + with patch.object(provider_module, "_LocalAuthProxy", return_value=proxy): + instance = HFSandboxProvider( + image="hf.co/spaces/openenv/echo", + flavor="a10g-small", + env_vars={"MODE": "certify"}, + mode="dedicated", + ) + instance.start_container() + + get_sandbox_cls.assert_called_once_with() + get_pool.assert_not_called() + sandbox_cls.create.assert_called_once_with( + image="hf.co/spaces/openenv/echo", + flavor="a10g-small", + env={"MODE": "certify"}, + ) + command = sandbox.run.call_args.args[0] + assert "unset SBX_PROXY_DIR" in command + assert sandbox.run.call_args.kwargs == {"shell": True, "background": True} + + +def test_hf_sandbox_provider_rejects_invalid_mode() -> None: + with pytest.raises(ValueError, match="mode"): + HFSandboxProvider(image="example", mode="shared") # type: ignore[arg-type] + + +def test_hf_sandbox_provider_requires_dedicated_mode_for_secrets() -> None: + with pytest.raises(ValueError, match="dedicated"): + HFSandboxProvider(image="example", secrets={"API_KEY": "secret"}) + + +def test_hf_sandbox_provider_uses_encrypted_dedicated_secrets() -> None: + sandbox_cls = MagicMock() + provider = HFSandboxProvider( + image="example", + mode="dedicated", + secrets={"API_KEY": "secret"}, + ) + + with patch.object(provider_module, "_get_sandbox_cls", return_value=sandbox_cls): + provider._create_sandbox(image="example", env_vars={"MODE": "test"}) + + sandbox_cls.create.assert_called_once_with( + image="example", + flavor="cpu-basic", + env={"MODE": "test"}, + secrets={"API_KEY": "secret"}, + ) + + +def test_hf_sandbox_provider_mode_is_read_only() -> None: + provider = HFSandboxProvider(image="example", mode="dedicated") + + with pytest.raises(AttributeError): + provider.mode = "pooled" # type: ignore[misc] + + +def test_stop_retains_sandbox_handle_when_kill_fails() -> None: + provider = HFSandboxProvider(image="example", mode="dedicated") + sandbox = MagicMock() + sandbox.kill.side_effect = RuntimeError("cancel failed") + provider._sandbox = sandbox + + with pytest.raises(RuntimeError, match="cancel failed"): + provider.stop_container() + + assert provider._sandbox is sandbox + + +def test_stop_retains_sandbox_when_sdk_swallows_kill_failure() -> None: + provider = HFSandboxProvider(image="example", mode="dedicated") + sandbox = MagicMock() + sandbox._killed = False + provider._sandbox = sandbox + + with pytest.raises(RuntimeError, match="did not confirm termination"): + provider.stop_container() + + assert provider._sandbox is sandbox + + +def test_stop_fails_closed_without_sdk_termination_state() -> None: + class SandboxWithoutTerminationState: + def kill(self) -> None: + return None + + provider = HFSandboxProvider(image="example", mode="dedicated") + sandbox = SandboxWithoutTerminationState() + provider._sandbox = sandbox + + with pytest.raises(RuntimeError, match="did not confirm termination"): + provider.stop_container() + + assert provider._sandbox is sandbox + + +def test_local_auth_proxy_rejects_upstream_redirects() -> None: + class _ServerWithoutSocket: + def __init__(self, config: object): + self.config = config + self.started = True + self.should_exit = False + + def run(self) -> None: + return None + + upstream = MagicMock() + upstream.status_code = 302 + upstream.headers = {"Location": "http://169.254.169.254/latest/meta-data"} + + with patch.object(provider_module, "_find_available_port", return_value=12345): + with patch.object(provider_module.uvicorn, "Server", _ServerWithoutSocket): + proxy = provider_module._LocalAuthProxy( + target_url="https://sandbox.example/proxy", + headers={"X-Sandbox-Token": "secret"}, + ) + proxy.start() + assert proxy._server is not None + app = proxy._server.config.app + with patch.object( + provider_module.requests, "request", return_value=upstream + ): + response = TestClient(app).get("/health", follow_redirects=False) + proxy.stop() + + assert response.status_code == 502 + assert response.headers.get("Location") is None + upstream.close.assert_called_once_with()