worker: enforce quotas during execution (claim/retry/fallback/followup/poll kill)
- claim-time gate: queued task whose quota was spent while waiting is marked interrupted with a quota reason, no cost row, resume kept - retry/fallback boundaries and every followup round check DB month usage plus this run's accumulated usage (priced via cost table) - local popen poll loop kills the process group on quota (30s throttle); Docker/sync paths get boundary checks only (documented) - quota stops finalize as interrupted + quota_exceeded error, keep the actual cost row, no push, followups restored to pending
This commit is contained in:
1 parent
c189df4a20
commit
ad014393ee
3 files changed
+474
-26
No files matched your search
+22
-8
@@ -913,14 +913,28 @@ interrupted 且终态任务 cancel 返回 400、心跳后 sweep 不扫刚心跳
|
||||
- 准入:建任务、会话分支、重试在权限校验通过后、写库之前查
|
||||
配额,本月用量 >= 配额时回 402 `quota_exceeded`(附 month、
|
||||
quota_usd_micros、used_usd_micros),不建任务、不改状态。
|
||||
老用户与未设配额的用户默认不限额,行为与过去一致。注意
|
||||
边界:只做准入校验,不做跑中掐断——已在跑的回合、后续消
|
||||
息触发的新一轮中途超限不会被停掉(worker 侧没有按金额中途
|
||||
停转的钩子);重试按操作人查配额、费用仍记到原任务创建者
|
||||
名下(建任务与分支的操作人就是创建者,没有这个错位)。
|
||||
系统任务(GitHub webhook 的 /fix-review,created_by
|
||||
是 'system')不记到任何用户名下、不占配额;要把它算到项
|
||||
目所有者头上需另做归属。
|
||||
老用户与未设配额的用户默认不限额,行为与过去一致。重试按
|
||||
操作人查配额、费用仍记到原任务创建者名下(建任务与分支的
|
||||
操作人就是创建者,没有这个错位)。系统任务(GitHub webhook
|
||||
的 /fix-review,created_by 是 'system')不记到任何用户名下、
|
||||
不占配额;要把它算到项目所有者头上需另做归属。
|
||||
- 执行时加固(worker 侧,`Worker._quota_exceeded`,口径与准入
|
||||
一致:当月已记账 + 本轮已累计用量,按价目表换算):认领后
|
||||
开跑前再查一次(排队期间配额被花完/调低时不跑,任务标
|
||||
interrupted + 配额原因、不记成本、续跑消息不消费);重试与
|
||||
fallback 换模型前各查一次;后续消息每轮开跑前查一次(超限
|
||||
时该条回 pending、任务 interrupted,不再排空剩余;建消息时
|
||||
不查,排队本身不花钱);Local popen 路径在轮询里每 30 秒查
|
||||
一次,超限时杀进程组(goal 长循环也一样,杀的是整个子进程
|
||||
组)。配额中断的任务一律 interrupted(可续跑:顶上来后重试
|
||||
/续跑继续),错误文本带 `quota_exceeded`,已花成本照记
|
||||
actual,不发终态推送(与取消同款)。诚实边界:Docker/同步
|
||||
路径没有轮询钩子,只有边界检查,单轮内的超支要到轮次结束
|
||||
才停;单轮本身花的钱无法中途按 token 流式截断(中转用量只
|
||||
在回合结束才回传),所以“掐断”粒度是轮次边界 + 长轮询中的
|
||||
30 秒检查,不是 token 级;并发下两个任务同时查到未超限再同
|
||||
时开跑会有少量超支(与限流器同款 best-effort,非精确配额);
|
||||
PR 任务(_run_pr_jobs)暂不查配额。
|
||||
- 接口:`GET /users/me/quota`(自查,返回 month、quota、
|
||||
used、remaining,不限额时 quota 与 remaining 为 null);
|
||||
`GET /admin/quotas?user_id=` 与 `POST /admin/quotas`
|
||||
|
||||
@@ -0,0 +1,248 @@
|
||||
"""cloud worker 配额加固单测:认领门、后续消息执行时查配额、轮询掐断。
|
||||
|
||||
运行(在 ~/workspace/codex-app-server 下):
|
||||
python3 -m unittest discover -s cloud/tests -p "test_worker_quota.py" -v
|
||||
|
||||
说明:只用注入的假 runner(假 CLI)与真 sleep 子进程,不跑真实
|
||||
模型回合、不联网、不登录任何服务;env 映射也是注入的普通
|
||||
dict,里面的值只是占位用的假名,绝不是真实密钥。
|
||||
|
||||
价目(cloud/cost.py,gpt-6.1-sol:input $2.00 / cached $0.10 /
|
||||
output $10.00 每百万 token):input 10000 + output 1000 =
|
||||
20000 + 10000 = 30000 微美元。配额 100000 微美元($0.10)。
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
try:
|
||||
from cloud.cost import CostStore
|
||||
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.server import init_db
|
||||
from cloud.worker import Worker
|
||||
|
||||
QUOTA = 100000
|
||||
|
||||
|
||||
def _insert_task(db_path, task_id, created_by="tester", status="queued",
|
||||
project_id="p1"):
|
||||
now = time.time()
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
conn.execute(
|
||||
"INSERT INTO tasks(id, project_id, prompt, status,"
|
||||
" idempotency_key, depends_on, priority, created_by,"
|
||||
" created_at, updated_at, last_heartbeat)"
|
||||
" VALUES(?, ?, ?, ?, NULL, '[]', 0, ?, ?, ?, NULL)",
|
||||
(task_id, project_id, "do work", status, created_by, now, now),
|
||||
)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _insert_quota(db_path, user_id, micros):
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
conn.execute(
|
||||
"INSERT INTO user_quotas(user_id, quota_usd_micros, updated_at)"
|
||||
" VALUES(?, ?, ?)",
|
||||
(user_id, micros, time.time()),
|
||||
)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _insert_donor_cost(db_path, micros):
|
||||
"""给 tester 记一笔已花成本:另起一个已完成任务并记 actual 行。"""
|
||||
_insert_task(db_path, "donor", status="completed")
|
||||
store = CostStore(db_path)
|
||||
try:
|
||||
store.record("donor", "p1", "gpt-6.1-sol", input_tokens=1,
|
||||
output_tokens=0, usd_micros=micros, source="actual")
|
||||
finally:
|
||||
try:
|
||||
store._conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _insert_followup(db_path, fu_id, task_id, body="follow up"):
|
||||
now = time.time()
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
conn.execute(
|
||||
"INSERT INTO task_followups(id, task_id, project_id, body,"
|
||||
" mode, status, created_by, created_at, updated_at)"
|
||||
" VALUES(?, ?, 'p1', ?, 'queued', 'pending', 'tester', ?, ?)",
|
||||
(fu_id, task_id, body, now, now),
|
||||
)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _status(db_path, task_id):
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
row = conn.execute(
|
||||
"SELECT status FROM tasks WHERE id = ?", (task_id,)).fetchone()
|
||||
finally:
|
||||
conn.close()
|
||||
return row[0] if row else None
|
||||
|
||||
|
||||
def _error_text(db_path, task_id):
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
row = conn.execute(
|
||||
"SELECT error_text FROM tasks WHERE id = ?",
|
||||
(task_id,)).fetchone()
|
||||
finally:
|
||||
conn.close()
|
||||
return row[0] if row else None
|
||||
|
||||
|
||||
def _followup_status(db_path, fu_id):
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
row = conn.execute(
|
||||
"SELECT status FROM task_followups WHERE id = ?",
|
||||
(fu_id,)).fetchone()
|
||||
finally:
|
||||
conn.close()
|
||||
return row[0] if row else None
|
||||
|
||||
|
||||
def _cost_rows(db_path, task_id):
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
rows = conn.execute(
|
||||
"SELECT usd_micros, source FROM cost_records"
|
||||
" WHERE task_id = ?", (task_id,)).fetchall()
|
||||
finally:
|
||||
conn.close()
|
||||
return rows
|
||||
|
||||
|
||||
def _usage_runner(usage, calls):
|
||||
"""假 CLI:dispatch/run 一律成功,stdout JSON 里带 usage。"""
|
||||
def fake(cmd):
|
||||
calls.append(list(cmd))
|
||||
return SimpleNamespace(
|
||||
returncode=0,
|
||||
stdout=json.dumps({"status": "completed", "text": "done",
|
||||
"usage": usage}),
|
||||
stderr="")
|
||||
return fake
|
||||
|
||||
|
||||
class WorkerQuotaTestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||||
init_db(self.db_path)
|
||||
self.env = {"XCM_API_KEY": "test-placeholder"}
|
||||
self.calls = []
|
||||
|
||||
def tearDown(self):
|
||||
self.tmp.cleanup()
|
||||
|
||||
def _worker(self, runner):
|
||||
return Worker(self.db_path, "/fake/host.py", runner=runner,
|
||||
env=self.env, model="gpt-6.1-sol",
|
||||
retry_backoff_seconds=0)
|
||||
|
||||
# -- 认领门 --------------------------------------------------------
|
||||
def test_claim_time_quota_stop(self):
|
||||
_insert_quota(self.db_path, "tester", QUOTA)
|
||||
_insert_donor_cost(self.db_path, QUOTA)
|
||||
_insert_task(self.db_path, "T1")
|
||||
worker = self._worker(_usage_runner(
|
||||
{"input_tokens": 10, "output_tokens": 1, "total_tokens": 11},
|
||||
self.calls))
|
||||
result = worker.run_one()
|
||||
self.assertEqual(result["status"], "interrupted")
|
||||
self.assertIn("quota_exceeded", result["error"])
|
||||
self.assertEqual(_status(self.db_path, "T1"), "interrupted")
|
||||
self.assertIn("quota_exceeded", _error_text(self.db_path, "T1"))
|
||||
# 还没跑:runner 一次没调,也不记成本。
|
||||
self.assertEqual(self.calls, [])
|
||||
self.assertEqual(_cost_rows(self.db_path, "T1"), [])
|
||||
|
||||
# -- 后续消息执行时查配额 ------------------------------------------
|
||||
def test_followup_quota_stop(self):
|
||||
_insert_quota(self.db_path, "tester", QUOTA)
|
||||
_insert_donor_cost(self.db_path, 90000)
|
||||
_insert_task(self.db_path, "T1")
|
||||
_insert_followup(self.db_path, "fu1", "T1")
|
||||
# 首轮花 30000 微美元:90000 + 30000 = 120000 >= 100000,
|
||||
# 排空时必须停,该条回 pending,任务 interrupted。
|
||||
worker = self._worker(_usage_runner(
|
||||
{"input_tokens": 10000, "cached_input_tokens": 0,
|
||||
"output_tokens": 1000, "total_tokens": 11000},
|
||||
self.calls))
|
||||
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(_followup_status(self.db_path, "fu1"), "pending")
|
||||
# 首轮只跑了一次(后续消息没跑),已花成本照记 actual。
|
||||
self.assertEqual(len(self.calls), 1)
|
||||
self.assertEqual(_cost_rows(self.db_path, "T1"),
|
||||
[(30000, "actual")])
|
||||
|
||||
def test_followup_runs_when_under_quota(self):
|
||||
_insert_quota(self.db_path, "tester", QUOTA)
|
||||
_insert_task(self.db_path, "T1")
|
||||
_insert_followup(self.db_path, "fu1", "T1")
|
||||
worker = self._worker(_usage_runner(
|
||||
{"input_tokens": 100, "cached_input_tokens": 0,
|
||||
"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")
|
||||
self.assertEqual(_followup_status(self.db_path, "fu1"), "done")
|
||||
self.assertEqual(len(self.calls), 2)
|
||||
|
||||
def test_no_quota_row_is_unlimited(self):
|
||||
_insert_task(self.db_path, "T1")
|
||||
worker = self._worker(_usage_runner(
|
||||
{"input_tokens": 100, "output_tokens": 10, "total_tokens": 110},
|
||||
self.calls))
|
||||
task = worker.get_task("T1")
|
||||
self.assertIsNone(worker._quota_exceeded(task, None))
|
||||
result = worker.run_one()
|
||||
self.assertEqual(result["status"], "completed")
|
||||
|
||||
# -- 轮询掐断 ------------------------------------------------------
|
||||
def test_poll_loop_quota_kill(self):
|
||||
_insert_task(self.db_path, "T1")
|
||||
worker = Worker(self.db_path, "/fake/host.py",
|
||||
env=self.env, model="gpt-6.1-sol",
|
||||
process_factory=lambda task: subprocess.Popen(
|
||||
["sleep", "30"], start_new_session=True,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE))
|
||||
task = worker.get_task("T1")
|
||||
result = worker._execute_via_process(
|
||||
task, kind="dispatch", quota_check=lambda: True)
|
||||
self.assertEqual(result["status"], "quota_exceeded")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+204
-18
@@ -121,6 +121,7 @@ Docker 下 repos_dir 必须放在 extra-mounts 已挂载的目录下,
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
from datetime import datetime, timezone
|
||||
import re
|
||||
import shutil
|
||||
import signal
|
||||
@@ -746,6 +747,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):
|
||||
"""配额加固的执行时检查:返回超限时的错误文本,否则 None。
|
||||
|
||||
口径与服务端准入(_quota_status)一致:按任务创建者算,
|
||||
当月(UTC)已记账用量 + 本轮已累计用量(按价目表换算,
|
||||
无价模型按 0 计,与服务端 SUM(usd NULL 不计)一致)。
|
||||
无配额行 = 不限额。查询本身异常时按不限额处理(fail-open
|
||||
并记一条警告):配额是第二道网,绝不能因查询抖动把全部
|
||||
任务挡在门外;服务端建任务/分支/重试时的准入仍是第一道。
|
||||
"""
|
||||
owner = (task or {}).get("created_by")
|
||||
if not owner:
|
||||
return None
|
||||
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:
|
||||
return None
|
||||
quota = row[0]
|
||||
if quota is None:
|
||||
return None
|
||||
# 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
|
||||
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 used >= quota:
|
||||
return ("quota_exceeded: 本月配额已用完(已用 %d / "
|
||||
"配额 %d 微美元,%s)" % (used, quota, month))
|
||||
return None
|
||||
except Exception as exc:
|
||||
try:
|
||||
self.logger.warning("quota check failed: %s", exc)
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
def _ready_tasks(self):
|
||||
"""当前就绪任务列表(pick_next 的批量版)。
|
||||
|
||||
@@ -977,14 +1045,20 @@ class Worker:
|
||||
pass
|
||||
|
||||
def _execute_via_process(self, task, kind="dispatch",
|
||||
deadline_seconds=None, seed_host_dir=None):
|
||||
deadline_seconds=None, seed_host_dir=None,
|
||||
quota_check=None):
|
||||
"""进程模式执行一轮:启动、轮询、取消时杀进程组。
|
||||
|
||||
返回与同步路径同款 result dict(status/stdout/stderr/
|
||||
host_status/returncode)。取消时返回 status='cancelled'
|
||||
的标记,调用方据此保持 interrupted;deadline_seconds
|
||||
到期时杀进程组并返回 status='timeout' 的标记,调用方
|
||||
据此置 failed(超时是失败,不是取消)。
|
||||
据此置 failed(超时是失败,不是取消)。quota_check 不为
|
||||
空时每 30 秒查一次配额(轮询本身 50 毫秒一圈,逐圈查会
|
||||
把库打爆),超限时同样杀进程组并返回 status=
|
||||
'quota_exceeded' 的内部标记——调用方据此走配额中断收尾,
|
||||
绝不能把这个标记原样写进 tasks.status(不在合法状态集
|
||||
里)。Docker/同步路径没有这个轮询钩子,只走边界检查。
|
||||
"""
|
||||
task_id = task["id"]
|
||||
if self.process_factory is not None:
|
||||
@@ -1002,7 +1076,9 @@ class Worker:
|
||||
pass
|
||||
cancelled = False
|
||||
timed_out = False
|
||||
quota_hit = False
|
||||
started = time.time()
|
||||
last_quota_check = 0.0
|
||||
try:
|
||||
while proc.poll() is None:
|
||||
if task_id in self._cancel_requested or self._loop_stop:
|
||||
@@ -1014,6 +1090,17 @@ class Worker:
|
||||
timed_out = True
|
||||
self._kill_proc_group(proc)
|
||||
break
|
||||
if quota_check is not None and \
|
||||
time.time() - last_quota_check >= 30:
|
||||
last_quota_check = time.time()
|
||||
try:
|
||||
over = quota_check()
|
||||
except Exception:
|
||||
over = False
|
||||
if over:
|
||||
quota_hit = True
|
||||
self._kill_proc_group(proc)
|
||||
break
|
||||
time.sleep(0.05)
|
||||
try:
|
||||
out, err = proc.communicate(timeout=5)
|
||||
@@ -1030,6 +1117,11 @@ class Worker:
|
||||
"status": "cancelled", "stdout": out or "",
|
||||
"stderr": err or "", "returncode": proc.returncode,
|
||||
"host_status": None}
|
||||
if quota_hit:
|
||||
return {"executor": "local", "task_id": task_id,
|
||||
"status": "quota_exceeded", "stdout": out or "",
|
||||
"stderr": err or "", "returncode": proc.returncode,
|
||||
"host_status": None}
|
||||
if timed_out:
|
||||
return {"executor": "local", "task_id": task_id,
|
||||
"status": "timeout", "stdout": out or "",
|
||||
@@ -1483,7 +1575,7 @@ class Worker:
|
||||
}
|
||||
return (True, tenant_mounts, tenant_network, proxy_env, None)
|
||||
|
||||
def _run_fork(self, task):
|
||||
def _run_fork(self, task, quota_check=None):
|
||||
"""分支任务执行:host.py fork 建分支 thread,再跑首轮。
|
||||
返回 (result, error),语义同 _run_attempt。fork 本身不
|
||||
跑模型;首轮走 _run_attempt(dispatch 撞已存在 → run
|
||||
@@ -1532,7 +1624,7 @@ class Worker:
|
||||
err = err.strip()
|
||||
return (fork_res,
|
||||
err[-1000:] if err else "fork failed")
|
||||
return self._run_attempt(task)
|
||||
return self._run_attempt(task, quota_check=quota_check)
|
||||
|
||||
def _claim_next_followup(self, task_id):
|
||||
"""取下一条后续消息:direct 优先、队内按创建时间先来
|
||||
@@ -1595,7 +1687,7 @@ class Worker:
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _run_attempt(self, task):
|
||||
def _run_attempt(self, task, quota_check=None):
|
||||
"""单次执行尝试(不含认领、key 门、重试、成本与最终审计)。
|
||||
|
||||
返回 (result, error):result 是执行器结果 dict,error 是
|
||||
@@ -1662,15 +1754,18 @@ class Worker:
|
||||
result = self._execute_via_process(
|
||||
task, kind=popen_kind,
|
||||
deadline_seconds=self.task_timeout,
|
||||
seed_host_dir=seed_host_dir)
|
||||
seed_host_dir=seed_host_dir,
|
||||
quota_check=quota_check)
|
||||
combined = (result.get("stdout", "") or "") + \
|
||||
(result.get("stderr", "") or "")
|
||||
if not goal_mode and result["status"] not in (
|
||||
"completed", "cancelled", "timeout") and \
|
||||
"completed", "cancelled", "timeout",
|
||||
"quota_exceeded") and \
|
||||
_ALREADY_EXISTS_MARKER in combined:
|
||||
result = self._execute_via_process(
|
||||
task, kind="run",
|
||||
deadline_seconds=self.task_timeout)
|
||||
deadline_seconds=self.task_timeout,
|
||||
quota_check=quota_check)
|
||||
else:
|
||||
if is_docker:
|
||||
(ok, tenant_mounts, tenant_network, proxy_env,
|
||||
@@ -1935,6 +2030,30 @@ class Worker:
|
||||
"stdout": "", "error": repo_error, "events": [],
|
||||
"cost_record_id": None}
|
||||
|
||||
# 配额认领门:排队期间配额被花完(或被调低)时,任务不能
|
||||
# 再开跑——直接标 interrupted + 配额原因,不记成本(还没花
|
||||
# 钱),续跑消息不消费,顶上来后重试/续跑可再跑。与 key
|
||||
# 门、仓库门同款提前返回形状,区别是状态为 interrupted
|
||||
#(可续跑)而非 failed。
|
||||
quota_error = self._quota_exceeded(task, None)
|
||||
if quota_error is not None:
|
||||
self._set_status(task_id, "interrupted")
|
||||
self._write_execution_result(
|
||||
task_id, error_text=quota_error, actual_model=self.model,
|
||||
started_at=started_at, finished_at=time.time())
|
||||
self._audit(task["project_id"], "worker_quota_stop",
|
||||
"%s: %s" % (task_id, quota_error))
|
||||
log_event(self.logger, "worker_quota_stop", task_id=task_id,
|
||||
project_id=task["project_id"],
|
||||
status="interrupted")
|
||||
return {"dispatched": True, "task_id": task_id,
|
||||
"host_name": self.executor._host_name(task)
|
||||
if hasattr(self.executor, "_host_name") else task_id,
|
||||
"status": "interrupted",
|
||||
"executor": self._executor_name(),
|
||||
"stdout": "", "error": quota_error, "events": [],
|
||||
"cost_record_id": None}
|
||||
|
||||
events = []
|
||||
fallback_count = 0
|
||||
result = None
|
||||
@@ -1963,18 +2082,47 @@ class Worker:
|
||||
# 各次尝试的真实用量都要算进成本(失败尝试花的钱也是钱):
|
||||
# 主尝试与 fallback 尝试的结果都先累加,最后只记一条汇总。
|
||||
usage_total = None
|
||||
# 配额中断状态:任一边界(重试/fallback/后续消息/轮询掐
|
||||
# 断)命中后置位,跳过后续轮次,走统一的配额中断收尾
|
||||
#(interrupted + 配额原因 + 已花成本照记 + 续跑消息保留)。
|
||||
quota_stop = False
|
||||
quota_error = None
|
||||
|
||||
def _quota_hit():
|
||||
# 闭包读当前值:usage_total 与 task 在重试/fallback
|
||||
# 时会被重绑,轮询掐断要按最新值判断。
|
||||
return self._quota_exceeded(task, usage_total) is not None
|
||||
|
||||
# 分支任务:先 fork 再跑首轮(见 _run_fork);首轮失败
|
||||
# 不自动重试/fallback,失败后走重试入口重跑首轮(已存在
|
||||
# 续跑),不重复 fork。range(0) 让下面的尝试循环跳过,
|
||||
# 直接进统一的收尾(超时/取消/成本/审计/通知)逻辑。
|
||||
if task.get("forked_from"):
|
||||
result, error = self._run_fork(task)
|
||||
result, error = self._run_fork(task, quota_check=_quota_hit)
|
||||
usage_total = sum_usage(usage_total,
|
||||
(result or {}).get("usage"))
|
||||
if (result or {}).get("status") == "quota_exceeded":
|
||||
quota_error = self._quota_exceeded(task, usage_total) or \
|
||||
"quota_exceeded: 本月配额已用完"
|
||||
quota_stop = True
|
||||
for attempt in range(0 if task.get("forked_from")
|
||||
else self.max_retries + 1):
|
||||
result, error = self._run_attempt(task)
|
||||
if quota_stop:
|
||||
break
|
||||
if attempt > 0:
|
||||
# 重试边界:上一轮已花钱,跑下一轮前再查一次。
|
||||
quota_error = self._quota_exceeded(task, usage_total)
|
||||
if quota_error is not None:
|
||||
quota_stop = True
|
||||
break
|
||||
result, error = self._run_attempt(task, quota_check=_quota_hit)
|
||||
usage_total = sum_usage(usage_total, (result or {}).get("usage"))
|
||||
if (result or {}).get("status") == "quota_exceeded":
|
||||
# 轮询掐断:进程已被杀,配额原因以现查为准。
|
||||
quota_error = self._quota_exceeded(task, usage_total) or \
|
||||
"quota_exceeded: 本月配额已用完"
|
||||
quota_stop = True
|
||||
break
|
||||
if self._was_cancelled(task_id) or \
|
||||
result.get("status") == "cancelled":
|
||||
break
|
||||
@@ -2011,9 +2159,19 @@ class Worker:
|
||||
actual_model = fb
|
||||
task = dict(task)
|
||||
task["model"] = fb
|
||||
result, error = self._run_attempt(task)
|
||||
# fallback 边界:换模型再跑一轮前先查配额,超限
|
||||
# 时不烧第二份钱,直接走配额中断收尾。
|
||||
quota_error = self._quota_exceeded(task, usage_total)
|
||||
if quota_error is not None:
|
||||
quota_stop = True
|
||||
break
|
||||
result, error = self._run_attempt(task, quota_check=_quota_hit)
|
||||
usage_total = sum_usage(usage_total,
|
||||
(result or {}).get("usage"))
|
||||
if (result or {}).get("status") == "quota_exceeded":
|
||||
quota_error = self._quota_exceeded(task, usage_total) or \
|
||||
"quota_exceeded: 本月配额已用完"
|
||||
quota_stop = True
|
||||
break
|
||||
if self.router is not None and host_status_try is not None \
|
||||
and is_fallbackable(host_status_try) and \
|
||||
@@ -2045,17 +2203,33 @@ class Worker:
|
||||
# 剩余 pending 留给重试入口(重试走已存在续跑,跑完继续
|
||||
# 排空剩余消息);取消时该条回 pending,任务走中断收尾。
|
||||
if (result is not None and result.get("status") == "completed"
|
||||
and not quota_stop
|
||||
and not task.get("verify_command")
|
||||
and not task.get("goal_mode")):
|
||||
while True:
|
||||
fu = self._claim_next_followup(task_id)
|
||||
if fu is None:
|
||||
break
|
||||
# 后续消息执行时查配额:超限时该条回 pending(留
|
||||
# 给重试入口),任务走配额中断收尾,不再排空剩余;
|
||||
# 建消息时不查(排队本身不花钱,顶上来后可跑)。
|
||||
quota_error = self._quota_exceeded(task, usage_total)
|
||||
if quota_error is not None:
|
||||
self._restore_followup(fu["id"])
|
||||
quota_stop = True
|
||||
break
|
||||
fu_task = dict(task)
|
||||
fu_task["prompt"] = fu["body"]
|
||||
fu_result, fu_error = self._run_attempt(fu_task)
|
||||
fu_result, fu_error = self._run_attempt(
|
||||
fu_task, quota_check=_quota_hit)
|
||||
usage_total = sum_usage(
|
||||
usage_total, (fu_result or {}).get("usage"))
|
||||
if (fu_result or {}).get("status") == "quota_exceeded":
|
||||
self._restore_followup(fu["id"])
|
||||
quota_error = self._quota_exceeded(task, usage_total) or \
|
||||
"quota_exceeded: 本月配额已用完"
|
||||
quota_stop = True
|
||||
break
|
||||
if self._was_cancelled(task_id) or \
|
||||
(fu_result or {}).get("status") == "cancelled":
|
||||
self._restore_followup(fu["id"])
|
||||
@@ -2097,7 +2271,16 @@ class Worker:
|
||||
stdout = result.get("stdout", "") or ""
|
||||
host_status = result.get("host_status")
|
||||
final = result["status"]
|
||||
if final != "completed":
|
||||
if quota_stop:
|
||||
# 配额中断收尾:状态一律 interrupted + 配额原因,不跑
|
||||
# 失败原因重算(否则 fallback 拼接会把配额原因盖掉),
|
||||
# 已花成本仍按 usage_total 照记(下面成本块),续跑
|
||||
# 消息保留(下面清消息时跳过),不发终态推送(与取消
|
||||
# 同款:interrupted 不在推送事件集里)。quota_exceeded
|
||||
# 这个内部标记绝不能原样写进 tasks.status。
|
||||
final = "interrupted"
|
||||
error = quota_error
|
||||
if final != "completed" and not quota_stop:
|
||||
if fallback_model is not None:
|
||||
# 真 fallback 后第二次也失败:错误文本必须同时写清
|
||||
# 两次原因,不能只留第二次的(超时转换已在前面把
|
||||
@@ -2175,11 +2358,13 @@ class Worker:
|
||||
source="estimated", fallback_count=fallback_count)
|
||||
if final_written:
|
||||
self._audit(task["project_id"],
|
||||
"worker_completed" if final == "completed"
|
||||
else "worker_failed", task_id)
|
||||
"worker_quota_stop" if quota_stop else
|
||||
("worker_completed" if final == "completed"
|
||||
else "worker_failed"), task_id)
|
||||
log_event(self.logger,
|
||||
"worker_completed" if final == "completed"
|
||||
else "worker_failed",
|
||||
"worker_quota_stop" if quota_stop else
|
||||
("worker_completed" if final == "completed"
|
||||
else "worker_failed"),
|
||||
task_id=task_id, project_id=task["project_id"],
|
||||
status=final)
|
||||
if result.get("goal_met") is False:
|
||||
@@ -2196,10 +2381,11 @@ class Worker:
|
||||
# 补一条 worker_completed/failed 审计误导排查。
|
||||
final = "interrupted"
|
||||
self._audit(task["project_id"], "worker_interrupted", task_id)
|
||||
if resume_message:
|
||||
if resume_message and not quota_stop:
|
||||
# 本轮已跑完(无论成败):续跑消息已消费,清掉;进程
|
||||
# 中途被杀时清不到,下次认领会重跑同一条消息——可接受
|
||||
# 的续跑语义,不会静默丢消息。清失败不挡结果返回。
|
||||
# 配额中断时不消费:顶上来后续跑还要说这句话。
|
||||
try:
|
||||
conn = _db_connect(self.db_path)
|
||||
try:
|
||||
|
||||
Reference in new issue
Block a user