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

osawa yuto

2019.01.31

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

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

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

Incoming Webhooks

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

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

SlackのIncoming Webhooksを使い倒す - Qiita

Lambda関数

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

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

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

コンソールで Lambda 関数を作成する - AWS 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に投稿することができます。

実行ロール

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

sagemaker:ListEndpoints
sagemaker:ListNotebookInstance

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

AWS サービスにアクセス許可を委任するロールの作成 - AWS Identity and Access Management

イベントトリガー

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

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

Rate または Cron を使用したスケジュール式 - AWS Lambda

確認

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

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

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

さいごに

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

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