作成した経緯
それは私が配属されてすぐの時でした...
生成AIを利用したアプリ開発するチームで、AWSのSageMakerやBedrockに様々なオープンソースの生成AIモデルをデプロイするというタスクを行っていたのですが、そこで事件は起きます。
上司 「昨日のAWS利用料金が高くて、Cost Explorerを確認したんだけど、SageMakerのインスタンスは止めた？」
私 「あれ？退勤する時にノートブックのインスタンスは止めたのですが」
上司 「エンドポイントは削除した?」
私 「それは削除してないかもです...」
料金を確認したところ、モデルをほぼ動かしていないのに1日で42ドル(日本円約6300円) も請求されていました
幸いにもチームの予算を超えることはなかったのですが、今後このようなことを防止するために、SageMakerの起動状態を監視するSlack Botを作成することになりました。
アーキテクチャ
今回作成したSlack Botの構成図を以下に示します。
LambdaからBoto3のAPIを叩き、SageMaker上で起動しているアプリの状態とエンドポイントを確認しています。EventBridgeのスケジューラでLambdaを実行することで、定期的にSageMakerの状態を監視する仕組みになっています。
SageMakerのレスポンスからLambdaで表示を整え、Slackにはwebhook経由でワークフローに表示しています。
ユーザーはSlackチャンネルでワークフローを動かしていれば、SageMaker上で起動しっぱなしのアプリやエンドポイントをチェックすることができます。
Slackのワークフロー設定
まずはSlackで新しいワークフローを作成します。
ワークフロービルダーが起動したら、右のメニューの"Webhookから"を選択します。
"Webhookを使って開始する"の鉛筆アイコンを選択して、データ変数にデータタイプがテキストの変数を設定して下さい。（キーは任意の名前）
また、編集ページの下部にある"ウェブリクエストのURL"はLambdaの環境変数に登録するため、メモ帳などにコピペしておきます。
Webhookの設定がすんだら、データ変数に追加したキーをチャンネルのメッセージとして表示するステップを追加すればSlackのワークフロー設定は完了です。
Slackワークフローの細かい作成方法などは、下のリンクを参考にすることをおすすめします。
Lambda
SageMakerの起動状態を確認し、Slackにレスポンスを返すLambdaの関数を作成します。
今回作成したLambdaのランタイムはPython3.12で動作を確認しています。
IAMロール
LambdaからSageMakerのAPIを叩くには、それを許可するIAMロールの作成が必要になります。
以下のアクセス許可ポリシーを追加したIAMロールを作成して下さい。
- AWSLambdaBasicExecutionRole
- AmazonSageMakerReadOnly
IAMロールはコンソールでも作成可能です。詳しくは以下のサイトを参照して下さい。
一般設定と環境変数
IAMロールを作成したら、設定タブから一般設定と環境変数を設定します。
一般設定(基本設定)
一般設定の編集ボタンをクリックすると基本設定の編集ページに遷移します。
まず、タイムアウトの設定を30秒にします。
デフォルトの3秒では、LambdaがSageMakerからの応答を待っている間にタイムアウトしてしまいます。
次に実行ロールを設定します。既存のロールを使用するを選択し、先ほど作成したIAMロールを選択します。
環境変数
設定タブの左のメニューから環境変数を選択し、編集ボタンをクリックします。
以下の環境変数を1つ追加します。
キーがSLACK_ENDPOINT_URL
値がSlackのウェブリクエストURL(ワークフロービルダーから確認)
以上でコードを動かす準備が完了しました。
コード
SageMakerのアプリとエンドポイントの状態を取得し、Slackに表示する出力を整えるソースコードを以下に示します。
import json
import boto3
import os
import urllib
from datetime import timezone, timedelta
SLACK_ENDPOINT_URL = os.environ['SLACK_ENDPOINT_URL']
# JSTを定義する
JST = timezone(timedelta(hours=9), 'JST')
def lambda_handler(event, context):
# SageMakerを使うリージョン名をlistの要素に追加する
regions = ["ap-northeast-1", "us-east-1"]
# リージョンごとに実行
apps_for_regions = {}
endpoints_for_regions = {}
for region in regions:
sm = boto3.client('sagemaker', region_name=region)
# 起動中のアプリを取得する
studio_app_list = sm.list_apps()['Apps']
running_apps = [app for app in studio_app_list if app['Status'] == 'InService']
if len(running_apps) > 0:
apps_for_regions[region] = running_apps
# エンドポイントを取得する
endpoints = sm.list_endpoints(StatusEquals = 'InService')['Endpoints']
if len(endpoints) > 0:
endpoints_for_regions[region] = endpoints
# 起動中のアプリやエンドポイントがない場合はこのまま終了する
if(len(apps_for_regions) == 0 and len(endpoints_for_regions) == 0):
return {
'statusCode': 200,
'body': json.dumps('No running apps and models')
}
# 表示を作成する
apps_view = create_view_for_regions(apps_for_regions, create_app_view)
endpoints_view = create_view_for_regions(endpoints_for_regions, create_endpoint_view)
view = '【Studioで起動中のアプリ】\n SageMaker StudioのRunning Instancesから起動しているインスタンスを停止して下さい。\n {}\n【起動中のエンドポイント】\n SageMaker StudioのDeploymentsのEndpointsからモデルを削除してください。\n {}'.format(apps_view, endpoints_view)
# Slackにメッセージを送る
send_slack_message(view, SLACK_ENDPOINT_URL)
return {
'statusCode': 200,
'body': json.dumps('done')
}
def create_view_for_regions(dic, func):
view = ''
for region, data_list in dic.items():
if len(data_list) == 0:
continue
view += '<リージョン名 : '+ region + '>' + '\n'
for data in data_list:
view += func(data) + '\n'
if len(view) == 0:
view = 'なし'
return view
def create_app_view(data):
view = ' ・ドメインID : {}\n ・スペース名 : {}\n ・アプリ名 : {}\n'.format(
data['DomainId'],
data['SpaceName'],
data['AppType']
)
return view
def create_endpoint_view(data):
view = ' ・{}({}起動)\n'.format(
data['EndpointName'],
data['CreationTime'].astimezone(JST).strftime('%Y/%m/%d %H:%M:%S')
)
return view
def send_slack_message(text, slack_endpoint_url):
data = {
'text':text
}
method = "POST"
headers = {"Content-Type" : "application/json"}
req = urllib.request.Request(slack_endpoint_url, method=method, data=json.dumps(data).encode(), headers=headers)
with urllib.request.urlopen(req) as res:
body = res.read()
return body
ポイントはBoto3のlist_apps()でSageMakerのアプリの起動状態をチェックし、list_endpoints()でエンドポイントの状態をチェックできることです。
以前よく使われていたlist_notebook_instances()では、Studioから起動されたノートブックの状態を取得できないので注意が必要です。
以下にBoto3のlist_apps()とlist_endpoints()の公式ドキュメントを添付したので、表示内容を制御したい場合はご参照下さい。
EventBridgeのスケジューラ設定
Lambdaを定期的に実行するためにはEventBridgeのスケジューラから、新しいスケジュールを作成する必要があります。
スケジュールのパターンは定期実行の場合cron式で記述する必要があります。
参考までに、毎日9:00と18:00にLambdaを実行するパターンを以下に示します。
cron ( 0 9,18 * * ? *)
分 時間 日 月 曜 年
詳しく設定したい方は以下のAWSドキュメントをご覧下さい。
ターゲットの選択はAWS Lambda Invokeにチェックし、作成したLambda関数を設定します。
次に、IAMロールからEventBridge SchedulerのLambdaの実行を許可するロールを作成し、アタッチします。
最後に、スケジュールの確認と保存から次のトリガー日が問題なさそうであれば、スケジュールを保存します。
以上でEventBridgeのスケジューラ設定は完了です。
実行例
実際にチームのSlackチャンネルで動かしているワークフローの例です。
SageMakerでアプリを起動していたり、エンドポイントにモデルが残っている時はチャンネルのメッセージ表示され、それ以外の時はメッセージを表示しないワークフローになっています。
おわりに
大学ではほとんどAWSに触る機会が無かったので、今回のSlack Bot作成を通じてLambdaやEventBridgeに初めて触れたのですが、とても便利で使い道の多いサービスだなと思いました。
これからもAWSについてスキルを深めて、日々のアプリケーション開発に活かしていきたいです。