SageMakerのスペースとエンドポイントの停止忘れをチェックするLambdaを作ってみた

SageMakerのスペースとエンドポイントの停止忘れをチェックするLambdaを作ってみた

#Amazon SageMaker
#SageMaker Studio
よなみね

よなみね

facebook logohatena logotwitter logo
Clock Icon2024.03.25

はじめに

みなさんこんにちは、先日SageMakerのエンドポイントを消し忘れて200ドルの請求がきたよなです。
今回はエンドポイントと合わせて、Studio内で起動したスペースの停止忘れをチェックするためのLambdaを作成してみました。

エンドポイントのチェックは他の記事を参考にすぐに実装できたのでコードのみ記載しています。
今回は特にスペースのステータスを確認することにフォーカスした記事を書いています。

構成図

今回作成する構成図はこのような感じです。

全体の流れはこんな感じです。

  1. EventBridgeで定期的にLambdaを実行
  2. Lambdaでリソースのステータスをチェック
  3. SNSでステータスをメール通知

SageMakerのリソースは学習を実行している可能性もあるので、削除ではなくステータスの確認のみを行います。

チェック用のリソースを作成するCloudFormationテンプレート

先にチェック用のリソースを作成するテンプレートを載せておきます。

AWSTemplateFormatVersion: '2010-09-09'
Description: 'AWS CloudFormation template for deploying the SageMaker Spaces status checker and SageMaker Endpoint checker Lambda function and EventBridge rule'

Parameters:
  ScheduleExpression:
    Type: String
    Default: 'cron(0 9 * * ? *)'
    Description: 'The schedule expression for running the Lambda function'
  EmailAddress:
    Type: String
    Default: '[email protected]'
    Description: 'The email address to send notifications to'

Resources:
  SageMakerResourceCheckerLambda:
    Type: 'AWS::Lambda::Function'
    Properties:
      FunctionName: 'SageMakerResourceChecker'
      Handler: 'index.lambda_handler'
      Runtime: 'python3.12'
      Timeout: 10
      Role: !GetAtt LambdaExecutionRole.Arn
      Environment:
        Variables:
          APP_STATUS_TOPIC_ARN: !Ref SageMakerAppStatusTopic
          ENDPOINT_STATUS_TOPIC_ARN: !Ref SageMakerEndpointStatusTopic
      Code:
        ZipFile: |
          import boto3
          import json
          import os

          def lambda_handler(event, context):
              sagemaker_client = boto3.client('sagemaker')
              sns_client = boto3.client('sns')

              # Check SageMaker space status
              app_response = sagemaker_client.list_apps()
              apps = app_response['Apps']

              app_status_list = []
              for app in apps:
                  space_name = app.get('SpaceName', '')
                  app_type = app['AppType']
                  app_name = app['AppName']
                  app_status = app['Status']

                  status = '起動中' if app_status == 'InService' else '停止中'
                  app_status_list.append(f"スペース名: {space_name}\nアプリケーションタイプ: {app_type}\nアプリケーション名: {app_name}\nステータス: {status}\n---")

              app_message = "現在のSageMakerアプリケーションの状況一覧:\n\n" + "\n".join(app_status_list)

              sns_client.publish(
                  TopicArn=os.environ['APP_STATUS_TOPIC_ARN'],
                  Message=app_message,
                  Subject='SageMakerアプリケーションの状況レポート'
              )

              # Check SageMaker endpoint status
              endpoint_response = sagemaker_client.list_endpoints()
              endpoints = endpoint_response['Endpoints']

              if endpoints:
                  endpoint_status_list = []
                  for endpoint in endpoints:
                      endpoint_name = endpoint['EndpointName']
                      endpoint_status = endpoint['EndpointStatus']
                      creation_time = endpoint['CreationTime']
                      last_modified_time = endpoint['LastModifiedTime']

                      endpoint_status_list.append(f"エンドポイント名: {endpoint_name}\nステータス: {endpoint_status}\n作成時刻: {creation_time}\n最終更新時刻: {last_modified_time}\n---")

                  endpoint_message = "現在のSageMakerエンドポイントの状況一覧:\n\n" + "\n".join(endpoint_status_list)
              else:
                  endpoint_message = "現在起動中のエンドポイントはありません。"

              sns_client.publish(
                  TopicArn=os.environ['ENDPOINT_STATUS_TOPIC_ARN'],
                  Message=endpoint_message,
                  Subject='SageMakerエンドポイントの状況レポート'
              )

              return {
                  'statusCode': 200,
                  'body': 'SageMaker application and endpoint status checked and email sent successfully.'
              }

  LambdaExecutionRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: 'sts:AssumeRole'
      Policies:
        - PolicyName: 'SageMakerResourceCheckerPolicy'
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - 'sagemaker:ListApps'
                  - 'sagemaker:ListEndpoints'
                Resource: '*'
              - Effect: Allow
                Action:
                  - 'sns:Publish'
                Resource:
                  - !Ref SageMakerAppStatusTopic
                  - !Ref SageMakerEndpointStatusTopic

  SageMakerAppStatusTopic:
    Type: 'AWS::SNS::Topic'
    Properties:
      TopicName: 'SageMakerAppStatusTopic'
      Subscription:
        - Endpoint: !Ref EmailAddress
          Protocol: 'email'

  SageMakerEndpointStatusTopic:
    Type: 'AWS::SNS::Topic'
    Properties:
      TopicName: 'SageMakerEndpointStatusTopic'
      Subscription:
        - Endpoint: !Ref EmailAddress
          Protocol: 'email'

  SageMakerResourceCheckerRule:
    Type: 'AWS::Events::Rule'
    Properties:
      Name: 'SageMakerResourceCheckerRule'
      ScheduleExpression: !Ref ScheduleExpression
      State: 'ENABLED'
      Targets:
        - Arn: !GetAtt SageMakerResourceCheckerLambda.Arn
          Id: 'SageMakerResourceCheckerLambdaTarget'

  LambdaInvokePermission:
    Type: 'AWS::Lambda::Permission'
    Properties:
      FunctionName: !Ref SageMakerResourceCheckerLambda
      Action: 'lambda:InvokeFunction'
      Principal: 'events.amazonaws.com'
      SourceArn: !GetAtt SageMakerResourceCheckerRule.Arn

デプロイ方法

CloudFormationでデプロイする際に以下を聞かれるので入力してデプロイします。

  • 停止忘れチェックの時間(cron式)
  • メールアドレス

その後、作成されたSNSから確認メールが2通（スペース用、エンドポイント用）届くので確認を行います。

あとは日次で以下のようなメールが届きます。

スペースのステータス通知メール

エンドポイントのステータス通知メール

なぜ作ったのか

リソース削除忘れによる不要なコストを防ぐため

SageMakerを使用していると、Studio内で起動したリソースを消し忘れてしまうことがあります。
その中でも私がよく止め忘れてしまうエンドポイントとノートブックインスタンス(本当はスペース)の停止忘れを確認するためにこの仕組みを作りました。

今回この仕組みを実装する際に、すでに誰かが作ってるだろう！と思い調べてみました。
すると、色々出てきました。
エンドポイントについては検索にヒットしたコードを参考にステータスを確認するものがすぐに実装できました。
しかし、ノートブックインスタンスのステータスを確認するのに苦戦しました。

詳細は割愛しますが、ノートブックインスタンスの自動終了や自動停止を行う内容が書かれた記事がたくさんありました。
これらを参考に、Studioで起動したノートブックインスタンスのステータスを確認する関数を作成してみました。
しかし...何を試しても起動したはずのノートブックインスタンスのステータスを確認することができませんでした。

そこでStudioでJupyterLabを触っている時に気づきました。
あれ？これノートブックインスタンスじゃなくて、スペースって書いてある、と。

最近のアップデートでStudioが新しくなり、起動時のプロセスが変わったことでノートブックインスタンスではなくスペースなるものが作成されていることを理解しました。

ちなみに去年の末にStudioが新しくなりました！

新しいStudioと古いStudioでは起動方式が違う

古いStudioでは起動時にノートブックインスタンスが起動していました。
しかし、新しいStudioは起動時にスペースというものが作成されるようです。

ということで、新旧でそれぞれStudioからアプリケーション(ノートブックとスペース)を立ち上げてみます。

SageMakerのコンソール画面を見てみましょう。
古いStudioでは起動後にアプリケーション一覧にJupyterServerが表示されています。

しかし新しいStudioで起動したアプリケーションは一覧に表示されていません。
そこでドメインの詳細から「スペースの管理」を確認してみました。

おお、ここにいるじゃないか！
ということで、新しいStudio内でアプリケーションを起動するとスペースが作成されることがわかりました。

つまり、起動プロセスが変わる＝APIも変わったということです。
当然ですが、今まで通りの方法ではノートブックインスタンスの確認はできてもスペースのステータスは分からないということです。

スペースのステータスを確認するAPIを探してみた

SpaceのAPI

そこで、Spaceの詳細を確認するAPIを探してみました。

以下のコマンドでSpaceの詳細が確認できます。

aws sagemaker describe-space --domain-id <domain-id> --space-name <space_name>

スペースを起動した状態で以下のようなレスポンスが返ってきました。

{
    "DomainId": "a-asdfghjklqwer",
    "SpaceArn": "arn:aws:sagemaker:ap-northeast-1:1234567890:space/a-asdfghjklqwer/demo-space",
    "SpaceName": "demo-space",
    "Status": "InService",
    "LastModifiedTime": "2024-03-24T03:01:36.084000+00:00",
    "CreationTime": "2024-03-24T03:01:19.382000+00:00",
    "SpaceSettings": {
        "JupyterLabAppSettings": {
            "DefaultResourceSpec": {
                "SageMakerImageArn": "arn:aws:sagemaker:ap-northeast-1:0109277902:image/sagemaker-distribution-cpu",
                "SageMakerImageVersionAlias": "1.5.2",
                "InstanceType": "ml.t3.medium"
            }
        },
        "AppType": "JupyterLab",
        "SpaceStorageSettings": {
            "EbsStorageSettings": {
                "EbsVolumeSizeInGb": 5
            }
        }
    },
    "OwnershipSettings": {
        "OwnerUserProfileName": "my-user-profile"
    },
    "SpaceSharingSettings": {
        "SharingType": "Private"
    },
    "Url": "https://fm4ipczf6wokccm.studio.ap-northeast-1.sagemaker.aws/jupyterlab/default"
}

スペースを停止して同じように確認しました。

{
    "DomainId": "a-asdfghjklqwer",
    "SpaceArn": "arn:aws:sagemaker:ap-northeast-1:1234567890:space/a-asdfghjklqwer/demo-space",
    "SpaceName": "demo-space",
    "Status": "InService",
    ...
}

あれ、Status一緒じゃん。
アプリケーションがrun状態でもstop状態でもStatusはInServiceと表示されていました。
これではスペースのステータスを確認することができません。

appのAPI

他にSpaceのステータスを確認できるものがないかと色々探しているとlist-appというAPIを発見しました。

試しにやってみました。

$ aws sagemaker list-apps

まずは起動状態で確認してみます。

{
    "Apps": [
        {
            "DomainId": "a-asdfghjklqwer",
            "SpaceName": "demo-space2",
            "AppType": "JupyterLab",
            "AppName": "default",
            "Status": "InService",
            ...
}

次に停止状態でも確認しました。

{
    "Apps": [
        {
            "DomainId": "a-asdfghjklqwer",
            "SpaceName": "demo-space2",
            "AppType": "JupyterLab",
            "AppName": "default",
            "Status": "Deleted",
            ...
}

おお！スペースの起動状態と停止状態でStatusが変わってるじゃないか！
と、いうことで、spaceではなくappのAPIでステータスの変化が確認できることがわかりました。

では、これをInServiceは起動中、Deletedは停止中と変換することでスペースの停止忘れに気づくことができます。
今回作成したCloudFormationテンプレートはこのようにしてスペースのステータスをチェックしています。

また、新しくなったStudioでも古いStudioを起動することができます。

そのため、念のために古いStudioから起動されたApplicationについても停止忘れを防止するため、ステータスを確認するコードも含めいています。
実際に新旧どちらも起動した状態で受け取ったメールはこんな感じです。（真ん中は関係ないので無視して下さい）

新旧でAppTypeが違うようなので、見分けることができます。
古いStudio→JupyterServer
新しいStudio→JupyterLab

まとめ

新しくなったStudioの情報がまだ少ないので色々と苦戦しました。
改めて記事を書いてナレッジを蓄えることの大切さを再認識しました。
また、私と同じように不要なコストによりダメージを受ける方が一人でも多く救われると嬉しいです。

Share this article

facebook logohatena logotwitter logo

関連記事

機械学習向けのノートブックをそのままパイプラインに組み込めるSageMaker Notebook JobをPython SDKから使ってみた

機械学習向けのノートブックをそのままパイプラインに組み込めるSageMaker Notebook JobをPython SDKから使ってみた

User avatar
nokomoro3
2024.11.26
CloudFormationでSageMaker Studioの自動シャットダウン設定を有効にする

CloudFormationでSageMaker Studioの自動シャットダウン設定を有効にする

User avatar
神野 雄大
2024.11.18
SageMaker Studio スペースにタグを付与する

SageMaker Studio スペースにタグを付与する

User avatar
神野 雄大
2024.11.18
まるクラ勉強会で「Amazon SageMaker MLFlow でスマートに機械学習モデルの実験を管理する方法」について話しました #まるクラ勉強会

まるクラ勉強会で「Amazon SageMaker MLFlow でスマートに機械学習モデルの実験を管理する方法」について話しました #まるクラ勉強会

User avatar
tomozou
2024.11.14
Classmethod's white logo
Facebook's logo
X's logo
YouTube's logo

© Classmethod, Inc. All rights reserved.