Orquestación del auto escalado de endpoints de SageMaker AI con AWS CDK

Actualizar la configuración de un endpoint en tiempo real de Amazon SageMaker AI (cambiar el tipo de instancia, el número de instancias o las variables de entorno) hace que las implementaciones a través de AWS CloudFormation o del AWS Cloud Development Kit (AWS CDK) fallen. SageMaker AI requiere que cancele el registro del scalable target antes de actualizar el endpoint, pero al cancelarlo, CloudFormation no puede rastrear los recursos de AWS Application Auto Scaling.

En esta publicación, le mostramos cómo usar un constructo de CDK reutilizable con recursos personalizados de AWS CloudFormation para orquestar la cancelación y el registro del auto escalado automáticamente durante las actualizaciones del endpoint. Puede actualizar la configuración de su endpoint y los parámetros de auto escalado en una sola implementación sin intervención manual.

Por qué ocurre este problema

Cuando implementa un endpoint de SageMaker AI con recursos nativos de auto escalado de CloudFormation y luego cambia la configuración del endpoint y vuelve a implementar, surgen varios modos de fallo. El fallo específico depende de lo que cambie y de cómo solucione el problema.

Cambio de la configuración del endpoint con auto escalado registrado

CloudFormation intenta actualizar el endpoint mientras el scalable target sigue registrado. SageMaker AI rechaza la actualización con el siguiente error:

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.

SageMaker AI bloquea la actualización del endpoint antes de que CloudFormation llegue a los recursos de auto escalado.

Cancelación manual del registro de auto escalado y luego cambio solo de la configuración del endpoint

Si cancela el registro del scalable target a través de la AWS Command Line Interface (AWS CLI) y vuelve a implementar con solo un cambio en la configuración del endpoint, la implementación se realiza correctamente. Sin embargo, CloudFormation no registra el auto escalado aunque esté en la plantilla de CloudFormation. Como los parámetros de auto escalado en la plantilla de CloudFormation no han cambiado, CloudFormation no registra el scalable target ni vuelve a aplicar la política de escalado. El endpoint ahora se ejecuta sin auto escalado y no hay ningún error que indique que algo salió mal.

Cancelación manual del registro de auto escalado y luego cambio de los parámetros de auto escalado

Si cancela el registro del scalable target a través de la AWS CLI y vuelve a implementar con cambios en los parámetros de auto escalado, CloudFormation intenta actualizar el recurso ScalableTarget que ya no existe:

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

Este fallo es notable porque la reversión también falla (UPDATE_ROLLBACK_FAILED), dejando el stack en un estado dañado que requiere intervención manual en la consola de CloudFormation para recuperarse.

Descripción general de la solución

Puede usar este constructo de CDK para reemplazar los recursos nativos de CloudFormation AWS::ApplicationAutoScaling::ScalableTarget y AWS::ApplicationAutoScaling::ScalingPolicy con tres recursos personalizados. Estos recursos personalizados cancelan y registran el auto escalado en torno a la actualización del endpoint con un orden de dependencias explícito, aplicado mediante llamadas node.add_dependency en CDK. La secuencia completa es:

1. Deregister (recurso personalizado): Elimina el scalable target antes de que comience la actualización del endpoint. En la primera implementación, aún no existe ningún scalable target, por lo que el handler gestiona el error de forma correcta y continúa.
2. Actualización del endpoint (administrada por CDK): CDK actualiza el recurso del endpoint normalmente, cambiando a la nueva configuración del endpoint.
3. Register Scalable Target (recurso personalizado): Espera a que el endpoint alcance el estado InService , luego registra el scalable target.
4. Attach Scaling Policy (recurso personalizado): Espera a que el scalable target sea consultable, luego aplica la política de target tracking.

Figura 1: Orquestación de recursos personalizados del auto escalado de endpoints de SageMaker AI en torno a una actualización del endpoint administrada por CDK.

Register Scalable Target y Attach Scaling Policy usan el patrón de finalización asíncrona del provider de CDK. Cada recurso personalizado tiene un on_event handler que devuelve inmediatamente y un is_complete handler. El waiter del proveedor vuelve a invocar el is_complete handler a intervalos (cada 30 segundos) hasta que su comprobación de preparación se supera. El total_timeout (1 hora para Register Scalable Target, 15 minutos para Attach Scaling Policy) limita la espera, de modo que no está limitado por el timeout de ninguna invocación individual de Lambda.

Una ConfigHash property en cada recurso personalizado obliga a CloudFormation a volver a invocar las funciones de Lambda cada vez que cambia un valor de configuración. La configuración del endpoint y los parámetros de auto scaling se definen juntos porque un solo ConfigHash impulsa la orquestación de auto scaling.

Si bien el constructo define tres recursos personalizados, la pila desplegada contiene más de tres funciones Lambda. Cada recurso personalizado está respaldado por un CDK cr.Provider, y el framework de proveedores genera sus propias funciones Lambda auxiliares. Un proveedor síncrono (Deregister) añade una función del framework. Un proveedor asíncrono (Register Scalable Target y Attach Scaling Policy) añade tres (onEvent, isCompletey onTimeout) más una máquina de estados de AWS Step Functions que controla el waiter. Como resultado, la pila desplegada contiene tres recursos personalizados, dos máquinas de estados de Step Functions y doce funciones Lambda. Cinco de los recursos contienen el código de handler propio del constructo (uno para Deregister, y un par de on_event y is_complete para cada recurso asíncrono), y los siete restantes son generados por el framework de proveedores. Usted no escribe ni mantiene las funciones generadas por el framework, pero están presentes en la pila y conviene tenerlas en cuenta al revisar los recursos desplegados.

Recorrido

En este recorrido, despliega un endpoint de SageMaker AI con autoescalado utilizando la pila de ejemplo proporcionada, y luego verifica que las actualizaciones de configuración funcionan correctamente.

El recorrido incluye los siguientes pasos:

1. Clonar el repositorio e instalar las dependencias.
2. Configurar el endpoint y los ajustes de autoescalado.
3. Desplegar y verificar.
4. Probar una actualización.

El código fuente está disponible en el repositorio de GitHub enlazado.

Requisitos previos

Para este recorrido, necesitas lo siguiente:

• Una cuenta de AWS con permisos para Amazon SageMaker AI, AWS Lambda, AWS Identity and Access Management (IAM), Application Auto Scaling y Amazon CloudWatch.
• AWS CDK v2 instalado.
• Python 3.12 o posterior.
• uv para la gestión de dependencias de Python.
• Un artefacto de modelo de SageMaker AI en Amazon Simple Storage Service (Amazon S3) y una imagen de contenedor en Amazon Elastic Container Registry (Amazon ECR), o un contenedor de un algoritmo integrado de SageMaker AI.

Clonar el repositorio e instalar las dependencias

1. Clona el repositorio:

   git clone https://github.com/aws-samples/sample-sagemaker-endpoint-autoscaling.git
   cd sample-sagemaker-endpoint-autoscaling

2. Instala las dependencias:

   uv sync

Configurar el endpoint y los ajustes de autoescalado

El endpoint y los ajustes de autoescalado se definen juntos en config/config.yaml porque un solo config_hash impulsa la orquestación del autoescalado. Un cambio en cualquiera de las secciones desencadena la secuencia completa.

Edite config/config.yaml con sus valores:

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

Los parámetros del endpoint corresponden a las API de SageMaker AI CreateModel y CreateEndpointConfig . Los parámetros de auto scaling corresponden a las API de Application Auto Scaling RegisterScalableTarget y PutScalingPolicy .

Despliegue y verifique

1. Inicialice (bootstrap) su entorno de AWS si aún no lo ha hecho:

   uv run cdk bootstrap

2. Despliegue el stack:

   uv run cdk deploy InferenceStack

CloudFormation crea el modelo de SageMaker AI, la configuración del endpoint y el endpoint, y luego ejecuta los tres recursos personalizados en secuencia para registrar el auto scaling.

1. Verifique que el despliegue registró el objetivo escalable:

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

2. Verifique que el despliegue aplicó la política de escalado:

   aws application-autoscaling describe-scaling-policies
   --service-namespace sagemaker
   --resource-id endpoint/my-model-endpoint/variant/AllTraffic

Pruebe una actualización

Cambie un valor en config/config.yaml, por ejemplo, actualice el tipo de instancia de ml.m5.large a ml.m5.xlarge, y vuelva a desplegar:

uv run cdk deploy InferenceStack

Una vez completado el despliegue, verifique que el auto scaling sigue registrado y tiene los valores actualizados:

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

Integración con una pila existente

Use la pila de ejemplo si está comenzando desde cero o si desea una implementación de referencia. Use el constructo directamente si ya tiene una pila de CDK con recursos de SageMaker AI.

Para integrarlo con una pila existente, copie el directorio sagemaker_autoscaling/constructs/ en su proyecto y agregue el constructo junto con los recursos de su endpoint:

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,
    ),
)

Cómo funciona el constructo

Esta sección repasa los detalles clave de implementación de SagemakerAutoscalingConstruct.

Deregister

El recurso personalizado Deregister llama a deregister_scalable_target antes de la actualización del endpoint. Gestiona ObjectNotFoundException para las primeras implementaciones donde no existe un destino escalable:

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}"}

La dependencia se establece en CDK:

props.endpoint.node.add_dependency(deregister)

Register Scalable Target

Aunque CloudFormation considera que el recurso del endpoint está completo una vez aceptada la actualización, el destino escalable solo puede registrarse después de que el endpoint alcance InService. En lugar de bloquear una única invocación de Lambda durante esa espera, Register Scalable Target utiliza el patrón de finalización asíncrona del proveedor. El controlador on_event retorna inmediatamente, y el waiter del proveedor vuelve a invocar el controlador is_complete en query_interval hasta que el endpoint esté InService. El total_timeout del proveedor (1 hora) limita la espera, no el timeout de ninguna invocación individual de 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}

Al eliminar la pila, el controlador on_event anula el registro del destino escalable para limpiar.

Attach Scaling Policy

Hay una breve ventana después de que RegisterScalableTarget retorna en la que el destino aún no es consultable. Attach Scaling Policy utiliza el mismo patrón de finalización asíncrona: el controlador on_event retorna inmediatamente, y el controlador is_complete informa la finalización solo cuando el destino escalable es consultable, momento en el que aplica la política. El total_timeout del proveedor (15 minutos) limita la espera:

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}

Al eliminar la pila, el controlador on_event elimina la política de escalado para limpiar.

Limpieza

Para evitar cargos futuros, elimine los recursos:

uv run cdk destroy InferenceStack

Esto elimina el endpoint de SageMaker AI, la configuración del endpoint, el modelo, las funciones de Lambda, los roles de IAM, las alarmas de Amazon CloudWatch creadas por la política de escalado de seguimiento de destino y anula el registro del destino escalable.

Conclusión

En esta publicación, le mostramos cómo usar un constructo reutilizable de CDK para desregistrar y registrar el auto scaling automáticamente con el fin de orquestar las actualizaciones de los endpoints de SageMaker AI. El constructo utiliza tres recursos personalizados de CloudFormation con un orden explícito de dependencias y una propiedad ConfigHash para mantener intacto el auto scaling entre implementaciones.

Para comenzar, clone el sample-sagemaker-endpoint-autoscaling repositorio y siga las instrucciones de configuración. El repositorio incluye una pila de ejemplo que puede implementar directamente, o bien puede integrar el constructo en su pila de CDK existente.

Para obtener más información sobre Application Auto Scaling para SageMaker AI, consulte Automatically scale Amazon SageMaker AI models. Para obtener más información sobre los recursos personalizados de CloudFormation, consulte Custom resources en la AWS CloudFormation User Guide.

Sobre el autor