[train][tune] fix WandbLoggerCallback to reuse actors upon restore (r…

…ay-project#47985)

This fixes an issue in the _restoration_ flow of Trials. The
`_WandbLoggingActor` actor is cleaned up as part of the trial completion
flow, but not appropriately handled during restoration, leading to
leaked actors. This fix reuses the same actor across restoration.

Signed-off-by: Matthew Deng <[email protected]>

python/ray/air/integrations/wandb.py

@@ -642,6 +642,11 @@ def log_trial_start(self, trial: "Trial"):
     def _start_logging_actor(
         self, trial: "Trial", exclude_results: List[str], **wandb_init_kwargs
     ):
+        # Reuse actor if one already exists.
+        # This can happen if the trial is restarted.
+        if trial in self._trial_logging_futures:
+            return
+
         if not self._remote_logger_class:
             env_vars = {}
             # API key env variable is not set if authenticating through `wandb login`

python/ray/air/tests/test_integration_wandb.py

@@ -483,6 +483,26 @@ def _handle_result(self, result):
             state = ray.get(actor.get_state.remote())
             assert [metrics["training_iteration"] for metrics in state.logs] == [4, 5]
 
+    def test_wandb_restart(self, trial):
+        """Test that the WandbLoggerCallback reuses actors for trial restarts."""
+
+        logger = WandbLoggerCallback(project="test_project", api_key="1234")
+        logger._logger_actor_cls = _MockWandbLoggingActor
+        logger.setup()
+
+        assert len(logger._trial_logging_futures) == 0
+        assert len(logger._logging_future_to_trial) == 0
+
+        logger.log_trial_start(trial)
+
+        assert len(logger._trial_logging_futures) == 1
+        assert len(logger._logging_future_to_trial) == 1
+
+        logger.log_trial_start(trial)
+
+        assert len(logger._trial_logging_futures) == 1
+        assert len(logger._logging_future_to_trial) == 1
+
 
 def test_wandb_logging_process_run_info_hook(monkeypatch):
     """