diff --git a/Makefile b/Makefile index 23aae4a..07201b1 100644 --- a/Makefile +++ b/Makefile @@ -11,15 +11,15 @@ install: # lint 带 --fix 会改工作区,check 只读。CI 用 check,人手修用 lint。 lint: - $(RUN) ruff check src/ tests/ --fix + $(RUN) ruff check src/ tests/ tools/ --fix $(RUN) lint-imports format: - $(RUN) ruff format src/ tests/ + $(RUN) ruff format src/ tests/ tools/ check: - $(RUN) ruff format --check src/ tests/ - $(RUN) ruff check src/ tests/ + $(RUN) ruff format --check src/ tests/ tools/ + $(RUN) ruff check src/ tests/ tools/ $(RUN) lint-imports test: diff --git a/tools/soak/__init__.py b/tools/soak/__init__.py new file mode 100644 index 0000000..b8d4ebf --- /dev/null +++ b/tools/soak/__init__.py @@ -0,0 +1,9 @@ +"""压测 harness:给 PolyLoop 造真实负载用的一次性工具,**不是库的一部分**。 + +`pyproject.toml` 的 `[tool.setuptools.packages.find]` 只收 `src/`,所以这个目录不会 +被打包、不会随 `pip install polyloop` 装到下游手里。它可以依赖 `src/polyloop/` 明确 +拒绝的东西(httpx、docker 命令行、某个具体 benchmark 的数据布局),因为它的失败只会 +影响我们自己的压测,不会击穿任何下游项目。 + +**反过来的方向是禁止的**:`src/polyloop/` 里任何一处都不许 import `tools.soak`。 +""" diff --git a/tools/soak/appworld.py b/tools/soak/appworld.py new file mode 100644 index 0000000..a381540 --- /dev/null +++ b/tools/soak/appworld.py @@ -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", +] diff --git a/tools/soak/check_appworld.py b/tools/soak/check_appworld.py new file mode 100644 index 0000000..2783653 --- /dev/null +++ b/tools/soak/check_appworld.py @@ -0,0 +1,134 @@ +"""AppWorld 环境层的离线冒烟:起 1 个容器,跑通一整条会话链路,**一次模型调用都不打**。 + +跑法(在仓库根目录):: + + PYTHONUNBUFFERED=1 conda run --live-stream -n PolyLoop \\ + python -m tools.soak.check_appworld --data-root /path/to/appworld + +用 `-m` 而不是直接给文件路径:直接跑文件时 `tools` 不在 `sys.path` 上, +`from tools.soak.appworld import ...` 会 ImportError。 + +它验的是环境层自己:容器起得来、题实例化得了、代码执行得了、评测调得通、容器删得干净。 +执行的三段代码是写死的,其中一段故意写错——那一段验的是「代码报错不抛异常,错误文本原样 +回来」这条行为,因为它是上层循环最依赖的一条,而正常路径验不出来。 + +退出码 0 表示全绿。任何一步失败都以非零退出并把原始报错打出来,**不做任何降级**。 +""" + +from __future__ import annotations + +import argparse +import asyncio +import logging +import sys +from pathlib import Path + +from tools.soak.appworld import AppWorldPool + +#: 冒烟用的题目。取 train 划分的第一条而不是写死一个 ID:题集换过一次(发布的 train 是 +#: 90 题,原论文写的是 105 题),写死的 ID 有一天会变成一句「找不到这道题」,而那时看起来 +#: 像是环境坏了。 +_SPLIT = "train" + +#: 三段写死的代码。顺序有意义:先确认 API 文档读得到(环境的数据挂载没问题),再确认 +#: 有状态执行器跨步保留变量(这是 AppWorld 与 docker-exec 型环境的关键区别),最后确认 +#: 报错不会被压成异常。 +_PROBES: tuple[tuple[str, str], ...] = ( + ("列出可用的 app", "print(apis.api_docs.show_app_descriptions())"), + ("定义一个变量", "soak_marker = 6 * 7"), + ("读回上一步的变量(验有状态执行器)", "print(soak_marker)"), + ("故意写错(验错误不抛异常)", "print(this_name_does_not_exist)"), +) + + +def _line(title: str) -> None: + print(f"\n=== {title} ===", flush=True) + + +async def _smoke(data_root: Path, port_base: int) -> int: + pool = AppWorldPool(data_root=data_root, size=1, port_base=port_base) + _line("池配置") + print(f"数据根目录:{data_root}(只读挂载)") + print(f"输出目录: {pool.outputs_dir}") + print(f"宿主端口: {pool.ports}") + + async with pool: + task_ids = pool.list_task_ids(_SPLIT) + task_id = task_ids[0] + _line(f"划分 {_SPLIT}") + print(f"共 {len(task_ids)} 道题,取第一道:{task_id}") + + async with pool.session(task_id) as session: + _line("初始化") + print(f"容器地址:{session.base_url}") + print(f"虚拟时间:{session.datetime}") + print(f"主管: {session.supervisor}") + print(f"题面: {session.instruction.strip()[:400]}") + + for title, code in _PROBES: + _line(f"执行:{title}") + print(f"$ {code}") + output = await session.execute(code) + print(output.strip()[:1200] or "<无输出>") + + _line("问一次完成没有") + print(f"is_done() = {await session.is_done()}(没做题,应为 False)") + + _line("评测") + score = await session.evaluate() + print(f"success = {score.success}(没做题,应为 False)") + print(f"执行次数 = {score.n_executions}") + print(f"detail 的键 = {sorted(score.detail)}") + + _line("容器残留检查") + return await _report_leftovers() + + +async def _report_leftovers() -> int: + """确认压测的容器一个都没剩下。剩了就非零退出——它会占着端口让下一次跑直接失败。""" + process = await asyncio.create_subprocess_exec( + "docker", + "ps", + "-a", + "--filter", + "name=polyloop-soak-appworld", + "--format", + "{{.Names}} {{.Status}}", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + stdout, stderr = await process.communicate() + if process.returncode != 0: + print(f"docker ps 失败(退出码 {process.returncode}):{stderr.decode(errors='replace')}") + return 1 + leftovers = stdout.decode(errors="replace").strip() + if leftovers: + print(f"仍有残留容器:\n{leftovers}") + return 1 + print("docker ps -a 里没有 polyloop-soak-appworld 开头的容器,干净。") + return 0 + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--data-root", + type=Path, + required=True, + help="AppWorld 数据根目录,下面应有 data/datasets 与 data/tasks", + ) + parser.add_argument( + "--port-base", + type=int, + default=8200, + help="容器池起始宿主端口(默认 8200,避开 dissect 的 8100)", + ) + args = parser.parse_args() + # 容器池的告警走 logging,默认级别是 WARNING 且没有 handler,会被静默丢掉——而 + # 「清理容器失败」正是这条路上唯一的线索。 + logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s") + return asyncio.run(_smoke(args.data_root.resolve(), args.port_base)) + + +if __name__ == "__main__": + sys.exit(main())