wandb 0.21.0__py3-none-win_amd64.whl → 0.21.2__py3-none-win_amd64.whl

wandb/sdk/wandb_setup.py

@@ -28,12 +28,13 @@ import logging
 import os
 import pathlib
 import sys
+import threading
 from typing import TYPE_CHECKING, Any, Union
 
 import wandb
 import wandb.integration.sagemaker as sagemaker
 from wandb.env import CONFIG_DIR
-from wandb.sdk.lib import import_hooks, wb_logging
+from wandb.sdk.lib import asyncio_manager, import_hooks, wb_logging
 
 from . import wandb_settings
 from .lib import config_util, server
@@ -88,19 +89,16 @@ Logger = Union[logging.Logger, _EarlyLogger]
 class _WandbSetup:
     """W&B library singleton."""
 
-    def __init__(
-        self,
-        pid: int,
-        settings: Settings | None = None,
-        environ: dict | None = None,
-    ) -> None:
+    def __init__(self, pid: int) -> None:
+        self._asyncer = asyncio_manager.AsyncioManager()
+        self._asyncer.start()
+
         self._connection: ServiceConnection | None = None
 
         self._active_runs: list[wandb_run.Run] = []
+        self._active_runs_lock = threading.Lock()
 
-        self._environ = environ or dict(os.environ)
         self._sweep_config: dict | None = None
-        self._config: dict | None = None
         self._server: server.Server | None = None
         self._pid = pid
 
@@ -108,11 +106,13 @@ class _WandbSetup:
         #            and logging is ready
         self._logger: Logger = _EarlyLogger()
 
-        self._settings = self._settings_setup(settings)
+        self._settings: Settings | None = None
+        self._settings_environ: dict[str, str] | None = None
 
-        wandb.termsetup(self._settings, None)
-
-        self._setup()
+    @property
+    def asyncer(self) -> asyncio_manager.AsyncioManager:
+        """The internal asyncio thread used by wandb."""
+        return self._asyncer
 
     def add_active_run(self, run: wandb_run.Run) -> None:
         """Append a run to the active runs list.
@@ -122,8 +122,9 @@ class _WandbSetup:
         Args:
             run: A newly initialized run.
         """
-        if run not in self._active_runs:
-            self._active_runs.append(run)
+        with self._active_runs_lock:
+            if run not in self._active_runs:
+                self._active_runs.append(run)
 
     def remove_active_run(self, run: wandb_run.Run) -> None:
         """Remove the run from the active runs list.
@@ -134,17 +135,19 @@ class _WandbSetup:
             run: A run that is finished or crashed.
         """
         try:
-            self._active_runs.remove(run)
+            with self._active_runs_lock:
+                self._active_runs.remove(run)
         except ValueError:
             pass  # Removing a run multiple times is not an error.
 
     @property
     def most_recent_active_run(self) -> wandb_run.Run | None:
         """The most recently initialized run that is not yet finished."""
-        if not self._active_runs:
-            return None
+        with self._active_runs_lock:
+            if not self._active_runs:
+                return None
 
-        return self._active_runs[-1]
+            return self._active_runs[-1]
 
     def finish_all_active_runs(self) -> None:
         """Finish all unfinished runs.
@@ -156,74 +159,131 @@ class _WandbSetup:
         default and only behavior, it does not seem worth optimizing.
         """
         # Take a snapshot as each call to `finish()` modifies `_active_runs`.
-        runs_copy = list(self._active_runs)
+        with self._active_runs_lock:
+            runs_copy = list(self._active_runs)
+
         for run in runs_copy:
             run.finish()
 
-    def _settings_setup(
+    def did_environment_change(self) -> bool:
+        """Check if os.environ has changed since settings were initialized."""
+        if not self._settings_environ:
+            return False
+
+        exclude_env_vars = {"WANDB_SERVICE", "WANDB_KUBEFLOW_URL"}
+        singleton_env = {
+            k: v
+            for k, v in self._settings_environ.items()
+            if k.startswith("WANDB_") and k not in exclude_env_vars
+        }
+        os_env = {
+            k: v
+            for k, v in os.environ.items()
+            if k.startswith("WANDB_") and k not in exclude_env_vars
+        }
+
+        return (
+            set(singleton_env.keys()) == set(os_env.keys())  #
+            and set(singleton_env.values()) == set(os_env.values())
+        )
+
+    def _load_settings(
         self,
-        settings: Settings | None,
-    ) -> wandb_settings.Settings:
-        s = wandb_settings.Settings()
+        *,
+        system_settings_path: str | None,
+        disable_sagemaker: bool,
+        overrides: Settings | None = None,
+    ) -> None:
+        """Load settings from environment variables, config files, etc.
+
+        Args:
+            system_settings_path: Location of system settings file to use.
+                If not provided, reads the WANDB_CONFIG_DIR environment
+                variable or uses the default location.
+            disable_sagemaker: If true, skips modifying settings based on
+                SageMaker.
+            overrides: Additional settings to apply to the global settings.
+        """
+        self._settings = wandb_settings.Settings()
 
         # the pid of the process to monitor for system stats
         pid = os.getpid()
         self._logger.info(f"Current SDK version is {wandb.__version__}")
         self._logger.info(f"Configure stats pid to {pid}")
-        s.x_stats_pid = pid
+        self._settings.x_stats_pid = pid
 
-        if settings and settings.settings_system:
-            s.settings_system = settings.settings_system
+        if system_settings_path:
+            self._settings.settings_system = system_settings_path
         elif config_dir_str := os.getenv(CONFIG_DIR, None):
             config_dir = pathlib.Path(config_dir_str).expanduser()
-            s.settings_system = str(config_dir / "settings")
+            self._settings.settings_system = str(config_dir / "settings")
         else:
-            s.settings_system = str(
+            self._settings.settings_system = str(
                 pathlib.Path("~", ".config", "wandb", "settings").expanduser()
             )
 
         # load settings from the system config
-        if s.settings_system:
-            self._logger.info(f"Loading settings from {s.settings_system}")
-        s.update_from_system_config_file()
+        if self._settings.settings_system:
+            self._logger.info(
+                f"Loading settings from {self._settings.settings_system}",
+            )
+        self._settings.update_from_system_config_file()
 
         # load settings from the workspace config
-        if s.settings_workspace:
-            self._logger.info(f"Loading settings from {s.settings_workspace}")
-        s.update_from_workspace_config_file()
+        if self._settings.settings_workspace:
+            self._logger.info(
+                f"Loading settings from {self._settings.settings_workspace}",
+            )
+        self._settings.update_from_workspace_config_file()
 
         # load settings from the environment variables
         self._logger.info("Loading settings from environment variables")
-        s.update_from_env_vars(self._environ)
+        self._settings_environ = os.environ.copy()
+        self._settings.update_from_env_vars(self._settings_environ)
 
         # infer settings from the system environment
-        s.update_from_system_environment()
+        self._settings.update_from_system_environment()
 
         # load SageMaker settings
-        check_sagemaker_env = not s.sagemaker_disable
-        if settings and settings.sagemaker_disable:
-            check_sagemaker_env = False
-        if check_sagemaker_env and sagemaker.is_using_sagemaker():
+        if (
+            not self._settings.sagemaker_disable
+            and not disable_sagemaker
+            and sagemaker.is_using_sagemaker()
+        ):
             self._logger.info("Loading SageMaker settings")
-            sagemaker.set_global_settings(s)
+            sagemaker.set_global_settings(self._settings)
 
         # load settings from the passed init/setup settings
-        if settings:
-            s.update_from_settings(settings)
+        if overrides:
+            self._settings.update_from_settings(overrides)
 
-        return s
+        wandb.termsetup(self._settings, None)
 
-    def _update(self, settings: Settings | None = None) -> None:
-        if not settings:
-            return
-        self._settings.update_from_settings(settings)
+    def _update(self, settings: Settings | None) -> None:
+        """Update settings, initializing them if necessary.
 
-    def _update_user_settings(self) -> None:
+        Args:
+            settings: Overrides to apply, if any.
+        """
+        if not self._settings:
+            system_settings_path = settings.settings_system if settings else None
+            disable_sagemaker = settings.sagemaker_disable if settings else False
+            self._load_settings(
+                system_settings_path=system_settings_path,
+                disable_sagemaker=disable_sagemaker,
+                overrides=settings,
+            )
+
+        # This is 'elif' because load_settings already applies overrides.
+        elif settings:
+            self._settings.update_from_settings(settings)
+
+    def update_user_settings(self) -> None:
         # Get rid of cached results to force a refresh.
         self._server = None
         user_settings = self._load_user_settings()
         if user_settings is not None:
-            self._settings.update_from_dict(user_settings)
+            self.settings.update_from_dict(user_settings)
 
     def _early_logger_flush(self, new_logger: Logger) -> None:
         if self._logger is new_logger:
@@ -238,6 +298,22 @@ class _WandbSetup:
 
     @property
     def settings(self) -> wandb_settings.Settings:
+        """The global wandb settings.
+
+        Initializes settings if they have not yet been loaded.
+        """
+        if not self._settings:
+            self._load_settings(
+                system_settings_path=None,
+                disable_sagemaker=False,
+            )
+            assert self._settings
+
+        return self._settings
+
+    @property
+    def settings_if_loaded(self) -> wandb_settings.Settings | None:
+        """The global wandb settings, or None if not yet loaded."""
         return self._settings
 
     def _get_entity(self) -> str | None:
@@ -262,7 +338,7 @@ class _WandbSetup:
     @property
     def viewer(self) -> dict[str, Any]:
         if self._server is None:
-            self._server = server.Server(settings=self._settings)
+            self._server = server.Server(settings=self.settings)
 
         return self._server.viewer
 
@@ -282,36 +358,36 @@ class _WandbSetup:
 
         return user_settings
 
-    def _setup(self) -> None:
-        sweep_path = self._settings.sweep_param_path
+    @property
+    def config(self) -> dict:
+        sweep_path = self.settings.sweep_param_path
         if sweep_path:
             self._sweep_config = config_util.dict_from_config_file(
                 sweep_path, must_exist=True
             )
 
+        config = {}
+
         # if config_paths was set, read in config dict
-        if self._settings.config_paths:
+        if self.settings.config_paths:
             # TODO(jhr): handle load errors, handle list of files
-            for config_path in self._settings.config_paths:
+            for config_path in self.settings.config_paths:
                 config_dict = config_util.dict_from_config_file(config_path)
-                if config_dict is None:
-                    continue
-                if self._config is not None:
-                    self._config.update(config_dict)
-                else:
-                    self._config = config_dict
+                if config_dict:
+                    config.update(config_dict)
+
+        return config
 
     def _teardown(self, exit_code: int | None = None) -> None:
         import_hooks.unregister_all_post_import_hooks()
 
-        if not self._connection:
-            return
+        if self._connection:
+            internal_exit_code = self._connection.teardown(exit_code or 0)
+        else:
+            internal_exit_code = None
 
-        # Reset to None so that setup() creates a new connection.
-        connection = self._connection
-        self._connection = None
+        self._asyncer.join()
 
-        internal_exit_code = connection.teardown(exit_code or 0)
         if internal_exit_code not in (None, 0):
             sys.exit(internal_exit_code)
 
@@ -322,7 +398,10 @@ class _WandbSetup:
 
         from wandb.sdk.lib.service import service_connection
 
-        self._connection = service_connection.connect_to_service(self._settings)
+        self._connection = service_connection.connect_to_service(
+            self._asyncer,
+            self.settings,
+        )
         return self._connection
 
     def assert_service(self) -> ServiceConnection:
@@ -343,6 +422,8 @@ _singleton: _WandbSetup | None = None
 The value is invalid and must not be used if `os.getpid() != _singleton._pid`.
 """
 
+_singleton_lock = threading.Lock()
+
 
 def singleton() -> _WandbSetup:
     """The W&B singleton for the current process.
@@ -351,29 +432,14 @@ def singleton() -> _WandbSetup:
     process) creates the singleton, and all subsequent calls return it
     until teardown(). This does not start the service process.
     """
-    return _setup(start_service=False)
-
-
-def singleton_if_setup() -> _WandbSetup | None:
-    """The W&B singleton for the current process or None if it isn't set up.
-
-    Always prefer singleton() over this function.
-
-    Unlike singleton(), this never creates the singleton and therefore never
-    initializes global settings from the environment. This is useful only
-    during tests, which may modify the environment after having imported wandb
-    and called certain functions.
-    """
-    if _singleton and _singleton._pid == os.getpid():
-        return _singleton
-    else:
-        return None
+    return _setup(start_service=False, load_settings=False)
 
 
 @wb_logging.log_to_all_runs()
 def _setup(
     settings: Settings | None = None,
     start_service: bool = True,
+    load_settings: bool = True,
 ) -> _WandbSetup:
     """Set up library context.
 
@@ -384,18 +450,31 @@ def _setup(
             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".
+        load_settings: Whether to load settings from the environment
+            if creating a new singleton. If False, then settings and
+            start_service must be None.
     """
     global _singleton
 
+    if not load_settings and settings:
+        raise ValueError("Cannot pass settings if load_settings is False.")
+    if not load_settings and start_service:
+        raise ValueError("Cannot use start_service if load_settings is False.")
+
     pid = os.getpid()
+    with _singleton_lock:
+        if _singleton and _singleton._pid == pid:
+            current_singleton = _singleton
+        else:
+            current_singleton = _WandbSetup(pid=pid)
+
+        if load_settings:
+            current_singleton._update(settings)
 
-    if _singleton and _singleton._pid == pid:
-        _singleton._update(settings=settings)
-    else:
-        _singleton = _WandbSetup(settings=settings, pid=pid)
+        if start_service and not current_singleton.settings._noop:
+            current_singleton.ensure_service()
 
-    if start_service and not _singleton.settings._noop:
-        _singleton.ensure_service()
+        _singleton = current_singleton
 
     return _singleton
 
@@ -419,49 +498,49 @@ def setup(settings: Settings | None = None) -> _WandbSetup:
             overridden by subsequent `wandb.init()` calls.
 
     Example:
-        ```python
-        import multiprocessing
+    ```python
+    import multiprocessing
 
-        import wandb
+    import wandb
 
 
-        def run_experiment(params):
-            with wandb.init(config=params):
-                # Run experiment
-                pass
+    def run_experiment(params):
+        with wandb.init(config=params):
+            # Run experiment
+            pass
 
 
-        if __name__ == "__main__":
-            # Start backend and set global config
-            wandb.setup(settings={"project": "my_project"})
+    if __name__ == "__main__":
+        # Start backend and set global config
+        wandb.setup(settings={"project": "my_project"})
 
-            # Define experiment parameters
-            experiment_params = [
-                {"learning_rate": 0.01, "epochs": 10},
-                {"learning_rate": 0.001, "epochs": 20},
-            ]
+        # Define experiment parameters
+        experiment_params = [
+            {"learning_rate": 0.01, "epochs": 10},
+            {"learning_rate": 0.001, "epochs": 20},
+        ]
 
-            # Start multiple processes, each running a separate experiment
-            processes = []
-            for params in experiment_params:
-                p = multiprocessing.Process(target=run_experiment, args=(params,))
-                p.start()
-                processes.append(p)
+        # Start multiple processes, each running a separate experiment
+        processes = []
+        for params in experiment_params:
+            p = multiprocessing.Process(target=run_experiment, args=(params,))
+            p.start()
+            processes.append(p)
 
-            # Wait for all processes to complete
-            for p in processes:
-                p.join()
+        # Wait for all processes to complete
+        for p in processes:
+            p.join()
 
-            # Optional: Explicitly shut down the backend
-            wandb.teardown()
-        ```
+        # Optional: Explicitly shut down the backend
+        wandb.teardown()
+    ```
     """
     return _setup(settings=settings)
 
 
 @wb_logging.log_to_all_runs()
 def teardown(exit_code: int | None = None) -> None:
-    """Waits for wandb to finish and frees resources.
+    """Waits for W&B to finish and frees resources.
 
     Completes any runs that were not explicitly finished
     using `run.finish()` and waits for all data to be uploaded.
@@ -473,8 +552,9 @@ def teardown(exit_code: int | None = None) -> None:
     """
     global _singleton
 
-    orig_singleton = _singleton
-    _singleton = None
+    with _singleton_lock:
+        orig_singleton = _singleton
+        _singleton = None
 
-    if orig_singleton:
-        orig_singleton._teardown(exit_code=exit_code)
+        if orig_singleton:
+            orig_singleton._teardown(exit_code=exit_code)

wandb/sdk/wandb_sweep.py

@@ -45,11 +45,12 @@ def sweep(
     Make note the unique identifier, `sweep_id`, that is returned.
     At a later step provide the `sweep_id` to a sweep agent.
 
+    See [Sweep configuration structure](https://docs.wandb.ai/guides/sweeps/define-sweep-configuration)
+    for information on how to define your sweep.
+
     Args:
       sweep: The configuration of a hyperparameter search.
-        (or configuration generator). See
-        [Sweep configuration structure](https://docs.wandb.ai/guides/sweeps/define-sweep-configuration)
-        for information on how to define your sweep.
+        (or configuration generator).
         If you provide a callable, ensure that the callable does
         not take arguments and that it returns a dictionary that
         conforms to the W&B sweep config spec.
@@ -64,7 +65,7 @@ def sweep(
       prior_runs: The run IDs of existing runs to add to this sweep.
 
     Returns:
-      sweep_id: str. A unique identifier for the sweep.
+      str: A unique identifier for the sweep.
     """
     if callable(sweep):
         sweep = sweep()
@@ -99,16 +100,16 @@ def controller(
 ) -> "_WandbController":
     """Public sweep controller constructor.
 
-    Usage:
-        ```python
-        import wandb
+    Examples:
+    ```python
+    import wandb
 
-        tuner = wandb.controller(...)
-        print(tuner.sweep_config)
-        print(tuner.sweep_id)
-        tuner.configure_search(...)
-        tuner.configure_stopping(...)
-        ```
+    tuner = wandb.controller(...)
+    print(tuner.sweep_config)
+    print(tuner.sweep_id)
+    tuner.configure_search(...)
+    tuner.configure_stopping(...)
+    ```
 
     """
     from ..wandb_controller import _WandbController

wandb/sdk/wandb_watch.py

@@ -17,15 +17,13 @@ from .lib import telemetry
 if TYPE_CHECKING:
     import torch  # type: ignore [import-not-found]
 
-    from wandb.sdk.wandb_run import Run
-
 logger = logging.getLogger("wandb")
 
 _global_watch_idx = 0
 
 
 def _watch(
-    run: Run,
+    run: wandb.Run,
     models: torch.nn.Module | Sequence[torch.nn.Module],
     criterion: torch.F | None = None,
     log: Literal["gradients", "parameters", "all"] | None = "gradients",
@@ -39,7 +37,7 @@ def _watch(
     extended to support arbitrary machine learning models in the future.
 
     Args:
-        run (wandb.sdk.wandb_run.Run): The run object to log to.
+        run (wandb.Run): The run object to log to.
         models (Union[torch.nn.Module, Sequence[torch.nn.Module]]):
             A single model or a sequence of models to be monitored.
         criterion (Optional[torch.F]):
@@ -122,12 +120,12 @@ def _watch(
 
 
 def _unwatch(
-    run: Run, models: torch.nn.Module | Sequence[torch.nn.Module] | None = None
+    run: wandb.Run, models: torch.nn.Module | Sequence[torch.nn.Module] | None = None
 ) -> None:
     """Remove pytorch model topology, gradient and parameter hooks.
 
     Args:
-        run (wandb.sdk.wandb_run.Run):
+        run (wandb.Run):
             The run object to log to.
         models (torch.nn.Module | Sequence[torch.nn.Module]):
             Optional list of pytorch models that have had watch called on them

wandb/sync/sync.py

@@ -54,6 +54,7 @@ class SyncThread(threading.Thread):
         log_path=None,
         append=None,
         skip_console=None,
+        replace_tags=None,
     ):
         threading.Thread.__init__(self)
         # mark this process as internal
@@ -71,6 +72,7 @@ class SyncThread(threading.Thread):
         self._log_path = log_path
         self._append = append
         self._skip_console = skip_console
+        self._replace_tags = replace_tags or {}
 
         self._tmp_dir = tempfile.TemporaryDirectory()
         atexit.register(self._tmp_dir.cleanup)
@@ -94,6 +96,11 @@ class SyncThread(threading.Thread):
                 pb.run.entity = self._entity
             if self._job_type:
                 pb.run.job_type = self._job_type
+            # Replace tags if specified
+            if self._replace_tags:
+                new_tags = [self._replace_tags.get(tag, tag) for tag in pb.run.tags]
+                pb.run.ClearField("tags")
+                pb.run.tags.extend(new_tags)
             pb.control.req_resp = True
         elif record_type in ("output", "output_raw") and self._skip_console:
             return pb, exit_pb, True
@@ -338,6 +345,7 @@ class SyncManager:
         log_path=None,
         append=None,
         skip_console=None,
+        replace_tags=None,
     ):
         self._sync_list = []
         self._thread = None
@@ -353,6 +361,7 @@ class SyncManager:
         self._log_path = log_path
         self._append = append
         self._skip_console = skip_console
+        self._replace_tags = replace_tags or {}
 
     def status(self):
         pass
@@ -376,6 +385,7 @@ class SyncManager:
             log_path=self._log_path,
             append=self._append,
             skip_console=self._skip_console,
+            replace_tags=self._replace_tags,
         )
         self._thread.start()