Merge branch 'master' into warn-on-train-step-returning-none

amorehead · web-flow · commit 82eb3f2ae9da · 2025-04-24T18:55:38.000-05:00
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
@@ -18,3 +18,6 @@
 /README.md                           @williamfalcon @lantiga
 /src/pytorch_lightning/__about__.py  @williamfalcon @lantiga @borda
 /src/lightning_fabric/__about__.py   @williamfalcon @lantiga @borda
+
+/src/lightning/fabric/loggers        @williamfalcon
+/src/lightning/pytorch/loggers       @williamfalcon
diff --git a/.github/workflows/_legacy-checkpoints.yml b/.github/workflows/_legacy-checkpoints.yml
@@ -113,6 +113,7 @@ jobs:
 
       - run: pip install -r requirements/ci.txt
       - name: Upload checkpoints to S3
+        if: secrets.AWS_REGION != ''
         working-directory: ${{ env.LEGACY_FOLDER }}
         env:
           AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY }}
diff --git a/docs/source-pytorch/visualize/loggers.rst b/docs/source-pytorch/visualize/loggers.rst
@@ -54,37 +54,3 @@ 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/requirements/ci.txt b/requirements/ci.txt
@@ -5,4 +5,4 @@ twine ==6.0.1
 importlib-metadata <8.0.0
 wget
 pkginfo ==1.12.0
-packaging <24.2
+packaging <25.1
diff --git a/requirements/docs.txt b/requirements/docs.txt
@@ -17,7 +17,7 @@ sphinx-rtd-dark-mode
 sphinxcontrib-video ==0.2.0
 jinja2 <3.2.0
 
-lightning-utilities >=0.11.1, <0.12.0
+lightning-utilities >=0.11.1, <0.15.0
 
 # installed from S3 location and fetched in advance
 lai-sphinx-theme
diff --git a/requirements/doctests.txt b/requirements/doctests.txt
@@ -1,2 +1,2 @@
-pytest ==7.4.0
-pytest-doctestplus ==1.0.0
+pytest ==8.3.5
+pytest-doctestplus ==1.4.0
diff --git a/requirements/fabric/base.txt b/requirements/fabric/base.txt
@@ -1,8 +1,8 @@
 # NOTE: the upper bound for the package version is only set for CI stability, and it is dropped while installing this package
 #  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
 
-torch >=2.1.0, <2.6.0
+torch >=2.1.0, <2.7.0
 fsspec[http] >=2022.5.0, <2024.4.0
-packaging >=20.0, <=23.1
+packaging >=20.0, <=25.0
 typing-extensions >=4.4.0, <4.11.0
-lightning-utilities >=0.10.0, <0.12.0
+lightning-utilities >=0.10.0, <0.15.0
diff --git a/requirements/fabric/examples.txt b/requirements/fabric/examples.txt
@@ -1,6 +1,6 @@
 # NOTE: the upper bound for the package version is only set for CI stability, and it is dropped while installing this package
 #  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
 
-torchvision >=0.16.0, <0.21.0
-torchmetrics >=0.10.0, <1.5.0
-lightning-utilities >=0.8.0, <0.12.0
+torchvision >=0.16.0, <0.22.0
+torchmetrics >=0.10.0, <1.8.0
+lightning-utilities >=0.8.0, <0.15.0
diff --git a/requirements/fabric/test.txt b/requirements/fabric/test.txt
@@ -1,6 +1,6 @@
 coverage ==7.3.1
 numpy >=1.17.2, <1.27.0
-pytest ==7.4.0
+pytest ==8.3.5
 pytest-cov ==4.1.0
 pytest-timeout ==2.1.0
 pytest-rerunfailures ==12.0
diff --git a/requirements/pytorch/base.txt b/requirements/pytorch/base.txt
@@ -1,11 +1,11 @@
 # NOTE: the upper bound for the package version is only set for CI stability, and it is dropped while installing this package
 #  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
 
-torch >=2.1.0, <2.6.0
+torch >=2.1.0, <2.7.0
 tqdm >=4.57.0, <4.67.0
 PyYAML >=5.4, <6.1.0
 fsspec[http] >=2022.5.0, <2024.4.0
-torchmetrics >=0.7.0, <1.5.0
-packaging >=20.0, <=23.1
+torchmetrics >=0.7.0, <1.8.0
+packaging >=20.0, <=25.0
 typing-extensions >=4.4.0, <4.11.0
-lightning-utilities >=0.10.0, <0.12.0
+lightning-utilities >=0.10.0, <0.15.0
diff --git a/requirements/pytorch/examples.txt b/requirements/pytorch/examples.txt
@@ -2,7 +2,7 @@
 #  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
 
 requests <2.32.0
-torchvision >=0.16.0, <0.21.0
-ipython[all] <8.15.0
-torchmetrics >=0.10.0, <1.5.0
-lightning-utilities >=0.8.0, <0.12.0
+torchvision >=0.16.0, <0.22.0
+ipython[all] <8.19.0
+torchmetrics >=0.10.0, <1.8.0
+lightning-utilities >=0.8.0, <0.15.0
diff --git a/requirements/pytorch/test.txt b/requirements/pytorch/test.txt
@@ -1,19 +1,19 @@
 coverage ==7.3.1
-pytest ==7.4.0
+pytest ==8.3.5
 pytest-cov ==4.1.0
 pytest-timeout ==2.1.0
 pytest-rerunfailures ==12.0
 pytest-random-order ==1.1.0
 
 # needed in tests
-cloudpickle >=1.3, <2.3.0
+cloudpickle >=1.3, <3.2.0
 scikit-learn >0.22.1, <1.4.0
 numpy >=1.17.2, <1.27.0
 onnx >=1.12.0, <1.17.0
-onnxruntime >=1.12.0, <1.19.0
+onnxruntime >=1.12.0, <1.21.0
 psutil <5.9.6 # for `DeviceStatsMonitor`
 pandas >1.0, <2.2.0  # needed in benchmarks
 fastapi  # for `ServableModuleValidator`  # not setting version as re-defined in App
 uvicorn  # for `ServableModuleValidator`  # not setting version as re-defined in App
 
-tensorboard >=2.9.1, <2.15.0  # for `TensorBoardLogger`
+tensorboard >=2.9.1, <2.20.0  # for `TensorBoardLogger`
diff --git a/requirements/typing.txt b/requirements/typing.txt
@@ -1,5 +1,5 @@
 mypy==1.11.0
-torch==2.5.1
+torch==2.6.0
 
 types-Markdown
 types-PyYAML
diff --git a/src/lightning/fabric/plugins/collectives/torch_collective.py b/src/lightning/fabric/plugins/collectives/torch_collective.py
@@ -50,7 +50,7 @@ def world_size(self) -> int:
 
     @override
     def broadcast(self, tensor: Tensor, src: int) -> Tensor:
-        dist.broadcast(tensor, src, group=self.group)
+        dist.broadcast(tensor, src, group=self.group)  # type: ignore[arg-type]
         return tensor
 
     @override
@@ -62,7 +62,7 @@ def all_reduce(self, tensor: Tensor, op: Union[str, ReduceOp, RedOpType] = "sum"
     @override
     def reduce(self, tensor: Tensor, dst: int, op: Union[str, ReduceOp, RedOpType] = "sum") -> Tensor:
         op = self._convert_to_native_op(op)
-        dist.reduce(tensor, dst, op=op, group=self.group)
+        dist.reduce(tensor, dst, op=op, group=self.group)  # type: ignore[arg-type]
         return tensor
 
     @override
@@ -72,12 +72,12 @@ def all_gather(self, tensor_list: list[Tensor], tensor: Tensor) -> list[Tensor]:
 
     @override
     def gather(self, tensor: Tensor, gather_list: list[Tensor], dst: int = 0) -> list[Tensor]:
-        dist.gather(tensor, gather_list, dst, group=self.group)
+        dist.gather(tensor, gather_list, dst, group=self.group)  # type: ignore[arg-type]
         return gather_list
 
     @override
     def scatter(self, tensor: Tensor, scatter_list: list[Tensor], src: int = 0) -> Tensor:
-        dist.scatter(tensor, scatter_list, src, group=self.group)
+        dist.scatter(tensor, scatter_list, src, group=self.group)  # type: ignore[arg-type]
         return tensor
 
     @override
@@ -109,27 +109,27 @@ def all_gather_object(self, object_list: list[Any], obj: Any) -> list[Any]:
     def broadcast_object_list(
         self, object_list: list[Any], src: int, device: Optional[torch.device] = None
     ) -> list[Any]:
-        dist.broadcast_object_list(object_list, src, group=self.group, device=device)
+        dist.broadcast_object_list(object_list, src, group=self.group, device=device)  # type: ignore[arg-type]
         return object_list
 
     def gather_object(self, obj: Any, object_gather_list: list[Any], dst: int = 0) -> list[Any]:
-        dist.gather_object(obj, object_gather_list, dst, group=self.group)
+        dist.gather_object(obj, object_gather_list, dst, group=self.group)  # type: ignore[arg-type]
         return object_gather_list
 
     def scatter_object_list(
         self, scatter_object_output_list: list[Any], scatter_object_input_list: list[Any], src: int = 0
     ) -> list[Any]:
-        dist.scatter_object_list(scatter_object_output_list, scatter_object_input_list, src, group=self.group)
+        dist.scatter_object_list(scatter_object_output_list, scatter_object_input_list, src, group=self.group)  # type: ignore[arg-type]
         return scatter_object_output_list
 
     @override
     def barrier(self, device_ids: Optional[list[int]] = None) -> None:
         if self.group == dist.GroupMember.NON_GROUP_MEMBER:
             return
-        dist.barrier(group=self.group, device_ids=device_ids)
+        dist.barrier(group=self.group, device_ids=device_ids)  # type: ignore[arg-type]
 
     def monitored_barrier(self, timeout: Optional[datetime.timedelta] = None, wait_all_ranks: bool = False) -> None:
-        dist.monitored_barrier(group=self.group, timeout=timeout, wait_all_ranks=wait_all_ranks)
+        dist.monitored_barrier(group=self.group, timeout=timeout, wait_all_ranks=wait_all_ranks)  # type: ignore[arg-type]
 
     @override
     def setup(self, main_address: Optional[str] = None, main_port: Optional[str] = None, **kwargs: Any) -> Self:
diff --git a/src/lightning/fabric/strategies/xla_fsdp.py b/src/lightning/fabric/strategies/xla_fsdp.py
@@ -295,6 +295,7 @@ def clip_gradients_norm(
     ) -> Tensor:
         """Clip gradients by norm."""
         self.precision.unscale_gradients(optimizer)
+        assert callable(module.clip_grad_norm_)
         return module.clip_grad_norm_(max_norm=max_norm, norm_type=norm_type)
 
     @override
diff --git a/src/lightning/fabric/utilities/init.py b/src/lightning/fabric/utilities/init.py
@@ -67,7 +67,8 @@ def _materialize(module: Module, device: _DEVICE) -> None:
             f"Materialization requires that the `{type(module).__name__}.reset_parameters` method is implemented."
             " This method is used to initialize any children parameters or buffers in this module."
         )
-    module.reset_parameters()
+    if callable(module.reset_parameters):
+        module.reset_parameters()
 
 
 def _materialize_meta_tensors(module: Module, device: _DEVICE) -> None:
diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md
@@ -35,7 +35,6 @@ 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))
 - Change `wandb` default x-axis to `tensorboard`'s `global_step` when `sync_tensorboard=True` ([#20611](https://github.com/Lightning-AI/pytorch-lightning/pull/20611))
-- 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 ([#20538](https://github.com/Lightning-AI/pytorch-lightning/pull/20538))
 - CometML logger was updated to support the recent Comet SDK ([#20275](https://github.com/Lightning-AI/pytorch-lightning/pull/20275))
 - bump: testing with latest `torch` 2.6 ([#20509](https://github.com/Lightning-AI/pytorch-lightning/pull/20509))
 
diff --git a/src/lightning/pytorch/callbacks/finetuning.py b/src/lightning/pytorch/callbacks/finetuning.py
@@ -133,7 +133,7 @@ def flatten_modules(modules: Union[Module, Iterable[Union[Module, Iterable]]]) -
 
         if isinstance(modules, Iterable):
             _flatten_modules = []
-            for m in modules:  # type: ignore[union-attr]
+            for m in modules:
                 _flatten_modules.extend(BaseFinetuning.flatten_modules(m))
 
             _modules = iter(_flatten_modules)
diff --git a/src/lightning/pytorch/callbacks/throughput_monitor.py b/src/lightning/pytorch/callbacks/throughput_monitor.py
@@ -140,7 +140,7 @@ def _update(self, trainer: "Trainer", pl_module: "LightningModule", batch: Any,
             # this assumes that all iterations used the same batch size
             samples=iter_num * batch_size,
             lengths=None if self.length_fn is None else self._lengths[stage],
-            flops=flops_per_batch,
+            flops=flops_per_batch,  # type: ignore[arg-type]
         )
 
     def _compute(self, trainer: "Trainer", iter_num: Optional[int] = None) -> None:
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,7 +121,6 @@ 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,
@@ -148,7 +147,6 @@ 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
 
@@ -363,7 +361,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(self._checkpoint_path_prefix, Path(p).stem).as_posix()
+            artifact_path = Path(p).stem
 
             # Log the checkpoint
             self.experiment.log_artifact(self._run_id, p, artifact_path)
diff --git a/src/lightning/pytorch/overrides/distributed.py b/src/lightning/pytorch/overrides/distributed.py
@@ -163,9 +163,7 @@ def _register_ddp_comm_hook(
 
 def _sync_module_states(module: torch.nn.Module) -> None:
     """Taken from https://github.com/pytorch/pytorch/blob/v2.0.0/torch/nn/parallel/distributed.py#L675-L682."""
-    parameters_to_ignore = (
-        set(module._ddp_params_and_buffers_to_ignore) if hasattr(module, "_ddp_params_and_buffers_to_ignore") else set()
-    )
+    parameters_to_ignore = set(getattr(module, "_ddp_params_and_buffers_to_ignore", []))
     from torch.distributed.distributed_c10d import _get_default_group
     from torch.distributed.utils import _sync_module_states as torch_sync_module_states
 
diff --git a/src/lightning/pytorch/plugins/precision/double.py b/src/lightning/pytorch/plugins/precision/double.py
@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from collections.abc import Generator
+from collections.abc import Generator, Iterable
 from contextlib import AbstractContextManager, contextmanager
 from typing import Any, Literal
 
@@ -72,6 +72,8 @@ class LightningDoublePrecisionModule(_DeviceDtypeModuleMixin, nn.Module):
 
     """
 
+    _ddp_params_and_buffers_to_ignore: Iterable[str]
+
     def __init__(self, pl_module: "pl.LightningModule") -> None:
         super().__init__()
         rank_zero_deprecation(
diff --git a/tests/tests_pytorch/loggers/test_mlflow.py b/tests/tests_pytorch/loggers/test_mlflow.py
@@ -427,33 +427,3 @@ 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")
