Can view Model Registry using a Service Principal, but cannot load the model for inference.

JC3
New Contributor II

‎06-04-2024 07:39 AM

I have a Service Principal (for M2M auth) with read access to a Databricks Model Registry. I can successfully search the registry (via the `WorkspaceClient`) and find the model that I want to load using (Python) APIs, but I cannot load the model for inference.

Loading for inference seems to require MLFlow, but each time I try to use `mlflow.MlflowClient` to load a model, it errors out with `InvalidConfigurationError: You haven't configured the CLI yet!`.

I need to be able to use M2M auth, and configuring the CLI is not possible in my workflow.

How can I use the M2M auth approach to load models with MLFlow, or is there a different model-loading workflow that needs to be followed?

 

0 Kudos

Kumaran
Databricks Employee
Databricks Employee

‎06-06-2024 11:27 AM

Hello @JC3,

Thank you for posting your question in the Databricks community.

Is it possible to share with us the minimum reproducible code?

0 Kudos

JC3
New Contributor II
In response to Kumaran

‎06-07-2024 08:58 AM

Thank you, @Kumaran.

The short version of this is below. Note that I have the following env vars loaded at runtime:

  • DATABRICKS_HOST
  • DATABRICKS_CLIENT_ID
  • DATABRICKS_CLIENT_SECRET
  • MLFLOW_TRACKING_URI

 

import os
import mlflow

if __name__ == "__main__":
    # attempt to use mlflow to search model registry and load model
    mlflow_client = mlflow.MlflowClient(
        registry_uri=os.getenv('MLFLOW_TRACKING_URI')
    )

    model_name = 'test-model'
    for model in mlflow_client.search_model_versions(filter_string=f"name='{model_name}'"):
        pymodel = mlflow.pyfunc.load_model(model_uri=f"models:/{model.name}/{model.version}")

 

If DATABRICKS_CLIENT_ID and DATABRICKS_CLIENT_SECRET were replaced by a personal access token at  DATABRICKS_TOKEN, then this works as expected. But not when using the Service principal's client id and secret, which is what I hope to use, so that auth is M2M.

Can you help here?

0 Kudos