PolyGateway/tools/soak/run_soak.py

"""压测入口(M2 设计 §8): 预算硬顶、隔离守卫、多进程 worker、跑后记分。

用法(conda 环境内,项目根目录):
    python tools/soak/run_soak.py --scenario P3 --budget-calls 6 \\
        --budget-tokens 50000 --workers 2 [--scope LLM] [--concurrency 8]

守卫(违反即拒跑): REDIS_URL 必须指向 db3(实验室 db0 有在用键);
PGW_TELEMETRY_PG_DSN 若配必须指向 polygateway 专用库;--workers>1 时
限流/熔断后端必须是 redis(内存后端不跨进程,多 worker 无共享闸即超发)。
预算双上限(请求数/token)任一命中优雅停;签字硬顶(设计 §8.1):
全程 token ≤ 2 亿。soak 与 pytest 不并跑(FLUSHDB 清 db3 测试状态)。
"""

from __future__ import annotations

import argparse
import asyncio
import json
import multiprocessing as mp
import os
import subprocess
import sys
import time
from pathlib import Path

_ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(_ROOT))
sys.path.insert(0, str(_ROOT / "src"))

TOKEN_HARD_CAP = 200_000_000  # 2 亿,2026-07-20 人类签字(设计 §8.1)


def _merged_env() -> dict[str, str]:
    from dotenv import dotenv_values

    return {k: v for k, v in {**dotenv_values(_ROOT / ".env"), **os.environ}.items() if v}


_SIGNED_MAX_CONCURRENCY = 100  # 网关保护签字值(设计 §8.1,2026-07-20 人类)
_SIGNED_MAX_RPM = 600


def _guard(env: dict[str, str], workers: int, scope: str) -> None:
    redis_url = env.get("REDIS_URL", "")
    if not redis_url.rstrip("/").endswith("/3"):
        raise SystemExit(f"拒跑: REDIS_URL 必须指向专用 db3,当前 {redis_url!r}")
    pg = env.get("PGW_TELEMETRY_PG_DSN", "")
    if pg and not pg.rstrip("/").endswith("/polygateway"):
        raise SystemExit("拒跑: PG DSN 必须指向 polygateway 专用库,当前库名不符")
    if workers > 1 and (
        env.get("PGW_LIMITER_BACKEND") != "redis" or env.get("PGW_BREAKER_BACKEND") != "redis"
    ):
        raise SystemExit("拒跑: --workers>1 需要 PGW_LIMITER_BACKEND/PGW_BREAKER_BACKEND=redis")
    # 网关保护(签字值 100/600): 压测 scope 必须配全局闸且不超签字上限
    conc = env.get(f"{scope}__GLOBAL__MAX_CONCURRENCY")
    rpm = env.get(f"{scope}__GLOBAL__RPM")
    if not conc or not rpm:
        raise SystemExit(
            f"拒跑: 压测须配网关保护 {scope}__GLOBAL__MAX_CONCURRENCY(≤{_SIGNED_MAX_CONCURRENCY})"
            f" 与 {scope}__GLOBAL__RPM(≤{_SIGNED_MAX_RPM})——设计 §8.1 签字值"
        )
    if int(conc) > _SIGNED_MAX_CONCURRENCY or int(rpm) > _SIGNED_MAX_RPM:
        raise SystemExit(
            f"拒跑: 全局闸 {conc}/{rpm} 超签字上限 {_SIGNED_MAX_CONCURRENCY}/{_SIGNED_MAX_RPM}"
        )


async def _flush_db3(redis_url: str) -> None:
    import redis.asyncio as aioredis

    client = aioredis.from_url(redis_url)
    try:
        await client.flushdb()
    finally:
        await client.aclose()


def _rss_mb() -> float:
    """当前 RSS(ps 口径,darwin/linux 皆 KB)。

    不用 ru_maxrss——那是历史峰值只增不减,2026-07-21 P6 曾把分发期
    瞬时缓冲峰值误判为 780MB "泄漏"(实测跑中稳定 36MB)。
    """
    out = subprocess.run(
        ["ps", "-o", "rss=", "-p", str(os.getpid())], capture_output=True, text=True, check=True
    )
    return int(out.stdout.strip()) / 1024.0


async def _paced_dispatch(
    generator, *, sem, spawn, should_stop, on_dispatched
) -> set[asyncio.Task]:
    """有界分发: 先占并发名额再建任务,名额由任务收尾释放。

    2026-07-21 P6 教训: 无界 create_task 曾在 15s 内入队全部预算,
    内存峰值 ~800MB,且 --max-hours 截止检查随分发结束而失效。
    """
    inflight: set[asyncio.Task] = set()

    async def _run(coro) -> None:
        try:
            await coro
        finally:
            sem.release()

    async for kind, kwargs in generator:
        if should_stop():
            break
        await sem.acquire()
        try:
            on_dispatched()
            task = asyncio.create_task(_run(spawn(kind, kwargs)))
        except BaseException:
            sem.release()  # 名额已占而任务未建(如进度采样失败),不泄漏
            raise
        inflight.add(task)
        task.add_done_callback(inflight.discard)
    return inflight


async def _worker_async(args: argparse.Namespace, worker_idx: int) -> None:
    """单 worker: 消费场景生成器,信号量限并发,双预算任一命中即停。"""
    from polygateway import GatewayClient
    from polygateway.telemetry.sqlite import SQLiteRecorder
    from tools.soak.scenarios import SCENARIOS, SoakCorpus

    env = _merged_env()
    run_id = args.run_id
    telemetry_path = _ROOT / f"data/soak/telemetry_{run_id}_{worker_idx}.db"
    recorder = SQLiteRecorder(telemetry_path)
    if args.scenario == "P7":
        from polygateway.ocr import OcrClient

        client = OcrClient.from_env(args.scope, telemetry=recorder, env=env)
    else:
        client = GatewayClient.from_env(args.scope, telemetry=recorder, env=env)
    corpus = SoakCorpus(
        harness_db=_ROOT / "data/soak/harness.db",
        telemetry_db=_ROOT / "data/soak/generate_questions_telemetry.db",
        frames_root=_ROOT / "data/soak/vt_frames",
        images_root=_ROOT / "data/soak/chs_images",
        # P7 只需 images: chains/replays 的 base64 负载曾把起跑 RSS 顶到 ~690MB(verifier I1)
        modalities=("images",) if args.scenario == "P7" else ("chains", "replays", "images"),
    )
    generator = SCENARIOS[args.scenario](corpus, f"{run_id}-w{worker_idx}")
    sem = asyncio.Semaphore(args.concurrency)
    stats = {"calls": 0, "ok": 0, "failed": 0, "cancelled": 0, "tokens": 0}
    rss_samples = [_rss_mb()]
    from tools.soak.scoreboard import capped_budget

    deadline = time.monotonic() + args.max_hours * 3600
    budget_calls = capped_budget(args.scenario, args.budget_calls) // args.workers
    budget_tokens = min(args.budget_tokens, TOKEN_HARD_CAP) // args.workers

    async def _one(kind: str, kwargs: dict) -> None:
        try:
            if kind == "ocr_text":
                await client.recognize_text(**kwargs)
            elif kind == "ocr_layout":
                await client.parse_layout(**kwargs)
            else:
                # run 级双层隔离: FLUSHDB(入口)+ 每 run 独立 namespace(此处)
                kwargs.setdefault("cache_namespace", run_id)
                resp = await client.chat(**kwargs)
                stats["tokens"] += resp.prompt_tokens + resp.completion_tokens
            stats["ok"] += 1
        except asyncio.CancelledError:
            stats["cancelled"] += 1
            raise
        except Exception as exc:  # 记分板从遥测读错误分布;这里只计数
            stats["failed"] += 1
            chain = getattr(exc, "per_source_reasons", None) or {}
            print(
                f"[w{worker_idx}] 调用失败: {type(exc).__name__}: {exc} 链={chain}",
                flush=True,
            )

    def _should_stop() -> bool:
        return (
            stats["calls"] >= budget_calls
            or stats["tokens"] >= budget_tokens
            or time.monotonic() > deadline
        )

    def _on_dispatched() -> None:
        stats["calls"] += 1
        if stats["calls"] % 20 == 0:
            rss_samples.append(_rss_mb())
            print(f"[w{worker_idx}] {stats}", flush=True)

    inflight = await _paced_dispatch(
        generator, sem=sem, spawn=_one, should_stop=_should_stop, on_dispatched=_on_dispatched
    )
    if inflight:
        await asyncio.gather(*inflight, return_exceptions=True)  # 优雅收尾 in-flight
    rss_samples.append(_rss_mb())
    await client.aclose()
    result = {"stats": stats, "rss_current_mb": rss_samples, "telemetry": str(telemetry_path)}
    (_ROOT / f"data/soak/result_{run_id}_{worker_idx}.json").write_text(
        json.dumps(result), encoding="utf-8"
    )
    print(f"[w{worker_idx}] 完成: {stats}", flush=True)


def _worker_main(args: argparse.Namespace, worker_idx: int) -> None:
    asyncio.run(_worker_async(args, worker_idx))


def _scoreboard(args: argparse.Namespace, env: dict[str, str]) -> None:
    """跑后记分: 合并 worker 遥测,断言硬不变量,产出报告。"""
    from tools.soak import scoreboard as sb

    results = []
    for w in range(args.workers):
        path = _ROOT / f"data/soak/result_{args.run_id}_{w}.json"
        results.append(json.loads(path.read_text(encoding="utf-8")))
    rows = sb.load_rows(*(r["telemetry"] for r in results))
    calls = sum(r["stats"]["calls"] for r in results)
    max_attempts = int(
        env.get(f"{args.scope}__RETRY__MAX_ATTEMPTS", env.get("LLM_MAX_RETRIES", "3"))
    )
    verdicts: list[tuple[str, str]] = []

    def _check(name: str, fn, *fargs, **fkwargs) -> None:
        try:
            fn(*fargs, **fkwargs)
            verdicts.append((name, "PASS"))
        except AssertionError as exc:
            verdicts.append((name, f"FAIL — {exc}"))

    # 行数下界=请求数(每请求至少 1 行);上界 = 逐次尝试放大(每请求 ≤ max_attempts 行)
    # + 终态行(scope 级失败/取消由 TelemetryMW 额外记一行,P5 全失败跑实证)
    terminal_rows = sum(r["stats"]["failed"] + r["stats"]["cancelled"] for r in results)
    _check(
        "遥测完备(行数≥请求数,重试容差内)",
        sb.inv_rows_match_calls,
        rows,
        expected_calls=calls,
        tolerance=calls * max(max_attempts - 1, 0) + terminal_rows,
    )
    _check("call_id 唯一", sb.inv_call_ids_unique, rows)
    rpm_conf = {}
    for key, val in env.items():
        parts = key.split("__")
        # 仅四段源键 {SCOPE}__{PROVIDER}__{N}__RPM(排除 {SCOPE}__GLOBAL__RPM)
        if len(parts) == 4 and parts[0] == args.scope and parts[3] == "RPM":
            rpm_conf[f"{parts[1].lower()}_{parts[2]}"] = int(val)
    # RPM 窗口口径 = 限流器所用时钟: redis 后端为服务器钟(需测偏移),memory 为本机钟
    clock_offset_s = 0.0
    if env.get("PGW_LIMITER_BACKEND") == "redis":
        clock_offset_s = asyncio.run(_redis_clock_offset_s(env["REDIS_URL"]))
        print(f"Redis 服务器钟偏移: {clock_offset_s:+.2f}s(RPM 窗口口径校正)")
    _check(
        "RPM 从未击穿(限流器窗口口径)",
        sb.inv_rpm_never_exceeded,
        rows,
        rpm_conf,
        clock_offset_s=clock_offset_s,
    )
    for r in results:
        samples = r.get("rss_current_mb")
        if samples is None:
            # 2026-07-21 前的旧结果只存 ru_maxrss 峰值序列,不可判泄漏
            verdicts.append(("RSS 平稳(w)", "SKIP — 旧格式峰值采样不可判(P6 伪击穿教训)"))
        else:
            _check("RSS 平稳(w)", sb.inv_rss_stable, samples, max_growth_mb=args.max_rss_growth_mb)
    # 不变量 1(记账归零 + gate 可再准入): 活后端检查,仅 redis 后端可跨进程复查
    if env.get("PGW_LIMITER_BACKEND") == "redis":
        asyncio.run(_live_checks(args, env, verdicts))
    else:
        verdicts.append(("记账归零/gate 可再准入", "SKIP — memory 后端跨进程不可查"))
    # 不变量 2c: P5/P6 故障混编池下错误必须真实发生。紧闸源被限流闸跳过
    # 不产错误行、故障名单无法从配置泛化推断 → 钉全局存在性,按源归因
    # 交报告"错误分布"人工核对(findings §4 条 2 的比例校准留待 P5 常态化)
    if args.scenario in ("P5", "P6"):
        _check("故障混编生效(P5/P6)", sb.inv_any_errors, rows)
    if args.scenario == "P7":
        _p7_checks(args, env, results, rows, calls, verdicts, _check)
    else:
        rate = sb.structured_success_rate(rows, session_suffix="-p3")
        verdicts.append(("P3 结构化成功率", f"{rate:.3f}(基线首跑建立)"))
    report = sb.render_report(args.run_id, rows, verdicts)
    path = sb.write_report(args.run_id, report)
    print(report)
    print(f"报告: {path}")
    if any(v.startswith("FAIL") for _, v in verdicts):
        raise SystemExit("硬不变量被击穿,见报告")


def _p7_checks(args, env, results, rows, calls, verdicts, check) -> None:
    """P7 专属裁决(M3 计划 T8): 成功率/坏源压制/故障生效/错误分类/取消泄漏/一致性。"""
    from tools.soak import scoreboard as sb

    ok_total = sum(r["stats"]["ok"] for r in results)
    check("成功率 ≥98%(P7)", sb.inv_success_rate, ok_total, calls, min_rate=0.98)
    fault_names = [
        s.strip() for s in env.get(f"{args.scope}_FAULT_SOURCES", "").split(",") if s.strip()
    ]
    check("坏源尝试占比 ≤15%(P7)", sb.inv_fault_share, rows, fault_names, max_share=0.15)
    check("故障混编生效(P7)", sb.inv_fault_errors_present, rows, fault_source_names=fault_names)
    check("错误全部可分类(P7)", sb.inv_errors_classified, rows, _KNOWN_ERROR_PREFIXES)
    for r in results:
        check("RSS 绝对值 <500MB(P7)", sb.inv_rss_absolute, r["rss_current_mb"], max_mb=500.0)
    cancelled = sum(r["stats"]["cancelled"] for r in results)
    verdicts.append(
        ("零取消泄漏(P7)", "PASS" if cancelled == 0 else f"FAIL — cancelled={cancelled}")
    )
    asyncio.run(_ocr_consistency_check(env, args.scope, verdicts))


# P7 已知错误前缀(metric ocr-call-success 口径: 异常类名 + 取消哨兵)
_KNOWN_ERROR_PREFIXES = (
    "TransientError",
    "SourceDeadError",
    "RequestRejectedError",
    "ResultInvalidError",
    "AllSourcesExhausted",
    "CircuitOpenError",
    "GatewayUnavailableError",
    "cancelled",
)


async def _ocr_consistency_check(
    env: dict[str, str], scope: str, verdicts: list[tuple[str, str]], n: int = 20
) -> None:
    """P7 不变量⑥: tables/para_blocks 一致性抽查(跑后活检查,设计 §1.2 护栏)。"""
    import io
    import random
    import zipfile

    import httpx

    from polygateway.transports.monkey_ocr import _parse_middle_json

    base = env[f"{scope}__MONKEY__1__BASE_URL"]
    images = sorted((_ROOT / "data/soak/chs_images").glob("chs_*.jpg"))
    sample = random.sample(images, min(n, len(images)))
    mismatches: list[str] = []
    async with httpx.AsyncClient(base_url=base, trust_env=False, timeout=300) as client:
        for path in sample:
            resp = await client.post(
                "/parse", files={"file": ("image.jpg", path.read_bytes(), "image/jpeg")}
            )
            resp.raise_for_status()
            zip_resp = await client.get(resp.json()["download_url"])
            zip_resp.raise_for_status()
            elements, _ = _parse_middle_json(zip_resp.content)
            lib = sorted(e.bbox for e in elements if e.type == "table")
            with zipfile.ZipFile(io.BytesIO(zip_resp.content)) as archive:
                member = next(m for m in archive.namelist() if m.endswith("_middle.json"))
                payload = json.load(archive.open(member))
            raw = sorted(
                tuple(float(v) for v in t["bbox"])
                for page in payload["pdf_info"]
                for t in page.get("tables", [])
            )
            if lib != raw:
                mismatches.append(path.name)
    verdict = "PASS" if not mismatches else f"FAIL — 不一致: {mismatches[:5]}"
    verdicts.append((f"tables/para_blocks 一致性抽查 n={len(sample)}", verdict))


async def _redis_clock_offset_s(redis_url: str) -> float:
    """Redis 服务器钟相对本机钟的偏移(秒);限流窗口 id 以服务器钟为准。"""
    import redis.asyncio as aioredis

    client = aioredis.from_url(redis_url)
    try:
        sec, usec = await client.time()
        return float(sec) + float(usec) / 1e6 - time.time()
    finally:
        await client.aclose()


async def _live_checks(
    args: argparse.Namespace, env: dict[str, str], verdicts: list[tuple[str, str]]
) -> None:
    """不变量 1 活后端复查(跑后): 限流 inflight 归零 + 熔断门可再准入。"""
    from polygateway.backends.redis.breaker import RedisGate
    from polygateway.backends.redis.limiter import RedisLimiter
    from polygateway.config import GatewaySettings
    from tools.soak import scoreboard as sb

    settings = GatewaySettings.from_env(args.scope, env=env)
    names = [s.name for s in settings.sources]
    limiter = RedisLimiter.from_url(
        env["REDIS_URL"],
        scope=settings.scope,
        sources={s.name: s for s in settings.sources},
        global_limits=settings.global_limits,
        lease_ttl_s=settings.lease_ttl_s,
    )
    gate = RedisGate.from_url(env["REDIS_URL"], config=settings.breaker, scope=settings.scope)
    checks = [
        ("记账归零(inflight)", sb.inv_accounting_zeroed(limiter, names)),
        ("gate 可再准入(探针不悬挂)", sb.inv_gate_reenterable(gate, names)),
    ]
    if args.scenario == "P7":
        fault = {
            s.strip() for s in env.get(f"{args.scope}_FAULT_SOURCES", "").split(",") if s.strip()
        }
        healthy = [n for n in names if n not in fault]
        checks.append(("真源门态 CLOSED(P7 零误熔)", sb.inv_healthy_gates_closed(gate, healthy)))
    try:
        for name, invariant in checks:
            try:
                await invariant
                verdicts.append((name, "PASS"))
            except AssertionError as exc:
                verdicts.append((name, f"FAIL — {exc}"))
    finally:
        await limiter.aclose()
        await gate.aclose()


def main() -> None:
    parser = argparse.ArgumentParser(description="PolyGateway 真实数据压测")
    parser.add_argument(
        "--scenario", required=True, choices=["P1", "P2", "P3", "P4", "P5", "P6", "P7"]
    )
    parser.add_argument("--budget-calls", type=int)
    parser.add_argument("--budget-tokens", type=int)
    parser.add_argument("--workers", type=int, default=1)
    parser.add_argument("--concurrency", type=int, default=8, help="单 worker 并发上限")
    parser.add_argument("--scope", default="SOAK")
    parser.add_argument("--run-id", default=None)
    parser.add_argument("--max-hours", type=float, default=3.0)
    parser.add_argument("--max-rss-growth-mb", type=float, default=500.0)
    parser.add_argument(
        "--rescore",
        metavar="RUN_ID",
        default=None,
        help="对已完成 run 离线重跑记分板(不产流量、不 FLUSHDB;须带原 env 覆盖)",
    )
    args = parser.parse_args()
    env = _merged_env()
    if args.rescore:
        args.run_id = args.rescore
        found = sorted((_ROOT / "data/soak").glob(f"result_{args.run_id}_*.json"))
        if not found:
            raise SystemExit(f"拒跑: 找不到 data/soak/result_{args.run_id}_*.json")
        args.workers = len(found)
        _guard(env, args.workers, args.scope)
        print(
            "rescore: 记账归零/gate 两项活检查反映**当前** db3 状态;若 db3 已被后续 run 复用请忽略"
        )
        _scoreboard(args, env)
        return
    if args.budget_calls is None or args.budget_tokens is None:
        parser.error("--budget-calls/--budget-tokens 为实跑必填")
    if args.run_id is None:
        args.run_id = time.strftime("soak_%Y%m%d_%H%M%S")
    _guard(env, args.workers, args.scope)
    print(f"run_id={args.run_id}: FLUSHDB db3 + namespace 隔离")
    asyncio.run(_flush_db3(env["REDIS_URL"]))
    if args.workers == 1:
        _worker_main(args, 0)
    else:
        ctx = mp.get_context("spawn")
        procs = [ctx.Process(target=_worker_main, args=(args, w)) for w in range(args.workers)]
        for p in procs:
            p.start()
        for p in procs:
            p.join()
        if any(p.exitcode != 0 for p in procs):
            raise SystemExit(f"worker 退出码异常: {[p.exitcode for p in procs]}")
    _scoreboard(args, env)


if __name__ == "__main__":
    main()