SageMaker Studio ノートブックの自動終了を設定してみた

消し忘れたSageMaker Studioノートブックを自動終了する拡張機能を紹介します。

#AWS

#Amazon SageMaker

#SageMaker Studio

八木優成

2023.03.13

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

こんにちは、データアナリティクス事業本部の八木です。

皆さん、SageMaker Studio使ってますか？
SageMaker StudioはWebベースの統合開発環境であり、機能の一部としてJupyterノートブックを簡単に実行することができる非常に便利なサービスです。
そんなサービスでも注意したいのが料金です。SageMaker Studio ノートブックは実行時に自動的にコンピュートリソースを確保しますが、削除は手動で行う必要があるため、うっかり削除し忘れてしまうと余分なコストがかかってしまいます。

このようなケースの対策として、今回はAWS公式から提供されている拡張機能を使って、アイドル状態になっているSageMaker Studioノートブックを自動終了する方法をご紹介します。

なお、SageMakerノートブックインスタンス（not SageMaker Studio ノートブック）の場合は設定が異なります。SageMakerノートブックインスタンスで自動終了の設定を行いたい場合は、こちらのGitHubリポジトリの「auto-stop-idle」をご参照ください。

自動終了拡張機能の仕組み

自動終了拡張機能の仕組みを説明するにあたり、まずSageMaker Studioのアーキテクチャについて理解する必要があります。
SageMaker Studioは「JupyterServer」と「KernelGateway」の２種類のDockerコンテナから成り立っています。

画像引用：Dive deep into Amazon SageMaker Studio Notebooks architecture | AWS Machine Learning Blog

JupyterServerはUIをホスティングするサーバとなっており、ユーザプロファイルごとに1つ割り当てられています。ユーザがウェブブラウザからアクセスすると、このコンテナで実行されているNotebook Serverに接続します。
KernelGatewayは実際にカーネルを実行する環境となっており、ユーザごとに複数割り当てることができます。

料金が発生するのはKernelGatewayをホスティングしているEC2インスタンスです。

拡張機能はJupyterServerからKernelGatewayを監視し、アイドル状態が指定時間以上経過しているKernelGatewayアプリケーションを停止することで、使われていないEC2インスタンスを開放する仕組みとなっています。この機能により、アプリケーションを削除し忘れたときでも余分なコストを抑えることができます。

監視を行うプログラムをJupyterServerにインストールする方法はいくつかありますが、ライフサイクル設定を利用することが推奨されています。ライフサイクル設定を利用することで、JupyterServerが再作成された場合や、新たなユーザプロファイルが作成された場合も、最小の手間で設定を適用することができます。

やってみた

ということで、実際に設定を行っていきます。

前提

今回の記事では以下の環境を前提にしています。
- SageMaker Studioドメインが作成済み - SageMaker Studioユーザプロファイルが作成済み - SageMaker Studioのネットワーク設定は「パブリックインターネットのみ」 - Jupyter Lab のバージョンは3.0

ライフサイクル設定の作成

まず、ライフサイクル設定を作成します。

KernelGatewayを監視するプログラムはJupyterServer上で動くため、設定タイプは「Jupyter Serverアプリケーション」を選択します。

スクリプトの欄に以下のスクリプトを設定します。オリジナルのコードからTIMEOUT_IN_MINS（タイムアウトするアイドル時間）の値だけ変更しています。今回は5分に設定しました。

#!/bin/bash
# This script installs the idle notebook auto-checker server extension to SageMaker Studio
# The original extension has a lab extension part where users can set the idle timeout via a Jupyter Lab widget.
# In this version the script installs the server side of the extension only. The idle timeout
# can be set via a command-line script which will be also created by this create and places into the
# user's home folder
#
# Installing the server side extension does not require Internet connection (as all the dependencies are stored in the
# install tarball) and can be done via VPCOnly mode.

set -eux

# timeout in minutes
export TIMEOUT_IN_MINS=5

# Should already be running in user home directory, but just to check:
cd /home/sagemaker-user

# By working in a directory starting with ".", we won't clutter up users' Jupyter file tree views
mkdir -p .auto-shutdown

# Create the command-line script for setting the idle timeout
cat > .auto-shutdown/set-time-interval.sh << EOF
#!/opt/conda/bin/python
import json
import requests
TIMEOUT=${TIMEOUT_IN_MINS}
session = requests.Session()
# Getting the xsrf token first from Jupyter Server
response = session.get("http://localhost:8888/jupyter/default/tree")
# calls the idle_checker extension's interface to set the timeout value
response = session.post("http://localhost:8888/jupyter/default/sagemaker-studio-autoshutdown/idle_checker",
            json={"idle_time": TIMEOUT, "keep_terminals": False},
            params={"_xsrf": response.headers['Set-Cookie'].split(";")[0].split("=")[1]})
if response.status_code == 200:
    print("Succeeded, idle timeout set to {} minutes".format(TIMEOUT))
else:
    print("Error!")
    print(response.status_code)
EOF
chmod +x .auto-shutdown/set-time-interval.sh

# "wget" is not part of the base Jupyter Server image, you need to install it first if needed to download the tarball
sudo yum install -y wget
# You can download the tarball from GitHub or alternatively, if you're using VPCOnly mode, you can host on S3
wget -O .auto-shutdown/extension.tar.gz https://github.com/aws-samples/sagemaker-studio-auto-shutdown-extension/raw/main/sagemaker_studio_autoshutdown-0.1.5.tar.gz

# Or instead, could serve the tarball from an S3 bucket in which case "wget" would not be needed:
# aws s3 --endpoint-url [S3 Interface Endpoint] cp s3://[tarball location] .auto-shutdown/extension.tar.gz

# Installs the extension
cd .auto-shutdown
tar xzf extension.tar.gz
cd sagemaker_studio_autoshutdown-0.1.5

# Activate studio environment just for installing extension
export AWS_SAGEMAKER_JUPYTERSERVER_IMAGE="${AWS_SAGEMAKER_JUPYTERSERVER_IMAGE:-'jupyter-server'}"
if [ "$AWS_SAGEMAKER_JUPYTERSERVER_IMAGE" = "jupyter-server-3" ] ; then
    eval "$(conda shell.bash hook)"
    conda activate studio
fi;
pip install --no-dependencies --no-build-isolation -e .
jupyter serverextension enable --py sagemaker_studio_autoshutdown
if [ "$AWS_SAGEMAKER_JUPYTERSERVER_IMAGE" = "jupyter-server-3" ] ; then
    conda deactivate
fi;

# Restarts the jupyter server
nohup supervisorctl -c /etc/supervisor/conf.d/supervisord.conf restart jupyterlabserver

# Waiting for 30 seconds to make sure the Jupyter Server is up and running
sleep 30

# Calling the script to set the idle-timeout and active the extension
/home/sagemaker-user/.auto-shutdown/set-time-interval.sh

以上の内容でライフサイクル設定を作成します。

ライフサイクル設定のアタッチ

続いて作成したライフサイクル設定をSageMaker Studioドメインにアタッチします。「ドメイン詳細」→「環境」→「個人用 Studio アプリのライフサイクル設定」で作成したライフサイクル設定をアタッチします。

また、アタッチするだけでなく、デフォルトとしても設定する必要があります。

以上で設定は完了しましたが、すでにJupyterServerがある場合は一度削除し、作成し直す必要がります。JupyterServer用のライフサイクル設定は起動時に実行される機能であり、すでに起動しているJupyterServerには適用されないためです。

ユーザプロファイルから詳細画面へ移動し、JupyterServerを削除します。

削除が完了したら、SageMaker Studioを起動します。

ユーザプロファイルに紐付いたJupyterServerが存在しない場合は自動的に作成されるため、手動での作成は不要です。

以上で自動終了の設定が完了しました。

自動終了の動作確認

実際に、タイムアウトに設定した時間を過ぎるとKernelGateway及びインスタンスが終了するか、確認してみます。

ノートブックを作成して適当なスクリプトを実行し、しばらく待機します。

5分ほど経過すると、KERNEL SESSIONが終了し、インスタンスも開放されました！

注意点

今回の記事では、作成したライフサイクル設定をドメインのデフォルトとして設定しましたが、ユーザプロファイルレベルでデフォルトのライフサイクル設定がされていた場合、ドメインのデフォルトライフサイクル設定は実行されず、自動終了処理が実行されません。この場合はユーザプロファイルに設定されているライフサイクル設定にも、自動終了設定のスクリプトを追加しましょう。

また、ライフサイクル設定で他の処理も行う場合は、実行順序に気をつけましょう。
以下の行で監視処理を開始していますが、この行のあとでrestart-jupyter-serverコマンドなどを実行すると、監視処理は行われなくなってしまいます。

/home/sagemaker-user/.auto-shutdown/set-time-interval.sh

また、紹介したドメインのデフォルト設定はユーザの個人環境用の設定となっており、共有スペースでも自動終了設定を適用するには、共有スペース用のデフォルト設定を行う必要があります。記事執筆時点ではコンソールからは設定できないため、CLIなどから設定を行ってください。

最後に

今回の記事では、SageMaker Studio ノートブックの自動終了設定についてご紹介しました。
リソースの消し忘れで思わぬ料金が発生してしまった！というケースはSageMaker Studioに限らず起こりがちです。今回ご紹介した設定を行うことで、消し忘れが発生した場合も影響を最小限にできるでしょう。

ミスを防ぐことも大切ですが、人はどうしてもミスをするもの。なんらかのミスがあった場合も影響が出にくい設計を心がけていきましょう！

以上、八木でした！