From ebfae7b9923a5eaa266b65adda5fc474ca1e12de Mon Sep 17 00:00:00 2001 From: host Date: Thu, 29 Oct 2026 05:37:14 +0800 Subject: [PATCH] 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). --- cloud/tests/test_run_all_pool.py | 143 +++++++++++++++++++++++++++++++ cloud/worker.py | 19 ++-- 2 files changed, 154 insertions(+), 8 deletions(-) create mode 100644 cloud/tests/test_run_all_pool.py diff --git a/cloud/tests/test_run_all_pool.py b/cloud/tests/test_run_all_pool.py new file mode 100644 index 0000000..8469f91 --- /dev/null +++ b/cloud/tests/test_run_all_pool.py @@ -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() diff --git a/cloud/worker.py b/cloud/worker.py index ffe3525..fb11313 100644 --- a/cloud/worker.py +++ b/cloud/worker.py @@ -4602,15 +4602,18 @@ class Worker: with results_lock: results.append(result) - try: - ready_count = len(self._ready_tasks()) - except Exception: - # 深度快照失败时退回开满线程，池线程自己的认领 - # 失败容错会处理，不让一次读库失败炸掉整轮。 - ready_count = self.capacity + # 池线程数恒为 capacity，绝不按本轮开始时的就绪快照缩减。 + # 快照缩减是 PM 等待死锁的根因：本轮开始时若只有 1 个就绪 + # 任务（如刚认领的 PM），只开 1 个线程；该线程跑 PM 模型回合 + # 后进入标记等待（最长 3600 秒），期间 PM 新建/重试的子任务 + # 入队，却再无空闲线程认领它——run_loop 又被本轮 run_all 的 + # join 挡住、开不了新一轮；若此时另一 worker 已缩容，就只剩 + # 一个被等待占住的 worker，子任务等满等待超时也无人跑。 + # 恒开 capacity 个线程没有空转代价：空闲线程在 task 为 None + # 且 active 为 0 时立刻退出，有在跑任务时才 0.2 秒轮询认领， + # 正好接住轮中新建的任务。 threads = [threading.Thread(target=_pool_thread, daemon=True) - for _ in range(min(self.capacity, - max(1, ready_count)))] + for _ in range(self.capacity)] for t in threads: t.start() for t in threads: