ストリーミング応答のCPUと取りこぼしを実測して長時間バッチを安定させる

import time
import psutil
from dataclasses import dataclass, field
from anthropic import Anthropic
 
client = Anthropic()  # ANTHROPIC_API_KEY を環境変数から読む
proc = psutil.Process()
 
 
@dataclass
class StreamStats:
    wall_seconds: float = 0.0       # 壁時計の経過秒数
    cpu_seconds: float = 0.0        # このプロセスが消費したCPU秒数
    output_tokens: int = 0          # 受信した出力トークン数
    chunks: int = 0                 # 受信したテキストデルタの個数
    text: str = field(default="")
 
    @property
    def cpu_cores(self) -> float:
        # 1.0 なら平均してコア1つを占有していた、という意味
        return self.cpu_seconds / self.wall_seconds if self.wall_seconds else 0.0
 
    @property
    def tokens_per_sec(self) -> float:
        return self.output_tokens / self.wall_seconds if self.wall_seconds else 0.0
 
 
def measured_stream(model: str, messages: list, max_tokens: int = 2048) -> StreamStats:
    stats = StreamStats()
    cpu_before = sum(proc.cpu_times()[:2])   # user + system のCPU秒
    wall_before = time.perf_counter()
 
    with client.messages.stream(
        model=model,
        max_tokens=max_tokens,
        messages=messages,
    ) as stream:
        for text in stream.text_stream:
            stats.text += text
            stats.chunks += 1
        final = stream.get_final_message()
        stats.output_tokens = final.usage.output_tokens
 
    stats.wall_seconds = time.perf_counter() - wall_before
    stats.cpu_seconds = sum(proc.cpu_times()[:2]) - cpu_before
    return stats

stats = measured_stream(
    "claude-sonnet-4-6",
    [{"role": "user", "content": "Summarize the benefits of streaming in 200 words."}],
)
print(f"wall={stats.wall_seconds:.1f}s "
      f"cpu={stats.cpu_seconds:.1f}s "
      f"cores={stats.cpu_cores:.2f} "
      f"tok/s={stats.tokens_per_sec:.1f}")
# 例: wall=6.3s cpu=1.4s cores=0.22 tok/s=33.5

import time
import httpx
from anthropic import Anthropic, APIStatusError
 
client = Anthropic()
 
 
def resilient_stream(
    model: str,
    messages: list,
    max_tokens: int = 2048,
    idle_timeout: float = 30.0,   # この秒数チャンクが来なければ中断扱い
    max_retries: int = 4,
) -> str:
    attempt = 0
    while True:
        attempt += 1
        last_recv = time.monotonic()
        buf = []
        try:
            with client.messages.stream(
                model=model, max_tokens=max_tokens, messages=messages,
            ) as stream:
                for text in stream.text_stream:
                    now = time.monotonic()
                    if now - last_recv > idle_timeout:
                        raise TimeoutError(
                            f"idle {now - last_recv:.0f}s > {idle_timeout}s"
                        )
                    last_recv = now
                    buf.append(text)
                stream.get_final_message()   # 完了確定（ここまで来れば成功）
            return "".join(buf)
 
        except (TimeoutError, httpx.HTTPError) as e:
            if attempt > max_retries:
                raise
            backoff = min(2 ** attempt, 30)
            print(f"[retry {attempt}] {type(e).__name__}: {e} -> {backoff}s")
            time.sleep(backoff)
 
        except APIStatusError as e:
            # 429（レート制限）と 5xx だけ再試行。4xx の論理エラーは即座に投げる
            if e.status_code not in (429, 500, 502, 503, 529) or attempt > max_retries:
                raise
            backoff = min(2 ** attempt, 30)
            print(f"[retry {attempt}] HTTP {e.status_code} -> {backoff}s")
            time.sleep(backoff)

import json
import os
from pathlib import Path
 
CKPT = Path("batch_progress.jsonl")
 
 
def load_done() -> set:
    if not CKPT.exists():
        return set()
    done = set()
    for line in CKPT.read_text().splitlines():
        if line.strip():
            done.add(json.loads(line)["id"])
    return done
 
 
def run_batch(jobs: list[dict], model: str = "claude-sonnet-4-6"):
    done = load_done()
    pending = [j for j in jobs if j["id"] not in done]
    print(f"total={len(jobs)} done={len(done)} pending={len(pending)}")
 
    with CKPT.open("a") as ck:
        for job in pending:
            stats = measured_stream(model, job["messages"])
            # 中断が起きうるならここを resilient_stream に差し替える
            ck.write(json.dumps({
                "id": job["id"],
                "tokens": stats.output_tokens,
                "wall": round(stats.wall_seconds, 2),
                "cores": round(stats.cpu_cores, 2),
            }, ensure_ascii=False) + "\n")
            ck.flush()
            os.fsync(ck.fileno())   # 電源断でも直前の完了を失わない
            print(f"done {job['id']} ({stats.tokens_per_sec:.0f} tok/s)")