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:
1 parent
d3e588e4b3
commit
306cdf0d0c
8 files changed
+1256
-34
No files matched your search
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
@@ -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 前先建表/迁移:全新库直接起循环不应因缺表而炸。
|
||||
|
||||
@@ -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"),
|
||||
}
|
||||
Reference in new issue
Block a user