wandb 0.19.6__py3-none-macosx_11_0_arm64.whl → 0.19.7__py3-none-macosx_11_0_arm64.whl

wandb/sdk/lib/mailbox.py

@@ -1,442 +0,0 @@
-"""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."""
-
-
-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 _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