wandb 0.22.1__py3-none-win_arm64.whl

wandb/sdk/lib/asyncio_compat.py

@@ -0,0 +1,278 @@
+"""Functions for compatibility with asyncio."""
+
+from __future__ import annotations
+
+import asyncio
+import concurrent
+import concurrent.futures
+import contextlib
+import threading
+from typing import Any, AsyncIterator, Callable, Coroutine, TypeVar
+
+_T = TypeVar("_T")
+
+
+def run(fn: Callable[[], Coroutine[Any, Any, _T]]) -> _T:
+    """Run `fn` in an asyncio loop in a new thread.
+
+    This must always be used instead of `asyncio.run` which fails if there is
+    an active `asyncio` event loop in the current thread. Since `wandb` was not
+    originally designed with `asyncio` in mind, using `asyncio.run` would break
+    users who were calling `wandb` methods from an `asyncio` loop.
+
+    Note that due to starting a new thread, this is slightly slow.
+    """
+    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+        runner = CancellableRunner()
+        future = executor.submit(runner.run, fn)
+
+        try:
+            return future.result()
+
+        finally:
+            runner.cancel()
+
+
+class RunnerCancelledError(Exception):
+    """The `CancellableRunner.run()` invocation was cancelled."""
+
+
+class CancellableRunner:
+    """Runs an asyncio event loop allowing cancellation.
+
+    The `run()` method is like `asyncio.run()`. The `cancel()` method may
+    be used in a different thread, for instance in a `finally` block, to cancel
+    all tasks, and it is a no-op if `run()` completed.
+
+    Without this, it is impossible to make `asyncio.run()` stop if it runs
+    in a non-main thread. In particular, a KeyboardInterrupt causes the
+    ThreadPoolExecutor above to block until the asyncio thread completes,
+    but there is no way to tell the asyncio thread to cancel its work.
+    A second KeyboardInterrupt makes ThreadPoolExecutor give up while the
+    asyncio thread still runs in the background, with terrible effects if it
+    prints to the user's terminal.
+    """
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+
+        self._is_cancelled = False
+        self._started = False
+        self._done = False
+
+        self._loop: asyncio.AbstractEventLoop | None = None
+        self._cancel_event: asyncio.Event | None = None
+
+    def run(self, fn: Callable[[], Coroutine[Any, Any, _T]]) -> _T:
+        """Run a coroutine in asyncio, cancelling it on `cancel()`.
+
+        Returns:
+            The result of the coroutine returned by `fn`.
+
+        Raises:
+            RunnerCancelledError: If `cancel()` is called.
+        """
+        return asyncio.run(self._run_or_cancel(fn))
+
+    async def _run_or_cancel(
+        self,
+        fn: Callable[[], Coroutine[Any, Any, _T]],
+    ) -> _T:
+        with self._lock:
+            if self._is_cancelled:
+                raise RunnerCancelledError()
+
+            self._loop = asyncio.get_running_loop()
+            self._cancel_event = asyncio.Event()
+            self._started = True
+
+        cancellation_task = asyncio.create_task(self._cancel_event.wait())
+        fn_task = asyncio.create_task(fn())
+
+        try:
+            await asyncio.wait(
+                [cancellation_task, fn_task],
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+
+            if fn_task.done():
+                return fn_task.result()
+            else:
+                raise RunnerCancelledError()
+
+        finally:
+            # NOTE: asyncio.run() cancels all tasks after the main task exits,
+            #   but this is not documented, so we cancel them explicitly here
+            #   as well. It also blocks until canceled tasks complete.
+            cancellation_task.cancel()
+            fn_task.cancel()
+
+            with self._lock:
+                self._done = True
+
+    def cancel(self) -> None:
+        """Cancel all asyncio work started by `run()`."""
+        with self._lock:
+            if self._is_cancelled:
+                return
+            self._is_cancelled = True
+
+            if self._done or not self._started:
+                # If the runner already finished, no need to cancel it.
+                #
+                # If the runner hasn't started the loop yet, then it will not
+                # as we already set _is_cancelled.
+                return
+
+            assert self._loop
+            assert self._cancel_event
+            self._loop.call_soon_threadsafe(self._cancel_event.set)
+
+
+class TaskGroup:
+    """Object that `open_task_group()` yields."""
+
+    def __init__(self) -> None:
+        self._tasks: list[asyncio.Task] = []
+
+    def start_soon(self, coro: Coroutine[Any, Any, Any]) -> None:
+        """Schedule a task in the group.
+
+        Args:
+            coro: The return value of the `async` function defining the task.
+        """
+        self._tasks.append(asyncio.create_task(coro))
+
+    async def _wait_all(self, *, race: bool, timeout: float | None) -> None:
+        """Block until tasks complete.
+
+        Args:
+            race: If true, blocks until the first task completes and then
+                cancels the rest. Otherwise, waits for all tasks or until
+                the first exception.
+            timeout: How long to wait.
+
+        Raises:
+            TimeoutError: If the timeout expires.
+            Exception: If one or more tasks raises an exception, one of these
+                is raised arbitrarily.
+        """
+        if not self._tasks:
+            return
+
+        if race:
+            return_when = asyncio.FIRST_COMPLETED
+        else:
+            return_when = asyncio.FIRST_EXCEPTION
+
+        done, pending = await asyncio.wait(
+            self._tasks,
+            timeout=timeout,
+            return_when=return_when,
+        )
+
+        if not done:
+            raise TimeoutError(f"Timed out after {timeout} seconds.")
+
+        # If any of the finished tasks raised an exception, pick the first one.
+        for task in done:
+            if exc := task.exception():
+                raise exc
+
+        # Wait for remaining tasks to clean up, then re-raise any exceptions
+        # that arise. Note that pending is only non-empty when race=True.
+        for task in pending:
+            task.cancel()
+        await asyncio.gather(*pending, return_exceptions=True)
+        for task in pending:
+            if task.cancelled():
+                continue
+            if exc := task.exception():
+                raise exc
+
+    async def _cancel_all(self) -> None:
+        """Cancel all tasks.
+
+        Blocks until cancelled tasks complete to allow them to clean up.
+        Ignores exceptions.
+        """
+        for task in self._tasks:
+            # NOTE: It is safe to cancel tasks that have already completed.
+            task.cancel()
+        await asyncio.gather(*self._tasks, return_exceptions=True)
+
+
+@contextlib.asynccontextmanager
+async def open_task_group(
+    *,
+    exit_timeout: float | None = None,
+    race: bool = False,
+) -> AsyncIterator[TaskGroup]:
+    """Create a task group.
+
+    `asyncio` gained task groups in Python 3.11.
+
+    This is an async context manager, meant to be used with `async with`.
+    On exit, it blocks until all subtasks complete. If any subtask fails, or if
+    the current task is cancelled, it cancels all subtasks in the group and
+    raises the subtask's exception. If multiple subtasks fail simultaneously,
+    one of their exceptions is chosen arbitrarily.
+
+    NOTE: Subtask exceptions do not propagate until the context manager exits.
+    This means that the task group cannot cancel code running inside the
+    `async with` block .
+
+    Args:
+        exit_timeout: An optional timeout in seconds. When exiting the
+            context manager, if tasks don't complete in this time,
+            they are cancelled and a TimeoutError is raised.
+        race: If true, all pending tasks are cancelled once any task
+            in the group completes. Prefer to use the race() function instead.
+
+    Raises:
+        TimeoutError: if exit_timeout is specified and tasks don't finish
+            in time.
+    """
+    task_group = TaskGroup()
+
+    try:
+        yield task_group
+        await task_group._wait_all(race=race, timeout=exit_timeout)
+    finally:
+        await task_group._cancel_all()
+
+
+@contextlib.asynccontextmanager
+async def cancel_on_exit(coro: Coroutine[Any, Any, Any]) -> AsyncIterator[None]:
+    """Schedule a task, cancelling it when exiting the context manager.
+
+    If the context manager exits successfully but the given coroutine raises
+    an exception, that exception is reraised. The exception is suppressed
+    if the context manager raises an exception.
+    """
+
+    async def stop_immediately():
+        pass
+
+    async with open_task_group(race=True) as group:
+        group.start_soon(stop_immediately())
+        group.start_soon(coro)
+        yield
+
+
+async def race(*coros: Coroutine[Any, Any, Any]) -> None:
+    """Wait until the first completed task.
+
+    After any coroutine completes, all others are cancelled.
+    If the current task is cancelled, all coroutines are cancelled too.
+
+    If coroutines complete simultaneously and any one of them raises
+    an exception, an arbitrary one is propagated. Similarly, if any coroutines
+    raise exceptions during cancellation, one of them propagates.
+
+    Args:
+        coros: Coroutines to race.
+    """
+    async with open_task_group(race=True) as tg:
+        for coro in coros:
+            tg.start_soon(coro)

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