Merge branch 'master' into cli-custom-parser-class

Borda · web-flow · commit 0cb6cce5bb47 · 2025-03-10T14:28:42.000+01:00
diff --git a/.github/workflows/call-clear-cache.yml b/.github/workflows/call-clear-cache.yml
@@ -23,18 +23,18 @@ on:
 jobs:
   cron-clear:
     if: github.event_name == 'schedule' || github.event_name == 'pull_request'
-    uses: Lightning-AI/utilities/.github/workflows/cleanup-caches.yml@v0.12.0
+    uses: Lightning-AI/utilities/.github/workflows/cleanup-caches.yml@v0.14.0
     with:
-      scripts-ref: v0.11.8
+      scripts-ref: v0.14.0
       dry-run: ${{ github.event_name == 'pull_request' }}
       pattern: "latest|docs"
       age-days: 7
 
   direct-clear:
     if: github.event_name == 'workflow_dispatch' || github.event_name == 'pull_request'
-    uses: Lightning-AI/utilities/.github/workflows/cleanup-caches.yml@v0.12.0
+    uses: Lightning-AI/utilities/.github/workflows/cleanup-caches.yml@v0.14.0
     with:
-      scripts-ref: v0.11.8
+      scripts-ref: v0.14.0
       dry-run: ${{ github.event_name == 'pull_request' }}
       pattern: ${{ inputs.pattern || 'pypi_wheels' }} # setting str in case of PR / debugging
       age-days: ${{ fromJSON(inputs.age-days) || 0 }} # setting 0 in case of PR / debugging
diff --git a/.github/workflows/ci-check-md-links.yml b/.github/workflows/ci-check-md-links.yml
@@ -14,7 +14,7 @@ on:
 
 jobs:
   check-md-links:
-    uses: Lightning-AI/utilities/.github/workflows/check-md-links.yml@v0.12.0
+    uses: Lightning-AI/utilities/.github/workflows/check-md-links.yml@v0.14.0
     with:
       config-file: ".github/markdown-links-config.json"
       base-branch: "master"
diff --git a/.github/workflows/ci-schema.yml b/.github/workflows/ci-schema.yml
@@ -8,7 +8,7 @@ on:
 
 jobs:
   check:
-    uses: Lightning-AI/utilities/.github/workflows/check-schema.yml@v0.12.0
+    uses: Lightning-AI/utilities/.github/workflows/check-schema.yml@v0.14.0
     with:
       # skip azure due to the wrong schema file by MSFT
       # https://github.com/Lightning-AI/lightning-flash/pull/1455#issuecomment-1244793607
diff --git a/docs/source-pytorch/visualize/loggers.rst b/docs/source-pytorch/visualize/loggers.rst
@@ -54,3 +54,37 @@ Track and Visualize Experiments
 
         </div>
     </div>
+
+.. _mlflow_logger:
+
+MLflow Logger
+-------------
+
+The MLflow logger in PyTorch Lightning now includes a `checkpoint_path_prefix` parameter. This parameter allows you to prefix the checkpoint artifact's path when logging checkpoints as artifacts.
+
+Example usage:
+
+.. code-block:: python
+
+    import lightning as L
+    from lightning.pytorch.loggers import MLFlowLogger
+
+    mlf_logger = MLFlowLogger(
+        experiment_name="lightning_logs",
+        tracking_uri="file:./ml-runs",
+        checkpoint_path_prefix="my_prefix"
+    )
+    trainer = L.Trainer(logger=mlf_logger)
+
+    # Your LightningModule definition
+    class LitModel(L.LightningModule):
+        def training_step(self, batch, batch_idx):
+            # example
+            self.logger.experiment.whatever_ml_flow_supports(...)
+
+        def any_lightning_module_function_or_hook(self):
+            self.logger.experiment.whatever_ml_flow_supports(...)
+
+    # Train your model
+    model = LitModel()
+    trainer.fit(model)
diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md
@@ -10,12 +10,25 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Allow LightningCLI to use a customized argument parser class ([#20596](https://github.com/Lightning-AI/pytorch-lightning/pull/20596))
 
+
 ### Changed
 
+- Added a new `checkpoint_path_prefix` parameter to the MLflow logger which can control the path to where the MLflow artifacts for the model checkpoints are stored.
+
+
 ### Removed
 
+-
+
+
 ### Fixed
 
+- Fix CSVLogger logging hyperparameter at every write which increase latency  ([#20594](https://github.com/Lightning-AI/pytorch-lightning/pull/20594))
+
+
+- Always call `WandbLogger.experiment` first in `_call_setup_hook` to ensure `tensorboard` logs can sync to `wandb` ([#20610](https://github.com/Lightning-AI/pytorch-lightning/pull/20610))
+
+
 ## [2.5.0] - 2024-12-19
 
 ### Added
diff --git a/src/lightning/pytorch/loggers/csv_logs.py b/src/lightning/pytorch/loggers/csv_logs.py
@@ -55,15 +55,10 @@ def __init__(self, log_dir: str) -> None:
         self.hparams: dict[str, Any] = {}
 
     def log_hparams(self, params: dict[str, Any]) -> None:
-        """Record hparams."""
+        """Record hparams and save into files."""
         self.hparams.update(params)
-
-    @override
-    def save(self) -> None:
-        """Save recorded hparams and metrics into files."""
         hparams_file = os.path.join(self.log_dir, self.NAME_HPARAMS_FILE)
         save_hparams_to_yaml(hparams_file, self.hparams)
-        return super().save()
 
 
 class CSVLogger(Logger, FabricCSVLogger):
@@ -144,7 +139,7 @@ def save_dir(self) -> str:
 
     @override
     @rank_zero_only
-    def log_hyperparams(self, params: Union[dict[str, Any], Namespace]) -> None:
+    def log_hyperparams(self, params: Optional[Union[dict[str, Any], Namespace]] = None) -> None:
         params = _convert_params(params)
         self.experiment.log_hparams(params)
 
diff --git a/src/lightning/pytorch/loggers/mlflow.py b/src/lightning/pytorch/loggers/mlflow.py
@@ -97,7 +97,7 @@ def any_lightning_module_function_or_hook(self):
               :paramref:`~lightning.pytorch.callbacks.Checkpoint.save_top_k` ``== -1``
               which also logs every checkpoint during training.
             * if ``log_model == False`` (default), no checkpoint is logged.
-
+        checkpoint_path_prefix: A string to prefix the checkpoint artifact's path.
         prefix: A string to put at the beginning of metric keys.
         artifact_location: The location to store run artifacts. If not provided, the server picks an appropriate
             default.
@@ -121,6 +121,7 @@ def __init__(
         tags: Optional[dict[str, Any]] = None,
         save_dir: Optional[str] = "./mlruns",
         log_model: Literal[True, False, "all"] = False,
+        checkpoint_path_prefix: str = "",
         prefix: str = "",
         artifact_location: Optional[str] = None,
         run_id: Optional[str] = None,
@@ -147,6 +148,7 @@ def __init__(
         self._artifact_location = artifact_location
         self._log_batch_kwargs = {} if synchronous is None else {"synchronous": synchronous}
         self._initialized = False
+        self._checkpoint_path_prefix = checkpoint_path_prefix
 
         from mlflow.tracking import MlflowClient
 
@@ -361,7 +363,7 @@ def _scan_and_log_checkpoints(self, checkpoint_callback: ModelCheckpoint) -> Non
             aliases = ["latest", "best"] if p == checkpoint_callback.best_model_path else ["latest"]
 
             # Artifact path on mlflow
-            artifact_path = Path(p).stem
+            artifact_path = Path(self._checkpoint_path_prefix) / Path(p).stem
 
             # Log the checkpoint
             self.experiment.log_artifact(self._run_id, p, artifact_path)
diff --git a/src/lightning/pytorch/trainer/call.py b/src/lightning/pytorch/trainer/call.py
@@ -21,6 +21,7 @@
 import lightning.pytorch as pl
 from lightning.fabric.utilities.device_dtype_mixin import _DeviceDtypeModuleMixin
 from lightning.pytorch.callbacks import Checkpoint, EarlyStopping
+from lightning.pytorch.loggers import WandbLogger
 from lightning.pytorch.strategies.launchers import _SubprocessScriptLauncher
 from lightning.pytorch.trainer.connectors.signal_connector import _get_sigkill_signal
 from lightning.pytorch.trainer.states import TrainerStatus
@@ -91,8 +92,12 @@ def _call_setup_hook(trainer: "pl.Trainer") -> None:
         if isinstance(module, _DeviceDtypeModuleMixin):
             module._device = trainer.strategy.root_device
 
+    # wandb.init must be called before any tensorboard writers are created in order to sync tensorboard logs to wandb:
+    # https://github.com/wandb/wandb/issues/1782#issuecomment-779161203
+    loggers = sorted(trainer.loggers, key=lambda logger: not isinstance(logger, WandbLogger))
+
     # Trigger lazy creation of experiment in loggers so loggers have their metadata available
-    for logger in trainer.loggers:
+    for logger in loggers:
         if hasattr(logger, "experiment"):
             _ = logger.experiment
 
diff --git a/tests/tests_fabric/utilities/test_seed.py b/tests/tests_fabric/utilities/test_seed.py
@@ -1,5 +1,6 @@
 import os
 import random
+import warnings
 from unittest import mock
 from unittest.mock import Mock
 
@@ -30,9 +31,9 @@ def test_seed_stays_same_with_multiple_seed_everything_calls():
         seed_everything()
     initial_seed = os.environ.get("PL_GLOBAL_SEED")
 
-    with pytest.warns(None) as record:
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
         seed_everything()
-    assert not record  # does not warn
     seed = os.environ.get("PL_GLOBAL_SEED")
 
     assert initial_seed == seed
diff --git a/tests/tests_pytorch/callbacks/test_lr_monitor.py b/tests/tests_pytorch/callbacks/test_lr_monitor.py
@@ -548,10 +548,10 @@ def finetune_function(self, pl_module, epoch: int, optimizer):
             """Called when the epoch begins."""
             if epoch == 1 and isinstance(optimizer, torch.optim.SGD):
                 self.unfreeze_and_add_param_group(pl_module.backbone[0], optimizer, lr=0.1)
-            if epoch == 2 and isinstance(optimizer, torch.optim.Adam):
+            if epoch == 2 and type(optimizer) is torch.optim.Adam:
                 self.unfreeze_and_add_param_group(pl_module.layer, optimizer, lr=0.1)
 
-            if epoch == 3 and isinstance(optimizer, torch.optim.Adam):
+            if epoch == 3 and type(optimizer) is torch.optim.Adam:
                 assert len(optimizer.param_groups) == 2
                 self.unfreeze_and_add_param_group(pl_module.backbone[1], optimizer, lr=0.1)
                 assert len(optimizer.param_groups) == 3
diff --git a/tests/tests_pytorch/core/test_results.py b/tests/tests_pytorch/core/test_results.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 from functools import partial
 
+import pytest
 import torch
 import torch.distributed as dist
 
@@ -48,6 +49,8 @@ def result_reduce_ddp_fn(strategy):
     assert actual.item() == dist.get_world_size()
 
 
+# flaky with "process 0 terminated with signal SIGABRT"
+@pytest.mark.flaky(reruns=3, only_rerun="torch.multiprocessing.spawn.ProcessExitedException")
 @RunIf(skip_windows=True)
 def test_result_reduce_ddp():
     spawn_launch(result_reduce_ddp_fn, [torch.device("cpu")] * 2)
diff --git a/tests/tests_pytorch/loggers/test_csv.py b/tests/tests_pytorch/loggers/test_csv.py
@@ -75,7 +75,6 @@ def test_named_version(tmp_path):
 
     logger = CSVLogger(save_dir=tmp_path, name=exp_name, version=expected_version)
     logger.log_hyperparams({"a": 1, "b": 2})
-    logger.save()
     assert logger.version == expected_version
     assert os.listdir(tmp_path / exp_name) == [expected_version]
     assert os.listdir(tmp_path / exp_name / expected_version)
@@ -85,7 +84,7 @@ def test_named_version(tmp_path):
 def test_no_name(tmp_path, name):
     """Verify that None or empty name works."""
     logger = CSVLogger(save_dir=tmp_path, name=name)
-    logger.save()
+    logger.log_hyperparams()
     assert os.path.normpath(logger.root_dir) == str(tmp_path)  # use os.path.normpath to handle trailing /
     assert os.listdir(tmp_path / "version_0")
 
@@ -116,7 +115,6 @@ def test_log_hyperparams(tmp_path):
         "layer": torch.nn.BatchNorm1d,
     }
     logger.log_hyperparams(hparams)
-    logger.save()
 
     path_yaml = os.path.join(logger.log_dir, ExperimentWriter.NAME_HPARAMS_FILE)
     params = load_hparams_from_yaml(path_yaml)
diff --git a/tests/tests_pytorch/loggers/test_mlflow.py b/tests/tests_pytorch/loggers/test_mlflow.py
@@ -427,3 +427,33 @@ def test_set_tracking_uri(mlflow_mock):
     mlflow_mock.set_tracking_uri.assert_not_called()
     _ = logger.experiment
     mlflow_mock.set_tracking_uri.assert_called_with("the_tracking_uri")
+
+
+@mock.patch("lightning.pytorch.loggers.mlflow._get_resolve_tags", Mock())
+def test_mlflow_log_model_with_checkpoint_path_prefix(mlflow_mock, tmp_path):
+    """Test that the logger creates the folders and files in the right place with a prefix."""
+    client = mlflow_mock.tracking.MlflowClient
+
+    # Get model, logger, trainer and train
+    model = BoringModel()
+    logger = MLFlowLogger("test", save_dir=str(tmp_path), log_model="all", checkpoint_path_prefix="my_prefix")
+    logger = mock_mlflow_run_creation(logger, experiment_id="test-id")
+
+    trainer = Trainer(
+        default_root_dir=tmp_path,
+        logger=logger,
+        max_epochs=2,
+        limit_train_batches=3,
+        limit_val_batches=3,
+    )
+    trainer.fit(model)
+
+    # Checkpoint log
+    assert client.return_value.log_artifact.call_count == 2
+    # Metadata and aliases log
+    assert client.return_value.log_artifacts.call_count == 2
+
+    # Check that the prefix is used in the artifact path
+    for call in client.return_value.log_artifact.call_args_list:
+        args, _ = call
+        assert str(args[2]).startswith("my_prefix")
diff --git a/tests/tests_pytorch/loggers/test_wandb.py b/tests/tests_pytorch/loggers/test_wandb.py
@@ -24,7 +24,7 @@
 from lightning.pytorch.callbacks import ModelCheckpoint
 from lightning.pytorch.cli import LightningCLI
 from lightning.pytorch.demos.boring_classes import BoringModel
-from lightning.pytorch.loggers import WandbLogger
+from lightning.pytorch.loggers import TensorBoardLogger, WandbLogger
 from lightning.pytorch.utilities.exceptions import MisconfigurationException
 from tests_pytorch.test_cli import _xfail_python_ge_3_11_9
 
@@ -133,6 +133,43 @@ def test_wandb_logger_init_before_spawn(wandb_mock):
     assert logger._experiment is not None
 
 
+def test_wandb_logger_experiment_called_first(wandb_mock, tmp_path):
+    wandb_experiment_called = False
+
+    def tensorboard_experiment_side_effect() -> mock.MagicMock:
+        nonlocal wandb_experiment_called
+        assert wandb_experiment_called
+        return mock.MagicMock()
+
+    def wandb_experiment_side_effect() -> mock.MagicMock:
+        nonlocal wandb_experiment_called
+        wandb_experiment_called = True
+        return mock.MagicMock()
+
+    with (
+        mock.patch.object(
+            TensorBoardLogger,
+            "experiment",
+            new_callable=lambda: mock.PropertyMock(side_effect=tensorboard_experiment_side_effect),
+        ),
+        mock.patch.object(
+            WandbLogger,
+            "experiment",
+            new_callable=lambda: mock.PropertyMock(side_effect=wandb_experiment_side_effect),
+        ),
+    ):
+        model = BoringModel()
+        trainer = Trainer(
+            default_root_dir=tmp_path,
+            log_every_n_steps=1,
+            limit_train_batches=0,
+            limit_val_batches=0,
+            max_steps=1,
+            logger=[TensorBoardLogger(tmp_path), WandbLogger(save_dir=tmp_path)],
+        )
+        trainer.fit(model)
+
+
 def test_wandb_pickle(wandb_mock, tmp_path):
     """Verify that pickling trainer with wandb logger works.
 
