fix(loggers): add best and latest aliases to wandb artifact in WandbLogger (#17121)

Co-authored-by: awaelchli <[email protected]>

Author: parambharat

src/lightning/pytorch/CHANGELOG.md

@@ -29,6 +29,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Fixed WandbLogger not showing "best" aliases for model checkpoints when `ModelCheckpoint(save_top_k>0)` is used ([#17121](https://github.com/Lightning-AI/lightning/pull/17121))
 
 
 

src/lightning/pytorch/loggers/wandb.py

@@ -599,6 +599,7 @@ def _scan_and_log_checkpoints(self, checkpoint_callback: ModelCheckpoint) -> Non
                 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")
-            self.experiment.log_artifact(artifact, aliases=[tag])
+            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)
             self._logged_model_time[p] = t

tests/tests_pytorch/loggers/test_wandb.py

@@ -277,6 +277,78 @@ def test_wandb_log_model(wandb, monkeypatch, tmpdir):
         },
     )
 
+    # Test wandb artifact with checkpoint_callback top_k logging latest
+    wandb.init().log_artifact.reset_mock()
+    wandb.init.reset_mock()
+    wandb.Artifact.reset_mock()
+    logger = WandbLogger(save_dir=tmpdir, log_model=True)
+    logger.experiment.id = "1"
+    logger.experiment.name = "run_name"
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        logger=logger,
+        max_epochs=3,
+        limit_train_batches=3,
+        limit_val_batches=3,
+        callbacks=[ModelCheckpoint(monitor="step", save_top_k=2)],
+    )
+    trainer.fit(model)
+    wandb.Artifact.assert_called_with(
+        name="model-1",
+        type="model",
+        metadata={
+            "score": 6,
+            "original_filename": "epoch=1-step=6-v5.ckpt",
+            "ModelCheckpoint": {
+                "monitor": "step",
+                "mode": "min",
+                "save_last": None,
+                "save_top_k": 2,
+                "save_weights_only": False,
+                "_every_n_train_steps": 0,
+            },
+        },
+    )
+    wandb.init().log_artifact.assert_called_with(wandb.Artifact(), aliases=["latest"])
+
+    # Test wandb artifact with checkpoint_callback top_k logging latest and best
+    wandb.init().log_artifact.reset_mock()
+    wandb.init.reset_mock()
+    wandb.Artifact.reset_mock()
+    logger = WandbLogger(save_dir=tmpdir, log_model=True)
+    logger.experiment.id = "1"
+    logger.experiment.name = "run_name"
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        logger=logger,
+        max_epochs=3,
+        limit_train_batches=3,
+        limit_val_batches=3,
+        callbacks=[
+            ModelCheckpoint(
+                monitor="step",
+            )
+        ],
+    )
+    trainer.fit(model)
+    wandb.Artifact.assert_called_with(
+        name="model-1",
+        type="model",
+        metadata={
+            "score": 3,
+            "original_filename": "epoch=0-step=3-v1.ckpt",
+            "ModelCheckpoint": {
+                "monitor": "step",
+                "mode": "min",
+                "save_last": None,
+                "save_top_k": 1,
+                "save_weights_only": False,
+                "_every_n_train_steps": 0,
+            },
+        },
+    )
+    wandb.init().log_artifact.assert_called_with(wandb.Artifact(), aliases=["latest", "best"])
+
 
 @mock.patch("lightning.pytorch.loggers.wandb.Run", new=mock.Mock)
 @mock.patch("lightning.pytorch.loggers.wandb.wandb")