feat: 跨 worker 真 kill(kill_requests 意图表 + 收殓器)与自动化触发入闸

机制 A 移植自旧 PR #101(分支 feat-residual-b,头 895e004),
不合旧 PR:旧基线 6e3eb7b 已 CONFLICTING,本批重做到现 main
6d4f87e(#151 合并后)上。

- kill_requests 意图表:cloud/db.py POSTGRES_SCHEMA + migrations
  清单补表;server.ensure_kill_requests_table / request_kill
  (同一任务只留一条 pending,写失败返回 False);init_db 的
  SQLite 分支补建表(PG 分支由 POSTGRES_SCHEMA 全量建表覆盖)。
- 取消/sweep 写请求 best-effort:_cancel_task 在 commit 后调
  request_kill(失败不挡取消);sweep_stale_tasks 在同一事务
  内为被扫任务写 system 请求。
- worker 收殓器 _reap_kill_requests:owner 活着(心跳停滞未超
  300 秒,COALESCE(last_heartbeat, updated_at))绝不抢杀;只
  有 owner 已死且任务仍 interrupted 时,按任务派生容器名
  (executor.container_name,codex-<安全字符任务 id>)在本机
  真杀,绝不按前缀扫、绝不杀自由传入的名字;本机无容器时
  请求照常终结;任务行失踪或已不在 interrupted 时请求直接
  终结;单条异常只记日志。run loop 每轮调一次,放在 sweep
  块外(旧 PR 只在 sweep 块内,sweep 关掉时请求会积压)。
- 自动化入闸只补两条旁路:worker._coord_sync_children 建子任
  务前过 _quota_exceeded(超限跳过建任务、记审计,同一 worker
  内同一分工只记一次);automations.run_due_automation
  (claim 抢单 -> 配额门 -> 建任务 -> 记产出,超限记
  skipped_quota 不进失败重试)。triage/webhook/续跑/后续/
  分叉未动。
- 单测:旧 PR 三份单测原样移植(Worker/Server 签名与现 main
  兼容,未改用例语义)。
This commit is contained in:
host committed 2026-10-30 01:32:35 +08:00
1 parent b306b56517
commit 3a06205d8b
8 files changed
+938 -12

No files matched your search

+207 -10
View File
@@ -12,16 +12,15 @@
数据库里(automations 表,带 project_id),不是 host state 数据库里(automations 表,带 project_id),不是 host state
文件。两者存储不相交,同一个自动化绝不会同时出现在两边, 文件。两者存储不相交,同一个自动化绝不会同时出现在两边,
所以本批不存在双重执行。 所以本批不存在双重执行。
- 本批不给云端自动化加执行器、不动 host.py 的 run-due 区: - 云端执行器入口是本模块的 run_due_automation(残留批次 B
host 的 run-due 跑的是 host 任务模型(codex 协议线程),云 落地,host.py 的 run-due 区不动):它先走 claim_due_automation
端任务模型(POST /tasks 建、worker 认领)是另一条链路,把 抢单,抢到后再过配额门、建云端任务、记 automation_runs 一
云端定义硬塞进 host run-due 会造出第二个调度器。云端执行 行;抢不到的执行器绝不能建任务。这个 claim 门就是防重复
器是后续一步,它必须走 claim_due_automation 先抢单(条件 执行的唯一关口,比 host run-due 的“跑完才写 last_run”更严
UPDATE,rowcount 为 1 才算抢到),抢到后再建云端任务、再 (中途被杀也不会在当间隔内重跑,代价是崩掉的那次要等到
记 automation_runs 一行;抢不到的执行器绝不能建任务。这 下个间隔才补跑,见 claim 文档)。配额门与服务端建任务口
个 claim 门就是防重复执行的唯一关口,比 host run-due 的 径一致(见 _automation_quota_exceeded):超配额时不建任务,
“跑完才写 last_run”更严(中途被杀也不会在当间隔内重跑, 产出记 skipped_quota(不是失败,不进失败重试循环),并记审计。
代价是崩掉的那次要等到下个间隔,见 claim 文档)。
- 分诊收件箱是只读视图:到期自动化、失败或中断的任务、待 - 分诊收件箱是只读视图:到期自动化、失败或中断的任务、待
处理的后续消息、最近的自动化产出记录。它不跑模型、不改 处理的后续消息、最近的自动化产出记录。它不跑模型、不改
任务状态,只把待处理项聚到一处。 任务状态,只把待处理项聚到一处。
@@ -44,8 +43,10 @@ import;表迁移写法仿 peer_messages(CREATE TABLE IF NOT
EXISTS + PRAGMA 查列补齐,SQLite 与 Postgres 适配层同款)。 EXISTS + PRAGMA 查列补齐,SQLite 与 Postgres 适配层同款)。
""" """
import os
import time import time
import uuid import uuid
from datetime import datetime, timezone
MIN_INTERVAL_SECONDS = 60 MIN_INTERVAL_SECONDS = 60
MAX_INTERVAL_SECONDS = 30 * 86400 MAX_INTERVAL_SECONDS = 30 * 86400
@@ -461,6 +462,202 @@ def record_automation_run(conn, automation_id, project_id, task_id=None,
} }
def _default_quota_usd_micros():
"""默认月配额(微美元):与服务端 _default_quota_usd_micros
同值同环境变量(CLOUD_DEFAULT_QUOTA_USD_MICROS,默认 $10)。"""
try:
return max(0, int(os.environ.get(
"CLOUD_DEFAULT_QUOTA_USD_MICROS", "10000000")))
except (TypeError, ValueError):
return 10000000
def _unpriced_token_cap():
"""无价模型月 token 上限:与 cost.unpriced_token_cap 同值
同环境变量(CLOUD_QUOTA_UNPRICED_TOKEN_CAP,默认 2000 万)。"""
try:
return max(0, int(os.environ.get(
"CLOUD_QUOTA_UNPRICED_TOKEN_CAP", "20000000")))
except (TypeError, ValueError):
return 20000000
def _automation_quota_exceeded(conn, owner):
"""自动化建任务前的配额门:超限返回错误文本,否则 None。
口径与服务端 _quota_status / worker _quota_exceeded 一致:
按创建人算当月(UTC)已记账 USD + 无价模型 actual token;
无配额行按默认额度(不是不限额),显式 NULL 行才不限额;
查库失败 fail-closed(返回 quota_check_failed,调用方按跳
过处理,绝不建任务)。cost_records 表不存在时视为零花费
(CostStore 首次记账时才建表,不存在即从未记过账),其
它查库异常才 fail-closed——全新库上不能因此误拒全部运行。
本模块只依赖标准库,不 import cloud 其它模块,故此处就
地实现 SQL 版口径,注释里标明三处同口径点,改口径时三
处要一起改(server._quota_status、worker._quota_exceeded、
本函数)。
"""
if not owner:
return ("quota_exceeded: 自动化缺少创建人,无法查配额,"
"拒绝建任务")
try:
conn.execute(
"CREATE TABLE IF NOT EXISTS user_quotas("
"user_id TEXT PRIMARY KEY, quota_usd_micros BIGINT, "
" updated_at DOUBLE PRECISION)")
row = conn.execute(
"SELECT quota_usd_micros FROM user_quotas"
" WHERE user_id = ?",
(owner,)).fetchone()
if row is None:
quota = _default_quota_usd_micros()
else:
quota = row_get(row, "quota_usd_micros")
if quota is None:
return None
month = datetime.now(timezone.utc).strftime("%Y-%m")
try:
conn.execute("SELECT COUNT(*) FROM cost_records").fetchone()
table_exists = True
except Exception:
# PG 上失败语句会毒化当前事务:先 rollback,后面的
# 任务 INSERT / 产出记录才能继续写;SQLite 上无害。
try:
conn.rollback()
except Exception:
pass
table_exists = False
if table_exists:
used_row = conn.execute(
"SELECT COALESCE(SUM(c.usd_micros), 0) AS total"
" FROM cost_records c JOIN tasks t"
" ON c.task_id = t.id"
" WHERE t.created_by = ? AND c.day LIKE ?",
(owner, month + "%")).fetchone()
used = row_get(used_row, "total") or 0
unpriced_row = conn.execute(
"SELECT COALESCE(SUM(c.input_tokens + c.output_tokens), 0)"
" AS total"
" FROM cost_records c JOIN tasks t"
" ON c.task_id = t.id"
" WHERE t.created_by = ? AND c.day LIKE ?"
" AND c.usd_micros IS NULL AND c.source = 'actual'",
(owner, month + "%")).fetchone()
unpriced_used = row_get(unpriced_row, "total") or 0
else:
used = 0
unpriced_used = 0
except Exception as exc:
# 同上:先 rollback 再回 fail-closed,调用方仍能记
# skipped_quota 产出(不记成失败重试循环)。
try:
conn.rollback()
except Exception:
pass
return "quota_check_failed: %s" % exc
if used >= quota:
return ("quota_exceeded: 本月配额已用完(已用 %d / "
"配额 %d 微美元,%s)" % (used, quota, month))
if unpriced_used >= _unpriced_token_cap():
return ("quota_exceeded: 本月无价模型 token 已用完"
"(已用 %d / 上限 %d,%s)"
% (unpriced_used, _unpriced_token_cap(), month))
return None
def _audit_row(conn, project_id, action, detail, now):
"""best-effort 写审计行:失败不影响自动化运行本身。"""
try:
conn.execute(
"CREATE TABLE IF NOT EXISTS audit_log("
"id TEXT PRIMARY KEY, "
"user_id TEXT, "
"project_id TEXT, "
"action TEXT, "
"detail TEXT, "
"created_at DOUBLE PRECISION)")
conn.execute(
"INSERT INTO audit_log(id, user_id, project_id, action,"
" detail, created_at) VALUES(?, 'system', ?, ?, ?, ?)",
(uuid.uuid4().hex, project_id, action, detail, now))
conn.commit()
except Exception:
try:
conn.rollback()
except Exception:
pass
def run_due_automation(conn, automation_id, now=None):
"""云端自动化执行器的唯一入口:抢单 -> 配额门 -> 建任务 -> 记产出。
抢不到(不存在、已取消、未到期、已被别的执行器抢走)时
返回 None,不记产出、不建任务。抢到后先过配额门:创建人
超配额时不建任务,产出记 skipped_quota(不是失败,不进
失败重试循环)并记审计,返回该产出;配额内才 INSERT 云
端任务(queued,executor/model 留 NULL 走 worker 默认,
与服务端建任务缺省行为一致),产出记 created 并记审计。
任务 INSERT 与产出记录由 record_automation_run 的 commit
一并落盘。调用方(未来的定时执行器)只调本函数,绝不
自己拼 claim + INSERT,否则配额门可被绕过。
"""
if now is None:
now = time.time()
ensure_automations_tables(conn)
if not claim_due_automation(conn, automation_id, now):
return None
auto = conn.execute(
"SELECT id, project_id, prompt, created_by FROM automations"
" WHERE id = ?",
(automation_id,)).fetchone()
if auto is None:
return None
project_id = row_get(auto, "project_id")
prompt = row_get(auto, "prompt")
owner = row_get(auto, "created_by")
gate = _automation_quota_exceeded(conn, owner)
if gate is not None:
run = record_automation_run(
conn, automation_id, project_id, task_id=None,
status="skipped_quota", detail=gate, now=now)
_audit_row(conn, project_id, "automation_run_quota_skipped",
"%s: %s" % (automation_id, gate), now)
return run
task_id = uuid.uuid4().hex
# 任务列补齐(与 ensure_automations_tables 的补列惯例同款):
# 裸库/老库的 tasks 表可能没有 executor 等列,INSERT 前补齐,
# 否则建任务直接报错;PG 适配层会翻译 PRAGMA 查询。
existing_tasks = _existing_columns(conn, "tasks")
for col, ddl in (
("executor", "TEXT"),
("model", "TEXT"),
("effort", "TEXT"),
("environment_id", "TEXT"),
("parent_task_id", "TEXT"),
("coordinator_part", "TEXT")):
if col not in existing_tasks:
try:
conn.execute(
"ALTER TABLE tasks ADD COLUMN %s %s" % (col, ddl))
except Exception:
pass
conn.execute(
"INSERT INTO tasks(id, project_id, prompt, status,"
" idempotency_key, depends_on, priority, created_by,"
" created_at, updated_at, last_heartbeat, executor,"
" model, effort, environment_id, parent_task_id,"
" coordinator_part)"
" VALUES(?, ?, ?, 'queued', NULL, '[]', 0, ?, ?, ?,"
" NULL, NULL, NULL, NULL, NULL, NULL, NULL)",
(task_id, project_id, prompt, owner, now, now))
run = record_automation_run(
conn, automation_id, project_id, task_id=task_id,
status="created", detail="", now=now)
_audit_row(conn, project_id, "automation_run_created",
"%s -> %s" % (automation_id, task_id), now)
return run
def list_automation_runs(conn, automation_id, limit=TRIAGE_LIMIT): def list_automation_runs(conn, automation_id, limit=TRIAGE_LIMIT):
"""列某自动化的最近产出记录,按时间倒序。""" """列某自动化的最近产出记录,按时间倒序。"""
ensure_automations_tables(conn) ensure_automations_tables(conn)
+13
View File
@@ -1030,6 +1030,19 @@ POSTGRES_SCHEMA = (
"memory_id TEXT, " "memory_id TEXT, "
"created_at DOUBLE PRECISION, " "created_at DOUBLE PRECISION, "
"processed_at DOUBLE PRECISION)", "processed_at DOUBLE PRECISION)",
# kill_requests:跨 worker 真 kill 的意图表。取消/sweep 只改库
# 状态杀不掉 owner 已死时遗留的容器;写一条 pending 请求后,
# 任何活着的 worker 在 run loop 的收殓器里按任务派生容器名
# 在本机 docker 杀掉并标 done。同一任务只留一条 pending
# (request_kill 的 SELECT 守卫,见 server.py;并发重复写
# 撞上时多出的行同样幂等终结,不会多杀——杀本身按名幂等)。
"CREATE TABLE IF NOT EXISTS kill_requests("
"id TEXT PRIMARY KEY, "
"task_id TEXT, "
"requested_by TEXT, "
"status TEXT, "
"requested_at DOUBLE PRECISION, "
"done_at DOUBLE PRECISION)",
) )
+4 -2
View File
@@ -57,17 +57,19 @@ MIGRATION_TABLES = (
# (本批后无行按默认额度处理,不再是无限额,但显式 # (本批后无行按默认额度处理,不再是无限额,但显式
# 管理员额度仍需原样迁过去,不能丢)。 # 管理员额度仍需原样迁过去,不能丢)。
"user_quotas", "user_quotas",
# P2-ops 对齐:以下 4 张表在 POSTGRES_SCHEMA 里有、 # P2-ops 对齐:以下 5 张表在 POSTGRES_SCHEMA 里有、
# 原清单里没有,补进后清单与 schema 表集合一致。 # 原清单里没有,补进后清单与 schema 表集合一致。
# uploads / task_attachments:附件包与任务附件; # uploads / task_attachments:附件包与任务附件;
# task_followups:任务间互发消息的后续消息队列; # task_followups:任务间互发消息的后续消息队列;
# environments:项目环境(分支/executor)。 # environments:项目环境(分支/executor)。
# 四张表主键均为 TEXT 或复合主键,没有身份列, # kill_requests:跨 worker 真 kill 的意图表(残留批次 B)。
# 五张表主键均为 TEXT 或复合主键,没有身份列,
# 不需要序列拨动(见 SEQUENCE_TABLES)。 # 不需要序列拨动(见 SEQUENCE_TABLES)。
"uploads", "uploads",
"task_attachments", "task_attachments",
"task_followups", "task_followups",
"environments", "environments",
"kill_requests",
# Cloud identity batch: TEXT primary keys, no sequence to sync. # Cloud identity batch: TEXT primary keys, no sequence to sync.
"user_profiles", "user_profiles",
"identity_memories", "identity_memories",
+75
View File
@@ -1694,6 +1694,7 @@ def init_db(db_path):
_ensure_uploads_tables(conn) _ensure_uploads_tables(conn)
_ensure_task_checkpoints_table(conn) _ensure_task_checkpoints_table(conn)
_ensure_task_steps_table(conn) _ensure_task_steps_table(conn)
ensure_kill_requests_table(conn)
# worker 每轮 poll 都按 status 取任务(见 load_tasks 的 # worker 每轮 poll 都按 status 取任务(见 load_tasks 的
# WHERE status IN (...)),无索引时是全表扫描;任务表随 # WHERE status IN (...)),无索引时是全表扫描;任务表随
# 时间变大后索引是必须的。IF NOT EXISTS 保证老库幂等。 # 时间变大后索引是必须的。IF NOT EXISTS 保证老库幂等。
@@ -1803,6 +1804,56 @@ def _project_to_dict(row):
} }
def ensure_kill_requests_table(conn):
"""建 kill_requests 表(幂等,SQLite/PG 通用写法)。"""
conn.execute(
"CREATE TABLE IF NOT EXISTS kill_requests("
"id TEXT PRIMARY KEY, "
"task_id TEXT, "
"requested_by TEXT, "
"status TEXT, "
"requested_at DOUBLE PRECISION, "
"done_at DOUBLE PRECISION)")
def request_kill(db_path, task_id, requested_by):
"""写一条 kill 请求:同一任务只留一条 pending,返回 True/False。
已有 pending 时不重复写(返回 False);否则写一行 pending
并返回 True。调用方(取消/sweep)必须 best-effort 调用:
写请求失败绝不能挡住取消本身——任务状态已是 interrupted,
owner 活着时的 1 秒看门狗仍会杀,收殓器只是补 owner 已死
的情形。表不存在的老库先建表(与 worker 侧同款守卫)。
"""
try:
conn = _connect(db_path)
except Exception:
return False
try:
ensure_kill_requests_table(conn)
row = conn.execute(
"SELECT id FROM kill_requests"
" WHERE task_id = ? AND status = 'pending' LIMIT 1",
(task_id,)).fetchone()
if row is not None:
return False
conn.execute(
"INSERT INTO kill_requests(id, task_id, requested_by,"
" status, requested_at, done_at)"
" VALUES(?, ?, ?, 'pending', ?, NULL)",
(uuid.uuid4().hex, task_id, requested_by, time.time()))
conn.commit()
return True
except Exception:
try:
conn.rollback()
except Exception:
pass
return False
finally:
conn.close()
def sweep_stale_tasks(db_path, threshold_seconds=300): def sweep_stale_tasks(db_path, threshold_seconds=300):
"""把 running 且心跳过期的任务改为 interrupted,返回被扫任务 id 列表。 """把 running 且心跳过期的任务改为 interrupted,返回被扫任务 id 列表。
@@ -1834,6 +1885,23 @@ def sweep_stale_tasks(db_path, threshold_seconds=300):
) )
if cur.rowcount: if cur.rowcount:
swept.append(row["id"]) swept.append(row["id"])
# 被扫任务的 owner 可能已死:光改库状态杀不掉它遗留的
# 容器,顺手写 kill 请求,活着的 worker 在收殓器里按
# 任务派生容器名在本机真杀(见 worker._reap_kill_requests)。
# 同一事务内写,状态与请求要么一起落盘、要么一起丢。
if swept:
ensure_kill_requests_table(conn)
for task_id in swept:
exists = conn.execute(
"SELECT id FROM kill_requests"
" WHERE task_id = ? AND status = 'pending' LIMIT 1",
(task_id,)).fetchone()
if exists is None:
conn.execute(
"INSERT INTO kill_requests(id, task_id,"
" requested_by, status, requested_at, done_at)"
" VALUES(?, ?, 'system', 'pending', ?, NULL)",
(uuid.uuid4().hex, task_id, now))
conn.commit() conn.commit()
stale = [(row["id"], row["project_id"]) for row in rows stale = [(row["id"], row["project_id"]) for row in rows
if row["id"] in swept] if row["id"] in swept]
@@ -6748,6 +6816,13 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
"SELECT * FROM tasks WHERE id = ?", (task_id,) "SELECT * FROM tasks WHERE id = ?", (task_id,)
).fetchone() ).fetchone()
) )
# owner 活着时它自己的 1 秒看门狗会杀;owner 已死时
# 靠这条 kill 请求让别的活着的 worker 在收殓器里杀
# 遗留容器。写请求失败绝不挡取消本身(best-effort)。
try:
request_kill(self.server.db_path, task_id, user_id)
except Exception:
pass
# 取消后 worker 不会再跑它:暂存包没人 seed 了,顺手删掉 # 取消后 worker 不会再跑它:暂存包没人 seed 了,顺手删掉
# (任务终局后 worker 侧也会删,两边删同一目录都容忍 # (任务终局后 worker 侧也会删,两边删同一目录都容忍
# 不存在;删失败不挡取消本身)。 # 不存在;删失败不挡取消本身)。
+112
View File
@@ -0,0 +1,112 @@
"""自动化执行器入口 run_due_automation:抢单 + 配额门单测。"""
import os
import sqlite3
import tempfile
import time
import unittest
from cloud.automations import (
create_automation,
ensure_automations_tables,
list_automation_runs,
run_due_automation,
)
from cloud.server import init_db
def _tasks(db_path):
conn = sqlite3.connect(db_path)
try:
return conn.execute("SELECT id, prompt, status FROM tasks").fetchall()
finally:
conn.close()
def _audit_count(db_path, action):
conn = sqlite3.connect(db_path)
try:
return conn.execute(
"SELECT COUNT(*) FROM audit_log WHERE action = ?",
(action,)).fetchone()[0]
finally:
conn.close()
def _set_quota(db_path, user_id, micros):
conn = sqlite3.connect(db_path)
try:
conn.execute(
"CREATE TABLE IF NOT EXISTS user_quotas("
"user_id TEXT PRIMARY KEY, quota_usd_micros BIGINT, "
" updated_at DOUBLE PRECISION)")
conn.execute(
"INSERT OR REPLACE INTO user_quotas(user_id,"
" quota_usd_micros, updated_at) VALUES(?, ?, ?)",
(user_id, micros, time.time()))
conn.commit()
finally:
conn.close()
class RunDueAutomationTestCase(unittest.TestCase):
def setUp(self):
self.tmp = tempfile.TemporaryDirectory()
self.db_path = os.path.join(self.tmp.name, "control.db")
init_db(self.db_path)
self.conn = sqlite3.connect(self.db_path)
self.conn.row_factory = sqlite3.Row
ensure_automations_tables(self.conn)
def tearDown(self):
self.conn.close()
self.tmp.cleanup()
def _auto(self, owner="owner1"):
return create_automation(self.conn, "proj", "daily", "do thing",
owner, interval_seconds=60)
def test_run_creates_task_within_quota(self):
auto = self._auto()
run = run_due_automation(self.conn, auto["id"])
self.assertIsNotNone(run)
self.assertEqual(run["status"], "created")
self.assertIsNotNone(run["task_id"])
tasks = _tasks(self.db_path)
self.assertEqual(len(tasks), 1)
self.assertEqual(tasks[0][1], "do thing")
self.assertEqual(tasks[0][2], "queued")
self.assertEqual(
_audit_count(self.db_path, "automation_run_created"), 1)
# 同一间隔内第二次跑:已抢过,返回 None,不重复建任务。
self.assertIsNone(run_due_automation(self.conn, auto["id"]))
self.assertEqual(len(_tasks(self.db_path)), 1)
def test_run_skips_when_over_quota(self):
auto = self._auto()
_set_quota(self.db_path, "owner1", 0)
run = run_due_automation(self.conn, auto["id"])
self.assertIsNotNone(run)
self.assertEqual(run["status"], "skipped_quota")
self.assertIsNone(run["task_id"])
self.assertEqual(_tasks(self.db_path), [])
self.assertEqual(
_audit_count(self.db_path, "automation_run_quota_skipped"), 1)
# 跳过仍记了产出(不是失败),同间隔内不再重跑。
runs = list_automation_runs(self.conn, auto["id"])
self.assertEqual(len(runs), 1)
self.assertEqual(runs[0]["status"], "skipped_quota")
def test_run_not_due_returns_none(self):
auto = self._auto()
self.conn.execute(
"UPDATE automations SET last_run_at = ? WHERE id = ?",
(time.time(), auto["id"]))
self.conn.commit()
self.assertIsNone(run_due_automation(self.conn, auto["id"]))
self.assertEqual(_tasks(self.db_path), [])
self.assertEqual(
list_automation_runs(self.conn, auto["id"]), [])
if __name__ == "__main__":
unittest.main()
+130
View File
@@ -0,0 +1,130 @@
"""协调员子任务创建时配额门单测。"""
import json
import os
import sqlite3
import tempfile
import unittest
from cloud.server import init_db
from cloud.worker import Worker
def _insert_parent(db_path, task_id="coord", created_by="owner1"):
conn = sqlite3.connect(db_path)
try:
parts = json.dumps([{"name": "后端", "prompt": "做后端"},
{"name": "前端", "prompt": "做前端"}])
conn.execute(
"INSERT INTO tasks(id, project_id, prompt, status,"
" created_by, created_at, updated_at, coordinator_parts)"
" VALUES(?, 'proj', 'coord', 'running', ?, ?, ?, ?)",
(task_id, created_by, 1000.0, 1000.0, parts))
conn.commit()
finally:
conn.close()
def _children(db_path, parent_id):
conn = sqlite3.connect(db_path)
try:
return conn.execute(
"SELECT id FROM tasks WHERE parent_task_id = ?",
(parent_id,)).fetchall()
finally:
conn.close()
def _audit_count(db_path, action):
conn = sqlite3.connect(db_path)
try:
try:
return conn.execute(
"SELECT COUNT(*) FROM audit_log WHERE action = ?",
(action,)).fetchone()[0]
except sqlite3.OperationalError:
# 审计是 best-effort:写审计的连接与任务写连接在
# SQLite 上撞写锁时审计会静默丢失(既有行为),
# 此时表可能根本不存在,计 0 而不是让测试炸掉。
return 0
finally:
conn.close()
def _set_quota(db_path, user_id, micros):
conn = sqlite3.connect(db_path)
try:
conn.execute(
"CREATE TABLE IF NOT EXISTS user_quotas("
"user_id TEXT PRIMARY KEY, quota_usd_micros BIGINT, "
" updated_at DOUBLE PRECISION)")
conn.execute(
"INSERT OR REPLACE INTO user_quotas(user_id,"
" quota_usd_micros, updated_at) VALUES(?, ?, ?)",
(user_id, micros, 1000.0))
conn.commit()
finally:
conn.close()
class CoordQuotaGateTestCase(unittest.TestCase):
def setUp(self):
self.tmp = tempfile.TemporaryDirectory()
self.db_path = os.path.join(self.tmp.name, "control.db")
init_db(self.db_path)
self.env = {"XCM_API_KEY": "test-placeholder"}
def tearDown(self):
self.tmp.cleanup()
def _worker(self):
return Worker(self.db_path, "/fake/host.py", env=self.env)
def _parts(self):
return [{"name": "后端", "prompt": "做后端"},
{"name": "前端", "prompt": "做前端"}]
def test_skip_when_over_quota(self):
_insert_parent(self.db_path)
_set_quota(self.db_path, "owner1", 0)
worker = self._worker()
task = worker.get_task("coord")
children = worker._coord_sync_children(
task, task["coordinator_parts"])
self.assertEqual(children, [])
self.assertEqual(_children(self.db_path, "coord"), [])
# 两个分工各记一次跳过审计;第二轮重调不再重复记。
self.assertEqual(
_audit_count(self.db_path, "coordinator_child_quota_skipped"),
2)
worker._coord_sync_children(task, task["coordinator_parts"])
self.assertEqual(
_audit_count(self.db_path, "coordinator_child_quota_skipped"),
2)
self.assertEqual(_children(self.db_path, "coord"), [])
def test_create_when_within_quota(self):
_insert_parent(self.db_path)
_set_quota(self.db_path, "owner1", 100000000)
worker = self._worker()
task = worker.get_task("coord")
children = worker._coord_sync_children(
task, task["coordinator_parts"])
self.assertEqual(len(children), 2)
self.assertEqual(len(_children(self.db_path, "coord")), 2)
self.assertEqual(
_audit_count(self.db_path, "coordinator_child_quota_skipped"),
0)
def test_create_with_default_quota_when_no_row(self):
# 无配额行按默认 $10 计,未花过即放行(不是不限额,
# 也不是零额度误拒)。
_insert_parent(self.db_path)
worker = self._worker()
task = worker.get_task("coord")
children = worker._coord_sync_children(
task, task["coordinator_parts"])
self.assertEqual(len(children), 2)
if __name__ == "__main__":
unittest.main()
+241
View File
@@ -0,0 +1,241 @@
"""跨 worker 真 kill:kill_requests 意图表 + 收殓器单测。"""
import http.client
import json
import os
import sqlite3
import tempfile
import threading
import time
import unittest
from cloud.auth import AuthStore
from cloud.server import (create_server, init_db, request_kill,
sweep_stale_tasks)
from cloud.worker import Worker
def _insert_task(db_path, task_id, status, heartbeat_age=1000):
conn = sqlite3.connect(db_path)
try:
now = time.time()
conn.execute(
"INSERT INTO tasks(id, project_id, prompt, status,"
" created_by, created_at, updated_at, last_heartbeat)"
" VALUES(?, 'proj', 'do', ?, 'u1', ?, ?, ?)",
(task_id, status, now, now, now - heartbeat_age))
conn.commit()
finally:
conn.close()
def _kill_rows(db_path, task_id=None):
conn = sqlite3.connect(db_path)
conn.row_factory = sqlite3.Row
try:
if task_id is None:
rows = conn.execute(
"SELECT * FROM kill_requests").fetchall()
else:
rows = conn.execute(
"SELECT * FROM kill_requests WHERE task_id = ?",
(task_id,)).fetchall()
return [dict(r) for r in rows]
finally:
conn.close()
def _audit_actions(db_path, action):
conn = sqlite3.connect(db_path)
try:
rows = conn.execute(
"SELECT * FROM audit_log WHERE action = ?", (action,)).fetchall()
return rows
finally:
conn.close()
class _FakeDockerExecutor:
"""假 docker 执行器:只记杀了哪个任务派生名,不碰真 docker。"""
def __init__(self, succeed=True):
self.killed = []
self.succeed = succeed
def container_name(self, task):
return "codex-" + str(task.get("id"))
def kill_container(self, task_or_name):
name = (self.container_name(task_or_name)
if isinstance(task_or_name, dict) else str(task_or_name))
self.killed.append(name)
return self.succeed
class KillRequestTestCase(unittest.TestCase):
def setUp(self):
self.tmp = tempfile.TemporaryDirectory()
self.db_path = os.path.join(self.tmp.name, "control.db")
init_db(self.db_path)
def tearDown(self):
self.tmp.cleanup()
def test_request_kill_idempotent(self):
self.assertTrue(request_kill(self.db_path, "t1", "u1"))
self.assertFalse(request_kill(self.db_path, "t1", "u1"))
rows = _kill_rows(self.db_path, "t1")
self.assertEqual(len(rows), 1)
self.assertEqual(rows[0]["status"], "pending")
def test_sweep_writes_kill_request(self):
_insert_task(self.db_path, "t-old", "running")
swept = sweep_stale_tasks(self.db_path, 300)
self.assertIn("t-old", swept)
rows = _kill_rows(self.db_path, "t-old")
self.assertEqual(len(rows), 1)
self.assertEqual(rows[0]["requested_by"], "system")
# 第二次 sweep:任务已不是 running,不再重复写。
sweep_stale_tasks(self.db_path, 300)
self.assertEqual(len(_kill_rows(self.db_path, "t-old")), 1)
class CancelHttpTestCase(unittest.TestCase):
def setUp(self):
self.tmp = tempfile.TemporaryDirectory()
self.db_path = os.path.join(self.tmp.name, "control.db")
init_db(self.db_path)
self.auth = AuthStore(self.db_path)
self.user_a, self.key_a = self.auth.create_user("user-a")
self.server = create_server(self.db_path, "127.0.0.1", 0)
self.port = self.server.server_address[1]
self.thread = threading.Thread(
target=self.server.serve_forever, daemon=True)
self.thread.start()
def tearDown(self):
self.server.shutdown()
self.server.server_close()
self.thread.join(timeout=5)
try:
self.auth.conn.close()
except Exception:
pass
self.tmp.cleanup()
def _req(self, method, path, body=None, key=None):
conn = http.client.HTTPConnection("127.0.0.1", self.port, timeout=10)
try:
headers = {}
payload = None
if body is not None:
payload = json.dumps(body)
headers["Content-Type"] = "application/json"
if key is not None:
headers["X-API-Key"] = key
conn.request(method, path, body=payload, headers=headers)
resp = conn.getresponse()
raw = resp.read().decode()
try:
parsed = json.loads(raw)
except ValueError:
parsed = raw
return resp.status, parsed
finally:
conn.close()
def test_cancel_writes_kill_request(self):
status, proj = self._req("POST", "/projects", {"name": "p"},
key=self.key_a)
self.assertEqual(status, 201)
pid = proj["id"] if isinstance(proj, dict) else proj["project"]["id"]
status, task = self._req("POST", "/tasks",
{"project_id": pid, "prompt": "do"},
key=self.key_a)
self.assertEqual(status, 201)
tid = task["id"] if isinstance(task, dict) else task["task"]["id"]
status, _ = self._req("POST", "/tasks/%s/cancel" % tid, {},
key=self.key_a)
self.assertEqual(status, 200)
rows = _kill_rows(self.db_path, tid)
self.assertEqual(len(rows), 1)
self.assertEqual(rows[0]["status"], "pending")
# 重复取消:任务已终态,400,不重复写请求。
status, _ = self._req("POST", "/tasks/%s/cancel" % tid, {},
key=self.key_a)
self.assertEqual(status, 400)
self.assertEqual(len(_kill_rows(self.db_path, tid)), 1)
class ReaperTestCase(unittest.TestCase):
def setUp(self):
self.tmp = tempfile.TemporaryDirectory()
self.db_path = os.path.join(self.tmp.name, "control.db")
init_db(self.db_path)
def tearDown(self):
self.tmp.cleanup()
def _worker(self, executor):
return Worker(self.db_path, "/fake/host.py",
env={"XCM_API_KEY": "test-placeholder"},
executor=executor)
def _seed_request(self, task_id):
conn = sqlite3.connect(self.db_path)
try:
conn.execute(
"INSERT INTO kill_requests(id, task_id, requested_by,"
" status, requested_at, done_at)"
" VALUES(?, ?, 'u1', 'pending', ?, NULL)",
("req-" + task_id, task_id, time.time()))
conn.commit()
finally:
conn.close()
def test_reap_kills_stale_owner_container(self):
_insert_task(self.db_path, "t1", "interrupted", heartbeat_age=1000)
self._seed_request("t1")
fake = _FakeDockerExecutor()
worker = self._worker(fake)
done = worker._reap_kill_requests(300)
self.assertEqual(done, ["req-t1"])
self.assertEqual(fake.killed, ["codex-t1"])
self.assertEqual(_kill_rows(self.db_path, "t1")[0]["status"], "done")
self.assertTrue(_audit_actions(self.db_path, "kill_request_reaped"))
def test_reap_skips_live_owner(self):
_insert_task(self.db_path, "t1", "interrupted", heartbeat_age=10)
self._seed_request("t1")
fake = _FakeDockerExecutor()
worker = self._worker(fake)
done = worker._reap_kill_requests(300)
self.assertEqual(done, [])
self.assertEqual(fake.killed, [])
self.assertEqual(_kill_rows(self.db_path, "t1")[0]["status"],
"pending")
def test_reap_stale_request_for_live_task(self):
# 任务已重回 queued(同 id 复活):旧请求直接终结,绝不杀新跑。
_insert_task(self.db_path, "t1", "queued", heartbeat_age=1000)
self._seed_request("t1")
fake = _FakeDockerExecutor()
worker = self._worker(fake)
done = worker._reap_kill_requests(300)
self.assertEqual(done, ["req-t1"])
self.assertEqual(fake.killed, [])
self.assertTrue(_audit_actions(self.db_path, "kill_request_stale"))
def test_reap_missing_task_and_missing_container(self):
self._seed_request("t-gone")
_insert_task(self.db_path, "t2", "interrupted", heartbeat_age=1000)
self._seed_request("t2")
fake = _FakeDockerExecutor(succeed=False)
worker = self._worker(fake)
done = worker._reap_kill_requests(300)
self.assertEqual(sorted(done), ["req-t-gone", "req-t2"])
# 本机无容器(kill 返回 False)也不报错,请求照常终结。
self.assertEqual(_kill_rows(self.db_path, "t2")[0]["status"], "done")
if __name__ == "__main__":
unittest.main()
+156
View File
@@ -614,6 +614,11 @@ class Worker:
# 配额中断就续跑下去了。受 _quota_hit_lock 保护。 # 配额中断就续跑下去了。受 _quota_hit_lock 保护。
self._quota_hit_tasks = set() self._quota_hit_tasks = set()
self._quota_hit_lock = threading.Lock() self._quota_hit_lock = threading.Lock()
# 协调员子任务因配额跳过建任务时已记过审计的分工
# ((parent_id, part_name)):协调员每轮都会重调
# _coord_sync_children,超配额期间每轮都记审计会刷屏,
# 同一 worker 进程内同一分工只记一次。
self._coord_quota_skipped = set()
def _code_changed(self): def _code_changed(self):
"""启动后关键代码是否被换过(部署换目录/手动改代码): """启动后关键代码是否被换过(部署换目录/手动改代码):
@@ -4016,6 +4021,24 @@ class Worker:
"%s -> %s" % (name, prev["id"])) "%s -> %s" % (name, prev["id"]))
children.append({"name": name, "id": prev["id"]}) children.append({"name": name, "id": prev["id"]})
continue continue
# 自动化入闸(创建时门):子任务认领时虽有配额门,
# 但不建任务才是扇出不烧钱的硬保证。父任务创建人
# 超配额时跳过建这个子任务、记审计,父任务本身不
# 失败(协调员按缺失子任务的既有语义等超时);同
# 一 worker 进程内同一分工只记一次审计,避免每轮
# 重调都刷屏。查库失败时 _quota_exceeded 同样返回
# 非空串(fail-closed),一并跳过。
gate = self._quota_exceeded(
{"created_by": task["created_by"]}, None)
if gate is not None:
key = (task["id"], name)
if key not in self._coord_quota_skipped:
self._coord_quota_skipped.add(key)
self._audit(
task["project_id"],
"coordinator_child_quota_skipped",
"%s: %s" % (name, gate))
continue
child_id = uuid.uuid4().hex child_id = uuid.uuid4().hex
conn.execute( conn.execute(
"INSERT INTO tasks(id, project_id, prompt, status," "INSERT INTO tasks(id, project_id, prompt, status,"
@@ -4124,6 +4147,125 @@ class Worker:
conn.close() conn.close()
return marked return marked
def _reap_kill_requests(self, sweep_threshold=300):
"""跨 worker 真 kill 的收殓器(run loop 每轮调一次)。
分工(与 owner 活着时的 1 秒看门狗互补,绝不抢杀):
owner 活着时,它的 _docker_cancel_watchdog 在任务变
interrupted 后 1 秒内按容器名真杀,本函数不碰;只有
owner 已死(任务心跳停滞超过 sweep_threshold,与陈旧
回收同一阈值、同一口径 COALESCE(last_heartbeat,
updated_at))时,本函数才动手——因为这时已没有任何
进程会来杀这个容器,不杀就是一直烧钱。
判定链(任一不满足即跳过,请求留待下轮;任务行失踪
或任务已不在 interrupted 时请求直接终结,避免复活的
任务(同 id 重回 queued)被旧请求误杀):
1. 请求 status 仍为 pending;
2. 任务存在且 status == 'interrupted';
3. 任务心跳停滞超过 sweep_threshold;
4. 本机 executor 有 container_name/kill_container
(docker 执行器才有;local 执行器无容器名,跳过)。
杀本身按任务派生容器名(executor.container_name,
codex-<安全字符任务 id>),绝不按前缀扫、绝不杀自由
传入的名字;本机无此容器时 kill_container 返回 False
也不报错,请求照常终结。终结时写 kill_request_reaped
审计;单条请求异常只记日志、不炸整轮。
返回终结的请求 id 列表。
"""
done = []
try:
conn = _db_connect(self.db_path)
except Exception as exc:
self.logger.warning("kill reap: db connect failed: %s", exc)
return done
try:
try:
conn.execute(
"CREATE TABLE IF NOT EXISTS kill_requests("
"id TEXT PRIMARY KEY, "
"task_id TEXT, "
"requested_by TEXT, "
"status TEXT, "
"requested_at DOUBLE PRECISION, "
"done_at DOUBLE PRECISION)")
conn.commit()
except Exception:
pass
try:
reqs = conn.execute(
"SELECT id, task_id FROM kill_requests"
" WHERE status = 'pending'").fetchall()
except Exception as exc:
self.logger.warning("kill reap: list failed: %s", exc)
return done
now = time.time()
kill = getattr(self.executor, "kill_container", None)
name_of = getattr(self.executor, "container_name", None)
for req in reqs:
req_id = req["id"]
task_id = req["task_id"]
try:
row = conn.execute(
"SELECT status, project_id, last_heartbeat,"
" updated_at FROM tasks WHERE id = ?",
(task_id,)).fetchone()
if row is None:
# 任务行失踪:请求已无意义,直接终结。
conn.execute(
"UPDATE kill_requests SET status = 'done',"
" done_at = ? WHERE id = ?",
(now, req_id))
conn.commit()
done.append(req_id)
continue
if row["status"] != "interrupted":
# 任务已不在 interrupted(重回 queued 等):
# 旧请求直接终结,绝不能杀新跑。
conn.execute(
"UPDATE kill_requests SET status = 'done',"
" done_at = ? WHERE id = ?",
(now, req_id))
conn.commit()
done.append(req_id)
self._audit(row["project_id"],
"kill_request_stale",
"%s not_interrupted" % task_id)
continue
hb = row["last_heartbeat"]
if hb is None:
hb = row["updated_at"]
if hb is not None and hb >= now - sweep_threshold:
# owner 还活着:留待下轮,绝不抢杀。
continue
# owner 已死:按任务派生容器名在本机真杀。
# 容器名只由任务 id 派生,不接受外部输入;
# 本机无此容器时 kill 返回 False 也不报错。
if kill is not None and name_of is not None:
try:
kill({"id": task_id})
except Exception:
pass
conn.execute(
"UPDATE kill_requests SET status = 'done',"
" done_at = ? WHERE id = ?",
(now, req_id))
conn.commit()
done.append(req_id)
self._audit(row["project_id"], "kill_request_reaped",
task_id)
except Exception as exc:
self.logger.warning(
"kill reap: request %s failed: %s", req_id, exc)
try:
conn.rollback()
except Exception:
pass
continue
finally:
conn.close()
return done
def sweep_orphaned_coordinator_children(self): def sweep_orphaned_coordinator_children(self):
"""常驻循环每轮的孤儿扫尾(P1):父已终态或父行失踪、子 """常驻循环每轮的孤儿扫尾(P1):父已终态或父行失踪、子
仍未终态时,把子标 interrupted 并记审计。 仍未终态时,把子标 interrupted 并记审计。
@@ -5495,6 +5637,20 @@ class Worker:
except Exception as exc: except Exception as exc:
log_event(self.logger, "worker_orphan_sweep_failed", log_event(self.logger, "worker_orphan_sweep_failed",
message="orphan sweep failed: %s" % exc) message="orphan sweep failed: %s" % exc)
# 跨 worker 真 kill 的收殓器:放在 sweep 块外、每轮
# 都跑(旧 PR 只在 sweep 块内跑,sweep 关掉时请
# 求会积压;本函数自带心跳停滞判定,sweep 关掉
# 也不抢杀活 owner)。陈旧回收刚把死 worker 的
# running 任务标 interrupted 并写了 kill 请求,
# 这里按请求杀本机遗留容器;owner 活着的任务心
# 跳新鲜,收殓器绝不抢杀。
try:
self._reap_kill_requests(
sweep_threshold
if sweep_threshold is not None else 300)
except Exception as exc:
log_event(self.logger, "worker_reap_failed",
message="worker reap failed: %s" % exc)
# 长任务自动恢复:接在 sweep 之后——刚被 sweep # 长任务自动恢复:接在 sweep 之后——刚被 sweep
# 标 interrupted 的系统中断 goal 任务,本轮即可 # 标 interrupted 的系统中断 goal 任务,本轮即可
# 回 queued、下一轮认领即断点续跑;用户暂停、 # 回 queued、下一轮认领即断点续跑;用户暂停、