wandb 0.19.6rc4__py3-none-macosx_11_0_arm64.whl → 0.19.8__py3-none-macosx_11_0_arm64.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/redirect.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 try:
     import fcntl
     import pty
@@ -17,8 +19,10 @@ import sys
 import threading
 import time
 from collections import defaultdict
+from typing import Callable, Iterable, Literal
 
 import wandb
+from wandb.sdk.lib import console_capture
 
 
 class _Numpy:  # fallback in case numpy is not available
@@ -55,8 +59,6 @@ except ImportError:
 
 logger = logging.getLogger("wandb")
 
-_redirects = {"stdout": None, "stderr": None}
-
 
 ANSI_CSI_RE = re.compile("\001?\033\\[((?:\\d|;)*)([a-zA-Z])\002?")
 ANSI_OSC_RE = re.compile("\001?\033\\]([^\a]*)(\a)\002?")
@@ -491,7 +493,11 @@ _MIN_CALLBACK_INTERVAL = 2  # seconds
 
 
 class RedirectBase:
-    def __init__(self, src, cbs=()):
+    def __init__(
+        self,
+        src: Literal["stdout", "stderr"],
+        cbs: Iterable[Callable[[str], None]] = (),
+    ) -> None:
         """# Arguments.
 
         `src`: Source stream to be redirected. "stdout" or "stderr".
@@ -499,7 +505,7 @@ class RedirectBase:
 
         """
         assert hasattr(sys, src)
-        self.src = src
+        self.src: Literal["stdout", "stderr"] = src
         self.cbs = cbs
 
     @property
@@ -514,71 +520,82 @@ class RedirectBase:
     def src_wrapped_stream(self):
         return getattr(sys, self.src)
 
-    def save(self):
+    def install(self) -> None:
         pass
 
-    def install(self):
-        curr_redirect = _redirects.get(self.src)
-        if curr_redirect and curr_redirect != self:
-            curr_redirect.uninstall()
-        _redirects[self.src] = self
-
-    def uninstall(self):
-        if _redirects[self.src] != self:
-            return
-        _redirects[self.src] = None
+    def uninstall(self) -> None:
+        pass
 
 
 class StreamWrapper(RedirectBase):
     """Patches the write method of current sys.stdout/sys.stderr."""
 
-    def __init__(self, src, cbs=()):
+    def __init__(
+        self,
+        src: Literal["stdout", "stderr"],
+        cbs: Iterable[Callable[[str], None]] = (),
+    ) -> None:
         super().__init__(src=src, cbs=cbs)
-        self._installed = False
+        self._uninstall: Callable[[], None] | None = None
         self._emulator = TerminalEmulator()
+        self._queue: queue.Queue[str] = queue.Queue()
+        self._stopped = threading.Event()
 
-    def _emulator_write(self):
+    def _emulator_write(self) -> None:
         while True:
             if self._queue.empty():
                 if self._stopped.is_set():
                     return
                 time.sleep(0.5)
                 continue
-            data = []
+
+            data: list[str] = []
             while not self._queue.empty():
                 data.append(self._queue.get())
+
             if self._stopped.is_set() and sum(map(len, data)) > 100000:
                 wandb.termlog("Terminal output too large. Logging without processing.")
                 self.flush()
-                [self.flush(line.encode("utf-8")) for line in data]
+
+                for line in data:
+                    self.flush(line)
+
                 return
+
             try:
                 self._emulator.write("".join(data))
             except Exception:
                 pass
 
-    def _callback(self):
+    def _callback(self) -> None:
         while not (self._stopped.is_set() and self._queue.empty()):
             self.flush()
             time.sleep(_MIN_CALLBACK_INTERVAL)
 
-    def install(self):
-        super().install()
-        if self._installed:
-            return
-        stream = self.src_wrapped_stream
-        old_write = stream.write
-        self._prev_callback_timestamp = time.time()
-        self._old_write = old_write
+    def _on_write(self, data: str | bytes, written: int, /) -> None:
+        if isinstance(data, bytes):
+            written_data = data[:written].decode("utf-8")
+        else:
+            written_data = data[:written]
 
-        def write(data):
-            self._old_write(data)
-            self._queue.put(data)
+        self._queue.put(written_data)
 
-        stream.write = write
+    def install(self) -> None:
+        if self._uninstall:
+            return
+
+        try:
+            if self.src == "stdout":
+                self._uninstall = console_capture.capture_stdout(self._on_write)
+            else:
+                self._uninstall = console_capture.capture_stderr(self._on_write)
+        except console_capture.CannotCaptureConsoleError:
+            logger.exception("failed to install %s hooks", self.src)
+            wandb.termwarn(
+                f"Failed to wrap {self.src}. Console logs will not be captured.",
+            )
+            return
 
-        self._queue = queue.Queue()
-        self._stopped = threading.Event()
         self._emulator_write_thread = threading.Thread(target=self._emulator_write)
         self._emulator_write_thread.daemon = True
         self._emulator_write_thread.start()
@@ -588,25 +605,25 @@ class StreamWrapper(RedirectBase):
             self._callback_thread.daemon = True
             self._callback_thread.start()
 
-        self._installed = True
-
-    def flush(self, data=None):
+    def flush(self, data: str | None = None) -> None:
         if data is None:
             try:
                 data = self._emulator.read().encode("utf-8")
             except Exception:
-                pass
+                logger.exception("exception reading TerminalEmulator")
+
         if data:
             for cb in self.cbs:
                 try:
                     cb(data)
                 except Exception:
-                    pass  # TODO(frz)
+                    logger.exception("exception in StreamWrapper callback")
 
-    def uninstall(self):
-        if not self._installed:
+    def uninstall(self) -> None:
+        if not self._uninstall:
             return
-        self.src_wrapped_stream.write = self._old_write
+
+        self._uninstall()
 
         self._stopped.set()
         self._emulator_write_thread.join(timeout=5)
@@ -616,9 +633,6 @@ class StreamWrapper(RedirectBase):
             wandb.termlog("Done.")
         self.flush()
 
-        self._installed = False
-        super().uninstall()
-
 
 class StreamRawWrapper(RedirectBase):
     """Patches the write method of current sys.stdout/sys.stderr.
@@ -626,40 +640,44 @@ class StreamRawWrapper(RedirectBase):
     Captures data in a raw form rather than using the emulator
     """
 
-    def __init__(self, src, cbs=()):
+    def __init__(
+        self,
+        src: Literal["stdout", "stderr"],
+        cbs: Iterable[Callable[[str], None]] = (),
+    ) -> None:
         super().__init__(src=src, cbs=cbs)
-        self._installed = False
+        self._uninstall: Callable[[], None] | None = None
 
-    def save(self):
-        stream = self.src_wrapped_stream
-        self._old_write = stream.write
+    def _on_write(self, data: str | bytes, written: int, /) -> None:
+        if isinstance(data, bytes):
+            written_data = data[:written].decode("utf-8")
+        else:
+            written_data = data[:written]
 
-    def install(self):
-        super().install()
-        if self._installed:
-            return
-        stream = self.src_wrapped_stream
-        self._prev_callback_timestamp = time.time()
+        for cb in self.cbs:
+            try:
+                cb(written_data)
+            except Exception:
+                logger.exception("error in %s callback", self.src)
 
-        def write(data):
-            self._old_write(data)
-            for cb in self.cbs:
-                try:
-                    cb(data)
-                except Exception:
-                    # TODO: Figure out why this was needed and log or error out appropriately
-                    # it might have been strange terminals? maybe shutdown cases?
-                    pass
+    def install(self) -> None:
+        if self._uninstall:
+            return
 
-        stream.write = write
-        self._installed = True
+        try:
+            if self.src == "stdout":
+                self._uninstall = console_capture.capture_stdout(self._on_write)
+            else:
+                self._uninstall = console_capture.capture_stderr(self._on_write)
+        except console_capture.CannotCaptureConsoleError:
+            logger.exception("failed to install %s hooks", self.src)
+            wandb.termwarn(
+                f"Failed to wrap {self.src}. Console logs will not be captured.",
+            )
 
-    def uninstall(self):
-        if not self._installed:
-            return
-        self.src_wrapped_stream.write = self._old_write
-        self._installed = False
-        super().uninstall()
+    def uninstall(self) -> None:
+        if self._uninstall:
+            self._uninstall()
 
 
 class _WindowSizeChangeHandler:
@@ -708,6 +726,8 @@ class _WindowSizeChangeHandler:
 
 _WSCH = _WindowSizeChangeHandler()
 
+_redirects: dict[str, Redirect | None] = {"stdout": None, "stderr": None}
+
 
 class Redirect(RedirectBase):
     """Redirect low level file descriptors."""
@@ -725,7 +745,11 @@ class Redirect(RedirectBase):
         return r, w
 
     def install(self):
-        super().install()
+        curr_redirect = _redirects.get(self.src)
+        if curr_redirect and curr_redirect != self:
+            curr_redirect.uninstall()
+        _redirects[self.src] = self
+
         if self._installed:
             return
         self._pipe_read_fd, self._pipe_write_fd = self._pipe()
@@ -776,7 +800,9 @@ class Redirect(RedirectBase):
         self.flush()
 
         _WSCH.remove_fd(self._pipe_read_fd)
-        super().uninstall()
+
+        if _redirects[self.src] == self:
+            _redirects[self.src] = None
 
     def flush(self, data=None):
         if data is None:

wandb/sdk/lib/service_connection.py

@@ -10,10 +10,11 @@ from wandb.proto import wandb_settings_pb2
 from wandb.sdk import wandb_settings
 from wandb.sdk.interface.interface import InterfaceBase
 from wandb.sdk.interface.interface_sock import InterfaceSock
+from wandb.sdk.interface.router_sock import MessageSockRouter
 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.lib.sock_client import SockClient, SockClientClosedError
+from wandb.sdk.mailbox import HandleAbandonedError, Mailbox, MailboxClosedError
 from wandb.sdk.service import service
 
 
@@ -115,6 +116,9 @@ class ServiceConnection:
         """Returns a new ServiceConnection.
 
         Args:
+            mailbox: The mailbox to use for all communication over the socket.
+            router: A handle to the thread that reads from the socket and
+                updates the mailbox.
             client: A socket that's connected to the service.
             proc: The service process if we own it, or None otherwise.
             cleanup: A callback to run on teardown before doing anything.
@@ -124,9 +128,12 @@ class ServiceConnection:
         self._torn_down = False
         self._cleanup = cleanup
 
-    def make_interface(self, mailbox: Mailbox, stream_id: str) -> InterfaceBase:
+        self._mailbox = Mailbox()
+        self._router = MessageSockRouter(self._client, self._mailbox)
+
+    def make_interface(self, stream_id: str) -> InterfaceBase:
         """Returns an interface for communicating with the service."""
-        return InterfaceSock(self._client, mailbox, stream_id=stream_id)
+        return InterfaceSock(self._client, self._mailbox, stream_id=stream_id)
 
     def send_record(self, record: pb.Record) -> None:
         """Sends data to the service."""
@@ -141,13 +148,13 @@ class ServiceConnection:
         request = spb.ServerInformInitRequest()
         request.settings.CopyFrom(settings)
         request._info.stream_id = run_id
-        self._client.send(inform_init=request)
+        self._client.send_server_request(spb.ServerRequest(inform_init=request))
 
     def inform_finish(self, run_id: str) -> None:
         """Sends an finish request to the service."""
         request = spb.ServerInformFinishRequest()
         request._info.stream_id = run_id
-        self._client.send(inform_finish=request)
+        self._client.send_server_request(spb.ServerRequest(inform_finish=request))
 
     def inform_attach(
         self,
@@ -157,18 +164,26 @@ class ServiceConnection:
 
         Raises a WandbAttachFailedError if attaching is not possible.
         """
-        request = spb.ServerInformAttachRequest()
-        request._info.stream_id = attach_id
+        request = spb.ServerRequest()
+        request.inform_attach._info.stream_id = attach_id
 
         try:
-            response = self._client.send_and_recv(inform_attach=request)
+            handle = self._mailbox.require_response(request)
+            self._client.send_server_request(request)
+            response = handle.wait_or(timeout=10)
             return response.inform_attach_response.settings
-        except SockClientTimeoutError:
+
+        except (MailboxClosedError, HandleAbandonedError, SockClientClosedError):
+            raise WandbAttachFailedError(
+                "Failed to attach: the service process is not running.",
+            ) from None
+
+        except TimeoutError:
             raise WandbAttachFailedError(
-                "Could not attach because the run does not belong to"
+                "Failed to attach because the run does not belong to"
                 " the current service process, or because the service"
                 " process is busy (unlikely)."
-            )
+            ) from None
 
     def inform_start(
         self,
@@ -179,7 +194,7 @@ class ServiceConnection:
         request = spb.ServerInformStartRequest()
         request.settings.CopyFrom(settings)
         request._info.stream_id = run_id
-        self._client.send(inform_start=request)
+        self._client.send_server_request(spb.ServerRequest(inform_start=request))
 
     def teardown(self, exit_code: int) -> int:
         """Shuts down the service process and returns its exit code.
@@ -207,10 +222,15 @@ class ServiceConnection:
         # Clear the service token to prevent new connections from being made.
         service_token.clear_service_token()
 
-        self._client.send(
-            inform_teardown=spb.ServerInformTeardownRequest(
-                exit_code=exit_code,
-            )
+        # Stop reading responses on the socket.
+        self._router.join()
+
+        self._client.send_server_request(
+            spb.ServerRequest(
+                inform_teardown=spb.ServerInformTeardownRequest(
+                    exit_code=exit_code,
+                )
+            ),
         )
 
         return self._proc.join()

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,63 +161,15 @@ 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,
-        )
-        # 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
-        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
         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 +208,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)