Amazon SageMakerノートブックインスタンスの長時間稼働を自動検知する

# はじめに
ご覧いただきありがとうございます。

[前回](https://qiita.com/gahirosan/items/257116510b133c1a2889)の記事で、EC2インスタンスの長時間稼働を検知して、自動通知を送る検証を行いました。

今回SageMakerノートブックインスタンスの長時間稼働を検知して、自動通知を送るようにします。
前回のようにRunCommandを使った検知はできなかったので、CloudWatchのログを利用した検知方法を検証します。

# 概要
(★)がついているセクションは、手を動かして頂く項目です。

1. 今回の構成(★)
2. SageMakerの環境構築(★)
3. SNSトピック/サブスクリプションの準備(★)
4. Lambda実装(連続稼働の検知)(★)
5. EventBridgeで検知を自動化する(★)
6. 挙動確認(★)
　

事前準備
- AWSアカウント作成
- AdministratorAccessを付与したIAMユーザーの作成

# 1.今回の構成
- Amazon SageMaker
- AWS Lambda
- Amazon SNS
- Amazon EventBridge




# 2.SageMakerの環境構築

#### S3バケットの作成

- 任意のバケット名("sagemaker"という文字列を含める)
- 任意のリージョン(今回は東京リージョンを選択)

#### SageMakerノートブックインスタンスの作成

ノートブックインスタンスを作成します。
今回自動検知の検証がメインなので、インスタンスタイプは最低限にします。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2697264/33dae478-5bc9-bfec-2e80-ea804ff130d3.png)

- 任意のノートブックインスタンス名
- インスタンスタイプ: ml.t2.medium
- IAMロール: 新しいロール

IAMロールで「新しいロール」を選ぶと、ロール作成画面が表示されます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2697264/fead229e-bd56-feaf-d6bc-895b0ebbe600.png)


先程作成したS3バケットを指定して、ロールを作成します。
その他はデフォルトでノートブックインスタンスを作成します。

作成したノートブックインスタンスのステータスが「InService」になるのを待ちます。

なお今回作成したIAMロールは以下の通りです。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2697264/18756694-6390-f9dc-1a1f-d9124868dc18.png)

以下はExecutionPolicy
```
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Action": [
                "s3:ListBucket"
            ],
            "Effect": "Allow",
            "Resource": [
                "arn:aws:s3:::xxxxxxxxxxxxxxxx"
            ]
        },
        {
            "Action": [
                "s3:GetObject",
                "s3:PutObject",
                "s3:DeleteObject"
            ],
            "Effect": "Allow",
            "Resource": [
                "arn:aws:s3:::xxxxxxxxxxxxxxxx/*"
            ]
        }
    ]
}
```

合わせてCloudWatchのページを確認します。

ロググループ ⇒ /aws/sagemaker/NotebookInstacesの配下を確認すると、「※ノートブックインタスタンス名/jupyter.log」という名前のログストリームがあります。

今回こちらのログを活用しますので、把握しておいてください。


# 3. SNSトピック/サブスクリプションの準備

Amazon SNSのページに移動します。

まずトピックを作成します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2697264/9130a634-5e4e-171f-f9be-918c862676f2.png)


- タイプ: スタンダード
- 名前: 任意の名前

次にサブスクリプションを作成します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2697264/054bb437-9754-69a1-d6b7-ffc6709b0fec.png)


- プロトコル: Eメール
- エンドポイント: ※通知先Eメールアドレス

サブスクリプションを作成すると、登録したメールアドレスに確認メールが届きます。

メール文中の「Confirm subscription」をクリックします。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2697264/d7e74b21-5f0f-a690-669a-6ff78008f83d.png)


これで通知メールの設定は完了です。
次の工程でLambdaと連携させます。


# 4. Lambda実装(連続稼働の検知)

#### やりたいこと

今回やりたいことは「SageMakerノートブックインスタンスが、使っていないのにも関わらず一定時間以上稼働していたときに、SNS通知を行うこと」ことです。

EventBridgeでLambdaを定期実行します。

LambdaはCloudWatch⇒ロググループ(/aws/sagemaker/NotebookInstaces/)⇒ログストリーム(※ノートブックインタスタンス名/jupyter.log)のログを確認します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2697264/cdaf3fa8-0642-2fde-f825-3de6df0cd7dc.png)

ノートブックインスタンスを「InService」のまま放置した状態です。
特にこれ以降の更新はありません。

もしSageMakerノートブックインスタンスの状態が「InService」であるにも関わらず、ログに更新がなければ、ノートブックインスタンスを停止し忘れていると判断してSNS通知を行います。
具体的には現在時刻とログ最終更新時間の差分をとります。

タスク実行中で長時間稼働している可能性があるので、あくまでSNS通知のみで、自動停止の処理は行いません。

#### Lambda関数の作成

Lambdaのページに移行します。

- 任意の関数名
- ランタイム: Python3.8
- アーキテクチャ: x86_64
- 実行ロール: 既存のロールを使用

今回は「AmazonSageMakerFullAccess」「CloudWatchLogsFullAccess」「AmazonSNSFullAccess」をつけたIAMロールを作成して、付与しました。

関数を作成したら、一般設定で

- タイムアウト: 3秒 ⇒ 1分

に変更します。

#### コードを記述

Parameterの部分は、各自の環境に合わせてください。
index_timeの数値を変えることで、SNS通知を送る判断をする基準となる時間を変更できます。今回は検証作業のため、1時間に設定しています。


```
import boto3
import json
from datetime import datetime, timezone, timedelta
jst = timezone(timedelta(hours=9), 'JST')

# 指標時間
index_time = 1

# 連続稼働の上限(秒)
limit_time = 60 * 60 * index_time

# Parameter
input_topic_arn = "xxxxxxxxxxxxxxxx"
address = "xxxxxxxxxxxxxxxx"
target_region = "ap-northeast-1"

# Client
sns_client = boto3.client('sns', target_region)
sm_client = boto3.client('sagemaker', target_region)
cw_logs_client = boto3.client('logs', target_region)


def lambda_handler(event, context):

    in_operation_longtime = []
    
    for instance in sm_client.list_notebook_instances()['NotebookInstances']:
        i_name = instance['NotebookInstanceName']
        last_modified_time = instance['LastModifiedTime']
            
        if instance['NotebookInstanceStatus'] != 'InService':
            continue
        try:
            logs = cw_logs_client.get_log_events(
                logGroupName="/aws/sagemaker/NotebookInstances",
                logStreamName=f"{i_name}/jupyter.log",
                limit=10,
                startFromHead=False,
                startTime=int(last_modified_time.timestamp()) * 1000,
                endTime=int(datetime.now().timestamp()) * 1000,
                )['events']
                
        except:
                continue

        if len(logs) == 0:
            timestamp = last_modified_time.timestamp()
            
        else:
            target_logs = logs[-1]
            timestamp = int(str(target_logs['timestamp'])[:10])
            
        diff = datetime.now() - datetime.fromtimestamp(timestamp)
        if diff.seconds >= limit_time:
            in_operation_longtime.append({
                "instance_name": i_name,
                "last_modified_time": (last_modified_time.astimezone(tz=jst)).strftime("%Y-%m-%d %H:%I:%S"),
                "seconds": diff.seconds
            })
            
    if len(in_operation_longtime) >= 1:
        topic_arn = input_topic_arn
        title = f"Amazon SageMakerノートブックが{index_time}時間以上起動中です"
        message = f"{index_time}時間以上起動中です"
        
        sns_client.publish(
        TopicArn=topic_arn,
        Message=message,
        Subject=title
    )
    
    return True
```

まず「in_operation_longtime」という空リストを作る。

[list_notebook_instances](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.list_notebook_instances)は、SageMaker notebookインスタンスの一覧を返します。

インスタンスの数だけ以下の処理を行います。

- [ノートブックインスタンスの情報](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.describe_notebook_instance)の中で、「NotebookInstanceName」「LastModifiedTime(ノートブックインスタンスが最後に変更された時刻)」を取得します。

- インスタンスが「InService」であるなら、以降の処理に進む。
- [get_log_events](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/logs.html#CloudWatchLogs.Client.get_log_events)で、指定されたログストリームからのログイベントを一覧表示します。リファレンスに書いてあるパラメータに沿って、「ロググループ」「ログストリーム」「ログを取得する時間範囲」を指定します。

- タイムスタンプを取得します。
- 現在時刻と最終更新時刻の差分を取得。
- 差分がlimit_timeより大きい場合は、「in_operation_longtime」に情報を加える。
- 「in_operation_longtime」に情報が入っていれば、メール通知を送ります。

上記のコードをLambdaに記述します。

Lambdaの実行します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2697264/c941ba16-2976-a381-a65f-18039bf5352a.png)

ノートブックインスタンスは1時間以上稼働しているため、メール通知が実行されました。



# 5.EventBridgeで検知を自動化する

Lambdaを定期実行します。

EventBridgeのページに移動してください。
「ルールを作成」をクリック。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2697264/4a793cf1-2508-deed-088b-0f319c1c36c2.png)

- 任意の名前
- イベントバス: default
- ルールタイプ: スケジュール

次のページに進みます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2697264/e977d2de-c81a-549a-5348-73fd0b4e1eb6.png)


- スケジュールパターン: 通常レート
- レート式: 30分

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2697264/d6ed87a1-98a4-f6a9-8707-514fe5a5d11d.png)

ターゲットに、今回作成したLambda関数を選択します。
あとはデフォルトで設定します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2697264/8f2301bc-192b-3c45-5bbd-a69bfb59e887.png)

Lambdaのトリガーとして、EventBridgeが設定されました。

# 6. 挙動確認

EventBridgeを設定後に、メール通知が自動でくるかを確認します。

今回はもう1台SageMakerノートブックインスタンスをたてました。

#### 検証①

現在のノートブックインスタンスの状態は以下の通りです。

- 1台目(InServiceのまま数時間放置)
- 2台目(InService 立ち上げてすぐの状態)

時間を追って、結果を確認していきます。

EventBridge設定後、すぐにメールが届きました。1台目が1時間以上起動しているからです。
メール確認後、1台目を停止して、再起動します。

#### 検証②
- 1台目(InService 立ち上げてすぐの状態)
- 2台目(InService 立ち上げて30分経過)

メールは届きませんでした。
ここでコードを一部追加してみました。

```
    if len(in_operation_longtime) >= 1:
        topic_arn = input_topic_arn
        title = f"Amazon SageMakerノートブックが{index_time}時間以上起動中です"
        message = f"{index_time}時間以上起動中です"
        
　　　　# コード追加
        for notification in in_operation_longtime:
            message += f"インスタンス名: {notification['instance_name']}\n"
        
        sns_client.publish(
        TopicArn=topic_arn,
        Message=message,
        Subject=title
    )
```

検知されたインスタンスの名前が分かるようにします。

#### 検証③
- 1台目(InService 立ち上げて30分経過)
- 2台目(InService 立ち上げて1時間経過)

メールが届きました。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2697264/61227280-fc0f-cc4f-5111-37f6d2515acc.png)

インスタンス名から2台目が起動しているのが分かります。

#### 検証④
- 1台目(InService 立ち上げて1時間以上経過)
- 2台目(InService 立ち上げて1時間以上経過)

変わらずメールが届きます。
インスタンス名から、2台とも起動しているのが分かります。

メールの確認ができたら、2台とも停止します。

#### 検証⑤
- 1台目(InService 停止)
- 2台目(InService 停止)

メールが届きませんでした。

狙い通りの挙動になっています。

# さいごに

以上、LambdaからCloudWatchログを参照することでノートブックインスタンスの稼働のチェックが可能です。

御覧いただき ありがとうございました！