spaces 0.52.0.dev1__tar.gz → 0.52.0.dev3__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: spaces
-Version: 0.52.0.dev1
+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.dev1 → spaces-0.52.0.dev3}/pyproject.toml

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

{spaces-0.52.0.dev1 → spaces-0.52.0.dev3}/spaces/__init__.py

@@ -19,8 +19,8 @@ if (gr := sys.modules.get('gradio')) is not None: # pragma: no cover
 from .zero.decorator import GPU
 
 
-def gradio_auto_wrap(*args, **kwargs): # pragma: no cover
-    warnings.warn("gradio_auto_wrap has no effect and is deprecated")
+def gradio_auto_wrap(task): # pragma: no cover
+    return task
 
 def disable_gradio_auto_wrap(*args, **kwargs): # pragma: no cover
     warnings.warn("gradio_auto_wrap has no effect and is deprecated")

{spaces-0.52.0.dev1 → 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.dev1 → 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 ctx_var
 from .context import ZeroGPUContext
 
 
@@ -104,7 +104,7 @@ def schedule(
     _first_attempt: bool = True,
 ) -> ScheduleResponse:
 
-    ctx = ctx_var.get()
+    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 = ctx_var.get()
+    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():
-    ctx_var.get().warning(
+    context.get().warning(
         title="ZeroGPU duration",
         message="GPU task is exceeding its requested duration and might be aborted",
     )

{spaces-0.52.0.dev1 → 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,4 +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
+
+
+def get():
+    if _context_getter is not None:
+        if (ctx := _context_getter()) is not None:
+            return ctx
+    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.dev1 → spaces-0.52.0.dev3}/spaces/zero/gradio.py

@@ -5,7 +5,7 @@ import os
 from functools import wraps
 from typing import Any
 from typing import Literal
-from typing import Type
+from typing import cast
 
 import gradio as gr
 from fastapi import FastAPI
@@ -18,8 +18,13 @@ from gradio.routes import App
 
 from ..utils import create_thread_daemon
 from ..utils import SimpleQueue
+from .context import ZeroGPUContext
+from .context import ZeroGPUErrorResponse
+from .context import ZeroGPULog
 from .lifecycle import zerogpu_mode_enabled
 from .lifecycle import zerogpu_startup
+from .middleware import ASGIScope
+from .middleware import get_server_port
 from .middleware import ZeroGPUMiddleware
 
 
@@ -33,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)
@@ -55,57 +109,33 @@ def disconnect_detector(): # pragma: no cover
     return False
 
 
-def patch():
+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 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()},
+        )
 
-    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 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=lambda err, exc: (
-                    setattr(exc, 'print_exception', False) or exc
-                    if isinstance(exc, gr.Error)
-                    else error(
-                        title=err['detail']['title'],
-                        message=err['detail']['message'],
-                    )
-                ),
-                log_emitter=lambda log: info(
-                    title=log['title'],
-                    message=log['message'],
-                    level=log['level'],
-                ),
-                worker_contextualizer=WorkerContextualizer,
-                disconnect_detector=disconnect_detector,
-                skip_zerogpu_startup=skip_zerogpu_startup,
-            )
+
+def patch():
 
     @wraps(orig_create_app)
     def create_app(*args, **kwargs):
@@ -113,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():

{spaces-0.52.0.dev1 → 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,23 +81,19 @@ 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,
             disconnect_detector=self.disconnect_detector,
             connection_state=connection_state,
-            server_port=_server_port(asgi_scope),
+            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 _server_port(scope: ASGIScope):
+def get_server_port(scope: ASGIScope):
     match scope.get('server'):
         case (_, int() as port):
             return port

{spaces-0.52.0.dev1 → 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 ctx_var
 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 = ctx_var.get()
+        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 = ctx_var.get()
+        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)