topic Re: Can view Model Registry using a Service Principal, but cannot load the model for inference. in Machine Learning

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

JC3 — Tue, 04 Jun 2024 14:39:01 GMT

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?

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

Kumaran — Thu, 06 Jun 2024 18:27:30 GMT

Hello @JC3,

Thank you for posting your question in the Databricks community.

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

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

JC3 — Fri, 07 Jun 2024 15:58:30 GMT

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?