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.
200 lines
7.7 KiB
Python
200 lines
7.7 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()
|
||
|
||
|
||
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()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|