wandb 0.19.6rc4__py3-none-any.whl → 0.19.7__py3-none-any.whl

wandb/sdk/service/server_sock.py

@@ -70,6 +70,7 @@ class SockServerInterfaceReaderThread(threading.Thread):
             sock_client = self._clients.get_client(sockid)
             assert sock_client
             sresp = spb.ServerResponse()
+            sresp.request_id = result.control.mailbox_slot
             sresp.result_communicate.CopyFrom(result)
             sock_client.send_server_response(sresp)
 
@@ -148,7 +149,10 @@ class SockServerReadThread(threading.Thread):
         inform_attach_response.settings.CopyFrom(
             self._mux._streams[stream_id]._settings._proto,
         )
-        response = spb.ServerResponse(inform_attach_response=inform_attach_response)
+        response = spb.ServerResponse(
+            request_id=sreq.request_id,
+            inform_attach_response=inform_attach_response,
+        )
         self._sock_client.send_server_response(response)
         iface = self._mux.get_stream(stream_id).interface
 

wandb/sdk/service/streams.py

@@ -8,13 +8,13 @@ StreamMux: Container for dictionary of stream threads per runid
 
 from __future__ import annotations
 
+import asyncio
 import functools
-import multiprocessing
 import queue
 import threading
 import time
 from threading import Event
-from typing import Any, Callable
+from typing import Any, Callable, NoReturn
 
 import psutil
 
@@ -22,14 +22,9 @@ 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 asyncio_compat, progress
 from wandb.sdk.lib import printer as printerlib
-from wandb.sdk.lib import progress
-from wandb.sdk.lib.mailbox import (
-    Mailbox,
-    MailboxProbe,
-    MailboxProgress,
-    MailboxProgressAll,
-)
+from wandb.sdk.mailbox import Mailbox, MailboxHandle, wait_all_with_progress
 from wandb.sdk.wandb_run import Run
 
 from ..interface.interface_relay import InterfaceRelay
@@ -61,19 +56,16 @@ class StreamRecord:
     _settings: SettingsStatic
     _started: bool
 
-    def __init__(self, settings: SettingsStatic, mailbox: Mailbox) -> None:
+    def __init__(self, settings: SettingsStatic) -> None:
         self._started = False
-        self._mailbox = mailbox
+        self._mailbox = Mailbox()
         self._record_q = queue.Queue()
         self._result_q = queue.Queue()
         self._relay_q = queue.Queue()
-        process = multiprocessing.current_process()
         self._iface = InterfaceRelay(
             record_q=self._record_q,
             result_q=self._result_q,
             relay_q=self._relay_q,
-            process=process,
-            process_check=False,
             mailbox=self._mailbox,
         )
         self._settings = settings
@@ -84,7 +76,7 @@ class StreamRecord:
         self._wait_thread_active()
 
     def _wait_thread_active(self) -> None:
-        self._iface.deliver_status().wait(timeout=-1)
+        self._iface.deliver_status().wait_or(timeout=None)
 
     def join(self) -> None:
         self._iface.join()
@@ -141,7 +133,6 @@ class StreamMux:
     _action_q: queue.Queue[StreamAction]
     _stopped: Event
     _pid_checked_ts: float | None
-    _mailbox: Mailbox
 
     def __init__(self) -> None:
         self._streams_lock = threading.Lock()
@@ -151,8 +142,6 @@ class StreamMux:
         self._stopped = Event()
         self._action_q = queue.Queue()
         self._pid_checked_ts = None
-        self._mailbox = Mailbox()
-        self._mailbox.enable_keepalive()
 
     def _get_stopped_event(self) -> Event:
         # TODO: clean this up, there should be a better way to abstract this
@@ -209,7 +198,7 @@ class StreamMux:
             return stream
 
     def _process_add(self, action: StreamAction) -> None:
-        stream = StreamRecord(action._data, mailbox=self._mailbox)
+        stream = StreamRecord(action._data)
         # run_id = action.stream_id  # will want to fix if a streamid != runid
         settings = action._data
         thread = StreamThread(
@@ -247,41 +236,51 @@ class StreamMux:
                 stream.drop()
                 stream.join()
 
-    def _on_probe_exit(self, probe_handle: MailboxProbe, stream: StreamRecord) -> 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)
-        handle = stream.interface.deliver_poll_exit()
-        probe_handle.set_mailbox_handle(handle)
-
-    def _on_progress_exit(self, progress_handle: MailboxProgress) -> None:
-        pass
-
-    def _on_progress_exit_all(
+    async def _finish_all_progress(
         self,
         progress_printer: progress.ProgressPrinter,
-        progress_all_handle: MailboxProgressAll,
+        streams_to_watch: dict[str, StreamRecord],
     ) -> 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())
-
-        assert probe_handles
-
-        if self._check_orphaned():
-            self._stopped.set()
-
-        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)
-
-        progress_printer.update(poll_exit_responses)
+        """Poll the streams and display statistics about them.
+
+        This never returns and must be cancelled.
+
+        Args:
+            progress_printer: Printer to use for displaying finish progress.
+            streams_to_watch: Streams to poll for finish progress.
+        """
+        results: dict[str, pb.Result | None] = {}
+
+        async def loop_poll_stream(
+            stream_id: str,
+            stream: StreamRecord,
+        ) -> NoReturn:
+            while True:
+                start_time = time.monotonic()
+
+                handle = stream.interface.deliver_poll_exit()
+                results[stream_id] = await handle.wait_async(timeout=None)
+
+                elapsed_time = time.monotonic() - start_time
+                if elapsed_time < 1:
+                    await asyncio.sleep(1 - elapsed_time)
+
+        async def loop_update_printer() -> NoReturn:
+            while True:
+                poll_exit_responses: list[pb.PollExitResponse] = []
+                for result in results.values():
+                    if not result or not result.response:
+                        continue
+                    if poll_exit_response := result.response.poll_exit_response:
+                        poll_exit_responses.append(poll_exit_response)
+
+                progress_printer.update(poll_exit_responses)
+                await asyncio.sleep(1)
+
+        async with asyncio_compat.open_task_group() as task_group:
+            for stream_id, stream in streams_to_watch.items():
+                task_group.start_soon(loop_poll_stream(stream_id, stream))
+            task_group.start_soon(loop_update_printer())
 
     def _finish_all(self, streams: dict[str, StreamRecord], exit_code: int) -> None:
         if not streams:
@@ -291,7 +290,7 @@ class StreamMux:
 
         # 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
-        exit_handles = []
+        exit_handles: list[MailboxHandle] = []
 
         # only finish started streams, non started streams failed early
         started_streams: dict[str, StreamRecord] = {}
@@ -302,27 +301,24 @@ class StreamMux:
 
         for stream in started_streams.values():
             handle = stream.interface.deliver_exit(exit_code)
-            handle.add_progress(self._on_progress_exit)
-            handle.add_probe(functools.partial(self._on_probe_exit, stream=stream))
             exit_handles.append(handle)
 
-            # this message is confusing, we should remove it
-            # Run._footer_exit_status_info(
-            #     exit_code, settings=stream._settings, printer=printer  # type: ignore
-            # )
-
-        with progress.progress_printer(printer) as progress_printer:
+        with progress.progress_printer(
+            printer,
+            default_text="Finishing up...",
+        ) 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,
+            wait_all_with_progress(
+                exit_handles,
+                timeout=None,
+                progress_after=1,
+                display_progress=functools.partial(
+                    self._finish_all_progress,
                     progress_printer,
+                    started_streams,
                 ),
             )
-            assert got_result
 
         # These could be done in parallel in the future
         for _sid, stream in started_streams.items():
@@ -332,20 +328,16 @@ class StreamMux:
             sampled_history_handle = stream.interface.deliver_request_sampled_history()
             internal_messages_handle = stream.interface.deliver_internal_messages()
 
-            result = internal_messages_handle.wait(timeout=-1)
-            assert result
+            result = internal_messages_handle.wait_or(timeout=None)
             internal_messages_response = result.response.internal_messages_response
 
-            result = poll_exit_handle.wait(timeout=-1)
-            assert result
+            result = poll_exit_handle.wait_or(timeout=None)
             poll_exit_response = result.response.poll_exit_response
 
-            result = sampled_history_handle.wait(timeout=-1)
-            assert result
+            result = sampled_history_handle.wait_or(timeout=None)
             sampled_history = result.response.sampled_history_response
 
-            result = final_summary_handle.wait(timeout=-1)
-            assert result
+            result = final_summary_handle.wait_or(timeout=None)
             final_summary = result.response.get_summary_response
 
             Run._footer(

wandb/sdk/verify/verify.py

@@ -1,6 +1,8 @@
 """Utilities for wandb verify."""
 
+import contextlib
 import getpass
+import io
 import os
 import time
 from functools import partial
@@ -163,8 +165,8 @@ def check_run(api: Api) -> bool:
         )
         print_results(failed_test_strings, False)
         return False
-    for key, value in prev_run.config.items():
-        if config[key] != value:
+    for key, value in config.items():
+        if prev_run.config.get(key) != value:
             failed_test_strings.append(
                 "Read config values don't match run config. Contact W&B for support."
             )
@@ -486,6 +488,56 @@ def check_wandb_version(api: Api) -> None:
     print_results(fail_string, warning)
 
 
+def check_sweeps(api: Api) -> bool:
+    print("Checking sweep creation and agent execution".ljust(72, "."), end="")  # noqa: T201
+    failed_test_strings: List[str] = []
+
+    sweep_config = {
+        "method": "random",
+        "metric": {"goal": "minimize", "name": "score"},
+        "parameters": {
+            "x": {"values": [0.01, 0.05, 0.1]},
+            "y": {"values": [1, 2, 3]},
+        },
+        "name": "verify_sweep",
+    }
+
+    try:
+        with contextlib.redirect_stdout(io.StringIO()):
+            sweep_id = wandb.sweep(
+                sweep=sweep_config, project=PROJECT_NAME, entity=api.default_entity
+            )
+    except Exception as e:
+        failed_test_strings.append(f"Failed to create sweep: {e}")
+        print_results(failed_test_strings, False)
+        return False
+
+    if not sweep_id:
+        failed_test_strings.append("Sweep creation returned an invalid ID.")
+        print_results(failed_test_strings, False)
+        return False
+
+    try:
+
+        def objective(config):
+            score = config.x**3 + config.y
+            return score
+
+        def main():
+            with wandb.init(project=PROJECT_NAME) as run:
+                score = objective(run.config)
+                run.log({"score": score})
+
+        wandb.agent(sweep_id, function=main, count=10)
+    except Exception as e:
+        failed_test_strings.append(f"Failed to run sweep agent: {e}")
+        print_results(failed_test_strings, False)
+        return False
+
+    print_results(failed_test_strings, False)
+    return len(failed_test_strings) == 0
+
+
 def retry_fn(fn: Callable) -> Any:
     ini_time = time.time()
     res = None

wandb/sdk/wandb_init.py

@@ -34,7 +34,7 @@ from wandb.errors import CommError, Error, UsageError
 from wandb.errors.links import url_registry
 from wandb.errors.util import ProtobufErrorHandler
 from wandb.integration import sagemaker
-from wandb.sdk.lib import runid
+from wandb.sdk.lib import progress, runid
 from wandb.sdk.lib.paths import StrPath
 from wandb.util import _is_artifact_representation
 
@@ -42,7 +42,7 @@ from . import wandb_login, wandb_setup
 from .backend.backend import Backend
 from .lib import SummaryDisabled, filesystem, module, paths, printer, telemetry
 from .lib.deprecate import Deprecated, deprecate
-from .lib.mailbox import Mailbox, MailboxProgress
+from .mailbox import Mailbox, wait_with_progress
 from .wandb_helper import parse_config
 from .wandb_run import Run, TeardownHook, TeardownStage
 from .wandb_settings import Settings
@@ -293,6 +293,16 @@ class _WandbInit:
 
         settings.x_start_time = time.time()
 
+        # In shared mode, generate a unique label if not provided.
+        # The label is used to distinguish between system metrics and console logs
+        # from different writers to the same run.
+        if settings._shared and not settings.x_label:
+            # TODO: If executed in a known distributed environment (e.g. Ray or SLURM),
+            #   use the env vars to generate a label (e.g. SLURM_JOB_ID or RANK)
+            prefix = settings.host or ""
+            label = runid.generate_id()
+            settings.x_label = f"{prefix}-{label}" if prefix else label
+
         return settings
 
     def _load_autoresume_run_id(self, resume_file: pathlib.Path) -> str | None:
@@ -747,11 +757,6 @@ class _WandbInit:
         )
         return drun
 
-    def _on_progress_init(self, handle: MailboxProgress) -> None:
-        line = "Waiting for wandb.init()...\r"
-        percent_done = handle.percent_done
-        self.printer.progress_update(line, percent_done=percent_done)
-
     def init(self, settings: Settings, config: _ConfigParts) -> Run:  # noqa: C901
         self._logger.info("calling init triggers")
         trigger.call("on_init")
@@ -763,28 +768,18 @@ class _WandbInit:
             f"\nconfig: {config.base_no_artifacts}"
         )
 
-        if (
-            settings.reinit or (settings._jupyter and settings.reinit is not False)
-        ) and len(self._wl._global_run_stack) > 0:
-            if len(self._wl._global_run_stack) > 1:
-                wandb.termwarn(
-                    "Launching multiple wandb runs using Python's threading"
-                    " module is not well-supported."
-                    " Please use multiprocessing instead."
-                    " Finishing previous run before initializing another."
-                )
-
-            latest_run = self._wl._global_run_stack[-1]
-            self._logger.info(f"found existing run on stack: {latest_run.id}")
-            latest_run.finish()
-        elif wandb.run is not None and os.getpid() == wandb.run._init_pid:
-            self._logger.info("wandb.init() called when a run is still active")
+        if wandb.run is not None and os.getpid() == wandb.run._init_pid:
+            if settings.reinit:
+                self._logger.info(f"finishing previous run: {wandb.run.id}")
+                wandb.run.finish()
+            else:
+                self._logger.info("wandb.init() called while a run is active")
 
-            # NOTE: Updates telemetry on the pre-existing run.
-            with telemetry.context() as tel:
-                tel.feature.init_return_run = True
+                # NOTE: Updates telemetry on the pre-existing run.
+                with telemetry.context() as tel:
+                    tel.feature.init_return_run = True
 
-            return wandb.run
+                return wandb.run
 
         self._logger.info("starting backend")
 
@@ -905,7 +900,6 @@ class _WandbInit:
         run._set_teardown_hooks(self._teardown_hooks)
 
         assert backend.interface
-        mailbox.enable_keepalive()
         backend.interface.publish_header()
 
         # Using GitRepo() blocks & can be slow, depending on user's current git setup.
@@ -921,16 +915,6 @@ class _WandbInit:
             )
         error: wandb.Error | None = None
 
-        # In shared mode, generate a unique label if not provided.
-        # The label is used to distinguish between system metrics and console logs
-        # from different writers to the same run.
-        if settings._shared and not settings.x_label:
-            # TODO: If executed in a known distributed environment (e.g. Ray or SLURM),
-            #   use the env vars to generate a label (e.g. SLURM_JOB_ID or RANK)
-            prefix = settings.host or ""
-            label = runid.generate_id()
-            settings.x_label = f"{prefix}-{label}" if prefix else label
-
         timeout = settings.init_timeout
 
         self._logger.info(
@@ -938,11 +922,18 @@ class _WandbInit:
         )
 
         run_init_handle = backend.interface.deliver_run(run)
-        result = run_init_handle.wait(
-            timeout=timeout,
-            on_progress=self._on_progress_init,
-            cancel=True,
-        )
+
+        async def display_init_message() -> None:
+            assert backend.interface
+
+            with progress.progress_printer(
+                self.printer,
+                default_text="Waiting for wandb.init()...",
+            ) as progress_printer:
+                await progress.loop_printing_operation_stats(
+                    progress_printer,
+                    backend.interface,
+                )
 
         # Raise an error if deliver_run failed.
         #
@@ -951,8 +942,16 @@ class _WandbInit:
         #
         # TODO: Remove try-except once x_disable_service is removed.
         try:
-            if not result or not result.run_result:
-                run_init_handle._cancel()
+            try:
+                result = wait_with_progress(
+                    run_init_handle,
+                    timeout=timeout,
+                    progress_after=1,
+                    display_progress=display_init_message,
+                )
+
+            except TimeoutError:
+                run_init_handle.cancel(backend.interface)
 
                 # This may either be an issue with the W&B server (a CommError)
                 # or a bug in the SDK (an Error). We cannot distinguish between
@@ -963,6 +962,8 @@ class _WandbInit:
                     " setting: `wandb.init(settings=wandb.Settings(init_timeout=120))`."
                 )
 
+            assert result.run_result
+
             if error := ProtobufErrorHandler.to_exception(result.run_result.error):
                 raise error
 
@@ -1004,13 +1005,13 @@ class _WandbInit:
         assert backend.interface
 
         run_start_handle = backend.interface.deliver_run_start(run)
-        # TODO: add progress to let user know we are doing something
-        run_start_result = run_start_handle.wait(timeout=30)
-        if run_start_result is None:
-            run_start_handle.abandon()
+        try:
+            # TODO: add progress to let user know we are doing something
+            run_start_handle.wait_or(timeout=30)
+        except TimeoutError:
+            pass
 
         assert self._wl is not None
-        self._wl._global_run_stack.append(run)
         self.run = run
 
         run._handle_launch_artifact_overrides()
@@ -1097,14 +1098,13 @@ def _attach(
     run._set_backend(backend)
     assert backend.interface
 
-    mailbox.enable_keepalive()
-
     attach_handle = backend.interface.deliver_attach(attach_id)
-    # TODO: add progress to let user know we are doing something
-    attach_result = attach_handle.wait(timeout=30)
-    if not attach_result:
-        attach_handle.abandon()
+    try:
+        # TODO: add progress to let user know we are doing something
+        attach_result = attach_handle.wait_or(timeout=30)
+    except TimeoutError:
         raise UsageError("Timeout attaching to run")
+
     attach_response = attach_result.response.attach_response
     if attach_response.error and attach_response.error.message:
         raise UsageError(f"Failed to attach to run: {attach_response.error.message}")
@@ -1222,10 +1222,10 @@ def init(  # noqa: C901
             on the system, such as checking the git root or the current program
             file. If we can't infer the project name, the project will default to
             `"uncategorized"`.
-        dir: An absolute path to the directory where metadata and downloaded
-            files will be stored. When calling `download()` on an artifact, files
-            will be saved to this directory. If not specified, this defaults to
-            the `./wandb` directory.
+        dir: The absolute path to the directory where experiment logs and
+            metadata files are stored. If not specified, this defaults
+            to the `./wandb` directory. Note that this does not affect the
+            location where artifacts are stored when calling `download()`.
         id: A unique identifier for this run, used for resuming. It must be unique
             within the project and cannot be reused once a run is deleted. The
             identifier must not contain any of the following special characters:
@@ -1426,7 +1426,7 @@ def init(  # noqa: C901
     wl: wandb_setup._WandbSetup | None = None
 
     try:
-        wl = wandb.setup()
+        wl = wandb_setup._setup(start_service=False)
 
         wi = _WandbInit(wl, init_telemetry)
 

wandb/sdk/wandb_login.py

@@ -162,8 +162,8 @@ class _WandbLogin:
             repeat=False,
         )
 
-    def configure_api_key(self, key: str) -> None:
-        """Saves the API key and updates the the global setup object."""
+    def try_save_api_key(self, key: str) -> None:
+        """Saves the API key to disk for future use."""
         if self._settings._notebook and not self._settings.silent:
             wandb.termwarn(
                 "If you're specifying your api key in code, ensure this "
@@ -172,7 +172,10 @@ class _WandbLogin:
                 "`wandb login` from the command line."
             )
         if key:
-            apikey.write_key(self._settings, key)
+            try:
+                apikey.write_key(self._settings, key)
+            except apikey.WriteNetrcError as e:
+                wandb.termwarn(str(e))
 
     def update_session(
         self,
@@ -305,7 +308,7 @@ def _login(
         wlogin._verify_login(key)
 
     if not key_is_pre_configured:
-        wlogin.configure_api_key(key)
+        wlogin.try_save_api_key(key)
         wlogin.update_session(key, status=key_status)
         wlogin._update_global_anonymous_setting()