wandb 0.21.1__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
 
@@ -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/util.py

@@ -702,7 +702,7 @@ def json_friendly_val(val: Any) -> Any:
         return converted
     if is_dataclass(val) and not isinstance(val, type):
         converted = asdict(val)
-        return converted
+        return json_friendly_val(converted)
     else:
         if val.__class__.__module__ not in ("builtins", "__builtin__"):
             val = str(val)

{wandb-0.21.1.dist-info → wandb-0.21.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: wandb
-Version: 0.21.1
+Version: 0.21.2
 Summary: A CLI and library for interacting with the Weights & Biases API.
 Project-URL: Source, https://github.com/wandb/wandb
 Project-URL: Bug Reports, https://github.com/wandb/wandb/issues