wandb 0.19.1__py3-none-musllinux_1_2_aarch64.whl

wandb/sdk/lib/retry.py

@@ -0,0 +1,289 @@
+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
+from wandb.util import CheckRetryFnType
+
+from .mailbox import ContextCancelledError
+
+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 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, **kwargs: Any) -> _R:  # noqa: C901
+        """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.
+        """
+        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
+
+        sleep_base: float = kwargs.pop("retry_sleep_base", 1)
+
+        # 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
+        )
+
+        sleep = sleep_base
+        now = NOW_FN()
+        start_time = now
+        start_time_triggered = None
+
+        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,
+                            "{} resolved after {}, resuming normal operation.".format(
+                                self._error_prefix, NOW_FN() - start_time
+                            ),
+                        )
+                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:
+                    raise
+
+                # always enforce num_retries no matter which type of exception was seen
+                if self._num_iter >= num_retries:
+                    raise
+
+                now = NOW_FN()
+
+                # 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
+
+                    # make sure that we haven't run out of time from secondary trigger
+                    if now - start_time_triggered >= retry_timedelta_triggered:
+                        raise
+
+                # always enforce the default timeout from start of retries
+                if now - start_time >= retry_timedelta:
+                    raise
+
+                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(
+                            "{} ({}), entering retry loop.".format(
+                                self._error_prefix, e.__class__.__name__
+                            )
+                        )
+                # if wandb.env.is_debug():
+                #     traceback.print_exc()
+            cancelled = self._sleep_check_cancelled(
+                sleep + random.random() * 0.25 * sleep, cancel_event=retry_cancel_event
+            )
+            if cancelled:
+                raise ContextCancelledError("retry timeout")
+            sleep *= 2
+            if sleep > self.MAX_SLEEP_SECONDS:
+                sleep = self.MAX_SLEEP_SECONDS
+            now = NOW_FN()
+
+            self._num_iter += 1
+
+
+_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,72 @@
+from dataclasses import dataclass
+from typing import Literal, Union, cast
+from urllib import parse
+
+_STEP = Literal["_step"]
+
+
+@dataclass
+class RunMoment:
+    """A moment in a run."""
+
+    run: str  # run name
+
+    # currently, the _step value to fork from. in future, this will be optional
+    value: Union[int, float]
+
+    # only step for now, in future this will be relaxed to be any metric
+    metric: _STEP = "_step"
+
+    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 ValueError(
+                f"Only int or float values are supported, got '{self.value}'."
+            )
+        if not isinstance(self.run, str):
+            raise ValueError(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=cast(_STEP, 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,38 @@
+"""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", "{}"))

wandb/sdk/lib/service_connection.py

@@ -0,0 +1,216 @@
+from __future__ import annotations
+
+import atexit
+import os
+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.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.service import service
+
+
+class WandbServiceNotOwnedError(Exception):
+    """Raised when the current process does not own the service process."""
+
+
+class WandbServiceConnectionError(Exception):
+    """Raised on failure to connect to the service process."""
+
+
+class WandbAttachFailedError(Exception):
+    """Raised if attaching to a run fails."""
+
+
+def connect_to_service(
+    settings: wandb_settings.Settings,
+) -> ServiceConnection:
+    """Connects to the service process, starting one up if necessary."""
+    conn = _try_connect_to_existing_service()
+    if conn:
+        return conn
+
+    return _start_and_connect_service(settings)
+
+
+def _try_connect_to_existing_service() -> ServiceConnection | None:
+    """Attemps to connect to an existing service process."""
+    token = service_token.get_service_token()
+    if not token:
+        return None
+
+    # Only localhost sockets are supported below.
+    assert token.host == "localhost"
+    client = SockClient()
+
+    try:
+        # TODO: This may block indefinitely if the service is unhealthy.
+        client.connect(token.port)
+
+    except Exception as e:
+        raise WandbServiceConnectionError(
+            "Failed to connect to internal service."
+        ) from e
+
+    return ServiceConnection(client=client, proc=None)
+
+
+def _start_and_connect_service(
+    settings: wandb_settings.Settings,
+) -> ServiceConnection:
+    """Starts a service process and returns a connection to it.
+
+    An atexit hook is registered to tear down the service process and wait for
+    it to complete. The hook does not run in processes started using the
+    multiprocessing module.
+    """
+    proc = service._Service(settings)
+    proc.start()
+
+    port = proc.sock_port
+    assert port
+    client = SockClient()
+    client.connect(port)
+
+    service_token.set_service_token(
+        parent_pid=os.getpid(),
+        transport="tcp",
+        host="localhost",
+        port=port,
+    )
+
+    hooks = ExitHooks()
+    hooks.hook()
+
+    def teardown_atexit():
+        conn.teardown(hooks.exit_code)
+
+    conn = ServiceConnection(
+        client=client,
+        proc=proc,
+        cleanup=lambda: atexit.unregister(teardown_atexit),
+    )
+
+    atexit.register(teardown_atexit)
+
+    return conn
+
+
+class ServiceConnection:
+    """A connection to the W&B internal service process."""
+
+    def __init__(
+        self,
+        client: SockClient,
+        proc: service._Service | None,
+        cleanup: Callable[[], None] | None = None,
+    ):
+        """Returns a new ServiceConnection.
+
+        Args:
+            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.
+        """
+        self._client = client
+        self._proc = proc
+        self._torn_down = False
+        self._cleanup = cleanup
+
+    def make_interface(self, mailbox: Mailbox) -> InterfaceBase:
+        """Returns an interface for communicating with the service."""
+        return InterfaceSock(self._client, mailbox)
+
+    def send_record(self, record: pb.Record) -> None:
+        """Sends data to the service."""
+        self._client.send_record_publish(record)
+
+    def inform_init(
+        self,
+        settings: wandb_settings_pb2.Settings,
+        run_id: str,
+    ) -> None:
+        """Sends an init request to the service."""
+        request = spb.ServerInformInitRequest()
+        request.settings.CopyFrom(settings)
+        request._info.stream_id = run_id
+        self._client.send(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)
+
+    def inform_attach(
+        self,
+        attach_id: str,
+    ) -> wandb_settings_pb2.Settings:
+        """Sends an attach request to the service.
+
+        Raises a WandbAttachFailedError if attaching is not possible.
+        """
+        request = spb.ServerInformAttachRequest()
+        request._info.stream_id = attach_id
+
+        try:
+            response = self._client.send_and_recv(inform_attach=request)
+            return response.inform_attach_response.settings
+        except SockClientTimeoutError:
+            raise WandbAttachFailedError(
+                "Could not attach because the run does not belong to"
+                " the current service process, or because the service"
+                " process is busy (unlikely)."
+            )
+
+    def inform_start(
+        self,
+        settings: wandb_settings_pb2.Settings,
+        run_id: str,
+    ) -> None:
+        """Sends a start request to the service."""
+        request = spb.ServerInformStartRequest()
+        request.settings.CopyFrom(settings)
+        request._info.stream_id = run_id
+        self._client.send(inform_start=request)
+
+    def teardown(self, exit_code: int) -> int:
+        """Shuts down the service process and returns its exit code.
+
+        This may only be called once.
+
+        Returns:
+            The exit code of the service process.
+
+        Raises:
+            WandbServiceNotOwnedError: If the current process did not start
+                the service process.
+        """
+        if not self._proc:
+            raise WandbServiceNotOwnedError(
+                "Cannot tear down service started by different process",
+            )
+
+        assert not self._torn_down
+        self._torn_down = True
+
+        if self._cleanup:
+            self._cleanup()
+
+        # 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,
+            )
+        )
+
+        return self._proc.join()