spaces 0.52.0.dev2__tar.gz → 0.52.0.dev3__tar.gz

{spaces-0.52.0.dev2 → spaces-0.52.0.dev3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: spaces
-Version: 0.52.0.dev2
+Version: 0.52.0.dev3
 Summary: Utilities for Hugging Face Spaces
 Author: Charles Bensimon
 Author-email: Charles Bensimon <charles@huggingface.co>

{spaces-0.52.0.dev2 → spaces-0.52.0.dev3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "spaces"
-version = "0.52.0.dev2"
+version = "0.52.0.dev3"
 description = "Utilities for Hugging Face Spaces"
 readme = "README.md"
 requires-python = ">=3.10,<3.15"

{spaces-0.52.0.dev2 → spaces-0.52.0.dev3}/spaces/zero/__init__.py

@@ -6,6 +6,7 @@ from importlib import metadata
 from packaging.version import Version
 
 from ..config import Config
+from . import context
 from .context import ZeroGPUError
 from .context import ZeroGPUErrorResponse
 from .context import ZeroGPUException
@@ -26,6 +27,7 @@ def _maybe_enable_gradio_compat():
         if Version(gradio_version) < Version('7.0.0.dev0'):
             from . import gradio
             gradio.patch()
+            context.set_context_getter(gradio.get_zerogpu_context)
 
 
 if Config.zero_gpu:
@@ -33,6 +35,10 @@ if Config.zero_gpu:
     _maybe_enable_gradio_compat()
 
 
+# Compat
+startup = zerogpu_startup
+
+
 __all__ = [
     'ZeroGPUError',
     'ZeroGPUErrorResponse',

{spaces-0.52.0.dev2 → spaces-0.52.0.dev3}/spaces/zero/client.py

@@ -11,6 +11,7 @@ from typing_extensions import assert_never
 
 from ..config import Config
 from . import config
+from . import context
 from . import default_gpu_size
 from . import utils
 from .api import APIClient
@@ -18,7 +19,6 @@ from .api import GPUSize
 from .api import AuthLevel
 from .api import QuotaInfos
 from .api import ScheduleResponse
-from .context import get_context
 from .context import ZeroGPUContext
 
 
@@ -104,7 +104,7 @@ def schedule(
     _first_attempt: bool = True,
 ) -> ScheduleResponse:
 
-    ctx = get_context()
+    ctx = context.get()
     token, payload = _get_token_and_payload(ctx.headers)
     if token is not None and (token_error := payload.get('error')):
         message = f"Falling back to IP-based quotas ({token_error})"
@@ -287,7 +287,7 @@ def release(
     allow_404: bool = False,
 ) -> None:
 
-    ctx = get_context()
+    ctx = context.get()
     res = api_client().release(
         allow_token=allow_token,
         fail=fail,
@@ -327,7 +327,7 @@ def _get_token_and_payload(headers: dict[str, str]) -> tuple[str | None, dict[st
 
 
 def duration_warning():
-    get_context().warning(
+    context.get().warning(
         title="ZeroGPU duration",
         message="GPU task is exceeding its requested duration and might be aborted",
     )

{spaces-0.52.0.dev2 → spaces-0.52.0.dev3}/spaces/zero/context.py

@@ -1,6 +1,7 @@
 """
 """
 
+from contextlib import contextmanager
 from contextvars import ContextVar
 from typing import Callable
 from typing import Literal
@@ -119,14 +120,26 @@ class ZeroGPUContext:
         return False
 
 
-ctx_var: ContextVar[ZeroGPUContext] = ContextVar('zerogpu_context', default=ZeroGPUContext())
+_context_var: ContextVar[ZeroGPUContext] = ContextVar('zerogpu_context', default=ZeroGPUContext())
+_context_getter: Callable[[], ZeroGPUContext | None] | None = None
 
 
-_context_resolver: Callable[[], ZeroGPUContext | None] | None = None
-
-def get_context() -> ZeroGPUContext:
-    if _context_resolver is not None:
-        ctx = _context_resolver()
-        if ctx is not None:
+def get():
+    if _context_getter is not None:
+        if (ctx := _context_getter()) is not None:
             return ctx
-    return ctx_var.get()
+    return _context_var.get()
+
+
+@contextmanager
+def set_context_var(value: ZeroGPUContext):
+    token = _context_var.set(value)
+    try:
+        yield
+    finally:
+        _context_var.reset(token)
+
+
+def set_context_getter(context_getter: Callable[[], ZeroGPUContext | None]):
+    global _context_getter
+    _context_getter = context_getter

{spaces-0.52.0.dev2 → spaces-0.52.0.dev3}/spaces/zero/gradio.py

@@ -18,7 +18,6 @@ from gradio.routes import App
 
 from ..utils import create_thread_daemon
 from ..utils import SimpleQueue
-from . import context
 from .context import ZeroGPUContext
 from .context import ZeroGPUErrorResponse
 from .context import ZeroGPULog
@@ -39,6 +38,55 @@ orig_queue_methods: dict[str, Any] = {
 }
 
 
+try:
+    Success = gr.Success # pyright: ignore[reportAttributeAccessIssue] (Gradio<5.10)
+except AttributeError: # pragma: no cover
+    Success = gr.Info
+
+
+def info(title: str, message: str, level: Literal['success', 'info', 'warning'] = 'info'): # pragma: no cover
+    params = inspect.signature(gr.Info).parameters
+    kwargs: dict[str, Any] = {}
+    if 'title' in params:
+        kwargs = {**kwargs, 'title': title}
+    match level:
+        case 'info':
+            return gr.Info(message, **kwargs)
+        case 'success':
+            return Success(message, **kwargs)
+        case 'warning':
+            return gr.Warning(message, **kwargs)
+
+
+def error(title: str, message: str): # pragma: no cover
+    params = inspect.signature(gr.Error).parameters
+    kwargs: dict[str, Any] = {}
+    if 'title' in params:
+        kwargs = {**kwargs, 'title': title}
+    if 'print_exception' in params:
+        kwargs = {**kwargs, 'print_exception': False}
+    return gr.Error(message, **kwargs)
+
+
+def log_emitter(log: ZeroGPULog):
+    return info(
+        title=log['title'],
+        message=log['message'],
+        level=log['level'],
+    )
+
+
+def exception_mapper(err: ZeroGPUErrorResponse, exc: Exception | None):
+    if isinstance(exc, gr.Error):
+        setattr(exc, 'print_exception', False)
+        return exc
+    detail = err['detail']
+    return error(
+        title=detail['title'],
+        message=detail['message'],
+    )
+
+
 class WorkerContextualizer:
     def __init__(self):
         self.event_id = LocalContext.event_id.get(None)
@@ -61,75 +109,33 @@ def disconnect_detector(): # pragma: no cover
     return False
 
 
-def patch():
-
-    try:
-        Success = gr.Success # pyright: ignore[reportAttributeAccessIssue] (Gradio<5.10)
-    except AttributeError: # pragma: no cover
-        Success = gr.Info
-
-    def error(title: str, message: str): # pragma: no cover
-        params = inspect.signature(gr.Error).parameters
-        kwargs: dict[str, Any] = {}
-        if 'title' in params:
-            kwargs = {**kwargs, 'title': title}
-        if 'print_exception' in params:
-            kwargs = {**kwargs, 'print_exception': False}
-        return gr.Error(message, **kwargs)
-
-    def info(title: str, message: str, level: Literal['success', 'info', 'warning'] = 'info'): # pragma: no cover
-        params = inspect.signature(gr.Info).parameters
-        kwargs: dict[str, Any] = {}
-        if 'title' in params:
-            kwargs = {**kwargs, 'title': title}
-        match level:
-            case 'info':
-                return gr.Info(message, **kwargs)
-            case 'success':
-                return Success(message, **kwargs)
-            case 'warning':
-                return gr.Warning(message, **kwargs)
-
-    def exception_mapper(err: ZeroGPUErrorResponse, exc: Exception | None):
-        if isinstance(exc, gr.Error):
-            setattr(exc, 'print_exception', False)
-            return exc
-        detail = err['detail']
-        return error(
-            title=detail['title'],
-            message=detail['message'],
+def setup_middleware(app: FastAPI, *, skip_zerogpu_startup: bool = False):
+    mounted = any(middleware.cls is ZeroGPUMiddleware for middleware in app.user_middleware)
+    if not mounted:
+        app.add_middleware(
+            ZeroGPUMiddleware,
+            exception_mapper=exception_mapper,
+            log_emitter=log_emitter,
+            worker_contextualizer=WorkerContextualizer,
+            disconnect_detector=disconnect_detector,
+            skip_zerogpu_startup=skip_zerogpu_startup,
         )
 
-    def log_emitter(log: ZeroGPULog):
-        return info(
-            title=log['title'],
-            message=log['message'],
-            level=log['level'],
+
+def get_zerogpu_context():
+    if zerogpu_mode_enabled() and (request := LocalContext.request.get(None)) is not None:
+        scope = cast(ASGIScope | None, request.request.scope if request.request is not None else None)
+        return ZeroGPUContext(
+            exception_mapper=exception_mapper,
+            log_emitter=log_emitter,
+            worker_contextualizer=WorkerContextualizer,
+            disconnect_detector=disconnect_detector,
+            server_port=get_server_port(scope) if scope is not None else None,
+            headers={key.lower(): value for key, value in dict(getattr(request, 'headers', {})).items()},
         )
 
-    def resolve_context():
-        if zerogpu_mode_enabled() and (request := LocalContext.request.get(None)) is not None:
-            scope = cast(ASGIScope | None, request.request.scope if request.request is not None else None)
-            return ZeroGPUContext(
-                exception_mapper=exception_mapper,
-                log_emitter=log_emitter,
-                worker_contextualizer=WorkerContextualizer,
-                disconnect_detector=disconnect_detector,
-                server_port=get_server_port(scope) if scope is not None else None,
-                headers={key.lower(): value for key, value in dict(getattr(request, 'headers', {})).items()},
-            )
-
-    def setup_middleware(app: FastAPI, *, skip_zerogpu_startup: bool = False):
-        mounted = any(middleware.cls is ZeroGPUMiddleware for middleware in app.user_middleware)
-        if not mounted:
-            app.add_middleware(
-                ZeroGPUMiddleware,
-                exception_mapper=exception_mapper,
-                log_emitter=log_emitter,
-                worker_contextualizer=WorkerContextualizer,
-                disconnect_detector=disconnect_detector,
-                skip_zerogpu_startup=skip_zerogpu_startup,
-            )
+
+def patch():
 
     @wraps(orig_create_app)
     def create_app(*args, **kwargs):
@@ -137,12 +143,14 @@ def patch():
         setup_middleware(app, skip_zerogpu_startup=True)
         return app
 
+
     @wraps(orig_mount_gradio_app)
     def mount_gradio_app(*args, **kwargs):
         app = orig_mount_gradio_app(*args, **kwargs)
         setup_middleware(app)
         return app
 
+
     @wraps(orig_start_server)
     def start_server(*args, **kwargs):
         if zerogpu_mode_enabled():
@@ -153,7 +161,6 @@ def patch():
     gr.mount_gradio_app = mount_gradio_app
     routes.mount_gradio_app = mount_gradio_app
     http_server.start_server = start_server
-    context._context_resolver = resolve_context
     _patch_queue()
 
 

{spaces-0.52.0.dev2 → spaces-0.52.0.dev3}/spaces/zero/middleware.py

@@ -15,7 +15,7 @@ from .._vendor.asgiref.typing import ASGIReceiveEvent
 from .._vendor.asgiref.typing import ASGISendCallable
 from .._vendor.asgiref.typing import Scope as ASGIScope
 
-from .context import ctx_var
+from . import context
 from .context import ConnectionState
 from .context import ZeroGPUErrorResponse
 from .context import ZeroGPUContext
@@ -81,7 +81,7 @@ class ZeroGPUMiddleware:
                 connection_state.alive = False
             return message
 
-        token = ctx_var.set(ZeroGPUContext(
+        with context.set_context_var(ZeroGPUContext(
             exception_mapper=self.exception_mapper,
             log_emitter=self.log_emitter,
             worker_contextualizer=self.worker_contextualizer,
@@ -89,12 +89,8 @@ class ZeroGPUMiddleware:
             connection_state=connection_state,
             server_port=get_server_port(asgi_scope),
             headers=_decode_headers(asgi_scope.get('headers', [])),
-        ))
-
-        try:
+        )):
             await asgi_app(asgi_scope, wrapped_receive, asgi_send)
-        finally:
-            ctx_var.reset(token)
 
 
 def get_server_port(scope: ASGIScope):

{spaces-0.52.0.dev2 → spaces-0.52.0.dev3}/spaces/zero/wrappers.py

@@ -29,12 +29,12 @@ from ..utils import debug
 from ..utils import drop_params
 from ..utils import SimpleQueue as Queue
 from . import client
+from . import context
 from . import torch
 from .api import AllowToken
 from .api import GPUSize
 from .api import NvidiaIndex
 from .api import NvidiaUUID
-from .context import get_context
 from .tqdm import remove_tqdm_multiprocessing_lock
 from .tqdm import tqdm
 from .utils import apply_cleanups
@@ -170,7 +170,7 @@ def regular_function_wrapper(
         if forked:
             return task(*args, **kwargs)
 
-        ctx = get_context()
+        ctx = context.get()
         duration_ = static_duration(duration, *args, **kwargs)
         duration_ = process_duration(duration_)
         schedule_response = client.schedule(task_id=task_id, duration=duration_, gpu_size=gpu_size)
@@ -297,7 +297,7 @@ def generator_function_wrapper(
             yield from task(*args, **kwargs)
             return
 
-        ctx = get_context()
+        ctx = context.get()
         duration_ = static_duration(duration, *args, **kwargs)
         duration_ = process_duration(duration_)
         schedule_response = client.schedule(task_id=task_id, duration=duration_, gpu_size=gpu_size)