将 MLflow 模型部署到 Amazon SageMaker

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

如果你是MLflow模型部署的新手,请先阅读 MLflow部署 以了解MLflow模型和部署的基本概念。

工作原理

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

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

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

备注

除了通用的 pyfunc 部署(默认),SageMaker 部署还支持 mleap 风格。对于这种部署风格,端点仅接受 JSON 序列化的 pandas DataFrame,采用 split 方向,例如 data = pandas_df.to_json(orient='split')。这种格式是通过 Content-Type 请求头值为 application/json 来指定的。

将模型部署到 SageMaker 端点

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

步骤 0:准备

安装工具

确保安装以下工具(如果尚未安装):

  • mlflow

  • awscli 是一个命令行界面工具,用于与Amazon Web Services进行交互。

  • docker

权限设置

正确设置AWS账户和权限。您需要一个具有创建SageMaker端点、访问S3存储桶和使用ECR仓库权限的IAM角色。该角色还应可被执行部署的用户承担。了解更多关于此设置的信息,请访问 在AWS CLI中使用IAM角色

创建一个 MLflow 模型

在部署之前,您必须拥有一个 MLflow 模型。如果您没有,可以按照 MLflow 跟踪快速入门 创建一个示例的 scikit-learn 模型。记得记下模型 URI,例如 runs:/<run_id>/<artifact_path> (或者如果您在 MLflow 模型注册表 中注册了模型,则为 models:/<model_name>/<model_version>)。

步骤 1:在本地测试您的模型

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

mlflow deployments run-local -t sagemaker -m runs:/<run_id>/model -p 5000

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

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

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

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

$ mlflow sagemaker build-and-push-container  -m runs:/<run_id>/model

或者,您可以使用 官方 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。

故障排除