wandb 0.19.8__py3-none-any.whl → 0.19.9__py3-none-any.whl

wandb/sdk/wandb_run.py

@@ -44,7 +44,7 @@ from wandb.proto.wandb_internal_pb2 import (
 )
 from wandb.sdk.artifacts.artifact import Artifact
 from wandb.sdk.internal import job_builder
-from wandb.sdk.lib import asyncio_compat
+from wandb.sdk.lib import asyncio_compat, wb_logging
 from wandb.sdk.lib.import_hooks import (
     register_post_import_hook,
     unregister_post_import_hook,
@@ -62,6 +62,7 @@ from wandb.util import (
 
 from . import wandb_config, wandb_metric, wandb_summary
 from .artifacts._validators import (
+    MAX_ARTIFACT_METADATA_KEYS,
     is_artifact_registry_project,
     validate_aliases,
     validate_tags,
@@ -165,11 +166,13 @@ class RunStatusChecker:
 
     def __init__(
         self,
+        run_id: str,
         interface: InterfaceBase,
         stop_polling_interval: int = 15,
         retry_polling_interval: int = 5,
         internal_messages_polling_interval: int = 10,
     ) -> None:
+        self._run_id = run_id
         self._interface = interface
         self._stop_polling_interval = stop_polling_interval
         self._retry_polling_interval = retry_polling_interval
@@ -275,19 +278,20 @@ class RunStatusChecker:
                         )
                     )
 
-        try:
-            self._loop_check_status(
-                lock=self._network_status_lock,
-                set_handle=lambda x: setattr(self, "_network_status_handle", x),
-                timeout=self._retry_polling_interval,
-                request=self._interface.deliver_network_status,
-                process=_process_network_status,
-            )
-        except BrokenPipeError:
-            self._abandon_status_check(
-                self._network_status_lock,
-                self._network_status_handle,
-            )
+        with wb_logging.log_to_run(self._run_id):
+            try:
+                self._loop_check_status(
+                    lock=self._network_status_lock,
+                    set_handle=lambda x: setattr(self, "_network_status_handle", x),
+                    timeout=self._retry_polling_interval,
+                    request=self._interface.deliver_network_status,
+                    process=_process_network_status,
+                )
+            except BrokenPipeError:
+                self._abandon_status_check(
+                    self._network_status_lock,
+                    self._network_status_handle,
+                )
 
     def check_stop_status(self) -> None:
         def _process_stop_status(result: Result) -> None:
@@ -299,19 +303,20 @@ class RunStatusChecker:
                     interrupt.interrupt_main()
                     return
 
-        try:
-            self._loop_check_status(
-                lock=self._stop_status_lock,
-                set_handle=lambda x: setattr(self, "_stop_status_handle", x),
-                timeout=self._stop_polling_interval,
-                request=self._interface.deliver_stop_status,
-                process=_process_stop_status,
-            )
-        except BrokenPipeError:
-            self._abandon_status_check(
-                self._stop_status_lock,
-                self._stop_status_handle,
-            )
+        with wb_logging.log_to_run(self._run_id):
+            try:
+                self._loop_check_status(
+                    lock=self._stop_status_lock,
+                    set_handle=lambda x: setattr(self, "_stop_status_handle", x),
+                    timeout=self._stop_polling_interval,
+                    request=self._interface.deliver_stop_status,
+                    process=_process_stop_status,
+                )
+            except BrokenPipeError:
+                self._abandon_status_check(
+                    self._stop_status_lock,
+                    self._stop_status_handle,
+                )
 
     def check_internal_messages(self) -> None:
         def _process_internal_messages(result: Result) -> None:
@@ -319,19 +324,20 @@ class RunStatusChecker:
             for msg in internal_messages.messages.warning:
                 wandb.termwarn(msg)
 
-        try:
-            self._loop_check_status(
-                lock=self._internal_messages_lock,
-                set_handle=lambda x: setattr(self, "_internal_messages_handle", x),
-                timeout=self._internal_messages_polling_interval,
-                request=self._interface.deliver_internal_messages,
-                process=_process_internal_messages,
-            )
-        except BrokenPipeError:
-            self._abandon_status_check(
-                self._internal_messages_lock,
-                self._internal_messages_handle,
-            )
+        with wb_logging.log_to_run(self._run_id):
+            try:
+                self._loop_check_status(
+                    lock=self._internal_messages_lock,
+                    set_handle=lambda x: setattr(self, "_internal_messages_handle", x),
+                    timeout=self._internal_messages_polling_interval,
+                    request=self._interface.deliver_internal_messages,
+                    process=_process_internal_messages,
+                )
+            except BrokenPipeError:
+                self._abandon_status_check(
+                    self._internal_messages_lock,
+                    self._internal_messages_handle,
+                )
 
     def stop(self) -> None:
         self._join_event.set()
@@ -355,6 +361,34 @@ class RunStatusChecker:
         self._internal_messages_thread.join()
 
 
+def _log_to_run(func: Callable) -> Callable:
+    """Decorate a Run method to set the run ID in the logging context.
+
+    Any logs during the execution of the method go to the run's log file
+    and not to other runs' log files.
+
+    This is meant for use on all public methods and some callbacks. Private
+    methods can be assumed to be called from some public method somewhere.
+    The general rule is to use it on methods that can be called from a
+    context that isn't specific to this run (such as all user code or
+    internal methods that aren't run-specific).
+    """
+
+    @functools.wraps(func)
+    def wrapper(self: Run, *args, **kwargs) -> Any:
+        # In "attach" usage, many properties of the Run are not initially
+        # populated.
+        if hasattr(self, "_settings"):
+            run_id = self._settings.run_id
+        else:
+            run_id = self._attach_id
+
+        with wb_logging.log_to_run(run_id):
+            return func(self, *args, **kwargs)
+
+    return wrapper
+
+
 class _run_decorator:  # noqa: N801
     _is_attaching: str = ""
 
@@ -811,29 +845,34 @@ class Run:
         return self._torch_history
 
     @property
+    @_log_to_run
     @_run_decorator._attach
     def settings(self) -> Settings:
         """A frozen copy of run's Settings object."""
         return self._settings.model_copy(deep=True)
 
     @property
+    @_log_to_run
     @_run_decorator._attach
     def dir(self) -> str:
         """The directory where files associated with the run are saved."""
         return self._settings.files_dir
 
     @property
+    @_log_to_run
     @_run_decorator._attach
     def config(self) -> wandb_config.Config:
         """Config object associated with this run."""
         return self._config
 
     @property
+    @_log_to_run
     @_run_decorator._attach
     def config_static(self) -> wandb_config.ConfigStatic:
         return wandb_config.ConfigStatic(self._config)
 
     @property
+    @_log_to_run
     @_run_decorator._attach
     def name(self) -> str | None:
         """Display name of the run.
@@ -846,6 +885,7 @@ class Run:
         return None
 
     @name.setter
+    @_log_to_run
     @_run_decorator._noop_on_finish()
     def name(self, name: str) -> None:
         with telemetry.context(run=self) as tel:
@@ -855,6 +895,7 @@ class Run:
             self._backend.interface.publish_run(self)
 
     @property
+    @_log_to_run
     @_run_decorator._attach
     def notes(self) -> str | None:
         """Notes associated with the run, if there are any.
@@ -865,6 +906,7 @@ class Run:
         return self._settings.run_notes
 
     @notes.setter
+    @_log_to_run
     @_run_decorator._noop_on_finish()
     def notes(self, notes: str) -> None:
         self._settings.run_notes = notes
@@ -872,12 +914,14 @@ class Run:
             self._backend.interface.publish_run(self)
 
     @property
+    @_log_to_run
     @_run_decorator._attach
     def tags(self) -> tuple | None:
         """Tags associated with the run, if there are any."""
         return self._settings.run_tags or ()
 
     @tags.setter
+    @_log_to_run
     @_run_decorator._noop_on_finish()
     def tags(self, tags: Sequence) -> None:
         with telemetry.context(run=self) as tel:
@@ -887,6 +931,7 @@ class Run:
             self._backend.interface.publish_run(self)
 
     @property
+    @_log_to_run
     @_run_decorator._attach
     def id(self) -> str:
         """Identifier for this run."""
@@ -895,6 +940,7 @@ class Run:
         return self._settings.run_id
 
     @property
+    @_log_to_run
     @_run_decorator._attach
     def sweep_id(self) -> str | None:
         """ID of the sweep associated with the run, if there is one."""
@@ -912,6 +958,7 @@ class Run:
         )
 
     @property
+    @_log_to_run
     @_run_decorator._attach
     def path(self) -> str:
         """Path to the run.
@@ -922,24 +969,28 @@ class Run:
         return self._get_path()
 
     @property
+    @_log_to_run
     @_run_decorator._attach
     def start_time(self) -> float:
         """Unix timestamp (in seconds) of when the run started."""
         return self._start_time
 
     @property
+    @_log_to_run
     @_run_decorator._attach
     def starting_step(self) -> int:
         """The first step of the run."""
         return self._starting_step
 
     @property
+    @_log_to_run
     @_run_decorator._attach
     def resumed(self) -> bool:
         """True if the run was resumed, False otherwise."""
         return self._settings.resumed
 
     @property
+    @_log_to_run
     @_run_decorator._attach
     def step(self) -> int:
         """Current value of the step.
@@ -949,6 +1000,7 @@ class Run:
         return self._step
 
     @property
+    @_log_to_run
     @_run_decorator._attach
     def mode(self) -> str:
         """For compatibility with `0.9.x` and earlier, deprecate eventually."""
@@ -963,16 +1015,19 @@ class Run:
         return "dryrun" if self._settings._offline else "run"
 
     @property
+    @_log_to_run
     @_run_decorator._attach
     def offline(self) -> bool:
         return self._settings._offline
 
     @property
+    @_log_to_run
     @_run_decorator._attach
     def disabled(self) -> bool:
         return self._settings._noop
 
     @property
+    @_log_to_run
     @_run_decorator._attach
     def group(self) -> str:
         """Name of the group associated with the run.
@@ -987,6 +1042,7 @@ class Run:
         return self._settings.run_group or ""
 
     @property
+    @_log_to_run
     @_run_decorator._attach
     def job_type(self) -> str:
         return self._settings.run_job_type or ""
@@ -996,12 +1052,14 @@ class Run:
         return self._settings.project or ""
 
     @property
+    @_log_to_run
     @_run_decorator._attach
     def project(self) -> str:
         """Name of the W&B project associated with the run."""
         return self.project_name()
 
     @_run_decorator._noop_on_finish()
+    @_log_to_run
     @_run_decorator._attach
     def log_code(
         self,
@@ -1088,6 +1146,7 @@ class Run:
 
         return self._log_artifact(art)
 
+    @_log_to_run
     def get_project_url(self) -> str | None:
         """Return the url for the W&B project associated with the run, if there is one.
 
@@ -1098,6 +1157,7 @@ class Run:
             return None
         return self._settings.project_url
 
+    @_log_to_run
     def get_sweep_url(self) -> str | None:
         """Return the url for the sweep associated with the run, if there is one."""
         if self._settings._offline:
@@ -1105,6 +1165,7 @@ class Run:
             return None
         return self._settings.sweep_url
 
+    @_log_to_run
     def get_url(self) -> str | None:
         """Return the url for the W&B run, if there is one.
 
@@ -1116,12 +1177,14 @@ class Run:
         return self._settings.run_url
 
     @property
+    @_log_to_run
     @_run_decorator._attach
     def url(self) -> str | None:
         """The W&B url associated with the run."""
         return self.get_url()
 
     @property
+    @_log_to_run
     @_run_decorator._attach
     def entity(self) -> str:
         """The name of the W&B entity associated with the run.
@@ -1216,6 +1279,7 @@ class Run:
         if lines:
             self._label_probe_lines(lines)
 
+    @_log_to_run
     @_run_decorator._attach
     def display(self, height: int = 420, hidden: bool = False) -> bool:
         """Display this run in jupyter."""
@@ -1235,6 +1299,7 @@ class Run:
             wandb.termwarn(".display() only works in jupyter environments")
             return False
 
+    @_log_to_run
     @_run_decorator._attach
     def to_html(self, height: int = 420, hidden: bool = False) -> str:
         """Generate HTML containing an iframe displaying the current run."""
@@ -1251,6 +1316,7 @@ class Run:
     ) -> dict[str, str]:
         return {"text/html": self.to_html(hidden=True)}
 
+    @_log_to_run
     @_run_decorator._noop_on_finish()
     def _config_callback(
         self,
@@ -1262,6 +1328,7 @@ class Run:
         if self._backend and self._backend.interface:
             self._backend.interface.publish_config(key=key, val=val, data=data)
 
+    @_log_to_run
     def _config_artifact_callback(
         self, key: str, val: str | Artifact | dict
     ) -> Artifact:
@@ -1301,6 +1368,7 @@ class Run:
     def _set_config_wandb(self, key: str, val: Any) -> None:
         self._config_callback(key=("_wandb", key), val=val)
 
+    @_log_to_run
     @_run_decorator._noop_on_finish()
     def _summary_update_callback(self, summary_record: SummaryRecord) -> None:
         with telemetry.context(run=self) as tel:
@@ -1308,6 +1376,7 @@ class Run:
         if self._backend and self._backend.interface:
             self._backend.interface.publish_summary(self, summary_record)
 
+    @_log_to_run
     def _summary_get_current_summary_callback(self) -> dict[str, Any]:
         if self._is_finished:
             # TODO: WB-18420: fetch summary from backend and stage it before run is finished
@@ -1325,10 +1394,12 @@ class Run:
         get_summary_response = result.response.get_summary_response
         return proto_util.dict_from_proto_list(get_summary_response.item)
 
+    @_log_to_run
     def _metric_callback(self, metric_record: MetricRecord) -> None:
         if self._backend and self._backend.interface:
             self._backend.interface._publish_metric(metric_record)
 
+    @_log_to_run
     def _datatypes_callback(self, fname: str) -> None:
         if not self._backend or not self._backend.interface:
             return
@@ -1399,6 +1470,7 @@ class Run:
 
         return data
 
+    @_log_to_run
     def _partial_history_callback(
         self,
         data: dict[str, Any],
@@ -1423,11 +1495,13 @@ class Run:
             publish_step=not_using_tensorboard,
         )
 
+    @_log_to_run
     def _console_callback(self, name: str, data: str) -> None:
         # logger.info("console callback: %s, %s", name, data)
         if self._backend and self._backend.interface:
             self._backend.interface.publish_output(name, data)
 
+    @_log_to_run
     @_run_decorator._noop_on_finish(only_warn=True)
     def _console_raw_callback(self, name: str, data: str) -> None:
         # logger.info("console callback: %s, %s", name, data)
@@ -1443,6 +1517,7 @@ class Run:
         if self._backend and self._backend.interface:
             self._backend.interface.publish_output_raw(name, data)
 
+    @_log_to_run
     def _tensorboard_callback(
         self, logdir: str, save: bool = True, root_logdir: str = ""
     ) -> None:
@@ -1476,9 +1551,6 @@ class Run:
         # Grab the config from resuming
         if run_obj.config:
             c_dict = config_util.dict_no_value_from_proto_list(run_obj.config.update)
-            # TODO: Windows throws a wild error when this is set...
-            if "_wandb" in c_dict:
-                del c_dict["_wandb"]
             # We update the config object here without triggering the callback
             self._config._update(c_dict, allow_val_change=True, ignore_locked=True)
         # Update the summary, this will trigger an un-needed graphql request :(
@@ -1610,6 +1682,7 @@ class Run:
         if (step is None and commit is None) or commit:
             self._step += 1
 
+    @_log_to_run
     @_run_decorator._noop
     @_run_decorator._noop_on_finish()
     @_run_decorator._attach
@@ -1872,6 +1945,7 @@ class Run:
             )
         self._log(data=data, step=step, commit=commit)
 
+    @_log_to_run
     @_run_decorator._noop_on_finish()
     @_run_decorator._attach
     def save(
@@ -2057,6 +2131,7 @@ class Run:
 
         return [str(f) for f in globbed_files]
 
+    @_log_to_run
     @_run_decorator._attach
     def restore(
         self,
@@ -2072,6 +2147,7 @@ class Run:
             root or self._settings.files_dir,
         )
 
+    @_log_to_run
     @_run_decorator._noop
     @_run_decorator._attach
     def finish(
@@ -2105,6 +2181,7 @@ class Run:
             )
         return self._finish(exit_code)
 
+    @_log_to_run
     def _finish(
         self,
         exit_code: int | None = None,
@@ -2145,6 +2222,7 @@ class Run:
             module.unset_globals()
             wandb._sentry.end_session()
 
+    @_log_to_run
     @_run_decorator._noop
     @_run_decorator._attach
     def join(self, exit_code: int | None = None) -> None:
@@ -2158,6 +2236,7 @@ class Run:
             )
         self._finish(exit_code=exit_code)
 
+    @_log_to_run
     @_run_decorator._noop_on_finish()
     @_run_decorator._attach
     def status(
@@ -2410,7 +2489,9 @@ class Run:
                 self._backend.interface.publish_python_packages(working_set())
 
         if self._backend and self._backend.interface and not self._settings._offline:
+            assert self._settings.run_id
             self._run_status_checker = RunStatusChecker(
+                self._settings.run_id,
                 interface=self._backend.interface,
             )
             self._run_status_checker.start()
@@ -2657,6 +2738,7 @@ class Run:
         for module_name in import_telemetry_set:
             unregister_post_import_hook(module_name, run_id)
 
+    @_log_to_run
     @_run_decorator._noop_on_finish()
     @_run_decorator._attach
     def define_metric(
@@ -2798,6 +2880,7 @@ class Run:
         m._commit()
         return m
 
+    @_log_to_run
     @_run_decorator._attach
     def watch(
         self,
@@ -2835,6 +2918,7 @@ class Run:
         """
         wandb.sdk._watch(self, models, criterion, log, log_freq, idx, log_graph)
 
+    @_log_to_run
     @_run_decorator._attach
     def unwatch(
         self, models: torch.nn.Module | Sequence[torch.nn.Module] | None = None
@@ -2885,6 +2969,7 @@ class Run:
     def _detach(self) -> None:
         pass
 
+    @_log_to_run
     @_run_decorator._noop_on_finish()
     @_run_decorator._attach
     def link_artifact(
@@ -2950,6 +3035,7 @@ class Run:
         if response.error_message:
             wandb.termerror(response.error_message)
 
+    @_log_to_run
     @_run_decorator._noop_on_finish()
     @_run_decorator._attach
     def use_artifact(
@@ -3021,6 +3107,8 @@ class Run:
                 artifact.id,
                 entity_name=self._settings.entity,
                 project_name=self._settings.project,
+                artifact_entity_name=artifact.entity,
+                artifact_project_name=artifact.project,
                 use_as=use_as or artifact_or_name,
             )
         else:
@@ -3053,7 +3141,10 @@ class Run:
                     )
                 artifact._use_as = use_as or artifact.name
                 api.use_artifact(
-                    artifact.id, use_as=use_as or artifact._use_as or artifact.name
+                    artifact.id,
+                    use_as=use_as or artifact._use_as or artifact.name,
+                    artifact_entity_name=artifact.entity,
+                    artifact_project_name=artifact.project,
                 )
             else:
                 raise ValueError(
@@ -3064,6 +3155,7 @@ class Run:
             self._backend.interface.publish_use_artifact(artifact)
         return artifact
 
+    @_log_to_run
     @_run_decorator._noop_on_finish()
     @_run_decorator._attach
     def log_artifact(
@@ -3106,6 +3198,7 @@ class Run:
             tags=tags,
         )
 
+    @_log_to_run
     @_run_decorator._noop_on_finish()
     @_run_decorator._attach
     def upsert_artifact(
@@ -3160,6 +3253,7 @@ class Run:
             finalize=False,
         )
 
+    @_log_to_run
     @_run_decorator._noop_on_finish()
     @_run_decorator._attach
     def finish_artifact(
@@ -3244,6 +3338,12 @@ class Run:
         artifact, aliases = self._prepare_artifact(
             artifact_or_path, name, type, aliases
         )
+
+        if len(artifact.metadata) > MAX_ARTIFACT_METADATA_KEYS:
+            raise ValueError(
+                f"Artifact must not have more than {MAX_ARTIFACT_METADATA_KEYS} metadata keys."
+            )
+
         artifact.distributed_id = distributed_id
         self._assert_can_log_artifact(artifact)
         if self._backend and self._backend.interface:
@@ -3355,6 +3455,7 @@ class Run:
         artifact.finalize()
         return artifact, _resolve_aliases(aliases)
 
+    @_log_to_run
     @_run_decorator._noop_on_finish()
     @_run_decorator._attach
     def log_model(
@@ -3406,6 +3507,7 @@ class Run:
             artifact_or_path=path, name=name, type="model", aliases=aliases
         )
 
+    @_log_to_run
     @_run_decorator._noop_on_finish()
     @_run_decorator._attach
     def use_model(self, name: str) -> FilePathStr:
@@ -3463,6 +3565,7 @@ class Run:
             return FilePathStr(os.path.join(path, dir_list[0]))
         return path
 
+    @_log_to_run
     @_run_decorator._noop_on_finish()
     @_run_decorator._attach
     def link_model(
@@ -3564,6 +3667,7 @@ class Run:
             )
         self.link_artifact(artifact=artifact, target_path=target_path, aliases=aliases)
 
+    @_log_to_run
     @_run_decorator._noop_on_finish()
     @_run_decorator._attach
     def alert(
@@ -3615,6 +3719,7 @@ class Run:
         self._finish(exit_code=exit_code)
         return not exception_raised
 
+    @_log_to_run
     @_run_decorator._noop_on_finish()
     @_run_decorator._attach
     def mark_preempting(self) -> None:
@@ -3626,6 +3731,7 @@ class Run:
             self._backend.interface.publish_preempting()
 
     @property
+    @_log_to_run
     @_run_decorator._noop_on_finish()
     @_run_decorator._attach
     def _system_metrics(self) -> dict[str, list[tuple[datetime, float]]]:
@@ -3674,6 +3780,7 @@ class Run:
                 return {}
 
     @property
+    @_log_to_run
     @_run_decorator._attach
     @_run_decorator._noop_on_finish()
     def _metadata(self) -> Metadata | None:
@@ -3712,6 +3819,7 @@ class Run:
 
         return None
 
+    @_log_to_run
     @_run_decorator._noop_on_finish()
     @_run_decorator._attach
     def _metadata_callback(