Revert "Fix: Make WandbLogger upload models from all ModelCheckpoint callbacks, not just one (#20191)"

This reverts commit bd9d114

Author: Borda

src/lightning/pytorch/loggers/wandb.py

@@ -278,7 +278,6 @@ def any_lightning_module_function_or_hook(self):
         prefix: A string to put at the beginning of metric keys.
         experiment: WandB experiment object. Automatically set when creating a run.
         checkpoint_name: Name of the model checkpoint artifact being logged.
-        add_file_policy: If "mutable", copies file to tempdirectory before upload.
         \**kwargs: Arguments passed to :func:`wandb.init` like `entity`, `group`, `tags`, etc.
 
     Raises:
@@ -305,7 +304,6 @@ def __init__(
         experiment: Union["Run", "RunDisabled", None] = None,
         prefix: str = "",
         checkpoint_name: Optional[str] = None,
-        add_file_policy: Literal["mutable", "immutable"] = "mutable",
         **kwargs: Any,
     ) -> None:
         if not _WANDB_AVAILABLE:
@@ -324,8 +322,7 @@ def __init__(
         self._prefix = prefix
         self._experiment = experiment
         self._logged_model_time: dict[str, float] = {}
-        self._checkpoint_callbacks: dict[int, ModelCheckpoint] = {}
-        self.add_file_policy = add_file_policy
+        self._checkpoint_callback: Optional[ModelCheckpoint] = None
 
         # paths are processed as strings
         if save_dir is not None:
@@ -594,7 +591,7 @@ def after_save_checkpoint(self, checkpoint_callback: ModelCheckpoint) -> None:
         if self._log_model == "all" or self._log_model is True and checkpoint_callback.save_top_k == -1:
             self._scan_and_log_checkpoints(checkpoint_callback)
         elif self._log_model is True:
-            self._checkpoint_callbacks[id(checkpoint_callback)] = checkpoint_callback
+            self._checkpoint_callback = checkpoint_callback
 
     @staticmethod
     @rank_zero_only
@@ -647,9 +644,8 @@ def finalize(self, status: str) -> None:
             # Currently, checkpoints only get logged on success
             return
         # log checkpoints as artifacts
-        if self._experiment is not None:
-            for checkpoint_callback in self._checkpoint_callbacks.values():
-                self._scan_and_log_checkpoints(checkpoint_callback)
+        if self._checkpoint_callback and self._experiment is not None:
+            self._scan_and_log_checkpoints(self._checkpoint_callback)
 
     def _scan_and_log_checkpoints(self, checkpoint_callback: ModelCheckpoint) -> None:
         import wandb
@@ -679,7 +675,7 @@ def _scan_and_log_checkpoints(self, checkpoint_callback: ModelCheckpoint) -> Non
             if not self._checkpoint_name:
                 self._checkpoint_name = f"model-{self.experiment.id}"
             artifact = wandb.Artifact(name=self._checkpoint_name, type="model", metadata=metadata)
-            artifact.add_file(p, name="model.ckpt", policy=self.add_file_policy)
+            artifact.add_file(p, name="model.ckpt")
             aliases = ["latest", "best"] if p == checkpoint_callback.best_model_path else ["latest"]
             self.experiment.log_artifact(artifact, aliases=aliases)
             # remember logged models - timestamp needed in case filename didn't change (lastkckpt or custom name)

tests/tests_pytorch/loggers/test_wandb.py

@@ -426,44 +426,6 @@ def test_wandb_log_model(wandb_mock, tmp_path):
     )
     wandb_mock.init().log_artifact.assert_called_with(wandb_mock.Artifact(), aliases=["latest", "best"])
 
-    # Test wandb artifact with two checkpoint_callbacks
-    wandb_mock.init().log_artifact.reset_mock()
-    wandb_mock.init.reset_mock()
-    wandb_mock.Artifact.reset_mock()
-    logger = WandbLogger(save_dir=tmp_path, log_model=True)
-    logger.experiment.id = "1"
-    logger.experiment.name = "run_name"
-    trainer = Trainer(
-        default_root_dir=tmp_path,
-        logger=logger,
-        max_epochs=3,
-        limit_train_batches=3,
-        limit_val_batches=3,
-        callbacks=[
-            ModelCheckpoint(monitor="epoch", save_top_k=2),
-            ModelCheckpoint(monitor="step", save_top_k=2),
-        ],
-    )
-    trainer.fit(model)
-    for name, val, version in [("epoch", 0, 2), ("step", 3, 3)]:
-        wandb_mock.Artifact.assert_any_call(
-            name="model-1",
-            type="model",
-            metadata={
-                "score": val,
-                "original_filename": f"epoch=0-step=3-v{version}.ckpt",
-                "ModelCheckpoint": {
-                    "monitor": name,
-                    "mode": "min",
-                    "save_last": None,
-                    "save_top_k": 2,
-                    "save_weights_only": False,
-                    "_every_n_train_steps": 0,
-                },
-            },
-        )
-        wandb_mock.init().log_artifact.assert_any_call(wandb_mock.Artifact(), aliases=["latest"])
-
 
 def test_wandb_log_model_with_score(wandb_mock, tmp_path):
     """Test to prevent regression on #15543, ensuring the score is logged as a Python number, not a scalar tensor."""