Amazon SageMaker ノートブックインスタンスで最新版の SageMaker Python SDK を使用する

2019 年 8 月より Amazon SageMaker のトレーニングジョブでスポットインスタンスが使用できるようになったため、早速試してみようとしたところ、SageMaker ノートブックインスタンスにプリインストールされる SageMaker Python SDK のバージョンが古く使用できなかった。そこで、本稿ではこの SageMaker Python SDK のアップデート方法について記述する。

関連記事

• 管理型スポットトレーニング: Amazon SageMaker トレーニングジョブで最大 90% を節約 | Amazon Web Services ブログ
• SageMakerのトレーニングジョブでマネージドスポットトレーニング機能がリリースされました！ | DevelopersIO

解決方法

ライフサイクル設定 を作成し、以下のスクリプトをノートブックインスタンスの開始時 (Start notebook) に実行するよう設定する。

#!/bin/bash
set -e

sudo -u ec2-user -i <<'EOF'
source /home/ec2-user/anaconda3/bin/activate mxnet_p36
pip install --upgrade sagemaker
source /home/ec2-user/anaconda3/bin/deactivate
EOF

initctl restart jupyter-server --no-wait

mxnet_p36 の部分はノートブックインスタンスで使用するカーネルに合わせて適宜書き換える。mxnet_p36 と記述した場合は、conda_mxnet_p36 カーネルでのみ最新版の SageMaker Python SDK が使用できるようになる。

なお、以下のように記述すればすべてのカーネルで最新版の SageMaker Python SDK が使用できるようになるが、スクリプトの実行に 5 分以上かかるとノートブックインスタンスの起動に失敗することに注意が必要だ。

#!/bin/bash
set -e

sudo -u ec2-user -i <<'EOF'
for env in base /home/ec2-user/anaconda3/envs/*; do
  source /home/ec2-user/anaconda3/bin/activate $(basename "$env")
  if [ $env = 'JupyterSystemEnv' ]; then continue; fi
  pip install --upgrade sagemaker
  source /home/ec2-user/anaconda3/bin/deactivate
done
EOF

initctl restart jupyter-server --no-wait

実際にスポットトレーニングを実行してみる

Estimator のパラメータに train_use_spot_instances , train_max_wait , checkpoint_s3_uri , checkpoint_local_path を追加する。最新版の SageMaker Python SDK がインストールされていれば正常に動作する。

od_model = sagemaker.estimator.Estimator(
    image_name,
    role,
    train_instance_count=1,
    train_instance_type='ml.p3.2xlarge',
    train_volume_size=8,
    train_max_run=3600,
    input_mode='File',
    output_path='s3://example-bucket/2019-09-15/output',
    sagemaker_session=session,
    train_use_spot_instances=True,
    train_max_wait=3600,
    checkpoint_s3_uri='s3://example-bucket/2019-09-15/checkpoints',
    checkpoint_local_path='/opt/ml/checkpoints')

参照： [Estimators — sagemaker 1.39.2 documentation] (https://sagemaker.readthedocs.io/en/stable/estimators.html)

知見など

• ノートブックインスタンスのターミナルから pip install --upgrade sagemaker を実行しても効果がなかった。
• 上記を実行後ノートブックインスタンスを再起動すると SageMaker Python SDK のバージョンが元に戻っていた。

参考

GitHub の以下のリポジトリに AWS 公式のライフサイクル設定例がいくつか紹介されている。
https://github.com/aws-samples/amazon-sagemaker-notebook-instance-lifecycle-config-samples