cloud: production gaps (worker loop + true cancel, task result columns + migration, DB interface migration, metrics + backup)
This commit is contained in:
1 parent
07ef869267
commit
631580beab
14 files changed
+956
-29
No files matched your search
+34
-1
@@ -274,4 +274,37 @@ interrupted 且终态任务 cancel 返回 400、心跳后 sweep 不扫刚心跳
|
||||
APNs / FCM 的真实凭证与发送端仍缺;单机 SQLite 在高并发写下
|
||||
仍有上限,生产量上来后需换 Postgres(接口层与骨架已就绪,
|
||||
但真实部署、驱动安装与迁移脚本仍缺,见上文“数据库后端选择”);
|
||||
worker 执行、云端执行、真实通知与真实模型回合均未接通。
|
||||
worker 执行已接通(见下文“worker 常驻与真取消”),云端执行、
|
||||
真实通知与真实模型回合仍受 key 供给与云账号卡点限制。
|
||||
|
||||
## worker 常驻与真取消(本轮新增)
|
||||
|
||||
- 单次:`python3 -m cloud.worker --db <库> --host-py ../host.py`;
|
||||
常驻:`python3 -m cloud.worker --db <库> --loop --interval 300`,
|
||||
每轮 run_all,SIGTERM/SIGINT 时优雅退出,正在跑的任务先标
|
||||
interrupted 再退出;CLI 起 worker 前先 init_db,全新库不再因
|
||||
缺表而炸,循环内一轮异常记日志后继续,不带走整个 worker。
|
||||
- 真取消:子进程以进程组启动(start_new_session),取消时先
|
||||
SIGTERM 给整个进程组、等 2 秒、仍活着再 SIGKILL,杀完任务
|
||||
保持 interrupted,执行结果不覆盖它;取消竞态(run_one 已先
|
||||
标 interrupted)时 cancel 仍返回 True,因为中断确由本次取消触发。
|
||||
- 任务结果列:tasks 表新增 result_text、error_text、actual_model、
|
||||
started_at、finished_at、child_pid,init_db 用 PRAGMA 检查后
|
||||
ALTER TABLE 幂等迁移,老库升级不丢数据;GET /tasks/{id} 返回
|
||||
它们,缺 key 任务的 error_text 含 Missing environment variable。
|
||||
|
||||
## 可观测性与备份(本轮新增)
|
||||
|
||||
- 结构化日志:cloud/logging_config.py,JSON 行日志(时间、级别、
|
||||
服务、任务 id、事件),worker 的认领/完成/失败/取消/fallback/
|
||||
缺 key 记事件日志,绝不写密钥与环境变量值。
|
||||
- 指标:`GET /metrics`(需 X-API-Key,任意项目成员可读,只统计
|
||||
调用者是成员的项目),返回任务总数、各状态计数、成功率、
|
||||
平均耗时、fallback 总次数、成本合计(actual 与 estimated 分开)。
|
||||
- 备份:`python3 -m cloud.backup --db <库> --out <目录> [--keep 3]`,
|
||||
用 SQLite backup API 在线备份,轮转只留最近 3 份(命名含时间戳),
|
||||
失败明确报错不静默;备份文件本身就是可打开的 SQLite 库。
|
||||
- DB 接口迁移:server 及 auth/memory/kb/cost/oauth/push/webhooks
|
||||
的 Store 连接已统一走 cloud.db 的 WAL + busy_timeout 连接
|
||||
(Postgres 后端仍是骨架,驱动缺失时启动即失败、不先建 SQLite
|
||||
文件,这一点迁移后未退化)。
|
||||
+2
-1
@@ -24,6 +24,7 @@ approve_merge -> maintainer,manage_project -> owner(仅 owner)。
|
||||
import hashlib
|
||||
import secrets
|
||||
import sqlite3
|
||||
from cloud.db import connect as _db_connect
|
||||
import time
|
||||
import uuid
|
||||
|
||||
@@ -44,7 +45,7 @@ class AuthStore:
|
||||
|
||||
def __init__(self, db_path):
|
||||
self.db_path = str(db_path)
|
||||
self.conn = sqlite3.connect(self.db_path)
|
||||
self.conn = _db_connect(self.db_path)
|
||||
self.conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS users("
|
||||
"id TEXT PRIMARY KEY, "
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
"""云库备份:SQLite 在线备份 + 轮转。
|
||||
|
||||
用法:python3 -m cloud.backup --db <库路径> --out <目录> [--keep 3]
|
||||
|
||||
- 用 sqlite3 的 backup API 做在线备份(源库正在被 server 读写时
|
||||
也能安全拷,不用先停服务、不用直接 cp 正在写的文件)。
|
||||
- 轮转:备份文件名含时间戳,写完后按时间倒序只保留最近 keep
|
||||
份(默认 3),多余的删掉;keep < 1 时报错,不删光。
|
||||
- 备份失败(源库不存在、目标目录建不出来、backup API 报错)
|
||||
明确抛错、不静默,CLI 退出码非零。
|
||||
- 备份文件本身就是 SQLite 库,可直接用 sqlite3 打开验证含
|
||||
tasks 表;备份里不含任何密钥(库里本来就只存 API key 哈希)。
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import os
|
||||
import sqlite3
|
||||
import time
|
||||
|
||||
from cloud.db import connect as _db_connect
|
||||
|
||||
|
||||
def backup_db(db_path, out_dir, keep=3, timestamp=None):
|
||||
"""执行一次备份并轮转,返回备份文件路径。
|
||||
|
||||
timestamp 可注入(单测用固定时间戳造多份);缺省用当前时间。
|
||||
"""
|
||||
if keep < 1:
|
||||
raise ValueError("keep 必须 >= 1,否则会删光所有备份")
|
||||
if not os.path.exists(db_path):
|
||||
raise FileNotFoundError("源库不存在:%s" % db_path)
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
stamp = timestamp or time.strftime("%Y%m%d%H%M%S", time.localtime())
|
||||
dest = os.path.join(out_dir, "cloud-backup-%s.db" % stamp)
|
||||
src = _db_connect(str(db_path))
|
||||
try:
|
||||
dst = sqlite3.connect(dest)
|
||||
try:
|
||||
src.backup(dst)
|
||||
finally:
|
||||
dst.close()
|
||||
finally:
|
||||
src.close()
|
||||
# 轮转:按文件名时间戳倒序(文件名即时间戳,排序即时间序),
|
||||
# 只留最近 keep 份。
|
||||
backups = sorted(glob.glob(os.path.join(out_dir, "cloud-backup-*.db")))
|
||||
for old in backups[:-keep] if len(backups) > keep else []:
|
||||
os.remove(old)
|
||||
return dest
|
||||
|
||||
|
||||
def main(argv=None):
|
||||
parser = argparse.ArgumentParser(description="cloud control-plane backup")
|
||||
parser.add_argument("--db", required=True)
|
||||
parser.add_argument("--out", required=True)
|
||||
parser.add_argument("--keep", type=int, default=3)
|
||||
args = parser.parse_args(argv)
|
||||
dest = backup_db(args.db, args.out, keep=args.keep)
|
||||
print(dest)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+2
-1
@@ -10,6 +10,7 @@ so estimates are never mixed into actual totals.
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
from cloud.db import connect as _db_connect
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
|
||||
@@ -56,7 +57,7 @@ class CostStore:
|
||||
if db_path != ":memory:":
|
||||
parent = os.path.dirname(os.path.abspath(db_path))
|
||||
os.makedirs(parent, exist_ok=True)
|
||||
self._conn = sqlite3.connect(db_path)
|
||||
self._conn = _db_connect(db_path)
|
||||
self._conn.execute(
|
||||
"""CREATE TABLE IF NOT EXISTS cost_records (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
|
||||
+2
-1
@@ -17,6 +17,7 @@ Design notes:
|
||||
"""
|
||||
|
||||
import sqlite3
|
||||
from cloud.db import connect as _db_connect
|
||||
import time
|
||||
|
||||
_KINDS = ("file", "readme", "doc", "task", "review")
|
||||
@@ -27,7 +28,7 @@ class KBStore:
|
||||
|
||||
def __init__(self, db_path):
|
||||
self.db_path = db_path
|
||||
self.conn = sqlite3.connect(db_path)
|
||||
self.conn = _db_connect(db_path)
|
||||
self.conn.row_factory = sqlite3.Row
|
||||
self.conn.execute(
|
||||
"""CREATE TABLE IF NOT EXISTS kb_docs(
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
"""结构化日志:JSON 行日志,含时间、级别、服务、任务 id、事件。
|
||||
|
||||
只用标准库 logging。日志绝不写任何密钥、API key、环境变量值:
|
||||
调用方只传任务 id、状态、事件名这类非敏感字段。
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import sys
|
||||
import time
|
||||
|
||||
|
||||
class JsonFormatter(logging.Formatter):
|
||||
"""把每条记录格式化成一行 JSON,方便后续按字段过滤聚合。"""
|
||||
|
||||
def format(self, record):
|
||||
payload = {
|
||||
"ts": time.strftime("%Y-%m-%dT%H:%M:%S", time.gmtime(record.created)),
|
||||
"level": record.levelname,
|
||||
"service": "cloud-control-plane",
|
||||
"logger": record.name,
|
||||
"message": record.getMessage(),
|
||||
}
|
||||
for key in ("task_id", "project_id", "event", "status"):
|
||||
if hasattr(record, key):
|
||||
payload[key] = getattr(record, key)
|
||||
if record.exc_info and record.exc_info[0] is not None:
|
||||
payload["error_type"] = record.exc_info[0].__name__
|
||||
return json.dumps(payload, ensure_ascii=False)
|
||||
|
||||
|
||||
def setup_logging(level=logging.INFO, stream=None):
|
||||
"""配置根 logger 为 JSON 行日志;重复调用不重复加 handler。"""
|
||||
handler = logging.StreamHandler(stream or sys.stderr)
|
||||
handler.setFormatter(JsonFormatter())
|
||||
root = logging.getLogger()
|
||||
root.handlers = [handler]
|
||||
root.setLevel(level)
|
||||
return root
|
||||
|
||||
|
||||
def get_logger(name="cloud"):
|
||||
return logging.getLogger(name)
|
||||
|
||||
|
||||
def log_event(logger, event, task_id=None, project_id=None, status=None,
|
||||
message=None):
|
||||
"""记一条结构化事件日志;message 默认就是事件名。"""
|
||||
logger.info(message or event, extra={
|
||||
"event": event,
|
||||
"task_id": task_id,
|
||||
"project_id": project_id,
|
||||
"status": status,
|
||||
})
|
||||
+2
-1
@@ -14,6 +14,7 @@ Design notes:
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
from cloud.db import connect as _db_connect
|
||||
import time
|
||||
|
||||
|
||||
@@ -22,7 +23,7 @@ class MemoryStore:
|
||||
|
||||
def __init__(self, db_path):
|
||||
self.db_path = db_path
|
||||
self.conn = sqlite3.connect(db_path)
|
||||
self.conn = _db_connect(db_path)
|
||||
self.conn.row_factory = sqlite3.Row
|
||||
self.conn.execute(
|
||||
"""CREATE TABLE IF NOT EXISTS summaries(
|
||||
|
||||
+2
-1
@@ -33,6 +33,7 @@ import json
|
||||
import os
|
||||
import secrets
|
||||
import sqlite3
|
||||
from cloud.db import connect as _db_connect
|
||||
import time
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
@@ -96,7 +97,7 @@ class OAuthStateStore:
|
||||
|
||||
def __init__(self, db_path):
|
||||
self.db_path = str(db_path)
|
||||
self.conn = sqlite3.connect(self.db_path)
|
||||
self.conn = _db_connect(self.db_path)
|
||||
self.conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS oauth_states("
|
||||
"state TEXT PRIMARY KEY, "
|
||||
|
||||
+2
-1
@@ -32,6 +32,7 @@
|
||||
"""
|
||||
|
||||
import sqlite3
|
||||
from cloud.db import connect as _db_connect
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import contextmanager
|
||||
@@ -78,7 +79,7 @@ def _connect(db):
|
||||
if _is_connection(db):
|
||||
yield db
|
||||
elif isinstance(db, (str, bytes)) or hasattr(db, "__fspath__"):
|
||||
conn = sqlite3.connect(str(db), timeout=30)
|
||||
conn = _db_connect(str(db))
|
||||
try:
|
||||
yield conn
|
||||
finally:
|
||||
|
||||
@@ -167,6 +167,32 @@ class LocalExecutor(Executor):
|
||||
return self._result_from_runner(
|
||||
task["id"], self.runner(cmd))
|
||||
|
||||
def popen_task(self, task, kind="dispatch"):
|
||||
"""进程组启动 host.py 子进程,返回 Popen(供真取消用)。
|
||||
|
||||
start_new_session=True 让子进程自成进程组,worker 取消时
|
||||
killpg 才能连同子孙一起杀掉。注意:只有默认 runner(真实
|
||||
subprocess)时才用本方法;单测注入的假 runner 走同步路径,
|
||||
不走这里。
|
||||
"""
|
||||
if not isinstance(task, dict) or task.get("id") is None:
|
||||
raise ValueError("任务缺少 id")
|
||||
if not task.get("prompt"):
|
||||
raise ValueError("任务缺少 prompt")
|
||||
cmd = self._base_command()
|
||||
if kind == "run":
|
||||
cmd += ["run", "--task", self._host_name(task),
|
||||
"--prompt", task["prompt"]]
|
||||
else:
|
||||
cmd += ["dispatch", "--name", self._host_name(task),
|
||||
"--task", task["prompt"]]
|
||||
repo = task.get("repo") or self.repo
|
||||
if repo:
|
||||
cmd += ["--repo", repo]
|
||||
return subprocess.Popen(cmd, stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE, text=True,
|
||||
start_new_session=True)
|
||||
|
||||
def execute(self, task) -> dict:
|
||||
task_id = task.get("id") if isinstance(task, dict) else None
|
||||
if task_id is None:
|
||||
|
||||
+110
@@ -87,6 +87,30 @@ def _connect(db_path):
|
||||
return _db_connect(str(db_path))
|
||||
|
||||
|
||||
def migrate_tasks_table(conn):
|
||||
"""幂等迁移:给老库 tasks 表补新列,不丢数据。
|
||||
|
||||
新列:result_text、error_text、actual_model、started_at、
|
||||
finished_at、child_pid。PRAGMA table_info 查现有列,缺哪个
|
||||
补哪个;已迁移的库再跑一次是 no-op。worker 写执行结果、
|
||||
server 读任务时都依赖这些列存在。
|
||||
"""
|
||||
cols = {row[1] for row in conn.execute(
|
||||
"PRAGMA table_info(tasks)").fetchall()}
|
||||
new_cols = {
|
||||
"result_text": "TEXT",
|
||||
"error_text": "TEXT",
|
||||
"actual_model": "TEXT",
|
||||
"started_at": "REAL",
|
||||
"finished_at": "REAL",
|
||||
"child_pid": "INTEGER",
|
||||
}
|
||||
for name, coltype in new_cols.items():
|
||||
if name not in cols:
|
||||
conn.execute(
|
||||
"ALTER TABLE tasks ADD COLUMN %s %s" % (name, coltype))
|
||||
|
||||
|
||||
def init_db(db_path):
|
||||
"""创建 server 自己的 projects/tasks 表(AuthStore 的表由其自建)。"""
|
||||
conn = _connect(db_path)
|
||||
@@ -113,6 +137,7 @@ def init_db(db_path):
|
||||
"last_heartbeat REAL, "
|
||||
"UNIQUE(project_id, idempotency_key))"
|
||||
)
|
||||
migrate_tasks_table(conn)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
@@ -131,6 +156,14 @@ def _auth_store(db_path):
|
||||
pass
|
||||
|
||||
|
||||
def _row_get(row, name):
|
||||
"""sqlite Row 的容错取值:老库行缺列时返回 None 不炸。"""
|
||||
try:
|
||||
return row[name]
|
||||
except (KeyError, IndexError, TypeError):
|
||||
return None
|
||||
|
||||
|
||||
def _task_to_dict(row):
|
||||
depends_on = row["depends_on"]
|
||||
try:
|
||||
@@ -151,6 +184,12 @@ def _task_to_dict(row):
|
||||
"created_at": row["created_at"],
|
||||
"updated_at": row["updated_at"],
|
||||
"last_heartbeat": row["last_heartbeat"],
|
||||
"result_text": _row_get(row, "result_text"),
|
||||
"error_text": _row_get(row, "error_text"),
|
||||
"actual_model": _row_get(row, "actual_model"),
|
||||
"started_at": _row_get(row, "started_at"),
|
||||
"finished_at": _row_get(row, "finished_at"),
|
||||
"child_pid": _row_get(row, "child_pid"),
|
||||
}
|
||||
|
||||
|
||||
@@ -306,6 +345,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
||||
self._create_project(user_id)
|
||||
elif path == "/projects" and method == "GET":
|
||||
self._list_projects(user_id)
|
||||
elif path == "/metrics" and method == "GET":
|
||||
self._metrics(user_id)
|
||||
elif path == "/models" and method == "GET":
|
||||
self._models_overview(user_id)
|
||||
elif path.startswith("/projects/"):
|
||||
@@ -989,6 +1030,75 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
||||
swept = sweep_stale_tasks(self.server.db_path, threshold)
|
||||
self._send_json(200, {"swept": swept})
|
||||
|
||||
def _metrics(self, user_id):
|
||||
"""可观测性:任务计数、成功率、平均耗时、fallback、成本合计。
|
||||
|
||||
权限:任意项目成员即可读(不要求 owner),但只统计调用者
|
||||
是成员的项目——不是全库数字,避免跨项目信息泄露。
|
||||
成本真数与估算分开:cost_records 的 source 列区分。
|
||||
"""
|
||||
conn = self._db()
|
||||
try:
|
||||
member_ids = {
|
||||
row["project_id"]
|
||||
for row in conn.execute(
|
||||
"SELECT project_id FROM project_members WHERE user_id = ?",
|
||||
(user_id,),
|
||||
).fetchall()
|
||||
}
|
||||
rows = conn.execute(
|
||||
"SELECT project_id, status, started_at, finished_at"
|
||||
" FROM tasks"
|
||||
).fetchall()
|
||||
by_status = {}
|
||||
durations = []
|
||||
for row in rows:
|
||||
if row["project_id"] not in member_ids:
|
||||
continue
|
||||
by_status[row["status"]] = by_status.get(row["status"], 0) + 1
|
||||
if row["started_at"] and row["finished_at"]:
|
||||
durations.append(row["finished_at"] - row["started_at"])
|
||||
total = sum(by_status.values())
|
||||
done = by_status.get("completed", 0) + by_status.get("merged", 0)
|
||||
failed = by_status.get("failed", 0) + by_status.get("merge-failed", 0)
|
||||
finished = done + failed
|
||||
try:
|
||||
cost_rows = conn.execute(
|
||||
"SELECT project_id, source, total_tokens, usd_micros,"
|
||||
" fallback_count FROM cost_records"
|
||||
).fetchall()
|
||||
except sqlite3.OperationalError:
|
||||
cost_rows = []
|
||||
actual_tokens = estimated_tokens = 0
|
||||
actual_usd = estimated_usd = 0
|
||||
fallback_total = 0
|
||||
for row in cost_rows:
|
||||
if row["project_id"] not in member_ids:
|
||||
continue
|
||||
if row["source"] == "actual":
|
||||
actual_tokens += row["total_tokens"] or 0
|
||||
actual_usd += row["usd_micros"] or 0
|
||||
else:
|
||||
estimated_tokens += row["total_tokens"] or 0
|
||||
estimated_usd += row["usd_micros"] or 0
|
||||
fallback_total += row["fallback_count"] or 0
|
||||
finally:
|
||||
conn.close()
|
||||
self._send_json(200, {
|
||||
"total_tasks": total,
|
||||
"by_status": by_status,
|
||||
"success_rate": (done / finished) if finished else None,
|
||||
"avg_duration_seconds": (
|
||||
sum(durations) / len(durations)) if durations else None,
|
||||
"fallback_total": fallback_total,
|
||||
"cost": {
|
||||
"actual_tokens": actual_tokens,
|
||||
"estimated_tokens": estimated_tokens,
|
||||
"actual_usd_micros": actual_usd,
|
||||
"estimated_usd_micros": estimated_usd,
|
||||
},
|
||||
})
|
||||
|
||||
# -- OAuth(pre-auth,靠 state + secret,不走 X-API-Key)-----------
|
||||
def _oauth_provider(self, path):
|
||||
parts = [p for p in path.split("/") if p]
|
||||
|
||||
@@ -0,0 +1,342 @@
|
||||
"""生产级缺口单测:worker 常驻/真取消、tasks 迁移、metrics、备份、日志。
|
||||
|
||||
只用标准库,不跑真实模型回合、不联网、不登录任何服务;
|
||||
process_factory 返回的长跑子进程只是 `sleep`,取消后必须消失。
|
||||
"""
|
||||
|
||||
import http.client
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sqlite3
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
try:
|
||||
from cloud.auth import AuthStore
|
||||
from cloud.backup import backup_db
|
||||
from cloud.cost import CostStore
|
||||
from cloud.db import PostgresDriverMissingError
|
||||
from cloud.logging_config import JsonFormatter, log_event, setup_logging
|
||||
from cloud.server import create_server, init_db
|
||||
from cloud.worker import Worker
|
||||
except ImportError:
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
from cloud.auth import AuthStore
|
||||
from cloud.backup import backup_db
|
||||
from cloud.cost import CostStore
|
||||
from cloud.db import PostgresDriverMissingError
|
||||
from cloud.logging_config import JsonFormatter, log_event, setup_logging
|
||||
from cloud.server import create_server, init_db
|
||||
from cloud.worker import Worker
|
||||
|
||||
|
||||
def _insert_task(db_path, task_id, prompt="do work", status="queued",
|
||||
project_id="p1", started=None, finished=None):
|
||||
now = time.time()
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
conn.execute(
|
||||
"INSERT INTO tasks(id, project_id, prompt, status,"
|
||||
" idempotency_key, depends_on, priority, created_by,"
|
||||
" created_at, updated_at, last_heartbeat)"
|
||||
" VALUES(?, ?, ?, ?, NULL, '[]', 0, 'tester', ?, ?, NULL)",
|
||||
(task_id, project_id, prompt, status, now, now),
|
||||
)
|
||||
if started is not None:
|
||||
conn.execute(
|
||||
"UPDATE tasks SET started_at = ?, finished_at = ?"
|
||||
" WHERE id = ?", (started, finished, task_id))
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _row(db_path, task_id):
|
||||
conn = sqlite3.connect(db_path)
|
||||
conn.row_factory = sqlite3.Row
|
||||
try:
|
||||
row = conn.execute(
|
||||
"SELECT * FROM tasks WHERE id = ?", (task_id,)).fetchone()
|
||||
return dict(row) if row else None
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _completed_runner(cmd):
|
||||
return SimpleNamespace(
|
||||
returncode=0,
|
||||
stdout=json.dumps({"task": "x", "status": "completed",
|
||||
"text": "done"}),
|
||||
stderr="")
|
||||
|
||||
|
||||
class MigrationTestCase(unittest.TestCase):
|
||||
def test_old_db_migrates_without_data_loss(self):
|
||||
tmp = tempfile.TemporaryDirectory()
|
||||
db_path = os.path.join(tmp.name, "old.db")
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
conn.execute(
|
||||
"CREATE TABLE projects(id TEXT PRIMARY KEY, name TEXT,"
|
||||
" owner_id TEXT, created_at REAL)")
|
||||
conn.execute(
|
||||
"CREATE TABLE tasks(id TEXT PRIMARY KEY, project_id TEXT,"
|
||||
" prompt TEXT, status TEXT, idempotency_key TEXT,"
|
||||
" depends_on TEXT, priority INTEGER, created_by TEXT,"
|
||||
" created_at REAL, updated_at REAL, last_heartbeat REAL,"
|
||||
" UNIQUE(project_id, idempotency_key))")
|
||||
conn.execute(
|
||||
"INSERT INTO tasks(id, project_id, prompt, status,"
|
||||
" depends_on, priority, created_by, created_at,"
|
||||
" updated_at, last_heartbeat)"
|
||||
" VALUES('old1', 'p1', 'old prompt', 'completed', '[]',"
|
||||
" 0, 'tester', 1, 1, NULL)")
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
init_db(db_path)
|
||||
cols = {r[1] for r in sqlite3.connect(db_path).execute(
|
||||
"PRAGMA table_info(tasks)").fetchall()}
|
||||
for col in ("result_text", "error_text", "actual_model",
|
||||
"started_at", "finished_at", "child_pid"):
|
||||
self.assertIn(col, cols)
|
||||
row = _row(db_path, "old1")
|
||||
self.assertEqual(row["prompt"], "old prompt")
|
||||
self.assertEqual(row["status"], "completed")
|
||||
# 再跑一次 init_db 是 no-op,不炸、不丢数据。
|
||||
init_db(db_path)
|
||||
self.assertEqual(_row(db_path, "old1")["prompt"], "old prompt")
|
||||
tmp.cleanup()
|
||||
|
||||
|
||||
class WorkerResultColumnsTestCase(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_completed_task_writes_result_columns(self):
|
||||
_insert_task(self.db_path, "t1")
|
||||
worker = Worker(self.db_path, "/fake/host.py",
|
||||
runner=_completed_runner,
|
||||
env={"XCM_API_KEY": "test-placeholder"})
|
||||
result = worker.run_one()
|
||||
self.assertEqual(result["status"], "completed")
|
||||
row = _row(self.db_path, "t1")
|
||||
self.assertIn("done", row["result_text"])
|
||||
self.assertIsNone(row["error_text"])
|
||||
self.assertEqual(row["actual_model"], worker.model)
|
||||
self.assertIsNotNone(row["started_at"])
|
||||
self.assertIsNotNone(row["finished_at"])
|
||||
self.assertGreaterEqual(row["finished_at"], row["started_at"])
|
||||
|
||||
def test_missing_key_writes_error_text_and_calls_runner_zero_times(self):
|
||||
calls = []
|
||||
|
||||
def counting(cmd):
|
||||
calls.append(cmd)
|
||||
return _completed_runner(cmd)
|
||||
|
||||
_insert_task(self.db_path, "t2")
|
||||
worker = Worker(self.db_path, "/fake/host.py",
|
||||
runner=counting, env={})
|
||||
result = worker.run_one()
|
||||
self.assertEqual(result["status"], "failed")
|
||||
self.assertEqual(calls, [])
|
||||
row = _row(self.db_path, "t2")
|
||||
self.assertIn("Missing environment variable", row["error_text"])
|
||||
self.assertIsNotNone(row["started_at"])
|
||||
|
||||
|
||||
class WorkerTrueCancelTestCase(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 _sleep_factory(self, task):
|
||||
return subprocess.Popen(
|
||||
[sys.executable, "-c", "import time; time.sleep(30)"],
|
||||
stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True,
|
||||
start_new_session=True)
|
||||
|
||||
def test_cancel_kills_process_group_and_keeps_interrupted(self):
|
||||
_insert_task(self.db_path, "long1")
|
||||
worker = Worker(self.db_path, "/fake/host.py",
|
||||
process_factory=self._sleep_factory,
|
||||
env={"XCM_API_KEY": "test-placeholder"})
|
||||
results = []
|
||||
thread = threading.Thread(
|
||||
target=lambda: results.append(worker.run_one()), daemon=True)
|
||||
thread.start()
|
||||
pgid = None
|
||||
deadline = time.time() + 5
|
||||
while time.time() < deadline:
|
||||
row = _row(self.db_path, "long1")
|
||||
if row and row["child_pid"]:
|
||||
pgid = row["child_pid"]
|
||||
break
|
||||
time.sleep(0.05)
|
||||
self.assertIsNotNone(pgid, "子进程 pid 未写进任务记录")
|
||||
# 进程组当前存在(killpg 0 不抛)。
|
||||
os.killpg(os.getpgid(pgid), 0)
|
||||
self.assertTrue(worker.cancel("long1"))
|
||||
thread.join(timeout=10)
|
||||
self.assertFalse(thread.is_alive())
|
||||
self.assertEqual(results[0]["status"], "interrupted")
|
||||
self.assertEqual(_row(self.db_path, "long1")["status"], "interrupted")
|
||||
# 进程组确实消失。
|
||||
with self.assertRaises((ProcessLookupError, PermissionError, OSError)):
|
||||
os.killpg(os.getpgid(pgid), 0)
|
||||
|
||||
def test_loop_stops_and_marks_running_interrupted(self):
|
||||
_insert_task(self.db_path, "loop1")
|
||||
worker = Worker(self.db_path, "/fake/host.py",
|
||||
process_factory=self._sleep_factory,
|
||||
env={"XCM_API_KEY": "test-placeholder"})
|
||||
stop = threading.Event()
|
||||
results = []
|
||||
thread = threading.Thread(
|
||||
target=lambda: results.extend(
|
||||
worker.run_loop(interval=0.2, stop_event=stop)),
|
||||
daemon=True)
|
||||
thread.start()
|
||||
time.sleep(1.0)
|
||||
stop.set()
|
||||
# 循环在任务进行中收到停止信号:_loop_stop 由 run_loop 的
|
||||
# finally 外的轮询感知——这里直接置位模拟 SIGTERM 处理。
|
||||
worker._loop_stop = True
|
||||
thread.join(timeout=10)
|
||||
self.assertFalse(thread.is_alive())
|
||||
self.assertEqual(_row(self.db_path, "loop1")["status"], "interrupted")
|
||||
|
||||
|
||||
class MetricsAndBackupTestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||||
self.auth = AuthStore(self.db_path)
|
||||
self.owner_id, self.owner_key = self.auth.create_user("owner")
|
||||
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)
|
||||
self.tmp.cleanup()
|
||||
|
||||
def _req(self, method, path, key=None, body=None):
|
||||
conn = http.client.HTTPConnection("127.0.0.1", self.port, timeout=10)
|
||||
hdrs = {}
|
||||
payload = None
|
||||
if body is not None:
|
||||
payload = json.dumps(body).encode()
|
||||
hdrs["Content-Type"] = "application/json"
|
||||
if key:
|
||||
hdrs["X-API-Key"] = key
|
||||
try:
|
||||
conn.request(method, path, body=payload, headers=hdrs)
|
||||
resp = conn.getresponse()
|
||||
data = resp.read()
|
||||
finally:
|
||||
conn.close()
|
||||
return resp.status, json.loads(data) if data else None
|
||||
|
||||
def test_metrics_counts_match_db(self):
|
||||
status, proj = self._req("POST", "/projects", self.owner_key,
|
||||
{"name": "m"})
|
||||
self.assertEqual(status, 201)
|
||||
pid = proj["id"]
|
||||
_insert_task(self.db_path, "m1", status="completed",
|
||||
project_id=pid, started=1000.0, finished=1002.0)
|
||||
_insert_task(self.db_path, "m2", status="failed",
|
||||
project_id=pid, started=1000.0, finished=1001.0)
|
||||
_insert_task(self.db_path, "other", status="completed",
|
||||
project_id="someone-else")
|
||||
costs = CostStore(self.db_path)
|
||||
costs.record("m1", pid, "gpt-6.1-sol", input_tokens=10,
|
||||
output_tokens=5, source="actual", fallback_count=1)
|
||||
costs.record("m2", pid, "gpt-6.1-sol", input_tokens=20,
|
||||
output_tokens=0, source="estimated")
|
||||
status, metrics = self._req("GET", "/metrics", self.owner_key)
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(metrics["total_tasks"], 2)
|
||||
self.assertEqual(metrics["by_status"]["completed"], 1)
|
||||
self.assertEqual(metrics["by_status"]["failed"], 1)
|
||||
self.assertAlmostEqual(metrics["success_rate"], 0.5)
|
||||
self.assertAlmostEqual(metrics["avg_duration_seconds"], 1.5)
|
||||
self.assertEqual(metrics["fallback_total"], 1)
|
||||
self.assertEqual(metrics["cost"]["actual_tokens"], 15)
|
||||
self.assertEqual(metrics["cost"]["estimated_tokens"], 20)
|
||||
# 无 key 401。
|
||||
status, _ = self._req("GET", "/metrics")
|
||||
self.assertEqual(status, 401)
|
||||
|
||||
def test_backup_rotation_keeps_three(self):
|
||||
init_db(self.db_path)
|
||||
_insert_task(self.db_path, "b1")
|
||||
out_dir = os.path.join(self.tmp.name, "backups")
|
||||
for stamp in ("20260101000000", "20260102000000", "20260103000000",
|
||||
"20260104000000", "20260105000000"):
|
||||
backup_db(self.db_path, out_dir, keep=3, timestamp=stamp)
|
||||
files = sorted(os.listdir(out_dir))
|
||||
self.assertEqual(len(files), 3)
|
||||
self.assertEqual(files[0], "cloud-backup-20260103000000.db")
|
||||
conn = sqlite3.connect(os.path.join(out_dir, files[-1]))
|
||||
try:
|
||||
row = conn.execute(
|
||||
"SELECT id FROM tasks WHERE id = 'b1'").fetchone()
|
||||
self.assertIsNotNone(row)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def test_postgres_backend_without_driver_fails_before_file(self):
|
||||
other_db = os.path.join(self.tmp.name, "pg.db")
|
||||
with self.assertRaises(PostgresDriverMissingError):
|
||||
create_server(other_db, "127.0.0.1", 0, backend="postgres")
|
||||
self.assertFalse(os.path.exists(other_db))
|
||||
|
||||
|
||||
class LoggingTestCase(unittest.TestCase):
|
||||
def test_json_log_contains_task_and_event_without_secret(self):
|
||||
stream = io.StringIO()
|
||||
handler = logging.StreamHandler(stream)
|
||||
handler.setFormatter(JsonFormatter())
|
||||
logger = logging.getLogger("cloud.test-prod-gaps")
|
||||
logger.handlers = [handler]
|
||||
logger.setLevel(logging.INFO)
|
||||
logger.propagate = False
|
||||
try:
|
||||
log_event(logger, "worker_claim", task_id="t9",
|
||||
project_id="p9", status="running")
|
||||
finally:
|
||||
logger.handlers = []
|
||||
payload = json.loads(stream.getvalue().strip())
|
||||
self.assertEqual(payload["task_id"], "t9")
|
||||
self.assertEqual(payload["event"], "worker_claim")
|
||||
self.assertEqual(payload["service"], "cloud-control-plane")
|
||||
# setup_logging 本身不重复加 handler。
|
||||
root = setup_logging(stream=io.StringIO())
|
||||
self.assertEqual(len(root.handlers), 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+3
-1
@@ -21,6 +21,8 @@ import hmac
|
||||
import sqlite3
|
||||
import time
|
||||
|
||||
from cloud.db import connect as _db_connect
|
||||
|
||||
__all__ = [
|
||||
"verify_github_signature",
|
||||
"parse_github_pr_comment",
|
||||
@@ -155,7 +157,7 @@ class SlackUserMap:
|
||||
|
||||
def __init__(self, db_path):
|
||||
self.db_path = db_path
|
||||
self._conn = sqlite3.connect(db_path)
|
||||
self._conn = _db_connect(db_path)
|
||||
self._conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS slack_users ("
|
||||
"slack_id TEXT PRIMARY KEY, "
|
||||
|
||||
+311
-21
@@ -38,13 +38,13 @@
|
||||
(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,不用执行结果覆盖它。
|
||||
取消:worker 的 cancel() 走 server 的 cancel 语义:把
|
||||
queued/running 任务直接标 interrupted(已终态的任务不动),
|
||||
且正在跑的子进程会真杀——子进程以进程组启动(start_new_session),
|
||||
取消时先 SIGTERM 给整个进程组、等 2 秒、仍活着再 SIGKILL,
|
||||
杀完后任务保持 interrupted,执行结果不覆盖它。run_one() 在
|
||||
执行返回后会先查任务是否已被取消,已被取消的一律保持
|
||||
interrupted,不用执行结果覆盖它。
|
||||
|
||||
可注入:runner(假 CLI,见 sandbox.LocalExecutor)与 env 映射
|
||||
均可注入,单测不跑真实模型回合、不联网、不登录任何服务。
|
||||
@@ -54,14 +54,20 @@ import argparse
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import signal
|
||||
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.sandbox import LocalExecutor
|
||||
from cloud.sandbox import _default_runner as _sandbox_default_runner
|
||||
|
||||
# 与 server.TERMINAL_STATUSES 对齐:这些状态的任务不再被认领或取消。
|
||||
TERMINAL_STATUSES = frozenset(
|
||||
@@ -152,7 +158,7 @@ class Worker:
|
||||
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):
|
||||
cost_store=None, process_factory=None):
|
||||
self.db_path = str(db_path)
|
||||
self.host_py_path = host_py_path
|
||||
self.config_path = config_path
|
||||
@@ -177,6 +183,14 @@ class Worker:
|
||||
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._loop_stop = False
|
||||
self.logger = get_logger("cloud.worker")
|
||||
|
||||
# -- 成本存储(惰性建表,避免只认领不执行时也建表)--------------
|
||||
@property
|
||||
@@ -241,6 +255,99 @@ class Worker:
|
||||
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"}.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()
|
||||
@@ -295,11 +402,23 @@ class Worker:
|
||||
|
||||
# -- 取消(server cancel 语义,见模块 docstring)-------------------
|
||||
def cancel(self, task_id):
|
||||
"""把 queued/running 任务标 interrupted;已终态或不存在返回 False。"""
|
||||
"""把 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
|
||||
conn = _db_connect(self.db_path)
|
||||
try:
|
||||
cur = conn.execute(
|
||||
@@ -313,8 +432,16 @@ class Worker:
|
||||
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 _was_cancelled(self, task_id):
|
||||
@@ -332,6 +459,67 @@ class Worker:
|
||||
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"):
|
||||
"""进程模式执行一轮:启动、轮询、取消时杀进程组。
|
||||
|
||||
返回与同步路径同款 result dict(status/stdout/stderr/
|
||||
host_status/returncode)。取消时返回 status='cancelled'
|
||||
的标记,调用方据此保持 interrupted。
|
||||
"""
|
||||
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
|
||||
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
|
||||
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}
|
||||
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 run_one(self, task=None):
|
||||
"""认领并执行一个任务,返回结果 dict(见模块 docstring)。
|
||||
@@ -351,13 +539,21 @@ class Worker:
|
||||
"task_id": task_id}
|
||||
task = self.get_task(task_id) or task
|
||||
self._audit(task["project_id"], "worker_claim", task_id)
|
||||
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,
|
||||
@@ -370,14 +566,20 @@ class Worker:
|
||||
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)
|
||||
if self._use_popen():
|
||||
result = self._execute_via_process(task)
|
||||
combined = (result.get("stdout", "") or "") + \
|
||||
(result.get("stderr", "") or "")
|
||||
if result["status"] not in ("completed", "cancelled") and \
|
||||
_ALREADY_EXISTS_MARKER in combined:
|
||||
result = self._execute_via_process(task, kind="run")
|
||||
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,
|
||||
@@ -386,8 +588,14 @@ class Worker:
|
||||
|
||||
# 取消优先:执行期间被 cancel 的任务保持 interrupted,
|
||||
# 绝不用执行结果覆盖它(见模块 docstring 的取消说明)。
|
||||
if self._was_cancelled(task_id):
|
||||
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=self.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,
|
||||
@@ -417,6 +625,11 @@ class Worker:
|
||||
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=self.model, started_at=started_at,
|
||||
finished_at=finished_at)
|
||||
self.usage_turns += 1
|
||||
# 成本:host 协议不暴露真实 token 数,用 cost.py 的估算结构
|
||||
# 记 estimated,绝不标 actual;fallback_count 复用同一记录。
|
||||
@@ -428,6 +641,11 @@ class Worker:
|
||||
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,
|
||||
@@ -449,14 +667,65 @@ class Worker:
|
||||
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。
|
||||
"""
|
||||
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)
|
||||
# 按 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 main(argv=None):
|
||||
"""单次派发循环:python3 -m cloud.worker --db <库> --host-py <路径>。
|
||||
"""worker CLI:默认单次 run_all;--loop 进常驻循环。
|
||||
|
||||
只打印任务结果摘要(状态、host 名、错误文本),绝不打印
|
||||
任何环境变量的值。
|
||||
"""
|
||||
parser = argparse.ArgumentParser(description="cloud worker (one pass)")
|
||||
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")
|
||||
@@ -468,7 +737,12 @@ def main(argv=None):
|
||||
parser.add_argument("--model", default=None)
|
||||
parser.add_argument("--max-turns", type=int, default=None)
|
||||
parser.add_argument("--max-tasks", type=int, default=None)
|
||||
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__))
|
||||
@@ -477,8 +751,24 @@ def main(argv=None):
|
||||
repo=args.repo, env_key=args.env_key,
|
||||
codex_config_path=args.codex_config, model=args.model,
|
||||
max_turns=args.max_turns)
|
||||
# CLI 起 worker 前先建表/迁移:全新库直接起循环不应因缺表而炸。
|
||||
try:
|
||||
results = worker.run_all(max_tasks=args.max_tasks)
|
||||
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(
|
||||
|
||||
Reference in new issue
Block a user