wandb 0.21.0__py3-none-win32.whl → 0.21.2__py3-none-win32.whl

wandb/sdk/lib/asyncio_manager.py

@@ -0,0 +1,252 @@
+"""Implements an asyncio thread suitable for internal wandb use."""
+
+from __future__ import annotations
+
+import asyncio
+import concurrent.futures
+import contextlib
+import logging
+import threading
+from typing import Any, Callable, Coroutine, TypeVar
+
+from . import asyncio_compat
+
+_T = TypeVar("_T")
+
+_logger = logging.getLogger(__name__)
+
+
+class RunCancelledError(Exception):
+    """A function passed to AsyncioManager.run() was cancelled."""
+
+
+class AlreadyJoinedError(Exception):
+    """AsyncioManager.run() used after join()."""
+
+
+class AsyncioManager:
+    """Manages a thread running an asyncio loop.
+
+    The thread must be started using start() and should be joined using
+    join(). The thread is a daemon thread, so if join() is not invoked,
+    the asyncio work could end abruptly when all non-daemon threads exit.
+
+    The run() method allows invoking an async function in the asyncio thread
+    and waiting until it completes. The run_soon() method allows running
+    an async function without waiting for it.
+
+    Note that although tempting, it is **not** possible to write a safe
+    run_in_loop() method that chooses whether to use run() or execute a function
+    directly based on whether it's called from the asyncio thread: Suppose a
+    function bad() holds a threading.Lock while using run_in_loop() and an
+    asyncio task calling bad() is scheduled. If bad() is then invoked in a
+    different thread that reaches run_in_loop(), the aforementioned asyncio task
+    will deadlock. It is unreasonable to require that run_in_loop() never be
+    called while holding a lock (which would apply to the callers of its
+    callers, and so on), so it cannot safely exist.
+    """
+
+    def __init__(self) -> None:
+        self._runner = asyncio_compat.CancellableRunner()
+        self._thread = threading.Thread(
+            target=self._main,
+            name="wandb-AsyncioManager-main",
+            daemon=True,
+        )
+        self._lock = threading.Lock()
+
+        self._ready_event = threading.Event()
+        """Whether asyncio primitives have been initialized."""
+
+        self._joined = False
+        """Whether join() has been called. Guarded by _lock."""
+
+        self._loop: asyncio.AbstractEventLoop
+        """A handle for interacting with the asyncio event loop."""
+
+        self._done_event: asyncio.Event
+        """Indicates to the asyncio loop that join() was called."""
+
+        self._remaining_tasks = 0
+        """The number of tasks remaining. Guarded by _lock."""
+
+        self._task_finished_cond: asyncio.Condition
+        """Signalled when _remaining_tasks is decremented."""
+
+    def start(self) -> None:
+        """Start the asyncio thread."""
+        self._thread.start()
+
+    def join(self) -> None:
+        """Stop accepting new asyncio tasks and wait for the remaining ones."""
+        try:
+            with self._lock:
+                # If join() was already called, block until the thread completes
+                # and then return.
+                if self._joined:
+                    self._thread.join()
+                    return
+
+                self._joined = True
+
+            # Wait until _loop and _done_event are initialized.
+            self._ready_event.wait()
+
+            # Set the done event. The main function will exit once all
+            # tasks complete.
+            self._loop.call_soon_threadsafe(self._done_event.set)
+
+            self._thread.join()
+
+        finally:
+            # Any of the above may get interrupted by Ctrl+C, in which case we
+            # should cancel all tasks, since join() can only be called once.
+            # This only matters if the KeyboardInterrupt is suppressed.
+            self._runner.cancel()
+
+    def run(self, fn: Callable[[], Coroutine[Any, Any, _T]]) -> _T:
+        """Run an async function to completion.
+
+        The function is called in the asyncio thread. Blocks until start()
+        is called. This raises an error if called inside an async function,
+        and as a consequence, the caller may also not be called inside an
+        async function.
+
+        Args:
+            fn: The function to run.
+
+        Returns:
+            The return value of fn.
+
+        Raises:
+            Exception: Any exception raised by fn.
+            RunCancelledError: If fn is cancelled, particularly when join()
+                is interrupted by Ctrl+C or if it otherwise cancels itself.
+            AlreadyJoinedError: If join() was already called.
+            ValueError: If called inside an async function.
+        """
+        self._ready_event.wait()
+
+        if threading.current_thread().ident == self._thread.ident:
+            raise ValueError("Cannot use run() inside async loop.")
+
+        future = self._schedule(fn, daemon=False)
+
+        try:
+            return future.result()
+
+        except concurrent.futures.CancelledError:
+            raise RunCancelledError from None
+
+        except KeyboardInterrupt:
+            # If we're interrupted here, we only cancel this task rather than
+            # cancelling all tasks like in join(). This only matters if the
+            # interrupt is then suppressed (or delayed) in which case we
+            # should let other tasks progress.
+            future.cancel()
+            raise
+
+    def run_soon(
+        self,
+        fn: Callable[[], Coroutine[Any, Any, None]],
+        *,
+        daemon: bool = False,
+        name: str | None = None,
+    ) -> None:
+        """Run an async function without waiting for it to complete.
+
+        The function is called in the asyncio thread. Note that since that's
+        a daemon thread, it will not get joined when the main thread exits,
+        so fn can stop abruptly.
+
+        Unlike run(), it is OK to call this inside an async function.
+
+        Blocks until start() is called.
+
+        Args:
+            fn: The function to run.
+            daemon: If true, join() will cancel fn after all non-daemon
+                tasks complete. By default, join() blocks until fn
+                completes.
+            name: An optional name to give to long-running tasks which can
+                appear in error traces and be useful to debugging.
+
+        Raises:
+            AlreadyJoinedError: If join() was already called.
+        """
+
+        # Wrap exceptions so that they're not printed to console.
+        async def fn_wrap_exceptions() -> None:
+            try:
+                await fn()
+            except Exception:
+                _logger.exception("Uncaught exception in run_soon callback.")
+
+        _ = self._schedule(fn_wrap_exceptions, daemon=daemon, name=name)
+
+    def _schedule(
+        self,
+        fn: Callable[[], Coroutine[Any, Any, _T]],
+        daemon: bool,
+        name: str | None = None,
+    ) -> concurrent.futures.Future[_T]:
+        # Wait for _loop to be initialized.
+        self._ready_event.wait()
+
+        with self._lock:
+            if self._joined:
+                raise AlreadyJoinedError
+
+            if not daemon:
+                self._remaining_tasks += 1
+
+        return asyncio.run_coroutine_threadsafe(
+            self._wrap(fn, daemon=daemon, name=name),
+            self._loop,
+        )
+
+    async def _wrap(
+        self,
+        fn: Callable[[], Coroutine[Any, Any, _T]],
+        daemon: bool,
+        name: str | None,
+    ) -> _T:
+        """Run fn to completion and possibly decrement _remaining tasks."""
+        try:
+            if name and (task := asyncio.current_task()):
+                task.set_name(name)
+
+            return await fn()
+        finally:
+            if not daemon:
+                async with self._task_finished_cond:
+                    with self._lock:
+                        self._remaining_tasks -= 1
+                    self._task_finished_cond.notify_all()
+
+    def _main(self) -> None:
+        """Run the asyncio loop until join() is called and all tasks finish."""
+        # A cancellation error is expected if join() is interrupted.
+        #
+        # Were it not suppressed, its stacktrace would get printed.
+        with contextlib.suppress(asyncio_compat.RunnerCancelledError):
+            self._runner.run(self._main_async)
+
+    async def _main_async(self) -> None:
+        """Wait until join() is called and all tasks finish."""
+        self._loop = asyncio.get_running_loop()
+        self._done_event = asyncio.Event()
+        self._task_finished_cond = asyncio.Condition()
+
+        self._ready_event.set()
+
+        # Wait until done.
+        await self._done_event.wait()
+
+        # Wait for all tasks to complete.
+        #
+        # Once we exit, asyncio will cancel any leftover tasks.
+        async with self._task_finished_cond:
+            await self._task_finished_cond.wait_for(
+                lambda: self._remaining_tasks <= 0,
+            )

wandb/sdk/lib/deprecate.py

@@ -1,20 +1,14 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 import wandb
 from wandb.proto.wandb_deprecated import DEPRECATED_FEATURES
 from wandb.sdk.lib import telemetry
 
-# Necessary to break import cycle.
-if TYPE_CHECKING:
-    from wandb import wandb_run
-
 
 def deprecate(
     field_name: DEPRECATED_FEATURES,
     warning_message: str,
-    run: wandb_run.Run | None = None,
+    run: wandb.Run | None = None,
 ) -> None:
     """Warn the user that a feature has been deprecated.
 

wandb/sdk/lib/disabled.py

@@ -26,5 +26,5 @@ class RunDisabled:
         deprecate.deprecate(
             field_name=Deprecated.run_disabled,
             warning_message="RunDisabled is deprecated and is a no-op. "
-            '`wandb.init(mode="disabled")` now returns and instance of `wandb.sdk.wandb_run.Run`.',
+            '`wandb.init(mode="disabled")` now returns an instance of `wandb.Run`.',
         )

wandb/sdk/lib/hashutil.py

@@ -2,18 +2,31 @@ from __future__ import annotations
 
 import base64
 import hashlib
+import logging
 import mmap
 import sys
-from typing import TYPE_CHECKING, NewType
+import time
+from typing import TYPE_CHECKING
+
+from typing_extensions import TypeAlias
 
 from wandb.sdk.lib.paths import StrPath
 
 if TYPE_CHECKING:
     import _hashlib  # type: ignore[import-not-found]
 
-ETag = NewType("ETag", str)
-HexMD5 = NewType("HexMD5", str)
-B64MD5 = NewType("B64MD5", str)
+
+logger = logging.getLogger(__name__)
+
+# In the future, consider relying on pydantic to validate these types via e.g.
+# - Base64Str: https://docs.pydantic.dev/latest/api/types/#pydantic.types.Base64Str
+# - a custom EncodedStr + Encoder impl: https://docs.pydantic.dev/latest/api/types/#pydantic.types.EncodedStr
+#
+# Note that so long as we continue to support Pydantic v1, the options above will require a compatible shim/backport
+# implementation, since those types are not in Pydantic v1.
+ETag: TypeAlias = str
+HexMD5: TypeAlias = str
+B64MD5: TypeAlias = str
 
 
 def _md5(data: bytes = b"") -> _hashlib.HASH:
@@ -44,7 +57,16 @@ def hex_to_b64_id(encoded_string: str | bytes) -> B64MD5:
 
 
 def md5_file_b64(*paths: StrPath) -> B64MD5:
-    return _b64_from_hasher(_md5_file_hasher(*paths))
+    start_time = time.monotonic()
+    digest = _b64_from_hasher(_md5_file_hasher(*paths))
+    hash_time_seconds = time.monotonic() - start_time
+    if hash_time_seconds > 1.0:
+        logger.debug(
+            "Computed MD5 hash for file. paths=%s, hashTimeMs=%d",
+            paths,
+            int(hash_time_seconds * 1000),
+        )
+    return digest
 
 
 def md5_file_hex(*paths: StrPath) -> HexMD5:

wandb/sdk/lib/module.py

@@ -57,22 +57,16 @@ 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.watch = preinit.PreInitCallable(
-        "wandb.watch", wandb.wandb_sdk.wandb_run.Run.watch
-    )
-    wandb.unwatch = preinit.PreInitCallable(
-        "wandb.unwatch", wandb.wandb_sdk.wandb_run.Run.unwatch
-    )
-    wandb.save = preinit.PreInitCallable(
-        "wandb.save", wandb.wandb_sdk.wandb_run.Run.save
-    )
+    wandb.log = preinit.PreInitCallable("wandb.log", wandb.Run.log)
+    wandb.watch = preinit.PreInitCallable("wandb.watch", wandb.Run.watch)
+    wandb.unwatch = preinit.PreInitCallable("wandb.unwatch", wandb.Run.unwatch)
+    wandb.save = preinit.PreInitCallable("wandb.save", wandb.Run.save)
     wandb.use_artifact = preinit.PreInitCallable(
-        "wandb.use_artifact", wandb.wandb_sdk.wandb_run.Run.use_artifact
+        "wandb.use_artifact", wandb.Run.use_artifact
     )
     wandb.log_artifact = preinit.PreInitCallable(
-        "wandb.log_artifact", wandb.wandb_sdk.wandb_run.Run.log_artifact
+        "wandb.log_artifact", wandb.Run.log_artifact
     )
     wandb.define_metric = preinit.PreInitCallable(
-        "wandb.define_metric", wandb.wandb_sdk.wandb_run.Run.define_metric
+        "wandb.define_metric", wandb.Run.define_metric
     )

wandb/sdk/lib/printer.py

@@ -102,8 +102,8 @@ def new_printer(settings: wandb.Settings | None = None) -> Printer:
             has been called, then global settings are used. Otherwise,
             settings (such as silent mode) are ignored.
     """
-    if not settings and (singleton := wandb_setup.singleton_if_setup()):
-        settings = singleton.settings
+    if not settings and (s := wandb_setup.singleton().settings_if_loaded):
+        settings = s
 
     if ipython.in_jupyter():
         return _PrinterJupyter(settings=settings)

wandb/sdk/lib/printer_asyncio.py

@@ -1,6 +1,7 @@
 import asyncio
 from typing import Callable, TypeVar
 
+from wandb.sdk import wandb_setup
 from wandb.sdk.lib import asyncio_compat, printer
 
 _T = TypeVar("_T")
@@ -43,4 +44,5 @@ def run_async_with_spinner(
             func_running.set()
             return res
 
-    return asyncio_compat.run(_loop_run_with_spinner)
+    asyncer = wandb_setup.singleton().asyncer
+    return asyncer.run(_loop_run_with_spinner)

wandb/sdk/lib/progress.py

@@ -14,25 +14,6 @@ from wandb.sdk.lib import asyncio_compat
 from . import printer as p
 
 
-def print_sync_dedupe_stats(
-    printer: p.Printer,
-    final_result: pb.PollExitResponse,
-) -> None:
-    """Print how much W&B sync reduced the amount of uploaded data.
-
-    Args:
-        final_result: The final PollExit result.
-    """
-    deduped_bytes = final_result.pusher_stats.deduped_bytes
-    total_bytes = final_result.pusher_stats.total_bytes
-
-    if total_bytes <= 0 or deduped_bytes <= 0:
-        return
-
-    frac = deduped_bytes / total_bytes
-    printer.display(f"W&B sync reduced upload amount by {frac:.1%}")
-
-
 async def loop_printing_operation_stats(
     progress: ProgressPrinter,
     interface: interface.InterfaceBase,