wandb 0.22.0__py3-none-win_amd64.whl → 0.22.1__py3-none-win_amd64.whl

wandb/sdk/interface/interface_sock.py

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any
 from typing_extensions import override
 
 from wandb.proto import wandb_server_pb2 as spb
+from wandb.sdk.lib import asyncio_manager
 
 from .interface_shared import InterfaceShared
 
@@ -21,10 +22,12 @@ logger = logging.getLogger("wandb")
 class InterfaceSock(InterfaceShared):
     def __init__(
         self,
+        asyncer: asyncio_manager.AsyncioManager,
         client: ServiceClient,
         stream_id: str,
     ) -> None:
         super().__init__()
+        self._asyncer = asyncer
         self._client = client
         self._stream_id = stream_id
 
@@ -37,13 +40,16 @@ class InterfaceSock(InterfaceShared):
         self._assign(record)
         request = spb.ServerRequest()
         request.record_publish.CopyFrom(record)
-        self._client.publish(request)
+        self._asyncer.run(lambda: self._client.publish(request))
 
-    @override
     def _deliver(self, record: pb.Record) -> MailboxHandle[pb.Result]:
+        return self._asyncer.run(lambda: self.deliver_async(record))
+
+    @override
+    async def deliver_async(self, record: pb.Record) -> MailboxHandle[pb.Result]:
         self._assign(record)
         request = spb.ServerRequest()
         request.record_publish.CopyFrom(record)
 
-        handle = self._client.deliver(request)
+        handle = await self._client.deliver(request)
         return handle.map(lambda response: response.result_communicate)

wandb/sdk/internal/_generated/__init__.py

@@ -1,15 +1,5 @@
 # Generated by ariadne-codegen
 
+__all__ = ["SERVER_FEATURES_QUERY_GQL", "ServerFeaturesQuery"]
 from .operations import SERVER_FEATURES_QUERY_GQL
-from .server_features_query import (
-    ServerFeaturesQuery,
-    ServerFeaturesQueryServerInfo,
-    ServerFeaturesQueryServerInfoFeatures,
-)
-
-__all__ = [
-    "SERVER_FEATURES_QUERY_GQL",
-    "ServerFeaturesQuery",
-    "ServerFeaturesQueryServerInfo",
-    "ServerFeaturesQueryServerInfoFeatures",
-]
+from .server_features_query import ServerFeaturesQuery

wandb/sdk/internal/sender.py

@@ -343,7 +343,7 @@ class SendManager:
         publish_interface = InterfaceQueue(record_q=record_q)
         context_keeper = context.ContextKeeper()
         return SendManager(
-            settings=SettingsStatic(settings.to_proto()),
+            settings=SettingsStatic(dict(settings)),
             record_q=record_q,
             result_q=result_q,
             interface=publish_interface,

wandb/sdk/internal/settings_static.py

@@ -2,9 +2,7 @@ from __future__ import annotations
 
 from typing import Any, Iterable
 
-from wandb.proto import wandb_settings_pb2
-from wandb.sdk.lib import RunMoment
-from wandb.sdk.wandb_settings import CLIENT_ONLY_SETTINGS, Settings
+from wandb.sdk.wandb_settings import Settings
 
 
 class SettingsStatic(Settings):
@@ -14,87 +12,9 @@ class SettingsStatic(Settings):
     attributes or items.
     """
 
-    def __init__(self, proto: wandb_settings_pb2.Settings) -> None:
-        data = self._proto_to_dict(proto)
+    def __init__(self, data: dict[str, Any]) -> None:
         super().__init__(**data)
 
-    def _proto_to_dict(self, proto: wandb_settings_pb2.Settings) -> dict:
-        data = {}
-
-        exclude_fields = {
-            "model_config",
-            "model_fields",
-            "model_fields_set",
-            "__fields__",
-            "__model_fields_set",
-            "__pydantic_self__",
-            "__pydantic_initialised__",
-        }
-
-        fields = (
-            Settings.model_fields
-            if hasattr(Settings, "model_fields")
-            else Settings.__fields__
-        )  # type: ignore [attr-defined]
-
-        fields = {k: v for k, v in fields.items() if k not in exclude_fields}  # type: ignore [union-attr]
-
-        forks_specified: list[str] = []
-        for key in fields:
-            if key in CLIENT_ONLY_SETTINGS:
-                continue
-
-            value: Any = None
-
-            field_info = fields[key]
-            annotation = str(field_info.annotation)
-
-            if key == "_stats_open_metrics_filters":
-                # todo: it's an underscored field, refactor into
-                #  something more elegant?
-                # I'm really about this. It's ugly, but it works.
-                # Do not try to repeat this at home.
-                value_type = getattr(proto, key).WhichOneof("value")
-                if value_type == "sequence":
-                    value = list(getattr(proto, key).sequence.value)
-                elif value_type == "mapping":
-                    unpacked_mapping = {}
-                    for outer_key, outer_value in getattr(
-                        proto, key
-                    ).mapping.value.items():
-                        unpacked_inner = {}
-                        for inner_key, inner_value in outer_value.value.items():
-                            unpacked_inner[inner_key] = inner_value
-                        unpacked_mapping[outer_key] = unpacked_inner
-                    value = unpacked_mapping
-            elif key == "fork_from" or key == "resume_from":
-                value = getattr(proto, key)
-                if value.run:
-                    value = RunMoment(
-                        run=value.run, value=value.value, metric=value.metric
-                    )
-                    forks_specified.append(key)
-                else:
-                    value = None
-            else:
-                if proto.HasField(key):  # type: ignore [arg-type]
-                    value = getattr(proto, key).value
-                    # Convert to list if the field is a sequence
-                    if any(t in annotation for t in ("tuple", "Sequence", "list")):
-                        value = list(value)
-                else:
-                    value = None
-
-            if value is not None:
-                data[key] = value
-
-        if len(forks_specified) > 1:
-            raise ValueError(
-                "Only one of fork_from or resume_from can be specified, not both"
-            )
-
-        return data
-
     def __setattr__(self, name: str, value: object) -> None:
         raise AttributeError("Error: SettingsStatic is a readonly object")
 

wandb/sdk/launch/runner/kubernetes_runner.py

@@ -27,7 +27,11 @@ from wandb.sdk.launch.runner.kubernetes_monitor import (
     CustomResource,
     LaunchKubernetesMonitor,
 )
-from wandb.sdk.launch.utils import recursive_macro_sub
+from wandb.sdk.launch.utils import (
+    recursive_macro_sub,
+    sanitize_identifiers_for_k8s,
+    yield_containers,
+)
 from wandb.sdk.lib.retry import ExponentialBackoff, retry_async
 from wandb.util import get_module
 
@@ -39,6 +43,7 @@ from ..utils import (
     MAX_ENV_LENGTHS,
     PROJECT_SYNCHRONOUS,
     get_kube_context_and_api_client,
+    make_k8s_label_safe,
     make_name_dns_safe,
 )
 from .abstract import AbstractRun, AbstractRunner
@@ -708,6 +713,7 @@ class KubernetesRunner(AbstractRunner):
         api_key_secret: Optional["V1Secret"] = None,
         wait_for_ready: bool = True,
         wait_timeout: int = 300,
+        auxiliary_resource_label_value: Optional[str] = None,
     ) -> None:
         """Prepare a service for launch.
 
@@ -725,6 +731,10 @@ class KubernetesRunner(AbstractRunner):
         config["metadata"].setdefault("labels", {})
         config["metadata"]["labels"][WANDB_K8S_RUN_ID] = run_id
         config["metadata"]["labels"]["wandb.ai/created-by"] = "launch-agent"
+        if auxiliary_resource_label_value:
+            config["metadata"]["labels"][WANDB_K8S_LABEL_AUXILIARY_RESOURCE] = (
+                auxiliary_resource_label_value
+            )
 
         env_vars = launch_project.get_env_vars_dict(
             self._api, MAX_ENV_LENGTHS[self.__class__.__name__]
@@ -734,6 +744,13 @@ class KubernetesRunner(AbstractRunner):
         }
         add_wandb_env(config, wandb_config_env)
 
+        if auxiliary_resource_label_value:
+            add_label_to_pods(
+                config,
+                WANDB_K8S_LABEL_AUXILIARY_RESOURCE,
+                auxiliary_resource_label_value,
+            )
+
         if api_key_secret:
             for cont in yield_containers(config):
                 env = cont.setdefault("env", [])
@@ -751,6 +768,8 @@ class KubernetesRunner(AbstractRunner):
                 cont["env"] = env
 
         try:
+            sanitize_identifiers_for_k8s(config)
+
             await kubernetes_asyncio.utils.create_from_dict(
                 api_client, config, namespace=namespace
             )
@@ -924,6 +943,9 @@ class KubernetesRunner(AbstractRunner):
         additional_services: List[Dict[str, Any]] = recursive_macro_sub(
             launch_project.launch_spec.get("additional_services", []), update_dict
         )
+        auxiliary_resource_label_value = make_k8s_label_safe(
+            f"aux-{launch_project.target_entity}-{launch_project.target_project}-{launch_project.run_id}"
+        )
         if additional_services:
             wandb.termlog(
                 f"{LOG_PREFIX}Creating additional services: {additional_services}"
@@ -943,6 +965,7 @@ class KubernetesRunner(AbstractRunner):
                         secret,
                         wait_for_ready,
                         wait_timeout,
+                        auxiliary_resource_label_value,
                     )
                     for resource in additional_services
                     if resource.get("config", {})
@@ -980,7 +1003,7 @@ class KubernetesRunner(AbstractRunner):
             job_name,
             namespace,
             secret,
-            f"aux-{launch_project.target_entity}-{launch_project.target_project}-{launch_project.run_id}",
+            auxiliary_resource_label_value,
         )
         if self.backend_config[PROJECT_SYNCHRONOUS]:
             await submitted_job.wait()
@@ -1131,24 +1154,6 @@ async def maybe_create_imagepull_secret(
         raise LaunchError(f"Exception when creating Kubernetes secret: {str(e)}\n")
 
 
-def yield_containers(root: Any) -> Iterator[dict]:
-    """Yield all container specs in a manifest.
-
-    Recursively traverses the manifest and yields all container specs. Container
-    specs are identified by the presence of a "containers" key in the value.
-    """
-    if isinstance(root, dict):
-        for k, v in root.items():
-            if k == "containers":
-                if isinstance(v, list):
-                    yield from v
-            elif isinstance(v, (dict, list)):
-                yield from yield_containers(v)
-    elif isinstance(root, list):
-        for item in root:
-            yield from yield_containers(item)
-
-
 def add_wandb_env(root: Union[dict, list], env_vars: Dict[str, str]) -> None:
     """Injects wandb environment variables into specs.
 

wandb/sdk/launch/utils.py

@@ -7,7 +7,17 @@ import re
 import subprocess
 import sys
 from collections import defaultdict
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+    Tuple,
+    Union,
+    cast,
+)
 
 import click
 
@@ -599,6 +609,32 @@ def make_name_dns_safe(name: str) -> str:
     return resp
 
 
+def make_k8s_label_safe(value: str) -> str:
+    """Return a Kubernetes label/identifier safe string (DNS-1123 label).
+
+    See:
+    https://kubernetes.io/docs/concepts/overview/working-with-objects/names/#dns-label-names
+
+    Rules:
+    - lowercase alphanumeric and '-'
+    - must start and end with an alphanumeric
+    - max length 63
+    """
+    # Normalize common separators first
+    safe = value.replace("_", "-").lower()
+    # Remove any invalid characters
+    safe = re.sub(r"[^a-z0-9\-]", "", safe)
+    # Collapse consecutive '-'
+    safe = re.sub(r"-+", "-", safe)
+    # Trim to 63 and strip leading/trailing '-'
+    safe = safe[:63].strip("-")
+
+    if not safe:
+        raise LaunchError(f"Invalid value for Kubernetes label: {value}")
+
+    return safe
+
+
 def warn_failed_packages_from_build_logs(
     log: str, image_uri: str, api: Api, job_tracker: Optional["JobAndRunStatusTracker"]
 ) -> None:
@@ -744,3 +780,48 @@ def get_current_python_version() -> Tuple[str, str]:
     major = full_version[0]
     version = ".".join(full_version[:2]) if len(full_version) >= 2 else major + ".0"
     return version, major
+
+
+def yield_containers(root: Union[dict, list]) -> Iterator[dict]:
+    """Yield all container specs in a manifest.
+
+    Recursively traverses the manifest and yields all container specs. Container
+    specs are identified by the presence of a "containers" key in the value.
+    """
+    if isinstance(root, dict):
+        for k, v in root.items():
+            if k == "containers":
+                if isinstance(v, list):
+                    yield from v
+            elif isinstance(v, (dict, list)):
+                yield from yield_containers(v)
+    elif isinstance(root, list):
+        for item in root:
+            yield from yield_containers(item)
+
+
+def sanitize_identifiers_for_k8s(root: Any) -> None:
+    if isinstance(root, list):
+        for item in root:
+            sanitize_identifiers_for_k8s(item)
+        return
+
+    # Only dicts have metadata and nested structures we need to sanitize.
+    if not isinstance(root, dict):
+        return
+
+    metadata = root.get("metadata")
+    if isinstance(metadata, dict):
+        if name := metadata.get("name"):
+            metadata["name"] = make_k8s_label_safe(str(name))
+
+    for container in yield_containers(root):
+        if name := container.get("name"):
+            container["name"] = make_k8s_label_safe(str(name))
+
+    # nested names
+    for key, value in root.items():
+        if isinstance(value, (dict, list)):
+            sanitize_identifiers_for_k8s(value)
+        elif key == "name" and isinstance(value, str):
+            root[key] = make_k8s_label_safe(value)

wandb/sdk/lib/progress.py

@@ -45,7 +45,11 @@ async def loop_printing_operation_stats(
         while True:
             start_time = time.monotonic()
 
-            handle = interface.deliver_operation_stats()
+            handle = await interface.deliver_async(
+                pb.Record(
+                    request=pb.Request(operations=pb.OperationStatsRequest()),
+                )
+            )
             result = await handle.wait_async(timeout=None)
             stats = result.response.operations_response.operation_stats
 
@@ -141,10 +145,9 @@ class ProgressPrinter:
             if extra_operations > 0:
                 line += f" (+ {extra_operations} more)"
 
-            if line != self._last_printed_line:
+            if line and line != self._last_printed_line:
                 self._printer.display(line)
-
-            self._last_printed_line = line
+                self._last_printed_line = line
 
     def _update_single_run(
         self,

wandb/sdk/lib/service/service_client.py

@@ -24,7 +24,6 @@ class ServiceClient:
         reader: asyncio.StreamReader,
         writer: asyncio.StreamWriter,
     ) -> None:
-        self._asyncer = asyncer
         self._reader = reader
         self._writer = writer
         self._mailbox = Mailbox(asyncer)
@@ -34,11 +33,11 @@ class ServiceClient:
             name="ServiceClient._forward_responses",
         )
 
-    def publish(self, request: spb.ServerRequest) -> None:
+    async def publish(self, request: spb.ServerRequest) -> None:
         """Send a request without waiting for a response."""
-        self._asyncer.run_soon(lambda: self._send_server_request(request))
+        await self._send_server_request(request)
 
-    def deliver(
+    async def deliver(
         self,
         request: spb.ServerRequest,
     ) -> MailboxHandle[spb.ServerResponse]:
@@ -52,7 +51,7 @@ class ServiceClient:
                 stopped due to an error.
         """
         handle = self._mailbox.require_response(request)
-        self._asyncer.run_soon(lambda: self._send_server_request(request))
+        await self._send_server_request(request)
         return handle
 
     async def _send_server_request(self, request: spb.ServerRequest) -> None:
@@ -64,11 +63,8 @@ class ServiceClient:
 
         await self._writer.drain()
 
-    def close(self) -> None:
+    async def close(self) -> None:
         """Flush and close the socket."""
-        self._asyncer.run_soon(self._close)
-
-    async def _close(self) -> None:
         self._writer.close()
         await self._writer.wait_closed()
 

wandb/sdk/lib/service/service_connection.py

@@ -31,6 +31,7 @@ def connect_to_service(
 
     if token:
         return ServiceConnection(
+            asyncer=asyncer,
             client=token.connect(asyncer=asyncer),
             proc=None,
         )
@@ -60,6 +61,7 @@ def _start_and_connect_service(
         conn.teardown(hooks.exit_code)
 
     conn = ServiceConnection(
+        asyncer=asyncer,
         client=client,
         proc=proc,
         cleanup=lambda: atexit.unregister(teardown_atexit),
@@ -71,10 +73,14 @@ def _start_and_connect_service(
 
 
 class ServiceConnection:
-    """A connection to the W&B internal service process."""
+    """A connection to the W&B internal service process.
+
+    None of the synchronous methods may be called in an asyncio context.
+    """
 
     def __init__(
         self,
+        asyncer: asyncio_manager.AsyncioManager,
         client: ServiceClient,
         proc: service_process.ServiceProcess | None,
         cleanup: Callable[[], None] | None = None,
@@ -82,13 +88,12 @@ class ServiceConnection:
         """Returns a new ServiceConnection.
 
         Args:
-            mailbox: The mailbox to use for all communication over the socket.
-            router: A handle to the thread that reads from the socket and
-                updates the mailbox.
-            client: A socket that's connected to the service.
+            asyncer: An asyncio runner.
+            client: A client for communicating with the service over a socket.
             proc: The service process if we own it, or None otherwise.
             cleanup: A callback to run on teardown before doing anything.
         """
+        self._asyncer = asyncer
         self._client = client
         self._proc = proc
         self._torn_down = False
@@ -96,9 +101,13 @@ class ServiceConnection:
 
     def make_interface(self, stream_id: str) -> InterfaceBase:
         """Returns an interface for communicating with the service."""
-        return InterfaceSock(self._client, stream_id=stream_id)
+        return InterfaceSock(
+            self._asyncer,
+            self._client,
+            stream_id=stream_id,
+        )
 
-    def init_sync(
+    async def init_sync(
         self,
         paths: set[pathlib.Path],
         settings: wandb_settings.Settings,
@@ -110,10 +119,10 @@ class ServiceConnection:
         )
         request = spb.ServerRequest(init_sync=init_sync)
 
-        handle = self._client.deliver(request)
+        handle = await self._client.deliver(request)
         return handle.map(lambda r: r.init_sync_response)
 
-    def sync(
+    async def sync(
         self,
         id: str,
         *,
@@ -123,10 +132,10 @@ class ServiceConnection:
         sync = wandb_sync_pb2.ServerSyncRequest(id=id, parallelism=parallelism)
         request = spb.ServerRequest(sync=sync)
 
-        handle = self._client.deliver(request)
+        handle = await self._client.deliver(request)
         return handle.map(lambda r: r.sync_response)
 
-    def sync_status(
+    async def sync_status(
         self,
         id: str,
     ) -> MailboxHandle[wandb_sync_pb2.ServerSyncStatusResponse]:
@@ -134,7 +143,7 @@ class ServiceConnection:
         sync_status = wandb_sync_pb2.ServerSyncStatusRequest(id=id)
         request = spb.ServerRequest(sync_status=sync_status)
 
-        handle = self._client.deliver(request)
+        handle = await self._client.deliver(request)
         return handle.map(lambda r: r.sync_status_response)
 
     def inform_init(
@@ -146,13 +155,17 @@ class ServiceConnection:
         request = spb.ServerInformInitRequest()
         request.settings.CopyFrom(settings)
         request._info.stream_id = run_id
-        self._client.publish(spb.ServerRequest(inform_init=request))
+        self._asyncer.run(
+            lambda: self._client.publish(spb.ServerRequest(inform_init=request))
+        )
 
     def inform_finish(self, run_id: str) -> None:
         """Send an finish request to the service."""
         request = spb.ServerInformFinishRequest()
         request._info.stream_id = run_id
-        self._client.publish(spb.ServerRequest(inform_finish=request))
+        self._asyncer.run(
+            lambda: self._client.publish(spb.ServerRequest(inform_finish=request))
+        )
 
     def inform_attach(
         self,
@@ -166,7 +179,7 @@ class ServiceConnection:
         request.inform_attach._info.stream_id = attach_id
 
         try:
-            handle = self._client.deliver(request)
+            handle = self._asyncer.run(lambda: self._client.deliver(request))
             response = handle.wait_or(timeout=10)
 
         except (MailboxClosedError, HandleAbandonedError):
@@ -210,13 +223,16 @@ class ServiceConnection:
         # Clear the service token to prevent new connections to the process.
         service_token.clear_service_in_env()
 
-        self._client.publish(
-            spb.ServerRequest(
-                inform_teardown=spb.ServerInformTeardownRequest(
-                    exit_code=exit_code,
-                )
-            ),
-        )
-        self._client.close()
+        async def publish_teardown_and_close() -> None:
+            await self._client.publish(
+                spb.ServerRequest(
+                    inform_teardown=spb.ServerInformTeardownRequest(
+                        exit_code=exit_code,
+                    )
+                ),
+            )
+            await self._client.close()
+
+        self._asyncer.run(publish_teardown_and_close)
 
         return self._proc.join()

wandb/sdk/mailbox/mailbox_handle.py

@@ -58,6 +58,8 @@ class MailboxHandle(abc.ABC, Generic[_T]):
     def cancel(self, iface: interface.InterfaceBase) -> None:
         """Cancel the handle, requesting any associated work to not complete.
 
+        It is an error to call this from an async function.
+
         This automatically abandons the handle, as a response is no longer
         guaranteed.
 

wandb/sdk/projects/_generated/__init__.py

@@ -1,47 +1,26 @@
 # Generated by ariadne-codegen
 
-from .delete_project import DeleteProject, DeleteProjectDeleteModel
-from .fetch_registry import FetchRegistry, FetchRegistryEntity
-from .fragments import (
-    RegistryFragment,
-    RegistryFragmentArtifactTypes,
-    RegistryFragmentArtifactTypesEdges,
-    RegistryFragmentArtifactTypesEdgesNode,
-)
-from .input_types import ArtifactTypeInput
-from .operations import (
-    DELETE_PROJECT_GQL,
-    FETCH_REGISTRY_GQL,
-    RENAME_PROJECT_GQL,
-    UPSERT_REGISTRY_PROJECT_GQL,
-)
-from .rename_project import (
-    RenameProject,
-    RenameProjectRenameProject,
-    RenameProjectRenameProjectProject,
-)
-from .upsert_registry_project import (
-    UpsertRegistryProject,
-    UpsertRegistryProjectUpsertModel,
-)
-
 __all__ = [
     "DELETE_PROJECT_GQL",
     "FETCH_REGISTRY_GQL",
     "RENAME_PROJECT_GQL",
     "UPSERT_REGISTRY_PROJECT_GQL",
     "FetchRegistry",
-    "FetchRegistryEntity",
     "RenameProject",
-    "RenameProjectRenameProject",
-    "RenameProjectRenameProjectProject",
     "UpsertRegistryProject",
-    "UpsertRegistryProjectUpsertModel",
     "DeleteProject",
-    "DeleteProjectDeleteModel",
     "ArtifactTypeInput",
     "RegistryFragment",
-    "RegistryFragmentArtifactTypes",
-    "RegistryFragmentArtifactTypesEdges",
-    "RegistryFragmentArtifactTypesEdgesNode",
 ]
+from .delete_project import DeleteProject
+from .fetch_registry import FetchRegistry
+from .fragments import RegistryFragment
+from .input_types import ArtifactTypeInput
+from .operations import (
+    DELETE_PROJECT_GQL,
+    FETCH_REGISTRY_GQL,
+    RENAME_PROJECT_GQL,
+    UPSERT_REGISTRY_PROJECT_GQL,
+)
+from .rename_project import RenameProject
+from .upsert_registry_project import UpsertRegistryProject