Merge pull request #117 from openqj/feat-autoscaler-coord-guard

autoscaler: guard scale-down against live coordinators and busy victims

cloud/autoscaler.py

@@ -22,6 +22,23 @@
   worker 不重复拉起，名字冲突的孤儿先 SIGTERM 清掉再重拉。
 - scaler 绝不凭猜测扩容：深度读库失败时保持现状并记日志，
   不扩也不缩（未知状态下乱动比不动更危险）。
+- 缩容护栏（两层 veto，见 _has_live_running_coordinator /
+  _victim_is_busy）：深度只算 queued 就绪任务，等待中的
+  协调员对深度贡献为 0，且协调员等齐循环不响应停机信号、
+  60 秒 drain 内等不完就会被 run_loop finally 标 interrupted，
+  所以缩容绝不能只看深度就停 worker。第一层是全局活协调员
+  veto：DB 里有 status=running、coordinator_parts 经
+  cloud.coordinator.parse_parts 解析为非空列表、且
+  last_heartbeat 在 120 秒内的任务时，本轮不执行任何缩容，
+  low_streak 冻结不重置（veto 解除后下一轮低深度即可缩，
+  不必再等 scale_down_after 轮）。第二层是受害者忙 veto：
+  按 LIFO 选到受害者后查 workers 表，alive 且 running>0 时
+  跳过该受害者；但注册表 running 有幽灵占位先例（报 busy
+  却库里无 running 任务），所以只有 DB 里同时存在活 running
+  任务时才拦，幽灵不拦。tasks 表没有 worker_id 列，无法把
+  协调员归属到某个 worker，因此第一层只能是全局 veto，
+  不能是按受害者 veto——保守一点，代价只是多留一个空闲
+  worker 到协调员跑完。
 - worker 命令由调用方以 worker_args 传入，scaler 强制追加
   --loop、--worker-id、--worker-capacity（argparse 后出现
   的值胜出，所以追加在末尾一定生效）；--db 若 worker_args
@@ -39,9 +56,11 @@ import subprocess
 import sys
 import time
 
+from cloud.coordinator import parse_parts as _coord_parse_parts
 from cloud.db import connect as _db_connect
 from cloud.logging_config import get_logger, log_event, setup_logging
 from cloud.queue import ready_tasks
+from cloud.workers import DEFAULT_ALIVE_SECONDS as _WORKER_ALIVE_SECONDS
 
 
 def queue_depth(db_path, executor=None):
@@ -117,13 +136,129 @@ def desired_workers(depth, per_worker_capacity,
     return max(min_workers, min(max_workers, want))
 
 
+def _is_live_heartbeat(last_heartbeat, now, alive_seconds):
+    """心跳是否新鲜：NULL 算不新鲜；未来时间戳（时钟/精度舍入）
+    也算新鲜，因为它一定是刚写过的。"""
+    if last_heartbeat is None:
+        return False
+    try:
+        age = now - float(last_heartbeat)
+    except (TypeError, ValueError):
+        return False
+    return age <= float(alive_seconds)
+
+
+def _running_task_rows(db_path):
+    """status=running 的任务行（coordinator_parts, last_heartbeat）。
+    调用方按需过滤；读库失败时抛给调用方决定 fail-open/closed。"""
+    conn = _db_connect(db_path)
+    try:
+        return conn.execute(
+            "SELECT coordinator_parts, last_heartbeat FROM tasks"
+            " WHERE status = 'running'").fetchall()
+    finally:
+        conn.close()
+
+
+def _has_live_running_coordinator(db_path, now=None,
+                                  alive_seconds=_WORKER_ALIVE_SECONDS):
+    """DB 里是否存在活着的等待中协调员。
+
+    判定与 worker 侧一致：coordinator_parts 经 parse_parts
+    解析为非空列表才算协调员（坏值解析为 []，worker 也不会
+    把它当协调员跑等齐循环）；last_heartbeat 必须新鲜，
+    否则是死 worker 遗留的 running 行，300 秒陈旧回收会去
+    处理它，不应由它永久拦住缩容。
+    """
+    if now is None:
+        now = time.time()
+    for row in _running_task_rows(db_path):
+        try:
+            parts = row["coordinator_parts"]
+            hb = row["last_heartbeat"]
+        except (KeyError, IndexError, TypeError):
+            continue
+        if not _is_live_heartbeat(hb, now, alive_seconds):
+            continue
+        if _coord_parse_parts(parts):
+            return True
+    return False
+
+
+def _has_live_running_task(db_path, now=None,
+                           alive_seconds=_WORKER_ALIVE_SECONDS):
+    """DB 里是否存在心跳新鲜的 running 任务（不限协调员）。"""
+    if now is None:
+        now = time.time()
+    for row in _running_task_rows(db_path):
+        try:
+            hb = row["last_heartbeat"]
+        except (KeyError, IndexError, TypeError):
+            continue
+        if _is_live_heartbeat(hb, now, alive_seconds):
+            return True
+    return False
+
+
+def _worker_registry_running(db_path, worker_id, now=None,
+                             alive_seconds=_WORKER_ALIVE_SECONDS):
+    """workers 表里该 worker 上报的 running 数；行缺失、stale、
+    或表缺失时一律按 0（空闲）处理——注册发生在 run_loop
+    启动、认领之前，所以已认领任务的 worker 一定已有行。"""
+    if now is None:
+        now = time.time()
+    try:
+        conn = _db_connect(db_path)
+        try:
+            row = conn.execute(
+                "SELECT running, last_seen FROM workers"
+                " WHERE worker_id = ?", (worker_id,)).fetchone()
+        finally:
+            conn.close()
+    except Exception:
+        return 0
+    if row is None:
+        return 0
+    try:
+        running = row["running"]
+        last_seen = row["last_seen"]
+    except (KeyError, IndexError, TypeError):
+        return 0
+    if not _is_live_heartbeat(last_seen, now, alive_seconds):
+        return 0
+    try:
+        return int(running) if running is not None else 0
+    except (TypeError, ValueError):
+        return 0
+
+
+def _victim_is_busy(db_path, worker_id, now=None,
+                    alive_seconds=_WORKER_ALIVE_SECONDS):
+    """受害者是否真忙：注册表报 busy 且 DB 里确实有活 running
+    任务时才算忙；注册表报 busy 但 DB 里一个活任务都没有时
+    是幽灵占位，不拦（否则幽灵会把缩容永久卡死）。"""
+    if now is None:
+        now = time.time()
+    if _worker_registry_running(db_path, worker_id, now,
+                                alive_seconds) <= 0:
+        return False
+    try:
+        return _has_live_running_task(db_path, now, alive_seconds)
+    except Exception:
+        # 第二层 veto 读库失败时 fail-open：第一层协调员 veto
+        # 仍是 fail-closed，见 reconcile；这里不因一次读失败
+        # 就把一个可能空闲的 worker 永久留着。
+        return False
+
+
 class Autoscaler:
     """管理 worker 子进程：按深度启停，带缩容冷却与优雅停机。"""
 
     def __init__(self, db_path, worker_args, min_workers=1, max_workers=2,
                  per_worker_capacity=2, interval=30, scale_down_after=3,
                  drain_timeout=60, state_file=None, log_dir=None,
-                 worker_cmd=None, executor=None):
+                 worker_cmd=None, executor=None,
+                 guard_alive_seconds=_WORKER_ALIVE_SECONDS):
         self.db_path = db_path
         self.worker_args = list(worker_args)
         self.min_workers = min_workers
@@ -132,6 +267,7 @@ class Autoscaler:
         self.interval = interval
         self.scale_down_after = scale_down_after
         self.drain_timeout = drain_timeout
+        self.guard_alive_seconds = guard_alive_seconds
         self.state_file = state_file
         self.log_dir = log_dir
         self.worker_cmd = worker_cmd
@@ -273,15 +409,51 @@ class Autoscaler:
             # 缩容延迟：连续 low 轮都用不了这么多才缩。
             self.low_streak += 1
             if self.low_streak >= self.scale_down_after:
-                self.low_streak = 0
+                # 缩容护栏（两层 veto）：被拦时 low_streak 不重置
-                # LIFO：停最年轻的（它的任务最新，排空代价最小）。
+                # （冻结），veto 解除后下一轮低深度即可缩，不必
-                victims = sorted(
+                # 再等 scale_down_after 轮；只有真正停了 worker
-                    self.workers,
+                # 才重置 streak。
-                    key=lambda wid: self.workers[wid]["started_at"],
+                try:
-                    reverse=True)[: actual - desired]
+                    coord_blocked = _has_live_running_coordinator(
-                for wid in victims:
+                        self.db_path,
-                    self._stop_worker(wid)
+                        alive_seconds=self.guard_alive_seconds)
-                action = "scaled_down_to_%d" % desired
+                except Exception as exc:
+                    # 第一层读库失败时 fail-closed：未知状态下
+                    # 停 worker 比不动更危险，本轮不缩。
+                    self.last_error = str(exc)
+                    log_event(self.logger, "autoscaler_guard_failed",
+                              message="autoscaler guard failed: %s"
+                              % exc)
+                    self.last_action = "scale_down_blocked_by_guard_error"
+                    self._write_state()
+                    return depth, desired, len(self.workers)
+                if coord_blocked:
+                    action = "scale_down_blocked_by_coordinator"
+                else:
+                    # LIFO：停最年轻的（它的任务最新，排空代价最小）。
+                    victims = sorted(
+                        self.workers,
+                        key=lambda wid: self.workers[wid]["started_at"],
+                        reverse=True)[: actual - desired]
+                    stopped = []
+                    blocked = []
+                    for wid in victims:
+                        if _victim_is_busy(
+                                self.db_path, wid,
+                                alive_seconds=self.guard_alive_seconds):
+                            blocked.append(wid)
+                            continue
+                        self._stop_worker(wid)
+                        stopped.append(wid)
+                    if not stopped:
+                        action = "scale_down_blocked_by_busy_worker"
+                    else:
+                        self.low_streak = 0
+                        if blocked:
+                            action = "scaled_down_to_%d" % len(
+                                self.workers)
+                        else:
+                            action = "scaled_down_to_%d" % desired
             else:
                 action = "scale_down_pending_%d/%d" % (
                     self.low_streak, self.scale_down_after)

cloud/tests/test_autoscaler.py

@@ -49,6 +49,35 @@ def _set_status(db_path, task_id, status):
         conn.close()
 
 
+def _insert_running_task(db_path, task_id, parts=None,
+                         heartbeat_offset=0):
+    """插入一个 running 任务；parts 为 coordinator_parts 列的
+    原始值（None=普通任务，JSON 文本=协调员分工）。"""
+    _insert_task(db_path, task_id, status="running")
+    now = time.time()
+    conn = sqlite3.connect(db_path)
+    try:
+        conn.execute(
+            "UPDATE tasks SET coordinator_parts = ?,"
+            " last_heartbeat = ? WHERE id = ?",
+            (parts, now + heartbeat_offset, task_id))
+        conn.commit()
+    finally:
+        conn.close()
+
+
+def _set_worker_running(db_path, worker_id, running):
+    """在 workers 表里给 worker 写一行新鲜心跳 + running 数。
+    用 cloud.workers 的注册函数建表，避免测试自己拼表结构。"""
+    try:
+        from cloud.workers import register_worker
+    except ImportError:
+        sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
+        from cloud.workers import register_worker
+    register_worker(db_path, worker_id, executor="docker",
+                    capacity=2, running=running)
+
+
 class DesiredWorkersTestCase(unittest.TestCase):
     def test_depth_zero_keeps_min(self):
         self.assertEqual(desired_workers(0, 2, 1, 2), 1)
@@ -195,5 +224,126 @@ class AutoscalerReconcileTestCase(unittest.TestCase):
         mock_loop.assert_called_once()
 
 
+class ScaleDownGuardTestCase(unittest.TestCase):
+    """缩容护栏单测：两层 veto 与 streak 冻结。
+
+    全部用临时 SQLite + 假 worker（sleep 进程），不跑真实
+    worker。注意：假 worker 没有 workers 表心跳，注册行
+    全部由 _set_worker_running 显式写；没写注册行的 worker
+    在护栏里按空闲处理（见 autoscaler docstring）。
+    """
+
+    def setUp(self):
+        self.tmp = tempfile.TemporaryDirectory()
+        self.db_path = os.path.join(self.tmp.name, "control.db")
+        init_db(self.db_path)
+        self.fake_cmd = [sys.executable, "-c",
+                          "import time; time.sleep(120)"]
+
+    def tearDown(self):
+        if hasattr(self, "scaler"):
+            for wid in list(self.scaler.workers):
+                try:
+                    self.scaler._stop_worker(wid)
+                except Exception:
+                    pass
+        self.tmp.cleanup()
+
+    def _scaler(self, **kwargs):
+        kwargs.setdefault("min_workers", 1)
+        kwargs.setdefault("per_worker_capacity", 2)
+        kwargs.setdefault("scale_down_after", 1)
+        kwargs.setdefault("drain_timeout", 5)
+        self.scaler = Autoscaler(
+            self.db_path, [], worker_cmd=self.fake_cmd, **kwargs)
+        return self.scaler
+
+    def _scale_to(self, scaler, count):
+        for i in range(count * 2):
+            _insert_task(self.db_path, "Q%d" % i)
+        scaler.reconcile()
+        self.assertEqual(len(scaler.workers), count)
+        for i in range(count * 2):
+            _set_status(self.db_path, "Q%d" % i, "completed")
+
+    def test_live_coordinator_blocks_scale_down(self):
+        scaler = self._scaler(max_workers=2)
+        self._scale_to(scaler, 2)
+        _insert_running_task(
+            self.db_path, "COORD",
+            parts='[{"name": "A", "prompt": "x"}]')
+        scaler.reconcile()
+        self.assertEqual(len(scaler.workers), 2)
+        self.assertEqual(scaler.last_action,
+                         "scale_down_blocked_by_coordinator")
+
+    def test_stale_coordinator_does_not_block(self):
+        scaler = self._scaler(max_workers=2)
+        self._scale_to(scaler, 2)
+        _insert_running_task(
+            self.db_path, "COORD",
+            parts='[{"name": "A", "prompt": "x"}]',
+            heartbeat_offset=-300)
+        scaler.reconcile()
+        self.assertEqual(len(scaler.workers), 1)
+
+    def test_non_coordinator_running_does_not_trigger_coordinator_veto(self):
+        scaler = self._scaler(max_workers=2)
+        self._scale_to(scaler, 2)
+        _insert_running_task(self.db_path, "PLAIN", parts=None)
+        _insert_running_task(self.db_path, "EMPTY", parts="[]")
+        scaler.reconcile()
+        self.assertEqual(len(scaler.workers), 1)
+        self.assertEqual(scaler.last_action, "scaled_down_to_1")
+
+    def test_invalid_coordinator_parts_not_a_coordinator(self):
+        # 与 worker 侧一致：parse_parts 把坏值解析为 []，
+        # worker 不会把它当协调员跑等齐循环，护栏也不拦。
+        scaler = self._scaler(max_workers=2)
+        self._scale_to(scaler, 2)
+        _insert_running_task(self.db_path, "BAD", parts="not-json")
+        scaler.reconcile()
+        self.assertEqual(len(scaler.workers), 1)
+
+    def test_busy_victim_skipped_idle_victim_stopped(self):
+        scaler = self._scaler(max_workers=3)
+        self._scale_to(scaler, 3)
+        # DB 里有一个活 running 任务（幽灵交叉校验的另一半）。
+        _insert_running_task(self.db_path, "LIVE", parts=None)
+        _set_worker_running(self.db_path, "autoscale-1", 0)
+        _set_worker_running(self.db_path, "autoscale-2", 0)
+        _set_worker_running(self.db_path, "autoscale-3", 1)
+        scaler.reconcile()
+        # LIFO 受害者是 [3, 2]：3 忙跳过，2 闲停掉，实际 3->2。
+        self.assertEqual(len(scaler.workers), 2)
+        self.assertIn("autoscale-3", scaler.workers)
+        self.assertNotIn("autoscale-2", scaler.workers)
+        self.assertEqual(scaler.last_action, "scaled_down_to_2")
+
+    def test_ghost_registry_does_not_block(self):
+        scaler = self._scaler(max_workers=2)
+        self._scale_to(scaler, 2)
+        # 注册表报 busy，但 DB 里一个活 running 任务都没有：
+        # 幽灵占位，不拦。
+        _set_worker_running(self.db_path, "autoscale-2", 2)
+        scaler.reconcile()
+        self.assertEqual(len(scaler.workers), 1)
+
+    def test_coordinator_veto_freezes_low_streak(self):
+        scaler = self._scaler(max_workers=2, scale_down_after=3)
+        self._scale_to(scaler, 2)
+        _insert_running_task(
+            self.db_path, "COORD",
+            parts='[{"name": "A", "prompt": "x"}]')
+        for _ in range(3):
+            scaler.reconcile()
+            self.assertEqual(len(scaler.workers), 2)
+        # veto 解除后不需要再等 3 轮：streak 已冻结在阈值上，
+        # 下一轮低深度立即缩。
+        _set_status(self.db_path, "COORD", "completed")
+        scaler.reconcile()
+        self.assertEqual(len(scaler.workers), 1)
+
+
 if __name__ == "__main__":
     unittest.main()