Cloud dispatch pool + auto PR (opt-in, skip clearly)

Worker pool: workers table + register/heartbeat (heartbeat is
registration, no 404), GET /workers with alive (120s default),
POST /workers/heartbeat (X-API-Key, missing worker_id 400).
Tasks gain executor (local/docker/cloud, None = any worker);
worker claim filters by executor in pick_next and in the CAS
WHERE clause, CAS semantics unchanged; CloudExecutor _name
now correctly returns cloud, not local.
CloudExecutor: generic HTTP runner interface (endpoint from
explicit arg or CLOUD_RUNNER_ENDPOINT, auth header name only,
token value only from env into the request header, never into
logs/body/errors); no endpoint configured = explicit failure,
never silent downgrade. Vendor adapters (ECS/Cloud Run) are
out of scope, documented as such.
Auto PR: projects.auto_pr (0/1, create-time only, no PATCH
route exists); POST /tasks/{id}/pr (maintainer+ via
approve_merge) only for merged tasks (else 409); skip reasons
in order: auto_pr_disabled, no_worktree, no_remote, no_gh,
gh_not_authenticated — all return 200 with pr_skipped, task
stays merged, no push attempted; push/pr-create failure
returns 502 with step, task stays merged; success writes
pr_url/pr_number and audits pr_created. No default auto-push
from the worker. Remote URLs in errors are scrubbed of
userinfo; gh credentials are managed by gh itself, never
stored by the control plane.
Tests: host 27 pass, cloud 323 pass (8 skipped are
environment-dependent real-Postgres tests).
This commit is contained in:
host committed 2026-10-26 04:57:37 +08:00
1 parent d3e588e4b3
commit 306cdf0d0c
8 files changed
+1256 -34

No files matched your search

+76
View File
@@ -83,6 +83,14 @@ X-API-Key,因为外部服务无法持本站 API key,改靠各自的签名认
本站 API key——它的认证靠随机 state(10 分钟过期、一次性)加
client_secret(只走环境变量),不是无认证后门。
本轮新增端点(同样要求 X-API-Key):`GET /workers`(注册 worker
列表,含 alive)、`POST /workers/heartbeat`(body 含 worker_id,
心跳即注册)、`POST /tasks/{id}/pr`(已 merged 任务的自动 PR,
默认关闭,缺条件时返回 pr_skipped 而不是失败;见下文“云端派发与
自动 PR”一节)。`POST /tasks` 新增可选 `executor`
(local/docker/cloud,未传表示任意 worker 可认领);
`POST /projects` 新增可选 `auto_pr`(0/1,默认 0)。
### OAuth 真实浏览器登录(已实现,部署前置仍需用户动手)
流程(标准 OAuth2 授权码流程):
@@ -626,3 +634,71 @@ interrupted 且终态任务 cancel 返回 400、心跳后 sweep 不扫刚心跳
恢复演练、monitor、多用户验证四项的主机实跑还没做(本轮只在 VM 上
跑了单测、deploy 干跑、monitor 自测与 VM 实跑)。
## 云端派发与自动 PR(本轮新增:多 worker 池、通用 runner、默认关闭的自动 PR)
- 多 worker 池(`cloud/workers.py` + worker 注册/心跳):worker 启动
时向控制面注册(workers 表:worker_id、executor、capacity、
last_seen),`run_loop` 每轮结束写注册表心跳;`GET /workers` 列出
全部 worker 及其 alive(last_seen 在 120 秒内,阈值可经环境变量
`CLOUD_WORKER_ALIVE_SECONDS` 调);`POST /workers/heartbeat`
要求 body 含 worker_id(缺了 400),心跳本身就是注册——worker_id
不存在时自动注册,绝不 404。worker_id 可经 CLI `--worker-id`
指定(生产部署应配稳定 id),未指定时默认“主机名-进程号”
(重启后进程号会变,旧行自然变 stale,不要用随机 uuid 让注册表
每次重启都多一行无法对应到同一台机器);capacity 经
`--worker-capacity` 指定(默认 1,注册表展示用,本轮 run_loop
仍是串行一轮一任务)。两个端点都要 X-API-Key(与其它已认证
端点一致);worker 直接写库注册时不需要 key,HTTP 心跳是给有
key 的外部 worker 用的另一条路。
- 任务按执行器池派发:tasks 表新增 `executor` 列(可空,老库幂等
补列);`POST /tasks` 接受可选 executor(local/docker/cloud,
非法值 400),未传时为 None 表示任意 worker 可认领;worker
认领时只认领 executor 为空或等于自己执行器类型的任务——
pick_next 先过滤,`_claim` 的 WHERE 里再加
`(executor IS NULL OR executor = ?)` 兜底,CAS 语义不变(仅当
queued 时标 running,两个同池 worker 抢同一任务时只有一个拿到
rowcount == 1);认领日志与审计里含 executor。注意
`_executor_name` 的修正:CloudExecutor 以前会被算成 local,
现在明确返回 cloud,否则 cloud worker 会认领 local 池的任务。
- CloudExecutor 通用 runner 接口(厂商无关,`cloud/sandbox.py`):
本轮不假装接通某个云厂商(ECS、Cloud Run 等专属适配明确不在
范围)。CloudExecutor 接受一个 runner endpoint(显式参数或
环境变量 `CLOUD_RUNNER_ENDPOINT`);未设时 execute 一律明确
失败“未配置云端 runner”,绝不静默降级到 Docker/Local;已设时
以 HTTP POST 把任务发给该 endpoint,请求体只有 id、prompt、
repo_url(task["repo"]),不含任何密钥;endpoint 返回的 status
为 completed 才算 completed,否则算 failed。认证只走 header
名:`CLOUD_RUNNER_AUTH_HEADER` 给 header 名,token 值只从环境
变量取(`CLOUD_RUNNER_AUTH_TOKEN_ENV` 指定的变量名,默认
`CLOUD_RUNNER_AUTH_TOKEN`),只放进请求头,绝不写进日志、
错误文本或返回结果。worker CLI 的 `--executor` 新增 cloud 选项,
构造 CloudExecutor;worker 的 run_one 走 dispatch_task,
CloudExecutor 的 dispatch_task/run_task 直接委托给 execute。
- 自动 PR(默认关闭,显式触发,`POST /tasks/{id}/pr`,需项目
maintainer 以上权限 approve_merge):仅当任务状态为 merged 时
可调(其它状态 409;任务不存在 404;权限不足 403)。逻辑按顺序
检查(见 `cloud/pr.py`):项目 auto_pr 不为 1 时返回 200 且
pr_skipped=true、reason=auto_pr_disabled,不失败、不推;本地
没有对应 host 任务 worktree(`<仓库根>/work/host-worktrees/
<host 任务名>`,host 任务名 = cloud 任务 id 经安全字符过滤,
分支名 = host/<host 任务名>,与 sandbox._host_name 同规则)时
reason=no_worktree;worktree 里没有 git remote origin 时
reason=no_remote;本机没有 gh 二进制时 reason=no_gh;gh 未登录
(gh auth status 非零)时 reason=gh_not_authenticated;以上全过
才 push 分支 + gh pr create,成功后把 pr_url/pr_number 写回任务
记录并记审计 pr_created,返回 200 含 pr_url。push 或 pr create
失败时返回 502 且错误文本写清是哪一步(push 还是 pr_create),
任务状态保持 merged 不变,绝不把任务标 failed。worker 不自动推:
自动 PR 只由显式调本端点触发,默认自动推会把测试任务也推到远端。
projects 表新增 `auto_pr` 列(默认 0,老库幂等补列);
`POST /projects` 接受可选 auto_pr(只接受 0/1,非法值 400);
本轮没有现成的 PATCH /projects/{id} 路由,auto_pr 只能在创建时
指定,建后改它需直接改库或重建项目,不要新造一个只为改 auto_pr
的路由。gh 的登录凭证由 gh 自己管,控制面不存、不打任何 token。
- 仍未做(不要宣称已做):厂商专属适配(ECS、Cloud Run 等,需
用户选厂商后另做);默认自动推(本轮明确不做);主机上 gh 的安装
与登录、远端写凭证(控制面不存,自动 PR 的真实成功路径仍需它们);
多 worker 的并发容量 enforcement(capacity 目前只是注册表展示,
run_loop 仍串行);worker 经 HTTP 发心跳的凭证供给(worker 直接
写库时不需要,外部 worker 走 HTTP 时需自带 API key)。
+14 -1
View File
@@ -781,7 +781,8 @@ POSTGRES_SCHEMA = (
"name TEXT, "
"owner_id TEXT, "
"created_at DOUBLE PRECISION, "
"repo_url TEXT)",
"repo_url TEXT, "
"auto_pr INTEGER)",
# tasks 含 server.migrate_tasks_table 补的那 6 列,建表时一次建齐;
# 老库迁移仍走 migrate_tasks_table(PRAGMA 翻译后在 PG 上是 no-op)。
"CREATE TABLE IF NOT EXISTS tasks("
@@ -802,6 +803,9 @@ POSTGRES_SCHEMA = (
"started_at DOUBLE PRECISION, "
"finished_at DOUBLE PRECISION, "
"child_pid BIGINT, "
"executor TEXT, "
"pr_url TEXT, "
"pr_number INTEGER, "
"UNIQUE(project_id, idempotency_key))",
"CREATE TABLE IF NOT EXISTS summaries("
"id BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, "
@@ -875,6 +879,15 @@ def init_postgres_schema(conn):
"PRAGMA table_info(projects)").fetchall()}
if "repo_url" not in cols:
conn.execute("ALTER TABLE projects ADD COLUMN repo_url TEXT")
if "auto_pr" not in cols:
conn.execute("ALTER TABLE projects ADD COLUMN auto_pr INTEGER")
tcols = {row[1] for row in conn.execute(
"PRAGMA table_info(tasks)").fetchall()}
for _name, _type in (("executor", "TEXT"), ("pr_url", "TEXT"),
("pr_number", "INTEGER")):
if _name not in tcols:
conn.execute(
"ALTER TABLE tasks ADD COLUMN %s %s" % (_name, _type))
conn.commit()
+158
View File
@@ -0,0 +1,158 @@
"""云端自动 PR:默认关闭、显式触发、缺条件时跳过不失败。
零第三方依赖,只用 Python 标准库;不改 host.py 的任务逻辑,
只复用它的规则:host 任务名 = cloud 任务 id 经安全字符过滤
(见 cloud/sandbox.py LocalExecutor._host_name),分支名是
host/<host 任务名>,worktree 在 <仓库根>/work/host-worktrees/
<host 任务名>。
触发规则(按顺序,任一步跳过即返回,不推、不改任务状态):
1. 任务所属项目的 auto_pr 不为 1 -> pr_skipped, auto_pr_disabled
2. 本地没有对应 worktree -> pr_skipped, no_worktree
3. worktree 里没有 git remote origin -> pr_skipped, no_remote
4. 本机没有 gh 二进制 -> pr_skipped, no_gh
5. gh 未登录(gh auth status 非零) -> pr_skipped,
gh_not_authenticated
只有以上全过才 push 分支 + gh pr create;push 或 pr create
失败时抛 PrStepError(调用方据此返回 502),任务状态保持
merged 不变,绝不把任务标 failed。
worker 不自动推:自动 PR 只由显式调 POST /tasks/{id}/pr
触发。默认自动推会把测试任务也推到远端,本模块不做。
本模块不存、不打任何密钥:gh 的登录凭证由 gh 自己管
(gh auth login),push/pr 的命令与错误文本里不出现任何
token;错误文本只含 git/gh 的原文摘要,不含远端 URL 里的
userinfo(若 remote URL 含 userinfo,摘要里把它抹掉)。
"""
import os
import re
import shutil
import subprocess
_SKIP_AUTO_PR = "auto_pr_disabled"
_SKIP_NO_WORKTREE = "no_worktree"
_SKIP_NO_REMOTE = "no_remote"
_SKIP_NO_GH = "no_gh"
_SKIP_NOT_AUTHED = "gh_not_authenticated"
class PrSkipped(Exception):
"""跳过:reason 见上列常量,调用方返回 200 + pr_skipped。"""
def __init__(self, reason):
super().__init__(reason)
self.reason = reason
class PrStepError(Exception):
"""push / pr create 哪一步失败:step 为 push 或 pr_create。"""
def __init__(self, step, message):
super().__init__(message)
self.step = step
self.message = message
def host_name_for_task(task_id):
"""cloud 任务 id -> host 任务名(与 sandbox._host_name 同规则)。"""
return re.sub(r"[^A-Za-z0-9_.-]", "-", str(task_id))
def default_worktrees_root():
"""默认 worktree 根:<仓库根>/work/host-worktrees。"""
here = os.path.dirname(os.path.abspath(__file__))
return os.path.join(os.path.dirname(here), "work", "host-worktrees")
def resolve_worktrees_root(override=None):
"""worktree 根优先级:显式参数 > 环境变量 > 默认。"""
if override:
return override
return os.environ.get("CLOUD_HOST_WORKTREES_DIR") or \
default_worktrees_root()
def _scrub_url(text):
"""抹掉 URL 里的 userinfo(https://user:pass@host -> https://host),
防止错误文本把凭证带出去。"""
if not text:
return text
return re.sub(r"(https?://)[^/@\s]+@", r"\1", text)
def _run(cmd, cwd=None, runner=None):
"""命令唯一入口:runner 供单测注入,返回带 returncode/stdout/
stderr 的对象;真实路径走 subprocess。"""
if runner is not None:
return runner(cmd, cwd)
return subprocess.run(cmd, cwd=cwd, capture_output=True, text=True)
def _git_remote_url(worktree_dir, git_runner=None):
r = _run(["git", "-C", worktree_dir, "remote", "get-url", "origin"],
cwd=worktree_dir, runner=git_runner)
url = (getattr(r, "stdout", "") or "").strip()
if getattr(r, "returncode", 1) != 0 or not url:
return None
return url
def check_pr_ready(worktree_dir, git_runner=None, gh_which=None,
gh_runner=None):
"""按顺序做跳过检查,全过返回 (worktree_dir, branch);任一步
不过抛 PrSkipped。branch 由调用方按 host 任务名算好传进来
的是 worktree 检查——本函数只查 worktree 与 remote/gh/登录,
分支名由 create_task_pr 组装。"""
if not os.path.isdir(worktree_dir):
raise PrSkipped(_SKIP_NO_WORKTREE)
if _git_remote_url(worktree_dir, git_runner) is None:
raise PrSkipped(_SKIP_NO_REMOTE)
gh_bin = gh_which if gh_which is not None else shutil.which("gh")
if not gh_bin:
raise PrSkipped(_SKIP_NO_GH)
r = _run(["gh", "auth", "status"], cwd=worktree_dir, runner=gh_runner)
if getattr(r, "returncode", 1) != 0:
raise PrSkipped(_SKIP_NOT_AUTHED)
return True
def create_task_pr(task_id, auto_pr, worktrees_root=None,
git_runner=None, gh_runner=None, gh_which=None):
"""给一个已 merged 的云端任务建 PR。
auto_pr 不为 1 时直接抛 PrSkipped(auto_pr_disabled),调用方
在查 worktree 之前就应先查项目 auto_pr——本函数也再查一次,
防止直接调用者漏查。成功返回 {pr_url, pr_number, branch}。
"""
if auto_pr != 1:
raise PrSkipped(_SKIP_AUTO_PR)
host_name = host_name_for_task(task_id)
worktree_dir = os.path.join(
resolve_worktrees_root(worktrees_root), host_name)
branch = "host/" + host_name
check_pr_ready(worktree_dir, git_runner=git_runner,
gh_which=gh_which, gh_runner=gh_runner)
push = _run(["git", "-C", worktree_dir, "push", "origin", branch],
cwd=worktree_dir, runner=git_runner)
if getattr(push, "returncode", 1) != 0:
raise PrStepError(
"push", "git push origin %s 失败:%s" % (
branch, _scrub_url(
(getattr(push, "stderr", "") or "")
or (getattr(push, "stdout", "") or ""))))
title = "cloud: %s" % task_id
r = _run(["gh", "pr", "create", "--title", title,
"--body", "cloud 任务 %s 的自动 PR。" % task_id,
"--head", branch],
cwd=worktree_dir, runner=gh_runner)
out = (getattr(r, "stdout", "") or "").strip()
if getattr(r, "returncode", 1) != 0 or not out:
raise PrStepError(
"pr_create", "gh pr create 失败:%s" % _scrub_url(
(getattr(r, "stderr", "") or "") or out))
pr_url = out.splitlines()[0].strip()
m = re.search(r"/(\d+)\s*$", pr_url)
pr_number = int(m.group(1)) if m else None
return {"pr_url": pr_url, "pr_number": pr_number, "branch": branch}
+91 -7
View File
@@ -593,14 +593,98 @@ def _docker_default_runner(cmd, timeout=None):
class CloudExecutor(Executor):
"""云端执行器(占位,未接通)。
"""云端执行器:可插拔的通用 HTTP runner 接口(厂商无关)。
execute(task) 一律 raise RuntimeError,绝不假装已接通云端。
接通前提:真实云账号 + 预算上限配置,届时由控制面注入凭证与
限额后再替换本占位实现。
本轮不假装接通某个云厂商(ECS、Cloud Run 等专属适配明确
不在范围)。它只做一件事:把任务 POST 到一个配置好的通用
runner endpoint,endpoint 返回任务状态,本类据此返回结果。
配置(只走参数或环境变量,不写文件):
- endpoint:显式参数 > CLOUD_RUNNER_ENDPOINT;未设时 execute
一律 raise RuntimeError,绝不静默降级到 Docker/Local。
- 认证:auth_header 参数或 CLOUD_RUNNER_AUTH_HEADER 给 header
名;token 值只从环境变量取(CLOUD_RUNNER_AUTH_TOKEN_ENV
指定的变量名,默认 CLOUD_RUNNER_AUTH_TOKEN),只放进请求头,
绝不写进日志、错误文本或返回结果。
- 请求体只有 id、prompt、repo_url(task["repo"]),不含任何密钥。
- http_client 供单测注入:http_client(url, payload, headers)
返回带 status 字段的 dict(或具 .status/.json 的对象);
真实路径走 urllib.request。
- dispatch_task/run_task:worker 的 run_one 走 dispatch_task,
这里直接委托给 execute,保证 cloud worker 池能真正派发。
"""
def __init__(self, endpoint=None, http_client=None, auth_header=None,
auth_token_env=None, env=None):
self._env = env if env is not None else os.environ
self.endpoint = endpoint or self._env.get("CLOUD_RUNNER_ENDPOINT")
self.http_client = http_client
self.auth_header = auth_header or self._env.get(
"CLOUD_RUNNER_AUTH_HEADER")
self.auth_token_env = auth_token_env or self._env.get(
"CLOUD_RUNNER_AUTH_TOKEN_ENV") or "CLOUD_RUNNER_AUTH_TOKEN"
def _headers(self):
headers = {"Content-Type": "application/json"}
if self.auth_header:
try:
token = self._env.get(self.auth_token_env)
except Exception:
token = None
if token:
headers[self.auth_header] = token
return headers
def _post(self, payload):
if self.http_client is not None:
return self.http_client(self.endpoint, payload, self._headers())
import urllib.request
req = urllib.request.Request(
self.endpoint,
data=json.dumps(payload).encode("utf-8"),
headers=self._headers(), method="POST")
with urllib.request.urlopen(req, timeout=30) as resp:
return json.loads(resp.read().decode("utf-8"))
def execute(self, task) -> dict:
raise RuntimeError(
"cloud executor not connected: missing cloud account "
"and budget limit")
if not isinstance(task, dict) or task.get("id") is None:
raise ValueError("任务缺少 id")
if not self.endpoint:
raise RuntimeError(
"cloud runner not configured: set CLOUD_RUNNER_ENDPOINT "
"to a generic runner endpoint; no cloud account/runner "
"configured, refusing to silently fall back to "
"Docker/Local")
payload = {
"id": task["id"],
"prompt": task.get("prompt") or "",
"repo_url": task.get("repo"),
}
result = self._post(payload)
if isinstance(result, dict):
status = result.get("status")
stdout = result.get("stdout", "")
else:
status = getattr(result, "status", None)
try:
stdout = result.json().get("stdout", "")
except Exception:
stdout = ""
if status == "completed":
final = "completed"
else:
final = "failed"
return {
"executor": "cloud",
"task_id": task["id"],
"status": final,
"stdout": _truncate_stdout(stdout),
"returncode": 0 if final == "completed" else 1,
"host_status": None,
}
def dispatch_task(self, task) -> dict:
return self.execute(task)
def run_task(self, task) -> dict:
return self.execute(task)
+176 -7
View File
@@ -59,6 +59,11 @@ from cloud.db import PostgresIntegrityError as _PgIntegrityError
from cloud.db import PostgresOperationalError as _PgOperationalError
from cloud.kb import KBStore
from cloud.memory import MemoryStore
from cloud.pr import PrSkipped, PrStepError, create_task_pr
from cloud.workers import (
heartbeat_worker as _registry_heartbeat,
list_workers as _registry_list,
)
from cloud.oauth import (
OAuthError,
OAuthStateStore,
@@ -227,6 +232,9 @@ def migrate_tasks_table(conn):
"started_at": "REAL",
"finished_at": "REAL",
"child_pid": "INTEGER",
"executor": "TEXT",
"pr_url": "TEXT",
"pr_number": "INTEGER",
}
for name, coltype in new_cols.items():
if name not in cols:
@@ -246,6 +254,8 @@ def migrate_projects_table(conn):
"PRAGMA table_info(projects)").fetchall()}
if "repo_url" not in cols:
conn.execute("ALTER TABLE projects ADD COLUMN repo_url TEXT")
if "auto_pr" not in cols:
conn.execute("ALTER TABLE projects ADD COLUMN auto_pr INTEGER")
def _is_postgres_backend():
@@ -275,7 +285,8 @@ def init_db(db_path):
"name TEXT, "
"owner_id TEXT, "
"created_at REAL, "
"repo_url TEXT)"
"repo_url TEXT, "
"auto_pr INTEGER)"
)
conn.execute(
"CREATE TABLE IF NOT EXISTS tasks("
@@ -346,6 +357,9 @@ def _task_to_dict(row):
"started_at": _row_get(row, "started_at"),
"finished_at": _row_get(row, "finished_at"),
"child_pid": _row_get(row, "child_pid"),
"executor": _row_get(row, "executor"),
"pr_url": _row_get(row, "pr_url"),
"pr_number": _row_get(row, "pr_number"),
}
@@ -358,6 +372,9 @@ def _project_to_dict(row):
# 老库行缺 repo_url 列时 _row_get 给 None,不炸(列本身由
# migrate_projects_table 在 init_db 时补,老行读到也是 None)。
"repo_url": _row_get(row, "repo_url"),
# auto_pr 默认 0(关闭):老行缺列时 None 按 0 计,绝不因
# 迁移就把老项目默认改成自动推。
"auto_pr": _row_get(row, "auto_pr") or 0,
}
@@ -569,6 +586,10 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
self._create_project(user_id)
elif path == "/projects" and method == "GET":
self._list_projects(user_id)
elif path == "/workers" and method == "GET":
self._workers_list(user_id)
elif path == "/workers/heartbeat" and method == "POST":
self._workers_heartbeat(user_id)
elif path == "/metrics" and method == "GET":
self._metrics(user_id)
elif path == "/models" and method == "GET":
@@ -595,6 +616,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
self._cancel_task(user_id, rest[: -len("/cancel")])
elif rest.endswith("/merge") and method == "POST":
self._merge_task(user_id, rest[: -len("/merge")])
elif rest.endswith("/pr") and method == "POST":
self._task_pr(user_id, rest[: -len("/pr")])
elif rest.endswith("/heartbeat") and method == "POST":
self._heartbeat_task(user_id, rest[: -len("/heartbeat")])
else:
@@ -626,6 +649,18 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
except ValueError as exc:
self._send_json(400, {"error": str(exc)})
return
# auto_pr 可选,默认 0(关闭):只接受 0/1(布尔 True/False
# 按 1/0 计),其它值 400 且不建项目。本轮没有现成的
# PATCH /projects/{id} 路由,auto_pr 只能在创建时指定,
# 建后改它需直接改库或重建项目(见 README)。
auto_pr_raw = body.get("auto_pr", 0)
if isinstance(auto_pr_raw, bool):
auto_pr = 1 if auto_pr_raw else 0
elif isinstance(auto_pr_raw, int) and auto_pr_raw in (0, 1):
auto_pr = auto_pr_raw
else:
self._send_json(400, {"error": "auto_pr must be 0 or 1"})
return
now = time.time()
conn = self._db()
try:
@@ -637,8 +672,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
return
conn.execute(
"INSERT INTO projects(id, name, owner_id, created_at,"
" repo_url) VALUES(?, ?, ?, ?, ?)",
(project_id, name, user_id, now, repo_url),
" repo_url, auto_pr) VALUES(?, ?, ?, ?, ?, ?)",
(project_id, name, user_id, now, repo_url, auto_pr),
)
conn.commit()
project = {
@@ -647,6 +682,7 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
"owner_id": user_id,
"created_at": now,
"repo_url": repo_url,
"auto_pr": auto_pr,
}
except (sqlite3.IntegrityError, _PgIntegrityError):
self._send_json(409, {"error": "project exists"})
@@ -671,7 +707,7 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
).fetchall()
}
rows = conn.execute(
"SELECT id, name, owner_id, created_at, repo_url"
"SELECT id, name, owner_id, created_at, repo_url, auto_pr"
" FROM projects ORDER BY created_at, id"
).fetchall()
finally:
@@ -1095,6 +1131,14 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
except (TypeError, ValueError):
self._send_json(400, {"error": "priority must be an integer"})
return
# executor 可选:None 表示任意 worker 可认领;指定时只能是
# local/docker/cloud,非法值 400 且不建任务。
executor = body.get("executor")
if executor is not None and executor not in (
"local", "docker", "cloud"):
self._send_json(400, {
"error": "executor must be local, docker or cloud"})
return
idempotency_key = body.get("idempotency_key")
if idempotency_key is not None and not isinstance(idempotency_key, str):
self._send_json(400, {"error": "idempotency_key must be a string"})
@@ -1120,10 +1164,11 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
conn.execute(
"INSERT INTO tasks(id, project_id, prompt, status,"
" idempotency_key, depends_on, priority, created_by,"
" created_at, updated_at, last_heartbeat)"
" VALUES(?, ?, ?, 'queued', ?, ?, ?, ?, ?, ?, NULL)",
" created_at, updated_at, last_heartbeat, executor)"
" VALUES(?, ?, ?, 'queued', ?, ?, ?, ?, ?, ?, NULL, ?)",
(task_id, project_id, prompt, idempotency_key,
json.dumps(depends_on), priority, user_id, now, now),
json.dumps(depends_on), priority, user_id, now, now,
executor),
)
conn.commit()
except (sqlite3.IntegrityError, _PgIntegrityError):
@@ -1270,6 +1315,124 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
store.audit(user_id, task["project_id"], "merge_task", task_id)
self._send_json(200, task)
# -- worker 池 -----------------------------------------------------
def _workers_list(self, user_id):
"""列出注册 worker:任何已认证用户可读(worker 是基础设施
信息,不属于某个项目的数据,不按项目成员过滤)。"""
workers = _registry_list(self.server.db_path)
self._send_json(200, {"workers": workers})
def _workers_heartbeat(self, user_id):
"""worker 心跳:body 含 worker_id(必填,缺了 400)与可选
executor/capacity。心跳本身就是注册:worker_id 不存在时
自动注册,绝不 404。"""
body, err = self._read_json()
if err:
self._send_json(400, {"error": err})
return
worker_id = body.get("worker_id")
if not worker_id or not isinstance(worker_id, str):
self._send_json(400, {"error": "worker_id required"})
return
executor = body.get("executor")
capacity = body.get("capacity")
if capacity is not None:
try:
capacity = int(capacity)
except (TypeError, ValueError):
self._send_json(400, {"error": "capacity must be an integer"})
return
try:
worker = _registry_heartbeat(
self.server.db_path, worker_id,
executor=executor, capacity=capacity)
except ValueError as exc:
self._send_json(400, {"error": str(exc)})
return
worker["alive"] = True
self._send_json(200, worker)
# -- 自动 PR(默认关闭,显式触发)----------------------------------
def _task_pr(self, user_id, task_id):
"""给已 merged 的任务建 PR(见 cloud/pr.py 的跳过规则)。
权限:项目 maintainer 以上(approve_merge,403);任务
不存在 404;非 merged 状态 409;跳过时返回 200 +
pr_skipped 与 reason,任务状态保持 merged 不变;push/pr
创建失败时返回 502 并写清哪一步失败,任务同样保持 merged。
成功时把 pr_url/pr_number 写回任务并记审计 pr_created。
"""
row = self._fetch_task(task_id)
if row is None:
self._send_json(404, {"error": "task not found"})
return
task = _task_to_dict(row)
with _auth_store(self.server.db_path) as store:
if not store.check_permission(
user_id, task["project_id"], "approve_merge"):
self._send_json(403, {"error": "forbidden"})
return
if task["status"] != "merged":
self._send_json(409, {
"error": "only merged tasks can create a PR",
"status": task["status"],
})
return
conn = self._db()
try:
prow = conn.execute(
"SELECT auto_pr FROM projects WHERE id = ?",
(task["project_id"],)).fetchone()
finally:
conn.close()
auto_pr = 0
if prow is not None:
try:
auto_pr = prow["auto_pr"] or 0
except (KeyError, IndexError, TypeError):
auto_pr = 0
try:
result = create_task_pr(
task_id, auto_pr,
worktrees_root=getattr(
self.server, "host_worktrees_dir", None),
git_runner=getattr(self.server, "pr_git_runner", None),
gh_runner=getattr(self.server, "pr_gh_runner", None),
gh_which=getattr(self.server, "pr_gh_which", None))
except PrSkipped as exc:
self._send_json(200, {
"task_id": task_id, "status": task["status"],
"pr_skipped": True, "reason": exc.reason,
})
return
except PrStepError as exc:
self._send_json(502, {
"error": exc.message, "step": exc.step,
"task_id": task_id, "status": task["status"],
})
return
now = time.time()
conn = self._db()
try:
conn.execute(
"UPDATE tasks SET pr_url = ?, pr_number = ?,"
" updated_at = ? WHERE id = ?",
(result["pr_url"], result["pr_number"], now, task_id),
)
conn.commit()
task = _task_to_dict(
conn.execute(
"SELECT * FROM tasks WHERE id = ?",
(task_id,)).fetchone())
finally:
conn.close()
with _auth_store(self.server.db_path) as store:
store.audit(user_id, task["project_id"], "pr_created",
result["pr_url"])
self._send_json(200, dict(task, pr_url=result["pr_url"],
pr_number=result["pr_number"],
branch=result["branch"]))
def _heartbeat_task(self, user_id, task_id):
row = self._fetch_task(task_id)
if row is None:
@@ -1857,6 +2020,12 @@ def create_server(db_path, host="127.0.0.1", port=8765, backend=None):
# 测试可注入 fake http_client / userinfo_client;绝不注入真实 secret。
server.oauth_http_client = None
server.oauth_userinfo_client = None
# 自动 PR 注入点:生产为 None(走真实 git/gh),测试可注入
# fake git/gh runner 与 worktree 根;绝不注入真实凭证。
server.pr_git_runner = None
server.pr_gh_runner = None
server.pr_gh_which = None
server.host_worktrees_dir = None
return server
+462
View File
@@ -0,0 +1,462 @@
"""多 worker 池派发 + 通用 CloudExecutor + 自动 PR 单测。
只用标准库,不跑真实模型回合、不联网、不登录任何服务;
git/gh 一律用假 runner 注入,不依赖本机是否装 gh、是否有
远端仓库;CloudExecutor 的 endpoint 一律用假 http_client。
"""
import http.client
import json
import os
import sqlite3
import sys
import tempfile
import threading
import time
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest import mock
try:
from cloud.auth import AuthStore
from cloud.pr import PrSkipped, create_task_pr
from cloud.sandbox import CloudExecutor, DockerExecutor
from cloud.server import create_server, init_db
from cloud.worker import Worker
from cloud.workers import (
heartbeat_worker,
list_workers,
register_worker,
)
except ImportError:
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from cloud.auth import AuthStore
from cloud.pr import PrSkipped, create_task_pr
from cloud.sandbox import CloudExecutor, DockerExecutor
from cloud.server import create_server, init_db
from cloud.worker import Worker
from cloud.workers import (
heartbeat_worker,
list_workers,
register_worker,
)
def _result(returncode=0, stdout="", stderr=""):
return SimpleNamespace(returncode=returncode, stdout=stdout,
stderr=stderr)
class _ServerCase(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.viewer_id, self.viewer_key = self.auth.create_user("viewer")
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()
status, self.project = self._req(
"POST", "/projects", {"name": "p1", "auto_pr": 1},
self.owner_key)
self.assertEqual(status, 201, self.project)
self.auth.set_member(
self.project["id"], self.owner_id, self.viewer_id, "viewer")
self.worktrees = os.path.join(self.tmp.name, "worktrees")
os.makedirs(self.worktrees, exist_ok=True)
self.server.host_worktrees_dir = self.worktrees
def tearDown(self):
self.server.shutdown()
self.server.server_close()
self.thread.join(timeout=5)
try:
self.auth.conn.close()
except Exception:
pass
self.tmp.cleanup()
def _req(self, method, path, body=None, key=None, headers=None):
conn = http.client.HTTPConnection(
"127.0.0.1", self.port, timeout=10)
hdrs = dict(headers or {})
payload = None
if body is not None:
payload = json.dumps(body).encode("utf-8")
hdrs.setdefault("Content-Type", "application/json")
if key is not None:
hdrs["X-API-Key"] = key
try:
conn.request(method, path, body=payload, headers=hdrs)
resp = conn.getresponse()
data = resp.read()
finally:
conn.close()
try:
parsed = json.loads(data.decode("utf-8")) if data else None
except ValueError:
parsed = None
return resp.status, parsed
def _insert_task(self, task_id, status, project_id=None):
conn = sqlite3.connect(self.db_path)
try:
conn.execute(
"INSERT INTO tasks(id, project_id, prompt, status,"
" depends_on, priority, created_by, created_at,"
" updated_at, last_heartbeat) VALUES(?, ?, ?, ?, '[]',"
" 0, 'system', 0, 0, NULL)",
(task_id, project_id or self.project["id"],
"do it", status))
conn.commit()
finally:
conn.close()
def _task_row(self, task_id):
conn = sqlite3.connect(self.db_path)
try:
row = conn.execute(
"SELECT status, pr_url, pr_number FROM tasks"
" WHERE id = ?", (task_id,)).fetchone()
finally:
conn.close()
return row
def _audit_actions(self):
conn = sqlite3.connect(self.db_path)
try:
return [r[0] for r in conn.execute(
"SELECT action FROM audit_log").fetchall()]
finally:
conn.close()
def _git_runner(self, push_fail=False, remote_fail=False):
def run(cmd, cwd):
if "remote" in cmd:
if remote_fail:
return _result(1, "", "no remote")
return _result(0, "https://example.com/repo.git\n")
if "push" in cmd:
if push_fail:
return _result(1, "", "push rejected")
return _result(0, "")
return _result(0, "")
return run
def _gh_runner(self, authed=True, pr_url=None):
def run(cmd, cwd):
if cmd[:2] == ["gh", "auth"]:
return _result(0 if authed else 1, "", "")
if cmd[:2] == ["gh", "pr"]:
return _result(0, (pr_url or
"https://example.com/repo/pull/42\n"))
return _result(0, "")
return run
def _worktree(self, task_id):
path = os.path.join(self.worktrees, task_id)
os.makedirs(path, exist_ok=True)
return path
class WorkerRegistryTestCase(_ServerCase):
def test_heartbeat_requires_worker_id(self):
status, body = self._req(
"POST", "/workers/heartbeat", {}, self.owner_key)
self.assertEqual(status, 400, body)
def test_heartbeat_registers_and_lists_alive(self):
status, body = self._req(
"POST", "/workers/heartbeat",
{"worker_id": "w-docker-1", "executor": "docker",
"capacity": 2},
self.owner_key)
self.assertEqual(status, 200, body)
self.assertTrue(body["alive"])
self.assertEqual(body["executor"], "docker")
status, body = self._req("GET", "/workers", None, self.owner_key)
self.assertEqual(status, 200, body)
self.assertEqual(len(body["workers"]), 1)
self.assertTrue(body["workers"][0]["alive"])
def test_heartbeat_preserves_executor_when_not_given(self):
self._req("POST", "/workers/heartbeat",
{"worker_id": "w1", "executor": "docker"}, self.owner_key)
status, body = self._req(
"POST", "/workers/heartbeat", {"worker_id": "w1"},
self.owner_key)
self.assertEqual(status, 200, body)
self.assertEqual(body["executor"], "docker")
def test_stale_worker_lists_not_alive(self):
register_worker(self.db_path, "w-old", executor="local")
conn = sqlite3.connect(self.db_path)
try:
conn.execute(
"UPDATE workers SET last_seen = ? WHERE worker_id = ?",
(time.time() - 3600, "w-old"))
conn.commit()
finally:
conn.close()
workers = list_workers(self.db_path)
self.assertEqual(len(workers), 1)
self.assertFalse(workers[0]["alive"])
def test_heartbeat_without_key_is_unauthorized(self):
status, body = self._req(
"POST", "/workers/heartbeat", {"worker_id": "w1"}, None)
self.assertEqual(status, 401, body)
class TaskExecutorFilterTestCase(unittest.TestCase):
def setUp(self):
self.tmp = tempfile.TemporaryDirectory()
init_db(os.path.join(self.tmp.name, "control.db"))
self.db_path = os.path.join(self.tmp.name, "control.db")
conn = sqlite3.connect(self.db_path)
try:
for tid, exc in (("t-local", "local"), ("t-none", None),
("t-docker", "docker")):
conn.execute(
"INSERT INTO tasks(id, project_id, prompt, status,"
" depends_on, priority, created_by, created_at,"
" updated_at, last_heartbeat, executor)"
" VALUES(?, 'proj', 'do it', 'queued', '[]', 0,"
" 'system', 0, 0, NULL, ?)", (tid, exc))
conn.commit()
finally:
conn.close()
def tearDown(self):
self.tmp.cleanup()
def _docker_worker(self):
return Worker(
self.db_path, "host.py", config_path=None, repo=None,
executor=DockerExecutor(runner=lambda cmd: _result()),
env={"XCM_API_KEY": "test-placeholder"}, env_key=None,
task_timeout=1800, max_retries=0, retry_backoff_seconds=0)
def test_claim_rejects_other_pool_task(self):
worker = self._docker_worker()
self.assertFalse(worker._claim("t-local", executor="docker"))
conn = sqlite3.connect(self.db_path)
try:
status = conn.execute(
"SELECT status FROM tasks WHERE id = 't-local'"
).fetchone()[0]
finally:
conn.close()
self.assertEqual(status, "queued")
def test_claim_accepts_unassigned_and_same_pool(self):
worker = self._docker_worker()
self.assertTrue(worker._claim("t-none", executor="docker"))
self.assertTrue(worker._claim("t-docker", executor="docker"))
def test_cas_race_only_one_worker_wins(self):
w1 = self._docker_worker()
w2 = self._docker_worker()
self.assertTrue(w1._claim("t-none", executor="docker"))
self.assertFalse(w2._claim("t-none", executor="docker"))
def test_pick_next_skips_other_pool(self):
# 只剩 local 池任务时,docker worker 无就绪任务。
conn = sqlite3.connect(self.db_path)
try:
conn.execute("DELETE FROM tasks WHERE id != 't-local'")
conn.commit()
finally:
conn.close()
worker = self._docker_worker()
self.assertIsNone(worker.pick_next())
def test_run_one_skips_other_pool_without_model(self):
conn = sqlite3.connect(self.db_path)
try:
conn.execute("DELETE FROM tasks WHERE id != 't-local'")
conn.commit()
finally:
conn.close()
worker = self._docker_worker()
result = worker.run_one()
self.assertFalse(result["dispatched"])
self.assertEqual(result["reason"], "no_ready_task")
class CloudExecutorTestCase(unittest.TestCase):
def test_unconfigured_fails_without_fallback(self):
ex = CloudExecutor(endpoint=None, env={})
with self.assertRaises(RuntimeError) as ctx:
ex.execute({"id": "t1", "prompt": "do it"})
self.assertIn("not configured", str(ctx.exception))
def test_configured_posts_only_id_prompt_repo_no_secrets(self):
seen = {}
def fake_client(url, payload, headers):
seen["url"] = url
seen["payload"] = payload
seen["headers"] = headers
return {"status": "completed", "stdout": "done"}
ex = CloudExecutor(
endpoint="https://runner.example.invalid/run",
http_client=fake_client,
auth_header="X-Runner-Token",
env={"CLOUD_RUNNER_AUTH_TOKEN": "test-token-value"})
result = ex.execute({"id": "t9", "prompt": "do it",
"repo": "https://example.com/r.git"})
self.assertEqual(result["status"], "completed")
self.assertEqual(seen["url"], "https://runner.example.invalid/run")
self.assertEqual(set(seen["payload"].keys()),
{"id", "prompt", "repo_url"})
self.assertEqual(seen["payload"]["repo_url"],
"https://example.com/r.git")
self.assertNotIn("test-token-value",
json.dumps(seen["payload"]))
self.assertEqual(seen["headers"]["X-Runner-Token"],
"test-token-value")
# dispatch_task 是 worker 走的入口,必须与 execute 同结果。
result2 = ex.dispatch_task({"id": "t9", "prompt": "do it"})
self.assertEqual(result2["status"], "completed")
def test_failed_status_maps_to_failed(self):
ex = CloudExecutor(
endpoint="https://runner.example.invalid/run",
http_client=lambda u, p, h: {"status": "failed",
"stdout": "boom"})
result = ex.execute({"id": "t1", "prompt": "x"})
self.assertEqual(result["status"], "failed")
class AutoPrTestCase(_ServerCase):
def _pr(self, task_id, key):
return self._req("POST", "/tasks/%s/pr" % task_id, {},
key)
def test_auto_pr_disabled_skips_and_stays_merged(self):
status, proj = self._req(
"POST", "/projects", {"name": "p-off"}, self.owner_key)
self.assertEqual(status, 201, proj)
self._insert_task("toff", "merged", project_id=proj["id"])
status, body = self._pr("toff", self.owner_key)
self.assertEqual(status, 200, body)
self.assertTrue(body["pr_skipped"])
self.assertEqual(body["reason"], "auto_pr_disabled")
self.assertEqual(self._task_row("toff")[0], "merged")
def test_auto_pr_illegal_value_rejected(self):
status, body = self._req(
"POST", "/projects", {"name": "p-bad", "auto_pr": 2},
self.owner_key)
self.assertEqual(status, 400, body)
def test_non_merged_returns_409(self):
self._insert_task("tq", "queued")
status, body = self._pr("tq", self.owner_key)
self.assertEqual(status, 409, body)
def test_viewer_cannot_pr(self):
self._insert_task("tm", "merged")
status, body = self._pr("tm", self.viewer_key)
self.assertEqual(status, 403, body)
def test_no_worktree_skips(self):
self._insert_task("tnw", "merged")
self.server.pr_git_runner = self._git_runner()
self.server.pr_gh_runner = self._gh_runner()
self.server.pr_gh_which = "gh"
status, body = self._pr("tnw", self.owner_key)
self.assertEqual(status, 200, body)
self.assertEqual(body["reason"], "no_worktree")
self.assertEqual(self._task_row("tnw")[0], "merged")
def test_no_remote_skips(self):
self._insert_task("tnr", "merged")
self._worktree("tnr")
self.server.pr_git_runner = self._git_runner(remote_fail=True)
self.server.pr_gh_runner = self._gh_runner()
self.server.pr_gh_which = "gh"
status, body = self._pr("tnr", self.owner_key)
self.assertEqual(status, 200, body)
self.assertEqual(body["reason"], "no_remote")
def test_no_gh_skips(self):
self._insert_task("tngh", "merged")
self._worktree("tngh")
self.server.pr_git_runner = self._git_runner()
self.server.pr_gh_runner = self._gh_runner()
self.server.pr_gh_which = ""
status, body = self._pr("tngh", self.owner_key)
self.assertEqual(status, 200, body)
self.assertEqual(body["reason"], "no_gh")
def test_gh_not_authenticated_skips(self):
self._insert_task("tna", "merged")
self._worktree("tna")
self.server.pr_git_runner = self._git_runner()
self.server.pr_gh_runner = self._gh_runner(authed=False)
self.server.pr_gh_which = "gh"
status, body = self._pr("tna", self.owner_key)
self.assertEqual(status, 200, body)
self.assertEqual(body["reason"], "gh_not_authenticated")
def test_push_failure_returns_502_and_stays_merged(self):
self._insert_task("tpf", "merged")
self._worktree("tpf")
self.server.pr_git_runner = self._git_runner(push_fail=True)
self.server.pr_gh_runner = self._gh_runner()
self.server.pr_gh_which = "gh"
status, body = self._pr("tpf", self.owner_key)
self.assertEqual(status, 502, body)
self.assertEqual(body["step"], "push")
self.assertEqual(self._task_row("tpf")[0], "merged")
def test_success_writes_back_pr_and_audit(self):
self._insert_task("tok", "merged")
self._worktree("tok")
self.server.pr_git_runner = self._git_runner()
self.server.pr_gh_runner = self._gh_runner()
self.server.pr_gh_which = "gh"
status, body = self._pr("tok", self.owner_key)
self.assertEqual(status, 200, body)
self.assertEqual(body["pr_url"],
"https://example.com/repo/pull/42")
self.assertEqual(body["pr_number"], 42)
self.assertEqual(body["branch"], "host/tok")
row = self._task_row("tok")
self.assertEqual(row[0], "merged")
self.assertEqual(row[1], "https://example.com/repo/pull/42")
self.assertEqual(row[2], 42)
self.assertIn("pr_created", self._audit_actions())
def test_pr_module_auto_pr_gate(self):
with self.assertRaises(PrSkipped) as ctx:
create_task_pr("x", 0)
self.assertEqual(ctx.exception.reason, "auto_pr_disabled")
class ProjectTaskFieldTestCase(_ServerCase):
def test_create_task_executor_validation(self):
status, body = self._req(
"POST", "/tasks",
{"project_id": self.project["id"], "prompt": "x",
"executor": "bogus"},
self.owner_key)
self.assertEqual(status, 400, body)
status, body = self._req(
"POST", "/tasks",
{"project_id": self.project["id"], "prompt": "x",
"executor": "docker"},
self.owner_key)
self.assertEqual(status, 201, body)
self.assertEqual(body["executor"], "docker")
if __name__ == "__main__":
unittest.main()
+113 -19
View File
@@ -101,6 +101,7 @@ import json
import os
import re
import signal
import socket
import subprocess
import sys
import time
@@ -113,8 +114,9 @@ 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 DockerExecutor, LocalExecutor
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(
@@ -178,6 +180,11 @@ def _task_from_row(row):
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"],
@@ -190,6 +197,9 @@ def _task_from_row(row):
"created_at": row["created_at"],
"updated_at": row["updated_at"],
"last_heartbeat": row["last_heartbeat"],
# 老库行缺 executor 列时给 None(表示任意 worker 可认领),
# 不炸;列本身由 _ensure_result_columns 在写回时补。
"executor": _col("executor"),
}
@@ -208,7 +218,8 @@ class Worker:
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):
repos_runner=None, heartbeat_file=None,
worker_id=None, capacity=1):
self.db_path = str(db_path)
self.host_py_path = host_py_path
self.config_path = config_path
@@ -259,6 +270,39 @@ class Worker:
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。
@@ -312,17 +356,34 @@ class Worker:
conn.close()
return _task_from_row(row) if row is not None else None
def _claim(self, task_id):
"""CAS 认领:仅当任务仍为 queued 时标 running 并记心跳。"""
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:
cur = conn.execute(
"UPDATE tasks SET status = 'running',"
" last_heartbeat = ?, updated_at = ?"
" WHERE id = ? AND status = 'queued'",
(now, now, task_id),
)
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:
@@ -355,7 +416,9 @@ class Worker:
for name, coltype in {
"result_text": "TEXT", "error_text": "TEXT",
"actual_model": "TEXT", "started_at": "REAL",
"finished_at": "REAL", "child_pid": "INTEGER"}.items():
"finished_at": "REAL", "child_pid": "INTEGER",
"executor": "TEXT", "pr_url": "TEXT",
"pr_number": "INTEGER"}.items():
if name not in cols:
try:
conn.execute(
@@ -484,6 +547,12 @@ class Worker:
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)-------------------
@@ -572,9 +641,13 @@ class Worker:
# -- 缺 key 门 ------------------------------------------------------
def _executor_name(self):
"""执行器名:DockerExecutor -> docker,其它 -> local。"""
"""执行器名: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;值本身绝不返回。"""
@@ -800,11 +873,15 @@ class Worker:
if task is None:
return {"dispatched": False, "reason": "no_ready_task"}
task_id = task["id"]
if not self._claim(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", task_id)
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()
@@ -982,6 +1059,9 @@ class Worker:
# 运维心跳:启动时先写一次,之后每轮结束再写,server 的
# /health 据此算 worker_ok(见 server.py 的健康检查)。
self._write_heartbeat()
# 多 worker 池注册:启动时注册一次,之后每轮结束写注册表
# 心跳(见本轮结束处);注册失败只记日志,不炸循环。
self.register()
started = time.time()
all_results = []
try:
@@ -1001,8 +1081,9 @@ class Worker:
results = []
all_results.extend(results)
# 每轮结束写心跳:哪怕本轮一个任务都没派,worker 活着
# 这件事本身也要被 server 看到。
# 这件事本身也要被 server 看到(文件心跳 + 注册表心跳)。
self._write_heartbeat()
self._registry_heartbeat()
# 按 0.2 秒切片睡,期间可响应 stop_event/_loop_stop。
deadline = time.time() + interval
while time.time() < deadline:
@@ -1032,7 +1113,11 @@ class Worker:
def _build_cli_executor(args, host_py):
"""CLI 的执行器构造:local 时返回 None(Worker 默认 LocalExecutor),
docker 时返回 DockerExecutor(env_key 名用 Worker 的解析规则,
值本身绝不从这里传——DockerExecutor 只查存在性,-e 只传名)。"""
值本身绝不从这里传——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(
@@ -1075,10 +1160,17 @@ def main(argv=None):
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"],
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")
@@ -1119,7 +1211,9 @@ def main(argv=None):
max_retries=args.max_retries,
retry_backoff_seconds=args.retry_backoff_seconds,
repos_dir=args.repos_dir,
heartbeat_file=args.heartbeat_file)
heartbeat_file=args.heartbeat_file,
worker_id=args.worker_id,
capacity=args.worker_capacity)
# CLI 启动清理兜底(常驻与单次都做):清掉上一轮残留容器。
worker.cleanup_stale_containers()
# CLI 起 worker 前先建表/迁移:全新库直接起循环不应因缺表而炸。
+166
View File
@@ -0,0 +1,166 @@
"""Worker 注册表:多 worker 池的注册、心跳与存活判断。
零第三方依赖,只用 Python 标准库。
设计:
- 注册表存在控制面数据库的 workers 表里(worker 与 server 共用
同一个库,所以 worker 直接写库注册,不需要为此再发 HTTP;
HTTP 侧的 GET /workers 与 POST /workers/heartbeat 只是同一张
表的 API 面,给有 API key 的外部 worker 用)。
- worker_id 默认由调用方给(生产 worker 应配稳定的 --worker-id);
未给时由 worker.py 用“主机名-进程号”生成——重启后进程号会变,
旧行会自然变 stale(alive=false),不要用随机 uuid 让注册表
每次重启都多一行无法对应到同一台机器。
- 心跳本身就是注册:heartbeat 时 worker_id 不存在就 INSERT,
存在就 UPDATE,绝不因“未注册”返回 404 报错。
- alive 定义:last_seen 在阈值内(默认 120 秒,可经环境变量
CLOUD_WORKER_ALIVE_SECONDS 调,测试可传参覆盖)。
- 本模块不存、不打任何密钥:注册信息只有 worker_id、executor、
capacity、last_seen,没有任何凭证字段。
"""
import os
import time
from cloud.db import connect as _db_connect
DEFAULT_ALIVE_SECONDS = 120
def alive_threshold_seconds(override=None):
"""存活阈值:显式参数 > 环境变量 > 默认 120 秒。"""
if override is not None:
try:
return float(override)
except (TypeError, ValueError):
return float(DEFAULT_ALIVE_SECONDS)
try:
return float(os.environ.get(
"CLOUD_WORKER_ALIVE_SECONDS") or DEFAULT_ALIVE_SECONDS)
except (TypeError, ValueError):
return float(DEFAULT_ALIVE_SECONDS)
def ensure_workers_table(db_path):
"""建 workers 表(不存在时建,已存在时不动)。"""
conn = _db_connect(str(db_path))
try:
conn.execute(
"CREATE TABLE IF NOT EXISTS workers("
"worker_id TEXT PRIMARY KEY, "
"executor TEXT, "
"capacity INTEGER, "
"last_seen REAL, "
"created_at REAL)"
)
conn.commit()
finally:
conn.close()
def register_worker(db_path, worker_id, executor=None, capacity=1):
"""注册(或刷新)一个 worker:不存在 INSERT,存在 UPDATE。
返回注册行 dict。worker_id 为空时抛 ValueError,调用方
(server 的 heartbeat 端点)据此返回 400。
"""
if not worker_id or not isinstance(worker_id, str):
raise ValueError("worker_id required")
ensure_workers_table(db_path)
now = time.time()
conn = _db_connect(str(db_path))
try:
cur = conn.execute(
"UPDATE workers SET executor = ?, capacity = ?,"
" last_seen = ? WHERE worker_id = ?",
(executor, capacity, now, worker_id),
)
if cur.rowcount == 0:
conn.execute(
"INSERT INTO workers(worker_id, executor, capacity,"
" last_seen, created_at) VALUES(?, ?, ?, ?, ?)",
(worker_id, executor, capacity, now, now),
)
conn.commit()
row = conn.execute(
"SELECT * FROM workers WHERE worker_id = ?",
(worker_id,)).fetchone()
finally:
conn.close()
return _worker_to_dict(row)
def heartbeat_worker(db_path, worker_id, executor=None, capacity=None):
"""心跳:就是带 last_seen 刷新的注册。
executor/capacity 未给时保留旧值(已注册的 worker 只刷时间,
不把 executor 洗成 None);未注册时按给定值(或默认)注册。
"""
if not worker_id or not isinstance(worker_id, str):
raise ValueError("worker_id required")
ensure_workers_table(db_path)
now = time.time()
conn = _db_connect(str(db_path))
try:
row = conn.execute(
"SELECT * FROM workers WHERE worker_id = ?",
(worker_id,)).fetchone()
if row is None:
conn.execute(
"INSERT INTO workers(worker_id, executor, capacity,"
" last_seen, created_at) VALUES(?, ?, ?, ?, ?)",
(worker_id, executor, capacity if capacity is not None
else 1, now, now),
)
else:
old = _worker_to_dict(row)
conn.execute(
"UPDATE workers SET executor = ?, capacity = ?,"
" last_seen = ? WHERE worker_id = ?",
(executor if executor is not None else old["executor"],
capacity if capacity is not None else old["capacity"],
now, worker_id),
)
conn.commit()
row = conn.execute(
"SELECT * FROM workers WHERE worker_id = ?",
(worker_id,)).fetchone()
finally:
conn.close()
return _worker_to_dict(row)
def list_workers(db_path, alive_seconds=None):
"""列出全部 worker,每行含 alive(last_seen 在阈值内)。"""
threshold = alive_threshold_seconds(alive_seconds)
ensure_workers_table(db_path)
now = time.time()
conn = _db_connect(str(db_path))
try:
rows = conn.execute(
"SELECT * FROM workers ORDER BY worker_id").fetchall()
finally:
conn.close()
out = []
for row in rows:
item = _worker_to_dict(row)
last = item.get("last_seen")
item["alive"] = bool(
last is not None and (now - last) <= threshold)
out.append(item)
return out
def _worker_to_dict(row):
def _get(name):
try:
return row[name]
except (KeyError, IndexError, TypeError):
return None
return {
"worker_id": _get("worker_id"),
"executor": _get("executor"),
"capacity": _get("capacity"),
"last_seen": _get("last_seen"),
"created_at": _get("created_at"),
}