Support remote Databricks model registries in mlflow.<flavor>.load_model

[sueann]

Signed-off-by: Sue Ann Hong sueann@databricks.com

What changes are proposed in this pull request?

Support remote Databricks model registries in mlflow.<flavor>.load_model when called with a models:/ URI.

Currently, the following two ways to call mlflow..load_model from a remote Databricks model registry work:

mlflow.pyfunc.load_model('models://profile@databricks/model_name/Staging')

mlflow.set_tracking_uri(registry_uri)
mlflow.pyfunc.load_model('models:/model_name/Staging')

But the following, more natural, way does not:

mlflow.set_registry_uri(registry_uri)
mlflow.pyfunc.load_model('models:/model_name/Staging')

This is because each flavor’s load_model
calls _download_artifact_from_uri
which calls get_artifact_repository(artifact_uri=root_uri).download_artifacts
and no registry server information is passed into the eventual DBFS artifact repository instance. So then the DBFS artifact repository uses the tracking URI from the context (i.e. set via the global variable or the environment variable).

To fix this, here we have ModelsArtifactRepository automatically use the Databricks model registry server information if specified in the context.

How is this patch tested?

• Unit tests
• Manual tests

mlflow.set_registry_uri('databricks://blah')  # invalid but ignored
mlflow.set_tracking_uri('databricks://blurgh')  # invalid but ignored
model = mlflow.pyfunc.load_model(f'models://{profile}@databricks/{model_name}/Staging')
model.predict(1)

mlflow.set_registry_uri(f'databricks://{profile}')  # valid registry URI
mlflow.set_tracking_uri('databricks://blurgh')  # invalid but ignored
model = mlflow.pyfunc.load_model(f'models:/{model_name}/Staging')
model.predict(1)

mlflow.set_registry_uri(None)
mlflow.set_tracking_uri(f'databricks://{profile}')  # valid registry URI
model = mlflow.pyfunc.load_model(f'models:/{model_name}/Staging')
model.predict(1)

Release Notes

Is this a user-facing change?

• No. You can skip the rest of this section.
• Yes. Give a description of this change to be included in the release notes for MLflow users.

mlflow.<flavor>.load_model methods will fetch the model from the Databricks model registry specified by mlflow.set_registry_uri if it is set to one.

What component(s), interfaces, languages, and integrations does this PR affect?

Components

• area/artifacts: Artifact stores and artifact logging
• area/build: Build and test infrastructure for MLflow
• area/docs: MLflow documentation pages
• area/examples: Example code
• area/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registry
• area/models: MLmodel format, model serialization/deserialization, flavors
• area/projects: MLproject format, project running backends
• area/scoring: Local serving, model deployment tools, spark UDFs
• area/server-infra: MLflow server, JavaScript dev server
• area/tracking: Tracking Service, tracking client APIs, autologging

Interface

• area/uiux: Front-end, user experience, JavaScript, plotting
• area/docker: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
• area/sqlalchemy: Use of SQLAlchemy in the Tracking Service or Model Registry
• area/windows: Windows support

Language

• language/r: R APIs and clients
• language/java: Java APIs and clients
• language/new: Proposals for new client languages

Integrations

• integrations/azure: Azure and Azure ML integrations
• integrations/sagemaker: SageMaker integrations
• integrations/databricks: Databricks integrations

How should the PR be classified in the release notes? Choose one:

• rn/breaking-change - The PR will be mentioned in the "Breaking Changes" section
• rn/none - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
• rn/feature - A new user-facing feature worth mentioning in the release notes
• rn/bug-fix - A user-facing bug fix worth mentioning in the release notes
• rn/documentation - A user-facing documentation change worth mentioning in the release notes

[ankit-db]

Tried it out, and it works! Small note about adding a test - I tried a few of the cases, but would be good to codify in the test

[ankit-db]

Could we also just verify that the right error is thrown if the registry URI is malformed?

[sueann]

@ankitmathur-db is something like this what you were thinking? please do let me know if there are other cases you were thinking of I didn't cover here. thanks!

[ankit-db]

Yeah this is what I was thinking. Have we tested stuff like databricks://scope:key. (with the extra spaces at the end)?

[sueann]

Ah we don't test it - but that should be handled at a lower level. Created a PR to address that: #3338.

[ankit-db]

Looks great!