Merge pull request #78 from openqj/feat-quota-hardening
worker: enforce quotas during execution
This commit is contained in:
3 files changed
+474
-26
No files matched your search
+22
-8
@@ -913,14 +913,28 @@ interrupted 且终态任务 cancel 返回 400、心跳后 sweep 不扫刚心跳
|
|||||||
- 准入:建任务、会话分支、重试在权限校验通过后、写库之前查
|
- 准入:建任务、会话分支、重试在权限校验通过后、写库之前查
|
||||||
配额,本月用量 >= 配额时回 402 `quota_exceeded`(附 month、
|
配额,本月用量 >= 配额时回 402 `quota_exceeded`(附 month、
|
||||||
quota_usd_micros、used_usd_micros),不建任务、不改状态。
|
quota_usd_micros、used_usd_micros),不建任务、不改状态。
|
||||||
老用户与未设配额的用户默认不限额,行为与过去一致。注意
|
老用户与未设配额的用户默认不限额,行为与过去一致。重试按
|
||||||
边界:只做准入校验,不做跑中掐断——已在跑的回合、后续消
|
操作人查配额、费用仍记到原任务创建者名下(建任务与分支的
|
||||||
息触发的新一轮中途超限不会被停掉(worker 侧没有按金额中途
|
操作人就是创建者,没有这个错位)。系统任务(GitHub webhook
|
||||||
停转的钩子);重试按操作人查配额、费用仍记到原任务创建者
|
的 /fix-review,created_by 是 'system')不记到任何用户名下、
|
||||||
名下(建任务与分支的操作人就是创建者,没有这个错位)。
|
不占配额;要把它算到项目所有者头上需另做归属。
|
||||||
系统任务(GitHub webhook 的 /fix-review,created_by
|
- 执行时加固(worker 侧,`Worker._quota_exceeded`,口径与准入
|
||||||
是 'system')不记到任何用户名下、不占配额;要把它算到项
|
一致:当月已记账 + 本轮已累计用量,按价目表换算):认领后
|
||||||
目所有者头上需另做归属。
|
开跑前再查一次(排队期间配额被花完/调低时不跑,任务标
|
||||||
|
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、
|
- 接口:`GET /users/me/quota`(自查,返回 month、quota、
|
||||||
used、remaining,不限额时 quota 与 remaining 为 null);
|
used、remaining,不限额时 quota 与 remaining 为 null);
|
||||||
`GET /admin/quotas?user_id=` 与 `POST /admin/quotas`
|
`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 argparse
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
|
from datetime import datetime, timezone
|
||||||
import re
|
import re
|
||||||
import shutil
|
import shutil
|
||||||
import signal
|
import signal
|
||||||
@@ -746,6 +747,73 @@ class Worker:
|
|||||||
return budget_exceeded(self.usage_turns, self.max_turns,
|
return budget_exceeded(self.usage_turns, self.max_turns,
|
||||||
self.usage_cost_usd, self.max_cost_usd)
|
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):
|
def _ready_tasks(self):
|
||||||
"""当前就绪任务列表(pick_next 的批量版)。
|
"""当前就绪任务列表(pick_next 的批量版)。
|
||||||
|
|
||||||
@@ -977,14 +1045,20 @@ class Worker:
|
|||||||
pass
|
pass
|
||||||
|
|
||||||
def _execute_via_process(self, task, kind="dispatch",
|
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/
|
返回与同步路径同款 result dict(status/stdout/stderr/
|
||||||
host_status/returncode)。取消时返回 status='cancelled'
|
host_status/returncode)。取消时返回 status='cancelled'
|
||||||
的标记,调用方据此保持 interrupted;deadline_seconds
|
的标记,调用方据此保持 interrupted;deadline_seconds
|
||||||
到期时杀进程组并返回 status='timeout' 的标记,调用方
|
到期时杀进程组并返回 status='timeout' 的标记,调用方
|
||||||
据此置 failed(超时是失败,不是取消)。
|
据此置 failed(超时是失败,不是取消)。quota_check 不为
|
||||||
|
空时每 30 秒查一次配额(轮询本身 50 毫秒一圈,逐圈查会
|
||||||
|
把库打爆),超限时同样杀进程组并返回 status=
|
||||||
|
'quota_exceeded' 的内部标记——调用方据此走配额中断收尾,
|
||||||
|
绝不能把这个标记原样写进 tasks.status(不在合法状态集
|
||||||
|
里)。Docker/同步路径没有这个轮询钩子,只走边界检查。
|
||||||
"""
|
"""
|
||||||
task_id = task["id"]
|
task_id = task["id"]
|
||||||
if self.process_factory is not None:
|
if self.process_factory is not None:
|
||||||
@@ -1002,7 +1076,9 @@ class Worker:
|
|||||||
pass
|
pass
|
||||||
cancelled = False
|
cancelled = False
|
||||||
timed_out = False
|
timed_out = False
|
||||||
|
quota_hit = False
|
||||||
started = time.time()
|
started = time.time()
|
||||||
|
last_quota_check = 0.0
|
||||||
try:
|
try:
|
||||||
while proc.poll() is None:
|
while proc.poll() is None:
|
||||||
if task_id in self._cancel_requested or self._loop_stop:
|
if task_id in self._cancel_requested or self._loop_stop:
|
||||||
@@ -1014,6 +1090,17 @@ class Worker:
|
|||||||
timed_out = True
|
timed_out = True
|
||||||
self._kill_proc_group(proc)
|
self._kill_proc_group(proc)
|
||||||
break
|
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)
|
time.sleep(0.05)
|
||||||
try:
|
try:
|
||||||
out, err = proc.communicate(timeout=5)
|
out, err = proc.communicate(timeout=5)
|
||||||
@@ -1030,6 +1117,11 @@ class Worker:
|
|||||||
"status": "cancelled", "stdout": out or "",
|
"status": "cancelled", "stdout": out or "",
|
||||||
"stderr": err or "", "returncode": proc.returncode,
|
"stderr": err or "", "returncode": proc.returncode,
|
||||||
"host_status": None}
|
"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:
|
if timed_out:
|
||||||
return {"executor": "local", "task_id": task_id,
|
return {"executor": "local", "task_id": task_id,
|
||||||
"status": "timeout", "stdout": out or "",
|
"status": "timeout", "stdout": out or "",
|
||||||
@@ -1483,7 +1575,7 @@ class Worker:
|
|||||||
}
|
}
|
||||||
return (True, tenant_mounts, tenant_network, proxy_env, None)
|
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,再跑首轮。
|
"""分支任务执行:host.py fork 建分支 thread,再跑首轮。
|
||||||
返回 (result, error),语义同 _run_attempt。fork 本身不
|
返回 (result, error),语义同 _run_attempt。fork 本身不
|
||||||
跑模型;首轮走 _run_attempt(dispatch 撞已存在 → run
|
跑模型;首轮走 _run_attempt(dispatch 撞已存在 → run
|
||||||
@@ -1532,7 +1624,7 @@ class Worker:
|
|||||||
err = err.strip()
|
err = err.strip()
|
||||||
return (fork_res,
|
return (fork_res,
|
||||||
err[-1000:] if err else "fork failed")
|
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):
|
def _claim_next_followup(self, task_id):
|
||||||
"""取下一条后续消息:direct 优先、队内按创建时间先来
|
"""取下一条后续消息:direct 优先、队内按创建时间先来
|
||||||
@@ -1595,7 +1687,7 @@ class Worker:
|
|||||||
except Exception:
|
except Exception:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
def _run_attempt(self, task):
|
def _run_attempt(self, task, quota_check=None):
|
||||||
"""单次执行尝试(不含认领、key 门、重试、成本与最终审计)。
|
"""单次执行尝试(不含认领、key 门、重试、成本与最终审计)。
|
||||||
|
|
||||||
返回 (result, error):result 是执行器结果 dict,error 是
|
返回 (result, error):result 是执行器结果 dict,error 是
|
||||||
@@ -1662,15 +1754,18 @@ class Worker:
|
|||||||
result = self._execute_via_process(
|
result = self._execute_via_process(
|
||||||
task, kind=popen_kind,
|
task, kind=popen_kind,
|
||||||
deadline_seconds=self.task_timeout,
|
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 "") + \
|
combined = (result.get("stdout", "") or "") + \
|
||||||
(result.get("stderr", "") or "")
|
(result.get("stderr", "") or "")
|
||||||
if not goal_mode and result["status"] not in (
|
if not goal_mode and result["status"] not in (
|
||||||
"completed", "cancelled", "timeout") and \
|
"completed", "cancelled", "timeout",
|
||||||
|
"quota_exceeded") and \
|
||||||
_ALREADY_EXISTS_MARKER in combined:
|
_ALREADY_EXISTS_MARKER in combined:
|
||||||
result = self._execute_via_process(
|
result = self._execute_via_process(
|
||||||
task, kind="run",
|
task, kind="run",
|
||||||
deadline_seconds=self.task_timeout)
|
deadline_seconds=self.task_timeout,
|
||||||
|
quota_check=quota_check)
|
||||||
else:
|
else:
|
||||||
if is_docker:
|
if is_docker:
|
||||||
(ok, tenant_mounts, tenant_network, proxy_env,
|
(ok, tenant_mounts, tenant_network, proxy_env,
|
||||||
@@ -1935,6 +2030,30 @@ class Worker:
|
|||||||
"stdout": "", "error": repo_error, "events": [],
|
"stdout": "", "error": repo_error, "events": [],
|
||||||
"cost_record_id": None}
|
"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 = []
|
events = []
|
||||||
fallback_count = 0
|
fallback_count = 0
|
||||||
result = None
|
result = None
|
||||||
@@ -1963,18 +2082,47 @@ class Worker:
|
|||||||
# 各次尝试的真实用量都要算进成本(失败尝试花的钱也是钱):
|
# 各次尝试的真实用量都要算进成本(失败尝试花的钱也是钱):
|
||||||
# 主尝试与 fallback 尝试的结果都先累加,最后只记一条汇总。
|
# 主尝试与 fallback 尝试的结果都先累加,最后只记一条汇总。
|
||||||
usage_total = None
|
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);首轮失败
|
# 分支任务:先 fork 再跑首轮(见 _run_fork);首轮失败
|
||||||
# 不自动重试/fallback,失败后走重试入口重跑首轮(已存在
|
# 不自动重试/fallback,失败后走重试入口重跑首轮(已存在
|
||||||
# 续跑),不重复 fork。range(0) 让下面的尝试循环跳过,
|
# 续跑),不重复 fork。range(0) 让下面的尝试循环跳过,
|
||||||
# 直接进统一的收尾(超时/取消/成本/审计/通知)逻辑。
|
# 直接进统一的收尾(超时/取消/成本/审计/通知)逻辑。
|
||||||
if task.get("forked_from"):
|
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,
|
usage_total = sum_usage(usage_total,
|
||||||
(result or {}).get("usage"))
|
(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")
|
for attempt in range(0 if task.get("forked_from")
|
||||||
else self.max_retries + 1):
|
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"))
|
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 \
|
if self._was_cancelled(task_id) or \
|
||||||
result.get("status") == "cancelled":
|
result.get("status") == "cancelled":
|
||||||
break
|
break
|
||||||
@@ -2011,9 +2159,19 @@ class Worker:
|
|||||||
actual_model = fb
|
actual_model = fb
|
||||||
task = dict(task)
|
task = dict(task)
|
||||||
task["model"] = fb
|
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,
|
usage_total = sum_usage(usage_total,
|
||||||
(result or {}).get("usage"))
|
(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
|
break
|
||||||
if self.router is not None and host_status_try is not None \
|
if self.router is not None and host_status_try is not None \
|
||||||
and is_fallbackable(host_status_try) and \
|
and is_fallbackable(host_status_try) and \
|
||||||
@@ -2045,17 +2203,33 @@ class Worker:
|
|||||||
# 剩余 pending 留给重试入口(重试走已存在续跑,跑完继续
|
# 剩余 pending 留给重试入口(重试走已存在续跑,跑完继续
|
||||||
# 排空剩余消息);取消时该条回 pending,任务走中断收尾。
|
# 排空剩余消息);取消时该条回 pending,任务走中断收尾。
|
||||||
if (result is not None and result.get("status") == "completed"
|
if (result is not None and result.get("status") == "completed"
|
||||||
|
and not quota_stop
|
||||||
and not task.get("verify_command")
|
and not task.get("verify_command")
|
||||||
and not task.get("goal_mode")):
|
and not task.get("goal_mode")):
|
||||||
while True:
|
while True:
|
||||||
fu = self._claim_next_followup(task_id)
|
fu = self._claim_next_followup(task_id)
|
||||||
if fu is None:
|
if fu is None:
|
||||||
break
|
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 = dict(task)
|
||||||
fu_task["prompt"] = fu["body"]
|
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 = sum_usage(
|
||||||
usage_total, (fu_result or {}).get("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 \
|
if self._was_cancelled(task_id) or \
|
||||||
(fu_result or {}).get("status") == "cancelled":
|
(fu_result or {}).get("status") == "cancelled":
|
||||||
self._restore_followup(fu["id"])
|
self._restore_followup(fu["id"])
|
||||||
@@ -2097,7 +2271,16 @@ class Worker:
|
|||||||
stdout = result.get("stdout", "") or ""
|
stdout = result.get("stdout", "") or ""
|
||||||
host_status = result.get("host_status")
|
host_status = result.get("host_status")
|
||||||
final = result["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:
|
if fallback_model is not None:
|
||||||
# 真 fallback 后第二次也失败:错误文本必须同时写清
|
# 真 fallback 后第二次也失败:错误文本必须同时写清
|
||||||
# 两次原因,不能只留第二次的(超时转换已在前面把
|
# 两次原因,不能只留第二次的(超时转换已在前面把
|
||||||
@@ -2175,11 +2358,13 @@ class Worker:
|
|||||||
source="estimated", fallback_count=fallback_count)
|
source="estimated", fallback_count=fallback_count)
|
||||||
if final_written:
|
if final_written:
|
||||||
self._audit(task["project_id"],
|
self._audit(task["project_id"],
|
||||||
"worker_completed" if final == "completed"
|
"worker_quota_stop" if quota_stop else
|
||||||
else "worker_failed", task_id)
|
("worker_completed" if final == "completed"
|
||||||
|
else "worker_failed"), task_id)
|
||||||
log_event(self.logger,
|
log_event(self.logger,
|
||||||
"worker_completed" if final == "completed"
|
"worker_quota_stop" if quota_stop else
|
||||||
else "worker_failed",
|
("worker_completed" if final == "completed"
|
||||||
|
else "worker_failed"),
|
||||||
task_id=task_id, project_id=task["project_id"],
|
task_id=task_id, project_id=task["project_id"],
|
||||||
status=final)
|
status=final)
|
||||||
if result.get("goal_met") is False:
|
if result.get("goal_met") is False:
|
||||||
@@ -2196,10 +2381,11 @@ class Worker:
|
|||||||
# 补一条 worker_completed/failed 审计误导排查。
|
# 补一条 worker_completed/failed 审计误导排查。
|
||||||
final = "interrupted"
|
final = "interrupted"
|
||||||
self._audit(task["project_id"], "worker_interrupted", task_id)
|
self._audit(task["project_id"], "worker_interrupted", task_id)
|
||||||
if resume_message:
|
if resume_message and not quota_stop:
|
||||||
# 本轮已跑完(无论成败):续跑消息已消费,清掉;进程
|
# 本轮已跑完(无论成败):续跑消息已消费,清掉;进程
|
||||||
# 中途被杀时清不到,下次认领会重跑同一条消息——可接受
|
# 中途被杀时清不到,下次认领会重跑同一条消息——可接受
|
||||||
# 的续跑语义,不会静默丢消息。清失败不挡结果返回。
|
# 的续跑语义,不会静默丢消息。清失败不挡结果返回。
|
||||||
|
# 配额中断时不消费:顶上来后续跑还要说这句话。
|
||||||
try:
|
try:
|
||||||
conn = _db_connect(self.db_path)
|
conn = _db_connect(self.db_path)
|
||||||
try:
|
try:
|
||||||
|
|||||||
Reference in new issue
Block a user