Unverified Commit e644b600 authored by Alexander Jipa's avatar Alexander Jipa Committed by GitHub
Browse files

fix: get mlflow version from mlflow-skinny (#29918)


Co-authored-by: default avatarAlexander Jipa <azzhipa@amazon.com>
parent 156d30da
...@@ -131,6 +131,13 @@ def is_mlflow_available(): ...@@ -131,6 +131,13 @@ def is_mlflow_available():
return importlib.util.find_spec("mlflow") is not None return importlib.util.find_spec("mlflow") is not None
def get_mlflow_version():
try:
return importlib.metadata.version("mlflow")
except importlib.metadata.PackageNotFoundError:
return importlib.metadata.version("mlflow-skinny")
def is_dagshub_available(): def is_dagshub_available():
return None not in [importlib.util.find_spec("dagshub"), importlib.util.find_spec("mlflow")] return None not in [importlib.util.find_spec("dagshub"), importlib.util.find_spec("mlflow")]
...@@ -1002,7 +1009,7 @@ class MLflowCallback(TrainerCallback): ...@@ -1002,7 +1009,7 @@ class MLflowCallback(TrainerCallback):
# "synchronous" flag is only available with mlflow version >= 2.8.0 # "synchronous" flag is only available with mlflow version >= 2.8.0
# https://github.com/mlflow/mlflow/pull/9705 # https://github.com/mlflow/mlflow/pull/9705
# https://github.com/mlflow/mlflow/releases/tag/v2.8.0 # https://github.com/mlflow/mlflow/releases/tag/v2.8.0
if packaging.version.parse(importlib.metadata.version("mlflow")) >= packaging.version.parse("2.8.0"): if packaging.version.parse(get_mlflow_version()) >= packaging.version.parse("2.8.0"):
self._async_log = True self._async_log = True
logger.debug( logger.debug(
f"MLflow experiment_name={self._experiment_name}, run_name={args.run_name}, nested={self._nested_run}," f"MLflow experiment_name={self._experiment_name}, run_name={args.run_name}, nested={self._nested_run},"
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment