"""Defines wandb.init() and associated classes and methods. `wandb.init()` indicates the beginning of a new run. In an ML training pipeline, you could add `wandb.init()` to the beginning of your training script as well as your evaluation script, and each step would be tracked as a run in W&B. For more on using `wandb.init()`, including code snippets, check out our [guide and FAQs](https://docs.wandb.ai/guides/track/launch). """ from __future__ import annotations import contextlib import dataclasses import functools import json import logging import os import pathlib import platform import sys import tempfile import time from collections.abc import Iterable, Iterator, Sequence from typing_extensions import Any, Literal, Protocol, Self import wandb import wandb.env from wandb import env, trigger from wandb.analytics import get_sentry 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, weave from wandb.proto.wandb_telemetry_pb2 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 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.deprecation import UNSET, DoNotSet, warn_and_record_deprecation from .mailbox import wait_with_progress from .wandb_helper import parse_config from .wandb_run import Run, TeardownHook, TeardownStage from .wandb_settings import Settings def _huggingface_version() -> str | None: if "transformers" in sys.modules: trans = wandb.util.get_module("transformers") if hasattr(trans, "__version__"): return str(trans.__version__) return None def _handle_launch_config(settings: Settings) -> dict[str, Any]: launch_run_config: dict[str, Any] = {} if not settings.launch: return launch_run_config if os.environ.get("WANDB_CONFIG") is not None: try: launch_run_config = json.loads(os.environ.get("WANDB_CONFIG", "{}")) except (ValueError, SyntaxError): wandb.termwarn("Malformed WANDB_CONFIG, using original config") elif settings.launch_config_path and os.path.exists(settings.launch_config_path): with open(settings.launch_config_path) as fp: launch_config = json.loads(fp.read()) launch_run_config = launch_config.get("overrides", {}).get("run_config") else: i = 0 chunks = [] while True: key = f"WANDB_CONFIG_{i}" if key in os.environ: chunks.append(os.environ[key]) i += 1 else: break if len(chunks) > 0: config_string = "".join(chunks) try: launch_run_config = json.loads(config_string) except (ValueError, SyntaxError): wandb.termwarn("Malformed WANDB_CONFIG, using original config") return launch_run_config @dataclasses.dataclass(frozen=True) class _ConfigParts: base_no_artifacts: dict[str, Any] """The run config passed to `init()` minus any artifact-valued keys.""" sweep_no_artifacts: dict[str, Any] """The config loaded as part of a sweep minus any artifact-valued keys.""" launch_no_artifacts: dict[str, Any] """The config loaded as part of Launch minus any artifact-valued keys.""" artifacts: dict[str, Any] """Artifact keys removed from config dictionaries. Due to implementation details of how a Run is constructed, artifacts must be inserted into its config after initialization. """ 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, wl: wandb_setup._WandbSetup, telemetry: telemetry.TelemetryRecord, ) -> None: self._wl = wl self._telemetry = telemetry """Telemetry gathered before creating a run. After the run is created, `telemetry.context()` is used instead. """ self.kwargs = None self.run: Run | None = None self.backend: Backend | None = None self._teardown_hooks: list[TeardownHook] = [] self.notebook: wandb.jupyter.Notebook | None = None self.deprecated_features_used: list[tuple[Deprecated, str]] = [] @property def _logger(self) -> wandb_setup.Logger: return self._wl._get_logger() def maybe_login(self, init_settings: Settings) -> None: """Log in if we are not creating an offline or disabled run. This may change the W&B singleton settings. Args: init_settings: Settings passed to `wandb.init()` or set via keyword arguments. """ # Allow settings passed to init() to override inferred values. # # Calling login() may change settings on the singleton, # so these may not be the final run settings. run_settings = self._wl.settings.model_copy() run_settings.update_from_settings(init_settings) # NOTE: _noop or _offline can become true after _login(). # _noop happens if _login hits a timeout. # _offline can be selected by the user at the login prompt. if run_settings._noop or run_settings._offline: return # Only pass an explicit key when the key was provided directly # to ensure correct messaging in _login(). explicit_key = init_settings.api_key wandb_login._login( host=run_settings.base_url, force=run_settings.force, _silent=run_settings.quiet or run_settings.silent, key=explicit_key, update_api_key=explicit_key is None, ) def warn_env_vars_change_after_setup(self) -> _PrinterCallback: """Warn if environment variables changed after `wandb.setup()`. Returns: A callback to print any generated warnings. """ if not self._wl.did_environment_change(): 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"{run_printer.link(url_registry.url('wandb-init'), 'the W&B docs')}." ) run_printer.display(line, level="warn") return print_warning def clear_run_path_if_sweep_or_launch( self, init_settings: Settings, ) -> _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 = "" if self._wl.settings.sweep_id: when_doing_thing = "when running a sweep" elif self._wl.settings.launch: when_doing_thing = "when running from a wandb launch context" if not when_doing_thing: return _noop_printer_callback() warnings = [] def warn(key: str, value: str) -> None: warnings.append(f"Ignoring {key} {value!r} {when_doing_thing}.") if init_settings.project is not None: warn("project", init_settings.project) init_settings.project = None if init_settings.entity is not None: warn("entity", init_settings.entity) init_settings.entity = None if init_settings.run_id is not None: warn("run_id", init_settings.run_id) init_settings.run_id = None 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. """ 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() # Apply settings from wandb.init() call. settings.update_from_settings(init_settings) # Infer the run ID from SageMaker. if ( (not settings.sagemaker_disable) and sagemaker.is_using_sagemaker() and sagemaker.set_run_id(settings) ): self._logger.info("set run ID and group based on SageMaker") self._telemetry.feature.sagemaker = True # get status of code saving before applying user settings save_code_pre_user_settings = settings.save_code if not settings._offline and not settings._noop: user_settings = self._wl._load_user_settings() if user_settings is not None: settings.update_from_dict(user_settings) # ensure that user settings don't set saving to true # if user explicitly set these to false in UI if save_code_pre_user_settings is False: settings.save_code = False # TODO: remove this once we refactor the client. This is a temporary # fix to make sure that we use the same project name for wandb-core. # The reason this is not going through the settings object is to # avoid failure cases in other parts of the code that will be # removed with the switch to wandb-core. if settings.project is None: settings.project = wandb.util.auto_project_name(settings.program) 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, _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. Returns `None` if the file does not exist or is not in a valid format. Args: resume_file: The file path to use for resume='auto' mode. """ if not resume_file.exists(): return None with resume_file.open() as f: try: return json.load(f)["run_id"] except json.JSONDecodeError as e: self._logger.exception( f"could not decode {resume_file}, ignoring", exc_info=e, ) return None except KeyError: self._logger.exception( f"resume file at {resume_file} did not store a run_id" ) return None def _save_autoresume_run_id( self, *, resume_file: pathlib.Path, run_id: str, ) -> None: """Write the run ID to the auto-resume file.""" resume_file.parent.mkdir(exist_ok=True) with resume_file.open("w") as f: json.dump({"run_id": run_id}, f) def set_run_id(self, settings: Settings) -> None: """Set the run ID and possibly save it to the auto-resume file. After this, `settings.run_id` is guaranteed to be set. If a `resume_from` is provided and `run_id` is not set, initialize `run_id` with the `resume_from` run's `run_id`. Args: settings: The run's settings derived from the environment and explicit values passed to `wandb.init()`. """ if settings.resume == "auto" and settings.resume_fname: resume_path = pathlib.Path(settings.resume_fname) else: resume_path = None if resume_path: previous_id = self._load_autoresume_run_id(resume_path) if not previous_id: pass elif settings.run_id is None: self._logger.info(f"loaded run ID from {resume_path}") settings.run_id = previous_id elif settings.run_id != previous_id: wandb.termwarn( f"Ignoring ID {previous_id} loaded due to resume='auto'" f" because the run ID is set to {settings.run_id}.", ) # If no run ID was inferred, explicitly set, or loaded from an # auto-resume file, then we generate a new ID. if settings.run_id is None: # If resume_from is provided and run_id is not already set, # initialize run_id with the value from resume_from. if settings.resume_from: settings.run_id = settings.resume_from.run else: settings.run_id = runid.generate_id() if resume_path: self._save_autoresume_run_id( resume_file=resume_path, run_id=settings.run_id, ) def set_sync_dir_suffix(self, settings: Settings) -> None: """Add a suffix to sync_dir if it already exists. The sync_dir uses a timestamp with second-level precision which can result in conflicts if a run with the same ID is initialized within the same second. This is most likely to happen in tests. This can't prevent conflicts from multiple processes attempting to create a wandb run simultaneously. Args: settings: Fully initialized settings other than the x_sync_dir_suffix setting which will be modified. """ index = 1 while pathlib.Path(settings.sync_dir).exists(): settings.x_sync_dir_suffix = f"{index}" index += 1 def make_run_config( self, settings: Settings, config: dict | str | None = None, config_exclude_keys: list[str] | None = None, config_include_keys: list[str] | None = None, ) -> _ConfigParts: """Construct the run's config. Args: settings: The run's finalized settings. config: The config passed to `init()`. config_exclude_keys: Deprecated. Keys to filter out from `config`. config_include_keys: Deprecated. Keys to include from `config`. Returns: Initial values for the run's config. """ if config_exclude_keys: self.deprecated_features_used.append( ( Deprecated(init__config_exclude_keys=True), "config_exclude_keys is deprecated. Use" " `config=wandb.helper.parse_config(config_object," " exclude=('key',))` instead.", ) ) if config_include_keys: self.deprecated_features_used.append( ( Deprecated(init__config_include_keys=True), "config_include_keys is deprecated. Use" " `config=wandb.helper.parse_config(config_object," " include=('key',))` instead.", ) ) config = parse_config( config or dict(), include=config_include_keys, exclude=config_exclude_keys, ) result = _ConfigParts( base_no_artifacts=dict(), sweep_no_artifacts=dict(), launch_no_artifacts=dict(), artifacts=dict(), ) if not settings.sagemaker_disable and sagemaker.is_using_sagemaker(): sagemaker_config = sagemaker.parse_sm_config() self._split_artifacts_from_config( sagemaker_config, config_target=result.base_no_artifacts, artifacts=result.artifacts, ) self._telemetry.feature.sagemaker = True if self._wl.config: self._split_artifacts_from_config( self._wl.config, config_target=result.base_no_artifacts, artifacts=result.artifacts, ) if config and isinstance(config, dict): self._split_artifacts_from_config( config, config_target=result.base_no_artifacts, artifacts=result.artifacts, ) if self._wl._sweep_config: self._split_artifacts_from_config( self._wl._sweep_config, config_target=result.sweep_no_artifacts, artifacts=result.artifacts, ) if launch_config := _handle_launch_config(settings): self._split_artifacts_from_config( launch_config, config_target=result.launch_no_artifacts, artifacts=result.artifacts, ) wandb_internal = result.base_no_artifacts.setdefault("_wandb", dict()) if settings.save_code and settings.program_relpath: wandb_internal["code_path"] = paths.LogicalPath( os.path.join("code", settings.program_relpath) ) if settings.fork_from is not None: wandb_internal["branch_point"] = { "run_id": settings.fork_from.run, "step": settings.fork_from.value, } if settings.resume_from is not None: wandb_internal["branch_point"] = { "run_id": settings.resume_from.run, "step": settings.resume_from.value, } return result def teardown(self) -> None: # TODO: currently this is only called on failed wandb.init attempts # normally this happens on the run object self._logger.info("tearing down wandb.init") for hook in self._teardown_hooks: hook.call() def _split_artifacts_from_config( self, config_source: dict, config_target: dict, artifacts: dict, ) -> None: for k, v in config_source.items(): if _is_artifact_representation(v): artifacts[k] = v else: config_target.setdefault(k, v) def _safe_symlink( self, base: str, target: str, name: str, delete: bool = False ) -> None: # TODO(jhr): do this with relpaths, but i can't figure it out on no sleep if not hasattr(os, "symlink"): return pid = os.getpid() tmp_name = os.path.join(base, f"{name}.{pid}") if delete: try: os.remove(os.path.join(base, name)) except OSError: pass target = os.path.relpath(target, base) try: os.symlink(target, tmp_name) os.rename(tmp_name, os.path.join(base, name)) except OSError: pass def _pre_run_cell_hook(self, *args, **kwargs) -> None: """Hook for the IPython pre_run_cell event. 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) if self.backend.interface is not None: self._logger.info("pausing backend") self.backend.interface.publish_pause() 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 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) 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) 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) ipython.display_pub._orig_publish = ipython.display_pub.publish 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: ipython.display_pub._orig_publish(data, metadata=metadata, **kwargs) assert self.notebook is not None self.notebook.save_display( ipython.execution_count, {"data": data, "metadata": metadata} ) ipython.display_pub.publish = publish @contextlib.contextmanager def setup_run_log_directory(self, settings: Settings) -> Iterator[None]: """Set up the run's log directory. This is a context manager that closes and unregisters the log handler in case of an uncaught exception, so that future logged messages do not modify this run's log file. """ filesystem.mkdir_exists_ok(os.path.dirname(settings.log_user)) filesystem.mkdir_exists_ok(os.path.dirname(settings.log_internal)) filesystem.mkdir_exists_ok(os.path.dirname(settings.sync_file)) filesystem.mkdir_exists_ok(settings.files_dir) filesystem.mkdir_exists_ok(settings._tmp_code_dir) if settings.symlink: self._safe_symlink( os.path.dirname(settings.sync_symlink_latest), os.path.dirname(settings.sync_file), os.path.basename(settings.sync_symlink_latest), delete=True, ) self._safe_symlink( os.path.dirname(settings.log_symlink_user), settings.log_user, os.path.basename(settings.log_symlink_user), delete=True, ) self._safe_symlink( os.path.dirname(settings.log_symlink_internal), settings.log_internal, os.path.basename(settings.log_symlink_internal), delete=True, ) assert settings.run_id handler = wb_logging.add_file_handler( settings.run_id, pathlib.Path(settings.log_user), ) if env.is_debug(): handler.setLevel(logging.DEBUG) disposed = False def dispose_handler() -> None: nonlocal disposed if not disposed: disposed = True logging.getLogger("wandb").removeHandler(handler) handler.close() try: self._teardown_hooks.append( TeardownHook( call=dispose_handler, stage=TeardownStage.LATE, ) ) self._wl._early_logger_flush(logging.getLogger("wandb")) self._logger.info(f"Logging user logs to {settings.log_user}") self._logger.info(f"Logging internal logs to {settings.log_internal}") yield except Exception: dispose_handler() raise def make_disabled_run(self, config: _ConfigParts) -> Run: """Returns a Run-like object where all methods are no-ops. This method is used when the `mode` setting is set to "disabled", such as by wandb.init(mode="disabled") or by setting the WANDB_MODE environment variable to "disabled". It creates a Run object that mimics the behavior of a normal Run but doesn't communicate with the W&B servers. The returned Run object has all expected attributes and methods, but they are no-op versions that don't perform any actual logging or communication. """ run_id = runid.generate_id() drun = Run( settings=Settings( mode="disabled", root_dir=tempfile.gettempdir(), run_id=run_id, run_tags=tuple(), run_notes=None, run_group=None, run_name=f"dummy-{run_id}", project="dummy", entity="dummy", ) ) # config, summary, and metadata objects drun._config = wandb.sdk.wandb_config.Config() drun._config.update(config.sweep_no_artifacts) drun._config.update(config.base_no_artifacts) drun.summary = SummaryDisabled() # type: ignore # methods drun.log = lambda data, *_, **__: drun.summary.update(data) # type: ignore[method-assign] drun.finish = lambda *_, **__: module.unset_globals() # type: ignore[method-assign] drun.join = drun.finish # type: ignore[method-assign] drun.define_metric = lambda *_, **__: wandb.sdk.wandb_metric.Metric("dummy") # type: ignore[method-assign] drun.save = lambda *_, **__: False # type: ignore[method-assign] for symbol in ( "alert", "finish_artifact", "get_project_url", "get_sweep_url", "get_url", "link_artifact", "link_model", "use_artifact", "log_code", "log_model", "use_model", "mark_preempting", "restore", "status", "watch", "unwatch", "upsert_artifact", "_finish", ): 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.""" def __getattr__(self, _: str) -> Self: return self def __call__(self, *_: Any, **__: Any) -> Self: return self class _ChainableNoOpField: # This is used to chain arbitrary attributes and method calls. # For example, `run.log_artifact().state` will work in disabled mode. def __init__(self) -> None: self._value = None def __set__(self, instance: Any, value: Any) -> None: self._value = value def __get__(self, instance: Any, owner: type) -> Any: return _ChainableNoOp() if (self._value is None) else self._value def __call__(self, *args: Any, **kwargs: Any) -> _ChainableNoOp: return _ChainableNoOp() drun.log_artifact = _ChainableNoOpField() # type: ignore # attributes drun._start_time = time.time() drun._starting_step = 0 drun._step = 0 drun._attach_id = None drun._backend = None # set the disabled run as the global run module.set_global( run=drun, config=drun.config, log=drun.log, summary=drun.summary, save=drun.save, use_artifact=drun.use_artifact, log_artifact=drun.log_artifact, define_metric=drun.define_metric, alert=drun.alert, watch=drun.watch, unwatch=drun.unwatch, ) return drun def init( # noqa: C901 self, settings: Settings, config: _ConfigParts, run_printer: printer.Printer, ) -> Run: self._logger.info("calling init triggers") trigger.call("on_init") assert self._wl is not None self._logger.info( f"wandb.init called with sweep_config: {config.sweep_no_artifacts}" f"\nconfig: {config.base_no_artifacts}" ) 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()) ): 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" ) elif settings.resume == "must": raise wandb.Error( "Cannot resume a run while another run is active." " You must either finish it using run.finish()," " or use reinit='create_new' when calling wandb.init()." ) else: 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." ) with telemetry.context(run=previous_run) as tel: tel.feature.init_return_run = True return previous_run self._logger.info("starting backend") service = self._wl.ensure_service() self._logger.info("sending inform_init request") service.inform_init( settings=settings.to_proto(), run_id=settings.run_id, # type: ignore ) backend = Backend(settings=settings, service=service) backend.ensure_launched() self._logger.info("backend started and connected") run = Run( config=config.base_no_artifacts, settings=settings, sweep_config=config.sweep_no_artifacts, launch_config=config.launch_no_artifacts, ) # Populate initial telemetry with telemetry.context(run=run, obj=self._telemetry) as tel: tel.cli_version = wandb.__version__ tel.python_version = platform.python_version() tel.platform = f"{platform.system()}-{platform.machine()}".lower() hf_version = _huggingface_version() if hf_version: tel.huggingface_version = hf_version if settings._jupyter: tel.env.jupyter = True if settings._ipython: tel.env.ipython = True if settings._colab: tel.env.colab = True if settings._kaggle: tel.env.kaggle = True if settings._windows: tel.env.windows = True if settings.launch: tel.feature.launch = True for module_name in telemetry.list_telemetry_imports(only_imported=True): setattr(tel.imports_init, module_name, True) if os.environ.get("PEX"): tel.env.pex = True if settings._aws_lambda: tel.env.aws_lambda = True if settings.x_flow_control_disabled: tel.feature.flow_control_disabled = True if settings.x_flow_control_custom: tel.feature.flow_control_custom = True if settings._shared: wandb.termwarn( "The `shared` mode feature is experimental and may change. " "Please contact support@wandb.com for guidance and to report any issues." ) 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 if not settings.label_disable: if self.notebook: run._label_probe_notebook(self.notebook) else: run._label_probe_main() for deprecated_feature, msg in self.deprecated_features_used: warn_and_record_deprecation( feature=deprecated_feature, message=msg, run=run, ) self._logger.info("updated telemetry") run._set_library(self._wl) run._set_backend(backend) run._set_teardown_hooks(self._teardown_hooks) assert backend.interface backend.interface.publish_header() # Using GitRepo() blocks & can be slow, depending on user's current git setup. # We don't want to block run initialization/start request, so populate run's git # info beforehand. if not (settings.disable_git or settings.x_disable_machine_info): run._populate_git_info() if settings._offline and settings.resume: wandb.termwarn( "`resume` will be ignored since W&B syncing is set to `offline`. " f"Starting a new run with run id {run.id}." ) error: wandb.Error | None = None timeout = settings.init_timeout self._logger.info( f"communicating run to backend with {timeout} second timeout", ) run_init_handle = backend.interface.deliver_run(run) try: with progress.progress_printer( run_printer, default_text="Waiting for wandb.init()...", ) as progress_printer: result = wait_with_progress( run_init_handle, timeout=timeout, display_progress=functools.partial( progress.loop_printing_operation_stats, progress_printer, backend.interface, ), ) except TimeoutError: # 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 # the two causes here. raise CommError( f"Run initialization has timed out after {timeout} sec." + " Please try increasing the timeout with the `init_timeout`" + " setting: `wandb.init(settings=wandb.Settings(init_timeout=120))`." ) from None assert result.run_result if error := ProtobufErrorHandler.to_exception(result.run_result.error): raise error if not result.run_result.HasField("run"): raise Error("Assertion failed: run_result is missing the run field") if result.run_result.run.resumed: self._logger.info("run resumed") with telemetry.context(run=run) as tel: tel.feature.resumed = result.run_result.run.resumed run._set_run_obj(result.run_result.run) self._logger.info("starting run threads in backend") assert backend.interface run_start_handle = backend.interface.deliver_run_start(run) try: # TODO: add progress to let user know we are doing something run_start_handle.wait_or(timeout=30) except TimeoutError: pass backend.interface.publish_probe_system_info() assert self._wl is not None self.run = run run._handle_launch_artifact_overrides() if ( settings.launch and settings.launch_config_path and os.path.exists(settings.launch_config_path) ): run.save(settings.launch_config_path) # put artifacts in run config here # since doing so earlier will cause an error # as the run is not upserted for k, v in config.artifacts.items(): run.config.update({k: v}, allow_val_change=True) job_artifact = run._launch_artifact_mapping.get( wandb.util.LAUNCH_JOB_ARTIFACT_SLOT_NAME ) if job_artifact: 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 def _attach( attach_id: str | None = None, run_id: str | None = None, *, run: Run | None = None, ) -> Run | None: """Attach to a run currently executing in another process/thread. Args: attach_id: (str, optional) The id of the run or an attach identifier that maps to a run. run_id: (str, optional) The id of the run to attach to. run: (Run, optional) The run instance to attach """ attach_id = attach_id or run_id if not ((attach_id is None) ^ (run is None)): raise UsageError("Either (`attach_id` or `run_id`) or `run` must be specified") attach_id = attach_id or (run._attach_id if run else None) if attach_id is None: raise UsageError( "Either `attach_id` or `run_id` must be specified or `run` must have `_attach_id`" ) _wl = wandb_setup.singleton() logger = _wl._get_logger() service = _wl.ensure_service() try: attach_settings = service.inform_attach(attach_id=attach_id) except Exception as e: raise UsageError(f"Unable to attach to run {attach_id}") from e settings = _wl.settings.model_copy() settings.update_from_dict( { "run_id": attach_id, "x_start_time": attach_settings.x_start_time.value, "mode": attach_settings.mode.value, } ) # TODO: consolidate this codepath with wandb.init() backend = Backend(settings=settings, service=service) backend.ensure_launched() logger.info("attach backend started and connected") if run is None: run = Run(settings=settings) else: run._init(settings=settings) run._set_library(_wl) run._set_backend(backend) assert backend.interface attach_handle = backend.interface.deliver_attach(attach_id) 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}") 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: return from wandb.integration import gym gym.monitor() def _monkeypatch_tensorboard() -> None: """Patch TensorBoard to log to the global `wandb.run`.""" if len(wandb.patched["tensorboard"]) > 0: return from wandb.integration import tensorboard as tb_module tb_module.patch() def try_create_root_dir(settings: Settings) -> None: """Try to create the root directory specified in settings. If creation fails due to permissions or other errors, falls back to using the system temp directory. Args: settings: The runs settings containing root_dir configuration. This function may update the root_dir to a temporary directory if the parent directory is not writable. """ fallback_to_temp_dir = False try: os.makedirs(settings.root_dir, exist_ok=True) except OSError: wandb.termwarn( f"Unable to create root directory {settings.root_dir}", repeat=False, ) 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 entity: str | None = None, project: str | None = None, dir: StrPath | None = None, id: str | None = None, name: str | None = None, notes: str | None = None, tags: Sequence[str] | None = None, config: dict[str, Any] | str | None = None, config_exclude_keys: list[str] | None = None, config_include_keys: list[str] | None = None, allow_val_change: bool | None = None, group: str | None = None, job_type: str | None = None, mode: Literal["online", "offline", "disabled", "shared"] | None = None, force: bool | None = None, reinit: ( bool | Literal[ None, "default", "return_previous", "finish_previous", "create_new", ] ) = None, resume: bool | Literal["allow", "never", "must", "auto"] | None = None, resume_from: str | None = None, fork_from: str | None = None, save_code: bool | None = None, tensorboard: bool | None = None, sync_tensorboard: bool | None = None, monitor_gym: bool | None = None, settings: Settings | dict[str, Any] | None = None, anonymous: DoNotSet = UNSET, ) -> Run: r"""Start a new run to track and log to W&B. In an ML training pipeline, you could add `wandb.init()` to the beginning of your training script as well as your evaluation script, and each piece would be tracked as a run in W&B. `wandb.init()` spawns a new background process to log data to a run, and it also syncs data to https://wandb.ai by default, so you can see your results in real-time. When you're done logging data, call `wandb.Run.finish()` to end the run. If you don't call `run.finish()`, the run will end when your script exits. Run IDs must not contain any of the following special characters `/ \ # ? % :` Args: entity: The username or team name the runs are logged to. The entity must already exist, so ensure you create your account or team in the UI before starting to log runs. If not specified, the run will default your default entity. To change the default entity, go to your settings and update the "Default location to create new projects" under "Default team". project: The name of the project under which this run will be logged. If not specified, we use a heuristic to infer the project name based 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: 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. For a short descriptive name, use the `name` field, or for saving hyperparameters to compare across runs, use `config`. name: A short display name for this run, which appears in the UI to help you identify it. By default, we generate a random two-word name allowing easy cross-reference runs from table to charts. Keeping these run names brief enhances readability in chart legends and tables. For saving hyperparameters, we recommend using the `config` field. notes: A detailed description of the run, similar to a commit message in Git. Use this argument to capture any context or details that may help you recall the purpose or setup of this run in the future. tags: A list of tags to label this run in the UI. Tags are helpful for organizing runs or adding temporary identifiers like "baseline" or "production." You can easily add, remove tags, or filter by tags in the UI. If resuming a run, the tags provided here will replace any existing tags. To add tags to a resumed run without overwriting the current tags, use `run.tags += ("new_tag",)` after calling `run = wandb.init()`. config: Sets `wandb.config`, a dictionary-like object for storing input parameters to your run, such as model hyperparameters or data preprocessing settings. The config appears in the UI in an overview page, allowing you to group, filter, and sort runs based on these parameters. Keys should not contain periods (`.`), and values should be smaller than 10 MB. If a dictionary, `argparse.Namespace`, or `absl.flags.FLAGS` is provided, the key-value pairs will be loaded directly into `wandb.config`. If a string is provided, it is interpreted as a path to a YAML file, from which configuration values will be loaded into `wandb.config`. config_exclude_keys: A list of specific keys to exclude from `wandb.config`. config_include_keys: A list of specific keys to include in `wandb.config`. allow_val_change: Controls whether config values can be modified after their initial set. By default, an exception is raised if a config value is overwritten. For tracking variables that change during training, such as a learning rate, consider using `wandb.log()` instead. By default, this is `False` in scripts and `True` in Notebook environments. group: Specify a group name to organize individual runs as part of a larger experiment. This is useful for cases like cross-validation or running multiple jobs that train and evaluate a model on different test sets. Grouping allows you to manage related runs collectively in the UI, making it easy to toggle and review results as a unified experiment. job_type: Specify the type of run, especially helpful when organizing runs within a group as part of a larger experiment. For example, in a group, you might label runs with job types such as "train" and "eval". Defining job types enables you to easily filter and group similar runs in the UI, facilitating direct comparisons. mode: Specifies how run data is managed, with the following options: - `"online"` (default): Enables live syncing with W&B when a network connection is available, with real-time updates to visualizations. - `"offline"`: Suitable for air-gapped or offline environments; data is saved locally and can be synced later. Ensure the run folder is preserved to enable future syncing. - `"disabled"`: Disables all W&B functionality, making the run’s methods no-ops. Typically used in testing to bypass W&B operations. - `"shared"`: (This is an experimental feature). Allows multiple processes, possibly on different machines, to simultaneously log to the same run. In this approach you use a primary node and one or more worker nodes to log data to the same run. Within the primary node you initialize a run. For each worker node, initialize a run using the run ID used by the primary node. force: Determines if a W&B login is required to run the script. If `True`, the user must be logged in to W&B; otherwise, the script will not proceed. If `False` (default), the script can proceed without a login, switching to offline mode if the user is not logged in. reinit: Shorthand for the "reinit" setting. Determines the behavior of `wandb.init()` when a run is active. resume: Controls the behavior when resuming a run with the specified `id`. Available options are: - `"allow"`: If a run with the specified `id` exists, it will resume from the last step; otherwise, a new run will be created. - `"never"`: If a run with the specified `id` exists, an error will be raised. If no such run is found, a new run will be created. - `"must"`: If a run with the specified `id` exists, it will resume from the last step. If no run is found, an error will be raised. - `"auto"`: Automatically resumes the previous run if it crashed on this machine; otherwise, starts a new run. - `True`: Deprecated. Use `"auto"` instead. - `False`: Deprecated. Use the default behavior (leaving `resume` unset) to always start a new run. If `resume` is set, `fork_from` and `resume_from` cannot be used. When `resume` is unset, the system will always start a new run. resume_from: Specifies a moment in a previous run to resume a run from, using the format `{run_id}?_step={step}`. This allows users to truncate the history logged to a run at an intermediate step and resume logging from that step. The target run must be in the same project. If an `id` argument is also provided, the `resume_from` argument will take precedence. `resume`, `resume_from` and `fork_from` cannot be used together, only one of them can be used at a time. Note that this feature is in beta and may change in the future. fork_from: Specifies a point in a previous run from which to fork a new run, using the format `{id}?_step={step}`. This creates a new run that resumes logging from the specified step in the target run’s history. The target run must be part of the current project. If an `id` argument is also provided, it must be different from the `fork_from` argument, an error will be raised if they are the same. `resume`, `resume_from` and `fork_from` cannot be used together, only one of them can be used at a time. Note that this feature is in beta and may change in the future. save_code: Enables saving the main script or notebook to W&B, aiding in experiment reproducibility and allowing code comparisons across runs in the UI. By default, this is disabled, but you can change the default to enable on your settings page. tensorboard: Deprecated. Use `sync_tensorboard` instead. sync_tensorboard: Enables automatic syncing of W&B logs from TensorBoard or TensorBoardX, saving relevant event files for viewing in the W&B UI. monitor_gym: Enables automatic logging of videos of the environment when using OpenAI Gym. settings: Specifies a dictionary or `wandb.Settings` object with advanced settings for the run. Returns: A `Run` object. Raises: Error: If some unknown or internal error happened during the run initialization. AuthenticationError: If the user failed to provide valid credentials. CommError: If there was a problem communicating with the WandB server. UsageError: If the user provided invalid arguments. KeyboardInterrupt: If user interrupts the run. Examples: `wandb.init()` returns a `Run` object. Use the run object to log data, save artifacts, and manage the run lifecycle. ```python import wandb config = {"lr": 0.01, "batch_size": 32} with wandb.init(config=config) as run: # Log accuracy and loss to the run acc = 0.95 # Example accuracy loss = 0.05 # Example loss run.log({"accuracy": acc, "loss": loss}) ``` """ init_telemetry = telemetry.TelemetryRecord() init_settings = Settings() if isinstance(settings, dict): init_settings = Settings(**settings) elif isinstance(settings, Settings): init_settings = settings # Explicit function arguments take precedence over settings if job_type is not None: init_settings.run_job_type = job_type if dir is not None: init_settings.root_dir = dir # type: ignore if project is not None: init_settings.project = project if entity is not None: init_settings.entity = entity if reinit is not None: init_settings.reinit = reinit if tags is not None: init_settings.run_tags = tuple(tags) if group is not None: init_settings.run_group = group if name is not None: init_settings.run_name = name if notes is not None: init_settings.run_notes = notes if anonymous is not UNSET: init_settings.anonymous = anonymous if mode is not None: init_settings.mode = mode # type: ignore if resume is not None: init_settings.resume = resume # type: ignore if force is not None: init_settings.force = force # TODO: deprecate "tensorboard" in favor of "sync_tensorboard" if tensorboard is not None: init_settings.sync_tensorboard = tensorboard if sync_tensorboard is not None: init_settings.sync_tensorboard = sync_tensorboard if save_code is not None: init_settings.save_code = save_code if id is not None: init_settings.run_id = id if fork_from is not None: init_settings.fork_from = fork_from # type: ignore if resume_from is not None: init_settings.resume_from = resume_from # type: ignore if config is not None: init_telemetry.feature.set_init_config = True wl: wandb_setup._WandbSetup | None = None try: wl = wandb_setup.singleton() wi = _WandbInit(wl, init_telemetry) wi.maybe_login(init_settings) run_settings, show_warnings = wi.make_run_settings(init_settings) if isinstance(run_settings.reinit, bool): wi.deprecated_features_used.append( ( Deprecated(run__reinit_bool=True), "Using a boolean value for 'reinit' is deprecated." " Use 'return_previous' or 'finish_previous' instead.", ) ) if run_settings.run_id is not None: init_telemetry.feature.set_init_id = True if run_settings.run_name is not None: init_telemetry.feature.set_init_name = True if run_settings.run_tags is not None: init_telemetry.feature.set_init_tags = True if run_settings._offline: init_telemetry.feature.offline = True if run_settings.fork_from is not None: init_telemetry.feature.fork_mode = True if run_settings.resume_from is not None: init_telemetry.feature.rewind_mode = True wi.set_run_id(run_settings) wi.set_sync_dir_suffix(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)) run_config = wi.make_run_config( settings=run_settings, config=config, config_exclude_keys=config_exclude_keys, config_include_keys=config_include_keys, ) if run_settings._noop: return wi.make_disabled_run(run_config) try_create_root_dir(run_settings) exit_stack.enter_context(wi.setup_run_log_directory(run_settings)) if run_settings._jupyter: wi.monkeypatch_ipython(run_settings) if monitor_gym: _monkeypatch_openai_gym() if wandb.patched["tensorboard"]: # NOTE: The user may have called the patch function directly. init_telemetry.feature.tensorboard_patch = True if run_settings.sync_tensorboard: _monkeypatch_tensorboard() init_telemetry.feature.tensorboard_sync = True if run_settings.x_server_side_derived_summary: init_telemetry.feature.server_side_derived_summary = True run = wi.init(run_settings, run_config, run_printer) # Set up automatic Weave integration if Weave is installed weave.setup(run_settings.entity, run_settings.project) return run except KeyboardInterrupt as e: if wl: wl._get_logger().warning("interrupted", exc_info=e) raise except Exception as e: if wl: wl._get_logger().exception("error in wandb.init()", exc_info=e) get_sentry().reraise(e)