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

wandb/sdk/lib/progress.py

@@ -2,12 +2,15 @@
 
 from __future__ import annotations
 
+import asyncio
 import contextlib
-from typing import Iterable, Iterator
+import time
+from typing import Iterable, Iterator, NoReturn
 
-import wandb
 from wandb import env
 from wandb.proto import wandb_internal_pb2 as pb
+from wandb.sdk.interface import interface
+from wandb.sdk.lib import asyncio_compat
 
 from . import printer as p
 
@@ -31,14 +34,67 @@ def print_sync_dedupe_stats(
     printer.display(f"W&B sync reduced upload amount by {frac:.1%}")
 
 
+async def loop_printing_operation_stats(
+    progress: ProgressPrinter,
+    interface: interface.InterfaceBase,
+) -> None:
+    """Poll and display ongoing tasks in the internal service process.
+
+    This never returns and must be cancelled. This is meant to be used with
+    `mailbox.wait_with_progress()`.
+
+    Args:
+        progress: The printer to update with operation stats.
+        interface: The interface to use to poll for updates.
+
+    Raises:
+        HandleAbandonedError: If the mailbox associated with the interface
+            becomes closed.
+        Exception: Any other problem communicating with the service process.
+    """
+    stats: pb.OperationStats | None = None
+
+    async def loop_update_screen() -> NoReturn:
+        while True:
+            if stats:
+                progress.update(stats)
+            await asyncio.sleep(0.1)
+
+    async def loop_poll_stats() -> NoReturn:
+        nonlocal stats
+        while True:
+            start_time = time.monotonic()
+
+            handle = interface.deliver_operation_stats()
+            result = await handle.wait_async(timeout=None)
+            stats = result.response.operations_response.operation_stats
+
+            elapsed_time = time.monotonic() - start_time
+            if elapsed_time < 0.5:
+                await asyncio.sleep(0.5 - elapsed_time)
+
+    async with asyncio_compat.open_task_group() as task_group:
+        task_group.start_soon(loop_update_screen())
+        task_group.start_soon(loop_poll_stats())
+
+
 @contextlib.contextmanager
 def progress_printer(
     printer: p.Printer,
-    settings: wandb.Settings | None = None,
+    default_text: str,
 ) -> Iterator[ProgressPrinter]:
-    """Context manager providing an object for printing run progress."""
+    """Context manager providing an object for printing run progress.
+
+    Args:
+        printer: The printer to use.
+        default_text: The text to show if no information is available.
+    """
     with printer.dynamic_text() as text_area:
-        yield ProgressPrinter(printer, text_area, settings)
+        yield ProgressPrinter(
+            printer,
+            text_area,
+            default_text=default_text,
+        )
         printer.progress_close()
 
 
@@ -49,28 +105,27 @@ class ProgressPrinter:
         self,
         printer: p.Printer,
         progress_text_area: p.DynamicText | None,
-        settings: wandb.Settings | None,
+        default_text: str,
     ) -> None:
-        self._show_operation_stats = (
-            settings
-            and settings.x_show_operation_stats
-            # Not implemented by the legacy service.
-            and not env.is_require_legacy_service()
-        )
+        # Not implemented by the legacy service.
+        self._show_operation_stats = not env.is_require_legacy_service()
         self._printer = printer
         self._progress_text_area = progress_text_area
+        self._default_text = default_text
         self._tick = 0
         self._last_printed_line = ""
 
     def update(
         self,
-        progress: list[pb.PollExitResponse],
+        progress: list[pb.PollExitResponse] | pb.OperationStats,
     ) -> None:
         """Update the displayed information."""
         if not progress:
             return
 
-        if self._show_operation_stats:
+        if isinstance(progress, pb.OperationStats):
+            self._update_operation_stats([progress])
+        elif self._show_operation_stats:
             self._update_operation_stats(
                 list(response.operation_stats for response in progress)
             )
@@ -88,6 +143,7 @@ class ProgressPrinter:
                 self._progress_text_area,
                 max_lines=6,
                 loading_symbol=self._printer.loading_symbol(self._tick),
+                default_text=self._default_text,
             ).display(stats_list)
 
         else:
@@ -159,6 +215,10 @@ class ProgressPrinter:
             self._update_progress_text(line, 1.0)
 
     def _update_progress_text(self, text: str, progress: float) -> None:
+        if text == self._last_printed_line:
+            return
+        self._last_printed_line = text
+
         if self._progress_text_area:
             self._progress_text_area.set_text(text)
         else:
@@ -174,11 +234,13 @@ class _DynamicOperationStatsPrinter:
         text_area: p.DynamicText,
         max_lines: int,
         loading_symbol: str,
+        default_text: str,
     ) -> None:
         self._printer = printer
         self._text_area = text_area
         self._max_lines = max_lines
         self._loading_symbol = loading_symbol
+        self._default_text = default_text
 
         self._lines: list[str] = []
         self._ops_shown = 0
@@ -204,9 +266,9 @@ class _DynamicOperationStatsPrinter:
 
         if len(self._lines) == 0:
             if self._loading_symbol:
-                self._text_area.set_text(f"{self._loading_symbol} Finishing up...")
+                self._text_area.set_text(f"{self._loading_symbol} {self._default_text}")
             else:
-                self._text_area.set_text("Finishing up...")
+                self._text_area.set_text(self._default_text)
         else:
             self._text_area.set_text("\n".join(self._lines))
 

wandb/sdk/lib/service_connection.py

@@ -12,8 +12,8 @@ from wandb.sdk.interface.interface import InterfaceBase
 from wandb.sdk.interface.interface_sock import InterfaceSock
 from wandb.sdk.lib import service_token
 from wandb.sdk.lib.exit_hooks import ExitHooks
-from wandb.sdk.lib.mailbox import Mailbox
 from wandb.sdk.lib.sock_client import SockClient, SockClientTimeoutError
+from wandb.sdk.mailbox import Mailbox
 from wandb.sdk.service import service
 
 

wandb/sdk/lib/sock_client.py

@@ -177,9 +177,9 @@ class SockClient:
             inform_finish=inform_finish,
             inform_teardown=inform_teardown,
         )
-        # TODO: this solution is fragile, but for checking attach
-        # it should be relatively stable.
-        # This pass would be solved as part of the fix in https://wandb.atlassian.net/browse/WB-8709
+
+        # 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:
@@ -213,11 +213,13 @@ class SockClient:
 
     def send_record_communicate(self, record: "pb.Record") -> None:
         server_req = spb.ServerRequest()
+        server_req.request_id = record.control.mailbox_slot
         server_req.record_communicate.CopyFrom(record)
         self.send_server_request(server_req)
 
     def send_record_publish(self, record: "pb.Record") -> None:
         server_req = spb.ServerRequest()
+        server_req.request_id = record.control.mailbox_slot
         server_req.record_publish.CopyFrom(record)
         self.send_server_request(server_req)
 
@@ -256,10 +258,8 @@ class SockClient:
                 data = self._sock.recv(self._bufsize)
             except socket.timeout:
                 break
-            except ConnectionResetError:
-                raise SockClientClosedError
-            except OSError:
-                raise SockClientClosedError
+            except OSError as e:
+                raise SockClientClosedError from e
             finally:
                 if timeout:
                     self._sock.settimeout(None)

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 .handles import HandleAbandonedError, MailboxHandle
+from .mailbox import Mailbox, MailboxClosedError
+from .wait_with_progress import wait_all_with_progress, wait_with_progress
+
+__all__ = [
+    "HandleAbandonedError",
+    "MailboxHandle",
+    "Mailbox",
+    "MailboxClosedError",
+    "wait_all_with_progress",
+    "wait_with_progress",
+]

wandb/sdk/mailbox/handles.py

@@ -0,0 +1,199 @@
+from __future__ import annotations
+
+import asyncio
+import math
+import threading
+from typing import TYPE_CHECKING
+
+from wandb.proto import wandb_internal_pb2 as pb
+
+# 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."""
+
+
+class MailboxHandle:
+    """A thread-safe handle that allows waiting for a response to a request."""
+
+    def __init__(self, address: str) -> None:
+        self._address = address
+        self._lock = threading.Lock()
+        self._event = threading.Event()
+
+        self._abandoned = False
+        self._result: pb.Result | None = None
+
+        self._asyncio_events: dict[asyncio.Event, _AsyncioEvent] = dict()
+
+    def deliver(self, result: pb.Result) -> 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._result:
+                raise ValueError(
+                    f"A response has already been delivered to {self._address}."
+                )
+
+            self._result = result
+            self._signal_done()
+
+    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()
+
+    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.
+
+        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()
+
+    def check(self) -> pb.Result | None:
+        """Returns the result if it's ready."""
+        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.
+
+        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.
+        """
+        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._result:
+                return self._result
+
+            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.
+        """
+        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._result:
+                    return self._result
+                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._result:
+                    return self._result
+
+                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 result is received.
+
+        If a result already exists, this notifies the event loop immediately.
+        """
+        asyncio_event = _AsyncioEvent(loop, event)
+
+        with self._lock:
+            if self._result 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."""
+        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/mailbox.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+import logging
+import secrets
+import string
+import threading
+
+from wandb.proto import wandb_internal_pb2 as pb
+
+from . import handles
+
+_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 Record 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_lock = threading.Lock()
+        self._closed = False
+
+    def require_response(self, request: pb.Record) -> handles.MailboxHandle:
+        """Set a response address on a request.
+
+        Args:
+            request: The request on which to set a 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 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)
+            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, result: pb.Result) -> 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
+        if not address:
+            _logger.error(
+                "Received response with no mailbox slot."
+                f" Kind: {result.WhichOneof('result_type')}"
+            )
+            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(result)
+
+    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/wait_with_progress.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import time
+from typing import Any, Callable, Coroutine, List, cast
+
+from wandb.proto import wandb_internal_pb2 as pb
+from wandb.sdk.lib import asyncio_compat
+
+from . import handles
+
+
+def wait_with_progress(
+    handle: handles.MailboxHandle,
+    *,
+    timeout: float | None,
+    progress_after: float,
+    display_progress: Callable[[], Coroutine[Any, Any, None]],
+) -> pb.Result:
+    """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[handles.MailboxHandle],
+    *,
+    timeout: float | None,
+    progress_after: float,
+    display_progress: Callable[[], Coroutine[Any, Any, None]],
+) -> list[pb.Result]:
+    """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[pb.Result]:
+        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[handles.MailboxHandle],
+    *,
+    timeout: float,
+) -> list[pb.Result]:
+    """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[pb.Result] = []
+
+    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[handles.MailboxHandle],
+    *,
+    timeout: float | None,
+) -> list[pb.Result]:
+    """Asynchronously wait for multiple mailbox handles.
+
+    Just like _wait_handles.
+    """
+    results: list[pb.Result | 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[pb.Result], results)