Amazon Web Services ブログ

Amazon SageMaker におけるカスタムコンテナ実装パターン詳説 〜学習編〜

みなさんこんにちは、AWS Japanのソリューションアーキテクトの辻です。このブログでは、Amazon SageMaker のカスタムコンテナの実装方法の選択肢が多すぎてどれが自分たちのユースケースに向いているのか分からない！、という意見にお答えして、カスタムコンテナの実装方法のパターンを整理・解説しています。
SageMaker ではカスタムコンテナが利用可能な場面はいくつかあるのですが、このブログでは学習 (SageMaker Training Job) で利用する場合について説明します。推論編もこちらで公開しているので、是非一読ください。

はじめに

このブログは、既に SageMaker の概要を知っている方を対象読者としています。SageMaker Training Job という単語を知っており、SageMaker Python SDK や Boto3 での学習開始の実装方法がイメージできていると躓かずに読み進められると思います。また Docker の概要や Dockerfile の記述方法についても基礎的な知識があるとなお良いです。

Amazon SageMaker は機械学習のあるゆるワークロードを支援するフルマネージドサービスです。SageMaker に数多ある機能の中で汎用性の高いものとしてモデルの学習環境をマネージドで提供する機能があります。このマネージドな学習環境は、コンテナでの処理が前提であり、ユーザーは学習開始時にコンテナイメージを指定します。ほとんどのユースケースではAWS が提供しているコンテナイメージをそのまま利用できますが、ビルドが必要なライブラリをインストールしたいなど、どうしてもカスタマイズが必要になってくる場合があります。その際にカスタムコンテナを定義することになりますが、このカスタムコンテナの実装方法は複数あり、どれを選択すれば良いかが判断が難しいという相談をよく受けていました。そのような相談にお答えする形で、このブログでは学習に絞った上で、カスタムコンテナの実装パターンとその詳細について説明します。

カスタムコンテナの実装パターン

SageMaker Training Job でカスタムコンテナを利用する方法としては、大きく3つのパターンがあります。

  1. AWS が提供しているコンテナイメージを拡張する方法
  2. 独自のコンテナイメージに SageMaker Training Toolkit をインストールする方法
  3. スクラッチでコンテナイメージを作成する方法

パターン 1 が最も実装が簡単で、続いてパターン 2, 3 と難易度が高くなって行きます。カスタムコンテナに関するブログやサンプルコードのほとんどはパターン 2 で実装されており、次点にパターン 1、そしてパターン 3 は稀にしかお目にかかれないかといった印象です。

それぞれのパターンのコンテナイメージで、どこまでのレイヤーが提供され、どこからユーザーが定義するのかを表したのが下記の図です。

SageMaker Custom Containers Pattern

パターン 1: AWS が提供しているコンテナイメージを拡張する方法

このパターンは、AWS が提供しているコンテナイメージをそのまま利用しつつも、独自のライブラリなどを追加したい場合などに選択します。このパターンは最も実装が簡単ですが、一方でカスタマイズの程度によってはそもそもカスタムコンテナが不要となる場合もあるので、まずは本当に必要かを検討した上で選択するのが良いです。例えば pip でインストール可能なライブラリであれば、SageMaker Python SDK の Estimator クラスのコンストラクタの source_dir で指定したディレクトリ内に requirements.txt ファイルを配置して置くことで、学習時にライブラリを利用することが可能なので、カスタムコンテナの利用は不要になります。

このパターンで重要となるのは、ベースとなるコンテナイメージを特定することです。特定しさえすれば、Dockerfile にて FROM でベースイメージを記述して、イメージのビルドとプッシュを行えば OK です。AWS が提供しているコンテナイメージの URI はこちらのドキュメントから辿ることができますが、少々手間がかかるため、オススメは一度目的のフレームワークのバージョンなどを指定して Estimator を作成し、URI を Estimator.training_image_uri() で取得してしまう方法です。この Estimator は URI 取得用としてのみ利用する形です。

from sagemaker.pytorch import PyTorch
estimator = PyTorch(py_version='py38', framework_version='1.10', ...)
estimator.training_image_uri()

あとは Dockerfile を記述し、ビルドとプッシュをし、その URI を Estimator クラスのコンストラクタに image_uri として渡せば学習を開始できます。ベースイメージとして AWS が提供しているものを利用しているので、ローカルの学習コードの指定や学習中のチェックポイントの自動アップロードなど SageMaker との連携機能は通常通り利用可能です。もちろん SageMaker Python SDK を利用せずに Boto3 の create_training_job() などでも開始できます。

サンプルコードとして、Amazon SageMaker Examples の PyTorch extending our containers がパターン 1 を実装したものになっています。

パターン 2: 独自のコンテナイメージに SageMaker Training Toolkit をインストールする方法

このパターンは、独自のベースイメージを利用したい場合や、フレームワークを独自のオプションでビルドしたい場合などに選択します。このパターンを実装する場合、SageMaker Training Toolkit がどのようなツールで、どのような機能を提供しているかを理解すると実装の見通しが非常よくなります。AWS が提供しているコンテナイメージも（一部古いバージョンの例外を除き）、このツールがプリインストールされており、パターン 1 を選択した場合でも実質的にこのツールの機能を利用して SageMaker と連携していることになります。まずはこのツールの概要について説明します。

SageMaker Training Toolkit (以下 Toolkit) は OSS のツールであり、SageMaker Training Job のコンテナ内で利用されます。SageMaker のサービスとの様々な連携機能が実装されており、スクラッチでこれらの実装をすることは実装工数が大幅に増えるため非推奨です。Toolkit をインストールすると、train というコマンドが追加されます（コード参照）。さらに SageMaker Training Job はコンテナを docker run <image> train で起動される（ドキュメント参照）ため、Dockerfile に ENTYRPOINT がなければ、この train コマンドが学習開始時に最初に実行されます。Toolkit 利用時は、逆に train コマンドを entry point として処理を開始したいので、Dockerfile には ENTRYPOINT を記述すべきではありません。
Toolkit はコンテナが SageMaker と連携する心臓部となっており、具体的には以下の機能は、独自に処理を記述しない場合、すべて Toolkit がなければ正しく動作しません。

  • ローカルの学習コードを entry_point として指定できる機能
  • Hyperparamters が学習コードの引数 (sys.argv) として渡される機能
  • 学習の正常終了や異常終了をサービスに通知する機能
  • などなど

逆に Toolkit を利用すれば、これらの機能は AWS 提供のコンテナイメージと同様に動作します。これらの機能については、また別の機会にブログで詳細を紹介する予定ですが、もし興味があれば是非 Toolkit のコードを読んでみてください。

このパターンの一番単純な実装方法は、Dockerfile に RUN pip3 install sagemaker-training を 1 行追加する方法です。学習コードをコンテナ内にコピーすることもできますが（GitHub の README.md はそのような例が記載されていますが）、上に記載した通り、SageMaker Python SDK の Estimator クラスの entry_point や source_dir でローカルやリモートの Git リポジトリ、Amazon S3 にあるコードを、学習開始時に動的に指定することが可能なので、必須ではありません。1 行追加した Dockerfile からイメージをビルドとプッシュをしたら、あとはパターン 1 同様に学習を開始できます。

サンプルコードとして、Amazon SageMaker Examples の custom training containers の basic-training-container 以外がパターン 2 を実装したものになっています（basic-training-container はパターン 3 のサンプルです）。また日本語では PyTorch の BYOC のサンプルがこちらに公開されています。

パターン 3: スクラッチでコンテナイメージを作成する方法

このパターンは、SageMaker Training Toolkit を何かしらの理由で利用できない場合に選択します。パターン 2 でも記載したように Toolkit の機能を再実装するのは非推奨ですが、割り切ってこれらの機能を利用しない場合はパターン 3 を選択しても良いでしょう。

このパターンでは、SageMaker Training Job がどのような仕様で外部とデータをやり取りするのかを正しく理解することが重要です。こちらのドキュメントに仕様が書かれているので、参考にしながら実装します。またクライアントが SageMaker Python SDK を利用している場合は、この SDK が CreateTrainingJob API に追加する各種パラメータを理解しておかないと、想定通りの動作にならないので注意が必要です。

サンプルコードとして、Amazon SageMaker Examples の custom training containers の basic-training-container がまさにパターン 3 を実装したものになっています。

おわりに

カスタムコンテナは、SageMaker のマネージドなインスタンスによって得られる運用負荷の低減と、独自のコンテナによって得られる柔軟性の向上を兼ね備えた、非常に強力な利用方法です。ドキュメントやサンプルコード、ブログなど様々な情報がありましたが、一歩引いた形で概要をまとめたものはあまりない状況でした。このブログはそこの溝を埋められるものと願っています。

著者について

Yohei Tsuji's profile picture辻 陽平 (Yohei Tsuji) は AWS Japan のエンタープライズ ソリューションアーキテクトで、製造領域や研究開発領域のお客様を中心にアーキテクチャ設計や構築をサポートしています。HPC や機械学習の領域を得意としており、特に分散深層学習は大好物です。最近はジム通いで身体を大きくすることコーヒーに夢中です。