feat: Update NeptuneScaleLogger to log model checkpoint paths instead of uploading checkpoints

Author: SiddhantSadangi

src/lightning/pytorch/loggers/neptune.py

@@ -669,18 +669,18 @@ def any_lightning_module_function_or_hook(self):
         neptune_logger.run.log_configs(data={"your/metadata/structure": metadata})
         neptune_logger.run.add_tags(["tag1", "tag2"])
 
-    **Log model checkpoints**
+    **Log model checkpoint paths**
 
     If you have :class:`~lightning.pytorch.callbacks.ModelCheckpoint` configured,
-    the Neptune logger automatically logs model checkpoints.
-    Model weights will be uploaded to the "model/checkpoints" namespace in the Neptune run.
+    the Neptune logger can log model checkpoint paths.
+    Paths will be logged to the "model/checkpoints" namespace in the Neptune run.
     You can disable this option with:
 
     .. code-block:: python
 
         neptune_logger = NeptuneScaleLogger(log_model_checkpoints=False)
 
-    Note: All model checkpoints will be uploaded. ``save_last`` and ``save_top_k`` are currently not supported.
+    Note: All model checkpoint paths will be logged. ``save_last`` and ``save_top_k`` are currently not supported.
 
     **Pass additional parameters to the Neptune run**
 
@@ -743,9 +743,9 @@ def any_lightning_module_function_or_hook(self):
         run: Optional. Default is ``None``. A Neptune ``Run`` object.
             If specified, this existing run will be used for logging, instead of a new run being created.
         prefix: Optional. Default is ``"training"``. Root namespace for all metadata logging.
-        log_model_checkpoints: Optional. Default is ``True``. Log model checkpoints to Neptune.
+        log_model_checkpoints: Optional. Default is ``True``. Log model checkpoint paths to Neptune.
             Works only if ``ModelCheckpoint`` is passed to the ``Trainer``.
-            NOTE: All model checkpoints will be uploaded.
+            NOTE: All model checkpoint paths will be logged.
             ``save_last`` and ``save_top_k`` are currently not supported.
         neptune_run_kwargs: Additional arguments like ``creation_time``, ``log_directory``,
             ``fork_run_id``, ``fork_step``, ``*_callback``, etc. used when a run is created.
@@ -1050,7 +1050,7 @@ def log_model_summary(self, model: "pl.LightningModule", max_depth: int = -1) ->
     @override
     @rank_zero_only
     def after_save_checkpoint(self, checkpoint_callback: Checkpoint) -> None:
-        """Automatically log checkpointed model. Called after model checkpoint callback saves a new checkpoint.
+        """Automatically log checkpointed model's path. Called after model checkpoint callback saves a new checkpoint.
 
         Args:
             checkpoint_callback: the model checkpoint callback instance
@@ -1066,7 +1066,7 @@ def after_save_checkpoint(self, checkpoint_callback: Checkpoint) -> None:
         if hasattr(checkpoint_callback, "last_model_path") and checkpoint_callback.last_model_path:
             model_last_name = self._get_full_model_name(checkpoint_callback.last_model_path, checkpoint_callback)
             file_names.add(model_last_name)
-            self.run.assign_files({
+            self.run.log_configs({
                 f"{checkpoints_namespace}/{model_last_name}": checkpoint_callback.last_model_path,
             })
 
@@ -1075,7 +1075,7 @@ def after_save_checkpoint(self, checkpoint_callback: Checkpoint) -> None:
             for key in checkpoint_callback.best_k_models:
                 model_name = self._get_full_model_name(key, checkpoint_callback)
                 file_names.add(model_name)
-                self.run.assign_files({
+                self.run.log_configs({
                     f"{checkpoints_namespace}/{model_name}": key,
                 })
 
@@ -1087,7 +1087,7 @@ def after_save_checkpoint(self, checkpoint_callback: Checkpoint) -> None:
 
             model_name = self._get_full_model_name(checkpoint_callback.best_model_path, checkpoint_callback)
             file_names.add(model_name)
-            self.run.assign_files({
+            self.run.log_configs({
                 f"{checkpoints_namespace}/{model_name}": checkpoint_callback.best_model_path,
             })
 

tests/tests_pytorch/loggers/test_neptune.py

@@ -461,7 +461,8 @@ def test_neptune_scale_logger_finalize(neptune_scale_logger):
     """Test finalize method sets status correctly."""
     logger, mock_run = neptune_scale_logger
     logger.finalize("success")
-    assert mock_run._status == "success"
+    expected_key = logger._construct_path_with_prefix("status")
+    mock_run.log_configs.assert_any_call({expected_key: "success"})
 
 
 @pytest.mark.skipif(not _NEPTUNE_SCALE_AVAILABLE, reason="Neptune-Scale is required for this test.")
@@ -472,24 +473,20 @@ def test_neptune_scale_logger_invalid_run():
 
 
 @pytest.mark.skipif(not _NEPTUNE_SCALE_AVAILABLE, reason="Neptune-Scale is required for this test.")
-def test_neptune_scale_logger_log_model_summary(neptune_scale_logger):
+def test_neptune_scale_logger_log_model_summary(neptune_scale_logger, monkeypatch):
     from neptune_scale.types import File
 
     model = BoringModel()
-    test_variants = [
-        ({}, "training/model/summary"),
-        ({"prefix": "custom_prefix"}, "custom_prefix/model/summary"),
-        ({"prefix": "custom/nested/prefix"}, "custom/nested/prefix/model/summary"),
-    ]
-
-    for prefix, model_summary_key in test_variants:
-        logger, run_instance_mock, _ = _get_logger_with_mocks(api_key="test", project="project", **prefix)
-
-        logger.log_model_summary(model)
-
-        assert run_instance_mock.__setitem__.call_count == 1
-        assert run_instance_mock.__getitem__.call_count == 0
-        run_instance_mock.__setitem__.assert_called_once_with(model_summary_key, File)
+    logger, mock_run = neptune_scale_logger
+    # Patch assign_files to track calls
+    assign_files_mock = mock.MagicMock()
+    monkeypatch.setattr(mock_run, "assign_files", assign_files_mock)
+    logger.log_model_summary(model)
+    # Check that assign_files was called with the correct key and a File instance
+    called_args = assign_files_mock.call_args[0][0]
+    assert list(called_args.keys())[0].endswith("model/summary")
+    file_val = list(called_args.values())[0]
+    assert isinstance(file_val, File)
 
 
 @pytest.mark.skipif(not _NEPTUNE_SCALE_AVAILABLE, reason="Neptune-Scale is required for this test.")
@@ -511,32 +508,28 @@ def test_neptune_scale_logger_with_prefix(neptune_scale_logger):
 
 @pytest.mark.skipif(not _NEPTUNE_SCALE_AVAILABLE, reason="Neptune-Scale is required for this test.")
 def test_neptune_scale_logger_after_save_checkpoint(neptune_scale_logger):
-    test_variants = [
-        ({}, "training/model"),
-        ({"prefix": "custom_prefix"}, "custom_prefix/model"),
-        ({"prefix": "custom/nested/prefix"}, "custom/nested/prefix/model"),
+    logger, mock_run = neptune_scale_logger
+    models_root_dir = os.path.join("path", "to", "models")
+    cb_mock = MagicMock(
+        dirpath=models_root_dir,
+        last_model_path=os.path.join(models_root_dir, "last"),
+        best_k_models={
+            f"{os.path.join(models_root_dir, 'model1')}": None,
+            f"{os.path.join(models_root_dir, 'model2/with/slashes')}": None,
+        },
+        best_model_path=os.path.join(models_root_dir, "best_model"),
+        best_model_score=None,
+    )
+    logger.after_save_checkpoint(cb_mock)
+    prefix = logger._prefix
+    model_key_prefix = f"{prefix}/model" if prefix else "model"
+    expected_calls = [
+        call.log_configs({f"{model_key_prefix}/checkpoints/model1": os.path.join(models_root_dir, "model1")}),
+        call.log_configs({
+            f"{model_key_prefix}/checkpoints/model2/with/slashes": os.path.join(models_root_dir, "model2/with/slashes")
+        }),
+        call.log_configs({f"{model_key_prefix}/checkpoints/last": os.path.join(models_root_dir, "last")}),
+        call.log_configs({f"{model_key_prefix}/checkpoints/best_model": os.path.join(models_root_dir, "best_model")}),
+        call.log_configs({f"{model_key_prefix}/best_model_path": os.path.join(models_root_dir, "best_model")}),
     ]
-
-    for prefix, model_key_prefix in test_variants:
-        logger, run_instance_mock, run_attr_mock = _get_logger_with_mocks(api_key="test", project="project", **prefix)
-        models_root_dir = os.path.join("path", "to", "models")
-        cb_mock = MagicMock(
-            dirpath=models_root_dir,
-            last_model_path=os.path.join(models_root_dir, "last"),
-            best_k_models={
-                f"{os.path.join(models_root_dir, 'model1')}": None,
-                f"{os.path.join(models_root_dir, 'model2/with/slashes')}": None,
-            },
-            best_model_path=os.path.join(models_root_dir, "best_model"),
-            best_model_score=None,
-        )
-
-        logger.after_save_checkpoint(cb_mock)
-
-        run_instance_mock.__getitem__.assert_any_call(f"{model_key_prefix}/checkpoints/model1")
-        run_instance_mock.__getitem__.assert_any_call(f"{model_key_prefix}/checkpoints/model2/with/slashes")
-
-        run_attr_mock.upload.assert_has_calls([
-            call(os.path.join(models_root_dir, "model1")),
-            call(os.path.join(models_root_dir, "model2/with/slashes")),
-        ])
+    mock_run.log_configs.assert_has_calls(expected_calls, any_order=True)