wandb 0.19.7__py3-none-win_amd64.whl → 0.19.9__py3-none-win_amd64.whl

wandb/sdk/lib/sock_client.py

@@ -150,10 +150,10 @@ class SockClient:
         with self._lock:
             self._sendall_with_error_handle(header + data)
 
-    def send_server_request(self, msg: Any) -> None:
+    def send_server_request(self, msg: spb.ServerRequest) -> None:
         self._send_message(msg)
 
-    def send_server_response(self, msg: Any) -> None:
+    def send_server_response(self, msg: spb.ServerResponse) -> None:
         try:
             self._send_message(msg)
         except BrokenPipeError:
@@ -161,56 +161,6 @@ class SockClient:
             #  things like network status poll loop, there might be a better way to quiesce
             pass
 
-    def send_and_recv(
-        self,
-        *,
-        inform_init: Optional[spb.ServerInformInitRequest] = None,
-        inform_start: Optional[spb.ServerInformStartRequest] = None,
-        inform_attach: Optional[spb.ServerInformAttachRequest] = None,
-        inform_finish: Optional[spb.ServerInformFinishRequest] = None,
-        inform_teardown: Optional[spb.ServerInformTeardownRequest] = None,
-    ) -> spb.ServerResponse:
-        self.send(
-            inform_init=inform_init,
-            inform_start=inform_start,
-            inform_attach=inform_attach,
-            inform_finish=inform_finish,
-            inform_teardown=inform_teardown,
-        )
-
-        # HACK: This assumes nothing else is reading on the socket, and that
-        # the next response is for this request.
-        response = self.read_server_response(timeout=1)
-
-        if response is None:
-            raise SockClientTimeoutError("No response after 1 second.")
-
-        return response
-
-    def send(
-        self,
-        *,
-        inform_init: Optional[spb.ServerInformInitRequest] = None,
-        inform_start: Optional[spb.ServerInformStartRequest] = None,
-        inform_attach: Optional[spb.ServerInformAttachRequest] = None,
-        inform_finish: Optional[spb.ServerInformFinishRequest] = None,
-        inform_teardown: Optional[spb.ServerInformTeardownRequest] = None,
-    ) -> None:
-        server_req = spb.ServerRequest()
-        if inform_init:
-            server_req.inform_init.CopyFrom(inform_init)
-        elif inform_start:
-            server_req.inform_start.CopyFrom(inform_start)
-        elif inform_attach:
-            server_req.inform_attach.CopyFrom(inform_attach)
-        elif inform_finish:
-            server_req.inform_finish.CopyFrom(inform_finish)
-        elif inform_teardown:
-            server_req.inform_teardown.CopyFrom(inform_teardown)
-        else:
-            raise Exception("unmatched")
-        self.send_server_request(server_req)
-
     def send_record_communicate(self, record: "pb.Record") -> None:
         server_req = spb.ServerRequest()
         server_req.request_id = record.control.mailbox_slot

wandb/sdk/lib/wb_logging.py

@@ -0,0 +1,161 @@
+"""Logging configuration for the "wandb" logger.
+
+Most log statements in wandb are made in the context of a run and should be
+redirected to that run's log file (usually named 'debug.log'). This module
+provides a context manager to temporarily set the current run ID and registers
+a global handler for the 'wandb' logger that sends log statements to the right
+place.
+
+All functions in this module are threadsafe.
+
+NOTE: The pytest caplog fixture will fail to capture logs from the wandb logger
+because they are not propagated to the root logger.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import contextvars
+import logging
+import pathlib
+from typing import Iterator
+
+
+class _NotRunSpecific:
+    """Sentinel for `not_run_specific()`."""
+
+
+_NOT_RUN_SPECIFIC = _NotRunSpecific()
+
+
+_run_id: contextvars.ContextVar[str | _NotRunSpecific | None] = contextvars.ContextVar(
+    "_run_id",
+    default=None,
+)
+
+_logger = logging.getLogger("wandb")
+
+
+def configure_wandb_logger() -> None:
+    """Configures the global 'wandb' logger.
+
+    The wandb logger is not intended to be customized by users. Instead, it is
+    used as a mechanism to redirect log messages into wandb run-specific log
+    files.
+
+    This function is idempotent: calling it multiple times has the same effect.
+    """
+    # Send all DEBUG and above messages to registered handlers.
+    #
+    # Per-run handlers can set different levels.
+    _logger.setLevel(logging.DEBUG)
+
+    # Do not propagate wandb logs to the root logger, which the user may have
+    # configured to point elsewhere. All wandb log messages should go to a run's
+    # log file.
+    _logger.propagate = False
+
+    # If no handlers are configured for the 'wandb' logger, don't activate the
+    # "lastResort" handler which sends messages to stderr with a level of
+    # WARNING by default.
+    #
+    # This occurs in wandb code that runs outside the context of a Run and
+    # not as part of the CLI.
+    #
+    # Most such code uses the `termlog` / `termwarn` / `termerror` methods
+    # to communicate with the user. When that code executes while a run is
+    # active, its logger messages go to that run's log file.
+    if not _logger.handlers:
+        _logger.addHandler(logging.NullHandler())
+
+
+@contextlib.contextmanager
+def log_to_run(run_id: str | None) -> Iterator[None]:
+    """Direct all wandb log messages to the given run.
+
+    Args:
+        id: The current run ID, or None if actions in the context manager are
+            not associated to a specific run. In the latter case, log messages
+            will go to all runs.
+
+    Usage:
+
+        with wb_logging.run_id(...):
+            ... # Log messages here go to the specified run's logger.
+    """
+    token = _run_id.set(run_id)
+    try:
+        yield
+    finally:
+        _run_id.reset(token)
+
+
+@contextlib.contextmanager
+def log_to_all_runs() -> Iterator[None]:
+    """Direct wandb log messages to all runs.
+
+    Unlike `log_to_run(None)`, this indicates an intentional choice.
+    This is often convenient to use as a decorator:
+
+        @wb_logging.log_to_all_runs()
+        def my_func():
+            ... # Log messages here go to the specified run's logger.
+    """
+    token = _run_id.set(_NOT_RUN_SPECIFIC)
+    try:
+        yield
+    finally:
+        _run_id.reset(token)
+
+
+def add_file_handler(run_id: str, filepath: pathlib.Path) -> logging.Handler:
+    """Direct log messages for a run to a file.
+
+    Args:
+        run_id: The run for which to create a log file.
+        filepath: The file to write log messages to.
+
+    Returns:
+        The added handler which can then be configured further or removed
+        from the 'wandb' logger directly.
+
+        The default logging level is INFO.
+    """
+    handler = logging.FileHandler(filepath)
+    handler.setLevel(logging.INFO)
+    handler.addFilter(_RunIDFilter(run_id))
+    handler.setFormatter(
+        logging.Formatter(
+            "%(asctime)s %(levelname)-7s %(threadName)-10s:%(process)d"
+            " [%(filename)s:%(funcName)s():%(lineno)s]%(run_id_tag)s"
+            " %(message)s"
+        )
+    )
+
+    _logger.addHandler(handler)
+    return handler
+
+
+class _RunIDFilter(logging.Filter):
+    """Filters out messages logged for a different run."""
+
+    def __init__(self, run_id: str) -> None:
+        """Create a _RunIDFilter.
+
+        Args:
+            run_id: Allows messages when the run ID is this or None.
+        """
+        self._run_id = run_id
+
+    def filter(self, record: logging.LogRecord) -> bool:
+        run_id = _run_id.get()
+
+        if run_id is None:
+            record.run_id_tag = " [no run ID]"
+            return True
+        elif isinstance(run_id, _NotRunSpecific):
+            record.run_id_tag = " [all runs]"
+            return True
+        else:
+            record.run_id_tag = ""
+            return run_id == self._run_id

wandb/sdk/mailbox/__init__.py

@@ -9,15 +9,15 @@ The Mailbox handles matching responses to requests. An internal thread
 continuously reads data from the service and passes it to the mailbox.
 """
 
-from .handles import HandleAbandonedError, MailboxHandle
 from .mailbox import Mailbox, MailboxClosedError
+from .mailbox_handle import HandleAbandonedError, MailboxHandle
 from .wait_with_progress import wait_all_with_progress, wait_with_progress
 
 __all__ = [
-    "HandleAbandonedError",
-    "MailboxHandle",
     "Mailbox",
     "MailboxClosedError",
+    "HandleAbandonedError",
+    "MailboxHandle",
     "wait_all_with_progress",
     "wait_with_progress",
 ]

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)