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

wandb/sdk/lib/console_capture.py

@@ -0,0 +1,172 @@
+"""Module for intercepting stdout/stderr.
+
+This patches the `write()` method of `stdout` and `stderr` on import.
+Once patched, it is not possible to unpatch or repatch, though individual
+callbacks can be removed.
+
+We assume that all other writing methods on the object delegate to `write()`,
+like `writelines()`. This is not guaranteed to be true, but it is true for
+common implementations. In particular, CPython's implementation of IOBase's
+`writelines()` delegates to `write()`.
+
+It is important to note that this technique interacts poorly with other
+code that performs similar patching if it also allows unpatching as this
+discards our modification. This is why we patch on import and do not support
+unpatching:
+
+    with contextlib.redirect_stderr(...):
+        from ... import console_capture
+        # Here, everything works fine.
+    # Here, callbacks are never called again.
+
+In particular, it does not work with some combinations of pytest's
+`capfd` / `capsys` fixtures and pytest's `--capture` option.
+"""
+
+from __future__ import annotations
+
+import sys
+import threading
+from typing import IO, AnyStr, Callable, Protocol
+
+
+class CannotCaptureConsoleError(Exception):
+    """The module failed to patch stdout or stderr."""
+
+
+class _WriteCallback(Protocol):
+    """A callback that receives intercepted bytes or string data."""
+
+    def __call__(
+        self,
+        data: bytes | str,
+        written: int,
+        /,
+    ) -> None:
+        """Intercept data passed to `write()`.
+
+        Args:
+            data: The object passed to stderr's or stdout's `write()`.
+            written: The number of bytes or characters written.
+                This is the return value of `write()`.
+        """
+
+
+_module_lock = threading.Lock()
+
+_patch_exception: CannotCaptureConsoleError | None = None
+
+_next_callback_id: int = 1
+
+_stdout_callbacks: dict[int, _WriteCallback] = {}
+_stderr_callbacks: dict[int, _WriteCallback] = {}
+
+
+def capture_stdout(callback: _WriteCallback) -> Callable[[], None]:
+    """Install a callback that runs after every write to sys.stdout.
+
+    Args:
+        callback: A callback to invoke after running `sys.stdout.write`.
+            This may be called from any thread, so it must be thread-safe.
+            Exceptions are propagated to the caller of `write`.
+            See `_WriteCallback` for the exact protocol.
+
+    Returns:
+        A function to uninstall the callback.
+
+    Raises:
+        CannotCaptureConsoleError: If patching failed on import.
+    """
+    with _module_lock:
+        if _patch_exception:
+            raise _patch_exception
+
+        return _insert_disposably(
+            _stdout_callbacks,
+            callback,
+        )
+
+
+def capture_stderr(callback: _WriteCallback) -> Callable[[], None]:
+    """Install a callback that runs after every write to sys.sdterr.
+
+    Args:
+        callback: A callback to invoke after running `sys.stderr.write`.
+            This may be called from any thread, so it must be thread-safe.
+            Exceptions are propagated to the caller of `write`.
+            See `_WriteCallback` for the exact protocol.
+
+    Returns:
+        A function to uninstall the callback.
+
+    Raises:
+        CannotCaptureConsoleError: If patching failed on import.
+    """
+    with _module_lock:
+        if _patch_exception:
+            raise _patch_exception
+
+        return _insert_disposably(
+            _stderr_callbacks,
+            callback,
+        )
+
+
+def _insert_disposably(
+    callback_dict: dict[int, _WriteCallback],
+    callback: _WriteCallback,
+) -> Callable[[], None]:
+    global _next_callback_id
+    id = _next_callback_id
+    _next_callback_id += 1
+
+    disposed = False
+
+    def dispose() -> None:
+        nonlocal disposed
+
+        with _module_lock:
+            if disposed:
+                return
+
+            del callback_dict[id]
+
+            disposed = True
+
+    callback_dict[id] = callback
+    return dispose
+
+
+def _patch(
+    stdout_or_stderr: IO[AnyStr],
+    callbacks: dict[int, _WriteCallback],
+) -> None:
+    orig_write: Callable[[AnyStr], int]
+
+    def write_with_callbacks(s: AnyStr, /) -> int:
+        n = orig_write(s)
+
+        # We make a copy here because callbacks could, in theory, modify
+        # the list of callbacks.
+        with _module_lock:
+            callbacks_copy = list(callbacks.values())
+
+        for cb in callbacks_copy:
+            cb(s, n)
+
+        return n
+
+    orig_write = stdout_or_stderr.write
+
+    # mypy==1.14.1 fails to type-check this:
+    #   Incompatible types in assignment (expression has type
+    #   "Callable[[bytes], int]", variable has type overloaded function)
+    stdout_or_stderr.write = write_with_callbacks  # type: ignore
+
+
+try:
+    _patch(sys.stdout, _stdout_callbacks)
+    _patch(sys.stderr, _stderr_callbacks)
+except Exception as _patch_exception_cause:
+    _patch_exception = CannotCaptureConsoleError()
+    _patch_exception.__cause__ = _patch_exception_cause

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.sock_client import SockClient, SockClientTimeoutError
-from wandb.sdk.mailbox import Mailbox
+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,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/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",
 ]