Jupyter NotebookのコードをSageMakerのJobとして動かす

概要

Jupyter Notebook(.ipynbファイル)をSageMakerのTraining JobやProcessing Jobのスクリプトとして使えるようにするためのライブラリの紹介です。

ライブラリ名: Sagemill (SageMaker+Papermill)

背景

SageMakerを使っていると大規模データに対する前処理や複雑なモデルの学習など、負荷のかかる処理は大体SageMakerのJobとしてハイスペックなサーバー上で動かす一方で、notebook instanceのスペックはそこまで高くないことが多いと思います。

ただnotebookに書いて動作確認したコードを、いちいちpythonファイルに書き写したりSageMakerのJobの環境設定(独自のdocker imageを使う場合はDockerfileや使用するライブラリの定義など)をするのは面倒です。
そこでnotebook上のコードをそのままSageMaker jobとして起動できたら特に初期の実験段階は便利だと思い作りました。

Sagemillの機能

詳細は次のサンプルコードにあるexample notebookを見て頂けたらと思います。

1. notebook(=.ipynbファイル)をpythonファイルに変換

notebookのコードをpythonファイルに変換することで、SageMaker jobで使う独自docker imageに加えたり、sagemaker.tensorflow.Tensorflowのentrypointや sagemaker.processing.ScriptProcessorのcodeとして指定することができます。

parameterize notebook

papermillのように"parameters"タグがついたセルで定義した変数は、jobの実行時に値を渡して上書きことができます。
用途としては、

「notebook上で動作確認するときはepoch数が小さいけどSageMaker Training Job上では100 epoch分学習させたい」

といったときに、notebook上では epochs = 2 にした上で、sagemaker.estimator.Estimatorのインスタンス作成時に hyperpameters={'epochs': 100} を指定することで上の要件が満たされます。

SageMaker関係のコードを、作成するpythonファイルから除外

notebookのコードには sagemaker.estimator.Estimator などSageMakerのjobを定義するためのコードも書くと思いますが、これらはもちろんjobで実行するスクリプト内に含めてはいけません。
そこでSagemillは"sagemaker"タグがついたセルは作成するpythonファイルからは除外するようにしています。

2. SageMaker jobで使う用のdocker imageの作成

これは独自のdocker imageを使うときのみ使用する機能です。

テンプレートとなるDockerfileの作成

指定したdirectoryにDockerfileのテンプレート(+Training Jobの場合はentrypoint.shも)を作成する関数があります。
テンプレートなので、もちろん作成後は自由に書き換えられます。

requirements.txtの作成

作成したpythonファイルに対して pipreqsを実行してrequirements.txtを作成する関数があります。

build image

notebookが動いているサーバーのAWS情報(アカウントIDとリージョン)を基にECRのpush先を決めてimageをbuildする関数があります。

ECRへのpush

• 作成したimageをECRにpush
  • notebookが動いているサーバーのIAM RoleからECRへの認証を自動で行う
  • 指定したrepositoryが存在しない場合は自動で作成

する関数があります。

サンプルコード

Github上ではnotebookのタグが見れないので、どのセルにタグがついてるかは実際に動かして確認してみて下さいmm

予め用意されているdocker imageでモデルの学習を行う例

• サンプルコード
• notebook上のコードをSageMaker Training Jobとして起動
• sagemaker.tensorflow.Tensorflowを使用し、tensorflowのモデルを学習

独自のdocker imageでモデルの学習を行う例

• サンプルコード
• notebook上のコードをSageMaker Training Jobとして起動
• Sagemillの関数で自動作成されたDockerfile, entrypoint.sh, requirements.txt, entrypoint.pyをそのまま使用してdocker imageを作成

予め用意されているdocker imageで前処理を行う例

• サンプルコード
• notebook上のコードをSageMaker Processing Jobとして起動
• sagemaker.sklearn.processing.SKLearnProcessorを使用し、csvファイルを処理

独自のdocker imageで前処理を行う例

• サンプルコード
• notebook上のコードをSageMaker Processing Jobとして起動
• sagemaker.processing.ScriptProcessorを使用し、csvファイルを処理
• Sagemillの関数で自動作成されたDockerfile, requirements.txt, entrypoint.pyをそのまま使用してdocker imageを作成