Amazon SageMaker Studio notebooks で独自のカスタムコンテナイメージを使用する

本記事は Bringing your own custom container image to Amazon SageMaker Studio notebooks を翻訳したものです。

Amazon SageMaker Studio は、機械学習 (ML) のための統合開発環境 (IDE) です。SageMaker Studio を使用すると、データサイエンティストは Studio ノートブックを起動して、データの探索、モデルの構築、Amazon SageMaker トレーニングジョブの起動、ホストされたエンドポイントにモデルのデプロイを行うことができます。Studio ノートブックには、Amazon SageMaker Python SDK と IPython ランタイムまたはカーネルの最新バージョンで構成された、あらかじめ構築されたイメージのセットが付属しています。この新機能により、Amazon SageMaker ノートブックに独自のカスタム画像を取り込むことができます。これらのイメージは、ドメインで認証されたすべてのユーザーが使用できます。この記事では、カスタムコンテナイメージを SageMaker Studio ノートブックに取り込む方法を共有します。

開発者やデータサイエンティストは、いくつかの異なるユースケースでカスタムイメージを必要とする場合があります。

TensorFlow、MXNet、PyTorch、またはその他の人気の ML フレームワークの特定のバージョンまたは最新バージョンを使用する
ローカルで開発されたカスタムコードまたはアルゴリズムをStudioノートブックに導入して、迅速な反復処理とモデルの学習する
API 経由でデータレイクまたはオンプレミスのデータストアにアクセスし、必要なドライバーをコンテナイメージに含める
R や Julia などの IPython 以外のカーネルとも呼ばれるバックエンドランタイムへのアクセス

大企業では、ML プラットフォーム管理者は、多くの場合、サードパーティ製のパッケージとコードがセキュリティチームによって事前に承認され、インターネットから直接ダウンロードされないようにする必要があります。一般的なワークフローとして、ML Platform チームが使用できるように一連のパッケージとフレームワークを承認し、これらのパッケージを使用してカスタムコンテナを構築し、コンテナに脆弱性がないかテストし、承認されたイメージを Amazon Elastic Container Registry（Amazon ECR）などのプライベートコンテナレジストリにプッシュするというワークフローがあります。この記事でご紹介する手順によって、ML プラットフォームチームは、承認されたイメージを Studio ドメインに直接添付できるようになりました（次のワークフロー図を参照）。Studioでは、承認済みのカスタムイメージを選択するだけで済みます。その後、Studioノートブックでカスタムイメージを使用できます。このリリースでは、1 つの Studio ドメインに最大 30 個のカスタムイメージを含めることができ、必要に応じて新しいバージョンを追加または削除するオプションを使用できます。

ここでは、この機能を使用して SageMaker Studio ノートブックにカスタムコンテナイメージを取り込む方法について説明します。インターネット上でのデフォルトのアプローチを実証しますが、プライベート Amazon Virtual Private Cloud（Amazon VPC）で動作するようにこれを変更する方法について詳しく説明します。

前提条件

開始する前に、次の前提条件を満たしていることを確認する必要があります。

AWS アカウントを持っている
Amazon SageMaker へのアクセスに使用する実行ロールに、以下の AWS Identity and Access Management（IAM）権限が付与されていることを確認します。これにより、SageMaker Studio はプレフィックスの smstudio を使用して Amazon ECR にリポジトリを作成し、このリポジトリからイメージをプッシュおよびプルするアクセス許可を付与します。既存のリポジトリを使用するには、`Resource` をリポジトリの ARN に置き換えます。コンテナーイメージを構築するには、ローカルの Docker クライアントを使用するか、ここで説明する SageMaker Studio から直接イメージを作成します。Amazon ECR でリポジトリを作成するには、SageMaker Studio は AWS CodeBuild を使用します。また、以下に示す CodeBuild アクセス許可を含める必要があります。

{
   "Version":"2012-10-17",
   "Statement":[
      {
         "Effect":"Allow",
         "Action":[
            "codebuild:DeleteProject",
            "codebuild:CreateProject",
            "codebuild:BatchGetBuilds",
            "codebuild:StartBuild"
         ],
         "Resource":"arn:aws:codebuild:::project/sagemaker-studio"
      },
      {
         "Effect":"Allow",
         "Action":"logs:CreateLogStream",
         "Resource":"arn:aws:logs:::log-group:/aws/codebuild/sagemaker-studio"
      },
      {
         "Effect":"Allow",
         "Action":[
            "logs:GetLogEvents",
            "logs:PutLogEvents"
         ],
         "Resource":"arn:aws:logs:::log-group:/aws/codebuild/sagemaker-studio:log-stream:"
      },
      {
         "Effect":"Allow",
         "Action":"logs:CreateLogGroup",
         "Resource":""
      },
      {
         "Effect":"Allow",
         "Action":[
            "ecr:CreateRepository",
            "ecr:BatchGetImage",
            "ecr:CompleteLayerUpload",
            "ecr:DescribeImages",
            "ecr:DescribeRepositories",
            "ecr:UploadLayerPart",
            "ecr:ListImages",
            "ecr:InitiateLayerUpload",
            "ecr:BatchCheckLayerAvailability",
            "ecr:PutImage"
         ],
         "Resource":"arn:aws:ecr:::repository/sagemaker-studio"
      },
      {
         "Effect":"Allow",
         "Action":"ecr:GetAuthorizationToken",
         "Resource":""
      },
      {
         "Effect":"Allow",
         "Action":[
            "s3:GetObject",
            "s3:DeleteObject",
            "s3:PutObject"
         ],
         "Resource":"arn:aws:s3:::sagemaker-/"
      },
      {
         "Effect":"Allow",
         "Action":[
            "s3:CreateBucket"
         ],
         "Resource":"arn:aws:s3:::sagemaker"
      },
      {
         "Effect":"Allow",
         "Action":[
            "iam:GetRole",
            "iam:ListRoles"
         ],
         "Resource":""
      },
      {
         "Effect":"Allow",
         "Action":"iam:PassRole",
         "Resource":"arn:aws:iam:::role/*",
         "Condition":{
            "StringLikeIfExists":{
               "iam:PassedToService":"codebuild.amazonaws.com"
            }
         }
      }
   ]
}

以下に示すように、SageMaker ロールには AWS CodeBuild の信頼ポリシーも必要です。詳細については、「Amazon SageMaker Studio イメージビルド CLI を使用して Studio ノートブックからコンテナイメージを構築する」を参照してください。

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Principal": {
        "Service": [
          "codebuild.amazonaws.com"
        ]
      },
      "Action": "sts:AssumeRole"
    }
  ]
}

AWS Command Line Interface（AWS CLI）をローカルマシンにインストールします。手順については、「AWS をインストールする」を参照してください。
SageMaker スタジオドメインを持っている。ドメインを作成するには、CreateDomain API または create-domain CLI コマンドを使用します。

プライベート VPC を使用してカスタムコンテナを安全に取り込む場合は、次のものも必要です。

プライベートサブネットを持つ VPC
次のサービスの VPC エンドポイント:
- Amazon Simple Storage Service (Amazon S3)
- Amazon SageMaker
- Amazon ECR
- AWS Security Token Service (AWS STS)
- Dockerコンテナを構築するためのCodeBuild

これらのリソースを設定するには、プライベート VPC を使用した Amazon SageMaker Studio 接続の保護と GitHub リポジトリを参照してください。

Dockerfile の作成

最新のフレームワークを試したいというデータサイエンティストの一般的なニーズを満たせることを実証するために、以下の Dockerfile を使用します。このDockerfile では、最新バージョンの TensorFlow 2.3 をベースイメージとして使用しています。この Dockerfile は、任意の Dockerfile に置き換えることができます。現在、SageMaker Studioでは、Ubuntu、Amazon Linux 2など、多数のベースイメージがサポートされています。Dockerfile は、Jupyter ノートブックを実行するのに必要な IPython ランタイムをインストールし、Amazon SageMaker Python SDK と boto3 をインストールします。

ノートブックに加えて、データサイエンティストと ML エンジニアは、Visual Studio Code や PyCharm などのさまざまな人気の IDE を使用して、ローカルのラップトップで実験を行うことがよくあります。スケーラブルな学習やデータ処理のために、これらのスクリプトをクラウドに持ち込むこともできます。これらのスクリプトを Docker コンテナの一部として含めて、SageMaker Studio のローカルストレージに表示することができます。次の Dockerfile では、train.py スクリプトをコピーします。これは、MNIST データセットで単純なディープラーニングモデルをトレーニングするための基本スクリプトです。このスクリプトは、独自のスクリプトまたはコードを含むパッケージに置き換えることができます。

FROM tensorflow/tensorflow:2.3.0
RUN apt-get update 
RUN apt-get install -y git
RUN pip install --upgrade pip
RUN pip install ipykernel && \
    python -m ipykernel install --sys-prefix && \
    pip install --quiet --no-cache-dir \
    'boto3>1.0<2.0' \
    'sagemaker>2.0<3.0'
COPY train.py /root/train.py #Replace with your own custom scripts or packages

以下のコードは train.py スクリプトです。

import tensorflow as tf
import os 

mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10, activation='softmax')
])
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])
model.fit(x_train, y_train, epochs=1)
model.evaluate(x_test, y_test)

カスタムスクリプトの代わりに、AWS Secrets Manager または AWS Systems Manager Parameter Store 経由で認証情報や環境変数にアクセスする Python ファイル、プライベート PyPI リポジトリとの接続を可能にする設定ファイル、またはその他のパッケージ管理ツールなど、他のファイルを含めることもできます。カスタムイメージを使用してスクリプトをコピーできますが、Dockerfile 内の ENTRYPOINT コマンドまたは CMD コマンドは実行されません。

インストールフォルダのセットアップ

ローカルマシン上にフォルダを作成し、そのフォルダに次のファイルを追加する必要があります。

前の手順で作成した Dockerfile
以下の内容で作成された app-image-config-input.json

"AppImageConfigName": "custom-tf2",
    "KernelGatewayImageConfig": {
        "KernelSpecs": [
            {
                "Name": "python3",
                "DisplayName": "Python 3"
            }
        ],
        "FileSystemConfig": {
            "MountPath": "/root/data",
            "DefaultUid": 0,
            "DefaultGid": 0
        }
    }
}

この Dockerfile のバックエンドカーネルを IPython カーネルとして設定し、Amazon Elastic File System (Amazon EFS) へのマウントパスを提供します。Amazon SageMaker は、Jupyter によって定義されたカーネルを認識します。たとえば、R カーネルの場合は、前のコードの Name を ir に設定します。Uid、Gid、およびカーネル名が Docker イメージのカーネル仕様およびユーザー情報と一致することを確認する必要があります。これらの値を抽出するには、このドキュメントを参照してください。

次の内容で default-user-settings.jsonという名前のファイルを作成します。複数のカスタムイメージを追加する場合は、CustomImages のリストに追加するだけです。

{
  "DefaultUserSettings": {
    "KernelGatewayAppSettings": {
      "CustomImages": [
          {
                   "ImageName": "tf2kernel",
                   "AppImageConfigName": "custom-tf2"
                }
            ]
        }
    }
}

イメージの作成と Studio ドメインへのアタッチ

既存のドメインがある場合は、新しいイメージでドメインを更新するだけです。このセクションでは、既存の Studio ユーザーがイメージをアタッチする方法について説明します。新しいユーザーのオンボーディングの手順については、「IAM を使用して Amazon SageMaker Studio にオンボードする」を参照してください。

まず、SageMaker Studio Docker ビルド CLI を使用して、Dockerファイルをビルドして Amazon ECR にプッシュします。ローカルの Docker クライアントや AWS CLI など、他の方法を使用して ECR にコンテナをプッシュできます。

ユーザープロファイルを使用して Studio にログイン
Dockerfile と、コンテナにコピーするその他のコードや依存関係を Studio ドメインにアップロード
Dockerファイルを含むフォルダに移動
ターミナルウィンドウまたはノートブック内で以下のコマンドを実行
```
!pip install sagemaker-studio-image-build
```
IMAGE_NAMEという変数をエクスポートし、default-user-settings.json で指定した値に設定
```
sm-docker build . --repository smstudio-custom:IMAGE_NAME
```
別のリポジトリを使用する場合は、前のコードの smstudio-custom をリポジトリ名に置き換え
SageMaker Studio は Docker イメージをビルドし、適切なイメージ名でタグ付けされた smstudio-custom という名前のリポジトリにある Amazon ECR にイメージをプッシュします。詳細なファイルパスやその他のオプションを指定するなど、これをさらにカスタマイズするには、「Amazon SageMaker Studio イメージビルド CLI を使用して Studio ノートブックからコンテナーイメージを構築する」を参照してください。上記の pip コマンドをプライベート VPC 環境で動作させるには、インターネットへのルート、またはプライベートリポジトリ内のこのパッケージへのアクセスが必要です。

以前のインストールフォルダに、create-and-update-image.sh という新しいファイルを作成

ACCOUNT_ID=AWS ACCT ID # Replace with your AWS account ID
REGION=us-east-2 #Replace with your region
DOMAINID=d-####### #Replace with your SageMaker Studio domain name.
IMAGE_NAME=tf2kernel #Replace with your Image name

# Using with SageMaker Studio
## Create SageMaker Image with the image in ECR (modify image name as required)
ROLE_ARN='The Execution Role ARN for the execution role you want to use'

aws --region ${REGION} sagemaker create-image \
    --image-name ${IMAGE_NAME} \
    --role-arn ${ROLE_ARN}

aws --region ${REGION} sagemaker create-image-version \
    --image-name ${IMAGE_NAME} \
    --base-image "${ACCOUNT_ID}.dkr.ecr.${REGION}.amazonaws.com/smstudio-custom:${IMAGE_NAME}"
    
## Create AppImageConfig for this image (modify AppImageConfigName and KernelSpecs in app-image-config-input.json as needed)
aws --region ${REGION} sagemaker create-app-image-config --cli-input-json file://app-image-config-input.json

## Update the Domain, providing the Image and AppImageConfig
aws --region ${REGION} sagemaker update-domain --domain-id ${DOMAINID} --cli-input-json file://default-user-settings.json

create-image API に渡すことができる引数の詳細については、AWS CLI を参照してください。ステータスを確認するには、Amazon SageMaker コンソールに移動し、ナビゲーションペインから Amazon SageMaker Studio を選択します。

Studio の UI を使ってイメージをアタッチ

GUI を使用して、Studioドメインにイメージを添付する最後の手順を実行することもできます。この場合、GUI ではイメージとイメージバージョンの作成、およびアタッチされたイメージによるドメインの更新を行います。

1. Amazon SageMaker コンソールで、Amazon SageMaker Studio を選択
  コントロールパネルページで、Studioドメインがプロビジョニングされたこと、および作成したユーザープロファイルが表示されます。
2. イメージをアタッチ を選択

1. 新しいイメージまたは既存のイメージのどちらを添付するかを選択
  1. 既存の画像 を選択した場合は、Amazon SageMaker イメージストアから画像を選択します。
  2. 新しいイメージ を選択した場合は、Docker イメージの Amazon ECR レジストリパスを指定します。パスは、Studio ドメインと同じリージョンにある必要があります。ECRリポジトリは、Studioドメインと同じアカウントにあるか、Studioのクロスアカウント権限を有効にする必要があります。
2. 次へを選択

1. イメージ名 に、tensorflowimageを入力
2. イメージ表示名 に、tf2kernelを入力
3. 説明に、ラベル定義を入力
4. IAM ロールについては、Amazon SageMaker がお客様に代わって Amazon ECR イメージを Amazon SageMaker イメージにアタッチするために必要な IAM ロールを選択します。
5. さらに、画像にタグを付けることも可能
6. 次へを選択

EFS mount pathに /root/data と入力
Kernel name に python3 と入力
Kernel display name に tf2kernel と入力
Advanced configuration セクションをクリックして展開
User ID (UID) に 0 と入力
Group ID (GUID) に 0 と入力
Submit をクリック

緑色のチェックボックスは、イメージがドメインに正常にアタッチされたことを示します。

設定が正しいかどうか tensorflowimage リンクをクリックして確認してください。

カーネルが python3 で User ID とGroup ID が None になっています

Amazon SageMaker イメージストアでは、イメージのバージョンが自動的に作成されます。アタッチ済みのイメージを選択し、デタッチ を選択してイメージとすべてのバージョンをアタッチ解除するか、イメージをアタッチ を選択して新しいバージョンをアタッチできます。イメージごとのバージョン数や、イメージをデタッチする機能に制限はありません。

カスタムイメージを使う

それでは、Studio のユーザーエクスペリエンスに飛び込んでみましょう。

1. ユーザープロファイルを使用して Studio にログイン
2. 新しいアクティビティを起動するには、Launcher を選択
3. Select a SageMaker image で、tf2kernel を選択
4. Notebook アイコンを選んでカスタムカーネルを使って新しいノートブックを開く

ノートブックカーネルは起動完了までに数分かかります。これで準備は完了です！

カスタムコンテナをノートブックで試す

カーネルが起動して実行されると、ノートブックでコードを実行できます。まず、Dockerfile で指定された TensorFlow の正しいバージョンが使用可能かどうかをテストしましょう。次のスクリーンショットでは、ノートブックがちょうど起動したtf2kernelを使用していることがわかります。

Amazon SageMaker ノートブックには、ローカル CPU とメモリ使用量も表示されます。

次に、ノートブックでカスタムトレーニングスクリプトを直接試してみましょう。トレーニングスクリプトをノートブックセルにコピーして実行します。このスクリプトは、tf.keras.datasets ユーティリティから mnist データセットをダウンロードし、データをトレーニングセットとテストセットに分割し、カスタムのディープニューラルネットワークアルゴリズムを定義し、トレーニングデータに対してアルゴリズムをトレーニングし、テストデータセットでアルゴリズムをテストします。

TensorFlow 2.3 フレームワークを試すには、Keras の新しい機能の前処理ユーティリティなど、新しくリリースされた API をテストすることができます。次のスクリーンショットでは、データの前処理のための新しいAPIが含まれているTensorFlow 2.3でリリースされた keras.layers.experimentalライブラリをインポートします。これらの API の 1 つをロードし、ノートブックでスクリプトを再実行します。

Amazon SageMaker は、コードの実行時に CPU とメモリ使用量を動的に変更します。カスタムコンテナとトレーニングスクリプトを利用することで、カスタムトレーニングスクリプトとアルゴリズムを Amazon SageMaker ノートブックで直接試すことができます。Studio ノートブックでの実験に満足したら、トレーニングジョブを開始します。

Dockerfile で COPY コマンドを使用してイメージの中に入れた Pythonファイルやそのほかのカスタムファイルはどうなっているでしょうか？ SageMaker Studioは、app-image-config-input.json で設定されたファイルパスに elastic file system (EFS) をマウント `/root/data` とします。Studio がカスタムファイルを上書きしないよう、COPY コマンドは train.py ファイルをパス /root にロードします。ターミナルまたはノートブックを開き、以下のコードを実行することでこのファイルにアクセスできます。

! cat /root/train.py

以下のスクリーンショットのように、出力が表示されます。

train.py ファイルが指定された場所に保存されていることがわかりました。

CloudWatch へのログイン

SageMaker Studio は Amazon CloudWatch にカーネルメトリックスを公開し、トラブルシューティングに使用できます。メトリックスは、名前空間 /aws/sagemaker/studio の下にキャプチャされます。

ログにアクセスするには、CloudWatch コンソールでログを選択します。ロググループ ページで名前空間を入力して、Jupyter サーバーとカーネルゲートウェイに関連付けられたログを表示します。

イメージやバージョンのデタッチ

サポートされなくなった場合は、ドメインからイメージまたはイメージバージョンをデタッチできます。

イメージとすべてのバージョンをデタッチするには、ドメインにアタッチされたカスタムイメージ テーブルからイメージを選択し、デタッチを選択します。

イメージとすべてのバージョンを削除することもできます。これは Amazon ECR のイメージには影響しません。

イメージバージョンをデタッチするには、イメージを選択します。イメージの詳細 ページで、ドメインにアタッチされたイメージバージョン テーブルからイメージのバージョン (または複数のバージョン) を選択し、デタッチを選択します。前述のフローと同様の警告とオプションが表示されます。

まとめ

SageMaker Studio を使用すると、合理化された方法で ML モデルをコラボレーション、実験、学習、デプロイすることができます。そのために、データサイエンティストは、多くの場合、パブリックおよびプライベートのコードリポジトリとパッケージ管理ツールから最新の ML フレームワーク、カスタムスクリプト、およびパッケージへのアクセスを必要とします。本記事の手順によって、関連するすべてのコードを含むカスタムイメージを作成し、Studio ノートブックを使用してこれらを起動できるようになりました。これらのイメージは、Studioドメイン内のすべてのユーザーが使用できます。この機能を使用して、R、Julia、Scala など、Python 以外の他の一般的な言語やランタイムを試すこともできます。サンプルファイルは GitHub リポジトリで利用できます。この機能の詳細については、「独自の SageMaker イメージを取り込む」を参照してください。

About the Authors

Stefan Natu is a Sr. Machine Learning Specialist at AWS. He is focused on helping financial services customers build end-to-end machine learning solutions on AWS. In his spare time, he enjoys reading machine learning blogs, playing the guitar, and exploring the food scene in New York City.

Jaipreet Singh is a Senior Software Engineer on the Amazon SageMaker Studio team. He has been working on Amazon SageMaker since its inception in 2017 and has contributed to various Project Jupyter open-source projects. In his spare time, he enjoys hiking and skiing in the Pacific Northwest.

Huong Nguyen is a Sr. Product Manager at AWS. She is leading the user experience for SageMaker Studio. She has 13 years’ experience creating customer-obsessed and data-driven products for both enterprise and consumer spaces. In her spare time, she enjoys reading, being in nature, and spending time with her family.

Amazon Web Services ブログ