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

wandb/sdk/mailbox/mailbox.py

@@ -7,6 +7,7 @@ import threading
 
 from wandb.proto import wandb_internal_pb2 as pb
 from wandb.proto import wandb_server_pb2 as spb
+from wandb.sdk.lib import asyncio_manager
 
 from .mailbox_handle import MailboxHandle
 from .response_handle import MailboxResponseHandle
@@ -27,7 +28,8 @@ class Mailbox:
     service process becomes unreachable.
     """
 
-    def __init__(self) -> None:
+    def __init__(self, asyncer: asyncio_manager.AsyncioManager) -> None:
+        self._asyncer = asyncer
         self._handles: dict[str, MailboxResponseHandle] = {}
         self._handles_lock = threading.Lock()
         self._closed = False
@@ -51,11 +53,17 @@ class Mailbox:
                 cannot be created.
         """
         if isinstance(request, spb.ServerRequest):
-            if address := request.request_id:
+            if (address := request.request_id) or (
+                address := request.record_publish.control.mailbox_slot
+            ):
                 raise ValueError(f"Request already has an address ({address})")
 
             address = self._new_address()
             request.request_id = address
+            if request.HasField("record_publish"):
+                request.record_publish.control.mailbox_slot = address
+            if request.HasField("record_communicate"):
+                request.record_communicate.control.mailbox_slot = address
         else:
             if address := request.control.mailbox_slot:
                 raise ValueError(f"Request already has an address ({address})")
@@ -67,7 +75,7 @@ class Mailbox:
             if self._closed:
                 raise MailboxClosedError()
 
-            handle = MailboxResponseHandle(address)
+            handle = MailboxResponseHandle(address, asyncer=self._asyncer)
             self._handles[address] = handle
 
         return handle
@@ -92,7 +100,7 @@ class Mailbox:
 
         return address
 
-    def deliver(self, response: spb.ServerResponse) -> None:
+    async def deliver(self, response: spb.ServerResponse) -> None:
         """Deliver a response from the service.
 
         If the response address is invalid, this does nothing.
@@ -116,7 +124,7 @@ class Mailbox:
         # It is not an error if there is no handle for the address:
         # handles can be abandoned if the result is no longer needed.
         if handle:
-            handle.deliver(response)
+            await handle.deliver(response)
 
     def close(self) -> None:
         """Indicate no further responses will be delivered.

wandb/sdk/mailbox/mailbox_handle.py

@@ -5,6 +5,8 @@ from typing import TYPE_CHECKING, Callable, Generic
 
 from typing_extensions import TypeVar, override
 
+from wandb.sdk.lib import asyncio_manager
+
 # Necessary to break an import loop.
 if TYPE_CHECKING:
     from wandb.sdk.interface import interface
@@ -19,7 +21,18 @@ class HandleAbandonedError(Exception):
 
 
 class MailboxHandle(abc.ABC, Generic[_T]):
-    """A thread-safe handle that allows waiting for a response to a request."""
+    """A handle for waiting on a response to a request."""
+
+    def __init__(self, asyncer: asyncio_manager.AsyncioManager) -> None:
+        self._asyncer = asyncer
+
+    @property
+    def asyncer(self) -> asyncio_manager.AsyncioManager:
+        """The asyncio thread to which the handle belongs.
+
+        The handle's async methods must be run using this object.
+        """
+        return self._asyncer
 
     def map(self, fn: Callable[[_T], _S]) -> MailboxHandle[_S]:
         """Returns a transformed handle.
@@ -35,7 +48,11 @@ class MailboxHandle(abc.ABC, Generic[_T]):
 
     @abc.abstractmethod
     def abandon(self) -> None:
-        """Abandon the handle, indicating it will not receive a response."""
+        """Abandon the handle, indicating it will not receive a response.
+
+        This may not happen immediately: it is possible for an existing
+        call to `wait_async` to complete successfully after this method returns.
+        """
 
     @abc.abstractmethod
     def cancel(self, iface: interface.InterfaceBase) -> None:
@@ -48,14 +65,12 @@ class MailboxHandle(abc.ABC, Generic[_T]):
             iface: The interface on which to publish the cancel request.
         """
 
-    @abc.abstractmethod
-    def check(self) -> _T | None:
-        """Returns the response if it's ready."""
-
     @abc.abstractmethod
     def wait_or(self, *, timeout: float | None) -> _T:
         """Wait for a response or a timeout.
 
+        It is an error to call this from an async function.
+
         Args:
             timeout: A finite number of seconds or None to never time out.
                 If less than or equal to zero, times out immediately unless
@@ -95,6 +110,7 @@ class _MailboxMappedHandle(Generic[_S], MailboxHandle[_S]):
         handle: MailboxHandle[_T],
         fn: Callable[[_T], _S],
     ) -> None:
+        super().__init__(handle.asyncer)
         self._handle = handle
         self._fn = fn
 
@@ -106,13 +122,6 @@ class _MailboxMappedHandle(Generic[_S], MailboxHandle[_S]):
     def cancel(self, iface: interface.InterfaceBase) -> None:
         self._handle.cancel(iface)
 
-    @override
-    def check(self) -> _S | None:
-        if response := self._handle.check():
-            return self._fn(response)
-        else:
-            return None
-
     @override
     def wait_or(self, *, timeout: float | None) -> _S:
         return self._fn(self._handle.wait_or(timeout=timeout))

wandb/sdk/mailbox/response_handle.py

@@ -2,12 +2,12 @@ from __future__ import annotations
 
 import asyncio
 import math
-import threading
 from typing import TYPE_CHECKING
 
 from typing_extensions import override
 
 from wandb.proto import wandb_server_pb2 as spb
+from wandb.sdk.lib import asyncio_manager
 
 from .mailbox_handle import HandleAbandonedError, MailboxHandle
 
@@ -19,33 +19,36 @@ if TYPE_CHECKING:
 class MailboxResponseHandle(MailboxHandle[spb.ServerResponse]):
     """A general handle for any ServerResponse."""
 
-    def __init__(self, address: str) -> None:
+    def __init__(
+        self,
+        address: str,
+        *,
+        asyncer: asyncio_manager.AsyncioManager,
+    ) -> None:
+        super().__init__(asyncer)
+
         self._address = address
-        self._lock = threading.Lock()
-        self._event = threading.Event()
 
         self._abandoned = False
         self._response: spb.ServerResponse | None = None
 
-        self._asyncio_events: dict[asyncio.Event, _AsyncioEvent] = dict()
+        # Initialized on first use in the asyncio thread.
+        self._done_event: asyncio.Event | None = None
 
-    def deliver(self, response: spb.ServerResponse) -> None:
-        """Deliver the response.
+    async def deliver(self, response: spb.ServerResponse) -> None:
+        if self._abandoned:
+            return
 
-        This may only be called once. It is an error to respond to the same
-        request more than once. It is a no-op if the handle has been abandoned.
-        """
-        with self._lock:
-            if self._abandoned:
-                return
+        if self._response:
+            raise ValueError(
+                f"A response has already been delivered to {self._address}."
+            )
 
-            if self._response:
-                raise ValueError(
-                    f"A response has already been delivered to {self._address}."
-                )
+        self._response = response
 
-            self._response = response
-            self._signal_done()
+        if not self._done_event:
+            self._done_event = asyncio.Event()
+        self._done_event.set()
 
     @override
     def cancel(self, iface: interface.InterfaceBase) -> None:
@@ -54,110 +57,43 @@ class MailboxResponseHandle(MailboxHandle[spb.ServerResponse]):
 
     @override
     def abandon(self) -> None:
-        with self._lock:
+        async def impl() -> None:
             self._abandoned = True
-            self._signal_done()
-
-    def _signal_done(self) -> None:
-        """Indicate that the handle either got a response or became abandoned.
 
-        The lock must be held.
-        """
-        # Unblock threads blocked on `wait_or`.
-        self._event.set()
+            if not self._done_event:
+                self._done_event = asyncio.Event()
+            self._done_event.set()
 
-        # Unblock asyncio loops blocked on `wait_async`.
-        for asyncio_event in self._asyncio_events.values():
-            asyncio_event.set_threadsafe()
-        self._asyncio_events.clear()
-
-    @override
-    def check(self) -> spb.ServerResponse | None:
-        with self._lock:
-            return self._response
+        self.asyncer.run_soon(impl)
 
     @override
     def wait_or(self, *, timeout: float | None) -> spb.ServerResponse:
-        if timeout is not None and not math.isfinite(timeout):
-            raise ValueError("Timeout must be finite or None.")
-
-        if not self._event.wait(timeout=timeout):
-            raise TimeoutError(
-                f"Timed out waiting for response on {self._address}",
-            )
-
-        with self._lock:
-            if self._response:
-                return self._response
-
-            assert self._abandoned
-            raise HandleAbandonedError()
+        return self.asyncer.run(lambda: self.wait_async(timeout=timeout))
 
     @override
     async def wait_async(self, *, timeout: float | None) -> spb.ServerResponse:
         if timeout is not None and not math.isfinite(timeout):
             raise ValueError("Timeout must be finite or None.")
 
-        evt = asyncio.Event()
-        self._add_asyncio_event(asyncio.get_event_loop(), evt)
+        if not self._done_event:
+            self._done_event = asyncio.Event()
 
         try:
-            await asyncio.wait_for(evt.wait(), timeout=timeout)
+            await asyncio.wait_for(self._done_event.wait(), timeout=timeout)
 
         except (asyncio.TimeoutError, TimeoutError) as e:
-            with self._lock:
-                if self._response:
-                    return self._response
-                elif self._abandoned:
-                    raise HandleAbandonedError()
-                else:
-                    raise TimeoutError(
-                        f"Timed out waiting for response on {self._address}"
-                    ) from e
-
-        else:
-            with self._lock:
-                if self._response:
-                    return self._response
-
-                assert self._abandoned
+            if self._response:
+                return self._response
+            elif self._abandoned:
                 raise HandleAbandonedError()
-
-        finally:
-            self._forget_asyncio_event(evt)
-
-    def _add_asyncio_event(
-        self,
-        loop: asyncio.AbstractEventLoop,
-        event: asyncio.Event,
-    ) -> None:
-        """Add an event to signal when a response is received.
-
-        If a response already exists, this notifies the event loop immediately.
-        """
-        asyncio_event = _AsyncioEvent(loop, event)
-
-        with self._lock:
-            if self._response or self._abandoned:
-                asyncio_event.set_threadsafe()
             else:
-                self._asyncio_events[event] = asyncio_event
-
-    def _forget_asyncio_event(self, event: asyncio.Event) -> None:
-        """Cancel signalling an event when a response is received."""
-        with self._lock:
-            self._asyncio_events.pop(event, None)
+                raise TimeoutError(
+                    f"Timed out waiting for response on {self._address}"
+                ) from e
 
+        else:
+            if self._response:
+                return self._response
 
-class _AsyncioEvent:
-    def __init__(
-        self,
-        loop: asyncio.AbstractEventLoop,
-        event: asyncio.Event,
-    ):
-        self._loop = loop
-        self._event = event
-
-    def set_threadsafe(self) -> None:
-        """Set the asyncio event in its own loop."""
-        self._loop.call_soon_threadsafe(self._event.set)
+            assert self._abandoned
+            raise HandleAbandonedError()

wandb/sdk/mailbox/wait_with_progress.py

@@ -14,7 +14,6 @@ def wait_with_progress(
     handle: MailboxHandle[_T],
     *,
     timeout: float | None,
-    progress_after: float,
     display_progress: Callable[[], Coroutine[Any, Any, None]],
 ) -> _T:
     """Wait for a handle, possibly displaying progress to the user.
@@ -24,7 +23,6 @@ def wait_with_progress(
     return wait_all_with_progress(
         [handle],
         timeout=timeout,
-        progress_after=progress_after,
         display_progress=display_progress,
     )[0]
 
@@ -33,7 +31,6 @@ def wait_all_with_progress(
     handle_list: list[MailboxHandle[_T]],
     *,
     timeout: float | None,
-    progress_after: float,
     display_progress: Callable[[], Coroutine[Any, Any, None]],
 ) -> list[_T]:
     """Wait for multiple handles, possibly displaying progress to the user.
@@ -42,18 +39,14 @@ def wait_all_with_progress(
         handle_list: The handles to wait for.
         timeout: A number of seconds after which to raise a TimeoutError,
             or None if this should never timeout.
-        progress_after: A number of seconds after which to start the
-            display_progress callback. Starting the callback creates a thread
-            and starts an asyncio loop, so we want to avoid doing it if
-            the handle is resolved quickly.
         display_progress: An asyncio function that displays progress to
-            the user. This function is executed on a new thread and cancelled
-            if the timeout is exceeded.
+            the user. This function runs using the handles' AsyncioManager.
 
     Returns:
         A list where the Nth item is the Nth handle's result.
 
     Raises:
+        ValueError: If the handles live in different asyncio threads.
         TimeoutError: If the overall timeout expires.
         HandleAbandonedError: If any handle becomes abandoned.
         Exception: Any exception from the display function is propagated.
@@ -61,16 +54,13 @@ def wait_all_with_progress(
     if not handle_list:
         return []
 
-    if timeout is not None and timeout <= progress_after:
-        return _wait_handles(handle_list, timeout=timeout)
+    asyncer = handle_list[0].asyncer
+    for handle in handle_list:
+        if handle.asyncer is not asyncer:
+            raise ValueError("Handles have different AsyncioManagers.")
 
     start_time = time.monotonic()
 
-    try:
-        return _wait_handles(handle_list, timeout=progress_after)
-    except TimeoutError:
-        pass
-
     async def progress_loop_with_timeout() -> list[_T]:
         with asyncio_compat.cancel_on_exit(display_progress()):
             if timeout is not None:
@@ -84,32 +74,7 @@ def wait_all_with_progress(
                 timeout=remaining_timeout,
             )
 
-    return asyncio_compat.run(progress_loop_with_timeout)
-
-
-def _wait_handles(
-    handle_list: list[MailboxHandle[_T]],
-    *,
-    timeout: float,
-) -> list[_T]:
-    """Wait for multiple mailbox handles.
-
-    Returns:
-        Each handle's result, in the same order as the given handles.
-
-    Raises:
-        TimeoutError: If the overall timeout expires.
-        HandleAbandonedError: If any handle becomes abandoned.
-    """
-    results: list[_T] = []
-
-    start_time = time.monotonic()
-    for handle in handle_list:
-        elapsed_time = time.monotonic() - start_time
-        remaining_timeout = timeout - elapsed_time
-        results.append(handle.wait_or(timeout=remaining_timeout))
-
-    return results
+    return asyncer.run(progress_loop_with_timeout)
 
 
 async def _wait_handles_async(

wandb/sdk/wandb_init.py

@@ -11,7 +11,6 @@ For more on using `wandb.init()`, including code snippets, check out our
 from __future__ import annotations
 
 import contextlib
-import copy
 import dataclasses
 import json
 import logging
@@ -31,7 +30,7 @@ from wandb import env, trigger
 from wandb.errors import CommError, Error, UsageError
 from wandb.errors.links import url_registry
 from wandb.errors.util import ProtobufErrorHandler
-from wandb.integration import sagemaker
+from wandb.integration import sagemaker, weave
 from wandb.proto.wandb_deprecated import Deprecated
 from wandb.sdk.lib import ipython as wb_ipython
 from wandb.sdk.lib import progress, runid, wb_logging
@@ -202,23 +201,7 @@ class _WandbInit:
         Returns:
             A callback to print any generated warnings.
         """
-        exclude_env_vars = {"WANDB_SERVICE", "WANDB_KUBEFLOW_URL"}
-        # check if environment variables have changed
-        singleton_env = {
-            k: v
-            for k, v in wandb_setup.singleton()._environ.items()
-            if k.startswith("WANDB_") and k not in exclude_env_vars
-        }
-        os_env = {
-            k: v
-            for k, v in os.environ.items()
-            if k.startswith("WANDB_") and k not in exclude_env_vars
-        }
-
-        if (
-            set(singleton_env.keys()) == set(os_env.keys())  #
-            and set(singleton_env.values()) == set(os_env.values())
-        ):
+        if not self._wl.did_environment_change():
             return _noop_printer_callback()
 
         def print_warning(run_printer: printer.Printer) -> None:
@@ -475,9 +458,9 @@ class _WandbInit:
             )
             self._telemetry.feature.sagemaker = True
 
-        if self._wl._config:
+        if self._wl.config:
             self._split_artifacts_from_config(
-                self._wl._config,
+                self._wl.config,
                 config_target=result.base_no_artifacts,
                 artifacts=result.artifacts,
             )
@@ -996,7 +979,6 @@ class _WandbInit:
             result = wait_with_progress(
                 run_init_handle,
                 timeout=timeout,
-                progress_after=1,
                 display_progress=display_init_message,
             )
 
@@ -1112,8 +1094,7 @@ def _attach(
     except Exception as e:
         raise UsageError(f"Unable to attach to run {attach_id}") from e
 
-    settings: Settings = copy.copy(_wl._settings)
-
+    settings = _wl.settings.model_copy()
     settings.update_from_dict(
         {
             "run_id": attach_id,
@@ -1578,7 +1559,12 @@ def init(  # noqa: C901
             if run_settings.x_server_side_derived_summary:
                 init_telemetry.feature.server_side_derived_summary = True
 
-            return wi.init(run_settings, run_config, run_printer)
+            run = wi.init(run_settings, run_config, run_printer)
+
+            # Set up automatic Weave integration if Weave is installed
+            weave.setup(run_settings.entity, run_settings.project)
+
+            return run
 
     except KeyboardInterrupt as e:
         if wl:

wandb/sdk/wandb_login.py

@@ -202,7 +202,7 @@ class _WandbLogin:
         # Whenever the key changes, make sure to pull in user settings
         # from server.
         if not self._wandb_setup.settings._offline:
-            self._wandb_setup._update_user_settings()
+            self._wandb_setup.update_user_settings()
 
     def _prompt_api_key(
         self, referrer: Optional[str] = None