Fallback to module available check for mlflow (#17467)

(cherry picked from commit b0af0ee)

src/lightning/pytorch/loggers/mlflow.py

@@ -25,7 +25,7 @@
 from typing import Any, Dict, List, Literal, Mapping, Optional, Union
 
 import yaml
-from lightning_utilities.core.imports import RequirementCache
+from lightning_utilities.core.imports import module_available, RequirementCache
 from torch import Tensor
 
 from lightning.fabric.utilities.logger import _add_prefix, _convert_params, _flatten_dict
@@ -36,9 +36,8 @@
 
 log = logging.getLogger(__name__)
 LOCAL_FILE_URI_PREFIX = "file:"
-_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
+_MLFLOW_AVAILABLE = RequirementCache("mlflow>=1.0.0") or module_available("mlflow")
+
 if _MLFLOW_AVAILABLE:
     from mlflow.entities import Metric, Param
     from mlflow.tracking import context, MlflowClient
@@ -149,9 +148,7 @@ def __init__(
         run_id: Optional[str] = None,
     ):
         if not _MLFLOW_AVAILABLE:
-            raise ModuleNotFoundError(
-                f"{_MLFLOW_FULL_AVAILABLE!s}. You can also try {_MLFLOW_SKINNY_AVAILABLE.requirement!r}"
-            )
+            raise ModuleNotFoundError(str(_MLFLOW_AVAILABLE))
         super().__init__()
         if not tracking_uri:
             tracking_uri = f"{LOCAL_FILE_URI_PREFIX}{save_dir}"