mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 23:11:58 +00:00
复核 R-b5、R-b6：转发运行事件、读缓存失效通知都从游标往回留一小段重读（运行 64 条、缓存 200 条），按发布过的 seq / 处理过的 id 去重——号先取后提交、号大的先提交时，号小的那一条后到也照样转、照样失效，不越过回放已经发过的那一段；合成之后超长的正文增量按字切成几条（二分找每段最长的前缀，不切坏 UTF-8），不再截掉。补测试：号小的后提交照样转出来、不重复；缓存失效 id 大的先到、小的后到都处理；超长增量切开拼回来一字不少
This commit is contained in:
ThearchyHelios
@@ -26,6 +26,10 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
_handlers: dict[str, Callable[[str | None], Any]] = {}
|
||||
_last_seen: int | None = None
|
||||
#: 往回多读几条（复核 R-b5）：id 是先取后提交的，id 大的先提交时游标会越过 id 小的那一条——每次从游标往回留这么多条重读，按处理过
|
||||
#: 的 id 去重
|
||||
LOOKBACK = 200
|
||||
_done: set[int] = set()
|
||||
_lock: asyncio.Lock | None = None
|
||||
|
||||
|
||||
@@ -58,12 +62,18 @@ async def broadcast(name: str, key: str | None = None) -> None:
|
||||
|
||||
async def start() -> None:
|
||||
"""起服务时：从现在的最后一条往后读（之前的变化，自己启动时读进来的已经是新的）。"""
|
||||
global _last_seen
|
||||
global _last_seen, _floor
|
||||
from app.db.models_cluster import CacheEvent
|
||||
|
||||
signals.on("cache", lambda _parts: asyncio.ensure_future(poll()))
|
||||
async with SessionLocal() as session:
|
||||
_last_seen = int(await session.scalar(select(func.coalesce(func.max(CacheEvent.id), 0))) or 0)
|
||||
_floor = _last_seen
|
||||
_done.clear()
|
||||
|
||||
|
||||
#: 起服务时已有的最后一条：往回重读不越过它（那之前的变化，自己启动时读进来的已经是新的）
|
||||
_floor = 0
|
||||
|
||||
|
||||
async def poll() -> int:
|
||||
@@ -77,10 +87,13 @@ async def poll() -> int:
|
||||
_lock = asyncio.Lock()
|
||||
async with _lock:
|
||||
async with SessionLocal() as session:
|
||||
rows = (await session.execute(select(CacheEvent).where(CacheEvent.id > _last_seen)
|
||||
rows = (await session.execute(select(CacheEvent).where(CacheEvent.id > max(_floor, _last_seen - LOOKBACK))
|
||||
.order_by(CacheEvent.id).limit(500))).scalars().all()
|
||||
done = 0
|
||||
for row in rows:
|
||||
if int(row.id) in _done:
|
||||
continue
|
||||
_done.add(int(row.id))
|
||||
_last_seen = max(_last_seen, int(row.id))
|
||||
if row.owner == identity.owner():
|
||||
continue
|
||||
@@ -94,6 +107,9 @@ async def poll() -> int:
|
||||
done += 1
|
||||
except Exception: # noqa: BLE001
|
||||
logger.exception("按别的进程的通知重读缓存「%s」出错", row.name)
|
||||
if len(_done) > 4 * LOOKBACK:
|
||||
low = _last_seen - 2 * LOOKBACK
|
||||
_done.difference_update({i for i in _done if i <= low})
|
||||
return done
|
||||
|
||||
|
||||
|
||||
@@ -43,6 +43,9 @@ EPHEMERAL_FLUSH_S = 0.05
|
||||
POLL_S = 2.0
|
||||
#: 一条载荷最多多少字节（PostgreSQL 的 NOTIFY 上限是 8000）
|
||||
PAYLOAD_MAX = 7000
|
||||
#: 往回多读几条（复核 R-b5）：号是先取后提交的，号大的先提交时游标会越过号小的那一条；每次从游标往回留这么多条重读，按已经
|
||||
#: 发布过的 seq 去重
|
||||
LOOKBACK = 64
|
||||
#: 不落库的事件（同 engine/runner._EPHEMERAL）
|
||||
EPHEMERAL = {EventType.LLM_TOKEN, EventType.LLM_THINKING_DELTA}
|
||||
_TERMINAL = ("succeeded", "failed", "cancelled")
|
||||
@@ -53,6 +56,10 @@ class Relay:
|
||||
#: 本进程里在跟随的运行：{run_id: 跟随的连接数}、游标（已经在本进程发布过的最大 seq）
|
||||
self._follow: dict[str, int] = {}
|
||||
self._cursor: dict[str, int] = {}
|
||||
#: 本进程已经发布过的 seq（游标往回 LOOKBACK 条以内的）：往回重读时去重
|
||||
self._seen: dict[str, set[int]] = {}
|
||||
#: 接上时回放到了哪一条：往回重读不越过它（那之前的回放已经发过）
|
||||
self._floor: dict[str, int] = {}
|
||||
self._pumping: dict[str, asyncio.Lock] = {}
|
||||
#: 本进程里在看的问数据一轮
|
||||
self._turns: dict[str, int] = {}
|
||||
@@ -106,6 +113,9 @@ class Relay:
|
||||
"""这条连接在看这次运行（WebSocket 接上、回放完之后）：在本进程里没在执行时从库里追。"""
|
||||
self._follow[run_id] = self._follow.get(run_id, 0) + 1
|
||||
self._cursor[run_id] = max(self._cursor.get(run_id, 0), int(after or 0))
|
||||
# 回放已经发到 after 为止（订阅位那边按 after 去重）：往回重读不越过它
|
||||
self._floor[run_id] = max(self._floor.get(run_id, 0), int(after or 0)) if run_id not in self._floor else \
|
||||
min(self._floor[run_id], int(after or 0))
|
||||
if signals.active:
|
||||
signals.publish("w", f"r:{run_id}")
|
||||
self._kick(run_id)
|
||||
@@ -116,6 +126,8 @@ class Relay:
|
||||
if left <= 0:
|
||||
self._follow.pop(run_id, None)
|
||||
self._cursor.pop(run_id, None)
|
||||
self._seen.pop(run_id, None)
|
||||
self._floor.pop(run_id, None)
|
||||
self._pumping.pop(run_id, None)
|
||||
else:
|
||||
self._follow[run_id] = left
|
||||
@@ -151,15 +163,19 @@ class Relay:
|
||||
lock = self._pumping.setdefault(run_id, asyncio.Lock())
|
||||
async with lock:
|
||||
after = self._cursor.get(run_id, 0)
|
||||
floor = self._floor.get(run_id, 0)
|
||||
seen = self._seen.setdefault(run_id, set())
|
||||
async with SessionLocal() as session:
|
||||
# 先读状态、再读事件：终态和终态事件是同一个事务提交的，读到终态时终态事件一定也读得到
|
||||
status = await session.scalar(select(Run.status).where(Run.id == run_id))
|
||||
rows = list((await session.execute(select(RunEvent).where(RunEvent.run_id == run_id,
|
||||
RunEvent.seq > after)
|
||||
.order_by(RunEvent.seq))).scalars())
|
||||
fetched = list((await session.execute(select(RunEvent).where(
|
||||
RunEvent.run_id == run_id, RunEvent.seq > max(floor, after - LOOKBACK))
|
||||
.order_by(RunEvent.seq))).scalars())
|
||||
rows = [r for r in fetched if int(r.seq or 0) not in seen and int(r.seq or 0) > floor]
|
||||
last_log = None
|
||||
for row in rows:
|
||||
self._cursor[run_id] = max(self._cursor.get(run_id, 0), int(row.seq or 0))
|
||||
seen.add(int(row.seq or 0))
|
||||
last_log = row
|
||||
try:
|
||||
event = RunEventModel(id=row.id, run_id=run_id, seq=row.seq, type=EventType(row.type),
|
||||
@@ -167,6 +183,9 @@ class Relay:
|
||||
except ValueError:
|
||||
continue # 不认识的事件种类（更新的版本写的）：回放时照样读得到，这里不转
|
||||
await bus.publish(event, relayed=True)
|
||||
low = self._cursor.get(run_id, 0) - 2 * LOOKBACK
|
||||
if len(seen) > 4 * LOOKBACK:
|
||||
self._seen[run_id] = {q for q in seen if q > low}
|
||||
ended = status is None or status in _TERMINAL or (
|
||||
status == "interrupted" and last_log is not None and last_log.type == EventType.LOG
|
||||
and (last_log.data or {}).get("code") == "server_shutdown")
|
||||
@@ -275,6 +294,7 @@ class Relay:
|
||||
return
|
||||
if event.run_id in self._follow:
|
||||
self._cursor[event.run_id] = max(self._cursor.get(event.run_id, 0), int(event.seq or 0))
|
||||
self._seen.setdefault(event.run_id, set()).add(int(event.seq or 0))
|
||||
signals.publish("r", event.run_id, event.seq)
|
||||
|
||||
async def _flush(self) -> None:
|
||||
@@ -292,8 +312,33 @@ class Relay:
|
||||
signals.publish("b", turn_id, text)
|
||||
|
||||
|
||||
def _size(item: dict[str, Any]) -> int:
|
||||
return len(json.dumps(item, ensure_ascii=False, default=str).encode())
|
||||
|
||||
|
||||
def _split(item: dict[str, Any], limit: int) -> list[dict[str, Any]]:
|
||||
"""太长的一条正文增量按字切成几条（复核 R-b6：以前截掉一截）：按字切，不会切坏 UTF-8；每条都在 limit 字节以内。切不开的
|
||||
（不是正文增量）交回空。"""
|
||||
data = dict(item.get("data") or {})
|
||||
text = data.get("delta")
|
||||
if not isinstance(text, str) or set(data) != {"delta"}:
|
||||
return []
|
||||
out: list[dict[str, Any]] = []
|
||||
while text:
|
||||
low, high = 1, len(text)
|
||||
while low < high: # 二分：最长的一段前缀，装得进 limit
|
||||
mid = (low + high + 1) // 2
|
||||
if _size({**item, "data": {"delta": text[:mid]}}) <= limit:
|
||||
low = mid
|
||||
else:
|
||||
high = mid - 1
|
||||
out.append({**item, "data": {"delta": text[:low]}})
|
||||
text = text[low:]
|
||||
return out
|
||||
|
||||
|
||||
def _chunks(items: list[dict[str, Any]]) -> Iterator[list[dict[str, Any]]]:
|
||||
"""按载荷上限切开（同一个节点接着的正文增量先合成一条）。太长的单条截断增量。"""
|
||||
"""按载荷上限切开（同一个节点接着的正文增量先合成一条；合成之后太长的再按字切成几条，不截）。"""
|
||||
merged: list[dict[str, Any]] = []
|
||||
for item in items:
|
||||
prev = merged[-1] if merged else None
|
||||
@@ -303,19 +348,17 @@ def _chunks(items: list[dict[str, Any]]) -> Iterator[list[dict[str, Any]]]:
|
||||
prev["seq"] = item.get("seq", prev.get("seq"))
|
||||
continue
|
||||
merged.append(dict(item))
|
||||
limit = PAYLOAD_MAX - 2
|
||||
pieces: list[dict[str, Any]] = []
|
||||
for item in merged:
|
||||
if _size(item) <= limit:
|
||||
pieces.append(item)
|
||||
else:
|
||||
pieces.extend(_split(item, limit)) # 不是正文增量又太长的（不该有）丢掉，丢了无妨：不落库的本来就不进回放
|
||||
chunk: list[dict[str, Any]] = []
|
||||
size = 2
|
||||
for item in merged:
|
||||
text = json.dumps(item, ensure_ascii=False, default=str)
|
||||
n = len(text.encode())
|
||||
if n > PAYLOAD_MAX - 2:
|
||||
data = dict(item.get("data") or {})
|
||||
if "delta" in data:
|
||||
data["delta"] = str(data["delta"])[: PAYLOAD_MAX // 4]
|
||||
item = {**item, "data": data}
|
||||
n = len(json.dumps(item, ensure_ascii=False, default=str).encode())
|
||||
if n > PAYLOAD_MAX - 2:
|
||||
continue
|
||||
for item in pieces:
|
||||
n = _size(item)
|
||||
if chunk and size + n + 1 > PAYLOAD_MAX:
|
||||
yield chunk
|
||||
chunk, size = [], 2
|
||||
|
||||
@@ -500,3 +500,66 @@ def test_the_worker_count_comes_from_the_backend_settings():
|
||||
assert ask(AGENTLAB_DATABASE_URL="", AGENTLAB_WORKERS="1").stdout.strip() == "1"
|
||||
pg = ask(AGENTLAB_DATABASE_URL="postgresql://u:p@db.internal:5432/agentlab", AGENTLAB_WORKERS="3")
|
||||
assert pg.returncode == 0 and pg.stdout.strip() == "3", pg.stderr[-1000:]
|
||||
|
||||
|
||||
async def test_the_relay_rereads_a_little_behind_so_a_late_commit_is_not_skipped():
|
||||
"""复核 R-b5：号先取后提交，号大的先提交时游标会越过号小的那一条——转发器往回留一小段重读、按发布过的 seq 去重。"""
|
||||
from app.cluster.relay import relay
|
||||
from app.core.bus import bus
|
||||
|
||||
async with SessionLocal() as session:
|
||||
run = Run(status="running", owner=PEER, graph={}, input={})
|
||||
session.add(run)
|
||||
await session.flush()
|
||||
session.add_all([RunEvent(id=new_id(), run_id=run.id, seq=s, type="log", ts=0, data={"n": s}) for s in (1, 3)])
|
||||
await session.commit()
|
||||
got: list[int] = []
|
||||
with bus.attach(run.id) as sub, relay.follow_run(run.id, 0):
|
||||
async def collect() -> None:
|
||||
async for event in sub.stream():
|
||||
got.append(event.seq)
|
||||
|
||||
reader = asyncio.create_task(collect())
|
||||
await relay.pump(run.id)
|
||||
async with SessionLocal() as session: # 号小的那一条后提交
|
||||
session.add(RunEvent(id=new_id(), run_id=run.id, seq=2, type="log", ts=0, data={"n": 2}))
|
||||
await session.commit()
|
||||
await relay.pump(run.id)
|
||||
await relay.pump(run.id)
|
||||
await asyncio.sleep(0.1)
|
||||
reader.cancel()
|
||||
assert sorted(got) == [1, 2, 3] and len(got) == 3, got
|
||||
|
||||
|
||||
async def test_cache_invalidations_committed_out_of_order_are_all_applied(monkeypatch):
|
||||
"""复核 R-b5：缓存失效的通知 id 大的先提交，id 小的后到也照样处理（往回留一小段重读、按 id 去重）。"""
|
||||
from app.cluster import caches
|
||||
from app.db.models_cluster import CacheEvent
|
||||
|
||||
seen: list[str | None] = []
|
||||
caches.on("test-cache", lambda key: seen.append(key))
|
||||
monkeypatch.setattr(caches, "_enabled", lambda: True)
|
||||
await caches.start()
|
||||
base = caches._last_seen or 0
|
||||
async with SessionLocal() as session:
|
||||
session.add(CacheEvent(id=base + 5, name="test-cache", key="late-big", owner=PEER))
|
||||
await session.commit()
|
||||
assert await caches.poll() == 1
|
||||
async with SessionLocal() as session:
|
||||
session.add(CacheEvent(id=base + 3, name="test-cache", key="early-small", owner=PEER))
|
||||
await session.commit()
|
||||
assert await caches.poll() == 1 and await caches.poll() == 0
|
||||
assert seen == ["late-big", "early-small"]
|
||||
|
||||
|
||||
def test_a_long_merged_delta_is_split_not_cut():
|
||||
"""复核 R-b6：几条正文增量合成一条之后超过载荷上限，按字切成几条（不切坏 UTF-8），拼回来一字不少。"""
|
||||
import json
|
||||
|
||||
from app.cluster.relay import PAYLOAD_MAX, _chunks
|
||||
|
||||
items = [{"type": "llm.token", "node_id": "w", "seq": i, "data": {"delta": "数据" * 500 + "😀"}} for i in range(1, 9)]
|
||||
chunks = list(_chunks(items))
|
||||
assert all(len(json.dumps(c, ensure_ascii=False).encode()) <= PAYLOAD_MAX for c in chunks)
|
||||
text = "".join(e["data"]["delta"] for c in chunks for e in c)
|
||||
assert text == ("数据" * 500 + "😀") * 8
|
||||
|
||||
Reference in New Issue
Block a user