wandb 0.21.2__py3-none-macosx_12_0_x86_64.whl

wandb/sdk/lib/retry.py

@@ -0,0 +1,395 @@
+import abc
+import asyncio
+import datetime
+import functools
+import logging
+import os
+import random
+import threading
+import time
+from typing import Any, Awaitable, Callable, Generic, Optional, Tuple, Type, TypeVar
+
+from requests import HTTPError
+
+import wandb
+import wandb.errors
+from wandb.util import CheckRetryFnType
+
+logger = logging.getLogger(__name__)
+
+
+# To let tests mock out the retry logic's now()/sleep() funcs, this file
+# should only use these variables, not call the stdlib funcs directly.
+NOW_FN = datetime.datetime.now
+SLEEP_FN = time.sleep
+SLEEP_ASYNC_FN = asyncio.sleep
+
+
+class RetryCancelledError(wandb.errors.Error):
+    """A retry did not occur because it was cancelled."""
+
+
+class TransientError(Exception):
+    """Exception type designated for errors that may only be temporary.
+
+    Can have its own message and/or wrap another exception.
+    """
+
+    def __init__(
+        self, msg: Optional[str] = None, exc: Optional[BaseException] = None
+    ) -> None:
+        super().__init__(msg)
+        self.message = msg
+        self.exception = exc
+
+
+_R = TypeVar("_R")
+
+
+class Retry(Generic[_R]):
+    """Create a retryable version of a function.
+
+    Calling this will call the passed function, retrying if any exceptions in
+    retryable_exceptions are caught, with exponential backoff.
+    """
+
+    MAX_SLEEP_SECONDS = 5 * 60
+
+    def __init__(
+        self,
+        call_fn: Callable[..., _R],
+        retry_timedelta: Optional[datetime.timedelta] = None,
+        retry_cancel_event: Optional[threading.Event] = None,
+        num_retries: Optional[int] = None,
+        check_retry_fn: CheckRetryFnType = lambda e: True,
+        retryable_exceptions: Optional[Tuple[Type[Exception], ...]] = None,
+        error_prefix: str = "Network error",
+        retry_callback: Optional[Callable[[int, str], Any]] = None,
+    ) -> None:
+        self._call_fn = call_fn
+        self._check_retry_fn = check_retry_fn
+        self._error_prefix = error_prefix
+        self._last_print = datetime.datetime.now() - datetime.timedelta(minutes=1)
+        self._retry_timedelta = retry_timedelta
+        self._retry_cancel_event = retry_cancel_event
+        self._num_retries = num_retries
+        if retryable_exceptions is not None:
+            self._retryable_exceptions = retryable_exceptions
+        else:
+            self._retryable_exceptions = (TransientError,)
+        self._index = 0
+        self.retry_callback = retry_callback
+
+    def _sleep_check_cancelled(
+        self, wait_seconds: float, cancel_event: Optional[threading.Event]
+    ) -> bool:
+        if not cancel_event:
+            SLEEP_FN(wait_seconds)
+            return False
+        cancelled = cancel_event.wait(wait_seconds)
+        return cancelled
+
+    @property
+    def num_iters(self) -> int:
+        """The number of iterations the previous __call__ retried."""
+        return self._num_iter
+
+    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:
+            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.
+        """
+        if os.environ.get("WANDB_TEST"):
+            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
+
+        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)
+
+        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,
+        )
+
+        start_time = NOW_FN()
+        self._num_iter = 0
+
+        while True:
+            try:
+                result = self._call_fn(*args, **kwargs)
+
+            except self._retryable_exceptions as e:
+                if not retry_loop.should_retry(e):
+                    raise
+
+                if self._num_iter == 2:
+                    logger.info("Retry attempt failed:", exc_info=e)
+                    self._print_entered_retry_loop(e)
+
+                retry_loop.wait_before_retry()
+                self._num_iter += 1
+
+            else:
+                if self._num_iter > 2:
+                    self._print_recovered(start_time)
+
+                return result
+
+    def _print_entered_retry_loop(self, exception: Exception) -> None:
+        """Emit a message saying we've begun retrying.
+
+        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("Cancelled while retrying.")
+        else:
+            SLEEP_FN(sleep_amount)
+
+        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)
+
+
+def retriable(*args: Any, **kargs: Any) -> Callable[[_F], _F]:
+    def decorator(fn: _F) -> _F:
+        retrier: Retry[Any] = Retry(fn, *args, **kargs)
+
+        @functools.wraps(fn)
+        def wrapped_fn(*args: Any, **kargs: Any) -> Any:
+            return retrier(*args, **kargs)
+
+        return wrapped_fn  # type: ignore
+
+    return decorator
+
+
+class Backoff(abc.ABC):
+    """A backoff strategy: decides whether to sleep or give up when an exception is raised."""
+
+    @abc.abstractmethod
+    def next_sleep_or_reraise(self, exc: Exception) -> datetime.timedelta:
+        raise NotImplementedError  # pragma: no cover
+
+
+class ExponentialBackoff(Backoff):
+    """Jittered exponential backoff: sleep times increase ~exponentially up to some limit."""
+
+    def __init__(
+        self,
+        initial_sleep: datetime.timedelta,
+        max_sleep: datetime.timedelta,
+        max_retries: Optional[int] = None,
+        timeout_at: Optional[datetime.datetime] = None,
+    ) -> None:
+        self._next_sleep = min(max_sleep, initial_sleep)
+        self._max_sleep = max_sleep
+        self._remaining_retries = max_retries
+        self._timeout_at = timeout_at
+
+    def next_sleep_or_reraise(self, exc: Exception) -> datetime.timedelta:
+        if self._remaining_retries is not None:
+            if self._remaining_retries <= 0:
+                raise exc
+            self._remaining_retries -= 1
+
+        if self._timeout_at is not None and NOW_FN() > self._timeout_at:
+            raise exc
+
+        result, self._next_sleep = (
+            self._next_sleep,
+            min(self._max_sleep, self._next_sleep * (1 + random.random())),
+        )
+
+        return result
+
+
+class FilteredBackoff(Backoff):
+    """Re-raise any exceptions that fail a predicate; delegate others to another Backoff."""
+
+    def __init__(self, filter: Callable[[Exception], bool], wrapped: Backoff) -> None:
+        self._filter = filter
+        self._wrapped = wrapped
+
+    def next_sleep_or_reraise(self, exc: Exception) -> datetime.timedelta:
+        if not self._filter(exc):
+            raise exc
+        return self._wrapped.next_sleep_or_reraise(exc)
+
+
+async def retry_async(
+    backoff: Backoff,
+    fn: Callable[..., Awaitable[_R]],
+    *args: Any,
+    on_exc: Optional[Callable[[Exception], None]] = None,
+    **kwargs: Any,
+) -> _R:
+    """Call `fn` repeatedly until either it succeeds, or `backoff` decides we should give up.
+
+    Each time `fn` fails, `on_exc` is called with the exception.
+    """
+    while True:
+        try:
+            return await fn(*args, **kwargs)
+        except Exception as e:
+            if on_exc is not None:
+                on_exc(e)
+            await SLEEP_ASYNC_FN(backoff.next_sleep_or_reraise(e).total_seconds())

wandb/sdk/lib/run_moment.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Union
+from urllib import parse
+
+
+@dataclass
+class RunMoment:
+    """A moment in a run.
+
+    Defines a branching point in a finished run to fork or resume from.
+    A run moment is identified by a run ID and a metric value.
+    Currently, only the metric '_step' is supported.
+    """
+
+    run: str
+    """run ID"""
+
+    value: Union[int, float]
+    """Value of the metric."""
+
+    metric: str = "_step"
+    """Metric to use to determine the moment in the run.
+
+    Currently, only the metric '_step' is supported.
+    In future, this will be relaxed to be any metric.
+    """
+
+    def __post_init__(self):
+        if self.metric != "_step":
+            raise ValueError(
+                f"Only the metric '_step' is supported, got '{self.metric}'."
+            )
+        if not isinstance(self.value, (int, float)):
+            raise TypeError(
+                f"Only int or float values are supported, got '{self.value}'."
+            )
+        if not isinstance(self.run, str):
+            raise TypeError(f"Only string run names are supported, got '{self.run}'.")
+
+    @classmethod
+    def from_uri(cls, uri: str) -> RunMoment:
+        parsable = "runmoment://" + uri
+        parse_err = ValueError(
+            f"Could not parse passed run moment string '{uri}', "
+            f"expected format '<run>?<metric>=<numeric_value>'. "
+            f"Currently, only the metric '_step' is supported. "
+            f"Example: 'ans3bsax?_step=123'."
+        )
+
+        try:
+            parsed = parse.urlparse(parsable)
+        except ValueError as e:
+            raise parse_err from e
+
+        if parsed.scheme != "runmoment":
+            raise parse_err
+
+        # extract run, metric, value from parsed
+        if not parsed.netloc:
+            raise parse_err
+
+        run = parsed.netloc
+
+        if parsed.path or parsed.params or parsed.fragment:
+            raise parse_err
+
+        query = parse.parse_qs(parsed.query)
+        if len(query) != 1:
+            raise parse_err
+
+        metric = list(query.keys())[0]
+        if metric != "_step":
+            raise parse_err
+        value: str = query[metric][0]
+        try:
+            num_value = int(value) if value.isdigit() else float(value)
+        except ValueError as e:
+            raise parse_err from e
+
+        return cls(run=run, metric=metric, value=num_value)

wandb/sdk/lib/runid.py

@@ -0,0 +1,12 @@
+"""runid util."""
+
+import secrets
+import string
+
+
+def generate_id(length: int = 8) -> str:
+    """Generate a random base-36 string of `length` digits."""
+    # There are ~2.8T base-36 8-digit strings. If we generate 210k ids,
+    # we'll have a ~1% chance of collision.
+    alphabet = string.ascii_lowercase + string.digits
+    return "".join(secrets.choice(alphabet) for _ in range(length))

wandb/sdk/lib/server.py

@@ -0,0 +1,58 @@
+"""module server."""
+
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING, Any
+
+from wandb import util
+from wandb.apis import InternalApi
+
+if TYPE_CHECKING:
+    from wandb.sdk.wandb_settings import Settings
+
+
+class Server:
+    def __init__(
+        self,
+        settings: Settings,
+    ) -> None:
+        self._api = InternalApi(default_settings=settings)
+        self._viewer: dict[str, Any] = {}
+        self._flags: dict[str, Any] = {}
+        self._settings: Settings = settings
+
+    def query_with_timeout(self, timeout: int | float = 5) -> None:
+        if self._settings.x_disable_viewer:
+            return
+
+        async_viewer = util.async_call(self._api.viewer_server_info, timeout=timeout)
+        try:
+            viewer_tuple, viewer_thread = async_viewer()
+        except Exception:
+            return
+        if viewer_thread.is_alive():
+            # this is likely a DNS hang
+            return
+        # TODO(jhr): should we kill the thread?
+        self._viewer, self._serverinfo = viewer_tuple
+        self._flags = json.loads(self._viewer.get("flags", "{}"))
+
+    @property
+    def viewer(self) -> dict[str, Any]:
+        """Returns information about the currently authenticated user.
+
+        If the API key is valid, the following is returned:
+        - id
+        - entity
+        - username
+        - flags
+        - teams
+
+        If the API key is not valid or the server is not reachable,
+        an empty dict is returned.
+        """
+        if not self._viewer and not self._settings._offline:
+            self.query_with_timeout()
+
+        return self._viewer

wandb/sdk/lib/service/ipc_support.py

@@ -0,0 +1,13 @@
+"""Constants determining what IPC methods are supported."""
+
+import socket
+
+SUPPORTS_UNIX = hasattr(socket, "AF_UNIX")
+"""Whether Unix sockets are supported.
+
+AF_UNIX is not supported on Windows:
+https://github.com/python/cpython/issues/77589
+
+Windows has supported Unix sockets since ~2017, but support in Python is
+missing as of 2025.
+"""

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