AWS Builder Center

Builder Center

Q-Bits: Auto Cleanup Sagemaker Endpoints with Amazon Q Developer

Q-Bits: Auto Cleanup Sagemaker Endpoints with Amazon Q Developer

The blog post show the power of Amazon Q Developer command line being able to save hundreds by cleaning up developer Sagemaker Endpoints

Michael Kaiser

Welcome to another installment of  Q-Bits , our regular series showcasing cool ways Amazon employees are leveraging Amazon Q Developer. Today, we're diving into how Q Developer can assist with building a solution for autoclean up Sagemaker Endpoints.
I work with customer to help them in their cloud journey. One of those calls I another AWS colleague had commented about sometimes forgetting to delete Sagemaker endpoints and it costing hundreds to thousands of dollars for idle endpoints. That gave me an idea to see if I can build a solution that would automatically delete endpoints after a set amount of time. I wanted to try out the  Amazon Q Developer CLI  and see what it could do with just a couple prompts.
I started thinking I would do a config rule with  RDK , to do reporting and it would auto remediate. So I loaded up a terminal and got into Q with
Bash
q chat -a
(the -a tells Q to auto accept running commands instead of asking me every time. I really wanted to let it go and quickly get a solution).
From here I gave it the prompt
plaintext
create rdk application that automatically cleans up Amazon SageMaker endpoints that are older than 48 hours
  1. Now Q starts going through:
  1. checking if I have the correct pip packages
  1. attempting to install them (it fails because don't have a python env setup and it was relying my system python. Here if I was confirming each command I could have told it to run and jump into my pyenv environment)
  1. It pivots to creating a solution manually using  AWS Serverless Application Model , sam cli.
  1. Creates the lambda python code with a function for getting the current Sagemaker endpoints, methods to check comparing the current time with the time that the endpoint was created to see if it is older than 48 hours.
  1. Creates a README.md, samconfig.toml, requirements.txt, and  Amazon CloudFormation  template.yaml
  1. Q wraps up by letting me all that it has done, how to run sam, and an overview of the solution.
We can see the main logic of app and at first glance it looks good. I wanted to deploy it to test it out, but I don't have sam cli installed.
Python
def lambda_handler(event, context):
    """
    AWS Lambda function to identify and delete SageMaker endpoints that are older than 48 hours.
    
    Parameters:
    event (dict): AWS Lambda event data
    context (object): AWS Lambda context object
    
    Returns:
    dict: Results of the cleanup operation
    """
    logger.info("Starting SageMaker endpoint cleanup process")
    
    # Initialize SageMaker client
    sagemaker_client = boto3.client('sagemaker')
    
    # Get current time
    current_time = datetime.datetime.now(datetime.timezone.utc)
    
    # Get all endpoints
    try:
        response = sagemaker_client.list_endpoints()
        endpoints = response['Endpoints']
        
        # Process paginated results if any
        while 'NextToken' in response:
            response = sagemaker_client.list_endpoints(NextToken=response['NextToken'])
            endpoints.extend(response['Endpoints'])
            
        logger.info(f"Found {len(endpoints)} SageMaker endpoints")
    except ClientError as e:
        logger.error(f"Error listing SageMaker endpoints: {e}")
        return {
            'statusCode': 500,
            'message': f"Error listing SageMaker endpoints: {str(e)}"
        }
    
    # Track endpoints processed
    endpoints_deleted = []
    endpoints_skipped = []
    errors = []
    
    # Check each endpoint
    for endpoint in endpoints:
        endpoint_name = endpoint['EndpointName']
        creation_time = endpoint['CreationTime']
        
        # Calculate age of endpoint in hours
        endpoint_age = (current_time - creation_time).total_seconds() / 3600
        
        # Check if endpoint is older than 48 hours
        if endpoint_age >= 48:
            logger.info(f"Endpoint {endpoint_name} is {endpoint_age:.2f} hours old and will be deleted")
            
            try:
                # Delete the endpoint
                sagemaker_client.delete_endpoint(EndpointName=endpoint_name)
                endpoints_deleted.append({
                    'name': endpoint_name,
                    'age_hours': endpoint_age
                })
                logger.info(f"Successfully deleted endpoint: {endpoint_name}")
            except ClientError as e:
                error_message = f"Error deleting endpoint {endpoint_name}: {str(e)}"
                logger.error(error_message)
                errors.append(error_message)
        else:
            logger.info(f"Endpoint {endpoint_name} is {endpoint_age:.2f} hours old and will be kept")
            endpoints_skipped.append({
                'name': endpoint_name,
                'age_hours': endpoint_age
            })
    
    # Prepare result
    result = {
        'statusCode': 200,
        'summary': {
            'total_endpoints': len(endpoints),
            'endpoints_deleted': len(endpoints_deleted),
            'endpoints_skipped': len(endpoints_skipped),
            'errors': len(errors)
        },
        'deleted_endpoints': endpoints_deleted,
        'skipped_endpoints': endpoints_skipped,
        'errors': errors
    }
    
    logger.info(f"SageMaker endpoint cleanup completed. Deleted: {len(endpoints_deleted)}, Skipped: {len(endpoints_skipped)}, Errors: {len(errors)}")
    return result

But I do have and I am very familiar with writing  CDK  apps and using  Amazon Q Developer  inside of VSCode, running cdk init, and then using /dev to have it get me majority if not a complete solution for what was looking for as a solution. Again I want to see if I can just stay in the CLI without opening an IDE. So I just tell Q to convert it to a CDK app using.
plaintext
convert app to cdk
With another round of it going through commands and creating folders and files, it had generated working cdk app that I could not tell the difference if I had hand created. With the main part of the cdk app being
TypeScript
   // Create IAM role for the Lambda function
    const lambdaRole = new iam.Role(this, 'SageMakerEndpointCleanupRole', {
      assumedBy: new iam.ServicePrincipal('lambda.amazonaws.com'),
      description: 'Role for SageMaker Endpoint Cleanup Lambda function',
    });

    // Add required policies
    lambdaRole.addManagedPolicy(iam.ManagedPolicy.fromAwsManagedPolicyName('AmazonSageMakerFullAccess'));
    lambdaRole.addManagedPolicy(iam.ManagedPolicy.fromAwsManagedPolicyName('CloudWatchLogsFullAccess'));

    // Create Lambda function
    const cleanupFunction = new lambda.Function(this, 'SageMakerEndpointCleanupFunction', {
      runtime: lambda.Runtime.PYTHON_3_9,
      handler: 'sagemaker_endpoint_cleanup.lambda_handler',
      code: lambda.Code.fromAsset(path.join(__dirname, '../lambda')),
      role: lambdaRole,
      timeout: cdk.Duration.seconds(300),
      memorySize: 128,
      description: 'Identifies and deletes SageMaker endpoints older than 48 hours',
    });

    // Create CloudWatch Events Rule to trigger the Lambda function daily
    const rule = new events.Rule(this, 'SageMakerEndpointCleanupSchedule', {
      schedule: events.Schedule.rate(cdk.Duration.minutes(10)),
      description: 'Triggers SageMaker endpoint cleanup once per day',
    });

    // Add the Lambda function as a target for the rule
    rule.addTarget(new targets.LambdaFunction(cleanupFunction));

    // Output the Lambda function ARN and role ARN
    new cdk.CfnOutput(this, 'SageMakerEndpointCleanupFunctionArn', {
      description: 'SageMaker Endpoint Cleanup Lambda Function ARN',
      value: cleanupFunction.functionArn,
    });

    new cdk.CfnOutput(this, 'SageMakerEndpointCleanupFunctionRoleArn', {
      description: 'IAM Role created for the cleanup function',
      value: lambdaRole.roleArn,
    });
Now I have something I can test. I changed line 52 to the following, so that I don't have to wait 48 hours for the endpoint to be deleted.
Python
 if endpoint_age >= 1:
So I deploy the solution with cdk deploy and it deploys successfully! I deploy a Sagemaker Endpoint and after 1 hour the Sagemaker Endpoint is deleted with the CloudWatch Logs logging it!
With Amazon Q Developer CLI I was able to create a solution with two prompts that helps auto clean up Sagemaker endpoints in Dev accounts with only two prompts using natural language. I was truly amazed how Q has been constantly improving and is established itself as a critical tool.  
# amazon-q-developer# q-bits# cdk# sagemaker
Any opinions in this article are those of the individual author and may not reflect the opinions of AWS.
Enjoyed reading this content? Let the author know!

Your likes, comments, shares, and saves help creators reach more builders.

Loading recommendations

Loading article