SageMaker Studio でローカル モードが利用できるようになりました

こんにちは、森田です。

以下のアップデートで、SageMaker Studio でローカルモードが利用できるようになりました。

また、上記と合わせて、SageMaker Studio でDockerの利用もサポートもされるようになりました。

本記事では、実際に今回のアップデートで追加された SageMaker Studio のローカルモードの使い方をご紹介いたします。

ローカルモードとは

ローカルモードとは、以下図のトレーニングインスタンスを使う前に、事前にStudio環境でモデルのトレーニングなどをコンテナ環境で実行させ、各処理の実行を試すことができるようになる機能です。

引用: Amazon SageMaker Training (座学編)

ローカルモードでは、他のコンピューティングリソースにアクセスすることが不要となるため、デバッグなどが素早くできるようになります。

なお、ローカルモードを使うためには、ドメインに対してEnableDockerAccess : ENABLEDの設定とDockerのインストールが必要となります。

やってみた

前提条件

Amazon SageMaker のドメインは作成済みとします。

ドメインの設定の変更

以下のコマンドを実行して、SageMaker StudioでDockerを利用できるように設定します。

# update domain
aws sagemaker update-domain \
    --domain-id ドメイン名 --region ap-northeast-1 \
    --domain-settings-for-update '{"DockerSettings": {"EnableDockerAccess": "ENABLED"}}'

なお、当初以下の環境で実行したところエラーが発生しましたので、新しいバージョンを利用しましょう。

error version

[cloudshell-user@ip-xxx ~]$ aws --version
aws-cli/2.15.4 Python/3.11.6 Linux/6.1.61-85.141.amzn2023.x86_64 exec-env/CloudShell exe/x86_64.amzn.2023 prompt/off

success version

[cloudshell-user@ip-xxx ~]$ aws --version
aws-cli/2.15.6 Python/3.11.6 Linux/6.1.61-85.141.amzn2023.x86_64 exec-env/CloudShell exe/x86_64.amzn.2023 prompt/off

Dockerのインストール

続いて、スペースに入ってDockerのインストールを行います。

今回は、JupyterLab を利用します。

JupyterLab スペースを開き、ターミナルを起動します。

ターミナルで以下のコマンドを実行し、Dockerのインストールを行います。

なお、ドキュメント上で20.10.Xのバージョンの指定があったため、以下のコマンドでは、20.10.13のインストールを行います。

# Add Docker's official GPG key:
sudo apt-get update
sudo apt-get install ca-certificates curl gnupg
sudo install -m 0755 -d /etc/apt/keyrings
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /etc/apt/keyrings/docker.gpg
sudo chmod a+r /etc/apt/keyrings/docker.gpg

# Add the repository to Apt sources:
echo \
  "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu \
  $(. /etc/os-release && echo "$VERSION_CODENAME") stable" | \
  sudo tee /etc/apt/sources.list.d/docker.list > /dev/null
sudo apt-get update

sudo apt-get install docker-ce=5:20.10.13~3-0~ubuntu-jammy \
docker-ce-cli=5:20.10.13~3-0~ubuntu-jammy \
containerd.io docker-buildx-plugin docker-compose-plugin

ローカルモードでモデルの学習

前準備

以下のリポジトリをターミナルでクローンし、コードを準備します。

git clone https://github.com/cH6noota/SageMakerSample.git

また、以下のコマンドを実行して、必要なライブラリのインストールを行います。

pip install 'sagemaker[local]' --upgrade
pip install -U sagemaker

コードの変更

クローンしたtrain.ipynb内のコードを修正します。

instance_typeをlocalに指定します。

model train

from sagemaker.tensorflow import TensorFlow
cifar_estimator = TensorFlow(
    entry_point="main.py",
    role=role,
    instance_count=1,
    framework_version="2.1.0",
    py_version="py3",
    metric_definitions=metric_definitions,
    enable_sagemaker_metrics=True,
    instance_type='local',
    distribution={"parameter_server": {"enabled": True}},
)

コードの実行

変更後、実行すると以下のようなログが表示され、ローカル環境上で実行されていることが確認できます。

実行結果

INFO:sagemaker.image_uris:image_uri is not presented, retrieving image_uri based on instance_type, framework etc.
INFO:sagemaker:Amazon SageMaker Debugger does not currently support Parameter Server distribution
INFO:sagemaker:Amazon SageMaker Debugger does not currently support Parameter Server distribution
INFO:sagemaker.image_uris:image_uri is not presented, retrieving image_uri based on instance_type, framework etc.
INFO:sagemaker:Creating training-job with name: tensorflow-training-2023-12-30-17-42-26-317
INFO:sagemaker.local.image:'Docker Compose' found using Docker CLI.
INFO:sagemaker.local.local_session:Starting training job
WARNING:sagemaker.local.image:Using the short-lived AWS credentials found in session. They might expire while running.
INFO:sagemaker.local.image:docker compose file: 
services:
  sagemaker-local:
    command: train
    container_name: fuggbi8r5k-sagemaker-local
    environment:
    - '[Masked]'
    - '[Masked]'
    - '[Masked]'
    - '[Masked]'
    - '[Masked]'
    - '[Masked]'
    - '[Masked]'
    image: 763104351884.dkr.ecr.ap-northeast-1.amazonaws.com/tensorflow-training:2.1.0-cpu-py3
    network_mode: sagemaker
    stdin_open: true
    tty: true
    volumes:
    - /home/sagemaker-user/tmp/tmpzq1vbba8/sagemaker-local/input:/opt/ml/input
    - /home/sagemaker-user/tmp/tmpzq1vbba8/sagemaker-local/output/data:/opt/ml/output/data
    - /home/sagemaker-user/tmp/tmpzq1vbba8/sagemaker-local/output:/opt/ml/output
    - /home/sagemaker-user/tmp/tmpzq1vbba8/model:/opt/ml/model
    - /opt/ml/metadata:/opt/ml/metadata
    - /home/sagemaker-user/SageMakerSample/data:/opt/ml/input/data/training
version: '2.3'

さいごに

先日もモデル学習時のコンテナ環境にアクセスできるようになったりと、SageMakerがより便利に使用できる機能が発表されています。

最初のセットアップについても慣れれば、そこまで大変ではありませんので、ぜひ活用してみてください。