wandb 0.21.1__py3-none-win32.whl → 0.21.2__py3-none-win32.whl

wandb/sdk/lib/retry.py

@@ -94,109 +94,216 @@ class Retry(Generic[_R]):
         """The number of iterations the previous __call__ retried."""
         return self._num_iter
 
-    def __call__(self, *args: Any, **kwargs: Any) -> _R:
+    def __call__(
+        self,
+        *args: Any,
+        num_retries: Optional[int] = None,
+        retry_timedelta: Optional[datetime.timedelta] = None,
+        retry_sleep_base: Optional[float] = None,
+        retry_cancel_event: Optional[threading.Event] = None,
+        check_retry_fn: Optional[CheckRetryFnType] = None,
+        **kwargs: Any,
+    ) -> _R:
         """Call the wrapped function, with retries.
 
         Args:
-           retry_timedelta (kwarg): amount of time to retry before giving up.
-           sleep_base (kwarg): amount of time to sleep upon first failure, all other sleeps
-               are derived from this one.
+            num_retries: The number of retries after which to give up.
+            retry_timedelta: An amount of time after which to give up.
+            retry_sleep_base: Number of seconds to sleep for the first retry.
+                This is used as the base for exponential backoff.
+            retry_cancel_event: An event that causes this to raise
+                a RetryCancelledException on the next attempted retry.
+            check_retry_fn: A custom check for deciding whether an exception
+                should be retried. Retrying is prevented if this returns a falsy
+                value, even if more retries are left. This may also return a
+                timedelta that represents a shorter timeout: retrying is
+                prevented if the value is less than the amount of time that has
+                passed since the last timedelta was returned.
         """
-        retry_timedelta = kwargs.pop("retry_timedelta", self._retry_timedelta)
-        if retry_timedelta is None:
-            retry_timedelta = datetime.timedelta(days=365)
-
-        retry_cancel_event = kwargs.pop("retry_cancel_event", self._retry_cancel_event)
-
-        num_retries = kwargs.pop("num_retries", self._num_retries)
-        if num_retries is None:
-            num_retries = 1000000
-
         if os.environ.get("WANDB_TEST"):
-            num_retries = 0
+            max_retries = 0
+        elif num_retries is not None:
+            max_retries = num_retries
+        elif self._num_retries is not None:
+            max_retries = self._num_retries
+        else:
+            max_retries = 1000000
 
-        sleep_base: float = kwargs.pop("retry_sleep_base", 1)
+        if retry_timedelta is not None:
+            timeout = retry_timedelta
+        elif self._retry_timedelta is not None:
+            timeout = self._retry_timedelta
+        else:
+            timeout = datetime.timedelta(days=365)
 
-        # an extra function to allow performing more logic on the filtered exception
-        check_retry_fn: CheckRetryFnType = kwargs.pop(
-            "check_retry_fn", self._check_retry_fn
+        if retry_sleep_base is not None:
+            initial_sleep = retry_sleep_base
+        else:
+            initial_sleep = 1
+
+        retry_loop = _RetryLoop(
+            max_retries=max_retries,
+            timeout=timeout,
+            initial_sleep=initial_sleep,
+            max_sleep=self.MAX_SLEEP_SECONDS,
+            cancel_event=retry_cancel_event or self._retry_cancel_event,
+            retry_check=check_retry_fn or self._check_retry_fn,
         )
 
-        sleep = sleep_base
-        now = NOW_FN()
-        start_time = now
-        start_time_triggered = None
-
+        start_time = NOW_FN()
         self._num_iter = 0
 
         while True:
             try:
                 result = self._call_fn(*args, **kwargs)
-                # Only print resolved attempts once every minute
-                if self._num_iter > 2 and now - self._last_print > datetime.timedelta(
-                    minutes=1
-                ):
-                    self._last_print = NOW_FN()
-                    if self.retry_callback:
-                        self.retry_callback(
-                            200,
-                            f"{self._error_prefix} resolved after {NOW_FN() - start_time}, resuming normal operation.",
-                        )
-                return result
+
             except self._retryable_exceptions as e:
-                # if the secondary check fails, re-raise
-                retry_timedelta_triggered = check_retry_fn(e)
-                if not retry_timedelta_triggered:
+                if not retry_loop.should_retry(e):
                     raise
 
-                # always enforce num_retries no matter which type of exception was seen
-                if self._num_iter >= num_retries:
-                    raise
+                if self._num_iter == 2:
+                    logger.info("Retry attempt failed:", exc_info=e)
+                    self._print_entered_retry_loop(e)
 
-                now = NOW_FN()
+                retry_loop.wait_before_retry()
+                self._num_iter += 1
 
-                # handle a triggered secondary check which could have a shortened timeout
-                if isinstance(retry_timedelta_triggered, datetime.timedelta):
-                    # save the time of the first secondary trigger
-                    if not start_time_triggered:
-                        start_time_triggered = now
+            else:
+                if self._num_iter > 2:
+                    self._print_recovered(start_time)
 
-                    # make sure that we haven't run out of time from secondary trigger
-                    if now - start_time_triggered >= retry_timedelta_triggered:
-                        raise
+                return result
 
-                # always enforce the default timeout from start of retries
-                if now - start_time >= retry_timedelta:
-                    raise
+    def _print_entered_retry_loop(self, exception: Exception) -> None:
+        """Emit a message saying we've begun retrying.
 
-                if self._num_iter == 2:
-                    logger.info("Retry attempt failed:", exc_info=e)
-                    if (
-                        isinstance(e, HTTPError)
-                        and e.response is not None
-                        and self.retry_callback is not None
-                    ):
-                        self.retry_callback(e.response.status_code, e.response.text)
-                    else:
-                        # todo: would like to catch other errors, eg wandb.errors.Error, ConnectionError etc
-                        # but some of these can be raised before the retry handler thread (RunStatusChecker) is
-                        # spawned in wandb_init
-                        wandb.termlog(
-                            f"{self._error_prefix} ({e.__class__.__name__}), entering retry loop."
-                        )
-                # if wandb.env.is_debug():
-                #     traceback.print_exc()
-            cancelled = self._sleep_check_cancelled(
-                sleep + random.random() * 0.25 * sleep, cancel_event=retry_cancel_event
+        Either calls the retry callback or prints a warning to console.
+
+        Args:
+            exception: The most recent exception we will retry.
+        """
+        if (
+            isinstance(exception, HTTPError)
+            and exception.response is not None
+            and self.retry_callback is not None
+        ):
+            self.retry_callback(
+                exception.response.status_code,
+                exception.response.text,
             )
+        else:
+            wandb.termlog(
+                f"{self._error_prefix}"
+                f" ({exception.__class__.__name__}), entering retry loop."
+            )
+
+    def _print_recovered(self, start_time: datetime.datetime) -> None:
+        """Emit a message saying we've recovered after retrying.
+
+        Args:
+            start_time: When we started retrying.
+        """
+        if not self.retry_callback:
+            return
+
+        now = NOW_FN()
+        if now - self._last_print < datetime.timedelta(minutes=1):
+            return
+        self._last_print = now
+
+        time_to_recover = now - start_time
+        self.retry_callback(
+            200,
+            (
+                f"{self._error_prefix} resolved after"
+                f" {time_to_recover}, resuming normal operation."
+            ),
+        )
+
+
+class _RetryLoop:
+    """An invocation of a Retry instance."""
+
+    def __init__(
+        self,
+        *,
+        max_retries: int,
+        timeout: datetime.timedelta,
+        initial_sleep: float,
+        max_sleep: float,
+        cancel_event: Optional[threading.Event],
+        retry_check: CheckRetryFnType,
+    ) -> None:
+        """Start a new call of a Retry instance.
+
+        Args:
+            max_retries: The number of retries after which to give up.
+            timeout: An amount of time after which to give up.
+            initial_sleep: Number of seconds to sleep for the first retry.
+                This is used as the base for exponential backoff.
+            max_sleep: Maximum number of seconds to sleep between retries.
+            cancel_event: An event that's set when the function is cancelled.
+            retry_check: A custom check for deciding whether an exception should
+                be retried. Retrying is prevented if this returns a falsy value,
+                even if more retries are left. This may also return a timedelta
+                that represents a shorter timeout: retrying is prevented if the
+                value is less than the amount of time that has passed since the
+                last timedelta was returned.
+        """
+        self._max_retries = max_retries
+        self._total_retries = 0
+
+        self._timeout = timeout
+        self._start_time = NOW_FN()
+
+        self._next_sleep_time = initial_sleep
+        self._max_sleep = max_sleep
+        self._cancel_event = cancel_event
+
+        self._retry_check = retry_check
+        self._last_custom_timeout: Optional[datetime.datetime] = None
+
+    def should_retry(self, exception: Exception) -> bool:
+        """Returns whether an exception should be retried."""
+        if self._total_retries >= self._max_retries:
+            return False
+        self._total_retries += 1
+
+        now = NOW_FN()
+        if now - self._start_time >= self._timeout:
+            return False
+
+        retry_check_result = self._retry_check(exception)
+        if not retry_check_result:
+            return False
+
+        if isinstance(retry_check_result, datetime.timedelta):
+            if not self._last_custom_timeout:
+                self._last_custom_timeout = now
+
+            if now - self._last_custom_timeout >= retry_check_result:
+                return False
+
+        return True
+
+    def wait_before_retry(self) -> None:
+        """Block until the next retry should happen.
+
+        Raises:
+            RetryCancelledError: If the operation is cancelled.
+        """
+        sleep_amount = self._next_sleep_time * (1 + random.random() * 0.25)
+
+        if self._cancel_event:
+            cancelled = self._cancel_event.wait(sleep_amount)
             if cancelled:
-                raise RetryCancelledError("retry timeout")
-            sleep *= 2
-            if sleep > self.MAX_SLEEP_SECONDS:
-                sleep = self.MAX_SLEEP_SECONDS
-            now = NOW_FN()
+                raise RetryCancelledError("Cancelled while retrying.")
+        else:
+            SLEEP_FN(sleep_amount)
 
-            self._num_iter += 1
+        self._next_sleep_time *= 2
+        if self._next_sleep_time > self._max_sleep:
+            self._next_sleep_time = self._max_sleep
 
 
 _F = TypeVar("_F", bound=Callable)

wandb/sdk/lib/service/service_client.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import struct
+
+from wandb.proto import wandb_server_pb2 as spb
+from wandb.sdk.lib import asyncio_manager
+from wandb.sdk.mailbox.mailbox import Mailbox
+from wandb.sdk.mailbox.mailbox_handle import MailboxHandle
+
+_logger = logging.getLogger(__name__)
+
+_HEADER_BYTE_INT_LEN = 5
+_HEADER_BYTE_INT_FMT = "<BI"
+
+
+class ServiceClient:
+    """Implements socket communication with the internal service."""
+
+    def __init__(
+        self,
+        asyncer: asyncio_manager.AsyncioManager,
+        reader: asyncio.StreamReader,
+        writer: asyncio.StreamWriter,
+    ) -> None:
+        self._asyncer = asyncer
+        self._reader = reader
+        self._writer = writer
+        self._mailbox = Mailbox(asyncer)
+        asyncer.run_soon(
+            self._forward_responses,
+            daemon=True,
+            name="ServiceClient._forward_responses",
+        )
+
+    def publish(self, request: spb.ServerRequest) -> None:
+        """Send a request without waiting for a response."""
+        self._asyncer.run_soon(lambda: self._send_server_request(request))
+
+    def deliver(
+        self,
+        request: spb.ServerRequest,
+    ) -> MailboxHandle[spb.ServerResponse]:
+        """Send a request and return a handle to wait for a response.
+
+        NOTE: This may mutate the request. The request should not be used
+        after.
+
+        Raises:
+            MailboxClosedError: If used after the client is closed or has
+                stopped due to an error.
+        """
+        handle = self._mailbox.require_response(request)
+        self._asyncer.run_soon(lambda: self._send_server_request(request))
+        return handle
+
+    async def _send_server_request(self, request: spb.ServerRequest) -> None:
+        header = struct.pack(_HEADER_BYTE_INT_FMT, ord("W"), request.ByteSize())
+        self._writer.write(header)
+
+        data = request.SerializeToString()
+        self._writer.write(data)
+
+        await self._writer.drain()
+
+    def close(self) -> None:
+        """Flush and close the socket."""
+        self._asyncer.run_soon(self._close)
+
+    async def _close(self) -> None:
+        self._writer.close()
+        await self._writer.wait_closed()
+
+    async def _forward_responses(self) -> None:
+        try:
+            while response := await self._read_server_response():
+                await self._mailbox.deliver(response)
+
+        except Exception:
+            _logger.exception("Error reading server response.")
+
+        else:
+            _logger.info("Reached EOF.")
+
+        finally:
+            self._mailbox.close()
+
+    async def _read_server_response(self) -> spb.ServerResponse | None:
+        try:
+            header = await self._reader.readexactly(_HEADER_BYTE_INT_LEN)
+        except asyncio.IncompleteReadError as e:
+            if e.partial:
+                raise
+            else:
+                return None
+
+        magic, length = struct.unpack(_HEADER_BYTE_INT_FMT, header)
+
+        if magic != ord("W"):
+            raise ValueError(f"Bad header: {header.hex()}")
+
+        data = await self._reader.readexactly(length)
+        response = spb.ServerResponse()
+        response.ParseFromString(data)
+        return response

wandb/sdk/lib/service/service_connection.py

@@ -3,16 +3,15 @@ from __future__ import annotations
 import atexit
 from typing import Callable
 
-from wandb.proto import wandb_internal_pb2 as pb
 from wandb.proto import wandb_server_pb2 as spb
 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 asyncio_manager
 from wandb.sdk.lib.exit_hooks import ExitHooks
-from wandb.sdk.lib.sock_client import SockClient, SockClientClosedError
-from wandb.sdk.mailbox import HandleAbandonedError, Mailbox, MailboxClosedError
+from wandb.sdk.lib.service.service_client import ServiceClient
+from wandb.sdk.mailbox import HandleAbandonedError, MailboxClosedError
 
 from . import service_process, service_token
 
@@ -22,18 +21,23 @@ class WandbAttachFailedError(Exception):
 
 
 def connect_to_service(
+    asyncer: asyncio_manager.AsyncioManager,
     settings: wandb_settings.Settings,
 ) -> ServiceConnection:
     """Connect to the service process, starting one up if necessary."""
     token = service_token.from_env()
 
     if token:
-        return ServiceConnection(client=token.connect(), proc=None)
+        return ServiceConnection(
+            client=token.connect(asyncer=asyncer),
+            proc=None,
+        )
     else:
-        return _start_and_connect_service(settings)
+        return _start_and_connect_service(asyncer, settings)
 
 
 def _start_and_connect_service(
+    asyncer: asyncio_manager.AsyncioManager,
     settings: wandb_settings.Settings,
 ) -> ServiceConnection:
     """Start a service process and returns a connection to it.
@@ -44,7 +48,7 @@ def _start_and_connect_service(
     """
     proc = service_process.start(settings)
 
-    client = proc.token.connect()
+    client = proc.token.connect(asyncer=asyncer)
     proc.token.save_to_env()
 
     hooks = ExitHooks()
@@ -69,7 +73,7 @@ class ServiceConnection:
 
     def __init__(
         self,
-        client: SockClient,
+        client: ServiceClient,
         proc: service_process.ServiceProcess | None,
         cleanup: Callable[[], None] | None = None,
     ):
@@ -88,16 +92,9 @@ class ServiceConnection:
         self._torn_down = False
         self._cleanup = cleanup
 
-        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, self._mailbox, stream_id=stream_id)
-
-    def send_record(self, record: pb.Record) -> None:
-        """Send data to the service."""
-        self._client.send_record_publish(record)
+        return InterfaceSock(self._client, stream_id=stream_id)
 
     def inform_init(
         self,
@@ -108,13 +105,13 @@ class ServiceConnection:
         request = spb.ServerInformInitRequest()
         request.settings.CopyFrom(settings)
         request._info.stream_id = run_id
-        self._client.send_server_request(spb.ServerRequest(inform_init=request))
+        self._client.publish(spb.ServerRequest(inform_init=request))
 
     def inform_finish(self, run_id: str) -> None:
         """Send an finish request to the service."""
         request = spb.ServerInformFinishRequest()
         request._info.stream_id = run_id
-        self._client.send_server_request(spb.ServerRequest(inform_finish=request))
+        self._client.publish(spb.ServerRequest(inform_finish=request))
 
     def inform_attach(
         self,
@@ -128,11 +125,10 @@ class ServiceConnection:
         request.inform_attach._info.stream_id = attach_id
 
         try:
-            handle = self._mailbox.require_response(request)
-            self._client.send_server_request(request)
+            handle = self._client.deliver(request)
             response = handle.wait_or(timeout=10)
 
-        except (MailboxClosedError, HandleAbandonedError, SockClientClosedError):
+        except (MailboxClosedError, HandleAbandonedError):
             raise WandbAttachFailedError(
                 "Failed to attach: the service process is not running.",
             ) from None
@@ -156,7 +152,7 @@ class ServiceConnection:
         request = spb.ServerInformStartRequest()
         request.settings.CopyFrom(settings)
         request._info.stream_id = run_id
-        self._client.send_server_request(spb.ServerRequest(inform_start=request))
+        self._client.publish(spb.ServerRequest(inform_start=request))
 
     def teardown(self, exit_code: int) -> int | None:
         """Close the connection.
@@ -178,21 +174,19 @@ class ServiceConnection:
         if self._cleanup:
             self._cleanup()
 
-        # Stop reading responses on the socket.
-        self._router.join()
-
         if not self._proc:
             return None
 
         # Clear the service token to prevent new connections to the process.
         service_token.clear_service_in_env()
 
-        self._client.send_server_request(
+        self._client.publish(
             spb.ServerRequest(
                 inform_teardown=spb.ServerInformTeardownRequest(
                     exit_code=exit_code,
                 )
             ),
         )
+        self._client.close()
 
         return self._proc.join()

wandb/sdk/lib/service/service_token.py

@@ -1,15 +1,17 @@
 from __future__ import annotations
 
 import abc
+import asyncio
 import os
 import re
-import socket
 
 from typing_extensions import final, override
 
 from wandb import env
+from wandb.sdk.lib import asyncio_manager
 from wandb.sdk.lib.service import ipc_support
-from wandb.sdk.lib.sock_client import SockClient
+
+from .service_client import ServiceClient
 
 _CURRENT_VERSION = "3"
 
@@ -53,9 +55,16 @@ class ServiceToken(abc.ABC):
     """A way of connecting to a running service process."""
 
     @abc.abstractmethod
-    def connect(self) -> SockClient:
+    def connect(
+        self,
+        *,
+        asyncer: asyncio_manager.AsyncioManager,
+    ) -> ServiceClient:
         """Connect to the service process.
 
+        Args:
+            asyncer: A started AsyncioManager for asyncio operations.
+
         Returns:
             A socket object for communicating with the service.
 
@@ -81,21 +90,25 @@ class UnixServiceToken(ServiceToken):
         self._path = path
 
     @override
-    def connect(self) -> SockClient:
+    def connect(
+        self,
+        *,
+        asyncer: asyncio_manager.AsyncioManager,
+    ) -> ServiceClient:
         if not ipc_support.SUPPORTS_UNIX:
             raise WandbServiceConnectionError("AF_UNIX socket not supported")
 
-        sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-
         try:
             # TODO: This may block indefinitely if the service is unhealthy.
-            sock.connect(self._path)
+            reader, writer = asyncer.run(
+                lambda: asyncio.open_unix_connection(self._path),
+            )
         except Exception as e:
             raise WandbServiceConnectionError(
                 f"Failed to connect to service on socket {self._path}",
             ) from e
 
-        return SockClient(sock)
+        return ServiceClient(asyncer, reader, writer)
 
     @override
     def _as_env_string(self):
@@ -128,18 +141,22 @@ class TCPServiceToken(ServiceToken):
         self._port = port
 
     @override
-    def connect(self) -> SockClient:
-        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-
+    def connect(
+        self,
+        *,
+        asyncer: asyncio_manager.AsyncioManager,
+    ) -> ServiceClient:
         try:
             # TODO: This may block indefinitely if the service is unhealthy.
-            sock.connect(("localhost", self._port))
+            reader, writer = asyncer.run(
+                lambda: asyncio.open_connection("localhost", self._port),
+            )
         except Exception as e:
             raise WandbServiceConnectionError(
                 f"Failed to connect to service on port {self._port}",
             ) from e
 
-        return SockClient(sock)
+        return ServiceClient(asyncer, reader, writer)
 
     @override
     def _as_env_string(self):