wandb 0.19.7__py3-none-musllinux_1_2_aarch64.whl → 0.19.8__py3-none-musllinux_1_2_aarch64.whl

wandb/sdk/mailbox/mailbox.py

@@ -6,8 +6,10 @@ import string
 import threading
 
 from wandb.proto import wandb_internal_pb2 as pb
+from wandb.proto import wandb_server_pb2 as spb
 
-from . import handles
+from .mailbox_handle import MailboxHandle
+from .response_handle import MailboxResponseHandle
 
 _logger = logging.getLogger(__name__)
 
@@ -19,22 +21,25 @@ class MailboxClosedError(Exception):
 class Mailbox:
     """Matches service responses to requests.
 
-    The mailbox can set an address on a Record and create a handle for
+    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, handles.MailboxHandle] = {}
+        self._handles: dict[str, MailboxResponseHandle] = {}
         self._handles_lock = threading.Lock()
         self._closed = False
 
-    def require_response(self, request: pb.Record) -> handles.MailboxHandle:
+    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 mailbox slot.
+            request: The request on which to set a request ID or mailbox slot.
                 This is mutated. An address must not already be set.
 
         Returns:
@@ -45,17 +50,24 @@ class Mailbox:
                 no new responses are expected to be delivered and new handles
                 cannot be created.
         """
-        if address := request.control.mailbox_slot:
-            raise ValueError(f"Request already has an address ({address})")
+        if isinstance(request, spb.ServerRequest):
+            if address := request.request_id:
+                raise ValueError(f"Request already has an address ({address})")
 
-        address = self._new_address()
-        request.control.mailbox_slot = 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 = handles.MailboxHandle(address)
+            handle = MailboxResponseHandle(address)
             self._handles[address] = handle
 
         return handle
@@ -80,18 +92,20 @@ class Mailbox:
 
         return address
 
-    def deliver(self, result: pb.Result) -> None:
+    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 = result.control.mailbox_slot
+        address = response.request_id
         if not address:
-            _logger.error(
-                "Received response with no mailbox slot."
-                f" Kind: {result.WhichOneof('result_type')}"
-            )
+            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:
@@ -102,7 +116,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(result)
+            handle.deliver(response)
 
     def close(self) -> None:
         """Indicate no further responses will be delivered.

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/{handles.py → response_handle.py}

@@ -2,22 +2,26 @@ from __future__ import annotations
 
 import asyncio
 import math
+import sys
 import threading
 from typing import TYPE_CHECKING
 
-from wandb.proto import wandb_internal_pb2 as pb
+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
 
-
-class HandleAbandonedError(Exception):
-    """The handle has no result and has been abandoned."""
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override
 
 
-class MailboxHandle:
-    """A thread-safe handle that allows waiting for a response to a request."""
+class MailboxResponseHandle(MailboxHandle[spb.ServerResponse]):
+    """A general handle for any ServerResponse."""
 
     def __init__(self, address: str) -> None:
         self._address = address
@@ -25,11 +29,11 @@ class MailboxHandle:
         self._event = threading.Event()
 
         self._abandoned = False
-        self._result: pb.Result | None = None
+        self._response: spb.ServerResponse | None = None
 
         self._asyncio_events: dict[asyncio.Event, _AsyncioEvent] = dict()
 
-    def deliver(self, result: pb.Result) -> None:
+    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
@@ -39,34 +43,27 @@ class MailboxHandle:
             if self._abandoned:
                 return
 
-            if self._result:
+            if self._response:
                 raise ValueError(
                     f"A response has already been delivered to {self._address}."
                 )
 
-            self._result = result
+            self._response = response
             self._signal_done()
 
+    @override
     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:
-            interface: The interface on which to publish the cancel request.
-        """
         iface.publish_cancel(self._address)
         self.abandon()
 
+    @override
     def abandon(self) -> None:
-        """Abandon the handle, indicating it will not receive a response."""
         with self._lock:
             self._abandoned = True
             self._signal_done()
 
     def _signal_done(self) -> None:
-        """Indicate that the handle either got a result or became abandoned.
+        """Indicate that the handle either got a response or became abandoned.
 
         The lock must be held.
         """
@@ -78,29 +75,13 @@ class MailboxHandle:
             asyncio_event.set_threadsafe()
         self._asyncio_events.clear()
 
-    def check(self) -> pb.Result | None:
-        """Returns the result if it's ready."""
+    @override
+    def check(self) -> spb.ServerResponse | None:
         with self._lock:
-            return self._result
-
-    def wait_or(self, *, timeout: float | None) -> pb.Result:
-        """Wait for a response or a timeout.
-
-        This is called `wait_or` because it replaces a method called `wait`
-        with different semantics.
-
-        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 result is available.
+            return self._response
 
-        Returns:
-            The result if it arrives before the timeout or has already arrived.
-
-        Raises:
-            TimeoutError: If the timeout is reached.
-            HandleAbandonedError: If the handle becomes abandoned.
-        """
+    @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.")
 
@@ -110,27 +91,14 @@ class MailboxHandle:
             )
 
         with self._lock:
-            if self._result:
-                return self._result
+            if self._response:
+                return self._response
 
             assert self._abandoned
             raise HandleAbandonedError()
 
-    async def wait_async(self, *, timeout: float | None) -> pb.Result:
-        """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 result if it arrives before the timeout or has already arrived.
-
-        Raises:
-            TimeoutError: If the timeout is reached.
-            HandleAbandonedError: If the handle becomes abandoned.
-        """
+    @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.")
 
@@ -142,8 +110,8 @@ class MailboxHandle:
 
         except (asyncio.TimeoutError, TimeoutError) as e:
             with self._lock:
-                if self._result:
-                    return self._result
+                if self._response:
+                    return self._response
                 elif self._abandoned:
                     raise HandleAbandonedError()
                 else:
@@ -153,8 +121,8 @@ class MailboxHandle:
 
         else:
             with self._lock:
-                if self._result:
-                    return self._result
+                if self._response:
+                    return self._response
 
                 assert self._abandoned
                 raise HandleAbandonedError()
@@ -167,20 +135,20 @@ class MailboxHandle:
         loop: asyncio.AbstractEventLoop,
         event: asyncio.Event,
     ) -> None:
-        """Add an event to signal when a result is received.
+        """Add an event to signal when a response is received.
 
-        If a result already exists, this notifies the event loop immediately.
+        If a response already exists, this notifies the event loop immediately.
         """
         asyncio_event = _AsyncioEvent(loop, event)
 
         with self._lock:
-            if self._result or self._abandoned:
+            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 result is received."""
+        """Cancel signalling an event when a response is received."""
         with self._lock:
             self._asyncio_events.pop(event, None)
 

wandb/sdk/mailbox/wait_with_progress.py

@@ -1,21 +1,22 @@
 from __future__ import annotations
 
 import time
-from typing import Any, Callable, Coroutine, List, cast
+from typing import Any, Callable, Coroutine, List, TypeVar, cast
 
-from wandb.proto import wandb_internal_pb2 as pb
 from wandb.sdk.lib import asyncio_compat
 
-from . import handles
+from .mailbox_handle import MailboxHandle
+
+_T = TypeVar("_T")
 
 
 def wait_with_progress(
-    handle: handles.MailboxHandle,
+    handle: MailboxHandle[_T],
     *,
     timeout: float | None,
     progress_after: float,
     display_progress: Callable[[], Coroutine[Any, Any, None]],
-) -> pb.Result:
+) -> _T:
     """Wait for a handle, possibly displaying progress to the user.
 
     Equivalent to passing a single handle to `wait_all_with_progress`.
@@ -29,12 +30,12 @@ def wait_with_progress(
 
 
 def wait_all_with_progress(
-    handle_list: list[handles.MailboxHandle],
+    handle_list: list[MailboxHandle[_T]],
     *,
     timeout: float | None,
     progress_after: float,
     display_progress: Callable[[], Coroutine[Any, Any, None]],
-) -> list[pb.Result]:
+) -> list[_T]:
     """Wait for multiple handles, possibly displaying progress to the user.
 
     Args:
@@ -70,7 +71,7 @@ def wait_all_with_progress(
     except TimeoutError:
         pass
 
-    async def progress_loop_with_timeout() -> list[pb.Result]:
+    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
@@ -87,10 +88,10 @@ def wait_all_with_progress(
 
 
 def _wait_handles(
-    handle_list: list[handles.MailboxHandle],
+    handle_list: list[MailboxHandle[_T]],
     *,
     timeout: float,
-) -> list[pb.Result]:
+) -> list[_T]:
     """Wait for multiple mailbox handles.
 
     Returns:
@@ -100,7 +101,7 @@ def _wait_handles(
         TimeoutError: If the overall timeout expires.
         HandleAbandonedError: If any handle becomes abandoned.
     """
-    results: list[pb.Result] = []
+    results: list[_T] = []
 
     start_time = time.monotonic()
     for handle in handle_list:
@@ -112,15 +113,15 @@ def _wait_handles(
 
 
 async def _wait_handles_async(
-    handle_list: list[handles.MailboxHandle],
+    handle_list: list[MailboxHandle[_T]],
     *,
     timeout: float | None,
-) -> list[pb.Result]:
+) -> list[_T]:
     """Asynchronously wait for multiple mailbox handles.
 
     Just like _wait_handles.
     """
-    results: list[pb.Result | None] = [None for _ in handle_list]
+    results: list[_T | None] = [None for _ in handle_list]
 
     async def wait_single(index: int) -> None:
         handle = handle_list[index]
@@ -131,4 +132,4 @@ async def _wait_handles_async(
             task_group.start_soon(wait_single(index))
 
     # NOTE: `list` is not subscriptable until Python 3.10, so we use List.
-    return cast(List[pb.Result], results)
+    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)

wandb/sdk/service/streams.py

@@ -21,16 +21,14 @@ import psutil
 import wandb
 import wandb.util
 from wandb.proto import wandb_internal_pb2 as pb
+from wandb.sdk.interface.interface_relay import InterfaceRelay
+from wandb.sdk.interface.router_relay import MessageRelayRouter
 from wandb.sdk.internal.settings_static import SettingsStatic
 from wandb.sdk.lib import asyncio_compat, progress
 from wandb.sdk.lib import printer as printerlib
 from wandb.sdk.mailbox import Mailbox, MailboxHandle, wait_all_with_progress
 from wandb.sdk.wandb_run import Run
 
-from ..interface.interface_relay import InterfaceRelay
-
-# from wandb.sdk.wandb_settings import Settings
-
 
 class StreamThread(threading.Thread):
     """Class to running internal process as a thread."""
@@ -62,6 +60,12 @@ class StreamRecord:
         self._record_q = queue.Queue()
         self._result_q = queue.Queue()
         self._relay_q = queue.Queue()
+        self._router = MessageRelayRouter(
+            request_queue=self._record_q,
+            response_queue=self._result_q,
+            relay_queue=self._relay_q,
+            mailbox=self._mailbox,
+        )
         self._iface = InterfaceRelay(
             record_q=self._record_q,
             result_q=self._result_q,
@@ -80,6 +84,7 @@ class StreamRecord:
 
     def join(self) -> None:
         self._iface.join()
+        self._router.join()
         if self._thread:
             self._thread.join()
 
@@ -290,7 +295,7 @@ class StreamMux:
 
         # fixme: for now we have a single printer for all streams,
         # and jupyter is disabled if at least single stream's setting set `_jupyter` to false
-        exit_handles: list[MailboxHandle] = []
+        exit_handles: list[MailboxHandle[pb.Result]] = []
 
         # only finish started streams, non started streams failed early
         started_streams: dict[str, StreamRecord] = {}

wandb/sdk/wandb_init.py

@@ -42,7 +42,7 @@ from . import wandb_login, wandb_setup
 from .backend.backend import Backend
 from .lib import SummaryDisabled, filesystem, module, paths, printer, telemetry
 from .lib.deprecate import Deprecated, deprecate
-from .mailbox import Mailbox, wait_with_progress
+from .mailbox import wait_with_progress
 from .wandb_helper import parse_config
 from .wandb_run import Run, TeardownHook, TeardownStage
 from .wandb_settings import Settings
@@ -682,11 +682,11 @@ class _WandbInit:
         drun._Run__metadata = wandb.sdk.wandb_metadata.Metadata()
 
         # methods
-        drun.log = lambda data, *_, **__: drun.summary.update(data)  # type: ignore
-        drun.finish = lambda *_, **__: module.unset_globals()  # type: ignore
-        drun.join = drun.finish  # type: ignore
-        drun.define_metric = lambda *_, **__: wandb.sdk.wandb_metric.Metric("dummy")  # type: ignore
-        drun.save = lambda *_, **__: False  # type: ignore
+        drun.log = lambda data, *_, **__: drun.summary.update(data)  # type: ignore[method-assign]
+        drun.finish = lambda *_, **__: module.unset_globals()  # type: ignore[method-assign]
+        drun.join = drun.finish  # type: ignore[method-assign]
+        drun.define_metric = lambda *_, **__: wandb.sdk.wandb_metric.Metric("dummy")  # type: ignore[method-assign]
+        drun.save = lambda *_, **__: False  # type: ignore[method-assign]
         for symbol in (
             "alert",
             "finish_artifact",
@@ -733,7 +733,7 @@ class _WandbInit:
             def __call__(self, *args: Any, **kwargs: Any) -> _ChainableNoOp:
                 return _ChainableNoOp()
 
-        drun.log_artifact = _ChainableNoOpField()
+        drun.log_artifact = _ChainableNoOpField()  # type: ignore[method-assign]
         # attributes
         drun._start_time = time.time()
         drun._starting_step = 0
@@ -793,12 +793,7 @@ class _WandbInit:
         else:
             service = None
 
-        mailbox = Mailbox()
-        backend = Backend(
-            settings=settings,
-            service=service,
-            mailbox=mailbox,
-        )
+        backend = Backend(settings=settings, service=service)
         backend.ensure_launched()
         self._logger.info("backend started and connected")
 
@@ -1085,8 +1080,7 @@ def _attach(
     )
 
     # TODO: consolidate this codepath with wandb.init()
-    mailbox = Mailbox()
-    backend = Backend(settings=settings, service=service, mailbox=mailbox)
+    backend = Backend(settings=settings, service=service)
     backend.ensure_launched()
     logger.info("attach backend started and connected")
 
@@ -1468,6 +1462,9 @@ def init(  # noqa: C901
             _monkeypatch_tensorboard()
             init_telemetry.feature.tensorboard_sync = True
 
+        if run_settings.x_server_side_derived_summary:
+            init_telemetry.feature.server_side_derived_summary = True
+
         return wi.init(run_settings, run_config)
 
     except KeyboardInterrupt as e: