超长时间异步推理解决方案 Async-Inference-Service

背景和问题

亚马逊云科技在 AI/ML 领域提供了 SageMaker 产品，为用户提供了众多功能强大的工具，帮助用户从训练、推理到模型上线后的跟踪等方面用很少的代价，即可获得高速的产出，而不必去花掉大量的时间来搭建基础设施。

对于模型的部署而言，SageMaker 提供了许多种方案供用户来选择，如实时推理部署，异步推理，批量转换等。但由于对场景定位的不同，所以不同类型的推理方式，会有各自的限制，比如对于实时同步推理，会被限制在 1 分钟就需要有回复，而异步推理，则最长不能超过 15 分钟等等。虽然这能覆盖很大一部分的场景，但总归有一些需要超长推理时间（小时级别甚至更长时间）的场合，目前产品不能满足一些特殊需要。

解决方案

为了满足超长时间推理的需求，且仍然可以利用 SageMaker 全托管环境带来的诸多优点，解决方案的核心思路在于充分利用 SageMaker Training 提供的能力，将其 Train Code 的部分替换为 Inference Code，从而在 Training 的触发和运行框架之下实现对长时间推理的支持。为了使整体架构更加简便易用和健壮，在具体实现中增加了许多工程化内容，如白名单、服务限制（Quota Limit）和一键部署。

Async-Inference-Service 架构

整体流程描述

Rest API – 整个系统，呈现给用户的是一个简单的 Restful API，由 APIGateway 来承担，用户发送标准的 HTTP 请求来触发推理。请求的例子如下，

curl -d '{"type":"sagemaker-xxx-deberta-inf"}' -H 'Content-Type: application/json'    https://xxxxxxxx.execute-api.cn-north-1.amazonaws.com.cn/dev/inference

请求发送之后，将会在 Apigateway 后连接的 Lambda 中进行处理。一切正常的情况下，将得到如下回应：

{"status": 200, "errMsg": "A new job [{jobName}] was scheduled successfully, and {default_jobs_limit - 1} slots reminded.."}

这个 response 的含义是，系统为你启动了一个 SageMaker 的推理 job，且将系统剩余的配额进行告知（如果是 0 slots reminded，意味着当前所有任务结束前，后续同样的 http 请求将会被拒绝。

此时，系统将启动一个 SageMaker Training Job，但实际上此 Trainning Job 的 Train Code 被替换成了 Inference Code，从业务上讲，它执行的是一段推理为目的代码。
当推理执行完毕之后，其结果可以写入/opt/ml/output 中，其会映射为一个 s3 的路径。
同时，一个规则被注册到了 Event Bridge 当中，这条规则的数据源被设置为 SageMaker 的 Training job status change，而目标被设置为 Lambda，意味着所有 Training Job 从创建到完成的每一个状态的变化，都会以一定的频率主动发送到 Lambda 中进行处理。用这种方式实现无侵入的异步通知机制。
在 Lambda 中，将会利用 Dynamo DB 进行查重处理，保证只有唯一一条标记为 Completed 的 Training Job 可以被插入 Dynamo DB，且同时去触发 SNS 相应的的 Topic。
在另一边，我们提前建好了 SNS，并将 Subscriber 配置为了邮件。如果收件人邮箱已经被验证无误，则收件人将收到来自通知邮件——邮件将包含间断但准确的信息，通知用户推理已经完成，并带有足够的参考信息让用户知道具体是哪个模，时间，以及到哪里去看到结果。

系统特点

整个系统的组件全部为托管服务，且均为无服务组件，按需计费，节约成本
可使用 Spot 实例，进一步优化成本。
Terraform 自动化部署和关键参数配置，部署和清理简单快捷。
白名单/限流机制，进一步保护客户资源不被意外滥开。

接口和部署模版

HTTP 接口说明

代码结构

请求	回应	场景
curl -d ‘{“type”:”sagemaker-kwm-deberta-inf”}’ -H ‘Content-Type: application/json’ https://xxxx.execute-api.cn-north-1.amazonaws.com.cn/dev/inference	{“status”: 200, “errMsg”: “A new job [{jobName}] was scheduled successfully, and 0 slots reminded.”}	假设 Limit 被配置为 1，且如果系统中也没有 Training Job 在运行，则系统如此 response，返回结果的含义是一个名为 JobName 的推理 Job 已经开始运行，系统当前有 Limit – 1 个空位，对于本例而言，系统当前的可供继续开资源的数字为 0.
	{“status”: 502, “errMsg”: “Failed to schedule inference at the moment, due to reaching limit 1.”}	假设 Limit 为 1，系统中已经有一个 type == “sagemaker-kwm-deberta-inf”的 job 在运行，则此时 reminded slots 为 0，则无法为该请求分配 SageMaker 资源，此时将如此 response，返回 502，拒绝继续为用户新开推理资源。
	{“status”: 501, “errMsg”: “missing “type” in your request, please check and try again.”}	传入的数据格式不正确，应为一个 json字符串，且格式为 “type”： “xxx”，对数据格式进行一定的规范。
	{“status”: 510, “errMsg”: “The type [xxx] of inference request is NOT supported..”}	传入的模型 type 的值，不在系统设定的模型白名单内，会如此 response。设置白名单的目的是进一步缩小用户可以激活系统资源的途径。

配置模版（系统参数）

Terraform 配置文件(vars.tf)

variable stage_vars{
 type = map
 default = {
   "s3_bucket" = "sagemaker-studio-bucket"
   "white_list" = "model_a,model_b,model_c"
   "role_arn" = "arn:aws:iam::123456789012:role/service-role/AmazonSageMaker-ExecutionRole-20221104T162296"
   "limit" = 1
 }
}

配置参数说明

配置项	含义
limit	设定，最小为 1
role_arn	SageMaker Role
S3_bucket	SageMaker 输入/输出需要映射到的 S3 桶
White_list	客户定义的模型名称列表，如果多于一个用逗号进行分割，形如：[model_a, mode_b]

配置项目以 ApiGateway Stage 变量存放，如图所示。

最后

本文描述了一个可以应对用户超长推理时间需求的，无服务化的异步推理实现服务。

亚马逊AWS官方博客