Kiro のクレジットが想定以上に消費されてしまった話

SCSKの畑です。

小ネタというか、Kiro そのものの話とは少し外れた話題なのですが、放っておくと大事になりかねない内容だったので備忘として残しておこうと思った次第です。

 

本エントリの前提

とある案件の PoC において、kiro-cli を使用して大量のファイルをバッチ処理するためのツールを Python で作成していました。内容の詳細は今回の内容に直接関連しないため触れませんが、対象のファイルから特定の情報を抽出する目的で使用しています。

他、ツールの作りは大まかに以下のようになっています。

  • コンテキストサイズの制約を考慮し、処理対象となるファイル単位でツールから kiro-cli を実行
  • 大量のファイルを処理する場合は相応の時間がかかるため、ツールから kiro-cli をマルチスレッドで実行可能に
  • 後述するようにサイズの大きいファイルを処理する場合は処理時間が想定以上に長くなってしまうことも起こり得るため、一定時間が経過した kiro-cli のプロセスはタイムアウト扱いとしてツール側から終了する

 

発生した事象

このツールを使用して PoC を進めていたのですが、ワンサイクルの処理で概ね 200-300 前後のクレジットを消費する重い処理（プロンプト）について、コンテキスト/クレジットの消費量を抑えるためにプロンプトやツール自体の見直しを実施していました。良く知られている通り、コンテキストサイズが上限近くに達してしまうと要約が走ってしまい目的の処理が無限ループしてしまったり、kiro-cli からの出力が途中で切れてしまうなどの弊害がある上、コンテキストの消費量を抑えることは消費クレジットの低減に直結するためです。

その試行錯誤をしている最中、現在のクレジット消費量がどれくらいかどうかを確認したくなったため、kiro-cli で /usage を叩いてみました。ツール側で kiro-cli の処理完了時に表示される消費クレジットを集計できるようにしており、概ね 2000 程度まで消費されていることは分かっていましたが、実際のクレジット消費量との突合せがしたいなと。

ところが、/usage に表示されたクレジットの消費量は、想定の 1.5 倍である約 3000 まで増加していました。しかも、継続的にクレジットの消費量を観察していたところ、少しずつゆっくりと値が増え続けているような状況でした。

さすがに血の気が引きつつも、これはマズいということで至急原因を調査する羽目になりました。

 

原因

そもそも何故消費量が 3000 に達したのかという点も疑問でしたが、まずクレジットの値が増え続けている状況を改善する必要がありました。クレジットの消費量は完全にリアルタイムで反映される訳ではなく、特にマルチスレッドで並列実行するようなケースだと反映が遅れるケースがあったので当初はその可能性を疑っていました。しかし数分程度経っても状況が変わらなかったため、この線はなさそうと判断しました。

次に Kiro のダッシュボードから手がかりとなる情報がないかを探したのですが、サブスクリプション利用状況の概観は確認できるものの、リアルタイムで利用状況を詳細に掘り下げる用途では使用できなさそうでした。機能的には S3 に出力するアクティビティレポートが唯一使えそうだったのですがその時点では有効化しておらず、このタイミングで有効化しても意味があるのかが分からなかったので一旦見送ることに。

この時点で最悪 AWS サポートに問い合わせるしかないかなとも考え始めていたのですが、ひょっとするとタイムアウトした kiro-cli のプロセスが残存しているのではないか？という可能性がふと頭に浮かびました。一連の試行錯誤において数回タイムアウトが発生していたのですが、そのプロセスが残存している可能性があるのではないかと考えたためです。そこで慌てて ps コマンドで kiro-cli プロセスの存在を確認したところ・・

何と 2 つのプロセスが残存していました。いずれのプロセスも入力（プロンプト）の内容がタイムアウトした処理に合致していた上、しかもその内の 1 個は数時間前から稼働していることに気づきました！

ということで速攻でそれらのプロセスを kill したところ、ようやくクレジットの増加を止めることができました。

また、この時点で消費量が 3000 に達した原因は、残存していた kiro-cli プロセスが単純に少しずつクレジットを消費していたためではないかという推測もつきました。タイムアウトの回数と残存していたプロセスの個数は一致しませんでしたが、プロセスによってはある程度の時間稼働した結果正常に終了した可能性も考えられるため、それも考えると十分に辻褄の合う程度の消費量だったためです。

 

解決策

ツール側でタイムアウトとして検知したプロセスが正常に kill されていないという事象はツール側の責任であるため、ツール側を改修することで解決しました。

厳密に言うとタイムアウトしたプロセスを kill する仕組みは実装されていたのですが、その対象となっていたのは kiro-cli プロセスのみでした。実際には、kiro-cli におけるチャットセッションは以下のように kiro-cli-chat という子プロセスとして起動されるようで、親プロセスの kiro-cli を kill するだけでは Linux(WSL) の仕様として子プロセスは kill されません。今回残存していたのはいずれもこの子プロセスであり、それが断続的なクレジット消費に繋がっていたと考えられるため、タイムアウト時に親プロセスと合わせて確実に kill する必要があります。

★kiro-cli-chatがkiro-cliの子プロセスとして稼働している
$ ps -ef | grep kiro-cli | grep -v grep
user 2322 644 0 17:52 pts/0 00:00:00 kiro-cli
user 2330 2322 0 17:52 pts/0 00:00:00 /home/user/.local/bin/kiro-cli-chat chat

★親プロセスのkiro-cliをkill
$ kill -9 2322

★子プロセスのkiro-cli-chatは親プロセスをkillしても残存している
$ ps -ef | grep kiro-cli| grep -v grep
user 2330 643 0 17:52 pts/0 00:00:00 /home/user/.local/bin/kiro-cli-chat chat

よって、以下のようにツール（Python）から kiro-cli を subprocess として起動する際、「preexec_fn=os.setsid」を指定することでプロセスグループを作成するように変更しました。この変更により、親プロセス（kiro-cli）とその子プロセス（kiro-cli-chat）の両方に対して SIGTERM や SIGKILL などのシグナルを送れるようになります。

process = subprocess.Popen(
            command,
            text=True,
            encoding='utf-8',
            preexec_fn=os.setsid,
        )

その上で、対象プロセスグループに対して kill 用のシグナルを送信するように、対象メソッドのロジックを修正しました。具体的な実装例は以下の通りです。

def kill_pt(process: subprocess.Popen) -> None:
  # プロセスグループにSIGTERMを送信
  pid = process.pid
  pgid = os.getpgid(pid)
  try:
    os.killpg(pgid, signal.SIGTERM)
  except ProcessLookupError:
    return

  # SIGTERM送信後、プロセスグループのgraceful shutdownを待機
  deadline = time.time() + 5
  while time.time() < deadline:
    try:
      os.killpg(pgid, 0) # Null Signalを使用した生存確認
    except ProcessLookupError:
      process.wait()
      return
    time.sleep(1)

  # SIGTERMで停止しなかった場合、プロセスグループにSIGKILLを送信して強制終了
  try:
    os.killpg(pgid, signal.SIGKILL)
  except ProcessLookupError:
    process.wait()

また、試行錯誤の最中に場合によっては該当処理がタイムアウトしそうと判断できた時点で Ctrl+C でツールを強制終了するようなこともしていたのですが、このタイミングでツール経由で起動していた kiro-cli を終了するような仕組みも不十分であったため、KeyboardInterrupt を kiro-cli プロセスの起動メソッドで catch して上記メソッドで kill するようにしています。（ひょっとすると主原因はこちらだったかもしれませんが・・）

それとは別の話として、先述したような「コンテキスト逼迫による無限ループ」については Kiro 側で検知して欲しさもあります。。IDE での実行だと一定回数以上のイテレーションがあった場合は続けるかどうかの選択肢が最近出るようになりましたが、検知の仕組みがインタラクティブである場合バッチ実行すると結局そこでスタックしてしまうので、他の仕組みがないものでしょうかね。ひょっとするともう何かありそうな気もしますが・・

 

まとめ

実はトータルのクレジット消費量が 3000 に達する前に 1 回確認しており、その時にも同じようにクレジット消費量がゆっくりと増加するような事象は確認できていたのですが、原因のセクションに書いた通り一過性の事象だとスルーしてしまっていました・・思い込みは危険なので気になった時点でちゃんと原因を探りましょう。本当に。

Kiro の仕組みを活用するという観点だと、S3 に出力するアクティビティレポートを有効化する以外の方策はなさそうなので、こちらも別の機会に試したいです。このへん、他の生成 AI エージェントだとどのように監視できるのかとかも気になりますね・・

本記事がどなたかの役に立てば幸いです。

タイトルとURLをコピーしました