从 Databricks 工作区下载模型工件

Question

如何从 databricks workspace 下载 mlflow 模型工件到 docker 容器中？

Answer 1

要从 Databricks 工作区下载模型，您需要做两件事：

使用 python 将 MLFlow 跟踪 URI 设置为数据块 API
设置数据块身份验证。我更喜欢通过设置以下环境变量来进行身份验证，您也可以使用 databricks CLI 进行身份验证：
```
DATABRICKS_HOST

DATABRICKS_TOKEN
```

这是从 Databricks 工作区模型注册表下载模型的基本代码片段：

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}')

运行以上 python 脚本将在模型目录中下载 ML 模型。

服务于 Docker

的容器化 MLFlow 模型

下一步是将下载的模型打包到 docker 图像中，并在运行图像时提供模型。

这是一个基本的 Docker 文件来执行相同的操作：

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}

有关更多信息，您可以关注来自 Akshay Milmile

的 article

从 Databricks 工作区下载模型工件

Download model artefact from Databricks workspace

docker

databricks

azure-databricks

mlflow