wandb 0.22.0__py3-none-macosx_12_0_arm64.whl → 0.22.2__py3-none-macosx_12_0_arm64.whl

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
 
@@ -89,7 +93,6 @@ class ProgressPrinter:
         progress_text_area: p.DynamicText | None,
         default_text: str,
     ) -> None:
-        self._show_operation_stats = True
         self._printer = printer
         self._progress_text_area = progress_text_area
         self._default_text = default_text
@@ -106,14 +109,10 @@ class ProgressPrinter:
 
         if isinstance(progress, pb.OperationStats):
             self._update_operation_stats([progress])
-        elif self._show_operation_stats:
+        else:
             self._update_operation_stats(
                 list(response.operation_stats for response in progress)
             )
-        elif len(progress) == 1:
-            self._update_single_run(progress[0])
-        else:
-            self._update_multiple_runs(progress)
 
         self._tick += 1
 
@@ -141,69 +140,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
-
-    def _update_single_run(
-        self,
-        progress: pb.PollExitResponse,
-    ) -> None:
-        stats = progress.pusher_stats
-        line = (
-            f"{_megabytes(stats.uploaded_bytes):.3f} MB"
-            f" of {_megabytes(stats.total_bytes):.3f} MB uploaded"
-        )
-
-        if stats.deduped_bytes > 0:
-            line += f" ({_megabytes(stats.deduped_bytes):.3f} MB deduped)"
-
-        if stats.total_bytes > 0:
-            self._update_progress_text(
-                line,
-                stats.uploaded_bytes / stats.total_bytes,
-            )
-        else:
-            self._update_progress_text(line, 1.0)
-
-    def _update_multiple_runs(
-        self,
-        progress_list: list[pb.PollExitResponse],
-    ) -> None:
-        total_files = 0
-        uploaded_bytes = 0
-        total_bytes = 0
-
-        for progress in progress_list:
-            total_files += progress.file_counts.wandb_count
-            total_files += progress.file_counts.media_count
-            total_files += progress.file_counts.artifact_count
-            total_files += progress.file_counts.other_count
-
-            uploaded_bytes += progress.pusher_stats.uploaded_bytes
-            total_bytes += progress.pusher_stats.total_bytes
-
-        line = (
-            f"Processing {len(progress_list)} runs with {total_files} files"
-            f" ({_megabytes(uploaded_bytes):.2f} MB"
-            f" / {_megabytes(total_bytes):.2f} MB)"
-        )
-
-        if total_bytes > 0:
-            self._update_progress_text(line, uploaded_bytes / total_bytes)
-        else:
-            self._update_progress_text(line, 1.0)
-
-    def _update_progress_text(self, text: str, progress: float) -> None:
-        if text == self._last_printed_line:
-            return
-        self._last_printed_line = text
-
-        if self._progress_text_area:
-            self._progress_text_area.set_text(text)
-        else:
-            self._printer.progress_update(text + "\r", progress)
+                self._last_printed_line = line
 
 
 class _DynamicOperationStatsPrinter:
@@ -315,8 +254,3 @@ def _time_to_string(seconds: float) -> str:
     hours = int(seconds / (60 * 60))
     minutes = int((seconds / 60) % 60)
     return f"{hours}h{minutes}m"
-
-
-def _megabytes(bytes: int) -> float:
-    """Returns the number of megabytes in `bytes`."""
-    return bytes / (1 << 20)

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

wandb/sdk/wandb_init.py

@@ -406,6 +406,25 @@ class _WandbInit:
                 run_id=settings.run_id,
             )
 
+    def set_sync_dir_suffix(self, settings: Settings) -> None:
+        """Add a suffix to sync_dir if it already exists.
+
+        The sync_dir uses a timestamp with second-level precision which can
+        result in conflicts if a run with the same ID is initialized within the
+        same second. This is most likely to happen in tests.
+
+        This can't prevent conflicts from multiple processes attempting
+        to create a wandb run simultaneously.
+
+        Args:
+            settings: Fully initialized settings other than the
+                x_sync_dir_suffix setting which will be modified.
+        """
+        index = 1
+        while pathlib.Path(settings.sync_dir).exists():
+            settings.x_sync_dir_suffix = f"{index}"
+            index += 1
+
     def make_run_config(
         self,
         settings: Settings,
@@ -708,7 +727,7 @@ class _WandbInit:
         drun = Run(
             settings=Settings(
                 mode="disabled",
-                x_files_dir=tempfile.gettempdir(),
+                root_dir=tempfile.gettempdir(),
                 run_id=run_id,
                 run_tags=tuple(),
                 run_notes=None,
@@ -1407,8 +1426,8 @@ def init(  # noqa: C901
             enable on your settings page.
         tensorboard: Deprecated. Use `sync_tensorboard` instead.
         sync_tensorboard: Enables automatic syncing of W&B logs from TensorBoard
-            or TensorBoardX, saving relevant event files for viewing in the W&B UI.
-            saving relevant event files for viewing in the W&B UI. (Default: `False`)
+            or TensorBoardX, saving relevant event files for viewing in
+            the W&B UI.
         monitor_gym: Enables automatic logging of videos of the environment when
             using OpenAI Gym.
         settings: Specifies a dictionary or `wandb.Settings` object with advanced
@@ -1524,6 +1543,7 @@ def init(  # noqa: C901
             init_telemetry.feature.rewind_mode = True
 
         wi.set_run_id(run_settings)
+        wi.set_sync_dir_suffix(run_settings)
         run_printer = printer.new_printer(run_settings)
         show_warnings(run_printer)
 

wandb/sdk/wandb_login.py

@@ -77,7 +77,7 @@ def login(
         UsageError: If `api_key` cannot be configured and no tty.
     """
     _handle_host_wandb_setting(host)
-    return _login(
+    logged_in, _ = _login(
         anonymous=anonymous,
         key=key,
         relogin=relogin,
@@ -87,6 +87,7 @@ def login(
         verify=verify,
         referrer=referrer,
     )
+    return logged_in
 
 
 class ApiKeyStatus(enum.Enum):
@@ -245,24 +246,6 @@ class _WandbLogin:
 
         return key, status
 
-    def _verify_login(self, key: str) -> None:
-        api = InternalApi(api_key=key)
-
-        try:
-            is_api_key_valid = api.validate_api_key()
-        except ConnectionError:
-            raise AuthenticationError(
-                "Unable to connect to server to verify API token."
-            )
-        except Exception:
-            raise AuthenticationError("An error occurred while verifying the API key.")
-
-        if not is_api_key_valid:
-            raise AuthenticationError(
-                f"API key verification failed for host {self._settings.base_url}."
-                " Make sure your API key is valid."
-            )
-
 
 def _login(
     *,
@@ -277,11 +260,30 @@ def _login(
     update_api_key: bool = True,
     _silent: Optional[bool] = None,
     _disable_warning: Optional[bool] = None,
-) -> bool:
+) -> (bool, Optional[str]):
+    """Logs in to W&B.
+
+    This is the internal implementation of wandb.login(),
+    with many of the same arguments as wandb.login().
+    Additional arguments are documented below.
+
+    Args:
+        update_api_key: If true, the api key will be saved or updated
+            in the users .netrc file.
+        _silent: If true, will not print any messages to the console.
+        _disable_warning: If true, no warning will be displayed
+            when calling wandb.login() after wandb.init().
+
+    Returns:
+        bool: If the login was successful
+            or the user is assumed to be already be logged in.
+        str: The API key used to log in,
+            or None if the api key was not verified during the login process.
+    """
     if wandb.run is not None:
         if not _disable_warning:
             wandb.termwarn("Calling wandb.login() after wandb.init() has no effect.")
-        return True
+        return True, None
 
     wlogin = _WandbLogin(
         anonymous=anonymous,
@@ -293,19 +295,19 @@ def _login(
     )
 
     if wlogin._settings._noop:
-        return True
+        return True, None
 
     if wlogin._settings._offline and not wlogin._settings.x_cli_only_mode:
         wandb.termwarn("Unable to verify login in offline mode.")
-        return False
+        return False, None
     elif wandb.util._is_kaggle() and not wandb.util._has_internet():
         wandb.termerror(
             "To use W&B in kaggle you must enable internet in the settings panel on the right."
         )
-        return False
+        return False, None
 
     if wlogin._settings.identity_token_file:
-        return True
+        return True, None
 
     key_is_pre_configured = False
     key_status = None
@@ -318,7 +320,7 @@ def _login(
             key, key_status = wlogin.prompt_api_key(referrer=referrer)
 
     if verify:
-        wlogin._verify_login(key)
+        _verify_login(key, wlogin._settings.base_url)
 
     if not key_is_pre_configured:
         if update_api_key:
@@ -329,4 +331,28 @@ def _login(
     if key and not _silent:
         wlogin._print_logged_in_message()
 
-    return key is not None
+    return key is not None, key
+
+
+def _verify_login(key: str, base_url: str) -> None:
+    api = InternalApi(
+        api_key=key,
+        default_settings={"base_url": base_url},
+    )
+
+    try:
+        is_api_key_valid = api.validate_api_key()
+    except ConnectionError:
+        raise AuthenticationError(
+            "Unable to connect to server to verify API token."
+        ) from None
+    except Exception as e:
+        raise AuthenticationError(
+            "An error occurred while verifying the API key."
+        ) from e
+
+    if not is_api_key_valid:
+        raise AuthenticationError(
+            f"API key verification failed for host {base_url}."
+            " Make sure your API key is valid."
+        )