LLMの並列処理で残る1タスクをTaskGroupで止める

端末にはエラーが出た。でも、もう片方は待機したまま。今日、LLMの並列呼び出しを模したコードで再現した。

自分は、両方の応答が必須ならTaskGroupでまとめたい。gather()で例外を捕まえても、残りは止まらないんだよね。外部APIは呼ばず、残る1タスクを0にする。

2タスクで再現する

Python 3.12.13で実行済み。cancel_demo.pyに保存して実行する。

slow()は応答待ち役、fail()は失敗役。Eventで待機開始をそろえた。

import asyncio

async def demo(grouped):
    started = asyncio.Event()

    async def slow():
        started.set()
        await asyncio.Event().wait()

    async def fail():
        await started.wait()
        raise RuntimeError("mock LLM failed")

    if grouped:
        try:
            async with asyncio.TaskGroup() as group:
                task = group.create_task(slow())
                group.create_task(fail())
        except* RuntimeError as errors:
            print("error:", errors.exceptions[0])
        print("TaskGroup pending:", int(not task.done()))
        print("TaskGroup cancelled:", task.cancelled())
    else:
        task = asyncio.create_task(slow())
        batch = asyncio.gather(task, fail())
        try:
            await batch
        except RuntimeError as error:
            print("error:", error)
        print("gather pending:", int(not task.done()))
        print("late cancel:", batch.cancel())
        task.cancel()
        await asyncio.gather(task, return_exceptions=True)

async def main():
    await demo(False)
    await demo(True)

asyncio.run(main())

error: mock LLM failed
gather pending: 1
late cancel: False
error: mock LLM failed
TaskGroup pending: 0
TaskGroup cancelled: True

pendingは待機役が未完了なら1。gather側の末尾2行で、残したタスクを後始末している。

例外のあとにcancelしても遅い

引っかかったのはlate cancel: False。例外を受け取った時点でFutureは完了扱い。そこへcancel()しても、待機役には届かない。

同梱のCPython 3.12.13の実装を読むと、gatherは子の例外を外側のFutureに設定して戻る。この経路に、ほかの子をキャンセルする処理はない。

TaskGroupは、子が通常の例外で失敗すると残りの終了を待つ。taskgroups.pyでも確認した。残りの子にキャンセルを要求し、終了後に例外グループを投げる。ここではexcept* RuntimeErrorで受けた。

終了時にどこまで待つかが違う。

asyncio.run()の終了時にも残りがキャンセルされる。今回はその前に数えた。常駐サーバーでは、要求のたびにループは閉じない。

LLM呼び出しに戻すとき

両方必須ならgroup.create_task()で登録する。片方でも使えるなら、残りを待つgather(return_exceptions=True)も選べる。

キャンセルは協調動作なので、子がCancelledErrorを握りつぶすと終了を待ち続ける場合がある。今回は伝播させた。

確認できたのは残るタスクが1から0になったこと。LLMサーバーが推論を止めるか、料金が減るかは、この実験では分からない。組み込むときも片方を失敗させ、要求を返す前にもう片方の終了を確認する。並列数を増やす前に、この小さい失敗経路を通しておきたい。