host
26d9edac33
Worker polls the DB every interval (production 30s) and every poll ran SELECT * FROM tasks: full rows including result_text and error_text (up to 4000 chars each) for every task ever, growing with history. load_tasks now selects only the columns pick_next actually needs (id/project/prompt/status/depends/ priority/timestamps/executor) and only queued/completed/merged rows; running/failed rows are neither ready nor completed, so excluding them cannot change pick_next. get_task still does SELECT * for the single claimed task. Plus CREATE INDEX IF NOT EXISTS idx_tasks_status on tasks(status) in both SQLite init_db and the Postgres schema, so the poll filter stops being a full table scan on old databases too. Test: load_tasks returns only queued+completed out of queued/completed/running/failed.
1443 lines
70 KiB
Python
1443 lines
70 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 的失败（404/429/5xx/timeout，401/400
|
||
不算）。默认（未配 --fallback-models）时 worker 只记录
|
||
fallback 事件（record_fallback() 记 model_fallback 事件，
|
||
fallback_count 写进成本记录），不自动换模型重试——自动重试
|
||
等于未经批准烧 relay 额度。显式配了 --fallback-models 时，
|
||
第一次失败且可 fallback 就真的换列表里的第一个模型再跑一次
|
||
（最多 1 次，同一任务、同一 worktree/分支，dispatch 报
|
||
“任务已存在”时走现有的 run 续跑路径，模型经 host.py 的
|
||
顶层 --model 参数传过去，不改 host.py 签名）；第二次的结果
|
||
就是最终结果，actual_model 写成最后跑的模型，fallback_count=1。
|
||
审计：认领/完成/失败/取消/缺 key 各记一条 AuthStore 审计
|
||
（user_id="system"，与 server.sweep_stale_tasks 的做法一致），
|
||
审计失败不影响任务状态（best-effort）。
|
||
|
||
取消：worker 的 cancel() 走 server 的 cancel 语义：把
|
||
queued/running 任务直接标 interrupted（已终态的任务不动），
|
||
且正在跑的子进程会真杀——子进程以进程组启动（start_new_session），
|
||
取消时先 SIGTERM 给整个进程组、等 2 秒、仍活着再 SIGKILL，
|
||
杀完后任务保持 interrupted，执行结果不覆盖它。run_one() 在
|
||
执行返回后会先查任务是否已被取消，已被取消的一律保持
|
||
interrupted，不用执行结果覆盖它。
|
||
|
||
任务级超时与失败重试（生产级，见 Worker.__init__ 参数）：
|
||
- task_timeout（默认 1800 秒，CLI --task-timeout 可调）是
|
||
worker 级的每任务截止时间，与下面两个参数不要混淆：
|
||
--docker-timeout 是执行器级的单次 docker run 超时（Docker
|
||
Executor 自身的 timeout），--max-seconds 是 run_loop 的
|
||
整体墙钟上限（不是每任务）。三者互不覆盖：Docker 路径下
|
||
worker 取 min(task_timeout, docker_timeout)（docker_timeout
|
||
未设时取 task_timeout）作为本次 docker run 的超时，超时后
|
||
真杀容器、任务置 failed（超时是失败，不是取消，绝不标
|
||
interrupted）；Local popen 路径下 worker 在轮询循环里到
|
||
截止时间就杀进程组、同样置 failed；Local 同步路径（假
|
||
runner / 无超时的 subprocess.run）无法在中途抢占一个已阻塞
|
||
的调用，超时只能靠执行器自身的 timeout，这是诚实边界，
|
||
不要假装同步路径也能中途杀掉。
|
||
- max_retries（默认 1 次重试，CLI --max-retries 可调，0 表示
|
||
不重试）是对模型任务失败的有限重试：取消不重试；可 fallback
|
||
的失败优先走真实 fallback（见下），不再做同模型重试（fallback
|
||
本身就是模型级的重试，再同模型重试会重复烧额度）；其它失败
|
||
（含 worker 超时）重试前按 retry_backoff_seconds（默认 5 秒，
|
||
测试可注入 0）等待，重试次数记在审计（worker_retry）与
|
||
最终 error_text（[retries=N] 前缀）里，不另起新表；重试
|
||
耗尽后任务才最终 failed，成本记录与 usage_turns 只记一次
|
||
- 真实 fallback（CLI --fallback-models，默认空 = 关闭）：一次
|
||
尝试失败且可 fallback（404/429/5xx/timeout，401/400 不算，
|
||
取消不算）时，若配了 fallback 列表，就记 model_fallback 事件
|
||
与 worker_fallback 审计，然后换列表里的第一个模型再跑一次
|
||
（最多 1 次，与 max_retries 无关，max_retries=0 时 fallback
|
||
照样跑）；第二次的结果就是最终结果，actual_model 写成最后
|
||
跑的模型，fallback_count=1，两次都失败时 error_text 同时写清
|
||
两次原因；未配列表时保持老行为（只记事件、直接 failed）。
|
||
（按最终结果记，不按尝试次数重复记）。
|
||
- 容器清理兜底：Worker.cleanup_stale_containers() 列出
|
||
codex- 前缀的残留容器并逐个 kill+rm，失败只记日志不炸
|
||
worker；run_loop 启动时与 CLI 启动时各调一次（run_all /
|
||
__init__ 里不调，避免单测构造 Worker 时多余 I/O）。
|
||
|
||
可注入：runner（假 CLI，见 sandbox.LocalExecutor）与 env 映射
|
||
均可注入，单测不跑真实模型回合、不联网、不登录任何服务。
|
||
|
||
执行器选择：Worker 默认用 LocalExecutor；传入 DockerExecutor
|
||
（或 CLI --executor docker）时，派发走容器内 host.py，取消时
|
||
按任务派生的容器名 docker kill，不只改数据库状态；本机无
|
||
docker 时 DockerExecutor 在调 runner 前就明确报错，任务标
|
||
failed 且 error_text 含 docker 缺失，绝不静默降级为本地执行。
|
||
|
||
项目仓库：run_one 派发前查任务所属项目的 repo_url（projects
|
||
表），非空时经 cloud.repos.resolve_repo 在 repos_dir（CLI
|
||
--repos-dir，默认 ./repos）下 clone/fetch 该项目仓库，设
|
||
task["repo"] 再交给执行器；Local 与 Docker 执行器都已读
|
||
task["repo"]，同一条路径。repo_url 为空时走原有默认 repo；
|
||
解析失败时任务直接 failed，错误文本写清是仓库解析失败。
|
||
Docker 下 repos_dir 必须放在 extra-mounts 已挂载的目录下，
|
||
否则容器内 repo 路径不存在（见 cloud/README.md）。
|
||
推送到远端不在本轮范围：任务在本地检出上执行、本地 worktree
|
||
合并仍按 host.py 现有流程，git push 需要仓库写凭证，控制面
|
||
不存、不推。
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
import re
|
||
import signal
|
||
import socket
|
||
import subprocess
|
||
import sys
|
||
import time
|
||
from types import SimpleNamespace
|
||
|
||
from cloud.auth import AuthStore
|
||
from cloud.cost import CostStore, estimate_tokens
|
||
from cloud.db import connect as _db_connect
|
||
from cloud.logging_config import get_logger, log_event, setup_logging
|
||
from cloud.models import ModelRouter, is_fallbackable, record_fallback
|
||
from cloud.queue import budget_exceeded, ready_tasks
|
||
from cloud.repos import RepoError, resolve_repo
|
||
from cloud.sandbox import CloudExecutor, DockerExecutor, LocalExecutor
|
||
from cloud.sandbox import _default_runner as _sandbox_default_runner
|
||
from cloud.workers import heartbeat_worker, register_worker
|
||
|
||
# 与 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]
|
||
def _col(name, default=None):
|
||
try:
|
||
return row[name]
|
||
except (KeyError, IndexError, TypeError):
|
||
return default
|
||
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"],
|
||
# 老库行缺 executor 列时给 None（表示任意 worker 可认领），
|
||
# 不炸；列本身由 _ensure_result_columns 在写回时补。
|
||
"executor": _col("executor"),
|
||
}
|
||
|
||
|
||
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, process_factory=None,
|
||
task_timeout=1800, max_retries=1,
|
||
retry_backoff_seconds=5.0, repos_dir="./repos",
|
||
repos_runner=None, heartbeat_file=None,
|
||
worker_id=None, capacity=1, fallback_models=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
|
||
# 真实 fallback 换模型列表：默认空 = 关闭，只记事件不换模型
|
||
# （老行为，生产 worker 未配该参数时绝不突然多花模型钱）。
|
||
# 归一化：去空白、去重保序、去掉与主模型相同的项（换成同一
|
||
# 个模型不叫 fallback，只会白烧一次额度）；fallback 总次数
|
||
# 最多 1 次，所以运行期只用列表的第一个元素，后面的只是
|
||
# 备用声明，不会被依次试下去。
|
||
self.fallback_models = []
|
||
for _m in (fallback_models or []):
|
||
if not isinstance(_m, str):
|
||
continue
|
||
_m = _m.strip()
|
||
if _m and _m != self.model and _m not in self.fallback_models:
|
||
self.fallback_models.append(_m)
|
||
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()
|
||
# 真取消用：正在跑的子进程（Popen 对象）与任务 id。
|
||
# process_factory 可注入：单测用它返回一个假长跑子进程，
|
||
# 真实执行时走 LocalExecutor.popen_task（进程组启动）。
|
||
self.process_factory = process_factory
|
||
self._running_proc = None
|
||
self._running_task_id = None
|
||
self._running_container = None
|
||
self._loop_stop = False
|
||
self.logger = get_logger("cloud.worker")
|
||
# 任务级超时 / 失败重试（见模块 docstring 的生产级说明）。
|
||
self.task_timeout = task_timeout
|
||
self.max_retries = max_retries
|
||
self.retry_backoff_seconds = retry_backoff_seconds
|
||
# 项目仓库检出根目录：CLI --repos-dir 可调（默认 ./repos）；
|
||
# Docker 执行器下它必须放在 extra-mounts 已挂载的目录下，
|
||
# 否则容器内 task["repo"] 路径不存在——见 cloud/README.md。
|
||
# repos_runner 可注入（单测用假 git runner），缺省走真实 git。
|
||
self.repos_dir = repos_dir
|
||
self.repos_runner = repos_runner
|
||
# 运维心跳：路径优先显式参数，其次环境变量，默认当前目录下
|
||
# 的 worker-heartbeat.json；主机路径不在代码里硬编码，主机
|
||
# 部署时用 --heartbeat-file 或环境变量传入（见 deploy/README）。
|
||
# 文件里只有 last_seen 与 pid，绝无任何密钥。
|
||
self.heartbeat_file = heartbeat_file or \
|
||
os.environ.get("CLOUD_WORKER_HEARTBEAT_FILE") or \
|
||
"./worker-heartbeat.json"
|
||
# 多 worker 池身份：显式 --worker-id 优先（生产部署应配稳定
|
||
# id）；未给时用“主机名-进程号”，重启后进程号会变，旧行
|
||
# 自然变 stale，不要用随机 uuid 让注册表无法对应同一台机器。
|
||
self.worker_id = worker_id or "%s-%d" % (
|
||
socket.gethostname(), os.getpid())
|
||
try:
|
||
self.capacity = int(capacity)
|
||
except (TypeError, ValueError):
|
||
self.capacity = 1
|
||
|
||
def register(self):
|
||
"""向控制面注册本 worker（写 workers 表，失败只记日志）。
|
||
|
||
注册与心跳共用同一张表：心跳就是带 last_seen 刷新的注册，
|
||
未注册的 worker 发心跳时自动注册，不 404。
|
||
"""
|
||
try:
|
||
register_worker(
|
||
self.db_path, self.worker_id,
|
||
executor=self._executor_name(), capacity=self.capacity)
|
||
except Exception as exc:
|
||
log_event(self.logger, "worker_register_failed",
|
||
message="worker register failed: %s" % exc)
|
||
|
||
def _registry_heartbeat(self):
|
||
"""向控制面写注册表心跳（与文件心跳并行，失败只记日志）。"""
|
||
try:
|
||
heartbeat_worker(
|
||
self.db_path, self.worker_id,
|
||
executor=self._executor_name(), capacity=self.capacity)
|
||
except Exception as exc:
|
||
log_event(self.logger, "worker_registry_heartbeat_failed",
|
||
message="worker registry heartbeat failed: %s" % exc)
|
||
|
||
def _write_heartbeat(self):
|
||
"""写一次运维心跳：写失败只记日志，绝不炸 worker。
|
||
|
||
先写临时文件再 os.replace，保证 server 读到的要么是旧的
|
||
完整文件、要么是新的完整文件，不会读到写一半的 JSON。
|
||
"""
|
||
try:
|
||
payload = {"last_seen": time.time(), "pid": os.getpid()}
|
||
tmp = "%s.tmp" % self.heartbeat_file
|
||
parent = os.path.dirname(os.path.abspath(self.heartbeat_file))
|
||
os.makedirs(parent, exist_ok=True)
|
||
with open(tmp, "w", encoding="utf-8") as f:
|
||
f.write(json.dumps(payload))
|
||
os.replace(tmp, self.heartbeat_file)
|
||
except Exception as exc:
|
||
log_event(self.logger, "worker_heartbeat_failed",
|
||
message="worker heartbeat write failed: %s" % exc)
|
||
|
||
# -- 成本存储（惰性建表，避免只认领不执行时也建表）--------------
|
||
@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):
|
||
"""pick_next 的输入：只取 queued/completed/merged 三类任务。
|
||
|
||
旧实现是 SELECT * 全表：prompt/result_text/error_text
|
||
（后两者各可达 4000 字）全量进内存，且每轮 poll 都来一次；
|
||
任务表随时间变大后这是纯浪费。pick_next 只需要 queued
|
||
任务和已完成 id（依赖判断），running/failed/interrupted
|
||
的任务对它既不是就绪、也不是已完成，排除它们不影响结果。
|
||
列清单必须覆盖 _task_from_row 直接取的全部字段（prompt、
|
||
idempotency_key 等），缺一列就会 KeyError；result_text、
|
||
error_text 等重字段不在此列，真正需要它们的 get_task
|
||
仍是 SELECT * 单行取。
|
||
"""
|
||
conn = _db_connect(self.db_path)
|
||
try:
|
||
rows = conn.execute(
|
||
"SELECT id, project_id, prompt, status, idempotency_key,"
|
||
" depends_on, priority, created_by, created_at, updated_at,"
|
||
" last_heartbeat, executor FROM tasks"
|
||
" WHERE status IN ('queued', 'completed', 'merged')"
|
||
" 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, executor=None):
|
||
"""CAS 认领：仅当任务仍为 queued 时标 running 并记心跳。
|
||
|
||
executor 过滤（多 worker 池）：executor 给定时，WHERE 里
|
||
再加 (executor IS NULL OR executor = ?)，即只认领未指定
|
||
执行器、或指定为本 worker 执行器类型的任务；CAS 语义不变，
|
||
仍是“仅当 queued 时标 running”，两个同池 worker 抢同一
|
||
任务时只有一个能拿到 rowcount == 1。
|
||
"""
|
||
now = time.time()
|
||
conn = _db_connect(self.db_path)
|
||
try:
|
||
self._ensure_result_columns(conn)
|
||
if executor is None:
|
||
cur = conn.execute(
|
||
"UPDATE tasks SET status = 'running',"
|
||
" last_heartbeat = ?, updated_at = ?"
|
||
" WHERE id = ? AND status = 'queued'",
|
||
(now, now, task_id),
|
||
)
|
||
else:
|
||
cur = conn.execute(
|
||
"UPDATE tasks SET status = 'running',"
|
||
" last_heartbeat = ?, updated_at = ?"
|
||
" WHERE id = ? AND status = 'queued'"
|
||
" AND (executor IS NULL OR executor = ?)",
|
||
(now, now, task_id, executor),
|
||
)
|
||
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 _ensure_result_columns(self, conn):
|
||
"""worker 侧的迁移兜底：与 server.migrate_tasks_table 同列。
|
||
|
||
worker 的库理论上已由 server.init_db 迁过；这里再查一次
|
||
是为了 worker 直连老库（测试里手工建的表）时也不炸。
|
||
"""
|
||
try:
|
||
cols = {row[1] for row in conn.execute(
|
||
"PRAGMA table_info(tasks)").fetchall()}
|
||
except Exception:
|
||
return
|
||
for name, coltype in {
|
||
"result_text": "TEXT", "error_text": "TEXT",
|
||
"actual_model": "TEXT", "started_at": "REAL",
|
||
"finished_at": "REAL", "child_pid": "INTEGER",
|
||
"executor": "TEXT", "pr_url": "TEXT",
|
||
"pr_number": "INTEGER"}.items():
|
||
if name not in cols:
|
||
try:
|
||
conn.execute(
|
||
"ALTER TABLE tasks ADD COLUMN %s %s"
|
||
% (name, coltype))
|
||
except Exception:
|
||
pass
|
||
|
||
def _write_execution_result(self, task_id, result_text=None,
|
||
error_text=None, actual_model=None,
|
||
started_at=None, finished_at=None,
|
||
child_pid=None):
|
||
"""把执行结果写进 tasks 新列；老库缺列时先迁移再写。"""
|
||
conn = _db_connect(self.db_path)
|
||
try:
|
||
self._ensure_result_columns(conn)
|
||
if child_pid is None:
|
||
conn.execute(
|
||
"UPDATE tasks SET result_text = ?, error_text = ?,"
|
||
" actual_model = ?, started_at = ?, finished_at = ?,"
|
||
" updated_at = ? WHERE id = ?",
|
||
((result_text or "")[:4000] if result_text else None,
|
||
(error_text or "")[:4000] if error_text else None,
|
||
actual_model, started_at, finished_at,
|
||
time.time(), task_id),
|
||
)
|
||
else:
|
||
conn.execute(
|
||
"UPDATE tasks SET result_text = ?, error_text = ?,"
|
||
" actual_model = ?, started_at = ?, finished_at = ?,"
|
||
" child_pid = ?, updated_at = ? WHERE id = ?",
|
||
((result_text or "")[:4000] if result_text else None,
|
||
(error_text or "")[:4000] if error_text else None,
|
||
actual_model, started_at, finished_at, child_pid,
|
||
time.time(), task_id),
|
||
)
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
def _set_child_pid(self, task_id, pid):
|
||
conn = _db_connect(self.db_path)
|
||
try:
|
||
self._ensure_result_columns(conn)
|
||
conn.execute(
|
||
"UPDATE tasks SET child_pid = ? WHERE id = ?",
|
||
(pid, task_id),
|
||
)
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
def _kill_proc_group(self, proc):
|
||
"""真杀：先 SIGTERM 给整个进程组，等 2 秒，仍活着再 SIGKILL。
|
||
|
||
proc 须是 start_new_session 启动的 Popen（自成进程组），
|
||
杀组才能连同它拉起的子孙一起杀掉，不只杀父进程。
|
||
"""
|
||
try:
|
||
pgid = os.getpgid(proc.pid)
|
||
except Exception:
|
||
return
|
||
try:
|
||
os.killpg(pgid, signal.SIGTERM)
|
||
except (ProcessLookupError, PermissionError, OSError):
|
||
return
|
||
deadline = time.time() + 2.0
|
||
while time.time() < deadline:
|
||
try:
|
||
os.killpg(pgid, 0)
|
||
except (ProcessLookupError, PermissionError, OSError):
|
||
return
|
||
time.sleep(0.05)
|
||
try:
|
||
os.killpg(pgid, signal.SIGKILL)
|
||
except (ProcessLookupError, PermissionError, OSError):
|
||
pass
|
||
|
||
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
|
||
# 多 worker 池过滤：只取未指定执行器、或指定为本 worker
|
||
# 执行器类型的任务；_claim 里还有一道 WHERE 过滤兜底，
|
||
# pick_next 这里是第一道，避免 worker 老去抢别人的任务。
|
||
mine = self._executor_name()
|
||
ready = [t for t in ready
|
||
if t.get("executor") in (None, "", mine)]
|
||
return ready[0] if ready else None
|
||
|
||
# -- 取消（server cancel 语义，见模块 docstring）-------------------
|
||
def cancel(self, task_id):
|
||
"""把 queued/running 任务标 interrupted；已终态或不存在返回 False。
|
||
|
||
正在跑的子进程要真杀：先杀进程组（见 _kill_proc_group），
|
||
再标状态；已终态或不存在返回 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)
|
||
if self._running_task_id == task_id and self._running_proc is not None:
|
||
self._kill_proc_group(self._running_proc)
|
||
try:
|
||
self._running_proc.wait(timeout=3)
|
||
except Exception:
|
||
pass
|
||
self._running_proc = None
|
||
self._running_task_id = None
|
||
# Docker 取消：按任务派生的容器名真杀，不只改数据库状态。
|
||
# 同步路径（runner 阻塞在 docker run 里）下 _running_proc
|
||
# 为空，此时靠 _running_container 杀容器；kill 尽力而为，
|
||
# 失败也不影响后续的状态标记。
|
||
if self._running_task_id == task_id and \
|
||
self._running_container is not None and \
|
||
hasattr(self.executor, "kill_container"):
|
||
try:
|
||
self.executor.kill_container(self._running_container)
|
||
except Exception:
|
||
pass
|
||
self._running_container = None
|
||
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 not claimed:
|
||
# 并发竞态：run_one 的轮询已看到本次取消的标记、先一步
|
||
# 把任务标成 interrupted——中断确实由本次取消触发，返回 True。
|
||
current = self.get_task(task_id)
|
||
if current is not None and current["status"] == "interrupted":
|
||
claimed = True
|
||
if claimed:
|
||
self._audit(task["project_id"], "worker_cancelled", task_id)
|
||
log_event(self.logger, "worker_cancelled", task_id=task_id,
|
||
project_id=task["project_id"], status="interrupted")
|
||
return claimed
|
||
|
||
def cleanup_stale_containers(self):
|
||
"""启动清理兜底：清掉上一轮残留的 codex- 容器。
|
||
|
||
只在执行器有 cleanup_stale_containers 时做（Docker）；
|
||
列出或删除失败只记日志，绝不抛异常炸掉 worker。
|
||
返回已清理的容器名列表（非 Docker 执行器返回 []）。
|
||
"""
|
||
cleanup = getattr(self.executor, "cleanup_stale_containers", None)
|
||
if cleanup is None:
|
||
return []
|
||
try:
|
||
cleaned = cleanup()
|
||
except Exception as exc:
|
||
log_event(self.logger, "worker_cleanup_failed",
|
||
message="cleanup stale containers failed: %s" % exc)
|
||
return []
|
||
if cleaned:
|
||
log_event(self.logger, "worker_cleanup",
|
||
message="cleaned stale containers: %s"
|
||
% ",".join(cleaned))
|
||
return cleaned
|
||
|
||
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 _executor_name(self):
|
||
"""执行器名：DockerExecutor -> docker，CloudExecutor -> cloud，
|
||
其它 -> local。cloud 绝不能回落成 local，否则 cloud worker
|
||
会认领 local 池的任务，池过滤就失效了。"""
|
||
if isinstance(self.executor, DockerExecutor):
|
||
return "docker"
|
||
if isinstance(self.executor, CloudExecutor):
|
||
return "cloud"
|
||
return "local"
|
||
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 _use_popen(self):
|
||
"""是否走进程模式：注入了 process_factory，或执行器是
|
||
默认 runner 的 LocalExecutor（真实 subprocess）。"""
|
||
if self.process_factory is not None:
|
||
return True
|
||
return (isinstance(self.executor, LocalExecutor)
|
||
and getattr(self.executor, "runner", None)
|
||
is _sandbox_default_runner
|
||
and hasattr(self.executor, "popen_task"))
|
||
|
||
def _execute_via_process(self, task, kind="dispatch",
|
||
deadline_seconds=None):
|
||
"""进程模式执行一轮：启动、轮询、取消时杀进程组。
|
||
|
||
返回与同步路径同款 result dict（status/stdout/stderr/
|
||
host_status/returncode）。取消时返回 status='cancelled'
|
||
的标记，调用方据此保持 interrupted；deadline_seconds
|
||
到期时杀进程组并返回 status='timeout' 的标记，调用方
|
||
据此置 failed（超时是失败，不是取消）。
|
||
"""
|
||
task_id = task["id"]
|
||
if self.process_factory is not None:
|
||
proc = self.process_factory(task)
|
||
else:
|
||
proc = self.executor.popen_task(task, kind=kind)
|
||
self._running_proc = proc
|
||
self._running_task_id = task_id
|
||
try:
|
||
self._set_child_pid(task_id, getattr(proc, "pid", None))
|
||
except Exception:
|
||
pass
|
||
cancelled = False
|
||
timed_out = False
|
||
started = time.time()
|
||
try:
|
||
while proc.poll() is None:
|
||
if task_id in self._cancel_requested or self._loop_stop:
|
||
cancelled = True
|
||
self._kill_proc_group(proc)
|
||
break
|
||
if deadline_seconds is not None and \
|
||
time.time() - started >= deadline_seconds:
|
||
timed_out = True
|
||
self._kill_proc_group(proc)
|
||
break
|
||
time.sleep(0.05)
|
||
try:
|
||
out, err = proc.communicate(timeout=5)
|
||
except Exception:
|
||
try:
|
||
proc.kill()
|
||
except Exception:
|
||
pass
|
||
out, err = proc.communicate()
|
||
finally:
|
||
self._running_proc = None
|
||
self._running_task_id = None
|
||
if cancelled:
|
||
return {"executor": "local", "task_id": task_id,
|
||
"status": "cancelled", "stdout": out or "",
|
||
"stderr": err or "", "returncode": proc.returncode,
|
||
"host_status": None}
|
||
if timed_out:
|
||
return {"executor": "local", "task_id": task_id,
|
||
"status": "timeout", "stdout": out or "",
|
||
"stderr": err or "", "returncode": proc.returncode,
|
||
"host_status": None}
|
||
fake = SimpleNamespace(returncode=proc.returncode,
|
||
stdout=out or "", stderr=err or "")
|
||
if hasattr(self.executor, "_result_from_runner"):
|
||
return self.executor._result_from_runner(task_id, fake)
|
||
return {"executor": "local", "task_id": task_id,
|
||
"status": "completed" if proc.returncode == 0 else "failed",
|
||
"stdout": out or "", "stderr": err or "",
|
||
"returncode": proc.returncode, "host_status": None}
|
||
|
||
# -- 项目仓库 ------------------------------------------------------
|
||
def _project_repo_url(self, project_id):
|
||
"""查任务所属项目的 repo_url；查不到项目、列缺失、表缺失时
|
||
一律返回 None（走原有默认 repo，不报错、不炸掉派发）。
|
||
repo_url 本身只在 projects 表里，绝不从任务 prompt 里猜。"""
|
||
try:
|
||
conn = _db_connect(self.db_path)
|
||
try:
|
||
row = conn.execute(
|
||
"SELECT repo_url FROM projects WHERE id = ?",
|
||
(project_id,)).fetchone()
|
||
finally:
|
||
conn.close()
|
||
except Exception:
|
||
return None
|
||
if row is None:
|
||
return None
|
||
try:
|
||
url = row["repo_url"]
|
||
except (KeyError, IndexError, TypeError):
|
||
return None
|
||
return url or None
|
||
|
||
def _apply_project_repo(self, task):
|
||
"""派发前按项目 repo_url 解析仓库，返回错误文本或 None。
|
||
|
||
repo_url 为空：不动 task，返回 None（走原有默认 repo）。
|
||
repo_url 非空：调 cloud.repos.resolve_repo 得到本地检出，
|
||
设 task["repo"] 再交给执行器——LocalExecutor 与
|
||
DockerExecutor 的 dispatch_task 都已读 task["repo"]，同一条
|
||
路径，不另起分支。resolve 失败返回错误文本，调用方据此把
|
||
任务置 failed（仓库解析失败绝不标 completed，也不在此重试，
|
||
避免对不可用的仓库反复 clone 烧时间）。
|
||
"""
|
||
repo_url = self._project_repo_url(task.get("project_id"))
|
||
if not repo_url:
|
||
return None
|
||
try:
|
||
task["repo"] = resolve_repo(
|
||
repo_url, self.repos_dir, task.get("project_id"),
|
||
runner=self.repos_runner)
|
||
except (RepoError, ValueError) as exc:
|
||
return "仓库解析失败: %s" % exc
|
||
except Exception as exc:
|
||
return "仓库解析失败: %s" % exc
|
||
return None
|
||
|
||
# -- 执行 ----------------------------------------------------------
|
||
def _run_attempt(self, task):
|
||
"""单次执行尝试（不含认领、key 门、重试、成本与最终审计）。
|
||
|
||
返回 (result, error)：result 是执行器结果 dict，error 是
|
||
runner 异常时的文本（无异常时 None）。Docker 路径下把
|
||
worker 级 task_timeout 与执行器级 docker_timeout 取较小
|
||
值作为本次 docker run 超时（docker_timeout 未设时取
|
||
task_timeout）；Local popen 路径下把 task_timeout 作为
|
||
轮询截止时间；Local 同步路径无法中途抢占，超时只能靠
|
||
执行器自身 timeout（见模块 docstring 的诚实边界）。
|
||
调用方负责：取消优先、超时转 failed、重试、成本与审计。
|
||
"""
|
||
task_id = task["id"]
|
||
result = None
|
||
error = None
|
||
# Docker 同步路径：dispatch 阻塞在 docker run 里时，取消
|
||
# 线程靠 _running_container 按名杀容器（见 cancel()）。
|
||
if hasattr(self.executor, "container_name"):
|
||
try:
|
||
self._running_container = \
|
||
self.executor.container_name(task)
|
||
self._running_task_id = task_id
|
||
except Exception:
|
||
self._running_container = None
|
||
# Docker 超时取较小值：两个超时都不应静默延长另一个。
|
||
saved_docker_timeout = None
|
||
is_docker = isinstance(self.executor, DockerExecutor)
|
||
if is_docker and self.task_timeout is not None:
|
||
saved_docker_timeout = self.executor.timeout
|
||
if saved_docker_timeout is None:
|
||
self.executor.timeout = self.task_timeout
|
||
else:
|
||
try:
|
||
self.executor.timeout = min(
|
||
float(self.task_timeout),
|
||
float(saved_docker_timeout))
|
||
except (TypeError, ValueError):
|
||
self.executor.timeout = self.task_timeout
|
||
try:
|
||
if self._use_popen():
|
||
result = self._execute_via_process(
|
||
task, deadline_seconds=self.task_timeout)
|
||
combined = (result.get("stdout", "") or "") + \
|
||
(result.get("stderr", "") or "")
|
||
if result["status"] not in (
|
||
"completed", "cancelled", "timeout") and \
|
||
_ALREADY_EXISTS_MARKER in combined:
|
||
result = self._execute_via_process(
|
||
task, kind="run",
|
||
deadline_seconds=self.task_timeout)
|
||
else:
|
||
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:
|
||
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)
|
||
finally:
|
||
if is_docker:
|
||
self.executor.timeout = saved_docker_timeout
|
||
# 同步路径执行结束：清掉取消用的容器标记（取消路径里
|
||
# 若已清过，这里再清一次无害）。
|
||
self._running_container = None
|
||
if self._running_task_id == task_id and \
|
||
self._running_proc is None:
|
||
self._running_task_id = None
|
||
# Docker 超时标记：执行器已在超时时 kill+rm 容器并返回
|
||
# failed（stdout 含 timeout），这里把状态统一成 timeout，
|
||
# 调用方据此置 failed 并记 worker_timeout 审计。
|
||
if result is not None and result.get("status") == "failed" and \
|
||
"docker run timeout" in (result.get("stdout", "") or ""):
|
||
result = dict(result, status="timeout")
|
||
return result, error
|
||
|
||
def _attempt_error(self, result, error, token):
|
||
"""一次尝试失败的可读原因（供 fallback 两次原因拼接用）。
|
||
|
||
超时用统一的超时文本；runner 异常用异常文本；同步路径的
|
||
stdout 是 host.py 的 JSON 时取它的 text 字段（比整段 JSON
|
||
可读）；最后退到 host_status token 或通用文本。绝不返回
|
||
任何密钥（stdout 里本就不应含密钥，text 字段同样不含）。
|
||
"""
|
||
if result is not None and result.get("status") == "timeout":
|
||
return "task timeout after %s seconds" % self.task_timeout
|
||
if error:
|
||
return error
|
||
stdout = (result or {}).get("stdout", "") or ""
|
||
try:
|
||
parsed = json.loads(stdout)
|
||
except (TypeError, ValueError):
|
||
parsed = None
|
||
if isinstance(parsed, dict) and parsed.get("text"):
|
||
return str(parsed["text"])
|
||
# dispatch 输出现在带 error 字段（host.py 从 thread_history
|
||
# 库里捞回的失败原因，如 404）；text 为空时优先用它，
|
||
# 否则两次失败的拼接里只剩 host_status=404 这种代号。
|
||
if isinstance(parsed, dict) and parsed.get("error"):
|
||
err = parsed["error"]
|
||
if isinstance(err, dict):
|
||
err = err.get("message") or str(err)
|
||
return str(err)[:1000]
|
||
# host.py 崩溃时 stdout 常为空、真正的报错在 stderr（例如
|
||
# “缺少 --repo”、Read-only file system）；只用 stdout 时
|
||
# 最终错误会被写成通用的 host execution failed，排障无从
|
||
# 下手，所以 stdout 为空时取 stderr 尾部（实际报错行在尾）。
|
||
stderr = (result or {}).get("stderr", "") or ""
|
||
if not stdout.strip() and stderr.strip():
|
||
return stderr.strip()[-1000:]
|
||
if token is not None:
|
||
return "host_status=%s" % token
|
||
return stdout.strip() or "host execution failed"
|
||
|
||
def _fallback_token(self, result, error):
|
||
"""一次失败尝试的可 fallback 信号。
|
||
|
||
host_status 的真实含义是 host.py 回合状态（completed /
|
||
failed / timeout / interrupted），不是 HTTP 状态码——
|
||
真实世界里 host.py 从不返回 status=\"404\"，只有单测里的
|
||
假 runner 才会这么写。所以这里不能把 host_status 原样
|
||
当 token：host_status 本身已是可 fallback 码（404/429/
|
||
5xx/timeout，含历史假 runner 形态）时直接返回；否则
|
||
（最常见：host_status=\"failed\"）从回合 JSON 的 text/error
|
||
字段、stdout、stderr、runner 异常里找显式 HTTP 状态码
|
||
或“model not found”类文本（按 404 计），找到且可
|
||
fallback 才返回；找不到显式信号时返回 None，绝不凭
|
||
“host execution failed”这类通用文本猜测，否则任何执行
|
||
器故障都会被误当成模型故障换模型烧额度。thread/start
|
||
崩溃（host.py 非 0 退出、不打印 JSON）时 host_status
|
||
为 None，同样走这条提取路径。
|
||
"""
|
||
host_status = (result or {}).get("host_status")
|
||
if host_status is not None and is_fallbackable(host_status):
|
||
return host_status
|
||
combined_parts = [
|
||
(result or {}).get("stdout", "") or "",
|
||
(result or {}).get("stderr", "") or "",
|
||
error or "",
|
||
]
|
||
combined = " ".join(combined_parts)
|
||
if not combined:
|
||
return None
|
||
lowered = combined.lower()
|
||
if "model not found" in lowered or "unknown model" in lowered \
|
||
or "no such model" in lowered:
|
||
return "404"
|
||
match = re.search(
|
||
r"(?:http/|status[\"':\s]{0,10}|code[\"':\s]{0,10}|error[\"':\s]{0,10})"
|
||
r"(404|429|500|502|503|504)\b",
|
||
combined, re.IGNORECASE)
|
||
if match:
|
||
return match.group(1)
|
||
match = re.search(r"\b(404|429|500|502|503|504)\b", combined)
|
||
if match:
|
||
return match.group(1)
|
||
return None
|
||
|
||
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, executor=self._executor_name()):
|
||
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",
|
||
"%s executor=%s" % (task_id, self._executor_name()))
|
||
# log_event 只接受 task_id/project_id/status/message 四个
|
||
# 字段，executor 只进审计 detail，不进日志调用（否则
|
||
# TypeError 会炸掉整个 run_one）。
|
||
log_event(self.logger, "worker_claim", task_id=task_id,
|
||
project_id=task["project_id"], status="running")
|
||
started_at = time.time()
|
||
|
||
# 缺 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._write_execution_result(
|
||
task_id, error_text=error, actual_model=self.model,
|
||
started_at=started_at, finished_at=time.time())
|
||
self._audit(task["project_id"], "worker_missing_key", error)
|
||
log_event(self.logger, "worker_missing_key", task_id=task_id,
|
||
project_id=task["project_id"], status="failed")
|
||
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": self._executor_name(),
|
||
"stdout": "", "error": error, "events": [],
|
||
"cost_record_id": None}
|
||
|
||
# 项目仓库解析：认领 + 缺 key 门之后、真正执行之前。
|
||
# repo_url 为空时不动（默认 repo）；解析失败时任务直接 failed，
|
||
# 错误文本写清是仓库解析失败，绝不标 completed，也不在此重试。
|
||
repo_error = self._apply_project_repo(task)
|
||
if repo_error is not None:
|
||
self._set_status(task_id, "failed")
|
||
self._write_execution_result(
|
||
task_id, error_text=repo_error, actual_model=self.model,
|
||
started_at=started_at, finished_at=time.time())
|
||
self._audit(task["project_id"], "worker_repo_failed", repo_error)
|
||
log_event(self.logger, "worker_repo_failed", task_id=task_id,
|
||
project_id=task["project_id"], status="failed")
|
||
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": self._executor_name(),
|
||
"stdout": "", "error": repo_error, "events": [],
|
||
"cost_record_id": None}
|
||
|
||
events = []
|
||
fallback_count = 0
|
||
result = None
|
||
error = None
|
||
retry_count = 0
|
||
# 最后真正跑过的模型：主模型起手，fallback 后换成 fallback 模型。
|
||
actual_model = self.model
|
||
fallback_model = None
|
||
first_error = None
|
||
# 模型经任务的 model 字段传给执行器（执行器再经 host.py 的
|
||
# 顶层 --model 参数传给 dispatch/run）；不设时 host.py 用
|
||
# 配置文件里的 model。这里显式设成 worker 的主模型，让
|
||
# --model 覆盖配置真正生效，fallback 时再换成新模型。
|
||
task = dict(task)
|
||
task["model"] = self.model
|
||
# 有限重试与真实 fallback 的优先级：取消不重试；一次尝试失败
|
||
# 且可 fallback、且显式配了 fallback_models 时，优先真换模型
|
||
# 再跑一次（最多 1 次，同一任务、同一 worktree/分支，dispatch
|
||
# 报“任务已存在”时 _run_attempt 内部走现有的 run 续跑路径），
|
||
# 第二次的结果就是最终结果，不再做同模型重试；未配
|
||
# fallback_models 时保持老行为（router 给出下一模型时只记事件
|
||
# 不重试）；其它失败（含 worker 超时）按 max_retries 做同模型
|
||
# 重试，退避等待后重跑一次 _run_attempt。
|
||
for attempt in range(self.max_retries + 1):
|
||
result, error = self._run_attempt(task)
|
||
if self._was_cancelled(task_id) or \
|
||
result.get("status") == "cancelled":
|
||
break
|
||
if result.get("status") == "completed":
|
||
break
|
||
host_status_try = self._fallback_token(result, error)
|
||
if host_status_try is None and \
|
||
result.get("status") == "timeout":
|
||
# 进程/容器超时路径的 host_status 为空，超时本身就是
|
||
# 可 fallback 的失败 token（models.FALLBACKABLE 含它）。
|
||
host_status_try = "timeout"
|
||
if fallback_model is None and self.fallback_models \
|
||
and host_status_try is not None \
|
||
and is_fallbackable(host_status_try):
|
||
fb = self.fallback_models[0]
|
||
record = dict(task)
|
||
record["events"] = events
|
||
# record_fallback 本身会 append 到 events，
|
||
# 这里不要再 append 一次，否则事件会重复。
|
||
record_fallback(
|
||
record, self.model, fb, str(host_status_try))
|
||
fallback_count = 1
|
||
self._audit(task["project_id"], "worker_fallback",
|
||
"%s -> %s" % (self.model, fb))
|
||
first_error = self._attempt_error(
|
||
result, error, host_status_try)
|
||
fallback_model = fb
|
||
actual_model = fb
|
||
task = dict(task)
|
||
task["model"] = fb
|
||
result, error = self._run_attempt(task)
|
||
break
|
||
if self.router is not None and host_status_try is not None \
|
||
and is_fallbackable(host_status_try) and \
|
||
self.router.next_after(
|
||
self.model, host_status_try) is not None:
|
||
break
|
||
if attempt < self.max_retries:
|
||
retry_count += 1
|
||
self._audit(
|
||
task["project_id"], "worker_retry",
|
||
"%s: attempt %d failed, retry %d/%d"
|
||
% (task_id, attempt + 1, retry_count,
|
||
self.max_retries))
|
||
if self.retry_backoff_seconds:
|
||
time.sleep(self.retry_backoff_seconds)
|
||
continue
|
||
break
|
||
# 超时转 failed：超时是失败，不是取消；最终状态必须是
|
||
# VALID_STATUSES 里的 failed，绝不能把 "timeout" 写进库。
|
||
if result is not None and result.get("status") == "timeout":
|
||
result = dict(result, status="failed")
|
||
error = "task timeout after %s seconds" % self.task_timeout
|
||
self._audit(task["project_id"], "worker_timeout", task_id)
|
||
|
||
# 取消优先：执行期间被 cancel 的任务保持 interrupted，
|
||
# 绝不用执行结果覆盖它（见模块 docstring 的取消说明）。
|
||
if self._was_cancelled(task_id) or result.get("status") == "cancelled":
|
||
self._set_status(task_id, "interrupted")
|
||
self._write_execution_result(
|
||
task_id, result_text=(result or {}).get("stdout", ""),
|
||
error_text="cancelled", actual_model=actual_model,
|
||
started_at=started_at, finished_at=time.time())
|
||
log_event(self.logger, "worker_cancelled", task_id=task_id,
|
||
project_id=task["project_id"], status="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": self._executor_name(),
|
||
"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 fallback_model is not None:
|
||
# 真 fallback 后第二次也失败：错误文本必须同时写清
|
||
# 两次原因，不能只留第二次的（超时转换已在前面把
|
||
# error 置成超时文本，那时 second_error 就是它）。
|
||
second_error = error or self._attempt_error(
|
||
result, None, host_status)
|
||
error = ("primary model %s failed: %s; "
|
||
"fallback model %s failed: %s"
|
||
% (self.model, first_error,
|
||
fallback_model, second_error))
|
||
else:
|
||
if error is None:
|
||
stderr = (result or {}).get("stderr", "") or ""
|
||
error = stdout.strip() or \
|
||
(stderr.strip()[-1000:] if stderr.strip() else "") \
|
||
or "host execution failed"
|
||
if retry_count:
|
||
error = "[retries=%d] %s" % (retry_count, error)
|
||
# 未配 fallback_models 时的老行为：只记事件、不重试。
|
||
# 注意用提取出的 token（真实回合失败时 host_status
|
||
# 只是 "failed"，真正的 404 藏在回合 JSON 里），
|
||
# 不能直接用原始 host_status 判断。
|
||
_token = self._fallback_token(result, error)
|
||
if self.router is not None and _token is not None \
|
||
and is_fallbackable(_token):
|
||
nxt = self.router.next_after(self.model, _token)
|
||
if nxt is not None:
|
||
record = dict(task)
|
||
record["events"] = events
|
||
# record_fallback 本身会 append 到 events，
|
||
# 这里不要再 append 一次，否则事件会重复。
|
||
record_fallback(
|
||
record, self.model, nxt, str(_token))
|
||
fallback_count = 1
|
||
self._audit(task["project_id"], "worker_fallback",
|
||
"%s -> %s" % (self.model, nxt))
|
||
self._set_status(task_id, final)
|
||
finished_at = time.time()
|
||
self._write_execution_result(
|
||
task_id, result_text=stdout, error_text=error,
|
||
actual_model=actual_model, started_at=started_at,
|
||
finished_at=finished_at)
|
||
self.usage_turns += 1
|
||
# 成本：host 协议不暴露真实 token 数，用 cost.py 的估算结构
|
||
# 记 estimated，绝不标 actual；fallback_count 复用同一记录；
|
||
# 模型记最后真正跑的那个（fallback 成功时就是 fallback 模型）。
|
||
total = estimate_tokens(task.get("prompt") or "", stdout)
|
||
cost_id = self.costs.record(
|
||
task_id, task["project_id"], actual_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)
|
||
log_event(self.logger,
|
||
"worker_completed" if final == "completed"
|
||
else "worker_failed",
|
||
task_id=task_id, project_id=task["project_id"],
|
||
status=final)
|
||
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": self._executor_name(),
|
||
"stdout": stdout, "error": error, "events": events,
|
||
"actual_model": actual_model,
|
||
"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 run_loop(self, interval=300, max_seconds=None, stop_event=None):
|
||
"""常驻循环：每 interval 秒跑一轮 run_all，空闲也不退出。
|
||
|
||
退出条件：stop_event 被置位、_loop_stop 被信号处理置位、
|
||
或 max_seconds 到期。退出时若有正在跑的任务（进程模式下
|
||
_running_task_id 非空），先杀进程组、再标 interrupted，
|
||
不让任务永远卡在 running。
|
||
"""
|
||
# 常驻启动清理兜底：清掉上一轮残留的 codex- 容器，
|
||
# 失败只记日志（见 cleanup_stale_containers），不炸循环。
|
||
self.cleanup_stale_containers()
|
||
# 运维心跳：启动时先写一次，之后每轮结束再写，server 的
|
||
# /health 据此算 worker_ok（见 server.py 的健康检查）。
|
||
self._write_heartbeat()
|
||
# 多 worker 池注册：启动时注册一次，之后每轮结束写注册表
|
||
# 心跳（见本轮结束处）；注册失败只记日志，不炸循环。
|
||
self.register()
|
||
started = time.time()
|
||
all_results = []
|
||
try:
|
||
while not self._loop_stop:
|
||
if stop_event is not None and stop_event.is_set():
|
||
break
|
||
if max_seconds is not None and \
|
||
time.time() - started >= max_seconds:
|
||
break
|
||
try:
|
||
results = self.run_all()
|
||
except Exception as exc:
|
||
# 常驻循环一轮炸了不带走整个 worker：记结构化日志，
|
||
# 睡完本轮间隔再试下一轮（缺表、锁冲突都属此类）。
|
||
log_event(self.logger, "worker_loop_error",
|
||
message="worker_loop_error: %s" % exc)
|
||
results = []
|
||
all_results.extend(results)
|
||
# 每轮结束写心跳：哪怕本轮一个任务都没派，worker 活着
|
||
# 这件事本身也要被 server 看到（文件心跳 + 注册表心跳）。
|
||
self._write_heartbeat()
|
||
self._registry_heartbeat()
|
||
# 按 0.2 秒切片睡，期间可响应 stop_event/_loop_stop。
|
||
deadline = time.time() + interval
|
||
while time.time() < deadline:
|
||
if self._loop_stop:
|
||
break
|
||
if stop_event is not None and stop_event.is_set():
|
||
break
|
||
if max_seconds is not None and \
|
||
time.time() - started >= max_seconds:
|
||
break
|
||
time.sleep(0.2)
|
||
finally:
|
||
if self._running_task_id is not None:
|
||
task_id = self._running_task_id
|
||
if self._running_proc is not None:
|
||
self._kill_proc_group(self._running_proc)
|
||
try:
|
||
self._running_proc.wait(timeout=3)
|
||
except Exception:
|
||
pass
|
||
self._running_proc = None
|
||
self._running_task_id = None
|
||
self._set_status(task_id, "interrupted")
|
||
return all_results
|
||
|
||
|
||
def _build_cli_executor(args, host_py):
|
||
"""CLI 的执行器构造：local 时返回 None（Worker 默认 LocalExecutor），
|
||
docker 时返回 DockerExecutor（env_key 名用 Worker 的解析规则，
|
||
值本身绝不从这里传——DockerExecutor 只查存在性，-e 只传名），
|
||
cloud 时返回 CloudExecutor（通用 HTTP runner，未配 endpoint
|
||
时执行即明确失败，绝不降级）。"""
|
||
if getattr(args, "executor", "local") == "cloud":
|
||
return CloudExecutor()
|
||
if getattr(args, "executor", "local") != "docker":
|
||
return None
|
||
return DockerExecutor(
|
||
host_py, config_path=args.config, image=args.docker_image,
|
||
cpus=args.docker_cpus, memory=args.docker_memory,
|
||
user=args.docker_user, worktrees_root=args.docker_worktrees_root,
|
||
env_key=resolve_env_key(args.env_key, args.codex_config),
|
||
timeout=args.docker_timeout,
|
||
extra_mounts=args.docker_extra_mounts)
|
||
|
||
|
||
def main(argv=None):
|
||
"""worker CLI：默认单次 run_all；--loop 进常驻循环。
|
||
|
||
只打印任务结果摘要（状态、host 名、错误文本），绝不打印
|
||
任何环境变量的值。
|
||
"""
|
||
parser = argparse.ArgumentParser(description="cloud worker")
|
||
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("--heartbeat-file", default=None,
|
||
help="运维心跳文件路径（默认：环境变量 "
|
||
"CLOUD_WORKER_HEARTBEAT_FILE，未设则为 "
|
||
"./worker-heartbeat.json；主机部署时用此"
|
||
"参数指到 /var/lib/codex-app-server/ 下，"
|
||
"与 server 的 health 读取路径保持一致）")
|
||
parser.add_argument("--repos-dir", default="./repos",
|
||
help="项目仓库检出根目录（默认 ./repos）；项目有"
|
||
" repo_url 时 worker 在此 clone/fetch 各项目"
|
||
"仓库；Docker 执行器下它必须放在"
|
||
" --docker-extra-mounts 已挂载的目录下，"
|
||
"否则容器内 task repo 路径不存在")
|
||
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("--fallback-models", default=None,
|
||
help="真实 fallback 换模型列表，逗号分隔"
|
||
"（默认空 = 关闭，只记事件不换模型；"
|
||
"配了时主模型可 fallback 失败后最多换"
|
||
"列表里的第一个模型再跑一次，401/400"
|
||
"绝不 fallback，取消绝不 fallback）")
|
||
parser.add_argument("--max-turns", type=int, default=None)
|
||
parser.add_argument("--max-tasks", type=int, default=None)
|
||
parser.add_argument("--executor", default="local",
|
||
choices=["local", "docker", "cloud"],
|
||
help="执行器：local 调 host.py 子进程；docker"
|
||
" 调 DockerExecutor（无 docker 时任务"
|
||
"明确失败，不静默降级）；cloud 调通用"
|
||
" HTTP runner（需 CLOUD_RUNNER_ENDPOINT）")
|
||
parser.add_argument("--worker-id", default=None,
|
||
help="多 worker 池身份（默认：主机名-进程号；"
|
||
"生产部署应配稳定 id，否则重启后旧行变 stale）")
|
||
parser.add_argument("--worker-capacity", type=int, default=1,
|
||
help="worker 并发容量声明（注册表里展示用，"
|
||
"本轮 run_loop 仍是串行一轮一任务）")
|
||
parser.add_argument("--docker-image", default="codex-app-server:latest")
|
||
parser.add_argument("--docker-cpus", default="2")
|
||
parser.add_argument("--docker-memory", default="2g")
|
||
parser.add_argument("--docker-user", default="1000:1000")
|
||
parser.add_argument("--docker-worktrees-root", default=None)
|
||
parser.add_argument("--docker-timeout", type=float, default=None)
|
||
parser.add_argument("--docker-extra-mounts", default=None,
|
||
help="docker 额外挂载，逗号分隔 host:container[:ro|rw]，"
|
||
"用于容器内 host.py 写 work/、host-state.json，"
|
||
"以及挂载原仓库和 codex-home（同路径挂载，"
|
||
"容器内配置无需改路径）")
|
||
parser.add_argument("--task-timeout", type=float, default=1800,
|
||
help="每任务截止时间（秒，默认 1800）；超时后真杀"
|
||
"容器/进程组、任务置 failed；与 --docker-timeout"
|
||
"（单次 docker run 超时）取较小值，与"
|
||
" --max-seconds（run_loop 整体墙钟上限）不同")
|
||
parser.add_argument("--max-retries", type=int, default=1,
|
||
help="模型任务失败重试次数（默认 1，0 表示不重试；"
|
||
"取消与已记 fallback 的失败不重试）")
|
||
parser.add_argument("--retry-backoff-seconds", type=float, default=5.0,
|
||
help="重试前等待秒数（默认 5）")
|
||
parser.add_argument("--loop", action="store_true",
|
||
help="常驻循环，每 --interval 秒跑一轮")
|
||
parser.add_argument("--interval", type=int, default=300)
|
||
parser.add_argument("--max-seconds", type=float, default=None)
|
||
args = parser.parse_args(argv)
|
||
setup_logging()
|
||
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,
|
||
fallback_models=(
|
||
[m.strip() for m in
|
||
args.fallback_models.split(",")]
|
||
if args.fallback_models else []),
|
||
max_turns=args.max_turns,
|
||
executor=_build_cli_executor(args, host_py),
|
||
task_timeout=args.task_timeout,
|
||
max_retries=args.max_retries,
|
||
retry_backoff_seconds=args.retry_backoff_seconds,
|
||
repos_dir=args.repos_dir,
|
||
heartbeat_file=args.heartbeat_file,
|
||
worker_id=args.worker_id,
|
||
capacity=args.worker_capacity)
|
||
# CLI 启动清理兜底（常驻与单次都做）：清掉上一轮残留容器。
|
||
worker.cleanup_stale_containers()
|
||
# CLI 起 worker 前先建表/迁移：全新库直接起循环不应因缺表而炸。
|
||
try:
|
||
from cloud.server import init_db as _init_db
|
||
_init_db(args.db)
|
||
except Exception as exc:
|
||
print("worker db 初始化失败：%s" % exc, file=sys.stderr)
|
||
raise SystemExit(1)
|
||
if args.loop:
|
||
def _on_signal(signum, frame):
|
||
worker._loop_stop = True
|
||
signal.signal(signal.SIGTERM, _on_signal)
|
||
signal.signal(signal.SIGINT, _on_signal)
|
||
try:
|
||
if args.loop:
|
||
results = worker.run_loop(interval=args.interval,
|
||
max_seconds=args.max_seconds)
|
||
else:
|
||
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()
|