How to download mlflow model artifacts from Azure Databricks workspace to local directory?

Question

I have trained a machine learning model in Databricks workspace using mlflow. Model is getting registered in databricks model registry and saved in databricks file share. Now I want to download model artifacts from workspace. Currently I am transferring model to azure machine learning workspace. There I am able to download all the artifacts. How to do it from databricks workspace?

Answer 1

@Sriram Reddy Gade Thanks for the question. To download a model from Databricks workspace you need to do two things:

Set MLFlow tracking URI to databricks using python API

Setup databricks authentication. I prefer authenticating by setting the following environment variables, you can also use databricks CLI to authenticate:

DATABRICKS_HOST

DATABRICKS_TOKEN
Here's a basic code snippet to download a model from Databricks workspace model registry:

import os  
import mlflow  
from mlflow.store.artifact.models_artifact_repo import ModelsArtifactRepository  
  
model_name = "example-model-name"  
model_stage = "Staging"  # Should be either 'Staging' or 'Production'  
  
mlflow.set_tracking_uri("databricks")  
  
os.makedirs("model", exist_ok=True)  
local_path = ModelsArtifactRepository(  
    f'models:/{model_name}/{model_stage}').download_artifacts("", dst_path="model")  
  
print(f'{model_stage} Model {model_name} is downloaded at {local_path}')  

Running above python script will download an ML model in the model directory.

Containerizing MLFlow model serving with Docker

For more information you can follow this article from Akshay Milmile