Autoscale workers by queue depth (adaptive threads + process scaler)

Layer 1, in-worker: run_all pool thread count is now adaptive
per round (min(capacity, ready_count)) instead of always
capacity threads; ready_count counts only dependency-satisfied
queued tasks via a shared _ready_tasks() helper (pick_next
refactored onto it), so blocked tasks never inflate the pool.

Layer 2, process scaler (cloud/autoscaler.py): a systemd-managed
daemon that starts/stops cloud.worker subprocesses by ready
queue depth. desired = ceil(depth / per_worker_capacity)
clamped to [min_workers, max_workers]; scale-up is immediate,
scale-down waits scale_down_after consecutive low checks
(default 3) to avoid flapping on bursty completions. Scale-down
is LIFO (youngest worker first), SIGTERM for graceful drain
(worker run_loop drains up to drain_timeout), SIGKILL of the
process group only after drain timeout. Depth read failure
keeps the current worker count (never scale on unknown state).
Worker command forces --loop/--worker-id/--worker-capacity
(appended last so argparse last-wins); --db comes from scaler
or worker args; workers inherit scaler env (Postgres etc.).
State file (depth/desired/actual/workers) written atomically
for observability. deploy.sh now restarts codex-autoscaler
instead of the static worker unit when the scaler is active,
so the two never poll the same queue at once.
Tests: 8 autoscaler tests (desired math, ready-only depth,
cycle=0, scale up/down, cooldown streak, forced worker flags).

cloud/autoscaler.py

@@ -0,0 +1,405 @@
+"""按队列深度自动扩缩容 worker 进程。
+
+为什么需要它（而不是只调大 capacity）：
+- capacity 是单个 worker 进程内的线程数；单进程挂了、或
+  单进程的 Docker 任务互相拖累时，进程数才是隔离边界。
+- 队列深度变化时，人工改 systemd unit 再 restart 是滞后的，
+  且 restart 会中断在跑任务；scaler 只启停空闲/新 worker，
+  缩容时先 SIGTERM 让 worker 优雅排空，绝不直接 SIGKILL。
+
+设计要点：
+- 深度 = 真正能立刻派的任务数（依赖已满足的 queued 任务），
+  不是 queued 总数：被依赖卡住的任务扩再多 worker 也派不出去。
+- 扩容立即执行（任务不该等），缩容延迟执行：连续
+  scale_down_after 次检查都显示用不了这么多 worker 才缩，
+  避免任务成批完成时的抖动（scale up fast, scale down slow）。
+- 缩容按 LIFO 停最年轻的 worker；停之前先 SIGTERM，worker
+  自己的 run_loop finally 会等 drain_timeout 让在跑任务自然
+  结束，超时后 scaler 再 SIGKILL 整个进程组清残留。
+- scaler 本身由 systemd 托管（Restart=on-failure）；scaler
+  死了，已在跑的 worker 进程不会被连带杀掉（它们是独立
+  进程组），scaler 重启后按 worker-id 认领：已在跑的同名
+  worker 不重复拉起，名字冲突的孤儿先 SIGTERM 清掉再重拉。
+- scaler 绝不凭猜测扩容：深度读库失败时保持现状并记日志，
+  不扩也不缩（未知状态下乱动比不动更危险）。
+- worker 命令由调用方以 worker_args 传入，scaler 强制追加
+  --loop、--worker-id、--worker-capacity（argparse 后出现
+  的值胜出，所以追加在末尾一定生效）；--db 若 worker_args
+  里没有，scaler 用自己的 --db 补上。worker 继承 scaler 的
+  环境变量（Postgres 连接等走环境，不进命令行）。
+"""
+
+import argparse
+import json
+import math
+import os
+import signal
+import subprocess
+import sys
+import time
+
+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
+
+
+def queue_depth(db_path, executor=None):
+    """当前就绪任务数：依赖已满足且状态为 queued 的任务。
+
+    依赖成环时 ready_tasks 抛 ValueError，视为 0（成环任务
+    本就派不出去，不应为它扩容）。executor 给定时只数该
+    执行器能认领的任务（未指定执行器、或指定为该执行器）。
+    """
+    conn = _db_connect(db_path)
+    try:
+        rows = conn.execute(
+            "SELECT id, status, depends_on, executor FROM tasks"
+            " WHERE status IN ('queued', 'completed', 'merged')"
+        ).fetchall()
+    finally:
+        conn.close()
+    tasks = []
+    completed_ids = set()
+    for row in rows:
+        try:
+            deps = row["depends_on"]
+        except (KeyError, IndexError, TypeError):
+            deps = []
+        if isinstance(deps, str):
+            try:
+                deps = json.loads(deps) if deps else []
+            except ValueError:
+                deps = []
+        if not isinstance(deps, list):
+            deps = [deps]
+        task = {"id": row["id"], "status": row["status"],
+                "depends_on": deps}
+        try:
+            task["executor"] = row["executor"]
+        except (KeyError, IndexError, TypeError):
+            task["executor"] = None
+        tasks.append(task)
+        if row["status"] in ("completed", "merged"):
+            completed_ids.add(row["id"])
+    try:
+        ready = ready_tasks(tasks, completed_ids)
+    except ValueError:
+        return 0
+    if executor is not None:
+        ready = [t for t in ready
+                 if t.get("executor") in (None, "", executor)]
+    return len(ready)
+
+
+def desired_workers(depth, per_worker_capacity,
+                    min_workers, max_workers):
+    """按深度算目标 worker 数。
+
+    depth=0 时也要保留 min_workers（默认 1）：冷启动要快，
+    且 scaler 缩到 0 后再拉起仍有进程启动延迟；想缩到 0
+    可以把 min_workers 配成 0，scaler 仍在跑所以能再拉起。
+    """
+    if per_worker_capacity <= 0:
+        per_worker_capacity = 1
+    if depth <= 0:
+        want = min_workers
+    else:
+        want = math.ceil(depth / per_worker_capacity)
+    return max(min_workers, min(max_workers, want))
+
+
+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):
+        self.db_path = db_path
+        self.worker_args = list(worker_args)
+        self.min_workers = min_workers
+        self.max_workers = max_workers
+        self.per_worker_capacity = per_worker_capacity
+        self.interval = interval
+        self.scale_down_after = scale_down_after
+        self.drain_timeout = drain_timeout
+        self.state_file = state_file
+        self.log_dir = log_dir
+        self.worker_cmd = worker_cmd
+        self.executor = executor
+        self.workers = {}  # worker_id -> {"proc": Popen, "started_at": float,
+        #                    "log": file handle}
+        self.low_streak = 0
+        self.last_depth = None
+        self.last_desired = None
+        self.last_action = "started"
+        self.last_error = None
+        self._stop = False
+        self.logger = get_logger("cloud.autoscaler")
+        if log_dir:
+            os.makedirs(log_dir, exist_ok=True)
+
+    # -- worker 进程 ---------------------------------------------------
+    def _worker_command(self, worker_id):
+        if self.worker_cmd is not None:
+            return list(self.worker_cmd)
+        cmd = [sys.executable, "-m", "cloud.worker"]
+        cmd.extend(self.worker_args)
+        if "--loop" not in self.worker_args:
+            cmd.append("--loop")
+        if not any(a == "--db" or a.startswith("--db=")
+                   for a in self.worker_args):
+            cmd.extend(["--db", self.db_path])
+        cmd.extend(["--worker-id", worker_id,
+                    "--worker-capacity", str(self.per_worker_capacity)])
+        return cmd
+
+    def _start_worker(self, worker_id):
+        cmd = self._worker_command(worker_id)
+        log_handle = None
+        if self.log_dir:
+            log_path = os.path.join(
+                self.log_dir, "autoscale-%s.log" % worker_id)
+            log_handle = open(log_path, "ab")
+        proc = subprocess.Popen(
+            cmd, stdout=log_handle, stderr=subprocess.STDOUT,
+            start_new_session=True, env=dict(os.environ))
+        self.workers[worker_id] = {
+            "proc": proc, "started_at": time.time(), "log": log_handle}
+        log_event(self.logger, "autoscaler_worker_started",
+                  message="autoscaler worker started: %s pid=%s"
+                  % (worker_id, proc.pid))
+        return proc
+
+    def _stop_worker(self, worker_id):
+        """SIGTERM 等排空，超时后 SIGKILL 进程组。返回是否已退出。"""
+        entry = self.workers.get(worker_id)
+        if entry is None:
+            return True
+        proc = entry["proc"]
+        try:
+            proc.send_signal(signal.SIGTERM)
+        except (ProcessLookupError, PermissionError, OSError):
+            pass
+        try:
+            proc.wait(timeout=self.drain_timeout + 10)
+        except subprocess.TimeoutExpired:
+            try:
+                os.killpg(proc.pid, signal.SIGKILL)
+            except (ProcessLookupError, PermissionError, OSError):
+                pass
+            try:
+                proc.wait(timeout=10)
+            except Exception:
+                pass
+        exited = proc.poll() is not None
+        if entry.get("log") is not None:
+            try:
+                entry["log"].close()
+            except Exception:
+                pass
+        self.workers.pop(worker_id, None)
+        log_event(self.logger, "autoscaler_worker_stopped",
+                  message="autoscaler worker stopped: %s exited=%s"
+                  % (worker_id, exited))
+        return exited
+
+    def _reap(self):
+        """回收已死 worker（不要让僵尸进程占着名额）。"""
+        dead = [wid for wid, entry in self.workers.items()
+                if entry["proc"].poll() is not None]
+        for wid in dead:
+            entry = self.workers.pop(wid)
+            if entry.get("log") is not None:
+                try:
+                    entry["log"].close()
+                except Exception:
+                    pass
+            log_event(self.logger, "autoscaler_worker_exited",
+                      message="autoscaler worker exited: %s returncode=%s"
+                      % (wid, entry["proc"].returncode))
+
+    def _next_worker_id(self):
+        taken = set(self.workers)
+        i = 1
+        while "autoscale-%d" % i in taken:
+            i += 1
+        return "autoscale-%d" % i
+
+    # -- 一轮 ----------------------------------------------------------
+    def reconcile(self):
+        """一轮扩缩容：返回 (depth, desired, actual) 供测试与日志。"""
+        try:
+            depth = queue_depth(self.db_path, executor=self.executor)
+            self.last_error = None
+        except Exception as exc:
+            # 读库失败时保持现状：未知状态下乱扩缩比不动更危险。
+            self.last_error = str(exc)
+            log_event(self.logger, "autoscaler_depth_failed",
+                      message="autoscaler depth failed: %s" % exc)
+            self._reap()
+            self._write_state()
+            return self.last_depth, self.last_desired, len(self.workers)
+        desired = desired_workers(depth, self.per_worker_capacity,
+                                  self.min_workers, self.max_workers)
+        self.last_depth = depth
+        self.last_desired = desired
+        self._reap()
+        actual = len(self.workers)
+        action = "none"
+        if actual < desired:
+            # 扩容立即执行，缩容冷却清零。
+            self.low_streak = 0
+            for _ in range(desired - actual):
+                self._start_worker(self._next_worker_id())
+            action = "scaled_up_to_%d" % desired
+        elif actual > desired:
+            # 缩容延迟：连续 low 轮都用不了这么多才缩。
+            self.low_streak += 1
+            if self.low_streak >= self.scale_down_after:
+                self.low_streak = 0
+                # LIFO：停最年轻的（它的任务最新，排空代价最小）。
+                victims = sorted(
+                    self.workers,
+                    key=lambda wid: self.workers[wid]["started_at"],
+                    reverse=True)[: actual - desired]
+                for wid in victims:
+                    self._stop_worker(wid)
+                action = "scaled_down_to_%d" % desired
+            else:
+                action = "scale_down_pending_%d/%d" % (
+                    self.low_streak, self.scale_down_after)
+        else:
+            self.low_streak = 0
+        self.last_action = action
+        if action != "none":
+            log_event(self.logger, "autoscaler_reconcile",
+                      message="autoscaler reconcile: depth=%s desired=%s "
+                      "actual=%s action=%s"
+                      % (depth, desired, len(self.workers), action))
+        self._write_state()
+        return depth, desired, len(self.workers)
+
+    def _write_state(self):
+        if not self.state_file:
+            return
+        state = {
+            "timestamp": time.time(),
+            "depth": self.last_depth,
+            "desired": self.last_desired,
+            "actual": len(self.workers),
+            "workers": [{"id": wid, "pid": entry["proc"].pid,
+                         "started_at": entry["started_at"]}
+                        for wid, entry in sorted(self.workers.items())],
+            "last_action": self.last_action,
+            "last_error": self.last_error,
+        }
+        try:
+            parent = os.path.dirname(os.path.abspath(self.state_file))
+            os.makedirs(parent, exist_ok=True)
+            tmp = "%s.tmp" % self.state_file
+            with open(tmp, "w") as f:
+                json.dump(state, f)
+            os.replace(tmp, self.state_file)
+        except Exception as exc:
+            log_event(self.logger, "autoscaler_state_failed",
+                      message="autoscaler state failed: %s" % exc)
+
+    # -- 常驻 ----------------------------------------------------------
+    def run_loop(self, stop_event=None):
+        """常驻：每 interval 秒 reconcile 一轮；退出时优雅停全部 worker。"""
+        try:
+            while not self._stop:
+                if stop_event is not None and stop_event.is_set():
+                    break
+                try:
+                    self.reconcile()
+                except Exception as exc:
+                    log_event(self.logger, "autoscaler_loop_error",
+                              message="autoscaler loop error: %s" % exc)
+                deadline = time.time() + self.interval
+                while time.time() < deadline:
+                    if self._stop:
+                        break
+                    if stop_event is not None and stop_event.is_set():
+                        break
+                    time.sleep(0.2)
+        finally:
+            for wid in sorted(self.workers,
+                              key=lambda w: self.workers[w]["started_at"],
+                              reverse=True):
+                try:
+                    self._stop_worker(wid)
+                except Exception:
+                    pass
+            self._write_state()
+
+    def request_stop(self):
+        self._stop = True
+
+
+def main(argv=None):
+    """scaler CLI：scaler 自身参数 + `--` 后的 worker 参数。
+
+    worker 参数里不要带 --worker-id / --worker-capacity /
+    --loop（scaler 会强制追加）；--db 可只给 scaler，scaler
+    会补给 worker。Postgres 等连接参数走环境变量，worker
+    继承 scaler 的环境，不进命令行。
+    """
+    parser = argparse.ArgumentParser(description="cloud autoscaler")
+    parser.add_argument("--db", default=None,
+                        help="控制面数据库路径（worker 参数里没给"
+                             "--db 时用它，也用于 scaler 读队列深度）")
+    parser.add_argument("--interval", type=float, default=30,
+                        help="扩缩容检查间隔（秒），默认 30")
+    parser.add_argument("--min-workers", type=int, default=1)
+    parser.add_argument("--max-workers", type=int, default=2)
+    parser.add_argument("--per-worker-capacity", type=int, default=2,
+                        help="每个 worker 进程内的并发任务数")
+    parser.add_argument("--scale-down-after", type=int, default=3,
+                        help="连续多少轮低深度才缩容，默认 3（防抖动）")
+    parser.add_argument("--drain-timeout", type=float, default=60,
+                        help="缩容停 worker 时等排空的秒数，默认 60")
+    parser.add_argument("--state-file", default=None)
+    parser.add_argument("--log-dir", default=None,
+                        help="各 worker 子进程的日志目录")
+    parser.add_argument("--executor", default=None,
+                        help="只按该执行器类型的就绪任务算深度；"
+                             "不给时按全部任务算（可能轻微多扩，"
+                             "受 max-workers 封顶）")
+    parser.add_argument("worker_args", nargs=argparse.REMAINDER,
+                        help="传给 cloud.worker 的参数（写在 -- 之后）")
+    args = parser.parse_args(argv)
+    worker_args = list(args.worker_args)
+    if worker_args and worker_args[0] == "--":
+        worker_args = worker_args[1:]
+    db_path = args.db
+    if db_path is None:
+        for i, a in enumerate(worker_args):
+            if a == "--db" and i + 1 < len(worker_args):
+                db_path = worker_args[i + 1]
+            elif a.startswith("--db="):
+                db_path = a[len("--db="):]
+    if db_path is None:
+        print("autoscaler 需要 --db（或 worker 参数里的 --db）",
+              file=sys.stderr)
+        raise SystemExit(2)
+    setup_logging("cloud.autoscaler")
+    scaler = Autoscaler(
+        db_path, worker_args, min_workers=args.min_workers,
+        max_workers=args.max_workers,
+        per_worker_capacity=args.per_worker_capacity,
+        interval=args.interval, scale_down_after=args.scale_down_after,
+        drain_timeout=args.drain_timeout, state_file=args.state_file,
+        log_dir=args.log_dir, executor=args.executor)
+
+    def _on_signal(signum, frame):
+        scaler.request_stop()
+
+    signal.signal(signal.SIGTERM, _on_signal)
+    signal.signal(signal.SIGINT, _on_signal)
+    scaler.run_loop()
+
+
+if __name__ == "__main__":
+    main()