"""按队列深度自动扩缩容 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 绝不凭猜测扩容：深度读库失败时保持现状并记日志， 不扩也不缩（未知状态下乱动比不动更危险）。 - 缩容护栏（两层 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 里没有，scaler 用自己的 --db 补上。worker 继承 scaler 的 环境变量（Postgres 连接等走环境，不进命令行）。 """ import argparse import json import math import os import signal import socket 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): """当前就绪任务数：依赖已满足且状态为 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)) 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, guard_alive_seconds=_WORKER_ALIVE_SECONDS): 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.guard_alive_seconds = guard_alive_seconds 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)]) # HA 阶段 1：宿主名强制追加（末尾胜出）。scaler 跑在宿主上， # socket.gethostname() 就是宿主名；容器内 worker 若用自己的 # hostname 会拿到容器名，worktree_host 就错了，所以这里必须 # 显式传，不能让 worker 自己猜。 if not any(a == "--host-name" or a.startswith("--host-name=") for a in self.worker_args): cmd.extend(["--host-name", socket.gethostname()]) 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: # 缩容护栏（两层 veto）：被拦时 low_streak 不重置 # （冻结），veto 解除后下一轮低深度即可缩，不必 # 再等 scale_down_after 轮；只有真正停了 worker # 才重置 streak。 try: coord_blocked = _has_live_running_coordinator( self.db_path, alive_seconds=self.guard_alive_seconds) 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) 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()