Amazon SageMakerのノートブックインスタンスを自動停止させるLambda関数を作ってみた

osawa yuto

2019.01.30

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

こんにちは、大阪DI部の大澤です。

みなさん、SageMakerのノートブックインスタンス使っていますか？僕は良く使ってるんですが、そこで起こる問題があります。インスタンスの停止忘れです。ノートブックインスタンス上で作業を行なって学習ジョブを回してる間とかにしばらく別の作業を行なってそのまま帰る...みたいな感じになるともれなくインスタンスを停止し忘れてしまいます。ml.t2.mediumのインスタンスを東京リージョンで、19時から次の日の朝の10時まで停止し忘れていた場合、およそ$0.9程度無駄になります。缶コーヒー買えます。1ヶ月放置すると三千円程度無駄になります...。

とはいえ、学習ジョブなどで待ちが発生するたびにノートブックインスタンスを停止させるのも億劫です。そこで消し忘れていてももう使わないだろう時間帯(20時くらい)に自動的にノートブックインスタンスを停止させるLambda関数を作ってみました。

Lambda関数

毎日20時ごろに自動実行
CloudWatch Eventsのスケジュールでトリガー
起動中のノートブックインスタンスを取得し、タグauto_stop=trueが付いていればインスタンスを停止
SageMakerが対応しているリージョンほぼ全てに対応
boto3.Session().get_available_regions('sagemaker')で取得できるリージョン
ランタイム: Python 3.7

Lambda関数の作成手順に関する詳細はドキュメントをご覧ください。

コンソールで Lambda 関数を作成する - AWS Lambda

関数コード

Lambda関数が実行する処理(関数コード)は次の通りです。

import json
import boto3
import os

# 環境変数を読み込む
TAG_KEY = os.environ['TAG_KEY']
TAG_VALUE = os.environ['TAG_VALUE']


def lambda_handler(event, context):
    # SageMakerが対応してるリージョン一覧を取得
    regions = boto3.Session().get_available_regions('sagemaker')

    # リージョンごとに実行
    for region in regions:
        sm = boto3.client('sagemaker', region_name=region)


        # 起動中(InService)のインスタンスを取得する
        nb_list = sm.list_notebook_instances(
            StatusEquals='InService'
        )['NotebookInstances']

        for nb in nb_list:
            # インスタンスのタグを取得する
            tags = sm.list_tags(ResourceArn=nb['NotebookInstanceArn'])['Tags']

            # タグから自動停止が有効かどうかを調べる
            stop_nb = False
            for tag in tags:
                if tag['Key'] == TAG_KEY and tag['Value'] == TAG_VALUE:
                    stop_nb = True
                    break

            # 自動停止が有効なノートブックインスタンスであれば停止させる
            if stop_nb:
                nb_name = nb['NotebookInstanceName']
                print('stop', nb_name+'@'+region)
                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
    )