wandb 0.19.9__py3-none-musllinux_1_2_aarch64.whl → 0.19.11__py3-none-musllinux_1_2_aarch64.whl

wandb/sdk/wandb_init.py

@@ -21,7 +21,7 @@ import platform
 import sys
 import tempfile
 import time
-from typing import Any, Iterator, Literal, Sequence
+from typing import TYPE_CHECKING, Any, Iterable, Iterator, Literal, Protocol, Sequence
 
 if sys.version_info >= (3, 11):
     from typing import Self
@@ -35,6 +35,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.proto.wandb_deprecated import Deprecated
 from wandb.sdk.lib import ipython as wb_ipython
 from wandb.sdk.lib import progress, runid, wb_logging
 from wandb.sdk.lib.paths import StrPath
@@ -43,12 +44,15 @@ from wandb.util import _is_artifact_representation
 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.deprecate import deprecate
 from .mailbox import wait_with_progress
 from .wandb_helper import parse_config
 from .wandb_run import Run, TeardownHook, TeardownStage
 from .wandb_settings import Settings
 
+if TYPE_CHECKING:
+    import wandb.jupyter
+
 
 def _huggingface_version() -> str | None:
     if "transformers" in sys.modules:
@@ -124,6 +128,34 @@ class _ConfigParts:
     """
 
 
+class _PrinterCallback(Protocol):
+    """A callback for displaying messages after a printer is configured.
+
+    This is used for a few messages that may be generated before run settings
+    are computed, which are necessary for creating a printer.
+    """
+
+    def __call__(self, run_printer: printer.Printer) -> None:
+        """Display information through the given printer."""
+
+
+def _noop_printer_callback() -> _PrinterCallback:
+    """A printer callback that does not print anything."""
+    return lambda _: None
+
+
+def _concat_printer_callbacks(
+    cbs: Iterable[_PrinterCallback],
+) -> _PrinterCallback:
+    """Returns a printer callback that runs the given callbacks in order."""
+
+    def do_callbacks(run_printer: printer.Printer) -> None:
+        for cb in cbs:
+            cb(run_printer)
+
+    return do_callbacks
+
+
 class _WandbInit:
     def __init__(
         self,
@@ -143,8 +175,7 @@ class _WandbInit:
         self.backend: Backend | None = None
 
         self._teardown_hooks: list[TeardownHook] = []
-        self.notebook: wandb.jupyter.Notebook | None = None  # type: ignore
-        self.printer = printer.new_printer()
+        self.notebook: wandb.jupyter.Notebook | None = None
 
         self.deprecated_features_used: dict[str, str] = dict()
 
@@ -176,20 +207,21 @@ class _WandbInit:
 
         wandb_login._login(
             anonymous=run_settings.anonymous,
+            host=run_settings.base_url,
             force=run_settings.force,
             _disable_warning=True,
             _silent=run_settings.quiet or run_settings.silent,
         )
 
-    def warn_env_vars_change_after_setup(self) -> None:
-        """Warn if environment variables change after wandb singleton is initialized.
+    def warn_env_vars_change_after_setup(self) -> _PrinterCallback:
+        """Warn if environment variables changed after `wandb.setup()`.
 
-        Any settings from environment variables set after the singleton is initialized
-        (via login/setup/etc.) will be ignored.
+        Returns:
+            A callback to print any generated warnings.
         """
         singleton = wandb_setup.singleton()
         if singleton is None:
-            return
+            return _noop_printer_callback()
 
         exclude_env_vars = {"WANDB_SERVICE", "WANDB_KUBEFLOW_URL"}
         # check if environment variables have changed
@@ -203,26 +235,36 @@ class _WandbInit:
             for k, v in os.environ.items()
             if k.startswith("WANDB_") and k not in exclude_env_vars
         }
-        if set(singleton_env.keys()) != set(os_env.keys()) or set(
-            singleton_env.values()
-        ) != set(os_env.values()):
+
+        if (
+            set(singleton_env.keys()) == set(os_env.keys())  #
+            and set(singleton_env.values()) == set(os_env.values())
+        ):
+            return _noop_printer_callback()
+
+        def print_warning(run_printer: printer.Printer) -> None:
             line = (
                 "Changes to your `wandb` environment variables will be ignored "
                 "because your `wandb` session has already started. "
                 "For more information on how to modify your settings with "
                 "`wandb.init()` arguments, please refer to "
-                f"{self.printer.link(url_registry.url('wandb-init'), 'the W&B docs')}."
+                f"{run_printer.link(url_registry.url('wandb-init'), 'the W&B docs')}."
             )
-            self.printer.display(line, level="warn")
+            run_printer.display(line, level="warn")
+
+        return print_warning
 
     def clear_run_path_if_sweep_or_launch(
         self,
         init_settings: Settings,
-    ) -> None:
+    ) -> _PrinterCallback:
         """Clear project/entity/run_id keys if in a Sweep or a Launch context.
 
         Args:
             init_settings: Settings specified in the call to `wandb.init()`.
+
+        Returns:
+            A callback to print any generated warnings.
         """
         when_doing_thing = ""
 
@@ -232,13 +274,12 @@ class _WandbInit:
             when_doing_thing = "when running from a wandb launch context"
 
         if not when_doing_thing:
-            return
+            return _noop_printer_callback()
+
+        warnings = []
 
         def warn(key: str, value: str) -> None:
-            self.printer.display(
-                f"Ignoring {key} {value!r} {when_doing_thing}.",
-                level="warn",
-            )
+            warnings.append(f"Ignoring {key} {value!r} {when_doing_thing}.")
 
         if init_settings.project is not None:
             warn("project", init_settings.project)
@@ -250,16 +291,26 @@ class _WandbInit:
             warn("run_id", init_settings.run_id)
             init_settings.run_id = None
 
-    def make_run_settings(self, init_settings: Settings) -> Settings:
-        """Returns the run's settings.
+        def print_warnings(run_printer: printer.Printer) -> None:
+            for warning in warnings:
+                run_printer.display(warning, level="warn")
+
+        return print_warnings
+
+    def make_run_settings(
+        self,
+        init_settings: Settings,
+    ) -> tuple[Settings, _PrinterCallback]:
+        """Returns the run's settings and any warnings.
 
         Args:
             init_settings: Settings passed to `wandb.init()` or set via
                 keyword arguments.
         """
-        self.warn_env_vars_change_after_setup()
-
-        self.clear_run_path_if_sweep_or_launch(init_settings)
+        warning_callbacks: list[_PrinterCallback] = [
+            self.warn_env_vars_change_after_setup(),
+            self.clear_run_path_if_sweep_or_launch(init_settings),
+        ]
 
         # Inherit global settings.
         settings = self._wl.settings.model_copy()
@@ -305,7 +356,7 @@ class _WandbInit:
             label = runid.generate_id()
             settings.x_label = f"{prefix}-{label}" if prefix else label
 
-        return settings
+        return settings, _concat_printer_callbacks(warning_callbacks)
 
     def _load_autoresume_run_id(self, resume_file: pathlib.Path) -> str | None:
         """Returns the run_id stored in the auto-resume file, if any.
@@ -533,62 +584,73 @@ class _WandbInit:
         except OSError:
             pass
 
-    def _pause_backend(self, *args: Any, **kwargs: Any) -> None:  #  noqa
-        if self.backend is None:
-            return None
+    def _pre_run_cell_hook(self, *args, **kwargs) -> None:
+        """Hook for the IPython pre_run_cell event.
 
-        # Attempt to save the code on every execution
-        if self.notebook.save_ipynb():  # type: ignore
+        This pauses a run, preventing system metrics from being collected
+        the run's runtime from increasing. It also uploads the notebook's code.
+        """
+        if not self.backend:
+            return
+
+        if self.notebook and self.notebook.save_ipynb():
             assert self.run is not None
             res = self.run.log_code(root=None)
-            self._logger.info("saved code: %s", res)  # type: ignore
+            self._logger.info("saved code: %s", res)
+
         if self.backend.interface is not None:
-            self._logger.info("pausing backend")  # type: ignore
+            self._logger.info("pausing backend")
             self.backend.interface.publish_pause()
 
-    def _resume_backend(self, *args: Any, **kwargs: Any) -> None:  #  noqa
-        if self.backend is not None and self.backend.interface is not None:
-            self._logger.info("resuming backend")  # type: ignore
-            self.backend.interface.publish_resume()
+    def _post_run_cell_hook(self, *args, **kwargs) -> None:
+        """Hook for the IPython post_run_cell event.
+
+        Resumes collection of system metrics and the run's timer.
+        """
+        if self.backend is None or self.backend.interface is None:
+            return
+
+        self._logger.info("resuming backend")
+        self.backend.interface.publish_resume()
 
     def _jupyter_teardown(self) -> None:
         """Teardown hooks and display saving, called with wandb.finish."""
         assert self.notebook
         ipython = self.notebook.shell
-        self.notebook.save_history()
+
+        if self.run:
+            self.notebook.save_history(self.run)
+
         if self.notebook.save_ipynb():
             assert self.run is not None
             res = self.run.log_code(root=None)
-            self._logger.info("saved code and history: %s", res)  # type: ignore
-        self._logger.info("cleaning up jupyter logic")  # type: ignore
-        # because of how we bind our methods we manually find them to unregister
-        for hook in ipython.events.callbacks["pre_run_cell"]:
-            if "_resume_backend" in hook.__name__:
-                ipython.events.unregister("pre_run_cell", hook)
-        for hook in ipython.events.callbacks["post_run_cell"]:
-            if "_pause_backend" in hook.__name__:
-                ipython.events.unregister("post_run_cell", hook)
+            self._logger.info("saved code and history: %s", res)
+        self._logger.info("cleaning up jupyter logic")
+
+        ipython.events.unregister("pre_run_cell", self._pre_run_cell_hook)
+        ipython.events.unregister("post_run_cell", self._post_run_cell_hook)
+
         ipython.display_pub.publish = ipython.display_pub._orig_publish
         del ipython.display_pub._orig_publish
 
     def monkeypatch_ipython(self, settings: Settings) -> None:
         """Add hooks, and session history saving."""
-        self.notebook = wandb.jupyter.Notebook(settings)  # type: ignore
+        self.notebook = wandb.jupyter.Notebook(settings)
         ipython = self.notebook.shell
 
         # Monkey patch ipython publish to capture displayed outputs
         if not hasattr(ipython.display_pub, "_orig_publish"):
-            self._logger.info("configuring jupyter hooks %s", self)  # type: ignore
+            self._logger.info("configuring jupyter hooks %s", self)
             ipython.display_pub._orig_publish = ipython.display_pub.publish
-            # Registering resume and pause hooks
 
-            ipython.events.register("pre_run_cell", self._resume_backend)
-            ipython.events.register("post_run_cell", self._pause_backend)
+            ipython.events.register("pre_run_cell", self._pre_run_cell_hook)
+            ipython.events.register("post_run_cell", self._post_run_cell_hook)
+
             self._teardown_hooks.append(
                 TeardownHook(self._jupyter_teardown, TeardownStage.EARLY)
             )
 
-        def publish(data, metadata=None, **kwargs) -> None:  # type: ignore
+        def publish(data, metadata=None, **kwargs) -> None:
             ipython.display_pub._orig_publish(data, metadata=metadata, **kwargs)
             assert self.notebook is not None
             self.notebook.save_display(
@@ -729,6 +791,10 @@ class _WandbInit:
         ):
             setattr(drun, symbol, lambda *_, **__: None)  # type: ignore
 
+        # set properties to None
+        for attr in ("url", "project_url", "sweep_url"):
+            setattr(type(drun), attr, property(lambda _: None))
+
         class _ChainableNoOp:
             """An object that allows chaining arbitrary attributes and method calls."""
 
@@ -753,7 +819,7 @@ class _WandbInit:
             def __call__(self, *args: Any, **kwargs: Any) -> _ChainableNoOp:
                 return _ChainableNoOp()
 
-        drun.log_artifact = _ChainableNoOpField()  # type: ignore[method-assign]
+        drun.log_artifact = _ChainableNoOpField()  # type: ignore
         # attributes
         drun._start_time = time.time()
         drun._starting_step = 0
@@ -777,7 +843,12 @@ class _WandbInit:
         )
         return drun
 
-    def init(self, settings: Settings, config: _ConfigParts) -> Run:  # noqa: C901
+    def init(  # noqa: C901
+        self,
+        settings: Settings,
+        config: _ConfigParts,
+        run_printer: printer.Printer,
+    ) -> Run:
         self._logger.info("calling init triggers")
         trigger.call("on_init")
 
@@ -788,23 +859,36 @@ class _WandbInit:
             f"\nconfig: {config.base_no_artifacts}"
         )
 
-        if wandb.run is not None and os.getpid() == wandb.run._init_pid:
+        if previous_run := self._wl.most_recent_active_run:
             if (
                 settings.reinit in (True, "finish_previous")
                 # calling wandb.init() in notebooks finishes previous runs
                 # by default for user convenience.
                 or (settings.reinit == "default" and wb_ipython.in_notebook())
             ):
-                self._logger.info("finishing previous run: %s", wandb.run.id)
-                wandb.run.finish()
+                run_printer.display(
+                    "Finishing previous runs because reinit is set"
+                    f" to {settings.reinit!r}."
+                )
+                self._wl.finish_all_active_runs()
+
+            elif settings.reinit == "create_new":
+                self._logger.info(
+                    "wandb.init() called while a run is active,"
+                    " and reinit is set to 'create_new', so continuing"
+                )
+
             else:
-                self._logger.info("wandb.init() called while a run is active")
+                run_printer.display(
+                    "wandb.init() called while a run is active and reinit is"
+                    f" set to {settings.reinit!r}, so returning the previous"
+                    " run."
+                )
 
-                # NOTE: Updates telemetry on the pre-existing run.
-                with telemetry.context() as tel:
+                with telemetry.context(run=previous_run) as tel:
                     tel.feature.init_return_run = True
 
-                return wandb.run
+                return previous_run
 
         self._logger.info("starting backend")
 
@@ -897,6 +981,12 @@ class _WandbInit:
                 )
                 tel.feature.shared_mode = True
 
+            if settings.x_label:
+                tel.feature.user_provided_label = True
+
+            if wandb.env.dcgm_profiling_enabled():
+                tel.feature.dcgm_profiling_enabled = True
+
             tel.env.maybe_mp = _maybe_mp_process(backend)
 
         if not settings.label_disable:
@@ -946,7 +1036,7 @@ class _WandbInit:
             assert backend.interface
 
             with progress.progress_printer(
-                self.printer,
+                run_printer,
                 default_text="Waiting for wandb.init()...",
             ) as progress_printer:
                 await progress.loop_printing_operation_stats(
@@ -1052,6 +1142,10 @@ class _WandbInit:
             run.use_artifact(job_artifact)
 
         self.backend = backend
+
+        if settings.reinit != "create_new":
+            _set_global_run(run)
+
         run._on_start()
         self._logger.info("run started, returning control to user process")
         return run
@@ -1128,10 +1222,36 @@ def _attach(
         raise UsageError(f"Failed to attach to run: {attach_response.error.message}")
 
     run._set_run_obj(attach_response.run)
+    _set_global_run(run)
     run._on_attach()
     return run
 
 
+def _set_global_run(run: Run) -> None:
+    """Set `wandb.run` and point some top-level functions to its methods.
+
+    Args:
+        run: The run to make global.
+    """
+    module.set_global(
+        run=run,
+        config=run.config,
+        log=run.log,
+        summary=run.summary,
+        save=run.save,
+        use_artifact=run.use_artifact,
+        log_artifact=run.log_artifact,
+        define_metric=run.define_metric,
+        alert=run.alert,
+        watch=run.watch,
+        unwatch=run.unwatch,
+        mark_preempting=run.mark_preempting,
+        log_model=run.log_model,
+        use_model=run.use_model,
+        link_model=run.link_model,
+    )
+
+
 def _monkeypatch_openai_gym() -> None:
     """Patch OpenAI gym to log to the global `wandb.run`."""
     if len(wandb.patched["gym"]) > 0:
@@ -1163,16 +1283,40 @@ def try_create_root_dir(settings: Settings) -> None:
             This function may update the root_dir to a temporary directory
             if the parent directory is not writable.
     """
+    fallback_to_temp_dir = False
+
     try:
-        if not os.path.exists(settings.root_dir):
-            os.makedirs(settings.root_dir, exist_ok=True)
+        os.makedirs(settings.root_dir, exist_ok=True)
     except OSError:
-        temp_dir = tempfile.gettempdir()
         wandb.termwarn(
-            f"Path {settings.root_dir} wasn't writable, using system temp directory {temp_dir}.",
+            f"Unable to create root directory {settings.root_dir}",
             repeat=False,
         )
-        settings.root_dir = temp_dir
+        fallback_to_temp_dir = True
+    else:
+        if not os.access(settings.root_dir, os.W_OK | os.R_OK):
+            wandb.termwarn(
+                f"Path {settings.root_dir} wasn't read/writable",
+                repeat=False,
+            )
+            fallback_to_temp_dir = True
+
+    if not fallback_to_temp_dir:
+        return
+
+    tmp_dir = tempfile.gettempdir()
+    if not os.access(tmp_dir, os.W_OK | os.R_OK):
+        raise ValueError(
+            f"System temp directory ({tmp_dir}) is not writable/readable, "
+            "please set the `dir` argument in `wandb.init()` to a writable/readable directory."
+        )
+
+    settings.root_dir = tmp_dir
+    wandb.termwarn(
+        f"Falling back to temporary directory {tmp_dir}.",
+        repeat=False,
+    )
+    os.makedirs(settings.root_dir, exist_ok=True)
 
 
 def init(  # noqa: C901
@@ -1199,6 +1343,7 @@ def init(  # noqa: C901
             "default",
             "return_previous",
             "finish_previous",
+            "create_new",
         ]
     ) = None,
     resume: bool | Literal["allow", "never", "must", "auto"] | None = None,
@@ -1476,7 +1621,7 @@ def init(  # noqa: C901
         wi = _WandbInit(wl, init_telemetry)
 
         wi.maybe_login(init_settings)
-        run_settings = wi.make_run_settings(init_settings)
+        run_settings, show_warnings = wi.make_run_settings(init_settings)
 
         if isinstance(run_settings.reinit, bool):
             wi.deprecated_features_used["run__reinit_bool"] = (
@@ -1494,6 +1639,8 @@ def init(  # noqa: C901
             init_telemetry.feature.offline = True
 
         wi.set_run_id(run_settings)
+        run_printer = printer.new_printer(run_settings)
+        show_warnings(run_printer)
 
         with contextlib.ExitStack() as exit_stack:
             exit_stack.enter_context(wb_logging.log_to_run(run_settings.run_id))
@@ -1527,7 +1674,7 @@ def init(  # noqa: C901
             if run_settings.x_server_side_derived_summary:
                 init_telemetry.feature.server_side_derived_summary = True
 
-            return wi.init(run_settings, run_config)
+            return wi.init(run_settings, run_config, run_printer)
 
     except KeyboardInterrupt as e:
         if wl:

wandb/sdk/wandb_login.py

@@ -45,6 +45,7 @@ def login(
     force: Optional[bool] = None,
     timeout: Optional[int] = None,
     verify: bool = False,
+    referrer: Optional[str] = None,
 ) -> bool:
     """Set up W&B login credentials.
 
@@ -64,6 +65,7 @@ def login(
         force: (bool, optional) If true, will force a relogin.
         timeout: (int, optional) Number of seconds to wait for user input.
         verify: (bool) Verify the credentials with the W&B server.
+        referrer: (string, optional) The referrer to use in the URL login request.
 
     Returns:
         bool: if key is configured
@@ -81,6 +83,7 @@ def login(
         force=force,
         timeout=timeout,
         verify=verify,
+        referrer=referrer,
     )
 
 
@@ -199,7 +202,9 @@ class _WandbLogin:
         if not self._wandb_setup.settings._offline:
             self._wandb_setup._update_user_settings()
 
-    def _prompt_api_key(self) -> Tuple[Optional[str], ApiKeyStatus]:
+    def _prompt_api_key(
+        self, referrer: Optional[str] = None
+    ) -> Tuple[Optional[str], ApiKeyStatus]:
         api = Api(self._settings)
         while True:
             try:
@@ -208,6 +213,7 @@ class _WandbLogin:
                     api=api,
                     no_offline=self._settings.force if self._settings else None,
                     no_create=self._settings.force if self._settings else None,
+                    referrer=referrer,
                 )
             except ValueError as e:
                 # invalid key provided, try again
@@ -222,9 +228,11 @@ class _WandbLogin:
                 return None, ApiKeyStatus.OFFLINE
             return key, ApiKeyStatus.VALID
 
-    def prompt_api_key(self) -> Tuple[Optional[str], ApiKeyStatus]:
+    def prompt_api_key(
+        self, referrer: Optional[str] = None
+    ) -> Tuple[Optional[str], ApiKeyStatus]:
         """Updates the global API key by prompting the user."""
-        key, status = self._prompt_api_key()
+        key, status = self._prompt_api_key(referrer)
         if status == ApiKeyStatus.NOTTY:
             directive = (
                 "wandb login [your_api_key]"
@@ -262,6 +270,7 @@ def _login(
     force: Optional[bool] = None,
     timeout: Optional[int] = None,
     verify: bool = False,
+    referrer: str = "models",
     _silent: Optional[bool] = None,
     _disable_warning: Optional[bool] = None,
 ) -> bool:
@@ -302,7 +311,7 @@ def _login(
         if key and not relogin:
             key_is_pre_configured = True
         else:
-            key, key_status = wlogin.prompt_api_key()
+            key, key_status = wlogin.prompt_api_key(referrer=referrer)
 
     if verify:
         wlogin._verify_login(key)