Autoscale workers by queue depth (adaptive threads + process scaler)
Layer 1, in-worker: run_all pool thread count is now adaptive per round (min(capacity, ready_count)) instead of always capacity threads; ready_count counts only dependency-satisfied queued tasks via a shared _ready_tasks() helper (pick_next refactored onto it), so blocked tasks never inflate the pool. Layer 2, process scaler (cloud/autoscaler.py): a systemd-managed daemon that starts/stops cloud.worker subprocesses by ready queue depth. desired = ceil(depth / per_worker_capacity) clamped to [min_workers, max_workers]; scale-up is immediate, scale-down waits scale_down_after consecutive low checks (default 3) to avoid flapping on bursty completions. Scale-down is LIFO (youngest worker first), SIGTERM for graceful drain (worker run_loop drains up to drain_timeout), SIGKILL of the process group only after drain timeout. Depth read failure keeps the current worker count (never scale on unknown state). Worker command forces --loop/--worker-id/--worker-capacity (appended last so argparse last-wins); --db comes from scaler or worker args; workers inherit scaler env (Postgres etc.). State file (depth/desired/actual/workers) written atomically for observability. deploy.sh now restarts codex-autoscaler instead of the static worker unit when the scaler is active, so the two never poll the same queue at once. Tests: 8 autoscaler tests (desired math, ready-only depth, cycle=0, scale up/down, cooldown streak, forced worker flags).
This commit is contained in:
host committed
1 parent
67e1d3eb24
commit
7e147be9ed
5 files changed
+643
-19
No files matched your search
@@ -0,0 +1,165 @@
|
||||
"""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_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)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in new issue
Block a user