commit 2fcb2c87fef05cd4c1646398e1b5f95c0c107f5e equwal <truex@equwal.com> 2026-09-22 16:45:22 -0700 Add the burst controller tools/burst_controller.py runs on the web server once a minute. It reads the paid queue and the burst workers in Redis, and decides the set of worker machines: it adds machines when paid jobs wait, up to a cap; it sends the rq shutdown command to a worker that is idle for 10 minutes while nothing waits, so the worker stops after its job; and it destroys the machine of a worker that stopped, or that did not start in 20 minutes. After a deploy, the workers of the old commit stop and new ones start. The controller writes the set to a Terraform variable file and runs terraform apply only when the set changes. It puts the app, at the commit of the web server, in the bucket for new workers, because the repository is private. Property tests check that a worker known to rq is never destroyed, that the cap holds, and that an idle worker stops only when no job waits.
tests/test_burst_controller.py | 251 +++++++++++++++++++++++++++++++++ tools/burst_controller.py | 313 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 564 insertions(+)
diff --git a/tests/test_burst_controller.py b/tests/test_burst_controller.py new file mode 100644 index 0000000..4613764 --- /dev/null +++ b/tests/test_burst_controller.py @@ -0,0 +1,251 @@ +"""The burst controller: when it adds, stops and destroys burst workers.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import fakeredis +import pytest +from hypothesis import given, settings +from hypothesis import strategies as st +from rq import Queue, Worker + +from tools import burst_controller as bc +from tools.burst_controller import Config, Memory, decide + +REF = "new" +NOW = 1_000_000.0 +MIN = 60.0 + + +def worker(created_ago: float, ref: str = REF) -> dict: + return {"created": NOW - created_ago, "ref": ref} + + +# --- invariants over any state ------------------------------------------------- + +names = st.sampled_from([f"{bc.NAME_PREFIX}{i}" for i in range(8)]) + + +@st.composite +def states(draw): + pool = draw(st.sets(names, max_size=8)) + workers = { + n: {"created": NOW - draw(st.floats(0, 60 * MIN)), "ref": draw(st.sampled_from(["old", REF]))} + for n in pool + } + live = {n: draw(st.sampled_from(["busy", "idle"])) for n in pool if draw(st.booleans())} + registered = {n for n in pool if n in live or draw(st.booleans())} + draining = {n for n in registered if draw(st.booleans())} + idle_since = {n: NOW - draw(st.floats(0, 30 * MIN)) for n in live if live[n] == "idle"} + memory = Memory(workers=workers, registered=registered, idle_since=idle_since, draining=draining) + return draw(st.integers(0, 12)), live, memory, Config(max_workers=draw(st.integers(0, 6))) + + +@settings(max_examples=500) +@given(states()) +def test_invariants(state): + waiting, live, memory, config = state + d, after = decide(NOW, waiting, live, memory, config, REF) + + # A worker that rq knows is never destroyed: it must stop by itself first. + assert not d.remove & set(live) + # Every old worker is kept or destroyed, never both; new ones are new. + assert set(after.workers) == (set(memory.workers) - d.remove) | set(d.add) + assert not d.remove & set(after.workers) + assert all(n.startswith(bc.NAME_PREFIX) and n not in memory.workers and n not in live + for n in d.add) + # Never more workers at work than the cap, and never more new ones than waiting jobs. + assert len([n for n in after.workers if n not in after.draining]) <= config.max_workers + assert len(d.add) <= waiting + # Only a worker that rq knows can get the shutdown command. + assert d.drain <= set(live) + + +@settings(max_examples=500) +@given(states()) +def test_a_fresh_busy_worker_under_the_cap_is_never_stopped(state): + waiting, live, memory, config = state + d, _ = decide(NOW, waiting, live, memory, config, REF) + at_work = [n for n in memory.workers if n not in memory.draining + and (n in live or n not in memory.registered)] + if len(at_work) <= config.max_workers: + for n, s in live.items(): + if s == "busy" and n in memory.workers and memory.workers[n]["ref"] == REF: + assert n not in d.drain + + +@settings(max_examples=500) +@given(states()) +def test_an_idle_worker_is_stopped_for_idleness_only_when_no_job_waits(state): + waiting, live, memory, config = state + d, _ = decide(NOW, waiting, live, memory, config, REF) + at_work = [n for n in memory.workers if n not in memory.draining + and (n in live or n not in memory.registered)] + if waiting > 0 and len(at_work) <= config.max_workers: + for n in d.drain: + assert live[n] == "busy" or memory.workers[n]["ref"] != REF + + +# --- scenarios ------------------------------------------------------------------- + +def test_a_backlog_adds_workers_up_to_the_cap(): + d, after = decide(NOW, 5, {}, Memory(), Config(max_workers=3), REF) + assert len(d.add) == 3 and not d.remove and not d.drain + assert all(after.workers[n] == {"created": NOW, "ref": REF} for n in d.add) + + +def test_booting_and_idle_workers_count_as_capacity(): + memory = Memory(workers={"subplz-burst-a": worker(2 * MIN), "subplz-burst-b": worker(15 * MIN)}, + registered={"subplz-burst-b"}) + d, _ = decide(NOW, 2, {"subplz-burst-b": "idle"}, memory, Config(max_workers=5), REF) + assert d.add == () and not d.drain and not d.remove + d, _ = decide(NOW, 3, {"subplz-burst-b": "idle"}, memory, Config(max_workers=5), REF) + assert len(d.add) == 1 + + +def test_an_idle_worker_stops_after_the_grace_time_then_its_droplet_goes(): + name = "subplz-burst-a" + memory = Memory(workers={name: worker(60 * MIN)}, registered={name}, + idle_since={name: NOW - 11 * MIN}) + d, after = decide(NOW, 0, {name: "idle"}, memory, Config(max_workers=3), REF) + assert d.drain == {name} and not d.remove + # It stopped: rq no longer lists it, so the droplet is destroyed. + d, after = decide(NOW + MIN, 0, {}, after, Config(max_workers=3), REF) + assert d.remove == {name} and after.workers == {} + + +def test_an_idle_worker_waits_for_new_jobs_during_the_grace_time(): + name = "subplz-burst-a" + memory = Memory(workers={name: worker(60 * MIN)}, registered={name}, + idle_since={name: NOW - 5 * MIN}) + d, _ = decide(NOW, 0, {name: "idle"}, memory, Config(max_workers=3), REF) + assert not d.drain and not d.remove + + +def test_a_droplet_that_never_starts_its_worker_is_destroyed(): + memory = Memory(workers={"subplz-burst-a": worker(21 * MIN)}) + d, _ = decide(NOW, 1, {}, memory, Config(max_workers=1), REF) + assert d.remove == {"subplz-burst-a"} + assert len(d.add) == 1 # and a new one takes its place + + +def test_after_a_deploy_the_workers_of_the_old_commit_are_replaced(): + memory = Memory( + workers={"subplz-burst-idle": worker(60 * MIN, "old"), + "subplz-burst-busy": worker(60 * MIN, "old"), + "subplz-burst-boot": worker(2 * MIN, "old")}, + registered={"subplz-burst-idle", "subplz-burst-busy"}, + ) + live = {"subplz-burst-idle": "idle", "subplz-burst-busy": "busy"} + d, after = decide(NOW, 2, live, memory, Config(max_workers=3), REF) + assert d.drain == {"subplz-burst-idle", "subplz-burst-busy"} # the busy one finishes first + assert d.remove == {"subplz-burst-boot"} + assert len(d.add) == 2 and all(after.workers[n]["ref"] == REF for n in d.add) + + +def test_a_cap_of_zero_stops_every_worker_and_starts_none(): + memory = Memory( + workers={"subplz-burst-boot": worker(MIN), "subplz-burst-idle": worker(60 * MIN), + "subplz-burst-busy": worker(60 * MIN)}, + registered={"subplz-burst-idle", "subplz-burst-busy"}, + ) + live = {"subplz-burst-idle": "idle", "subplz-burst-busy": "busy"} + d, _ = decide(NOW, 4, live, memory, Config(max_workers=0), REF) + assert d.remove == {"subplz-burst-boot"} + assert d.drain == {"subplz-burst-idle", "subplz-burst-busy"} + assert d.add == () + + +# --- Redis side ----------------------------------------------------------------- + +def test_the_controller_counts_the_paid_queue_of_the_app(): + from backend.queue import PAID_QUEUE + assert bc.QUEUE_NAME == PAID_QUEUE + + +def register(conn, rq_name: str, host: str, state: str) -> None: + w = Worker([Queue(bc.QUEUE_NAME, connection=conn)], connection=conn, name=rq_name) + w.hostname = host + w.register_birth() + w.set_state(state) + + +def test_read_queue_reports_burst_workers_only(): + conn = fakeredis.FakeRedis() + register(conn, "rq-1", "subplz-burst-1-0", "busy") + register(conn, "rq-2", "subplz-burst-1-1", "idle") + register(conn, "rq-web", "racknerd-655a1e9", "busy") # the web server's own worker + Queue(bc.QUEUE_NAME, connection=conn).enqueue("os.getcwd") + waiting, live, rq_names = bc.read_queue(conn) + assert waiting == 1 + assert live == {"subplz-burst-1-0": "busy", "subplz-burst-1-1": "idle"} + assert rq_names == {"subplz-burst-1-0": "rq-1", "subplz-burst-1-1": "rq-2"} + + +def test_send_drain_sends_the_rq_shutdown_command(monkeypatch): + sent = [] + monkeypatch.setattr("rq.command.send_shutdown_command", lambda conn, name: sent.append(name)) + bc.send_drain(None, {"subplz-burst-a": "rq-a", "subplz-burst-b": "rq-b"}, + frozenset({"subplz-burst-a"})) + assert sent == ["rq-a"] + + +# --- one run, with a fake Terraform ------------------------------------------------- + +@pytest.fixture +def world(tmp_path, monkeypatch): + conn = fakeredis.FakeRedis() + monkeypatch.setattr("redis.Redis.from_url", lambda url: conn) + released = [] + monkeypatch.setattr(bc, "ensure_release", lambda app_dir, ref, bucket: released.append(ref)) + calls = tmp_path / "calls.txt" + fake = tmp_path / "fake_terraform.py" + fake.write_text( + "import sys, pathlib\n" + f"pathlib.Path({str(calls)!r}).open('a').write(' '.join(sys.argv[1:]) + '\\n')\n" + f"sys.exit(int(pathlib.Path({str(tmp_path / 'exit_code')!r}).read_text() or 0))\n" + ) + (tmp_path / "exit_code").write_text("0") + + def run(extra=()): + argv = ["--redis-url", "redis://x", "--bucket", "b", "--max-workers", "2", + "--app-dir", str(Path(__file__).resolve().parents[1]), + "--module-dir", str(tmp_path / "workers"), + "--state", str(tmp_path / "state.json"), + "--var-file", str(tmp_path / "vars.json"), + "--terraform", sys.executable, str(fake), *extra] + return bc.main(argv) + + return conn, run, tmp_path, calls, released + + +def test_a_run_writes_the_workers_and_applies(world): + conn, run, tmp, calls, released = world + for _ in range(3): + Queue(bc.QUEUE_NAME, connection=conn).enqueue("os.getcwd") + assert run() == 0 + workers = json.loads((tmp / "vars.json").read_text())["workers"] + assert len(workers) == 2 # the cap + assert "apply -auto-approve" in calls.read_text() + assert len(released) == 1 # the app went to the bucket for the new workers + state = json.loads((tmp / "state.json").read_text()) + assert set(state["workers"]) == set(workers) + + +def test_a_failed_apply_keeps_the_old_state(world): + conn, run, tmp, calls, released = world + Queue(bc.QUEUE_NAME, connection=conn).enqueue("os.getcwd") + (tmp / "exit_code").write_text("1") + assert run() == 1 + assert json.loads((tmp / "vars.json").read_text()) == {"workers": {}} + assert not (tmp / "state.json").exists() + + +def test_a_dry_run_changes_nothing(world): + conn, run, tmp, calls, released = world + Queue(bc.QUEUE_NAME, connection=conn).enqueue("os.getcwd") + assert run(["--dry-run"]) == 0 + assert not calls.exists() and not (tmp / "vars.json").exists() and not released diff --git a/tools/burst_controller.py b/tools/burst_controller.py new file mode 100644 index 0000000..84e453b --- /dev/null +++ b/tools/burst_controller.py @@ -0,0 +1,313 @@ +"""Burst workers: add them when jobs wait, remove them when they are idle. + +The web server runs this once a minute (a systemd timer, see +infra/burst/README.md). One run does this: + +1. Read the Redis queue: how many jobs wait, and which burst workers are busy + or idle. +2. Decide the set of burst workers (`decide`). +3. Send the rq shutdown command to each worker that is no longer needed. rq + lets a worker finish its job before it stops, so a scale-in never kills a + paid job. +4. Write the set to the Terraform variables of infra/burst/workers and run + `terraform apply` when the set changed. Terraform makes the new droplets, + and destroys each droplet whose worker has stopped. + +The host name of a burst worker is its droplet name, and rq records the host +name of each worker. That is how a droplet and its rq worker are matched. +Workers run the app at the git commit of this server (`ref`). After a deploy, +the workers of the old commit stop and new ones take their place. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import os +import subprocess +import sys +import tempfile +import time +from dataclasses import dataclass, field +from pathlib import Path + +log = logging.getLogger("subplz.burst") + +NAME_PREFIX = "subplz-burst-" +QUEUE_NAME = "subplz-jobs-paid" # backend.queue.PAID_QUEUE: burst workers take paid jobs only +RELEASE_PREFIX = "releases/" + + +@dataclass(frozen=True) +class Config: + max_workers: int + # A new droplet installs the speech stack before its worker starts. + boot_timeout_s: float = 20 * 60 + # An idle worker waits this long for a new job before it stops. + idle_grace_s: float = 10 * 60 + + +@dataclass +class Memory: + """What the controller keeps between runs, in a JSON file.""" + + # name -> {"created": unix time, "ref": git commit of the app} + workers: dict[str, dict] = field(default_factory=dict) + # Names that had an rq worker at least once. + registered: set[str] = field(default_factory=set) + # name -> unix time from which the worker was idle without a break. + idle_since: dict[str, float] = field(default_factory=dict) + # Names that received the rq shutdown command. + draining: set[str] = field(default_factory=set) + + def to_json(self) -> dict: + return { + "workers": self.workers, + "registered": sorted(self.registered), + "idle_since": self.idle_since, + "draining": sorted(self.draining), + } + + @classmethod + def from_json(cls, data: dict) -> Memory: + return cls( + workers=dict(data.get("workers", {})), + registered=set(data.get("registered", [])), + idle_since=dict(data.get("idle_since", {})), + draining=set(data.get("draining", [])), + ) + + +@dataclass(frozen=True) +class Decision: + add: tuple[str, ...] # new droplets + drain: frozenset[str] # send these workers the rq shutdown command + remove: frozenset[str] # destroy these droplets + + @property + def changes_droplets(self) -> bool: + return bool(self.add or self.remove) + + +def decide( + now: float, + waiting: int, + live: dict[str, str], + memory: Memory, + config: Config, + ref: str, +) -> tuple[Decision, Memory]: + """The next set of burst workers. Pure: no I/O. + + `waiting` is the number of queued jobs that no worker has taken. `live` + maps the host name of each burst worker that rq knows to "busy" or "idle". + `ref` is the app commit for new workers. + """ + keep: set[str] = set() + remove: set[str] = set() + drain: set[str] = set() + registered = set(memory.registered) + idle_since = dict(memory.idle_since) + draining = set(memory.draining) + booting: list[str] = [] + idle_ready: list[str] = [] + busy: list[str] = [] + + for name, info in memory.workers.items(): + state = live.get(name) + stale = info.get("ref") != ref + if state is None: + if name in registered or stale: + remove.add(name) # it stopped, or it would run an old commit + elif now - info["created"] >= config.boot_timeout_s: + remove.add(name) # it never started + else: + keep.add(name) + booting.append(name) + continue + + keep.add(name) + registered.add(name) + if state == "busy": + idle_since.pop(name, None) + if stale and name not in draining: + drain.add(name) # it finishes the job, then stops + elif name not in draining: + busy.append(name) + continue + + idle_since.setdefault(name, now) + if name in draining: + continue # it stops soon + if stale or (waiting == 0 and now - idle_since[name] >= config.idle_grace_s): + drain.add(name) + else: + idle_ready.append(name) + + # Above the cap (the operator lowered it): stop the cheapest workers first. + active = len(booting) + len(idle_ready) + len(busy) + excess = active - config.max_workers + while excess > 0 and booting: + name = booting.pop() + keep.discard(name) + remove.add(name) + excess -= 1 + for group in (idle_ready, busy): + while excess > 0 and group: + drain.add(group.pop()) + excess -= 1 + + active = len(booting) + len(idle_ready) + len(busy) + need = waiting - len(idle_ready) - len(booting) + room = config.max_workers - active + count = max(0, min(need, room)) + taken = set(memory.workers) | set(live) + new_names: list[str] = [] + i = 0 + while len(new_names) < count: + name = f"{NAME_PREFIX}{int(now)}-{i}" + i += 1 + if name not in taken: + new_names.append(name) + add = tuple(new_names) + + workers = {n: memory.workers[n] for n in keep} + workers.update({n: {"created": now, "ref": ref} for n in add}) + draining = (draining | drain) & keep + new_memory = Memory( + workers=workers, + registered=registered & keep, + idle_since={n: t for n, t in idle_since.items() if n in keep and live.get(n) == "idle"}, + draining=draining, + ) + return Decision(add=add, drain=frozenset(drain), remove=frozenset(remove)), new_memory + + +# --- I/O --------------------------------------------------------------------- + +def read_queue(conn) -> tuple[int, dict[str, str], dict[str, str]]: + """Queued jobs, the state of each burst worker, and its rq worker name.""" + from rq import Queue, Worker + + waiting = Queue(QUEUE_NAME, connection=conn).count + live: dict[str, str] = {} + rq_names: dict[str, str] = {} + for worker in Worker.all(connection=conn): + host = worker.hostname or "" + if not host.startswith(NAME_PREFIX): + continue # the web server's own worker, or one this code does not manage + live[host] = "busy" if worker.get_state() == "busy" else "idle" + rq_names[host] = worker.name + return waiting, live, rq_names + + +def send_drain(conn, rq_names: dict[str, str], names: frozenset[str]) -> None: + from rq.command import send_shutdown_command + + for name in sorted(names): + if name in rq_names: + send_shutdown_command(conn, rq_names[name]) + log.info("told %s to stop after its current job", name) + + +def load_memory(path: Path) -> Memory: + if not path.exists(): + return Memory() + return Memory.from_json(json.loads(path.read_text(encoding="utf-8"))) + + +def save_json(path: Path, data: dict) -> None: + """Write `data` so that a crash never leaves half a file.""" + path.parent.mkdir(parents=True, exist_ok=True) + fd, tmp = tempfile.mkstemp(dir=path.parent, prefix=path.name, suffix=".tmp") + with os.fdopen(fd, "w", encoding="utf-8") as out: + json.dump(data, out, indent=2, sort_keys=True) + os.replace(tmp, path) + + +def terraform_vars(memory: Memory) -> dict: + return {"workers": {n: {"ref": info["ref"]} for n, info in sorted(memory.workers.items())}} + + +def terraform_apply(terraform: list[str], module_dir: Path, var_file: Path) -> bool: + cmd = [*terraform, f"-chdir={module_dir}", "apply", "-auto-approve", + "-input=false", "-no-color", f"-var-file={var_file}"] + result = subprocess.run(cmd, capture_output=True, text=True) + if result.returncode != 0: + log.error("terraform apply failed (%s):\n%s", result.returncode, result.stderr[-4000:]) + return False + return True + + +def ensure_release(app_dir: Path, ref: str, bucket: str) -> None: + """Put the app at `ref` in the bucket, where new workers download it.""" + import boto3 + from botocore.exceptions import ClientError + + s3 = boto3.client("s3") + key = f"{RELEASE_PREFIX}{ref}.tar.gz" + try: + s3.head_object(Bucket=bucket, Key=key) + return + except ClientError as exc: + if exc.response.get("Error", {}).get("Code") not in ("404", "NoSuchKey", "NotFound"): + raise + with tempfile.TemporaryDirectory() as tmp: + archive = Path(tmp) / "app.tar.gz" + subprocess.run(["git", "-C", str(app_dir), "archive", "--format=tar.gz", + "-o", str(archive), ref], check=True) + s3.upload_file(str(archive), bucket, key) + log.info("uploaded %s", key) + + +def run(args: argparse.Namespace) -> int: + from redis import Redis + + config = Config(max_workers=args.max_workers) + ref = subprocess.run(["git", "-C", str(args.app_dir), "rev-parse", "HEAD"], + check=True, capture_output=True, text=True).stdout.strip() + conn = Redis.from_url(args.redis_url) + waiting, live, rq_names = read_queue(conn) + memory = load_memory(args.state) + decision, new_memory = decide(time.time(), waiting, live, memory, config, ref) + log.info("waiting=%d live=%s add=%s drain=%s remove=%s", waiting, live, + list(decision.add), sorted(decision.drain), sorted(decision.remove)) + if args.dry_run: + return 0 + + send_drain(conn, rq_names, decision.drain) + if decision.changes_droplets: + if decision.add: + ensure_release(args.app_dir, ref, args.bucket) + save_json(args.var_file, terraform_vars(new_memory)) + if not terraform_apply(args.terraform, args.module_dir, args.var_file): + # Keep the old memory: the next run decides again from the facts. + save_json(args.var_file, terraform_vars(memory)) + return 1 + save_json(args.state, new_memory.to_json()) + return 0 + + +def main(argv: list[str] | None = None) -> int: + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") + p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + p.add_argument("--max-workers", type=int, + default=int(os.environ.get("SUBPLZ_BURST_MAX_WORKERS", "3"))) + p.add_argument("--redis-url", default=os.environ.get("SUBPLZ_WEB_REDIS_URL", "")) + p.add_argument("--bucket", default=os.environ.get("SUBPLZ_WEB_S3_BUCKET", "")) + p.add_argument("--app-dir", type=Path, default=Path("/opt/subplz-web")) + p.add_argument("--module-dir", type=Path, + default=Path("/opt/subplz-web/infra/burst/workers")) + p.add_argument("--state", type=Path, default=Path("/var/lib/subplz-burst/controller.json")) + p.add_argument("--var-file", type=Path, default=Path("/var/lib/subplz-burst/workers.tfvars.json")) + p.add_argument("--terraform", nargs="+", default=["terraform"]) + p.add_argument("--dry-run", action="store_true", help="decide and log, change nothing") + args = p.parse_args(argv) + if not args.redis_url or not args.bucket: + p.error("set SUBPLZ_WEB_REDIS_URL and SUBPLZ_WEB_S3_BUCKET (or --redis-url, --bucket)") + return run(args) + + +if __name__ == "__main__": + sys.exit(main())