Skip to content

Commit 7ef96d3

Browse files
committed
feat(hf): add dedicated sandbox runtime mode
1 parent 2dcab80 commit 7ef96d3

5 files changed

Lines changed: 351 additions & 27 deletions

File tree

docs/source/guides/runtime-providers.md

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ is a one-line change.
1414
| `UVProvider` | Local process via `uv` (no container) | core ||
1515
| `DaytonaProvider` | Daytona cloud sandboxes | `pip install openenv[daytona]` ||
1616
| `ACASandboxProvider` | Azure Container Apps Sandboxes | `pip install openenv[aca]` ||
17+
| `HFSandboxProvider` | Hugging Face Sandboxes | `pip install openenv[hf-sandbox]` ||
1718
| `ModalProvider` | Modal sandboxes | `pip install openenv[modal]` ||
1819
| `KubernetesProvider` | Kubernetes cluster | core | 🚧 planned |
1920

@@ -159,6 +160,26 @@ from openenv.core.containers.runtime import DockerSwarmProvider
159160
provider = DockerSwarmProvider()
160161
```
161162

163+
### HFSandboxProvider
164+
165+
Runs a registry image in a Hugging Face Sandbox. Pooled mode is the default and
166+
is appropriate for trusted same-user rollout fan-out. Dedicated mode allocates
167+
one VM per sandbox and is suited to untrusted or GPU workloads.
168+
169+
```python
170+
from openenv.core.containers.runtime.hf_sandbox_provider import HFSandboxProvider
171+
172+
provider = HFSandboxProvider(
173+
image="hf.co/spaces/openenv/coding_env",
174+
mode="dedicated",
175+
)
176+
base_url = provider.start_container()
177+
provider.wait_for_ready(base_url)
178+
```
179+
180+
Plain `env_vars` are visible as Job environment metadata. Use the dedicated
181+
mode's `secrets=` argument for encrypted runtime secrets.
182+
162183
### KubernetesProvider
163184

164185
🚧 Not yet implemented. The class exists as a placeholder for the planned

docs/source/reference/core.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -240,4 +240,6 @@ For a high-level explanation of how MCP-backed environments move through `step()
240240

241241
[[autodoc]] openenv.core.containers.runtime.aca_provider.ACASandboxProvider
242242

243+
[[autodoc]] openenv.core.containers.runtime.hf_sandbox_provider.HFSandboxProvider
244+
243245
[[autodoc]] openenv.core.containers.runtime.modal_provider.ModalProvider

pyproject.toml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,9 @@ modal = [
5555
inspect = [
5656
"inspect-ai>=0.3.0",
5757
]
58+
hf-sandbox = [
59+
"huggingface_hub>=1.22.0",
60+
]
5861

5962
[project.scripts]
6063
openenv = "openenv.cli.__main__:main"

src/openenv/core/containers/runtime/hf_sandbox_provider.py

Lines changed: 135 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
import threading
1515
import time
1616
from contextlib import suppress
17-
from typing import Any
17+
from typing import Any, Literal
1818

1919
import requests
2020
import uvicorn
@@ -27,12 +27,18 @@
2727

2828

2929
_DEFAULT_PORT = 8000
30-
_SERVER_COMMAND = (
30+
_POOLED_SERVER_COMMAND = (
3131
'export SBX_PROXY_DIR="${SBX_PROXY_DIR:-$HOME/.sbx/proxy}"; '
3232
'mkdir -p "$SBX_PROXY_DIR"; '
33-
'nohup server >"$HOME/openenv-server.log" 2>&1 &'
33+
'exec server >"$HOME/openenv-server.log" 2>&1'
3434
)
35+
_DEDICATED_SERVER_COMMAND = (
36+
'unset SBX_PROXY_DIR; exec server >"$HOME/openenv-server.log" 2>&1'
37+
)
38+
# Compatibility alias for callers that imported the old private constant.
39+
_SERVER_COMMAND = _POOLED_SERVER_COMMAND
3540
_MAX_WS_MESSAGE_SIZE = 100 * 1024 * 1024
41+
_MAX_HTTP_RESPONSE_SIZE = 100 * 1024 * 1024
3642
_HOP_BY_HOP_HEADERS = {
3743
"connection",
3844
"content-encoding",
@@ -50,6 +56,11 @@
5056
_POOL_LOCK = threading.Lock()
5157

5258

59+
class _NoRedirectWebSocketConnect(ws_connect):
60+
def process_redirect(self, exc: Exception) -> Exception | str:
61+
return exc
62+
63+
5364
def _find_available_port() -> int:
5465
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
5566
sock.bind(("127.0.0.1", 0))
@@ -81,6 +92,22 @@ def _get_sandbox_pool_cls() -> Any:
8192
return SandboxPool
8293

8394

95+
def _get_sandbox_cls() -> Any:
96+
try:
97+
from huggingface_hub import Sandbox
98+
except ImportError as exc:
99+
raise RuntimeError(
100+
"Dedicated HFSandboxProvider execution requires huggingface_hub>=1.22.0. "
101+
"Install it with `pip install 'openenv[hf-sandbox]'`."
102+
) from exc
103+
if not hasattr(Sandbox, "create"):
104+
raise RuntimeError(
105+
"Dedicated HFSandboxProvider execution requires huggingface_hub>=1.22.0. "
106+
"Install it with `pip install 'openenv[hf-sandbox]'`."
107+
)
108+
return Sandbox
109+
110+
84111
def _get_pool(image: str, flavor: str) -> Any:
85112
key = (image, flavor)
86113
with _POOL_LOCK:
@@ -133,23 +160,40 @@ async def proxy_http(path: str, request: Request) -> Response:
133160
data=body,
134161
headers=headers,
135162
timeout=60.0,
136-
allow_redirects=True,
163+
allow_redirects=False,
164+
stream=True,
137165
)
138166
except requests.RequestException:
139167
return Response(
140168
content=b"upstream HF job unreachable",
141169
status_code=502,
142170
)
143-
response_headers = {
144-
key: value
145-
for key, value in upstream.headers.items()
146-
if key.lower() not in _HOP_BY_HOP_HEADERS
147-
}
148-
return Response(
149-
content=upstream.content,
150-
status_code=upstream.status_code,
151-
headers=response_headers,
152-
)
171+
try:
172+
if 300 <= upstream.status_code < 400:
173+
return Response(
174+
content=b"upstream HF job redirects are not allowed",
175+
status_code=502,
176+
)
177+
content = upstream.raw.read(
178+
_MAX_HTTP_RESPONSE_SIZE + 1, decode_content=True
179+
)
180+
if len(content) > _MAX_HTTP_RESPONSE_SIZE:
181+
return Response(
182+
content=b"upstream HF job response exceeded proxy limit",
183+
status_code=502,
184+
)
185+
response_headers = {
186+
key: value
187+
for key, value in upstream.headers.items()
188+
if key.lower() not in _HOP_BY_HOP_HEADERS
189+
}
190+
return Response(
191+
content=content,
192+
status_code=upstream.status_code,
193+
headers=response_headers,
194+
)
195+
finally:
196+
upstream.close()
153197

154198
@app.websocket("/{path:path}")
155199
async def proxy_websocket(path: str, websocket: WebSocket) -> None:
@@ -202,7 +246,7 @@ async def _connect_upstream_websocket(
202246
last_error: Exception | None = None
203247
for _ in range(5):
204248
try:
205-
return await ws_connect(
249+
return await _NoRedirectWebSocketConnect(
206250
target,
207251
additional_headers=headers,
208252
max_size=_MAX_WS_MESSAGE_SIZE,
@@ -243,14 +287,59 @@ def __init__(
243287
*,
244288
image: str,
245289
env_vars: dict[str, str] | None = None,
290+
secrets: dict[str, str] | None = None,
246291
flavor: str = "cpu-basic",
292+
mode: Literal["pooled", "dedicated"] = "pooled",
247293
):
294+
if mode not in {"pooled", "dedicated"}:
295+
raise ValueError("HFSandboxProvider mode must be 'pooled' or 'dedicated'")
296+
if mode == "pooled" and secrets:
297+
raise ValueError("Encrypted secrets require dedicated HF sandbox mode")
248298
self.image = image
249-
self.env_vars = env_vars
299+
self._env_vars = dict(env_vars) if env_vars is not None else None
300+
self._secrets = dict(secrets) if secrets is not None else None
250301
self.flavor = flavor
302+
self._mode = mode
251303
self._sandbox: Any = None
304+
self._server_process: Any = None
252305
self._proxy: _LocalAuthProxy | None = None
253306

307+
@property
308+
def mode(self) -> Literal["pooled", "dedicated"]:
309+
return self._mode
310+
311+
@property
312+
def env_vars(self) -> dict[str, str] | None:
313+
return dict(self._env_vars) if self._env_vars is not None else None
314+
315+
@property
316+
def secrets(self) -> dict[str, str] | None:
317+
return dict(self._secrets) if self._secrets is not None else None
318+
319+
@property
320+
def isolation_mode(self) -> Literal["pooled", "dedicated", "unknown"]:
321+
if self._sandbox is None:
322+
return self._mode
323+
try:
324+
host_id = self._sandbox.host_id
325+
except (AttributeError, RuntimeError):
326+
return "unknown"
327+
return "dedicated" if host_id is None else "pooled"
328+
329+
def _create_sandbox(self, *, image: str, env_vars: dict[str, str] | None) -> Any:
330+
if self.mode == "dedicated":
331+
Sandbox = _get_sandbox_cls()
332+
create_kwargs: dict[str, Any] = {
333+
"image": image,
334+
"flavor": self.flavor,
335+
"env": env_vars,
336+
}
337+
if self._secrets is not None:
338+
create_kwargs["secrets"] = dict(self._secrets)
339+
return Sandbox.create(**create_kwargs)
340+
pool = _get_pool(image, self.flavor)
341+
return pool.create(env=env_vars)
342+
254343
def start_container(
255344
self,
256345
image: str | None = None,
@@ -272,9 +361,10 @@ def start_container(
272361
)
273362

274363
effective_image = self.image if image is None else image
275-
effective_env = self.env_vars if env_vars is None else env_vars
276-
pool = _get_pool(effective_image, self.flavor)
277-
self._sandbox = pool.create(env=effective_env)
364+
effective_env = self._env_vars if env_vars is None else env_vars
365+
self._sandbox = self._create_sandbox(
366+
image=effective_image, env_vars=effective_env
367+
)
278368
try:
279369
if not hasattr(self._sandbox, "proxy_url_for") or not hasattr(
280370
self._sandbox, "proxy_headers"
@@ -283,35 +373,53 @@ def start_container(
283373
"HFSandboxProvider requires a huggingface_hub version with "
284374
"Sandbox.proxy_url_for and Sandbox.proxy_headers support."
285375
)
286-
self._sandbox.run(
287-
_SERVER_COMMAND,
376+
self._server_process = self._sandbox.run(
377+
(
378+
_DEDICATED_SERVER_COMMAND
379+
if self.mode == "dedicated"
380+
else _POOLED_SERVER_COMMAND
381+
),
288382
shell=True,
383+
background=True,
289384
)
290385
self._proxy = _LocalAuthProxy(
291386
target_url=self._sandbox.proxy_url_for(_DEFAULT_PORT, "/"),
292387
headers=self._sandbox.proxy_headers,
293388
)
294389
return self._proxy.start()
295390
except Exception:
296-
self.stop_container()
391+
with suppress(Exception):
392+
self.stop_container()
297393
raise
298394

299395
def stop_container(self) -> None:
396+
proxy_error: Exception | None = None
300397
if self._proxy is not None:
301-
self._proxy.stop()
302-
self._proxy = None
398+
try:
399+
self._proxy.stop()
400+
except Exception as exc:
401+
proxy_error = exc
402+
finally:
403+
self._proxy = None
303404
if self._sandbox is not None:
304405
sandbox = self._sandbox
406+
sandbox.kill()
407+
if getattr(sandbox, "_killed", None) is not True:
408+
raise RuntimeError(
409+
"Hugging Face did not confirm termination of the sandbox; "
410+
"the provider retained its handle so cleanup can be retried."
411+
)
305412
self._sandbox = None
306-
with suppress(Exception):
307-
sandbox.kill()
413+
self._server_process = None
414+
if proxy_error is not None:
415+
raise proxy_error
308416

309417
def wait_for_ready(self, base_url: str, timeout_s: float = 120.0) -> None:
310418
deadline = time.time() + timeout_s
311419
health_url = f"{base_url}/health"
312420
while time.time() < deadline:
313421
try:
314-
response = requests.get(health_url, timeout=5.0)
422+
response = requests.get(health_url, timeout=5.0, allow_redirects=False)
315423
if response.status_code == 200:
316424
return
317425
except requests.exceptions.RequestException:

0 commit comments

Comments
 (0)