はじめに
データアナリティクス事業本部のおざわです。在宅勤務で毎日エアコンをつけっぱなしの生活なので今月分の請求が恐ろしいです。。
さて、いつもお世話になっているSageMakerのノートブックインスタンスですが、ときどき起動に時間がかかることがあります。できればあまり待たずに使えたら嬉しいですよね。また、ノートブックインスタンスは、起動している間は課金されますので節約のためにも停止忘れを防ぎたいものです。
ということで、指定した時間になったら自動でインスタンスを起動・停止してみたいと思います。おまけでインスタンスのアイドル状態が一定時間続いたら自動で停止するように設定します。
今回は、こちらのブログを参考にLambdaやインスタンスのタグ設定を行っています。
指定した時間になったらLambdaを実行するためAmazon EventBridge Schedulerを使用します。
一定時間アイドル状態のインスタンスを停止するには、以下のリポジトリで公開されているライフサイクル設定のスクリプトを使用できます。
1. ノートブックインスタンスのタグ設定
ノートブックインスタンスはすでに存在するものとします。インスタンスにタグを設定して、自動起動・停止の対象かどうかを判断させます。今回は、タグのキーとして「auto_start_stop」 、値に「true」を設定しています。このタグが存在しなかったり、値がTrueでない場合は無視されます。
2. Lambda関数を作成
このブログにあるコードを少し変更したものを使いました。次に設定するAmazon EventBridge Schedulerからcommandパラメータを受け取って起動と停止の処理を分けています。
import os
import json
import boto3
# 環境変数を読み込む
TAG_KEY = os.environ["TAG_KEY"]
TAG_VALUE = os.environ["TAG_VALUE"]
REGION = os.environ["REGION"]
def lambda_handler(event, context):
command = event["command"]
status_mapping = {
"start": "Stopped",
"stop": "InService"
}
if command not in status_mapping:
return {
"statusCode": 400,
"body": json.dumps("Invalid command")
}
current_status = status_mapping[command]
sm = boto3.client("sagemaker", region_name=REGION)
# インスタンスを取得する
nb_list = sm.list_notebook_instances(
StatusEquals=current_status
)['NotebookInstances']
for nb in nb_list:
# インスタンスのタグを取得する
tags = sm.list_tags(ResourceArn=nb['NotebookInstanceArn'])['Tags']
# タグから自動停止・起動が有効かどうかを調べる
tagged_nb = False
for tag in tags:
if tag['Key'] == TAG_KEY and tag['Value'] == TAG_VALUE:
tagged_nb = True
break
nb_name = nb["NotebookInstanceName"]
if tagged_nb and command == "start":
print(f"starting instance: {nb_name}")
start_notebook_instance(nb_name, sm)
elif tagged_nb and command == "stop":
print(f"stopping instance: {nb_name}")
stop_notebook_instance(nb_name, sm)
return {
"statusCode": 200,
"body": json.dumps("done")
}
def stop_notebook_instance(nb_name, sm_client):
# 指定されたノートブックインスタンスを停止する
sm_client.stop_notebook_instance(
NotebookInstanceName=nb_name
)
def start_notebook_instance(nb_name, sm_client):
# 指定されたノートブックインスタンスを起動する
sm_client.start_notebook_instance(
NotebookInstanceName=nb_name
)環境変数
環境変数にREGION、TAG_KEY、TAG_VALUEを設定しています。 REGIONにはノートブックインスタンスがあるリージョンを設定してください。
Lambdaのロール設定
Lambda関数の実行ロールにノートブックインスタンスの起動、停止、一覧の取得、またインスタンスに付与されたタグの取得を許可するポリシーを追加します。以下のActionに対する権限を与えておきます。
- sagemaker:StopNotebookInstance
- sagemaker:StartNotebookInstance
- sagemaker:ListNotebookInstances
- sagemaker:ListTags
3. EventBridge Schedulerの設定
続いて、時間になったらLambda関数を呼び出すためにEventBridge Schedulerで設定を行います。 下の画面では、月曜日から金曜日の朝8時30分に実行するスケジュールを設定しています。
ターゲットにはAWS Lambda - Invokeを選択します。
先ほど作成したLambda関数を呼び出すようにして、起動か停止かを判断させるJSONを渡します。
停止用には、同じ手順でもう1つ別のスケジュールを作成して"stop"を渡すようにすれば設定完了です。
4. ライフサイクル設定
最後に、2時間アイドル状態が続いたら自動で落とすようにします。 ライフサイクルの設定は、Amazon SageMakerのライフサイクル設定画面に移動して、以下のノートブック開始時の設定画面にスクリプトとしてauto-stop-idleのスクリプトを貼り付けます。
17行目のIDLE_TIMEをお好みの長さに設定してください。
おわりに
さっそく試してみたところ、指定した時間に自動でノートブックインスタンスの起動と停止をすることができました。また、インスタンスをほったらかしにしておくと、ライフサイクル設定で停止してくれることも確認できました。これからはインスタンスを停止したか気になって夜中に起きることもなくなりそうです。
どなたかの参考になれば幸いです。
2
