Amazon SageMaker Studioのアーキテクチャ

初めに

これまで、SageMakerノートブックインスタンスを使うことが多かったのですが、「SageMaker Studioを利用すればSageMakerが提供するほぼ全ての機能をGUI上で実現でき、生産性が高まりそう！」と考え、SageMaker Studioを触り始めました。
少し使ってみて確かに便利だなとは思ったのですが、どのような仕組みになっており、何にどのように課金されるのかいまいちよくわからず、アーキテクチャをちゃんと理解して安心安全に使いたいなと思い、アーキテクチャについて調べてみました。
公式の開発者ドキュメントよりも、以下の記事が参考になったので、この内容をもとに整理しようと思います。
https://aws.amazon.com/jp/blogs/machine-learning/dive-deep-into-amazon-sagemaker-studio-notebook-architecture/

※本投稿内の各画像も一部上記記事から拝借しています。

Amazon SageMaker Studio　概要

Amazon SageMaker Studioとは、AWSが提供する、Amazon SageMakerの機能の1つです。
ウェブベースの機械学習用の統合開発環境 (IDE) となっており、機械学習モデルを構築、トレーニング、デバッグ、デプロイ、モニタリングする全てのツールが用意されています。
機械学習開発運用のすべてのフェーズを効率的に実現することが可能です。
概要は、以下に説明があります。

• https://aws.amazon.com/jp/sagemaker/studio/
• https://aws.amazon.com/jp/blogs/news/amazon-sagemaker-studio-the-first-fully-integrated-development-environment-for-machine-learning/

Amazon SageMaker Studio　アーキテクチャ

Single-host Jupyter architecture（SageMakerノートブックインスタンス）

違いを理解するために、まずは、SageMakerノートブックインスタンスが採用している、アーキテクチャの説明です。
SageMakerノートブックインスタンスはJupyterノートブックをマネージドで提供してくれているサービスで、アーキテクチャはJupyterノートブックを普通に設定した場合と同様です。
Jupyterノートブックは1台のマシンにホストされ、どのWebブラウザからでもHTTPS/WSSプロトコルを利用して、アクセスできるようになっています。以下の図は、Amazon Elastic Compute Cloud（Amazon EC2インスタンス）上にセットアップした場合の動作を示しています。

一般的なWebアプリケーションと同じようにWebブラウザのみですぐ利用できるので、このアーキテクチャでも十分便利につかうことができます。
しかし、チームが大きくなったり、機械学習を商用提供する段階になってくると、新たな要件が発生することが多いです。
例えば、以下のような要件が考えられます。

• 他の人の作業に影響を与えないように、自身の作業だけで使うパッケージをインストールしたい。
• 開発フェーズによって、CPUやメモリなどマシンリソースの要求に差異があるため、柔軟にスケールさせたい。
  • データ準備：大量のメモリが必要
  • モデル構築：大量のCPUやGPUが必要
• カーネル（何かの処理）を実行せずに、サンプルノートブックを読んだり、スクリプトやデータファイルを閲覧したりするためだけに使用している時でも、カーネルに対する使用料がかかっており、使わないときは最小限のコストにしたい。
• 機密情報を扱うことも多く、セキュリティを保護するために、チームが使用するすべての環境に定期的にパッチを当て、保護する必要がある。
• gitなどのバージョン管理システムは、ノートブックファイル（.ipynb）をレンダリングする機能はサポートしてらず、ノートブックで行っている作業や成果物を簡単に共同作業して共有する仕組みが欲しい。
• 商用提供段階になると、MLモデルを自動化された方法でデプロイ、監視、再トレーニングする必要があり、さまざまなツールやサービスを切り替えることなく、単一のツールやサービスでシームレスに実現したい。

SageMaker Studio architecture

SageMaker Studioは、機械学習開発に必要なすべてのツールを単一のIDEで提供してくれます。開発者は、コードの記述、実験の追跡、データの視覚化、デバッグと監視をすべて単一のビジュアルインターフェイス内で実行できるため、開発者の生産性が大幅に向上します。

SageMaker Studioの構成要素

SageMaker Studioは大きく3つの要素で構成されています。

1. Domain
   • 各AWSアカウント、1リージョンごとのSageMaker Studioに関する全般的な設定情報。（Amazon Elastic File System (Amazon EFS) ボリューム、UserProfile、セキュリティ設定、App、など）
2. UserProfile
   • ドメインで利用するユーザーの情報。1つのUserProfileはドメイン内の1ユーザを表す。
3. App
   • SageMaker StudioでJupyterノートブックやコンソールなどを実行するための環境（Dockerコンテナー）。以下2種類のAppが用意されている。
     • JupyterServer：ノートブックを表示するUI環境。各ユーザーは、ドメイン内で一意の専用JupyterServerが用意される。
     • KernelGateway：実行中のSageMakerイメージコンテナー。Jupyterノートブックでの各処理や、SageMakerの各機能、ジョブ（TrainigJobやTransformJobなど）を実際に実行する環境。各ユーザーは、ドメイン内で一度に複数のKernelGatewayアプリを実行できる。

UI（JupyterServer）とカーネル実行環境（KernelGateway）を分離するアーキテクチャになっており、カーネルを使いたいときだけ使う仕組みになっています。

図に表すと以下のようなアーキテクチャとなっています。

詳細アーキテクチャ

各アーキテクチャについて詳しく説明します。

StudioUI接続方式

ユーザーがWebブラウザを使用してStudioにアクセスすると、SageMakerによって管理されるEC2インスタンスで実行されているJupyterServerコンテナ内で実行されているノートブックサーバーとのHTTPS/WSS接続が確立されます。これで、StudioのUIにアクセスできます。

Kernelとの通信

Studioは、KernelGatewayアーキテクチャを使用して、ノートブックサーバーがリモートホストで実行されているカーネルと通信できるようにします。ノートブックから実行する各処理は、ノートブックサーバーが存在するホストでは実行されずに、別のホストのDockerコンテナーで実行されます。（上記図だと、KernelGateway App SM Image: datascience-1.0や、KernelGateway App SM Image: mxnet-1.8-gpu-py37に該当）

アプリの割り当て

各ユーザーは、特定のタイプ（ml.t3.mediumなど）のインスタンスを1つだけ実行でき、各インスタンスに最大4つのアプリを割り当てることができます。つまり、ユーザーは、複数のノートブックや処理を同時に実行できます。
同じインスタンスで4つを超えるアプリを実行する必要がある場合は、異なるタイプのインスタンスを実行することで実現できます。
また、アプリ間にリソースの制約はなく、各アプリはすべてのコンピューティングリソースを使用することも可能です。

CPU,GPUの選択

CPUだけでなく、GPUのインスタンスも実行可能です。当然ですが、CPUとGPUでは、ハードウェア要件が異なるため、同じインスタンス上では実行できないという点には注意が必要となります。

スケーリング

複数のインスタンスタイプが用意されているため、必要なコンピューティングリソースを柔軟に用意することができます。Studioでノートブックを開くと、ノートブックが実行されているEC2インスタンス（黄色で強調表示）のvCPUとメモリが表示されます。

黄色で強調表示された領域を選択すると、別のインスタンスタイプに変更することが可能です。

サポートされているインスタンスタイプは以下で確認することができます。
https://aws.amazon.com/jp/sagemaker/pricing/

ストレージ

作業目的に応じて、SageMaker Studioから各種AWSのストレージサービスを活用することが可能です。

Amazon EFS

各ユーザーは、ドメインに関連付けられたAmazon EFSボリューム上に個別のプライベートホームディレクトリを利用できます。ユーザーごとに、一意のID（UID / GID）を自動的に関連付けて、ファイルシステム上のホームディレクトリにのみアクセスできるように制御されます。
Amazon EFSボリュームはすべてのKernelGatewayおよびJupyterServerに自動的にマウントされるので、簡単に実行状況や各種成果物を各コンテナ同士で共有することが可能です。
例えば、以下のようにKernelGateway側で作成した、ディレクトリtest_on_notebookをすぐにUI側（画面左の表示領域）で確認することが可能で、ファイル共有できることで非常に生産性が高まります。

AmazonEFSはもともと様々なクライアントにより、マウントすることができる仕組みになっています。そのため、ファイルシステムをEC2インスタンスにマウントし、ホームディレクトリに対して脆弱性スキャンを実行するといった、セキュリティ要件にも柔軟に対応できます。

Amazon S3

ノートブックのスナップショットとメタデータをS3に保存することで、ノートブックをチームメンバーと共有を可能することができます。
実行方法はとても簡単です。

1. 作業中のノートブック画面の右上の「Share」を押下。
   
2. 出力したい内容を選択（中間成果物や、gitリポジトリ情報も含むかのオプション）し、「Create」を押下。
3. スナップショットと、メタデータが指定したS3バケットに出力。
4. 以下の共有リンク（URL）が表示されるため、コピーして、チームメンバーに共有
   
5. チームメンバーは、URLをブラウザで開くと、共有されたノートブックを開くことができる。
   
6. 作業状況を確認したり、共有されたノートブックを使って自身の作業に活用することができる。

Amazon EBS

各カーネル実行環境（KernelGateway）のインスタンスにAmazonEBSボリュームがアタッチされます。EBSは、コンテナ実行時の一時領域（エフェメラルストレージ）として使用されます。そのため、インスタンスを終了すると、EBSボリュームは削除されます。
最終的には消えても良い中間生成物を管理する領域だと思われるので、利用時はあまり意識する必要がないと思います。

ターミナル

JupyterServer（システムターミナル）とKernelGateway（イメージターミナル）の両方でターミナルセッションを確立することで、各種Linuxコマンドの操作も実施可能です。前者は、ノートブックサーバー拡張機能をインストールしたり、ファイルシステム操作を実行したりできます。後者は、コンテナに特定のライブラリをインストールしたり、コマンドラインからスクリプトを実行したりできます。

余談ですが、前述している通り、各ターミナルの実態はDockerコンテナのため、デフォルトの状態だと使えないコマンドが結構あります。自分でカスタマイズして使いやすくする必要があります。（※以下のようにターミナルでroot直下を見るとコンテナ内にいることを示す.dockerenvがあります。）

ネットワーク

SageMaker Studioは2つのネットワーク環境を選択して、利用することができます。

• SageMakerマネージドVPC（デフォルト）：SageMakerサービスが管理するVPC。インターネット通信が可能となっており、Studioからの各種通信（SageMakerAPIの実行、S3へのアクセスなど）はすべてインターネット経由で実行される。
• CustomerVPC：各利用者が管理するVPC。Studioからの各種通信（SageMakerAPIの実行、S3へのアクセスなど）を利用者が管理するVPC経由で実行できるため、インターネットを介さない通信を実現できる。インターネット経由の通信を許容しない高セキュリティなユースケースではこちらを利用する。

詳細は以下でも解説されています。
https://aws.amazon.com/jp/blogs/machine-learning/securing-amazon-sagemaker-studio-connectivity-using-a-private-vpc/

セキュリティ

基本的には、他のSageMakerサービスと同じセキュリティレベルを提供していますが、SageMaker Studio特有のセキュリティ機能もあるようです。量がかなり多くなりそうなため、別記事でまとめたいと思います。

• https://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/security.html

Amazon SageMaker Studio コスト

利用者側が安心に使うために、気になるポイントとなるコストについてです。
課金要素を簡単にまとめると以下です。

• KernelGatewayのコンピューティングリソース（以下料金ドキュメントの「Amazon SageMaker Studio ノートブック）
  https://aws.amazon.com/jp/sagemaker/pricing/
• AmazonEFSのストレージ

UI（JupyterServer）とカーネル実行環境（KernelGateway）が分離されたアーキテクチャとなっていると説明した通り、UIを開くだけなら無料ということです。何人ユーザを追加しても同様に無料です。
各ユーザは、コストをかけずに、ファイルの参照、ノートブックの読み取り、Studioのシステム端末やその他のUIコンポーネントを使用できます。

実行中のKernelGatewayは、UI上のリソースブラウザ（画面左側領域）に表示され、簡単に終了することができるため、未使用のKernelGatewayは終了して無駄な料金をかけないようにしましょう。

図でいうと、「RUNNING APPS」に実行中のKernelGatewayが表示されます。

また、Sagemaker-Studio-Autoshutdownが提供されており、一定時間操作がないアイドル状態のKernelGatewayのシャットダウンを自動化することもできます。

まとめ

長くなりましたが、アーキテクチャと、アーキテクチャがもたらす各種恩恵について理解することができました。
UI（JupyterServer）とカーネル実行環境（KernelGateway）が分離されていることがアーキテクチャの最もコアな要素となっており、それによって各種機械学習チームの様々な要求を満たしていることが分かります。これだけでも理解すると、見通しが良くなると思います。