wandb 0.18.3__py3-none-any.whl → 0.18.5__py3-none-any.whl

wandb/sdk/lib/progress.py

@@ -0,0 +1,274 @@
+"""Defines an object for printing run progress at the end of a script."""
+
+from __future__ import annotations
+
+import contextlib
+from typing import Iterable, Iterator
+
+import wandb
+from wandb.proto import wandb_internal_pb2 as pb
+
+from . import printer as p
+
+
+def print_sync_dedupe_stats(
+    printer: p.Printer,
+    final_result: pb.PollExitResponse,
+) -> None:
+    """Print how much W&B sync reduced the amount of uploaded data.
+
+    Args:
+        final_result: The final PollExit result.
+    """
+    deduped_bytes = final_result.pusher_stats.deduped_bytes
+    total_bytes = final_result.pusher_stats.total_bytes
+
+    if total_bytes <= 0 or deduped_bytes <= 0:
+        return
+
+    frac = deduped_bytes / total_bytes
+    printer.display(f"W&B sync reduced upload amount by {frac:.1%}")
+
+
+@contextlib.contextmanager
+def progress_printer(
+    printer: p.Printer,
+    settings: wandb.Settings | None = None,
+) -> Iterator[ProgressPrinter]:
+    """Context manager providing an object for printing run progress."""
+    with printer.dynamic_text() as text_area:
+        yield ProgressPrinter(printer, text_area, settings)
+        printer.progress_close()
+
+
+class ProgressPrinter:
+    """Displays PollExitResponse results to the user."""
+
+    def __init__(
+        self,
+        printer: p.Printer,
+        progress_text_area: p.DynamicText | None,
+        settings: wandb.Settings | None,
+    ) -> None:
+        self._show_operation_stats = settings and settings._show_operation_stats
+        self._printer = printer
+        self._progress_text_area = progress_text_area
+        self._tick = 0
+        self._last_printed_line = ""
+
+    def update(
+        self,
+        progress: list[pb.PollExitResponse],
+    ) -> None:
+        """Update the displayed information."""
+        if not progress:
+            return
+
+        if 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)
+
+        self._tick += 1
+
+    def _update_operation_stats(self, stats_list: list[pb.OperationStats]) -> None:
+        if self._progress_text_area:
+            _DynamicOperationStatsPrinter(
+                self._printer,
+                self._progress_text_area,
+                max_lines=6,
+                loading_symbol=self._printer.loading_symbol(self._tick),
+            ).display(stats_list)
+
+        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 != 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 self._progress_text_area:
+            self._progress_text_area.set_text(text)
+        else:
+            self._printer.progress_update(text + "\r", progress)
+
+
+class _DynamicOperationStatsPrinter:
+    """Single-use object that writes operation stats into a text area."""
+
+    def __init__(
+        self,
+        printer: p.Printer,
+        text_area: p.DynamicText,
+        max_lines: int,
+        loading_symbol: str,
+    ) -> None:
+        self._printer = printer
+        self._text_area = text_area
+        self._max_lines = max_lines
+        self._loading_symbol = loading_symbol
+
+        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
+
+        if self._ops_shown < total_operations:
+            # NOTE: In Python 3.8, we'd use a chained comparison here.
+            if 1 <= self._max_lines and self._max_lines <= len(self._lines):
+                self._lines.pop()
+
+            remaining = 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} Finishing up...")
+            else:
+                self._text_area.set_text("Finishing up...")
+        else:
+            self._text_area.set_text("\n".join(self._lines))
+
+    def _add_operation(self, op: pb.Operation, is_subtask: bool, indent: str) -> None:
+        """Add the operation to `self._lines`."""
+        if len(self._lines) >= self._max_lines:
+            return
+
+        if not is_subtask:
+            self._ops_shown += 1
+
+        parts = []
+
+        # Subtask indicator.
+        if is_subtask and self._printer.supports_unicode:
+            parts.append("↳")
+
+        # Loading symbol.
+        if self._loading_symbol:
+            parts.append(self._loading_symbol)
+
+        # Task name.
+        parts.append(op.desc)
+
+        # Progress information.
+        if op.progress:
+            parts.append(f"{op.progress}")
+
+        # Task duration.
+        parts.append(f"({_time_to_string(seconds=op.runtime_seconds)})")
+
+        # Error status.
+        self._lines.append(indent + " ".join(parts))
+        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 ""
+            self._lines.append(
+                f"{indent}{subtask_indent}  {error_word} {error_desc}",
+            )
+
+        # Subtasks.
+        if op.subtasks:
+            subtask_indent = indent + "  "
+            for task in op.subtasks:
+                self._add_operation(
+                    task,
+                    is_subtask=True,
+                    indent=subtask_indent,
+                )
+
+
+def _time_to_string(seconds: float) -> str:
+    """Returns a short string representing the duration."""
+    if seconds < 10:
+        return f"{seconds:.1f}s"
+    if seconds < 60:
+        return f"{seconds:.0f}s"
+    if seconds < 60 * 60:
+        minutes = seconds / 60
+        return f"{minutes:.1f}m"
+
+    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/service/streams.py

@@ -20,6 +20,7 @@ import wandb
 import wandb.util
 from wandb.proto import wandb_internal_pb2 as pb
 from wandb.sdk.internal.settings_static import SettingsStatic
+from wandb.sdk.lib import progress
 from wandb.sdk.lib.mailbox import (
     Mailbox,
     MailboxProbe,
@@ -257,8 +258,12 @@ class StreamMux:
     def _on_progress_exit(self, progress_handle: MailboxProgress) -> None:
         pass
 
-    def _on_progress_exit_all(self, progress_all_handle: MailboxProgressAll) -> None:
-        probe_handles = []
+    def _on_progress_exit_all(
+        self,
+        progress_printer: progress.ProgressPrinter,
+        progress_all_handle: MailboxProgressAll,
+    ) -> None:
+        probe_handles: list[MailboxProbe] = []
         progress_handles = progress_all_handle.get_progress_handles()
         for progress_handle in progress_handles:
             probe_handles.extend(progress_handle.get_probe_handles())
@@ -268,13 +273,13 @@ class StreamMux:
         if self._check_orphaned():
             self._stopped.set()
 
-        poll_exit_responses: List[Optional[pb.PollExitResponse]] = []
+        poll_exit_responses: List[pb.PollExitResponse] = []
         for probe_handle in probe_handles:
             result = probe_handle.get_probe_result()
             if result:
                 poll_exit_responses.append(result.response.poll_exit_response)
 
-        Run._footer_file_pusher_status_info(poll_exit_responses, printer=self._printer)
+        progress_printer.update(poll_exit_responses)
 
     def _finish_all(self, streams: Dict[str, StreamRecord], exit_code: int) -> None:
         if not streams:
@@ -283,7 +288,6 @@ class StreamMux:
         printer = get_printer(
             all(stream._settings._jupyter for stream in streams.values())
         )
-        self._printer = printer
 
         # fixme: for now we have a single printer for all streams,
         # and jupyter is disabled if at least single stream's setting set `_jupyter` to false
@@ -307,12 +311,18 @@ class StreamMux:
             #     exit_code, settings=stream._settings, printer=printer  # type: ignore
             # )
 
-        # todo: should we wait for the max timeout (?) of all exit handles or just wait forever?
-        # timeout = max(stream._settings._exit_timeout for stream in streams.values())
-        got_result = self._mailbox.wait_all(
-            handles=exit_handles, timeout=-1, on_progress_all=self._on_progress_exit_all
-        )
-        assert got_result
+        with progress.progress_printer(printer) as progress_printer:
+            # todo: should we wait for the max timeout (?) of all exit handles or just wait forever?
+            # timeout = max(stream._settings._exit_timeout for stream in streams.values())
+            got_result = self._mailbox.wait_all(
+                handles=exit_handles,
+                timeout=-1,
+                on_progress_all=functools.partial(
+                    self._on_progress_exit_all,
+                    progress_printer,
+                ),
+            )
+            assert got_result
 
         # These could be done in parallel in the future
         for _sid, stream in started_streams.items():

wandb/sdk/wandb_init.py

@@ -8,6 +8,8 @@ For more on using `wandb.init()`, including code snippets, check out our
 [guide and FAQs](https://docs.wandb.ai/guides/track/launch).
 """
 
+from __future__ import annotations
+
 import copy
 import json
 import logging
@@ -16,7 +18,7 @@ import platform
 import sys
 import tempfile
 import time
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Union
+from typing import TYPE_CHECKING, Any, Sequence
 
 import wandb
 import wandb.env
@@ -43,7 +45,7 @@ from .wandb_settings import Settings, Source
 if TYPE_CHECKING:
     from wandb.proto import wandb_internal_pb2 as pb
 
-logger: Optional[logging.Logger] = None  # logger configured during wandb.init()
+logger: logging.Logger | None = None  # logger configured during wandb.init()
 
 
 def _set_logger(log_object: logging.Logger) -> None:
@@ -52,7 +54,7 @@ def _set_logger(log_object: logging.Logger) -> None:
     logger = log_object
 
 
-def _huggingface_version() -> Optional[str]:
+def _huggingface_version() -> str | None:
     if "transformers" in sys.modules:
         trans = wandb.util.get_module("transformers")
         if hasattr(trans, "__version__"):
@@ -74,8 +76,8 @@ def _maybe_mp_process(backend: Backend) -> bool:
     return False
 
 
-def _handle_launch_config(settings: "Settings") -> Dict[str, Any]:
-    launch_run_config: Dict[str, Any] = {}
+def _handle_launch_config(settings: Settings) -> dict[str, Any]:
+    launch_run_config: dict[str, Any] = {}
     if not settings.launch:
         return launch_run_config
     if os.environ.get("WANDB_CONFIG") is not None:
@@ -112,22 +114,22 @@ class _WandbInit:
 
     def __init__(self) -> None:
         self.kwargs = None
-        self.settings: Optional[Settings] = None
-        self.sweep_config: Dict[str, Any] = {}
-        self.launch_config: Dict[str, Any] = {}
-        self.config: Dict[str, Any] = {}
-        self.run: Optional[Run] = None
-        self.backend: Optional[Backend] = None
-
-        self._teardown_hooks: List[TeardownHook] = []
-        self._wl: Optional[wandb_setup._WandbSetup] = None
-        self._reporter: Optional[wandb.sdk.lib.reporting.Reporter] = None
-        self.notebook: Optional[wandb.jupyter.Notebook] = None  # type: ignore
-        self.printer: Optional[Printer] = None
+        self.setting: Settings | None = None
+        self.sweep_config: dict[str, Any] = {}
+        self.launch_config: dict[str, Any] = {}
+        self.config: dict[str, Any] = {}
+        self.run: Run | None = None
+        self.backend: Backend | None = None
+
+        self._teardown_hooks: list[TeardownHook] = []
+        self._wl: wandb_setup._WandbSetup | None = None
+        self._reporter: wandb.sdk.lib.reporting.Reporter | None = None
+        self.notebook: wandb.jupyter.Notebook | None = None  # type: ignore
+        self.printer: Printer | None = None
 
         self._init_telemetry_obj = telemetry.TelemetryRecord()
 
-        self.deprecated_features_used: Dict[str, str] = dict()
+        self.deprecated_features_used: dict[str, str] = dict()
 
     def _setup_printer(self, settings: Settings) -> None:
         if self.printer:
@@ -203,7 +205,7 @@ class _WandbInit:
         self._setup_printer(settings)
         self._reporter = reporting.setup_reporter(settings=settings)
 
-        sagemaker_config: Dict = (
+        sagemaker_config: dict = (
             dict() if settings.sagemaker_disable else sagemaker.parse_sm_config()
         )
         if sagemaker_config:
@@ -254,7 +256,7 @@ class _WandbInit:
         self.sweep_config = dict()
         sweep_config = self._wl._sweep_config or dict()
         self.config = dict()
-        self.init_artifact_config: Dict[str, Any] = dict()
+        self.init_artifact_config: dict[str, Any] = dict()
         for config_data in (
             sagemaker_config,
             self._wl._config,
@@ -368,7 +370,7 @@ class _WandbInit:
             else:
                 config_target.setdefault(k, v)
 
-    def _enable_logging(self, log_fname: str, run_id: Optional[str] = None) -> None:
+    def _enable_logging(self, log_fname: str, run_id: str | None = None) -> None:
         """Enable logging to the global debug log.
 
         This adds a run_id to the log, in case of multiple processes on the same machine.
@@ -602,6 +604,8 @@ class _WandbInit:
             define_metric=drun.define_metric,
             plot_table=drun.plot_table,
             alert=drun.alert,
+            watch=drun.watch,
+            unwatch=drun.unwatch,
         )
         return drun
 
@@ -719,7 +723,7 @@ class _WandbInit:
                 setattr(tel.imports_init, module_name, True)
 
             # probe the active start method
-            active_start_method: Optional[str] = None
+            active_start_method: str | None = None
             if self.settings.start_method == "thread":
                 active_start_method = self.settings.start_method
             else:
@@ -794,7 +798,7 @@ class _WandbInit:
         if not self.settings.disable_git:
             run._populate_git_info()
 
-        run_result: Optional[pb.RunUpdateResult] = None
+        run_result: pb.RunUpdateResult | None = None
 
         if self.settings._offline:
             with telemetry.context(run=run) as tel:
@@ -805,7 +809,7 @@ class _WandbInit:
                     "`resume` will be ignored since W&B syncing is set to `offline`. "
                     f"Starting a new run with run id {run.id}."
                 )
-        error: Optional[wandb.errors.Error] = None
+        error: wandb.Error | None = None
 
         timeout = self.settings.init_timeout
 
@@ -909,11 +913,11 @@ class _WandbInit:
 
 
 def _attach(
-    attach_id: Optional[str] = None,
-    run_id: Optional[str] = None,
+    attach_id: str | None = None,
+    run_id: str | None = None,
     *,
-    run: Optional["Run"] = None,
-) -> Optional[Run]:
+    run: Run | None = None,
+) -> Run | None:
     """Attach to a run currently executing in another process/thread.
 
     Arguments:
@@ -995,32 +999,32 @@ def _attach(
 
 
 def init(
-    job_type: Optional[str] = None,
-    dir: Optional[StrPath] = None,
-    config: Union[Dict, str, None] = None,
-    project: Optional[str] = None,
-    entity: Optional[str] = None,
-    reinit: Optional[bool] = None,
-    tags: Optional[Sequence] = None,
-    group: Optional[str] = None,
-    name: Optional[str] = None,
-    notes: Optional[str] = None,
-    magic: Optional[Union[dict, str, bool]] = None,
-    config_exclude_keys: Optional[List[str]] = None,
-    config_include_keys: Optional[List[str]] = None,
-    anonymous: Optional[str] = None,
-    mode: Optional[str] = None,
-    allow_val_change: Optional[bool] = None,
-    resume: Optional[Union[bool, str]] = None,
-    force: Optional[bool] = None,
-    tensorboard: Optional[bool] = None,  # alias for sync_tensorboard
-    sync_tensorboard: Optional[bool] = None,
-    monitor_gym: Optional[bool] = None,
-    save_code: Optional[bool] = None,
-    id: Optional[str] = None,
-    fork_from: Optional[str] = None,
-    resume_from: Optional[str] = None,
-    settings: Union[Settings, Dict[str, Any], None] = None,
+    job_type: str | None = None,
+    dir: StrPath | None = None,
+    config: dict | str | None = None,
+    project: str | None = None,
+    entity: str | None = None,
+    reinit: bool | None = None,
+    tags: Sequence | None = None,
+    group: str | None = None,
+    name: str | None = None,
+    notes: str | None = None,
+    magic: dict | str | bool | None = None,
+    config_exclude_keys: list[str] | None = None,
+    config_include_keys: list[str] | None = None,
+    anonymous: str | None = None,
+    mode: str | None = None,
+    allow_val_change: bool | None = None,
+    resume: bool | str | None = None,
+    force: bool | None = None,
+    tensorboard: bool | None = None,  # alias for sync_tensorboard
+    sync_tensorboard: bool | None = None,
+    monitor_gym: bool | None = None,
+    save_code: bool | None = None,
+    id: str | None = None,
+    fork_from: str | None = None,
+    resume_from: str | None = None,
+    settings: Settings | dict[str, Any] | None = None,
 ) -> Run:
     r"""Start a new run to track and log to W&B.