429を出さずにバルク処理を速くする：Gemini APIの適応的並行度制御

import asyncio
import time
 
class AdaptiveLimiter:
    def __init__(self, start=8, floor=2, ceil=96, cooldown=4.0):
        self.limit = float(start)   # 現在の許容並行数
        self.floor = floor          # 下限（停止防止）
        self.ceil = ceil            # 上限（TPM保護）
        self.cooldown = cooldown    # 減少後に増加を止める秒数
        self._last_drop = 0.0
        self._sem = asyncio.Semaphore(start)
        self._inflight = 0
        self._lock = asyncio.Lock()
 
    def _resize(self, new_limit):
        new_limit = max(self.floor, min(self.ceil, new_limit))
        delta = int(new_limit) - int(self.limit)
        self.limit = new_limit
        # セマフォの許可枚数を差分だけ増減する
        if delta > 0:
            for _ in range(delta):
                self._sem.release()
        # 縮小は acquire 済みトークンの自然返却で吸収（即時 acquire はしない）
 
    async def on_success(self):
        async with self._lock:
            if time.monotonic() - self._last_drop < self.cooldown:
                return  # クールダウン中は増やさない
            self._resize(self.limit + 1)            # Additive Increase
 
    async def on_429(self):
        async with self._lock:
            self._last_drop = time.monotonic()
            self._resize(self.limit * 0.5)          # Multiplicative Decrease
 
    async def acquire(self):
        await self._sem.acquire()
 
    def release(self):
        self._sem.release()

async def run_bulk(items, handler, limiter, workers=128):
    queue = asyncio.Queue()
    for it in items:
        queue.put_nowait(it)
    results, dead = [], []
 
    async def worker():
        while True:
            try:
                item = queue.get_nowait()
            except asyncio.QueueEmpty:
                return
            await limiter.acquire()
            try:
                out = await handler(item)      # 1件のGemini呼び出し
                await limiter.on_success()
                results.append((item, out))
            except RateLimited:
                await limiter.on_429()
                queue.put_nowait(item)         # 429は並行数を下げて再投入
            except PermanentError as e:
                dead.append((item, str(e)))    # 復旧不能はデッドレターへ
            finally:
                limiter.release()
                queue.task_done()
 
    await asyncio.gather(*[worker() for _ in range(workers)])
    return results, dead

import json, pathlib
 
CKPT = pathlib.Path("checkpoint.jsonl")
 
def load_done():
    if not CKPT.exists():
        return set()
    return {json.loads(l)["id"] for l in CKPT.read_text().splitlines() if l}
 
def mark_done(item_id, payload):
    with CKPT.open("a") as f:
        f.write(json.dumps({"id": item_id, "out": payload}) + "\n")
 
# run_bulk 起動前
done = load_done()
pending = [it for it in all_items if it["id"] not in done]