diff --git a/Dockerfile b/Dockerfile index 9137ed6c..c3bd9ba7 100644 --- a/Dockerfile +++ b/Dockerfile @@ -93,6 +93,8 @@ HEALTHCHECK --interval=30s --timeout=5s --start-period=60s --start-interval=2s - # 版本号放在最后：它每次打包都不同，放前面会让装依赖那几层的缓存全部失效 ARG VERSION=dev +# 进程登记时带上构建号（app/cluster/version.py：库里还有活着的、版本不同的进程时新进程不许加入） +ENV AGENTLAB_BUILD=${VERSION} LABEL org.opencontainers.image.title="AgentLab" \ org.opencontainers.image.description="可视化 Agent 编排实验台" \ org.opencontainers.image.version="${VERSION}" \ @@ -100,4 +102,5 @@ LABEL org.opencontainers.image.title="AgentLab" \ # 工作进程数照 AGENTLAB_WORKERS（缺省 1）。大于 1 要应用库是 PostgreSQL（AGENTLAB_DATABASE_URL）：SQLite 只支持单进程， # 服务启动时拒绝第二个进程（app/cluster/guard.py）。exec：uvicorn 直接收 SIGTERM，走完整的优雅退出 -CMD ["sh", "-c", "exec uvicorn app.main:app --host 0.0.0.0 --port 8000 --workers \"${AGENTLAB_WORKERS:-1}\""] +# 进程数问后端（python -m app.cluster workers，环境变量和 .env 一起读；SQLite 而大于 1 时照实报错、不起） +CMD ["sh", "-c", "n=\"$(python -m app.cluster workers)\" || exit 2; exec uvicorn app.main:app --host 0.0.0.0 --port 8000 --workers \"$n\""] diff --git a/README.md b/README.md index 5ac9d204..a9062297 100644 --- a/README.md +++ b/README.md @@ -1054,9 +1054,8 @@ AGENTLAB_DATABASE_URL=postgresql+asyncpg://agentlab:<密码>@127.0.0.1:5432/agen - 断点（LangGraph 的 checkpoints、checkpoint_blobs、checkpoint_writes、checkpoint_migrations 四张表，不和应用的表重名） 缺省放在同一个库；要放别处设 `AGENTLAB_CHECKPOINT_URL=postgresql://…`（这一项是 libpq 的写法，`sslrootcert` 等可以直接写）。 - 连接池：每个进程常驻 `AGENTLAB_DB_POOL_SIZE`（缺省 10）条、高峰再多开 `AGENTLAB_DB_MAX_OVERFLOW`（缺省 20）条，取连接前探活， - 30 分钟换一次；断点库另有一个小池子。前面有 PgBouncer（事务模式）时：`AGENTLAB_DB_POOL_SIZE` 设成 0（不用连接池）、 - 地址后加 `?prepared_statement_cache_size=0`；断点库用预备语句，要 PgBouncer 1.21 以上开 `max_prepared_statements`， - 或者用 `AGENTLAB_CHECKPOINT_URL` 让它直连 PostgreSQL。 + 30 分钟换一次；断点库另有一个小池子。前面有 PgBouncer 时照下面「直连与 PgBouncer」的三种部署写（事务模式要开 + `AGENTLAB_DB_PGBOUNCER=1` 并另给直连地址）。 - 连接数的账：每个进程最多 `AGENTLAB_DB_POOL_SIZE + AGENTLAB_DB_MAX_OVERFLOW`（缺省 10 + 20 = 30）条应用库连接，断点库的池子 最多 `min(AGENTLAB_MAX_CONCURRENT_RUNS, 20)` 条（至少 4，缺省 20），合起来一个进程最多 50 条；导入命令另占几条。PostgreSQL 缺省 `max_connections=100`，还要留几条给 `psql`、备份和超级用户（`superuser_reserved_connections` 缺省 3）：一个进程够用， @@ -1141,8 +1140,24 @@ AGENTLAB_WORKERS=4 … | `AGENTLAB_SCHEDULER_TICK_S` | 15 | 调度循环一拍的间隔 | 连接数照上一节「连接数的账」按进程数乘，另外每个进程多一条专门收发进程间通知（LISTEN / NOTIFY）的连接，不占连接池；跨进程的锁 -（咨询锁）持锁期间占连接池里的一条。前面放 PgBouncer（事务模式）时，通知那一条要直连 PostgreSQL（事务模式下 LISTEN 收不到通知）， -会话级的咨询锁也要求直连或会话模式——多进程部署建议应用库直连。 +（会话级咨询锁）持锁期间占一条。 + +**直连与 PgBouncer。** 有几样连接依赖会话状态：收发进程间通知的那一条（LISTEN）、跨进程的锁（会话级咨询锁，取锁和放锁要在同一条 +会话上）、断点库（预备语句，建表时取会话锁）。普通查询走 `AGENTLAB_DATABASE_URL`；设了 `AGENTLAB_DATABASE_DIRECT_URL`（写法同它） +时，上面这几样走直连地址（断点库缺省也走它，`AGENTLAB_CHECKPOINT_URL` 另给的除外）；不设就全用同一个地址。三种部署： + +| 部署 | 怎么设 | +| --- | --- | +| 直连 PostgreSQL | 只设 `AGENTLAB_DATABASE_URL`。最简单，多进程建议这样 | +| PgBouncer 会话模式（`pool_mode=session`） | `AGENTLAB_DATABASE_URL` 指 PgBouncer 就行：会话模式下一个客户端连接独占一条服务端连接，会话状态都在 | +| PgBouncer 事务模式（`pool_mode=transaction`） | `AGENTLAB_DATABASE_URL` 指 PgBouncer，另设 `AGENTLAB_DB_PGBOUNCER=1` 和 `AGENTLAB_DATABASE_DIRECT_URL`（直连 PostgreSQL，或指 PgBouncer 的会话模式那个库）。没给直连地址时拒绝启动 | + +`AGENTLAB_DB_PGBOUNCER=1` 做的事（只对普通地址）：不用客户端的连接池（交给 PgBouncer）；关掉 SQLAlchemy 和 asyncpg 的预备语句缓存 +（事务模式下下一条语句可能落在另一条服务端连接上，缓存着的预备语句在那边不存在）；asyncpg 每条语句仍先预备再执行（SQLAlchemy 的 +asyncpg 方言一律这样做，同一个事务里、在同一条服务端连接上），名字用随机的，不和别的客户端留下的撞名——PgBouncer 那边要配 +`server_reset_query = DISCARD ALL` 加 `server_reset_query_always = 1` 收掉留下的预备语句，或用 1.21 以上开 `max_prepared_statements`。 +事务级的咨询锁（发名额、同源队、建表时用）在事务模式下照常成立：取锁、放锁在同一个事务里。这几条照 SQLAlchemy、asyncpg、psycopg +的文档和代码核过，没有在真的 PgBouncer 上跑过（本机没有它的镜像）。 **认领和租约。** 每个进程在库里登记一行（主机、pid、启动时的随机串），隔一拍心跳续租约。运行、问数据的一轮、后台任务（探查、起草、 剖析、评测、影子起草）、处理中的文档、执行中的发布单都记着认领它的进程；租约按进程算，进程没了它认领的行一起到期。平稳关停的进程 @@ -1161,11 +1176,13 @@ AGENTLAB_WORKERS=4 … **停止。** 运行、问数据的一轮、后台任务在别的进程上跑时，「停止」写一条给属主的请求（PostgreSQL 上 NOTIFY 叫醒它，最迟一拍心跳）， 属主照本进程的规矩停、收尾（`run.cancelled` 带上是谁停的），界面上和单进程时一样；属主不在了的按上表接手。 -**运行并发。** 每段执行（发起、审批之后接着跑、继续运行）开始前要一个名额，按来源分四类，各有全站上限和优先级，几个进程合起来算： -交互（问数据、画布试运行、手动发起的正式运行、蓝图试运行、个人应用刷新）、应用（应用入口、按需视图、钩子、接口）、定时（定时任务、 -事件触发）、系统（评测、口径召回、对账、补算、按口径正式出具）。名额满了时先给优先级高的一类，已经在跑的不打断。缺省：全站上限取 -`AGENTLAB_MAX_CONCURRENT_RUNS`（20），交互 = 全站上限、应用 2/5、定时 1/5、系统 1/5，优先级 40 / 30 / 20 / 10——评测再多也只占得到 -系统那几个名额，不挤占问数据。在「设置 → 偏好设置 → 运行并发」里改（管理员），改了最迟几秒各进程按新的发。单进程时也照这套发名额。 +**运行并发。** 每段执行（发起、审批之后接着跑、继续运行）开始前要一个名额，按来源分四类，几个进程合起来算：交互（问数据、画布 +试运行、手动发起的正式运行、蓝图试运行、个人应用刷新）、应用（应用入口、按需视图、钩子、接口）、定时（定时任务、事件触发）、系统 +（评测、口径召回、对账、补算、按口径正式出具）。全站上限取 `AGENTLAB_MAX_CONCURRENT_RUNS`（缺省 20），其中一部分是**给交互的留底** +（缺省全站上限的 1/4、至少 1、至多全站上限减 1）：应用、定时、系统三类合起来永远用不到留底，有人问数据、试运行时总有名额；留底以外的 +它们可以**借用**——三类各自的上限缺省就是「全站上限 − 留底」，没人问数据时评测也能开满，设置里可以调小。交互自己的上限缺省等于全站 +上限。等名额时按优先级（缺省交互 40 > 应用 30 > 定时 20 > 系统 10，同级先来先得），已经在执行的不打断。在「设置 → 偏好设置 → 运行 +并发」里改（管理员），改了最迟几秒各进程按新的发。单进程时也照这套发名额。 **调度当班。** 只有一个进程当班（库里的一份租约）：定时认领、事件派发、Webhook 投递、流程的时限和核实、每天的汇总、各类周期清理、 口径比对的补比和自动召回的续发、接手不在了的进程留下的行，都只在当班的进程里做；其余进程待命，当班的没了待命的接班（平稳关停的马上 @@ -1191,7 +1208,15 @@ SQL 形状。各进程各自的并发闸门（口径比对、自动对比、起 另一台读，可能读到不一致的页。所以多台机器时，记录表写得多的部署建议用一台机器多进程，或者把数据目录放在支持可靠文件锁的存储上。 沙箱会话（microVM）在各进程自己手上，不跨进程共享。 -**升级。** 升级要先停掉全部进程再起新版本：不支持新旧版本的进程同时连着一个库。 +**升级。** 不支持新旧版本的进程同时连着一个库，代码里拦着：每个进程登记时带上代码版本（包版本、构建号或 git 提交、源码指纹）和 +库结构版本，库里还有活着的、版本不同的进程时，新版本的进程在补列之前就拒绝启动，报错里写明是哪几个进程、什么版本；同版本的照常 +加入。升级步骤： + +1. 停掉全部旧版本的进程（平稳关停的马上删掉自己的登记；强杀的等它的租约过期，新进程启动时会等）； +2. 备份（`pg_dump` 加数据目录，见上一节）； +3. 起新版本：第一个进程补列补索引，之后的照常加入。 + +只停了一部分就起新版本，新版本起不来，旧版本照常跑，不会两边各按各的写同一批行。 **测试。** `scripts/test-multiprocess.sh` 用本机 docker 的 `postgres:16` 起临时库，在上面起两个工作进程，用演示模型跑端到端 （`backend/tests/test_multiprocess_e2e.py`，缺省跳过，`AGENTLAB_MULTIPROCESS_E2E=1` 时跑）：运行由任一进程认领、另一个进程上的 diff --git a/backend/app/apps/resolve.py b/backend/app/apps/resolve.py index d20bb73a..95bb9d58 100644 --- a/backend/app/apps/resolve.py +++ b/backend/app/apps/resolve.py @@ -388,7 +388,7 @@ _LOCKS: dict[str, asyncio.Lock] = {} def _lock(key: str) -> Any: - """按需视图的单飞锁（app/cluster/mutex，多进程之间也是一把）：按应用取，同一个键正在运行的那次接上、不另发，每日上限也数得准。""" + """按需视图的单飞锁（app/cluster/mutex，多进程之间也是一把）：按缓存键（视图 + 参数）取，同一个键正在运行的那次接上、不另发。""" from app.cluster import mutex return mutex.lock(f"app-view:{key}") @@ -471,25 +471,35 @@ async def run_view(session: AsyncSession, *, app: Any, version: AppVersion, view raise FormRejected(errors) inputs = normalize_params(form, params) key = await cache_key(session, state, inputs) - # 按应用一把锁（不按缓存键）：同一个键正在运行的那次接上不另发，每日上限的「数一数再发」几个进程、几个键之间也不越过 - async with _lock(app.id): + # 锁按缓存键（这个视图、这组参数）：管的是去重——同一个键正在运行的那次接上，不另算一遍。每日上限不靠它：发起之前在库里先占 + # 一份这个应用今天的额度（app/cluster/quota，带条件的加一），几个键、几个进程同时刷新也正好到上限为止，发起不成退回 + from app.cluster import quota + + async with _lock(key): run = await _cached(session, key) if run is not None: return run, run.status == "succeeded" - # 「每天」按服务端所在时区的自然日算 - start = datetime.now().astimezone().replace(hour=0, minute=0, second=0, microsecond=0).astimezone(timezone.utc) - used = await session.scalar(select(func.count(Run.id)).where( - Run.app_id == app.id, Run.origin_kind == "app_view", Run.created_at >= start)) limit = await daily_limit(session) - if int(used or 0) >= limit: + + async def used_today(s: AsyncSession) -> int: + # 「每天」按服务端所在时区的自然日算 + start = datetime.now().astimezone().replace(hour=0, minute=0, second=0, microsecond=0).astimezone(timezone.utc) + return int(await s.scalar(select(func.count(Run.id)).where( + Run.app_id == app.id, Run.origin_kind == "app_view", Run.created_at >= start)) or 0) + + if not await quota.take(f"app_view:{app.id}", limit, seed=used_today): raise ViewRefused(429, f"这个应用今天的按需运行已达到上限（{limit} 次），请明天再试，或请管理员调整上限", DAILY_LIMIT) - run = await launch_run( - session, workflow_id=state.pin["workflow_id"], run_class="formal", version=state.pin["version"], - expect_hash=state.pin.get("graph_hash"), input_payload=inputs, actor=actor, origin_kind="app_view", - app_id=app.id, app_version=version.version, - origin={"view": view_id, "spec_hash": version.spec_hash, "actor_verified": principal.verified, - "auth_mode": principal.mode, **_key_trace(principal)}) + try: + run = await launch_run( + session, workflow_id=state.pin["workflow_id"], run_class="formal", version=state.pin["version"], + expect_hash=state.pin.get("graph_hash"), input_payload=inputs, actor=actor, origin_kind="app_view", + app_id=app.id, app_version=version.version, + origin={"view": view_id, "spec_hash": version.spec_hash, "actor_verified": principal.verified, + "auth_mode": principal.mode, **_key_trace(principal)}) + except BaseException: + await quota.give_back(f"app_view:{app.id}") + raise ttl = source.get("ttl_s") if ttl is None and await reads_live(session, state.graph): ttl = TTL_DEFAULT diff --git a/backend/app/chat/job.py b/backend/app/chat/job.py index b29b3a8c..438f6fba 100644 --- a/backend/app/chat/job.py +++ b/backend/app/chat/job.py @@ -1234,7 +1234,9 @@ async def startup(*, resume: bool | None = None, live: Any = None) -> dict[str, await session.commit() for turn_id in gone: turn_log.notify(turn_id) - for row in legacy: + # 老轮次（没有任务行）只在启动时整理：当班进程隔一会儿来的那一次（live.trust_mine）不碰——老页面开的轮次（不占锁的写入） + # 正在建图时也是这个样子 + for row in [] if live.trust_mine else legacy: if turn_run_id(row): continue # 有运行的老轮次：页面打开时请服务端接着收尾 async with SessionLocal() as session: diff --git a/backend/app/chat/log.py b/backend/app/chat/log.py index c0e8a37c..1f151097 100644 --- a/backend/app/chat/log.py +++ b/backend/app/chat/log.py @@ -221,12 +221,13 @@ async def _stream(turn_id: str, after: int, *, active: Callable[[], Awaitable[st yield PING -async def purge(now: datetime, *, retain_days: float = RETAIN_DAYS, limit: int = PURGE_BATCH) -> list[str]: +async def purge(now: datetime | None = None, *, retain_days: float = RETAIN_DAYS, limit: int = PURGE_BATCH) -> list[str]: """结束超过 retain_days 天的轮次整轮删事件日志（见模块说明），交回清了哪几轮。任务行、conversation_turns 那一行不动。 结束的时刻按任务的 finished_at，没记的按 updated_at。""" + from app.db.base import utcnow from app.db.models_chat import ChatTurnEvent, ChatTurnJob - cutoff = now - timedelta(days=retain_days) + cutoff = (now or utcnow()) - timedelta(days=retain_days) async with SessionLocal() as session: ended = select(ChatTurnJob.id).where(ChatTurnJob.status.in_(FINISHED), func.coalesce(ChatTurnJob.finished_at, ChatTurnJob.updated_at) < cutoff) @@ -242,16 +243,13 @@ async def purge(now: datetime, *, retain_days: float = RETAIN_DAYS, limit: int = _last_purge = 0.0 -async def purge_tick(now: datetime | None) -> list[str]: - """调度循环每一拍调（scheduler/loop.process_due）：一小时最多清一次。""" +async def purge_tick(now: datetime | None = None) -> list[str]: + """调度循环每一拍调（scheduler/loop.process_due，那一拍不一定给 now）：一小时最多清一次，不给 now 按当前时间算。""" global _last_purge if time.monotonic() - _last_purge < PURGE_EVERY_S and _last_purge: return [] _last_purge = time.monotonic() - # 调度循环那一拍不带时刻（process_due() 的 now 缺省为 None）：按现在算 - from app.db.base import utcnow - - return await purge(now or utcnow()) + return await purge(now) __all__ = ["FINISHED", "HEARTBEAT_S", "PING", "READ_BATCH", "RETAIN_DAYS", "STREAM_END", "add", "append", "beat", "forget", diff --git a/backend/app/cluster/__main__.py b/backend/app/cluster/__main__.py new file mode 100644 index 00000000..04824c79 --- /dev/null +++ b/backend/app/cluster/__main__.py @@ -0,0 +1,24 @@ +"""python -m app.cluster workers：照后端自己的配置（环境变量和 .env 一起读）说起几个工作进程，scripts/prod.sh 和 Docker 的入口按它给 +uvicorn --workers（复核 R-b4：写在 .env 里的 AGENTLAB_DATABASE_URL、AGENTLAB_WORKERS，shell 脚本自己读不到）。应用库是 SQLite 而进程数 +大于 1 时照实报错、退出码 2（SQLite 只支持单进程）。""" +from __future__ import annotations + +import sys + + +def main(argv: list[str]) -> int: + if argv[:1] != ["workers"]: + print("用法：python -m app.cluster workers", file=sys.stderr) + return 64 + from app.cluster.guard import SINGLE_PROCESS + from app.core.config import settings + + if settings.db_dialect == "sqlite" and settings.workers > 1: + print(SINGLE_PROCESS.format(dir=settings.data_dir, hint=f"AGENTLAB_WORKERS={settings.workers}"), file=sys.stderr) + return 2 + print(settings.workers) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:])) diff --git a/backend/app/cluster/caches.py b/backend/app/cluster/caches.py index 5f04de72..d9232492 100644 --- a/backend/app/cluster/caches.py +++ b/backend/app/cluster/caches.py @@ -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 diff --git a/backend/app/cluster/fence.py b/backend/app/cluster/fence.py new file mode 100644 index 00000000..60da513b --- /dev/null +++ b/backend/app/cluster/fence.py @@ -0,0 +1,116 @@ +"""自我隔离（fencing）：本进程续不上租约的时候，当作自己认领的行可能已经被别的进程接手，停下手上的执行、不再写。 + +**为什么。** 进程卡住了但没死（SIGSTOP、虚拟机暂停、事件循环被长同步代码卡住、库一时连不上）：别的进程看它的租约过期了，把它在跑 +的运行标中断、交给别人继续运行。它一缓过来，手上的执行接着往下走：照样写事件（和接手的那一边撞号）、照样写终态（把接手的结果 +盖掉）。所以： + +- 续租约在单独的任务、单独的连接上做（cluster/service.py），不被同一拍别的活拖住；最近一次续上的时刻记在这里（renewed）。 +- 本进程的执行在动手写之前问一句 check()：离上次续上已经超过一个租约时长（别的进程按两个租约时长才接手，中间留着一段），或者这段 + 执行是在上一次隔离之前开始的（世代不对），就抛 Fenced——写事件、写断点、写终态、写待办和动作、发外部写入之前都问。 +- 一旦发现续晚了（trip）：世代加一，手上的执行全部停掉（停的时候不写终态，那些行交给接手的进程按「属主不在了」处理），换一个新的 + 身份重新登记，之后发起的照常。 + +Fenced 是 BaseException：节点、工具里常见的 except Exception 拦不住它，一路退到执行的最外层。没起协调（测试、脚本）时什么都不查。 +""" +from __future__ import annotations + +import contextvars +import logging +import time +from collections.abc import Callable +from typing import Any + +logger = logging.getLogger(__name__) + + +class Fenced(BaseException): + """本进程续不上租约了（或者这段执行属于上一次隔离之前）：不许再写。""" + + +_generation = 0 +_renewed: float | None = None +_active = False +_handlers: list[Callable[[], Any]] = [] +_bound: contextvars.ContextVar[int | None] = contextvars.ContextVar("agentlab_fence_generation", default=None) + + +def activate(renewed_at: float | None = None) -> None: + """起协调、续上了第一次租约时（PostgreSQL；SQLite 只有一个进程，不用隔离）。""" + global _active, _renewed + _active = True + _renewed = time.monotonic() if renewed_at is None else renewed_at + + +def deactivate() -> None: + global _active, _renewed + _active, _renewed = False, None + + +def renewed(at: float | None = None) -> None: + """续上了一次租约。""" + global _renewed + _renewed = time.monotonic() if at is None else at + + +def generation() -> int: + return _generation + + +def bind() -> None: + """这段执行属于现在这一世代（运行的一段执行、后台任务、问数据的一轮开始时调）。""" + _bound.set(_generation) + + +def stale(now: float | None = None) -> bool: + """离上次续上租约是不是已经超过一个租约时长。""" + if not _active or _renewed is None: + return False + from app.cluster.workers import lease_s + + return (time.monotonic() if now is None else now) - _renewed > lease_s() + + +def fenced_here() -> bool: + """当前这段执行还能不能写：世代不对（隔离之前开始的）或者租约续晚了。""" + if not _active: + return False + bound = _bound.get() + return (bound is not None and bound != _generation) or stale() + + +def check() -> None: + """写之前问一句：不许写就抛 Fenced（续晚了的话顺手隔离）。""" + if not _active: + return + if stale(): + trip("执行要写的时候发现租约已经续晚了") + raise Fenced() + bound = _bound.get() + if bound is not None and bound != _generation: + raise Fenced() + + +def on_trip(handler: Callable[[], Any]) -> None: + if handler not in _handlers: + _handlers.append(handler) + + +def trip(why: str) -> bool: + """隔离：世代加一，叫各处停掉手上的执行（cluster/service 再换身份重新登记）。已经在这一世代隔离过的不重复。交回这一次有没有动手。""" + global _generation, _renewed + if not _active or _renewed is None: + return False + _generation += 1 + _renewed = None + logger.error("本进程的租约续晚了（%s）：当作认领的运行、问数据、后台任务可能已被别的进程接手，停掉手上的执行、不写终态，" + "换一个身份重新登记", why) + for handler in list(_handlers): + try: + handler() + except Exception: # noqa: BLE001 + logger.exception("隔离时停执行出错") + return True + + +__all__ = ["Fenced", "activate", "bind", "check", "deactivate", "fenced_here", "generation", "on_trip", "renewed", + "stale", "trip"] diff --git a/backend/app/cluster/identity.py b/backend/app/cluster/identity.py index 1a0d6695..7a554012 100644 --- a/backend/app/cluster/identity.py +++ b/backend/app/cluster/identity.py @@ -33,6 +33,14 @@ def owner() -> str: return OWNER +def regenerate() -> str: + """换一个身份（隔离之后重新登记用，cluster/fence.py）：同一个 pid、新的随机串。之前认领的行算「属主不在了」。""" + global OWNER, _PID + _PID = os.getpid() + OWNER = _make(_PID) + return OWNER + + def host() -> str: return _HOST @@ -45,4 +53,4 @@ def describe(value: str | None) -> str: return f"{parts[0]} 上的进程 {parts[1]}" if len(parts) >= 2 else value -__all__ = ["OWNER", "describe", "host", "owner"] +__all__ = ["OWNER", "describe", "host", "owner", "regenerate"] diff --git a/backend/app/cluster/leader.py b/backend/app/cluster/leader.py index ab2d800b..fe169232 100644 --- a/backend/app/cluster/leader.py +++ b/backend/app/cluster/leader.py @@ -32,10 +32,16 @@ class Leadership: self.until: datetime | None = None self.epoch = 0 + def single(self) -> None: + """只有本进程（SQLite）：直接算当班，不写库。""" + self.until = datetime.max.replace(tzinfo=None) + def held(self, now: datetime | None = None) -> bool: """本地看：还在当班（租约离到期还有余量）。""" if self.until is None: return False + if self.until.tzinfo is None: + return True # single()：只有本进程 # 余量不超过租约的四分之一：租约调得很短（验收脚本）时也留得出当班的时间 margin = min(MARGIN_S, lease_s() / 4) return (now or utcnow()) + timedelta(seconds=margin) < self.until @@ -85,7 +91,10 @@ class Leadership: if self.until is None: return + single = self.until.tzinfo is None self.until = None + if single: + return async with SessionLocal() as session: await session.execute(update(ClusterLease).where( ClusterLease.name == self.name, ClusterLease.owner == identity.owner()) diff --git a/backend/app/cluster/mutex.py b/backend/app/cluster/mutex.py index 2de45376..96214ba3 100644 --- a/backend/app/cluster/mutex.py +++ b/backend/app/cluster/mutex.py @@ -55,9 +55,9 @@ class ClusterLock: try: from sqlalchemy import text - from app.db.base import engine + from app.db.base import session_engine - conn = await engine.connect() + conn = await session_engine().connect() try: conn = await conn.execution_options(isolation_level="AUTOCOMMIT") got = bool(await conn.scalar(text("SELECT pg_try_advisory_lock(:k)"), {"k": self.key})) @@ -86,9 +86,11 @@ class ClusterLock: try: from sqlalchemy import text - from app.db.base import engine + from app.db.base import session_engine - conn = await engine.connect() + # 会话级的锁：走直连（AGENTLAB_DATABASE_DIRECT_URL，没设就是应用库）——PgBouncer 事务模式下取锁、放锁可能落在 + # 两条不同的服务端连接上 + conn = await session_engine().connect() try: conn = await conn.execution_options(isolation_level="AUTOCOMMIT") await conn.execute(text("SELECT pg_advisory_lock(:k)"), {"k": self.key}) diff --git a/backend/app/cluster/quota.py b/backend/app/cluster/quota.py new file mode 100644 index 00000000..9a726966 --- /dev/null +++ b/backend/app/cluster/quota.py @@ -0,0 +1,70 @@ +"""每日额度：先占一份额度再做事，做不成退回（ledger/issue 的出具上限、apps/resolve 的按需视图上限）。 + +**为什么不靠锁。** 以前是「数一数今天做了几次，没到上限就发起」，数和发起之间靠一把锁串行；多进程时那把锁要么只管本进程（几个 +进程一起越过上限），要么全站一把（几份出具只能一个接一个发起）。现在额度记在库里（rate_counters 里按「名字 + 本地日期」一行）， +占额度是一句带条件的 UPDATE：还没到上限才加一，否则什么都不改——两种库上这一句都是原子的，几个进程、几个协程同时来也正好 +占到上限为止，发起本身可以同时做。发起失败（被拒、出错）把那一份退回。 + +这一天第一次占的时候按库里已有的记录起算（seed：比如今天已经出具过几次），之后只看这一行。进程在占了额度、发起之前没了的， +那一份不退：宁可这一天少出一次，不会多出。 +""" +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from datetime import datetime, timedelta +from typing import Any + +from sqlalchemy import select, update + +from app.db.base import SessionLocal, utcnow + + +def today() -> str: + """本地日期（每日上限按服务端所在时区的自然日算，同以前按本地零点数）。""" + return datetime.now().astimezone().date().isoformat() + + +def _key(name: str, day: str) -> str: + return f"quota:{name}:{day}" + + +async def take(name: str, limit: int, *, seed: Callable[[Any], Awaitable[int]] | None = None, + day: str | None = None) -> bool: + """占一份额度：这一天占过的还不到 limit 就加一、交回 True，到了交回 False（什么都不改）。seed(session)：这一天第一次占时 + 已经用掉了几份（按库里的记录数）。""" + from app.db.models_cluster import RateCounter as Row + from app.db.portable import insert_ignore + + key = _key(name, day or today()) + async with SessionLocal() as session: + if await session.scalar(select(Row.key).where(Row.key == key)) is None: + used = int(await seed(session)) if seed is not None else 0 + await session.execute(insert_ignore(session, Row.__table__).values( + key=key, count=used, value=0.0, at=0.0, expires_at=utcnow() + timedelta(days=3)) + .on_conflict_do_nothing(index_elements=["key"])) + done = await session.execute(update(Row).where(Row.key == key, Row.count < int(limit)) + .values(count=Row.count + 1).execution_options(synchronize_session=False)) + await session.commit() + return done.rowcount == 1 + + +async def give_back(name: str, *, day: str | None = None) -> None: + """退回一份（占了额度、却没发起成）。""" + from app.db.models_cluster import RateCounter as Row + + key = _key(name, day or today()) + async with SessionLocal() as session: + await session.execute(update(Row).where(Row.key == key, Row.count > 0).values(count=Row.count - 1) + .execution_options(synchronize_session=False)) + await session.commit() + + +async def used(name: str, *, day: str | None = None) -> int | None: + """这一天占了几份（没占过为 None）。""" + from app.db.models_cluster import RateCounter as Row + + async with SessionLocal() as session: + return await session.scalar(select(Row.count).where(Row.key == _key(name, day or today()))) + + +__all__ = ["give_back", "take", "today", "used"] diff --git a/backend/app/cluster/relay.py b/backend/app/cluster/relay.py index df0873f8..6d0ed68f 100644 --- a/backend/app/cluster/relay.py +++ b/backend/app/cluster/relay.py @@ -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 diff --git a/backend/app/cluster/requests.py b/backend/app/cluster/requests.py index 09332874..34fb7e8f 100644 --- a/backend/app/cluster/requests.py +++ b/backend/app/cluster/requests.py @@ -80,14 +80,32 @@ async def _drain() -> int: if fn is None: logger.warning("收到不认识的进程间请求 %s/%s，不办", row.kind, row.target) continue - try: - await fn(row.target_id, row.actor, dict(row.payload or {})) - done += 1 - except Exception: # noqa: BLE001 - logger.exception("办进程间请求 %s/%s %s 出错", row.kind, row.target, row.target_id) + # 各办各的、另起任务：停止要等执行收尾（最多十来秒），不能拖住心跳和后面几件 + task = asyncio.ensure_future(_handle(fn, row.kind, row.target, row.target_id, row.actor, dict(row.payload or {}))) + _running.add(task) + task.add_done_callback(_running.discard) + done += 1 return done +#: 正在办的请求（另起的任务要有人拿着引用） +_running: set[asyncio.Future[Any]] = set() + + +async def _handle(fn: Callable[..., Awaitable[Any]], kind: str, target: str, target_id: str, actor: str | None, + payload: dict[str, Any]) -> None: + try: + await fn(target_id, actor, payload) + except Exception: # noqa: BLE001 + logger.exception("办进程间请求 %s/%s %s 出错", kind, target, target_id) + + +async def settle(timeout: float = 5.0) -> None: + """等手上正在办的请求办完（测试、关停用）。""" + if _running: + await asyncio.wait(list(_running), timeout=timeout) + + async def discard_for(owners: set[str]) -> int: """收件人已经不在了的请求删掉（接手的进程做：那些行按属主不在了的规矩处理）。""" from app.db.models_cluster import ClusterRequest diff --git a/backend/app/cluster/service.py b/backend/app/cluster/service.py index 424996c8..ff9cf7b3 100644 --- a/backend/app/cluster/service.py +++ b/backend/app/cluster/service.py @@ -1,7 +1,11 @@ """本进程在协调里的那一份：启动时登记、起心跳，关停时交班、删登记。 -心跳每一拍（AGENTLAB_WORKER_HEARTBEAT_S，缺省 5 秒）： -1. 续自己的租约（登记不见了就重新登记）； +续租约在单独的任务、单独的连接上（不走连接池，不被同一拍别的活拖住）；续晚了（离上次续上超过一个租约时长）就隔离：停掉手上的 +执行、不写终态，换一个身份重新登记（cluster/fence.py）。SQLite 只有一个进程：登记只写一次、不续租、本进程直接当班、接手只在 +启动时做，心跳只做本进程自己攒着的活。 + +心跳每一拍（AGENTLAB_WORKER_HEARTBEAT_S，缺省 5 秒；PostgreSQL）： +1. （续租约另有任务）； 2. 刷新「别的进程谁还活着」； 3. 取走发给本进程的请求照办（停止、往运行里记事件，cluster/requests.py）； 4. 本进程自己的周期活（on_beat 登记的：访问密钥的请求记录写库……——这些是每个进程攒在自己内存里的，只能自己写）； @@ -52,36 +56,48 @@ def leading() -> bool: class _Service: def __init__(self) -> None: self.task: asyncio.Task[None] | None = None + self.lease_task: asyncio.Task[None] | None = None self.wake: asyncio.Event | None = None self.reap_now = False self.last_reap = 0.0 self.others: frozenset[str] = frozenset() self.started = False + #: SQLite：守卫已经确认只有本进程——登记只写一次、不续租、本进程直接当班、接手只在启动时做（不为协调多写库） + self.single = False + self._lease_conn: Any = None async def start(self) -> None: """登记、开始收发通知、起心跳。SQLite 时先取单进程的守卫（拿不到抛 guard.SingleProcessOnly）。""" if self.started: return guard.acquire() - exclusive = settings.db_dialect == "sqlite" - await workers.register(exclusive=exclusive) + self.single = settings.db_dialect == "sqlite" + await workers.register(exclusive=self.single) live = await workers.liveness() state.remember(live) self.others = live.others await signals.start() signals.on("q", self._on_request) signals.on("gone", self._on_gone) - if settings.db_dialect == "postgresql": + if not self.single: await signals.ready(timeout=5.0) self.wake = asyncio.Event() self.started = True - from app.cluster import caches + from app.cluster import caches, fence await caches.start() - await self._lead() - self.task = asyncio.create_task(self._loop(), name="agentlab-cluster-heartbeat") + if self.single: + # 只有本进程：当班不用争（不写库），心跳只做本进程自己攒着的活（访问密钥的请求记录） + leader.scheduler.single() + self.task = asyncio.create_task(self._light_loop(), name="agentlab-cluster-beats") + else: + await self._lead() + fence.activate() + fence.on_trip(self._on_fenced) + self.lease_task = asyncio.create_task(self._lease_loop(), name="agentlab-cluster-lease") + self.task = asyncio.create_task(self._loop(), name="agentlab-cluster-heartbeat") logger.info("工作进程 %s 已登记（%s，租约 %g 秒、心跳 %g 秒；现在另有 %d 个进程）", identity.describe(identity.owner()), - "PostgreSQL，可以多进程" if settings.db_dialect == "postgresql" else "SQLite，单进程", + "SQLite，单进程" if self.single else "PostgreSQL，可以多进程", workers.lease_s(), workers.heartbeat_s(), len(live.others)) def _on_request(self, parts: list[str]) -> None: @@ -93,6 +109,75 @@ class _Service: if self.wake is not None: self.wake.set() + # ---------------- 续租约：单独的任务、单独的连接 ---------------- + + async def _lease_conn_open(self) -> Any: + import asyncpg + + from app.cluster.signal import listen_dsn + from app.db.base import APPLICATION_NAME + + if self._lease_conn is None or self._lease_conn.is_closed(): + self._lease_conn = await asyncpg.connect(listen_dsn(), server_settings={"application_name": APPLICATION_NAME}, + timeout=max(1.0, workers.heartbeat_s())) + return self._lease_conn + + async def renew(self) -> bool: + """续一拍租约（在专用连接上，不走连接池：池子被占满时取连接要等，等得比租约还长）。续晚了（离上次续上超过一个租约时长） + 就不续这个身份了：隔离（cluster/fence.py），换身份重新登记。交回续没续上。""" + from datetime import timedelta + + from app.cluster import fence + from app.db.base import utcnow + + if fence.stale(): + fence.trip("续租约的时候发现已经续晚了") + return False + now = utcnow() + until = now + timedelta(seconds=workers.lease_s()) + try: + conn = await self._lease_conn_open() + done = await asyncio.wait_for(conn.execute( + "UPDATE cluster_workers SET heartbeat_at = $2, lease_until = $3 WHERE id = $1", + identity.owner(), now.replace(tzinfo=None), until.replace(tzinfo=None)), timeout=workers.heartbeat_s()) + except Exception: # noqa: BLE001 - 库一时连不上：这一拍没续上，下一拍再试（续晚了的话下一拍隔离） + logger.warning("续租约没续上，下一拍再试", exc_info=True) + with contextlib.suppress(Exception): + if self._lease_conn is not None: + await self._lease_conn.close() + self._lease_conn = None + return False + if done.endswith(" 0"): + logger.warning("本进程的登记不见了（可能被当成没了），重新登记") + await workers.register() + fence.renewed() + return True + + async def _lease_loop(self) -> None: + while True: + await self.renew() + await asyncio.sleep(workers.heartbeat_s()) + + def _on_fenced(self) -> None: + """隔离了：手上的执行由各处自己停（runner、job_runtime 登记的处理），这里换身份重新登记。""" + asyncio.ensure_future(self._reidentify()) + + async def _reidentify(self) -> None: + from app.cluster import fence + + old = identity.owner() + new = identity.regenerate() + leader.scheduler.until = None + try: + await workers.register() + fence.renewed() + logger.warning("隔离之后换了身份重新登记：%s → %s（之前认领的行由当班进程按「属主不在了」接手）", + identity.describe(old), identity.describe(new)) + except Exception: # noqa: BLE001 + logger.exception("隔离之后重新登记失败，下一拍再试") + + # ---------------- 当班、心跳的其余活 ---------------- + async def _lead(self) -> bool: was = leader.scheduler.held() try: @@ -113,10 +198,10 @@ class _Service: return now_leading async def beat(self) -> None: - """心跳的一拍（测试里直接调）。""" - if not await workers.renew(): - logger.warning("本进程的登记不见了（可能被当成没了），重新登记") - await workers.register() + """心跳的一拍（续租约另在 _lease_loop；测试里直接调）。SQLite 只做本进程自己攒着的活。""" + if self.single: + await self._own_work() + return live = await workers.liveness() state.remember(live) if self.others - live.others: @@ -132,14 +217,17 @@ class _Service: await caches.poll() except Exception: # noqa: BLE001 logger.exception("读别的进程的缓存变化时出错") + await self._own_work() + if await self._lead() and (self.reap_now or time.monotonic() - self.last_reap >= workers.lease_s()): + # 隔一会儿的这一次信得过自己：属主是本进程的行（刚建、还没交给后台任务的）不当成没人管 + await self.reap(live.trusting()) + + async def _own_work(self) -> None: for name, fn in list(_beats): try: await fn() except Exception: # noqa: BLE001 logger.exception("心跳里的「%s」出错，下一拍再试", name) - if await self._lead() and (self.reap_now or time.monotonic() - self.last_reap >= workers.lease_s()): - # 隔一会儿的这一次信得过自己：属主是本进程的行（刚建、还没交给后台任务的）不当成没人管 - await self.reap(live.trusting()) async def reap(self, live: workers.Liveness | None = None) -> dict[str, Any]: self.reap_now = False @@ -152,10 +240,11 @@ class _Service: with contextlib.suppress(Exception): await workers.prune() await requests.prune() - if settings.db_dialect == "postgresql": - from app.cluster import caches, counters + from app.cluster import caches, counters - await counters.prune() + # 到期的计数、额度（每日额度两种库都记在库里） + await counters.prune() + if settings.db_dialect == "postgresql": await caches.prune() return out @@ -172,20 +261,33 @@ class _Service: except Exception: # noqa: BLE001 - 一拍出错（库一时不通）不该让心跳停下 logger.exception("心跳这一拍出错，下一拍照常进行") + async def _light_loop(self) -> None: + while True: + await asyncio.sleep(workers.heartbeat_s()) + await self._own_work() + async def stop(self) -> None: """关停：停心跳、交班、删登记（别的进程马上接手本进程留下的行），停收发通知。""" if not self.started: return self.started = False - if self.task is not None: - self.task.cancel() - with contextlib.suppress(asyncio.CancelledError, Exception): - await self.task - self.task = None + from app.cluster import fence + + for task in (self.lease_task, self.task): + if task is not None: + task.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await task + self.task = self.lease_task = None + fence.deactivate() with contextlib.suppress(Exception): await leader.scheduler.release() with contextlib.suppress(Exception): await workers.unregister() + if self._lease_conn is not None: + with contextlib.suppress(Exception): + await self._lease_conn.close() + self._lease_conn = None signals.publish("gone", identity.owner()) await signals.stop() state.forget() diff --git a/backend/app/cluster/signal.py b/backend/app/cluster/signal.py index a2592d97..46785494 100644 --- a/backend/app/cluster/signal.py +++ b/backend/app/cluster/signal.py @@ -37,6 +37,14 @@ RETRY_S = 2.0 Handler = Callable[[list[str]], Any] +def listen_dsn() -> str: + """收发通知那条专用连接的地址：LISTEN 要一条一直开着的会话，设了直连地址（AGENTLAB_DATABASE_DIRECT_URL）就走它——普通地址 + 可能是 PgBouncer 的事务模式，那里收不到通知。""" + from app.db.urls import libpq_conninfo + + return libpq_conninfo(settings.session_db_url) + + def _tag() -> str: return identity.owner().rsplit(":", 1)[-1] @@ -126,9 +134,8 @@ class Signals: import asyncpg from app.db.base import APPLICATION_NAME - from app.db.urls import libpq_conninfo - dsn = libpq_conninfo(settings.database_url) + dsn = listen_dsn() first = True while not self._stopping: lost = asyncio.Event() diff --git a/backend/app/cluster/slots.py b/backend/app/cluster/slots.py index 06e2312f..d363dedb 100644 --- a/backend/app/cluster/slots.py +++ b/backend/app/cluster/slots.py @@ -10,12 +10,15 @@ 定时 schedule：定时任务（含事件触发） 系统 system：评测、口径召回、对账、补算、按口径正式出具——服务端替人跑的 -每类有全站的并发上限和优先级（设置里「运行并发」一组，管理员可调），另有全站总上限。名额按库里的计数发：等着的按 -优先级高的先、同一优先级先来先得，只要它那一类没到上限、全站没到总上限就发给它。所以评测再多也只占得到系统那一类的 -几个名额，问数据不用等它；全站满了的时候，空出来的名额先给交互。已经在跑的不抢。 +每类有全站的并发上限和优先级（设置里「运行并发」一组，管理员可调），另有全站总上限，和一个**交互留底**： -缺省（总上限取 AGENTLAB_MAX_CONCURRENT_RUNS）：交互 = 总上限、应用 2/5、定时 1/5、系统 1/5（至少各 1），优先级 40 / 30 / 20 / 10。 -应用、定时、系统三类加起来不到总上限：交互总留得出名额。单进程时总上限和以前的信号量一样大，问数据、试跑不再被评测挤占。 +- 留底：全站名额里永远留给交互类的那几个（缺省全站上限的 1/4、至少 1；至多全站上限减 1）。应用、定时、系统三类合起来 + 永远用不到这一部分——有人问数据、试跑时总有名额，不用等评测、召回跑完。 +- 借用：没人问数据时，别的类可以借满到「全站上限 − 留底」（三类各自的上限缺省就是这个数，设置里可以调小）。评测照旧跑得快。 +- 交互自己的上限缺省 = 全站上限：留底之外空着的名额它也能用。 + +等着的按优先级高的先、同一优先级先来先得（缺省交互 40 > 应用 30 > 定时 20 > 系统 10）；已经在跑的不抢。单进程时总上限和以前的 +信号量一样大。 **怎么发。** 要名额的一方先在本进程登记、再插一行「等着」，然后自己去发一次（grant：把所有等着的按规矩过一遍，谁够格就 改成「已发」）；没轮到就等叫醒（本进程放名额时叫醒，PostgreSQL 上别的进程放名额时 NOTIFY slot 叫醒），另隔一会儿自己再 @@ -70,35 +73,47 @@ def class_of(origin_kind: str | None) -> str: return ORIGIN_CLASS.get(origin_kind or "", "interactive") -def defaults(total: int | None = None) -> dict[str, Any]: - """缺省的设置（总上限取 AGENTLAB_MAX_CONCURRENT_RUNS）。""" +def _reserve_of(total: int, reserve: int | None) -> int: + """生效的交互留底：缺省全站上限的 1/4、至少 1；至多全站上限减 1（全站只有 1 个名额时没有留底，别的类才跑得起来）。""" + wanted = max(1, total // 4) if reserve is None else int(reserve) + return max(0, min(wanted, total - 1)) + + +def defaults(total: int | None = None, reserve: int | None = None) -> dict[str, Any]: + """缺省的设置（总上限取 AGENTLAB_MAX_CONCURRENT_RUNS）：交互上限 = 全站上限，别的三类 = 全站上限 − 交互留底。""" total = max(1, int(total if total is not None else settings.max_concurrent_runs)) - share = {"interactive": total, "app": max(1, total * 2 // 5), "schedule": max(1, total // 5), - "system": max(1, total // 5)} - return {"max_total": total, - "classes": {k: {"limit": share[k], "priority": DEFAULT_PRIORITY[k]} for k in CLASSES}} + kept = _reserve_of(total, reserve) + others = max(1, total - kept) + return {"max_total": total, "interactive_reserve": kept, + "classes": {k: {"limit": total if k == "interactive" else others, "priority": DEFAULT_PRIORITY[k]} + for k in CLASSES}} + + +def _int_in(value: Any, low: int, high: int) -> bool: + return isinstance(value, int) and not isinstance(value, bool) and low <= value <= high def values(stored: Any) -> dict[str, Any]: - """存着的（可能缺项、存坏了）→ 生效的：缺的、坏的补缺省。""" - base = defaults() - out = {"max_total": base["max_total"], "classes": {k: dict(v) for k, v in base["classes"].items()}} - if not isinstance(stored, dict): - return out - total = stored.get("max_total") - if isinstance(total, int) and not isinstance(total, bool) and 1 <= total <= LIMIT_MAX: - out["max_total"] = total + """存着的（可能缺项、存坏了）→ 生效的：缺的、坏的补缺省（缺省按存着的全站上限、留底算）。""" + stored = stored if isinstance(stored, dict) else {} + total = stored.get("max_total") if _int_in(stored.get("max_total"), 1, LIMIT_MAX) else None + reserve = stored.get("interactive_reserve") if _int_in(stored.get("interactive_reserve"), 0, LIMIT_MAX) else None + out = defaults(total, reserve) classes = stored.get("classes") if isinstance(stored.get("classes"), dict) else {} for k in CLASSES: item = classes.get(k) if isinstance(classes.get(k), dict) else {} - limit, prio = item.get("limit"), item.get("priority") - if isinstance(limit, int) and not isinstance(limit, bool) and 1 <= limit <= LIMIT_MAX: - out["classes"][k]["limit"] = limit - if isinstance(prio, int) and not isinstance(prio, bool) and 0 <= prio <= PRIORITY_MAX: - out["classes"][k]["priority"] = prio + if _int_in(item.get("limit"), 1, LIMIT_MAX): + out["classes"][k]["limit"] = item["limit"] + if _int_in(item.get("priority"), 0, PRIORITY_MAX): + out["classes"][k]["priority"] = item["priority"] return out +def borrowable(pol: dict[str, Any]) -> int: + """应用、定时、系统三类合起来最多用几个名额：全站上限 − 交互留底。""" + return max(0, int(pol["max_total"]) - int(pol.get("interactive_reserve") or 0)) + + def problem(sent: Any) -> str | None: """设置接口写进来的这一组有没有问题（说明给人看）。只核对给了的项。""" if not isinstance(sent, dict): @@ -106,6 +121,11 @@ def problem(sent: Any) -> str | None: total = sent.get("max_total") if total is not None and (not isinstance(total, int) or isinstance(total, bool) or not 1 <= total <= LIMIT_MAX): return f"全站同时运行的上限需要是 1 到 {LIMIT_MAX} 之间的整数" + reserve = sent.get("interactive_reserve") + if reserve is not None and not _int_in(reserve, 0, LIMIT_MAX): + return f"给交互留底的名额需要是 0 到 {LIMIT_MAX} 之间的整数" + if reserve is not None and isinstance(total, int) and total > 1 and reserve >= total: + return "给交互留底的名额要小于全站同时运行的上限：留底占满了，应用、定时、系统就一个都运行不了" classes = sent.get("classes") if classes is None: return None @@ -125,14 +145,18 @@ def problem(sent: Any) -> str | None: def merged(saved: Any, sent: dict[str, Any]) -> dict[str, Any]: - """只给了一部分的（只改一类的上限）补上存着的另几项。""" - out = values(saved) - if "max_total" in sent and sent["max_total"] is not None: - out["max_total"] = sent["max_total"] + """只给了一部分的（只改一类的上限）补上存着的另几项。存的是原样（没给的类别上限不写，按全站上限、留底现算缺省）。""" + base = dict(saved) if isinstance(saved, dict) else {} + out: dict[str, Any] = {k: base[k] for k in ("max_total", "interactive_reserve") if k in base} + for key in ("max_total", "interactive_reserve"): + if sent.get(key) is not None: + out[key] = sent[key] + classes = {k: dict(v) for k, v in (base.get("classes") or {}).items() if k in CLASSES and isinstance(v, dict)} for k, item in (sent.get("classes") or {}).items(): for field_ in ("limit", "priority"): if item.get(field_) is not None: - out["classes"][k][field_] = item[field_] + classes.setdefault(k, {})[field_] = item[field_] + out["classes"] = classes return out @@ -170,6 +194,8 @@ class _Holder: run_id: str loop: asyncio.AbstractEventLoop woke: asyncio.Event = field(default_factory=asyncio.Event) + #: 还在排队等名额（发下来之前） + queued: bool = True _holders: dict[str, _Holder] = {} @@ -201,37 +227,83 @@ def _grant_lock() -> asyncio.Lock: return lock +@dataclass +class _MemRow: + run_id: str + klass: str + priority: int + requested_at: Any + owner: str = "" + granted: bool = False + + +#: SQLite 时的名额（复核 R-b3）：只有一个进程，规矩照同一套、只是不落库——一段运行不为要名额多写三次库 +_mem: dict[str, _MemRow] = {} + + +def _in_memory() -> bool: + return settings.db_dialect == "sqlite" + + +def _pick(rows: list[Any], pol: dict[str, Any]) -> list[str]: + """按规矩挑这一次发给谁（库里的行、进程里的行同一套）：优先级高的先、同级先来先得，类别上限、全站上限、交互留底。""" + used = Counter(r.klass for r in rows if r.granted) + total = sum(used.values()) + # 交互以外的三类合起来用了几个：到了「全站上限 − 交互留底」就不再发给它们（留底永远给交互） + others = total - used["interactive"] + others_cap = borrowable(pol) + limits = {k: int(v["limit"]) for k, v in pol["classes"].items()} + granted: list[str] = [] + waiting = sorted((r for r in rows if not r.granted), key=lambda r: (-int(r.priority or 0), r.requested_at, r.run_id)) + for row in waiting: + if total >= int(pol["max_total"]): + break + if used[row.klass] >= limits.get(row.klass, int(pol["max_total"])): + continue + if row.klass != "interactive" and others >= others_cap: + continue + used[row.klass] += 1 + total += 1 + if row.klass != "interactive": + others += 1 + granted.append(row.run_id) + return granted + + async def grant() -> list[str]: """把等着的按规矩过一遍，够格的改成已发；属主不在了的行删掉。交回这一次发出去的。全站串行（见模块说明）。""" + if _in_memory(): + pol = await policy() + for run_id in [r for r in _mem if not _local(r)]: + _mem.pop(run_id, None) + granted = _pick(list(_mem.values()), pol) + for run_id in granted: + _mem[run_id].granted = True + if granted: + wake_local() + return granted from app.cluster import state from app.db.models_cluster import RunSlot pol = await policy() + # 拿锁之前先备好「谁还活着」（必要时这时候查库），持锁期间只用持锁的那条连接 + live = await state.current() + known_dead = await _known_dead(live) async with _grant_lock(): async with SessionLocal() as session: if session.get_bind().dialect.name == "postgresql": await session.execute(text("SELECT pg_advisory_xact_lock(:k)"), {"k": _GRANT_KEY}) rows = list((await session.execute(select(RunSlot))).scalars()) - live = await state.live_for([r.owner for r in rows]) - dead = [r.run_id for r in rows if not live.alive(r.owner, local=_local(r.run_id))] + # 持着全站的锁时不另取连接查库（复核：连接池用尽时会卡到取连接超时）：用锁之前拿好的那一份「谁还活着」，不认识的属主 + # 这一次当活着（不删它的行），下一次手上那份更新了再说 + me = identity.owner() + dead = [r.run_id for r in rows + if (r.owner == me and not _local(r.run_id)) or (r.owner != me and r.owner in known_dead)] if dead: await session.execute(delete(RunSlot).where(RunSlot.run_id.in_(dead))) rows = [r for r in rows if r.run_id not in set(dead)] - used = Counter(r.klass for r in rows if r.granted) - total = sum(used.values()) - limits = {k: int(v["limit"]) for k, v in pol["classes"].items()} - granted: list[str] = [] - waiting = sorted((r for r in rows if not r.granted), - key=lambda r: (-int(r.priority or 0), r.requested_at, r.run_id)) + granted = _pick(rows, pol) now = utcnow() - for row in waiting: - if total >= int(pol["max_total"]): - break - if used[row.klass] >= limits.get(row.klass, int(pol["max_total"])): - continue - used[row.klass] += 1 - total += 1 - granted.append(row.run_id) if granted: await session.execute(update(RunSlot).where(RunSlot.run_id.in_(granted)) .values(granted=True, granted_at=now).execution_options(synchronize_session=False)) @@ -244,8 +316,25 @@ async def grant() -> list[str]: return granted +async def _known_dead(live: Any) -> set[str]: + """名额表里属主不在手上那份「活着」里的那几个，拿锁之前当场再查一次登记：仍不在的算没了。""" + from app.cluster import state + from app.db.models_cluster import RunSlot + + async with SessionLocal() as session: + owners = set((await session.execute(select(RunSlot.owner).distinct())).scalars()) + unknown = {o for o in owners if o and o != identity.owner() and o not in live.others} + if not unknown: + return set() + fresh = await state.current(fresh=True) + return {o for o in unknown if o not in fresh.others} + + async def _granted(run_id: str) -> bool | None: """这一行发了没有：True 发了，False 还在等，None 这一行不见了（被当成属主不在了删掉）。""" + if _in_memory(): + row = _mem.get(run_id) + return None if row is None else row.granted from app.db.models_cluster import RunSlot async with SessionLocal() as session: @@ -253,6 +342,9 @@ async def _granted(run_id: str) -> bool | None: async def _request(run_id: str, klass: str, priority: int) -> None: + if _in_memory(): + _mem[run_id] = _MemRow(run_id, klass, priority, utcnow(), identity.owner()) + return from app.db.models_cluster import RunSlot async with SessionLocal() as session: @@ -263,6 +355,9 @@ async def _request(run_id: str, klass: str, priority: int) -> None: async def _release(run_id: str) -> None: + if _in_memory(): + _mem.pop(run_id, None) + return from app.db.models_cluster import RunSlot async with SessionLocal() as session: @@ -286,6 +381,7 @@ async def hold(run_id: str, origin_kind: str | None) -> AsyncIterator[str]: await grant() got = await _granted(run_id) if got: + holder.queued = False break if got is None: await _request(run_id, klass, priority) @@ -304,17 +400,32 @@ async def hold(run_id: str, origin_kind: str | None) -> AsyncIterator[str]: signals.publish("slot") +async def waiting(run_id: str, *, here: bool = False) -> bool: + """这次运行是不是正排着队等名额（评测每题的时限、算期望的时限不算这一段，eval/runner、eval/gold）。 + here：这次运行的执行任务在本进程里（还没开始要、或者已经放掉了名额的就不是在等，不用查库）。""" + holder = _holders.get(run_id) + if holder is not None: + return holder.queued + if here: + return False + return (await _granted(run_id)) is False + + async def snapshot() -> dict[str, Any]: """现在各类在跑、在等几个（系统信息、验收脚本看）。""" - from app.db.models_cluster import RunSlot + if _in_memory(): + rows = [(r.klass, r.granted) for r in _mem.values() if _local(r.run_id)] + else: + from app.db.models_cluster import RunSlot - async with SessionLocal() as session: - rows = list((await session.execute(select(RunSlot.klass, RunSlot.granted))).all()) - out = {k: {"running": 0, "waiting": 0} for k in CLASSES} + async with SessionLocal() as session: + rows = list((await session.execute(select(RunSlot.klass, RunSlot.granted))).all()) + out = {k: {"running": 0, "waiting": 0} for k in CLASSES} # noqa: C420 for klass, granted in rows: out.setdefault(klass, {"running": 0, "waiting": 0})["running" if granted else "waiting"] += 1 return {"policy": await policy(), "classes": out} -__all__ = ["CLASSES", "CLASS_LABEL", "ORIGIN_CLASS", "SETTING_KEY", "class_of", "defaults", "forget_policy", "grant", +__all__ = ["CLASSES", "CLASS_LABEL", "ORIGIN_CLASS", "SETTING_KEY", "borrowable", "class_of", "defaults", "forget_policy", "grant", + "waiting", "hold", "merged", "policy", "problem", "snapshot", "values", "wake_local"] diff --git a/backend/app/cluster/version.py b/backend/app/cluster/version.py new file mode 100644 index 00000000..7020bd5a --- /dev/null +++ b/backend/app/cluster/version.py @@ -0,0 +1,189 @@ +"""这个进程是哪一版代码、要的是哪一版库结构；库里还有活着的、版本不同的进程时不许加入（不支持新旧版本混跑）。 + +**为什么要拦。** 升级时只停了一部分进程、或者新旧两批容器同时连着一个库：新版的进程启动时补列、补索引，旧版的进程不认识新的列和 +规矩（认领、租约、接手的说法），两边各按各的写同一批行。以前只在 README 里写「升级要先停掉全部进程」，现在代码里拦住。 + +- 代码版本 code_version：包版本 + 构建号（AGENTLAB_BUILD，Docker 镜像构建时写进去）或 git 提交 + 源码指纹（app/ 下全部 .py + 的内容哈希）。比的是源码指纹：同一份代码一定相同，改了一行就不同，取不到 git、没有构建号时也区分得开。 +- 库结构版本 schema_version：模型里各表各列（名字、类型、可空）加迁移清单的哈希。 + +怎么拦：PostgreSQL 时，补列之前先看一眼（preflight）——还有活着的、版本不同的进程就不启动；它的心跳停着（多半是强杀之后留下的、 +租约还没到期）就等它过期再起，心跳还在走就照实报错：哪几个进程、什么版本、要先停掉再升级。登记那一下再在事务级咨询锁里查一遍 +（cluster/workers.register），两个不同版本同时启动也只有先登记上的那一个起得来。SQLite 只有一个进程，不用看。 +""" +from __future__ import annotations + +import asyncio +import hashlib +import logging +import os +import subprocess +from functools import lru_cache +from pathlib import Path +from typing import Any + +from app.cluster import identity + +logger = logging.getLogger(__name__) + +_APP_DIR = Path(__file__).resolve().parents[1] + + +class VersionMismatch(RuntimeError): + """库里还有活着的、版本不同的进程。""" + + +@lru_cache(maxsize=1) +def source_fingerprint() -> str: + digest = hashlib.sha256() + for path in sorted(_APP_DIR.rglob("*.py")): + if "__pycache__" in path.parts: + continue + digest.update(str(path.relative_to(_APP_DIR)).encode()) + try: + digest.update(path.read_bytes()) + except OSError: + continue + return digest.hexdigest()[:12] + + +def _package() -> str: + try: + from importlib.metadata import version + + return version("agentlab-backend") + except Exception: # noqa: BLE001 + return "0" + + +def _build() -> str: + if build := (os.environ.get("AGENTLAB_BUILD") or "").strip(): + return build[:40] + try: + out = subprocess.run(["git", "rev-parse", "--short=10", "HEAD"], cwd=str(_APP_DIR), capture_output=True, + text=True, timeout=3) + return out.stdout.strip() if out.returncode == 0 else "" + except (OSError, subprocess.SubprocessError): + return "" + + +@lru_cache(maxsize=1) +def code_version() -> str: + """包版本+构建号或 git 提交+源码指纹，例如 0.1.0+e6a3ee8f12+3f9a…（最后一段是比较用的）。""" + return "+".join(p for p in (_package(), _build(), source_fingerprint()) if p) + + +def fingerprint_of(version: str | None) -> str: + return (version or "").rsplit("+", 1)[-1] + + +@lru_cache(maxsize=1) +def schema_version() -> str: + from app.db import models # noqa: F401 模型都登记进 metadata + from app.db.base import _COLUMN_MIGRATIONS, _INDEX_MIGRATIONS, Base + + digest = hashlib.sha256() + for name in sorted(Base.metadata.tables): + table = Base.metadata.tables[name] + digest.update(name.encode()) + for col in table.columns: + digest.update(f"{col.name}:{col.type!r}:{col.nullable}".encode()) + digest.update(repr(_COLUMN_MIGRATIONS).encode() + repr(_INDEX_MIGRATIONS).encode()) + return digest.hexdigest()[:12] + + +def same(code: str | None, schema: str | None) -> bool: + return fingerprint_of(code) == fingerprint_of(code_version()) and (schema or "") == schema_version() + + +async def _peers(conn: Any) -> list[dict[str, Any]] | None: + """库里登记着的别的进程（id、版本、租约、心跳）；表还不在（新库）交回 None。在一条连接上查（preflight 在建表之前）。""" + from sqlalchemy import inspect, text + + def has(sync_conn: Any) -> tuple[bool, set[str]]: + insp = inspect(sync_conn) + if not insp.has_table("cluster_workers"): + return False, set() + return True, {c["name"] for c in insp.get_columns("cluster_workers")} + + exists, cols = await conn.run_sync(has) + if not exists: + return None + extra = ", ".join(c for c in ("code_version", "schema_version") if c in cols) + rows = (await conn.execute(text(f"SELECT id, lease_until, heartbeat_at{', ' + extra if extra else ''} " + "FROM cluster_workers WHERE id <> :me"), {"me": identity.owner()})).mappings().all() + return [dict(r) for r in rows] + + +def _describe(peer: dict[str, Any]) -> str: + code = peer.get("code_version") or "（升级前的版本，没有记版本）" + schema = peer.get("schema_version") or "未知" + return f"{identity.describe(peer['id'])}（代码 {code}，库结构 {schema}）" + + +async def preflight(*, observe_s: float | None = None) -> None: + """PostgreSQL 时、建表补列之前：库里还有活着的、版本不同的进程就不启动（见模块说明）。observe_s：看它心跳还在不在走的时长 + （缺省两拍心跳）。""" + from app.core.config import settings + + if settings.db_dialect != "postgresql": + return + from app.cluster.workers import heartbeat_s, lease_s + from app.db.base import utcnow + + observe = 2 * heartbeat_s() + 1.0 if observe_s is None else observe_s + loop = asyncio.get_running_loop() + started = loop.time() + first_beats: dict[str, Any] = {} + warned = False + while True: + mismatched = await _mismatched_alive(utcnow()) + if not mismatched: + return + for peer in mismatched: + first_beats.setdefault(peer["id"], peer.get("heartbeat_at")) + # 心跳还在走：真活着，照实报错。心跳停着的（强杀、断电留下的登记）等它们的租约过期 + beating = [p for p in mismatched if p.get("heartbeat_at") != first_beats.get(p["id"])] + waited = loop.time() - started + if beating or waited >= lease_s() * 2 + observe: + raise VersionMismatch(_message(beating or mismatched)) + if waited >= observe and not warned: + warned = True + logger.warning("库里登记着 %d 个版本不同的进程，心跳停着（多半是强杀之后留下的）：等它们的租约过期再启动", + len(mismatched)) + await asyncio.sleep(min(0.5, max(0.1, heartbeat_s() / 2))) + + +async def _mismatched_alive(now: Any) -> list[dict[str, Any]]: + from app.db.base import engine + + async with engine.connect() as conn: + peers = await _peers(conn) + if not peers: + return [] + naive = now.replace(tzinfo=None) + alive = [p for p in peers if p.get("lease_until") is not None and _naive(p["lease_until"]) > naive] + return [p for p in alive if not same(p.get("code_version"), p.get("schema_version"))] + + +def _naive(moment: Any) -> Any: + return moment.replace(tzinfo=None) if getattr(moment, "tzinfo", None) else moment + + +def _message(peers: list[dict[str, Any]]) -> str: + listed = "；".join(_describe(p) for p in peers) + return (f"库里还有 {len(peers)} 个版本不同的 AgentLab 进程活着：{listed}。本进程是代码 {code_version()}、库结构 " + f"{schema_version()}。不支持新旧版本同时运行：先停掉这些进程再启动新版本（升级步骤见 README「多进程」的「升级」）") + + +def check_rows(rows: list[Any], now: Any) -> None: + """登记那一下（cluster/workers.register，持着事务级咨询锁）再查一遍：别的活着的进程里有版本不同的就抛 VersionMismatch。""" + naive = now.replace(tzinfo=None) + bad = [{"id": r.id, "code_version": r.code_version, "schema_version": r.schema_version} for r in rows + if r.lease_until is not None and _naive(r.lease_until) > naive and r.id != identity.owner() + and not same(r.code_version, r.schema_version)] + if bad: + raise VersionMismatch(_message(bad)) + + +__all__ = ["VersionMismatch", "check_rows", "code_version", "preflight", "same", "schema_version", "source_fingerprint"] diff --git a/backend/app/cluster/workers.py b/backend/app/cluster/workers.py index db9dd624..adbd6ef2 100644 --- a/backend/app/cluster/workers.py +++ b/backend/app/cluster/workers.py @@ -2,7 +2,7 @@ - register / renew / unregister：启动时登记、隔一拍心跳续租约、平稳关停时删掉自己那一行（别的进程马上就知道它不在了， 不必等租约到期）。 -- liveness()：现在哪些别的进程还活着（租约没过期）。各种「启动整理」「接手」都只碰属主不在了的行： +- liveness()：现在哪些别的进程还活着（租约过期不到一个租约时长：再晚才算不在了，见 cluster/fence.py）。各种「启动整理」「接手」都只碰属主不在了的行： Liveness.alive(owner, local=…) —— 属主是别的进程：看它的租约；属主是自己：看这一行在本进程里是不是真的在跑（local）， 同一个进程里前一次起停留下的（测试里一个进程起停多次服务）照样算不在了；没有属主（升级前的行）算不在了。 @@ -25,6 +25,10 @@ from app.db.base import SessionLocal, utcnow logger = logging.getLogger(__name__) +#: 登记时全站串行用的事务级咨询锁的号 +_REGISTER_KEY = 0x4167576F726B6572 # "AgWorker" + + def lease_s() -> float: """租约多长（秒）：至少是心跳间隔的三倍——库一时慢了、一拍没续上，不算进程没了。""" return max(float(settings.worker_lease_s), 3.0 * float(settings.worker_heartbeat_s)) @@ -61,9 +65,12 @@ async def liveness(now: datetime | None = None) -> Liveness: now = now or utcnow() me = identity.owner() + # 租约过期之后再多等一个租约时长才算不在了：卡住了但没死的进程（SIGSTOP、虚拟机暂停、库一时连不上）自己在一个租约时长时就 + # 停手（cluster/fence.py），两边中间留着一段，不会它还在写、这边已经接手 + floor = now - timedelta(seconds=lease_s()) async with SessionLocal() as session: ids = (await session.execute(select(ClusterWorker.id).where( - ClusterWorker.lease_until > now, ClusterWorker.id != me))).scalars().all() + ClusterWorker.lease_until > floor, ClusterWorker.id != me))).scalars().all() return Liveness(frozenset(ids)) @@ -83,15 +90,25 @@ async def register(*, exclusive: bool = False) -> None: 守卫已经确认没有别的进程在用这个数据目录，见 cluster/guard.py）。""" from app.db.models_cluster import ClusterWorker + from app.cluster import version + me = identity.owner() now = utcnow() async with SessionLocal() as session: if exclusive: await session.execute(delete(ClusterWorker).where(ClusterWorker.id != me)) + elif session.get_bind().dialect.name == "postgresql": + # 查「还有没有活着的、版本不同的进程」和登记在同一个事务里、全站串行：两个不同版本同时启动只有先登记上的起得来 + from sqlalchemy import text + + await session.execute(text("SELECT pg_advisory_xact_lock(:k)"), {"k": _REGISTER_KEY}) + others = (await session.execute(select(ClusterWorker).where(ClusterWorker.id != me))).scalars().all() + version.check_rows(list(others), now) row = await session.get(ClusterWorker, me) if row is None: session.add(ClusterWorker(id=me, host=identity.host(), pid=os.getpid(), started_at=now, heartbeat_at=now, - lease_until=now + timedelta(seconds=lease_s()), info=_info())) + lease_until=now + timedelta(seconds=lease_s()), info=_info(), + code_version=version.code_version(), schema_version=version.schema_version())) else: row.heartbeat_at, row.lease_until = now, now + timedelta(seconds=lease_s()) await session.commit() @@ -137,6 +154,7 @@ async def listing() -> list[dict[str, Any]]: async with SessionLocal() as session: rows = (await session.execute(select(ClusterWorker).order_by(ClusterWorker.started_at, ClusterWorker.id))).scalars().all() return [{"id": r.id, "host": r.host, "pid": r.pid, "alive": r.lease_until > now, "me": r.id == identity.owner(), + "code_version": r.code_version, "schema_version": r.schema_version, "started_at": r.started_at.isoformat() if r.started_at else None, "heartbeat_at": r.heartbeat_at.isoformat() if r.heartbeat_at else None} for r in rows] diff --git a/backend/app/core/config.py b/backend/app/core/config.py index 890db795..3fce235d 100644 --- a/backend/app/core/config.py +++ b/backend/app/core/config.py @@ -56,6 +56,13 @@ class Settings(BaseSettings): # LangGraph 的断点也放进同一个库；checkpoint_url 另给就放那里。数据目录照旧要有：工件、上传表格、沙箱等仍是文件 database_url: str = "" checkpoint_url: str = "" + # 直连地址（可选，写法同 database_url）：普通查询走 database_url（可以是 PgBouncer），依赖会话状态的连接走这里——进程间通知的 + # LISTEN 那一条、跨进程锁（会话级咨询锁）、断点库（预备语句、建表时的会话锁）。不设就全用 database_url + database_direct_url: str = "" + # database_url 是 PgBouncer 的事务模式（AGENTLAB_DB_PGBOUNCER=1）：应用库不用客户端连接池（交给 PgBouncer）、关掉 asyncpg 和 + # SQLAlchemy 的预备语句缓存、预备语句用不会撞名的名字。这时必须另给 database_direct_url（事务模式下 LISTEN、会话级咨询锁不成立）， + # 启动时检查 + db_pgbouncer: bool = False # PostgreSQL 的连接池（SQLite 不用这几项）。pool_size=0 表示不用连接池、每次新建连接（前面有 PgBouncer 时用， # 测试里也这么用：每个测试一个事件循环，asyncpg 的连接不能跨循环） db_pool_size: int = 10 @@ -224,7 +231,7 @@ class Settings(BaseSettings): return True return value - @field_validator("database_url", "checkpoint_url", mode="before") + @field_validator("database_url", "checkpoint_url", "database_direct_url", mode="before") @classmethod def _valid_postgres_url(cls, value: object, info: Any) -> object: # 写错了启动时就报出来，别等到第一次查库。只认 PostgreSQL：SQLite 就是不设 @@ -235,7 +242,7 @@ class Settings(BaseSettings): from app.db.urls import normalize_database_url, normalize_checkpoint_url name = f"AGENTLAB_{(info.field_name or '').upper()}" - normalize = normalize_database_url if info.field_name == "database_url" else normalize_checkpoint_url + normalize = normalize_checkpoint_url if info.field_name == "checkpoint_url" else normalize_database_url try: return normalize(value.strip()) except ValueError as e: @@ -296,15 +303,21 @@ class Settings(BaseSettings): """LangGraph checkpointer 用独立的库文件，避免和业务表抢写锁。只在 SQLite 时用。""" return self.data_dir / "checkpoints.db" + @property + def session_db_url(self) -> str: + """依赖会话状态的连接（LISTEN、会话级咨询锁、断点库）用的应用库地址：设了直连地址用它，否则同 database_url。""" + return self.database_direct_url or self.database_url + @property def checkpoint_conninfo(self) -> str | None: - """PostgreSQL 断点库的 libpq 连接串（psycopg 用）：checkpoint_url 优先，否则由 database_url 换算。SQLite 时 None。""" + """PostgreSQL 断点库的 libpq 连接串（psycopg 用）：checkpoint_url 优先，否则由直连地址（没设就是 database_url）换算。 + SQLite 时 None。断点库用预备语句、建表时取会话级的锁，所以缺省走直连。""" if self.checkpoint_url: return self.checkpoint_url if self.database_url: from app.db.urls import libpq_conninfo - return libpq_conninfo(self.database_url) + return libpq_conninfo(self.session_db_url) return None @property diff --git a/backend/app/core/job_runtime.py b/backend/app/core/job_runtime.py index e860e0d7..b050ee83 100644 --- a/backend/app/core/job_runtime.py +++ b/backend/app/core/job_runtime.py @@ -78,6 +78,11 @@ class SerialWriter: self._lock = asyncio.Lock() async def __call__(self, fn: Callable[..., Awaitable[Any]], *args: Any) -> Any: + from app.cluster import fence + + # 本进程续不上租约了（cluster/fence.py）：这个任务可能已经被别的进程接手，不再写 + fence.check() + async def go() -> Any: async with self._lock: async with SessionLocal() as session: @@ -236,16 +241,40 @@ class SourceLanes: await session.execute(text("SELECT pg_advisory_xact_lock(:k)"), {"k": advisory_key(f"lane:{source_id}")}) - async def _rows(self, session: AsyncSession, source_id: str) -> list[Any]: - """这个源队里活着的各行（号小的在前）；属主不在了的删掉（不提交）。""" - from app.cluster import state + async def _dead_owners(self, source_id: str) -> frozenset[str]: + """拿这个源的锁之前备好：队里别的进程的属主哪些确实不在了（手上那份不认识的当场再查一次登记）。持着咨询锁时不另取连接 + 查库（复核：连接池用尽时会卡到取连接超时）；拿锁之后才出现的属主不在这里面，这一次当活着，下一次再说。""" + from app.cluster import identity, state + from app.db.models_cluster import LaneTicket as Row + + async with SessionLocal() as session: + owners = set((await session.execute(select(Row.owner).where(Row.source_id == source_id).distinct())) + .scalars()) + me = identity.owner() + others = {o for o in owners if o and o != me} + if not others: + return frozenset() + live = await state.live_for(others) + return frozenset(o for o in others if o not in live.others) + + async def _rows(self, session: AsyncSession, source_id: str, dead_owners: frozenset[str]) -> list[Any]: + """这个源队里活着的各行（号小的在前）；属主不在了的删掉（不提交）。dead_owners 是拿锁之前备好的那一份：属主是别的进程的 + 只认它；属主是自己的看本进程里登记着没有（正有人排进来时一律算活着）。""" + from app.cluster import identity from app.db.models_cluster import LaneTicket as Row rows = list((await session.execute(select(Row).where(Row.source_id == source_id).order_by(Row.id))).scalars()) - live = await state.live_for([r.owner for r in rows]) - if self._entering.get(source_id): - live = live.trusting() - dead = [r.id for r in rows if not live.alive(r.owner, local=self._local_alive(r.id))] + me = identity.owner() + entering = bool(self._entering.get(source_id)) + + def alive(row: Any) -> bool: + if not row.owner: + return False + if row.owner == me: + return entering or self._local_alive(row.id) + return row.owner not in dead_owners + + dead = [r.id for r in rows if not alive(r)] if dead: from sqlalchemy import delete @@ -257,9 +286,10 @@ class SourceLanes: from app.db.models_cluster import LaneTicket as Row granted_now: int | None = None + dead_owners = await self._dead_owners(source_id) async with self._lock(), SessionLocal() as session: await self._locked(session, source_id) - rows = await self._rows(session, source_id) + rows = await self._rows(session, source_id, dead_owners) head = next((r for r in rows if r.granted), None) if head is None and rows: head = rows[0] @@ -313,9 +343,10 @@ class SourceLanes: if kind not in _LANE_KINDS: raise ValueError(f"同源队里没有登记「{kind}」这种事，先 register_lane_kind") self._listen() + dead_owners = await self._dead_owners(source_id) async with self._lock(), SessionLocal() as session: await self._locked(session, source_id) - rows = await self._rows(session, source_id) + rows = await self._rows(session, source_id, dead_owners) if rows: head = next((r for r in rows if r.granted), rows[0]) await session.commit() @@ -332,8 +363,9 @@ class SourceLanes: async def holders(self, source_id: str) -> list[LaneHolder]: """队里的各位，排在最前面（正在做）的在前。""" + dead_owners = await self._dead_owners(source_id) async with SessionLocal() as session: - rows = await self._rows(session, source_id) + rows = await self._rows(session, source_id, dead_owners) await session.commit() rows.sort(key=lambda r: (not r.granted, r.id)) return [LaneHolder(r.kind, r.job_id, r.owner) for r in rows] @@ -509,14 +541,22 @@ class JobRun: await self.writer.settle() async def run(self) -> None: + from app.cluster import fence + + fence.bind() try: try: await self._queue() await self._run() + except fence.Fenced: + # 本进程续不上租约了：停手、不写收尾（这个任务由接手的进程按「属主不在了」处理） + logger.error("%s %s：本进程续不上租约，停手、不写收尾", self.kind, self.job_id) except asyncio.CancelledError: current = asyncio.current_task() if current is not None: current.uncancel() + if fence.fenced_here(): + return await self._stopped() except JobFailed as e: await self._failed(str(e)) @@ -568,9 +608,18 @@ class JobManager: self._factory = factory self.label = label self.name = name - from app.cluster import requests + from app.cluster import fence, requests requests.handle("stop", name)(self._stop_requested) + fence.on_trip(self._stop_fenced) + + def _stop_fenced(self) -> None: + """本进程隔离了（cluster/fence.py）：停掉手上的任务，停的时候不写收尾。""" + for job_id in list(self._runs): + run = self._runs.get(job_id) + if run is not None and run.task is not None and not run.task.done(): + run.stop_reason = run.stop_reason or "fenced" + run.task.cancel() def model(self) -> Any: """任务表（停止时读属主用）：缺省取 factory 上的 model。""" diff --git a/backend/app/db/base.py b/backend/app/db/base.py index b8224fb7..85434064 100644 --- a/backend/app/db/base.py +++ b/backend/app/db/base.py @@ -188,12 +188,45 @@ def pg_json_dumps(value: Any) -> str: APPLICATION_NAME = "agentlab" -def make_engine(url: str | None = None, *, application_name: str = APPLICATION_NAME) -> AsyncEngine: +def pg_engine_kwargs(*, application_name: str = APPLICATION_NAME, pgbouncer: bool = False, + pool_size: int | None = None) -> dict[str, Any]: + """PostgreSQL 引擎的连接池和连接参数。pgbouncer：这个地址是 PgBouncer 的事务模式（AGENTLAB_DB_PGBOUNCER）—— + + - 不用客户端的连接池（NullPool，连接池交给 PgBouncer；SQLAlchemy 的说明里同样这么建议）； + - SQLAlchemy 的预备语句缓存（prepared_statement_cache_size）和 asyncpg 自己的（statement_cache_size）都关掉：事务模式下同一个 + 客户端连接的下一条语句可能落在另一条服务端连接上，缓存着的预备语句在那边不存在； + - asyncpg 每条语句仍要先预备再执行（SQLAlchemy 的 asyncpg 方言一律这么做，同一个事务里，落在同一条服务端连接上），名字用随机的 + （prepared_statement_name_func），不和别的客户端在同一条服务端连接上留下的撞名。PgBouncer 要配 server_reset_query=DISCARD ALL + （且 server_reset_query_always=1）收掉留下的预备语句，或用 1.21 以上开 max_prepared_statements。 + 事务级的咨询锁（pg_advisory_xact_lock，发名额、同源队、建表时用）在事务模式下照常成立：取锁和放锁都在同一个事务里。""" + connect_args: dict[str, Any] = { + # 会话时区定成 UTC：时间列是不带时区的 UTC（UTCDateTime），库里 now() 之类的也按 UTC 算 + "server_settings": {"application_name": application_name, "timezone": "UTC"}} + size = settings.db_pool_size if pool_size is None else pool_size + if pgbouncer: + connect_args.update(statement_cache_size=0, prepared_statement_cache_size=0, + prepared_statement_name_func=_unique_statement_name) + return {"poolclass": NullPool, "connect_args": connect_args} + if size <= 0: + return {"poolclass": NullPool, "connect_args": connect_args} + return {"pool_size": size, "max_overflow": max(settings.db_max_overflow, 0), "pool_timeout": settings.db_pool_timeout, + "pool_recycle": settings.db_pool_recycle, "connect_args": connect_args} + + +def _unique_statement_name() -> str: + return f"__agentlab_{uuid.uuid4().hex}__" + + +def make_engine(url: str | None = None, *, application_name: str = APPLICATION_NAME, pgbouncer: bool | None = None, + pool_size: int | None = None) -> AsyncEngine: """应用库的引擎。SQLite：WAL、外键、忙等，单个文件；PostgreSQL：连接池、方言调整（db/postgres.py）。 copy-to-postgres 建目标库的引擎也走这里，两边的列类型、排序规则是同一套。 """ url = url or settings.db_url + if pgbouncer is None: + # 只有应用库的普通地址按 AGENTLAB_DB_PGBOUNCER 处理：直连地址、copy-to-postgres 给的目标库照常 + pgbouncer = bool(settings.db_pgbouncer) and url == settings.database_url if url.startswith("sqlite"): # pool_pre_ping：取连接前先探一下活。取消异常打进 SQLAlchemy 的清理流程时 # （_close_connection 记完日志会把 CancelledError 原样抛出，跳过了把连接记录 @@ -206,15 +239,9 @@ def make_engine(url: str | None = None, *, application_name: str = APPLICATION_N from app.db import postgres _require_postgres_driver("asyncpg") - pool: dict[str, Any] = ( - {"poolclass": NullPool} if settings.db_pool_size <= 0 else { - "pool_size": settings.db_pool_size, "max_overflow": max(settings.db_max_overflow, 0), - "pool_timeout": settings.db_pool_timeout, "pool_recycle": settings.db_pool_recycle}) - made = create_async_engine( - url, echo=False, future=True, pool_pre_ping=True, json_serializer=pg_json_dumps, - # 会话时区定成 UTC：时间列是不带时区的 UTC（UTCDateTime），库里 now() 之类的也按 UTC 算 - connect_args={"server_settings": {"application_name": application_name, "timezone": "UTC"}}, - **pool) + made = create_async_engine(url, echo=False, future=True, pool_pre_ping=True, json_serializer=pg_json_dumps, + **pg_engine_kwargs(application_name=application_name, pgbouncer=pgbouncer, + pool_size=pool_size)) postgres.install(made.sync_engine.dialect) event.listen(made.sync_engine, "before_cursor_execute", postgres.strip_nul, retval=True) return made @@ -250,6 +277,35 @@ engine = make_engine() SessionLocal = async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession) +_session_engine: AsyncEngine | None = None + + +def session_engine() -> AsyncEngine: + """依赖会话状态的连接用的引擎（跨进程锁的会话级咨询锁，app/cluster/mutex）：设了 AGENTLAB_DATABASE_DIRECT_URL 时是直连那个库的 + 一个小池子，否则就是应用库的引擎。""" + global _session_engine + if not settings.database_direct_url or settings.database_direct_url == settings.database_url: + return engine + if _session_engine is None: + _session_engine = make_engine(settings.database_direct_url, pgbouncer=False, pool_size=max(2, min(settings.db_pool_size or 5, 10))) + return _session_engine + + +async def dispose_session_engine() -> None: + """关停时关掉直连那个小池子（没另建就什么都不做）。""" + global _session_engine + if _session_engine is not None: + await _session_engine.dispose() + _session_engine = None + + +def check_pgbouncer() -> None: + """AGENTLAB_DB_PGBOUNCER=1 却没给直连地址：拒绝启动（事务模式下 LISTEN、会话级咨询锁不成立，跨进程的锁会锁不住、放不掉）。""" + if settings.db_pgbouncer and settings.db_dialect == "postgresql" and not settings.database_direct_url: + raise RuntimeError("AGENTLAB_DB_PGBOUNCER=1（应用库地址是 PgBouncer 的事务模式）时要另给 AGENTLAB_DATABASE_DIRECT_URL：" + "进程间通知（LISTEN）、跨进程的锁（会话级咨询锁）、断点库要直连 PostgreSQL（或 PgBouncer 的会话模式）。" + "见 README「用 PostgreSQL」里的三种部署") + async def get_session() -> AsyncIterator[AsyncSession]: """FastAPI 依赖注入用的会话。""" @@ -398,6 +454,9 @@ _COLUMN_MIGRATIONS: list[tuple[str, str, str]] = [ ("publish_plans", "owner", "VARCHAR(120)"), # 问数据：等不及的那次复核还在后台跑（实时流在它回来之前不关；多进程时别的进程上的流也看得到） ("chat_turn_jobs", "late_review_at", "DATETIME"), + # 工作进程登记的代码版本、库结构版本（app/cluster/version.py：不许新旧版本混跑） + ("cluster_workers", "code_version", "VARCHAR(200)"), + ("cluster_workers", "schema_version", "VARCHAR(64)"), ] #: 迁移补上的列要建的索引：(索引名, 表, 列或多列)。create_all 不碰已有的表，表上新列的索引也就不会建。 @@ -512,6 +571,7 @@ def _migrate(conn) -> None: # pragma: no cover - 同步回调 cols = columns.get(table) or set() if all(c in cols for c in wanted): conn.execute(text(f"CREATE INDEX IF NOT EXISTS {name} ON {table} ({', '.join(wanted)})")) + _unique_run_event_seq(conn, columns) for name, table, wanted, where, dedupe in _UNIQUE_INDEX_MIGRATIONS: cols = columns.get(table) or set() if not all(c in cols for c in wanted): @@ -524,6 +584,29 @@ def _migrate(conn) -> None: # pragma: no cover - 同步回调 _backfill_not_null(conn, columns) +def _unique_run_event_seq(conn: Any, columns: dict[str, set[str]]) -> None: # pragma: no cover - 同步回调 + """run_events 的 (run_id, seq) 补唯一索引（ux_run_events_run_seq，换掉原来不唯一的 ix_run_events_run_seq）。存量库里可能已经 + 有撞了号的（多进程之前的进程内计数、按需裁判的老写法）：先数一数，有就照实记日志、这一次不建（照旧用不唯一的索引），不让启动 + 失败——那几次运行的封存核对本来就对不上，删哪一条都不对，要人看。""" + from sqlalchemy import inspect, text + + if "run_events" not in columns: + return + names = {ix["name"] for ix in inspect(conn).get_indexes("run_events")} + if "ux_run_events_run_seq" in names: + return + dupes = conn.execute(text("SELECT run_id, seq, COUNT(*) FROM run_events GROUP BY run_id, seq HAVING COUNT(*) > 1 " + "LIMIT 20")).all() + if dupes: + runs = sorted({str(r[0]) for r in dupes}) + logger.warning("run_events 里有撞了号的事件（%d 处以上，涉及运行 %s）：这一次不建 (run_id, seq) 的唯一索引。" + "这几次运行的封存核对对不上，请人工核对后处理", len(dupes), "、".join(runs[:5])) + return + conn.execute(text("CREATE UNIQUE INDEX IF NOT EXISTS ux_run_events_run_seq ON run_events (run_id, seq)")) + if "ix_run_events_run_seq" in names: + conn.execute(text("DROP INDEX IF EXISTS ix_run_events_run_seq")) + + #: init_db 在 PostgreSQL 上拿的会话级咨询锁的号：几个进程同时启动时，建表和迁移一个一个来（同一个号别处别用） _INIT_LOCK = 0x4167656E744C6162 # "AgentLab" diff --git a/backend/app/db/models.py b/backend/app/db/models.py index b010432a..bf7f41ce 100644 --- a/backend/app/db/models.py +++ b/backend/app/db/models.py @@ -721,7 +721,9 @@ class RunEvent(Base): """完整事件流落库 —— 这既是 trace，也是页面刷新后重建时间线的依据。""" __tablename__ = "run_events" - __table_args__ = (Index("ix_run_events_run_seq", "run_id", "seq"),) + # (run_id, seq) 唯一（编排完善第 2 波 W2-b）：撞号时报错，不悄悄写两条——封存按 seq 排序算哈希，撞了号核对时对时不对。 + # 存量库由 db/base._migrate 补建（先查有没有重复，有就照实记日志、不建） + __table_args__ = (Index("ux_run_events_run_seq", "run_id", "seq", unique=True),) id: Mapped[str] = mapped_column(String(32), primary_key=True, default=new_id) run_id: Mapped[str] = mapped_column(ForeignKey("runs.id", ondelete="CASCADE")) diff --git a/backend/app/db/models_cluster.py b/backend/app/db/models_cluster.py index 9555be21..6e2fe63f 100644 --- a/backend/app/db/models_cluster.py +++ b/backend/app/db/models_cluster.py @@ -38,6 +38,9 @@ class ClusterWorker(Base): lease_until: Mapped[datetime] = mapped_column(UTCDateTime, default=utcnow, index=True) #: 版本、数据目录之类给人看的信息 info: Mapped[dict[str, Any]] = mapped_column(default=dict) + #: 代码版本、库结构版本（app/cluster/version.py）：库里有活着的、版本不同的进程时新进程不许加入 + code_version: Mapped[str | None] = mapped_column(String(200), default=None) + schema_version: Mapped[str | None] = mapped_column(String(64), default=None) class ClusterLease(Base): diff --git a/backend/app/engine/runner.py b/backend/app/engine/runner.py index 552af462..6599ca5b 100644 --- a/backend/app/engine/runner.py +++ b/backend/app/engine/runner.py @@ -15,7 +15,8 @@ from langgraph.errors import GraphRecursionError from langgraph.types import Command from sqlalchemy import case, func, select, update -from app.cluster import identity +from app.cluster import fence, identity +from app.cluster.fence import Fenced from app.core.bus import bus from app.core.config import settings from app.core.errors import describe_exception, raw_detail @@ -92,7 +93,7 @@ class RunManager: from app.engine.checkpoints import open_saver self._cm = open_saver() - self._checkpointer = await self._cm.__aenter__() + self._checkpointer = _fenced_saver(await self._cm.__aenter__()) # 上次进程被强杀时留下的记录。running 和 queued 要分开处理： # running 至少跑过一步、有 checkpoint，可以续；queued 还在排队， # 一个 checkpoint 都没写过，"恢复"它只会从 START 用默认值重跑一遍， @@ -209,6 +210,8 @@ class RunManager: persist: bool = True, ts: float | None = None, ) -> None: + # 执行这次运行的那段（cluster/fence 绑定过世代的）写之前问一句：租约续晚了就不写（可能已经被别的进程接手） + fence.check() event = self._event(run_id, event_type, node_id=node_id, data=data, ts=ts) # 先落库，再广播。顺序反过来会漏事件，而且是结构性的漏： # @@ -791,6 +794,11 @@ class RunManager: # 等待从 run.interrupted 算起——以前执行一直计到 finally，收尾那一截执行和等待两边都算， # active + wait 比墙钟还长（机器一忙就差出几十毫秒） paused: float | None = None + # 这段执行属于现在这一世代（cluster/fence.py）：本进程隔离之后，它不再写 + fence.bind() + mine = identity.owner() + #: 被隔离了（续不上租约）：停手，不写终态（这次运行由接手的进程按「属主不在了」处理） + fenced = False # 先要并发名额（app/cluster/slots.py）。等名额的时候被停止（或者关停）也要照实收尾，见 _stopped_in_queue slot = slots.hold(run_id, origin_kind) try: @@ -806,7 +814,7 @@ class RunManager: row = await session.get(Run, run_id) if row is not None: row.status = "running" - row.owner = identity.owner() + row.owner = mine # started_at 只在第一次开始时写。以前每段 _drive 都覆盖它， # 审批恢复、接着跑之后，运行的"开始时间"就成了最后一段的开始 if row.started_at is None: @@ -891,7 +899,13 @@ class RunManager: elif snapshot.next: status = "interrupted" + except Fenced: + fenced = True + logger.error("运行 %s：本进程续不上租约，停手、不写终态（由接手的进程按「属主不在了」处理）", run_id) except asyncio.CancelledError: + if fence.fenced_here(): + fenced = True + raise if self._closing: # 服务关停（dev.sh 的 --reload 改一个文件就是一次）不是用户取消。 # checkpoint 完好，记成中断、重启后从断点接着跑——和强杀后启动 @@ -956,11 +970,12 @@ class RunManager: elapsed = int((stop - started) * 1000) self._finalizing.add(run_id) try: - await self._finalize(run_id, status, error, final_state, elapsed, - terminal=terminal, error_node=error_node, end_reason=end_reason) + if not fenced: + await self._finalize(run_id, status, error, final_state, elapsed, terminal=terminal, + error_node=error_node, end_reason=end_reason, owner=mine) # 记 action.unknown、安排核实也算收尾：做完才把自己摘掉。以前摘掉之后才做，等运行「不在跑了」 # 再去读事件的一方（接口、测试）有一小段读不到这一条；库慢一点（PostgreSQL 满载时）就读不到 - if pending_action is not None: + if pending_action is not None and not fenced: from app.process.actions import on_unknown try: @@ -978,7 +993,7 @@ class RunManager: # 不收的话每跑一次带代码节点的图就多一台常驻 microVM（几百 MB）， # 而且没有任何自动回收路径会碰它——闲置回收只在"下次有人执行 # 代码"时才触发。interrupted 要留着，人回来还得接着跑。 - if status in ("succeeded", "failed", "cancelled"): + if status in ("succeeded", "failed", "cancelled") and not fenced: with contextlib.suppress(Exception): from app.sandbox.manager import sandbox_manager @@ -992,11 +1007,13 @@ class RunManager: try: if self._closing: return + if fence.fenced_here(): + return actor = self._stopped_by.pop(run_id, None) self._finalizing.add(run_id) await self._finalize(run_id, "cancelled", "用户取消", {}, 0, terminal=(EventType.RUN_CANCELLED, {"reason": END_STOPPED, "actor": actor}), - end_reason=END_STOPPED) + end_reason=END_STOPPED, owner=identity.owner()) finally: self._finalizing.discard(run_id) self._tasks.pop(run_id, None) @@ -1021,6 +1038,7 @@ class RunManager: interrupts = chunk.get("__interrupt__") if not interrupts: return False + fence.check() for item in interrupts: value = getattr(item, "value", item) @@ -1082,7 +1100,9 @@ class RunManager: terminal: tuple[EventType, dict[str, Any]] | None = None, error_node: str | None = None, end_reason: str | None = None, + owner: str | None = None, ) -> None: + """收尾。owner：执行这一段的属主——运行那一行的属主已经不是它（被别的进程当成没了接手了，cluster/fence.py）就什么都不写。""" usage = dict(state.get("usage") or {}) output = dict(state.get("output") or {}) now = datetime.now(timezone.utc) @@ -1100,7 +1120,13 @@ class RunManager: run_class: str | None = None origin_kind: str | None = None async with SessionLocal() as session: - run = await session.get(Run, run_id) + # 先占住这一行再看属主（PostgreSQL 上 FOR UPDATE；SQLite 只有一个进程）：接手的进程改属主用的是带条件的 UPDATE， + # 两边一个接一个，被接手了的这边不写 + run = (await session.execute(select(Run).where(Run.id == run_id).with_for_update())).scalar_one_or_none() + if run is not None and owner is not None and run.owner != owner: + logger.error("运行 %s 已经被别的进程接手（属主 %s），本进程不写终态", run_id, identity.describe(run.owner)) + await session.rollback() + return if run: app_id = run.app_id run_class, origin_kind = run.run_class, run.origin_kind @@ -1759,6 +1785,31 @@ async def _timing(session: Any, run: Run, elapsed_ms: int, now: datetime) -> dic run_manager = RunManager() +def _fenced_saver(saver: BaseCheckpointSaver) -> BaseCheckpointSaver: + """断点库的写也要过 cluster/fence：本进程续不上租约之后，手上的执行不再写断点（接手的进程会从断点接着跑，两边都写就分叉了）。""" + for name in ("aput", "aput_writes"): + original = getattr(saver, name, None) + if original is None: + continue + + def guarded(*args: Any, __original: Any = original, **kwargs: Any) -> Any: + fence.check() + return __original(*args, **kwargs) + + setattr(saver, name, guarded) + return saver + + +def _stop_fenced() -> None: + """隔离了：停掉本进程在执行的运行（停的时候不写终态，见 _drive 的 fenced）。""" + for task in list(run_manager._tasks.values()): + if not task.done(): + task.cancel() + + +fence.on_trip(_stop_fenced) + + def _register_requests() -> None: """别的进程发来的：停止本进程在执行的运行、往它里面记一条事件（app/cluster/requests.py）。""" from app.cluster import requests diff --git a/backend/app/eval/answer.py b/backend/app/eval/answer.py index 6acc25cf..2842b27e 100644 --- a/backend/app/eval/answer.py +++ b/backend/app/eval/answer.py @@ -24,7 +24,7 @@ app/chat/job.py），生成过程是 api/copilot.answer_ops，几个「决定」 用量按角色记（做法 3、11）：助手 = 生成（answer_ops 末尾的 usage）+ 复核（review_of_run 的 usage）；写作 = 答案运行里的 llm.end（按模型分）。重来的话两次都算——那都是这一题花的钱。 -时限、停止由调用方管（runner.py 用 asyncio.wait_for 包住 drive）：这里随时把正在跑的那次运行记在 current_run，调用方 +时限、停止由调用方管（runner.py 用 within 包住 drive，排队等名额的那一段不算）：这里随时把正在跑的那次运行记在 current_run，调用方 超时、停止时据此取消它（stop_run）。 """ from __future__ import annotations diff --git a/backend/app/eval/gold.py b/backend/app/eval/gold.py index 872aaf32..4b584f1b 100644 --- a/backend/app/eval/gold.py +++ b/backend/app/eval/gold.py @@ -209,11 +209,19 @@ async def _launch_and_wait(graph: dict[str, Any], payload: dict[str, Any], *, na run_id = run.id except LaunchRefused as e: return _Ran("failed", None, f"算期望的运行没有发起：{e.message}", GOLD_FAILED) + from app.cluster import slots + loop = asyncio.get_running_loop() deadline = loop.time() + max(0.0, float(ctx.timeout_s)) + last = loop.time() while True: # 执行任务收完尾（含封存）才回来；任务已经不在了（刚建好还没排上、或者已经结束）就立刻回来，再按库里的状态判 - await run_manager.wait_idle(run_id, timeout=max(0.05, min(5.0, deadline - loop.time()))) + await run_manager.wait_idle(run_id, timeout=max(0.05, min(1.0, deadline - loop.time()))) + now = loop.time() + if await slots.waiting(run_id, here=run_manager.is_active(run_id)): + # 排队等名额的那一段不算时限（复核：满载时排队不该算进每题时限） + deadline += now - last + last = now async with SessionLocal() as session: row = await session.get(Run, run_id) status, error = (row.status, row.error) if row is not None else ("failed", "运行记录不见了") diff --git a/backend/app/eval/runner.py b/backend/app/eval/runner.py index 441430d7..a552a528 100644 --- a/backend/app/eval/runner.py +++ b/backend/app/eval/runner.py @@ -15,7 +15,7 @@ snapshot 视图）：起草、剖析照常进行，也不影响这次评测的 **并发。** 同时最多 max_parallel 题，且所有臂合起来不超过全局运行并发的四分之一。评测发起的作答、算期望的运行归「系统」类 （app/cluster/slots.py）：全站的名额按类别有上限和优先级，评测再多也只占得到系统那一类的名额，问数据、试跑优先，不被挤占。 A/B 的同一题几臂同时开跑。每题时限 per_case_timeout_s：作答一臂超时就取消它正在跑的运行、记超时； -期望探针每次也按这个时限。 +期望探针每次也按这个时限。排队等运行名额的那一段不算时限（within、gold._launch_and_wait）：满载时排队不该被记成超时、触发止损。 **预算**（评测-D9；单独一本账 engine/spend.eval_spend，不挤占裁判、核查的额度）：开跑前按上一次的每题平均花费估算（estimate）。 每题开评前按估算预留——加上会超过这次评测的金额上限、或评测每日上限就不再开新题，剩下的记「已到上限」；评完按实际找补。 @@ -33,6 +33,7 @@ A/B 的同一题几臂同时开跑。每题时限 per_case_timeout_s：作答一 from __future__ import annotations import asyncio +import contextlib import copy import logging from collections.abc import Mapping, Sequence @@ -519,9 +520,9 @@ class EvalExecution(runtime.JobRun): try: if case.get("kind") == "workflow": graph, inputs = await self._workflow_graph(case, values) - await asyncio.wait_for(driver.drive_workflow(graph, inputs), timeout=timeout) + await within(driver.drive_workflow(graph, inputs), timeout, lambda: _queued(driver)) else: - await asyncio.wait_for(driver.drive(), timeout=timeout) + await within(driver.drive(), timeout, lambda: _queued(driver)) except asyncio.TimeoutError: turn = driver.turn turn.timed_out = True @@ -964,6 +965,45 @@ def running(run: Any) -> bool: return manager.running(leader_of(run)) +async def _queued(driver: Any) -> bool: + """这一题的作答运行正排着队等名额（app/cluster/slots）：这一段不算每题的时限。""" + from app.cluster import slots + from app.engine.runner import run_manager + + run_id = getattr(driver.turn, "current_run", None) + return bool(run_id) and await slots.waiting(run_id, here=run_manager.is_active(run_id)) + + +async def within(coro: Any, budget: float, paused: Any, *, tick: float = 0.5) -> Any: + """同 asyncio.wait_for(coro, budget)，只是 paused() 为真的那几段不算时间（复核：满载时排队等名额不该算进每题时限、触发止损）。 + 超时取消 coro、抛 TimeoutError。每隔一小段（tick，且不超过时限的 1/20）看一次；一段的头尾有一头在暂停就整段不算—— + 宁可少算一点，不把排队的时间算进去。""" + loop = asyncio.get_running_loop() + task = asyncio.ensure_future(coro) + step_s = max(0.005, min(tick, budget / 20)) + spent, last = 0.0, loop.time() + was_paused = await paused() + try: + while True: + done, _ = await asyncio.wait({task}, timeout=max(0.005, min(step_s, budget - spent))) + if task in done: + return task.result() + now = loop.time() + step, last = now - last, now + is_paused = await paused() + if not (was_paused or is_paused): + spent += step + was_paused = is_paused + if spent >= budget: + task.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await task + raise asyncio.TimeoutError() + except asyncio.CancelledError: + task.cancel() + raise + + def leader_of_id(rows: Sequence[EvalRun], job_id: str) -> str: """rows 里 id 是 job_id 的那一臂由哪一臂的任务在跑（A/B 几臂共用领头那一臂的后台任务）。""" row = next((r for r in rows if r.id == job_id), None) diff --git a/backend/app/ledger/issue.py b/backend/app/ledger/issue.py index e7c9181d..e2bddd08 100644 --- a/backend/app/ledger/issue.py +++ b/backend/app/ledger/issue.py @@ -329,11 +329,16 @@ _locks: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, dict[tuple[str, st def _lock(key: tuple[str, str, int]) -> Any: - """跨进程的锁（app/cluster/mutex）：多进程时「有没有出具过、今天出了几次」的核对和发起照样一个接一个。全站一把（不按键）： - 每日上限是全站的，按键分锁时几个不同的答案同时出具会一起越过上限；出具不频繁，发起本身只是建一条运行。""" + """同一（运行、工作流、版本）的出具一个接一个（app/cluster/mutex，多进程之间也是一把）：管的是「这一次出具过没有」的去重， + 同时点了几下只发起一次。每日上限不靠它：发起之前先在库里占一份额度（app/cluster/quota），几份不同的出具可以同时发起，上限 + 照样数得准。""" from app.cluster import mutex - return mutex.lock("ledger-issue") + return mutex.lock(f"ledger-issue:{key[0]}:{key[1]}:{key[2]}") + + +#: 每日出具额度在 cluster/quota 里的名字 +QUOTA = "ledger_issue" async def _hider(run_id: str) -> Any: @@ -369,12 +374,23 @@ async def issue(run_id: str, workflow_id: str, version: int, *, principal: Any, "how": plan.how, "binding_hash": input_key(plan.binding), "match": plan.match_artifact, "auto": auto, **({"granted_by_app": plan.granted_by_app} if plan.granted_by_app else {})} actor = run.started_by if auto else principal.username + # 先占一份今天的额度（库里带条件的加一，几个进程同时出具也正好到上限为止），发起不成退回 + from app.api.settings import ledger_settings + from app.cluster import quota + + limit = int((await ledger_settings(session))["issue_daily_limit"]) + if not await quota.take(QUOTA, limit, seed=issued_today): + raise IssueRefused(429, OVER_LIMIT.format(limit=limit), DAILY_LIMIT) try: launched = await launch_run(session, workflow_id=workflow_id, run_class="formal", version=int(version), input_payload=plan.inputs, actor=actor, origin_kind=ORIGIN, origin=origin, expect_hash=plan.expect_hash) except LaunchRefused as e: + await quota.give_back(QUOTA) raise IssueRefused(e.status_code, e.message, e.code or REFUSED) from e + except BaseException: + await quota.give_back(QUOTA) + raise return launched, True diff --git a/backend/app/main.py b/backend/app/main.py index b926852f..5aa95a8e 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -58,6 +58,14 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]: # SQLite 只支持单进程：同一个数据目录已经有一个进程在用（uvicorn --workers、两个服务指着同一个目录）就拒绝启动、照实说 # （app/cluster/guard.py）。PostgreSQL 可以多进程，什么都不做 guard.acquire() + # PgBouncer 事务模式（AGENTLAB_DB_PGBOUNCER=1）要另给直连地址，没给就拒绝启动（db/base.check_pgbouncer） + from app.db.base import check_pgbouncer + + check_pgbouncer() + # 不许新旧版本混跑（app/cluster/version.py）：库里还有活着的、版本不同的进程就不启动——在补列之前看 + from app.cluster import version + + await version.preflight() await init_db() # 多进程协调（app/cluster）：登记这个工作进程、开始心跳和进程间通知、争调度当班。要在下面的启动整理之前：整理只碰属主不在了的 # 行，要先知道别的进程谁还活着 @@ -150,9 +158,10 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]: await cluster.stop() # 应用库的连接全关掉：SQLite 最后一个连接关掉时才把 -wal 并回、删掉（copy-to-postgres 按它认服务停没停）， # PostgreSQL 的连接池也正常关 - from app.db.base import engine as db_engine + from app.db.base import dispose_session_engine, engine as db_engine await db_engine.dispose() + await dispose_session_engine() logger.info("已停止") diff --git a/backend/app/process/actions.py b/backend/app/process/actions.py index 6231d047..f8c8e2b2 100644 --- a/backend/app/process/actions.py +++ b/backend/app/process/actions.py @@ -399,6 +399,10 @@ async def send(prep: Prepared, capability: Any, *, run_id: str, node_id: str, ex else "演练：探索运行里端点没有测试环境，没有发送请求") return _output(action_id, prep, "simulated", None, receipt) + # 本进程续不上租约了（cluster/fence.py）就不发：这次运行可能已经被别的进程接手 + from app.cluster import fence + + fence.check() # 先记 sent 再发：发到一半进程没了，重启后看到的是 sent，只能核实，不会当成没发过再发一次 if not await _set(action_id, ("pending",), state="sent", request_at=utcnow()): raise ActionPending(action_id) @@ -725,7 +729,8 @@ async def startup(*, live: Any = None) -> int: """启动时：上次进程没了的时候已经发出（sent）、还没拿到结果的动作收成结果未知，安排核实；结果未知却还没安排核实的 （挂起之后、安排之前进程没了）补上。返回处理的条数。 - 多进程：所属的运行还在别的活着的进程里执行的不碰（那边正在发、正在等结果）。""" + 多进程：所属的运行还在别的活着的进程里执行的不碰（那边正在发、正在等结果）；当班进程隔一会儿来的那一次只收拾停了两个租约 + 时长以上的。""" from app.cluster import state from app.db.models import Run from app.engine.runner import run_manager @@ -737,6 +742,16 @@ async def startup(*, live: Any = None) -> int: ActionRecord.state.in_(("sent", "unknown")), ActionRecord.resolution.is_(None)))).all() ids = [action_id for action_id, run_id, status, owner in rows if not (status in ("running", "queued") and live.alive(owner, local=run_manager.is_active(run_id)))] + if live.trust_mine and ids: + # 当班进程隔一会儿来的那一次：刚挂起、执行它的进程正在安排核实的（收尾和安排核实之间）不碰，只收拾停了两个租约时长以上的 + from datetime import timedelta + + from app.cluster.workers import lease_s + + floor = utcnow() - timedelta(seconds=2 * lease_s()) + async with SessionLocal() as session: + ids = list((await session.execute(select(ActionRecord.id).where( + ActionRecord.id.in_(ids), ActionRecord.updated_at < floor))).scalars()) for action_id in ids: try: await on_unknown(action_id) diff --git a/backend/app/process/tasks.py b/backend/app/process/tasks.py index e1690f40..a8bbbb4c 100644 --- a/backend/app/process/tasks.py +++ b/backend/app/process/tasks.py @@ -1027,14 +1027,15 @@ async def startup(*, live: Any = None) -> int: """启动时：存量的待审批补建 Task（派给任何人，幂等）；上次停在 resolving 的（判定之后、恢复之前进程没了）按 Approval 的状态收拾——还在等的退回 pending、撤回那一票，已经交给引擎的记成已处理。返回补建的条数。 - 多进程：别的进程活着时，resolving 可能正是它在判定、恢复（几秒的事），只收拾停在 resolving 超过两个租约时长的。""" + 多进程：别的进程活着时，resolving 可能正是它在判定、恢复（几秒的事），只收拾停在 resolving 超过两个租约时长的；当班进程隔一会儿 + 来的那一次（live.trust_mine）不论有没有别的进程都这样——本进程自己也可能正在判定。只有启动时、又没有别的进程的那一次全收拾。""" from datetime import timedelta from app.cluster import state from app.cluster.workers import lease_s live = live if live is not None else await state.current(fresh=True) - stale_before = utcnow() - timedelta(seconds=2 * lease_s()) if live.others else None + stale_before = utcnow() - timedelta(seconds=2 * lease_s()) if (live.others or live.trust_mine) else None async with SessionLocal() as session: orphans = list((await session.execute( select(Approval).outerjoin(Task, Task.approval_id == Approval.id) diff --git a/backend/tests/test_chat_turns.py b/backend/tests/test_chat_turns.py index f5f851b3..77104fe7 100644 --- a/backend/tests/test_chat_turns.py +++ b/backend/tests/test_chat_turns.py @@ -1155,3 +1155,17 @@ async def test_a_downgraded_account_that_can_still_chat_finishes_as_who_they_are row = await _row(turn_id) assert row.status == "done" and row.answer == "华东 300。" assert seen and seen[-1].username == "bob" and seen[-1].role == "user" + + +async def test_the_cleanup_tick_without_a_time_uses_now(monkeypatch): + """调度循环那一拍不一定给 now（process_due(None)）：清理按当前时间算，不报错。""" + from app.scheduler import loop + + old = await _finished_turn(days_ago=turn_log.RETAIN_DAYS + 2) + monkeypatch.setattr(turn_log, "_last_purge", 0.0) + assert old in await turn_log.purge_tick() + assert await _events(old) == [] + again = await _finished_turn(days_ago=turn_log.RETAIN_DAYS + 2) + monkeypatch.setattr(turn_log, "_last_purge", 0.0) + await loop.process_due(None) + assert await _events(again) == [] diff --git a/backend/tests/test_cluster.py b/backend/tests/test_cluster.py index e0f6f34d..89c708cd 100644 --- a/backend/tests/test_cluster.py +++ b/backend/tests/test_cluster.py @@ -14,6 +14,7 @@ import pytest from sqlalchemy import delete, select from app.cluster import identity, leader as leader_mod, slots, state, workers +from app.core.config import settings as app_settings from app.db.base import SessionLocal, new_id, utcnow from app.db.models import Run, RunEvent from app.db.models_cluster import ClusterLease, ClusterRequest, ClusterWorker, LaneTicket, RunSlot @@ -47,9 +48,11 @@ async def _peer(owner: str = PEER, *, alive: bool = True) -> None: state.forget() -def _policy(monkeypatch, total: int, **classes: tuple[int, int]) -> None: - value = slots.values({"max_total": total, "classes": {k: {"limit": lim, "priority": pri} - for k, (lim, pri) in classes.items()}}) +def _policy(monkeypatch, total: int, reserve: int | None = None, **classes: tuple[int, int]) -> None: + stored: dict = {"max_total": total, "classes": {k: {"limit": lim, "priority": pri} for k, (lim, pri) in classes.items()}} + if reserve is not None: + stored["interactive_reserve"] = reserve + value = slots.values(stored) async def fixed(*, fresh: bool = False): # noqa: ARG001 return value @@ -112,11 +115,66 @@ async def test_a_class_at_its_cap_does_not_hold_up_the_others(monkeypatch): await _until(lambda: "chat" in got) assert len({"e1", "e2"} & set(got)) == 1, "系统类到了上限：另一个等着" snapshot = await slots.snapshot() - assert snapshot["classes"]["system"] == {"running": 1, "waiting": 1} + assert snapshot["classes"]["system"] == {"running": 1, "waiting": 1}, snapshot gate.set() await asyncio.wait_for(asyncio.gather(*tasks), timeout=10) +def test_the_defaults_keep_a_reserve_for_interactive_and_lend_the_rest(): + pol = slots.defaults(20) + assert pol["interactive_reserve"] == 5 and pol["classes"]["interactive"]["limit"] == 20 + assert {pol["classes"][k]["limit"] for k in ("app", "schedule", "system")} == {15}, "别的类缺省借满到 全站 − 留底" + assert slots.defaults(2)["interactive_reserve"] == 1 and slots.defaults(1)["interactive_reserve"] == 0, \ + "至少留 1 个；全站只有 1 个名额时不留（别的类才跑得起来）" + assert slots.values({"max_total": 8})["classes"]["system"]["limit"] == 6, "缺省按存着的全站上限现算" + assert slots.problem({"max_total": 4, "interactive_reserve": 4}) is not None + + +async def _holders(names_origins: list[tuple[str, str]], got: list[str], gate: asyncio.Event) -> list[asyncio.Task]: + async def run(name: str, origin: str) -> None: + async with slots.hold(f"run-{name}", origin): + got.append(name) + await gate.wait() + + tasks = [] + for name, origin in names_origins: + tasks.append(asyncio.create_task(run(name, origin))) + await asyncio.sleep(0.05) + return tasks + + +async def test_idle_system_runs_borrow_all_but_the_interactive_reserve(monkeypatch): + """全站 4 个、留底 1 个：没人问数据时评测能同时跑 3 个（全站 − 留底），第 4 个等着。""" + _policy(monkeypatch, 4, reserve=1) + got: list[str] = [] + gate = asyncio.Event() + tasks = await _holders([(f"e{i}", "eval") for i in range(4)], got, gate) + await _until(lambda: len(got) == 3) + await asyncio.sleep(0.3) + assert len(got) == 3, "留底那一个不借给评测" + assert (await slots.snapshot())["classes"]["system"] == {"running": 3, "waiting": 1} + gate.set() + await asyncio.wait_for(asyncio.gather(*tasks), timeout=10) + + +async def test_interactive_takes_the_reserve_at_once_and_it_is_never_lent(monkeypatch): + """评测借满了（3 个在跑），应用类的再来也借不到留底（全站还空着 1 个）；问数据一来马上拿到留底，不用等评测跑完。""" + _policy(monkeypatch, 4, reserve=1) + got: list[str] = [] + gate = asyncio.Event() + tasks = await _holders([("e1", "eval"), ("e2", "eval"), ("e3", "eval")], got, gate) + await _until(lambda: len(got) == 3) + tasks += await _holders([("app", "app_view")], got, gate) + await asyncio.sleep(0.3) + assert "app" not in got, "留底不借给应用类" + tasks += await _holders([("chat", "chat")], got, gate) + await _until(lambda: "chat" in got) + assert "app" not in got and set(got) == {"e1", "e2", "e3", "chat"} + gate.set() + await asyncio.wait_for(asyncio.gather(*tasks), timeout=10) + + +@pytest.mark.skipif(app_settings.db_dialect != "postgresql", reason="SQLite 只有一个进程，名额在进程里（复核 R-b3），没有别的进程的名额行") async def test_slots_held_by_a_process_that_is_gone_are_taken_back(monkeypatch): _policy(monkeypatch, 1, interactive=(1, 40)) async with SessionLocal() as session: @@ -128,6 +186,7 @@ async def test_slots_held_by_a_process_that_is_gone_are_taken_back(monkeypatch): assert rows == ["mine"], "没了的进程占着的名额收回了" +@pytest.mark.skipif(app_settings.db_dialect != "postgresql", reason="SQLite 只有一个进程，名额在进程里（复核 R-b3），没有别的进程的名额行") async def test_slots_of_a_live_peer_are_left_alone(monkeypatch): _policy(monkeypatch, 1, interactive=(1, 40)) await _peer(PEER) @@ -171,6 +230,54 @@ async def test_the_source_lane_is_shared_and_skips_tickets_of_a_gone_process(): await other.release() +async def test_the_source_lane_does_not_look_up_liveness_while_holding_its_lock(monkeypatch): + """复核：持着这个源的咨询锁时不另取连接查「谁还活着」（连接池用尽时会卡到取连接超时）——拿锁之前备好；拿锁之后才出现的 + 别的属主这一次当活着，不删它的行。""" + from app.core import job_runtime as runtime + + import app.data.catalog_jobs # noqa: F401 - 登记 draft、profile + + lanes = runtime.lanes + holding = {"now": False, "looked_inside": 0} + real_locked, real_rows, real_live_for = lanes._locked, lanes._rows, state.live_for + + async def locked(session, source_id): + await real_locked(session, source_id) + holding["now"] = True + + async def rows(session, source_id, dead_owners): + try: + return await real_rows(session, source_id, dead_owners) + finally: + holding["now"] = False + + async def live_for(owners): + if holding["now"]: + holding["looked_inside"] += 1 + return await real_live_for(owners) + + monkeypatch.setattr(lanes, "_locked", locked) + monkeypatch.setattr(lanes, "_rows", rows) + monkeypatch.setattr(state, "live_for", live_for) + async with SessionLocal() as session: + session.add(LaneTicket(source_id="src-lock", kind="profile", job_id="ghost-job", owner=GHOST, granted=True)) + await session.commit() + mine = await lanes.enter("src-lock", "draft", "job-lock") + assert mine.granted.is_set(), "没了的进程占着的位置照旧让出来（拿锁之前就查清楚了）" + assert holding["looked_inside"] == 0 + await mine.release() + + async with SessionLocal() as session: + session.add(LaneTicket(source_id="src-late", kind="profile", job_id="late-job", owner=GHOST, granted=True)) + await session.commit() + kept = await real_rows(session, "src-late", frozenset()) + await session.commit() + assert [r.owner for r in kept] == [GHOST], "拿锁之前那份里没有的属主这一次当活着" + async with SessionLocal() as session: + await session.execute(delete(LaneTicket).where(LaneTicket.source_id.in_(["src-lock", "src-late"]))) + await session.commit() + + async def test_reaping_runs_only_touches_the_ones_whose_owner_is_gone(): from app.engine.runner import run_manager @@ -285,3 +392,223 @@ async def _until(check, timeout: float = 10.0) -> None: if loop.time() > deadline: raise AssertionError("等超时了") await asyncio.sleep(0.05) + + +async def test_a_daily_quota_is_taken_up_to_the_limit_by_concurrent_callers(): + """先占额度再发起（cluster/quota）：10 个协程同时占、上限 5，正好 5 个占到；这一天第一次占按已有的记录起算；退回一份又能占。 + （同时占的个数不超过连接池能直接给的连接数：SQLite 的池子是各测试共用的，排队等连接会把它绑到这个测试的事件循环上。）""" + from app.cluster import quota + + name = f"test-quota-{new_id()[:8]}" + + async def seed(_session) -> int: + return 2 # 今天已经做过 2 次 + + got = await asyncio.gather(*(quota.take(name, 5, seed=seed) for _ in range(10))) + assert sum(got) == 3 and await quota.used(name) == 5 + assert not await quota.take(name, 5) + await quota.give_back(name) + assert await quota.take(name, 5) and not await quota.take(name, 5) + + +def test_a_daily_quota_holds_across_two_processes(): + """两个进程同时各占 15 次、上限 10：合计正好 10 次。""" + name = f"test-quota-{new_id()[:8]}" + backend = Path(__file__).resolve().parents[1] + script = ("import asyncio, sys\n" + "from app.cluster import quota\n" + "async def main():\n" + f" got = await asyncio.gather(*(quota.take({name!r}, 10) for _ in range(15)))\n" + " print(sum(got))\n" + "asyncio.run(main())\n") + env = {**os.environ, "PYTHONPATH": str(backend), "AGENTLAB_DB_POOL_SIZE": "5"} + procs = [subprocess.Popen([sys.executable, "-c", script], cwd=str(backend), env=env, stdout=subprocess.PIPE, + stderr=subprocess.PIPE, text=True) for _ in range(2)] + outs = [p.communicate(timeout=120) for p in procs] + assert all(p.returncode == 0 for p in procs), [o[1][-2000:] for o in outs] + assert sum(int(o[0].strip().splitlines()[-1]) for o in outs) == 10 + + +# -------------------------------------------------------------------------- +# 不许新旧版本混跑（cluster/version.py） +# -------------------------------------------------------------------------- + +async def _peer_version(owner: str, code: str | None, *, lease_s: float = 60.0) -> None: + now = utcnow() + async with SessionLocal() as session: + session.add(ClusterWorker(id=owner, host="peer-host", pid=4242, started_at=now, heartbeat_at=now, + lease_until=now + timedelta(seconds=lease_s), code_version=code, + schema_version=None if code is None else _schema())) + await session.commit() + + +def _schema() -> str: + from app.cluster import version + + return version.schema_version() + + +async def _rows(): + async with SessionLocal() as session: + return list((await session.execute(select(ClusterWorker))).scalars()) + + +async def test_a_live_process_of_another_version_is_named_and_refused(): + from app.cluster import version + + await _peer_version(PEER, "0.0.9+old+aaaaaaaaaaaa") + with pytest.raises(version.VersionMismatch) as e: + version.check_rows(await _rows(), utcnow()) + assert "peer-host 上的进程 4242" in str(e.value) and "0.0.9+old+aaaaaaaaaaaa" in str(e.value) and "先停掉" in str(e.value) + + +async def test_the_same_version_joins_and_expired_or_unversioned_dead_ones_do_not_count(): + from app.cluster import version + + await _peer_version(PEER, version.code_version()) + await _peer_version(GHOST, "0.0.9+old+aaaaaaaaaaaa", lease_s=-5) + version.check_rows(await _rows(), utcnow()) + assert version.same("0.1.0+other-build+" + version.source_fingerprint(), version.schema_version()), \ + "比的是源码指纹：构建号写法不同、代码相同算同一版" + + +@pytest.mark.skipif(app_settings.db_dialect != "postgresql", reason="只在 PostgreSQL 上拦（SQLite 只有一个进程）") +async def test_preflight_refuses_a_beating_peer_and_waits_out_a_stale_one(): + from app.cluster import version + + await _peer_version(PEER, "0.0.9+old+aaaaaaaaaaaa") + + async def beat() -> None: + for _ in range(20): + await asyncio.sleep(0.1) + async with SessionLocal() as session: + row = await session.get(ClusterWorker, PEER) + row.heartbeat_at = utcnow() + await session.commit() + + beating = asyncio.create_task(beat()) + with pytest.raises(version.VersionMismatch): + await version.preflight(observe_s=0.5) + await beating + async with SessionLocal() as session: + await session.execute(delete(ClusterWorker).where(ClusterWorker.id == PEER)) + await session.commit() + # 强杀之后留下的旧版登记（心跳停着、租约还有 1 秒）：等它过期就起 + await _peer_version(GHOST, "0.0.9+old+aaaaaaaaaaaa", lease_s=1.0) + loop = asyncio.get_running_loop() + t0 = loop.time() + await version.preflight(observe_s=0.3) + assert 0.5 < loop.time() - t0 < 5 + with pytest.raises(version.VersionMismatch): + await _peer_version(PEER, "0.0.9+old+aaaaaaaaaaaa") + await workers.register() + + +@pytest.mark.sqlite_only +async def test_an_idle_single_process_writes_nothing_for_coordination(monkeypatch): + """复核 R-b3：SQLite（只有一个进程）空闲时协调不写库——不续租约、不写当班、不周期接手；要名额、放名额也不落库。用另一条连接 + 读 PRAGMA data_version：别的连接提交过写事务它就会变。""" + import sqlite3 + + from app import cluster + from app.core.config import settings + + monkeypatch.setattr(settings, "worker_heartbeat_s", 0.05) + await cluster.start() + try: + await asyncio.sleep(0.2) + probe = sqlite3.connect(settings.db_path) + try: + before = probe.execute("PRAGMA data_version").fetchone()[0] + await asyncio.sleep(0.6) # 十来拍心跳 + async with slots.hold("idle-check", "chat"): # 一段运行要一个名额、用完放掉 + pass + after = probe.execute("PRAGMA data_version").fetchone()[0] + finally: + probe.close() + assert before == after, "空闲的单进程为协调写了库" + assert cluster.leading(), "只有本进程：直接当班" + finally: + await cluster.stop() + + +def test_the_worker_count_comes_from_the_backend_settings(): + """复核 R-b4：prod.sh、Docker 的入口按 python -m app.cluster workers 起几个进程（后端自己读环境变量和 .env）；SQLite 而大于 1 时 + 照实报错、退出码 2。""" + backend = Path(__file__).resolve().parents[1] + base = {k: v for k, v in os.environ.items() if not k.startswith("AGENTLAB_")} + base.update(PYTHONPATH=str(backend), AGENTLAB_DATA_DIR=os.environ["AGENTLAB_DATA_DIR"], + AGENTLAB_SECRET_KEY="test-only-not-a-real-key") + + def ask(**env: str) -> subprocess.CompletedProcess: + return subprocess.run([sys.executable, "-m", "app.cluster", "workers"], cwd=str(backend), env={**base, **env}, + capture_output=True, text=True, timeout=60) + + sqlite = ask(AGENTLAB_DATABASE_URL="", AGENTLAB_WORKERS="3") + assert sqlite.returncode == 2 and "只支持单进程" in sqlite.stderr, sqlite.stderr[-1000:] + 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 diff --git a/backend/tests/test_cluster_fence.py b/backend/tests/test_cluster_fence.py new file mode 100644 index 00000000..f963e9c6 --- /dev/null +++ b/backend/tests/test_cluster_fence.py @@ -0,0 +1,161 @@ +"""自我隔离（cluster/fence.py，复核 R-b1）：本进程卡住了、续不上租约，别的进程已经把它在跑的运行接手了——它缓过来之后不再写： +不写事件（不和接手的一边撞号）、不写断点、不写终态（不把接手的结果盖掉）。在一个进程里装出来：把「上次续上租约」拨到一个租约 +时长以前，再把运行改成接手之后的样子。两种库都跑。""" +from __future__ import annotations + +import asyncio +import time + +import pytest +from sqlalchemy import func, select +from sqlalchemy.exc import IntegrityError + +from app.cluster import fence +from app.cluster.workers import lease_s +from app.db.base import SessionLocal, new_id +from app.db.models import Run, RunEvent +from app.engine import compiler +from app.engine.runner import run_manager +from app.engine.schema import NodeType + + +@pytest.fixture(autouse=True) +async def engine_up(): + await run_manager.setup() + fence.activate() + yield + fence.deactivate() + await run_manager.shutdown() + + +def node(nid, ntype, **config): + return {"id": nid, "type": ntype, "position": {"x": 0, "y": 0}, "data": {"label": nid, "config": config}} + + +GRAPH = {"nodes": [node("start", "input", fields=[{"name": "question"}]), + node("slow", "transform", mode="template", template="慢的这步"), + node("next", "transform", mode="template", template="后面这步"), + node("out", "output", fields=[{"name": "r", "value": "{{ nodes.next }}"}])], + "edges": [{"source": "start", "target": "slow"}, {"source": "slow", "target": "next"}, + {"source": "next", "target": "out"}]} + + +@pytest.fixture +def slow(monkeypatch): + """slow 那一步：第一次跑时卡住，等测试放行（好在它卡着的时候装出「续不上租约、被接手」）；之后立刻返回。""" + calls = {"n": 0} + gate = asyncio.Event() + entered = asyncio.Event() + original = compiler.RUNNERS[NodeType.TRANSFORM] + + async def runner(state, ctx): + if ctx.node.id == "slow": + calls["n"] += 1 + if calls["n"] == 1: + entered.set() + await gate.wait() + return await original(state, ctx) + + monkeypatch.setitem(compiler.RUNNERS, NodeType.TRANSFORM, runner) + return {"calls": calls, "gate": gate, "entered": entered} + + +async def _taken_over(run_id: str) -> None: + """别的进程看它的租约过期了、接手：标中断、属主清空（同 run_manager.reap 做的）。""" + from sqlalchemy import update + + async with SessionLocal() as session: + await session.execute(update(Run).where(Run.id == run_id).values( + status="interrupted", owner=None, error="服务重启，运行已挂起，可从断点继续运行")) + await session.commit() + + +async def _events(run_id: str) -> list[RunEvent]: + async with SessionLocal() as session: + return list((await session.execute(select(RunEvent).where(RunEvent.run_id == run_id) + .order_by(RunEvent.seq))).scalars()) + + +async def _no_repeated_seq(run_id: str) -> None: + async with SessionLocal() as session: + rows = (await session.execute(select(RunEvent.seq, func.count()).where(RunEvent.run_id == run_id) + .group_by(RunEvent.seq).having(func.count() > 1))).all() + assert rows == [], f"撞号：{rows}" + + +async def _checkpointed_before_slow(run_id: str) -> None: + """等 slow 前面那一步的断点落了库（LangGraph 存断点和下一步并行：慢一点的库上 slow 已经开始、断点还在写）。""" + config = {"configurable": {"thread_id": run_id}} + for _ in range(300): + tup = await run_manager.checkpointer.aget_tuple(config) + if tup is not None and int((tup.metadata or {}).get("step", -1)) >= 1: + return + await asyncio.sleep(0.02) + raise AssertionError("slow 前面那一步的断点一直没落库") + + +async def _stall_and_take_over(slow) -> str: + run = await run_manager.start(graph=GRAPH, input_payload={"question": "q"}) + await asyncio.wait_for(slow["entered"].wait(), timeout=15) + await _checkpointed_before_slow(run.id) + before = fence.generation() + fence.renewed(time.monotonic() - lease_s() - 1) # 卡了一个多租约时长 + await _taken_over(run.id) + slow["gate"].set() # 缓过来，手上的执行接着往下走 + for _ in range(200): + if not run_manager.is_active(run.id): + break + await asyncio.sleep(0.05) + assert not run_manager.is_active(run.id), "执行停了" + assert fence.generation() == before + 1, "发现续晚了：隔离了一次" + return run.id + + +async def test_a_stalled_executor_does_not_overwrite_the_takeover(slow): + run_id = await _stall_and_take_over(slow) + async with SessionLocal() as session: + row = await session.get(Run, run_id) + assert row.status == "interrupted" and row.owner is None, "接手之后的样子没被盖掉" + events = await _events(run_id) + assert not [e for e in events if e.type in ("run.finished", "run.failed", "run.cancelled")] + assert not [e for e in events if e.node_id == "next"], "缓过来之后一条都没写" + await _no_repeated_seq(run_id) + + +async def test_a_run_continued_elsewhere_finishes_once_and_the_stalled_side_stays_quiet(slow): + run_id = await _stall_and_take_over(slow) + fence.renewed() # 换了身份重新登记之后，这一世代照常 + await run_manager.continue_failed(run_id) # 接手的那一边接着跑 + for _ in range(300): + async with SessionLocal() as session: + row = await session.get(Run, run_id) + if row.status in ("succeeded", "failed") and row.manifest_hash: + break + await asyncio.sleep(0.05) + assert row.status == "succeeded", row.error + events = await _events(run_id) + assert len([e for e in events if e.type == "run.finished"]) == 1 + assert len([e for e in events if e.type == "node.started" and e.node_id == "next"]) == 1, "后面那一步只跑了一次" + await _no_repeated_seq(run_id) + + +async def test_the_final_write_checks_the_owner(): + async with SessionLocal() as session: + run = Run(status="running", owner="peer-host:1:abcd", graph={}, input={}) + session.add(run) + await session.commit() + await run_manager._finalize(run.id, "succeeded", None, {}, 10, owner="mine:1:ffff") + async with SessionLocal() as session: + row = await session.get(Run, run.id) + assert row.status == "running" and row.finished_at is None, "属主已经不是自己：不写终态" + + +async def test_a_repeated_seq_is_refused_by_the_database(): + async with SessionLocal() as session: + run = Run(status="running", graph={}, input={}) + session.add(run) + await session.flush() + session.add_all([RunEvent(id=new_id(), run_id=run.id, seq=1, type="log", ts=0, data={}), + RunEvent(id=new_id(), run_id=run.id, seq=1, type="log", ts=0, data={})]) + with pytest.raises(IntegrityError): + await session.commit() diff --git a/backend/tests/test_db_direct.py b/backend/tests/test_db_direct.py new file mode 100644 index 00000000..552a0c40 --- /dev/null +++ b/backend/tests/test_db_direct.py @@ -0,0 +1,69 @@ +"""直连地址和 PgBouncer 事务模式（AGENTLAB_DATABASE_DIRECT_URL、AGENTLAB_DB_PGBOUNCER）：地址怎么解析、哪些连接走直连、事务模式下 +关掉了哪些东西。只建引擎、不连库，两种库都跑。""" +from __future__ import annotations + +import pytest +from sqlalchemy.pool import NullPool + +from app.core.config import Settings, settings + +APP = "postgres://agentlab:secret@pgbouncer.internal:6432/agentlab" +DIRECT = "postgresql://agentlab:secret@db.internal:5432/agentlab" + + +def test_the_direct_url_is_normalized_and_used_for_session_bound_connections(): + s = Settings(database_url=APP, database_direct_url=DIRECT, _env_file=None) + assert s.database_url.startswith("postgresql+asyncpg://") and "pgbouncer.internal:6432" in s.database_url + assert s.database_direct_url.startswith("postgresql+asyncpg://") and "db.internal:5432" in s.database_direct_url + assert s.session_db_url == s.database_direct_url + assert "db.internal" in s.checkpoint_conninfo and s.checkpoint_conninfo.startswith("postgresql://"), "断点库缺省走直连" + plain = Settings(database_url=APP, _env_file=None) + assert plain.session_db_url == plain.database_url and "pgbouncer.internal" in plain.checkpoint_conninfo, "不设就全用同一个" + with pytest.raises(ValueError, match="AGENTLAB_DATABASE_DIRECT_URL"): + Settings(database_url=APP, database_direct_url="mysql://x@y/z", _env_file=None) + + +def test_transaction_mode_turns_off_client_pooling_and_statement_caches(): + from app.db.base import pg_engine_kwargs + + kw = pg_engine_kwargs(pgbouncer=True) + assert kw["poolclass"] is NullPool + args = kw["connect_args"] + assert args["statement_cache_size"] == 0 and args["prepared_statement_cache_size"] == 0 + names = {args["prepared_statement_name_func"]() for _ in range(3)} + assert len(names) == 3, "预备语句的名字每次都不一样，不和别的客户端撞名" + direct = pg_engine_kwargs(pgbouncer=False, pool_size=5) + assert direct["pool_size"] == 5 and "statement_cache_size" not in direct["connect_args"] + + +def test_listen_and_cluster_locks_go_direct(monkeypatch): + from app.cluster.signal import listen_dsn + from app.db import base + + monkeypatch.setattr(settings, "database_url", Settings(database_url=APP, _env_file=None).database_url) + monkeypatch.setattr(settings, "database_direct_url", Settings(database_url=APP, database_direct_url=DIRECT, + _env_file=None).database_direct_url) + monkeypatch.setattr(base, "_session_engine", None) + assert "db.internal:5432" in listen_dsn() and listen_dsn().startswith("postgresql://") + made = base.session_engine() + try: + assert made is not base.engine and made.url.host == "db.internal" and made.url.port == 5432 + assert base.session_engine() is made, "直连那个池子只建一个" + finally: + monkeypatch.setattr(base, "_session_engine", None) + made.sync_engine.dispose() + monkeypatch.setattr(settings, "database_direct_url", "") + assert base.session_engine() is base.engine and "pgbouncer.internal" in listen_dsn() + + +def test_transaction_mode_without_a_direct_url_refuses_to_start(monkeypatch): + from app.db import base + + monkeypatch.setattr(settings, "database_url", Settings(database_url=APP, _env_file=None).database_url) + monkeypatch.setattr(settings, "db_pgbouncer", True) + monkeypatch.setattr(settings, "database_direct_url", "") + with pytest.raises(RuntimeError, match="AGENTLAB_DATABASE_DIRECT_URL"): + base.check_pgbouncer() + monkeypatch.setattr(settings, "database_direct_url", Settings(database_url=APP, database_direct_url=DIRECT, + _env_file=None).database_direct_url) + base.check_pgbouncer() diff --git a/backend/tests/test_eval_runner.py b/backend/tests/test_eval_runner.py index a9685e62..6aaaa484 100644 --- a/backend/tests/test_eval_runner.py +++ b/backend/tests/test_eval_runner.py @@ -289,6 +289,87 @@ async def test_a_case_over_its_time_limit_stops_its_run(monkeypatch): assert item.status == "ok" and item.path == "build_failed" and "每题时限" in item.error +async def test_time_queued_for_a_run_slot_does_not_count_against_the_case_limit(monkeypatch): + """复核：满载时作答运行排队等名额的那一段不算每题时限（排了 2 秒、时限 1 秒，拿到名额后很快答完：不记超时）。""" + from app.cluster import slots + + stored = slots.values({"max_total": 2, "interactive_reserve": 1}) + + async def fixed(*, fresh: bool = False): # noqa: ARG001 + return stored + + monkeypatch.setattr(slots, "policy", fixed) + slots.forget_policy() + gate = asyncio.Event() + + async def blocker() -> None: + async with slots.hold("eval-slot-blocker", "eval"): + await gate.wait() + + async def drive(self): + self.turn.current_run = "eval-answer-in-queue" + await self._phase("build") + async with slots.hold(self.turn.current_run, "eval"): + await asyncio.sleep(0.05) + self.turn.path = "asked" + self.turn.usage["assistant"] = {"by_model": {}, "input_tokens": 0, "output_tokens": 0, "calls": 1, + "cost_usd": 0.0, "unpriced_calls": 0} + return self.turn + + monkeypatch.setattr(answer_mod.Driver, "drive", drive) + original = runner.EvalExecution._prepare + + async def prepare(self): + units = await original(self) + self.budget["per_case_timeout_s"] = 1.0 + return units + + monkeypatch.setattr(runner.EvalExecution, "_prepare", prepare) + held = asyncio.create_task(blocker()) + try: + while not (await slots.snapshot())["classes"]["system"]["running"]: + await asyncio.sleep(0.02) + suite_id = await es.make_suite("none", cases=ask_cases(1)) + [run] = await start(suite_id) + while not await slots.waiting("eval-answer-in-queue"): + await asyncio.sleep(0.02) + await asyncio.sleep(2.0) + gate.set() + await es.wait_eval(run.id) + finally: + gate.set() + await held + [item] = await results(run.id) + assert item.status == "ok" and item.path == "asked" and not item.error, item.error + + +async def test_within_does_not_count_the_paused_stretches(): + paused = {"now": True} + + async def is_paused() -> bool: + return paused["now"] + + async def work() -> str: + await asyncio.sleep(0.5) + paused["now"] = False + await asyncio.sleep(0.05) + return "done" + + assert await runner.within(work(), 0.2, is_paused, tick=0.02) == "done", "暂停的那 0.5 秒不算" + + async def never(): + await asyncio.sleep(10) + + async def awake() -> bool: + return False + + loop = asyncio.get_running_loop() + t0 = loop.time() + with pytest.raises(asyncio.TimeoutError): + await runner.within(never(), 0.2, awake, tick=0.02) + assert loop.time() - t0 < 1.0 + + async def test_a_run_stopped_at_an_approval_is_cancelled_not_approved(client, monkeypatch, tmp_path): from app.eval.answer import APPROVAL_REASON diff --git a/backend/tests/test_ledger_issue.py b/backend/tests/test_ledger_issue.py index 19803753..00c7f4df 100644 --- a/backend/tests/test_ledger_issue.py +++ b/backend/tests/test_ledger_issue.py @@ -297,6 +297,28 @@ async def test_users_issue_only_versions_pinned_by_apps_they_can_use(monkeypatch assert created and formal.started_by == "carol" and formal.origin["granted_by_app"] == board_id +async def test_different_answers_issued_at_once_stop_exactly_at_the_daily_limit(monkeypatch): + """几份不同的答案同时出具（不再一把全站锁一个接一个）：先在库里占今天的额度再发起，正好到上限为止，多出来的回 429。""" + wf = await S.publish(level="published") + rows = [await equivalent_run(monkeypatch) for _ in range(3)] + from app.cluster import quota + + async with SessionLocal() as session: + issued = await ledger_issue.issued_today(session) + # 今天的额度在库里记着（同一个测试进程里前面的测试也出具过）；还没记过的按今天已出具的起算 + already = await quota.used(ledger_issue.QUOTA) + already = issued if already is None else already + await set_ledger(issue_daily_limit=max(already, issued) + 2) + results = await asyncio.gather(*(ledger_issue.issue(r.id, wf, 1, principal=ADMIN) for r in rows), + return_exceptions=True) + made = [r for r in results if not isinstance(r, BaseException)] + refused = [r for r in results if isinstance(r, ledger_issue.IssueRefused)] + assert len(made) == 2 and all(created for _, created in made), results + assert len(refused) == 1 and refused[0].status_code == 429 and refused[0].code == ledger_issue.DAILY_LIMIT + async with SessionLocal() as session: + assert await ledger_issue.issued_today(session) == issued + 2 + + async def test_concurrent_requests_issue_once(monkeypatch): wf = await S.publish(level="published") row = await equivalent_run(monkeypatch) diff --git a/backend/tests/test_multiprocess_e2e.py b/backend/tests/test_multiprocess_e2e.py index 09469f96..7181b10e 100644 --- a/backend/tests/test_multiprocess_e2e.py +++ b/backend/tests/test_multiprocess_e2e.py @@ -6,11 +6,14 @@ 进程只执行一次；两边的事件流都按 seq 递增、不重不漏，封存核对得上； 3. 在 A 上跑的运行，从 B 停止：照本进程的停止办（run.cancelled 带上是谁停的）； 4. 定时任务到点只触发一次（两个进程都起着调度循环，只有当班的做）； -5. 评测和问数据并发时优先级生效：全站只给 1 个名额，评测的作答运行（系统类，由本测试进程作为第三个登记的工作进程以评测的 - 来源发起——工作进程里发起评测要真的去调模型建图）排着队时，后来的问数据运行（交互类，经 A 的接口发起）先拿到名额； +5. 评测和问数据并发时优先级生效：全站只给 1 个名额（没有留底），评测的作答运行（系统类，由本测试进程作为第三个登记的工作进程 + 以评测的来源发起——工作进程里发起评测要真的去调模型建图）排着队时，后来的问数据运行（交互类，经 A 的接口发起）先拿到名额； + 全站 2 个、给交互留底 1 个：评测借满了别的名额、还有题排着，问数据一来马上拿到留底，不等评测跑完； 6. 每日花费上限跨进程不越过：两个进程同时往同一本账上记，记下的合计不超过上限； 7. 强杀 B：租约过期后，B 上在跑的运行被照实标中断（可从断点继续），停在审批上的运行照常能在 A 上批、接着跑完；B 当着调度班的 - 话 A 接班。 + 话 A 接班； +8. 卡住但没死（复核 R-b1）：A 在跑一次运行时被 SIGSTOP，B 等它的租约过期两倍时长后接手、标中断；A 缓过来（SIGCONT）之后发现续晚了， + 停手、不写终态、不写事件，换一个身份重新登记——不管 B 有没有在这期间把那次运行接着跑完，都不双写、不盖掉终态、序号不重。 缺省不跑：要 AGENTLAB_MULTIPROCESS_E2E=1 和 PostgreSQL 测试库（scripts/test-multiprocess.sh 起临时容器、设好这两样）。 """ @@ -87,8 +90,16 @@ class Worker: self.proc.send_signal(signal.SIGKILL) self.proc.wait(timeout=10) + def pause(self) -> None: + self.proc.send_signal(signal.SIGSTOP) + + def resume(self) -> None: + if self.proc.poll() is None: + self.proc.send_signal(signal.SIGCONT) + def stop(self) -> None: if self.proc.poll() is None: + self.proc.send_signal(signal.SIGCONT) self.proc.send_signal(signal.SIGTERM) try: self.proc.wait(timeout=15) @@ -97,7 +108,7 @@ class Worker: def tail(self) -> str: try: - return self.log.read_text(encoding="utf-8", errors="replace")[-4000:] + return self.log.read_text(encoding="utf-8", errors="replace")[-12000:] except OSError: return "" @@ -404,6 +415,55 @@ async def test_5_chat_goes_before_queued_evaluation_runs(workers, monkeypatch): await client.put(f"{a.base}/api/settings", json={"values": {"concurrency": slots.defaults()}}) +async def test_5b_the_interactive_reserve_is_there_when_evaluations_fill_the_rest(workers, monkeypatch): + """全站 2 个名额、留底 1 个：评测（系统类）只借得到 1 个，第二题排着；问数据一来马上拿到留底，评测的第一题还在跑。""" + from app import cluster + from app.cluster import slots + from app.cluster.relay import relay + from app.core.config import settings + 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": 2, "interactive_reserve": 1, "classes": {}}}}) + assert put.status_code == 200, put.text + got = put.json()["concurrency"] + assert got["max_total"] - got["interactive_reserve"] == 1, "别的类最多借得到 全站 − 留底 = 1 个" + 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(3), input_payload={"question": f"评测 {i}"}, + origin_kind="eval", origin={"kind": "eval"}) for i in range(2)] + + 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="第一题的作答运行拿到名额") + 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 + + e1, e2, chat = await _until(all_done, timeout=60, what="都跑完") + assert chat.started_at < e1.finished_at, "问数据拿到的是留底：评测的第一题还没跑完它就开始了" + assert e2.started_at >= e1.finished_at, "评测的第二题借不到留底，等第一题跑完" + 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 = ( @@ -431,6 +491,88 @@ async def _spent() -> float: return float((await Ledger("e2e_spend").daily())["usd"]) +async def _stall_a_during_a_run(client, a: Worker, b: Worker, seconds: float) -> str: + """A 上发起一次运行，跑起来之后 SIGSTOP A，等 B 接手（标中断）。交回运行 id（A 还停着，调用方负责 SIGCONT）。""" + run_id = (await client.post(f"{a.base}/api/runs", json={"graph": slow_graph(seconds), "input": {"question": "q"}, + "origin": "canvas"})).json()["id"] + await _until(lambda: _owner_while_active(run_id), what="A 认领了运行") + await asyncio.sleep(0.5) + a.pause() + + async def taken_over() -> Any: + row = await _run(client, b, run_id) + return row if row["status"] == "interrupted" else None + + await _until(taken_over, timeout=LEASE_S * 6 + 10, what="B 接手了卡住的 A 上的运行") + return run_id + + +async def _a_reidentified(client, a: Worker) -> None: + """A 缓过来之后换了身份重新登记：fixture 里记的 owner 跟着换。""" + old = a.owner + + async def fresh() -> Any: + try: + me = (await client.get(f"{a.base}/api/system")).json()["cluster"]["me"] + except Exception: # noqa: BLE001 + return None + return me if me != old else None + + a.owner = await _until(fresh, timeout=15, what="A 换身份重新登记") + + +async def _settled_events(client, worker: Worker, run_id: str) -> list[dict[str, Any]]: + return (await client.get(f"{worker.base}/api/runs/{run_id}/events")).json() + + +def _seqs_unique(events: list[dict[str, Any]]) -> None: + seqs = [e["seq"] for e in events] + assert len(seqs) == len(set(seqs)), f"撞号：{sorted(seqs)}" + + +async def test_6b_a_stalled_worker_does_not_write_after_it_was_taken_over(workers): + a, b = workers["A"], workers["B"] + async with await _client() as client: + try: + run_id = await _stall_a_during_a_run(client, a, b, 6) + finally: + a.resume() + await _a_reidentified(client, a) + await asyncio.sleep(max(4.0, 3 * HEARTBEAT_S)) # A 手上那一步本来再过几秒就跑完：给它机会（它不该写） + row = await _run(client, b, run_id) + assert row["status"] == "interrupted", f"接手之后的状态被盖掉了：{row['status']}" + events = await _settled_events(client, b, run_id) + assert not [e for e in events if e["type"] in ("run.finished", "run.failed", "run.cancelled")] + assert any(e["type"] == "log" and e["data"].get("code") == "server_shutdown" for e in events) + _seqs_unique(events) + + +async def test_6c_a_run_continued_by_the_other_worker_is_not_written_twice(workers): + a, b = workers["A"], workers["B"] + async with await _client() as client: + try: + run_id = await _stall_a_during_a_run(client, a, b, 3) + cont = await client.post(f"{b.base}/api/runs/{run_id}/continue", json={}) + assert cont.status_code == 200, cont.text + finally: + a.resume() + await _a_reidentified(client, a) + + async def done() -> Any: + row = await _run(client, b, run_id) + return row if row["status"] in ("succeeded", "failed") and row.get("manifest_hash") else None + + row = await _until(done, timeout=60, what="B 接着跑完") + await asyncio.sleep(2.0) + assert row["status"] == "succeeded", row + assert (await _run(client, b, run_id))["status"] == "succeeded", "A 缓过来之后没有把终态改掉" + events = await _settled_events(client, b, run_id) + assert len([e for e in events if e["type"] == "run.finished"]) == 1 + _seqs_unique(events) + verified = (await client.get(f"{b.base}/api/runs/{run_id}/verify")).json() + assert verified["ok"] is True, verified + + async def test_7_kill_a_worker(workers): a, b = workers["A"], workers["B"] async with await _client() as client: diff --git a/backend/tests/test_task_resume_errors.py b/backend/tests/test_task_resume_errors.py index db9b885e..721a7fd2 100644 --- a/backend/tests/test_task_resume_errors.py +++ b/backend/tests/test_task_resume_errors.py @@ -84,3 +84,32 @@ async def test_any_resume_error_puts_the_task_back(monkeypatch): out = await tasks.respond(task.id, ANN, decision="approve") assert out["resolved"] is True assert (await _wait(run.id, ("succeeded", "failed"))).status == "succeeded" + + +async def test_the_periodic_takeover_leaves_a_task_that_is_being_resolved_alone(monkeypatch): + """复核 R-b2：当班进程隔一会儿的那一次接手（live.trust_mine）撞上一次正在判定、恢复的待办（停在 resolving 的那几秒）， + 单进程时也不能把它退回 pending、撤回那一票——只收拾停了两个租约时长以上的。启动时、又没有别的进程的那一次照旧全收拾。""" + from app.cluster.workers import Liveness + + run = await run_manager.start(graph=FLOW, input_payload={"q": "x"}) + await _wait(run.id, ("interrupted",)) + task = await _pending(run.id) + real = run_manager.resume + reaped = asyncio.Event() + + async def slow(*a, **k): + await asyncio.sleep(0.3) + await tasks.startup(live=Liveness().trusting()) # 判定到一半，当班进程来了一次 + reaped.set() + return await real(*a, **k) + + monkeypatch.setattr(run_manager, "resume", slow) + out = await tasks.respond(task.id, ANN, decision="approve") + assert reaped.is_set() and out["resolved"] is True + assert (await _wait(run.id, ("succeeded", "failed"))).status == "succeeded" + async with SessionLocal() as session: + row = await session.get(Task, task.id) + mine = (await session.execute(select(TaskAssignment).where(TaskAssignment.task_id == task.id))).scalars().all() + assert row.status == "resolved", row.status + assert mine and all(a.responded_by == "ann" and a.status not in ("pending", "withdrawn") for a in mine), \ + [(a.status, a.responded_by) for a in mine] diff --git a/frontend/src/lib/termsConcurrency.ts b/frontend/src/lib/termsConcurrency.ts index 335f1b8c..0b192c6c 100644 --- a/frontend/src/lib/termsConcurrency.ts +++ b/frontend/src/lib/termsConcurrency.ts @@ -6,6 +6,10 @@ export const CONCURRENCY_TEXT = { title: '运行并发', hint: '全站同时执行的运行按来源分四类，各有上限和优先级：名额满了时先给优先级高的一类，已经在执行的不打断。多进程部署时几个进程合起来算', maxTotal: '全站同时运行的上限', + reserve: '给交互留底', + reserveHint: '全站名额里永远留给交互类的个数，应用、定时、系统用不到这一部分：有人问数据、试运行时总有名额。至多全站上限减 1', + borrow: (n: number) => `应用、定时、系统合起来最多借用 ${n} 个名额（全站上限减留底）：没人问数据时评测可以同时多开几题`, + limitHint: '各类的上限可以调小；应用、定时、系统再大也借不到留底那一部分', limit: '上限', priority: '优先级', priorityHint: '数字大的先拿到空出来的名额（0 到 100）', diff --git a/frontend/src/pages/ConcurrencySettings.tsx b/frontend/src/pages/ConcurrencySettings.tsx index fbcfc18a..7bcc18f7 100644 --- a/frontend/src/pages/ConcurrencySettings.tsx +++ b/frontend/src/pages/ConcurrencySettings.tsx @@ -8,17 +8,19 @@ import { atLeast, usePrincipal } from '../store/auth' const T = CONCURRENCY_TEXT type Klass = (typeof CONCURRENCY_CLASSES)[number] -interface Draft { total: string; limits: Record; priorities: Record } +interface Draft { total: string; reserve: string; limits: Record; priorities: Record } type Now = Partial> const draftOf = (group: Record | undefined): Draft => ({ total: group?.max_total == null ? '' : String(group.max_total), + reserve: group?.interactive_reserve == null ? '' : String(group.interactive_reserve), limits: Object.fromEntries(CONCURRENCY_CLASSES.map((k) => [k, String(group?.classes?.[k]?.limit ?? '')])) as Record, priorities: Object.fromEntries(CONCURRENCY_CLASSES.map((k) => [k, String(group?.classes?.[k]?.priority ?? '')])) as Record, }) const changedCount = (a: Draft, b: Draft): number => (a.total.trim() !== b.total ? 1 : 0) + + (a.reserve.trim() !== b.reserve ? 1 : 0) + CONCURRENCY_CLASSES.filter((k) => a.limits[k].trim() !== b.limits[k]).length + CONCURRENCY_CLASSES.filter((k) => a.priorities[k].trim() !== b.priorities[k]).length @@ -62,6 +64,7 @@ export function ConcurrencySettingsSection() { try { const value = { max_total: Number(cur.total.trim()), + interactive_reserve: Number(cur.reserve.trim()), classes: Object.fromEntries(CONCURRENCY_CLASSES.map((k) => [k, { limit: Number(cur.limits[k].trim()), priority: Number(cur.priorities[k].trim()), }])), @@ -80,6 +83,8 @@ export function ConcurrencySettingsSection() { } if (loadError || !cur) return null + // 借用的上限按存着的算（服务端生效的那份），改到一半的输入不跟着跳 + const borrow = saved && saved.total !== '' && saved.reserve !== '' ? Math.max(0, Number(saved.total) - Number(saved.reserve)) : null const edit = (patch: Partial) => { setError(null) setDraft({ ...cur, ...patch, limits: { ...cur.limits, ...(patch.limits ?? {}) }, @@ -89,12 +94,21 @@ export function ConcurrencySettingsSection() {
- - {(p) => ( - edit({ total: e.target.value })} data-concurrency-setting="max_total" /> - )} - +
+ + {(p) => ( + edit({ total: e.target.value })} data-concurrency-setting="max_total" /> + )} + + + {(p) => ( + edit({ reserve: e.target.value })} data-concurrency-setting="interactive_reserve" /> + )} + +
+ {borrow != null &&

{T.borrow(borrow)}

}
{CONCURRENCY_CLASSES.map((k) => { const info = T.classes[k] @@ -106,7 +120,7 @@ export function ConcurrencySettingsSection() {
{info.hint}
{live &&
{T.now(live.running ?? 0, live.waiting ?? 0)}
}
- + {(p) => ( edit({ limits: { [k]: e.target.value } as Record })} diff --git a/scripts/prod.sh b/scripts/prod.sh index fe688b39..c3e9be81 100755 --- a/scripts/prod.sh +++ b/scripts/prod.sh @@ -26,6 +26,7 @@ USAGE HOST="127.0.0.1" PORT="${AGENTLAB_PORT:-8000}" WORKERS="${AGENTLAB_WORKERS:-1}" +WORKERS_CLI="" BUILD="auto" while [ $# -gt 0 ]; do case "$1" in @@ -36,8 +37,8 @@ while [ $# -gt 0 ]; do --port=*) PORT="${1#--port=}" ;; --workers) [ $# -ge 2 ] || { echo "--workers 后面要跟进程数"; echo; usage; exit 2; } - WORKERS="$2"; shift ;; - --workers=*) WORKERS="${1#--workers=}" ;; + WORKERS="$2"; WORKERS_CLI=1; shift ;; + --workers=*) WORKERS="${1#--workers=}"; WORKERS_CLI=1 ;; --build) BUILD="always" ;; --no-build) BUILD="never" ;; -h|--help) usage; exit 0 ;; @@ -57,12 +58,7 @@ esac if [ "$WORKERS" -lt 1 ]; then echo "进程数至少是 1：$WORKERS"; exit 2 fi -# SQLite 只支持单进程（服务自己也会拒绝启动第二个进程，app/cluster/guard.py）：在这里先说清楚，免得 uvicorn 一遍遍重起被拒的进程 -if [ "$WORKERS" -gt 1 ] && [ -z "${AGENTLAB_DATABASE_URL:-}" ]; then - echo "应用库是 SQLite 时只支持单进程：要 --workers $WORKERS，先设 AGENTLAB_DATABASE_URL 换成 PostgreSQL（README「用 PostgreSQL」「多进程」）" - exit 2 -fi -export AGENTLAB_WORKERS="$WORKERS" +# 命令行给了 --workers 才写进环境变量（见下面问后端的那一句）；没给的话照后端自己的配置（环境变量和 .env） ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" DIST="$ROOT/frontend/dist" @@ -169,6 +165,10 @@ fi echo "==> 启动 $URLS" cd "$ROOT/backend" +# 起几个进程问后端（python -m app.cluster workers）：它照环境变量和 .env 一起读应用库、进程数——写在 .env 里的 AGENTLAB_DATABASE_URL、 +# AGENTLAB_WORKERS 这里也算数；应用库是 SQLite 而进程数大于 1 时它照实报错（SQLite 只支持单进程），这里不起 +[ -n "$WORKERS_CLI" ] && export AGENTLAB_WORKERS="$WORKERS" +WORKERS="$("$PY_BIN" -m app.cluster workers)" || exit 2 # exec：让 uvicorn 直接收 Ctrl-C / SIGTERM，一个信号就走完整的优雅退出（收尾运行、关沙箱）。 # 要是包一层 shell 再转发，终端的 Ctrl-C 会让 uvicorn 收到两次信号，第二次直接强退、跳过收尾 # 多进程（--workers 大于 1，PostgreSQL）：uvicorn 起几个工作进程共用这个端口，各自在库里登记、续租约（app/cluster）