Amazon SageMakerの起動しているノートブックインスタンスとエンドポイントをSlackに通知するBotを作ってみた

こんにちは、大阪DI部の大澤です。

SageMakerのノートブックインスタンスとエンドポイントを起動しっぱなしにしないように、現在起動しているノートブックインスタンスとエンドポイントを毎朝Slackに通知するBot(Lambda関数)を作成しました。今回はその内容を紹介します。

Incoming Webhooks

Slackへのメッセージの投稿はIncoming Webhooksを使うことで実現できます。
SlackのワークスペースにIncoming Webhooksを追加し、メッセージを送信するために使用するエンドポイントURLを取得します。取得したエンドポイントURLに対してメッセージや送り先のチャンネルなどのデータをPOSTすることで、対象のSlackのチャンネルに投稿することができます。

Incoming Webhooksについては以下の記事が詳しいです。

Lambda関数

実際の処理はLambda関数で動かします。今回作るLambda関数の概要は以下の通りです。

  • 毎日9時ごろに自動実行
    • CloudWatch Eventsのスケジュールでトリガー
  • 起動中のノートブックインスタンスとエンドポイントを取得し、Slackの特定のチャンネルに通知
  • SageMakerが使えるリージョンほぼ全てに対応
    • boto3.Session().get_available_regions('sagemaker')で取得できるリージョン
  • ランタイム: Python 3.7

Lambda関数の作成方法の詳細についてはドキュメントをご覧ください。

関数コード

Lambda関数が実行する処理(関数コード)は次の通りです。

import json
import boto3
import os
import urllib
from datetime import timezone, timedelta

# Slack関連の設定を環境変数から読み込む
BOT_USERNAME = os.environ['BOT_USERNAME'] 
SLACK_ENDPOINT_URL = os.environ['SLACK_ENDPOINT_URL'] 
SLACK_CHANNEL = os.environ['SLACK_CHANNEL'] 

# JSTを定義する
JST = timezone(timedelta(hours=9), 'JST')

def lambda_handler(event, context):
    # SageMakerが対応してるリージョン一覧を取得
    regions = boto3.Session().get_available_regions('sagemaker')

    # リージョンごとに実行
    notebooks_for_regions = {}
    endpoints_for_regions = {}

    for region in regions:
        sm = boto3.client('sagemaker', region_name=region)


        # 起動中(InService)のインスタンスを取得する
        notebooks = sm.list_notebook_instances(
            StatusEquals='InService'
        )['NotebookInstances']
        if len(notebooks) > 0:
            notebooks_for_regions[region] = notebooks
        
        
        # 起動中のエンドポイントを取得する
        endpoints = sm.list_endpoints(
            StatusEquals = 'InService'
        )['Endpoints']
        if len(endpoints) > 0:
            endpoints_for_regions[region] = endpoints
        
    
    # ノートブックインスタンスとエンドポイントのそれぞれの表示内容を作成する
    notebooks_view = create_view_for_regions(notebooks_for_regions, create_notebook_view)
    endpoints_view = create_view_for_regions(endpoints_for_regions, create_endpoint_view)
    
    # 表示を作る
    view = '起動中のノートブックインスタンス\n{}\n起動中のエンドポイント\n{}'.format(notebooks_view, endpoints_view)
    
    # Slackにメッセージを送る
    send_slack_message(view, BOT_USERNAME, SLACK_CHANNEL, 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_notebook_view(data):
    # ノートブックの表示を作る
    
    view = ' ・{}({}起動)'.format(
        data['NotebookInstanceName'],
        data['LastModifiedTime'].astimezone(JST).strftime('%Y/%m/%d %H:%M:%S')
    )
    return view

def create_endpoint_view(data):
    # エンドポイントの表示を作る
    
    view = ' ・{}({}起動)'.format(
        data['EndpointName'],
        data['CreationTime'].astimezone(JST).strftime('%Y/%m/%d %H:%M:%S')
    )
    return view

def send_slack_message(text, username, channel, slack_endpoint_url):
    # Slackにメッセージを送信する
    
    data = {
        'username':username, # 表示名
        'text':text, # 内容
        'channel': SLACK_CHANNEL # 送信先チャンネル
    }
    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

環境変数

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

  • BOT_USERNAME : Slack投稿時のBotの表示名です。 例)SageMakerインスタンス放置警察bot
  • SLACK_CHANNEL : Slackのチャンネル名。エンドポイントとインスタンス情報を通知したいチャンネルです。
  • SLACK_ENDPOINT_URL : Incoming Webhooks設定時に取得したエンドポイントURLです。このURLにデータをPOSTすることでSlackに投稿することができます。

実行ロール

管理ポリシーのAWSLambdaBasicExecutionRoleAmazonSageMakerReadOnlyをアタッチしたIAMロールを作成し、実行ロールに設定します。 最低限の権限だけであればAWSLambdaBasicExecutionRoleをベースにして、今回の処理で使用する権限をもつインラインポリシーを追加します。以下のActionに対する権限を有効にする必要があります。

  • sagemaker:ListEndpoints
  • sagemaker:ListNotebookInstance

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

イベントトリガー

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

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

確認

SageMakerでus-east-1とap-northeast-1でノートブックインスタンスやエンドポイントを作成します。

作成したインスタンスやエンドポイントの起動が完了している状態で、先ほど作成したLambda関数を動かします。空のテストイベントを作成しテストを実行します。

作成したLambda関数が無事実行され、Slackに通知されました。

さいごに

今回はSageMakerの起動中のノートブックインスタンスとエンドポイントをSlackのチャンネルへお知らせするBotを紹介しました。エンドポイントやノートブックインスタンスが立ち上げっぱなしになるのは、1日とかだったらまだ安いですが1週間1ヶ月となるとそれなりにコストが膨れ上がります。変なところでお金を浪費するのはよくないですし、これを機に試してみてはいかがでしょうか。

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