Download model artifacts from Databricks workspace

Question

How can I download mlflow model artifacts in a docker container from databricks workspace?

Answer 1

To download a model from Databricks workspace you need to do two things:

1. Set MLFlow tracking URI to databricks using python API

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

   DATABRICKS_HOST
   
   DATABRICKS_TOKEN
   

3. 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

   The next step is to package this downloaded model in a docker image and serve a model when you run the image.

Here's a basic Dockerfile to do the same:

FROM continuumio/miniconda3

ENV MLFLOW_HOME /opt/mlflow
ENV MLFLOW_VERSION 1.12.1
ENV PORT 5000

RUN conda install -c conda-forge mlflow=${MLFLOW_VERSION}

COPY model/ ${MLFLOW_HOME}/model

WORKDIR ${MLFLOW_HOME}

RUN mlflow models prepare-env -m ${MLFLOW_HOME}/model

RUN useradd -d ${MLFLOW_HOME} mlflow
RUN chown mlflow: ${MLFLOW_HOME}
USER mlflow

CMD mlflow models serve -m ${MLFLOW_HOME}/model --host 0.0.0.0 --port ${PORT}

For more information you can follow this article from Akshay Milmile