Lambda + EventBridgeでSageMaker StudioのKernelGatewayを自動削除して不要なコストを削減する

データアナリティクス事業本部の貞松です。 最近は日々SageMakerと格闘しております。

機械学習のWebベースIDEであるSageMaker Studioは、Jupyter Notebookによる開発時にKernelという単位で管理される環境を利用することで、Notebookを実行する為のインスタンスや各種設定やライブラリが搭載されたコンテナ環境を柔軟に切り替えながら開発を行うことができます。

実行されているKernelはKernelGateWayというタイプのアプリケーションとして、SageMaker Studioの実行ユーザー単位で作成されるのですが、このKernelGatewayが起動している間は、Kernelを実行しているインスタンスタイプに応じて時間単位で課金が発生します。

また、KernelGatewayは、SageMaker Studioの画面を閉じたり、SageMaker Studioの本体(Jupyter Server)をシャットダウンしても停止しない為、ユーザーが明示的にKernelのインスタンスをシャットダウンする、あるいは該当のKernelGatewayタイプのアプリケーションを削除する必要があります。

AWSマネジメントコンソールの画面やSageMaker Studioの画面から手動で削除しても良いのですが、毎回手動で削除するのは手間がかかりますし、何より削除忘れによる無駄な課金が発生してしまう可能性があります。

そこで今回は、AWS SDK経由でSageMaker StudioのKernelGatewayを削除する処理をLambda関数として実装し、その関数をEventBridgeからスケジュール実行する仕組みを作成します。

SageMaker StudioのKernelGatewayを削除するLambda関数を作成する

開発環境

今回作成したLambda関数のランタイムはPython3.8を使用しています。

作成した関数の処理仕様

今回作成する処理の仕様は以下の通りです。

• SageMaker StudioのDomainIdを環境変数に設定して、該当のDomainIdを持つSageMaker StudioのAppのリストを取得する。
• 取得したAppのリストからAppType=KernelGateweyかつStatus=InServiceのAppのUserProfileNameとAppNameを取得して別のリストに詰め直す。
• 詰め直したリストに含まれるApp(KernelGateway)を全て削除する。

実行に必要なIAMロールの権限設定

上記の仕様に沿って処理を実行する為に、Lambdaの実行ロールに以下の権限を追加します。

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "sagemaker:ListApps",
                "sagemaker:DeleteApp"
            ],
            "Resource": "*"
        }
    ]
}

環境変数の設定

Lambda関数の環境変数としてDOMAIN_IDを設定します。

SageMaker StudioのDomainIdはAWSマネジメントコンソールのSageMaker画面上ではStudio IDとして表示されているので、この値を設定します。

Pythonコードの作成

AWS SDK for Python (Boto3)を使用して、SagaMaker StudioからAppのリストを取得し、KernelGatewayを削除する処理を実装します。
今回使用するBoto3のSageMakerクライアントの仕様については下記URLのページをご参照ください。

Boto3 - SageMaker.Client.list_apps

Boto3 - SageMaker.Client.delete_app

以下、作成したLambda関数のPythonコードです。
前述で設定した環境変数を先頭で読み込んで使用しています。

# sagemaker_studio_delete_apps.ｐｙ

import boto3
import json
import os

domain_id = os.environ['DOMAIN_ID']

def delete_apps(event, context):
    client = boto3.Session().client('sagemaker')

    # 指定されたのDomainIdのSageMaker StudioからAppのリストを取得
    response = client.list_apps(
        DomainIdEquals=domain_id
    )
    
    # 取得したAppのリストから
    # AppTypeがKernelGateweyかつStatusがInServiceのAppの
    # UserProfileNameとAppNameを取得してリストに詰め直す
    app_list = []
    apps = response['Apps']
    for app in apps:
        if app['AppType'] == 'KernelGateway' and app['Status'] == 'InService':
            app_list.append((app['UserProfileName'], app['AppName']))
    
    # 詰め直したリストに含まれるApp(KernelGateway)を削除する
    for app in app_list:
        client.delete_app(
            DomainId = domain_id,
            UserProfileName = app[0],
            AppType = 'KernelGateway',
            AppName = app[1]
        )

    return {
        'statusCode': 200,
        'body': json.dumps('Started deleting apps.')
    }

今回作成したPythonのファイル名はsagemaker_studio_delete_appsで関数名をdelete_appsとしているので、ランタイム設定のハンドラはsagemaker_studio_delete_apps.delete_appsとなっています。ここは適宜作成したファイルの内容に合わせて書き換えてください。

作成したLambda関数をスケジュール実行するEventBridgeルールを作成する

あとはEventBridgeで、作成したLambda関数を定時スケジュール実行するためのルールを作成するだけです。

ルールの作成画面で、パターン定義をCron式で設定したスケジュール実行とします。

そして、ターゲットに前述で作成したLambda関数を指定します。

これでEventBridgeのルールも作成完了です。

作成した仕組みの動作を確認する

作成した仕組みの動作を検証する為にSageMaker Studio上のNotebookで適当なKernelを一つ起動しておきます。

EventBridgeルールで設定した時間に再度Appの状態を確認すると、KernelGatewayのみDeletingになっており、処理が正常に動作していることが確認できました。

まとめ

SageMaker StudioのKernelGatewayを自動で削除する為の仕組みをLambda + EventBridgeで作成することにより、KernelGatewayの削除忘れによる余分なコストを削減することができそうです。

SageMaker Studio自体のアップデートによりこういった機能が追加されると嬉しいのですが、当面はこういった仕組みを活用して手間を省きつつコスト削減するのが良いと思います。