Zenn
🏃‍♀️

SageMaker Training Toolkit - ENVIRONMENT_VARIABLES.md 日本語版

2022/04/11に公開
AWS
機械学習
SageMaker
tech

Amazon SageMaker 使っていますか？

SageMaker のメイン機能のひとつである SageMaker Training Job では、バッチ的にサーバーインスタンスを起動し、指定のコンテナ上で任意のスクリプトを実行することができます。

SageMaker Python SDK の Estimator クラスをセットアップして .fit メソッド（もしくは直接 CreateTrainingJob API）を呼んであげれば使えます。

この時、SageMaker Python SDK \leftrightarrow 実行する (主に Python) スクリプト \leftrightarrow コンテナ、の間で連携する情報（どこにデータを置いて、どのハイパーパラメータをセットして、どこにモデルファイルを吐き出すか、など）はコンテナ上で 環境変数 として展開されます。

SageMaker のビルトインコンテナ (TensorFlow / PyTorch / MXNet / HuggingFace / Scikit-Learn など) は SageMaker Training Toolkit というライブラリをベースに構築されてます。

https://aws.amazon.com/jp/blogs/news/sagemaker-custom-containers-pattern-training/

Sagemaker Training Toolkit のドキュメントに、 Estimator から渡される情報が格納される環境変数のリストが掲載されているのですが、日本語の資料は世に出回ってなさそうなので、訳してみます。

https://github.com/aws/sagemaker-training-toolkit/blob/master/ENVIRONMENT_VARIABLES.md

  • SM_MODEL_DIR
  • SM_CHANNELS
  • SM_CHANNEL_{channel_name}
  • SM_HPS
  • SM_HP_{hyperparameter_name}
  • SM_CURRENT_HOST
  • SM_HOSTS
  • SM_NUM_GPUS
  • SM_NUM_CPUS
  • SM_LOG_LEVEL
  • SM_NETWORK_INTERFACE_NAME
  • SM_USER_ARGS
  • SM_INPUT_DIR
  • SM_INPUT_CONFIG_DIR
  • SM_RESOURCE_CONFIG
  • SM_INPUT_DATA_CONFIG
  • SM_TRAINING_ENV

SM_MODEL_DIR

SM_MODEL_DIR=/opt/ml/model

トレーニングジョブが終了すると、コンテナとそのファイルシステムは、/opt/ml/model と /opt/ml/output ディレクトリを除き、削除されます。モデルのチェックポイントを保存するには、/opt/ml/model を使用します。これらのチェックポイントは、デフォルトの Amazon S3 バケットにアップロードされます。

import os

# using it in argparse
parser.add_argument('model_dir', type=str, default=os.environ['SM_MODEL_DIR'])

# using it as variable
model_dir = os.environ['SM_MODEL_DIR']

# saving checkpoints to model dir in chainer
serializers.save_npz(os.path.join(os.environ['SM_MODEL_DIR'], 'model.npz'), model)

詳しくは、Amazon SageMaker がトレーニング出力を処理する方法をご覧ください。

SM_CHANNELS

SM_CHANNELS='["testing","training"]'

コンテナ内の入力データチャンネルのリストが格納されています。

トレーニングを実行する際に、トレーニングデータを異なる論理的な「チャンネル」に分割することができます。問題に応じて、 "training"、"testing"、"evaluation"、"image"、"labels" など、いくつかの一般的なチャンネルのアイデアがあります。

SM_CHANNELS は、コンテナ内で利用可能なチャネルの名前を JSON エンコードしたリストとして格納します。

import os
import json

# using it in argparse
parser.add_argument('channel_names', default=json.loads(os.environ['SM_CHANNELS'])))

# using it as variable
channel_names = json.loads(os.environ['SM_CHANNELS']))

詳しくは、チャンネルをご覧ください。

SM_CHANNEL_{channel_name}

SM_CHANNEL_TRAINING='/opt/ml/input/data/training'
SM_CHANNEL_TESTING='/opt/ml/input/data/testing'

コンテナ内の channel_name という名前のチャンネルが配置されているディレクトリが含まれます。

import os
import json

parser.add_argument('--train', type=str, default=os.environ['SM_CHANNEL_TRAINING'])
parser.add_argument('--test', type=str, default=os.environ['SM_CHANNEL_TESTING'])

args = parser.parse_args()

train_file = np.load(os.path.join(args.train, 'train.npz'))
test_file = np.load(os.path.join(args.test, 'test.npz'))

SM_HPS

SM_HPS='{"batch-size": "256", "learning-rate": "0.0001","communicator": "pure_nccl"}'

ユーザーが提供するハイパーパラメータを JSON エンコードした辞書が含まれます。

import os
import json

hyperparameters = json.loads(os.environ['SM_HPS'])
# {"batch-size": 256, "learning-rate": 0.0001, "communicator": "pure_nccl"}

SM_HP_{hyperparameter_name}

SM_HP_LEARNING-RATE=0.0001
SM_HP_BATCH-SIZE=10000
SM_HP_COMMUNICATOR=pure_nccl

hyperparameter_name という名前のハイパーパラメータの値が格納されます。

learning_rate = float(os.environ['SM_HP_LEARNING-RATE'])
batch_size = int(os.environ['SM_HP_BATCH-SIZE'])
comminicator = os.environ['SM_HP_COMMUNICATOR']

SM_CURRENT_HOST

SM_CURRENT_HOST=algo-1

コンテナネットワーク上の現在のコンテナの名前が格納されます。

import os

# using it in argparse
parser.add_argument('current_host', type=str, default=os.environ['SM_CURRENT_HOST'])

# using it as variable
current_host = os.environ['SM_CURRENT_HOST']

SM_HOSTS

SM_HOSTS='["algo-1","algo-2"]'

すべてのホストを含む JSON エンコードされたリストが格納されます。

import os
import json

# using it in argparse
parser.add_argument('hosts', type=str, default=json.loads(os.environ['SM_HOSTS']))

# using it as variable
hosts = json.loads(os.environ['SM_HOSTS'])

SM_NUM_GPUS

SM_NUM_GPUS=1

現在のコンテナで利用可能な GPU の数が格納されます。

import os

# using it in argparse
parser.add_argument('num_gpus', type=int, default=os.environ['SM_NUM_GPUS'])

# using it as variable
num_gpus = int(os.environ['SM_NUM_GPUS'])

SM_NUM_CPUS

SM_NUM_CPUS=32

現在のコンテナで利用可能な CPU の数が格納されます。

# using it in argparse
parser.add_argument('num_cpus', type=int, default=os.environ['SM_NUM_CPUS'])

# using it as variable
num_cpus = int(os.environ['SM_NUM_CPUS'])

SM_LOG_LEVEL

SM_LOG_LEVEL=20

コンテナ内の現在のログレベルが格納されます。

import os
import logging

logger = logging.getLogger(__name__)

logger.setLevel(int(os.environ.get('SM_LOG_LEVEL', logging.INFO)))

SM_NETWORK_INTERFACE_NAME

SM_NETWORK_INTERFACE_NAME=ethwe

ネットワークインターフェイスの名前が格納されます（分散型トレーニングを行う際に有用）。

# using it in argparse
parser.add_argument('network_interface', type=str, default=os.environ['SM_NETWORK_INTERFACE_NAME'])

# using it as variable
network_interface = os.environ['SM_NETWORK_INTERFACE_NAME']

SM_USER_ARGS

SM_USER_ARGS='["--batch-size","256","--learning_rate","0.0001","--communicator","pure_nccl"]'

トレーニング用に提供されたスクリプトの引数を JSON でエンコードしたリストが格納されます。

SM_INPUT_DIR

SM_INPUT_DIR=/opt/ml/input/

入力ディレクトリのパス（例、/opt/ml/input/）を格納します。入力ディレクトリとは、SageMakerがトレーニング前やトレーニング中に入力データや設定ファイルを保存するディレクトリのことを指します。

SM_INPUT_CONFIG_DIR

SM_INPUT_CONFIG_DIR=/opt/ml/input/config

SageMaker の標準的な設定ファイルが置かれているディレクトリで、例えば /opt/ml/input/config/ などを指します。

SageMaker のトレーニングでは、トレーニング開始時にこのフォルダに以下のファイルが作成されます。

  • hyperparameters.json: Amazon SageMaker は CreateTrainingJob リクエストに含まれるハイパーパラメータをこのファイルに書き込む。
  • inputdataconfig.json: CreateTrainingJob リクエストの InputDataConfig パラメータでデータチャネル情報を指定する。Amazon SageMaker はこの情報をこのファイルに書き込む。
  • resourceconfig.json: 現在のホストとトレーニング内のすべてのホストコンテナの名前が書き込まれる。

これらのファイルの詳細については、Amazon SageMakerがトレーニング情報を提供する方法を参照してください。

SM_RESOURCE_CONFIG

SM_RESOURCE_CONFIG='{"current_host":"algo-1","hosts":["algo-1","algo-2"]}'

opt/ml/input/config/resourceconfig.json から取得した内容が格納されます。以下のようなキーを持っています。

  • current_host: コンテナネットワーク上の現在のコンテナ名。例えば、'algo-1' のようになる。
  • hosts: コンテナネットワーク上のすべてのコンテナの名前のリスト。例えば、3ノードクラスタの場合は ['algo-1', 'algo-2', 'algo-3'] となる。

resourceconfig.json の詳細については、分散学習の設定を参照してください。

SM_INPUT_DATA_CONFIG

SM_INPUT_DATA_CONFIG='{
    "testing": {
        "RecordWrapperType": "None",
        "S3DistributionType": "FullyReplicated",
        "TrainingInputMode": "File"
    },
    "training": {
        "RecordWrapperType": "None",
        "S3DistributionType": "FullyReplicated",
        "TrainingInputMode": "File"
    }
}'

/opt/ml/input/config/inputdataconfig.json から取得した入力データの設定が格納されます。

inpudataconfig.json の詳細については、入力データの設定を参照してください。

SM_TRAINING_ENV

SM_TRAINING_ENV='
{
    "channel_input_dirs": {
        "test": "/opt/ml/input/data/testing",
        "train": "/opt/ml/input/data/training"
    },
    "current_host": "algo-1",
    "framework_module": "sagemaker_chainer_container.training:main",
    "hosts": [
        "algo-1",
        "algo-2"
    ],
    "hyperparameters": {
        "batch-size": 10000,
        "epochs": 1
    },
    "input_config_dir": "/opt/ml/input/config",
    "input_data_config": {
        "test": {
            "RecordWrapperType": "None",
            "S3DistributionType": "FullyReplicated",
            "TrainingInputMode": "File"
        },
        "train": {
            "RecordWrapperType": "None",
            "S3DistributionType": "FullyReplicated",
            "TrainingInputMode": "File"
        }
    },
    "input_dir": "/opt/ml/input",
    "job_name": "preprod-chainer-2018-05-31-06-27-15-511",
    "log_level": 20,
    "model_dir": "/opt/ml/model",
    "module_dir": "s3://sagemaker-{aws-region}-{aws-id}/{training-job-name}/source/sourcedir.tar.gz",
    "module_name": "user_script",
    "network_interface_name": "ethwe",
    "num_cpus": 4,
    "num_gpus": 1,
    "output_data_dir": "/opt/ml/output/data/algo-1",
    "output_dir": "/opt/ml/output",
    "resource_config": {
        "current_host": "algo-1",
        "hosts": [
            "algo-1",
            "algo-2"
        ]
    }
}'

すべてのトレーニング情報を JSON エンコードされた辞書として提供します。

Discussion