wandb 0.21.1__py3-none-win_amd64.whl → 0.21.2__py3-none-win_amd64.whl

wandb/sdk/launch/runner/kubernetes_runner.py

@@ -7,7 +7,6 @@ import json
 import logging
 import os
 import time
-import uuid
 from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
 
 import yaml
@@ -28,6 +27,7 @@ from wandb.sdk.launch.runner.kubernetes_monitor import (
     CustomResource,
     LaunchKubernetesMonitor,
 )
+from wandb.sdk.launch.utils import recursive_macro_sub
 from wandb.sdk.lib.retry import ExponentialBackoff, retry_async
 from wandb.util import get_module
 
@@ -62,6 +62,9 @@ from kubernetes_asyncio.client.api.core_v1_api import (  # type: ignore # noqa:
 from kubernetes_asyncio.client.api.custom_objects_api import (  # type: ignore # noqa: E402
     CustomObjectsApi,
 )
+from kubernetes_asyncio.client.api.networking_v1_api import (  # type: ignore # noqa: E402
+    NetworkingV1Api,
+)
 from kubernetes_asyncio.client.models.v1_secret import (  # type: ignore # noqa: E402
     V1Secret,
 )
@@ -85,6 +88,7 @@ class KubernetesSubmittedRun(AbstractRun):
         batch_api: "BatchV1Api",
         core_api: "CoreV1Api",
         apps_api: "AppsV1Api",
+        network_api: "NetworkingV1Api",
         name: str,
         namespace: Optional[str] = "default",
         secret: Optional["V1Secret"] = None,
@@ -103,6 +107,7 @@ class KubernetesSubmittedRun(AbstractRun):
         Arguments:
             batch_api: Kubernetes BatchV1Api object.
             core_api: Kubernetes CoreV1Api object.
+            network_api: Kubernetes NetworkV1Api object.
             name: Name of the job.
             namespace: Kubernetes namespace.
             secret: Kubernetes secret.
@@ -113,6 +118,7 @@ class KubernetesSubmittedRun(AbstractRun):
         self.batch_api = batch_api
         self.core_api = core_api
         self.apps_api = apps_api
+        self.network_api = network_api
         self.name = name
         self.namespace = namespace
         self._fail_count = 0
@@ -207,11 +213,9 @@ class KubernetesSubmittedRun(AbstractRun):
                 (self.core_api, "service"),
                 (self.batch_api, "job"),
                 (self.core_api, "pod"),
-                (self.core_api, "config_map"),
                 (self.core_api, "secret"),
                 (self.apps_api, "deployment"),
-                (self.apps_api, "replica_set"),
-                (self.apps_api, "daemon_set"),
+                (self.network_api, "network_policy"),
             ]
 
             for api_client, resource_type in resource_cleanups:
@@ -700,7 +704,6 @@ class KubernetesRunner(AbstractRunner):
         config: Dict[str, Any],
         namespace: str,
         run_id: str,
-        auxiliary_resource_label_key: str,
         launch_project: LaunchProject,
         api_key_secret: Optional["V1Secret"] = None,
         wait_for_ready: bool = True,
@@ -713,7 +716,6 @@ class KubernetesRunner(AbstractRunner):
             config: The resource configuration to prepare.
             namespace: The namespace to create the resource in.
             run_id: The run ID to label the resource with.
-            auxiliary_resource_label_key: The key of the auxiliary resource label.
             launch_project: The launch project to get environment variables from.
             api_key_secret: The API key secret to inject.
             wait_for_ready: Whether to wait for the resource to be ready after creation.
@@ -722,25 +724,8 @@ class KubernetesRunner(AbstractRunner):
         config.setdefault("metadata", {})
         config["metadata"].setdefault("labels", {})
         config["metadata"]["labels"][WANDB_K8S_RUN_ID] = run_id
-        config["metadata"]["labels"][WANDB_K8S_LABEL_AUXILIARY_RESOURCE] = (
-            auxiliary_resource_label_key
-        )
         config["metadata"]["labels"]["wandb.ai/created-by"] = "launch-agent"
 
-        if config.get("kind") == "Service" or config.get("kind") == "Deployment":
-            config.setdefault("metadata", {})
-            original_name = config["metadata"].get("name", config.get("kind"))
-            safe_name = make_name_dns_safe(original_name)
-            safe_entity = make_name_dns_safe(launch_project.target_entity or "")
-            safe_project = make_name_dns_safe(launch_project.target_project or "")
-            safe_run_id = make_name_dns_safe(run_id or "")
-
-            new_name = f"{safe_name}-{safe_entity}-{safe_project}-{safe_run_id}"
-            config["metadata"]["name"] = new_name
-            wandb.termlog(
-                f"{LOG_PREFIX}Modified {config.get('kind')} name from '{original_name}' to '{new_name}'"
-            )
-
         env_vars = launch_project.get_env_vars_dict(
             self._api, MAX_ENV_LENGTHS[self.__class__.__name__]
         )
@@ -920,19 +905,29 @@ class KubernetesRunner(AbstractRunner):
         batch_api = kubernetes_asyncio.client.BatchV1Api(api_client)
         core_api = kubernetes_asyncio.client.CoreV1Api(api_client)
         apps_api = kubernetes_asyncio.client.AppsV1Api(api_client)
+        network_api = kubernetes_asyncio.client.NetworkingV1Api(api_client)
 
         namespace = self.get_namespace(resource_args, context)
         job, secret = await self._inject_defaults(
             resource_args, launch_project, image_uri, namespace, core_api
         )
 
-        additional_services = launch_project.launch_spec.get("additional_services", [])
-        auxiliary_resource_label_key = None
+        update_dict = {
+            "project_name": launch_project.target_project,
+            "entity_name": launch_project.target_entity,
+            "run_id": launch_project.run_id,
+            "run_name": launch_project.name,
+            "image_uri": image_uri,
+            "author": launch_project.author,
+        }
+        update_dict.update(os.environ)
+        additional_services: List[Dict[str, Any]] = recursive_macro_sub(
+            launch_project.launch_spec.get("additional_services", []), update_dict
+        )
         if additional_services:
             wandb.termlog(
                 f"{LOG_PREFIX}Creating additional services: {additional_services}"
             )
-            auxiliary_resource_label_key = f"aux-{uuid.uuid4()}"
 
             wait_for_ready = resource_args.get("wait_for_ready", True)
             wait_timeout = resource_args.get("wait_timeout", 300)
@@ -941,10 +936,9 @@ class KubernetesRunner(AbstractRunner):
                 *[
                     self._prepare_resource(
                         api_client,
-                        resource.get("config"),
+                        resource.get("config", {}),
                         namespace,
                         launch_project.run_id,
-                        auxiliary_resource_label_key,
                         launch_project,
                         secret,
                         wait_for_ready,
@@ -982,10 +976,11 @@ class KubernetesRunner(AbstractRunner):
             batch_api,
             core_api,
             apps_api,
+            network_api,
             job_name,
             namespace,
             secret,
-            auxiliary_resource_label_key,
+            f"aux-{launch_project.target_entity}-{launch_project.target_project}-{launch_project.run_id}",
         )
         if self.backend_config[PROJECT_SYNCHRONOUS]:
             await submitted_job.wait()

wandb/sdk/lib/asyncio_compat.py

@@ -23,7 +23,7 @@ def run(fn: Callable[[], Coroutine[Any, Any, _T]]) -> _T:
     Note that due to starting a new thread, this is slightly slow.
     """
     with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
-        runner = _Runner()
+        runner = CancellableRunner()
         future = executor.submit(runner.run, fn)
 
         try:
@@ -33,15 +33,16 @@ def run(fn: Callable[[], Coroutine[Any, Any, _T]]) -> _T:
             runner.cancel()
 
 
-class _RunnerCancelledError(Exception):
-    """The `_Runner.run()` invocation was cancelled."""
+class RunnerCancelledError(Exception):
+    """The `CancellableRunner.run()` invocation was cancelled."""
 
 
-class _Runner:
+class CancellableRunner:
     """Runs an asyncio event loop allowing cancellation.
 
-    This is like `asyncio.run()`, except it provides a `cancel()` method
-    meant to be called in a `finally` block.
+    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
@@ -69,7 +70,7 @@ class _Runner:
             The result of the coroutine returned by `fn`.
 
         Raises:
-            _RunnerCancelledError: If `cancel()` is called.
+            RunnerCancelledError: If `cancel()` is called.
         """
         return asyncio.run(self._run_or_cancel(fn))
 
@@ -79,7 +80,7 @@ class _Runner:
     ) -> _T:
         with self._lock:
             if self._is_cancelled:
-                raise _RunnerCancelledError()
+                raise RunnerCancelledError()
 
             self._loop = asyncio.get_running_loop()
             self._cancel_event = asyncio.Event()
@@ -97,7 +98,7 @@ class _Runner:
             if fn_task.done():
                 return fn_task.result()
             else:
-                raise _RunnerCancelledError()
+                raise RunnerCancelledError()
 
         finally:
             # NOTE: asyncio.run() cancels all tasks after the main task exits,
@@ -157,11 +158,9 @@ class TaskGroup:
         )
 
         for task in done:
-            try:
+            with contextlib.suppress(asyncio.CancelledError):
                 if exc := task.exception():
                     raise exc
-            except asyncio.CancelledError:
-                pass
 
     def _cancel_all(self) -> None:
         """Cancel all tasks."""
@@ -199,15 +198,16 @@ async def open_task_group() -> AsyncIterator[TaskGroup]:
 def cancel_on_exit(coro: Coroutine[Any, Any, Any]) -> Iterator[None]:
     """Schedule a task, cancelling it when exiting the context manager.
 
-    If the given coroutine raises an exception, that exception is raised
-    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.
     """
     task = asyncio.create_task(coro)
 
     try:
         yield
-    finally:
+
         if task.done() and (exception := task.exception()):
             raise exception
-
+    finally:
         task.cancel()

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/hashutil.py

@@ -6,19 +6,28 @@ import logging
 import mmap
 import sys
 import time
-from typing import TYPE_CHECKING, NewType
+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:
     """Allow FIPS-compliant md5 hash when supported."""

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)