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

wandb/sdk/launch/inputs/schema.py

@@ -7,6 +7,9 @@ META_SCHEMA = {
         },
         "title": {"type": "string"},
         "description": {"type": "string"},
+        "label": {"type": "string"},
+        "placeholder": {"type": "string"},
+        "required": {"type": "boolean"},
         "format": {"type": "string"},
         "enum": {"type": "array", "items": {"type": ["integer", "number", "string"]}},
         "properties": {"type": "object", "patternProperties": {".*": {"$ref": "#"}}},
@@ -19,24 +22,24 @@ META_SCHEMA = {
     },
     "allOf": [
         {
-            "if": {"properties": {"type": {"const": "number"}}},
+            "if": {"properties": {"type": {"const": "integer"}}},
             "then": {
                 "properties": {
-                    "minimum": {"type": ["integer", "number"]},
-                    "maximum": {"type": ["integer", "number"]},
-                    "exclusiveMinimum": {"type": ["integer", "number"]},
-                    "exclusiveMaximum": {"type": ["integer", "number"]},
+                    "minimum": {"type": "integer"},
+                    "maximum": {"type": "integer"},
+                    "exclusiveMinimum": {"type": "integer"},
+                    "exclusiveMaximum": {"type": "integer"},
                 }
             },
         },
         {
-            "if": {"properties": {"type": {"const": "integer"}}},
+            "if": {"properties": {"type": {"const": "number"}}},
             "then": {
                 "properties": {
-                    "minimum": {"type": "integer"},
-                    "maximum": {"type": "integer"},
-                    "exclusiveMinimum": {"type": "integer"},
-                    "exclusiveMaximum": {"type": "integer"},
+                    "minimum": {"type": ["integer", "number"]},
+                    "maximum": {"type": ["integer", "number"]},
+                    "exclusiveMinimum": {"type": ["integer", "number"]},
+                    "exclusiveMaximum": {"type": ["integer", "number"]},
                 }
             },
         },

wandb/sdk/lib/apikey.py

@@ -136,12 +136,6 @@ def prompt_api_key(  # noqa: C901
     if (jupyter and not settings.login_timeout) or no_create:
         choices.remove(LOGIN_CHOICE_NEW)
 
-    if jupyter and "google.colab" in sys.modules:
-        log_string = term.LOG_STRING_NOCOLOR
-        key = wandb.jupyter.attempt_colab_login(app_url)  # type: ignore
-        if key is not None:
-            return key  # type: ignore
-
     if anon_mode == "must":
         result = LOGIN_CHOICE_ANON
     # If we're not in an interactive environment, default to dry-run.
@@ -236,9 +230,9 @@ def check_netrc_access(
 def write_netrc(host: str, entity: str, key: str):
     """Add our host and key to .netrc."""
     _, key_suffix = key.split("-", 1) if "-" in key else ("", key)
-    if len(key_suffix) != 40:
+    if len(key_suffix) < 40:
         raise ValueError(
-            f"API-key must be exactly 40 characters long: {key_suffix} ({len(key_suffix)} chars)"
+            f"API-key must be at least 40 characters long: {key_suffix} ({len(key_suffix)} chars)"
         )
 
     normalized_host = urlparse(host).netloc
@@ -305,12 +299,14 @@ def write_key(
     # TODO(jhr): api shouldn't be optional or it shouldn't be passed, clean up callers
     api = api or InternalApi()
 
-    # Normal API keys are 40-character hex strings. On-prem API keys have a
-    # variable-length prefix, a dash, then the 40-char string.
+    # API keys are strings of at least 40 characters. On-prem API keys have a
+    # variable-length prefix, a dash, then the string of at least 40 chars.
     _, suffix = key.split("-", 1) if "-" in key else ("", key)
 
-    if len(suffix) != 40:
-        raise ValueError(f"API key must be 40 characters long, yours was {len(key)}")
+    if len(suffix) < 40:
+        raise ValueError(
+            f"API key must be at least 40 characters long, yours was {len(key)}"
+        )
 
     write_netrc(settings.base_url, "user", key)
 

wandb/sdk/lib/asyncio_compat.py

@@ -133,7 +133,7 @@ class TaskGroup:
     """Object that `open_task_group()` yields."""
 
     def __init__(self) -> None:
-        self._tasks: list[asyncio.Task] = []
+        self._tasks: list[asyncio.Task[None]] = []
 
     def start_soon(self, coro: Coroutine[Any, Any, Any]) -> None:
         """Schedule a task in the group.

wandb/sdk/lib/asyncio_manager.py

@@ -7,7 +7,7 @@ import concurrent.futures
 import contextlib
 import logging
 import threading
-from typing import Any, Callable, Coroutine, TypeVar
+from typing import Awaitable, Callable, TypeVar
 
 from . import asyncio_compat
 
@@ -104,7 +104,7 @@ class AsyncioManager:
             # This only matters if the KeyboardInterrupt is suppressed.
             self._runner.cancel()
 
-    def run(self, fn: Callable[[], Coroutine[Any, Any, _T]]) -> _T:
+    def run(self, fn: Callable[[], Awaitable[_T]]) -> _T:
         """Run an async function to completion.
 
         The function is called in the asyncio thread. Blocks until start()
@@ -148,7 +148,7 @@ class AsyncioManager:
 
     def run_soon(
         self,
-        fn: Callable[[], Coroutine[Any, Any, None]],
+        fn: Callable[[], Awaitable[None]],
         *,
         daemon: bool = False,
         name: str | None = None,
@@ -186,7 +186,7 @@ class AsyncioManager:
 
     def _schedule(
         self,
-        fn: Callable[[], Coroutine[Any, Any, _T]],
+        fn: Callable[[], Awaitable[_T]],
         daemon: bool,
         name: str | None = None,
     ) -> concurrent.futures.Future[_T]:
@@ -207,7 +207,7 @@ class AsyncioManager:
 
     async def _wrap(
         self,
-        fn: Callable[[], Coroutine[Any, Any, _T]],
+        fn: Callable[[], Awaitable[_T]],
         daemon: bool,
         name: str | None,
     ) -> _T:

wandb/sdk/lib/console_capture.py

@@ -75,9 +75,12 @@ class _WriteCallback(Protocol):
         """
 
 
-# A reentrant lock is used to catch callbacks that write to stderr/stdout.
-_module_rlock = threading.RLock()
+_module_lock = threading.Lock()
 _is_writing = False
+"""Prevents infinite print-capture loops.
+
+If a capture callback prints, that output is not captured.
+"""
 
 _patch_exception: CannotCaptureConsoleError | None = None
 
@@ -99,7 +102,7 @@ def capture_stdout(callback: _WriteCallback) -> Callable[[], None]:
     Raises:
         CannotCaptureConsoleError: If patching failed on import.
     """
-    with _module_rlock:
+    with _module_lock:
         if _patch_exception:
             raise _patch_exception
 
@@ -121,7 +124,7 @@ def capture_stderr(callback: _WriteCallback) -> Callable[[], None]:
     Raises:
         CannotCaptureConsoleError: If patching failed on import.
     """
-    with _module_rlock:
+    with _module_lock:
         if _patch_exception:
             raise _patch_exception
 
@@ -144,7 +147,7 @@ def _insert_disposably(
     def dispose() -> None:
         nonlocal disposed
 
-        with _module_rlock:
+        with _module_lock:
             if disposed:
                 return
 
@@ -167,38 +170,43 @@ def _patch(
         global _is_writing
         n = orig_write(s)
 
-        # NOTE: Since _module_rlock is reentrant, this is safe. It will not
-        # deadlock if a callback invokes write() again.
-        with _module_rlock:
+        with _module_lock:
             if _is_writing:
                 return n
-
             _is_writing = True
-            try:
-                for cb in callbacks.values():
-                    cb(s, n)
-
-            except BaseException as e:
-                # Clear all callbacks on any exception to avoid infinite loops:
-                #
-                # * If we re-raise, an exception handler is likely to print
-                #   the exception to the console and trigger callbacks again
-                # * If we log, we can't guarantee that this doesn't print
-                #   to console.
-                #
-                # This is especially important for KeyboardInterrupt.
+
+            # Invoke callbacks outside of the lock to avoid deadlocks.
+            # 1. A callback may print, invoking this again.
+            # 2. A callback may block on a different thread which then prints.
+            callback_list = list(callbacks.values())
+
+        try:
+            for cb in callback_list:
+                cb(s, n)
+
+        except BaseException as e:
+            # Clear all callbacks on any exception to avoid infinite loops:
+            #
+            # * If we re-raise, an exception handler is likely to print
+            #   the exception to the console and trigger callbacks again
+            # * If we log, we can't guarantee that this doesn't print
+            #   to console.
+            #
+            # This is especially important for KeyboardInterrupt.
+            with _module_lock:
                 _stderr_callbacks.clear()
                 _stdout_callbacks.clear()
 
-                if isinstance(e, Exception):
-                    # We suppress Exceptions so that bugs in W&B code don't
-                    # cause the user's print() statements to raise errors.
-                    _logger.exception("Error in console callback, clearing all!")
-                else:
-                    # Re-raise errors like KeyboardInterrupt.
-                    raise
+            if isinstance(e, Exception):
+                # We suppress Exceptions so that bugs in W&B code don't
+                # cause the user's print() statements to raise errors.
+                _logger.exception("Error in console callback, clearing all!")
+            else:
+                # Re-raise errors like KeyboardInterrupt.
+                raise
 
-            finally:
+        finally:
+            with _module_lock:
                 _is_writing = False
 
         return n

wandb/sdk/lib/progress.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 import asyncio
 import contextlib
 import time
-from typing import Iterable, Iterator, NoReturn
+from typing import Iterator, NoReturn
 
 from wandb.proto import wandb_internal_pb2 as pb
 from wandb.sdk.interface import interface
@@ -13,6 +13,10 @@ from wandb.sdk.lib import asyncio_compat
 
 from . import printer as p
 
+_INDENT = "  "
+_MAX_LINES_TO_PRINT = 6
+_MAX_OPS_TO_PRINT = 5
+
 
 async def loop_printing_operation_stats(
     progress: ProgressPrinter,
@@ -96,101 +100,189 @@ class ProgressPrinter:
         self._printer = printer
         self._progress_text_area = progress_text_area
         self._default_text = default_text
-        self._tick = 0
+        self._tick = -1
         self._last_printed_line = ""
 
     def update(
         self,
-        progress: list[pb.PollExitResponse] | pb.OperationStats,
+        stats_or_groups: pb.OperationStats | dict[str, pb.OperationStats],
     ) -> None:
-        """Update the displayed information."""
-        if not progress:
+        """Update the displayed information.
+
+        Args:
+            stats_or_groups: A single group of operations, or zero or more
+                labeled operation groups.
+        """
+        self._tick += 1
+
+        if not self._progress_text_area:
+            line = self._to_static_text(stats_or_groups)
+            if line and line != self._last_printed_line:
+                self._printer.display(line)
+                self._last_printed_line = line
             return
 
-        if isinstance(progress, pb.OperationStats):
-            self._update_operation_stats([progress])
-        else:
-            self._update_operation_stats(
-                list(response.operation_stats for response in progress)
-            )
+        lines = self._to_dynamic_text(stats_or_groups)
+        if not lines:
+            loading_symbol = self._printer.loading_symbol(self._tick)
+            if loading_symbol:
+                lines = [f"{loading_symbol} {self._default_text}"]
+            else:
+                lines = [self._default_text]
 
-        self._tick += 1
+        self._progress_text_area.set_text("\n".join(lines))
 
-    def _update_operation_stats(self, stats_list: list[pb.OperationStats]) -> None:
-        if self._progress_text_area:
-            _DynamicOperationStatsPrinter(
+    def _to_dynamic_text(
+        self,
+        stats_or_groups: pb.OperationStats | dict[str, pb.OperationStats],
+    ) -> list[str]:
+        """Returns text to show in a dynamic text area."""
+        loading_symbol = self._printer.loading_symbol(self._tick)
+
+        if isinstance(stats_or_groups, dict):
+            return _GroupedOperationStatsPrinter(
                 self._printer,
-                self._progress_text_area,
-                max_lines=6,
-                loading_symbol=self._printer.loading_symbol(self._tick),
-                default_text=self._default_text,
-            ).display(stats_list)
+                _MAX_LINES_TO_PRINT,
+                loading_symbol,
+            ).render(stats_or_groups)
 
         else:
-            top_level_operations: list[str] = []
-            extra_operations = 0
-            for stats in stats_list:
-                for op in stats.operations:
-                    if len(top_level_operations) < 5:
-                        top_level_operations.append(op.desc)
-                    else:
-                        extra_operations += 1
-
-            line = "; ".join(top_level_operations)
-            if extra_operations > 0:
-                line += f" (+ {extra_operations} more)"
+            return _OperationStatsPrinter(
+                self._printer,
+                _MAX_LINES_TO_PRINT,
+                loading_symbol,
+            ).render(stats_or_groups)
 
-            if line and line != self._last_printed_line:
-                self._printer.display(line)
-                self._last_printed_line = line
+    def _to_static_text(
+        self,
+        stats_or_groups: pb.OperationStats | dict[str, pb.OperationStats],
+    ) -> str:
+        """Returns a single line of text to print out."""
+        if isinstance(stats_or_groups, dict):
+            sorted_prefixed_stats = list(
+                (f"[{group}] ", stats)  #
+                for group, stats in sorted(stats_or_groups.items())
+            )
+        else:
+            sorted_prefixed_stats = [("", stats_or_groups)]
+
+        group_strs: list[str] = []
+        total_operations = 0
+        total_printed = 0
 
+        for prefix, stats in sorted_prefixed_stats:
+            total_operations += stats.total_operations
+            if not stats.operations:
+                continue
+
+            group_ops: list[str] = []
+            i = 0
+            while total_printed < _MAX_OPS_TO_PRINT and i < len(stats.operations):
+                group_ops.append(stats.operations[i].desc)
+                total_printed += 1
+                i += 1
+
+            if group_ops:
+                group_strs.append(prefix + "; ".join(group_ops))
+
+        line = "; ".join(group_strs)
+        remaining = total_operations - total_printed
+        if total_printed > 0 and remaining > 0:
+            line += f" (+ {remaining} more)"
 
-class _DynamicOperationStatsPrinter:
-    """Single-use object that writes operation stats into a text area."""
+        return line
+
+
+class _GroupedOperationStatsPrinter:
+    """Renders a list of labeled operation stats groups into lines of text."""
 
     def __init__(
         self,
         printer: p.Printer,
-        text_area: p.DynamicText,
         max_lines: int,
         loading_symbol: str,
-        default_text: str,
     ) -> None:
         self._printer = printer
-        self._text_area = text_area
         self._max_lines = max_lines
         self._loading_symbol = loading_symbol
-        self._default_text = default_text
 
-        self._lines: list[str] = []
-        self._ops_shown = 0
+    def render(self, groups: dict[str, pb.OperationStats]) -> list[str]:
+        """Convert labeled operation stats groups into text to display.
+
+        Args:
+            groups: A mapping from group labels to stats for that group.
+
+        Returns:
+            The lines of text to print. The lines do not end with the newline
+            character. Returns an empty list if there are no operations.
+        """
+        lines: list[str] = []
+
+        for key, stats in sorted(groups.items()):
+            # Don't display empty groups.
+            if not stats.operations:
+                continue
+
+            # Ensure enough space left for the group header and at least
+            # one line of content.
+            remaining_lines = self._max_lines - len(lines)
+            if remaining_lines < 2:
+                break
+
+            # Group header.
+            lines.append(key)
 
-    def display(
+            # Group content.
+            stats_lines = _OperationStatsPrinter(
+                printer=self._printer,
+                max_lines=remaining_lines - 1,  # minus one for the header
+                loading_symbol=self._loading_symbol,
+            ).render(stats)
+            for line in stats_lines:
+                lines.append(f"{_INDENT}{line}")
+
+        return lines
+
+
+class _OperationStatsPrinter:
+    """Renders operation stats into lines of text."""
+
+    def __init__(
         self,
-        stats_list: Iterable[pb.OperationStats],
+        printer: p.Printer,
+        max_lines: int,
+        loading_symbol: str,
     ) -> None:
-        """Show the given stats in the text area."""
-        total_operations = 0
-        for stats in stats_list:
-            for op in stats.operations:
-                self._add_operation(op, is_subtask=False, indent="")
-            total_operations += stats.total_operations
+        self._printer = printer
+        self._max_lines = max_lines
+        self._loading_symbol = loading_symbol
+
+        self._lines: list[str] = []
+        self._ops_shown = 0
+
+    def render(self, stats: pb.OperationStats) -> list[str]:
+        """Convert the stats into a list of lines to display.
+
+        Args:
+            stats: Collection of operations to display.
 
-        if self._ops_shown < total_operations:
+        Returns:
+            The lines of text to print. The lines do not end with the newline
+            character. Returns an empty list if there are no operations.
+        """
+        for op in stats.operations:
+            self._add_operation(op, is_subtask=False, indent="")
+
+        if self._ops_shown < stats.total_operations:
             if 1 <= self._max_lines <= len(self._lines):
+                self._ops_shown -= 1
                 self._lines.pop()
 
-            remaining = total_operations - self._ops_shown
+            remaining = stats.total_operations - self._ops_shown
 
             self._lines.append(f"+ {remaining} more task(s)")
 
-        if len(self._lines) == 0:
-            if self._loading_symbol:
-                self._text_area.set_text(f"{self._loading_symbol} {self._default_text}")
-            else:
-                self._text_area.set_text(self._default_text)
-        else:
-            self._text_area.set_text("\n".join(self._lines))
+        return self._lines
 
     def _add_operation(self, op: pb.Operation, is_subtask: bool, indent: str) -> None:
         """Add the operation to `self._lines`."""
@@ -200,14 +292,17 @@ class _DynamicOperationStatsPrinter:
         if not is_subtask:
             self._ops_shown += 1
 
-        parts = []
+        status_indent_level = 0  # alignment for the status message, if any
+        parts: list[str] = []
 
         # Subtask indicator.
         if is_subtask and self._printer.supports_unicode:
+            status_indent_level += 2  # +1 for space
             parts.append("↳")
 
         # Loading symbol.
         if self._loading_symbol:
+            status_indent_level += 2  # +1 for space
             parts.append(self._loading_symbol)
 
         # Task name.
@@ -225,14 +320,14 @@ class _DynamicOperationStatsPrinter:
         if op.error_status:
             error_word = self._printer.error("ERROR")
             error_desc = self._printer.secondary_text(op.error_status)
-            subtask_indent = "  " if is_subtask else ""
+            status_indent = " " * status_indent_level
             self._lines.append(
-                f"{indent}{subtask_indent}  {error_word} {error_desc}",
+                f"{indent}{status_indent}{error_word} {error_desc}",
             )
 
         # Subtasks.
         if op.subtasks:
-            subtask_indent = indent + "  "
+            subtask_indent = indent + _INDENT
             for task in op.subtasks:
                 self._add_operation(
                     task,

wandb/sdk/lib/retry.py

@@ -77,9 +77,10 @@ class Retry(Generic[_R]):
             self._retryable_exceptions = retryable_exceptions
         else:
             self._retryable_exceptions = (TransientError,)
-        self._index = 0
         self.retry_callback = retry_callback
 
+        self._num_iter = 0
+
     def _sleep_check_cancelled(
         self, wait_seconds: float, cancel_event: Optional[threading.Event]
     ) -> bool:
@@ -194,7 +195,7 @@ class Retry(Generic[_R]):
         else:
             wandb.termlog(
                 f"{self._error_prefix}"
-                f" ({exception.__class__.__name__}), entering retry loop."
+                + f" ({exception.__class__.__name__}), entering retry loop."
             )
 
     def _print_recovered(self, start_time: datetime.datetime) -> None:

wandb/sdk/lib/service/service_connection.py

@@ -190,8 +190,8 @@ class ServiceConnection:
         except TimeoutError:
             raise WandbAttachFailedError(
                 "Failed to attach because the run does not belong to"
-                " the current service process, or because the service"
-                " process is busy (unlikely)."
+                + " the current service process, or because the service"
+                + " process is busy (unlikely)."
             ) from None
 
         else:

wandb/sdk/lib/wb_logging.py

@@ -136,7 +136,7 @@ def add_file_handler(run_id: str, filepath: pathlib.Path) -> logging.Handler:
     return handler
 
 
-class _RunIDFilter(logging.Filter):
+class _RunIDFilter:
     """Filters out messages logged for a different run."""
 
     def __init__(self, run_id: str) -> None:
@@ -148,6 +148,7 @@ class _RunIDFilter(logging.Filter):
         self._run_id = run_id
 
     def filter(self, record: logging.LogRecord) -> bool:
+        """Modify a log record and return whether it matches the run."""
         run_id = _run_id.get()
 
         if run_id is None:

wandb/sdk/mailbox/mailbox.py

@@ -89,7 +89,7 @@ class Mailbox:
         def generate():
             return "".join(
                 secrets.choice(string.ascii_lowercase + string.digits)
-                for i in range(12)
+                for _ in range(12)
             )
 
         address = generate()

wandb/sdk/wandb_init.py

@@ -12,6 +12,7 @@ from __future__ import annotations
 
 import contextlib
 import dataclasses
+import functools
 import json
 import logging
 import os
@@ -988,25 +989,21 @@ class _WandbInit:
 
         run_init_handle = backend.interface.deliver_run(run)
 
-        async def display_init_message() -> None:
-            assert backend.interface
-
+        try:
             with progress.progress_printer(
                 run_printer,
                 default_text="Waiting for wandb.init()...",
             ) as progress_printer:
-                await progress.loop_printing_operation_stats(
-                    progress_printer,
-                    backend.interface,
+                result = wait_with_progress(
+                    run_init_handle,
+                    timeout=timeout,
+                    display_progress=functools.partial(
+                        progress.loop_printing_operation_stats,
+                        progress_printer,
+                        backend.interface,
+                    ),
                 )
 
-        try:
-            result = wait_with_progress(
-                run_init_handle,
-                timeout=timeout,
-                display_progress=display_init_message,
-            )
-
         except TimeoutError:
             run_init_handle.cancel(backend.interface)