wandb 0.22.1__py3-none-musllinux_1_2_aarch64.whl → 0.22.3__py3-none-musllinux_1_2_aarch64.whl

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,
@@ -93,168 +97,192 @@ 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
-        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])
-        elif self._show_operation_stats:
-            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)
+        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)"
-
-            if line and line != self._last_printed_line:
-                self._printer.display(line)
-                self._last_printed_line = line
+            return _OperationStatsPrinter(
+                self._printer,
+                _MAX_LINES_TO_PRINT,
+                loading_symbol,
+            ).render(stats_or_groups)
 
-    def _update_single_run(
+    def _to_static_text(
         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,
+        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:
-            self._update_progress_text(line, 1.0)
+            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)"
+
+        return line
+
+
+class _GroupedOperationStatsPrinter:
+    """Renders a list of labeled operation stats groups into lines of text."""
 
-    def _update_multiple_runs(
+    def __init__(
         self,
-        progress_list: list[pb.PollExitResponse],
+        printer: p.Printer,
+        max_lines: int,
+        loading_symbol: str,
     ) -> 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)
+        self._printer = printer
+        self._max_lines = max_lines
+        self._loading_symbol = loading_symbol
 
-    def _update_progress_text(self, text: str, progress: float) -> None:
-        if text == self._last_printed_line:
-            return
-        self._last_printed_line = text
+    def render(self, groups: dict[str, pb.OperationStats]) -> list[str]:
+        """Convert labeled operation stats groups into text to display.
 
-        if self._progress_text_area:
-            self._progress_text_area.set_text(text)
-        else:
-            self._printer.progress_update(text + "\r", progress)
+        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)
 
+            # 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}")
 
-class _DynamicOperationStatsPrinter:
-    """Single-use object that writes operation stats into a text area."""
+        return lines
+
+
+class _OperationStatsPrinter:
+    """Renders operation stats 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 display(
-        self,
-        stats_list: Iterable[pb.OperationStats],
-    ) -> 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
+    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`."""
@@ -264,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.
@@ -289,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,
@@ -318,8 +349,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/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
@@ -727,7 +728,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,
@@ -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)
 

wandb/sdk/wandb_run.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import asyncio
 import functools
 import glob
 import json
@@ -33,7 +32,6 @@ from wandb.errors import CommError, UsageError
 from wandb.errors.links import url_registry
 from wandb.integration.torch import wandb_torch
 from wandb.plot import CustomChart, Visualize
-from wandb.proto import wandb_internal_pb2 as pb
 from wandb.proto.wandb_deprecated import Deprecated
 from wandb.proto.wandb_internal_pb2 import (
     MetricRecord,
@@ -44,7 +42,7 @@ from wandb.proto.wandb_internal_pb2 import (
 from wandb.sdk.artifacts._internal_artifact import InternalArtifact
 from wandb.sdk.artifacts.artifact import Artifact
 from wandb.sdk.internal import job_builder
-from wandb.sdk.lib import asyncio_compat, wb_logging
+from wandb.sdk.lib import wb_logging
 from wandb.sdk.lib.import_hooks import (
     register_post_import_hook,
     unregister_post_import_hook,
@@ -489,7 +487,7 @@ class Run:
 
     You can log data to a run with `wandb.Run.log()`. Anything you log using
     `wandb.Run.log()` is sent to that run. See
-    [Create an experiment](https://docs.wandb.ai/guides/track/launch) or
+    [Create an experiment](https://docs.wandb.ai/guides/track/create-an-experiment/) or
     [`wandb.init`](https://docs.wandb.ai/ref/python/init/) API reference page
     or more information.
 
@@ -1574,15 +1572,13 @@ class Run:
 
     @_log_to_run
     def _console_callback(self, name: str, data: str) -> None:
-        # logger.info("console callback: %s, %s", name, data)
         if self._backend and self._backend.interface:
-            self._backend.interface.publish_output(name, data)
+            # nowait=True so that this can be called from an asyncio context.
+            self._backend.interface.publish_output(name, data, nowait=True)
 
     @_log_to_run
     @_raise_if_finished
     def _console_raw_callback(self, name: str, data: str) -> None:
-        # logger.info("console callback: %s, %s", name, data)
-
         # NOTE: console output is only allowed on the process which installed the callback
         # this will prevent potential corruption in the socket to the service.  Other methods
         # are protected by the _attach run decorator, but this callback was installed on the
@@ -1592,7 +1588,8 @@ class Run:
             return
 
         if self._backend and self._backend.interface:
-            self._backend.interface.publish_output_raw(name, data)
+            # nowait=True so that this can be called from an asyncio context.
+            self._backend.interface.publish_output_raw(name, data, nowait=True)
 
     @_log_to_run
     def _tensorboard_callback(
@@ -1861,7 +1858,7 @@ class Run:
             run.log({"accuracy": 0.8}, step=current_step)
             current_step += 1
             run.log({"train-loss": 0.4}, step=current_step)
-            run.log({"accuracy": 0.9}, step=current_step)
+            run.log({"accuracy": 0.9}, step=current_step, commit=True)
         ```
 
         Args:
@@ -2051,6 +2048,9 @@ class Run:
         When given an absolute path or glob and no `base_path`, one
         directory level is preserved as in the example above.
 
+        Files are automatically deduplicated: calling `save()` multiple times
+        on the same file without modifications will not re-upload it.
+
         Args:
             glob_str: A relative or absolute path or Unix glob.
             base_path: A path to use to infer a directory structure; see examples.
@@ -2075,10 +2075,10 @@ class Run:
         run.save("these/are/myfiles/*", base_path="these")
         # => Saves files in an "are/myfiles/" folder in the run.
 
-        run.save("/User/username/Documents/run123/*.txt")
+        run.save("/Users/username/Documents/run123/*.txt")
         # => Saves files in a "run123/" folder in the run. See note below.
 
-        run.save("/User/username/Documents/run123/*.txt", base_path="/User")
+        run.save("/Users/username/Documents/run123/*.txt", base_path="/Users")
         # => Saves files in a "username/Documents/run123/" folder in the run.
 
         run.save("files/*/saveme.txt")
@@ -2679,41 +2679,6 @@ class Run:
         else:
             return artifact
 
-    async def _display_finish_stats(
-        self,
-        progress_printer: progress.ProgressPrinter,
-    ) -> None:
-        last_result: Result | None = None
-
-        async def loop_update_printer() -> None:
-            while True:
-                if last_result:
-                    progress_printer.update(
-                        [last_result.response.poll_exit_response],
-                    )
-                await asyncio.sleep(0.1)
-
-        async def loop_poll_exit() -> None:
-            nonlocal last_result
-            assert self._backend and self._backend.interface
-
-            while True:
-                handle = await self._backend.interface.deliver_async(
-                    pb.Record(request=pb.Request(poll_exit=pb.PollExitRequest()))
-                )
-
-                time_start = time.monotonic()
-                last_result = await handle.wait_async(timeout=None)
-
-                # Update at most once a second.
-                time_elapsed = time.monotonic() - time_start
-                if time_elapsed < 1:
-                    await asyncio.sleep(1 - time_elapsed)
-
-        async with asyncio_compat.open_task_group() as task_group:
-            task_group.start_soon(loop_update_printer())
-            task_group.start_soon(loop_poll_exit())
-
     def _on_finish(self) -> None:
         trigger.call("on_finished")
 
@@ -2738,8 +2703,9 @@ class Run:
                 exit_handle,
                 timeout=None,
                 display_progress=functools.partial(
-                    self._display_finish_stats,
+                    progress.loop_printing_operation_stats,
                     progress_printer,
+                    self._backend.interface,
                 ),
             )