feat(soak): 压测的环境层——AppWorld 容器池与薄 HTTP 客户端
⑥ 的验收要自己造负载压，环境用 AppWorld：733 个任务与 197M 数据都在本机、docker 镜像已 拉好、而且它自带评测端点做程序化判分——换成让另一个模型判对错，等于往验收里再塞一个非 确定源，验收本身就不可复现了。 **不 import dissect**（§1.2 禁止反向 import），照它那两个文件当协议文档自己写了一份。 副作用是这反而更强：证明一个不认识 dissect 的第三方只用公共 API 就能驱动真实环境。 复刻了 dissect 记下的几个坑：httpx 必须 trust_env=False，否则本机代理会把 127.0.0.1 的 请求也劫走、表现成 502；健康检查失败先抓容器日志再删；关闭失败按端口分开计数，用全局 计数器的话坏容器的失败会被别的容器的成功清零。 放在仓库根的 tools/ 下而不是 tests/ 下：打包只收 src/，所以它不进 wheel 也不进 sdist （实测两个产物里 soak 命中数都是 0）；而 §1.9 的四层是按「依赖什么」分的，压测不属于其中 任何一层，塞进 tests/ 要么破坏分层要么和 e2e 共用同一道花钱的闸。Makefile 的检查目标 跟着加上 tools/。 冒烟真跑通：起容器、初始化真题、跨调用保持变量、错误代码返回 traceback 而不抛异常、 评测、无容器残留；另验 2 容器并发的租借与归还。 Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,833 @@
|
||||
"""压测 harness 的环境层：一个 AppWorld 容器池，加一个薄 HTTP 客户端。
|
||||
|
||||
**这一层完全不碰模型。** 它只提供「起容器 → 实例化一道题 → 执行代码 → 问完成没有 →
|
||||
评分 → 关闭」这一串动作，谁来决定执行什么代码是上层的事。所以它可以离线验证——
|
||||
`check_appworld.py` 用一段写死的 Python 就能把整条链路走通，一次模型调用都不用打。
|
||||
|
||||
**为什么走 HTTP 而不是直接 import appworld。** appworld 包钉在 pydantic 1.x，而
|
||||
PolyGateway 要求 pydantic ≥ 2.8，两者装不进同一个解释器。AppWorld 官方把它的
|
||||
`AppWorld` 类整个包成了一个 environment server 跑在官方镜像里，所以本模块只是个瘦
|
||||
客户端，本进程一行 appworld 代码都不导入。
|
||||
|
||||
**为什么需要容器池而不是一个容器。** environment server 用一个模块级变量存「当前任务」，
|
||||
请求换题会被直接拒绝，也就是说一个容器同一时刻只能跑一道题。要并发跑 N 道题就得起
|
||||
N 个容器，各占一个宿主端口，用完归还。
|
||||
|
||||
**协议来自 dissect 的 `harness/envs/`，代码是独立写的。** PolyLoop 是被 dissect 依赖的
|
||||
库，反向 import 下游是硬约束（CLAUDE.md §1.2，由 import-linter 断言）。dissect 那两个
|
||||
文件在这里只当协议文档看——端点形状、初始化参数取值、两个坑的成因都出自那里，实现则
|
||||
只保留压测用得着的最小子集，不带 dissect 自己的 Episode/TaskEnv/ScoredEpisode 抽象、
|
||||
不带评分聚合、不带快照参数。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
import tempfile
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Self
|
||||
|
||||
import httpx
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import AsyncIterator, Callable, Mapping, Sequence
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
#: environment server 的官方镜像。用滚动标签而不是 digest：压测只关心「跑得动多少
|
||||
#: 并发」，不做跨时间可比的成绩，上游换镜像不影响结论。真做实验的那一侧（dissect）
|
||||
#: 才需要把 digest 钉进快照。
|
||||
DEFAULT_IMAGE = "ghcr.io/stonybrooknlp/appworld:latest"
|
||||
|
||||
#: 起始宿主端口，池占用 [PORT_BASE, PORT_BASE + size)。
|
||||
#:
|
||||
#: **刻意避开 8100**：dissect 的默认值就是 8100，而这台机器上两边可能同时在跑。撞了
|
||||
#: 端口的表现不是一句「端口被占用」——先起的那个容器活着，后起的 `docker run` 失败，
|
||||
#: 于是压测报「容器起不来」，而真正在受害的是另一个项目的实验。
|
||||
DEFAULT_PORT_BASE = 8200
|
||||
|
||||
#: 容器名前缀。名字里带端口号，`docker ps` 的输出就能直接对上出问题的那个端口。
|
||||
CONTAINER_NAME_PREFIX = "polyloop-soak-appworld"
|
||||
|
||||
#: 健康检查路径。返回 2xx 即视为就绪。
|
||||
_READINESS_PATH = "/"
|
||||
|
||||
#: 数据目录的挂载模式。**只读**：这份数据是 dissect 的实验数据，733 个任务目录，压测
|
||||
#: 没有任何理由改它，而一次误写会污染另一个项目的实验输入且无人察觉。环境自己要写的
|
||||
#: 东西全部落在输出目录那个挂载点上。
|
||||
_DATA_MOUNT_MODE = "ro"
|
||||
|
||||
#: 连续多少次会话关闭失败就让整轮压测停下来。取 3 而不是 1，是因为一次网络抖动不该
|
||||
#: 毁掉整批；取 3 而不是 50，是因为泄漏是累积的，等到几十次时环境早就不干净了。
|
||||
_MAX_CLOSE_FAILURES = 3
|
||||
|
||||
#: 允许的挂载模式。docker 会把认不得的模式当成一个额外的挂载选项，拼错了不一定报错，
|
||||
#: 所以在拼命令行之前先自己拦一道。
|
||||
_MOUNT_MODES = frozenset({"ro", "rw"})
|
||||
|
||||
|
||||
class ContainerPoolError(RuntimeError):
|
||||
"""容器池的启动、健康检查或清理失败。
|
||||
|
||||
一律带上 docker 的 stderr 或容器日志。容器起不来时只报一句「超时」，调试就只能靠猜。
|
||||
"""
|
||||
|
||||
|
||||
class AppWorldError(RuntimeError):
|
||||
"""AppWorld 环境返回了错误，或返回了我们无法解释的内容。
|
||||
|
||||
**它不表示「模型写的代码跑挂了」**——那是正常观察，由 `AppWorldSession.execute`
|
||||
原样返回。这个异常只在环境本身坏了时抛：连不上、HTTP 非 2xx、返回体不是约定的形状。
|
||||
"""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Mount:
|
||||
"""一条 docker 卷挂载。
|
||||
|
||||
模式单独成字段而不是拼进字符串，是为了让「这个挂载是只读的」在调用处看得见——
|
||||
`-v a:b:ro` 里那两个字母混在路径中间，改错了不会有人发现。
|
||||
"""
|
||||
|
||||
host: Path
|
||||
container: str
|
||||
mode: str
|
||||
|
||||
def to_arg(self) -> str:
|
||||
"""拼成 `docker run -v` 的取值。"""
|
||||
if self.mode not in _MOUNT_MODES:
|
||||
raise ValueError(f"挂载模式只能是 {sorted(_MOUNT_MODES)}，收到 {self.mode!r}")
|
||||
return f"{self.host}:{self.container}:{self.mode}"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TaskScore:
|
||||
"""官方评测器对一道题的判定。
|
||||
|
||||
`success` 是过与不过，AppWorld 没有连续分。`detail` 是评测器返回的完整结构，里面
|
||||
逐条 requirement 的通过情况在排查「是模型不行还是环境坏了」时是唯一线索。
|
||||
"""
|
||||
|
||||
task_id: str
|
||||
success: bool
|
||||
n_executions: int
|
||||
detail: Mapping[str, Any]
|
||||
|
||||
|
||||
class ContainerPool:
|
||||
"""一组同构容器与其宿主端口的租借池。
|
||||
|
||||
池里第 i 个容器占宿主端口 `port_base + i`，容器内监听同一个端口号——`-p` 的两侧
|
||||
一致，容器里的服务就不需要知道自己被映射到了哪里。
|
||||
|
||||
用法::
|
||||
|
||||
async with ContainerPool(...) as pool:
|
||||
async with pool.lease() as port:
|
||||
... # 独占这个端口上的服务
|
||||
|
||||
池不清理租借者在容器里留下的状态。归还前复位状态是租借者的责任——对 AppWorld
|
||||
而言就是调 `/close`。
|
||||
|
||||
**它刻意不知道容器里跑的是什么。** 启动命令由 `container_args` 给出，健康检查路径
|
||||
由 `readiness_path` 给出，本类只管容器的生死与端口的租借。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
image: str,
|
||||
size: int,
|
||||
port_base: int,
|
||||
name_prefix: str,
|
||||
container_args: Callable[[int], Sequence[str]],
|
||||
readiness_path: str,
|
||||
mounts: Sequence[Mount] = (),
|
||||
startup_timeout_s: float,
|
||||
) -> None:
|
||||
"""构造容器池。不启动任何东西，启动在 `start()`。
|
||||
|
||||
Args:
|
||||
image: 镜像名，需已在本地或可拉取。
|
||||
size: 容器数量，即并发上限。
|
||||
port_base: 起始宿主端口。
|
||||
name_prefix: 容器名前缀，实际名字是 `{prefix}-{端口号}`。
|
||||
container_args: 端口 → 容器内启动命令。每个容器的命令都要带上自己的端口号，
|
||||
所以这里是函数而不是一份固定列表。
|
||||
readiness_path: 健康检查的 HTTP 路径。
|
||||
mounts: 卷挂载。
|
||||
startup_timeout_s: 单个容器从 `docker run` 到健康检查通过的时限。
|
||||
"""
|
||||
if size < 1:
|
||||
raise ValueError(f"池大小必须至少为 1，收到 {size}")
|
||||
if port_base < 1024:
|
||||
raise ValueError(f"起始端口须在非特权区间，收到 {port_base}")
|
||||
if startup_timeout_s <= 0:
|
||||
raise ValueError(f"启动时限必须为正，收到 {startup_timeout_s}")
|
||||
self._image = image
|
||||
self._size = size
|
||||
self._ports = [port_base + i for i in range(size)]
|
||||
self._name_prefix = name_prefix
|
||||
self._container_args = container_args
|
||||
self._readiness_path = readiness_path
|
||||
self._mounts = tuple(mounts)
|
||||
self._startup_timeout_s = startup_timeout_s
|
||||
self._free_ports: asyncio.Queue[int] = asyncio.Queue()
|
||||
self._started = False
|
||||
self._stopped = False
|
||||
self._leased = 0
|
||||
|
||||
@property
|
||||
def ports(self) -> tuple[int, ...]:
|
||||
"""池占用的全部宿主端口，无论此刻是否空闲。"""
|
||||
return tuple(self._ports)
|
||||
|
||||
async def start(self) -> None:
|
||||
"""启动全部容器并等它们就绪。任何一个起不来就把已起的全部拆掉再抛错。
|
||||
|
||||
留下半个池比干净地失败更难排查：下一次跑会撞上同名容器，而报出来的是
|
||||
`docker run` 的重名错误，跟真正的起因隔了一整轮。
|
||||
"""
|
||||
if self._started:
|
||||
raise ContainerPoolError("容器池已启动，不要重复调用 start()")
|
||||
# 池是一次性的：停过就不能再起。允许重启的话，`stop()` 之后迟到归还的租约会把
|
||||
# 端口 put 进队列，紧接着的 `start()` 再 put 一遍，队列里就有了重复端口——两个
|
||||
# 任务会同时租到同一个容器，而 environment server 是单例，后者会被「当前活动任务
|
||||
# 是另一道题」打挂。压测没有重启的用例，就不留这个坑。
|
||||
if self._stopped:
|
||||
raise ContainerPoolError("容器池已停止，不支持重启；需要的话新建一个池")
|
||||
# 上一次运行若被强杀会留下同名容器，`docker run` 会因重名直接失败。先无条件清
|
||||
# 一遍，让重跑不需要人工干预。
|
||||
await self._remove_all()
|
||||
# 用 TaskGroup 而不是 gather：gather 在第一个子协程抛错时立刻把异常抛给调用方，
|
||||
# 但**不取消其余子协程**，它们会一直跑到 startup_timeout_s。那意味着下面的清理
|
||||
# 可能跑在某些 `docker run` 完成之前，于是那些容器活了下来占着端口——正是「要么
|
||||
# 全起要么全拆」要防的半个池。TaskGroup 会取消其余任务并等它们收敛之后才抛。
|
||||
try:
|
||||
async with asyncio.TaskGroup() as group:
|
||||
for port in self._ports:
|
||||
group.create_task(self._start_one(port))
|
||||
except BaseExceptionGroup as failures:
|
||||
await self._remove_all()
|
||||
# TaskGroup 把子任务的异常打包成组。这里拆开只抛第一个，让调用方拿到一个
|
||||
# 平常的异常而不是需要 `except*` 才能接的组；其余的记进日志，不然多个容器
|
||||
# 同时出问题时只看得见一个。
|
||||
first, *rest = failures.exceptions
|
||||
for extra in rest:
|
||||
logger.error("容器池启动时的另一处失败：%s", extra)
|
||||
# 不写 `from ...`：`first` 自己的 __cause__ 已经指着真正的故障原因（见
|
||||
# `_start_one` 的 `from exc`），再显式挂一个 from 会把它顶掉，异常链上就只剩
|
||||
# 一句「启动过程中出错」。B904 要防的是丢掉上下文，这里上下文本来就在。
|
||||
raise first # noqa: B904
|
||||
except BaseException:
|
||||
# 取消之类不经由 TaskGroup 打包的路径，同样要保证不留下半个池。
|
||||
await self._remove_all()
|
||||
raise
|
||||
for port in self._ports:
|
||||
self._free_ports.put_nowait(port)
|
||||
self._started = True
|
||||
logger.info("容器池就绪：%d 个 %s 实例，端口 %s", self._size, self._image, self._ports)
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""停止并删除全部容器。幂等，但停过之后不能再 `start()`。"""
|
||||
if self._leased:
|
||||
# 容器马上就没了，那些租约手里的端口随后会连不上。这说明调用方在还有任务在跑
|
||||
# 的时候就退出了 `async with`，属于调用方的 bug；但沉默地让它表现成一堆网络
|
||||
# 错误更糟——那时看到的是「AppWorld 连不上」，查的方向完全错了。
|
||||
logger.error("停止容器池时仍有 %d 个租约在途，它们的请求将会失败", self._leased)
|
||||
# 状态更新必须放 finally：`docker rm` 也会失败（docker 二进制不在时
|
||||
# create_subprocess_exec 直接抛 FileNotFoundError）。失败之后如果状态还停在
|
||||
# 「已启动」、队列里还剩着端口，接下来的 lease() 会把端口发给一批已经不确定还在
|
||||
# 不在的容器，故障从「停不掉」变成「跑在幽灵容器上」。
|
||||
try:
|
||||
await self._remove_all()
|
||||
finally:
|
||||
self._started = False
|
||||
self._stopped = True
|
||||
while not self._free_ports.empty():
|
||||
self._free_ports.get_nowait()
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def lease(self) -> AsyncIterator[int]:
|
||||
"""租借一个空闲容器的宿主端口，退出时归还。
|
||||
|
||||
池满时阻塞等待，**不设排队超时**——排队时间是正常的并发背压，把它变成异常只会
|
||||
让上层写一堆重试逻辑，而重试改变不了「容器就这么多」这个事实。
|
||||
"""
|
||||
# 先判已停止：停过之后 `_started` 也是 False，只报「尚未启动，先调用 start()」
|
||||
# 会把人引去调一个必然报「不支持重启」的方法。
|
||||
if self._stopped:
|
||||
raise ContainerPoolError("容器池已停止，它的容器都没了；需要的话新建一个池")
|
||||
if not self._started:
|
||||
raise ContainerPoolError("容器池尚未启动，先调用 start() 或用 async with")
|
||||
port = await self._free_ports.get()
|
||||
self._leased += 1
|
||||
try:
|
||||
yield port
|
||||
finally:
|
||||
self._leased -= 1
|
||||
self._free_ports.put_nowait(port)
|
||||
|
||||
async def __aenter__(self) -> Self:
|
||||
await self.start()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc_info: object) -> None:
|
||||
await self.stop()
|
||||
|
||||
# -- 内部 ----------------------------------------------------------------
|
||||
|
||||
def _container_name(self, port: int) -> str:
|
||||
return f"{self._name_prefix}-{port}"
|
||||
|
||||
async def _start_one(self, port: int) -> None:
|
||||
"""起一个容器并等它健康。"""
|
||||
name = self._container_name(port)
|
||||
command = ["docker", "run", "-d", "--name", name, "-p", f"{port}:{port}"]
|
||||
for mount in self._mounts:
|
||||
command += ["-v", mount.to_arg()]
|
||||
command += [self._image, *self._container_args(port)]
|
||||
|
||||
code, _, stderr = await _run(command)
|
||||
if code != 0:
|
||||
raise ContainerPoolError(f"容器 {name} 启动失败（docker run 退出码 {code}）：{stderr}")
|
||||
try:
|
||||
await self._await_ready(port, name)
|
||||
except asyncio.CancelledError:
|
||||
# 取消必须原样传播。转成普通异常会让这个任务不再处于「已取消」状态，事件循环
|
||||
# 随后报一次 unhandled exception，而现场看起来像是一次超时——真正发生的事
|
||||
# （有人按了 Ctrl-C 或上层撤销了这一批）就此消失。
|
||||
await self._remove_noisily(name)
|
||||
raise
|
||||
except Exception as exc:
|
||||
# 健康检查失败时容器日志是唯一的线索，**必须先抓出来再删**。删掉之后
|
||||
# `docker logs` 就没有东西可读了，而这类失败往往只在负载高的时候偶发。
|
||||
# 保留 `from exc`：能进来的不止超时，readiness_path 写错导致的 httpx.InvalidURL
|
||||
# 之类也会走到这儿，抹掉原因就只能靠猜。
|
||||
reason = "未就绪（超时）" if isinstance(exc, TimeoutError) else "启动过程中出错"
|
||||
logs = await self._container_logs(name)
|
||||
await self._remove_noisily(name)
|
||||
raise ContainerPoolError(
|
||||
f"容器 {name} {reason}（等待上限 {self._startup_timeout_s} 秒）。容器日志：\n{logs}"
|
||||
) from exc
|
||||
|
||||
async def _await_ready(self, port: int, name: str) -> None:
|
||||
"""轮询健康检查端点直到返回 2xx 或超时。"""
|
||||
url = f"http://127.0.0.1:{port}{self._readiness_path}"
|
||||
deadline = asyncio.get_running_loop().time() + self._startup_timeout_s
|
||||
# trust_env=False 是必须的：这台机器设了 http_proxy，而 httpx 默认会读环境变量，
|
||||
# 于是连 127.0.0.1 的请求也被送去代理。表现是 HTTP 502——看起来像容器里的服务
|
||||
# 出错，实际上请求根本没到过容器。
|
||||
async with httpx.AsyncClient(timeout=5.0, trust_env=False) as client:
|
||||
while asyncio.get_running_loop().time() < deadline:
|
||||
# 容器刚起时连接被拒是常态，只有超时才算失败，所以这里不区分具体的网络
|
||||
# 错误类型；非网络异常（配置错误之类）仍会照常冒泡。
|
||||
try:
|
||||
response = await client.get(url)
|
||||
except httpx.TransportError:
|
||||
await asyncio.sleep(0.5)
|
||||
continue
|
||||
if response.is_success:
|
||||
logger.debug("容器 %s 就绪", name)
|
||||
return
|
||||
await asyncio.sleep(0.5)
|
||||
raise TimeoutError(f"容器 {name} 健康检查超时")
|
||||
|
||||
async def _container_logs(self, name: str) -> str:
|
||||
"""取容器日志用于报错。取不到就说明取不到，不返回空串冒充「日志为空」。"""
|
||||
code, stdout, stderr = await _run(["docker", "logs", "--tail", "50", name])
|
||||
if code != 0:
|
||||
return f"<无法读取容器日志，docker logs 退出码 {code}：{stderr}>"
|
||||
return (stdout + stderr).strip() or "<容器未输出任何日志>"
|
||||
|
||||
async def _remove_noisily(self, name: str) -> None:
|
||||
"""尽力删掉一个容器；删不掉只记日志，不改变正在传播的异常。
|
||||
|
||||
这里不用 `contextlib.suppress`：清理失败意味着一个容器占着端口活了下来，下一次
|
||||
跑会以「docker run 重名」的形式失败，而那时已经没人记得是这一次没删干净。
|
||||
"""
|
||||
try:
|
||||
code, _, stderr = await _run(["docker", "rm", "-f", name])
|
||||
except OSError as exc:
|
||||
logger.error("清理容器 %s 时连 docker 都没跑起来：%s", name, exc)
|
||||
return
|
||||
if code != 0:
|
||||
logger.error("清理容器 %s 失败（退出码 %d）：%s", name, code, stderr.strip())
|
||||
|
||||
async def _remove_all(self) -> None:
|
||||
"""强删池内全部容器。
|
||||
|
||||
不检查退出码：绝大多数非零都是「容器本来就不存在」，而这正是启动前清残留时的常态。
|
||||
真正的故障（docker 守护进程没了）会在紧接着的 `docker run` 上以更清楚的形式暴露。
|
||||
"""
|
||||
names = [self._container_name(port) for port in self._ports]
|
||||
await _run(["docker", "rm", "-f", *names])
|
||||
|
||||
|
||||
def container_args_for_port(port: int) -> Sequence[str]:
|
||||
"""官方镜像的启动参数：第一个位置参数选服务类型，然后是端口。
|
||||
|
||||
单独提出来是为了不必起一个真容器就能核对这串参数。
|
||||
"""
|
||||
return ["environment", "--port", str(port), "--no-show-usage"]
|
||||
|
||||
|
||||
class AppWorldSession:
|
||||
"""一道题的一次会话，绑定在池里的某一个容器上。
|
||||
|
||||
生命周期由 `AppWorldPool.session()` 管理，不要直接构造——直接构造出来的会话没有
|
||||
对应的 `/close`，容器里的资源会一直挂着。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
client: httpx.AsyncClient,
|
||||
base_url: str,
|
||||
task_id: str,
|
||||
instruction: str,
|
||||
supervisor: Mapping[str, str],
|
||||
datetime: str,
|
||||
) -> None:
|
||||
self._client = client
|
||||
self._base_url = base_url
|
||||
self._task_id = task_id
|
||||
self._instruction = instruction
|
||||
self._supervisor = dict(supervisor)
|
||||
self._datetime = datetime
|
||||
self._n_executions = 0
|
||||
|
||||
@property
|
||||
def task_id(self) -> str:
|
||||
"""本题的题目 ID。"""
|
||||
return self._task_id
|
||||
|
||||
@property
|
||||
def base_url(self) -> str:
|
||||
"""本次会话所在容器的地址。排查时用来定位是哪个容器出的事。"""
|
||||
return self._base_url
|
||||
|
||||
@property
|
||||
def instruction(self) -> str:
|
||||
"""题面：主管交给 agent 的自然语言指令。"""
|
||||
return self._instruction
|
||||
|
||||
@property
|
||||
def supervisor(self) -> Mapping[str, str]:
|
||||
"""主管的身份信息（姓名、邮箱、电话）。agent 要靠它调 API。"""
|
||||
return dict(self._supervisor)
|
||||
|
||||
@property
|
||||
def datetime(self) -> str:
|
||||
"""任务发生的虚拟时间，题目的一部分。"""
|
||||
return self._datetime
|
||||
|
||||
@property
|
||||
def n_executions(self) -> int:
|
||||
"""到目前为止执行过多少次代码。"""
|
||||
return self._n_executions
|
||||
|
||||
async def execute(self, code: str) -> str:
|
||||
"""在环境里执行一段 Python 代码，返回它的输出。
|
||||
|
||||
这是 AppWorld 唯一的动作接口——它不是 JSON 形式的工具调用，agent 通过写代码调
|
||||
``apis.<应用>.<接口>(...)`` 来操作各个 app，通过 ``apis.supervisor.complete_task()``
|
||||
声明做完了。执行器是有状态的（环境侧是一个常驻 IPython shell），变量、import、
|
||||
打开的句柄都跨步存活。
|
||||
|
||||
**代码本身报错、超时、语法错，都不抛异常。** 环境会把 traceback 或超时提示放在
|
||||
输出里返回，那是给模型看的正常观察，压成异常就等于把「模型写错了」和「环境坏了」
|
||||
混成同一件事，而上层对这两者的处理完全不同。只有环境自己坏了——连不上、HTTP 非
|
||||
2xx、返回体不是约定形状——才抛 `AppWorldError`。
|
||||
"""
|
||||
payload = await self._post("/execute", {"task_id": self._task_id, "code": code})
|
||||
self._n_executions += 1
|
||||
if not isinstance(payload, str):
|
||||
raise AppWorldError(f"execute 期望返回字符串，收到 {type(payload).__name__}")
|
||||
return payload
|
||||
|
||||
async def is_done(self) -> bool:
|
||||
"""环境是否看到了完成信号，即 agent 调过 `apis.supervisor.complete_task()`。
|
||||
|
||||
这只是停机条件之一。步数耗尽、连续解析失败之类由上层自己判断，不经过环境。
|
||||
"""
|
||||
payload = await self._post("/task_completed", {"task_id": self._task_id})
|
||||
return bool(payload)
|
||||
|
||||
async def evaluate(self) -> TaskScore:
|
||||
"""跑官方评测器给这道题打分。**必须在会话关闭之前调用**，关闭之后环境状态就没了。
|
||||
|
||||
`suppress_errors=True` 是正常打分模式，不是「把错误藏起来」。AppWorld 用异常表达
|
||||
「这一条 requirement 没通过」：评测脚本把每条 requirement 包在一个上下文管理器里，
|
||||
断言失败时 `__exit__` 把它记进 failures 再返回 `suppress_errors` 决定是否抑制。
|
||||
取 False 的话第一条没过的 requirement 就会中断整个评测，连分数都拿不到。
|
||||
|
||||
代价是它无差别地吞掉所有异常，包括评测基础设施自己的故障——两者都表现成 failure，
|
||||
区别只在 `detail` 里那条 trace 的内容。所以成绩集体为零时要去看 trace 是断言失败
|
||||
还是别的东西，别把环境坏了当成模型不行。
|
||||
"""
|
||||
payload = await self._post(
|
||||
"/evaluate",
|
||||
{"task_id": self._task_id, "suppress_errors": True, "report": False},
|
||||
)
|
||||
if not isinstance(payload, dict):
|
||||
raise AppWorldError(f"evaluate 期望返回对象，收到 {type(payload).__name__}")
|
||||
return TaskScore(
|
||||
task_id=self._task_id,
|
||||
success=_success_of(payload, self._task_id),
|
||||
n_executions=self._n_executions,
|
||||
detail=payload,
|
||||
)
|
||||
|
||||
async def close(self) -> None:
|
||||
"""关闭会话，释放环境侧资源。由 `AppWorldPool.session()` 在退出时调用。"""
|
||||
await self._post("/close", {"task_id": self._task_id})
|
||||
|
||||
async def _post(self, path: str, body: Mapping[str, Any]) -> Any:
|
||||
return await _post_json(self._client, self._base_url, path, body)
|
||||
|
||||
|
||||
class AppWorldPool:
|
||||
"""AppWorld 环境的入口：管容器池、管 HTTP 连接、开会话。
|
||||
|
||||
用法::
|
||||
|
||||
pool = AppWorldPool(data_root=Path(...), size=2)
|
||||
async with pool:
|
||||
task_id = pool.list_task_ids("train")[0]
|
||||
async with pool.session(task_id) as session:
|
||||
print(await session.execute("print(1 + 1)"))
|
||||
score = await session.evaluate()
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
data_root: Path | str,
|
||||
size: int,
|
||||
outputs_dir: Path | str | None = None,
|
||||
experiment_name: str = "polyloop-soak",
|
||||
image: str = DEFAULT_IMAGE,
|
||||
port_base: int = DEFAULT_PORT_BASE,
|
||||
max_interactions: int = 40,
|
||||
execution_timeout_s: int = 100,
|
||||
startup_timeout_s: float = 180.0,
|
||||
) -> None:
|
||||
"""构造入口。不启动容器也不建连接，那些都在 `start()`。
|
||||
|
||||
Args:
|
||||
data_root: AppWorld 的数据根目录，即 `appworld download data --root` 指定的
|
||||
那个。它下面应有 `data/datasets/` 与 `data/tasks/`。**这个目录以只读方式
|
||||
挂进容器**，压测不会改它。
|
||||
size: 容器数量，即并发上限。
|
||||
outputs_dir: 环境侧写日志的目录，挂到容器的 `/run/experiments/outputs`。
|
||||
默认落在系统临时目录下我们自己的一个位置——不复用数据根目录下的
|
||||
`experiments/outputs`，因为那是别的项目的实验产物，压测不该往里掺东西。
|
||||
experiment_name: 环境用它作为输出子目录名。
|
||||
image: environment server 的镜像。
|
||||
port_base: 容器池的起始宿主端口。
|
||||
max_interactions: 环境侧允许的最大执行次数，当成上层步数预算的双保险——
|
||||
上层循环失控时环境会兜住。40 与 dissect 的取值一致（官方默认是 1000）。
|
||||
execution_timeout_s: 单次代码执行的超时，传给环境侧。
|
||||
startup_timeout_s: 单个容器启动并就绪的时限。镜像已在本地时通常几秒就够，
|
||||
留 180 秒是为了容得下一次冷拉取。
|
||||
"""
|
||||
self._data_root = Path(data_root).resolve()
|
||||
self._datasets_dir = self._data_root / "data" / "datasets"
|
||||
self._tasks_dir = self._data_root / "data" / "tasks"
|
||||
self._outputs_dir = (
|
||||
Path(outputs_dir).resolve()
|
||||
if outputs_dir is not None
|
||||
else Path(tempfile.gettempdir()) / "polyloop-soak-appworld-outputs"
|
||||
)
|
||||
self._experiment_name = experiment_name
|
||||
self._max_interactions = max_interactions
|
||||
self._execution_timeout_s = execution_timeout_s
|
||||
self._client: httpx.AsyncClient | None = None
|
||||
#: 每个容器端口各自的连续关闭失败次数。**必须按端口分开数**：池里有多个容器，
|
||||
#: 用一个全局计数器的话，坏掉那个容器每次失败都会被其他容器的成功清零，阈值
|
||||
#: 永远到不了，安全网等于不存在——而资源泄漏恰恰是发生在单个容器上的。
|
||||
self._close_failures: dict[int, int] = {}
|
||||
|
||||
self._pool = ContainerPool(
|
||||
image=image,
|
||||
size=size,
|
||||
port_base=port_base,
|
||||
name_prefix=CONTAINER_NAME_PREFIX,
|
||||
container_args=container_args_for_port,
|
||||
readiness_path=_READINESS_PATH,
|
||||
mounts=[
|
||||
Mount(host=self._data_root / "data", container="/run/data", mode=_DATA_MOUNT_MODE),
|
||||
Mount(
|
||||
host=self._outputs_dir,
|
||||
container="/run/experiments/outputs",
|
||||
mode="rw",
|
||||
),
|
||||
],
|
||||
startup_timeout_s=startup_timeout_s,
|
||||
)
|
||||
|
||||
@property
|
||||
def ports(self) -> tuple[int, ...]:
|
||||
"""池占用的宿主端口。"""
|
||||
return self._pool.ports
|
||||
|
||||
@property
|
||||
def outputs_dir(self) -> Path:
|
||||
"""环境侧日志落在宿主上的哪里。
|
||||
|
||||
**容器里的进程是 root，写出来的文件也是 root 的**，所以这个目录事后不由本模块
|
||||
删除——删不掉，而一次删不掉的清理会以「压测收尾报错」的形式盖住真正的结果。
|
||||
"""
|
||||
return self._outputs_dir
|
||||
|
||||
async def start(self) -> None:
|
||||
"""检查数据、建连接、启动容器池。
|
||||
|
||||
任何一步失败都要把已经建起来的东西收回去——`async with` 在 `__aenter__` 抛错时
|
||||
**不会**调用 `__aexit__`，不自己收就是一次泄漏。
|
||||
"""
|
||||
# 重入守卫。没有它的话第二次 start() 会先覆盖掉 self._client（旧连接就此泄漏），
|
||||
# 再被容器池的「已启动」守卫打回，然后 except 里关掉刚建的新连接并置 None——结果
|
||||
# 是容器还在跑、旧连接漏着、而对外声称尚未启动。
|
||||
if self._client is not None:
|
||||
raise AppWorldError("环境已启动，不要重复调用 start()")
|
||||
self._verify_data_layout()
|
||||
self._outputs_dir.mkdir(parents=True, exist_ok=True)
|
||||
# HTTP 超时要盖过环境侧的执行超时，否则代码还在跑我们就先断了连接，表现成一次
|
||||
# 假的网络故障。留 30 秒余量给序列化与调度。
|
||||
#
|
||||
# trust_env=False 同样是必须的，理由见 ContainerPool._await_ready。
|
||||
self._client = httpx.AsyncClient(timeout=self._execution_timeout_s + 30.0, trust_env=False)
|
||||
try:
|
||||
await self._pool.start()
|
||||
except BaseException:
|
||||
await self._client.aclose()
|
||||
self._client = None
|
||||
raise
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""停止容器池并关闭 HTTP 连接。
|
||||
|
||||
try/finally 保证连接一定被关掉：`pool.stop()` 也会失败（docker 二进制不在时
|
||||
`create_subprocess_exec` 直接抛 FileNotFoundError），那时连接不能跟着漏掉。
|
||||
"""
|
||||
try:
|
||||
await self._pool.stop()
|
||||
finally:
|
||||
if self._client is not None:
|
||||
await self._client.aclose()
|
||||
self._client = None
|
||||
|
||||
async def __aenter__(self) -> Self:
|
||||
await self.start()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc_info: object) -> None:
|
||||
await self.stop()
|
||||
|
||||
def list_task_ids(self, split: str) -> list[str]:
|
||||
"""列出某个数据划分下的全部题目 ID。
|
||||
|
||||
纯本地读文件，不需要容器——挑题、数题的场合不该为了拿一串 ID 去起一堆容器。
|
||||
"""
|
||||
dataset_file = self._datasets_dir / f"{split}.txt"
|
||||
if not dataset_file.exists():
|
||||
available = sorted(p.stem for p in self._datasets_dir.glob("*.txt"))
|
||||
raise FileNotFoundError(f"找不到划分文件 {dataset_file}；现有划分：{available}")
|
||||
task_ids: list[str] = []
|
||||
for line in dataset_file.read_text(encoding="utf-8").splitlines():
|
||||
entry = line.strip()
|
||||
if not entry:
|
||||
continue
|
||||
# 划分文件里的条目可能带 ":标签" 后缀，标签不是题目 ID 的一部分。
|
||||
task_ids.append(entry.split(":")[0])
|
||||
return task_ids
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def session(self, task_id: str) -> AsyncIterator[AppWorldSession]:
|
||||
"""开一次会话：租一个容器、实例化这道题，退出时关闭会话并归还容器。
|
||||
|
||||
**评测必须在退出这个上下文之前做**，退出后环境状态就销毁了。
|
||||
"""
|
||||
client = self._require_client()
|
||||
async with self._pool.lease() as port:
|
||||
base_url = f"http://127.0.0.1:{port}"
|
||||
payload = await _post_json(
|
||||
client,
|
||||
base_url,
|
||||
"/initialize",
|
||||
{"task_id": task_id, **self._init_params()},
|
||||
)
|
||||
if not isinstance(payload, dict):
|
||||
raise AppWorldError(f"initialize 期望返回对象，收到 {type(payload).__name__}")
|
||||
session = AppWorldSession(
|
||||
client=client,
|
||||
base_url=base_url,
|
||||
task_id=task_id,
|
||||
instruction=_require(payload, "instruction", task_id),
|
||||
supervisor=_require(payload, "supervisor", task_id),
|
||||
datetime=_require(payload, "datetime", task_id),
|
||||
)
|
||||
try:
|
||||
yield session
|
||||
finally:
|
||||
await self._close_quietly(session, port)
|
||||
|
||||
# -- 内部 ----------------------------------------------------------------
|
||||
|
||||
async def _close_quietly(self, session: AppWorldSession, port: int) -> None:
|
||||
"""关闭会话；失败只记账不抛，但同一个容器连续失败到阈值就让整轮压测停下来。
|
||||
|
||||
不抛的理由：走到这里时这道题的结果通常已经拿到手了，为一次清理失败丢掉整道题的
|
||||
数据不划算。
|
||||
|
||||
但「不抛」需要一张真的安全网，因为**这个故障不会自己暴露**。AppWorld 的
|
||||
`/initialize` 拿到请求后第一件事就是覆盖全局的 world 变量，既不检查旧的还开着
|
||||
没有、也不关它。所以下一个租户不会报错，只会让上一个 world 的资源永久留在容器
|
||||
里。一个容器在压测里要连续跑成百上千个会话，这是会累积到 fd 耗尽的慢性故障，而
|
||||
唯一的线索只有一行日志。
|
||||
|
||||
计数按端口分开，理由见 `_close_failures` 的定义。
|
||||
"""
|
||||
try:
|
||||
await session.close()
|
||||
except asyncio.CancelledError:
|
||||
# 取消要原样传播：它不是「关闭失败」，记进失败计数会让一次 Ctrl-C 把安全网
|
||||
# 的阈值推高，而真正的泄漏还没发生。
|
||||
raise
|
||||
except Exception as exc:
|
||||
failures = self._close_failures.get(port, 0) + 1
|
||||
self._close_failures[port] = failures
|
||||
logger.error(
|
||||
"题目 %s 的会话关闭失败（容器端口 %d，该容器已连续失败 %d 次）：%s",
|
||||
session.task_id,
|
||||
port,
|
||||
failures,
|
||||
exc,
|
||||
)
|
||||
if failures >= _MAX_CLOSE_FAILURES:
|
||||
raise AppWorldError(
|
||||
f"端口 {port} 上的容器会话关闭已连续失败 {failures} 次。它内部的资源"
|
||||
f"正在泄漏（AppWorld 的 /initialize 不会替我们清理旧会话），继续跑下去，"
|
||||
f"落在这个容器上的任务都会产出在一个不健康的环境里"
|
||||
) from exc
|
||||
else:
|
||||
self._close_failures.pop(port, None)
|
||||
|
||||
def _init_params(self) -> dict[str, Any]:
|
||||
"""传给 `/initialize` 的参数。
|
||||
|
||||
**会影响 agent 行为的开关全部显式给值**，不沿用 AppWorld 的默认——默认值会随它
|
||||
的版本变化，而这些开关直接改变 agent 能做什么。除 `max_interactions` 外，其余取值
|
||||
都等于 0.1.3 的默认值；写出来是为了防版本漂移，不是为了偏离官方设置。
|
||||
|
||||
官方标注「仅供测试、不应改动」的六个参数（`raise_on_extra_parameters`、
|
||||
`import_utils`、`parse_datetimes`、`allow_datetime_change`、`add_login_shortcut`、
|
||||
`munchify_response`）不传，沿用它们的默认值。
|
||||
|
||||
**刻意不传 `ground_truth_mode`。** AppWorld 0.1.3 的 `AppWorldInitDefaults` 里那个
|
||||
字段写的是 ``Literal["full" "minimal"]``，两个字面量之间漏了逗号，被 Python 拼接成
|
||||
单一取值 ``"fullminimal"``。默认值 ``"minimal"`` 并不在这个 Literal 里，只是因为
|
||||
pydantic 不校验默认值才没暴露；一旦显式传 ``"minimal"``，请求体就会走校验然后被拒。
|
||||
"""
|
||||
return {
|
||||
"experiment_name": self._experiment_name,
|
||||
"max_interactions": self._max_interactions,
|
||||
"max_api_calls_per_interaction": 1000,
|
||||
"raise_on_unsafe_syntax": True,
|
||||
"null_patch_unsafe_execution": True,
|
||||
"load_ground_truth": True, # evaluate 需要它
|
||||
"raise_on_failure": True,
|
||||
# 环境自身的随机种子固定不动。它控制的是环境初始状态与 API 响应里的随机成分，
|
||||
# 也就是「题目本身」——跟着别的什么东西变的话，两次压测面对的就不是同一道题，
|
||||
# 而「同一道题这次慢了」正是压测要看的东西。100 是 AppWorld 的默认值。
|
||||
"random_seed": 100,
|
||||
"timeout_seconds": self._execution_timeout_s,
|
||||
"show_api_response_schemas": True,
|
||||
"gc_threshold": 500000,
|
||||
}
|
||||
|
||||
def _require_client(self) -> httpx.AsyncClient:
|
||||
"""取 HTTP 连接；没 start 过就用是调用方的顺序错误，直接说清楚。"""
|
||||
if self._client is None:
|
||||
raise AppWorldError("环境尚未启动，先调用 start() 或用 async with")
|
||||
return self._client
|
||||
|
||||
def _verify_data_layout(self) -> None:
|
||||
"""启动前确认数据在位，免得每个容器各自失败一次才发现是数据没下。"""
|
||||
for path in (self._datasets_dir, self._tasks_dir):
|
||||
if not path.is_dir():
|
||||
raise FileNotFoundError(
|
||||
f"AppWorld 数据缺失：{path} 不存在。在装有 appworld 的环境里跑："
|
||||
f"appworld install && appworld download data --root {self._data_root}"
|
||||
)
|
||||
|
||||
|
||||
def _success_of(tracker: Mapping[str, Any], task_id: str) -> bool:
|
||||
"""从评测结果里取这道题过没过。
|
||||
|
||||
官方评测器先看 `success`、再回落到 `passes_fully`（字段迁移中的兼容写法）。两个都没有
|
||||
就报错而不是当成失败——把「读不出分数」记成「没做对」会让一次接口变更表现成成绩暴跌，
|
||||
而且查不出原因。
|
||||
"""
|
||||
for key in ("success", "passes_fully"):
|
||||
if key in tracker:
|
||||
return bool(tracker[key])
|
||||
raise AppWorldError(
|
||||
f"题目 {task_id} 的评测结果里既没有 success 也没有 passes_fully，"
|
||||
f"实有字段：{sorted(tracker)}"
|
||||
)
|
||||
|
||||
|
||||
def _require(payload: Mapping[str, Any], key: str, task_id: str) -> Any:
|
||||
"""从环境返回里取一个必需字段，缺了就报错。"""
|
||||
if key not in payload:
|
||||
raise AppWorldError(f"题目 {task_id} 的环境返回缺字段 {key!r}，实有字段：{sorted(payload)}")
|
||||
return payload[key]
|
||||
|
||||
|
||||
async def _post_json(
|
||||
client: httpx.AsyncClient,
|
||||
base_url: str,
|
||||
path: str,
|
||||
body: Mapping[str, Any],
|
||||
) -> Any:
|
||||
"""打一个 POST，检查状态码，剥掉 environment server 统一的 output 包装。"""
|
||||
response = await client.post(f"{base_url}{path}", json=dict(body))
|
||||
if response.is_error:
|
||||
raise AppWorldError(f"{path} 返回 HTTP {response.status_code}：{response.text[:2000]}")
|
||||
payload = response.json()
|
||||
if not isinstance(payload, dict) or "output" not in payload:
|
||||
raise AppWorldError(f"{path} 的返回不是 {{'output': ...}} 的形态：{str(payload)[:500]}")
|
||||
return payload["output"]
|
||||
|
||||
|
||||
async def _run(command: Sequence[str]) -> tuple[int, str, str]:
|
||||
"""跑一条 docker 命令，返回 (退出码, stdout, stderr)。"""
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
*command,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
stdout, stderr = await process.communicate()
|
||||
assert process.returncode is not None # communicate() 返回后必然已退出
|
||||
return process.returncode, stdout.decode(errors="replace"), stderr.decode(errors="replace")
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CONTAINER_NAME_PREFIX",
|
||||
"DEFAULT_IMAGE",
|
||||
"DEFAULT_PORT_BASE",
|
||||
"AppWorldError",
|
||||
"AppWorldPool",
|
||||
"AppWorldSession",
|
||||
"ContainerPool",
|
||||
"ContainerPoolError",
|
||||
"Mount",
|
||||
"TaskScore",
|
||||
"container_args_for_port",
|
||||
]
|
||||
Reference in New Issue
Block a user