画像を多用するエージェントのターンをメモリ内に収める

本番環境で、あるエージェントのターンが、カタログのスキャン6枚から商品を切り抜き、その切り抜きを確認している最中に、952 MBでメモリ不足になりました。原因は1つのリークではなく、3つありました。モデルへの呼び出しのたびに、そのターンの画像をフルサイズで再送していたこと。モデルを呼び出すライブラリが、ターンが終わるまですべてのリクエストを保持し続けていたこと。そして、エージェントフレームワークのトレースが、画像を含むすべての呼び出しの入力を、ターン全体にわたって保持していたことです。3つすべてを修正しました。送る前に画像を縮小し、各呼び出しの画像は、その呼び出しが終わったらすぐ解放するようにしました。

エージェントが画像をどう見るか

Mythexで作業すると、エージェントはプロジェクトのクラウドサンドボックスの中で動きます。画像を見ることができます。添付したスクリーンショット、モックアップ、プロジェクト内のファイル、あるいはエージェント自身が作った画像などです。会話と一緒にそれらをモデルに送ることで見ています。

ターンとは、あなたの1つのメッセージに応えてエージェントが行うすべての作業を指します。ターンの中で、エージェントは何度もモデルを呼び出します。考え、ツールを実行し、結果を見て、また考えます。それぞれの呼び出しでは、これまでの会話が送られます。そこには、そのターンの画像も含まれます。画像はテキスト（base64）としてエンコードされるため、約3分の1大きくなります。

クラッシュしたターンは、妥当なことをしていました。2.4 MBほどのカタログスキャンを6枚取り、そこから個々の商品を切り抜き、その切り抜きを何度も読み直して確認していました。大きな画像を複数含んだモデル呼び出しを、多数行っていたわけです。

原因1：呼び出しのたびに、フルサイズの画像を再送していた

各モデル呼び出しは、そのターンの画像をフルサイズのまま、リクエストボディを組み立て直していました。1回の呼び出しあたり、約20 MBのbase64です。ボディはエージェントのメモリ内で1つの文字列として組み立てられるので、呼び出しのたびに約20 MBを新たに確保していました。

しかも、フルサイズで送っても何も得られていませんでした。モデルは、長辺が約1,300ピクセルになるように自分で画像を縮小します。余分なピクセルは、メモリとアップロード時間を無駄にするだけでした。

修正：一度縮小して、小さく送る

400 KBを超える画像は、送る前に、長辺が最大1,568ピクセルのJPEGに変換されます。小さい画像はそのまま送られます。
変換には、すべてのMythexサンドボックスに入っているImageMagickを使い、ファイルごとに1回だけ実行します。毎回の呼び出しで同じスキャンを見るターンは、縮小済みのコピーを再利用します。キャッシュは、ファイルのパス、サイズ、更新日時をキーにしているので、編集されたファイルは再変換されます。また、キャッシュ自体にもサイズ上限があり、それがメモリを埋めることはありません。
1回のモデル呼び出しで画像に使える予算は、16 MBから6 MBに減りました。

この結果、同じ6枚のスキャンは、呼び出しあたり19.7 MBではなく4.4 MBで済むようになりました。

原因2：モデルのクライアントがリクエストを解放しなかった

リクエストのサイズは、問題の一部にすぎませんでした。メモリは、呼び出しを重ねるごとに増えていきました。ステージング環境では、同じ種類のターンが、完全なガベージコレクションの後でも369 MBの生きたメモリを保持していました。ヒープのサンプリングによると、増加分はすべて、リクエストボディをエンコードするモデルクライアントライブラリの部分に集中していました。

原因は、アボートリスナーでした。Stopを押すと、ターンのアボートシグナルが発火し、それを監視しているすべてのリクエストがキャンセルされます。私たちは、そのターン全体で1つのシグナルを、すべてのリクエストに渡していました。クライアントライブラリは、渡されたシグナルに「abort」のリスナーを追加しますが、それを削除することはありません。

そのため、ターンのすべてのリクエストは、ターンが続く限り、ターンのシグナルから到達可能なままでした。リスナーがコントローラーを保持し、そのコントローラーがリクエストを保持し、そのリクエストが画像を含むリクエストボディ全体を保持していたのです。

修正：リクエストごとに1つのシグナル

各リクエストは、自分専用のシグナルを持つようになりました。それは、そのリクエストが実行中の間だけ、ターンのシグナルにリンクされます。ストリーミングされたレスポンスなら、最後のチャンクを読み終えるまでです。その直後にリンクは解除されます。Stopは、実行中のどのリクエストにも届きますが、ターンのシグナルが同時に保持するリスナーは1つだけになります。

テストでは、この3つの部分すべてを確認しています。

12回の呼び出しの後、ターンのシグナルにはリスナーが1つも残らないことを確認します。修正前は12個ありました。
8 MBのリクエストを15回送ると、メモリは横ばいのままです。修正前は128 MB増えていました。
Stopを押すと、ストリームの途中であっても、実行中のリクエストは中断されます。

原因3：トレースが、各呼び出しの入力を保持していた

ヒープのスナップショットには、3つ目の保持者が写っていました。モデルの呼び出し1回あたり、約6 MBが保持されていたのです。これはエージェントフレームワークのイベントストリーミングに由来していました。ターン内のモデル呼び出しを、入力も含めて、ターンの子として記録し、ターンが終わるまで保持していたのです。それぞれの呼び出しに添付されたすべての画像が、一緒にメモリに残っていました。

修正：呼び出しが返ったら画像を空にする

画像は、特定のモデル呼び出しのために組み立てられたメッセージにだけ添付されます。会話の保存された状態に入れられることはありません。そのため、呼び出しが返ったら、それらのメッセージ内の画像データを空にしても安全です。トレースには呼び出しのテキストの記録が残り、ピクセルデータは消えます。

重要だったのは、画像をその場で空にすることです。トレーサーは、同じ画像データに、メッセージ上の複数の参照を通じて到達できる場合があります。そのため、新しい空のリストに差し替えてしまうと、元のピクセルデータが依然として到達可能なまま残ってしまいます。保存された会話の状態は、一切変更されません。

テストでは、エージェントのグラフを、同じイベントストリーミングの下で、毎回の呼び出しに4 MBの画像を付けて実行しています。修正前は、16回の呼び出しで77 MBメモリが増えました。修正後は、メモリは横ばいのままです。

この件から学んだこと

メモリの問題は層になっていることがあります。 3つの別々の原因がそれぞれ画像を長く保持しており、ターンがメモリ内に収まるには3つすべてを解消する必要がありました。
ライブラリに渡したものは、保持されるかもしれません。 シグナルもトレースも、渡された呼び出しより長く生き続けていました。呼び出しごとにメモリが増えるなら、その呼び出しが触れている長寿命のオブジェクトを確認してください。
相手が捨てるものは送らないでください。 モデルがどうせ捨てるピクセルは、毎回の呼び出しでメモリを消費していました。
メモリは直接テストしてください。 修正のそれぞれに、ヒープの増加やリスナー数を測るテストを付けました。そうすれば、回帰はクラッシュではなく数字として現れます。

エージェントがメモリを使う話は、短期間に2つ目です。1つ目は、進捗ラベルが毎回41,000回も再送されていた件です。エージェントとは何か、どう動くのかを知りたい場合は、AIエージェントのガイドと、チャットのドキュメントをご覧ください。