Amazon SageMaker でモデル学習時のコンテナ環境にアクセスできるようになりました

こんにちは、森田です。

以下のアップデートで Amazon SageMaker でモデル学習時のコンテナ環境にアクセスできるようになりました。

本記事では、実際に SageMaker のコンテナ環境にアクセスを行ってみたいと思います。

やってみた

モデルの学習部分については、以前の記事をベースに作成します。

学習データの準備

以下のコードを実行して、学習させるデータを作成します。

データの作成

# tensorflow のインストール
! conda install tensorflow -y 
import tensorflow as tf
# データの読み込み
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar100.load_data()
# 格納するディレクトリ作成
!mkdir data
# npy形式で保存
import numpy as np
np.save('data/train_data', x_train)
np.save('data/train_labels', y_train)

np.save('data/eval_data', x_test)
np.save('data/eval_labels', y_test)

import boto3

region = "ap-northeast-1"

# バケット名
bucket_name = "sagemaker-datasets-sample-hogefuga"

s3_client = boto3.client('s3',region_name = region)
location = {'LocationConstraint' : region}
s3_client.create_bucket(Bucket = bucket_name , CreateBucketConfiguration = location)

import glob
for filepath in glob.glob("./data/*.npy"):
    s3_client.upload_file(filepath, bucket_name, filepath.replace("./data/", ""))

IAMロールの作成

SageMaker でモデルを学習させる際に利用するロールを作成します。

通常 Sagemakerの権限に加えてSSM エージェントを利用するため、以下の権限が必要となります。

{
    "Version": "2012-10-17",            
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "ssmmessages:CreateControlChannel",
                "ssmmessages:CreateDataChannel",
                "ssmmessages:OpenControlChannel",
                "ssmmessages:OpenDataChannel"
            ],
            "Resource": "*"    
        }
    ]
 }

モデルの学習

では、実際にモデルの学習を行います。

なお、Amazon DLCsまたは、自身所有のイメージのみが対象となっているようですので、tensorflowを使ってモデルの学習を行います。

モデルの学習

from sagemaker.tensorflow import TensorFlow

role = ""

metric_definitions = [
        {"Name": "loss", "Regex": "loss: ([0-9\\.]+)"},
        {"Name": "accuracy", "Regex": "accuracy: ([0-9\\.]+)"},
        {"Name": "val_loss", "Regex": "val_loss: ([0-9\\.]+)"},
        {"Name": "val_accuracy", "Regex": "val_accuracy: ([0-9\\.]+)"},
    ]

cifar_estimator = TensorFlow(
    entry_point="main.py",
    role=role,
    instance_count=2,
    instance_type="ml.p3.2xlarge",
    framework_version="2.13.0",
    py_version="py310",
    metric_definitions=metric_definitions,
    enable_sagemaker_metrics=True,
    distribution={"parameter_server": {"enabled": True}},
)

bucket_name = "バケット名"
training_data_uri = "s3://{}/".format(bucket_name)

cifar_estimator.fit(
    training_data_uri
)

実行後、ジョブ名が必要となるため、メモしておきます。

コンテナアクセスの設定

以下のコマンドを実行することで、コンテナへアクセスを行うことができるようになります。

RemoteDebugの有効化

aws sagemaker update-training-job  --training-job-name ジョブ名 \
--remote-debug-config EnableRemoteDebug=true

なお、アクセスが可能となるのは、ステータスがTraining（学習時）のみとなります。

ステータス確認

aws sagemaker describe-training-job \
--training-job-name ジョブ名 | jq .SecondaryStatus

コンテナへアクセス

ステータスがTrainingになりましたら、以下のコマンドを実行することでコンテナ環境へアクセス可能です。

aws ssm start-session --target sagemaker-training-job:ジョブ名_algo-インスタンス番号
# 例: aws ssm start-session --target sagemaker-training-job:tensorflow-training-xxx_algo-1

接続後

Starting session with SessionId: xxxxx
# cat /etc/os-release
NAME="Ubuntu"
VERSION="20.04.6 LTS (Focal Fossa)"
ID=ubuntu
ID_LIKE=debian
PRETTY_NAME="Ubuntu 20.04.6 LTS"
VERSION_ID="20.04"
HOME_URL="https://www.ubuntu.com/"
SUPPORT_URL="https://help.ubuntu.com/"
BUG_REPORT_URL="https://bugs.launchpad.net/ubuntu/"
PRIVACY_POLICY_URL="https://www.ubuntu.com/legal/terms-and-policies/privacy-policy"
VERSION_CODENAME=focal
UBUNTU_CODENAME=focal

さいごに

ドキュメントでは、sagemaker.estimator.Estimatorを使ってましたが、本記事では、sagemaker.tensorflowを使い、ジョブ作成後に、AWS CLIからEnableRemoteDebugを有効化しました。

コンテナ環境でアクセスすることができるようになったため、デバッグの際に非常に便利ですので、ぜひご活用ください。