跳到正文
AWS Machine Learning Blog· Lauren Eterno·· 16 小时前精选AI 评分59

如何用 AWS CDK 编排 SageMaker AI 端点自动扩缩容

Orchestrating SageMaker AI endpoint auto scaling with AWS CDK

AI 导读

文章介绍如何用可复用的 AWS CDK construct，在 SageMaker AI 实时端点更新前注销、更新后重新注册自动扩缩容，避免 CloudFormation 部署失败或扩缩容配置丢失。

推荐理由

文章拆解了 SageMaker AI 端点更新与自动扩缩容资源之间的部署冲突，并给出可复用的 CDK 编排方案及其资源开销。

正文 · AI 翻译

更新 Amazon SageMaker AI 实时端点配置（更改实例类型、实例数量或环境变量）会导致通过 AWS CloudFormation 或 AWS Cloud Development Kit (AWS CDK) 进行的部署失败。SageMaker AI 要求您在更新端点之前注销可扩展目标，但注销后，CloudFormation 无法跟踪 AWS Application Auto Scaling 资源。

在本文中，我们将介绍如何使用可复用的 CDK 构造和 AWS CloudFormation 自定义资源，在端点更新期间自动编排自动扩缩的注销和注册。您可以在一次部署中更新端点配置和自动扩缩参数，无需人工干预。

问题出现的原因

使用原生 CloudFormation 自动扩缩资源部署 SageMaker AI 端点后，如果更改端点配置并重新部署，就会出现多种故障模式。具体故障取决于您更改的内容以及采用的解决方法。

在已注册自动扩缩的情况下更改端点配置

CloudFormation 会尝试在可扩展目标仍处于注册状态时更新端点。SageMaker AI 会拒绝此更新，并返回以下错误：

The variant(s) "[AllTraffic]" must be deregistered as scalable targets with Application Auto Scaling before they can be removed or have their instance type updated.

在 CloudFormation 处理自动扩缩资源之前，SageMaker AI 就会阻止端点更新。

手动注销自动扩缩，然后仅更改端点配置

如果您通过 AWS Command Line Interface (AWS CLI) 注销可扩展目标，然后仅更改端点配置并重新部署，部署会成功。但是，即使 CloudFormation 模板中包含自动扩缩配置，CloudFormation 也不会注册自动扩缩。由于 CloudFormation 模板中的自动扩缩参数没有变化，CloudFormation 不会注册可扩展目标，也不会重新应用扩缩策略。此时端点已在运行，但没有自动扩缩，而且不会出现任何错误提示问题。

手动注销自动扩缩，然后更改自动扩缩参数

如果您通过 AWS CLI 注销可扩展目标，然后更改自动扩缩参数并重新部署，CloudFormation 会尝试更新已不存在的 ScalableTarget 资源：

Resource of type 'AWS::ApplicationAutoScaling::ScalableTarget' with identifier 'endpoint/<name>/variant/AllTraffic|sagemaker:variant:DesiredInstanceCount|sagemaker' was not found.

值得注意的是，此故障还会导致回滚失败 (UPDATE_ROLLBACK_FAILED)，使堆栈处于损坏状态，需要在 CloudFormation 控制台中手动干预才能恢复。

解决方案概述

您可以使用此 CDK 构造，将原生的 AWS::ApplicationAutoScaling::ScalableTarget 和 AWS::ApplicationAutoScaling::ScalingPolicy CloudFormation 资源替换为三个自定义资源。这些自定义资源会在端点更新前后注销和注册自动扩缩，并通过 CDK 中的 node.add_dependency 调用来强制执行显式依赖顺序。完整流程如下：

  1. 注销（自定义资源）：在端点更新开始前移除可扩展目标。首次部署时，可扩展目标尚不存在，因此处理程序会妥善处理错误并继续执行。
  2. 端点更新（由 CDK 管理）：CDK 正常更新端点资源，切换到新的端点配置。
  3. 注册可扩展目标（自定义资源）：等待端点达到 InService 状态，然后注册可扩展目标。
  4. 附加扩缩策略（自定义资源）：等待可扩展目标可供查询，然后应用目标跟踪策略。

Four-step orchestration the construct runs during an endpoint update: deregister the scalable target, update the endpoint, register the scalable target, and attach the scaling policy

图 1：围绕由 CDK 管理的端点更新，通过自定义资源编排 SageMaker AI 端点自动扩缩。

Register Scalable Target 和 Attach Scaling Policy 使用 CDK 提供程序的异步完成模式。每个自定义资源都有一个 on_event 处理程序，该处理程序会立即返回，以及一个 is_complete 处理程序。提供程序的等待程序会按固定间隔（每 30 秒）重新调用 is_complete 处理程序，直到其就绪检查通过。提供程序的 total_timeout（Register Scalable Target 为 1 小时，Attach Scaling Policy 为 15 分钟）限制了等待时间，因此不会受任何单次 Lambda 调用超时的限制。

每个自定义资源上的 ConfigHash 属性会强制 CloudFormation 在配置值发生变化时重新调用 Lambda 函数。端点配置和自动扩缩参数一同定义，这是因为单个 ConfigHash 会驱动自动扩缩编排。

虽然构造包含三个自定义资源，但部署的堆栈包含三个以上的 Lambda 函数。每个自定义资源都由 CDK cr.Provider 提供支持，而提供程序框架会生成自己的辅助 Lambda 函数。同步提供程序（Deregister）会添加一个框架函数。异步提供程序（Register Scalable Target 和 Attach Scaling Policy）会添加三个函数（onEvent、isComplete 和 onTimeout），以及一个驱动等待程序的 AWS Step Functions 状态机。因此，部署的堆栈包含三个自定义资源、两个 Step Functions 状态机和十二个 Lambda 函数。其中五个资源包含构造自身的处理程序代码（Deregister 一个，两个异步资源各包含一个 on_event 和 is_complete 对），其余七个由提供程序框架生成。你无需编写或维护框架生成的函数，但它们会存在于堆栈中，在检查已部署资源时值得纳入考虑。

演练

在本演练中，你将使用提供的示例堆栈部署一个带自动扩缩功能的 SageMaker AI 端点，然后验证配置更新是否正常工作。

本演练包括以下步骤：

  1. 克隆存储库并安装依赖项。
  2. 配置端点和自动扩缩设置。
  3. 部署并验证。
  4. 测试更新。

源代码位于链接的 GitHub 存储库中。

先决条件

在本演练中，你需要满足以下条件：

  • 一个具有 Amazon SageMaker AI、AWS Lambda、AWS Identity and Access Management (IAM)、Application Auto Scaling 和 Amazon CloudWatch 权限的 AWS 账户。
  • 已安装AWS CDK v2。
  • Python 3.12 或更高版本。
  • 用于 Python 依赖项管理的 uv。
  • Amazon Simple Storage Service (Amazon S3) 中的 SageMaker AI 模型构件和 Amazon Elastic Container Registry (Amazon ECR) 中的容器镜像，或 SageMaker AI 内置算法容器。

克隆存储库并安装依赖项

  1. Clone the repository: 
    git clone https://github.com/aws-samples/sample-sagemaker-endpoint-autoscaling.git
cd sample-sagemaker-endpoint-autoscaling
  2. Install dependencies: 
    uv sync

配置端点和自动扩缩设置

端点和自动扩缩设置在 config/config.yaml 中一同定义，这是因为单个 config_hash 会驱动自动扩缩编排。任一部分发生变化都会触发完整流程。

使用你的值编辑 config/config.yaml：

endpoint:
  endpoint_name: "my-model-endpoint"
  image: "123456789012.dkr.ecr.us-east-1.amazonaws.com/my-model:latest"
  model_data_url: "s3://my-bucket/models/model.tar.gz"
  instance_type: "ml.m5.large"
  initial_instance_count: 1
  environment: {}
autoscaling:
  min_capacity: 1
  max_capacity: 10
  target_value: 100
  scale_in_cooldown: 600
  scale_out_cooldown: 120

端点参数映射到 SageMaker AI CreateModel 和 CreateEndpointConfig API。自动扩缩参数映射到 Application Auto Scaling RegisterScalableTarget 和 PutScalingPolicy API。

部署并验证

  1. Bootstrap your AWS environment if you haven’t already: 
    uv run cdk bootstrap
  2. Deploy the stack: 
    uv run cdk deploy InferenceStack

CloudFormation 会创建 SageMaker AI 模型、端点配置和端点，然后按顺序运行三个自定义资源来注册自动扩缩。

  1. Verify that the deployment registered the scalable target: 
    aws application-autoscaling describe-scalable-targets
--service-namespace sagemaker
--resource-ids endpoint/my-model-endpoint/variant/AllTraffic
  2. Verify that the deployment applied the scaling policy: 
    aws application-autoscaling describe-scaling-policies
--service-namespace sagemaker
--resource-id endpoint/my-model-endpoint/variant/AllTraffic

测试更新

更改 config/config.yaml 中的一个值，例如，将实例类型从 ml.m5.large 更新为 ml.m5.xlarge，然后重新部署：

uv run cdk deploy InferenceStack

部署完成后，验证自动扩展仍已注册，并且已更新为新值：

aws application-autoscaling describe-scalable-targets
--service-namespace sagemaker
--resource-ids endpoint/my-model-endpoint/variant/AllTraffic

与现有堆栈集成

如果你是从头开始，或希望获取参考实现，请使用示例堆栈。如果你已经拥有包含 SageMaker AI 资源的 CDK 堆栈，请直接使用该构造体。

要与现有堆栈集成，请将 sagemaker_autoscaling/constructs/ 目录复制到你的项目中，并将该构造体添加到端点资源旁边：

import json
from sagemaker_autoscaling.constructs.sagemaker_autoscaling_construct import (
    AutoscalingConfig,
    SagemakerAutoscalingConstruct,
    SagemakerAutoscalingProps,
)

# Both endpoint and auto scaling values must be in config_hash so that
# a change to either triggers the auto scaling orchestration.
config_hash = json.dumps({
    "endpoint": {
        "image": "123456789012.dkr.ecr.us-east-1.amazonaws.com/my-model:latest",
        "modelDataUrl": "s3://my-bucket/models/model.tar.gz",
        "instanceType": "ml.m5.large",
        "initialInstanceCount": 1,
    },
    "autoscaling": {
        "minCapacity": 1,
        "maxCapacity": 10,
        "targetValue": 100,
        "scaleInCooldown": 600,
        "scaleOutCooldown": 120,
    },
})

SagemakerAutoscalingConstruct(
    self, "Autoscaling",
    props=SagemakerAutoscalingProps(
        endpoint=your_endpoint,  # your CfnEndpoint
        autoscaling=AutoscalingConfig(
            min_capacity=1,
            max_capacity=10,
            target_value=100,
            scale_in_cooldown=600,
            scale_out_cooldown=120,
        ),
        config_hash=config_hash,
    ),
)

构造体的工作原理

本节将介绍 SagemakerAutoscalingConstruct 的关键实现细节。

取消注册

Deregister 自定义资源会在端点更新之前调用 deregister_scalable_target。对于首次部署时不存在可扩展目标的情况，它会处理 ObjectNotFoundException：

def handler(event, context):
    resource_id = event["ResourceProperties"]["ResourceId"]
    if event["RequestType"] in ("Create", "Update"):
        try:
            client.deregister_scalable_target(
                ServiceNamespace="sagemaker",
                ResourceId=resource_id,
                ScalableDimension="sagemaker:variant:DesiredInstanceCount",
            )
        except client.exceptions.ObjectNotFoundException:
            pass  # No scalable target on first deploy
    return {"PhysicalResourceId": f"deregister-{resource_id}"}

依赖关系在 CDK 中建立：

props.endpoint.node.add_dependency(deregister)

注册可扩展目标

尽管 CloudFormation 会在端点更新请求被接受后将端点资源视为已完成，但只有当端点达到 InService 后，才能注册可扩展目标。Register Scalable Target 不会让单次 Lambda 调用一直等待，而是使用提供程序的异步完成模式。on_event 处理程序会立即返回，提供程序的等待器会在 query_interval 上重新调用 is_complete 处理程序，直到端点处于 InService 状态。提供程序的 total_timeout（1 小时）限定了等待时间，而不是任何单次 Lambda 调用的超时时间：

def handler(event, context):
    props = event["ResourceProperties"]
    resource_id = props["ResourceId"]
    endpoint_name = resource_id.split("/")[1]
    if event["RequestType"] == "Delete":
        return {"IsComplete": True}
    status = sm_client.describe_endpoint(EndpointName=endpoint_name)["EndpointStatus"]
    if status == "Failed":
        raise Exception("Endpoint in Failed state")
    if status != "InService":
        return {"IsComplete": False}
    as_client.register_scalable_target(
        ServiceNamespace="sagemaker",
        ResourceId=resource_id,
        ScalableDimension="sagemaker:variant:DesiredInstanceCount",
        MinCapacity=int(props["MinCapacity"]),
        MaxCapacity=int(props["MaxCapacity"]),
    )
    return {"IsComplete": True}

删除堆栈时，on_event 处理程序会取消注册可扩展目标以完成清理。

附加扩展策略

RegisterScalableTarget 返回后会有一个很短的时间窗口，在此期间无法查询目标。Attach Scaling Policy 使用相同的异步完成模式：on_event 处理程序会立即返回，is_complete 处理程序仅在可扩展目标可查询后报告完成，然后应用该策略。提供程序的 total_timeout（15 分钟）限定了等待时间：

def handler(event, context):
    props = event["ResourceProperties"]
    resource_id = props["ResourceId"]
    if event["RequestType"] == "Delete":
        return {"IsComplete": True}
    resp = client.describe_scalable_targets(
        ServiceNamespace="sagemaker",
        ResourceIds=[resource_id],
        ScalableDimension="sagemaker:variant:DesiredInstanceCount",
    )
    if not resp.get("ScalableTargets"):
        return {"IsComplete": False}
    client.put_scaling_policy(
        PolicyName=props["PolicyName"],
        ServiceNamespace="sagemaker",
        ResourceId=resource_id,
        ScalableDimension="sagemaker:variant:DesiredInstanceCount",
        PolicyType="TargetTrackingScaling",
        TargetTrackingScalingPolicyConfiguration=json.loads(props["ScalingPolicyConfig"]),
    )
    return {"IsComplete": True}

删除堆栈时，on_event 处理程序会删除扩展策略以完成清理。

清理

为避免产生后续费用，请删除以下资源：

uv run cdk destroy InferenceStack

这会移除 SageMaker AI 端点、端点配置、模型、Lambda 函数、IAM 角色，以及由目标跟踪扩展策略创建的 Amazon CloudWatch 警报，并取消注册可扩展目标。

结语

在本文中，我们向你展示了如何使用可复用的 CDK 构造体自动取消注册和注册自动扩展，从而协调 SageMaker AI 端点更新。该构造体使用三个具有明确依赖顺序的 CloudFormation 自定义资源，以及一个 ConfigHash 属性，以确保在部署过程中自动扩展保持不变。

要开始使用，请克隆 sample-sagemaker-endpoint-autoscaling 存储库并按照设置说明操作。该存储库包含一个可以直接部署的示例堆栈，你也可以将该构造体集成到现有的 CDK 堆栈中。

有关 SageMaker AI 的 Application Auto Scaling 的更多信息，请参阅自动扩展 Amazon SageMaker AI 模型。有关 CloudFormation 自定义资源的更多信息，请参阅AWS CloudFormation 用户指南中的“自定义资源”。

关于作者

Lauren Eterno

Lauren Eterno

Lauren 是 Amazon Web Services 的助理交付顾问，专注于 AI/ML。她与客户团队合作，设计、构建、部署和扩展机器学习及生成式 AI 解决方案。她拥有 Case Western Reserve University 的数据科学与分析理学学士学位（B.S.）和计算机科学理学硕士学位（M.S.）。

来源：AWS Machine Learning Blog · aws.amazon.com

部署工程教程实践#教程/实践#部署/工程