wandb 0.21.2__py3-none-macosx_12_0_arm64.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/capped_dict.py

@@ -0,0 +1,26 @@
+import collections
+from typing import Any, Optional
+
+
+class CappedDict(collections.OrderedDict):
+    default_max_size = 50
+
+    def __init__(self, max_size: Optional[int] = None) -> None:
+        self.max_size = max_size or self.default_max_size
+        super().__init__()
+
+    def __setitem__(self, key: str, val: Any) -> None:
+        if key not in self:
+            max_size = self.max_size - 1
+            self._prune_dict(max_size)
+        super().__setitem__(key, val)
+
+    def update(self, **kwargs: Any) -> None:  # type: ignore[override]
+        super().update(**kwargs)
+        self._prune_dict(self.max_size)
+
+    def _prune_dict(self, max_size: int) -> None:
+        if len(self) >= max_size:
+            diff = len(self) - max_size
+            for k in list(self.keys())[:diff]:
+                del self[k]

wandb/sdk/lib/config_util.py

@@ -0,0 +1,101 @@
+import json
+import logging
+import os
+from typing import Any, Dict, Optional
+
+import yaml
+
+import wandb
+from wandb.errors import Error
+from wandb.util import load_yaml
+
+from . import filesystem
+
+logger = logging.getLogger("wandb")
+
+
+class ConfigError(Error):
+    pass
+
+
+def dict_from_proto_list(obj_list):
+    d = dict()
+    for item in obj_list:
+        d[item.key] = dict(desc=None, value=json.loads(item.value_json))
+    return d
+
+
+def dict_strip_value_dict(config_dict):
+    d = dict()
+    for k, v in config_dict.items():
+        d[k] = v["value"]
+    return d
+
+
+def dict_no_value_from_proto_list(obj_list):
+    d = dict()
+    for item in obj_list:
+        possible_dict = json.loads(item.value_json)
+        if not isinstance(possible_dict, dict) or "value" not in possible_dict:
+            continue
+        d[item.key] = possible_dict["value"]
+
+    return d
+
+
+# TODO(jhr): these functions should go away once we merge jobspec PR
+def save_config_file_from_dict(config_filename, config_dict):
+    s = b"wandb_version: 1"
+    if config_dict:  # adding an empty dictionary here causes a parse error
+        s += b"\n\n" + yaml.dump(
+            config_dict,
+            Dumper=yaml.SafeDumper,
+            default_flow_style=False,
+            allow_unicode=True,
+            encoding="utf-8",
+            sort_keys=False,
+        )
+    data = s.decode("utf-8")
+    filesystem.mkdir_exists_ok(os.path.dirname(config_filename))
+    with open(config_filename, "w") as conf_file:
+        conf_file.write(data)
+
+
+def dict_from_config_file(
+    filename: str, must_exist: bool = False
+) -> Optional[Dict[str, Any]]:
+    if not os.path.exists(filename):
+        if must_exist:
+            raise ConfigError(f"config file {filename} doesn't exist")
+        logger.debug(f"no default config file found in {filename}")
+        return None
+    try:
+        conf_file = open(filename)
+    except OSError:
+        raise ConfigError(f"Couldn't read config file: {filename}")
+    try:
+        loaded = load_yaml(conf_file)
+    except yaml.parser.ParserError:
+        raise ConfigError("Invalid YAML in config yaml")
+    if loaded is None:
+        wandb.termwarn(
+            "Found an empty default config file (config-defaults.yaml). Proceeding with no defaults."
+        )
+        return None
+    config_version = loaded.pop("wandb_version", None)
+    if config_version is not None and config_version != 1:
+        raise ConfigError("Unknown config version")
+    data = dict()
+    for k, v in loaded.items():
+        data[k] = v["value"]
+    return data
+
+
+def merge_dicts(dest: dict, src: dict) -> dict:
+    """Recursively merge two dictionaries. Similar to Lodash's _.merge()."""
+    for key, value in src.items():
+        if isinstance(value, dict) and key in dest and isinstance(dest[key], dict):
+            merge_dicts(dest[key], value)
+        else:
+            dest[key] = value
+    return dest

wandb/sdk/lib/console_capture.py

@@ -0,0 +1,219 @@
+"""Module for intercepting stdout/stderr.
+
+This patches the `write()` method of `stdout` and `stderr` on import.
+Once patched, it is not possible to unpatch or repatch, though individual
+callbacks can be removed.
+
+We assume that all other writing methods on the object delegate to `write()`,
+like `writelines()`. This is not guaranteed to be true, but it is true for
+common implementations. In particular, CPython's implementation of IOBase's
+`writelines()` delegates to `write()`.
+
+It is important to note that this technique interacts poorly with other
+code that performs similar patching if it also allows unpatching as this
+discards our modification. This is why we patch on import and do not support
+unpatching:
+
+    with contextlib.redirect_stderr(...):
+        from ... import console_capture
+        # Here, everything works fine.
+    # Here, callbacks are never called again.
+
+In particular, it does not work with some combinations of pytest's
+`capfd` / `capsys` fixtures and pytest's `--capture` option.
+"""
+
+from __future__ import annotations
+
+import logging
+import sys
+import threading
+from typing import IO, AnyStr, Callable, Protocol
+
+from . import wb_logging
+
+_logger = logging.getLogger(__name__)
+
+
+class CannotCaptureConsoleError(Exception):
+    """The module failed to patch stdout or stderr."""
+
+
+class _WriteCallback(Protocol):
+    """A callback that receives intercepted bytes or string data.
+
+    This may be called from any thread, but is only called from one thread
+    at a time.
+
+    Note on errors: Any error raised during the callback will clear all
+    callbacks. This means that if a user presses Ctrl-C at an unlucky time
+    during a run, we will stop uploading console output---but it's not
+    likely to be a problem unless something catches the KeyboardInterrupt.
+
+    Regular Exceptions are caught and logged instead of bubbling up to the
+    user's print() statements; other exceptions like KeyboardInterrupt are
+    re-raised.
+
+    Callbacks should handle all exceptions---a callback that raises any
+    Exception is considered buggy.
+    """
+
+    def __call__(
+        self,
+        data: bytes | str,
+        written: int,
+        /,
+    ) -> None:
+        """Intercept data passed to `write()`.
+
+        See the protocol docstring for information about exceptions.
+
+        Args:
+            data: The object passed to stderr's or stdout's `write()`.
+            written: The number of bytes or characters written.
+                This is the return value of `write()`.
+        """
+
+
+# A reentrant lock is used to catch callbacks that write to stderr/stdout.
+_module_rlock = threading.RLock()
+_is_writing = False
+
+_patch_exception: CannotCaptureConsoleError | None = None
+
+_next_callback_id: int = 1
+
+_stdout_callbacks: dict[int, _WriteCallback] = {}
+_stderr_callbacks: dict[int, _WriteCallback] = {}
+
+
+def capture_stdout(callback: _WriteCallback) -> Callable[[], None]:
+    """Install a callback that runs after every write to sys.stdout.
+
+    Args:
+        callback: A callback to invoke after running `sys.stdout.write`.
+
+    Returns:
+        A function to uninstall the callback.
+
+    Raises:
+        CannotCaptureConsoleError: If patching failed on import.
+    """
+    with _module_rlock:
+        if _patch_exception:
+            raise _patch_exception
+
+        return _insert_disposably(
+            _stdout_callbacks,
+            callback,
+        )
+
+
+def capture_stderr(callback: _WriteCallback) -> Callable[[], None]:
+    """Install a callback that runs after every write to sys.sdterr.
+
+    Args:
+        callback: A callback to invoke after running `sys.stderr.write`.
+
+    Returns:
+        A function to uninstall the callback.
+
+    Raises:
+        CannotCaptureConsoleError: If patching failed on import.
+    """
+    with _module_rlock:
+        if _patch_exception:
+            raise _patch_exception
+
+        return _insert_disposably(
+            _stderr_callbacks,
+            callback,
+        )
+
+
+def _insert_disposably(
+    callback_dict: dict[int, _WriteCallback],
+    callback: _WriteCallback,
+) -> Callable[[], None]:
+    global _next_callback_id
+    id = _next_callback_id
+    _next_callback_id += 1
+
+    disposed = False
+
+    def dispose() -> None:
+        nonlocal disposed
+
+        with _module_rlock:
+            if disposed:
+                return
+
+            callback_dict.pop(id, None)
+
+            disposed = True
+
+    callback_dict[id] = callback
+    return dispose
+
+
+def _patch(
+    stdout_or_stderr: IO[AnyStr],
+    callbacks: dict[int, _WriteCallback],
+) -> None:
+    orig_write: Callable[[AnyStr], int]
+
+    @wb_logging.log_to_all_runs()
+    def write_with_callbacks(s: AnyStr, /) -> int:
+        global _is_writing
+        n = orig_write(s)
+
+        # NOTE: Since _module_rlock is reentrant, this is safe. It will not
+        # deadlock if a callback invokes write() again.
+        with _module_rlock:
+            if _is_writing:
+                return n
+
+            _is_writing = True
+            try:
+                for cb in callbacks.values():
+                    cb(s, n)
+
+            except BaseException as e:
+                # Clear all callbacks on any exception to avoid infinite loops:
+                #
+                # * If we re-raise, an exception handler is likely to print
+                #   the exception to the console and trigger callbacks again
+                # * If we log, we can't guarantee that this doesn't print
+                #   to console.
+                #
+                # This is especially important for KeyboardInterrupt.
+                _stderr_callbacks.clear()
+                _stdout_callbacks.clear()
+
+                if isinstance(e, Exception):
+                    # We suppress Exceptions so that bugs in W&B code don't
+                    # cause the user's print() statements to raise errors.
+                    _logger.exception("Error in console callback, clearing all!")
+                else:
+                    # Re-raise errors like KeyboardInterrupt.
+                    raise
+
+            finally:
+                _is_writing = False
+
+        return n
+
+    orig_write = stdout_or_stderr.write
+
+    # mypy==1.14.1 fails to type-check this:
+    #   Incompatible types in assignment (expression has type
+    #   "Callable[[bytes], int]", variable has type overloaded function)
+    stdout_or_stderr.write = write_with_callbacks  # type: ignore
+
+
+try:
+    _patch(sys.stdout, _stdout_callbacks)
+    _patch(sys.stderr, _stderr_callbacks)
+except Exception as _patch_exception_cause:
+    _patch_exception = CannotCaptureConsoleError()
+    _patch_exception.__cause__ = _patch_exception_cause