feat: P1 补齐（自动评审队列、自动化+分诊、并发 caps、运行时审批队列）
This commit is contained in:
host committed
1 parent
29efc6c61f
commit
e0cecd70aa
18 files changed
+3474
-27
No files matched your search
@@ -0,0 +1,223 @@
|
||||
"""并发上限单测（caps）。
|
||||
|
||||
覆盖：配置解析优先级与默认值、认领门的未满 / 已满 / 释放后
|
||||
恢复、启动中任务计入占用、CAS 抢输后占位不泄漏、缺 key 门
|
||||
提前返回后占位释放、双线程同时认领只赢一个。
|
||||
|
||||
只用注入的假环境与本地 SQLite，不跑真实模型、不联网。
|
||||
"""
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
from cloud.caps import (
|
||||
can_claim,
|
||||
effective_capacity,
|
||||
resolve_capacity,
|
||||
running_count,
|
||||
)
|
||||
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.caps import (
|
||||
can_claim,
|
||||
effective_capacity,
|
||||
resolve_capacity,
|
||||
running_count,
|
||||
)
|
||||
from cloud.server import init_db
|
||||
from cloud.worker import Worker
|
||||
|
||||
|
||||
def _insert_task(db_path, task_id, 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(?, ?, ?, 'queued', NULL, '[]', 0, 'tester', ?, ?, NULL)",
|
||||
(task_id, project_id, "do work", 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
|
||||
|
||||
|
||||
class CapsResolveTestCase(unittest.TestCase):
|
||||
def test_explicit_beats_env(self):
|
||||
env = {"CLOUD_WORKER_CAPACITY": "3"}
|
||||
self.assertEqual(resolve_capacity(2, env=env), 2)
|
||||
|
||||
def test_env_used_when_no_explicit(self):
|
||||
env = {"CLOUD_WORKER_CAPACITY": "3"}
|
||||
self.assertEqual(resolve_capacity(None, env=env), 3)
|
||||
|
||||
def test_default_is_1(self):
|
||||
self.assertEqual(resolve_capacity(None, env={}), 1)
|
||||
self.assertEqual(resolve_capacity(), 1)
|
||||
|
||||
def test_invalid_values_fall_back_to_1(self):
|
||||
self.assertEqual(resolve_capacity(0, env={}), 1)
|
||||
self.assertEqual(resolve_capacity(-2, env={}), 1)
|
||||
self.assertEqual(resolve_capacity("abc", env={}), 1)
|
||||
self.assertEqual(resolve_capacity(None, env={
|
||||
"CLOUD_WORKER_CAPACITY": "abc"}), 1)
|
||||
self.assertEqual(resolve_capacity(None, env={
|
||||
"CLOUD_WORKER_CAPACITY": "0"}), 1)
|
||||
self.assertEqual(resolve_capacity(None, env={
|
||||
"CLOUD_WORKER_CAPACITY": " "}), 1)
|
||||
|
||||
def test_effective_capacity(self):
|
||||
self.assertEqual(effective_capacity(2), 2)
|
||||
self.assertEqual(effective_capacity(0), 1)
|
||||
self.assertEqual(effective_capacity(None), 1)
|
||||
|
||||
def test_can_claim_boundaries(self):
|
||||
self.assertTrue(can_claim(0, 1))
|
||||
self.assertFalse(can_claim(1, 1))
|
||||
self.assertTrue(can_claim(1, 2))
|
||||
self.assertFalse(can_claim(2, 2))
|
||||
|
||||
def test_running_count_counts_starting_entries(self):
|
||||
# 空条目 = 已认领未起进程的启动中任务，同样占一个槽。
|
||||
self.assertEqual(running_count({"a": {}, "b": {"proc": 1}}), 2)
|
||||
self.assertEqual(running_count({}), 0)
|
||||
|
||||
|
||||
class WorkerClaimGateTestCase(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"}
|
||||
|
||||
def tearDown(self):
|
||||
self.tmp.cleanup()
|
||||
|
||||
def _worker(self, **kwargs):
|
||||
kwargs.setdefault("env", self.env)
|
||||
kwargs.setdefault("retry_backoff_seconds", 0)
|
||||
kwargs.setdefault("max_retries", 0)
|
||||
return Worker(self.db_path, "/fake/host.py", **kwargs)
|
||||
|
||||
def test_default_capacity_is_1(self):
|
||||
worker = self._worker()
|
||||
self.assertEqual(worker.capacity, 1)
|
||||
|
||||
def test_env_capacity_and_explicit_priority(self):
|
||||
env = {"CLOUD_WORKER_CAPACITY": "3"}
|
||||
self.assertEqual(self._worker(env=env).capacity, 3)
|
||||
self.assertEqual(self._worker(env=env, capacity=2).capacity, 2)
|
||||
|
||||
def test_full_worker_refuses_claim_and_task_stays_queued(self):
|
||||
_insert_task(self.db_path, "A")
|
||||
_insert_task(self.db_path, "B")
|
||||
worker = self._worker(capacity=1)
|
||||
self.assertTrue(worker._claim("A"))
|
||||
self.assertEqual(_status(self.db_path, "A"), "running")
|
||||
self.assertFalse(worker._claim("B"))
|
||||
self.assertEqual(_status(self.db_path, "B"), "queued")
|
||||
self.assertEqual(set(worker._running_snapshot()), {"A"})
|
||||
|
||||
def test_release_restores_claim(self):
|
||||
_insert_task(self.db_path, "A")
|
||||
_insert_task(self.db_path, "B")
|
||||
worker = self._worker(capacity=1)
|
||||
self.assertTrue(worker._claim("A"))
|
||||
self.assertFalse(worker._claim("B"))
|
||||
worker._running_clear("A")
|
||||
self.assertTrue(worker._claim("B"))
|
||||
self.assertEqual(_status(self.db_path, "B"), "running")
|
||||
|
||||
def test_starting_task_counts_toward_capacity(self):
|
||||
_insert_task(self.db_path, "B")
|
||||
worker = self._worker(capacity=1)
|
||||
# 空条目模拟“已认领、执行器还没起进程”的启动中任务。
|
||||
worker._running_set("STARTING")
|
||||
self.assertEqual(worker._running_count(), 1)
|
||||
self.assertFalse(worker._claim("B"))
|
||||
self.assertEqual(_status(self.db_path, "B"), "queued")
|
||||
worker._running_clear("STARTING")
|
||||
self.assertTrue(worker._claim("B"))
|
||||
|
||||
def test_capacity_2_allows_two_blocks_third(self):
|
||||
for task_id in ("A", "B", "C"):
|
||||
_insert_task(self.db_path, task_id)
|
||||
worker = self._worker(capacity=2)
|
||||
self.assertTrue(worker._claim("A"))
|
||||
self.assertTrue(worker._claim("B"))
|
||||
self.assertFalse(worker._claim("C"))
|
||||
self.assertEqual(_status(self.db_path, "C"), "queued")
|
||||
worker._running_clear("A")
|
||||
self.assertTrue(worker._claim("C"))
|
||||
|
||||
def test_lost_cas_race_releases_slot(self):
|
||||
_insert_task(self.db_path, "A")
|
||||
_insert_task(self.db_path, "B")
|
||||
first = self._worker(capacity=1)
|
||||
second = self._worker(capacity=1)
|
||||
self.assertTrue(first._claim("A"))
|
||||
# second 的 CAS 抢输：返回 False，且它自己的占位已清，
|
||||
# 不能留下一个空槽把自己堵死。
|
||||
self.assertFalse(second._claim("A"))
|
||||
self.assertEqual(second._running_snapshot(), {})
|
||||
self.assertTrue(second._claim("B"))
|
||||
|
||||
def test_two_threads_claim_only_one_wins(self):
|
||||
_insert_task(self.db_path, "A")
|
||||
_insert_task(self.db_path, "B")
|
||||
worker = self._worker(capacity=1)
|
||||
barrier = threading.Barrier(2, timeout=5)
|
||||
results = {}
|
||||
|
||||
def _try(task_id):
|
||||
barrier.wait()
|
||||
results[task_id] = worker._claim(task_id)
|
||||
|
||||
threads = [threading.Thread(target=_try, args=(t,))
|
||||
for t in ("A", "B")]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join(timeout=10)
|
||||
self.assertEqual(sorted(results.values()), [False, True])
|
||||
self.assertEqual(worker._running_count(), 1)
|
||||
|
||||
def test_missing_key_gate_releases_slot(self):
|
||||
_insert_task(self.db_path, "A")
|
||||
_insert_task(self.db_path, "B")
|
||||
# env 里无 key：run_one 在缺 key 门提前返回 failed，
|
||||
# 占位必须随之释放，否则 capacity=1 的 worker 会被堵死。
|
||||
worker = self._worker(env={}, capacity=1)
|
||||
result = worker.run_one()
|
||||
self.assertTrue(result["dispatched"])
|
||||
self.assertEqual(result["status"], "failed")
|
||||
self.assertEqual(_status(self.db_path, "A"), "failed")
|
||||
self.assertEqual(worker._running_snapshot(), {})
|
||||
self.assertTrue(worker._claim("B"))
|
||||
self.assertEqual(_status(self.db_path, "B"), "running")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in new issue
Block a user