AWS BatchでGPU対応の深層学習・数理最適化バッチ処理を実現した話

はじめに

本記事では、「Lambdaの15分制限では処理が終わらない」「GPUを使った深層学習や数理最適化のバッチ処理を定期実行したい」といった要件に対して、AWS Batchを使ってどのように定期処理基盤を構築したかを紹介します。

AWS BatchとAWS Lambdaの違い（実行時間制限やGPU対応など）や、「AWS Batch GPU 設定」「AWS Batch Lambda 違い」といったキーワードで悩んでいる方の参考になるよう、実際に1時間・30分クラスのジョブを運用した事例ベースで整理しています。

背景

弊社で受託した2件の案件において、AWS Lambdaなどの関数型サービスでは実現が難しい処理を実装する必要がありました。

それぞれ「衛星画像データ」「数理最適化」といった、比較的処理時間のかかるアプリケーションであり、画像解析では深層学習の利用も想定されていたため、GPUの利用も見込んでいました。

とくに、30分〜1時間程度かかるバッチ処理では、Lambdaの最大15分という実行時間制限では対応が難しく、「Lambdaの制限を回避しつつGPUも使える実行環境」が必要でした。

解決したい課題

• 処理時間をあまり気にせずに済むリクエスト・レスポンス型の処理を実現したい
• GPUを使って深層学習の処理を行いたい

• AWS Lambda

• 実行時間を気にせずに利用できるか
• GPUを使った処理に対応可能か（拡張性）

• 実行時間に制限がなかったこと
• マシンサイズ（ジョブ環境）を自由に選択できること

数理最適化を伴ったバッチ処理の実現

AWS Batchを利用することで、数理最適化を伴うバッチ処理を実現できた。2案件分(衛星画像、数理最適化)を実装し、それぞれの実行時間はおおよそ「1時間」「30分」程度でした。

これらはいずれもLambdaの最大実行時間（15分）を大きく超えるため、「AWS Lambdaの制限を回避したい」「GPUを使った深層学習・数理最適化ジョブを実行したい」といったケースでは、AWS Batchが有力な選択肢になりうると感じています。

AWS Batchの実行には「ジョブ」「ジョブ定義」「ジョブ環境」「ジョブキュー」といった4つの設定項目があり、初見ではわかりづらく戸惑いました。

また、Batch処理を実行する際に起動するEC2インスタンスと、Batch処理のオーケストレーターとの通信がうまくいかず、Batchが正常に実行開始するまでに2〜3日かかってしまいました。

上述の現象の原因について述べると、ジョブ環境に紐づくEC2インスタンス側のセキュリティグループやサブネットの設定が不適切で、Batchの管理プレーンからステータスが正しく取得できない状態になっていました。ネットワーク周り（VPC／サブネット／セキュリティグループ）を一つずつ見直し、インバウンド／アウトバウンドルールを整理・修正することで、原因となっていた設定ミスを特定し、無事にジョブが起動するようになりました。

詳細は、以下のZenn記事にまとめています。
https://zenn.dev/team_nishika/articles/34390ef5ad0004

AWS Batchにおける4つの要素の関係イメージ

• ジョブ: 実際に実行される処理の単位（「このコンテナイメージをこのパラメータで実行する」といった指示）
• ジョブ定義: ジョブで使うコンテナイメージやコマンド、リソース要求（vCPU・メモリ・GPUなど）を定義したテンプレート
• ジョブ環境（コンピュート環境）: ジョブを実行するEC2インスタンス群や、そのスケーリングルールを定義する場所
• ジョブキュー: 投げられたジョブをキューイングし、どのジョブ環境にどの順番で割り当てるかを制御するキュー

最初はこれら4つの役割が直感的に掴みにくいのですが、「ジョブ＝実行依頼」「ジョブ定義＝実行レシピ」「ジョブ環境＝実行するマシン群」「ジョブキュー＝待合室」とイメージすると整理しやすくなります。

今回の案件では、VPCやサブネット、セキュリティグループ、ジョブ環境の設定をコンソールベースで行っており、その結果として「どの変更が原因で通信エラーを引き起こしたのか」を追いづらい場面がありました。

後から振り返ると、TerraformなどのIaCツールを用いてクラウド環境の設定をコードベースで管理しておけば、変更履歴をGitで追いやすくなり、今回のようなトラブルシューティングももう少しスムーズにできたのではないかと感じています。

ざっくりとした目安としては、

• 衛星画像の解析や深層学習モデルの推論、数理最適化のように「1回あたりの処理が長い」「GPUや比較的大きなリソースを使いたい」ケースでは、AWS Batchのようなバッチ基盤が向いています。
• 数秒〜数十秒程度で終わる処理や、イベントドリブンで軽量なAPIバックエンドを作りたい場合は、AWS Lambdaで十分なことが多いです。

といった使い分けになります。「Lambdaの時間制限やリソース制限に悩み始めたらAWS Batchを検討する」という基準で考えると、設計の判断がしやすくなると思います。

また、今回の反省としては、「最初からTerraformなどでインフラ構成をコード化しておけばよかった」という点が大きいです。BatchやVPC周りで何度か設定を試行錯誤したので、同様の構成を今後も扱う見込みがある場合は、早めにIaC化しておくと将来的な運用負荷を抑えられると感じました。