ローカルで動くのにSageMakerで動かない！？ TensorFlowとSageMakerエンドポイント構築の苦闘記

はじめに

こんにちは。レバレジーズデータ戦略室、データサイエンスグループの徳田です。早速ですが、機械学習エンジニアのみなさん、ローカルで開発したモデルをSageMakerエンドポイントにデプロイしようとすると、予期せぬエラーに遭遇して時間を溶かしてしまうこと、ありますよね？今回は、筆者が実際に経験したSageMakerエンドポイント構築における苦悩と、その解決策についてお話します。

今回のモデル開発におけるアーキテクチャ図

ローカル環境でTensorFlowを使ってまずモデルの試行錯誤および開発を行いました。その後にSageMakerの学習パイプラインを開発し、そこで学習させたモデルを、推論エンドポイントにデプロイする構成で開発を進めていました。ローカルと同じ環境で学習パイプラインを構築し、そこで学習したモデルを推論エンドポイントにデプロイするような構成です。

ローカル環境での開発は順調だったのに…立ちはだかる壁：ローカル環境/学習パイプライン・推論エンドポイントの環境差異

ローカル環境でのモデルの試行錯誤と開発は順調に進み、「よし、いよいよSageMakerで実装だ！」と意気込んでいました。

しかし、現実は甘くありませんでした。ローカル環境では問題なく開発できたモデルが、学習パイプラインで学習し、SageMakerの推論エンドポイントにデプロイすると、エラーを吐き出してしまったのです。原因は、学習パイプラインと推論エンドポイントで使用されていた環境の差異でした。具体的には、以下の2点で環境が異なっていました。

TensorFlowのバージョン: 学習パイプラインで使用していたTensorFlowのバージョンと、推論エンドポイントで利用可能なTensorFlowのバージョンが異なっていました。具体的には、学習パイプラインでは最新バージョンのTensorFlowを使用していましたが、推論エンドポイントではバージョン2.12.1が利用されていました。その結果、学習時には問題なく動作していたモデルが、推論エンドポイントでエラーを引き起こすことになりました。
その他のライブラリのバージョン: 学習パイプラインでインストールしていたライブラリと、推論エンドポイントにインストールされているライブラリのバージョンがTensorFlowの依存関係により異なっていました。これにより、学習パイプラインと推論エンドポイント間で互換性の問題が生じ、予期しない動作やエラーが発生しました。

泥沼にハマる：試行錯誤の解決策

ローカル環境で使用していたライブラリの中には、AWSが用意しているDocker環境では提供されていないものや、バージョンが異なるものがありました。

これらの環境差異を解消するために、私は以下の手順で試行錯誤を重ねました。一見すると複雑で解決が難しいように思えるかもしれません。しかし、落ち着いてエラーメッセージを読み解き、問題の原因を特定することが解決への近道です。

SageMakerのドキュメントでバージョンを確認
- まずはSageMakerのドキュメントで、AWSが用意しているDocker環境で利用可能なTensorFlowとライブラリのバージョンを確認しました。
- 参考：AWSが提供しているコンテナ一覧
ローカル環境の再構築
- 次に、AWSが提供しているDocker環境を基に、ローカルで再度ライブラリの依存関係をインストールし直しました。この際に、Poetryを用いて依存関係のバージョン管理を行い、SageMaker環境に合わせた形で環境を再構築しました。これにより、ローカルとクラウド間での互換性が高まり、エラー発生のリスクを減らすことができました。
Dockerイメージのカスタマイズ
- さらに、学習パイプラインや推論エンドポイントで使用するDockerイメージに、ローカルで再構築したバージョンの依存関係から必要なライブラリを追加しました。AWSが提供するDockerイメージをベースイメージとして、必要なライブラリやツールを組み込むことで、カスタムDockerイメージを作成し、ローカルとの環境の差分がでないようにしました。
- 参考：Amazon SageMaker におけるカスタムコンテナ実装パターン詳説

これらの対応策によって、最終的には推論エンドポイントで正常にモデルを動作させることができました。

教訓：備えあれば憂いなし

今回の経験を通して、SageMakerエンドポイント構築における重要な教訓を得ることができました。それは、「事前に環境の差異を意識し、必要な準備をしておくこと」です。具体的には、以下の点に注意して開発を進めることをおすすめします。

SageMakerのドキュメントをよく読む: SageMakerのドキュメントには、利用可能なTensorFlowやライブラリのバージョン情報が記載されています。こまめにドキュメントを参照しましょう。
バージョンを固定する: requirements.txtなどを活用して、TensorFlowやライブラリのバージョンを固定することで、環境差異による問題を減らすことができます。
Dockerイメージをカスタマイズする: 環境差異が大きい場合は、Dockerイメージをカスタマイズして、必要なライブラリをインストールする必要があります。

これらの点を事前に確認しておくことで、デプロイ時のトラブルを未然に防ぎ、スムーズな開発を進めることができます。