"""按队列深度自动扩缩容 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 deps is None: # depends_on 为 NULL（裸 SQL 插的任务、或老数据）表示 # 无依赖，绝不能变成 [None]——那样 ready_tasks 会判定 # 它依赖一个永远不存在的 None，任务永远不算就绪， # scaler 也就永远看不到深度（2026-10-26 实测教训： # worker 的 _task_from_row 对 NULL 处理正确所以能 # 正常认领，只有 scaler 的深度计算被这个 bug 致盲）。 deps = [] elif 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, "%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() 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()