Orchestrer l'auto scaling des endpoints SageMaker AI avec AWS CDK

La mise à jour de la configuration d'un endpoint temps réel Amazon SageMaker AI (modification du type d'instance, du nombre d'instances ou des variables d'environnement) fait échouer les déploiements via AWS CloudFormation ou l'AWS Cloud Development Kit (AWS CDK). SageMaker AI vous oblige à désinscrire la cible évolutive avant de mettre à jour l'endpoint, mais lorsque vous la désinscrivez, CloudFormation ne peut plus suivre les ressources AWS Application Auto Scaling.

Dans cet article, nous vous montrons comment utiliser un construct CDK réutilisable avec des ressources personnalisées AWS CloudFormation pour orchestrer automatiquement la désinscription et l'inscription de l'auto scaling lors des mises à jour de l'endpoint. Vous pouvez mettre à jour la configuration de votre endpoint et les paramètres d'auto scaling en un seul déploiement, sans intervention manuelle.

Pourquoi ce problème se produit

Lorsque vous déployez un endpoint SageMaker AI avec les ressources d'auto scaling natives de CloudFormation, puis que vous modifiez la configuration de l'endpoint et redéployez, plusieurs modes d'échec apparaissent. L'échec spécifique dépend de ce que vous modifiez et de la façon dont vous contournez le problème.

Modification de la configuration de l'endpoint avec l'auto scaling enregistré

CloudFormation tente de mettre à jour l'endpoint alors que la cible évolutive est toujours inscrite. SageMaker AI rejette la mise à jour avec l'erreur suivante :

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 bloque la mise à jour de l'endpoint avant que CloudFormation n'atteigne les ressources d'auto scaling.

Désinscription manuelle de l'auto scaling, puis modification de la seule configuration de l'endpoint

Si vous désinscrivez la cible évolutive via l'AWS Command Line Interface (AWS CLI) et que vous redéployez avec uniquement une modification de la configuration de l'endpoint, le déploiement réussit. Cependant, CloudFormation n'inscrit pas l'auto scaling même s'il figure dans le modèle CloudFormation. Comme les paramètres d'auto scaling du modèle CloudFormation n'ont pas changé, CloudFormation n'inscrit pas la cible évolutive ni ne réapplique la politique de mise à l'échelle. L'endpoint s'exécute désormais sans auto scaling, et aucune erreur n'indique qu'un problème est survenu.

Désinscription manuelle de l'auto scaling, puis modification des paramètres d'auto scaling

Si vous désinscrivez la cible évolutive via l'AWS CLI et que vous redéployez avec des modifications des paramètres d'auto scaling, CloudFormation tente de mettre à jour la ressource ScalableTarget qui n'existe plus :

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

Cet échec est notable car le rollback échoue également (UPDATE_ROLLBACK_FAILED), laissant la pile dans un état défaillant qui nécessite une intervention manuelle dans la console CloudFormation pour être récupéré.

Présentation de la solution

Vous pouvez utiliser ce construct CDK pour remplacer les ressources CloudFormation natives AWS::ApplicationAutoScaling::ScalableTarget et AWS::ApplicationAutoScaling::ScalingPolicy par trois ressources personnalisées. Ces ressources personnalisées désinscrivent et inscrivent l'auto scaling autour de la mise à jour de l'endpoint avec un ordre de dépendance explicite, appliqué par des appels node.add_dependency dans CDK. La séquence complète est :

1. Désinscrire (ressource personnalisée) : supprime la cible évolutive avant le début de la mise à jour de l'endpoint. Lors du premier déploiement, aucune cible évolutive n'existe encore, donc le gestionnaire gère l'erreur correctement et continue.
2. Mise à jour de l'endpoint (gérée par CDK) : CDK met à jour la ressource de l'endpoint normalement, en passant à la nouvelle configuration d'endpoint.
3. Inscrire la cible évolutive (ressource personnalisée) : attend que l'endpoint atteigne le statut InService , puis inscrit la cible évolutive.
4. Attacher la politique de mise à l'échelle (ressource personnalisée) : attend que la cible évolutive soit interrogeable, puis applique la politique de target tracking.

Figure 1 : orchestration par ressources personnalisées de l'auto scaling d'un endpoint SageMaker AI autour d'une mise à jour d'endpoint gérée par CDK.

Register Scalable Target et Attach Scaling Policy utilisent le modèle de complétion asynchrone du provider CDK. Chaque ressource personnalisée possède un gestionnaire on_event qui renvoie immédiatement et un is_complete handler. Le waiter du fournisseur ré-invoque le is_complete handler à intervalles (toutes les 30 secondes) jusqu'à ce que son contrôle de préparation réussisse. Le total_timeout (1 heure pour Register Scalable Target, 15 minutes pour Attach Scaling Policy) borne l'attente, de sorte qu'elle n'est pas limitée par le délai d'expiration d'une seule invocation Lambda.

Une ConfigHash propriété sur chaque ressource personnalisée force CloudFormation à ré-invoquer les fonctions Lambda chaque fois qu'une valeur de configuration change. La configuration du point de terminaison et les paramètres d'auto scaling sont définis ensemble car une seule ConfigHash pilote l'orchestration de l'auto scaling.

Bien que le construct définisse trois ressources personnalisées, la pile déployée contient plus de trois fonctions Lambda. Chaque ressource personnalisée est adossée à un CDK cr.Provider, et le framework de fournisseurs génère ses propres fonctions Lambda auxiliaires. Un fournisseur synchrone (Deregister) ajoute une fonction de framework. Un fournisseur asynchrone (Register Scalable Target et Attach Scaling Policy) en ajoute trois (onEvent, isComplete, et onTimeout) plus une machine à états AWS Step Functions qui pilote le waiter. Par conséquent, la pile déployée contient trois ressources personnalisées, deux machines à états Step Functions et douze fonctions Lambda. Cinq des ressources contiennent le code de gestionnaire propre au construct (une pour Deregister, et une paire on_event et is_complete pour chaque ressource asynchrone), les sept autres étant générées par le framework de fournisseurs. Vous n’écrivez ni ne maintenez les fonctions générées par le framework, mais elles sont présentes dans la pile et il convient de les prendre en compte lorsque vous examinez les ressources déployées.

Procédure pas à pas

Dans cette procédure pas à pas, vous déployez un point de terminaison SageMaker AI avec mise à l’échelle automatique à l’aide de la pile d’exemple fournie, puis vous vérifiez que les mises à jour de configuration fonctionnent correctement.

La procédure comprend les étapes suivantes :

1. Cloner le dépôt et installer les dépendances.
2. Configurer votre point de terminaison et vos paramètres de mise à l’échelle automatique.
3. Déployer et vérifier.
4. Tester une mise à jour.

Le code source est disponible dans le dépôt GitHub lié.

Prérequis

Pour cette procédure, vous avez besoin des éléments suivants :

• Un compte AWS avec des autorisations pour Amazon SageMaker AI, AWS Lambda, AWS Identity and Access Management (IAM), Application Auto Scaling et Amazon CloudWatch.
• AWS CDK v2 installé.
• Python 3.12 ou version ultérieure.
• uv pour la gestion des dépendances Python.
• Un artefact de modèle SageMaker AI dans Amazon Simple Storage Service (Amazon S3) et une image de conteneur dans Amazon Elastic Container Registry (Amazon ECR), ou un conteneur d’algorithme intégré SageMaker AI.

Cloner le dépôt et installer les dépendances

1. Clonez le dépôt :

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

2. Installez les dépendances :

   uv sync

Configurer votre point de terminaison et vos paramètres de mise à l’échelle automatique

Les paramètres du point de terminaison et de la mise à l’échelle automatique sont définis ensemble dans config/config.yaml parce qu’un seul config_hash pilote l’orchestration de la mise à l’échelle automatique. Une modification de l’une ou l’autre section déclenche la séquence complète.

Modifiez config/config.yaml avec vos valeurs :

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

Les paramètres d'endpoint correspondent aux API SageMaker AI CreateModel et CreateEndpointConfig . Les paramètres d'auto scaling correspondent aux API Application Auto Scaling RegisterScalableTarget et PutScalingPolicy .

Déployer et vérifier

1. Initialisez votre environnement AWS (bootstrap) si ce n'est pas déjà fait :

   uv run cdk bootstrap

2. Déployez la pile (stack) :

   uv run cdk deploy InferenceStack

CloudFormation crée le modèle SageMaker AI, la configuration de l'endpoint et l'endpoint, puis exécute les trois ressources personnalisées en séquence pour enregistrer l'auto scaling.

1. Vérifiez que le déploiement a enregistré la cible scalable :

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

2. Vérifiez que le déploiement a appliqué la politique de mise à l'échelle :

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

Tester une mise à jour

Modifiez une valeur dans config/config.yaml, par exemple, mettez à jour le type d'instance de ml.m5.large vers ml.m5.xlarge, puis redéployez :

uv run cdk deploy InferenceStack

Une fois le déploiement terminé, vérifiez que l'auto scaling est toujours enregistré et qu'il contient les valeurs mises à jour :

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

Intégration à une pile existante

Utilisez l’exemple de pile si vous démarrez de zéro ou si vous souhaitez une implémentation de référence. Utilisez directement le construct si vous disposez déjà d’une pile CDK avec des ressources SageMaker AI.

Pour intégrer une pile existante, copiez le répertoire sagemaker_autoscaling/constructs/ dans votre projet et ajoutez le construct à côté de vos ressources de point de terminaison :

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

Fonctionnement du construct

Cette section décrit les détails clés de l’implémentation de SagemakerAutoscalingConstruct.

Deregister

La ressource personnalisée Deregister appelle deregister_scalable_target avant la mise à jour du point de terminaison. Elle gère ObjectNotFoundException pour les premiers déploiements où aucune cible évolutive n’existe :

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 dépendance est établie dans CDK :

props.endpoint.node.add_dependency(deregister)

Register Scalable Target

Bien que CloudFormation considère la ressource du point de terminaison comme terminée dès que la mise à jour est acceptée, la cible évolutive ne peut être enregistrée qu’une fois que le point de terminaison atteint InService. Plutôt que de bloquer une seule invocation Lambda sur cette attente, Register Scalable Target utilise le modèle d’achèvement asynchrone du provider. Le gestionnaire on_event retourne immédiatement, et le waiter du provider réinvoque le gestionnaire is_complete à intervalles de query_interval jusqu’à ce que le point de terminaison soit InService. Le total_timeout du provider (1 heure) borne l’attente, et non le délai d’attente d’une invocation Lambda individuelle :

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}

Lors de la suppression de la pile, le gestionnaire on_event désenregistre la cible évolutive pour nettoyer.

Attach Scaling Policy

Il existe une brève fenêtre après le retour de RegisterScalableTarget pendant laquelle la cible n’est pas encore interrogeable. Attach Scaling Policy utilise le même modèle d’achèvement asynchrone : le gestionnaire on_event retourne immédiatement, et le gestionnaire is_complete signale l’achèvement uniquement lorsque la cible évolutive est interrogeable, moment auquel il applique la stratégie. Le total_timeout du provider (15 minutes) borne l’attente :

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}

Lors de la suppression de la pile, le gestionnaire on_event supprime la stratégie de mise à l’échelle pour nettoyer.

Nettoyage

Pour éviter des frais futurs, supprimez les ressources :

uv run cdk destroy InferenceStack

Cela supprime le point de terminaison SageMaker AI, la configuration du point de terminaison, le modèle, les fonctions Lambda, les rôles IAM, les alarmes Amazon CloudWatch créées par la stratégie de mise à l’échelle par suivi de cible, et désenregistre la cible évolutive.

Conclusion

Dans cet article, nous vous avons montré comment utiliser un construct CDK réutilisable pour dé-enregistrer et ré-enregistrer automatiquement l'auto scaling afin d'orchestrer les mises à jour de points de terminaison SageMaker AI. Le construct utilise trois ressources personnalisées CloudFormation avec un ordre de dépendance explicite et une propriété ConfigHash pour conserver l'auto scaling intact entre les déploiements.

Pour commencer, clonez le sample-sagemaker-endpoint-autoscaling dépôt et suivez les instructions de configuration. Le dépôt inclut une pile d'exemple que vous pouvez déployer directement, ou vous pouvez intégrer le construct dans votre pile CDK existante.

Pour plus d'informations sur Application Auto Scaling pour SageMaker AI, consultez Automatically scale Amazon SageMaker AI models. Pour plus d'informations sur les ressources personnalisées CloudFormation, consultez Custom resources dans le AWS CloudFormation User Guide.

À propos de l'auteur