Allow MLFlowLogger to work with mlflow-skinny (#16513)

Co-authored-by: Carlos Mocholí <[email protected]>
Fixes #16486

Author: dconathan

src/pytorch_lightning/CHANGELOG.md

@@ -15,6 +15,8 @@ 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))
+
 - Added a `Trainer.received_sigterm` property to check whether a SIGTERM signal was received ([#16501](https://github.com/Lightning-AI/lightning/pull/16501))
 
 ### Changed
@@ -153,7 +155,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))
 
 
-
 ## [1.9.0] - 2023-01-17
 
 ### Added

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")
+_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}"