wandb 0.21.0__py3-none-win_amd64.whl → 0.21.2__py3-none-win_amd64.whl

wandb/sdk/mailbox/mailbox.py

@@ -7,6 +7,7 @@ import threading
 
 from wandb.proto import wandb_internal_pb2 as pb
 from wandb.proto import wandb_server_pb2 as spb
+from wandb.sdk.lib import asyncio_manager
 
 from .mailbox_handle import MailboxHandle
 from .response_handle import MailboxResponseHandle
@@ -27,7 +28,8 @@ class Mailbox:
     service process becomes unreachable.
     """
 
-    def __init__(self) -> None:
+    def __init__(self, asyncer: asyncio_manager.AsyncioManager) -> None:
+        self._asyncer = asyncer
         self._handles: dict[str, MailboxResponseHandle] = {}
         self._handles_lock = threading.Lock()
         self._closed = False
@@ -51,11 +53,17 @@ class Mailbox:
                 cannot be created.
         """
         if isinstance(request, spb.ServerRequest):
-            if address := request.request_id:
+            if (address := request.request_id) or (
+                address := request.record_publish.control.mailbox_slot
+            ):
                 raise ValueError(f"Request already has an address ({address})")
 
             address = self._new_address()
             request.request_id = address
+            if request.HasField("record_publish"):
+                request.record_publish.control.mailbox_slot = address
+            if request.HasField("record_communicate"):
+                request.record_communicate.control.mailbox_slot = address
         else:
             if address := request.control.mailbox_slot:
                 raise ValueError(f"Request already has an address ({address})")
@@ -67,7 +75,7 @@ class Mailbox:
             if self._closed:
                 raise MailboxClosedError()
 
-            handle = MailboxResponseHandle(address)
+            handle = MailboxResponseHandle(address, asyncer=self._asyncer)
             self._handles[address] = handle
 
         return handle
@@ -92,7 +100,7 @@ class Mailbox:
 
         return address
 
-    def deliver(self, response: spb.ServerResponse) -> None:
+    async def deliver(self, response: spb.ServerResponse) -> None:
         """Deliver a response from the service.
 
         If the response address is invalid, this does nothing.
@@ -116,7 +124,7 @@ class Mailbox:
         # It is not an error if there is no handle for the address:
         # handles can be abandoned if the result is no longer needed.
         if handle:
-            handle.deliver(response)
+            await handle.deliver(response)
 
     def close(self) -> None:
         """Indicate no further responses will be delivered.

wandb/sdk/mailbox/mailbox_handle.py

@@ -5,6 +5,8 @@ from typing import TYPE_CHECKING, Callable, Generic
 
 from typing_extensions import TypeVar, override
 
+from wandb.sdk.lib import asyncio_manager
+
 # Necessary to break an import loop.
 if TYPE_CHECKING:
     from wandb.sdk.interface import interface
@@ -19,7 +21,18 @@ class HandleAbandonedError(Exception):
 
 
 class MailboxHandle(abc.ABC, Generic[_T]):
-    """A thread-safe handle that allows waiting for a response to a request."""
+    """A handle for waiting on a response to a request."""
+
+    def __init__(self, asyncer: asyncio_manager.AsyncioManager) -> None:
+        self._asyncer = asyncer
+
+    @property
+    def asyncer(self) -> asyncio_manager.AsyncioManager:
+        """The asyncio thread to which the handle belongs.
+
+        The handle's async methods must be run using this object.
+        """
+        return self._asyncer
 
     def map(self, fn: Callable[[_T], _S]) -> MailboxHandle[_S]:
         """Returns a transformed handle.
@@ -35,7 +48,11 @@ class MailboxHandle(abc.ABC, Generic[_T]):
 
     @abc.abstractmethod
     def abandon(self) -> None:
-        """Abandon the handle, indicating it will not receive a response."""
+        """Abandon the handle, indicating it will not receive a response.
+
+        This may not happen immediately: it is possible for an existing
+        call to `wait_async` to complete successfully after this method returns.
+        """
 
     @abc.abstractmethod
     def cancel(self, iface: interface.InterfaceBase) -> None:
@@ -48,14 +65,12 @@ class MailboxHandle(abc.ABC, Generic[_T]):
             iface: The interface on which to publish the cancel request.
         """
 
-    @abc.abstractmethod
-    def check(self) -> _T | None:
-        """Returns the response if it's ready."""
-
     @abc.abstractmethod
     def wait_or(self, *, timeout: float | None) -> _T:
         """Wait for a response or a timeout.
 
+        It is an error to call this from an async function.
+
         Args:
             timeout: A finite number of seconds or None to never time out.
                 If less than or equal to zero, times out immediately unless
@@ -95,6 +110,7 @@ class _MailboxMappedHandle(Generic[_S], MailboxHandle[_S]):
         handle: MailboxHandle[_T],
         fn: Callable[[_T], _S],
     ) -> None:
+        super().__init__(handle.asyncer)
         self._handle = handle
         self._fn = fn
 
@@ -106,13 +122,6 @@ class _MailboxMappedHandle(Generic[_S], MailboxHandle[_S]):
     def cancel(self, iface: interface.InterfaceBase) -> None:
         self._handle.cancel(iface)
 
-    @override
-    def check(self) -> _S | None:
-        if response := self._handle.check():
-            return self._fn(response)
-        else:
-            return None
-
     @override
     def wait_or(self, *, timeout: float | None) -> _S:
         return self._fn(self._handle.wait_or(timeout=timeout))

wandb/sdk/mailbox/response_handle.py

@@ -2,12 +2,12 @@ from __future__ import annotations
 
 import asyncio
 import math
-import threading
 from typing import TYPE_CHECKING
 
 from typing_extensions import override
 
 from wandb.proto import wandb_server_pb2 as spb
+from wandb.sdk.lib import asyncio_manager
 
 from .mailbox_handle import HandleAbandonedError, MailboxHandle
 
@@ -19,33 +19,36 @@ if TYPE_CHECKING:
 class MailboxResponseHandle(MailboxHandle[spb.ServerResponse]):
     """A general handle for any ServerResponse."""
 
-    def __init__(self, address: str) -> None:
+    def __init__(
+        self,
+        address: str,
+        *,
+        asyncer: asyncio_manager.AsyncioManager,
+    ) -> None:
+        super().__init__(asyncer)
+
         self._address = address
-        self._lock = threading.Lock()
-        self._event = threading.Event()
 
         self._abandoned = False
         self._response: spb.ServerResponse | None = None
 
-        self._asyncio_events: dict[asyncio.Event, _AsyncioEvent] = dict()
+        # Initialized on first use in the asyncio thread.
+        self._done_event: asyncio.Event | None = None
 
-    def deliver(self, response: spb.ServerResponse) -> None:
-        """Deliver the response.
+    async def deliver(self, response: spb.ServerResponse) -> None:
+        if self._abandoned:
+            return
 
-        This may only be called once. It is an error to respond to the same
-        request more than once. It is a no-op if the handle has been abandoned.
-        """
-        with self._lock:
-            if self._abandoned:
-                return
+        if self._response:
+            raise ValueError(
+                f"A response has already been delivered to {self._address}."
+            )
 
-            if self._response:
-                raise ValueError(
-                    f"A response has already been delivered to {self._address}."
-                )
+        self._response = response
 
-            self._response = response
-            self._signal_done()
+        if not self._done_event:
+            self._done_event = asyncio.Event()
+        self._done_event.set()
 
     @override
     def cancel(self, iface: interface.InterfaceBase) -> None:
@@ -54,110 +57,43 @@ class MailboxResponseHandle(MailboxHandle[spb.ServerResponse]):
 
     @override
     def abandon(self) -> None:
-        with self._lock:
+        async def impl() -> None:
             self._abandoned = True
-            self._signal_done()
-
-    def _signal_done(self) -> None:
-        """Indicate that the handle either got a response or became abandoned.
 
-        The lock must be held.
-        """
-        # Unblock threads blocked on `wait_or`.
-        self._event.set()
+            if not self._done_event:
+                self._done_event = asyncio.Event()
+            self._done_event.set()
 
-        # Unblock asyncio loops blocked on `wait_async`.
-        for asyncio_event in self._asyncio_events.values():
-            asyncio_event.set_threadsafe()
-        self._asyncio_events.clear()
-
-    @override
-    def check(self) -> spb.ServerResponse | None:
-        with self._lock:
-            return self._response
+        self.asyncer.run_soon(impl)
 
     @override
     def wait_or(self, *, timeout: float | None) -> spb.ServerResponse:
-        if timeout is not None and not math.isfinite(timeout):
-            raise ValueError("Timeout must be finite or None.")
-
-        if not self._event.wait(timeout=timeout):
-            raise TimeoutError(
-                f"Timed out waiting for response on {self._address}",
-            )
-
-        with self._lock:
-            if self._response:
-                return self._response
-
-            assert self._abandoned
-            raise HandleAbandonedError()
+        return self.asyncer.run(lambda: self.wait_async(timeout=timeout))
 
     @override
     async def wait_async(self, *, timeout: float | None) -> spb.ServerResponse:
         if timeout is not None and not math.isfinite(timeout):
             raise ValueError("Timeout must be finite or None.")
 
-        evt = asyncio.Event()
-        self._add_asyncio_event(asyncio.get_event_loop(), evt)
+        if not self._done_event:
+            self._done_event = asyncio.Event()
 
         try:
-            await asyncio.wait_for(evt.wait(), timeout=timeout)
+            await asyncio.wait_for(self._done_event.wait(), timeout=timeout)
 
         except (asyncio.TimeoutError, TimeoutError) as e:
-            with self._lock:
-                if self._response:
-                    return self._response
-                elif self._abandoned:
-                    raise HandleAbandonedError()
-                else:
-                    raise TimeoutError(
-                        f"Timed out waiting for response on {self._address}"
-                    ) from e
-
-        else:
-            with self._lock:
-                if self._response:
-                    return self._response
-
-                assert self._abandoned
+            if self._response:
+                return self._response
+            elif self._abandoned:
                 raise HandleAbandonedError()
-
-        finally:
-            self._forget_asyncio_event(evt)
-
-    def _add_asyncio_event(
-        self,
-        loop: asyncio.AbstractEventLoop,
-        event: asyncio.Event,
-    ) -> None:
-        """Add an event to signal when a response is received.
-
-        If a response already exists, this notifies the event loop immediately.
-        """
-        asyncio_event = _AsyncioEvent(loop, event)
-
-        with self._lock:
-            if self._response or self._abandoned:
-                asyncio_event.set_threadsafe()
             else:
-                self._asyncio_events[event] = asyncio_event
-
-    def _forget_asyncio_event(self, event: asyncio.Event) -> None:
-        """Cancel signalling an event when a response is received."""
-        with self._lock:
-            self._asyncio_events.pop(event, None)
+                raise TimeoutError(
+                    f"Timed out waiting for response on {self._address}"
+                ) from e
 
+        else:
+            if self._response:
+                return self._response
 
-class _AsyncioEvent:
-    def __init__(
-        self,
-        loop: asyncio.AbstractEventLoop,
-        event: asyncio.Event,
-    ):
-        self._loop = loop
-        self._event = event
-
-    def set_threadsafe(self) -> None:
-        """Set the asyncio event in its own loop."""
-        self._loop.call_soon_threadsafe(self._event.set)
+            assert self._abandoned
+            raise HandleAbandonedError()

wandb/sdk/mailbox/wait_with_progress.py

@@ -14,7 +14,6 @@ def wait_with_progress(
     handle: MailboxHandle[_T],
     *,
     timeout: float | None,
-    progress_after: float,
     display_progress: Callable[[], Coroutine[Any, Any, None]],
 ) -> _T:
     """Wait for a handle, possibly displaying progress to the user.
@@ -24,7 +23,6 @@ def wait_with_progress(
     return wait_all_with_progress(
         [handle],
         timeout=timeout,
-        progress_after=progress_after,
         display_progress=display_progress,
     )[0]
 
@@ -33,7 +31,6 @@ def wait_all_with_progress(
     handle_list: list[MailboxHandle[_T]],
     *,
     timeout: float | None,
-    progress_after: float,
     display_progress: Callable[[], Coroutine[Any, Any, None]],
 ) -> list[_T]:
     """Wait for multiple handles, possibly displaying progress to the user.
@@ -42,18 +39,14 @@ def wait_all_with_progress(
         handle_list: The handles to wait for.
         timeout: A number of seconds after which to raise a TimeoutError,
             or None if this should never timeout.
-        progress_after: A number of seconds after which to start the
-            display_progress callback. Starting the callback creates a thread
-            and starts an asyncio loop, so we want to avoid doing it if
-            the handle is resolved quickly.
         display_progress: An asyncio function that displays progress to
-            the user. This function is executed on a new thread and cancelled
-            if the timeout is exceeded.
+            the user. This function runs using the handles' AsyncioManager.
 
     Returns:
         A list where the Nth item is the Nth handle's result.
 
     Raises:
+        ValueError: If the handles live in different asyncio threads.
         TimeoutError: If the overall timeout expires.
         HandleAbandonedError: If any handle becomes abandoned.
         Exception: Any exception from the display function is propagated.
@@ -61,16 +54,13 @@ def wait_all_with_progress(
     if not handle_list:
         return []
 
-    if timeout is not None and timeout <= progress_after:
-        return _wait_handles(handle_list, timeout=timeout)
+    asyncer = handle_list[0].asyncer
+    for handle in handle_list:
+        if handle.asyncer is not asyncer:
+            raise ValueError("Handles have different AsyncioManagers.")
 
     start_time = time.monotonic()
 
-    try:
-        return _wait_handles(handle_list, timeout=progress_after)
-    except TimeoutError:
-        pass
-
     async def progress_loop_with_timeout() -> list[_T]:
         with asyncio_compat.cancel_on_exit(display_progress()):
             if timeout is not None:
@@ -84,32 +74,7 @@ def wait_all_with_progress(
                 timeout=remaining_timeout,
             )
 
-    return asyncio_compat.run(progress_loop_with_timeout)
-
-
-def _wait_handles(
-    handle_list: list[MailboxHandle[_T]],
-    *,
-    timeout: float,
-) -> list[_T]:
-    """Wait for multiple mailbox handles.
-
-    Returns:
-        Each handle's result, in the same order as the given handles.
-
-    Raises:
-        TimeoutError: If the overall timeout expires.
-        HandleAbandonedError: If any handle becomes abandoned.
-    """
-    results: list[_T] = []
-
-    start_time = time.monotonic()
-    for handle in handle_list:
-        elapsed_time = time.monotonic() - start_time
-        remaining_timeout = timeout - elapsed_time
-        results.append(handle.wait_or(timeout=remaining_timeout))
-
-    return results
+    return asyncer.run(progress_loop_with_timeout)
 
 
 async def _wait_handles_async(