host
ebfae7b992
run_all previously sized its pool from the ready snapshot at round start (min(capacity, ready_count)). If a round started with only 1 ready task (e.g. a just-claimed PM), only 1 thread ran; that thread entered a marker wait (up to 3600s) after its model turn, and children created/retried mid-round had no free thread to claim them. run_loop is blocked joining that round, so no new round starts; if the other worker had already scaled down, the queued child waited out the full wait timeout unclaimed. Fix: always start capacity threads. Idle threads exit immediately when no task is ready and nothing is active, and poll every 0.2s while something is active, so mid-round arrivals are claimed. Test: run_all with capacity=2, 1 initial task whose mocked run_one creates a child then polls for it; without the fix the PM poll times out in 5s (1 result), with the fix both tasks dispatch in the same round (0.2s).
144 lines
5.6 KiB
Python
144 lines
5.6 KiB
Python
"""run_all 池线程数单测：轮中新建任务不能死锁。
|
|
|
|
根因：run_all 曾按本轮开始时的就绪快照开线程（min(capacity,
|
|
ready_count)）。本轮开始时若只有 1 个就绪任务（如刚认领的 PM），
|
|
只开 1 个线程；该线程跑 PM 模型回合后进入标记等待（最长 3600 秒），
|
|
期间 PM 新建/重试的子任务入队，却再无空闲线程认领它——run_loop
|
|
又被本轮 run_all 的 join 挡住、开不了新一轮；若此时另一 worker
|
|
已缩容，就只剩一个被等待占住的 worker，子任务等满等待超时也
|
|
无人跑。修法：池线程数恒为 capacity。
|
|
|
|
不跑真实模型、不联网：用 init_db 建的临时 SQLite 库 + 一个模拟
|
|
run_one 的 Worker 子类——PM 任务先建子任务、再轮询等子任务完成；
|
|
子任务任务一被认领就立刻完成。无修法时 PM 轮询 5 秒超时（子任务
|
|
在本轮内无人认领，PM 失败后下一轮才跑子任务）；有修法时第二线程
|
|
在本轮内认领子任务，PM 轮询成功，两个任务都在本轮内派发。
|
|
"""
|
|
|
|
import os
|
|
import sqlite3
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
try:
|
|
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.server import init_db
|
|
from cloud.worker import Worker
|
|
|
|
FAKE_ENV = {"XCM_API_KEY": "test-placeholder"}
|
|
|
|
|
|
class _WaitWorker(Worker):
|
|
"""模拟 PM 等待：PM 建子任务后轮询等它完成，子任务立刻完成。"""
|
|
|
|
def _claim(self, task_id):
|
|
"""模拟真实 run_one 的 CAS 认领：只有第一个线程能拿到任务。"""
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
cur = conn.execute(
|
|
"UPDATE tasks SET status = 'running', updated_at = ?"
|
|
" WHERE id = ? AND status = 'queued'",
|
|
(time.time(), task_id))
|
|
conn.commit()
|
|
return cur.rowcount == 1
|
|
finally:
|
|
conn.close()
|
|
|
|
def run_one(self, task=None):
|
|
if task is None:
|
|
return {"dispatched": False}
|
|
task_id = task["id"]
|
|
if not self._claim(task_id):
|
|
return {"dispatched": False}
|
|
if task_id == "CHILD":
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
conn.execute(
|
|
"UPDATE tasks SET status = 'completed',"
|
|
" updated_at = ? WHERE id = 'CHILD'",
|
|
(time.time(),))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
return {"dispatched": True}
|
|
if task_id == "PM":
|
|
now = time.time()
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
conn.execute(
|
|
"INSERT OR IGNORE INTO tasks(id, project_id, prompt,"
|
|
" status, idempotency_key, depends_on, priority,"
|
|
" created_by, created_at, updated_at, last_heartbeat)"
|
|
" VALUES('CHILD', 'p1', 'do work', 'queued', NULL,"
|
|
" '[]', 0, 'tester', ?, ?, NULL)",
|
|
(now, now))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
deadline = time.time() + 5
|
|
while time.time() < deadline:
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
row = conn.execute(
|
|
"SELECT status FROM tasks WHERE id = 'CHILD'"
|
|
).fetchone()
|
|
finally:
|
|
conn.close()
|
|
if row is not None and row[0] == "completed":
|
|
return {"dispatched": True}
|
|
time.sleep(0.1)
|
|
# 超时 = 本轮内无人认领子任务，即死锁症状。
|
|
return {"dispatched": False}
|
|
return {"dispatched": False}
|
|
|
|
|
|
class RunAllPoolTestCase(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)
|
|
now = time.time()
|
|
conn = sqlite3.connect(self.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('PM', 'p1', 'do work', 'queued', NULL,"
|
|
" '[]', 0, 'tester', ?, ?, NULL)",
|
|
(now, now))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
def tearDown(self):
|
|
self.tmp.cleanup()
|
|
|
|
def test_mid_round_child_claimed_in_same_run_all(self):
|
|
worker = _WaitWorker(self.db_path, "/fake/host.py",
|
|
env=FAKE_ENV, capacity=2)
|
|
# max_tasks=2：模拟任务本身不标 completed（真实 run_one 会标），
|
|
# 用派发上限让 run_all 在 PM 与子任务都派发后收尾，避免测试
|
|
# 因 PM 仍 queued 而被池线程反复认领。
|
|
results = worker.run_all(max_tasks=2)
|
|
# 两个任务都在本轮内派发成功：PM 轮询等到了子任务完成。
|
|
self.assertEqual(len(results), 2)
|
|
self.assertTrue(all(r["dispatched"] for r in results))
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
row = conn.execute(
|
|
"SELECT status FROM tasks WHERE id = 'CHILD'").fetchone()
|
|
finally:
|
|
conn.close()
|
|
self.assertEqual(row[0], "completed")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|