493 lines
21 KiB
Python
493 lines
21 KiB
Python
"""云端 worker：把控制面的 queued 任务接通到 host.py 执行。
|
||
|
||
零第三方依赖，只用 Python 标准库；不改 host.py 行为、不动
|
||
~/.codex 配置、不碰 wechat-bridge。
|
||
|
||
流程（与 server.py 的任务状态机对齐）：
|
||
queued -> running -> completed / failed / interrupted
|
||
|
||
1. 认领：pick_next() 用 queue.ready_tasks() 按依赖拓扑 + 优先级
|
||
排序，跳过依赖未完成（未 completed/merged）的任务；再过
|
||
queue.budget_exceeded() 预算门，超限不派发、任务保持 queued。
|
||
claim() 用 `WHERE status='queued'` 的 CAS 更新把任务标 running
|
||
并记 last_heartbeat / updated_at，避免两个 worker 重复认领。
|
||
2. 缺 key 门（执行前必查）：只检查配置的 env_key 名在进程环境
|
||
里是否存在且非空，绝不读出、打印、写进文件或日志 key 的值。
|
||
缺 key 时任务直接标 failed，错误固定为
|
||
`Missing environment variable: <env_key 名>`，与之前
|
||
real-nightly-tests automation 的失败模式一致，绝不静默假装跑了。
|
||
env_key 名的来源：显式参数 > codex 配置文件里的 env_key 行
|
||
（只读名字）> 默认 XCM_API_KEY（中转站 relay 的约定名）。
|
||
3. 执行：经 LocalExecutor 调 host.py 真实子命令。dispatch
|
||
（建 host 任务并跑一轮 turn）是首次派发；dispatch 报
|
||
“任务已存在”时改用 run（在已有 host 任务上再跑一轮）续跑。
|
||
host 任务名由 cloud 任务 id 派生（见 sandbox._host_name），
|
||
同一任务每次算出同一名字，续跑才能找到同一个 host 任务。
|
||
4. 写回：server 的 tasks 表没有 result/error 列，所以执行结果
|
||
的持久化分三处，各复用已有结构、不另起字段名：
|
||
- tasks 表：最终状态 + updated_at（取消时为 interrupted）；
|
||
- cost.py CostStore：每轮实际执行记一条 source="estimated"
|
||
的成本记录（host 协议不暴露真实 token 数，估算绝不冒充
|
||
actual；total_tokens 由 CostStore.record 按 input+output 算）；
|
||
- models.py：可 fallback 的失败（如 host 状态 timeout）且
|
||
router 给出下一个模型时，用 record_fallback() 在任务记录上
|
||
记 model_fallback 事件，并把 fallback_count 写进成本记录。
|
||
注意：worker 只记录 fallback 事件，不自动用下一个模型重试——
|
||
自动重试等于未经批准烧 relay 额度，重试与否由调用方决定。
|
||
审计：认领/完成/失败/取消/缺 key 各记一条 AuthStore 审计
|
||
（user_id="system"，与 server.sweep_stale_tasks 的做法一致），
|
||
审计失败不影响任务状态（best-effort）。
|
||
|
||
取消：host.py 的 argparse 子命令表里没有 interrupt（turn/interrupt
|
||
只存在于 host.py Server 类内部的 app-server 协议调用，未暴露为
|
||
CLI），所以 worker 的 cancel() 走 server 的 cancel 语义：把
|
||
queued/running 任务直接标 interrupted（已终态的任务不动）。
|
||
进行中的 host.py 子进程不会因此被杀（中途杀进程可能把 host
|
||
state 写一半）；run_one() 在执行返回后会先查任务是否已被取消，
|
||
已被取消的一律保持 interrupted，不用执行结果覆盖它。
|
||
|
||
可注入：runner（假 CLI，见 sandbox.LocalExecutor）与 env 映射
|
||
均可注入，单测不跑真实模型回合、不联网、不登录任何服务。
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
import re
|
||
import time
|
||
|
||
from cloud.auth import AuthStore
|
||
from cloud.cost import CostStore, estimate_tokens
|
||
from cloud.db import connect as _db_connect
|
||
from cloud.models import ModelRouter, is_fallbackable, record_fallback
|
||
from cloud.queue import budget_exceeded, ready_tasks
|
||
from cloud.sandbox import LocalExecutor
|
||
|
||
# 与 server.TERMINAL_STATUSES 对齐：这些状态的任务不再被认领或取消。
|
||
TERMINAL_STATUSES = frozenset(
|
||
{"completed", "failed", "interrupted", "merged", "merge-failed"})
|
||
COMPLETED_STATUSES = frozenset({"completed", "merged"})
|
||
DEFAULT_ENV_KEY = "XCM_API_KEY"
|
||
DEFAULT_MODEL = "gpt-6.1-sol"
|
||
# host.py dispatch 的原话是 "task 已存在: <name>，先换名或 discard 它"
|
||
#（英文 task + 已存在），取共有子串做判断，不依赖标点。
|
||
_ALREADY_EXISTS_MARKER = "已存在"
|
||
|
||
|
||
def _missing_key_error(env_key):
|
||
"""缺 key 的标准错误：只含 env_key 名，与 relay 侧的报错一致。"""
|
||
return "Missing environment variable: %s" % env_key
|
||
|
||
|
||
def resolve_env_key(env_key=None, codex_config_path=None):
|
||
"""确定模型 key 对应的环境变量名（只返回名字，从不返回值）。
|
||
|
||
显式参数优先；其次读 codex 配置文件（如 relay 的 config.toml）
|
||
里的 `env_key = "..."` 行，只取名字；文件不存在或无此行时
|
||
用默认名。绝不读 ~/.codex/config.toml 之外的默认位置——
|
||
调用方要读哪个配置文件，必须显式传 codex_config_path。
|
||
"""
|
||
if env_key:
|
||
return env_key
|
||
if codex_config_path:
|
||
try:
|
||
with open(codex_config_path, encoding="utf-8") as f:
|
||
text = f.read()
|
||
except (OSError, UnicodeDecodeError):
|
||
text = ""
|
||
match = re.search(
|
||
r'^\s*env_key\s*=\s*"([^"]+)"', text, re.MULTILINE)
|
||
if match:
|
||
return match.group(1)
|
||
return DEFAULT_ENV_KEY
|
||
|
||
|
||
def _resolve_model(model=None, config_path=None):
|
||
"""模型名：显式参数 > host 配置文件里的 model > host 默认值。"""
|
||
if model:
|
||
return model
|
||
if config_path:
|
||
try:
|
||
with open(config_path, encoding="utf-8") as f:
|
||
cfg = json.load(f)
|
||
except (OSError, ValueError):
|
||
cfg = {}
|
||
if isinstance(cfg, dict) and cfg.get("model"):
|
||
return cfg["model"]
|
||
return DEFAULT_MODEL
|
||
|
||
|
||
def _task_from_row(row):
|
||
depends_on = row["depends_on"]
|
||
try:
|
||
depends_on = json.loads(depends_on) if depends_on else []
|
||
except (TypeError, ValueError):
|
||
depends_on = []
|
||
if not isinstance(depends_on, list):
|
||
depends_on = [depends_on]
|
||
return {
|
||
"id": row["id"],
|
||
"project_id": row["project_id"],
|
||
"prompt": row["prompt"],
|
||
"status": row["status"],
|
||
"idempotency_key": row["idempotency_key"],
|
||
"depends_on": depends_on,
|
||
"priority": row["priority"],
|
||
"created_by": row["created_by"],
|
||
"created_at": row["created_at"],
|
||
"updated_at": row["updated_at"],
|
||
"last_heartbeat": row["last_heartbeat"],
|
||
}
|
||
|
||
|
||
class Worker:
|
||
"""单机 worker：一轮认领一个 queued 任务并执行到底。
|
||
|
||
前置条件：db_path 须已由 cloud.server.init_db() 建好 tasks 表；
|
||
本模块不自建 server 的表（成本表与审计表由 CostStore/AuthStore
|
||
在首次使用时自建，与 server 的做法一致）。
|
||
"""
|
||
|
||
def __init__(self, db_path, host_py_path, config_path=None, repo=None,
|
||
runner=None, env=None, env_key=None,
|
||
codex_config_path=None, model=None, router=None,
|
||
max_turns=None, max_cost_usd=None, executor=None,
|
||
cost_store=None):
|
||
self.db_path = str(db_path)
|
||
self.host_py_path = host_py_path
|
||
self.config_path = config_path
|
||
self.repo = repo
|
||
# env 只用于 key 存在性检查：默认接进程环境，单测可注入
|
||
# 一个普通 dict；worker 永远不把 env 里的值写出。
|
||
self._env = env if env is not None else os.environ
|
||
self.env_key = resolve_env_key(env_key, codex_config_path)
|
||
self.model = _resolve_model(model, config_path)
|
||
if router is not None and not isinstance(router, ModelRouter):
|
||
raise ValueError("router 必须是 cloud.models.ModelRouter")
|
||
self.router = router
|
||
self.max_turns = max_turns
|
||
self.max_cost_usd = max_cost_usd
|
||
self.usage_turns = 0
|
||
self.usage_cost_usd = 0.0
|
||
if executor is not None:
|
||
self.executor = executor
|
||
else:
|
||
self.executor = LocalExecutor(
|
||
host_py_path, config_path=config_path, runner=runner,
|
||
repo=repo)
|
||
self._cost_store = cost_store
|
||
self._cancel_requested = set()
|
||
|
||
# -- 成本存储（惰性建表，避免只认领不执行时也建表）--------------
|
||
@property
|
||
def costs(self):
|
||
if self._cost_store is None:
|
||
self._cost_store = CostStore(self.db_path)
|
||
return self._cost_store
|
||
|
||
def close(self):
|
||
"""关闭惰性建的成本存储连接（注入的 cost_store 由调用方管）。"""
|
||
if self._cost_store is not None:
|
||
try:
|
||
self._cost_store._conn.close()
|
||
except Exception:
|
||
pass
|
||
|
||
# -- DB 读写 ------------------------------------------------------
|
||
def load_tasks(self):
|
||
conn = _db_connect(self.db_path)
|
||
try:
|
||
rows = conn.execute(
|
||
"SELECT * FROM tasks ORDER BY created_at, id").fetchall()
|
||
finally:
|
||
conn.close()
|
||
return [_task_from_row(r) for r in rows]
|
||
|
||
def get_task(self, task_id):
|
||
conn = _db_connect(self.db_path)
|
||
try:
|
||
row = conn.execute(
|
||
"SELECT * FROM tasks WHERE id = ?", (task_id,)).fetchone()
|
||
finally:
|
||
conn.close()
|
||
return _task_from_row(row) if row is not None else None
|
||
|
||
def _claim(self, task_id):
|
||
"""CAS 认领：仅当任务仍为 queued 时标 running 并记心跳。"""
|
||
now = time.time()
|
||
conn = _db_connect(self.db_path)
|
||
try:
|
||
cur = conn.execute(
|
||
"UPDATE tasks SET status = 'running',"
|
||
" last_heartbeat = ?, updated_at = ?"
|
||
" WHERE id = ? AND status = 'queued'",
|
||
(now, now, task_id),
|
||
)
|
||
conn.commit()
|
||
return cur.rowcount == 1
|
||
finally:
|
||
conn.close()
|
||
|
||
def _set_status(self, task_id, status):
|
||
now = time.time()
|
||
conn = _db_connect(self.db_path)
|
||
try:
|
||
conn.execute(
|
||
"UPDATE tasks SET status = ?, last_heartbeat = ?,"
|
||
" updated_at = ? WHERE id = ?",
|
||
(status, now, now, task_id),
|
||
)
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
def heartbeat(self, task_id):
|
||
"""刷新 running 任务的心跳；任务不在 running 时返回 False。"""
|
||
now = time.time()
|
||
conn = _db_connect(self.db_path)
|
||
try:
|
||
cur = conn.execute(
|
||
"UPDATE tasks SET last_heartbeat = ?, updated_at = ?"
|
||
" WHERE id = ? AND status = 'running'",
|
||
(now, now, task_id),
|
||
)
|
||
conn.commit()
|
||
return cur.rowcount == 1
|
||
finally:
|
||
conn.close()
|
||
|
||
def _audit(self, project_id, action, detail=""):
|
||
"""best-effort 审计：失败不影响任务状态。"""
|
||
try:
|
||
store = AuthStore(self.db_path)
|
||
try:
|
||
store.audit("system", project_id, action, detail)
|
||
finally:
|
||
try:
|
||
store.conn.close()
|
||
except Exception:
|
||
pass
|
||
except Exception:
|
||
pass
|
||
|
||
# -- 认领 ----------------------------------------------------------
|
||
def _budget_exceeded(self):
|
||
return budget_exceeded(self.usage_turns, self.max_turns,
|
||
self.usage_cost_usd, self.max_cost_usd)
|
||
|
||
def pick_next(self):
|
||
"""下一个可派任务，或 None（无就绪任务 / 预算超限 / 依赖成环）。
|
||
|
||
依赖满足的定义：依赖任务状态为 completed 或 merged。
|
||
依赖成环时 ready_tasks 抛 ValueError，worker 视为无就绪任务
|
||
（成环任务本就永远派不出去，run_all 不应因此崩溃）。
|
||
"""
|
||
if self._budget_exceeded():
|
||
return None
|
||
tasks = self.load_tasks()
|
||
completed_ids = {t["id"] for t in tasks
|
||
if t["status"] in COMPLETED_STATUSES}
|
||
try:
|
||
ready = ready_tasks(tasks, completed_ids)
|
||
except ValueError:
|
||
return None
|
||
return ready[0] if ready else None
|
||
|
||
# -- 取消（server cancel 语义，见模块 docstring）-------------------
|
||
def cancel(self, task_id):
|
||
"""把 queued/running 任务标 interrupted；已终态或不存在返回 False。"""
|
||
task = self.get_task(task_id)
|
||
if task is None or task["status"] in TERMINAL_STATUSES:
|
||
return False
|
||
self._cancel_requested.add(task_id)
|
||
conn = _db_connect(self.db_path)
|
||
try:
|
||
cur = conn.execute(
|
||
"UPDATE tasks SET status = 'interrupted', updated_at = ?"
|
||
" WHERE id = ? AND status NOT IN"
|
||
" ('completed', 'failed', 'interrupted', 'merged',"
|
||
" 'merge-failed')",
|
||
(time.time(), task_id),
|
||
)
|
||
conn.commit()
|
||
claimed = cur.rowcount == 1
|
||
finally:
|
||
conn.close()
|
||
if claimed:
|
||
self._audit(task["project_id"], "worker_cancelled", task_id)
|
||
return claimed
|
||
|
||
def _was_cancelled(self, task_id):
|
||
if task_id in self._cancel_requested:
|
||
return True
|
||
task = self.get_task(task_id)
|
||
return task is not None and task["status"] == "interrupted"
|
||
|
||
# -- 缺 key 门 ------------------------------------------------------
|
||
def key_available(self):
|
||
"""只查存在性：名在环境里且值非空即 True；值本身绝不返回。"""
|
||
try:
|
||
value = self._env.get(self.env_key)
|
||
except Exception:
|
||
return False
|
||
return value is not None and value != ""
|
||
|
||
# -- 执行 ----------------------------------------------------------
|
||
def run_one(self, task=None):
|
||
"""认领并执行一个任务，返回结果 dict（见模块 docstring）。
|
||
|
||
未派发时返回 {"dispatched": False, "reason": ...} 且不动任务；
|
||
派发后（即使缺 key 失败）返回 dispatched=True 与最终状态。
|
||
"""
|
||
if self._budget_exceeded():
|
||
return {"dispatched": False, "reason": "budget_exceeded"}
|
||
if task is None:
|
||
task = self.pick_next()
|
||
if task is None:
|
||
return {"dispatched": False, "reason": "no_ready_task"}
|
||
task_id = task["id"]
|
||
if not self._claim(task_id):
|
||
return {"dispatched": False, "reason": "already_claimed",
|
||
"task_id": task_id}
|
||
task = self.get_task(task_id) or task
|
||
self._audit(task["project_id"], "worker_claim", task_id)
|
||
|
||
# 缺 key 门：认领之后、调 host.py 之前；缺 key 直接 failed，
|
||
# 不调 runner（假 CLI 单测可断言 runner 未被调用）。
|
||
if not self.key_available():
|
||
error = _missing_key_error(self.env_key)
|
||
self._set_status(task_id, "failed")
|
||
self._audit(task["project_id"], "worker_missing_key", error)
|
||
return {"dispatched": True, "task_id": task_id,
|
||
"host_name": self.executor._host_name(task)
|
||
if hasattr(self.executor, "_host_name") else task_id,
|
||
"status": "failed", "executor": "local",
|
||
"stdout": "", "error": error, "events": [],
|
||
"cost_record_id": None}
|
||
|
||
events = []
|
||
fallback_count = 0
|
||
result = None
|
||
error = None
|
||
try:
|
||
result = self.executor.dispatch_task(task)
|
||
combined = (result.get("stdout", "") or "") + \
|
||
(result.get("stderr", "") or "")
|
||
if result["status"] != "completed" and \
|
||
_ALREADY_EXISTS_MARKER in combined:
|
||
# host 任务已存在（上轮 worker 在 dispatch 后、写回前
|
||
# 崩溃）：改用 run 在已有任务上续跑，不重复建 worktree。
|
||
result = self.executor.run_task(task)
|
||
except Exception as exc: # runner 异常也如实写回 failed
|
||
result = {"executor": "local", "task_id": task_id,
|
||
"status": "failed", "stdout": "", "returncode": None,
|
||
"host_status": None, "stderr": ""}
|
||
error = str(exc)
|
||
|
||
# 取消优先：执行期间被 cancel 的任务保持 interrupted，
|
||
# 绝不用执行结果覆盖它（见模块 docstring 的取消说明）。
|
||
if self._was_cancelled(task_id):
|
||
self._set_status(task_id, "interrupted")
|
||
return {"dispatched": True, "task_id": task_id,
|
||
"host_name": self.executor._host_name(task)
|
||
if hasattr(self.executor, "_host_name") else task_id,
|
||
"status": "interrupted", "executor": "local",
|
||
"stdout": (result or {}).get("stdout", ""),
|
||
"error": "cancelled", "events": [],
|
||
"cost_record_id": None}
|
||
|
||
stdout = result.get("stdout", "") or ""
|
||
host_status = result.get("host_status")
|
||
final = result["status"]
|
||
if final != "completed":
|
||
if error is None:
|
||
error = stdout.strip() or "host execution failed"
|
||
# fallback 事件：只记、不自动重试（见模块 docstring）。
|
||
if self.router is not None and host_status is not None \
|
||
and is_fallbackable(host_status):
|
||
nxt = self.router.next_after(self.model, host_status)
|
||
if nxt is not None:
|
||
record = dict(task)
|
||
record["events"] = events
|
||
# record_fallback 本身会 append 到 events，
|
||
# 这里不要再 append 一次，否则事件会重复。
|
||
record_fallback(
|
||
record, self.model, nxt, str(host_status))
|
||
fallback_count = 1
|
||
self._audit(task["project_id"], "worker_fallback",
|
||
"%s -> %s" % (self.model, nxt))
|
||
self._set_status(task_id, final)
|
||
self.usage_turns += 1
|
||
# 成本：host 协议不暴露真实 token 数，用 cost.py 的估算结构
|
||
# 记 estimated，绝不标 actual；fallback_count 复用同一记录。
|
||
total = estimate_tokens(task.get("prompt") or "", stdout)
|
||
cost_id = self.costs.record(
|
||
task_id, task["project_id"], self.model,
|
||
input_tokens=total, output_tokens=0, usd_micros=None,
|
||
source="estimated", fallback_count=fallback_count)
|
||
self._audit(task["project_id"],
|
||
"worker_completed" if final == "completed"
|
||
else "worker_failed", task_id)
|
||
return {"dispatched": True, "task_id": task_id,
|
||
"host_name": self.executor._host_name(task)
|
||
if hasattr(self.executor, "_host_name") else task_id,
|
||
"status": final, "executor": "local",
|
||
"stdout": stdout, "error": error, "events": events,
|
||
"cost_record_id": cost_id}
|
||
|
||
def run_all(self, max_tasks=None):
|
||
"""循环认领执行，直到无就绪任务、预算超限或达到 max_tasks。
|
||
|
||
每轮都重新读库，所以上一轮刚 completed 的任务会解锁它的
|
||
依赖任务；每轮至少推进一个任务到终态，不会空转死循环。
|
||
"""
|
||
results = []
|
||
while max_tasks is None or len(results) < max_tasks:
|
||
result = self.run_one()
|
||
if not result["dispatched"]:
|
||
break
|
||
results.append(result)
|
||
return results
|
||
|
||
|
||
def main(argv=None):
|
||
"""单次派发循环：python3 -m cloud.worker --db <库> --host-py <路径>。
|
||
|
||
只打印任务结果摘要（状态、host 名、错误文本），绝不打印
|
||
任何环境变量的值。
|
||
"""
|
||
parser = argparse.ArgumentParser(description="cloud worker (one pass)")
|
||
parser.add_argument("--db", default="cloud-control-plane.db")
|
||
parser.add_argument("--host-py", default=None,
|
||
help="host.py 路径，缺省为 cloud/ 上一级的 host.py")
|
||
parser.add_argument("--config", default=None)
|
||
parser.add_argument("--repo", default=None)
|
||
parser.add_argument("--env-key", default=None)
|
||
parser.add_argument("--codex-config", default=None,
|
||
help="含 env_key 行的 codex 配置文件（只读名字）")
|
||
parser.add_argument("--model", default=None)
|
||
parser.add_argument("--max-turns", type=int, default=None)
|
||
parser.add_argument("--max-tasks", type=int, default=None)
|
||
args = parser.parse_args(argv)
|
||
host_py = args.host_py
|
||
if host_py is None:
|
||
here = os.path.dirname(os.path.abspath(__file__))
|
||
host_py = os.path.join(os.path.dirname(here), "host.py")
|
||
worker = Worker(args.db, host_py, config_path=args.config,
|
||
repo=args.repo, env_key=args.env_key,
|
||
codex_config_path=args.codex_config, model=args.model,
|
||
max_turns=args.max_turns)
|
||
try:
|
||
results = worker.run_all(max_tasks=args.max_tasks)
|
||
finally:
|
||
worker.close()
|
||
print(json.dumps(
|
||
[{"task_id": r.get("task_id"), "status": r.get("status"),
|
||
"reason": r.get("reason"), "error": r.get("error")}
|
||
for r in results],
|
||
indent=2, ensure_ascii=False))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|