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:
host committed
1 parent
c189df4a20
commit
ad014393ee
3 files changed
+474
-26
No files matched your search
@@ -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()
|
||||
Reference in new issue
Block a user