diff --git a/mlflow_export_import/common/iterators.py b/mlflow_export_import/common/iterators.py index 38b3d621..ebfd6131 100644 --- a/mlflow_export_import/common/iterators.py +++ b/mlflow_export_import/common/iterators.py @@ -1,4 +1,5 @@ import mlflow +from packaging import version class BaseIterator(): @@ -14,7 +15,7 @@ def __init__(self, search_method, max_results=None, filter=None): self.kwargs = { "max_results": self.max_results } if self.max_results else {} def _call_iter(self): - if mlflow.__version__ < "2.2.1": + if version.parse(mlflow.__version__) < version.parse("2.2.1"): return self.search_method(filter_string=self.filter) #7623 - https://mlflow.org/docs/2.1.1/python_api/mlflow.client.html else: return self.search_method(filter_string=self.filter, **self.kwargs) # https://mlflow.org/docs/latest/python_api/mlflow.client.html