Bedrockのバッチ推論は段階的に出力されます

こんにちは！
株式会社リブセンス アルバイト事業部 開発グループのこりんです！

弊社ではアドベントカレンダーを実施しており、8日目は私が担当します！
今回は、少しクセのあるBedrockのバッチ推論について書こうと思います。

ちなみに、学生時代からぼちぼちと書き続けているQiitaですが、25新卒としてリブセンスに入社し、所属を明らかにして投稿するのは初めてです。
ちょっとドキドキしますが、よろしくお願いします。

TL; DR

1. Bedrockのバッチ推論は、推論結果がS3に段階的に出力されることがある
2. 全ての推論が完了しているかはmanifest.json.outで確認可能

背景

まずは簡単にことの経緯を。

弊社リブセンスは「あたりまえを、発明しよう。」をビジョンに掲げ、社会の不均衡を解消し、新しい「あたりまえ」を生み出す複数のサービスを運営しています。
アルバイト事業部の運営する「マッハバイト」もその一つで、「求人を探す」という体験を刷新すべく様々な挑戦をしています。

私はその中で、主に社内オペレーションの改善を担当しています。
求人掲載に関わる様々な業務の効率化を目指しているのですが、直近の急速なAIの成長を受け、AIを活用した改善に挑戦することとなりました。

この記事では、その検証段階で詰まったところを紹介したいと思います！

システムの概要

AIを使った業務改善は初めてだったため、まずはPoCとして以下のようなシステムを構築しました。

• 推論件数が多いため、コストとクォータの観点からバッチ推論を利用
• 推論の前後処理やBedrock APIの呼び出しにはLambdaを利用
• 実装コストを削減するために、後処理のLambdaはS3Putのイベントで起動

Bedrockのバッチ推論は、プロンプトをJSONL形式でS3にアップロードし、APIで推論を開始すると、完了後に推論結果がS3に出力される仕組みです。

詰まったところ

問題は後処理のLambdaで発生しました。
150件ほどのプロンプトを入力したのですが、後処理のLambdaの出力には1/3ほどしか含まれていませんでした。

まずは後処理のLambdaのバグを疑いましたが、ログを詳しく調査した結果、Lambdaを起動した時点でS3上の推論結果ファイル自体が欠損していると判明しました。

どうやら、推論結果のファイルが作成された時には全件の推論が終わっておらず、段階的に出力されていそうです。
しかし、この動作については公式ドキュメントや検索でも情報が見当たりませんでした。

そこで、イベントトリガーではなく、バッチ推論の完了を確認した上で、手動でLambdaを起動したところ、全件が正しく出力されました。
これは段階的な書き出しが原因と確定して良さそうです。

解決方法

バッチ推論に続けて処理を行う場合、バッチ推論ジョブの状態を確認する必要があると分かりました。
この確認には、二つの方法が考えられます。

1. Bedrock APIでバッチ推論ジョブの状態を取得する
2. manifest.json.outファイルを確認する

manifest.json.outファイルとは？

Bedrockにより自動で作成される、推論結果のメタ情報を記録したファイルです。
バッチ推論のレコード数や、トークン数が記載されています。
この中に、入力の総レコード数（totalRecordCount）と処理の完了したレコード数（processedRecordCount）が書かれているので、この一致を確認すればバッチ推論ジョブの状態を確認できます。

今回は、IAMの権限を増やさなくて良い後者を選択することにしました。
既存のS3に対するアクセス権だけで対応できるので、変更が少なく実装コストを抑えられます。

PoCということから、簡易的にLambdaの中でポーリングを行います。
ポーリングの処理はおおよそ以下の通りです。

1. manifest.json.outを取得
2. Pydanticで取得したテキスト（JSON）をバリデート
3. totalRecordCountとprocessedRecordCountを比較
4. 不一致なら1分待機して再試行

今回のケースでは、約5分ほどで待機が完了しました。
出力も問題なく、全行が取得できました。

エビデンス

実装していた時は見つけられなかったのですが、後日、段階的な書き出しに関する情報を見つけました。
AWS Black Belt Online Seminar の資料
こちらの27ページにしっかりと書いてありました。

バッチ推論ジョブの途中でも定期的に更新される

また、バッチ推論の状態確認についても以下の記載がありました。

GetModelInvocationJob API でジョブの状態を、manifest.json.out ファイルで処理されたレコード等の数値を確認すれば、バッチ推論ジョブの実⾏途中でも進捗を追える

どちらも結果的に正しい対応ができていたようです。
よかったよかった。

現状

現在は前述したシステムの本実装を進めています。

本実装ではStep Functionsを利用し、推論の完了はTaskからGetModelInvocationJobを用いてポーリングで確認しています。
EventBridgeでも推論完了は取得できるようですが、少し煩雑になりそうだったので見送りました。

まとめ

今回はBedrockバッチ推論を利用して詰まったところを記事にしました。
段階的な書き出し（分割書き出し？）の情報はパッと見つからなかったので、同じく迷える方のところに届けば幸いです。

Bedrockはとても便利です。
便利ですが、急成長な分、ドキュメント整備が追いついていない部分があると思います。

色々手探りで実装を進める中で、Qiitaをはじめ様々な先人の知恵のお世話になりました。
この記事も、どこかで同じような方のお力になれれば嬉しいです！

Nova2も出たことですし、これからもBedrockはますます楽しみですね。
それでは、メリークリスマス！