Orchestrating SageMaker AI endpoint auto scaling with AWS CDK

Updating the Amazon Sagemaker AI real-time endpoint configuration (changing the instance type, instance count, or environment variables) causes deployments through AWS CloudFormation or the AWS Cloud Development Kit (AWS CDK) to fail. SageMaker AI requires you to deregister the scalable target before updating the endpoint, but when you deregister it, CloudFormation can’t track the AWS Application Auto Scaling resources. In this post, we show you how to use a reusable CDK construct with AWS CloudFormation custom resources to orchestrate deregister and register auto scaling automatically during endpoint updates. You can update your endpoint configuration and auto scaling parameters in a single deployment without manual intervention. Why this problem occurs When you deploy a SageMaker AI endpoint with native CloudFormation auto scaling resources, then change the endpoint configuration and redeploy, several failure modes emerge. The specific failure depends on what you change and how you work around the issue. Changing the endpoint configuration with auto scaling registered CloudFormation tries to update the endpoint while the scalable target is still registered. SageMaker AI rejects the update with the following 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 blocks the endpoint update before CloudFormation reaches the auto scaling resources. Manually deregistering auto scaling, then changing only the endpoint configuration If you deregister the scalable target through the AWS Command Line Interface (AWS CLI) and redeploy with only an endpoint configuration change, the deployment succeeds. However, CloudFormation doesn’t register auto scaling even though it is in the CloudFormation template. Because the auto scaling parameters in the CloudFormation template haven’t changed, CloudFormation doesn’t register the scalable target or re-apply the scaling policy. The endpoint is now running without auto scaling, and there is no error to indicate anything went wrong. Manually deregistering auto scaling, then changing auto scaling parameters If you deregister the scalable target through the AWS CLI and redeploy with auto scaling parameter changes, CloudFormation tries to update the ScalableTarget resource that no longer exists: Resource of type 'AWS::ApplicationAutoScaling::ScalableTarget' with identifier 'endpoint//variant/AllTraffic|sagemaker:variant:DesiredInstanceCount|sagemaker' was not found. This failure is notable because the rollback also fails (UPDATE_ROLLBACK_FAILED), leaving the stack in a broken state that requires manual intervention in the CloudFormation console to recover. Solution overview You can use this CDK construct to replace the native AWS::ApplicationAutoScaling::ScalableTarget and AWS::ApplicationAutoScaling::ScalingPolicy CloudFormation resources with three custom resources. These custom resources deregister and register auto scaling around the endpoint update with explicit dependency ordering, enforced by node.add_dependency calls in CDK. The full sequence is: Deregister (custom resource): Removes the scalable target before the endpoint update begins. On first deploy, no scalable target exists yet, so the handler handles the error gracefully and continues. Endpoint update (CDK-managed): CDK updates the endpoint resource normally, swapping to the new endpoint configuration. Register Scalable Target (custom resource): Waits for the endpoint to reach InService status, then registers the scalable target. Attach Scaling Policy (custom resource): Waits for the scalable target to be queryable, then applies the target tracking policy. Figure 1: Custom resource orchestration of SageMaker AI endpoint auto scaling around a CDK-managed endpoint update. Register Scalable Target and Attach Scaling Policy use the CDK provider’s asynchronous completion pattern. Each custom resource has an on_event handler that returns immediately and an is_complete handler. The provider’s waiter re-invokes the is_complete handler on an interval (every 30 seconds) until its readiness check passes. The provider’s total_timeout (1 hour for Register Scalable Target, 15 minutes for Attach Scaling Policy) bounds the wait, so it isn’t limited by the timeout of any single Lambda invocation. A ConfigHash property on each custom resource forces CloudFormation to re-invoke the Lambda functions whenever a configuration value changes. The endpoint configuration and auto scaling parameters are defined together because a single ConfigHash drives the auto scaling orchestration. While the construct defines three custom resources, the deployed stack contains more than three Lambda functions. Each custom resource is backed by a CDK cr.Provider , and the provider framework generates its own helper Lambda functions. A synchronous provider (Deregister) adds one framework function. An asynchronous provider (Register Scalable Target and Attach Scaling Policy) adds three ( onEvent , isComplete , and onTimeout ) plus an AWS Step Functions state machine that drives the waiter. As a result, the deployed stack contains three custom resources, two Step Functions state machines, and twelve Lambda functions. Five of the resources hold the construct’s own handler code (one for Deregister, and an on_event and is_complete pair for each asynchronous resource), with the remaining seven generated by the provider framework. You don’t write or maintain the framework-generated functions, but they are present in the stack and worth accounting for when you review the deployed resources. Walkthrough In this walkthrough, you deploy a SageMaker AI endpoint with auto scaling using the provided example stack, then verify that configuration updates work correctly. The walkthrough includes the following steps: Clone the repository and install dependencies. Configure your endpoint and auto scaling settings. Deploy and verify. Test an update. The source code is available in the linked GitHub repository . Prerequisites For this walkthrough, you need the following: An AWS account with permissions for Amazon SageMaker AI, AWS Lambda, AWS Identity and Access Management (IAM), Application Auto Scaling, and Amazon CloudWatch. AWS CDK v2 installed. Python 3.12 or later. uv for Python dependency management. A SageMaker AI model artifact in Amazon Simple Storage Service (Amazon S3) and a container image in Amazon Elastic Container Registry (Amazon ECR), or a SageMaker AI built-in algorithm container. Clone the repository and install dependencies Clone the repository: git clone https://github.com/aws-samples/sample-sagemaker-endpoint-autoscaling.git cd sample-sagemaker-endpoint-autoscaling Install dependencies: uv sync Configure your endpoint and auto scaling settings The endpoint and auto scaling settings are defined together in config/ config.yaml because a single config_hash drives the auto scaling orchestration. A change to either section triggers the full sequence. Edit config/ config.yaml with your values: 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 The endpoint parameters map to the SageMaker AI CreateModel and CreateEndpointConfig APIs. The auto scaling parameters map to the Application Auto Scaling RegisterScalableTarget and PutScalingPolicy APIs. Deploy and verify Bootstrap your AWS environment if you haven’t already: uv run cdk bootstrap Deploy the stack: uv run cdk deploy InferenceStack CloudFormation creates the SageMaker AI model, endpoint configuration, and endpoint, then runs the three custom resources in sequence to register auto scaling. Verify that the deployment regi