Amazon SageMakerのノートブックインスタンスをアイドル時のみ自動停止させる

WRITER: 黒崎 優太
このエントリーをはてなブックマークに追加

AI事業本部 Dynalystのエンジニアの黒崎( @kuro_m88 )です。

AI事業本部のDynalystではPythonを使ったデータ分析等のノートブック環境としてAmazon Sagemakerのノートブックインスタンスを導入しはじめました。

Amazon SagemakerのノートブックインスタンスはJupyterLabがインストールされた環境がAWS上に簡単に構築できるサービスです。クラウド上のマネージドサービスですので、好きなサイズのインスタンスを必要なときに起動して使うことができます。また、JupyterLabにはブラウザ経由でアクセスしますが、アクセス制御がAWSコンソールの認証情報を使ってくれるので、自前でIP制限や認証プロキシを挟む必要がなく、安全にインターネット経由でアクセスすることができます。

気軽に立てられるのはよいのですが、長時間アイドル状態で放置してしまうとインスタンス料金が無駄になってしまうので、インスタンスの停止忘れ対策を入れて日常的に安心して使えるようにしたいと考え、アイドル状態が続いたらSlackに通知し、それでも対応しなければインスタンスを停止する方法を考えました。

タイトルではAmazon Sagemakerのノートブックインスタンスという書き方をしていますが、実際にはJupyterLabが動いている環境全般に応用可能な手法です。

実現したいこと

起動してから一定時間経過しただけで急にインスタンスが停止されたりすると困るので、

  • ノートブックインスタンス上のすべてのセッションが非アクティブになってから
    • 一定時間後にSlackに非アクティブになったことを通知
    • 一定時間後にSlackにインスタンスを停止することを通知、インスタンスの停止

というように二段階の手順を踏んでインスタンスが自動停止されるようにしたいです。

必要な情報

セッションの状態

JupyterLabはブラウザ経由で利用するソフトウェアなので何かしらAPIがあるだろうということで探していたところ、Sessions APIというものがあり、これが利用できそうです。

Jupyter Notebook Server API · jupyter/jupyter Wiki

実際にノートブックインスタンス上でcurlコマンドで情報を取得してみました。自己署名証明書が使われているようなので-kコマンドで証明書の検証エラーを無視します。

[ec2-user@ip-172-16-80-20 ~]$ curl -s -k https://localhost:8443/api/sessions | jq .
[
  {
    "id": "169de785-f1aa-4f3f-a498-e7d27e94a699",
    "path": "Untitled.ipynb",
    "name": "Untitled.ipynb",
    "type": "notebook",
    "kernel": {
      "id": "49d801c3-9627-471a-97b7-372b3cc87b8d",
      "name": "conda_python3",
      "last_activity": "2020-02-03T01:50:50.446732Z",
      "execution_state": "idle",
      "connections": 1
    },
    "notebook": {
      "path": "Untitled.ipynb",
      "name": "Untitled.ipynb"
    }
  },
  {
    "id": "947d5cbf-0a98-442b-a4c0-4ea37a87a85d",
    "path": "Untitled1.ipynb",
    "name": "Untitled1.ipynb",
    "type": "notebook",
    "kernel": {
      "id": "2f90ff3c-e5f9-44cc-b3ff-b4e301025f13",
      "name": "conda_python3",
      "last_activity": "2020-02-03T01:51:40.538698Z",
      "execution_state": "idle",
      "connections": 1
    },
    "notebook": {
      "path": "Untitled1.ipynb",
      "name": "Untitled1.ipynb"
    }
  }
]

JupyterLab上でコードを実行するとlast_activityという項目が更新されたので、この項目利用できそうなことがわかりました。また、何も実行していない時はexecution_stateがidleになり、何か実行している時はbusyになりました。

ノートブックインスタンスのリソース名

この情報の取得がすこしハマりました。
ノートブックインスタンスの実態はEC2インスタンスなのでEC2のインスタンスメタデータサービス(169.254.169.254)にアクセスしてインスタンスIDを取得すればSageMaker上のノートブックインスタンスの情報と紐付けられるのではないかと思い、インスタンスIDを取得しました。

$ curl http://169.254.169.254/latest/meta-data/instance-id
i-0846b3zf4a0b0c12c

このインスタンスIDは自分のAWSアカウント上には存在しませんでした。

よくわからなかったので他にもいくつかのメタデータを取得して見ていたところ、インスタンスプロファイルの情報を見ていたときにあることに気づきました。(出力は適当な値に置き換えています)

$ curl http://169.254.169.254/latest/meta-data/iam/info
{
  "Code" : "Success",
  "LastUpdated" : "2020-02-03T01:00:09Z",
  "InstanceProfileArn" : "arn:aws:iam::012345678901:instance-profile/BaseNotebookInstanceEc2InstanceRole",
  "InstanceProfileId" : "AIPACADQMTABCSZFMBIW2"
}

InstanceProfileArnの項目にはAWSのアカウントIDが含まれるのですが、このアカウントIDが普段自分が使っているアカウントIDとは異なるものでした。これらのことから、SageMakerのノートブックインスタンスは自分が使っているのとは別のAmazon Sage Maker側で用意されたAWSアカウントで動いているインスタンスなのではないかと推測しました。

このままだとインスタンス上から取得できる値と自分のAWSアカウントのSageMakerのAPIから得られる情報とが紐付かずどうしようかと困っていたところ、ノートブックインスタンスのメタデータというものがある事を知りました。

ノートブックインスタンスのメタデータを取得する

インスタンス起動時に /opt/ml/metadata/resource-metadata.json にJSON形式でリソース名とリソースARNを記述しておいてくれているようです。これを使えばインスタンス内部から自身がどのSageMakerノートブックインスタンスと紐付けられているのかがわかりそうです。

コード

コードは長くなるのでGistにアップロードしました。
やっていることとしては

  • 現在アイドル状態なのかどうかを判定
  • 最後にアクティブだった時間を取得
  • アイドル状態が一定時間経過していればSlackに通知
  • アイドル状態がさらに一定時間経過していればSlackに通知してからAWSのAPI経由でインスタンスをシャットダウン

という流れです。

セットアップ

前述のGistにアップロードされている notebook_idle_monitor.py を自身で管理しているS3にアップロードします。

その後、Amazon SageMakerノートブックのライフサイクル設定のノートブックの作成という欄に以下のコードを記述します。S3のパスは自身のものに書き替えてください。

pip install uptime
aws s3 cp s3://path/to/notebook_idle_monitor.py /home/ec2-user/notebook_idle_monitor.py
(crontab -l 2>/dev/null; echo "1 * * * * /usr/bin/python /home/ec2-user/notebook_idle_monitor.py") | crontab -

ライフサイクル設定

前述の内容でライフサイクルを登録し、アイドル時のみ自動停止されるように管理したいインスタンスのライフサイクル設定に今回作成したものを指定すれば準備は完了です。
注意が必要な点として、ノートブックインスタンス自身がS3からスクリプトを取得し、条件を満たした場合にAWSのAPI経由でインスタンスを停止するのでノートブックインスタンスに割り当てるIAM Roleは以下の操作ができる必要があります。

  • s3:GetObject
  • sagemaker:StopNotebookInstance

動作確認

cronで1分おきにチェックするようにしてみたので、ノートブックインスタンスを起動してから30分放置してみます。アイドル状態であることが通知されました。

アイドル状態のslack通知

さらに30分放置(合計60分)してみます。

インスタンス停止のSlack通知

停止通知が流れ、AWSコンソールからノートブックインスタンスの一覧を確認したところ、無事インスタンスが停止されていることが確認されました。

必要になったときは再度インスタンスを起動させれば続きから作業ができます。

課題

Slack上からインスタンスを起動させたり、停止期限の延長ができるようにしたいなど改善の余地はありそうです。
これに加えてひとつ問題があり、カーネル以外のアクティブセッションは通知できません。
どういうことかというと、Python等のノートブックのセッションはカーネルセッション種類なのですが、ブラウザ上でのターミナル操作のセッションはターミナルセッションという種類のようです。

JupyterLabの画面上では全てのセッションが見えているのですが、APIにて取得する方法がわからず、ターミナル操作のアクティブセッションの検知については諦めています。

JupyterLabのUIでのセッション一覧表示

ターミナルセッションではシェルが立ち上がりっぱなしで、子プロセスの有無などを検知しないとアクティブかどうかはわからなさそうですし、この方法でも監視した瞬間だけ子プロセスがなかった場合に検知漏れが発生してしまいそうなのでうまくはいかなさそうです。
今の所、ノートブックなので停止されて困るものをシェルで長時間使うことはないだろうという想定なのでそこまで問題にはならなさそうですが、もしも良いアイデアをご存知の方がいらっしゃいましたらTwitter等で教えていただけますと嬉しいです。

以上、Amazon SageMakerのノートブックインスタンスをアイドル時のみ自動停止させる方法でした。

KubeflowとJupyter Enterprise GatewayでJupyter Notebook環境をさらに便利に
次の記事
CA × atmaCupを開催しました & 開催準備悲喜こもごも
前の記事