Sagemaker model CI/CD

Intro

On our journey to build cheap, scalable GPU inference at Tonar we have selected to go with Sagemaker Async Inference. We manage our infrastructure with Terraform and this is a small how-to on actually getting CI/CD to work when you are continuously updating models behind Sagemaker Endpoints. We will use a custom Whisper model as the base reference here.

Expected behaviour...

Reading the documentation of both AWS and Terraform you would end up with something like the below code to have an autoscaling endpoint.

resource "aws_sagemaker_model" "whisper" {
  execution_role_arn = aws_iam_role.whisper.arn
  name               = "whisper"

  primary_container {
    image          = data.aws_ssm_parameter.whisper_digest_id.value
    model_data_url = "s3://..."
    environment = {
      STAGE                 = terraform.workspace
      CONTAINER_ID          = data.aws_ssm_parameter.whisper_digest_id.value
      SENTRY_DSN            = data.aws_ssm_parameter.sentry_dsn.value
      SENTRY_ENVIRONMENT    = terraform.workspace
      INSTANCE_TYPE         = "ml.g4dn.xlarge"
    }
  }
}

resource "aws_sagemaker_endpoint_configuration" "whisper" {
  production_variants {
    model_name                                        = aws_sagemaker_model.whisper.name
    variant_name                                      = "AllTraffic"
    initial_instance_count                            = 1
    instance_type                                     = "ml.g4dn.xlarge"
    container_startup_health_check_timeout_in_seconds = 1800
    model_data_download_timeout_in_seconds            = 1200
  }

  async_inference_config {
    output_config {
      s3_output_path  = "s3://..."
      s3_failure_path = "s3://..."
      notification_config {
        include_inference_response_in = ["SUCCESS_NOTIFICATION_TOPIC", "ERROR_NOTIFICATION_TOPIC"]
        success_topic                 = aws_sns_topic.whisper_success_topic.arn
        error_topic                   = aws_sns_topic.whisper_error_topic.arn
      }
    }
    client_config {
      max_concurrent_invocations_per_instance = 1
    }
  }
}


resource "aws_sagemaker_endpoint" "whisper" {
  name                 = "whisper"
  endpoint_config_name = aws_sagemaker_endpoint_configuration.whisper.name
  deployment_config {
    rolling_update_policy {
      maximum_batch_size {
        type  = "CAPACITY_PERCENT"
        value = 50
      }
      maximum_execution_timeout_in_seconds = 900
      wait_interval_in_seconds             = 180
    }
  }
  depends_on = [aws_sagemaker_model.whisper, aws_sagemaker_endpoint_configuration.whisper]
}

resource "aws_appautoscaling_target" "sagemaker_target" {
  max_capacity       = 30
  min_capacity       = 0
  resource_id        = "endpoint/${aws_sagemaker_endpoint.whisper.name}/variant/AllTraffic"
  role_arn           = aws_iam_role.whisper.arn
  scalable_dimension = "sagemaker:variant:DesiredInstanceCount"
  service_namespace  = "sagemaker"
  depends_on         = [aws_sagemaker_endpoint.whisper, aws_sagemaker_endpoint_configuration.whisper]
}

resource "aws_appautoscaling_policy" "sagemaker_policy_regular" {
  name               = "whisper-invocations-scaling-policy"
  policy_type        = "TargetTrackingScaling"
  resource_id        = aws_appautoscaling_target.sagemaker_target.resource_id
  scalable_dimension = aws_appautoscaling_target.sagemaker_target.scalable_dimension
  service_namespace  = aws_appautoscaling_target.sagemaker_target.service_namespace

  target_tracking_scaling_policy_configuration {

    customized_metric_specification {
      metric_name = "ApproximateBacklogSizePerInstance"
      namespace   = "AWS/SageMaker"
      dimensions {
        name  = "EndpointName"
        value = aws_sagemaker_endpoint.whisper.name
      }
      statistic = "Average"
    }
    target_value       = 10
    scale_in_cooldown  = 30
    scale_out_cooldown = 120
  }
}

// Scales from 0 to 1 without waiting for queue to fill up
resource "aws_appautoscaling_policy" "sagemaker_policy_zero_to_one" {
  name               = "whisper-backlog-without-capacity-scaling-policy"
  policy_type        = "StepScaling"
  resource_id        = aws_appautoscaling_target.sagemaker_target.resource_id
  scalable_dimension = aws_appautoscaling_target.sagemaker_target.scalable_dimension
  service_namespace  = aws_appautoscaling_target.sagemaker_target.service_namespace
  step_scaling_policy_configuration {
    adjustment_type         = "ChangeInCapacity"
    cooldown                = 300
    metric_aggregation_type = "Average"
    step_adjustment {
      metric_interval_lower_bound = 0
      scaling_adjustment          = 1
    }
  }
}

resource "aws_cloudwatch_metric_alarm" "sagemaker_policy_zero_to_one" {
  alarm_name          = "whisper-backlog-without-capacity-scaling-policy"
  metric_name         = "HasBacklogWithoutCapacity"
  namespace           = "AWS/SageMaker"
  statistic           = "Average"
  evaluation_periods  = 2
  datapoints_to_alarm = 2
  comparison_operator = "GreaterThanOrEqualToThreshold"
  threshold           = 1
  treat_missing_data  = "missing"
  dimensions = {
    EndpointName = aws_sagemaker_endpoint.whisper.name
  }
  period            = 60
  alarm_description = "This metric is used to trigger the scaling policy that scales from 0 to 1 without waiting for queue to fill up"
  alarm_actions     = [aws_appautoscaling_policy.sagemaker_policy_zero_to_one.arn]
}

The problem

Why this is not working has to do with the inner workings of Terraform not fully understanding the update-chain that must happen to satisfy Sagemaker. Basically updating the model will result in one of two scenarios:

• The model and endpoint configuration is updated, but the changes aren't applied to the endpoint, so you have to destroy the endpoint in order to apply the changes (not good for production...)
• The endpoint apply fails because it's always referring to the previous version of the endpoint configuration (which terraform has destroyed), which leads to deleting the endpoint again...

The dirty solution

So if you really want to just be able to update the underlying code of the model (like any other container), push it to ECR and expect it to roll out to your endpoint, this is the only solution I've come up with so far.

First we will change the terraform code to:

resource "aws_sagemaker_model" "whisper" {
  execution_role_arn = aws_iam_role.whisper.arn
  primary_container {
    image          = data.aws_ssm_parameter.whisper_digest_id.value
    model_data_url = "s3://..."
    environment = {
      STAGE                 = terraform.workspace
      CONTAINER_ID          = data.aws_ssm_parameter.whisper_digest_id.value
      SENTRY_DSN            = data.aws_ssm_parameter.sentry_dsn.value
      SENTRY_ENVIRONMENT    = terraform.workspace
      INSTANCE_TYPE         = "ml.g4dn.xlarge"
    }
  }
}

resource "aws_cloudwatch_metric_alarm" "sagemaker_endpoint_error_rate" {
  alarm_name          = "EndToEndDeploymentHighErrorRateAlarm"
  alarm_description   = "Monitors the error rate of 4xx errors"
  metric_name         = "Invocation4XXErrors"
  namespace           = "AWS/SageMaker"
  statistic           = "Average"
  evaluation_periods  = 2
  comparison_operator = "GreaterThanThreshold"
  threshold           = 1
  period              = 600
  treat_missing_data  = "notBreaching"
  dimensions = {
    EndpointName = "whisper"
    VariantName  = "AllTraffic"
  }
}

data "external" "deploy_model" {
  program = ["python", "${path.module}/deploy_model.py"]

  query = {
    deploy_action                 = var.DEPLOY_ACTION
    aws_access_key_id             = var.ACCESS_KEY
    aws_secret_access_key         = var.SECRET_KEY
    endpoint_name                 = "whisper"
    endpoint_alarm_name           = "EndToEndDeploymentHighErrorRateAlarm"
    endpoint_config_model_name    = aws_sagemaker_model.whisper.name
    endpoint_config_model_image   = data.aws_ssm_parameter.whisper_digest_id.value
    endpoint_config_instance_type = "ml.g4dn.xlarge"
    endpoint_config_output_path   = "s3://..."
    endpoint_config_error_path    = "s3://..."
    endpoint_config_success_topic = aws_sns_topic.whisper_success_topic.arn
    endpoint_config_error_topic   = aws_sns_topic.whisper_error_topic.arn
  }
  depends_on = [aws_sagemaker_model.whisper]
}

output "model_deployment" {
  description = "Model deployment"
  value       = data.external.deploy_model.result
}

resource "aws_appautoscaling_target" "sagemaker_target" {
  max_capacity       = 30
  min_capacity       = 0
  resource_id        = "endpoint/${data.external.deploy_model.result["endpoint_name"]}/variant/AllTraffic"
  role_arn           = aws_iam_role.whisper.arn
  scalable_dimension = "sagemaker:variant:DesiredInstanceCount"
  service_namespace  = "sagemaker"
  depends_on         = [data.external.deploy_model]
}

resource "aws_appautoscaling_policy" "sagemaker_policy_regular" {
  name               = "whisper-invocations-scaling-policy"
  policy_type        = "TargetTrackingScaling"
  resource_id        = aws_appautoscaling_target.sagemaker_target.resource_id
  scalable_dimension = aws_appautoscaling_target.sagemaker_target.scalable_dimension
  service_namespace  = aws_appautoscaling_target.sagemaker_target.service_namespace

  target_tracking_scaling_policy_configuration {

    customized_metric_specification {
      metric_name = "ApproximateBacklogSizePerInstance"
      namespace   = "AWS/SageMaker"
      dimensions {
        name  = "EndpointName"
        value = data.external.deploy_model.result["endpoint_name"]
      }
      statistic = "Average"
    }
    target_value       = 10
    scale_in_cooldown  = 30
    scale_out_cooldown = 120
  }
}

// Scales from 0 to 1 without waiting for queue to fill up
resource "aws_appautoscaling_policy" "sagemaker_policy_zero_to_one" {
  name               = "whisper-backlog-without-capacity-scaling-policy"
  policy_type        = "StepScaling"
  resource_id        = aws_appautoscaling_target.sagemaker_target.resource_id
  scalable_dimension = aws_appautoscaling_target.sagemaker_target.scalable_dimension
  service_namespace  = aws_appautoscaling_target.sagemaker_target.service_namespace
  step_scaling_policy_configuration {
    adjustment_type         = "ChangeInCapacity"
    cooldown                = 300
    metric_aggregation_type = "Average"
    step_adjustment {
      metric_interval_lower_bound = 0
      scaling_adjustment          = 1
    }
  }
}

resource "aws_cloudwatch_metric_alarm" "sagemaker_policy_zero_to_one" {
  alarm_name          = "whisper-backlog-without-capacity-scaling-policy"
  metric_name         = "HasBacklogWithoutCapacity"
  namespace           = "AWS/SageMaker"
  statistic           = "Average"
  evaluation_periods  = 2
  datapoints_to_alarm = 2
  comparison_operator = "GreaterThanOrEqualToThreshold"
  threshold           = 1
  treat_missing_data  = "missing"
  dimensions = {
    EndpointName = data.external.deploy_model.result["endpoint_name"]
  }
  period            = 60
  alarm_description = "This metric is used to trigger the scaling policy that scales from 0 to 1 without waiting for queue to fill up"
  alarm_actions     = [aws_appautoscaling_policy.sagemaker_policy_zero_to_one.arn]
}

Basically we are moving the entire management of the CRUD of the endpoint/config/model to a Python script that runs boto3. In this case I've hard coded some values as you can see below, and some values I keep as parameters passed into the script. For this scenario I'm also passing in the terraform action (plan, apply etc.) since I don't want to run the script if it's not an apply action. Maybe not the best way, but it fits well with our Github actions.

import boto3
import json
import sys
from datetime import datetime
from time import sleep


def endpoint_exists(client, endpoint_name: str) -> tuple[str | None, str | None]:
    try:
        res = client.describe_endpoint(EndpointName=endpoint_name)
        return (
            res["EndpointConfigName"],
            res["ProductionVariants"][0]["DeployedImages"][0]["SpecifiedImage"],
        )
    except Exception:
        return None, None


def create_endpoint_config(
    client,
    name: str,
    model_name: str,
    instance_type: str,
    output_path: str,
    error_path: str,
    success_topic: str,
    error_topic: str,
):
    config_name = f"{name}-{datetime.now().strftime('%Y%m%d%H%M%S')}"
    res = client.create_endpoint_config(
        EndpointConfigName=config_name,
        ProductionVariants=[
            {
                "ModelName": model_name,
                "VariantName": "AllTraffic",
                "InstanceType": instance_type,
                "InitialInstanceCount": 1,
                "ContainerStartupHealthCheckTimeoutInSeconds": 1800,
                "ModelDataDownloadTimeoutInSeconds": 1200,
            }
        ],
        AsyncInferenceConfig={
            "OutputConfig": {
                "S3OutputPath": output_path,
                "S3FailurePath": error_path,
                "NotificationConfig": {
                    "IncludeInferenceResponseIn": [
                        "SUCCESS_NOTIFICATION_TOPIC",
                        "ERROR_NOTIFICATION_TOPIC",
                    ],
                    "SuccessTopic": success_topic,
                    "ErrorTopic": error_topic,
                },
            },
            "ClientConfig": {
                "MaxConcurrentInvocationsPerInstance": 1,
            },
        },
    )
    return config_name, res["EndpointConfigArn"]


def create_endpoint(
    client, endpoint_name: str, endpoint_config_name: str, alarm_name: str
):
    client.create_endpoint(
        EndpointName=endpoint_name,
        EndpointConfigName=endpoint_config_name,
        DeploymentConfig={
            "BlueGreenUpdatePolicy": {
                "TrafficRoutingConfiguration": {
                    "Type": "ALL_AT_ONCE",
                    "WaitIntervalInSeconds": 0,
                },
            },
            "AutoRollbackConfiguration": {
                "Alarms": [
                    {"AlarmName": alarm_name},
                ]
            },
        },
    )
    waiter = client.get_waiter("endpoint_in_service")
    waiter.wait(
        EndpointName=endpoint_name, WaiterConfig={"Delay": 10, "MaxAttempts": 60}
    )


def delete_endpoint_config(client, endpoint_config_name: str):
    client.delete_endpoint_config(EndpointConfigName=endpoint_config_name)


def update_endpoint(
    client, endpoint_name: str, endpoint_config_name: str, alarm_name: str
):
    retry_delay = 10
    for _ in range(60):
        try:
            client.update_endpoint(
                EndpointName=endpoint_name,
                EndpointConfigName=endpoint_config_name,
                RetainAllVariantProperties=True,
                RetainDeploymentConfig=True,
            )
            break
        except Exception as e:
            if "Cannot update in-progress endpoint" in str(e):
                sleep(retry_delay)
            else:
                raise e
    waiter = client.get_waiter("endpoint_in_service")
    waiter.wait(
        EndpointName=endpoint_name, WaiterConfig={"Delay": 10, "MaxAttempts": 120}
    )


def main(input: dict):
    deploy_action = input["deploy_action"]
    if deploy_action != "apply":
        print(
            json.dumps({"type": "no_change", "endpoint_name": input["endpoint_name"]})
        )
        return
    session = boto3.Session(
        aws_access_key_id=input["aws_access_key_id"],
        aws_secret_access_key=input["aws_secret_access_key"],
    )
    ec_model_name = input["endpoint_config_model_name"]
    ec_model_image = input["endpoint_config_model_image"]
    ec_instance_type = input["endpoint_config_instance_type"]
    ec_output_path = input["endpoint_config_output_path"]
    ec_error_path = input["endpoint_config_error_path"]
    ec_success_topic = input["endpoint_config_success_topic"]
    ec_error_topic = input["endpoint_config_error_topic"]
    endpoint_name = input["endpoint_name"]
    endpoint_alarm_name = input["endpoint_alarm_name"]

    client = session.client("sagemaker", region_name="eu-north-1")
    config_name, image_name = endpoint_exists(client, endpoint_name)
    if config_name:
        if ec_model_image == image_name:
            print(json.dumps({"type": "no_change", "endpoint_name": endpoint_name}))
        else:
            new_config_name, config_arn = create_endpoint_config(
                client,
                endpoint_name,
                ec_model_name,
                ec_instance_type,
                ec_output_path,
                ec_error_path,
                ec_success_topic,
                ec_error_topic,
            )
            update_endpoint(client, endpoint_name, new_config_name, endpoint_alarm_name)
            delete_endpoint_config(client, config_name)
            print(
                json.dumps(
                    {
                        "type": "update",
                        "endpoint_name": endpoint_name,
                        "endpoint_config_name": new_config_name,
                        "model_name": ec_model_name,
                        "old_endpoint_config_name": config_name,
                        "endpoint_config_arn": config_arn,
                    }
                )
            )
    else:
        new_config_name, config_arn = create_endpoint_config(
            client,
            endpoint_name,
            ec_model_name,
            ec_instance_type,
            ec_output_path,
            ec_error_path,
            ec_success_topic,
            ec_error_topic,
        )
        create_endpoint(client, endpoint_name, new_config_name, endpoint_alarm_name)
        print(
            json.dumps(
                {
                    "type": "new",
                    "endpoint_name": endpoint_name,
                    "endpoint_config_name": new_config_name,
                    "endpoint_config_arn": config_arn,
                }
            )
        )


if __name__ == "__main__":
    input = sys.stdin.read()
    input_json = json.loads(input)
    main(input_json)

Conclusion

It's unfortunate that this has been a forum problem for >4 years with no fix. But at least this solution allows you to correctly handle the Sagemaker logic for replacing the underlying containers in a safe way with progressive rollout and rollback.