worker: always start capacity pool threads in run_all
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).
This commit is contained in:
1 parent
ab0aaf4a5e
commit
ebfae7b992
2 files changed
+154
-8
No files matched your search
@@ -0,0 +1,143 @@
|
||||
"""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()
|
||||
Reference in new issue
Block a user