From a426cbbd21f8ba7cf21f4b50e2c7ada10d4a21d7 Mon Sep 17 00:00:00 2001 From: iomgaa Date: Tue, 11 Aug 2026 09:02:53 -0400 Subject: [PATCH] =?UTF-8?q?feat(soak):=20=E5=8E=8B=E6=B5=8B=E5=85=A5?= =?UTF-8?q?=E5=8F=A3=E2=80=94=E2=80=94=E9=A2=84=E7=AE=97=E6=8A=A4=E6=A0=8F?= =?UTF-8?q?=E3=80=81=E5=B9=B6=E5=8F=91=E4=B8=8A=E9=99=90=E3=80=81=E5=B9=B2?= =?UTF-8?q?=E8=B7=91=EF=BC=8C=E4=BB=A5=E5=8F=8A=E8=AE=B0=E5=88=86=E6=9D=BF?= =?UTF-8?q?=E8=A6=81=E7=9A=84=E5=9B=9B=E4=B8=AA=E4=BA=A7=E7=89=A9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --budget-calls 与 --concurrency 都没有默认值：一个能跑飞的压测入口迟早会跑飞。达到调用 上限时停止派发新任务，但已经在跑的让它跑完——半路砍断会制造一批没有结束记录的日志，而 那和崩溃长得一模一样，会污染故障注入那边的判定。 --dry-run 不打模型，只把任务集列出来、把每个 RunRequest 真的装配一遍，确认容器起得来、 语料读得进、脱敏闸过得了。它也刻意不往 runs-dir 写东西，写了的话紧接着的全量会撞上 RunIdentityError。 两个场景同时跑时任务轮流排开而不是拼接：共用一份预算，拼接的话排在前面的场景会把预算 吃光，而报告看起来只是「因预算停在第 N 个任务」——一次只压了一半的跑长得像一次正常的跑。 测试里最有价值的一条是真的把产物喂给记分板：手工搭两个 GovDoc 任务共六次运行（真的走 polyloop.session.run，模型是写死的替身），再 import 记分板判定，验十一条不变量全绿。 两边的 sidecar 约定对不上的话这条会当场红。 实跑过一次干跑：两道 AppWorld 题各装配出 12.4k 字符上下文，GovDoc 两个任务六次运行全部 装配成功，脱敏替换 64 处，跑完零残留容器。 Co-Authored-By: Claude Opus 5 (1M context) --- tools/soak/run_soak.py | 1020 +++++++++++++++++++++++++++++ tools/soak/soak.sh | 144 ++++ tools/soak/tests/test_run_soak.py | 599 +++++++++++++++++ 3 files changed, 1763 insertions(+) create mode 100644 tools/soak/run_soak.py create mode 100644 tools/soak/soak.sh create mode 100644 tools/soak/tests/test_run_soak.py diff --git a/tools/soak/run_soak.py b/tools/soak/run_soak.py new file mode 100644 index 0000000..b1f7b45 --- /dev/null +++ b/tools/soak/run_soak.py @@ -0,0 +1,1020 @@ +"""压测入口：把 AppWorld 与 GovDoc 两个场景跑成一批真实负载，并留下记分板要读的四个文件。 + +跑法（在仓库根目录）:: + + PYTHONUNBUFFERED=1 conda run --live-stream -n PolyLoop \\ + python -m tools.soak.run_soak --scenario both --budget-calls 400 --concurrency 4 \\ + --runs-dir soak-out/runs --report soak-out/soak.md + +用 `-m` 而不是直接给文件路径：直接跑文件时 `tools` 不在 `sys.path` 上， +`from tools.soak.appworld import ...` 会 ImportError。 + +每次运行落四个文件，文件名与字段由 `tools/soak/scoreboard.py` 的模块 docstring 定义，这里 +只负责写全。三个后缀常量直接从记分板 import，不在这边写第二份。 + +**预算与并发上限都没有默认值。** 一个能跑飞的压测入口迟早会跑飞：模型调用要花真钱，容器要 +占别人也在用的机器，而这两个数字是唯一能拦住它的东西。缺了直接拒跑，不猜。 + +**`--dry-run` 是全量之前的必经一步。** 它不打模型，但把每个任务的 `RunRequest` 真的装配 +一遍——容器起不起得来、语料读不读得进、脱敏闸过不过得了、提示词模板渲不渲染得出来，这几件 +事全都在装配这一步暴露，而在全量里暴露的代价是已经花掉的那部分调用费。 +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import logging +import sys +import time +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +from polyloop.serialization import encode +from polyloop.session import AgentDefinition +from polyloop.session import run as session_run +from polyloop.stores import JsonlRunStore +from tools.soak.appworld import AppWorldPool +from tools.soak.scenarios import appworld as appworld_scenario +from tools.soak.scenarios import govdoc as govdoc_scenario +from tools.soak.scoreboard import EVENTS_SUFFIX, META_SUFFIX, RESULT_SUFFIX + +if TYPE_CHECKING: + from polyloop.ports import Event, ModelCall + from polyloop.session import RunRequest + from polyloop.types import ModelReply, RunResult + +#: 场景名。它原样进 `.meta.json` 的 `scenario`，记分板按它分组统计。 +APPWORLD = "appworld" +GOVDOC = "govdoc" + +#: 传给每次模型调用的项目侧标识。**空的**：网关适配器只往下转发 `session_id` 与 +#: `parent_call_id`，而这两样每批都不同；进了参数快照，故障注入那一步续跑时的逐字段比对 +#: 就会报一次假的参数漂移。 +MODEL_BINDING: Mapping[str, str] = {} + +#: GovDoc 的工作区落在 runs 目录下的这个子目录里。记分板枚举 run 用的是 +#: `runs_dir.glob("*.jsonl")`，那不是递归的，所以子目录里的东西它看不见。 +WORKSPACES_DIRNAME = "workspaces" + +_LOG = logging.getLogger("polyloop.soak") + + +def _say(message: str) -> None: + """往 stdout 打一行。压测是长跑，缓冲住的输出等于没有输出。""" + print(message, flush=True) + + +# --------------------------------------------------------------------------- +# 一、预算护栏 +# --------------------------------------------------------------------------- + + +class BudgetGuard: + """包住一个 `ModelClient`，数每次调用；达到上限之后不再派发新任务。 + + **计数在真正发起调用之前加。** 数的是「尝试」而不是「成功」：一次失败的调用同样占了时间、 + 可能也已经在网关那边计了费，按成功数会让上限形同虚设。 + + **达到上限不砍断在跑的运行。** 半路砍断会制造一批没有结束记录的日志，而那和进程崩溃留下 + 的日志长得一模一样——故障注入那一步正是靠「有没有结束记录」判定的，混进来的话两边分不开。 + + `parameters()` 原样转发内层客户端的返回。它进参数快照，包一层不该改变快照的内容，否则 + 续跑时逐字段比对会报一次假漂移。 + """ + + def __init__(self, *, inner: object, limit: int) -> None: + """Args: + inner: 真正打模型的客户端，满足 `polyloop.ports.ModelClient`。 + limit: 整批允许的模型调用次数上限，必须 ≥ 1。 + """ + if limit < 1: + raise ValueError(f"模型调用预算必须 ≥ 1，收到 {limit}") + self._inner = inner + self._limit = limit + #: 并发安全靠这把锁。asyncio 单线程下 `+= 1` 本身不会被打断，但「读计数、判上限、 + #: 写计数」这三步之间有 await 点时就会，所以计数的更新整体放进锁里。 + self._lock = asyncio.Lock() + self._total = 0 + self._per_run: dict[str, int] = {} + + @property + def limit(self) -> int: + return self._limit + + @property + def total(self) -> int: + """整批已经发起的模型调用次数。""" + return self._total + + @property + def exhausted(self) -> bool: + return self._total >= self._limit + + def calls_for(self, run_id: str) -> int: + """某一次运行发起了几次模型调用。它进 `.meta.json` 的 `model_calls`。""" + return self._per_run.get(run_id, 0) + + def parameters(self) -> Mapping[str, str]: + return self._inner.parameters() # type: ignore[attr-defined] + + async def call(self, call: ModelCall) -> ModelReply: + async with self._lock: + self._total += 1 + self._per_run[call.run_id] = self._per_run.get(call.run_id, 0) + 1 + return await self._inner.call(call) # type: ignore[attr-defined] + + +# --------------------------------------------------------------------------- +# 二、事件出口 +# --------------------------------------------------------------------------- + + +class JsonlEventSink: + """把事件逐行写进 `.events.jsonl`，并自己数投递失败次数。 + + **自己数是有意义的**：库也数一份（`RunResult.event_delivery_failures`），记分板比的就是 + 这两份对不对得上。两份都由库来数的话，这条不变量验的是库和它自己一致。 + + 一次运行一个实例、一个文件，所以不需要跨运行的锁。 + """ + + __slots__ = ("_failures", "_path") + + def __init__(self, *, path: Path) -> None: + self._path = Path(path) + self._failures = 0 + + @property + def failures(self) -> int: + return self._failures + + def parameters(self) -> Mapping[str, str]: + return {"kind": "jsonl_events"} + + async def emit(self, event: Event) -> None: + """写一行。失败先记账再原样抛出——库会接住它并计数，两边的数才对得上。 + + `CancelledError` 继承 `BaseException`，下面那个 `except` 接不到它（CLAUDE.md §1.6）。 + """ + line = json.dumps( + { + "kind": event.kind.value, + "run_id": event.run_id, + "step_idx": event.step.step_idx, + }, + ensure_ascii=False, + ) + try: + with self._path.open("a", encoding="utf-8") as handle: + handle.write(line + "\n") + except Exception: + self._failures += 1 + raise + + +# --------------------------------------------------------------------------- +# 三、一次运行：跑完并落四个文件 +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True, slots=True, kw_only=True) +class RunOutcome: + """一次运行在报告里的那一行。产物的权威是磁盘上那四个文件，这个只进报告。""" + + run_id: str + scenario: str + task_id: str + phase: str | None + stop_reason: str | None + steps: int + model_calls: int + wall_ms: int + error: str | None = None + + +#: 收尾取证：返回 `(env_executions, success)`。两样都由环境侧给，取不到就是 `None`。 +Collector = Callable[[], Awaitable[tuple[int | None, bool | None]]] + + +async def execute_run( + *, + runs_dir: Path, + request: RunRequest, + model_client: object, + decision_parser: object, + synthetic_observations: object, + scenario: str, + task_id: str, + phase: str | None, + count_model_calls: Callable[[], int], + collect: Collector | None = None, + fault: str | None = None, + resumed_from_step: int = 0, +) -> RunOutcome: + """跑一次运行，并把记分板要的四个文件写全。 + + 库自己写 `.jsonl`，另外三个由这里写。**`.meta.json` 无论跑成什么样都会写**： + 运行抛异常时它记下的是「这次跑到哪儿、花了多少次调用」，而缺文件在记分板那边只会变成 + 一条「无法判定」，等于什么都没记下来。 + + Args: + runs_dir: 四个文件落在哪儿。 + request: 已经装配好的运行请求，`run_id` 从它取。 + model_client: 满足 `polyloop.ports.ModelClient` 的客户端，通常是 `BudgetGuard`。 + decision_parser: 场景自己的决策解释器。 + synthetic_observations: 场景自己的三段合成观察。 + scenario: 场景名，原样进 `.meta.json`。 + task_id: 任务标识，原样进 `.meta.json`。 + phase: 阶段名，没有阶段的场景填 `None`。 + count_model_calls: 跑完之后问一次「这次运行发起了几次模型调用」。 + collect: 跑完之后的收尾取证。AppWorld 那一路必须在会话上下文退出之前调用，所以它是 + 个回调而不是返回值——调用时机由这里定，取什么由场景定。 + fault: 注入了哪个故障。正常负载填 `None`。 + resumed_from_step: 从第几步续跑。正常负载填 0。 + """ + runs_dir = Path(runs_dir) + runs_dir.mkdir(parents=True, exist_ok=True) + run_id = request.run_id + sink = JsonlEventSink(path=runs_dir / f"{run_id}{EVENTS_SUFFIX}") + definition = AgentDefinition( + model_client=model_client, # type: ignore[arg-type] + decision_parser=decision_parser, # type: ignore[arg-type] + store=JsonlRunStore(directory=runs_dir), + event_sink=sink, + synthetic_observations=synthetic_observations, # type: ignore[arg-type] + ) + + started = time.monotonic() + result: RunResult | None = None + error: str | None = None + try: + result = await session_run(definition, request) + except Exception as exc: + # 一次运行炸掉不该带走整批。它进报告、进 `.meta.json`，然后接着跑下一个任务。 + # `CancelledError` 不在这里（它是 BaseException），Ctrl-C 照样穿得过去。 + error = f"{type(exc).__name__}: {exc}" + wall_ms = int((time.monotonic() - started) * 1000) + + env_executions: int | None = None + success: bool | None = None + if collect is not None: + try: + env_executions, success = await collect() + except Exception as exc: + note = f"收尾取证失败 {type(exc).__name__}: {exc}" + error = note if error is None else f"{error}；{note}" + + if result is not None: + (runs_dir / f"{run_id}{RESULT_SUFFIX}").write_text( + json.dumps(encode(result), ensure_ascii=False), + encoding="utf-8", + ) + + meta = { + "scenario": scenario, + "task_id": task_id, + "phase": phase, + "wall_ms": wall_ms, + "model_calls": count_model_calls(), + "sink_failures": sink.failures, + "env_executions": env_executions, + "fault": fault, + "success": success, + "resumed_from_step": resumed_from_step, + } + (runs_dir / f"{run_id}{META_SUFFIX}").write_text( + json.dumps(meta, ensure_ascii=False, indent=2), + encoding="utf-8", + ) + + return RunOutcome( + run_id=run_id, + scenario=scenario, + task_id=task_id, + phase=phase, + stop_reason=result.stop_reason.value if result is not None else None, + steps=len(result.steps) if result is not None else 0, + model_calls=meta["model_calls"], # type: ignore[arg-type] + wall_ms=wall_ms, + error=error, + ) + + +# --------------------------------------------------------------------------- +# 四、两个场景各自的任务运行器 +# --------------------------------------------------------------------------- + + +def appworld_run_id(task_id: str) -> str: + """`appworld-<题目 ID>`。必须匹配 `[A-Za-z0-9._-]+`（`JsonlRunStore` 的约束）。""" + return f"appworld-{task_id}" + + +async def run_appworld_task( + *, + pool: AppWorldPool, + task_id: str, + app_descriptions: str, + runs_dir: Path, + guard: BudgetGuard, +) -> list[RunOutcome]: + """跑一道 AppWorld 题：租一个容器、装配、跑完、评分。 + + **`evaluate()` 在会话上下文退出之前调**：退出之后环境状态就销毁了，那时再问分数只会拿到 + 一个异常。所以它被放进 `collect` 回调里，由 `execute_run` 在 `run` 返回之后立刻调用，而 + 整段都还在 `async with pool.session(...)` 里面。 + """ + run_id = appworld_run_id(task_id) + async with pool.session(task_id) as session: + request = appworld_scenario.build_run_request( + run_id=run_id, + session=session, + app_descriptions=app_descriptions, + model_binding=MODEL_BINDING, + ) + executor = request.action_executor + + async def collect() -> tuple[int | None, bool | None]: + # 执行次数取自环境自己的计数器，不是本地计数器——见 `AppWorldExecutor` 那条注释。 + executions = ( + executor.env_executions + if isinstance(executor, appworld_scenario.AppWorldExecutor) + else None + ) + score = await session.evaluate() + return executions, score.success + + outcome = await execute_run( + runs_dir=runs_dir, + request=request, + model_client=guard, + decision_parser=appworld_scenario.AppWorldParser(), + synthetic_observations=appworld_scenario.build_synthetic_observations(), + scenario=APPWORLD, + task_id=task_id, + phase=None, + count_model_calls=lambda: guard.calls_for(run_id), + collect=collect, + ) + return [outcome] + + +async def run_govdoc_task( + *, + task: govdoc_scenario.AuditTask, + runs_dir: Path, + workspace_root: Path, + guard: BudgetGuard, + phases: Sequence[str] = govdoc_scenario.PHASES, +) -> list[RunOutcome]: + """跑一个 GovDoc 任务：同一条审核点上依次跑三个阶段。 + + **三个阶段必须串行。** execute 要读 plan 写下的 `plan.md`，summarize 要读 execute 写下的 + `evidence.md`；并发跑的话后一阶段读到的是一个还不存在的文件，而表现是模型「找不到计划」 + 然后自己瞎编一个——那不是压测想看的形态。任务之间才是并发的，由派发器管。 + """ + workspace = Path(workspace_root) / f"govdoc-{task.index}" + workspace.mkdir(parents=True, exist_ok=True) + outcomes: list[RunOutcome] = [] + for phase in phases: + run_id = govdoc_scenario.make_run_id(task_index=task.index, phase=phase) + # 审计账是三个阶段共用的，所以这一次运行的执行次数是它的增量，不是总行数。 + before = len(govdoc_scenario.read_audit_lines(workspace)) + request = govdoc_scenario.build_run_request( + task=task, + phase=phase, + run_id=run_id, + workspace=workspace, + model_binding=MODEL_BINDING, + ) + + async def collect(baseline: int = before) -> tuple[int | None, bool | None]: + # GovDoc 没有程序化判分，`success` 恒为 None——填一个我们自己算的分数，会让记分板 + # 的成功率变成「我们和我们自己一致」。 + return len(govdoc_scenario.read_audit_lines(workspace)) - baseline, None + + outcomes.append( + await execute_run( + runs_dir=runs_dir, + request=request, + model_client=guard, + decision_parser=govdoc_scenario.GovDocParser(), + synthetic_observations=govdoc_scenario.SYNTHETIC_OBSERVATIONS, + scenario=GOVDOC, + task_id=task.checkpoint.checkpoint_id, + phase=phase, + count_model_calls=lambda run_id=run_id: guard.calls_for(run_id), + collect=collect, + ) + ) + return outcomes + + +# --------------------------------------------------------------------------- +# 五、派发：并发上限与预算护栏 +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True, slots=True, kw_only=True) +class Unit: + """一个派发单位。AppWorld 一道题是一个单位（一次运行），GovDoc 一个任务是一个单位 + （三次运行）。并发上限数的是单位，预算护栏也按单位为粒度决定派不派。""" + + label: str + scenario: str + run: Callable[[], Awaitable[Sequence[RunOutcome]]] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class BatchReport: + """一批跑完之后的全部事实。""" + + outcomes: tuple[RunOutcome, ...] = () + #: 单位级的失败：`(单位名, 错误)`。运行级的失败记在 `RunOutcome.error` 上。 + failures: tuple[tuple[str, str], ...] = () + total_units: int = 0 + dispatched: int = 0 + #: 因为预算耗尽而没有派出去的第一个单位的序号（从 1 数）。没停就是 `None`。 + stopped_at: int | None = None + + +async def dispatch( + units: Sequence[Unit], + *, + concurrency: int, + guard: BudgetGuard | None = None, +) -> BatchReport: + """按并发上限派发，按预算护栏决定还派不派。 + + **信号量在派发之前拿。** 拿到槽位之后才判预算，判的是「此刻」的调用数——先建好全部任务 + 再让它们自己抢槽位的话，预算判定会在一瞬间全部通过，护栏等于不存在。 + + **单个单位抛异常不打断整批**：记进 `failures`，接着派下一个。`CancelledError` 是 + `BaseException`，下面的 `except Exception` 接不到它；外层那个 `except BaseException` 只做 + 一件事——把还在跑的任务收干净，然后原样抛出去，Ctrl-C 因此穿得过整个编排（CLAUDE.md §1.6）。 + """ + if concurrency < 1: + raise ValueError(f"并发上限必须 ≥ 1，收到 {concurrency}") + + semaphore = asyncio.Semaphore(concurrency) + outcomes: list[RunOutcome] = [] + failures: list[tuple[str, str]] = [] + tasks: list[asyncio.Task[None]] = [] + stopped_at: int | None = None + + async def worker(unit: Unit) -> None: + try: + produced = await unit.run() + outcomes.extend(produced) + for item in produced: + mark = "失败" if item.error else item.stop_reason or "无结果" + _say(f" ✓ {item.run_id} {mark} {item.steps} 步 {item.wall_ms} ms") + except Exception as exc: + failures.append((unit.label, f"{type(exc).__name__}: {exc}")) + _say(f" ✗ {unit.label} {type(exc).__name__}: {exc}") + finally: + semaphore.release() + + try: + for index, unit in enumerate(units, start=1): + await semaphore.acquire() + if guard is not None and guard.exhausted: + semaphore.release() + stopped_at = index + _say( + f"预算用完（已发起 {guard.total} 次模型调用，上限 {guard.limit}），" + f"停在第 {index} 个任务；已经在跑的会跑完" + ) + break + _say(f"派发 {index}/{len(units)}：{unit.label}") + tasks.append(asyncio.create_task(worker(unit), name=f"soak-{unit.label}")) + if tasks: + await asyncio.gather(*tasks) + except BaseException: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + raise + + return BatchReport( + outcomes=tuple(sorted(outcomes, key=lambda item: item.run_id)), + failures=tuple(failures), + total_units=len(units), + dispatched=len(tasks), + stopped_at=stopped_at, + ) + + +# --------------------------------------------------------------------------- +# 六、任务集 +# --------------------------------------------------------------------------- + + +def select_appworld_tasks( + pool: AppWorldPool, *, splits: Sequence[str], limit: int | None +) -> list[str]: + """按给定顺序依次取每个划分的题目，跨划分去重，取到 `limit` 个为止。 + + 去重是必要的：`train` 与 `dev` 各自的题目 ID 不重叠，但划分列表是人给的，同一个划分给 + 两次的话没有去重就会把同一道题跑两遍——而两次的 `run_id` 一样，第二次会撞上「这个运行 + 标识已经有日志了」直接失败。 + """ + ordered: list[str] = [] + seen: set[str] = set() + for split in splits: + for task_id in pool.list_task_ids(split): + if task_id in seen: + continue + seen.add(task_id) + ordered.append(task_id) + if limit is not None and len(ordered) >= limit: + return ordered + return ordered + + +def load_govdoc_tasks( + *, db_path: Path, corpus_dir: Path, count: int +) -> tuple[tuple[govdoc_scenario.AuditTask, ...], tuple[govdoc_scenario.RedactionResult, ...]]: + """读前 `count` 条审核点、读语料、脱敏、过闸，装成任务列表。 + + **不走 `govdoc.build_audit_tasks`**：那个函数只收一个 `data_root`，从它推出库文件与语料 + 目录两个位置，而这个入口的 `--govdoc-db` 与 `--govdoc-corpus` 是两条独立的路径，表达不了。 + 做的事情逐字相同，脱敏与过闸都由这里调的那两个函数完成。 + """ + redactor = govdoc_scenario.Redactor() + checkpoints = govdoc_scenario.load_checkpoints(db_path=Path(db_path), limit=count) + loaded = govdoc_scenario.load_documents(prepared_dir=Path(corpus_dir), redactor=redactor) + documents = tuple(document for document, _ in loaded) + reports = tuple(report for _, report in loaded) + tasks = tuple( + govdoc_scenario.AuditTask(index=index, checkpoint=checkpoint, documents=documents) + for index, checkpoint in enumerate(checkpoints) + ) + return tasks, reports + + +def interleave(groups: Sequence[Sequence[Unit]]) -> list[Unit]: + """把几个场景的单位轮流排开。 + + **不是简单拼接。** 拼接的话整批共用的模型调用预算会被排在前面的场景吃光，排在后面的那个 + 一个任务都跑不到，而报告看起来只是「因预算停在第 N 个任务」——一次只压了一半的跑，长得 + 像一次正常的跑。 + """ + merged: list[Unit] = [] + for row in range(max((len(group) for group in groups), default=0)): + for group in groups: + if row < len(group): + merged.append(group[row]) + return merged + + +# --------------------------------------------------------------------------- +# 七、干跑 +# --------------------------------------------------------------------------- + + +async def dry_run_appworld(args: argparse.Namespace) -> list[str]: + """起容器、取 app 清单、逐题装配一次 `RunRequest`。装配失败原样抛出去。""" + notes: list[str] = [] + pool = AppWorldPool( + data_root=args.appworld_data_root, + size=args.containers, + port_base=args.port_base, + ) + task_ids = select_appworld_tasks(pool, splits=args.split, limit=args.limit) + _say(f"AppWorld：划分 {list(args.split)} 去重后取 {len(task_ids)} 道题") + if not task_ids: + raise SystemExit("AppWorld 的题目集是空的，检查 --split 与 --limit") + notes.append(f"AppWorld 任务 {len(task_ids)} 道，容器 {args.containers} 个") + + async with pool: + descriptions = await appworld_scenario.load_app_descriptions(pool, task_id=task_ids[0]) + notes.append(f"app 清单 {len(descriptions)} 字符") + _say(f"AppWorld：app 清单 {len(descriptions)} 字符") + + semaphore = asyncio.Semaphore(args.concurrency) + + async def assemble(task_id: str) -> None: + async with semaphore, pool.session(task_id) as session: + request = appworld_scenario.build_run_request( + run_id=appworld_run_id(task_id), + session=session, + app_descriptions=descriptions, + model_binding=MODEL_BINDING, + ) + chars = sum( + len(block.text) + for message in (*request.context.run_level, *request.context.goal_level) + for block in message.content + ) + _say(f" ✓ {request.run_id} 上下文 {chars} 字符") + + await asyncio.gather(*(assemble(task_id) for task_id in task_ids)) + return notes + + +def dry_run_govdoc(args: argparse.Namespace) -> list[str]: + """读语料、脱敏、过闸，逐任务逐阶段装配一次 `RunRequest`。""" + if args.limit is None: + raise SystemExit("GovDoc 场景必须给 --limit：审核点库有多少条不该由压测入口替人决定") + tasks, reports = load_govdoc_tasks( + db_path=args.govdoc_db, corpus_dir=args.govdoc_corpus, count=args.limit + ) + notes = [f"GovDoc 任务 {len(tasks)} 个 × {len(govdoc_scenario.PHASES)} 阶段"] + for report in reports: + summary = "、".join(f"{name} {count}" for name, count in sorted(report.counts.items())) + notes.append(f"脱敏替换：{summary or '无'}") + _say(f"GovDoc：脱敏替换 {summary or '无'}") + + workspace_root = Path(args.runs_dir) / WORKSPACES_DIRNAME + for task in tasks: + for phase in govdoc_scenario.PHASES: + request = govdoc_scenario.build_run_request( + task=task, + phase=phase, + run_id=govdoc_scenario.make_run_id(task_index=task.index, phase=phase), + workspace=workspace_root / f"govdoc-{task.index}", + model_binding=MODEL_BINDING, + ) + _say(f" ✓ {request.run_id} 工具 {list(request.tools.names())}") + return notes + + +async def dry_run(args: argparse.Namespace) -> int: + """干跑：列任务集、装配全部 `RunRequest`，一次模型都不打。 + + **不建网关客户端。** 干跑要能在一台没配 `.env` 的机器上跑起来——它验的是容器、语料、 + 脱敏闸与提示词装配，网关凭据不在这几件事里。凭据的问题会在全量的第一次调用上暴露。 + """ + _say("=== 干跑：不打模型 ===") + notes: list[str] = [] + if args.scenario in (APPWORLD, "both"): + notes.extend(await dry_run_appworld(args)) + if args.scenario in (GOVDOC, "both"): + notes.extend(dry_run_govdoc(args)) + + _say("\n=== 干跑通过 ===") + for note in notes: + _say(f" {note}") + _say("网关客户端没有装配（干跑不打模型），凭据要到全量的第一次调用才会被验证。") + + if args.report is not None: + body = "# 压测干跑\n\n" + "\n".join(f"- {note}" for note in notes) + "\n" + _write_report(Path(args.report), body) + return 0 + + +# --------------------------------------------------------------------------- +# 八、全量 +# --------------------------------------------------------------------------- + + +def build_model_client() -> tuple[object, Callable[[], Awaitable[None]]]: + """装配一个连着真实网关的模型客户端，并交回关它的办法。 + + **在函数里 import 网关**：`polyloop.adapters` 要 `polyloop[gateway]`，而干跑与这个模块的 + 测试都用不着它。放在模块顶层的话，没装网关的机器上连 `--help` 都跑不起来。 + """ + from polygateway import GatewayClient, GatewaySettings + + from polyloop.adapters import GatewayModelClient + + client = GatewayClient.from_env() + + async def close() -> None: + await client.aclose() + + return GatewayModelClient(client=client, settings=GatewaySettings.from_env()), close + + +async def full_run(args: argparse.Namespace) -> int: + """跑正常负载。""" + runs_dir = Path(args.runs_dir) + runs_dir.mkdir(parents=True, exist_ok=True) + workspace_root = runs_dir / WORKSPACES_DIRNAME + + model_client, close_client = build_model_client() + guard = BudgetGuard(inner=model_client, limit=args.budget_calls) + started = time.monotonic() + + pool: AppWorldPool | None = None + groups: list[list[Unit]] = [] + notes: list[str] = [] + try: + if args.scenario in (APPWORLD, "both"): + candidate = AppWorldPool( + data_root=args.appworld_data_root, + size=args.containers, + port_base=args.port_base, + ) + task_ids = select_appworld_tasks(candidate, splits=args.split, limit=args.limit) + if not task_ids: + raise SystemExit("AppWorld 的题目集是空的，检查 --split 与 --limit") + # `start()` 自己保证「要么全起要么全拆」，所以只有起成功之后才把池交给 finally + # 去停——没起过的池不需要停，而对它调 stop() 会打出一串看起来像故障的清理日志。 + await candidate.start() + pool = candidate + descriptions = await appworld_scenario.load_app_descriptions(pool, task_id=task_ids[0]) + notes.append(f"AppWorld：{len(task_ids)} 道题，容器 {args.containers} 个") + groups.append( + [ + Unit( + label=f"appworld/{task_id}", + scenario=APPWORLD, + run=_appworld_unit( + pool=pool, + task_id=task_id, + app_descriptions=descriptions, + runs_dir=runs_dir, + guard=guard, + ), + ) + for task_id in task_ids + ] + ) + + if args.scenario in (GOVDOC, "both"): + if args.limit is None: + raise SystemExit("GovDoc 场景必须给 --limit：一个任务是三次运行，预算要按它算") + tasks, reports = load_govdoc_tasks( + db_path=args.govdoc_db, corpus_dir=args.govdoc_corpus, count=args.limit + ) + for report in reports: + summary = "、".join( + f"{name} {count}" for name, count in sorted(report.counts.items()) + ) + notes.append(f"GovDoc 脱敏替换：{summary or '无'}") + notes.append( + f"GovDoc：{len(tasks)} 个任务 × {len(govdoc_scenario.PHASES)} 阶段 = " + f"{len(tasks) * len(govdoc_scenario.PHASES)} 次运行" + ) + groups.append( + [ + Unit( + label=f"govdoc/{task.index}", + scenario=GOVDOC, + run=_govdoc_unit( + task=task, + runs_dir=runs_dir, + workspace_root=workspace_root, + guard=guard, + ), + ) + for task in tasks + ] + ) + + units = interleave(groups) + _say(f"=== 正常负载：{len(units)} 个任务，并发 {args.concurrency}，预算 {guard.limit} ===") + report = await dispatch(units, concurrency=args.concurrency, guard=guard) + finally: + if pool is not None: + await pool.stop() + await close_client() + + body = render_report( + report, + guard=guard, + notes=notes, + runs_dir=runs_dir, + wall_ms=int((time.monotonic() - started) * 1000), + ) + _say("\n" + body) + if args.report is not None: + _write_report(Path(args.report), body) + # **退出码只看任务级失败，不看运行级失败。** 一次以 `llm_error` 或 `step_budget` 结束的 + # 运行是压测的正常产出，判它是好是坏是记分板的事；任务级失败则是异常逃出了整次运行 + # （容器租不到、请求装配不出来），那说明这套 harness 本身有问题，值得让 soak.sh 停下来。 + return 1 if report.failures else 0 + + +def _appworld_unit( + *, + pool: AppWorldPool, + task_id: str, + app_descriptions: str, + runs_dir: Path, + guard: BudgetGuard, +) -> Callable[[], Awaitable[Sequence[RunOutcome]]]: + async def start() -> Sequence[RunOutcome]: + return await run_appworld_task( + pool=pool, + task_id=task_id, + app_descriptions=app_descriptions, + runs_dir=runs_dir, + guard=guard, + ) + + return start + + +def _govdoc_unit( + *, + task: govdoc_scenario.AuditTask, + runs_dir: Path, + workspace_root: Path, + guard: BudgetGuard, +) -> Callable[[], Awaitable[Sequence[RunOutcome]]]: + async def start() -> Sequence[RunOutcome]: + return await run_govdoc_task( + task=task, + runs_dir=runs_dir, + workspace_root=workspace_root, + guard=guard, + ) + + return start + + +# --------------------------------------------------------------------------- +# 九、报告 +# --------------------------------------------------------------------------- + + +def _cell(text: str) -> str: + """markdown 表格单元格：竖线要转义，换行要压掉，否则整张表塌了。异常消息里两样都有。""" + return text.replace("|", "\\|").replace("\n", " ") + + +def _write_report(path: Path, body: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(body, encoding="utf-8") + _say(f"报告写到了 {path}") + + +def render_report( + report: BatchReport, + *, + guard: BudgetGuard, + notes: Iterable[str], + runs_dir: Path, + wall_ms: int, +) -> str: + """渲染一份 markdown 报告。判定不在这里——那是记分板的事，这里只报事实。""" + lines = ["# 压测正常负载", ""] + lines.append(f"产物目录：`{runs_dir}`") + lines.append("") + lines.append(f"- 任务：派发 {report.dispatched} / 共 {report.total_units}") + lines.append(f"- 运行：{len(report.outcomes)} 次") + lines.append(f"- 模型调用：{guard.total} / 预算 {guard.limit}") + lines.append(f"- 墙钟：{wall_ms} ms") + for note in notes: + lines.append(f"- {note}") + if report.stopped_at is not None: + lines.append( + f"- **因预算停在第 {report.stopped_at} 个任务**：" + f"它和它后面的 {report.total_units - report.stopped_at + 1} 个任务没有派发，" + "已经在跑的都跑完了" + ) + else: + lines.append("- 预算没有用完，全部任务都派发了") + + lines.extend( + [ + "", + "## 逐次运行", + "", + "| run_id | 场景 | 阶段 | 停止原因 | 步数 | 调用 | 墙钟 ms | 错误 |", + ] + ) + lines.append("|---|---|---|---|---:|---:|---:|---|") + for item in report.outcomes: + lines.append( + f"| `{item.run_id}` | {item.scenario} | {item.phase or '—'} | " + f"{item.stop_reason or '—'} | {item.steps} | {item.model_calls} | " + f"{item.wall_ms} | {_cell(item.error) if item.error else '—'} |" + ) + + lines.extend(["", "## 任务级失败", ""]) + if report.failures: + for label, error in report.failures: + lines.append(f"- `{label}`：{error}") + else: + lines.append("没有。") + lines.append("") + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# 十、命令行 +# --------------------------------------------------------------------------- + + +def _positive(text: str) -> int: + value = int(text) + if value < 1: + raise argparse.ArgumentTypeError(f"必须 ≥ 1，收到 {value}") + return value + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="run_soak", + description="跑一批正常负载，并留下记分板要读的四个文件。", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=( + "GovDoc 的 --limit 数的是**任务**，不是运行：一个任务 = 一条审核点 + 一份文书，" + "跑 plan / execute / summarize 三个阶段，也就是三次运行。--limit 20 是 60 次运行，" + "预算要按 60 次运行去估，不是 20 次。\n\n" + "AppWorld 的 --limit 数的是题目，一道题就是一次运行。\n\n" + "全量的形状：--split train --split dev --limit 100" + ), + ) + parser.add_argument( + "--scenario", + required=True, + choices=[APPWORLD, GOVDOC, "both"], + help="压哪个场景。both 时两个场景的任务轮流排开，共用同一份模型调用预算", + ) + parser.add_argument( + "--budget-calls", + required=True, + type=_positive, + help=( + "整批允许的模型调用次数上限。**没有默认值**：一个能跑飞的压测入口迟早会跑飞。" + "达到上限就不再派发新任务，已经在跑的会跑完，所以实际调用数会略微超出" + ), + ) + parser.add_argument( + "--concurrency", + required=True, + type=_positive, + help="同时在跑的任务数上限。**没有默认值**，缺了拒跑", + ) + parser.add_argument( + "--containers", type=_positive, default=4, help="AppWorld 的容器数（默认 4）" + ) + parser.add_argument( + "--split", + action="append", + default=None, + metavar="NAME", + help="AppWorld 的数据划分，可给多次；按给的顺序依次取，跨划分去重（默认 train）", + ) + parser.add_argument( + "--limit", + type=_positive, + default=None, + help="每个场景最多跑几个任务。AppWorld 不给就是整个划分；GovDoc 必须给", + ) + parser.add_argument("--runs-dir", required=True, type=Path, help="四个产物文件落在哪儿") + parser.add_argument("--report", type=Path, default=None, help="报告写到哪儿，不给就只打屏") + parser.add_argument( + "--dry-run", + action="store_true", + help=( + "不打模型：列出任务集、把每个任务的 RunRequest 真的装配一遍，" + "确认容器起得来、语料读得进、脱敏闸过得了，然后退出。全量之前的必经一步" + ), + ) + parser.add_argument( + "--appworld-data-root", + type=Path, + default=None, + help="AppWorld 数据根目录，下面应有 data/datasets 与 data/tasks（跑 AppWorld 时必填）", + ) + parser.add_argument( + "--port-base", + type=int, + default=8200, + help="AppWorld 容器池的起始宿主端口（默认 8200，避开 dissect 的 8100）", + ) + parser.add_argument( + "--govdoc-db", + type=Path, + default=govdoc_scenario.DEFAULT_DATA_ROOT / "app.sqlite", + help="GovDoc 的审核点 sqlite（只读打开）", + ) + parser.add_argument( + "--govdoc-corpus", + type=Path, + default=govdoc_scenario.DEFAULT_DATA_ROOT / "storage" / "prepared", + help="GovDoc 的语料目录", + ) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + parser = build_parser() + args = parser.parse_args(argv) + if args.split is None: + args.split = ["train"] + if args.scenario in (APPWORLD, "both") and args.appworld_data_root is None: + parser.error("--scenario 含 appworld 时必须给 --appworld-data-root") + + # 容器池的告警走 logging，默认级别是 WARNING 且没有 handler，会被静默丢掉——而 + # 「清理容器失败」正是那条路上唯一的线索。 + logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s") + _LOG.debug("参数：%s", vars(args)) + + if args.dry_run: + return asyncio.run(dry_run(args)) + return asyncio.run(full_run(args)) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/soak/soak.sh b/tools/soak/soak.sh new file mode 100644 index 0000000..fac843d --- /dev/null +++ b/tools/soak/soak.sh @@ -0,0 +1,144 @@ +#!/usr/bin/env bash +# +# 压测的四步：干跑 → 正常负载 → 故障注入 → 记分板。 +# +# 这个脚本只做「在当前 shell 里按顺序跑」，自己不建 tmux 会话、也不 attach。**长跑要在 tmux +# 里跑**（CLAUDE.md §4），所以正确的用法是人先开会话再调它： +# +# tmux new -s polyloop-soak +# SOAK_BUDGET_CALLS=4000 SOAK_FAULT_BUDGET_CALLS=600 SOAK_CONCURRENCY=4 \ +# SOAK_LIMIT=100 SOAK_APPWORLD_DATA_ROOT=/path/to/appworld \ +# bash tools/soak/soak.sh +# +# 会话名约定是 polyloop-soak。跑完不要急着 kill，留着给人复查。 +# +# **末尾一律不接管道。** `pytest ... | tail` 的退出码来自管道最后一节，于是一次失败的跑会 +# 报成 exit 0。要留日志就设 SOAK_LOG_DIR，那一路用重定向，不用 tee。 +# +# 全部配置都是环境变量，没有位置参数： +# +# 必填 +# SOAK_BUDGET_CALLS 正常负载的模型调用次数上限 +# SOAK_FAULT_BUDGET_CALLS 故障注入那一步的模型调用次数上限 +# SOAK_CONCURRENCY 同时在跑的任务数上限 +# SOAK_LIMIT 每个场景最多跑几个任务（GovDoc 一个任务是三次运行） +# SOAK_APPWORLD_DATA_ROOT AppWorld 数据根目录（SOAK_SCENARIO 含 appworld 时） +# +# 可选 +# SOAK_OUT 产物根目录，默认 soak-out/<时间戳> +# SOAK_SCENARIO appworld / govdoc / both，默认 both +# SOAK_CONTAINERS AppWorld 容器数，默认 4 +# SOAK_SPLITS AppWorld 划分，空格分隔，默认 "train dev" +# SOAK_GOVDOC_DB GovDoc 的审核点 sqlite，不给就用场景里的默认路径 +# SOAK_GOVDOC_CORPUS GovDoc 的语料目录，同上 +# SOAK_LOG_DIR 每一步的输出重定向到这里；不给就直接打屏（tmux 里能实时看） +# SOAK_SKIP_FAULTS 非空则跳过故障注入那一步 +# SOAK_CONDA_ENV conda 环境名，默认 PolyLoop + +set -euo pipefail + +STEP="启动" +trap 'echo "!!! 第「${STEP}」步失败（退出码 $?）" >&2' ERR + +CONDA_ENV="${SOAK_CONDA_ENV:-PolyLoop}" +SCENARIO="${SOAK_SCENARIO:-both}" +CONTAINERS="${SOAK_CONTAINERS:-4}" +SPLITS="${SOAK_SPLITS:-train dev}" +LOG_DIR="${SOAK_LOG_DIR:-}" +OUT="${SOAK_OUT:-soak-out/$(date +%Y%m%d-%H%M%S)}" + +BUDGET_CALLS="${SOAK_BUDGET_CALLS:?必须给 SOAK_BUDGET_CALLS：正常负载的模型调用次数上限}" +FAULT_BUDGET_CALLS="${SOAK_FAULT_BUDGET_CALLS:?必须给 SOAK_FAULT_BUDGET_CALLS：故障注入那一步的上限}" +CONCURRENCY="${SOAK_CONCURRENCY:?必须给 SOAK_CONCURRENCY：同时在跑的任务数上限}" +LIMIT="${SOAK_LIMIT:?必须给 SOAK_LIMIT：每个场景最多跑几个任务}" + +# conda 和 Python 各缓冲一层，两层都得拆——只加其中一个，长跑命令仍然全程无输出。 +RUN=(env PYTHONUNBUFFERED=1 conda run --live-stream -n "$CONDA_ENV" python) + +RUNS_DIR="$OUT/runs" +FAULT_RUNS_DIR="$OUT/fault-runs" + +# 场景相关的参数拼成数组。**用数组不用字符串**：路径里有空格时字符串会在展开时被切开， +# 而表现是「找不到这个目录」，看起来像数据没准备好。 +SCENARIO_ARGS=(--scenario "$SCENARIO" --limit "$LIMIT" --containers "$CONTAINERS") +if [[ "$SCENARIO" == "appworld" || "$SCENARIO" == "both" ]]; then + APPWORLD_DATA_ROOT="${SOAK_APPWORLD_DATA_ROOT:?SOAK_SCENARIO 含 appworld 时必须给 SOAK_APPWORLD_DATA_ROOT}" + SCENARIO_ARGS+=(--appworld-data-root "$APPWORLD_DATA_ROOT") + for split in $SPLITS; do + SCENARIO_ARGS+=(--split "$split") + done +fi +# 故障注入那一步的两个 GovDoc 路径没有默认值，而它们的权威在场景模块里。**问 Python 要， +# 不在这里写第二份**：两处各写一份路径，迟早有一处被改、另一处没改，而表现是「数据不在」。 +GOVDOC_DB="${SOAK_GOVDOC_DB:-$("${RUN[@]}" -c 'from tools.soak.scenarios.govdoc import DEFAULT_DATA_ROOT as R; print(R / "app.sqlite")')}" +GOVDOC_CORPUS="${SOAK_GOVDOC_CORPUS:-$("${RUN[@]}" -c 'from tools.soak.scenarios.govdoc import DEFAULT_DATA_ROOT as R; print(R / "storage" / "prepared")')}" +if [[ "$SCENARIO" == "govdoc" || "$SCENARIO" == "both" ]]; then + SCENARIO_ARGS+=(--govdoc-db "$GOVDOC_DB" --govdoc-corpus "$GOVDOC_CORPUS") +fi + +# 故障注入的参数是另一套：它自己的 `--split` 只收一个值，AppWorld 数据根目录那一项叫 +# `--data-root`。这里按 `python -m tools.soak.faults --help` 的形状拼。 +FAULT_ARGS=(--govdoc-db "$GOVDOC_DB" --govdoc-corpus "$GOVDOC_CORPUS") +if [[ "$SCENARIO" == "appworld" || "$SCENARIO" == "both" ]]; then + FAULT_ARGS+=(--data-root "$APPWORLD_DATA_ROOT" --split "${SPLITS%% *}") +fi + +mkdir -p "$OUT" + +step() { + local name="$1" + shift + STEP="$name" + echo "" + echo "=== [$name] $* ===" + if [[ -n "$LOG_DIR" ]]; then + mkdir -p "$LOG_DIR" + "$@" >"$LOG_DIR/$name.log" 2>&1 + else + "$@" + fi +} + +step 01-干跑 "${RUN[@]}" -m tools.soak.run_soak \ + "${SCENARIO_ARGS[@]}" \ + --budget-calls "$BUDGET_CALLS" \ + --concurrency "$CONCURRENCY" \ + --runs-dir "$RUNS_DIR" \ + --report "$OUT/dry-run.md" \ + --dry-run + +step 02-正常负载 "${RUN[@]}" -m tools.soak.run_soak \ + "${SCENARIO_ARGS[@]}" \ + --budget-calls "$BUDGET_CALLS" \ + --concurrency "$CONCURRENCY" \ + --runs-dir "$RUNS_DIR" \ + --report "$OUT/normal-load.md" + +# 正常负载不该有缺文件，所以这一轮不开 --allow-undetermined：判不了就是要人去看一眼。 +step 03-记分板-正常负载 "${RUN[@]}" -m tools.soak.scoreboard \ + --runs-dir "$RUNS_DIR" \ + --report "$OUT/scoreboard-normal.md" \ + --completing-tool submit_finding + +if [[ -n "${SOAK_SKIP_FAULTS:-}" ]]; then + echo "" + echo "=== [04-故障注入] 按 SOAK_SKIP_FAULTS 跳过 ===" +else + step 04-故障注入 "${RUN[@]}" -m tools.soak.faults \ + "${FAULT_ARGS[@]}" \ + --runs-dir "$FAULT_RUNS_DIR" \ + --budget-calls "$FAULT_BUDGET_CALLS" + + # 崩溃注入那一类天然会缺文件，这一轮才该开 --allow-undetermined。 + step 05-记分板-故障注入 "${RUN[@]}" -m tools.soak.scoreboard \ + --runs-dir "$FAULT_RUNS_DIR" \ + --report "$OUT/scoreboard-faults.md" \ + --completing-tool submit_finding \ + --allow-undetermined +fi + +STEP="收尾" +echo "" +echo "=== 全部步骤通过 ===" +echo "产物：$OUT" +ls -1 "$OUT" diff --git a/tools/soak/tests/test_run_soak.py b/tools/soak/tests/test_run_soak.py new file mode 100644 index 0000000..4968822 --- /dev/null +++ b/tools/soak/tests/test_run_soak.py @@ -0,0 +1,599 @@ +"""压测入口驱动的测试。**不打真实模型、不起容器**，模型那一侧全是写死脚本的替身。 + +分五块：参数校验、预算护栏与并发上限、产物与记分板的对接、错误隔离与取消、GovDoc 的阶段串行。 + +最有价值的一条是「产物喂给记分板判成全绿」。驱动写四个文件、记分板读四个文件，两边的字段 +约定只写在记分板的模块 docstring 里，没有任何机器约束把它们钉在一起——那条测试就是那个约束： +少写一个字段、类型写错一个，它当场红。 + +GovDoc 那几块用手工搭的 `AuditTask`，不读磁盘上的真实语料：那份数据是另一个项目的工作副本， +换一台机器就没有，而这里要验的是编排（阶段串行、任务并发），跟语料内容无关。 +""" + +from __future__ import annotations + +import asyncio +import json +from typing import TYPE_CHECKING + +import pytest + +from polyloop.types import ModelReply +from tools.soak import scoreboard +from tools.soak.run_soak import ( + BudgetGuard, + RunOutcome, + Unit, + dispatch, + execute_run, + interleave, + main, + run_govdoc_task, + select_appworld_tasks, +) +from tools.soak.scenarios import govdoc as govdoc_scenario + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + from pathlib import Path + + from polyloop.ports import ModelCall + +# --------------------------------------------------------------------------- +# 替身 +# --------------------------------------------------------------------------- + + +class _ScriptedModel: + """按 run_id 给回复的模型替身。每次调用都让出一次事件循环，好让并发真的发生。""" + + def __init__(self, *, render: Callable[[ModelCall], str]) -> None: + self._render = render + #: 调用到达的顺序，按 run_id 记。阶段串行与任务并发都从它上面判。 + self.seen: list[str] = [] + + async def call(self, call: ModelCall) -> ModelReply: + await asyncio.sleep(0.001) + self.seen.append(call.run_id) + return ModelReply( + call_id=f"call-{len(self.seen)}", + content=self._render(call), + thinking="", + ) + + def parameters(self): # noqa: ANN201 - 替身，签名由 Protocol 定 + return {"kind": "scripted"} + + +def _guard(*, render: Callable[[ModelCall], str], limit: int = 1000) -> BudgetGuard: + return BudgetGuard(inner=_ScriptedModel(render=render), limit=limit) + + +_SUBMIT = json.dumps( + { + "tool": "submit_finding", + "arguments": { + "verdict": "存疑", + "evidence": "tender.md 第 1 行：示例甲公司", + "reasoning": "测试替身给的固定结论", + }, + }, + ensure_ascii=False, +) + +#: 解析不出来的输出。GovDoc 的预算里连续解析失败上限是 3，所以三步就停——plan 与 execute +#: 两个阶段没有结束通路（提交工具只在 summarize 那一阶段），不这样它们会一路走到 50 步。 +_GARBAGE = "我先想一想这道题。" + + +def _audit_task(index: int) -> govdoc_scenario.AuditTask: + checkpoint = govdoc_scenario.Checkpoint( + checkpoint_id=f"cp-{index}", + category="资格条件", + title="供应商资格要求", + description="不得以注册地设置差别待遇。", + legal_basis=("政府采购法第五条",), + severity="高", + ) + document = govdoc_scenario.CorpusDocument.from_text( + logical_name="tender.md", + text="第一行：示例甲公司参与投标。\n第二行：投标截止时间为示例日期。", + ) + return govdoc_scenario.AuditTask(index=index, checkpoint=checkpoint, documents=(document,)) + + +def _fake_unit( + label: str, + *, + body: Callable[[], object], +) -> Unit: + async def start() -> Sequence[RunOutcome]: + result = body() + if asyncio.iscoroutine(result): + await result + return [ + RunOutcome( + run_id=label, + scenario="fake", + task_id=label, + phase=None, + stop_reason="task_completed", + steps=1, + model_calls=1, + wall_ms=0, + ) + ] + + return Unit(label=label, scenario="fake", run=start) + + +# --------------------------------------------------------------------------- +# 一、参数校验 +# --------------------------------------------------------------------------- + + +def test_missing_budget_calls_refuses(tmp_path: Path) -> None: + with pytest.raises(SystemExit) as caught: + main(["--scenario", "govdoc", "--concurrency", "2", "--runs-dir", str(tmp_path)]) + assert caught.value.code != 0 + + +def test_missing_concurrency_refuses(tmp_path: Path) -> None: + with pytest.raises(SystemExit) as caught: + main(["--scenario", "govdoc", "--budget-calls", "10", "--runs-dir", str(tmp_path)]) + assert caught.value.code != 0 + + +def test_non_positive_budget_refuses(tmp_path: Path) -> None: + with pytest.raises(SystemExit): + main( + [ + "--scenario", + "govdoc", + "--budget-calls", + "0", + "--concurrency", + "1", + "--runs-dir", + str(tmp_path), + ] + ) + + +def test_appworld_without_data_root_refuses(tmp_path: Path) -> None: + with pytest.raises(SystemExit): + main( + [ + "--scenario", + "appworld", + "--budget-calls", + "5", + "--concurrency", + "1", + "--runs-dir", + str(tmp_path), + "--dry-run", + ] + ) + + +# --------------------------------------------------------------------------- +# 二、预算护栏与并发上限 +# --------------------------------------------------------------------------- + + +async def test_budget_counts_are_per_run_and_total() -> None: + guard = _guard(render=lambda call: "ok") + calls = [_call("run-a"), _call("run-a"), _call("run-b")] + for item in calls: + await guard.call(item) + assert guard.total == 3 + assert guard.calls_for("run-a") == 2 + assert guard.calls_for("run-b") == 1 + assert guard.exhausted is False + + +async def test_budget_stops_dispatch_and_lets_running_tasks_finish() -> None: + guard = _guard(render=lambda call: "ok", limit=2) + + async def slow() -> None: + await guard.call(_call("slow")) + await asyncio.sleep(0.05) + + async def quick() -> None: + await guard.call(_call("quick")) + + units = [ + _fake_unit("unit-1", body=slow), + _fake_unit("unit-2", body=quick), + _fake_unit("unit-3", body=quick), + _fake_unit("unit-4", body=quick), + ] + report = await dispatch(units, concurrency=2, guard=guard) + + assert report.dispatched == 2 + # 「第 N 个任务」从 1 数：第 3 个是第一个没派出去的。 + assert report.stopped_at == 3 + # 慢的那个是在预算耗尽时正在跑的，它必须跑完并留下结果，不许被砍断。 + assert {item.run_id for item in report.outcomes} == {"unit-1", "unit-2"} + assert guard.total == 2 + + +async def test_concurrency_never_exceeds_the_cap() -> None: + live = 0 + peak = 0 + + async def body() -> None: + nonlocal live, peak + live += 1 + peak = max(peak, live) + await asyncio.sleep(0.005) + live -= 1 + + units = [_fake_unit(f"unit-{index}", body=body) for index in range(12)] + report = await dispatch(units, concurrency=3) + + assert peak == 3 + assert report.dispatched == 12 + assert len(report.outcomes) == 12 + + +async def test_zero_concurrency_is_rejected() -> None: + with pytest.raises(ValueError, match="并发上限"): + await dispatch([], concurrency=0) + + +def test_interleave_alternates_between_scenarios() -> None: + left = [_fake_unit(f"L{index}", body=lambda: None) for index in range(3)] + right = [_fake_unit(f"R{index}", body=lambda: None) for index in range(2)] + assert [unit.label for unit in interleave([left, right])] == ["L0", "R0", "L1", "R1", "L2"] + + +def test_select_appworld_tasks_dedupes_across_splits() -> None: + class _Pool: + def list_task_ids(self, split: str) -> list[str]: + return {"train": ["a", "b"], "dev": ["b", "c", "d"]}[split] + + assert select_appworld_tasks(_Pool(), splits=["train", "dev"], limit=None) == [ + "a", + "b", + "c", + "d", + ] + assert select_appworld_tasks(_Pool(), splits=["train", "dev"], limit=3) == ["a", "b", "c"] + + +# --------------------------------------------------------------------------- +# 三、错误隔离与取消 +# --------------------------------------------------------------------------- + + +async def test_one_failing_unit_does_not_stop_the_batch() -> None: + def boom() -> None: + raise RuntimeError("这一道题炸了") + + units = [ + _fake_unit("ok-1", body=lambda: None), + _fake_unit("boom", body=boom), + _fake_unit("ok-2", body=lambda: None), + ] + report = await dispatch(units, concurrency=1) + + assert {item.run_id for item in report.outcomes} == {"ok-1", "ok-2"} + assert len(report.failures) == 1 + label, message = report.failures[0] + assert label == "boom" + assert "RuntimeError" in message and "这一道题炸了" in message + + +async def test_cancellation_passes_through_and_cleans_up() -> None: + running = asyncio.Event() + cleaned = False + + async def body() -> None: + nonlocal cleaned + running.set() + try: + await asyncio.sleep(10) + finally: + cleaned = True + + units = [_fake_unit("slow", body=body)] + task = asyncio.create_task(dispatch(units, concurrency=1)) + await running.wait() + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + assert cleaned is True + + +# --------------------------------------------------------------------------- +# 四、产物：字段齐全，且记分板判得动 +# --------------------------------------------------------------------------- + + +async def _run_summarize(tmp_path: Path, *, index: int = 0) -> None: + """跑一次 summarize 阶段：模型第一步就提交结论，运行以 task_completed 结束。""" + guard = _guard(render=lambda call: _SUBMIT) + await run_govdoc_task( + task=_audit_task(index), + runs_dir=tmp_path / "runs", + workspace_root=tmp_path / "workspaces", + guard=guard, + phases=("summarize",), + ) + + +async def test_sidecar_fields_are_complete_and_typed(tmp_path: Path) -> None: + await _run_summarize(tmp_path) + runs = tmp_path / "runs" + run_id = govdoc_scenario.make_run_id(task_index=0, phase="summarize") + + assert (runs / f"{run_id}.jsonl").is_file() + assert (runs / f"{run_id}.result.json").is_file() + assert (runs / f"{run_id}.events.jsonl").is_file() + + meta = json.loads((runs / f"{run_id}.meta.json").read_text(encoding="utf-8")) + assert set(meta) == { + "scenario", + "task_id", + "phase", + "wall_ms", + "model_calls", + "sink_failures", + "env_executions", + "fault", + "success", + "resumed_from_step", + } + assert meta["scenario"] == "govdoc" + assert meta["task_id"] == "cp-0" + assert meta["phase"] == "summarize" + assert isinstance(meta["wall_ms"], int) + assert meta["model_calls"] == 1 + assert meta["sink_failures"] == 0 + # submit_finding 往审计账里追加一行，那是环境自己记的账。 + assert meta["env_executions"] == 1 + assert meta["fault"] is None + assert meta["success"] is None + assert meta["resumed_from_step"] == 0 + + events = [ + json.loads(line) + for line in (runs / f"{run_id}.events.jsonl").read_text(encoding="utf-8").splitlines() + ] + assert events == [{"kind": "step_finished", "run_id": run_id, "step_idx": 0}] + + result = json.loads((runs / f"{run_id}.result.json").read_text(encoding="utf-8")) + assert result["run_id"] == run_id + assert result["stop_reason"] == "task_completed" + assert result["event_delivery_failures"] == 0 + + +async def test_artifacts_pass_the_scoreboard(tmp_path: Path) -> None: + """把驱动写出来的产物直接喂给记分板，要它报全绿。 + + 两边的字段约定没有任何机器约束把它们钉在一起，这条测试就是那个约束。 + """ + guard = _guard(render=lambda call: _SUBMIT if "summarize" in call.run_id else _GARBAGE) + for index in range(2): + await run_govdoc_task( + task=_audit_task(index), + runs_dir=tmp_path / "runs", + workspace_root=tmp_path / "workspaces", + guard=guard, + ) + + board = scoreboard.evaluate( + tmp_path / "runs", + completing_tools=frozenset({"submit_finding"}), + ) + trouble = [ + f"{item.name}: {[evidence.describe() for evidence in (*item.breaches, *item.undetermined)]}" + for item in (*board.breached, *board.undetermined) + ] + assert board.verdict is scoreboard.Verdict.PASSED, trouble + assert len(board.summaries) == 6 + assert {item.scenario for item in board.summaries} == {"govdoc"} + + +async def test_sink_failures_match_what_the_library_counted(tmp_path: Path) -> None: + """事件文件写不出去时，出口自己数的失败次数与结果里那份必须一致。 + + 这条不变量是记分板的一条判定，而两边的计数由不同的代码写：出口在自己的 `except` 里加一， + 库在接住异常之后加一。只要出口漏加或多加，记分板当场报击穿。 + """ + runs = tmp_path / "runs" + runs.mkdir() + run_id = govdoc_scenario.make_run_id(task_index=0, phase="summarize") + # 在事件文件该在的位置放一个目录，追加写就必然失败。 + (runs / f"{run_id}.events.jsonl").mkdir() + + guard = _guard(render=lambda call: _SUBMIT) + await run_govdoc_task( + task=_audit_task(0), + runs_dir=runs, + workspace_root=tmp_path / "workspaces", + guard=guard, + phases=("summarize",), + ) + + meta = json.loads((runs / f"{run_id}.meta.json").read_text(encoding="utf-8")) + result = json.loads((runs / f"{run_id}.result.json").read_text(encoding="utf-8")) + assert meta["sink_failures"] == 1 + assert result["event_delivery_failures"] == meta["sink_failures"] + + +async def test_a_failing_run_still_leaves_meta(tmp_path: Path) -> None: + """`run` 抛异常也要留下 `.meta.json`——缺文件在记分板那边只是「无法判定」。 + + 这里用同一个运行标识跑第二次来触发失败：库对已经有日志的标识直接报 `RunIdentityError`， + 而那是压测里最可能真的撞上的一种失败（任务集去重漏了一处）。 + """ + runs = tmp_path / "runs" + await _run_summarize(tmp_path) + run_id = govdoc_scenario.make_run_id(task_index=0, phase="summarize") + (runs / f"{run_id}.result.json").unlink() + + guard = _guard(render=lambda call: _SUBMIT) + task = _audit_task(0) + request = govdoc_scenario.build_run_request( + task=task, + phase="summarize", + run_id=run_id, + workspace=tmp_path / "ws", + model_binding={}, + ) + outcome = await execute_run( + runs_dir=runs, + request=request, + model_client=guard, + decision_parser=govdoc_scenario.GovDocParser(), + synthetic_observations=govdoc_scenario.SYNTHETIC_OBSERVATIONS, + scenario="govdoc", + task_id=task.checkpoint.checkpoint_id, + phase="summarize", + count_model_calls=lambda: guard.calls_for(run_id), + ) + + assert outcome.error is not None + assert "RunIdentityError" in outcome.error + assert outcome.stop_reason is None + meta = json.loads((runs / f"{run_id}.meta.json").read_text(encoding="utf-8")) + assert meta["model_calls"] == 0 + assert meta["success"] is None + # 这一次没有结果，所以没有 `.result.json` 可写——记分板会把它记成一条要人看见的观察。 + assert not (runs / f"{run_id}.result.json").exists() + + +async def test_a_model_error_is_a_stop_reason_not_a_crash(tmp_path: Path) -> None: + """模型调用失败不会把异常抛出循环，它是一次以 `llm_error` 结束的正常运行。 + + 压测的报告里这两件事必须分得开：`error` 是驱动这边出的事，`llm_error` 是库判定的停止原因。 + """ + + class _Exploding: + async def call(self, call: ModelCall) -> ModelReply: + raise RuntimeError("网关炸了") + + def parameters(self): # noqa: ANN202 - 替身 + return {"kind": "exploding"} + + guard = BudgetGuard(inner=_Exploding(), limit=5) + task = _audit_task(0) + request = govdoc_scenario.build_run_request( + task=task, + phase="summarize", + run_id="govdoc-0-summarize", + workspace=tmp_path / "ws", + model_binding={}, + ) + outcome = await execute_run( + runs_dir=tmp_path / "runs", + request=request, + model_client=guard, + decision_parser=govdoc_scenario.GovDocParser(), + synthetic_observations=govdoc_scenario.SYNTHETIC_OBSERVATIONS, + scenario="govdoc", + task_id=task.checkpoint.checkpoint_id, + phase="summarize", + count_model_calls=lambda: guard.calls_for("govdoc-0-summarize"), + ) + + assert outcome.error is None + assert outcome.stop_reason == "llm_error" + meta = json.loads( + (tmp_path / "runs" / "govdoc-0-summarize.meta.json").read_text(encoding="utf-8") + ) + assert meta["model_calls"] == 1 + assert (tmp_path / "runs" / "govdoc-0-summarize.result.json").is_file() + + +# --------------------------------------------------------------------------- +# 五、GovDoc：阶段串行、任务并发 +# --------------------------------------------------------------------------- + + +async def test_govdoc_phases_run_in_order_within_a_task(tmp_path: Path) -> None: + model = _ScriptedModel(render=lambda call: _GARBAGE) + guard = BudgetGuard(inner=model, limit=1000) + await run_govdoc_task( + task=_audit_task(0), + runs_dir=tmp_path / "runs", + workspace_root=tmp_path / "workspaces", + guard=guard, + ) + + phases = [run_id.rsplit("-", 1)[1] for run_id in model.seen] + # 每个阶段的调用连成一段，段与段之间不交错。 + assert phases == sorted(phases, key=["plan", "execute", "summarize"].index) + assert set(phases) == {"plan", "execute", "summarize"} + + +async def test_govdoc_tasks_overlap_while_phases_do_not(tmp_path: Path) -> None: + model = _ScriptedModel(render=lambda call: _GARBAGE) + guard = BudgetGuard(inner=model, limit=1000) + units = [ + Unit( + label=f"govdoc/{index}", + scenario="govdoc", + run=_govdoc_runner( + index=index, + runs_dir=tmp_path / "runs", + workspace_root=tmp_path / "workspaces", + guard=guard, + ), + ) + for index in range(2) + ] + report = await dispatch(units, concurrency=2, guard=guard) + assert len(report.outcomes) == 6 + assert report.failures == () + + # 阶段串行：同一个任务里，后一阶段的第一次调用晚于前一阶段的最后一次调用。 + for index in range(2): + own = [ + position + for position, run_id in enumerate(model.seen) + if run_id.startswith(f"govdoc-{index}-") + ] + for earlier, later in (("plan", "execute"), ("execute", "summarize")): + last_earlier = max( + position for position in own if model.seen[position].endswith(earlier) + ) + first_later = min(position for position in own if model.seen[position].endswith(later)) + assert last_earlier < first_later + + # 任务并发：两个任务的调用在时间上交错。 + owners = [run_id.split("-")[1] for run_id in model.seen] + assert any(left != right for left, right in zip(owners, owners[1:], strict=False)) + + +def _govdoc_runner(*, index: int, runs_dir: Path, workspace_root: Path, guard: BudgetGuard): # noqa: ANN202 + async def start() -> Sequence[RunOutcome]: + return await run_govdoc_task( + task=_audit_task(index), + runs_dir=runs_dir, + workspace_root=workspace_root, + guard=guard, + ) + + return start + + +# --------------------------------------------------------------------------- +# 小工具 +# --------------------------------------------------------------------------- + + +def _call(run_id: str) -> ModelCall: + from polyloop.ports import ModelCall as _ModelCall + + return _ModelCall( + messages=(), + call_index=0, + run_id=run_id, + result_id=f"{run_id}-0", + binding={}, + )