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

import * as cdk from 'aws-cdk-lib';
import * as sagemaker from 'aws-cdk-lib/aws-sagemaker';
import * as appscaling from 'aws-cdk-lib/aws-applicationautoscaling';
import * as lambda from 'aws-cdk-lib/aws-lambda';
import * as iam from 'aws-cdk-lib/aws-iam';
import * as cr from 'aws-cdk-lib/custom-resources';
import { Construct } from 'constructs';

export interface AutoScalingSageMakerEndpointProps {
  readonly endpointName: string;
  readonly endpointConfigName: string;
  readonly variantName: string;
  readonly minCapacity: number;
  readonly maxCapacity: number;
  readonly targetValue: number;
  readonly scaleInCooldown?: number;
  readonly scaleOutCooldown?: number;
}

export class AutoScalingSageMakerEndpoint extends Construct {
  public readonly endpoint: sagemaker.CfnEndpoint;

  constructor(scope: Construct, id: string, props: AutoScalingSageMakerEndpointProps) {
    super(scope, id);

    // 1. 定义标准的SageMaker端点配置和端点
    const endpointConfig = new sagemaker.CfnEndpointConfig(this, 'EndpointConfig', {
      endpointConfigName: props.endpointConfigName,
      productionVariants: [
        {
          initialInstanceCount: props.minCapacity,
          instanceType: 'ml.m5.xlarge',
          variantName: props.variantName,
          modelName: 'your-model-name',
        },
      ],
    });

    this.endpoint = new sagemaker.CfnEndpoint(this, 'Endpoint', {
      endpointName: props.endpointName,
      endpointConfigName: endpointConfig.attrEndpointConfigName,
    });

    this.endpoint.addDependency(endpointConfig);

    // 2. 用于管理扩缩注册生命周期的Lambda处理程序
    const scalingManagerLambda = new lambda.SingletonFunction(this, 'ScalingManagerLambda', {
      uuid: 'a1b2c3d4-e5f6-7890-abcd-ef0123456789',
      code: lambda.Code.fromInline(`
        const { ApplicationAutoScalingClient, RegisterScalableTargetCommand, DeregisterScalableTargetCommand, PutScalingPolicyCommand } = require('@aws-sdk/client-application-auto-scaling');
        const client = new ApplicationAutoScalingClient();

        exports.handler = async (event) => {
          const props = event.ResourceProperties;
          const resourceId = \`endpoint/\${props.EndpointName}/variant/\${props.VariantName}\`;
          
          if (event.RequestType === 'Delete') {
            await client.send(new DeregisterScalableTargetCommand({
              ServiceNamespace: 'sagemaker',
              ResourceId: resourceId,
              ScalableDimension: 'sagemaker:variant:DesiredInstanceCount',
            }));
            return { PhysicalResourceId: resourceId };
          }

          if (event.RequestType === 'Update' || event.RequestType === 'Create') {
            // 抢先注销以避免更新期间的锁定冲突
            try {
              await client.send(new DeregisterScalableTargetCommand({
                ServiceNamespace: 'sagemaker',
                ResourceId: resourceId,
                ScalableDimension: 'sagemaker:variant:DesiredInstanceCount',
              }));
            } catch (e) {
              console.log('目标尚未注册，继续执行。');
            }

            // 注册可扩缩目标
            await client.send(new RegisterScalableTargetCommand({
              ServiceNamespace: 'sagemaker',
              ResourceId: resourceId,
              ScalableDimension: 'sagemaker:variant:DesiredInstanceCount',
              MinCapacity: parseInt(props.MinCapacity),
              MaxCapacity: parseInt(props.MaxCapacity),
            }));

            // 设置扩缩策略
            await client.send(new PutScalingPolicyCommand({
              PolicyName: \`sagemaker-target-tracking-\${props.VariantName}\`,
              ServiceNamespace: 'sagemaker',
              ResourceId: resourceId,
              ScalableDimension: 'sagemaker:variant:DesiredInstanceCount',
              PolicyType: 'TargetTrackingScaling',
              TargetTrackingScalingPolicyConfiguration: {
                TargetValue: parseFloat(props.TargetValue),
                PredefinedMetricSpecification: {
                  PredefinedMetricType: 'SageMakerVariantInvocationsPerInstance',
                },
              },
            }));
          }

          return { PhysicalResourceId: resourceId };
        };
      `),
      handler: 'index.handler',
      runtime: lambda.Runtime.NODEJS_18_X,
      timeout: cdk.Duration.seconds(30),
    });

    scalingManagerLambda.addToRolePolicy(new iam.PolicyStatement({
      actions: [
        'application-autoscaling:RegisterScalableTarget',
        'application-autoscaling:DeregisterScalableTarget',
        'application-autoscaling:PutScalingPolicy',
        'application-autoscaling:DeleteScalingPolicy',
      ],
      resources: ['*'],
    }));

    // 3. 自定义资源提供程序绑定
    const provider = new cr.Provider(this, 'ScalingProvider', {
      onEventHandler: scalingManagerLambda,
    });

    const customResource = new cdk.CustomResource(this, 'ScalingCustomResource', {
      serviceToken: provider.serviceToken,
      properties: {
        EndpointName: this.endpoint.endpointName,
        VariantName: props.variantName,
        MinCapacity: props.minCapacity,
        MaxCapacity: props.maxCapacity,
        TargetValue: props.targetValue,
        // 强制在配置更改时重新执行
        EndpointConfigName: endpointConfig.attrEndpointConfigName,
      },
    });

    customResource.node.addDependency(this.endpoint);
  }
}

# 1. 综合CDK堆栈以检查模板有效性
cdk synth

# 2. 部署包含SageMaker端点和自动扩缩规则的初始堆栈
cdk deploy --require-approval never

# 3. 修改端点配置（例如更改实例数量或模型参数）并重新部署
cdk deploy --require-approval never