host
37b47cadc6
Two veto layers in reconcile(), worker.py untouched: - Global live-coordinator veto: a running task whose coordinator_parts parse (via cloud.coordinator.parse_parts) to a non-empty list and whose last_heartbeat is within 120s blocks all scale-down this round; low_streak is frozen, not reset, so scale-down resumes on the next low-depth round once the veto lifts. - Per-victim busy veto: a LIFO victim whose workers row is alive with running>0 is skipped, but only when the DB also has a live running task (ghost registry rows with no live DB task do not block). - Coordinator-veto query failure fails closed (no stop, guard error recorded); victim-veto read failure fails open for that layer only. - A missing/stale workers row is treated as idle. Tests: 7 new ScaleDownGuardTestCase cases (live/stale coordinator, non-coordinator, invalid parts, busy/idle victims, ghost registry, streak freeze); cloud suite 1077 tests OK (4 skipped).
350 lines
14 KiB
Python
350 lines
14 KiB
Python
"""Autoscaler 单测：深度计算、目标 worker 数、启停与缩容冷却。
|
||
|
||
只用临时 SQLite 与假 worker 命令（sleep 进程），不跑真实
|
||
worker、不联网、不碰生产库。假 worker 命令通过 worker_cmd
|
||
注入，scaler 的 worker_args 为空——真实 worker 命令构造
|
||
（强制 --loop/--worker-id/--worker-capacity）由构造测试覆盖。
|
||
"""
|
||
|
||
import os
|
||
import sqlite3
|
||
import sys
|
||
import tempfile
|
||
import time
|
||
import unittest
|
||
from pathlib import Path
|
||
|
||
try:
|
||
from cloud.autoscaler import Autoscaler, desired_workers, queue_depth
|
||
from cloud.server import init_db
|
||
except ImportError:
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||
from cloud.autoscaler import Autoscaler, desired_workers, queue_depth
|
||
from cloud.server import init_db
|
||
|
||
|
||
def _insert_task(db_path, task_id, status="queued", depends_on="[]"):
|
||
now = time.time()
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
conn.execute(
|
||
"INSERT INTO tasks(id, project_id, prompt, status,"
|
||
" idempotency_key, depends_on, priority, created_by,"
|
||
" created_at, updated_at, last_heartbeat)"
|
||
" VALUES(?, 'p1', 'x', ?, NULL, ?, 0, 'tester', ?, ?, NULL)",
|
||
(task_id, status, depends_on, now, now),
|
||
)
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def _set_status(db_path, task_id, status):
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
conn.execute("UPDATE tasks SET status = ? WHERE id = ?",
|
||
(status, task_id))
|
||
conn.commit()
|
||
finally:
|
||
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)
|
||
|
||
def test_depth_zero_min_zero_scales_to_zero(self):
|
||
self.assertEqual(desired_workers(0, 2, 0, 2), 0)
|
||
|
||
def test_depth_divides_by_capacity_and_caps(self):
|
||
self.assertEqual(desired_workers(1, 2, 1, 2), 1)
|
||
self.assertEqual(desired_workers(3, 2, 1, 2), 2)
|
||
self.assertEqual(desired_workers(10, 2, 1, 2), 2)
|
||
self.assertEqual(desired_workers(5, 1, 1, 3), 3)
|
||
|
||
|
||
class QueueDepthTestCase(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||
init_db(self.db_path)
|
||
|
||
def tearDown(self):
|
||
self.tmp.cleanup()
|
||
|
||
def test_only_ready_tasks_count(self):
|
||
# A 就绪；B 依赖未完成的 A，不算；C 已完成，不算；
|
||
# D 在跑，不算；E 失败，不算。
|
||
_insert_task(self.db_path, "A")
|
||
_insert_task(self.db_path, "B", depends_on='["A"]')
|
||
_insert_task(self.db_path, "C", status="completed")
|
||
_insert_task(self.db_path, "D", status="running")
|
||
_insert_task(self.db_path, "E", status="failed")
|
||
self.assertEqual(queue_depth(self.db_path), 1)
|
||
# A 完成后，B 解锁，深度仍是 1（不是 2）。
|
||
_set_status(self.db_path, "A", "completed")
|
||
self.assertEqual(queue_depth(self.db_path), 1)
|
||
|
||
def test_null_depends_on_counts_as_ready(self):
|
||
# 回归：depends_on 为 NULL 时必须算无依赖，不能变成
|
||
# [None] 导致任务永远不算就绪（scaler 致盲）。
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
now = time.time()
|
||
conn.execute(
|
||
"INSERT INTO tasks(id, project_id, prompt, status,"
|
||
" idempotency_key, depends_on, priority, created_by,"
|
||
" created_at, updated_at, last_heartbeat)"
|
||
" VALUES('N', 'p1', 'x', 'queued', NULL, NULL, 0,"
|
||
" 'tester', ?, ?, NULL)",
|
||
(now, now))
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
self.assertEqual(queue_depth(self.db_path), 1)
|
||
|
||
def test_dependency_cycle_counts_zero(self):
|
||
_insert_task(self.db_path, "A", depends_on='["B"]')
|
||
_insert_task(self.db_path, "B", depends_on='["A"]')
|
||
self.assertEqual(queue_depth(self.db_path), 0)
|
||
|
||
|
||
class AutoscalerReconcileTestCase(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||
init_db(self.db_path)
|
||
# 假 worker：睡 120 秒的 python 进程，scaler 用 SIGTERM
|
||
# 停它（默认行为即退出，不用等 drain 超时）。
|
||
self.fake_cmd = [sys.executable, "-c",
|
||
"import time; time.sleep(120)"]
|
||
|
||
def tearDown(self):
|
||
# 测试结束时残留的假 worker 必须清掉，否则泄漏进程。
|
||
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("max_workers", 2)
|
||
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 test_scales_up_on_depth_and_down_after_cooldown(self):
|
||
scaler = self._scaler()
|
||
depth, desired, actual = scaler.reconcile()
|
||
self.assertEqual((depth, desired, actual), (0, 1, 1))
|
||
for tid in ("A", "B", "C", "D"):
|
||
_insert_task(self.db_path, tid)
|
||
depth, desired, actual = scaler.reconcile()
|
||
self.assertEqual((depth, desired, actual), (4, 2, 2))
|
||
# 任务全完成：scale_down_after=1，下一轮直接缩回 1。
|
||
for tid in ("A", "B", "C", "D"):
|
||
_set_status(self.db_path, tid, "completed")
|
||
depth, desired, actual = scaler.reconcile()
|
||
self.assertEqual((depth, desired, actual), (0, 1, 1))
|
||
|
||
def test_scale_down_waits_for_cooldown_streak(self):
|
||
scaler = self._scaler(scale_down_after=3)
|
||
for tid in ("A", "B", "C", "D"):
|
||
_insert_task(self.db_path, tid)
|
||
scaler.reconcile()
|
||
self.assertEqual(len(scaler.workers), 2)
|
||
for tid in ("A", "B", "C", "D"):
|
||
_set_status(self.db_path, tid, "completed")
|
||
scaler.reconcile()
|
||
self.assertEqual(len(scaler.workers), 2)
|
||
scaler.reconcile()
|
||
self.assertEqual(len(scaler.workers), 2)
|
||
scaler.reconcile()
|
||
self.assertEqual(len(scaler.workers), 1)
|
||
|
||
def test_worker_command_forces_loop_id_capacity(self):
|
||
scaler = Autoscaler(self.db_path, ["--interval", "10"])
|
||
cmd = scaler._worker_command("autoscale-1")
|
||
self.assertIn("--loop", cmd)
|
||
self.assertIn("autoscale-1", cmd)
|
||
self.assertIn("--worker-capacity", cmd)
|
||
self.assertEqual(cmd[cmd.index("--worker-capacity") + 1], "2")
|
||
self.assertIn("--db", cmd)
|
||
self.assertEqual(cmd[cmd.index("--db") + 1], self.db_path)
|
||
|
||
def test_main_parses_args_and_starts_loop(self):
|
||
# 回归：main() 里 setup_logging 曾错传 logger 名当 level，
|
||
# 直接 ValueError，scaler 在主机上 crash-loop、一个 worker
|
||
# 都没拉起来。本测试 mock 掉 run_loop 与 signal，只验证
|
||
# main() 能走到启动循环（即参数解析与日志初始化没炸）。
|
||
from unittest import mock
|
||
|
||
import cloud.autoscaler as autoscaler_mod
|
||
with mock.patch.object(
|
||
autoscaler_mod.Autoscaler, "run_loop") as mock_loop, \
|
||
mock.patch("signal.signal"):
|
||
autoscaler_mod.main([
|
||
"--db", self.db_path, "--interval", "30",
|
||
"--", "--interval", "10"])
|
||
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()
|