host
ad014393ee
- 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
249 lines
8.8 KiB
Python
249 lines
8.8 KiB
Python
"""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()
|