wandb 0.19.6rc4__py3-none-any.whl → 0.19.8__py3-none-any.whl

wandb/sdk/mailbox/__init__.py

@@ -0,0 +1,23 @@
+"""A message protocol for the internal service process.
+
+The core of W&B is implemented by a side process that asynchronously uploads
+data. The client process (such as this Python code) sends requests to the
+service, and for some requests, the service eventually sends a response.
+
+The client can send multiple requests before the service provides a response.
+The Mailbox handles matching responses to requests. An internal thread
+continuously reads data from the service and passes it to the mailbox.
+"""
+
+from .mailbox import Mailbox, MailboxClosedError
+from .mailbox_handle import HandleAbandonedError, MailboxHandle
+from .wait_with_progress import wait_all_with_progress, wait_with_progress
+
+__all__ = [
+    "Mailbox",
+    "MailboxClosedError",
+    "HandleAbandonedError",
+    "MailboxHandle",
+    "wait_all_with_progress",
+    "wait_with_progress",
+]

wandb/sdk/mailbox/mailbox.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+import logging
+import secrets
+import string
+import threading
+
+from wandb.proto import wandb_internal_pb2 as pb
+from wandb.proto import wandb_server_pb2 as spb
+
+from .mailbox_handle import MailboxHandle
+from .response_handle import MailboxResponseHandle
+
+_logger = logging.getLogger(__name__)
+
+
+class MailboxClosedError(Exception):
+    """The mailbox has been closed and cannot be used."""
+
+
+class Mailbox:
+    """Matches service responses to requests.
+
+    The mailbox can set an address on a server request and create a handle for
+    waiting for a response to that record. Responses are delivered by calling
+    `deliver()`. The `close()` method abandons all handles in case the
+    service process becomes unreachable.
+    """
+
+    def __init__(self) -> None:
+        self._handles: dict[str, MailboxResponseHandle] = {}
+        self._handles_lock = threading.Lock()
+        self._closed = False
+
+    def require_response(
+        self,
+        request: spb.ServerRequest | pb.Record,
+    ) -> MailboxHandle[spb.ServerResponse]:
+        """Set a response address on a request.
+
+        Args:
+            request: The request on which to set a request ID or mailbox slot.
+                This is mutated. An address must not already be set.
+
+        Returns:
+            A handle for waiting for the response to the request.
+
+        Raises:
+            MailboxClosedError: If the mailbox has been closed, in which case
+                no new responses are expected to be delivered and new handles
+                cannot be created.
+        """
+        if isinstance(request, spb.ServerRequest):
+            if address := request.request_id:
+                raise ValueError(f"Request already has an address ({address})")
+
+            address = self._new_address()
+            request.request_id = address
+        else:
+            if address := request.control.mailbox_slot:
+                raise ValueError(f"Request already has an address ({address})")
+
+            address = self._new_address()
+            request.control.mailbox_slot = address
+
+        with self._handles_lock:
+            if self._closed:
+                raise MailboxClosedError()
+
+            handle = MailboxResponseHandle(address)
+            self._handles[address] = handle
+
+        return handle
+
+    def _new_address(self) -> str:
+        """Returns an unused address for a request.
+
+        Assumes `_handles_lock` is held.
+        """
+
+        def generate():
+            return "".join(
+                secrets.choice(string.ascii_lowercase + string.digits)
+                for i in range(12)
+            )
+
+        address = generate()
+
+        # Being extra cautious. This loop will almost never be entered.
+        while address in self._handles:
+            address = generate()
+
+        return address
+
+    def deliver(self, response: spb.ServerResponse) -> None:
+        """Deliver a response from the service.
+
+        If the response address is invalid, this does nothing.
+        It is a no-op if the mailbox has been closed.
+        """
+        address = response.request_id
+        if not address:
+            kind: str | None = response.WhichOneof("server_response_type")
+            if kind == "result_communicate":
+                result_type = response.result_communicate.WhichOneof("result_type")
+                kind = f"result_communicate.{result_type}"
+
+            _logger.error(f"Received response with no mailbox slot: {kind}")
+            return
+
+        with self._handles_lock:
+            # NOTE: If the mailbox is closed, this returns None because
+            # we clear the dict.
+            handle = self._handles.pop(address, None)
+
+        # 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)
+
+    def close(self) -> None:
+        """Indicate no further responses will be delivered.
+
+        Abandons all handles.
+        """
+        with self._handles_lock:
+            self._closed = True
+
+            _logger.info(
+                f"Closing mailbox, abandoning {len(self._handles)} handles.",
+            )
+
+            for handle in self._handles.values():
+                handle.abandon()
+            self._handles.clear()

wandb/sdk/mailbox/mailbox_handle.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+import abc
+import sys
+from typing import TYPE_CHECKING, Callable, Generic, TypeVar
+
+# Necessary to break an import loop.
+if TYPE_CHECKING:
+    from wandb.sdk.interface import interface
+
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override
+
+
+_T = TypeVar("_T")
+_S = TypeVar("_S")
+
+
+class HandleAbandonedError(Exception):
+    """The handle has no response and has been abandoned."""
+
+
+class MailboxHandle(abc.ABC, Generic[_T]):
+    """A thread-safe handle that allows waiting for a response to a request."""
+
+    def map(self, fn: Callable[[_T], _S]) -> MailboxHandle[_S]:
+        """Returns a transformed handle.
+
+        Methods on the returned handle call methods on this handle, but the
+        response type is derived using the given function.
+
+        Args:
+            fn: A function to apply to this handle's result to get the new
+                handle's result. The function should be pure and fast.
+        """
+        return _MailboxMappedHandle(self, fn)
+
+    @abc.abstractmethod
+    def abandon(self) -> None:
+        """Abandon the handle, indicating it will not receive a response."""
+
+    @abc.abstractmethod
+    def cancel(self, iface: interface.InterfaceBase) -> None:
+        """Cancel the handle, requesting any associated work to not complete.
+
+        This automatically abandons the handle, as a response is no longer
+        guaranteed.
+
+        Args:
+            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.
+
+        Args:
+            timeout: A finite number of seconds or None to never time out.
+                If less than or equal to zero, times out immediately unless
+                the response is available.
+
+        Returns:
+            The response if it arrives before the timeout or has already arrived.
+
+        Raises:
+            TimeoutError: If the timeout is reached.
+            HandleAbandonedError: If the handle becomes abandoned.
+        """
+
+    @abc.abstractmethod
+    async def wait_async(self, *, timeout: float | None) -> _T:
+        """Wait for a response or timeout.
+
+        This must run in an `asyncio` event loop.
+
+        Args:
+            timeout: A finite number of seconds or None to never time out.
+
+        Returns:
+            The response if it arrives before the timeout or has already arrived.
+
+        Raises:
+            TimeoutError: If the timeout is reached.
+            HandleAbandonedError: If the handle becomes abandoned.
+        """
+
+
+class _MailboxMappedHandle(Generic[_S], MailboxHandle[_S]):
+    """A mailbox handle whose result is derived from another handle."""
+
+    def __init__(
+        self,
+        handle: MailboxHandle[_T],
+        fn: Callable[[_T], _S],
+    ) -> None:
+        self._handle = handle
+        self._fn = fn
+
+    @override
+    def abandon(self) -> None:
+        self._handle.abandon()
+
+    @override
+    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))
+
+    @override
+    async def wait_async(self, *, timeout: float | None) -> _S:
+        response = await self._handle.wait_async(timeout=timeout)
+        return self._fn(response)

wandb/sdk/mailbox/response_handle.py

@@ -0,0 +1,167 @@
+from __future__ import annotations
+
+import asyncio
+import math
+import sys
+import threading
+from typing import TYPE_CHECKING
+
+from wandb.proto import wandb_server_pb2 as spb
+
+from .mailbox_handle import HandleAbandonedError, MailboxHandle
+
+# Necessary to break an import loop.
+if TYPE_CHECKING:
+    from wandb.sdk.interface import interface
+
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override
+
+
+class MailboxResponseHandle(MailboxHandle[spb.ServerResponse]):
+    """A general handle for any ServerResponse."""
+
+    def __init__(self, address: str) -> None:
+        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()
+
+    def deliver(self, response: spb.ServerResponse) -> None:
+        """Deliver the response.
+
+        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}."
+                )
+
+            self._response = response
+            self._signal_done()
+
+    @override
+    def cancel(self, iface: interface.InterfaceBase) -> None:
+        iface.publish_cancel(self._address)
+        self.abandon()
+
+    @override
+    def abandon(self) -> None:
+        with self._lock:
+            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()
+
+        # 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
+
+    @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()
+
+    @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)
+
+        try:
+            await asyncio.wait_for(evt.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
+                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)
+
+
+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)

wandb/sdk/mailbox/wait_with_progress.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+import time
+from typing import Any, Callable, Coroutine, List, TypeVar, cast
+
+from wandb.sdk.lib import asyncio_compat
+
+from .mailbox_handle import MailboxHandle
+
+_T = TypeVar("_T")
+
+
+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.
+
+    Equivalent to passing a single handle to `wait_all_with_progress`.
+    """
+    return wait_all_with_progress(
+        [handle],
+        timeout=timeout,
+        progress_after=progress_after,
+        display_progress=display_progress,
+    )[0]
+
+
+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.
+
+    Args:
+        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.
+
+    Returns:
+        A list where the Nth item is the Nth handle's result.
+
+    Raises:
+        TimeoutError: If the overall timeout expires.
+        HandleAbandonedError: If any handle becomes abandoned.
+        Exception: Any exception from the display function is propagated.
+    """
+    if not handle_list:
+        return []
+
+    if timeout is not None and timeout <= progress_after:
+        return _wait_handles(handle_list, timeout=timeout)
+
+    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:
+                elapsed_time = time.monotonic() - start_time
+                remaining_timeout = timeout - elapsed_time
+            else:
+                remaining_timeout = None
+
+            return await _wait_handles_async(
+                handle_list,
+                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
+
+
+async def _wait_handles_async(
+    handle_list: list[MailboxHandle[_T]],
+    *,
+    timeout: float | None,
+) -> list[_T]:
+    """Asynchronously wait for multiple mailbox handles.
+
+    Just like _wait_handles.
+    """
+    results: list[_T | None] = [None for _ in handle_list]
+
+    async def wait_single(index: int) -> None:
+        handle = handle_list[index]
+        results[index] = await handle.wait_async(timeout=timeout)
+
+    async with asyncio_compat.open_task_group() as task_group:
+        for index in range(len(handle_list)):
+            task_group.start_soon(wait_single(index))
+
+    # NOTE: `list` is not subscriptable until Python 3.10, so we use List.
+    return cast(List[_T], results)

wandb/sdk/service/server_sock.py

@@ -44,7 +44,10 @@ class SockServerInterfaceReaderThread(threading.Thread):
     _stopped: "Event"
 
     def __init__(
-        self, clients: ClientDict, iface: "InterfaceRelay", stopped: "Event"
+        self,
+        clients: ClientDict,
+        iface: "InterfaceRelay",
+        stopped: "Event",
     ) -> None:
         self._iface = iface
         self._clients = clients
@@ -53,7 +56,6 @@ class SockServerInterfaceReaderThread(threading.Thread):
         self._stopped = stopped
 
     def run(self) -> None:
-        assert self._iface.relay_q
         while not self._stopped.is_set():
             try:
                 result = self._iface.relay_q.get(timeout=1)
@@ -70,6 +72,7 @@ class SockServerInterfaceReaderThread(threading.Thread):
             sock_client = self._clients.get_client(sockid)
             assert sock_client
             sresp = spb.ServerResponse()
+            sresp.request_id = result.control.mailbox_slot
             sresp.result_communicate.CopyFrom(result)
             sock_client.send_server_response(sresp)
 
@@ -148,7 +151,10 @@ class SockServerReadThread(threading.Thread):
         inform_attach_response.settings.CopyFrom(
             self._mux._streams[stream_id]._settings._proto,
         )
-        response = spb.ServerResponse(inform_attach_response=inform_attach_response)
+        response = spb.ServerResponse(
+            request_id=sreq.request_id,
+            inform_attach_response=inform_attach_response,
+        )
         self._sock_client.send_server_response(response)
         iface = self._mux.get_stream(stream_id).interface