Make MLFlow version detection more robust and handles mlflow-skinny

src/transformers/integrations/integration_utils.py

@@ -131,6 +131,19 @@ def is_mlflow_available():
     return importlib.util.find_spec("mlflow") is not None
 
 
+def get_mlflow_version():
+    # MLFlow can also be provided from mlflow-skinny package, which has the same versioning of mlflow package.
+    for mlflow_package_name in ["mlflow", "mlflow-skinny"]:
+        try:
+            return importlib.metadata.version(mlflow_package_name)
+        except importlib.metadata.PackageNotFoundError:
+            # We will try different mlflow package candidates
+            pass
+
+    # Unable to determine the mlflow version
+    return None
+
+
 def is_dagshub_available():
     return None not in [importlib.util.find_spec("dagshub"), importlib.util.find_spec("mlflow")]
 
@@ -1002,7 +1015,8 @@ def setup(self, args, state, model):
         # "synchronous" flag is only available with mlflow version >= 2.8.0
         # https://github.com/mlflow/mlflow/pull/9705
         # https://github.com/mlflow/mlflow/releases/tag/v2.8.0
-        if packaging.version.parse(importlib.metadata.version("mlflow")) >= packaging.version.parse("2.8.0"):
+        mlflow_version = get_mlflow_version()
+        if mlflow_version is not None and packaging.version.parse(mlflow_version) >= packaging.version.parse("2.8.0"):
             self._async_log = True
         logger.debug(
             f"MLflow experiment_name={self._experiment_name}, run_name={args.run_name}, nested={self._nested_run},"