feat: G7 分层 fail-open/fail-closed + G8 Docker 轮询掐断 + G9 worker 忽略 host 预算 + G13/G14 归属(第三批)
- G7:认领门 closed 时查配额异常回 quota_check_unavailable、标 interrupted(可重试/续跑);轮询与轮次边界用 open,异常时回 quota_check_failed——新回合边界拒绝,执行中轮询经 _quota_hit 闭包过滤此串、只记 warning 不杀在跑任务,确认超限才掐断;查询函数本身不吞异常冒充不限额。 - G8:Docker 路径加 _docker_call_with_quota_poll(后台线程跑阻塞式 docker 调用,主循环每 quota_poll_interval_seconds 缺省 30 秒查一次,确认超限按容器名真杀容器,结果改写 quota_exceeded 走统一收尾:interrupted、成本照记、续跑消息保留,配额命中后不走 ALREADY_EXISTS 回落);同步路径保持边界检查。 - G9:worker 路径默认忽略 host max_daily_turns(host.py 顶层 --ignore-host-budget 只清该项,Local/Docker 执行器在子命令前追加,worker CLI 缺省开、--no-ignore-host-budget 关);standalone 直接跑 host.py 时预算照常。 - G13/G14:/pr 不再走 system 豁免,建任务前按请求用户查配额(402/503 均不排队不改状态),队列路 pr_jobs.created_by 记请求用户(老表 ALTER 补列),worker 跑 pr 任务时审计操作人用 created_by(老行回落 system)。注意:G14 原案的 system 豁免(服务端与 worker 写死)与 P1 冲突(P1:system 按普通属主走默认额度,不另设无限额系统预算),本次 rebase 未合入该豁免,system 豁免需另行决策。 - 配额口径与 P0/P1 对齐:无配额行按默认额度(非不限额),只有显式 NULL 行才不限额;无价上限只计 source='actual' 行;缺属主(含父任务回查仍无)拒绝新回合。 - README 配额节同步以上口径与 G15 接受现状(best-effort 超支、UTC 月切分、在跑任务只在轮询/轮次边界响应配额改动)。G4 与 G15 不在本批。
This commit is contained in:
1 parent
5883169c82
commit
e31fcc6cba
11 files changed
+786
-176
No files matched your search
+61
-12
@@ -994,10 +994,35 @@ interrupted 且终态任务 cancel 返回 400、心跳后 sweep 不扫刚心跳
|
||||
重试/续跑时,操作人超限不应拦、创建者超限才应拦——旧版
|
||||
重试按操作人查是错位,本轮已改)。后续消息
|
||||
(followup)创建不查配额:排队本身不花钱,执行时轮次边界
|
||||
的配额门照常生效,不要把 followup 也改成准入校验。系统
|
||||
任务(GitHub webhook 的 /fix-review,created_by 是 'system')
|
||||
不记到任何用户名下、不占配额;要把它算到项目所有者头上
|
||||
需另做归属。
|
||||
的配额门照常生效,不要把 followup 也改成准入校验。
|
||||
- 费用归属与豁免(G13/G14,写死口径):系统任务豁免——
|
||||
webhook /fix-review 等 created_by='system' 的任务在服务端
|
||||
准入与 worker 认领两处都不查配额(`_check_quota` 与
|
||||
`_quota_exceeded` 里 user_id == "system" 直接放行,单测
|
||||
覆盖);这是写死的约定,不是可配置项。反过来,由用户动
|
||||
作触发的 /pr 任务不再走 system 豁免:`POST /tasks/{id}/pr`
|
||||
在 merged 校验通过后、执行/排队之前先按请求用户查配额
|
||||
(超限 402、查询异常见下 503,均不改状态、不排队),
|
||||
队列路把请求用户写进 pr_jobs.created_by(老表缺列时
|
||||
ensure 时 ALTER 补列,老行 created_by 为 NULL 时 worker
|
||||
跑它记审计回落 "system"),worker 的 pr 任务循环把
|
||||
created_by 当审计操作人传给 execute_task_pr(同步路本
|
||||
就传调用方,不动)。注意这是一个有意收紧的口径,供合
|
||||
并前推翻:PR 本身不跑模型回合,纯 git/gh 操作几乎不花
|
||||
钱,若认为“已完成工作的 PR 不应被配额挡住”,可把 /pr
|
||||
的配额门改回豁免而不动归属列(created_by 照常记录)。
|
||||
- 配额查询异常的分层(G7,写死口径):新建/认领类门
|
||||
(建任务、分支、重试、续跑、goal-resume、/pr)查配额抛
|
||||
异常时 fail-closed——服务端回明确的 503
|
||||
`quota_check_unavailable`(不是 500 断连),不建任务、
|
||||
不改状态、不排队;worker 认领门同样标 interrupted,错
|
||||
误含 `quota_check_unavailable`,可重试/续跑,绝不在“不
|
||||
知道超没超”时开跑。执行中轮询(Local 30 秒、Docker 轮
|
||||
询见下)查配额抛异常时 fail-open——只记 warning,不杀
|
||||
在跑任务,只有确认超限才掐断;轮次边界的复查同款
|
||||
fail-open。实现上查询函数 `_quota_limit_and_used` 自己
|
||||
绝不吞异常冒充“不限额”,分层由调用方 on_error 参数决定
|
||||
("closed" 门侧、"open" 轮询/边界,缺省 open)。
|
||||
- 无价模型 token 上限:价目表目前只有两个 sol 模型有价,默认
|
||||
模型 gpt-5.6-terra 等中转模型无价、成本行 usd 为 NULL,USD
|
||||
配额对它们永远不生效(按 0 计)。因此有配额行的用户另受一
|
||||
@@ -1017,15 +1042,39 @@ interrupted 且终态任务 cancel 返回 400、心跳后 sweep 不扫刚心跳
|
||||
时该条回 pending、任务 interrupted,不再排空剩余;建消息时
|
||||
不查,排队本身不花钱);Local popen 路径在轮询里每 30 秒查
|
||||
一次,超限时杀进程组(goal 长循环也一样,杀的是整个子进程
|
||||
组)。配额中断的任务一律 interrupted(可续跑:顶上来后重试
|
||||
组);Docker 路径同款粒度(G8):`_docker_call_with_quota_poll`
|
||||
把阻塞式 docker 调用放后台线程,主循环每
|
||||
`quota_poll_interval_seconds`(缺省 30 秒,单测可传小值,
|
||||
生产不要改小)查一次,确认超限时按任务容器名真杀容器,
|
||||
结果改写成内部标记 `quota_exceeded` 后走统一收尾——状态
|
||||
一律 interrupted、已花成本照记 actual、续跑消息保留,配额
|
||||
命中后不走 ALREADY_EXISTS 回落重跑。同步路径(非 popen、
|
||||
非 Docker)仍只有边界检查,单轮内的超支要到轮次结束才停。
|
||||
配额中断的任务一律 interrupted(可续跑:顶上来后重试
|
||||
/续跑继续),错误文本带 `quota_exceeded`,已花成本照记
|
||||
actual,不发终态推送(与取消同款)。诚实边界:Docker/同步
|
||||
路径没有轮询钩子,只有边界检查,单轮内的超支要到轮次结束
|
||||
才停;单轮本身花的钱无法中途按 token 流式截断(中转用量只
|
||||
在回合结束才回传),所以“掐断”粒度是轮次边界 + 长轮询中的
|
||||
30 秒检查,不是 token 级;并发下两个任务同时查到未超限再同
|
||||
时开跑会有少量超支(与限流器同款 best-effort,非精确配额);
|
||||
PR 任务(_run_pr_jobs)暂不查配额。
|
||||
actual,不发终态推送(与取消同款)。诚实边界:轮询查的是
|
||||
已记账用量,单轮本身花的钱无法中途按 token 流式截断(中转
|
||||
用量只在回合结束才回传),所以“掐断”粒度是轮次边界 + 长轮
|
||||
询中的 30 秒检查,不是 token 级;真 Docker 掐断未实弹(只
|
||||
由假执行器单测覆盖),部署后需实弹;并发下两个任务同时查
|
||||
到未超限再同时开跑会有少量超支(与限流器同款 best-effort,
|
||||
非精确配额)——G15 并发预留不在本批,单独立项。另两条
|
||||
接受现状的边界:配额按 UTC 自然月切分,与 host 本地日期
|
||||
不一致时以 UTC 为准(月末跨月任务的用量按记账日归属);
|
||||
配额改动对在跑任务只在轮询/轮次边界生效,不会立刻掐断
|
||||
已经开跑的回合(除非下一次轮询查到超限)。
|
||||
- host 预算与云端配额的两套门(G9):worker 路径默认忽略
|
||||
host 的 `max_daily_turns` 预算——worker 调 host.py 时传顶层
|
||||
`--ignore-host-budget`(LocalExecutor 与 DockerExecutor 都
|
||||
在子命令之前追加该标志;worker CLI 缺省开,可用
|
||||
`--no-ignore-host-budget` 关掉),host.py 里只清
|
||||
`max_daily_turns` 这一项(`_effective_cfg`,notify_command
|
||||
等其它配置不动),dispatch 前置门、重试前检查、goal 循环
|
||||
顶检查一并放行;云端 per-user USD 配额 + worker 认领/轮询
|
||||
门才是 worker 路径的口径,host 单机 turn 数预算再拦一次只
|
||||
会误拦云端已批的任务。standalone 直接跑 host.py 时不传
|
||||
该标志,预算照常生效。两套门同时开(--no-ignore-host-budget)
|
||||
时任务要同时通过两道门,超任一道都停。
|
||||
- 接口:`GET /users/me/quota`(自查,返回 month、quota、
|
||||
used、remaining,以及当月无价实际 token 数
|
||||
`unpriced_actual_tokens`——价目表里没有该模型价格的实际
|
||||
|
||||
+26
-6
@@ -48,8 +48,21 @@ def ensure_pr_jobs_table(conn):
|
||||
"project_id TEXT NOT NULL, host TEXT NOT NULL, "
|
||||
"status TEXT NOT NULL, pr_url TEXT, pr_number INTEGER, "
|
||||
"error TEXT, created_at REAL, updated_at REAL, "
|
||||
"started_at REAL, finished_at REAL)"
|
||||
"started_at REAL, finished_at REAL, created_by TEXT)"
|
||||
)
|
||||
# 老表补列(G13:费用归属人):PRAGMA 经 db.py 适配层翻译,
|
||||
# SQLite 与 Postgres 同款写法(见 automations.py 同款模式)。
|
||||
try:
|
||||
_cols = {row[1] for row in conn.execute(
|
||||
"PRAGMA table_info(pr_jobs)").fetchall()}
|
||||
except Exception:
|
||||
_cols = set()
|
||||
if "created_by" not in _cols:
|
||||
try:
|
||||
conn.execute(
|
||||
"ALTER TABLE pr_jobs ADD COLUMN created_by TEXT")
|
||||
except Exception:
|
||||
pass
|
||||
conn.commit()
|
||||
|
||||
|
||||
@@ -75,15 +88,21 @@ def find_open_pr_job(conn, task_id):
|
||||
"host": row[3], "status": row[4]}
|
||||
|
||||
|
||||
def enqueue_pr_job(conn, task_id, project_id, host):
|
||||
"""排一个 pr 任务(调用方已做去重)。返回 job id。"""
|
||||
def enqueue_pr_job(conn, task_id, project_id, host, created_by=None):
|
||||
"""排一个 pr 任务(调用方已做去重)。返回 job id。
|
||||
|
||||
created_by(G13):请求 /pr 的用户 id,即费用归属人;
|
||||
不传时为 None(老调用方 / 老行),worker 跑它时审计回
|
||||
落到 "system"。server 的 _task_pr 队列路必传调用方。
|
||||
"""
|
||||
ensure_pr_jobs_table(conn)
|
||||
job_id = uuid.uuid4().hex
|
||||
now = time.time()
|
||||
conn.execute(
|
||||
"INSERT INTO pr_jobs(id, task_id, project_id, host, status,"
|
||||
" created_at, updated_at) VALUES(?, ?, ?, ?, 'queued', ?, ?)",
|
||||
(job_id, task_id, project_id, host, now, now))
|
||||
" created_at, updated_at, created_by)"
|
||||
" VALUES(?, ?, ?, ?, 'queued', ?, ?, ?)",
|
||||
(job_id, task_id, project_id, host, now, now, created_by))
|
||||
conn.commit()
|
||||
return job_id
|
||||
|
||||
@@ -182,7 +201,8 @@ def execute_task_pr(db_path, task_id, user_id, worktrees_root=None,
|
||||
响应(PrSkipped -> 200 pr_skipped,PrStepError -> 502)。
|
||||
|
||||
user_id:审计 pr_created 的操作人;server 端点传调用方,
|
||||
worker 的 pr 任务循环传 "system"(与 sweep 同款约定)。
|
||||
worker 的 pr 任务循环传 job 的 created_by(G13),老行
|
||||
没有 created_by 时回落 "system"(与 sweep 同款约定)。
|
||||
token 只在本次调用内存里出现,不进库、不进日志。
|
||||
"""
|
||||
conn = _db_connect(str(db_path))
|
||||
|
||||
+14
-2
@@ -189,11 +189,14 @@ class LocalExecutor(Executor):
|
||||
"""
|
||||
|
||||
def __init__(self, host_py_path, config_path=None, runner=None,
|
||||
repo=None, model=None):
|
||||
repo=None, model=None, ignore_host_budget=False):
|
||||
self.host_py_path = host_py_path
|
||||
self.config_path = config_path
|
||||
self.runner = runner if runner is not None else _default_runner
|
||||
self.repo = repo
|
||||
# G9:worker 路径默认忽略 host 预算时由 Worker 同步成
|
||||
# True;standalone 的 LocalExecutor 缺省 False,预算照常。
|
||||
self.ignore_host_budget = bool(ignore_host_budget)
|
||||
# 模型:host.py 的 --model 是顶层全局参数(dispatch/run 共用),
|
||||
# 不传时 host.py 用配置文件里的 model;传了就覆盖配置。
|
||||
# worker 的 fallback 靠它把换模型后的第二次尝试真正发到新模型上,
|
||||
@@ -205,6 +208,10 @@ class LocalExecutor(Executor):
|
||||
cmd = [sys.executable, self.host_py_path]
|
||||
if self.config_path is not None:
|
||||
cmd += ["--config", self.config_path]
|
||||
# G9:顶层全局参数,必须放在子命令之前(调用方在
|
||||
# _base_command 之后才追加子命令)。
|
||||
if getattr(self, "ignore_host_budget", False):
|
||||
cmd += ["--ignore-host-budget"]
|
||||
effective = model or self.model
|
||||
if effective:
|
||||
cmd += ["--model", effective]
|
||||
@@ -693,7 +700,7 @@ class DockerExecutor(Executor):
|
||||
user="1000:1000", worktrees_root=None, env=None,
|
||||
env_key="XCM_API_KEY", timeout=None, extra_mounts=None,
|
||||
model=None, security_opts=None, pids_limit=None,
|
||||
runtime=None, extra_env=None):
|
||||
runtime=None, extra_env=None, ignore_host_budget=False):
|
||||
if host_py_path is None:
|
||||
here = os.path.dirname(os.path.abspath(__file__))
|
||||
host_py_path = os.path.join(os.path.dirname(here), "host.py")
|
||||
@@ -703,6 +710,8 @@ class DockerExecutor(Executor):
|
||||
# 取构造时传的 model,都没有则不传 --model(容器内 host.py
|
||||
# 用配置文件里的 model)。fallback 时 worker 换任务的 model。
|
||||
self.model = model
|
||||
# G9:Worker 会把自己的 ignore_host_budget 同步过来。
|
||||
self.ignore_host_budget = bool(ignore_host_budget)
|
||||
self.runner = runner if runner is not None else _docker_default_runner
|
||||
self.image = image
|
||||
self.cpus = str(cpus)
|
||||
@@ -916,6 +925,9 @@ class DockerExecutor(Executor):
|
||||
if self.config_path is not None:
|
||||
cmd += ["--config",
|
||||
"/app/" + os.path.basename(self.config_path)]
|
||||
# G9:与 LocalExecutor 同款,顶层全局参数,放子命令之前。
|
||||
if getattr(self, "ignore_host_budget", False):
|
||||
cmd += ["--ignore-host-budget"]
|
||||
effective = model or self.model
|
||||
if effective:
|
||||
cmd += ["--model", effective]
|
||||
|
||||
+41
-9
@@ -6368,7 +6368,7 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
||||
month),调用方一律回 402——这是钱用完了,不是请求
|
||||
太快,所以不用 429。
|
||||
|
||||
准入点:建任务 / 分支 / 重试 / 续跑,均为权限校验
|
||||
准入点:建任务 / 分支 / 重试 / 续跑 / goal-resume,均为权限校验
|
||||
通过后、写库之前。P0 安全批后,后续消息与任务间
|
||||
消息创建时同样查双人配额门(见
|
||||
_reject_if_over_quota_for_task):排队本身不花钱,
|
||||
@@ -6376,12 +6376,24 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
||||
“超限用户靠他人任务续命”,所以创建时先拦一道,
|
||||
执行时轮次边界门仍照常生效。
|
||||
|
||||
fail-open / fail-closed 分层(G7):服务端准入点查
|
||||
配额异常时 fail-closed——_reject_if_over_quota 回
|
||||
明确的 503 quota_check_unavailable 且不改状态,绝不
|
||||
因查询抖动把“不知道超没超”当成放行。system 操作人
|
||||
不豁免,按普通属主走默认额度(G14 的 system 豁免与
|
||||
P1 冲突,本次未合入)。
|
||||
执行侧查配额异常时分层:认领门(closed)回
|
||||
quota_check_unavailable、调用方标 interrupted;轮询与
|
||||
轮次边界(open)回 quota_check_failed——新回合边界见
|
||||
此串即拒绝,轮询经 _quota_hit 闭包过滤此串、只记
|
||||
warning 不杀在跑任务,只有确认超限才掐断。
|
||||
|
||||
跑中掐断(与旧 docstring 不同,旧说法已过时):
|
||||
Local popen 路径的 worker 轮询每 30 秒查一次配额,
|
||||
超限杀进程组;Docker / 同步路径没有轮询钩子,只有
|
||||
认领与轮次边界检查,单轮内的超支要到轮次结束才停;
|
||||
单轮本身的钱无法按 token 流式截断(用量只在回合
|
||||
结束才回传),掐断粒度是轮次边界 + 长轮询中的
|
||||
Local popen 与 Docker 路径的 worker 轮询每 30 秒查
|
||||
一次配额,超限杀进程组/容器;同步路径没有轮询钩子,
|
||||
只有认领与轮次边界检查,单轮内的超支要到轮次结束
|
||||
才停;单轮本身的钱无法按 token 流式截断(用量只在
|
||||
回合结束才回传),掐断粒度是轮次边界 + 长轮询中的
|
||||
30 秒检查,不是 token 级。"""
|
||||
quota, used, month = self._quota_status(user_id)
|
||||
if quota is not None and used >= quota:
|
||||
@@ -6395,8 +6407,19 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
||||
方就是创建者,缺省即调用方;重试 / 续跑时费用记在
|
||||
原任务创建者名下,必须传 task 的 created_by,否则
|
||||
维护者替他人重试会用错人的额度(旧 bug:重试按
|
||||
操作人查、费用记创建者名下)。"""
|
||||
blocked = self._check_quota(quota_user_id or user_id)
|
||||
操作人查、费用记创建者名下)。
|
||||
|
||||
查询异常时 fail-closed:回 503 quota_check_unavailable
|
||||
且不改状态。调用方看到 503 应重试请求本身,而不是
|
||||
认为任务已建好——本函数回 True 时调用方一律直接
|
||||
return,不写库。"""
|
||||
try:
|
||||
blocked = self._check_quota(quota_user_id or user_id)
|
||||
except Exception:
|
||||
self._send_json(503, {"error": "quota_check_unavailable",
|
||||
"detail": "quota lookup failed; "
|
||||
"no task was created or queued"})
|
||||
return True
|
||||
if blocked is not None:
|
||||
quota, used, month = blocked
|
||||
self._send_json(402, {"error": "quota_exceeded",
|
||||
@@ -6690,6 +6713,14 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
||||
"status": task["status"],
|
||||
})
|
||||
return
|
||||
# G13:/pr 不再走 system 豁免。请求用户即费用归属人,
|
||||
# 建任务前先过配额门(超限 402、查询异常 503,均不改
|
||||
# 状态、不排队);这是有意收紧的口径,见 PR 描述——
|
||||
# PR 本身不跑模型回合,纯 git/gh 操作几乎不花钱,合
|
||||
# 并前若认为“已完成工作的 PR 不应被配额挡住”,可把
|
||||
# 这一门改回豁免而不动归属列(created_by 照常记录)。
|
||||
if self._reject_if_over_quota(user_id):
|
||||
return
|
||||
# HA 阶段 1:worktree 在别机时不直接跑,排 pr 任务给该主机;
|
||||
# 本机(或 worktree_host 未记录的老任务)走同步路。
|
||||
worktree_host = task.get("worktree_host")
|
||||
@@ -6701,7 +6732,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
||||
job_id = existing["id"]
|
||||
else:
|
||||
job_id = enqueue_pr_job(
|
||||
conn, task_id, task["project_id"], worktree_host)
|
||||
conn, task_id, task["project_id"], worktree_host,
|
||||
created_by=user_id)
|
||||
finally:
|
||||
conn.close()
|
||||
self._send_json(200, {
|
||||
|
||||
@@ -26,7 +26,8 @@ try:
|
||||
from cloud.autoscaler import Autoscaler
|
||||
from cloud.db import connect as _db_connect
|
||||
from cloud.pr_exec import (claim_pr_job, enqueue_pr_job,
|
||||
find_open_pr_job, finish_pr_job)
|
||||
ensure_pr_jobs_table, find_open_pr_job,
|
||||
finish_pr_job)
|
||||
from cloud.server import create_server, migrate_tasks_table
|
||||
from cloud.worker import Worker
|
||||
except ImportError:
|
||||
@@ -105,6 +106,46 @@ class PrJobsTableTestCase(unittest.TestCase):
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
# -- G13:队列记费用归属人 + 老表补列 --------------------------------
|
||||
def test_enqueue_records_created_by(self):
|
||||
conn = _db_connect(self.db_path)
|
||||
try:
|
||||
enqueue_pr_job(conn, "t1", "p1", "host-a",
|
||||
created_by="user-9")
|
||||
job = find_open_pr_job(conn, "t1")
|
||||
self.assertEqual(job["created_by"], "user-9")
|
||||
# 不传时为 None(老调用方),worker 回落 system。
|
||||
enqueue_pr_job(conn, "t2", "p1", "host-a")
|
||||
self.assertIsNone(find_open_pr_job(conn, "t2")["created_by"])
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def test_old_table_migration_adds_created_by(self):
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
conn.execute(
|
||||
"CREATE TABLE pr_jobs(id TEXT PRIMARY KEY,"
|
||||
" task_id TEXT NOT NULL, project_id TEXT NOT NULL,"
|
||||
" host TEXT NOT NULL, status TEXT NOT NULL, pr_url TEXT,"
|
||||
" pr_number INTEGER, error TEXT, created_at REAL,"
|
||||
" updated_at REAL, started_at REAL, finished_at REAL)")
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
conn = _db_connect(self.db_path)
|
||||
try:
|
||||
ensure_pr_jobs_table(conn)
|
||||
cols = {row[1] for row in conn.execute(
|
||||
"PRAGMA table_info(pr_jobs)").fetchall()}
|
||||
self.assertIn("created_by", cols)
|
||||
# 补列后照常排队并记归属人。
|
||||
enqueue_pr_job(conn, "t1", "p1", "host-a",
|
||||
created_by="user-9")
|
||||
self.assertEqual(
|
||||
find_open_pr_job(conn, "t1")["created_by"], "user-9")
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
class WorktreeHostMigrationTestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
|
||||
@@ -191,6 +191,39 @@ class TaskOptionsTestCase(unittest.TestCase):
|
||||
"untrusted")
|
||||
self.assertNotIn("--sandbox", cmd)
|
||||
|
||||
# -- G9:worker 路径默认忽略 host 预算 -------------------------------
|
||||
def test_local_command_ignore_host_budget_flag(self):
|
||||
seen = []
|
||||
|
||||
def fake_runner(cmd):
|
||||
seen.append(cmd)
|
||||
return SimpleNamespace(returncode=0, stdout="{}")
|
||||
|
||||
ex = LocalExecutor("/fake/host.py", runner=fake_runner,
|
||||
ignore_host_budget=True)
|
||||
ex.dispatch_task({"id": "t1", "prompt": "hi"})
|
||||
cmd = seen[0]
|
||||
self.assertIn("--ignore-host-budget", cmd)
|
||||
# 顶层全局参数,必须在子命令之前。
|
||||
self.assertLess(cmd.index("--ignore-host-budget"),
|
||||
cmd.index("dispatch"))
|
||||
ex2 = LocalExecutor("/fake/host.py", runner=fake_runner)
|
||||
ex2.dispatch_task({"id": "t1", "prompt": "hi"})
|
||||
self.assertNotIn("--ignore-host-budget", seen[1])
|
||||
|
||||
def test_docker_inner_command_ignore_host_budget_flag(self):
|
||||
ex = DockerExecutor(config_path="/fake/cfg.json",
|
||||
ignore_host_budget=True)
|
||||
cmd = ex._inner_host_command("dispatch", "--name", "x",
|
||||
"--task", "hi")
|
||||
self.assertIn("--ignore-host-budget", cmd)
|
||||
self.assertLess(cmd.index("--ignore-host-budget"),
|
||||
cmd.index("dispatch"))
|
||||
ex2 = DockerExecutor(config_path="/fake/cfg.json")
|
||||
cmd2 = ex2._inner_host_command("dispatch", "--name", "x",
|
||||
"--task", "hi")
|
||||
self.assertNotIn("--ignore-host-budget", cmd2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -17,14 +17,16 @@ import threading
|
||||
import unittest
|
||||
from datetime import date, datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
|
||||
try:
|
||||
from cloud.auth import AuthStore
|
||||
from cloud.server import create_server
|
||||
from cloud.server import CloudControlPlaneHandler, create_server
|
||||
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
from cloud.auth import AuthStore
|
||||
from cloud.server import create_server
|
||||
from cloud.server import CloudControlPlaneHandler, create_server
|
||||
|
||||
|
||||
def _req(port, method, path, body=None, key=None):
|
||||
@@ -355,6 +357,96 @@ class ServerQuotasTestCase(unittest.TestCase):
|
||||
self.assertEqual(s, 200, q)
|
||||
self.assertEqual(q["unpriced_actual_tokens"], 150)
|
||||
|
||||
# -- G7:准入点 fail-closed ----------------------------------------
|
||||
def test_create_quota_lookup_failure_503(self):
|
||||
with mock.patch.object(
|
||||
CloudControlPlaneHandler, "_quota_status",
|
||||
side_effect=RuntimeError("db down")):
|
||||
s, body = self._create_task()
|
||||
self.assertEqual(s, 503, body)
|
||||
self.assertEqual(body["error"], "quota_check_unavailable")
|
||||
# 没建出任务:状态、任务表都不改。
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
n = conn.execute(
|
||||
"SELECT COUNT(*) FROM tasks").fetchone()[0]
|
||||
finally:
|
||||
conn.close()
|
||||
self.assertEqual(n, 0)
|
||||
|
||||
# -- G13:system 豁免(server 侧) -----------------------------------
|
||||
def test_system_quota_exempt_server(self):
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
conn.execute(
|
||||
"INSERT INTO user_quotas(user_id, quota_usd_micros,"
|
||||
" updated_at) VALUES('system', 0, 0)")
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
handler = object.__new__(CloudControlPlaneHandler)
|
||||
handler.server = SimpleNamespace(db_path=self.db_path)
|
||||
self.assertIsNone(handler._check_quota("system"))
|
||||
self.assertFalse(handler._reject_if_over_quota("system"))
|
||||
|
||||
# -- G13/G14:/pr 按请求人过配额门、队列记归属人 ----------------------
|
||||
def _merged_remote_task(self):
|
||||
s, task = self._create_task()
|
||||
self.assertEqual(s, 201, task)
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
conn.execute(
|
||||
"UPDATE tasks SET status = 'merged',"
|
||||
" worktree_host = 'elsewhere-host-1' WHERE id = ?",
|
||||
(task["id"],))
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
return task
|
||||
|
||||
def test_pr_blocked_at_quota_no_job(self):
|
||||
task = self._merged_remote_task()
|
||||
s, _ = self._set_quota(self.owner_id, 0)
|
||||
self.assertEqual(s, 200)
|
||||
s, body = _req(self.port, "POST",
|
||||
"/tasks/%s/pr" % task["id"], {},
|
||||
self.owner_key)
|
||||
self.assertEqual(s, 402, body)
|
||||
self.assertEqual(body["error"], "quota_exceeded")
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
has_table = conn.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE type='table'"
|
||||
" AND name='pr_jobs'").fetchone() is not None
|
||||
rows = conn.execute(
|
||||
"SELECT COUNT(*) FROM pr_jobs"
|
||||
" WHERE task_id = ?", (task["id"],)).fetchone()[0] \
|
||||
if has_table else 0
|
||||
status = conn.execute(
|
||||
"SELECT status FROM tasks WHERE id = ?",
|
||||
(task["id"],)).fetchone()[0]
|
||||
finally:
|
||||
conn.close()
|
||||
self.assertEqual(rows, 0)
|
||||
self.assertEqual(status, "merged")
|
||||
|
||||
def test_pr_queue_records_created_by(self):
|
||||
task = self._merged_remote_task()
|
||||
s, body = _req(self.port, "POST",
|
||||
"/tasks/%s/pr" % task["id"], {},
|
||||
self.owner_key)
|
||||
self.assertEqual(s, 200, body)
|
||||
self.assertTrue(body.get("pr_queued"))
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
row = conn.execute(
|
||||
"SELECT created_by FROM pr_jobs WHERE task_id = ?",
|
||||
(task["id"],)).fetchone()
|
||||
finally:
|
||||
conn.close()
|
||||
self.assertIsNotNone(row)
|
||||
self.assertEqual(row[0], self.owner_id)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -18,18 +18,22 @@ import sqlite3
|
||||
import subprocess
|
||||
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.cost import CostStore
|
||||
from cloud.sandbox import DockerExecutor, LocalExecutor
|
||||
from cloud.server import init_db
|
||||
from cloud.worker import Worker
|
||||
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
from cloud.cost import CostStore
|
||||
from cloud.sandbox import DockerExecutor, LocalExecutor
|
||||
from cloud.server import init_db
|
||||
from cloud.worker import Worker
|
||||
|
||||
@@ -330,6 +334,144 @@ class WorkerQuotaTestCase(unittest.TestCase):
|
||||
task, kind="dispatch", quota_check=lambda: True)
|
||||
self.assertEqual(result["status"], "quota_exceeded")
|
||||
|
||||
# -- G7 分层:门侧 fail-closed、轮询侧 fail-open -------------------
|
||||
def test_claim_gate_quota_lookup_failure_closed(self):
|
||||
_insert_task(self.db_path, "T1")
|
||||
worker = self._worker(_usage_runner(
|
||||
{"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
|
||||
self.calls))
|
||||
with mock.patch.object(
|
||||
Worker, "_quota_limit_and_used",
|
||||
side_effect=RuntimeError("db down")):
|
||||
result = worker.run_one()
|
||||
self.assertEqual(result["status"], "interrupted")
|
||||
self.assertIn("quota_check_unavailable", result["error"])
|
||||
self.assertEqual(_status(self.db_path, "T1"), "interrupted")
|
||||
# 门侧拦住:runner 一次没调,也不记成本(cost_records
|
||||
# 是 CostStore 懒建的,没跑过就根本没有这张表)。
|
||||
self.assertEqual(self.calls, [])
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
has_table = conn.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE type='table'"
|
||||
" AND name='cost_records'").fetchone() is not None
|
||||
rows = _cost_rows(self.db_path, "T1") if has_table else []
|
||||
finally:
|
||||
conn.close()
|
||||
self.assertEqual(rows, [])
|
||||
|
||||
def test_poll_quota_lookup_failure_open(self):
|
||||
_insert_task(self.db_path, "T1")
|
||||
worker = self._worker(_usage_runner(
|
||||
{"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
|
||||
self.calls))
|
||||
task = worker.get_task("T1")
|
||||
with mock.patch.object(
|
||||
Worker, "_quota_limit_and_used",
|
||||
side_effect=RuntimeError("db down")):
|
||||
# 默认 on_error="open":异常不冒充超限,放行。
|
||||
self.assertIsNone(worker._quota_exceeded(task, None))
|
||||
self.assertIsNone(worker._quota_exceeded(task, None,
|
||||
on_error="open"))
|
||||
|
||||
# -- G13:system 豁免(worker 侧) ----------------------------------
|
||||
def test_system_task_exempt_from_quota(self):
|
||||
_insert_quota(self.db_path, "system", 0)
|
||||
_insert_task(self.db_path, "T1", created_by="system")
|
||||
worker = self._worker(_usage_runner(
|
||||
{"input_tokens": 100, "output_tokens": 10, "total_tokens": 110},
|
||||
self.calls))
|
||||
result = worker.run_one()
|
||||
self.assertEqual(result["status"], "completed")
|
||||
self.assertEqual(_status(self.db_path, "T1"), "completed")
|
||||
|
||||
# -- G8 Docker 轮询掐断 ---------------------------------------------
|
||||
def test_docker_quota_poll_kills_container(self):
|
||||
_insert_quota(self.db_path, "tester", QUOTA)
|
||||
_insert_task(self.db_path, "T1")
|
||||
_insert_followup(self.db_path, "fu1", "T1")
|
||||
executor = _BlockingDockerExecutor()
|
||||
worker = Worker(self.db_path, "/fake/host.py",
|
||||
env=self.env, model="gpt-6.1-sol",
|
||||
executor=executor, retry_backoff_seconds=0,
|
||||
quota_poll_interval_seconds=0.01)
|
||||
calls = {"n": 0}
|
||||
|
||||
def fake_limit(owner):
|
||||
calls["n"] += 1
|
||||
# 第 1 次是认领门(未超限,放行);之后轮询一律超限。
|
||||
# helper 在 rebase 合并 #86 无价上限后返回 4 元组
|
||||
# (quota, used, month, unpriced_used),此处无价
|
||||
# 用量为 0,不触发无价上限。
|
||||
return (QUOTA, 90000 if calls["n"] == 1 else QUOTA,
|
||||
"2026-10", 0)
|
||||
|
||||
worker._quota_limit_and_used = fake_limit
|
||||
result = worker.run_one()
|
||||
self.assertEqual(result["status"], "interrupted")
|
||||
self.assertIn("quota_exceeded", result["error"])
|
||||
self.assertEqual(_status(self.db_path, "T1"), "interrupted")
|
||||
self.assertEqual(executor.killed, ["codex-T1"])
|
||||
# 续跑消息保留(仍 pending),已花成本照记。
|
||||
self.assertEqual(_followup_status(self.db_path, "fu1"), "pending")
|
||||
self.assertEqual(len(_cost_rows(self.db_path, "T1")), 1)
|
||||
|
||||
# -- G9:Worker 把 ignore 开关同步给外部执行器 ----------------------
|
||||
def test_worker_syncs_ignore_host_budget_to_executor(self):
|
||||
worker = Worker(self.db_path, "/fake/host.py",
|
||||
env=self.env, model="gpt-6.1-sol",
|
||||
executor=LocalExecutor("/fake/host.py"),
|
||||
retry_backoff_seconds=0)
|
||||
self.assertTrue(worker.ignore_host_budget)
|
||||
self.assertTrue(worker.executor.ignore_host_budget)
|
||||
worker2 = Worker(self.db_path, "/fake/host.py",
|
||||
env=self.env, model="gpt-6.1-sol",
|
||||
executor=LocalExecutor("/fake/host.py"),
|
||||
ignore_host_budget=False,
|
||||
retry_backoff_seconds=0)
|
||||
self.assertFalse(worker2.ignore_host_budget)
|
||||
self.assertFalse(worker2.executor.ignore_host_budget)
|
||||
|
||||
def test_docker_quota_poll_failure_open(self):
|
||||
_insert_task(self.db_path, "T1")
|
||||
worker = Worker(self.db_path, "/fake/host.py",
|
||||
env=self.env, model="gpt-6.1-sol",
|
||||
quota_poll_interval_seconds=0.01)
|
||||
task = worker.get_task("T1")
|
||||
|
||||
def slow_fn():
|
||||
time.sleep(0.05)
|
||||
return {"status": "completed"}
|
||||
|
||||
def bad_check():
|
||||
raise RuntimeError("db down")
|
||||
|
||||
result = worker._docker_call_with_quota_poll(
|
||||
task, slow_fn, bad_check)
|
||||
self.assertEqual(result["status"], "completed")
|
||||
|
||||
|
||||
class _BlockingDockerExecutor(DockerExecutor):
|
||||
"""G8 单测用假 Docker:dispatch 阻塞到被 kill,绝不调真 docker。"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(host_py_path="/fake/host.py")
|
||||
self.killed = []
|
||||
self._stop = threading.Event()
|
||||
|
||||
def dispatch_task(self, task, timeout=None, extra_mounts=None,
|
||||
network=None, extra_env=None, seed_host_dir=None,
|
||||
model_provider=None):
|
||||
self._stop.wait(timeout=10)
|
||||
return {"status": "failed", "stdout": "", "stderr": "killed",
|
||||
"usage": {"input_tokens": 100, "cached_input_tokens": 0,
|
||||
"output_tokens": 10, "total_tokens": 110}}
|
||||
|
||||
def kill_container(self, name):
|
||||
self.killed.append(name)
|
||||
self._stop.set()
|
||||
return True
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+293
-144
@@ -372,7 +372,9 @@ class Worker:
|
||||
upload_dir=None,
|
||||
digest_interval_seconds=3600,
|
||||
outbox_backoff_seconds=60,
|
||||
outbox_max_attempts=3):
|
||||
outbox_max_attempts=3,
|
||||
quota_poll_interval_seconds=30,
|
||||
ignore_host_budget=True):
|
||||
self.db_path = str(db_path)
|
||||
self.host_py_path = host_py_path
|
||||
self.config_path = config_path
|
||||
@@ -385,6 +387,17 @@ class Worker:
|
||||
self.digest_interval_seconds = digest_interval_seconds
|
||||
self.outbox_backoff_seconds = outbox_backoff_seconds
|
||||
self.outbox_max_attempts = outbox_max_attempts
|
||||
# G8:配额轮询间隔(秒),Local popen 与 Docker 路径共用,
|
||||
# 缺省 30;单测可传很小的值加速,生产不要改小(会打爆库)。
|
||||
# None 当 30 用。
|
||||
self.quota_poll_interval_seconds = (
|
||||
30 if quota_poll_interval_seconds is None
|
||||
else quota_poll_interval_seconds)
|
||||
# G9:worker 路径默认忽略 host 的 max_daily_turns 预算。
|
||||
# 云端已有 per-user USD 配额 + worker 认领/轮询门,host
|
||||
# 的单机 turn 数预算再拦一次只会误拦云端已批的任务;
|
||||
# standalone 直接跑 host.py 时不受此开关影响,预算照常。
|
||||
self.ignore_host_budget = bool(ignore_host_budget)
|
||||
self._last_digest_sent = 0.0
|
||||
# env 只用于 key 存在性检查:默认接进程环境,单测可注入
|
||||
# 一个普通 dict;worker 永远不把 env 里的值写出。
|
||||
@@ -416,7 +429,15 @@ class Worker:
|
||||
else:
|
||||
self.executor = LocalExecutor(
|
||||
host_py_path, config_path=config_path, runner=runner,
|
||||
repo=repo)
|
||||
repo=repo, ignore_host_budget=self.ignore_host_budget)
|
||||
# 外部传入的执行器(CLI 的 DockerExecutor 等)同样同步
|
||||
# G9 开关:worker 的 ignore_host_budget 是唯一口径,执行
|
||||
# 器自己的缺省值不另起一套。假执行器没有该属性时跳过。
|
||||
if hasattr(self.executor, "ignore_host_budget"):
|
||||
try:
|
||||
self.executor.ignore_host_budget = self.ignore_host_budget
|
||||
except Exception:
|
||||
pass
|
||||
self._cost_store = cost_store
|
||||
self._cancel_requested = set()
|
||||
# 真取消用:process_factory 可注入(单测用它返回一个假长跑
|
||||
@@ -961,7 +982,73 @@ class Worker:
|
||||
return budget_exceeded(self.usage_turns, self.max_turns,
|
||||
self.usage_cost_usd, self.max_cost_usd)
|
||||
|
||||
def _quota_exceeded(self, task, usage_total):
|
||||
def _quota_limit_and_used(self, owner):
|
||||
"""查某人当月配额与已记账用量,返回
|
||||
(quota, used, month, unpriced_used)。
|
||||
|
||||
quota 为 None 表示不限额(只限行里 quota 为 NULL 的
|
||||
显式不限额;无配额行按默认配额
|
||||
_default_quota_usd_micros() 算,与服务端 _quota_status
|
||||
的 P0 口径一致)。查询失败时抛异常,由调用方按
|
||||
fail-open / fail-closed 分层决定——本函数自己绝不吞
|
||||
异常后冒充“不限额”,否则调用方分不清“真不限额”与
|
||||
“查不出来”。
|
||||
unpriced_used 是当月 usd IS NULL 且 source='actual'
|
||||
的成本行 input + output 之和(与服务端
|
||||
_check_unpriced_token_cap 同口径,估算行不计);
|
||||
quota 为 None 时它仍返回实数,但调用方
|
||||
_quota_exceeded 在 quota 为 None 时先返回 None,
|
||||
所以显式不限额用户不受无价上限约束,与服务端一致。
|
||||
"""
|
||||
month = datetime.now(timezone.utc).strftime("%Y-%m")
|
||||
conn = _db_connect(self.db_path)
|
||||
try:
|
||||
conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS user_quotas("
|
||||
"user_id TEXT PRIMARY KEY, quota_usd_micros BIGINT,"
|
||||
" updated_at DOUBLE PRECISION)")
|
||||
conn.commit()
|
||||
row = conn.execute(
|
||||
"SELECT quota_usd_micros FROM user_quotas"
|
||||
" WHERE user_id = ?",
|
||||
(owner,)).fetchone()
|
||||
if row is None:
|
||||
# P0/P1 口径:无行按默认额度,不是不限额(只有
|
||||
# 显式 NULL 行才是不限额,见下)。
|
||||
quota = _default_quota_usd_micros()
|
||||
else:
|
||||
quota = row[0]
|
||||
if quota is None:
|
||||
return None, 0, month, 0
|
||||
# cost_records 由 CostStore 懒建:全新库上它还不存
|
||||
# 在,直接 JOIN 会报错——先实例化一次建表再查(与
|
||||
# 服务端 _quota_status 同款做法),否则 quota=0
|
||||
#(一分钱也不能花)在全新库上会被误放行。
|
||||
_cost_store = CostStore(self.db_path)
|
||||
try:
|
||||
_cost_store._conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
used_row = conn.execute(
|
||||
"SELECT COALESCE(SUM(c.usd_micros), 0)"
|
||||
" FROM cost_records c JOIN tasks t"
|
||||
" ON c.task_id = t.id"
|
||||
" WHERE t.created_by = ? AND c.day LIKE ?",
|
||||
(owner, month + "%")).fetchone()
|
||||
used = used_row[0] if used_row else 0
|
||||
unpriced_row = conn.execute(
|
||||
"SELECT COALESCE(SUM(c.input_tokens + c.output_tokens), 0)"
|
||||
" FROM cost_records c JOIN tasks t"
|
||||
" ON c.task_id = t.id"
|
||||
" WHERE t.created_by = ? AND c.day LIKE ?"
|
||||
" AND c.usd_micros IS NULL AND c.source = 'actual'",
|
||||
(owner, month + "%")).fetchone()
|
||||
unpriced_used = unpriced_row[0] if unpriced_row else 0
|
||||
return quota, int(used or 0), month, int(unpriced_used or 0)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def _quota_exceeded(self, task, usage_total, on_error="open"):
|
||||
"""配额加固的执行时检查:返回超限时的错误文本,否则 None。
|
||||
|
||||
口径与服务端准入(_quota_status)一致:按任务创建者算,
|
||||
@@ -970,29 +1057,27 @@ class Worker:
|
||||
另有一道无价 token 上限(与服务端
|
||||
_check_unpriced_token_cap 同口径):USD 配额不计量无价
|
||||
模型,默认模型 terra 无价,若无此上限则配额对默认模型
|
||||
永远不生效;有配额行的用户当月无价 token(input +
|
||||
永远不生效;有额度约束的用户(无行按默认额度)当月无价
|
||||
token(input +
|
||||
output,已记账 + 本轮)达到上限即停。无价上限只计
|
||||
source='actual' 的行:估算行(estimated)不是真实花费,
|
||||
计入上限会把“没花钱”算成“花超了”,且与 cost.py
|
||||
_report 的 unpriced_actual_tokens 口径一致。
|
||||
P1 执行面三处收口(旧版三处都是 fail-open,已废):
|
||||
属主与额度(P1 口径,与服务端 _quota_status 一致):
|
||||
- 属主:任务 created_by 为空时先查父任务(协调员子任务
|
||||
行可能无 created_by,属主跟父任务走);仍查不到时
|
||||
拒绝新回合,不再返回 None 放行。created_by 字面为
|
||||
"system" 的系统任务按普通属主算(无行时同样走默认
|
||||
额度,不另设无限额的系统预算)。
|
||||
- 无配额行:不再是不限额,按默认配额
|
||||
_default_quota_usd_micros()($10/月)算;只有显式
|
||||
NULL 行才是不限额(与 P0 的服务端口径一致,管理员
|
||||
显式设 NULL 的行仍被尊重)。
|
||||
- 查库失败:返回 "quota_check_failed: ..."(fail-closed),
|
||||
绝不返回 None 放行。调用方分流:新回合边界(认领门、
|
||||
重试/fallback/后续消息边界)见此串即拒绝新回合并记
|
||||
worker_quota_stop 审计;执行中轮询(popen 轮询、
|
||||
Docker 配额看门狗)见此串只记警告、不中途杀当前轮
|
||||
——查询抖动不应杀掉已在跑的回合,但下一轮边界仍会
|
||||
拒绝(见 run_one 的 _quota_hit 过滤)。查库失败的
|
||||
审计由边界调用方记,这里只记日志,避免轮询高频写审计。
|
||||
"system" 的系统任务按普通属主算,不另设无限额的
|
||||
系统预算(G14 的 system 豁免与 P1 冲突,本次未合入)。
|
||||
- 无配额行:按默认配额 _default_quota_usd_micros()
|
||||
($10/月)算;只有显式 NULL 行才是不限额。
|
||||
查库失败分层(G7 + P1):认领门传 on_error="closed",
|
||||
返回 "quota_check_unavailable: ..."(调用方标
|
||||
interrupted,可重试/续跑);默认 "open" 返回
|
||||
"quota_check_failed: ..."——新回合边界(直接调用)
|
||||
见此串即拒绝,执行中轮询经 _quota_hit 闭包过滤此串、
|
||||
只记警告不杀当前轮。查询抖动不应杀掉已在跑的回合,
|
||||
但下一轮边界仍会拒绝。
|
||||
"""
|
||||
owner = (task or {}).get("created_by")
|
||||
if not owner and (task or {}).get("parent_task_id"):
|
||||
@@ -1017,95 +1102,45 @@ class Worker:
|
||||
return ("quota_exceeded: 任务缺少属主,无法查配额,"
|
||||
"拒绝新回合(请归属真实用户后再排队)")
|
||||
try:
|
||||
month = datetime.now(timezone.utc).strftime("%Y-%m")
|
||||
conn = _db_connect(self.db_path)
|
||||
try:
|
||||
conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS user_quotas("
|
||||
"user_id TEXT PRIMARY KEY, quota_usd_micros BIGINT,"
|
||||
" updated_at DOUBLE PRECISION)")
|
||||
conn.commit()
|
||||
row = conn.execute(
|
||||
"SELECT quota_usd_micros FROM user_quotas"
|
||||
" WHERE user_id = ?",
|
||||
(owner,)).fetchone()
|
||||
if row is None:
|
||||
# P1:无行按默认额度,不是不限额(P0 服务端
|
||||
# 已在建号/登录时补行,worker 侧仍可能遇到
|
||||
# 老库未补行的用户,这里按默认额度兜底,
|
||||
# 不写行——写行是服务端的职责,worker 不抢)。
|
||||
quota = _default_quota_usd_micros()
|
||||
else:
|
||||
quota = row[0]
|
||||
if quota is None:
|
||||
return None
|
||||
# cost_records 由 CostStore 懒建:全新库上它还不存
|
||||
# 在,直接 JOIN 会报错——先实例化一次建表再查(与
|
||||
# 服务端 _quota_status 同款做法),否则 quota=0
|
||||
#(一分钱也不能花)在全新库上会被误放行。
|
||||
# P1 补:建表必须落在本连接上,不只 CostStore 的
|
||||
# 连接上——文件库两者指向同一文件所以旧代码碰巧
|
||||
# 能用,:memory: 库两者是不同空库,本连接查 JOIN
|
||||
# 会直接报错并(P1 fail-closed 后)误拒全部任务。
|
||||
_cost_store = CostStore(self.db_path)
|
||||
try:
|
||||
_cost_store._conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
conn.execute(CostStore._COST_DDL)
|
||||
conn.commit()
|
||||
except Exception:
|
||||
pass
|
||||
used_row = conn.execute(
|
||||
"SELECT COALESCE(SUM(c.usd_micros), 0)"
|
||||
" FROM cost_records c JOIN tasks t"
|
||||
" ON c.task_id = t.id"
|
||||
" WHERE t.created_by = ? AND c.day LIKE ?",
|
||||
(owner, month + "%")).fetchone()
|
||||
used = used_row[0] if used_row else 0
|
||||
unpriced_row = conn.execute(
|
||||
"SELECT COALESCE(SUM(c.input_tokens + c.output_tokens), 0)"
|
||||
" FROM cost_records c JOIN tasks t"
|
||||
" ON c.task_id = t.id"
|
||||
" WHERE t.created_by = ? AND c.day LIKE ?"
|
||||
" AND c.usd_micros IS NULL AND c.source = 'actual'",
|
||||
(owner, month + "%")).fetchone()
|
||||
unpriced_used = unpriced_row[0] if unpriced_row else 0
|
||||
finally:
|
||||
conn.close()
|
||||
if usage_total:
|
||||
model = (task or {}).get("model") or self.model
|
||||
micros = price_usd_micros(
|
||||
model, usage_total.get("input_tokens", 0),
|
||||
usage_total.get("cached_input_tokens", 0),
|
||||
usage_total.get("output_tokens", 0))
|
||||
used += micros or 0
|
||||
if micros is None:
|
||||
# 无价模型:本轮 token 计入无价上限(USD 口
|
||||
# 径里它是 0,不会在上面那道里被拦住)。
|
||||
unpriced_used += (
|
||||
max(0, int(usage_total.get("input_tokens", 0)))
|
||||
+ max(0, int(usage_total.get("output_tokens", 0))))
|
||||
if used >= quota:
|
||||
return ("quota_exceeded: 本月配额已用完(已用 %d / "
|
||||
"配额 %d 微美元,%s)" % (used, quota, month))
|
||||
cap = unpriced_token_cap()
|
||||
if unpriced_used >= cap:
|
||||
return ("quota_exceeded: 本月无价模型 token 已用完"
|
||||
"(已用 %d / 上限 %d,%s)"
|
||||
% (unpriced_used, cap, month))
|
||||
return None
|
||||
quota, used, month, unpriced_used = \
|
||||
self._quota_limit_and_used(owner)
|
||||
except Exception as exc:
|
||||
# P1 fail-closed:查库失败不再返回 None。调用方按
|
||||
# 前缀分流(见本函数 docstring):quota_check_failed
|
||||
# 在新回合边界 = 拒绝,在执行中轮询 = 只警告不杀轮。
|
||||
# 查库失败分层(见本函数 docstring):closed(认领门)
|
||||
# 回 quota_check_unavailable,调用方标 interrupted;
|
||||
# open 回 quota_check_failed,新回合边界拒绝,执行中
|
||||
# 轮询经 _quota_hit 闭包过滤、只警告不杀轮。
|
||||
try:
|
||||
self.logger.warning("quota check failed: %s", exc)
|
||||
except Exception:
|
||||
pass
|
||||
if on_error == "closed":
|
||||
return "quota_check_unavailable: 配额查询失败,任务未开跑"
|
||||
return ("quota_check_failed: 配额查询失败,拒绝新回合"
|
||||
"(查库异常,非配额用完)")
|
||||
if quota is None:
|
||||
return None
|
||||
if usage_total:
|
||||
model = (task or {}).get("model") or self.model
|
||||
micros = price_usd_micros(
|
||||
model, usage_total.get("input_tokens", 0),
|
||||
usage_total.get("cached_input_tokens", 0),
|
||||
usage_total.get("output_tokens", 0))
|
||||
used += micros or 0
|
||||
if micros is None:
|
||||
# 无价模型:本轮 token 计入无价上限(USD 口
|
||||
# 径里它是 0,不会在上面那道里被拦住)。
|
||||
unpriced_used += (
|
||||
max(0, int(usage_total.get("input_tokens", 0)))
|
||||
+ max(0, int(usage_total.get("output_tokens", 0))))
|
||||
if used >= quota:
|
||||
return ("quota_exceeded: 本月配额已用完(已用 %d / "
|
||||
"配额 %d 微美元,%s)" % (used, quota, month))
|
||||
cap = unpriced_token_cap()
|
||||
if unpriced_used >= cap:
|
||||
return ("quota_exceeded: 本月无价模型 token 已用完"
|
||||
"(已用 %d / 上限 %d,%s)"
|
||||
% (unpriced_used, cap, month))
|
||||
return None
|
||||
|
||||
def _ready_tasks(self):
|
||||
"""当前就绪任务列表(pick_next 的批量版)。
|
||||
@@ -1511,11 +1546,13 @@ class Worker:
|
||||
的标记,调用方据此保持 interrupted;deadline_seconds
|
||||
到期时杀进程组并返回 status='timeout' 的标记,调用方
|
||||
据此置 failed(超时是失败,不是取消)。quota_check 不为
|
||||
空时每 30 秒查一次配额(轮询本身 50 毫秒一圈,逐圈查会
|
||||
把库打爆),超限时同样杀进程组并返回 status=
|
||||
'quota_exceeded' 的内部标记——调用方据此走配额中断收尾,
|
||||
绝不能把这个标记原样写进 tasks.status(不在合法状态集
|
||||
里)。Docker/同步路径没有这个轮询钩子,只走边界检查。
|
||||
空时每 quota_poll_interval_seconds(缺省 30 秒)查一次
|
||||
配额(轮询本身 50 毫秒一圈,逐圈查会把库打爆),超限时
|
||||
同样杀进程组并返回 status='quota_exceeded' 的内部标记
|
||||
——调用方据此走配额中断收尾,绝不能把这个标记原样写
|
||||
进 tasks.status(不在合法状态集里)。Docker 路径见
|
||||
_docker_call_with_quota_poll(同款轮询);同步路径没有
|
||||
轮询钩子,只走边界检查。
|
||||
"""
|
||||
task_id = task["id"]
|
||||
gateway_provider, gateway_env = self._ensure_gateway(task)
|
||||
@@ -1539,6 +1576,10 @@ class Worker:
|
||||
started = time.time()
|
||||
last_quota_check = 0.0
|
||||
last_db_cancel_check = 0.0
|
||||
quota_interval = getattr(
|
||||
self, "quota_poll_interval_seconds", 30)
|
||||
if quota_interval is None:
|
||||
quota_interval = 30
|
||||
try:
|
||||
while proc.poll() is None:
|
||||
if task_id in self._cancel_requested or self._loop_stop:
|
||||
@@ -1567,7 +1608,7 @@ class Worker:
|
||||
self._kill_proc_group(proc)
|
||||
break
|
||||
if quota_check is not None and \
|
||||
time.time() - last_quota_check >= 30:
|
||||
time.time() - last_quota_check >= quota_interval:
|
||||
last_quota_check = time.time()
|
||||
try:
|
||||
over = quota_check()
|
||||
@@ -2205,6 +2246,80 @@ class Worker:
|
||||
detail += " (" + "; ".join(rejected) + ")"
|
||||
self._audit(task["project_id"], "worker_peer_messages", detail)
|
||||
|
||||
def _docker_call_with_quota_poll(self, task, fn, quota_check=None):
|
||||
"""Docker 同步调用包一层配额轮询(G8),与 popen 路径同款。
|
||||
|
||||
fn 在后台线程里跑(DockerExecutor 的调用本身是阻塞的
|
||||
docker run,无法像 popen 那样在主循环里轮询);主循环
|
||||
每 quota_poll_interval_seconds(缺省 30 秒)调一次
|
||||
quota_check,确认超限时按任务容器名真杀容器,fn 返回
|
||||
后把结果状态改写成内部标记 'quota_exceeded',调用方
|
||||
(run_one)据此走统一的配额中断收尾:状态 interrupted、
|
||||
已花成本照记、续跑消息保留。quota_check 为 None 时
|
||||
直接调 fn,不起线程。quota_check 本身抛异常时按
|
||||
fail-open 处理(只记 warning,不杀容器),只有确认
|
||||
超限才杀;取消时 cancel() 已按容器名杀过容器,线程
|
||||
随之结束,轮询循环不再多杀一次。fn 抛异常时在主线程
|
||||
原样重抛,由 _run_attempt 外层的 except 写成 failed,
|
||||
与不包轮询时的行为一致。注意:轮询查的是已记账用量
|
||||
(与 popen 路径同款),在跑回合的 token 只在回合结束
|
||||
才回传,所以掐断粒度仍是轮询间隔,不是 token 级。
|
||||
"""
|
||||
if quota_check is None:
|
||||
return fn()
|
||||
box = {}
|
||||
|
||||
def _target():
|
||||
try:
|
||||
box["result"] = fn()
|
||||
except Exception as exc: # noqa: BLE001 -- 主线程重抛
|
||||
box["error"] = exc
|
||||
|
||||
thread = threading.Thread(target=_target, daemon=True)
|
||||
thread.start()
|
||||
interval = getattr(self, "quota_poll_interval_seconds", 30)
|
||||
if interval is None:
|
||||
interval = 30
|
||||
last_check = 0.0
|
||||
quota_hit = False
|
||||
while thread.is_alive():
|
||||
if task["id"] in self._cancel_requested or self._loop_stop:
|
||||
break
|
||||
now = time.time()
|
||||
if now - last_check >= interval:
|
||||
last_check = now
|
||||
try:
|
||||
over = bool(quota_check())
|
||||
except Exception as exc:
|
||||
try:
|
||||
self.logger.warning(
|
||||
"docker quota poll failed: %s", exc)
|
||||
except Exception:
|
||||
pass
|
||||
over = False
|
||||
if over:
|
||||
quota_hit = True
|
||||
container = None
|
||||
try:
|
||||
container = self.executor.container_name(task)
|
||||
except Exception:
|
||||
container = None
|
||||
if container is not None and hasattr(
|
||||
self.executor, "kill_container"):
|
||||
try:
|
||||
self.executor.kill_container(container)
|
||||
except Exception:
|
||||
pass
|
||||
break
|
||||
thread.join(timeout=0.05)
|
||||
thread.join()
|
||||
if "error" in box:
|
||||
raise box["error"]
|
||||
result = box.get("result")
|
||||
if quota_hit and isinstance(result, dict):
|
||||
result = dict(result, status="quota_exceeded")
|
||||
return result
|
||||
|
||||
def _run_attempt(self, task, quota_check=None):
|
||||
"""单次执行尝试(不含认领、key 门、重试、成本与最终审计)。
|
||||
|
||||
@@ -2337,12 +2452,16 @@ class Worker:
|
||||
if task.get("_goal_resume_breakpoint"):
|
||||
# G10 断点:Docker 路用无 --goal-text 的
|
||||
# goal-resume-breakpoint 命令。
|
||||
result = self.executor.goal_resume_breakpoint_task(
|
||||
result = self._docker_call_with_quota_poll(
|
||||
task,
|
||||
timeout=effective_docker_timeout,
|
||||
extra_mounts=tenant_mounts,
|
||||
network=tenant_network,
|
||||
extra_env=proxy_env)
|
||||
lambda: self.executor.goal_resume_breakpoint_task(
|
||||
task,
|
||||
timeout=effective_docker_timeout,
|
||||
extra_mounts=tenant_mounts,
|
||||
network=tenant_network,
|
||||
extra_env=_merge_env(proxy_env,
|
||||
gateway_env)),
|
||||
quota_check)
|
||||
elif task.get("_resume"):
|
||||
# P1 断点续跑传 None(不传 --goal-text,
|
||||
# host 沿用原 goal+verify);开新目标
|
||||
@@ -2350,31 +2469,41 @@ class Worker:
|
||||
_resume_text = None if task.get(
|
||||
"_resume_kind") == "breakpoint" \
|
||||
else task["prompt"]
|
||||
result = self.executor.goal_resume_task(
|
||||
task, _resume_text,
|
||||
timeout=effective_docker_timeout,
|
||||
extra_mounts=tenant_mounts,
|
||||
network=tenant_network,
|
||||
extra_env=_merge_env(proxy_env,
|
||||
gateway_env),
|
||||
model_provider=gateway_provider)
|
||||
result = self._docker_call_with_quota_poll(
|
||||
task,
|
||||
lambda: self.executor.goal_resume_task(
|
||||
task, _resume_text,
|
||||
timeout=effective_docker_timeout,
|
||||
extra_mounts=tenant_mounts,
|
||||
network=tenant_network,
|
||||
extra_env=_merge_env(proxy_env,
|
||||
gateway_env),
|
||||
model_provider=gateway_provider),
|
||||
quota_check)
|
||||
else:
|
||||
result = self.executor.goal_task(
|
||||
result = self._docker_call_with_quota_poll(
|
||||
task,
|
||||
lambda: self.executor.goal_task(
|
||||
task, timeout=effective_docker_timeout,
|
||||
extra_mounts=tenant_mounts,
|
||||
network=tenant_network,
|
||||
extra_env=_merge_env(proxy_env,
|
||||
gateway_env),
|
||||
seed_host_dir=seed_host_dir,
|
||||
model_provider=gateway_provider),
|
||||
quota_check)
|
||||
else:
|
||||
result = self._docker_call_with_quota_poll(
|
||||
task,
|
||||
lambda: self.executor.dispatch_task(
|
||||
task, timeout=effective_docker_timeout,
|
||||
extra_mounts=tenant_mounts,
|
||||
network=tenant_network,
|
||||
extra_env=_merge_env(proxy_env,
|
||||
gateway_env),
|
||||
seed_host_dir=seed_host_dir,
|
||||
model_provider=gateway_provider)
|
||||
else:
|
||||
result = self.executor.dispatch_task(
|
||||
task, timeout=effective_docker_timeout,
|
||||
extra_mounts=tenant_mounts,
|
||||
network=tenant_network,
|
||||
extra_env=_merge_env(proxy_env, gateway_env),
|
||||
seed_host_dir=seed_host_dir,
|
||||
model_provider=gateway_provider)
|
||||
model_provider=gateway_provider),
|
||||
quota_check)
|
||||
else:
|
||||
# seed 参数只传给 LocalExecutor:CloudExecutor
|
||||
# 基类与测试假执行器的 goal_task/dispatch_task
|
||||
@@ -2434,19 +2563,23 @@ class Worker:
|
||||
result = self.executor.dispatch_task(task)
|
||||
combined = (result.get("stdout", "") or "") + \
|
||||
(result.get("stderr", "") or "")
|
||||
if not goal_mode and result["status"] != "completed" and \
|
||||
if not goal_mode and result["status"] not in (
|
||||
"completed", "quota_exceeded") and \
|
||||
_ALREADY_EXISTS_MARKER in combined:
|
||||
if is_docker:
|
||||
result = self.executor.run_task(
|
||||
task, timeout=effective_docker_timeout,
|
||||
extra_mounts=tenant_mounts
|
||||
if self.tenant_isolation else None,
|
||||
network=tenant_network
|
||||
if self.tenant_isolation else None,
|
||||
extra_env=_merge_env(
|
||||
proxy_env if self.egress_proxy_port
|
||||
else None, gateway_env),
|
||||
model_provider=gateway_provider)
|
||||
result = self._docker_call_with_quota_poll(
|
||||
task,
|
||||
lambda: self.executor.run_task(
|
||||
task, timeout=effective_docker_timeout,
|
||||
extra_mounts=tenant_mounts
|
||||
if self.tenant_isolation else None,
|
||||
network=tenant_network
|
||||
if self.tenant_isolation else None,
|
||||
extra_env=_merge_env(
|
||||
proxy_env if self.egress_proxy_port
|
||||
else None, gateway_env),
|
||||
model_provider=gateway_provider),
|
||||
quota_check)
|
||||
else:
|
||||
if gateway_provider and not isinstance(
|
||||
self.executor, LocalExecutor):
|
||||
@@ -3092,8 +3225,11 @@ class Worker:
|
||||
# 再开跑——直接标 interrupted + 配额原因,不记成本(还没花
|
||||
# 钱),续跑消息不消费,顶上来后重试/续跑可再跑。与 key
|
||||
# 门、仓库门同款提前返回形状,区别是状态为 interrupted
|
||||
#(可续跑)而非 failed。
|
||||
quota_error = self._quota_exceeded(task, None)
|
||||
#(可续跑)而非 failed。G7 分层:认领门用 closed——
|
||||
# 查询异常时同样标 interrupted(quota_check_unavailable,
|
||||
# 可重试/续跑),绝不在“不知道超没超”时开跑;轮询与
|
||||
# 轮次边界仍用默认 open,查询抖动不杀在跑任务。
|
||||
quota_error = self._quota_exceeded(task, None, on_error="closed")
|
||||
if quota_error is not None:
|
||||
self._set_status(task_id, "interrupted")
|
||||
self._write_execution_result(
|
||||
@@ -3696,8 +3832,11 @@ class Worker:
|
||||
job_id = job["id"]
|
||||
task_id = job["task_id"]
|
||||
try:
|
||||
# G13:审计操作人用 job 的费用归属人(请求 /pr 的
|
||||
# 用户),老行没有 created_by 时回落 "system"。
|
||||
result = execute_task_pr(
|
||||
self.db_path, task_id, "system")
|
||||
self.db_path, task_id,
|
||||
job.get("created_by") or "system")
|
||||
conn = _db_connect(self.db_path)
|
||||
try:
|
||||
finish_pr_job(conn, job_id, "completed",
|
||||
@@ -4121,6 +4260,15 @@ def main(argv=None):
|
||||
"列表里的第一个模型再跑一次,401/400"
|
||||
"绝不 fallback,取消绝不 fallback)")
|
||||
parser.add_argument("--max-turns", type=int, default=None)
|
||||
parser.add_argument("--ignore-host-budget", dest="ignore_host_budget",
|
||||
action="store_true", default=True,
|
||||
help="G9(默认开):worker 路径忽略 host 的 "
|
||||
"max_daily_turns 预算,云端 per-user USD "
|
||||
"配额 + worker 门才是口径")
|
||||
parser.add_argument("--no-ignore-host-budget",
|
||||
dest="ignore_host_budget", action="store_false",
|
||||
help="G9:关闭忽略,worker 调 host.py 时 host "
|
||||
"的 max_daily_turns 预算照常生效(双重门)")
|
||||
parser.add_argument("--max-tasks", type=int, default=None)
|
||||
parser.add_argument("--executor", default="local",
|
||||
choices=["local", "docker", "cloud"],
|
||||
@@ -4219,6 +4367,7 @@ def main(argv=None):
|
||||
args.fallback_models.split(",")]
|
||||
if args.fallback_models else []),
|
||||
max_turns=args.max_turns,
|
||||
ignore_host_budget=args.ignore_host_budget,
|
||||
executor=_build_cli_executor(args, host_py),
|
||||
task_timeout=args.task_timeout,
|
||||
max_retries=args.max_retries,
|
||||
|
||||
Reference in new issue
Block a user