diff --git a/ingestion/setup.py b/ingestion/setup.py index b13c0aebf96..93d70089b16 100644 --- a/ingestion/setup.py +++ b/ingestion/setup.py @@ -197,7 +197,7 @@ plugins: Dict[str, Set[str]] = { "kinesis": {VERSIONS["boto3"]}, "ldap-users": {"ldap3==2.9.1"}, "looker": {"looker-sdk>=22.20.0", "lkml~=1.3"}, - "mlflow": {"mlflow-skinny~=1.30", "alembic~=1.10.2"}, + "mlflow": {"mlflow-skinny>=2.3.0", "alembic~=1.10.2"}, "mongo": {"pymongo~=4.3", VERSIONS["pandas"]}, "couchbase": {"couchbase~=4.1"}, "mssql": {"sqlalchemy-pytds~=0.3"}, diff --git a/ingestion/src/metadata/ingestion/source/mlmodel/mlflow/connection.py b/ingestion/src/metadata/ingestion/source/mlmodel/mlflow/connection.py index a43bc067fa7..b12c3a22176 100644 --- a/ingestion/src/metadata/ingestion/source/mlmodel/mlflow/connection.py +++ b/ingestion/src/metadata/ingestion/source/mlmodel/mlflow/connection.py @@ -47,7 +47,7 @@ def test_connection( of a metadata workflow or during an Automation Workflow """ - test_fn = {"GetModels": client.list_registered_models} + test_fn = {"GetModels": client.search_registered_models} test_connection_steps( metadata=metadata, diff --git a/ingestion/src/metadata/ingestion/source/mlmodel/mlflow/metadata.py b/ingestion/src/metadata/ingestion/source/mlmodel/mlflow/metadata.py index c8e49bf916c..dd433fcbe57 100644 --- a/ingestion/src/metadata/ingestion/source/mlmodel/mlflow/metadata.py +++ b/ingestion/src/metadata/ingestion/source/mlmodel/mlflow/metadata.py @@ -69,7 +69,7 @@ class MlflowSource(MlModelServiceSource): """ List and filters models from the registry """ - for model in cast(RegisteredModel, self.client.list_registered_models()): + for model in cast(RegisteredModel, self.client.search_registered_models()): if filter_by_mlmodel( self.source_config.mlModelFilterPattern, mlmodel_name=model.name ): @@ -92,8 +92,8 @@ class MlflowSource(MlModelServiceSource): self.status.failed( StackTraceError( name=model.name, - error="Invalid Version", - stack_trace=f"Cannot find latest version from version list {model.latest_versions}", + error="Version not found", + stack_trace=f"Unable to ingest model {model.name} due to missing version from version list {model.latest_versions}", # pylint: disable=line-too-long ) ) continue