1414import threading
1515import time
1616from contextlib import suppress
17- from typing import Any
17+ from typing import Any , Literal
1818
1919import requests
2020import uvicorn
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" ,
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+
5364def _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+
84111def _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