将 MLflow 模型部署到 Amazon SageMaker

Amazon SageMaker 是一项完全托管的服务，专为扩展 ML 推理容器而设计。MLflow 通过提供易于使用的命令，无需编写容器定义，从而简化了部署过程。

如果您不熟悉 MLflow 模型部署，请先阅读MLflow 部署，以了解 MLflow 模型和部署的基本概念。

工作原理

SageMaker 提供了一项名为自带容器 (BYOC) 的功能，允许您在推理端点上运行自定义 Docker 容器。这些容器必须满足特定要求，例如运行公开特定 REST 端点的 Web 服务器、具有指定的容器入口点、设置环境变量等。编写满足这些要求的 Dockerfile 和服务脚本可能是一项繁琐的任务。

MLflow 会代表您从 MLflow 模型构建 Docker 镜像，从而自动化此过程。随后，它会将镜像推送到 Elastic Container Registry (ECR)，并使用此镜像创建 SageMaker 端点。它还将模型工件上传到 S3 存储桶，并配置端点从该存储桶下载模型。

该容器提供与本地推理服务器相同的 REST 端点。例如，/invocations 端点接受 CSV 和 JSON 输入数据并返回预测结果。有关端点的更多详细信息，请参阅本地推理服务器。

将模型部署到 SageMaker 端点

本节概述了使用 MLflow CLI 将模型部署到 SageMaker 的过程。有关 Python API 参考和教程，请参阅有用链接部分。

步骤 0：准备

安装工具

确保已安装以下工具（如果尚未安装）：

• mlflow
• awscli
• docker

权限设置

正确设置 AWS 账户和权限。您需要一个 IAM 角色，该角色具有创建 SageMaker 端点、访问 S3 存储桶和使用 ECR 存储库的权限。执行部署的用户也应该能够承担此角色。有关此设置的详细信息，请参阅在 AWS CLI 中使用 IAM 角色。

创建 MLflow 模型

在部署之前，您必须拥有一个 MLflow 模型。如果您没有，可以按照MLflow Tracking 快速入门创建示例 scikit-learn 模型。请记住记下模型 URI，例如 models:/<model_id>（如果您在MLflow 模型注册表中注册了模型，则为 models:/<model_name>/<model_version>）。

步骤 1：在本地测试模型

建议在将模型部署到生产环境之前在本地进行测试。mlflow deployments run-local 命令会在具有相同镜像和环境配置的 Docker 容器中部署模型，因此非常适合部署前测试。

mlflow deployments run-local -t sagemaker -m models:/<model_id> -p 5000

然后，您可以通过向端点发送 POST 请求来测试模型

curl -X POST -H "Content-Type:application/json; format=pandas-split" --data '{"columns":["a","b"],"data":[[1,2]]}' https://:5000/invocations

步骤 2：构建 Docker 镜像并推送到 ECR

mlflow sagemaker build-and-push-container 命令会构建一个与 SageMaker 兼容的 Docker 镜像并将其上传到 ECR。

$ mlflow sagemaker build-and-push-container  -m models:/<model_id>

或者，您可以使用官方 MLflow Docker 镜像创建自定义 Docker 镜像，并手动将其推送到 ECR。

步骤 3：部署到 SageMaker 端点

mlflow deployments create 命令会将模型部署到 Amazon SageMaker 端点。MLflow 会将 Python 函数模型上传到 S3，并自动启动一个服务该模型的 Amazon SageMaker 端点。

有多种命令行选项可用于自定义部署，例如实例类型、数量、IAM 角色等。有关完整选项列表，请参阅CLI 参考。

$ mlflow deployments create -t sagemaker -m runs:/<run_id>/model \
    -C region_name=<your-region> \
    -C instance-type=ml.m4.xlarge \
    -C instance-count=1 \
    -C env='{"DISABLE_NGINX": "true"}''

API 参考

您可以通过两种方式将模型部署到 SageMaker：使用 CLI 或 Python API。

• CLI 参考
• Python API 文档

有用链接

• MLflow 快速入门第二部分：使用 Amazon SageMaker 部署模型 - 此分步教程演示了如何从 Databricks 笔记本中使用 MLflow Python API 将模型部署到 SageMaker。
• 使用 MLflow 和 Amazon SageMaker 管理您的机器学习生命周期 - 此综合教程涵盖了从模型训练到部署的整个 MLflow 生命周期与 SageMaker 的集成。