進行状況のラベルに潜んでいたメモリ不足によるクラッシュ

9月下旬のある晩、ユーザーのワークスペースでAIエージェントが、ターンの途中でメモリ不足のエラーにより停止しました。疑わしいとされたのは、同時に動くサブエージェントが多すぎることと、コンテキストが大きくなりすぎたことでした。しかしどちらも的外れでした。本当の原因はステータスラベルでした。モデルが大きなシェルコマンドを数文字ずつストリーミングしている間、私たちは、断片が届くたびにそれまでのコマンド全体を進行状況のラベルとして再送していたのです。124 KBのコマンドが、約41,000件のイベントと、約690 MBのキューに入った文字列になっていました。

エージェントの作業のストリーミング方式

Mythexでアプリを作る場合、エージェントはプロジェクトのプライベートなクラウドサンドボックスの中で動きます。モデルが書くと、その出力がブラウザにストリーミングされます。テキストは一語ずつ表示され、シェルコマンドのようなツールを実行すると、そのラベルがライブで更新されます。そのため、何をしているかが分かります。

モデルはデータを小さな断片でストリーミングします。使っていたモデルでは、その断片はしばしばわずか3文字でした。

最初の仮説は間違っていた

クラッシュは、規模の問題のように見えたので、まずそこから調べました。

並列のサブエージェント？ そのターンで使われていたのは2つだけで、それぞれ約1.5 MBの状態を保持していました。それでは足りません。
保持されたトレース？ トレースは実行ツリーをメモリに保持していましたが、クラッシュ時のヒープの大きさにはまったく届きませんでした。

どちらももっともらしい話で、どちらも信じるのは簡単でした。しかし、どちらでも失敗を再現できませんでした。何度も学び直すことになる教訓があります。失敗を再現できるまで、仮説は原因ではないのです。

本当に再現する

本番と同じメモリ上限を持つ使い捨てのサンドボックスを作り、現実的な3文字の断片で大きなファイルをストリーミングする、偽のモデルを用意しました。さらに、ソケットの反対側の受信側を、実際のネットワークのように遅くしました。

それで再現しました。ヒープが増え続け、同じエラーでプロセスが停止したのです。

仕組みは、いったん観察できれば単純でした。モデルが大きなコマンドを書いている間、この場合は124 KBのファイルを書き込むheredocでしたが、ストリーミングのコードは断片ごとにツールのラベルを更新していました。そのラベルは新しい3文字ではなく、それまでのコマンド全体でした。そのためイベントは、3文字、次に6文字、そして9文字と、124 KBに達するまで約41,000回、増え続けていたのです。ソケットは、生成される速さに送信が追いつかず、書き込みバッファーに文字列として積み上がりました。その合計が約690 MBでした。

修正

ラベルに上限を設けました。 ステータスラベルは、何が起きているかを示せばよいものです。今では200文字に制限し、部分的なコマンドを再スキャンする頻度も抑えています。
二乗の連結をやめました。 ストリーミングされた断片を、文字列の繰り返し連結でつなぐのは、大きな出力に対してO(n²)になります。線形時間で追加する、チャンクベースの方式に切り替えました。

その後、大量の出力を受けて無制限に増えうる箇所を探し、それぞれに上限を設けました。

チャットに表示されるライブのコマンド出力は、まとめて処理し、最後の8,000文字だけを保持します。
個々のイベント行は4 MBで上限を設けます。
ブラウザが再接続後に追いつくためのリプレイバッファーは、プロジェクトごとに8 MBで上限を設けます。
古いチェックポイントは、ターンが進むにつれて削除されます。ただし、最新の状態と、以前のターンが必要とするアンカーは残します。

大量出力を正直にテストする

このテストが想定よりずっと難しかった理由は2つあります。

実際のモデルは協力してくれません。 巨大なファイルを出力するよう頼んでも、たいてい拒否するか、要約するか、より穏当な書き方をします。大量出力をテストするには、指示したとおりに動き、実際のエージェントのコードを通って動く、スクリプト化されたモデルが必要です。
スクリプトはモデルらしく振る舞わなければなりません。 最初のスクリプト化したモデルは、すべての応答で同じレスポンスIDを再利用していました。そのためエージェントのフレームワークは、各応答を前の応答の置き換えとみなし、ターンを途中で終えてしまいました。UIの更新のバグのように見えましたが、そうではありませんでした。テストのハーネスが私たちに嘘をついていたのです。また、メモリの数値が非現実的にならないよう、スクリプトでモデルの実際のコンテキスト上限も強制する必要がありました。

得られた教訓

ストリーミングUIは小さなミスを増幅します。 「断片ごと」に送るものは、何万回も送られることになります。差分を送るか、ペイロードに上限を設けてください。
遅い境界をまたぐすべてのバッファーに上限を設けてください。 ソケット、リプレイログ、チェックポイントには、どれも天井が必要です。
修正する前に再現してください。 最初の2つの説明はもっともらしく聞こえましたが、間違っていました。

最近Mythexのチャットで、エージェントが止まることなく非常に長いファイルが表示されるのを見たことがあれば、これがその理由の一部です。