SageMaker Notebook instances and SageMaker Studios are expensive. Insanely expensive especially if you leave them unterminated. Fortunately, there are ways to set up auto-shutdown of both SageMaker Notebook and SageMaker Studio instances when they are idling.

SageMaker Notebook Instance

AWS generously provided the scripts to facilitate stopping SageMaker Notebook instance automatically using a background cron job.

Lifecycle Configuration

We will start fixing SageMaker Notebook Instances by specifying the on-start.sh shell script. This is a standard bash script that runs everytime SageMaker Notebook instances boots up.

#!/bin/bash

set -e

# PARAMETERS
IDLE_TIME=3600

echo "Fetching the autostop script"
wget https://raw.githubusercontent.com/aws-samples/amazon-sagemaker-notebook-instance-lifecycle-config-samples/master/scripts/auto-stop-idle/autostop.py

echo "Starting the SageMaker autostop script in cron"

(crontab -l 2>/dev/null; echo "*/5 * * * * /usr/bin/python $PWD/autostop.py --time $IDLE_TIME --ignore-connections") | crontab -

To setup the on-start script, do the following:

Go to Amazon SageMaker -> Notebook -> Lifecycle Configurations
Click on Create Configuration
On the next page, name your configuration, e.g., auto-stop-idle
Copy the above bash script to the Start notebook tab.
Click on Create Configuration.
Now the configuration should be created. Click on Notebook instances on the left hand side panel.
Click on the Notebook instances that you want to auto stop when idle.
Click on Edit, and at the Lifestyle configuration - optional dropdown, select auto-stop-idle or your named configuration.
Click on Update notebook instances.

Permissions and Policy

In addition to setting up the on-start.sh script, we will have to add permissions to sageMaker:StopNotebookInstance to stop the notebook instance and sageMaker:DescribeNotebookInstance to describe the notebook instance. We can do it with the following:

Go to the Notebook Instance, click on IAM role ARN
Click on Attached Policies
Click Create Policy

Past the following JSON to allow describe and stop notebook instance action

{
"Version": "2012-10-17",
"Statement": [
    {
        "Sid": "VisualEditor0",
        "Effect": "Allow",
        "Action": [
            "sagemaker:StopNotebookInstance",
            "sagemaker:DescribeNotebookInstance"
        ],
        "Resource": "*"
    }
]
}

Save the policy with a name, e.g. Sagemaker-Auto-Stop
Attach the newly created policy to the IAM role

This is easy enough and should let you use SageMaker Notebook Instances safely.

SageMaker Studio Auto-Shutdown Lambda Function

SageMaker Studio auto-shutdown is slightly more complicated as it hides the booted up instances under the containers running on top of them. AWS offers an example repo here for setting up auto-shutdown of SageMaker Studio instances via an AWS Lambda function that monitors the instance.

Because AWS Lambda is involved, it’s best to have both aws cli version 2 and SAM CLI installed to facilitate the process.

To download, package, and deploy the auto-shutdown Lambda function, run:

git clone https://github.com/aws-samples/sagemaker-studio-auto-shutdown-extension.git
cd auto-installer

export DEPLOYMENT_BUCKET_NAME=YOUR_BUCKET_NAME
export DEPLOYMENT_BUCKET_PREFIX=IF_YOU_WANT_ONE
export DEPLOYMENT_REGION=LAMBDA REGION

aws s3api create-bucket --bucket $DEPLOYMENT_BUCKET_NAME --region $DEPLOYMENT_REGION  --create-bucket-configuration LocationConstraint=$DEPLOYMENT_REGION

sam build --use-container --template template.sam.yaml
sam package --s3-bucket $DEPLOYMENT_BUCKET_NAME --s3-prefix $DEPLOYMENT_BUCKET_PREFIX --output-template-file template.tmp.yaml
sam deploy --template-file template.tmp.yaml --capabilities CAPABILITY_IAM --no-fail-on-empty-changeset --stack-name smstudio-setup

This Lambda function to install auto-stop instance only when ‘default’ JupyterServer app is started in SageMaker Studio. And the ‘default’ app only starts when brand new SageMaker Studio user is created. Thus, to retro-fit existing SageMaker Studio user, please do:

Click on the username to retrofit
Delete the ‘default’ app
Return to the SageMaker Studio Control Panel
Click on Open Studio to launch a brand new ‘default’ app.

Once launched, you will see the prompt asking you to build Jupyter Extension.

Once it’s done, you will see the prompt asking you to reload the instance to update the changes.

After reloading, an auto-shutdown button should appear at the left hand side control panel. You can change the idle time limits. However, due to the design, the instances won’t get shut down on the dot.

References:

AWS Samples Github

Save costs by automatically shutting down idle resources within Amazon SageMaker Studio

To cite this content, please use:

@article{
    leehanchung,
    author = {Lee, Hanchung},
    title = {Sagemaker: Automatically Stop Notebook Instances When Idle},
    year = {2021},
    howpublished = {\url{https://leehanchung.github.io}},
    url = {https://leehanchung.github.io/2021-06-11-sagemaker-auto-terminate/}
}