wandb 0.18.2__py3-none-musllinux_1_2_x86_64.whl

wandb/sdk/lib/mailbox.py

@@ -0,0 +1,460 @@
+"""mailbox."""
+
+import secrets
+import string
+import threading
+import time
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple
+
+from wandb.errors import Error
+from wandb.proto import wandb_internal_pb2 as pb
+
+if TYPE_CHECKING:
+    from wandb.sdk.interface.interface_shared import InterfaceShared
+
+
+def _generate_address(length: int = 12) -> str:
+    address = "".join(
+        secrets.choice(string.ascii_lowercase + string.digits) for i in range(length)
+    )
+    return address
+
+
+class MailboxError(Error):
+    """Generic Mailbox Exception."""
+
+    pass
+
+
+class ContextCancelledError(MailboxError):
+    """Context cancelled Exception."""
+
+    pass
+
+
+class _MailboxWaitAll:
+    _event: threading.Event
+    _lock: threading.Lock
+    _handles: List["MailboxHandle"]
+    _failed_handles: int
+
+    def __init__(self) -> None:
+        self._event = threading.Event()
+        self._lock = threading.Lock()
+        self._handles = []
+        self._failed_handles = 0
+
+    def notify(self) -> None:
+        with self._lock:
+            self._event.set()
+
+    def _add_handle(self, handle: "MailboxHandle") -> None:
+        handle._slot._set_wait_all(self)
+        self._handles.append(handle)
+
+        # set wait_all event if an event has already been set before added to wait_all
+        if handle._slot._event.is_set():
+            self._event.set()
+
+    @property
+    def active_handles(self) -> List["MailboxHandle"]:
+        return [h for h in self._handles if not h._is_failed]
+
+    @property
+    def active_handles_count(self) -> int:
+        return len(self.active_handles)
+
+    @property
+    def failed_handles_count(self) -> int:
+        return self._failed_handles
+
+    def _mark_handle_failed(self, handle: "MailboxHandle") -> None:
+        handle._mark_failed()
+        self._failed_handles += 1
+
+    def clear_handles(self) -> None:
+        for handle in self._handles:
+            handle._slot._clear_wait_all()
+        self._handles = []
+
+    def _wait(self, timeout: float) -> bool:
+        return self._event.wait(timeout=timeout)
+
+    def _get_and_clear(self, timeout: float) -> List["MailboxHandle"]:
+        found: List[MailboxHandle] = []
+        if self._wait(timeout=timeout):
+            with self._lock:
+                remove_handles = []
+
+                # Look through handles for triggered events
+                for handle in self._handles:
+                    if handle._slot._event.is_set():
+                        found.append(handle)
+                        remove_handles.append(handle)
+
+                for handle in remove_handles:
+                    self._handles.remove(handle)
+
+                self._event.clear()
+        return found
+
+
+class _MailboxSlot:
+    _result: Optional[pb.Result]
+    _event: threading.Event
+    _lock: threading.Lock
+    _wait_all: Optional[_MailboxWaitAll]
+    _address: str
+    _abandoned: bool
+
+    def __init__(self, address: str) -> None:
+        self._result = None
+        self._event = threading.Event()
+        self._lock = threading.Lock()
+        self._address = address
+        self._wait_all = None
+        self._abandoned = False
+
+    def _set_wait_all(self, wait_all: _MailboxWaitAll) -> None:
+        assert not self._wait_all, "Only one caller can wait_all for a slot at a time"
+        self._wait_all = wait_all
+
+    def _clear_wait_all(self) -> None:
+        self._wait_all = None
+
+    def _wait(self, timeout: float) -> bool:
+        return self._event.wait(timeout=timeout)
+
+    def _get_and_clear(self, timeout: float) -> Tuple[Optional[pb.Result], bool]:
+        found = None
+        if self._wait(timeout=timeout):
+            with self._lock:
+                found = self._result
+                self._event.clear()
+        abandoned = self._abandoned
+        return found, abandoned
+
+    def _deliver(self, result: pb.Result) -> None:
+        with self._lock:
+            self._result = result
+            self._event.set()
+
+        if self._wait_all:
+            self._wait_all.notify()
+
+    def _notify_abandon(self) -> None:
+        self._abandoned = True
+        with self._lock:
+            self._event.set()
+
+        if self._wait_all:
+            self._wait_all.notify()
+
+
+class MailboxProbe:
+    _result: Optional[pb.Result]
+    _handle: Optional["MailboxHandle"]
+
+    def __init__(self) -> None:
+        self._handle = None
+        self._result = None
+
+    def set_probe_result(self, result: pb.Result) -> None:
+        self._result = result
+
+    def get_probe_result(self) -> Optional[pb.Result]:
+        return self._result
+
+    def get_mailbox_handle(self) -> Optional["MailboxHandle"]:
+        return self._handle
+
+    def set_mailbox_handle(self, handle: "MailboxHandle") -> None:
+        self._handle = handle
+
+
+class MailboxProgress:
+    _percent_done: float
+    _handle: "MailboxHandle"
+    _probe_handles: List[MailboxProbe]
+    _stopped: bool
+
+    def __init__(self, _handle: "MailboxHandle") -> None:
+        self._handle = _handle
+        self._percent_done = 0.0
+        self._probe_handles = []
+        self._stopped = False
+
+    @property
+    def percent_done(self) -> float:
+        return self._percent_done
+
+    def set_percent_done(self, percent_done: float) -> None:
+        self._percent_done = percent_done
+
+    def add_probe_handle(self, probe_handle: MailboxProbe) -> None:
+        self._probe_handles.append(probe_handle)
+
+    def get_probe_handles(self) -> List[MailboxProbe]:
+        return self._probe_handles
+
+    def wait_stop(self) -> None:
+        self._stopped = True
+
+    @property
+    def _is_stopped(self) -> bool:
+        return self._stopped
+
+
+class MailboxProgressAll:
+    _progress_handles: List[MailboxProgress]
+
+    def __init__(self) -> None:
+        self._progress_handles = []
+
+    def add_progress_handle(self, progress_handle: MailboxProgress) -> None:
+        self._progress_handles.append(progress_handle)
+
+    def get_progress_handles(self) -> List[MailboxProgress]:
+        # only return progress handles for not failed handles
+        return [ph for ph in self._progress_handles if not ph._handle._is_failed]
+
+
+class MailboxHandle:
+    _mailbox: "Mailbox"
+    _slot: _MailboxSlot
+    _on_probe: Optional[Callable[[MailboxProbe], None]]
+    _on_progress: Optional[Callable[[MailboxProgress], None]]
+    _interface: Optional["InterfaceShared"]
+    _keepalive: bool
+    _failed: bool
+
+    def __init__(self, mailbox: "Mailbox", slot: _MailboxSlot) -> None:
+        self._mailbox = mailbox
+        self._slot = slot
+        self._on_probe = None
+        self._on_progress = None
+        self._interface = None
+        self._keepalive = False
+        self._failed = False
+
+    def add_probe(self, on_probe: Callable[[MailboxProbe], None]) -> None:
+        self._on_probe = on_probe
+
+    def add_progress(self, on_progress: Callable[[MailboxProgress], None]) -> None:
+        self._on_progress = on_progress
+
+    def _time(self) -> float:
+        return time.monotonic()
+
+    def wait(  # noqa: C901
+        self,
+        *,
+        timeout: float,
+        on_probe: Optional[Callable[[MailboxProbe], None]] = None,
+        on_progress: Optional[Callable[[MailboxProgress], None]] = None,
+        release: bool = True,
+        cancel: bool = False,
+    ) -> Optional[pb.Result]:
+        probe_handle: Optional[MailboxProbe] = None
+        progress_handle: Optional[MailboxProgress] = None
+        found: Optional[pb.Result] = None
+        start_time = self._time()
+        percent_done = 0.0
+        progress_sent = False
+        wait_timeout = 1.0
+        if timeout >= 0:
+            wait_timeout = min(timeout, wait_timeout)
+
+        on_progress = on_progress or self._on_progress
+        if on_progress:
+            progress_handle = MailboxProgress(_handle=self)
+
+        on_probe = on_probe or self._on_probe
+        if on_probe:
+            probe_handle = MailboxProbe()
+            if progress_handle:
+                progress_handle.add_probe_handle(probe_handle)
+
+        while True:
+            if self._keepalive and self._interface:
+                if self._interface._transport_keepalive_failed():
+                    raise MailboxError("transport failed")
+
+            found, abandoned = self._slot._get_and_clear(timeout=wait_timeout)
+            if found:
+                # Always update progress to 100% when done
+                if on_progress and progress_handle and progress_sent:
+                    progress_handle.set_percent_done(100)
+                    on_progress(progress_handle)
+                break
+            if abandoned:
+                break
+            now = self._time()
+            if timeout >= 0:
+                if now >= start_time + timeout:
+                    # todo: communicate that we timed out
+                    break
+            if on_probe and probe_handle:
+                on_probe(probe_handle)
+            if on_progress and progress_handle:
+                if timeout > 0:
+                    percent_done = min((now - start_time) / timeout, 1.0)
+                progress_handle.set_percent_done(percent_done)
+                on_progress(progress_handle)
+                if progress_handle._is_stopped:
+                    break
+                progress_sent = True
+        if not found and cancel:
+            self._cancel()
+        if release:
+            self._release()
+        return found
+
+    def _cancel(self) -> None:
+        mailbox_slot = self.address
+        if self._interface:
+            self._interface.publish_cancel(mailbox_slot)
+
+    def _release(self) -> None:
+        self._mailbox._release_slot(self.address)
+
+    def abandon(self) -> None:
+        self._slot._notify_abandon()
+        self._release()
+
+    @property
+    def _is_failed(self) -> bool:
+        return self._failed
+
+    def _mark_failed(self) -> None:
+        self._failed = True
+
+    @property
+    def address(self) -> str:
+        return self._slot._address
+
+
+class Mailbox:
+    _slots: Dict[str, _MailboxSlot]
+    _keepalive: bool
+
+    def __init__(self) -> None:
+        self._slots = {}
+        self._keepalive = False
+
+    def enable_keepalive(self) -> None:
+        self._keepalive = True
+
+    def wait(
+        self,
+        handle: MailboxHandle,
+        *,
+        timeout: float,
+        on_progress: Optional[Callable[[MailboxProgress], None]] = None,
+        cancel: bool = False,
+    ) -> Optional[pb.Result]:
+        return handle.wait(timeout=timeout, on_progress=on_progress, cancel=cancel)
+
+    def _time(self) -> float:
+        return time.monotonic()
+
+    def wait_all(
+        self,
+        handles: List[MailboxHandle],
+        *,
+        timeout: float,
+        on_progress_all: Optional[Callable[[MailboxProgressAll], None]] = None,
+    ) -> bool:
+        progress_all_handle: Optional[MailboxProgressAll] = None
+
+        if on_progress_all:
+            progress_all_handle = MailboxProgressAll()
+
+        wait_all = _MailboxWaitAll()
+        for handle in handles:
+            wait_all._add_handle(handle)
+            if progress_all_handle and handle._on_progress:
+                progress_handle = MailboxProgress(_handle=handle)
+                if handle._on_probe:
+                    probe_handle = MailboxProbe()
+                    progress_handle.add_probe_handle(probe_handle)
+                progress_all_handle.add_progress_handle(progress_handle)
+
+        start_time = self._time()
+
+        while wait_all.active_handles_count > 0:
+            # Make sure underlying interfaces are still up
+            if self._keepalive:
+                for handle in wait_all.active_handles:
+                    if not handle._interface:
+                        continue
+                    if handle._interface._transport_keepalive_failed():
+                        wait_all._mark_handle_failed(handle)
+
+                # if there are no valid handles left, either break or raise exception
+                if not wait_all.active_handles_count:
+                    if wait_all.failed_handles_count:
+                        wait_all.clear_handles()
+                        raise MailboxError("transport failed")
+                    break
+
+            # wait for next event
+            wait_all._get_and_clear(timeout=1)
+
+            # TODO: we can do more careful timekeeping and not run probes and progress
+            # indications until a full second elapses in the case where we found a wait_all
+            # event.  Extra probes should be ok for now.
+
+            if progress_all_handle and on_progress_all:
+                # Run all probe handles
+                for progress_handle in progress_all_handle.get_progress_handles():
+                    for probe_handle in progress_handle.get_probe_handles():
+                        if (
+                            progress_handle._handle
+                            and progress_handle._handle._on_probe
+                        ):
+                            progress_handle._handle._on_probe(probe_handle)
+
+                on_progress_all(progress_all_handle)
+
+            now = self._time()
+            if timeout >= 0 and now >= start_time + timeout:
+                break
+
+        return wait_all.active_handles_count == 0
+
+    def deliver(self, result: pb.Result) -> None:
+        mailbox = result.control.mailbox_slot
+        slot = self._slots.get(mailbox)
+        if not slot:
+            return
+        slot._deliver(result)
+
+    def _allocate_slot(self) -> _MailboxSlot:
+        address = _generate_address()
+        slot = _MailboxSlot(address=address)
+        self._slots[address] = slot
+        return slot
+
+    def _release_slot(self, address: str) -> None:
+        self._slots.pop(address, None)
+
+    def get_handle(self) -> MailboxHandle:
+        slot = self._allocate_slot()
+        handle = MailboxHandle(mailbox=self, slot=slot)
+        return handle
+
+    def _deliver_record(
+        self, record: pb.Record, interface: "InterfaceShared"
+    ) -> MailboxHandle:
+        handle = self.get_handle()
+        handle._interface = interface
+        handle._keepalive = self._keepalive
+        record.control.mailbox_slot = handle.address
+        try:
+            interface._publish(record)
+        except Exception:
+            interface._transport_mark_failed()
+            raise
+        interface._transport_mark_success()
+        return handle

wandb/sdk/lib/module.py

@@ -0,0 +1,69 @@
+#
+import wandb
+
+from . import preinit
+
+
+def set_global(
+    run=None,
+    config=None,
+    log=None,
+    summary=None,
+    save=None,
+    use_artifact=None,
+    log_artifact=None,
+    define_metric=None,
+    alert=None,
+    plot_table=None,
+    mark_preempting=None,
+    log_model=None,
+    use_model=None,
+    link_model=None,
+):
+    if run:
+        wandb.run = run
+    if config is not None:
+        wandb.config = config
+    if log:
+        wandb.log = log
+    if summary is not None:
+        wandb.summary = summary
+    if save:
+        wandb.save = save
+    if use_artifact:
+        wandb.use_artifact = use_artifact
+    if log_artifact:
+        wandb.log_artifact = log_artifact
+    if define_metric:
+        wandb.define_metric = define_metric
+    if plot_table:
+        wandb.plot_table = plot_table
+    if alert:
+        wandb.alert = alert
+    if mark_preempting:
+        wandb.mark_preempting = mark_preempting
+    if log_model:
+        wandb.log_model = log_model
+    if use_model:
+        wandb.use_model = use_model
+    if link_model:
+        wandb.link_model = link_model
+
+
+def unset_globals():
+    wandb.run = None
+    wandb.config = preinit.PreInitObject("wandb.config")
+    wandb.summary = preinit.PreInitObject("wandb.summary")
+    wandb.log = preinit.PreInitCallable("wandb.log", wandb.wandb_sdk.wandb_run.Run.log)
+    wandb.save = preinit.PreInitCallable(
+        "wandb.save", wandb.wandb_sdk.wandb_run.Run.save
+    )
+    wandb.use_artifact = preinit.PreInitCallable(
+        "wandb.use_artifact", wandb.wandb_sdk.wandb_run.Run.use_artifact
+    )
+    wandb.log_artifact = preinit.PreInitCallable(
+        "wandb.log_artifact", wandb.wandb_sdk.wandb_run.Run.log_artifact
+    )
+    wandb.define_metric = preinit.PreInitCallable(
+        "wandb.define_metric", wandb.wandb_sdk.wandb_run.Run.define_metric
+    )

wandb/sdk/lib/paths.py

@@ -0,0 +1,106 @@
+import os
+import platform
+from functools import wraps
+from pathlib import PurePath, PurePosixPath
+from typing import Any, NewType, Union
+
+# Path _inputs_ should generally accept any kind of path. This is named the same and
+# modeled after the hint defined in the Python standard library's `typeshed`:
+# https://github.com/python/typeshed/blob/0b1cd5989669544866213807afa833a88f649ee7/stdlib/_typeshed/__init__.pyi#L56-L65
+StrPath = Union[str, "os.PathLike[str]"]
+
+# A native path to a file on a local filesystem.
+FilePathStr = NewType("FilePathStr", str)
+
+URIStr = NewType("URIStr", str)
+
+
+class LogicalPath(str):
+    """A string that represents a path relative to an artifact or run.
+
+    The format of the string is always as a POSIX path, e.g. "foo/bar.txt".
+
+    A neat trick is that you can use this class as if it were a PurePosixPath. E.g.:
+    ```
+    >>> path = LogicalPath("foo/bar.txt")
+    >>> path.parts
+    ('foo', 'bar.txt')
+    >>> path.parent / "baz.txt"
+    'foo/baz.txt'
+    >>> type(path.relative_to("foo"))
+    LogicalPath
+    ```
+    """
+
+    # It should probably always be a relative path, but that would be a behavior change.
+    #
+    # These strings used to be the output of `to_forward_slash_path`, which only works
+    # with strings and whose behavior is pretty simple:
+    # ```
+    # if platform.system() == "Windows":
+    #     path = path.replace("\\", "/")
+    # ```
+    #
+    # This results in some weird things, such as backslashes being allowed from
+    # non-Windows platforms (which would probably break if such an artifact was used
+    # from Windows) and anchors or absolute paths being allowed. E.g., the Windows path
+    # "C:\foo\bar.txt" becomes "C:/foo/bar.txt", which then would mount as
+    # "./artifacts/artifact_name:v0/C:/foo/bar.txt" on MacOS and as
+    # "./artifacts/artifact_name-v0/C-/foo/bar.txt" on Windows.
+    #
+    # This implementation preserves behavior for strings but attempts to sanitize other
+    # formerly unsupported inputs more aggressively. It uses the `.as_posix()` form of
+    # pathlib objects rather than the `str()` form to reduce how often identical inputs
+    # will result in different outputs on different platforms; however, it doesn't alter
+    # absolute paths or check for prohibited characters etc.
+
+    def __new__(cls, path: StrPath) -> "LogicalPath":
+        if isinstance(path, LogicalPath):
+            return super().__new__(cls, path)
+        if hasattr(path, "as_posix"):
+            path = PurePosixPath(path.as_posix())
+            return super().__new__(cls, str(path))
+        if hasattr(path, "__fspath__"):
+            path = path.__fspath__()  # Can be str or bytes.
+        if isinstance(path, bytes):
+            path = os.fsdecode(path)
+        # For historical reasons we have to convert backslashes to forward slashes, but
+        # only on Windows, and need to do it before any pathlib operations.
+        if platform.system() == "Windows":
+            path = path.replace("\\", "/")
+        # This weird contortion and the one above are because in some unusual cases
+        # PurePosixPath(path.as_posix()).as_posix() != path.as_posix().
+        path = PurePath(path).as_posix()
+        return super().__new__(cls, str(PurePosixPath(path)))
+
+    def to_path(self) -> PurePosixPath:
+        """Convert this path to a PurePosixPath."""
+        return PurePosixPath(self)
+
+    def __getattr__(self, attr: str) -> Any:
+        """Act like a subclass of PurePosixPath for all methods not defined on str."""
+        try:
+            result = getattr(self.to_path(), attr)
+        except AttributeError as e:
+            raise AttributeError(f"LogicalPath has no attribute {attr!r}") from e
+
+        if isinstance(result, PurePosixPath):
+            return LogicalPath(result)
+
+        # If the result is a callable (a method), wrap it so that it has the same
+        # behavior: if the call result returns a PurePosixPath, return a LogicalPath.
+        if callable(result):
+
+            @wraps(result)
+            def wrapper(*args: Any, **kwargs: Any) -> Any:
+                inner_result = result(*args, **kwargs)
+                if isinstance(inner_result, PurePosixPath):
+                    return LogicalPath(inner_result)
+                return inner_result
+
+            return wrapper
+        return result
+
+    def __truediv__(self, other: StrPath) -> "LogicalPath":
+        """Act like a PurePosixPath for the / operator, but return a LogicalPath."""
+        return LogicalPath(self.to_path() / LogicalPath(other))

wandb/sdk/lib/preinit.py

@@ -0,0 +1,42 @@
+from typing import Any, Callable, Optional
+
+import wandb
+
+
+class PreInitObject:
+    def __init__(self, name: str, destination: Optional[Any] = None) -> None:
+        self._name = name
+
+        if destination is not None:
+            self.__doc__ = destination.__doc__
+
+    def __getitem__(self, key: str) -> None:
+        raise wandb.Error(f"You must call wandb.init() before {self._name}[{key!r}]")
+
+    def __setitem__(self, key: str, value: Any) -> Any:
+        raise wandb.Error(f"You must call wandb.init() before {self._name}[{key!r}]")
+
+    def __setattr__(self, key: str, value: Any) -> Any:
+        if not key.startswith("_"):
+            raise wandb.Error(f"You must call wandb.init() before {self._name}.{key}")
+        else:
+            return object.__setattr__(self, key, value)
+
+    def __getattr__(self, key: str) -> Any:
+        if not key.startswith("_"):
+            raise wandb.Error(f"You must call wandb.init() before {self._name}.{key}")
+        else:
+            raise AttributeError
+
+
+def PreInitCallable(  # noqa: N802
+    name: str, destination: Optional[Any] = None
+) -> Callable:
+    def preinit_wrapper(*args: Any, **kwargs: Any) -> Any:
+        raise wandb.Error(f"You must call wandb.init() before {name}()")
+
+    preinit_wrapper.__name__ = str(name)
+    if destination:
+        preinit_wrapper.__wrapped__ = destination  # type: ignore
+        preinit_wrapper.__doc__ = destination.__doc__
+    return preinit_wrapper