iomgaa
c0d9d7b66e
CHANGELOG 落在「未发布」段，按那一段自己的规矩不提前写版本号。写清了四条会抛 ValueError 的 情形与迁移写法，因为这是一次破坏性变更。 migrations/govdoc-saas.md 加一节讲租户命名空间怎么传，含迁完算不算数的四条判据，最终判据是 「同一段文本由两个租户各提交一次，各自拿到自己的那份输出」。参数一律指向 0017 不复述。这份 文档原来说 GovDoc 侧「还没有可迁移的东西」，现在有了第一条能逐条验的接入动作，文件头那句 状态说明跟着补了一句例外。 tools/soak/run_soak.py 那份空绑定上方的注释在复述旧转发规则，顺手改对。它给的理由本来就不准 ——让绑定留空的真正原因是绑定的全部键值都进参数快照，每批都不同的值会让故障注入那一步续跑时 报一次假的参数漂移，和转发哪些键无关。三份压测绑定常量里都没有裸键，所以这次变更打不到压测。 Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
1021 lines
40 KiB
Python
1021 lines
40 KiB
Python
"""压测入口：把 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"
|
||
|
||
#: 传给每次模型调用的项目侧标识。**空的**：绑定的全部键值都进参数快照，而每批都不同的值会让
|
||
#: 故障注入那一步续跑时的逐字段比对报一次假的参数漂移。这里跟转发无关——网关适配器只转发带
|
||
#: `gateway.` 前缀的键，而这份一个都没有。
|
||
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:
|
||
"""把事件逐行写进 `<run_id>.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:
|
||
"""跑一次运行，并把记分板要的四个文件写全。
|
||
|
||
库自己写 `<run_id>.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())
|