host
b74deffb26
queue_depth turned a NULL depends_on into [None], and ready_tasks then required None to be in completed_ids, so such tasks never counted as ready. Worker pick_next was unaffected (_task_from_row maps NULL to []), which is why tasks still ran while the scaler always saw depth 0 and never scaled up (proven live: depth probe with 1 queued task returned 0). NULL now means no dependencies, with a regression test.
414 lines
17 KiB
Python
414 lines
17 KiB
Python
"""按队列深度自动扩缩容 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()
|