SageMaker Training にジョブを投げる君を作った
プライベートで機械学習実験をする際、特に一時的に大量のメモリや GPU が必要な場合、皆さんはどんな環境を使っているだろうか？Google colaboratory や Kaggle のような構築済みのノートブック環境を使う人もいれば、AWS・Google Cloud や RunPod のような何らかのクラウド環境にインスタンスを自前で立てている人もいるだろう。リッチな人は自宅に GPU の乗った強いマシンを持っているかもしれない。
自分は最近そのような用途に AWS の SageMaker Training が便利だなと思っている。
SageMaker Training は機械学習の実験ジョブを AWS 上で実行するサービスである。が、実態としては指定したコンテナ上で任意の Python スクリプトを実行できるサービスと言うのが正しく、起動するインスタンスの種類や動かすコンテナイメージを指定し、実験スクリプトをアップロードすると、SageMaker が自動でインスタンスを立て、そこにスクリプトをダウンロードして実行し、成果物を S3 などの外部ストレージに保存してくれる。終了後はインスタンスを勝手に落としてくれるので、消し忘れで高額請求が来るような心配はない。AWS なので RunPod などに比べると GPU は高いが、一応スポットインスタンスも使えるので安めに済ませることもできる。
コンテナをバッチ的に実行してくれるサービスであれば、AWS Batch や Vertex Pipelines など色々な選択肢がある。その中でも SageMaker Training が使いやすいのは以下の点にあると思う
- コンテナイメージの指定と Python ライブラリの指定を分離できる。実験スクリプトと一緒に
requirements.txtを入れておくと、SageMaker がコンテナを起動した後に自動で必要なライブラリをインストールしてくれる。ライブラリを追加・変更するたびに、docker build して push するというようなことをしなくて良い。 - 同様なことは Vertex AI のカスタムトレーニングでもできるが、以前試した感じだとこちらは使えるイメージに縛りがあったりして思ったほど柔軟ではなかった。SageMaker Training は SDK が入ったイメージであれば基本何でも良いので便利。
- 必要なクラウドリソースは AWS 上の S3 バケットと SageMaker の実行ロールくらいなので準備が簡単。例えば AWS Batch は事前に ECS クラスタ（コンピューティング環境）を作る必要があり、その辺の適切な IAM ロールを準備するのが結構面倒。
欠点はドキュメントがわかりづらく、入出力データのパスの渡し方などの独特の convention に慣れるのに時間がかかること。この点に関しては以下の資料が参考になった。
Amazon SageMaker Training (座学編) - 機械学習のモデル開発の試行錯誤を簡単にする
SageMaker Training を使う際は基本的には Python SDK を経由するのが便利。そこで最近は SDK をラップしたシンプルな CLI を作ってジョブを投げている。
開発のワークフローとしては以下のような感じになる
- SageMaker Training に必要なリソース（S3 バケット・IAM ロール・Docker イメージ）を準備する。
- ローカルで実験スクリプトを開発する。小規模なデータセットを用いてローカルで動作確認をする。依存管理には
uvを使う。 - 使用するライブラリに変更があった場合は、
uv export -o requirements.txt --no-hashesでrequirements.txtを作る。 - 上記CLI
smtを実行する。するとカレントディレクトリがtar.gzが固められて S3 にアップロードされ、 SageMaker Training のジョブがキックされる。
ローカルで開発したスクリプトをそのままリモートのコンテナ環境で実行できるので結構便利な仕組みだと思っている。