spaces 0.50.dev2__tar.gz → 0.50.dev3__tar.gz

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

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

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

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

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

@@ -60,6 +60,7 @@ def html_string(html_contents: str, text_contents: str): # pragma: no cover
     return HTMLString(html_contents)
 
 
+# TODO: Remove (make it like ZeroGPU quota exceeded errors)
 def _toast_action(
     auth: AuthLevel | None,
     pro_message: str,
@@ -126,7 +127,7 @@ def schedule(
                 f"is larger than the maximum allowed"
             )
         elif token is None:
-            title = "ZeroGPU quota exceeded"
+            title = "ZeroGPU quota exceeded" # TODO: different title?
             message = (
                 f"Space app has reached its GPU limit. "
                 f"{EXAMPLES_RETRY_MESSAGE}"
@@ -189,11 +190,13 @@ def schedule(
                     raise RuntimeError("Error while re-trying after queue disconnect")
                 return schedule(task_id, duration, _first_attempt=False)
             if event.event == 'ping':
-                if connection_state is not None and not connection_state.alive:
+                print('ping', getattr(connection_state, 'id', None))
+                if ctx.disconnected:
                     res.close()
-                    raise RuntimeError("Connection closed by visitor while queueing")
+                    raise RuntimeError(f"Connection closed by visitor while queueing ({getattr(connection_state, 'id', None)})")
                 continue
             if event.event == 'failed':
+                print('failed', getattr(connection_state, 'id', None))
                 if token is None:
                     message = f"{NO_GPU_MESSAGE_INQUEUE}. {EXAMPLES_RETRY_MESSAGE}"
                 else:
@@ -205,16 +208,18 @@ def schedule(
                         ending="priority in ZeroGPU queues",
                     )
                     message = f"{NO_GPU_MESSAGE_INQUEUE}. {details}"
+                # TODO: API vs GUI message commme pour quota exceeded non ?
                 raise ctx.exception(
                     title="ZeroGPU queue timeout",
                     message=message,
                     status_code=503,
                 )
             if event.event == 'succeeded':
+                print('succeeded', getattr(connection_state, 'id', None))
                 assert event.data is not None
-                if connection_state is not None and not connection_state.alive:
+                if ctx.disconnected:
                     release(event.data.allowToken)
-                    raise RuntimeError("Connection closed by visitor on queue success")
+                    raise RuntimeError(f"Connection closed by visitor on queue success ({getattr(connection_state, 'id', None)})")
                 ctx.success(
                     title="ZeroGPU queue",
                     message="Successfully acquired a GPU",

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

@@ -1,6 +1,7 @@
 """
 """
 
+import secrets
 from contextvars import ContextVar
 from typing import Callable
 from typing import Literal
@@ -54,6 +55,7 @@ class WorkerContextualizerDefault:
 
 class ConnectionState:
     def __init__(self) -> None:
+        self.id = secrets.token_hex(2)
         self.alive = True
 
 
@@ -63,6 +65,7 @@ class ZeroGPUContext:
         exception_mapper: Callable[[ZeroGPUErrorResponse, Exception | None], Exception] | None = None,
         log_emitter: Callable[[ZeroGPULog], None] | None = None,
         worker_contextualizer: Callable[[], Picklable[Callable[[], None]]] | None = None,
+        disconnect_detector: Callable[[], bool] | None = None,
         connection_state: ConnectionState | None = None,
         server_port: int | None = None,
         headers: dict[str, str] | None = None,
@@ -70,6 +73,7 @@ class ZeroGPUContext:
         self.exception_mapper = exception_mapper or exception_mapper_default
         self.log_emitter = log_emitter or log_emitter_default
         self.worker_contextualizer = worker_contextualizer or WorkerContextualizerDefault
+        self.disconnect_detector = disconnect_detector
         self.connection_state = connection_state
         self.server_port = server_port
         self.headers = headers or {}
@@ -110,5 +114,13 @@ class ZeroGPUContext:
         res = {self.server_port} if self.server_port is not None else set()
         return res | set(GlobalContext.server_ports.values())
 
+    @property
+    def disconnected(self):
+        if self.disconnect_detector is not None:
+            return self.disconnect_detector()
+        if self.connection_state is not None:
+            return not self.connection_state.alive
+        return False
+
 
 ctx_var: ContextVar[ZeroGPUContext] = ContextVar('zerogpu_context', default=ZeroGPUContext())

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

@@ -40,11 +40,14 @@ class ZeroGPUMiddleware:
         exception_mapper: Callable[[ZeroGPUErrorResponse, Exception | None], Exception] | None = None,
         log_emitter: Callable[[ZeroGPULog], None] | None = None,
         worker_contextualizer: Callable[[], Picklable[Callable[[], None]]] | None = None,
+        disconnect_detector: Callable[[], bool] | None = None,
     ):
+        print("ZeroGPUMiddleware.__init__")
         self.app = app
         self.exception_mapper = exception_mapper
         self.log_emitter = log_emitter
         self.worker_contextualizer = worker_contextualizer
+        self.disconnect_detector = disconnect_detector
 
     async def __call__(self, scope: Scope, receive: Receive, send: Send):
 
@@ -69,27 +72,33 @@ class ZeroGPUMiddleware:
         connection_state = ConnectionState()
 
         async def wrapped_receive() -> ASGIReceiveEvent:
+            print('wrapped_receive', connection_state.id)
             message = await asgi_receive()
+            print('wrapped_receive', connection_state.id, message['type'])
             if (False
                 or message['type'] == 'http.disconnect'
                 or message['type'] == 'websocket.disconnect'
             ):
+                print("\n==== disconnect =====\n", connection_state.id)
                 connection_state.alive = False
             return message
 
-        ctx_var.set(ZeroGPUContext(
+        token = ctx_var.set(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),
             headers=_decode_headers(asgi_scope.get('headers', [])),
         ))
 
         try:
+            print("\n==== await =====\n", connection_state.id)
             await asgi_app(asgi_scope, wrapped_receive, asgi_send)
         finally:
-            connection_state.alive = False
+            print("\n==== finally =====\n", connection_state.id)
+            ctx_var.reset(token)
 
 
 def _server_port(scope: ASGIScope):