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

こんにちは、大阪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関数が実行する処理(関数コード)は次の通りです。

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
    )

環境変数

自動停止の有効かどうかを判定するために用いるタグキーとタグ値は環境変数から取得しています。以下のように環境変数を登録します。

  • TAG_KEY : auto_stop
  • TAG_VALUE : true

実行ロール

管理ポリシーのAWSLambdaBasicExecutionRoleをベースにして、今回の処理で使用する権限を追加します。以下のActionに対する権限を有効にする必要があります。

  • sagemaker:ListTags
  • sagemaker:StopNotebookInstance
  • sagemaker:ListNotebookInstances

IAMロールの作成手順についてはドキュメントをご覧ください。

イベントトリガー

Designerのトリガーの追加からCloudWatch Eventsを選択し、新たにトリガールールを作成します。

  • ルールタイプ: スケジュール式
  • スケジュール式: cron(0 11 * * ? *)
    • cronの書き方でいつトリガーするかを定義します。
    • 今回の場合だと毎日20時にトリガーする設定です。
    • スケジュール式の内容はUTC基準なので注意が必要です。
      • JSTはUTCと9時間のズレがあるので、20時(JST)とするためには11時(UTC)で設定する必要があります。
  • Rate または Cron を使用したスケジュール式 - AWS Lambda

確認

SageMakerでノートブックインスタンスを作成します。そのインスタンスにKey=auto_stopValue=trueとなるタグをつけます。

インスタンスの起動が完了している状態で、先ほど作成したLambda関数を動かします。Lambda関数編集画面の上部から空のテストイベントを作成しテスト実行します。

SageMakerの対応リージョン全てに対する処理のため、実行時間はだいたい10秒くらいかかります。

作成したLambda関数が無事実行され、インスタンスの停止が開始されました。

さいごに

Lambdaを使ってSageMakerのノートブックインスタンスを自動停止できるようにしました。インスタンスが自動停止されると困る時は自動停止用のタグを変更することで自動停止から逃れることができます。そのままにするとまた停止忘れにつながるので注意が必要です。

お読みくださり、ありがとうございました〜!