wandb 0.19.6rc4__py3-none-macosx_11_0_arm64.whl → 0.19.8__py3-none-macosx_11_0_arm64.whl

wandb/sdk/wandb_metadata.py

@@ -236,39 +236,6 @@ class Metadata(BaseModel, validate_assignment=True):
 
     NOTE: Definitions must be kept in sync with wandb_internal.proto::MetadataRequest.
 
-    Attributes:
-        os: Operating system.
-        python: Python version.
-        heartbeat_at: Timestamp of last heartbeat.
-        started_at: Timestamp of run start.
-        docker: Docker image.
-        cuda: CUDA version.
-        args: Command-line arguments.
-        state: Run state.
-        program: Program name.
-        code_path: Path to code.
-        git: Git repository information.
-        email: Email address.
-        root: Root directory.
-        host: Host name.
-        username: Username.
-        executable: Python executable path.
-        code_path_local: Local code path.
-        colab: Colab URL.
-        cpu_count: CPU count.
-        cpu_count_logical: Logical CPU count.
-        gpu_type: GPU type.
-        disk: Disk information.
-        memory: Memory information.
-        cpu: CPU information.
-        apple: Apple silicon information.
-        gpu_nvidia: NVIDIA GPU information.
-        gpu_amd: AMD GPU information.
-        slurm: Slurm environment information.
-        cuda_version: CUDA version.
-        trainium: Trainium information.
-        tpu: TPU information.
-
     Examples:
         Update Run metadata:
 
@@ -296,44 +263,108 @@ class Metadata(BaseModel, validate_assignment=True):
         ```
     """
 
-    # TODO: Pydantic configuration.
     model_config = ConfigDict(
         extra="ignore",  # ignore extra fields
         validate_default=True,  # validate default values
+        use_attribute_docstrings=True,  # for field descriptions
+        revalidate_instances="always",
     )
 
     os: str | None = None
+    """Operating system."""
+
     python: str | None = None
+    """Python version."""
+
     heartbeat_at: datetime | None = Field(default=None, alias="heartbeatAt")
+    """Timestamp of last heartbeat."""
+
     started_at: datetime | None = Field(default=None, alias="startedAt")
+    """Timestamp of run start."""
+
     docker: str | None = None
+    """Docker image."""
+
     cuda: str | None = None
+    """CUDA version."""
+
     args: list[str] = Field(default_factory=list)
+    """Command-line arguments."""
+
     state: str | None = None
+    """Run state."""
+
     program: str | None = None
+    """Program name."""
+
     code_path: str | None = Field(default=None, alias="codePath")
+    """Path to code."""
+
     git: GitRepoRecord | None = None
+    """Git repository information."""
+
     email: str | None = None
+    """Email address."""
+
     root: str | None = None
+    """Root directory."""
+
     host: str | None = None
+    """Host name."""
+
     username: str | None = None
+    """Username."""
+
     executable: str | None = None
+    """Python executable path."""
+
     code_path_local: str | None = Field(default=None, alias="codePathLocal")
+    """Local code path."""
+
     colab: str | None = None
+    """Colab URL."""
+
     cpu_count: int | None = Field(default=None, alias="cpuCount")
+    """CPU count."""
+
     cpu_count_logical: int | None = Field(default=None, alias="cpuCountLogical")
+    """Logical CPU count."""
+
     gpu_type: str | None = Field(default=None, alias="gpuType")
+    """GPU type."""
+
     gpu_count: int | None = Field(default=None, alias="gpuCount")
+    """GPU count."""
+
     disk: dict[str, DiskInfo] = Field(default_factory=dict)
+    """Disk information."""
+
     memory: MemoryInfo | None = None
+    """Memory information."""
+
     cpu: CpuInfo | None = None
+    """CPU information."""
+
     apple: AppleInfo | None = None
+    """Apple silicon information."""
+
     gpu_nvidia: list[GpuNvidiaInfo] = Field(default_factory=list, alias="gpuNvidia")
+    """NVIDIA GPU information."""
+
     gpu_amd: list[GpuAmdInfo] = Field(default_factory=list, alias="gpuAmd")
+    """AMD GPU information."""
+
     slurm: dict[str, str] = Field(default_factory=dict)
+    """Slurm environment information."""
+
     cuda_version: str | None = Field(default=None, alias="cudaVersion")
+    """CUDA version."""
+
     trainium: TrainiumInfo | None = None
+    """Trainium information."""
+
     tpu: TPUInfo | None = None
+    """TPU information."""
 
     def __init__(self, **data):
         super().__init__(**data)

wandb/sdk/wandb_require.py

@@ -9,8 +9,10 @@ Example:
     wandb.require("incremental-artifacts@beta")
 """
 
+from __future__ import annotations
+
 import os
-from typing import Optional, Sequence, Union
+from typing import Iterable
 
 import wandb
 from wandb.env import _REQUIRE_LEGACY_SERVICE
@@ -21,9 +23,9 @@ from wandb.sdk import wandb_run
 class _Requires:
     """Internal feature class."""
 
-    _features: Sequence[str]
+    _features: tuple[str, ...]
 
-    def __init__(self, features: Union[str, Sequence[str]]) -> None:
+    def __init__(self, features: str | Iterable[str]) -> None:
         self._features = (
             tuple([features]) if isinstance(features, str) else tuple(features)
         )
@@ -67,17 +69,21 @@ class _Requires:
 
 
 def require(
-    requirement: Optional[Union[str, Sequence[str]]] = None,
-    experiment: Optional[Union[str, Sequence[str]]] = None,
+    requirement: str | Iterable[str] | None = None,
+    experiment: str | Iterable[str] | None = None,
 ) -> None:
     """Indicate which experimental features are used by the script.
 
+    This should be called before any other `wandb` functions, ideally right
+    after importing `wandb`.
+
     Args:
-        requirement: (str or list) Features to require
-        experiment: (str or list) Features to require
+        requirement: The name of a feature to require or an iterable of
+            feature names.
+        experiment: An alias for `requirement`.
 
     Raises:
-        wandb.errors.UnsupportedError: if not supported
+        wandb.errors.UnsupportedError: If a feature name is unknown.
     """
     features = requirement or experiment
     if not features:

wandb/sdk/wandb_run.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import atexit
 import functools
 import glob
@@ -43,6 +44,7 @@ from wandb.proto.wandb_internal_pb2 import (
 )
 from wandb.sdk.artifacts.artifact import Artifact
 from wandb.sdk.internal import job_builder
+from wandb.sdk.lib import asyncio_compat
 from wandb.sdk.lib.import_hooks import (
     register_post_import_hook,
     unregister_post_import_hook,
@@ -82,7 +84,12 @@ from .lib import (
     telemetry,
 )
 from .lib.exit_hooks import ExitHooks
-from .lib.mailbox import MailboxError, MailboxHandle, MailboxProbe, MailboxProgress
+from .mailbox import (
+    HandleAbandonedError,
+    MailboxClosedError,
+    MailboxHandle,
+    wait_with_progress,
+)
 from .wandb_alerts import AlertLevel
 from .wandb_metadata import Metadata
 from .wandb_settings import Settings
@@ -150,11 +157,11 @@ class RunStatusChecker:
     """
 
     _stop_status_lock: threading.Lock
-    _stop_status_handle: MailboxHandle | None
+    _stop_status_handle: MailboxHandle[Result] | None
     _network_status_lock: threading.Lock
-    _network_status_handle: MailboxHandle | None
+    _network_status_handle: MailboxHandle[Result] | None
     _internal_messages_lock: threading.Lock
-    _internal_messages_handle: MailboxHandle | None
+    _internal_messages_handle: MailboxHandle[Result] | None
 
     def __init__(
         self,
@@ -202,7 +209,7 @@ class RunStatusChecker:
     @staticmethod
     def _abandon_status_check(
         lock: threading.Lock,
-        handle: MailboxHandle | None,
+        handle: MailboxHandle[Result] | None,
     ):
         with lock:
             if handle:
@@ -217,35 +224,36 @@ class RunStatusChecker:
         request: Any,
         process: Any,
     ) -> None:
-        local_handle: MailboxHandle | None = None
+        local_handle: MailboxHandle[Result] | None = None
         join_requested = False
         while not join_requested:
             time_probe = time.monotonic()
             if not local_handle:
-                local_handle = request()
+                try:
+                    local_handle = request()
+                except MailboxClosedError:
+                    # This can happen if the service process dies.
+                    break
             assert local_handle
 
             with lock:
                 if self._join_event.is_set():
                     break
                 set_handle(local_handle)
+
             try:
-                result = local_handle.wait(timeout=timeout, release=False)
-            except MailboxError:
-                # background threads are oportunistically getting results
-                # from the internal process but the internal process could
-                # be shutdown at any time.  In this case assume that the
-                # thread should exit silently.   This is possible
-                # because we do not have an atexit handler for the user
-                # process which quiesces active threads.
+                result = local_handle.wait_or(timeout=timeout)
+            except HandleAbandonedError:
+                # This can happen if the service process dies.
                 break
+            except TimeoutError:
+                result = None
+
             with lock:
                 set_handle(None)
 
             if result:
                 process(result)
-                # if request finished, clear the handle to send on the next interval
-                local_handle.abandon()
                 local_handle = None
 
             time_elapsed = time.monotonic() - time_probe
@@ -534,7 +542,7 @@ class Run:
 
     _sampled_history: SampledHistoryResponse | None
     _final_summary: GetSummaryResponse | None
-    _poll_exit_handle: MailboxHandle | None
+    _poll_exit_handle: MailboxHandle[Result] | None
     _poll_exit_response: PollExitResponse | None
     _internal_messages_response: InternalMessagesResponse | None
 
@@ -1300,12 +1308,6 @@ class Run:
         if self._backend and self._backend.interface:
             self._backend.interface.publish_summary(self, summary_record)
 
-    def _on_progress_get_summary(self, handle: MailboxProgress) -> None:
-        pass
-        # TODO(jhr): enable printing for get_summary in later mailbox dev phase
-        # line = "Waiting for run.summary data..."
-        # self._printer.display(line)
-
     def _summary_get_current_summary_callback(self) -> dict[str, Any]:
         if self._is_finished:
             # TODO: WB-18420: fetch summary from backend and stage it before run is finished
@@ -1314,12 +1316,12 @@ class Run:
         if not self._backend or not self._backend.interface:
             return {}
         handle = self._backend.interface.deliver_get_summary()
-        result = handle.wait(
-            timeout=self._settings.summary_timeout,
-            on_progress=self._on_progress_get_summary,
-        )
-        if not result:
+
+        try:
+            result = handle.wait_or(timeout=self._settings.summary_timeout)
+        except TimeoutError:
             return {}
+
         get_summary_response = result.response.get_summary_response
         return proto_util.dict_from_proto_list(get_summary_response.item)
 
@@ -1330,7 +1332,7 @@ class Run:
     def _datatypes_callback(self, fname: str) -> None:
         if not self._backend or not self._backend.interface:
             return
-        files: FilesDict = dict(files=[(GlobStr(glob.escape(fname)), "now")])
+        files: FilesDict = dict(files=[(GlobStr(fname), "now")])
         self._backend.interface.publish_files(files)
 
     def _pop_all_charts(
@@ -1503,7 +1505,7 @@ class Run:
         if run_obj.notes:
             self._settings.run_notes = run_obj.notes
         if run_obj.tags:
-            self._settings.run_tags = run_obj.tags
+            self._settings.run_tags = tuple(run_obj.tags)
         if run_obj.sweep_id:
             self._settings.sweep_id = run_obj.sweep_id
         if run_obj.host:
@@ -1643,7 +1645,7 @@ class Run:
         [guides to logging](https://docs.wandb.ai/guides/track/log) for examples,
         from 3D molecular structures and segmentation masks to PR curves and histograms.
         You can use `wandb.Table` to log structured data. See our
-        [guide to logging tables](https://docs.wandb.ai/guides/tables/tables-walkthrough)
+        [guide to logging tables](https://docs.wandb.ai/guides/models/tables/tables-walkthrough)
         for details.
 
         The W&B UI organizes metrics with a forward slash (`/`) in their name
@@ -2111,15 +2113,6 @@ class Run:
         with telemetry.context(run=self) as tel:
             tel.feature.finish = True
 
-        # Pop this run (hopefully) from the run stack, to support the "reinit"
-        # functionality of wandb.init().
-        #
-        # TODO: It's not clear how _global_run_stack could have length other
-        # than 1 at this point in the code. If you're reading this, consider
-        # refactoring this thing.
-        if self._wl and len(self._wl._global_run_stack) > 0:
-            self._wl._global_run_stack.pop()
-
         # Run hooks that need to happen before the last messages to the
         # internal service, like Jupyter hooks.
         for hook in self._teardown_hooks:
@@ -2175,8 +2168,7 @@ class Run:
             return RunStatus()
 
         handle_run_status = self._backend.interface.deliver_request_run_status()
-        result = handle_run_status.wait(timeout=-1)
-        assert result
+        result = handle_run_status.wait_or(timeout=None)
         sync_data = result.response.run_status_response
 
         sync_time = None
@@ -2316,8 +2308,6 @@ class Run:
             raise ValueError("unhandled console")
         try:
             # save stdout and stderr before installing new write functions
-            out_redir.save()
-            err_redir.save()
             out_redir.install()
             err_redir.install()
             self._out_redir = out_redir
@@ -2563,31 +2553,38 @@ class Run:
         else:
             return artifact
 
-    def _on_probe_exit(self, probe_handle: MailboxProbe) -> None:
-        handle = probe_handle.get_mailbox_handle()
-        if handle:
-            result = handle.wait(timeout=0, release=False)
-            if not result:
-                return
-            probe_handle.set_probe_result(result)
-        assert self._backend and self._backend.interface
-        handle = self._backend.interface.deliver_poll_exit()
-        probe_handle.set_mailbox_handle(handle)
-
-    def _on_progress_exit(
+    async def _display_finish_stats(
         self,
         progress_printer: progress.ProgressPrinter,
-        progress_handle: MailboxProgress,
     ) -> None:
-        probe_handles = progress_handle.get_probe_handles()
-        if not probe_handles or len(probe_handles) != 1:
-            return
+        last_result: Result | None = None
 
-        result = probe_handles[0].get_probe_result()
-        if not result:
-            return
+        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 = self._backend.interface.deliver_poll_exit()
+
+                time_start = time.monotonic()
+                last_result = await handle.wait_async(timeout=None)
 
-        progress_printer.update([result.response.poll_exit_response])
+                # 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")
@@ -2604,25 +2601,24 @@ class Run:
         else:
             exit_handle = self._backend.interface.deliver_finish_without_exit()
 
-        exit_handle.add_probe(on_probe=self._on_probe_exit)
-
         with progress.progress_printer(
             self._printer,
-            self._settings,
+            default_text="Finishing up...",
         ) as progress_printer:
             # Wait for the run to complete.
-            _ = exit_handle.wait(
-                timeout=-1,
-                on_progress=functools.partial(
-                    self._on_progress_exit,
+            wait_with_progress(
+                exit_handle,
+                timeout=None,
+                progress_after=1,
+                display_progress=functools.partial(
+                    self._display_finish_stats,
                     progress_printer,
                 ),
             )
 
         # Print some final statistics.
         poll_exit_handle = self._backend.interface.deliver_poll_exit()
-        result = poll_exit_handle.wait(timeout=-1)
-        assert result
+        result = poll_exit_handle.wait_or(timeout=None)
         progress.print_sync_dedupe_stats(
             self._printer,
             result.response.poll_exit_response,
@@ -2630,8 +2626,7 @@ class Run:
 
         self._poll_exit_response = result.response.poll_exit_response
         internal_messages_handle = self._backend.interface.deliver_internal_messages()
-        result = internal_messages_handle.wait(timeout=-1)
-        assert result
+        result = internal_messages_handle.wait_or(timeout=None)
         self._internal_messages_response = result.response.internal_messages_response
 
         # dispatch all our final requests
@@ -2641,12 +2636,10 @@ class Run:
             self._backend.interface.deliver_request_sampled_history()
         )
 
-        result = sampled_history_handle.wait(timeout=-1)
-        assert result
+        result = sampled_history_handle.wait_or(timeout=None)
         self._sampled_history = result.response.sampled_history_response
 
-        result = final_summary_handle.wait(timeout=-1)
-        assert result
+        result = final_summary_handle.wait_or(timeout=None)
         self._final_summary = result.response.get_summary_response
 
         if self._backend:
@@ -2952,13 +2945,10 @@ class Run:
             wandb.termwarn(
                 "Artifact TTL will be disabled for source artifacts that are linked to portfolios."
             )
-        result = handle.wait(timeout=-1)
-        if result is None:
-            handle.abandon()
-        else:
-            response = result.response.link_artifact_response
-            if response.error_message:
-                wandb.termerror(response.error_message)
+        result = handle.wait_or(timeout=None)
+        response = result.response.link_artifact_response
+        if response.error_message:
+            wandb.termerror(response.error_message)
 
     @_run_decorator._noop_on_finish()
     @_run_decorator._attach
@@ -3258,7 +3248,7 @@ class Run:
         self._assert_can_log_artifact(artifact)
         if self._backend and self._backend.interface:
             if not self._settings._offline:
-                future = self._backend.interface.communicate_artifact(
+                handle = self._backend.interface.deliver_artifact(
                     self,
                     artifact,
                     aliases,
@@ -3268,7 +3258,7 @@ class Run:
                     is_user_created=is_user_created,
                     use_after_commit=use_after_commit,
                 )
-                artifact._set_save_future(future, self._public_api().client)
+                artifact._set_save_handle(handle, self._public_api().client)
             else:
                 self._backend.interface.publish_artifact(
                     self,
@@ -3670,16 +3660,18 @@ class Run:
             return {}
 
         handle = self._backend.interface.deliver_get_system_metrics()
-        result = handle.wait(timeout=1)
 
-        if result:
+        try:
+            result = handle.wait_or(timeout=1)
+        except TimeoutError:
+            return {}
+        else:
             try:
                 response = result.response.get_system_metrics_response
-                if response:
-                    return pb_to_dict(response)
+                return pb_to_dict(response) if response else {}
             except Exception as e:
                 logger.error("Error getting system metrics: %s", e)
-        return {}
+                return {}
 
     @property
     @_run_decorator._attach
@@ -3698,10 +3690,11 @@ class Run:
             self.__metadata._set_callback(self._metadata_callback)
 
         handle = self._backend.interface.deliver_get_system_metadata()
-        result = handle.wait(timeout=1)
 
-        if not result:
-            logger.error("Error getting run metadata: no result")
+        try:
+            result = handle.wait_or(timeout=1)
+        except TimeoutError:
+            logger.error("Error getting run metadata: timeout")
             return None
 
         try:

wandb/sdk/wandb_settings.py

@@ -165,6 +165,9 @@ class Settings(BaseModel, validate_assignment=True):
     entity: str | None = None
     """The W&B entity, such as a user or a team."""
 
+    organization: str | None = None
+    """The W&B organization."""
+
     force: bool = False
     """Whether to pass the `force` flag to `wandb.login()`."""
 
@@ -558,22 +561,25 @@ class Settings(BaseModel, validate_assignment=True):
     x_save_requirements: bool = True
     """Flag to save the requirements file."""
 
+    x_server_side_derived_summary: bool = False
+    """Flag to delegate automatic computation of summary from history to the server.
+
+    This does not disable user-provided summary updates.
+    """
+
     x_service_transport: str | None = None
     """Transport method for communication with the wandb service."""
 
     x_service_wait: float = 30.0
     """Time in seconds to wait for the wandb-core internal service to start."""
 
-    x_show_operation_stats: bool = True
-    """Whether to show statistics about internal operations such as data uploads."""
-
     x_start_time: float | None = None
     """The start time of the run in seconds since the Unix epoch."""
 
     x_stats_pid: int = os.getpid()
     """PID of the process that started the wandb-core process to collect system stats for."""
 
-    x_stats_sampling_interval: float = Field(default=10.0)
+    x_stats_sampling_interval: float = Field(default=15.0)
     """Sampling interval for the system monitor in seconds."""
 
     x_stats_neuron_monitor_config_path: str | None = None

wandb/sdk/wandb_setup.py

@@ -28,7 +28,6 @@ from .lib import config_util, server
 
 if TYPE_CHECKING:
     from wandb.sdk.lib.service_connection import ServiceConnection
-    from wandb.sdk.wandb_run import Run
     from wandb.sdk.wandb_settings import Settings
 
 
@@ -90,9 +89,6 @@ class _WandbSetup:
         self._server: server.Server | None = None
         self._pid = pid
 
-        # keep track of multiple runs, so we can unwind with join()s
-        self._global_run_stack: list[Run] = []
-
         # TODO(jhr): defer strict checks until settings are fully initialized
         #            and logging is ready
         self._logger: Logger = _EarlyLogger()
@@ -298,18 +294,33 @@ def singleton() -> _WandbSetup | None:
         return None
 
 
-def _setup(settings: Settings | None = None) -> _WandbSetup:
-    """Set up library context."""
+def _setup(
+    settings: Settings | None = None,
+    start_service: bool = True,
+) -> _WandbSetup:
+    """Set up library context.
+
+    Args:
+        settings: Global settings to set, or updates to the global settings
+            if the singleton has already been initialized.
+        start_service: Whether to start up the service process.
+            NOTE: A service process will only be started if allowed by the
+            global settings (after the given updates). The service will not
+            start up if the mode resolves to "disabled".
+    """
     global _singleton
 
     pid = os.getpid()
 
     if _singleton and _singleton._pid == pid:
         _singleton._update(settings=settings)
-        return _singleton
     else:
         _singleton = _WandbSetup(settings=settings, pid=pid)
-        return _singleton
+
+    if start_service and not _singleton.settings._noop:
+        _singleton.ensure_service()
+
+    return _singleton
 
 
 def setup(settings: Settings | None = None) -> _WandbSetup: