wandb 0.21.0__py3-none-musllinux_1_2_aarch64.whl → 0.21.1__py3-none-musllinux_1_2_aarch64.whl

wandb/sdk/launch/runner/kubernetes_runner.py

@@ -6,6 +6,8 @@ import datetime
 import json
 import logging
 import os
+import time
+import uuid
 from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
 
 import yaml
@@ -20,6 +22,7 @@ from wandb.sdk.launch.registry.local_registry import LocalRegistry
 from wandb.sdk.launch.runner.abstract import Status
 from wandb.sdk.launch.runner.kubernetes_monitor import (
     WANDB_K8S_LABEL_AGENT,
+    WANDB_K8S_LABEL_AUXILIARY_RESOURCE,
     WANDB_K8S_LABEL_MONITOR,
     WANDB_K8S_RUN_ID,
     CustomResource,
@@ -47,6 +50,9 @@ get_module(
 
 import kubernetes_asyncio  # type: ignore # noqa: E402
 from kubernetes_asyncio import client  # noqa: E402
+from kubernetes_asyncio.client.api.apps_v1_api import (  # type: ignore # noqa: E402
+    AppsV1Api,
+)
 from kubernetes_asyncio.client.api.batch_v1_api import (  # type: ignore # noqa: E402
     BatchV1Api,
 )
@@ -78,9 +84,11 @@ class KubernetesSubmittedRun(AbstractRun):
         self,
         batch_api: "BatchV1Api",
         core_api: "CoreV1Api",
+        apps_api: "AppsV1Api",
         name: str,
         namespace: Optional[str] = "default",
         secret: Optional["V1Secret"] = None,
+        auxiliary_resource_label_key: Optional[str] = None,
     ) -> None:
         """Initialize a KubernetesSubmittedRun.
 
@@ -104,10 +112,12 @@ class KubernetesSubmittedRun(AbstractRun):
         """
         self.batch_api = batch_api
         self.core_api = core_api
+        self.apps_api = apps_api
         self.name = name
         self.namespace = namespace
         self._fail_count = 0
         self.secret = secret
+        self.auxiliary_resource_label_key = auxiliary_resource_label_key
 
     @property
     def id(self) -> str:
@@ -149,6 +159,7 @@ class KubernetesSubmittedRun(AbstractRun):
             await asyncio.sleep(5)
 
         await self._delete_secret()
+        await self._delete_auxiliary_resources_by_label()
         return (
             status.state == "finished"
         )  # todo: not sure if this (copied from aws runner) is the right approach? should we return false on failure
@@ -157,6 +168,7 @@ class KubernetesSubmittedRun(AbstractRun):
         status = LaunchKubernetesMonitor.get_status(self.name)
         if status in ["stopped", "failed", "finished", "preempted"]:
             await self._delete_secret()
+            await self._delete_auxiliary_resources_by_label()
         return status
 
     async def cancel(self) -> None:
@@ -167,6 +179,7 @@ class KubernetesSubmittedRun(AbstractRun):
                 name=self.name,
             )
             await self._delete_secret()
+            await self._delete_auxiliary_resources_by_label()
         except ApiException as e:
             raise LaunchError(
                 f"Failed to delete Kubernetes Job {self.name} in namespace {self.namespace}: {str(e)}"
@@ -181,6 +194,52 @@ class KubernetesSubmittedRun(AbstractRun):
             )
             self.secret = None
 
+    async def _delete_auxiliary_resources_by_label(self) -> None:
+        if self.auxiliary_resource_label_key is None:
+            return
+
+        label_selector = (
+            f"{WANDB_K8S_LABEL_AUXILIARY_RESOURCE}={self.auxiliary_resource_label_key}"
+        )
+
+        try:
+            resource_cleanups = [
+                (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"),
+            ]
+
+            for api_client, resource_type in resource_cleanups:
+                try:
+                    list_method = getattr(
+                        api_client, f"list_namespaced_{resource_type}"
+                    )
+                    delete_method = getattr(
+                        api_client, f"delete_namespaced_{resource_type}"
+                    )
+
+                    # List resources with our label
+                    resources = await list_method(
+                        namespace=self.namespace, label_selector=label_selector
+                    )
+
+                    # Delete each resource
+                    for resource in resources.items:
+                        await delete_method(
+                            name=resource.metadata.name, namespace=self.namespace
+                        )
+
+                except (AttributeError, ApiException) as e:
+                    wandb.termwarn(f"Could not clean up {resource_type}: {e}")
+
+        except Exception as e:
+            wandb.termwarn(f"Failed to clean up some auxiliary resources: {e}")
+
 
 class CrdSubmittedRun(AbstractRun):
     """Run submitted to a CRD backend, e.g. Volcano."""
@@ -366,6 +425,7 @@ class KubernetesRunner(AbstractRunner):
             job_metadata["generateName"] = make_name_dns_safe(
                 f"launch-{launch_project.target_entity}-{launch_project.target_project}-"
             )
+        job_metadata["namespace"] = namespace
 
         for i, cont in enumerate(containers):
             if "name" not in cont:
@@ -489,6 +549,235 @@ class KubernetesRunner(AbstractRunner):
 
         return job, api_key_secret
 
+    async def _wait_for_resource_ready(
+        self,
+        api_client: kubernetes_asyncio.client.ApiClient,
+        config: Dict[str, Any],
+        namespace: str,
+        timeout_seconds: int = 300,
+    ) -> None:
+        """Wait for a Kubernetes resource to be ready.
+
+        Arguments:
+            api_client: The Kubernetes API client.
+            config: The resource configuration.
+            namespace: The namespace where the resource was created.
+            timeout_seconds: Maximum time to wait for readiness.
+        """
+        resource_kind = config.get("kind")
+        resource_name = config.get("metadata", {}).get("name")
+
+        if not resource_kind or not resource_name:
+            wandb.termerror(
+                f"{LOG_PREFIX}Cannot wait for resource without kind or name"
+            )
+            return
+
+        wandb.termlog(
+            f"{LOG_PREFIX}Waiting for {resource_kind} '{resource_name}' to be ready..."
+        )
+
+        start_time = time.time()
+
+        if resource_kind == "Deployment":
+            await self._wait_for_deployment_ready(
+                api_client, resource_name, namespace, timeout_seconds
+            )
+        elif resource_kind == "Service":
+            await self._wait_for_service_ready(
+                api_client, resource_name, namespace, timeout_seconds
+            )
+        elif resource_kind == "Pod":
+            await self._wait_for_pod_ready(
+                api_client, resource_name, namespace, timeout_seconds
+            )
+        else:
+            wandb.termlog(
+                f"{LOG_PREFIX}No specific readiness check for {resource_kind}, waiting 5 seconds..."
+            )
+            await asyncio.sleep(5)
+
+        elapsed = time.time() - start_time
+        wandb.termlog(
+            f"{LOG_PREFIX}{resource_kind} '{resource_name}' is ready after {elapsed:.1f}s"
+        )
+
+    async def _wait_for_deployment_ready(
+        self,
+        api_client: kubernetes_asyncio.client.ApiClient,
+        name: str,
+        namespace: str,
+        timeout_seconds: int,
+    ) -> None:
+        """Wait for a Deployment to be ready."""
+        apps_api = kubernetes_asyncio.client.AppsV1Api(api_client)
+
+        async def check_deployment_ready():
+            deployment = await apps_api.read_namespaced_deployment(
+                name=name, namespace=namespace
+            )
+            status = deployment.status
+
+            if status.ready_replicas and status.replicas:
+                return status.ready_replicas >= status.replicas
+
+            return False
+
+        await self._wait_with_timeout(check_deployment_ready, timeout_seconds, name)
+
+    async def _wait_for_service_ready(
+        self,
+        api_client: kubernetes_asyncio.client.ApiClient,
+        name: str,
+        namespace: str,
+        timeout_seconds: int,
+    ) -> None:
+        """Wait for a Service to have endpoints."""
+        core_api = kubernetes_asyncio.client.CoreV1Api(api_client)
+
+        async def check_service_ready():
+            endpoints = await core_api.read_namespaced_endpoints(
+                name=name, namespace=namespace
+            )
+            if endpoints.subsets:
+                for subset in endpoints.subsets:
+                    if subset.addresses:  # These are ready pod addresses
+                        return True
+            return False
+
+        await self._wait_with_timeout(check_service_ready, timeout_seconds, name)
+
+    async def _wait_for_pod_ready(
+        self,
+        api_client: kubernetes_asyncio.client.ApiClient,
+        name: str,
+        namespace: str,
+        timeout_seconds: int,
+    ) -> None:
+        """Wait for a Pod to be ready."""
+        core_api = kubernetes_asyncio.client.CoreV1Api(api_client)
+
+        async def check_pod_ready():
+            pod = await core_api.read_namespaced_pod(name=name, namespace=namespace)
+            if pod.status.phase == "Running":
+                if pod.status.container_statuses:
+                    return all(status.ready for status in pod.status.container_statuses)
+                return True
+            return False
+
+        await self._wait_with_timeout(check_pod_ready, timeout_seconds, name)
+
+    async def _wait_with_timeout(
+        self, check_func, timeout_seconds: int, name: str
+    ) -> None:
+        """Generic timeout wrapper for readiness checks."""
+        start_time = time.time()
+
+        while time.time() - start_time < timeout_seconds:
+            try:
+                if await check_func():
+                    return
+            except kubernetes_asyncio.client.ApiException as e:
+                if e.status == 404:
+                    pass
+                else:
+                    wandb.termerror(
+                        f"{LOG_PREFIX}Error waiting for resource '{name}': {e}"
+                    )
+                    raise
+            except Exception as e:
+                wandb.termerror(f"{LOG_PREFIX}Error waiting for resource '{name}': {e}")
+                raise
+            await asyncio.sleep(2)
+
+        raise LaunchError(
+            f"Resource '{name}' not ready within {timeout_seconds} seconds"
+        )
+
+    async def _prepare_resource(
+        self,
+        api_client: kubernetes_asyncio.client.ApiClient,
+        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,
+        wait_timeout: int = 300,
+    ) -> None:
+        """Prepare a service for launch.
+
+        Arguments:
+            api_client: The Kubernetes API client.
+            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.
+            wait_timeout: Maximum time in seconds to wait for resource readiness.
+        """
+        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__]
+        )
+        wandb_config_env = {
+            "WANDB_CONFIG": env_vars.get("WANDB_CONFIG", "{}"),
+        }
+        add_wandb_env(config, wandb_config_env)
+
+        if api_key_secret:
+            for cont in yield_containers(config):
+                env = cont.setdefault("env", [])
+                env.append(
+                    {
+                        "name": "WANDB_API_KEY",
+                        "valueFrom": {
+                            "secretKeyRef": {
+                                "name": api_key_secret.metadata.name,
+                                "key": "password",
+                            }
+                        },
+                    }
+                )
+                cont["env"] = env
+
+        try:
+            await kubernetes_asyncio.utils.create_from_dict(
+                api_client, config, namespace=namespace
+            )
+
+            if wait_for_ready:
+                await self._wait_for_resource_ready(
+                    api_client, config, namespace, wait_timeout
+                )
+        except Exception as e:
+            wandb.termerror(f"{LOG_PREFIX}Failed to create Kubernetes resource: {e}")
+            raise LaunchError(f"Failed to create Kubernetes resource: {e}")
+
     async def run(
         self, launch_project: LaunchProject, image_uri: str
     ) -> Optional[AbstractRun]:
@@ -630,10 +919,42 @@ 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)
+
         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
+        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)
+
+            await asyncio.gather(
+                *[
+                    self._prepare_resource(
+                        api_client,
+                        resource.get("config"),
+                        namespace,
+                        launch_project.run_id,
+                        auxiliary_resource_label_key,
+                        launch_project,
+                        secret,
+                        wait_for_ready,
+                        wait_timeout,
+                    )
+                    for resource in additional_services
+                    if resource.get("config", {})
+                ]
+            )
+
         msg = "Creating Kubernetes job"
         if "name" in resource_args:
             msg += f": {resource_args['name']}"
@@ -658,7 +979,13 @@ class KubernetesRunner(AbstractRunner):
         job_name = job_response.metadata.name
         LaunchKubernetesMonitor.monitor_namespace(namespace)
         submitted_job = KubernetesSubmittedRun(
-            batch_api, core_api, job_name, namespace, secret
+            batch_api,
+            core_api,
+            apps_api,
+            job_name,
+            namespace,
+            secret,
+            auxiliary_resource_label_key,
         )
         if self.backend_config[PROJECT_SYNCHRONOUS]:
             await submitted_job.wait()

wandb/sdk/launch/sweeps/scheduler.py

@@ -36,7 +36,6 @@ if TYPE_CHECKING:
     import wandb.apis.public as public
     from wandb.apis.internal import Api
     from wandb.apis.public import QueuedRun, Run
-    from wandb.sdk.wandb_run import Run as SdkRun
 
 
 _logger = logging.getLogger(__name__)
@@ -255,10 +254,10 @@ class Scheduler(ABC):
             _id: w for _id, w in self._workers.items() if _id not in self.busy_workers
         }
 
-    def _init_wandb_run(self) -> "SdkRun":
+    def _init_wandb_run(self) -> "wandb.Run":
         """Controls resume or init logic for a scheduler wandb run."""
         settings = wandb.Settings(disable_job_creation=True)
-        run: SdkRun = wandb.init(  # type: ignore
+        run: wandb.Run = wandb.init(  # type: ignore
             name=f"Scheduler.{self._sweep_id}",
             resume="allow",
             config=self._kwargs,  # when run as a job, this sets config

wandb/sdk/lib/asyncio_compat.py

@@ -100,6 +100,9 @@ class _Runner:
                 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()
 

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,8 +2,10 @@ from __future__ import annotations
 
 import base64
 import hashlib
+import logging
 import mmap
 import sys
+import time
 from typing import TYPE_CHECKING, NewType
 
 from wandb.sdk.lib.paths import StrPath
@@ -15,6 +17,8 @@ ETag = NewType("ETag", str)
 HexMD5 = NewType("HexMD5", str)
 B64MD5 = NewType("B64MD5", str)
 
+logger = logging.getLogger(__name__)
+
 
 def _md5(data: bytes = b"") -> _hashlib.HASH:
     """Allow FIPS-compliant md5 hash when supported."""
@@ -44,7 +48,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/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,

wandb/sdk/lib/sock_client.py

@@ -110,10 +110,6 @@ class SockClient:
     def shutdown(self, val: int) -> None:
         self._sock.shutdown(val)
 
-    def set_socket(self, sock: socket.socket) -> None:
-        self._sock = sock
-        self._detect_bufsize()
-
     def _sendall_with_error_handle(self, data: bytes) -> None:
         # This is a helper function for sending data in a retry fashion.
         # Similar to the sendall() function in the socket module, but with