Suggestions

Author: carmocca

src/pytorch_lightning/CHANGELOG.md

@@ -11,6 +11,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added an argument `include_cuda` in `pytorch_lightning.utilities.seed.isolate_rng` to disable managing `torch.cuda`'s rng ([#16423](https://github.com/Lightning-AI/lightning/pull/16423))
 
+- Added support for running the `MLFlowLogger` with the `mlflow-skinny` package ([16513](https://github.com/Lightning-AI/lightning/pull/16513))
 
 ### Changed
 
@@ -152,8 +153,6 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed strict availability check for `torch_xla` requirement ([#16476](https://github.com/Lightning-AI/lightning/pull/16476))
 
-- Fixed being able to use `MLFlowLogger` with the `mlflow-skinny` dependency installed ([16513](https://github.com/Lightning-AI/lightning/pull/16513))
-
 
 ## [1.9.0] - 2023-01-17
 

src/pytorch_lightning/loggers/mlflow.py

@@ -37,7 +37,9 @@
 
 log = logging.getLogger(__name__)
 LOCAL_FILE_URI_PREFIX = "file:"
-_MLFLOW_AVAILABLE = RequirementCache("mlflow>=1.0.0") or RequirementCache("mlflow-skinny>=1.0.0")
+_MLFLOW_FULL_AVAILABLE = RequirementCache("mlflow>=1.0.0")
+_MLFLOW_SKINNY_AVAILABLE = RequirementCache("mlflow-skinny>=1.0.0")
+_MLFLOW_AVAILABLE = _MLFLOW_FULL_AVAILABLE or _MLFLOW_SKINNY_AVAILABLE
 if _MLFLOW_AVAILABLE:
     from mlflow.entities import Metric, Param
     from mlflow.tracking import context, MlflowClient
@@ -78,7 +80,7 @@ class MLFlowLogger(Logger):
 
     .. code-block:: bash
 
-        pip install mlflow
+        pip install mlflow  # or mlflow-skinny
 
     .. code-block:: python
 
@@ -148,7 +150,9 @@ def __init__(
         run_id: Optional[str] = None,
     ):
         if not _MLFLOW_AVAILABLE:
-            raise ModuleNotFoundError(str(_MLFLOW_AVAILABLE))
+            raise ModuleNotFoundError(
+                f"{_MLFLOW_FULL_AVAILABLE!s}. You can also try {_MLFLOW_SKINNY_AVAILABLE.requirement!r}"
+            )
         super().__init__()
         if not tracking_uri:
             tracking_uri = f"{LOCAL_FILE_URI_PREFIX}{save_dir}"