はじめに

こんにちは。先日 SageMaker で推論エンドポイント作って放置していたら 2 日で 1 万円ぐらい請求金額が跳ね上がって泣きました。
テストで 2・3 個、推論エンドポイントを作ったのですが、できたことに満足して削除を怠ってしまいました・・・。
というわけで、このような悲劇を防ぐために、Lambda と EventBridge で推論エンドポイントを自動削除する仕組みを作ってみました。

推論エンドポイントを自動削除する Lambda

本記事では、あくまで開発時の削除漏れを目的としているので、用意した Lambda は、推論エンドポイントを全て削除しています。

import boto3
from typing import Dict, Any


def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
    """
    AWS Lambda関数のハンドラで、指定されたSageMakerエンドポイントを一括削除します。

    :param event: Lambda関数に渡されるイベントデータ
    :param context: Lambda関数のランタイム情報
    :return: ステータスコード
    """
    # SageMakerクライアントの作成
    sagemaker_client = boto3.client('sagemaker')

    # 既存のエンドポイントの一覧を取得
    endpoints = sagemaker_client.list_endpoints()

    # エンドポイントを一つずつ削除
    for endpoint in endpoints['Endpoints']:
        endpoint_name: str = endpoint['EndpointName']
        try:
            sagemaker_client.delete_endpoint(EndpointName=endpoint_name)
            print(f"エンドポイント {endpoint_name} を削除しました")
        except Exception as ex:
            print(f"エンドポイント {endpoint_name} の削除中にエラーが発生しました: {str(ex)}")

    return {
        'statusCode': 200
    }

Lambda に必要なロール

以下の AWS のマネージポリシーを付与したロールを用意し、Lambda にアタッチします。
SageMaker のフルアクセスまではいらないと言えばいらないのですが、サクッと作りたいので今回はこのポリシーを使います。

"arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
"arn:aws:iam::aws:policy/AmazonSageMakerFullAccess"

Lambda を EventBridge から実行可能にする

ここは忘れやすいので要注意です。
Lambda の画面の設定 → 　アクセス権限 → 実行ロール → リソースベースのポリシーステートメントです。

設定内容はこちら。

{
  "Version": "2012-10-17",
  "Id": "default",
  "Statement": [
    {
      "Sid": "{Sid}",
      "Effect": "Allow",
      "Principal": {
        "Service": "events.amazonaws.com"
      },
      "Action": "lambda:InvokeFunction",
      "Resource": "{LambdaのARN}",
      "Condition": {
        "ArnLike": {
          "AWS:SourceArn": "{LambdaのARN}"
        }
      }
    }
  ]
}

EventBridge で定期実行する

あとは EventBridge ルールで定期実行するだけです。
ルールを作成し、ターゲットに作成した Lambda を指定します。

これで推論エンドポイントの自動削除ができるようになりました。

ノートブックインスタンスの自動停止

実は推論エンドポイントだけでなく、ノートブックインスタンスも忘れがちなので、こちらも自動停止の Lambda を用意しました。
こちらも EventBridge で定期実行するだけです。

import boto3
from typing import Any, Dict


def lambda_handler(event: Dict[Any, Any], context: Any) -> Dict[str, Any]:
    """
    既存のノートブックインスタンスを一覧で取得し、'InService'状態のものを停止します。

    :param event: Lambda関数に渡されるイベントデータ
    :param context: Lambda関数の実行コンテキスト
    :return: statusCodeと停止操作の結果を含む辞書
    """

    # SageMakerクライアントの作成
    sagemaker_client = boto3.client('sagemaker')

    # 既存のノートブックインスタンスの一覧を取得
    notebook_instances: Dict[str, Any] = sagemaker_client.list_notebook_instances()

    # ノートブックインスタンスを一つずつ停止
    for notebook_instance in notebook_instances['NotebookInstances']:
        notebook_instance_name: str = notebook_instance['NotebookInstanceName']
        status: str = notebook_instance['NotebookInstanceStatus']

        # インスタンスが 'InService' 状態の場合のみ停止
        if status == 'InService':
            try:
                sagemaker_client.stop_notebook_instance(NotebookInstanceName=notebook_instance_name)
                print(f"ノートブックインスタンス {notebook_instance_name} を停止しました")
            except Exception as e:
                print(f"ノートブックインスタンス {notebook_instance_name} の停止中にエラーが発生しました: {str(e)}")

    return {
        'statusCode': 200
    }

本記事の Lambda と EventBridge を一気に構築する CloudFormation テンプレート

本記事で記述した 2 つの Lambda と、それを定期実行する EventBridge ルールを一気に構築する CloudFormation テンプレートを用意しました。
よろしければ活用してください。
2 つの Lambda を平日 21 時に実行するようにしています。

テンプレートのポイントは、AWS::Lambda::Permissionリソースを作っている部分です。
こちら、あまり見かけないリソースだと思います。
Lambda のリソースベースのポリシーステートメントに相当するリソースです。