オープンソースのAmazon SageMaker Distributionを使い始める

このブログ記事はGet started with the open-source Amazon SageMaker Distributionを翻訳したものです。

データサイエンティストは、機械学習（ML）およびデータサイエンスのワークロードに対して、セキュア且つ一貫性のある、依存関係の管理と再現可能の環境が必要です。AWS Deep Learning Containersは、既にTensorFlow、PyTorch、MXNetなどのモデルトレーニングのフレームワークを事前構築されたDockerイメージで提供しています。利用者の体験を向上させるために、2023年のJupyterConでオープンソースのAmazon SageMaker Distributionのパブリックベータ版を発表しました。これにより、さまざまなレベルの専門知識を持つML開発者に統合されたエンドツーエンドのML体験が提供されます。開発者は、実験のために異なるフレームワークコンテナを切り替えたり、ローカルのJupyterLab環境とSageMakerノートブックからSageMaker上の本番ジョブに移行したりする必要がもはやありません。オープンソースのAmazon SageMaker Distributionは、TensorFlow、PyTorch、Scikit-learn、Pandas、Matplotlibなどのデータサイエンス、ML、可視化のための最も一般的なパッケージとライブラリをサポートしています。本日からAmazon ECR Public Galleryで公開されたコンテナを使用できます。

この記事では、オープンソースのAmazon SageMaker Distributionを使用して、ローカル環境で迅速に実験し、それらを簡単にSageMaker上のジョブに移行する方法を紹介します。

ソリューションの概要

この例では、PyTorchを使用して画像分類モデルのトレーニングを示します。公開されているPyTorchで利用可能なKMNISTデータセットを使用します。ニューラルネットワークモデルをトレーニングし、モデルの性能をテストしてトレーニングとテストの損失を表示します。この例のノートブックは、SageMaker Studio Labのサンプルリポジトリで利用できます。オープンソースディストリビューションを使用してローカルのノートパソコンで実験を開始し、より大きなインスタンスを使用するためにAmazon SageMaker Studioに移動し、最後にノートブックをノートブックジョブとしてスケジュールします。

前提条件

次の前提条件が必要です：

• Dockerがインストールされていること。
• 管理者権限を持つアクティブなAWSアカウント。
• AWS Command Line Interface（AWS CLI）とDockerがインストールされた環境。
• 既存のSageMakerドメイン。ドメインを作成するには、「Amazon SageMakerドメインへのオンボード」を参照してください。

ローカル環境のセットアップ

SageMakerオープンソースディストリビューションを直接ローカルのノートパソコンで使用できます。JupyterLabを起動するには、ターミナルで次のコマンドを実行してください：

export ECR_IMAGE_ID='public.ecr.aws/sagemaker/sagemaker-distribution:latest-cpu'
docker run -it \
    -p 8888:8888 \
    --user `id -u`:`id -g` \
    -v `pwd`/sample-notebooks:/home/sagemaker-user/sample-notebooks \
    $ECR_IMAGE_ID jupyter-lab --no-browser --ip=0.0.0.0

このコマンドはJupyterLabを起動し、ターミナルにhttp://127.0.0.1:8888/lab?token=<token>のようなURLを提供します。リンクをコピーして、お好きなブラウザに入力してJupyterLabを起動してください。

Studioのセットアップ

Studioは、開発者やデータサイエンティストがMLモデルをビルド、トレーニング、デプロイ、監視できる、スケールできるエンドツーエンドの統合開発環境（IDE）です。Studioには、Data Science、TensorFlow、PyTorch、Sparkなどの共通のフレームワークとパッケージを含む、幅広いファーストパーティのイメージが提供されています。これらのイメージを使用することで、データサイエンティストはフレームワークとインスタンスタイプを選ぶだけでMLを簡単に始められます。

SageMakerオープンソースディストリビューションをStudioで使用するために、Studioの「bring your own image」機能を使用できます。SageMakerドメインにオープンソースディストリビューションを追加するには、以下の手順を完了してください：

1. ターミナルで次のコマンドを実行して、オープンソースディストリビューションをアカウントのAmazon Elastic Container Registry（Amazon ECR）リポジトリに追加します：

# Use the latest-cpu or latest-gpu tag based on your requirements
export ECR_GALLERY_IMAGE_ID='sagemaker-distribution:latest-cpu'
export SAGEMAKER_IMAGE_NAME='sagemaker-distribution'
export SAGEMAKER_STUDIO_DOMAIN_ID='d-xxxx'
export SAGEMAKER_STUDIO_IAM_ROLE_ARN='<studio-default-execution-role-arn>'

docker pull public.ecr.aws/sagemaker/$ECR_GALLERY_IMAGE_ID

export ECR_PRIVATE_REPOSITORY_NAME='sm-distribution'
export ECR_IMAGE_TAG='sagemaker-distribution-cpu'
export AWS_ACCOUNT_ID='0123456789'
export AWS_ECR_REPOSITORY_REGION='us-east-1'

# create repository
aws --region ${AWS_ECR_REPOSITORY_REGION} ecr create-repository --repository-name $ECR_PRIVATE_REPOSITORY_NAME
aws --region ${AWS_ECR_REPOSITORY_REGION} ecr get-login-password | docker login --username AWS --password-stdin ${AWS_ACCOUNT_ID}.dkr.ecr.${AWS_ECR_REPOSITORY_REGION}.amazonaws.com
export ECR_IMAGE_URI=$AWS_ACCOUNT_ID.dkr.ecr.$AWS_ECR_REPOSITORY_REGION.amazonaws.com/$ECR_PRIVATE_REPOSITORY_NAME:$ECR_IMAGE_TAG

# Tag
docker tag public.ecr.aws/sagemaker/$ECR_GALLERY_IMAGE_ID $ECR_IMAGE_URI
# Push the image to your private repository
docker push $ECR_IMAGE_URI

2. SageMakerイメージを作成し、そのイメージをStudioドメインにアタッチします。

# Create a SageMaker image
aws sagemaker create-image \
    --image-name $SAGEMAKER_IMAGE_NAME \
    --role-arn $SAGEMAKER_STUDIO_IAM_ROLE_ARN
# Create a SageMaker Image Version.
aws sagemaker create-image-version \
    --image-name $SAGEMAKER_IMAGE_NAME \
    --base-image $ECR_IMAGE_URI

# Optionally, describe the image version to ensure it's succesfully created
aws sagemaker describe-image-version \
    --image-name $SAGEMAKER_IMAGE_NAME \
    --version-number 1
    
# Create the app image configuration file
cat > /tmp/app-config.json << EOF
{
   "AppImageConfigName": "app-image-config-$SAGEMAKER_IMAGE_NAME",
   "KernelGatewayImageConfig": { 
      "FileSystemConfig": { 
         "DefaultGid": 100,
         "DefaultUid": 1000,
         "MountPath": "/home/sagemaker-user"
      },
      "KernelSpecs": [ 
         { 
            "DisplayName": "Python 3 (ipykernel)",
            "Name": "python3"
         }
      ]
   }
}
EOF

# Create an Amazon SageMaker App Image Config.
aws sagemaker create-app-image-config \
    --cli-input-json file:///tmp/app-config.json
    
# Create a default user settings file
# Update the file with your existing settings if you have additional custom images
cat > /tmp/default-user-settings.json << EOF
{
    "DefaultUserSettings": {
        "KernelGatewayAppSettings": {
            "CustomImages": [
                {
                    "ImageName": "$SAGEMAKER_IMAGE_NAME",
                    "AppImageConfigName": "app-image-config-$SAGEMAKER_IMAGE_NAME",
                    "ImageVersionNumber": 1
                }
            ]
        }
    }
}
EOF

# Update Amazon SageMaker Domain with the new default User Settings.
aws sagemaker update-domain \
    --domain-id $SAGEMAKER_STUDIO_DOMAIN_ID \
    --cli-input-json file:///tmp/default-user-settings.json

3. SageMakerコンソールで、ドメインを選択し、既存のユーザープロファイルを選んでStudioを起動します。
4. 必要に応じて、SageMaker Studioを再起動するには、「SageMaker Studioのシャットダウンとアップデート」の手順に従ってください。

ノートブックのダウンロード

GitHubリポジトリからサンプルノートブックをローカルにダウンロードします。

好きなIDEでノートブックを開き、ノートブックの最初にtorchsummaryをインストールするセルを追加します。torchsummaryパッケージはディストリビューションの一部ではないため、このパッケージをノートブックにインストールすることで、ノートブックが終端まで実行されることが保証されます。環境と依存関係を管理するためにcondaまたはmicromambaを使用することをおすすめします。以下のセルをノートブックに追加して、ノートブックを保存してください。

%pip install torchsummary

ローカルのノートブックで実験する

次のスクリーンショットで示すように、アップロードアイコンを選択して、起動したJupyterLab UIにノートブックをアップロードします。

アップロードが完了したら、cv-kmnist.ipynbノートブックを起動します。torch、matplotlib、ipywidgetsなどの依存関係をインストールする必要なく、すぐにセルの実行を開始できます。

前の手順に従っている場合、ローカルのノートパソコンからディストリビューションを使用できることがわかるでしょう。次のステップでは、Studioの機能を活用するために、同じディストリビューションをStudioで使用します。

Studioへの移行（オプション）

これからは実験をStudioに移行しましょう。Studioの利点の1つは、基盤となるコンピュートリソースが完全に弾力的であるため、利用可能なリソースを容易に増減でき、作業が中断されることなく変更は自動的にバックグラウンドで行われます。以前のノートブックを大規模なデータセットとコンピュートインスタンスで実行したい場合、Studioに移行できます。
既に起動したStudio UIに移動し、アップロードアイコンを選択してノートブックをアップロードします。

ノートブックを起動した後、イメージとインスタンスタイプを選択するようにプロンプトが表示されます。カーネルランチャーで、イメージとしてsagemaker-distribution、インスタンスとしてml.t3.mediumを選択し、Selectをクリックしてください。

今、ローカルの開発環境からStudioのノートブックまで、ノートブックに変更を加えることなく、ノートブックを最初から最後まで実行できます！

ノートブックをジョブとしてスケジュールする

実験が完了したら、SageMakerはトレーニングジョブやSageMakerパイプラインなど、ノートブックを本番環境で利用するための複数のオプションを提供しています。その1つは、SageMakerノートブックジョブを使用して、ノートブック自体を非対話型のスケジュールされたノートブックジョブとして直接実行することです。例えば、定期的にモデルを再トレーニングしたり、定期的に入力データに対して推論を行い、ステークホルダー向けにレポートを生成します。

Studioから、ノートブックジョブアイコンを選択してノートブックジョブを起動します。ローカルのノートブックジョブ拡張機能をノートパソコンにインストールしている場合は、ノートパソコンから直接ノートブックをスケジュールすることもできます。ノートブックジョブ拡張機能をローカルにセットアップするには、インストールガイドを参照してください。

ノートブックジョブは自動的にオープンソースディストリビューションのECRイメージURIを使用するため、ノートブックジョブを直接スケジュールすることができます。

Run on scheduleを選択し、スケジュール（例: 毎週土曜日）を選択して、Createを選択します。また、結果をすぐに表示したい場合は、Run nowを選択することもできます。

最初のノートブックジョブが完了したら、Studio UIからOutput filesの下にあるNotebookを選択して、ノートブックの出力を直接表示できます。

追加の考慮事項

公開されているECRイメージをMLワークロードに直接使用することに加えて、オープンソースディストリビューションには以下の利点があります：

• イメージをビルドするために使用されるDockerfileは、開発者が探索して独自のイメージをビルドすることができるように公開されています。このイメージをベースイメージとして継承し、カスタムライブラリをインストールして再現可能な環境を作成することもできます。
• Dockerに慣れていない場合や、JupyterLab環境でConda環境を使用することを好む場合、各バージョンごとにenv.outファイルを提供しています。このファイルの指示に従って、同じ環境を模倣するための独自のConda環境を作成できます。たとえば、CPU環境ファイルcpu.env.outを参照してください。
• GPUバージョンのイメージを使用して、ディープラーニングや画像処理などのGPU互換のワークロードを実行できます。

クリーンアップ

リソースをクリーンアップするには、以下の手順を実行してください：

1. ノートブックをスケジュールで実行するように設定している場合、ジョブの追加料金を防ぐために、Notebook Job Definitionsタブでスケジュールを一時停止または削除してください。



2. すべてのStudioアプリをシャットダウンして、コンピュートの追加料金を防ぎましょう。手順については、「Studioアプリのシャットダウンとアップデート」を参照してください。
3. (オプション) 作成したドメインを削除します。

結論

MLライフサイクルのさまざまなステージで再現可能な環境を維持することは、データサイエンティストや開発者にとって最も大きな課題の1つです。SageMakerオープンソースディストリビューションは、最も一般的なMLフレームワークとパッケージの相互互換性のあるバージョンを含むイメージを提供しています。また、このディストリビューションはオープンソースであり、パッケージやビルドプロセスに対する透明性が提供され、カスタマイズが容易になっています。

この記事では、ディストリビューションの使用方法を、ローカル環境やStudio、トレーニングジョブのコンテナとして紹介しました。この機能は現在パブリックベータ版です。ぜひ試してみて、フィードバックや問題点を公開のGitHubリポジトリで共有してください！

著者について

Durga Sury は、Amazon SageMakerサービスチームのMLソリューションアーキテクトです。彼女は機械学習を誰にでも利用可能にすることに情熱を持って取り組んでいます。AWSでの4年間の経験を通じて、エンタープライズ顧客向けにAI/MLプラットフォームの設定を支援してきました。仕事以外では、バイクのライド、ミステリー小説、そして5歳のハスキーとの長い散歩が大好きです。

Ketan Vijayvargiya は、AWSのシニアソフトウェア開発エンジニアです。彼の専門分野は機械学習、分散システム、およびオープンソースです。仕事以外では、ウェブサービスの運用や自然を楽しむことに時間を費やしています。