host
ccb22c960d
Marker waits run after _run_attempt has cleared _running, so the run_loop heartbeat thread no longer refreshes the waiting parent. sweep_stale_tasks (default 300s) then kills a parent waiting up to the 3600s marker timeout, as seen in the marker live test (PM2 swept mid wait-threads). Poll and wait loops now heartbeat the parent each iteration (interval clamped to <=60s, well under the 300s threshold), best-effort, without touching _running, so no new capacity leak and cancel still works (empty _running entry has no proc/container to kill). Tests: 2 new sweep-survival tests (stale heartbeat at -299s, 2s wait, sweep must spare the parent); both fail 2/2 on pre-fix code and pass with the fix.
119 lines
4.8 KiB
Python
119 lines
4.8 KiB
Python
"""标记等待中父任务心跳单测：sweep 误杀回归。
|
|
|
|
根因：标记等待在 _run_attempt 返回后才跑，那时 _running 已被清掉，
|
|
run_loop 的心跳线程刷不到等候中的父任务；sweep 默认阈值 300 秒，
|
|
而标记等待标称可到 3600 秒，父任务等超 300 秒就被当死任务回收。
|
|
修法：轮询 / 等待循环每次迭代自己刷一次心跳（间隔至多 interval，
|
|
钳制后 <=60 秒，远小于 300 秒阈值），不依赖心跳线程。
|
|
|
|
不跑真实模型、不联网：用 init_db 建的临时 SQLite 库 + 最小任务行，
|
|
直接调 _deliver_marker_tools，不走 run_one 全量流程。
|
|
"""
|
|
|
|
import os
|
|
import sqlite3
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
try:
|
|
from cloud.server import init_db, sweep_stale_tasks
|
|
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, sweep_stale_tasks
|
|
from cloud.worker import Worker
|
|
|
|
FAKE_ENV = {"XCM_API_KEY": "test-placeholder"}
|
|
|
|
|
|
class MarkerWaitHeartbeatTestCase(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)
|
|
# 父任务 PM：running，但心跳已陈旧到离 sweep 阈值只差 1 秒。
|
|
# 不修时，等 2 秒后心跳就过期（301 > 300），sweep 会误杀；
|
|
# 修了时，等待循环每次迭代刷心跳，sweep 应放过它。
|
|
stale = time.time() - 299
|
|
self._insert("PM", status="running", last_heartbeat=stale,
|
|
updated_at=stale)
|
|
self._insert("E1", status="queued", parent="PM")
|
|
|
|
def tearDown(self):
|
|
self.tmp.cleanup()
|
|
|
|
def _insert(self, task_id, status="queued", parent=None,
|
|
last_heartbeat=None, updated_at=None):
|
|
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(?, 'p1', 'do work', ?, NULL, '[]', 0,"
|
|
" 'tester', ?, ?, ?)",
|
|
(task_id, status, now,
|
|
updated_at if updated_at is not None else now,
|
|
last_heartbeat))
|
|
if parent is not None:
|
|
conn.execute(
|
|
"UPDATE tasks SET parent_task_id = ? WHERE id = ?",
|
|
(parent, task_id))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
def _worker(self):
|
|
return Worker(self.db_path, "/fake/host.py", env=FAKE_ENV)
|
|
|
|
def _heartbeat(self, task_id):
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
row = conn.execute(
|
|
"SELECT status, last_heartbeat FROM tasks WHERE id = ?",
|
|
(task_id,)).fetchone()
|
|
return row[0], row[1]
|
|
finally:
|
|
conn.close()
|
|
|
|
def test_poll_wait_refreshes_heartbeat_and_survives_sweep(self):
|
|
worker = self._worker()
|
|
worker._deliver_marker_tools(worker.get_task("PM"), {
|
|
"status": "completed", "task_creates": [],
|
|
"task_polls": [{"task_id": "E1", "interval_seconds": 1,
|
|
"timeout_seconds": 2, "chars": 2000}],
|
|
"wait_threads": [], "task_cancels": [], "task_retries": []})
|
|
status, heartbeat = self._heartbeat("PM")
|
|
self.assertEqual(status, "running")
|
|
# 等待循环刷过心跳：心跳应比陈旧的 -299 秒新得多。
|
|
self.assertGreater(heartbeat, time.time() - 60)
|
|
swept = sweep_stale_tasks(self.db_path, 300)
|
|
self.assertNotIn("PM", swept)
|
|
status, _ = self._heartbeat("PM")
|
|
self.assertEqual(status, "running")
|
|
|
|
def test_wait_threads_refreshes_heartbeat_and_survives_sweep(self):
|
|
worker = self._worker()
|
|
worker._deliver_marker_tools(worker.get_task("PM"), {
|
|
"status": "completed", "task_creates": [],
|
|
"task_polls": [],
|
|
"wait_threads": [{"task_ids": ["E1"],
|
|
"interval_seconds": 1,
|
|
"timeout_seconds": 2, "chars": 2000}],
|
|
"task_cancels": [], "task_retries": []})
|
|
status, heartbeat = self._heartbeat("PM")
|
|
self.assertEqual(status, "running")
|
|
self.assertGreater(heartbeat, time.time() - 60)
|
|
swept = sweep_stale_tasks(self.db_path, 300)
|
|
self.assertNotIn("PM", swept)
|
|
status, _ = self._heartbeat("PM")
|
|
self.assertEqual(status, "running")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|