Problem

In Azure Databricks, you can store a model registered in Unity Catalog within an external storage account. Databricks provides a hosted version of the MLflow Model Registry as part of Unity Catalog and you can load models from this registry using the MLflow client. 

However, when the storage account uses a Private Link connection, you receive an authorization error while loading the model.

Example

In the following Python code example, a model is already registered in Unity Catalog with the alias Model_Alias. After importing the MLflow library, the client is configured to access models in Unity Catalog and load the model using the mlflow.pyfunc.load_model() function. 

import mlflow
mlflow.set_registry_uri("databricks-uc")  #configure MLflow to access models in Unity Catalog
model_version_uri = "models:/catalog.schema.model_name@Model_Alias"
model_version = mlflow.pyfunc.load_model(model_version_uri)  #load the model from registry

This results in the following error. 

RequestId: 83eacd9-d01e-00c-20ef-25a202000000
Time: 2024-10-24T05:50:08.7165446Z</Message></Error>
2024/10/24 05:50:08 INFO mflow.store.artifact.cloud_artifact_repo: Failed to complete request, possibly due to credential expiration. Refreshing credentials and trying again…(Error: This request is not authorized to perform this operation.) 
ErrorCode: AuthorizationFailure

Cause

The Azure Storage account only has a DFS private endpoint, but both DFS and Blob storage endpoints are required for proper authorization. 

Solution

Add private endpoints for both DFS and blob services on the Azure Storage account hosting the model. This ensures the ABFS driver can authenticate and perform read and write operations correctly.

1. In the Azure Portal, go to the Azure Storage account hosting the model.
2. Under Networking > Private Endpoint Connections, create two private endpoints with the following sub-resources:
   • Target sub-resource: Blob
   • Target sub-resource: DFS
3. Verify the endpoints resolve correctly within the Databricks workspace by executing the following shell commands in a notebook. 

%sh
nslookup storage_account_name.blob.core.windows.net
nslookup storage_account_name.dfs.core.windows.net

For more information, refer to Microsoft’s Use private endpoints for Azure Storage documentation.