Amazon SageMakerを活用した推論パイプライン運用　ディー・エヌ・エーのエンジニアが語る構成とツール検討の試行錯誤

自己紹介

大谷拓海氏（以下、大谷）：では「SageMakerで試行錯誤する推論パイプライン」という話をします。たぶん界隈だとGCPのほうが流行っているので、AWS使っている人は少ないのかなという気もしますが、AWSサービスの「SageMaker」を使って、推論パイプラインをいい感じに作る試行錯誤をした話をしていきたいと思います。

自己紹介です。moajoというハンドルネームでやっています。新卒2年目で、DeNAのAIシステム部のMLエンジニアリンググループで、ふだんはMLOps的なことをいろいろしています。趣味はクライミングで、これは断崖絶壁を登ったときの写真です。

プロジェクトの過程と課題

今日お話しする内容ですが、何のプロジェクトかは具体的には言えないんですけども、DeNAの“試合“の動画を撮影して、その“試合“の動画をいい感じに解析して、チームにそれをフィードバックすることでチームを強化していきましょう、というプロジェクトになっています。

このプロジェクトはけっこう前からやっているプロジェクトですが、初期段階はPoCというかたちで、かなり少人数で進められていました。そのため、モデルの開発自体は普通にEC2インスタンスを生で立てて、そこの中でローカルのファイルシステムを使いながら、モデルを作って検証して、みたいなことをしていました。

徐々にそれが運用フェーズに乗ってきて、推論パイプラインを作りましょう、運用しましょう、というような話になっていきます。当初の実装では、推論JobがSQS（Amazon Simple Queue Service）にキューイングされるので、そのSQSのキューに対して設定されたオートスケールによってカスタムAMIを指定したEC2インスタンスが起動して、SQSからメッセージを取り出して処理するだけの、非常にシンプルな構成になっていました。

もちろん最初のうちはこれでぜんぜん動いていましたが、これだけシンプルな構成だと、徐々にいろいろと限界が出てきました。どういう限界かというと、“パイプラインのステップの密結合”というのは、推論のパイプライン全体が1つのAMIのエントリーポイントに指定されていて、1個の巨大なPythonスクリプトで実装されているような状態だったので、当然その各推論ステップを個別にバージョン管理したりとか、個別に入れ替えたりがけっこう難しい、バージョン管理が難しかったというのがあります。

さらにけっこうクリティカルな問題で、インスタンス選択もリソースを最適化することができない。というのは、パイプラインが1つのプロセスで動くので、各ステップに応じて適切なインスタンスタイプを割り当てることができないわけです。なので当然、一番重い処理に合わせてインスタンスタイプを選ぶ必要があります。

ステップによっては、GPUを使うものと使わないものがあります。GPUがあるインスタンスを使わなければいけませんが、そうするとGPUを使わない処理をしているときはまったくGPUが無駄になってしまう、ということもありました。

さらに、1つのそのPythonのプロセスで動くので、そのPython環境の中のライブラリのバージョンも全部共通にせざるを得なくて、そのあたりにつらみがありました。

このいろいろつらみのあるパイプラインを、SageMakerを使っていい感じにしていきましょう、ということが今日お話しする内容になります。GCPだとたぶんみなさん「AI platform」などを使っていると思いますが、それのAWSにあるやつだと思ってください。

SageMakerを導入してみた

SageMakerを導入してみました。まずとりあえず作ってみた構成が、こんな感じになっています。SageMakerの中に「Batch Transform」というサービスがあります。名前のとおりこれはバッチで推論Jobを走らせるものですが、これを使って構成しました。

バッチ推論によって、そのパイプラインの各推論ステップを1つのJobとして定義します。そのJobを別途立てたAirflowからキックする感じ。依存関係やリトライなどの処理はAirflow側で管理して、一つひとつの処理をBatch Transformにする感じの構成にしています。

GCPだとAirflow、マネージドでComposerなどがあると思いますが、AWSの場合はないので、自前でECS上にAirflowのインスタンスを立てて使っています。

SageMaker導入後の構成を使ってみた感想

この構成をいろいろ使ってみた感じどうだったか、という感想になります。まず1つは、Airflowの管理がけっこう面倒くさい。まあそれはそうなんですが、わざわざ管理のためだけにAirflowをECSに定義して運用するというのはちょっと面倒くさかったです。このあたりGCPだと、けっこう簡単にできると思うので、GCPいいなあ、という感想です。

あともう1つは細かいところですが、書き出し先のpathが分散していたという話で。書き出し先はS3上に出力していくんですが、model_nameの後にtimestampみたいなpathを切ってそこに書き込んでいましたが、そうすると1回のワークフローで書き出されたオブジェクトがいろいろなところに分散してしまうことになってしまって、それを一覧できないのが辛かったという話です。これは言われてみればそれはそうだ、という話ですが。

一番大きかったのは、Batch Transformへの書き換えが辛い問題があって。このBatch TransformはPythonのスクリプトで書きますが、model_fnやpredict_fnみたいな特殊な名前の関数を定義すると、SageMakerがそれをいい感じに呼び出してくれるというような仕組みになっています。なのでその形式に書き換えなければいけないんですが、今までベタ書きで書いていたスクリプトをその形式に書き換えるのはけっこう面倒くさくて、辛いという問題がありました。

Batch Transform自体も意外とイケてない部分がいっぱいありました。まず1つに「SageMaker Endpoint」という別のサービスがあって、（別の登壇者の発表で）SavedModel Predictionみたいな話がありましたが、それと似たようなサービスで、httpのAPIとしてモデルをサービングしてくれる機能です。実はこのBatch Transformは、（SageMaker）Endpointと同一のコードで利用できるようなサービスになっています。

このBatch Transformは内部的にはEndpointを建てて、そこにhttp経由でひたすらバッチでデータを飛ばして推論するような仕組みになっています。本来必要ないはずのhttpレイヤーが中に存在しているので、そこからエラーが出てきたりするわけですね。これがちょっと面倒くさいところの1つです。

たぶんこのあたりの仕様から察するに、ふだんEndpointを使うけれど、たまにバッチ推論したいときにコードを使い回せるようなユースケースを想定しているんじゃないかと思いました。

今回のケースの場合はそもそもEndpoint建てなくて、完全にバッチ推論しかしないので、実はユースケースとしてはあまりハマってないんじゃないかな、というのが感想です。

Batch Transformの代わりとなるサービスの検討

そこで、Batch Transformの代わりにProcessingという別のサービスを使いました。これは「SageMaker Processing」と言いますが、前処理、後処理みたいなのをするためのサービスです。任意のコードを任意のDockerコンテナ上で実行できるような、非常に汎用性の高い、シンプルなサービスになっています。

このサービスは前処理、後処理はモデルを使わない処理なので、そのモデル管理などがマネージドではなくなりますが、その代わり汎用的な任意のコードが書けます。

あと、書き出し先は変えたら快適になりました。1回のトリガーで走ったパイプラインの書き出し先は1箇所にまとめるといいです。これは当たり前ですね。

Processingを使うことで任意のコードが動かせるようになったので、今まで使っていたコードをほとんど修正せず、そのまま走らせられるようになりました。書き換えとか移行コストがかなり少なくて済んだので、非常によかったなという感じになっています。

これによって既存のパイプラインからの移行はかなり簡単にできていい感じになったんですが、また別の問題としてコスト的な問題がありました。Processingがスポットインスタンスに対応していないので普通にGPUインスタンスを立ち上げないといけないので、けっこうお金がかかってしまうという問題です。

コスト最適化の検討

そのあたりを改善することをさらにやっていきました。次の最適化はいくつか選択肢があると思いますが、まずAWS Batchというものを検討しました。これはECS上でバッチ処理するという、これはもうSageMakerではないんですが、こういうサービスもありました。これはスポットインスタンスも使えるので、今回のユースケースだと普通に使えるものになります。

もう1つはSageMaker Training job、TrainJobです。これは名前のとおり学習Jobを回すやつで、学習 Jobの中で推論してしまう使い方をします。これはどう考えても想定されている使い方じゃないんですが、意外といい感じに動いていて（笑）。スポットインスタンスも使えるし、インスタンスタイプも他のサービスより使えるタイプが多いです。そのため、汎用計算環境としてはかなり使い勝手がいいものになってます。

一番よかったのはTrainJob

いろいろ検討して試しましたが、一番よかったのがTrainJobでした。TrainJobですが、いくつかクセがあって、出力が強制的に圧縮されてしまいます。圧縮ファイルに圧縮されてしまい、その設定が変更できないのがすごい使いにくいところです。そのため、マネージドな出力として出力するのではなく、その推論プロセスの中に自前でアップロード処理を書くみたいなことをしています。ちょっとダーティハックな感じがしてしまいますが、そういうことをしています。スポットインスタンスが使えるので、7割引きとかになって非常にお得感があります。

あと最後、データ入力です。プロセス走らすときにS3のオブジェクトをファイルシステムにマウントできる仕組みがあり、それをmanifest fileという形式のファイルで定義しますが、それに関してはProcessingと同じものが使い回せたので、ProcessingとTrainJobの移行コストはそんなに多くなかったと思います。

最後にまだ残っている辛さがあるので、ここだけ話させてもらうとmanifest fileが、実はローカルモードに対応してないという問題があります。ローカルモードというTrainJobをAWS上ではなくローカルで動かす機能が一応あるんですが、これがけっこうまだ不完全で、いろいろ対応していない機能があるので実はあんまり動かないみたいなことがあって。ローカルデバッグがしにくい問題があります。圧縮のオプションもどうにかしてほしいです。

すべて従量課金でマネージドにできた

ということで、まとめです。SageMaker、いろいろ使ってパイプラインを作ってみました。Batch Transformは実はバッチ推論にはあまり適していない場合があるので、気をつけてください。Processingは汎用的な計算環境としては使い勝手がいいです。

Training jobはある程度制約がつく代わり、スポットインスタンスが使えて実際は便利です。

結果的に、全部従量課金でマネージドにできたので、いい感じに動いています。発表内容は、以上になります。ありがとうございました。