Files
AgentLab/backend/tests/test_multiprocess_e2e.py
T

481 lines
22 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""多进程端到端(docs/design/orchestration-round/00-总方案.md §5 W2-b 的验收):PostgreSQL 上起两个工作进程(uvicorn,A、B),
指着同一个库、同一个数据目录,用演示模型(Mock)跑:
1. 两个进程都登记了、只有一个当调度班;
2. 运行由接到请求的那个进程认领;另一个进程上的 WebSocket 收得到实时事件(含不落库的流式增量);同一个审批决定同时发给两个
进程只执行一次;两边的事件流都按 seq 递增、不重不漏,封存核对得上;
3. 在 A 上跑的运行,从 B 停止:照本进程的停止办(run.cancelled 带上是谁停的);
4. 定时任务到点只触发一次(两个进程都起着调度循环,只有当班的做);
5. 评测和问数据并发时优先级生效:全站只给 1 个名额,评测的作答运行(系统类,由本测试进程作为第三个登记的工作进程以评测的
来源发起——工作进程里发起评测要真的去调模型建图)排着队时,后来的问数据运行(交互类,经 A 的接口发起)先拿到名额;
6. 每日花费上限跨进程不越过:两个进程同时往同一本账上记,记下的合计不超过上限;
7. 强杀 B:租约过期后,B 上在跑的运行被照实标中断(可从断点继续),停在审批上的运行照常能在 A 上批、接着跑完;B 当着调度班的
话 A 接班。
缺省不跑:要 AGENTLAB_MULTIPROCESS_E2E=1 和 PostgreSQL 测试库(scripts/test-multiprocess.sh 起临时容器、设好这两样)。
"""
from __future__ import annotations
import asyncio
import json
import os
import signal
import socket
import subprocess
import sys
import time
from pathlib import Path
from typing import Any
import pytest
pytestmark = pytest.mark.skipif(
os.environ.get("AGENTLAB_MULTIPROCESS_E2E") != "1" or not os.environ.get("AGENTLAB_TEST_DATABASE_URL"),
reason="多进程端到端缺省不跑:用 scripts/test-multiprocess.sh(要本机 docker 的 postgres:16)")
BACKEND = Path(__file__).resolve().parents[1]
#: 工作进程的租约和心跳:调小,强杀之后几秒就能看到接手
HEARTBEAT_S = 0.5
LEASE_S = 2.0
def _free_port() -> int:
with socket.socket() as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
def _node(nid: str, ntype: str, **config: Any) -> dict[str, Any]:
return {"id": nid, "type": ntype, "position": {"x": 0, "y": 0}, "data": {"label": nid, "config": config}}
def _chain(*nodes: dict[str, Any]) -> dict[str, Any]:
return {"nodes": list(nodes), "edges": [{"source": a["id"], "target": b["id"]} for a, b in zip(nodes, nodes[1:])]}
def _sleep_code(seconds: float) -> dict[str, Any]:
return _node("wait", "code", language="python", code=f"import time\ntime.sleep({seconds})\nprint('ok')")
def slow_graph(seconds: float) -> dict[str, Any]:
return _chain(_node("start", "input", fields=[{"name": "question"}]), _sleep_code(seconds),
_node("out", "output", fields=[{"name": "r", "value": "{{ nodes.wait.text }}"}]))
APPROVAL_GRAPH = _chain(
_node("start", "input", fields=[{"name": "question"}]),
_sleep_code(1.0),
_node("gate", "human", mode="approve", title="看一眼再接着"),
_node("write", "llm", prompt="用三句话写一段说明"),
_node("out", "output", fields=[{"name": "结果", "value": "{{ nodes.write.text }}"}]),
)
class Worker:
def __init__(self, name: str, env: dict[str, str], log_dir: Path) -> None:
self.name = name
self.port = _free_port()
self.base = f"http://127.0.0.1:{self.port}"
self.log = log_dir / f"worker-{name}.log"
self.proc = subprocess.Popen(
[sys.executable, "-m", "uvicorn", "app.main:app", "--host", "127.0.0.1", "--port", str(self.port)],
cwd=str(BACKEND), env=env, stdout=open(self.log, "wb"), stderr=subprocess.STDOUT)
self.owner: str | None = None
def kill(self) -> None:
if self.proc.poll() is None:
self.proc.send_signal(signal.SIGKILL)
self.proc.wait(timeout=10)
def stop(self) -> None:
if self.proc.poll() is None:
self.proc.send_signal(signal.SIGTERM)
try:
self.proc.wait(timeout=15)
except subprocess.TimeoutExpired:
self.kill()
def tail(self) -> str:
try:
return self.log.read_text(encoding="utf-8", errors="replace")[-4000:]
except OSError:
return ""
def _worker_env() -> dict[str, str]:
env = {k: v for k, v in os.environ.items()
if k not in ("ANTHROPIC_API_KEY", "ANTHROPIC_BASE_URL", "ANTHROPIC_AUTH_TOKEN", "OPENAI_API_KEY",
"PYTEST_XDIST_WORKER", "PYTEST_CURRENT_TEST")}
env.update({
"PYTHONPATH": str(BACKEND),
# 应用库、数据目录:和本测试进程同一个(conftest 已经设好)
"AGENTLAB_DATABASE_URL": os.environ["AGENTLAB_DATABASE_URL"],
"AGENTLAB_DATA_DIR": os.environ["AGENTLAB_DATA_DIR"],
"AGENTLAB_DB_POOL_SIZE": "5",
"AGENTLAB_SCHEDULER": "1",
"AGENTLAB_SCHEDULER_TICK_S": "0.5",
"AGENTLAB_WORKER_HEARTBEAT_S": str(HEARTBEAT_S),
"AGENTLAB_WORKER_LEASE_S": str(LEASE_S),
"AGENTLAB_SANDBOX_BACKEND": "local",
"AGENTLAB_AUTH_MODE": "off",
})
return env
@pytest.fixture(scope="module")
def workers(tmp_path_factory):
import httpx
log_dir = tmp_path_factory.mktemp("multiprocess")
env = _worker_env()
started = [Worker("A", env, log_dir), Worker("B", env, log_dir)]
try:
deadline = time.monotonic() + 90
for w in started:
while True:
if w.proc.poll() is not None:
raise AssertionError(f"工作进程 {w.name} 没起来:\n{w.tail()}")
try:
if httpx.get(f"{w.base}/api/health", timeout=2).status_code == 200:
break
except httpx.HTTPError:
pass
if time.monotonic() > deadline:
raise AssertionError(f"工作进程 {w.name} 90 秒没起来:\n{w.tail()}")
time.sleep(0.3)
w.owner = httpx.get(f"{w.base}/api/system", timeout=10).json()["cluster"]["me"]
yield {w.name: w for w in started}
finally:
for w in started:
w.stop()
for w in started:
print(f"---- 工作进程 {w.name} 的日志(末尾)----\n{w.tail()}")
async def _client():
import httpx
return httpx.AsyncClient(timeout=30)
async def _until(check, timeout: float = 20.0, what: str = "") -> Any:
deadline = time.monotonic() + timeout
last = None
while time.monotonic() < deadline:
last = await check()
if last:
return last
await asyncio.sleep(0.1)
raise AssertionError(f"等超时了:{what}(最后一次:{last!r})")
async def _run(client, worker: Worker, run_id: str) -> dict[str, Any]:
return (await client.get(f"{worker.base}/api/runs/{run_id}")).json()
async def _db(sql: str, *args: Any) -> list[Any]:
from sqlalchemy import text
from app.db.base import SessionLocal
async with SessionLocal() as session:
return list((await session.execute(text(sql), _params(args))).all())
def _params(args: tuple[Any, ...]) -> dict[str, Any]:
return {f"p{i}": v for i, v in enumerate(args)}
class Stream:
"""一条 WebSocket(/api/runs/{id}/stream)收到的全部消息。"""
def __init__(self, worker: Worker, run_id: str) -> None:
self.url = f"ws://127.0.0.1:{worker.port}/api/runs/{run_id}/stream?after=0"
self.messages: list[dict[str, Any]] = []
self.task: asyncio.Task[None] | None = None
async def start(self) -> Stream:
import websockets
ready = asyncio.Event()
async def pump() -> None:
async with websockets.connect(self.url) as ws:
ready.set()
async for raw in ws:
msg = json.loads(raw)
self.messages.append(msg)
if msg.get("type") == "stream.end":
return
self.task = asyncio.create_task(pump())
await asyncio.wait_for(ready.wait(), timeout=10)
return self
def types(self) -> list[str]:
return [m.get("type") for m in self.messages]
def seqs(self) -> list[int]:
"""落库事件的 seq(不落库的流式增量另算)。"""
return [int(m["seq"]) for m in self.messages if m.get("seq") and m.get("type") not in ("llm.token",
"llm.thinking.delta")]
async def ended(self, timeout: float = 30.0) -> None:
assert self.task is not None
await asyncio.wait_for(self.task, timeout=timeout)
# --------------------------------------------------------------------------
async def test_1_two_workers_one_leader(workers):
a, b = workers["A"], workers["B"]
async with await _client() as client:
async def both() -> Any:
info = (await client.get(f"{a.base}/api/system")).json()["cluster"]
alive = {w["id"] for w in info["workers"] if w["alive"]}
return info if {a.owner, b.owner} <= alive else None
info = await _until(both, what="两个进程都登记")
assert info["database"] == "postgresql"
assert info["leader"]["owner"] in (a.owner, b.owner)
other = (await client.get(f"{b.base}/api/system")).json()["cluster"]
assert other["leader"]["owner"] == info["leader"]["owner"], "两边看到的当班是同一个"
async def test_2_live_events_across_workers_and_one_decision(workers):
a, b = workers["A"], workers["B"]
async with await _client() as client:
made = (await client.post(f"{a.base}/api/runs", json={"graph": APPROVAL_GRAPH, "input": {"question": "q"},
"origin": "canvas"})).json()
run_id = made["id"]
# 两边各接一条流(趁它还在跑:停在审批上之后才接的流回放完就结束,要等恢复之后重新接)
on_b = await Stream(b, run_id).start()
on_a = await Stream(a, run_id).start()
owner = await _until(lambda: _owner_while_active(run_id), what="运行被认领")
assert owner == a.owner, "运行由接到请求的进程认领"
async def paused() -> Any:
return (await _run(client, b, run_id))["status"] == "interrupted"
await _until(paused, what="停在审批上")
await _until(lambda: _has(on_b, "run.interrupted"), what="B 上的流收到 run.interrupted")
# 同一个决定同时发给两个进程:只执行一次
decided = await asyncio.gather(
client.post(f"{a.base}/api/runs/{run_id}/resume", json={"response": {"approved": True}}),
client.post(f"{b.base}/api/runs/{run_id}/resume", json={"response": {"approved": True}}))
assert sorted(r.status_code for r in decided)[0] == 200, [r.text for r in decided]
async def finished() -> Any:
return (await _run(client, a, run_id))["status"] in ("succeeded", "failed")
await _until(finished, timeout=60, what="运行跑完")
assert (await _run(client, a, run_id))["status"] == "succeeded"
await on_a.ended()
await on_b.ended()
for stream in (on_a, on_b):
assert "run.finished" in stream.types(), stream.types()
seqs = stream.seqs()
assert seqs == sorted(set(seqs)), f"事件流有重复或乱序:{seqs}"
started = [m for m in on_a.messages if m.get("type") == "node.started" and m.get("node_id") == "write"]
assert len(started) == 1, "审批之后的那一步只执行了一次"
# 执行后半段的进程不是谁、另一个进程上的流照样收到了流式增量(不落库,只转给在看的进程)
assert any(m.get("type") == "llm.token" for m in on_a.messages) and \
any(m.get("type") == "llm.token" for m in on_b.messages), "两边都收到了正文增量"
rows = await _db("SELECT seq, count(*) FROM run_events WHERE run_id = :p0 GROUP BY seq HAVING count(*) > 1", run_id)
assert rows == [], f"库里的事件撞了号:{rows}"
verified = (await client.get(f"{b.base}/api/runs/{run_id}/verify")).json()
assert verified["ok"] is True, verified
async def _owner_while_active(run_id: str) -> str | None:
rows = await _db("SELECT owner, status FROM runs WHERE id = :p0", run_id)
if rows and rows[0][1] in ("running", "queued") and rows[0][0]:
return rows[0][0]
return None
async def _has(stream: Stream, kind: str) -> bool:
return kind in stream.types()
async def test_3_stop_from_the_other_worker(workers):
a, b = workers["A"], workers["B"]
async with await _client() as client:
run_id = (await client.post(f"{a.base}/api/runs", json={"graph": slow_graph(20), "input": {"question": "q"},
"origin": "canvas"})).json()["id"]
await _until(lambda: _owner_while_active(run_id), what="运行被认领")
t0 = time.monotonic()
stop = await client.post(f"{b.base}/api/runs/{run_id}/cancel", headers={"X-Actor": "e2e-stopper"})
assert stop.status_code == 200 and stop.json()["status"] == "stopping", stop.text
async def cancelled() -> Any:
row = await _run(client, b, run_id)
return row if row["status"] == "cancelled" else None
row = await _until(cancelled, timeout=15, what="停止生效")
assert row["end_reason"] == "stopped"
assert time.monotonic() - t0 < 10, "不是等它跑完才停的"
events = (await client.get(f"{a.base}/api/runs/{run_id}/events")).json()
stopped = [e for e in events if e["type"] == "run.cancelled"]
assert stopped and stopped[-1]["data"].get("actor") == "e2e-stopper", stopped
async def test_4_a_schedule_fires_once(workers):
a = workers["A"]
async with await _client() as client:
wf = (await client.post(f"{a.base}/api/workflows", json={"name": "定时验收", "graph": slow_graph(0.1)})).json()
made = await client.post(f"{a.base}/api/schedules", json={
"name": "每分钟一次", "workflow_id": wf["id"], "cron": "* * * * *", "timezone": "Asia/Shanghai", "run_class": "exploratory",
"input_template": {"question": "定时"}})
assert made.status_code == 201, made.text
schedule_id = made.json()["id"]
# 拨到「刚过点」:两个进程的调度循环都在转,只有当班的那个做
from sqlalchemy import text
from app.db.base import SessionLocal, utcnow
due = utcnow().replace(microsecond=0)
async with SessionLocal() as session:
await session.execute(text("UPDATE schedules SET next_run_at = :t WHERE id = :i"), {"t": due.replace(tzinfo=None),
"i": schedule_id})
await session.commit()
async def fired() -> Any:
rows = await _db("SELECT id FROM schedule_fires WHERE schedule_id = :p0", schedule_id)
return rows or None
await _until(fired, timeout=15, what="定时触发")
await asyncio.sleep(3) # 再等几拍:不会再发一次
fires = await _db("SELECT kind, status, run_id FROM schedule_fires WHERE schedule_id = :p0", schedule_id)
runs = await _db("SELECT id FROM runs WHERE schedule_id = :p0", schedule_id)
assert len(fires) == 1 and len(runs) == 1, (fires, runs)
await client.patch(f"{a.base}/api/schedules/{schedule_id}", json={"enabled": False})
async def test_5_chat_goes_before_queued_evaluation_runs(workers, monkeypatch):
"""全站 1 个名额:评测的作答运行(系统类)排着两个,问数据的运行(交互类)后来,先拿到名额。"""
from app import cluster
from app.core.config import settings
from app.cluster import slots
from app.cluster.relay import relay
from app.db.base import SessionLocal
from app.db.models import Run
from app.engine.runner import run_manager
a = workers["A"]
async with await _client() as client:
put = await client.put(f"{a.base}/api/settings", json={"values": {"concurrency": {
"max_total": 1, "classes": {"interactive": {"limit": 1, "priority": 40}, "system": {"limit": 1, "priority": 10}}}}})
assert put.status_code == 200, put.text
slots.forget_policy()
monkeypatch.setattr(settings, "sandbox_backend", "local")
await cluster.start() # 本测试进程登记成第三个工作进程,替评测发起作答运行
await relay.start()
await run_manager.setup()
try:
evals = [await run_manager.start(graph=slow_graph(1.5), input_payload={"question": f"评测 {i}"},
origin_kind="eval", origin={"kind": "eval"}) for i in range(3)]
async def first_running() -> Any:
async with SessionLocal() as session:
return (await session.get(Run, evals[0].id)).status == "running"
await _until(first_running, what="第一题的作答运行拿到名额")
await asyncio.sleep(0.3)
chat_id = (await client.post(f"{a.base}/api/runs", json={
"graph": slow_graph(0.2), "input": {"question": "问数据"}, "origin": "chat"})).json()["id"]
async def all_done() -> Any:
async with SessionLocal() as session:
rows = [await session.get(Run, rid, populate_existing=True) for rid in [*[e.id for e in evals], chat_id]]
return rows if all(r.status in ("succeeded", "failed") for r in rows) else None
rows = await _until(all_done, timeout=60, what="都跑完")
assert all(r.status == "succeeded" for r in rows), [(r.id, r.status, r.error) for r in rows]
e1, e2, e3, chat = rows
assert e2.created_at < chat.created_at and e3.created_at < chat.created_at, "评测的两题先来、排着队"
assert e1.started_at < chat.started_at < e2.started_at < e3.started_at, \
[(r.origin_kind, r.created_at, r.started_at) for r in rows]
finally:
await run_manager.shutdown()
await relay.stop()
await cluster.stop()
await client.put(f"{a.base}/api/settings", json={"values": {"concurrency": slots.defaults()}})
def test_6_the_daily_cap_holds_across_processes(workers):
"""两个进程同时往评测那本账上各记 20 次、每次 0.1 美元,上限 1 美元:合计记下的正好 10 次、不超过上限。"""
script = (
"import asyncio, json\n"
"from app.engine.spend import Ledger\n"
"async def main():\n"
" book = Ledger('e2e_spend')\n"
" got = await asyncio.gather(*(book.charge(0.1, limit=1.0, priced=True) for _ in range(20)))\n"
" print(json.dumps(sum(1 for d in got if d)))\n"
"asyncio.run(main())\n")
env = {**_worker_env(), "AGENTLAB_SCHEDULER": "0"}
procs = [subprocess.Popen([sys.executable, "-c", script], cwd=str(BACKEND), env=env, stdout=subprocess.PIPE,
stderr=subprocess.PIPE) for _ in range(2)]
outs = [p.communicate(timeout=120) for p in procs]
assert all(p.returncode == 0 for p in procs), [o[1].decode()[-2000:] for o in outs]
charged = sum(int(o[0].decode().strip().splitlines()[-1]) for o in outs)
assert charged == 10, f"两个进程合计记下 {charged} 次(上限只容得下 10 次)"
total = asyncio.run(_spent())
assert total <= 1.0 + 1e-9, total
async def _spent() -> float:
from app.engine.spend import Ledger
return float((await Ledger("e2e_spend").daily())["usd"])
async def test_7_kill_a_worker(workers):
a, b = workers["A"], workers["B"]
async with await _client() as client:
leader_before = (await client.get(f"{a.base}/api/system")).json()["cluster"]["leader"]["owner"]
# B 上:一次在跑(会被照实标中断),一次停在审批上(不属于任何进程,B 没了照样能在 A 上批)
busy = (await client.post(f"{b.base}/api/runs", json={"graph": slow_graph(30), "input": {"question": "q"},
"origin": "canvas"})).json()["id"]
paused = (await client.post(f"{b.base}/api/runs", json={"graph": APPROVAL_GRAPH, "input": {"question": "q"},
"origin": "canvas"})).json()["id"]
await _until(lambda: _owner_while_active(busy), what="B 认领了运行")
async def waiting() -> Any:
return (await _run(client, a, paused))["status"] == "interrupted"
await _until(waiting, what="停在审批上")
b.kill()
t0 = time.monotonic()
async def taken_over() -> Any:
row = await _run(client, a, busy)
return row if row["status"] == "interrupted" else None
row = await _until(taken_over, timeout=LEASE_S * 6 + 10, what="租约过期后被接手")
assert "可从断点继续运行" in (row["error"] or "")
events = (await client.get(f"{a.base}/api/runs/{busy}/events")).json()
assert any(e["type"] == "log" and e["data"].get("code") == "server_shutdown" for e in events)
print(f"强杀 B 之后 {time.monotonic() - t0:.1f} 秒被接手")
info = (await client.get(f"{a.base}/api/system")).json()["cluster"]
assert b.owner not in {w["id"] for w in info["workers"] if w["alive"]}
if leader_before == b.owner:
async def a_leads() -> Any:
got = (await client.get(f"{a.base}/api/system")).json()["cluster"]["leader"]
return got["owner"] == a.owner
await _until(a_leads, timeout=LEASE_S * 6, what="A 接班")
# 停在审批上的那一次:在 A 上批,接着跑完
ok = await client.post(f"{a.base}/api/runs/{paused}/resume", json={"response": {"approved": True}})
assert ok.status_code == 200, ok.text
async def done() -> Any:
return (await _run(client, a, paused))["status"] == "succeeded"
await _until(done, timeout=60, what="批了之后在 A 上跑完")
# 被标中断的那一次:在 A 上能接着跑(这里接上就停掉,不等它睡完)
cont = await client.post(f"{a.base}/api/runs/{busy}/continue", json={})
assert cont.status_code == 200, cont.text
await client.post(f"{a.base}/api/runs/{busy}/cancel")