使用Amazon SageMaker异步端点优化Amazon SageMaker JumpStart基础模型的部署成本
使用Amazon SageMaker异步端点优化部署成本
生成式AI应用在各个行业取得的成功吸引了全球公司的关注和兴趣，这些公司希望复制和超越竞争对手的成就，或者解决新颖而令人兴奋的用例。这些客户正在研究基础模型，如TII Falcon、稳定扩散XL或OpenAI的GPT-3.5，作为推动生成式AI创新的引擎。
基础模型是一类生成式AI模型，由于它们经过大量非结构化数据的训练，能够理解和生成类似人类的内容。这些模型已经在各种计算机视觉（CV）和自然语言处理（NLP）任务中引起了革命，包括图像生成、翻译和问题回答。它们是许多AI应用的构建模块，并且已成为高级智能系统开发中至关重要的组成部分。
然而，部署基础模型可能面临重大挑战，特别是在成本和资源需求方面。这些模型以其规模而闻名，参数数量通常在数亿到数十亿之间。它们庞大的规模需要大量的计算资源，包括强大的硬件和大容量的内存。事实上，部署基础模型通常需要至少一个（通常更多）GPU来高效处理计算负载。例如，TII Falcon-40B Instruct模型至少需要一个ml.g5.12xlarge实例才能成功加载到内存中，但在较大的实例上性能最佳。因此，部署和维护这些模型的投资回报率（ROI）可能过低，尤其是在开发周期或突发工作负载期间。这是因为在长时间会话中拥有基于GPU的实例的运行成本可能是24/7。
今年早些时候，我们宣布了Amazon Bedrock，这是一个无服务器API，用于从亚马逊和我们的生成式AI合作伙伴那里访问基础模型。尽管目前处于私有预览阶段，它的无服务器API允许您使用来自亚马逊、Anthropic、Stability AI和AI21的基础模型，而无需自行部署任何端点。然而，来自Hugging Face等社区的开源模型增长很快，并非所有模型都可以通过Amazon Bedrock获得。
在本文中，我们针对这些情况解决了通过从Amazon SageMaker JumpStart将大型基础模型部署到Amazon SageMaker异步端点而导致高成本的问题。这可以帮助减少架构的成本，允许端点仅在队列中有请求时运行，并且具有短的存活时间，在没有等待服务的请求时缩减到零。这对于许多用例来说听起来很棒；然而，缩减到零的端点将在能够提供推理之前引入冷启动时间。
解决方案概述
以下图示说明了我们的解决方案架构。
我们部署的架构非常简单：
- 用户界面是一个笔记本，可以用Streamlit或类似技术构建的Web UI替换。在我们的案例中，笔记本是一个运行在ml.m5.large实例上的Amazon SageMaker Studio笔记本，使用PyTorch 2.0 Python 3.10 CPU内核。
- 笔记本以三种方式查询端点：SageMaker Python SDK、Python的AWS SDK（Boto3）和LangChain。
- 端点在SageMaker上异步运行，我们在端点上部署了Falcon-40B Instruct模型。它目前是指导模型方面的最新技术，并且在SageMaker JumpStart中可用。通过单个API调用，我们可以将模型部署到端点上。
SageMaker异步推理是什么
SageMaker异步推理是SageMaker中的四种部署选项之一，其他选项包括实时端点、批量推理和无服务器推理。要了解有关不同部署选项的更多信息，请参阅部署推理模型。
SageMaker异步推理将传入的请求排队并以异步方式处理，使其成为处理负载大小高达1 GB、处理时间长且需要近实时延迟的请求的理想选择。然而，它在处理大型基础模型时提供的主要优势，特别是在概念验证（POC）或开发过程中，是能够配置异步推理以在没有请求需要处理时缩减到实例计数为零，从而节省成本。有关SageMaker异步推理的更多信息，请参阅异步推理。以下图示说明了这种架构。
要部署一个异步推理端点，您需要创建一个AsyncInferenceConfig对象。如果您创建AsyncInferenceConfig而不指定其参数，则默认的S3OutputPath将为s3://sagemaker-{REGION}-{ACCOUNTID}/async-endpoint-outputs/{UNIQUE-JOB-NAME}，S3FailurePath将为s3://sagemaker-{REGION}-{ACCOUNTID}/async-endpoint-failures/{UNIQUE-JOB-NAME}。
什么是SageMaker JumpStart
我们的模型来自SageMaker JumpStart，这是SageMaker的一个功能，通过提供预训练模型、解决方案模板和示例笔记本，加速机器学习（ML）的学习过程。它提供了各种类型的预训练模型，用于不同的问题类型，使您能够在坚实的基础上开始您的ML任务。SageMaker JumpStart还提供了常见用例的解决方案模板和学习的示例笔记本。借助SageMaker JumpStart，您可以通过一键式解决方案启动和全面资源的实践ML经验，减少启动ML项目所需的时间和精力。
以下截图显示了SageMaker JumpStart UI上可用的模型示例。
部署模型
我们的第一步是将模型部署到SageMaker。为此，我们可以使用SageMaker JumpStart的UI或SageMaker Python SDK。后者提供了一个API，我们可以使用它来将模型部署到异步端点：
%%time
from sagemaker.jumpstart.model import JumpStartModel, AsyncInferenceConfig
from sagemaker.serializers import JSONSerializer
from sagemaker.deserializers import JSONDeserializer
model_id, model_version = "huggingface-llm-falcon-40b-instruct-bf16", "*"
my_model = JumpStartModel(model_id=model_id)
predictor = my_model.deploy(
initial_instance_count=0,
instance_type="ml.g5.12xlarge",
async_inference_config=AsyncInferenceConfig()
)此调用可能需要大约10分钟才能完成。在此期间，将启动端点，将容器与模型工件一起下载到端点，从SageMaker JumpStart加载模型配置，然后通过DNS端点公开异步端点。为了确保我们的端点可以缩减到零，我们需要使用Application Auto Scaling在异步端点上配置自动缩放。您需要首先使用Application Auto Scaling注册您的端点变体，定义一个缩放策略，然后应用该缩放策略。在此配置中，我们使用CustomizedMetricSpecification创建了一个自定义指标，名为ApproximateBacklogSizePerInstance，如下所示。有关可用于异步推理端点的Amazon CloudWatch指标的详细列表，请参阅使用CloudWatch进行监控。
import boto3
client = boto3.client("application-autoscaling")
resource_id = "endpoint/" + my_model.endpoint_name + "/variant/" + "AllTraffic"
# 配置异步端点的自动缩放至零实例
response = client.register_scalable_target(
ServiceNamespace="sagemaker",
ResourceId=resource_id,
ScalableDimension="sagemaker:variant:DesiredInstanceCount",
MinCapacity=0, # 我们希望缩减到的最小实例数-缩减到0以停止产生成本
MaxCapacity=1, # 我们希望扩展到的最大实例数-扩展到1个最大实例对于开发来说足够好
)
response = client.put_scaling_policy(
PolicyName="Invocations-ScalingPolicy",
ServiceNamespace="sagemaker", # 提供资源的AWS服务的命名空间。
ResourceId=resource_id, # 端点名称
ScalableDimension="sagemaker:variant:DesiredInstanceCount", # SageMaker仅支持实例数
PolicyType="TargetTrackingScaling", # 'StepScaling'|'TargetTrackingScaling'
TargetTrackingScalingPolicyConfiguration={
"TargetValue": 5.0, # 指标的目标值。-此处指标为- SageMakerVariantInvocationsPerInstance
"CustomizedMetricSpecification": {
"MetricName": "ApproximateBacklogSizePerInstance",
"Namespace": "AWS/SageMaker",
"Dimensions": [{"Name": "EndpointName", "Value": my_model.endpoint_name}],
"Statistic": "Average",
},
"ScaleInCooldown": 600, # 在缩小活动完成后，再次开始缩小活动之前的时间量，以秒为单位。
"ScaleOutCooldown": 300, # ScaleOutCooldown - 在扩展活动完成后，再次开始扩展活动之前的时间量，以秒为单位。
# 'DisableScaleIn': True|False - 指示是否禁用目标跟踪策略的缩小。
# 如果值为true，则禁用缩小，目标跟踪策略将不会从可扩展资源中移除容量。
},
)您可以通过导航到SageMaker控制台，选择导航窗格中的推理下的终端节点，并查找刚刚部署的终端节点，以验证此策略是否已成功设置。
调用异步终端节点
要调用终端节点，您需要将请求有效载荷放置在Amazon Simple Storage Service (Amazon S3)中，并将此有效载荷的指针作为InvokeEndpointAsync请求的一部分提供。调用时，SageMaker将请求排队进行处理，并返回一个标识符和输出位置作为响应。处理完成后，SageMaker将结果放置在Amazon S3位置中。您可以选择使用Amazon Simple Notification Service (Amazon SNS)接收成功或错误通知。
SageMaker Python SDK
部署完成后，它将返回一个AsyncPredictor对象。要执行异步推理，您需要将数据上传到Amazon S3，并使用S3 URI作为输入使用predict_async()方法。它将返回一个AsyncInferenceResponse对象，您可以使用get_response()方法检查结果。
或者，如果您想定期检查结果并在生成后返回结果，请使用predict()方法。我们在以下代码中使用了这种第二种方法:
import time
# 使用SageMaker Python SDK调用异步终端节点
def query_endpoint(payload):
"""查询终端节点并打印响应"""
response = predictor.predict_async(
data=payload,
input_path="s3://{}/{}".format(bucket, prefix),
)
while True:
try:
response = response.get_result()
break
except:
print("推理尚未准备好...")
time.sleep(5)
print(f"\033[1m 输入:\033[0m {payload['inputs']}")
print(f"\033[1m 输出:\033[0m {response[0]['generated_text']}")
query_endpoint(payload)Boto3
现在让我们探索Boto3的sagemaker-runtime客户端中的invoke_endpoint_async方法。它使开发人员能够异步调用SageMaker终端节点，并提供用于跟踪进度和稍后检索响应的令牌。Boto3不提供等待异步推理完成的方法，例如SageMaker Python SDK的get_result()操作。因此，我们利用Boto3将推理输出存储在Amazon S3中的事实，可以使用以下函数等待将推理文件写入Amazon S3:
import json
import time
import boto3
from botocore.exceptions import ClientError
s3_client = boto3.client("s3")
# 等待生成预测的文件
def wait_inference_file(bucket, prefix):
while True:
try:
response = s3_client.get_object(Bucket=bucket, Key=prefix)
break
except ClientError as ex:
if ex.response['Error']['Code'] == 'NoSuchKey':
print("等待文件生成中...")
time.sleep(5)
next
else:
raise
except Exception as e:
print(e.__dict__)
raise
return response有了这个函数，我们现在可以查询终端节点:
# 使用Boto3 SDK调用异步终端节点
import boto3
sagemaker_client = boto3.client("sagemaker-runtime")
# 查询终端节点函数
def query_endpoint_boto3(payload):
"""查询终端节点并打印响应"""
response = sagemaker_client.invoke_endpoint_async(
EndpointName=my_model.endpoint_name,
InputLocation="s3://{}/{}".format(bucket, prefix),
ContentType="application/json",
Accept="application/json"
)
output_url = response["OutputLocation"]
output_prefix = "/".join(output_url.split("/")[3:])
# 使用Boto3从output_url中的Amazon S3中读取文件的字节
output = wait_inference_file(bucket, output_prefix)
output = json.loads(output['Body'].read())[0]['generated_text']
# 输出结果
print(f"\033[1m 输入:\033[0m {payload['inputs']}")
print(f"\033[1m 输出:\033[0m {output}")
query_endpoint_boto3(payload)LangChain
LangChain是由Harrison Chase于2022年10月推出的开源框架。它通过与各种系统和数据源集成，简化了使用大型语言模型（LLMs）开发应用程序的过程。LangChain可以进行文档分析、摘要、聊天机器人创建、代码分析等操作。它受到了广大开发者的贡献和风险投资公司的大量资金支持，因此变得非常流行。LangChain可以连接LLMs与外部资源，从而可以创建动态、数据响应的应用程序。它提供了库、API和文档，以简化开发过程。
LangChain提供了库和示例，用于在其框架中使用SageMaker端点，使得使用托管在SageMaker上的ML模型作为“链”的“大脑”变得更加容易。要了解LangChain如何与SageMaker集成，请参考LangChain文档中的SageMaker端点部分。
当前LangChain的一个限制是它不原生支持异步端点。要在LangChain中使用异步端点，我们需要定义一个新的类SagemakerAsyncEndpoint，它扩展了LangChain中已经存在的SagemakerEndpoint类。此外，我们还提供以下信息：
- 异步推断将存储输入（和输出）的S3存储桶和前缀
- 等待超时的最大秒数
- 一个
updated _call()函数，用于使用invoke_endpoint_async()而不是invoke_endpoint()查询端点 - 唤醒异步端点的方法（如果它处于冷启动状态，即被缩减到零）
要查看新创建的SagemakerAsyncEndpoint，请查看GitHub上可用的sagemaker_async_endpoint.py文件。
from typing import Dict
from langchain import PromptTemplate
from langchain.llms.sagemaker_endpoint import LLMContentHandler
from langchain.chains import LLMChain
from sagemaker_async_endpoint import SagemakerAsyncEndpoint
class ContentHandler(LLMContentHandler):
content_type:str = "application/json"
accepts:str = "application/json"
len_prompt:int = 0
def transform_input(self, prompt: str, model_kwargs: Dict) -> bytes:
self.len_prompt = len(prompt)
input_str = json.dumps({"inputs": prompt, "parameters": {"max_new_tokens": 100, "do_sample": False, "repetition_penalty": 1.1}})
return input_str.encode('utf-8')
def transform_output(self, output: bytes) -> str:
response_json = output.read()
res = json.loads(response_json)
ans = res[0]['generated_text']
return ans
chain = LLMChain(
llm=SagemakerAsyncEndpoint(
input_bucket=bucket,
input_prefix=prefix,
endpoint_name=my_model.endpoint_name,
region_name=sagemaker.Session().boto_region_name,
content_handler=ContentHandler(),
),
prompt=PromptTemplate(
input_variables=["query"],
template="{query}",
),
)
print(chain.run(payload['inputs']))清理
当您完成了对端点进行推理生成的测试后，请记得删除端点以避免产生额外的费用：
predictor.delete_endpoint()结论
在部署像TII Falcon这样的大型基础模型时，优化成本非常重要。这些模型需要强大的硬件和大量的内存容量，导致基础设施成本很高。SageMaker异步推断是一种处理请求的异步部署选项，通过在没有挂起请求时将实例数量缩减到零来降低开销。在本文中，我们演示了如何将大型SageMaker JumpStart基础模型部署到SageMaker异步端点。我们提供了使用SageMaker Python SDK、Boto3和LangChain的代码示例，以说明调用异步端点和检索结果的不同方法。这些技术使开发人员和研究人员能够在使用基础模型的能力进行高级语言理解系统时优化成本。
要了解有关异步推断和SageMaker JumpStart的更多信息，请查看以下文章：
- 使用Amazon Kendra、LangChain和大型语言模型在企业数据上快速构建高准确性生成式AI应用程序
- 在大型视频上运行计算机视觉推断，使用Amazon SageMaker异步端点
