commit 29e67302f4d53c7ad7a1d7febf792e5cf8fed1a5 equwal <truex@equwal.com> 2026-09-22 17:08:43 -0700 Start each burst worker on the cloud that is cheapest now The burst controller becomes the package tools/burst, and replaces tools/burst_controller.py. Each time paid jobs wait, it reads the live offers of each cloud that has credentials: the newest AWS spot prices, and the Hetzner prices of the server types in stock. A new worker goes to the offer with the lowest cost per book (the hourly price times the hours of a book, plus any transfer fee). When the apply of one cloud fails, that cloud rests for 30 minutes and the next offer gets the worker. A spend cap limits the workers. The cap is 30% of the net server revenue of the last 30 days (packs for 30 days, active subscriptions, after the Stripe fees), at least $5 and at most $100. The controller counts what each machine cost, and keeps 2.5 hours of each running worker in reserve. At the cap it starts no worker and stops no job. A machine that a cloud bills by the started hour stops only late in its paid hour. Property tests check that new workers go to the cheapest offer that fits, that the spend cap and the worker cap hold, and that a worker known to rq is never destroyed. Three deliberate bugs in the offer choice were caught.
tests/test_burst.py | 403 +++++++++++++++++++++++++++++++++++++++++ tests/test_burst_controller.py | 251 ------------------------- tools/burst/__init__.py | 1 + tools/burst/__main__.py | 303 +++++++++++++++++++++++++++++++ tools/burst/core.py | 252 ++++++++++++++++++++++++++ tools/burst/offers.py | 95 ++++++++++ tools/burst/revenue.py | 55 ++++++ tools/burst_controller.py | 313 -------------------------------- 8 files changed, 1109 insertions(+), 564 deletions(-)
diff --git a/tests/test_burst.py b/tests/test_burst.py new file mode 100644 index 0000000..eb407d2 --- /dev/null +++ b/tests/test_burst.py @@ -0,0 +1,403 @@ +"""The burst controller (tools/burst): which workers it starts, where, and when it stops.""" + +from __future__ import annotations + +import json +import math +import sys +import time +from datetime import datetime, timedelta, timezone +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 backend.db import Account, Purchase, SessionLocal +from tools.burst import __main__ as run_mod +from tools.burst import offers as offers_mod +from tools.burst.core import ( + NAME_PREFIX, Config, Memory, Offer, decide, machine_cost, spend_30d, spend_cap_usd, +) +from tools.burst.revenue import net_pack_cents, net_subscription_cents, revenue_30d_usd + +REF = "new" +NOW = 1_000_000.0 +MIN = 60.0 +AWS = Offer("aws", "c6a.xlarge", "us-west-2a", 4, 0.07) +HETZNER = Offer("hetzner", "cx33", "fsn1", 4, 0.0164, per_started_hour=True) + + +def worker(created_ago: float, ref: str = REF, offer: Offer = AWS) -> dict: + return {"created": NOW - created_ago, "ref": ref, "provider": offer.provider, + "machine": offer.machine, "zone": offer.zone, "usd_per_hour": offer.usd_per_hour, + "per_started_hour": offer.per_started_hour} + + +# --- invariants over any state ------------------------------------------------- + +names = st.sampled_from([f"{NAME_PREFIX}{i}" for i in range(8)]) +offer_st = st.builds( + Offer, provider=st.sampled_from(["aws", "hetzner"]), machine=st.sampled_from(["a", "b"]), + zone=st.sampled_from(["z1", "z2"]), vcpus=st.sampled_from([2, 4, 8]), + usd_per_hour=st.floats(0.005, 0.2), per_started_hour=st.booleans(), + capacity=st.integers(0, 3), +) + + +@st.composite +def states(draw): + pool = draw(st.sets(names, max_size=8)) + workers = {n: worker(draw(st.floats(0, 60 * MIN)), draw(st.sampled_from(["old", REF])), + draw(st.sampled_from([AWS, HETZNER]))) 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, cooldown=draw(st.sampled_from([{}, {"aws": NOW + 60}]))) + offers = draw(st.lists(offer_st, max_size=4, unique=True)) + cap = draw(st.sampled_from([math.inf, 0.0, 0.5, 5.0])) + return (draw(st.integers(0, 12)), live, memory, Config(max_workers=draw(st.integers(0, 6))), + offers, cap) + + +@settings(max_examples=500) +@given(states()) +def test_invariants(state): + waiting, live, memory, config, offers, cap = state + d, after = decide(NOW, waiting, live, memory, config, REF, offers, cap) + added = [n for n, _ in d.add] + + # A worker that rq knows is never destroyed: it must stop by itself first. + assert not d.remove & set(live) + assert set(after.workers) == (set(memory.workers) - d.remove) | set(added) + assert all(n.startswith(NAME_PREFIX) and n not in memory.workers and n not in live for n in added) + # The cap on workers holds, and no more new workers than waiting jobs. + assert len([n for n in after.workers if n not in after.draining]) <= config.max_workers + assert len(added) <= waiting + # Only a worker that rq knows gets the shutdown command. + assert d.drain <= set(live) + # New workers come from offers that have room, from a cloud that is not resting. + for offer in {o for _, o in d.add}: + assert offer in offers and memory.cooldown.get(offer.provider, 0) <= NOW + assert sum(1 for _, o in d.add if o == offer) <= offer.capacity + # Each machine that ends goes to the ledger, once. + assert sorted(e["name"] for e in after.ledger if e["end"] == NOW) == sorted(d.remove) + + +@settings(max_examples=500) +@given(states()) +def test_a_new_worker_goes_to_the_cheapest_offer_that_fits(state): + waiting, live, memory, config, offers, cap = state + d, _ = decide(NOW, waiting, live, memory, config, REF, offers, cap) + if not d.add: + return + first = d.add[0][1] + fitting = [o for o in offers if o.capacity > 0 and memory.cooldown.get(o.provider, 0) <= NOW] + cheaper = [o for o in fitting if o.usd_per_book(config) < first.usd_per_book(config)] + # A cheaper offer is left out only when the spend cap has no room for its + # reserve, which is its hourly price: so it must cost more by the hour. + # (A counterexample of 2026-09-22: 8 vCPU at $0.30/h is cheaper by the book + # than 2 vCPU at $0.125/h, and a small cap takes only the second.) + for o in cheaper: + assert o.usd_per_hour > first.usd_per_hour + if cap == math.inf: + assert not cheaper + + +@settings(max_examples=300) +@given(states()) +def test_the_spend_cap_holds(state): + waiting, live, memory, config, offers, cap = state + d, after = decide(NOW, waiting, live, memory, config, REF, offers, cap) + if d.add: + reserve = sum(config.reserve_hours * w["usd_per_hour"] for n, w in after.workers.items() + if n not in after.draining) + assert spend_30d(NOW, after) + reserve <= cap + 1e-9 + + +# --- scenarios ------------------------------------------------------------------- + +def test_a_backlog_adds_workers_on_the_cheapest_cloud_up_to_the_cap(): + d, after = decide(NOW, 5, {}, Memory(), Config(max_workers=3), REF, [AWS, HETZNER]) + assert [o for _, o in d.add] == [HETZNER, HETZNER, HETZNER] + assert all(after.workers[n]["provider"] == "hetzner" for n, _ in d.add) + + +def test_when_the_cheapest_cloud_is_full_the_next_one_gets_the_worker(): + small = Offer("hetzner", "cx33", "fsn1", 4, 0.0164, per_started_hour=True, capacity=1) + d, _ = decide(NOW, 3, {}, Memory(), Config(max_workers=3), REF, [AWS, small]) + assert [o.provider for _, o in d.add] == ["hetzner", "aws", "aws"] + + +def test_a_resting_cloud_gets_no_worker(): + memory = Memory(cooldown={"hetzner": NOW + 10 * MIN}) + d, _ = decide(NOW, 1, {}, memory, Config(max_workers=3), REF, [AWS, HETZNER]) + assert [o.provider for _, o in d.add] == ["aws"] + + +def test_a_small_cap_takes_the_offer_with_the_lower_hourly_price(): + # Found by the property test: the big machine is cheaper by the book, but + # the cap has room only for the reserve of the small one. + small = Offer("aws", "small", "z1", 2, 0.125) # $0.52 a book, $0.31 reserve + big = Offer("hetzner", "big", "z1", 8, 0.35) # $0.40 a book, $0.88 reserve + d, _ = decide(NOW, 1, {}, Memory(), Config(max_workers=1), REF, [small, big], cap_usd=0.5) + assert [o for _, o in d.add] == [small] + d, _ = decide(NOW, 1, {}, Memory(), Config(max_workers=1), REF, [small, big]) + assert [o for _, o in d.add] == [big] + + +def test_egress_can_make_the_cheaper_machine_the_dearer_book(): + aws_with_egress = Offer("aws", "c6a.xlarge", "us-west-2a", 4, 0.03, extra_usd_per_book=0.2) + d, _ = decide(NOW, 1, {}, Memory(), Config(max_workers=1), REF, [aws_with_egress, HETZNER]) + assert d.add[0][1] == HETZNER + + +def test_the_spend_cap_stops_new_workers_but_not_running_ones(): + memory = Memory(workers={"subplz-burst-a": worker(60 * MIN)}, registered={"subplz-burst-a"}, + ledger=[{"name": "old", "provider": "aws", "start": NOW - 50 * 3600, + "end": NOW - 10 * 3600, "usd_per_hour": 0.07, "per_started_hour": False}]) + # 40 h at $0.07 = $2.80 spent, plus 1 h of the running worker, and its reserve. + d, after = decide(NOW, 3, {"subplz-burst-a": "busy"}, memory, Config(max_workers=3), REF, + [AWS], cap_usd=3.2) + assert d.add == () and not d.drain and "subplz-burst-a" in after.workers + d, _ = decide(NOW, 3, {"subplz-burst-a": "busy"}, memory, Config(max_workers=3), REF, + [AWS], cap_usd=10.0) + assert len(d.add) == 2 + + +def test_an_hourly_machine_stops_late_in_its_paid_hour(): + name = "subplz-burst-h" + for age, stops in ((70 * MIN, False), (115 * MIN, True)): + memory = Memory(workers={name: worker(age, offer=HETZNER)}, registered={name}, + idle_since={name: NOW - 20 * MIN}) + d, _ = decide(NOW, 0, {name: "idle"}, memory, Config(max_workers=3), REF) + assert (name in d.drain) is stops, age + + +def test_an_idle_worker_stops_after_the_grace_time_then_its_machine_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 + d, after = decide(NOW + MIN, 0, {}, after, Config(max_workers=3), REF) + assert d.remove == {name} and after.workers == {} + assert after.ledger[-1]["name"] == name and after.ledger[-1]["end"] == NOW + MIN + + +def test_a_machine_that_never_starts_its_worker_is_destroyed_and_replaced(): + memory = Memory(workers={"subplz-burst-a": worker(21 * MIN)}) + d, _ = decide(NOW, 1, {}, memory, Config(max_workers=1), REF, [AWS]) + assert d.remove == {"subplz-burst-a"} and len(d.add) == 1 + + +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, [AWS]) + assert d.drain == {"subplz-burst-idle", "subplz-burst-busy"} + 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, [AWS]) + assert d.remove == {"subplz-burst-boot"} + assert d.drain == {"subplz-burst-idle", "subplz-burst-busy"} and d.add == () + + +def test_machine_cost_and_spend(): + assert machine_cost(0, 5400, 0.1, False) == pytest.approx(0.15) + assert machine_cost(0, 5400, 0.1, True) == pytest.approx(0.2) # 2 started hours + memory = Memory(workers={"w": worker(3600)}, + ledger=[{"name": "x", "provider": "aws", "start": 0, "end": NOW - 31 * 86400, + "usd_per_hour": 9.0, "per_started_hour": False}]) + assert spend_30d(NOW, memory) == pytest.approx(0.07) # the old machine is out of the window + + +def test_spend_cap(): + assert spend_cap_usd(0.0) == 5.0 # the floor serves the first buyer + assert spend_cap_usd(115.85) == pytest.approx(34.755) + assert spend_cap_usd(1000.0) == 100.0 + + +# --- live offers -------------------------------------------------------------------- + +def test_aws_offers_take_the_newest_price_and_skip_a_price_above_the_maximum(): + items = [ + {"InstanceType": "c6a.xlarge", "AvailabilityZone": "us-west-2a", "SpotPrice": "0.090", + "Timestamp": "2026-09-22T10:00:00Z"}, + {"InstanceType": "c6a.xlarge", "AvailabilityZone": "us-west-2a", "SpotPrice": "0.050", + "Timestamp": "2026-09-22T11:00:00Z"}, + {"InstanceType": "c7a.xlarge", "AvailabilityZone": "us-west-2b", "SpotPrice": "0.080", + "Timestamp": "2026-09-22T11:00:00Z"}, + {"InstanceType": "x9.huge", "AvailabilityZone": "us-west-2b", "SpotPrice": "0.001", + "Timestamp": "2026-09-22T11:00:00Z"}, + ] + found = offers_mod.aws_offers(items, max_price=0.0625, extra_usd_per_hour=0.0072, + extra_usd_per_book=0.0, capacity=1) + assert found == [Offer("aws", "c6a.xlarge", "us-west-2a", 4, pytest.approx(0.0572), False, 0.0, 1)] + + +def test_hetzner_offers_come_from_stock_at_the_net_price(): + api = { + "/server_types?per_page=50": {"server_types": [ + {"id": 1, "name": "cx33", "cores": 4, "architecture": "x86", "deprecation": None, + "prices": [{"location": "fsn1", "price_hourly": {"net": "0.0136", "gross": "0.0162"}}, + {"location": "nbg1", "price_hourly": {"net": "0.0136", "gross": "0.0162"}}]}, + {"id": 2, "name": "cax21", "cores": 4, "architecture": "arm", "deprecation": None, + "prices": [{"location": "fsn1", "price_hourly": {"net": "0.0100", "gross": "0.0119"}}]}, + {"id": 3, "name": "cpx31", "cores": 4, "architecture": "x86", "deprecation": None, + "prices": [{"location": "hil", "price_hourly": {"net": "0.0997", "gross": "0.0997"}}]}, + ]}, + "/datacenters?per_page=50": {"datacenters": [ + {"location": {"name": "fsn1"}, "server_types": {"available": [1, 2]}}, + {"location": {"name": "nbg1"}, "server_types": {"available": []}}, # sold out + {"location": {"name": "hil"}, "server_types": {"available": [3]}}, + ]}, + } + found = offers_mod.hetzner_offers(api.__getitem__, machines={"cx33", "cax21", "cpx31"}, + usd_per_eur=1.15, ipv4_eur_per_hour=0.0007, capacity=2) + assert [(o.machine, o.zone) for o in found] == [("cpx31", "hil"), ("cx33", "fsn1")] + assert found[1].usd_per_hour == pytest.approx((0.0136 + 0.0007) * 1.15) + assert all(o.per_started_hour and o.capacity == 2 for o in found) + + +# --- revenue ------------------------------------------------------------------------- + +def test_net_after_stripe(): + assert [net_pack_cents(a) for a in (499, 3999, 17499)] == [454, 3853, 16961] + assert [net_subscription_cents(a) for a in (499, 999, 5000)] == [451, 933, 4790] + + +def test_revenue_counts_recent_packs_and_active_subscriptions(): + now = datetime(2026, 9, 22, tzinfo=timezone.utc) + tag = str(time.time_ns()) + with SessionLocal() as s: + base = revenue_30d_usd(s, now) + s.add(Account(id=f"acct_a{tag}", device_token=f"a{tag}")) + s.add(Account(id=f"acct_b{tag}", device_token=f"b{tag}", subscription_plan_id="month30", + subscription_status="active", subscription_period_end=now + timedelta(days=9))) + s.add(Account(id=f"acct_c{tag}", device_token=f"c{tag}", subscription_plan_id="unlimited", + subscription_status="canceled", subscription_period_end=now + timedelta(days=9))) + s.add(Purchase(account_id=f"acct_a{tag}", plan_id="pack100", credits=100, amount_cents=3999, + stripe_session_id=f"cs_new{tag}", created_at=now - timedelta(days=3))) + s.add(Purchase(account_id=f"acct_a{tag}", plan_id="pack10", credits=10, amount_cents=499, + stripe_session_id=f"cs_old{tag}", created_at=now - timedelta(days=40))) + s.commit() + assert revenue_30d_usd(s, now) - base == pytest.approx((3853 + 933) / 100) + + +# --- one run, with fake clouds and a fake Terraform ------------------------------------ + +def register(conn, rq_name: str, host: str, state: str) -> None: + w = Worker([Queue(run_mod.QUEUE_NAME, connection=conn)], connection=conn, name=rq_name) + w.hostname = host + w.register_birth() + w.set_state(state) + + +def test_the_controller_counts_the_paid_queue_of_the_app(): + from backend.queue import PAID_QUEUE + assert run_mod.QUEUE_NAME == PAID_QUEUE + + +def test_read_queue_reports_burst_workers_only(): + conn = fakeredis.FakeRedis() + register(conn, "rq-1", "subplz-burst-1-0", "busy") + register(conn, "rq-web", "racknerd-655a1e9", "busy") # the web server's own worker + Queue(run_mod.QUEUE_NAME, connection=conn).enqueue("os.getcwd") + waiting, live, rq_names = run_mod.read_queue(conn) + assert (waiting, live, rq_names) == (1, {"subplz-burst-1-0": "busy"}, {"subplz-burst-1-0": "rq-1"}) + + +def test_terraform_gets_the_burst_keys_not_the_storage_keys(monkeypatch): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "r2-key") + monkeypatch.setenv("AWS_ENDPOINT_URL_S3", "https://r2.example") + monkeypatch.setenv("SUBPLZ_WEB_DATABASE_URL", "postgresql://secret") + monkeypatch.setenv("SUBPLZ_BURST_AWS_ACCESS_KEY_ID", "ec2-key") + monkeypatch.setenv("HCLOUD_TOKEN", "hz") + aws = run_mod.terraform_env("aws") + assert aws["AWS_ACCESS_KEY_ID"] == "ec2-key" and "AWS_ENDPOINT_URL_S3" not in aws + assert "SUBPLZ_WEB_DATABASE_URL" not in aws + hetzner = run_mod.terraform_env("hetzner") + assert hetzner["HCLOUD_TOKEN"] == "hz" and "AWS_ACCESS_KEY_ID" not in hetzner + + +@pytest.fixture +def world(tmp_path, monkeypatch): + conn = fakeredis.FakeRedis() + monkeypatch.setattr("redis.Redis.from_url", lambda url: conn) + released = [] + monkeypatch.setattr(run_mod, "ensure_release", lambda app_dir, ref, bucket: released.append(ref)) + monkeypatch.setattr(run_mod, "revenue_30d", lambda: 0.0) # the floor: $5 + offers = [AWS, HETZNER] + monkeypatch.setattr(run_mod, "gather_offers", lambda memory: list(offers)) + 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"bad = pathlib.Path({str(tmp_path / 'fail')!r}).read_text()\n" + "sys.exit(1 if bad and bad in ' '.join(sys.argv) else 0)\n" + ) + (tmp_path / "fail").write_text("") + + def run(extra=()): + argv = ["--redis-url", "redis://x", "--bucket", "b", "--max-workers", "2", + "--app-dir", str(Path(__file__).resolve().parents[1]), + "--modules-dir", str(tmp_path / "modules"), "--state-dir", str(tmp_path / "state"), + "--terraform", sys.executable, str(fake), *extra] + return run_mod.main(argv) + + return conn, run, tmp_path, calls, released + + +def enqueue(conn, n: int) -> None: + for _ in range(n): + Queue(run_mod.QUEUE_NAME, connection=conn).enqueue("os.getcwd") + + +def test_a_run_starts_workers_on_the_cheapest_cloud(world): + conn, run, tmp, calls, released = world + enqueue(conn, 3) + assert run() == 0 + hetzner = json.loads((tmp / "state" / "workers-hetzner.tfvars.json").read_text())["workers"] + assert len(hetzner) == 2 and all(w["machine"] == "cx33" for w in hetzner.values()) + assert "modules" in calls.read_text() and "hetzner" in calls.read_text() + assert not (tmp / "state" / "workers-aws.tfvars.json").exists() # nothing changed there + assert len(released) == 1 + + +def test_a_failed_apply_rests_the_cloud_and_the_next_run_uses_another(world): + conn, run, tmp, calls, released = world + enqueue(conn, 1) + (tmp / "fail").write_text("hetzner") + assert run() == 1 + state = json.loads((tmp / "state" / "controller.json").read_text()) + assert state["workers"] == {} and state["cooldown"]["hetzner"] > time.time() + assert json.loads((tmp / "state" / "workers-hetzner.tfvars.json").read_text()) == {"workers": {}} + assert run() == 0 # Hetzner rests: AWS gets the worker + aws = json.loads((tmp / "state" / "workers-aws.tfvars.json").read_text())["workers"] + assert [w["machine"] for w in aws.values()] == ["c6a.xlarge"] + + +def test_a_dry_run_changes_nothing(world): + conn, run, tmp, calls, released = world + enqueue(conn, 1) + assert run(["--dry-run"]) == 0 + assert not calls.exists() and not (tmp / "state").exists() and not released diff --git a/tests/test_burst_controller.py b/tests/test_burst_controller.py deleted file mode 100644 index 4613764..0000000 --- a/tests/test_burst_controller.py +++ /dev/null @@ -1,251 +0,0 @@ -"""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/__init__.py b/tools/burst/__init__.py new file mode 100644 index 0000000..904de78 --- /dev/null +++ b/tools/burst/__init__.py @@ -0,0 +1 @@ +"""The burst controller: see __main__.py and infra/burst/README.md.""" diff --git a/tools/burst/__main__.py b/tools/burst/__main__.py new file mode 100644 index 0000000..31757a1 --- /dev/null +++ b/tools/burst/__main__.py @@ -0,0 +1,303 @@ +"""Burst controller: add workers on the cheapest cloud when paid jobs wait. + +The web server runs `python -m tools.burst` once a minute (a systemd timer, +see infra/burst/README.md). One run does this: + +1. Read the paid queue in Redis, and the state of each burst worker. +2. Read the live offers of each cloud that has credentials (AWS spot prices, + Hetzner prices and stock), and the net revenue of the last 30 days. +3. Decide (core.decide): new workers go to the offer with the lowest cost per + book, inside the spend cap. Idle workers get the rq shutdown command: rq + lets a worker finish its job first, so a scale-in never kills a paid job. +4. For each cloud whose set of workers changed: write its Terraform variable + file and run `terraform apply` on infra/burst/workers/<cloud>. If the apply + fails, that cloud rests for 30 minutes, and the next run uses the next + cheapest offer. + +The host name of a worker is its machine name, and rq records the host name of +each worker. That is how a machine and its rq worker are matched. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import os +import subprocess +import sys +import tempfile +import time +from datetime import datetime, timedelta, timezone +from pathlib import Path + +from .core import NAME_PREFIX, Config, Memory, Offer, decide, spend_30d, spend_cap_usd + +log = logging.getLogger("subplz.burst") + +QUEUE_NAME = "subplz-jobs-paid" # backend.queue.PAID_QUEUE: burst workers take paid jobs only +RELEASE_PREFIX = "releases/" +PROVIDERS = ("aws", "hetzner") +COOLDOWN_S = 30 * 60 + + +# --- Redis --------------------------------------------------------------------- + +def read_queue(conn) -> tuple[int, dict[str, str], dict[str, str]]: + """Waiting paid 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 that 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) + + +# --- files --------------------------------------------------------------------- + +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 provider_vars(memory: Memory, provider: str) -> dict: + """The Terraform variables of one cloud: its workers.""" + return {"workers": { + name: {"ref": w["ref"], "machine": w["machine"], "zone": w["zone"]} + for name, w in sorted(memory.workers.items()) if w["provider"] == provider + }} + + +# --- clouds -------------------------------------------------------------------- + +def env_float(name: str, default: float) -> float: + return float(os.environ.get(name, default)) + + +def env_list(name: str, default: str) -> list[str]: + return [v.strip() for v in os.environ.get(name, default).split(",") if v.strip()] + + +def vcpus_in_use(memory: Memory, provider: str) -> int: + from .offers import AWS_VCPUS + + return sum(AWS_VCPUS.get(w["machine"], 4) for w in memory.workers.values() + if w["provider"] == provider) + + +def gather_offers(memory: Memory) -> list[Offer]: + """The live offers of each cloud that has credentials. A cloud whose API + fails gives no offers this run; the others still do.""" + from . import offers + + found: list[Offer] = [] + if os.environ.get("SUBPLZ_BURST_AWS_ACCESS_KEY_ID"): + try: + import boto3 + + ec2 = boto3.client( + "ec2", region_name=os.environ.get("SUBPLZ_BURST_AWS_REGION", "us-west-2"), + aws_access_key_id=os.environ["SUBPLZ_BURST_AWS_ACCESS_KEY_ID"], + aws_secret_access_key=os.environ["SUBPLZ_BURST_AWS_SECRET_ACCESS_KEY"], + ) + machines = env_list("SUBPLZ_BURST_AWS_TYPES", "c6a.xlarge,c7a.xlarge,m6a.xlarge") + since = datetime.now(timezone.utc) - timedelta(hours=1) + limit = int(os.environ.get("SUBPLZ_BURST_AWS_VCPU_LIMIT", "5")) + capacity = max(0, limit - vcpus_in_use(memory, "aws")) // 4 + found += offers.aws_offers( + offers.fetch_aws_spot_prices(ec2, machines, since), + max_price=env_float("SUBPLZ_BURST_AWS_MAX_PRICE", 0.0625), + # A public IPv4 address ($0.005/h) and a 20 GB disk ($0.0022/h). + extra_usd_per_hour=env_float("SUBPLZ_BURST_AWS_EXTRA_USD_PER_HOUR", 0.0072), + # AWS gives 100 GB of transfer out each month for free. Above it, + # each result that a worker sends to R2 costs $0.09/GB. + extra_usd_per_book=env_float("SUBPLZ_BURST_AWS_EGRESS_USD_PER_BOOK", 0.0), + capacity=capacity, + ) + except Exception: # noqa: BLE001 - one cloud failing must not stop the others + log.exception("AWS offers failed") + if os.environ.get("HCLOUD_TOKEN"): + try: + import requests + + def get_json(path: str) -> dict: + r = requests.get(f"https://api.hetzner.cloud/v1{path}", timeout=20, + headers={"Authorization": f"Bearer {os.environ['HCLOUD_TOKEN']}"}) + r.raise_for_status() + return r.json() + + limit = int(os.environ.get("SUBPLZ_BURST_HETZNER_SERVER_LIMIT", "5")) + in_use = sum(1 for w in memory.workers.values() if w["provider"] == "hetzner") + found += offers.hetzner_offers( + get_json, + machines=set(env_list("SUBPLZ_BURST_HETZNER_TYPES", "cx33,cx43,cpx31,cpx32,ccx13")), + usd_per_eur=env_float("SUBPLZ_BURST_USD_PER_EUR", 1.15), + # A primary IPv4 address: 0.50 EUR a month. + ipv4_eur_per_hour=env_float("SUBPLZ_BURST_HETZNER_IPV4_EUR_PER_HOUR", 0.0007), + capacity=max(0, limit - in_use), + ) + except Exception: # noqa: BLE001 + log.exception("Hetzner offers failed") + return found + + +def revenue_30d() -> float: + from backend.db import SessionLocal + + from .revenue import revenue_30d_usd + + with SessionLocal() as session: + return revenue_30d_usd(session, datetime.now(timezone.utc)) + + +def terraform_env(provider: str) -> dict[str, str]: + """The environment of one `terraform apply`. The storage keys of the app are + also AWS_* variables (R2 speaks the S3 protocol), so they must not reach the + AWS provider: it gets the burst keys instead.""" + env = {k: v for k, v in os.environ.items() + if not (k.startswith("AWS_") or k.startswith("SUBPLZ_WEB_"))} + if provider == "aws": + env["AWS_ACCESS_KEY_ID"] = os.environ.get("SUBPLZ_BURST_AWS_ACCESS_KEY_ID", "") + env["AWS_SECRET_ACCESS_KEY"] = os.environ.get("SUBPLZ_BURST_AWS_SECRET_ACCESS_KEY", "") + env["AWS_REGION"] = os.environ.get("SUBPLZ_BURST_AWS_REGION", "us-west-2") + return env + + +def terraform_apply(terraform: list[str], module_dir: Path, var_file: Path, env: dict) -> 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, env=env) + if result.returncode != 0: + log.error("terraform apply failed in %s (%s):\n%s", module_dir, 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: the + repository is private, and a worker must run the web server's commit.""" + 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 rollback(provider: str, old: Memory, new: Memory, now: float) -> Memory: + """Undo one cloud's part of a decision after its apply failed, and rest it.""" + workers = {n: w for n, w in new.workers.items() if w["provider"] != provider} + workers.update({n: w for n, w in old.workers.items() if w["provider"] == provider}) + new.workers = workers + new.ledger = [e for e in new.ledger + if not (e.get("provider") == provider and e["end"] == now)] + new.draining &= set(workers) + new.cooldown = {**new.cooldown, provider: now + COOLDOWN_S} + return new + + +# --- one run --------------------------------------------------------------------- + +def run(args: argparse.Namespace) -> int: + from redis import Redis + + now = time.time() + 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) + state_file = args.state_dir / "controller.json" + memory = load_memory(state_file) + offers = gather_offers(memory) if waiting else [] + cap = spend_cap_usd(revenue_30d(), k=args.cap_share, floor=args.cap_floor, ceiling=args.cap_ceiling) + decision, new_memory = decide(now, waiting, live, memory, config, ref, offers, cap) + log.info("waiting=%d live=%s spend30=$%.2f cap=$%.2f add=%s drain=%s remove=%s", + waiting, live, spend_30d(now, memory), cap, + [(n, o.provider, o.machine, o.zone, round(o.usd_per_book(config), 3)) + for n, o in decision.add], + sorted(decision.drain), sorted(decision.remove)) + if waiting and not decision.add and not live: + log.warning("paid jobs wait, and no burst worker can start (spend cap, or no offer)") + if args.dry_run: + return 0 + + send_drain(conn, rq_names, decision.drain) + if decision.add: + ensure_release(args.app_dir, ref, args.bucket) + status = 0 + for provider in PROVIDERS: + if provider_vars(memory, provider) == provider_vars(new_memory, provider): + continue + var_file = args.state_dir / f"workers-{provider}.tfvars.json" + save_json(var_file, provider_vars(new_memory, provider)) + if not terraform_apply(args.terraform, args.modules_dir / provider, var_file, + terraform_env(provider)): + new_memory = rollback(provider, memory, new_memory, now) + save_json(var_file, provider_vars(new_memory, provider)) + status = 1 + save_json(state_file, new_memory.to_json()) + return status + + +def main(argv: list[str] | None = None) -> int: + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") + p = argparse.ArgumentParser(prog="python -m tools.burst", description=__doc__.splitlines()[0]) + p.add_argument("--max-workers", type=int, + default=int(os.environ.get("SUBPLZ_BURST_MAX_WORKERS", "1"))) + p.add_argument("--cap-share", type=float, default=env_float("SUBPLZ_BURST_CAP_SHARE", 0.30), + help="share of the net revenue of 30 days that burst workers may spend") + p.add_argument("--cap-floor", type=float, default=env_float("SUBPLZ_BURST_CAP_FLOOR_USD", 5.0)) + p.add_argument("--cap-ceiling", type=float, default=env_float("SUBPLZ_BURST_CAP_CEILING_USD", 100.0)) + 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("--modules-dir", type=Path, default=Path("/opt/subplz-web/infra/burst/workers")) + p.add_argument("--state-dir", type=Path, default=Path("/var/lib/subplz-burst")) + 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()) diff --git a/tools/burst/core.py b/tools/burst/core.py new file mode 100644 index 0000000..cdd7949 --- /dev/null +++ b/tools/burst/core.py @@ -0,0 +1,252 @@ +"""The decisions of the burst controller. Pure: no I/O, no clock, no network. + +`decide` takes what rq reports about the queue and the workers, the live +offers of the clouds, and the spend cap. It returns which workers to start (on +the cheapest offer), which to stop, and which machines to destroy. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field + +NAME_PREFIX = "subplz-burst-" +HOUR = 3600.0 +WINDOW = 30 * 86400.0 # the spend cap and the revenue both look back 30 days + + +@dataclass(frozen=True) +class Config: + max_workers: int + # A new machine 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 + # A 10-hour book on 4 vCPU: the web server does it at about 5x real time. + book_hours_4vcpu: float = 2.0 + boot_hours: float = 0.15 + # The cap keeps this much work of each worker in reserve. + reserve_hours: float = 2.5 + # A cloud that bills each started hour gets its worker stopped only late in + # the hour that is paid already. + hour_drain_after_s: float = 50 * 60 + + +@dataclass(frozen=True) +class Offer: + """A machine that one cloud can start now, at its live price.""" + + provider: str # "aws" or "hetzner": the Terraform module infra/burst/workers/<provider> + machine: str # instance type or server type + zone: str # AWS availability zone or Hetzner location + vcpus: int + usd_per_hour: float # all that the machine costs per hour: price, IPv4 address, disk + per_started_hour: bool = False # Hetzner bills each started hour in full + extra_usd_per_book: float = 0.0 # for example the transfer fee of the results + capacity: int = 1000 # machines of this offer that the account can start now + + def usd_per_book(self, config: Config) -> float: + hours = config.book_hours_4vcpu * 4 / self.vcpus + config.boot_hours + return self.usd_per_hour * hours + self.extra_usd_per_book + + +@dataclass +class Memory: + """What the controller keeps between runs, in a JSON file.""" + + # name -> {"created", "ref", "provider", "machine", "zone", "usd_per_hour", "per_started_hour"} + 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 -> 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) + # Machines that ended in the last 30 days: + # {"name", "provider", "start", "end", "usd_per_hour", "per_started_hour"}. + ledger: list[dict] = field(default_factory=list) + # provider -> time until which the controller does not use it (an apply failed). + cooldown: dict[str, float] = field(default_factory=dict) + + def to_json(self) -> dict: + return { + "workers": self.workers, + "registered": sorted(self.registered), + "idle_since": self.idle_since, + "draining": sorted(self.draining), + "ledger": self.ledger, + "cooldown": self.cooldown, + } + + @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", [])), + ledger=list(data.get("ledger", [])), + cooldown=dict(data.get("cooldown", {})), + ) + + +@dataclass(frozen=True) +class Decision: + add: tuple[tuple[str, Offer], ...] # new machines: (name, offer) + drain: frozenset[str] # send these workers the rq shutdown command + remove: frozenset[str] # destroy these machines + + +def machine_cost(start: float, end: float, usd_per_hour: float, per_started_hour: bool) -> float: + hours = max(0.0, end - start) / HOUR + if per_started_hour: + hours = math.ceil(hours) + return hours * usd_per_hour + + +def spend_30d(now: float, memory: Memory) -> float: + """What the machines of the last 30 days cost: the ended ones and the running ones.""" + total = 0.0 + for entry in memory.ledger: + if entry["end"] >= now - WINDOW: + total += machine_cost(entry["start"], entry["end"], entry["usd_per_hour"], + entry["per_started_hour"]) + for info in memory.workers.values(): + total += machine_cost(info["created"], now, info["usd_per_hour"], info["per_started_hour"]) + return total + + +def spend_cap_usd(revenue_30d_usd: float, k: float = 0.30, floor: float = 5.0, + ceiling: float = 100.0) -> float: + """The burst budget of the last 30 days: a share of the net revenue, with a floor + so that the first buyer gets a burst worker, and a ceiling.""" + return min(ceiling, max(floor, k * revenue_30d_usd)) + + +def decide( + now: float, + waiting: int, + live: dict[str, str], + memory: Memory, + config: Config, + ref: str, + offers: list[Offer] = (), + cap_usd: float = math.inf, +) -> tuple[Decision, Memory]: + """The next set of burst workers. + + `waiting` is the number of paid 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. `offers` are the machines that the clouds + can start now. `cap_usd` is the spend cap of the last 30 days. + """ + 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 + late_in_paid_hour = (not info.get("per_started_hour") + or (now - info["created"]) % HOUR >= config.hour_drain_after_s) + if stale or (waiting == 0 and now - idle_since[name] >= config.idle_grace_s + and late_in_paid_hour): + drain.add(name) + else: + idle_ready.append(name) + + # Above the cap on workers (the operator lowered it): stop the cheapest 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 + + # The machines that end now go to the ledger of the spend cap. + ledger = [e for e in memory.ledger if e["end"] >= now - WINDOW] + for name in remove: + info = memory.workers[name] + ledger.append({"name": name, "provider": info["provider"], "start": info["created"], + "end": now, "usd_per_hour": info["usd_per_hour"], + "per_started_hour": info["per_started_hour"]}) + + kept = {n: memory.workers[n] for n in keep} + budget = cap_usd - spend_30d(now, Memory(workers=kept, ledger=ledger)) + reserved = sum(config.reserve_hours * info["usd_per_hour"] + for n, info in kept.items() if n not in draining | drain) + + active = len(booting) + len(idle_ready) + len(busy) + need = waiting - len(idle_ready) - len(booting) + room = config.max_workers - active + usable = sorted( + (o for o in offers if memory.cooldown.get(o.provider, 0.0) <= now and o.capacity > 0), + key=lambda o: (o.usd_per_book(config), o.provider, o.machine, o.zone), + ) + left = {o: o.capacity for o in usable} + taken = set(memory.workers) | set(live) + add: list[tuple[str, Offer]] = [] + i = 0 + while len(add) < min(need, room): + offer = next((o for o in usable if left[o] > 0 + and reserved + config.reserve_hours * o.usd_per_hour <= budget), None) + if offer is None: + break # no cloud can start one, or the spend cap is reached + name = f"{NAME_PREFIX}{int(now)}-{i}" + i += 1 + if name in taken: + continue + left[offer] -= 1 + reserved += config.reserve_hours * offer.usd_per_hour + add.append((name, offer)) + + workers = dict(kept) + for name, offer in add: + workers[name] = { + "created": now, "ref": ref, "provider": offer.provider, "machine": offer.machine, + "zone": offer.zone, "usd_per_hour": offer.usd_per_hour, + "per_started_hour": offer.per_started_hour, + } + 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 | drain) & keep, + ledger=ledger, + cooldown={p: t for p, t in memory.cooldown.items() if t > now}, + ) + return Decision(add=tuple(add), drain=frozenset(drain), remove=frozenset(remove)), new_memory diff --git a/tools/burst/offers.py b/tools/burst/offers.py new file mode 100644 index 0000000..ba75655 --- /dev/null +++ b/tools/burst/offers.py @@ -0,0 +1,95 @@ +"""Live offers of the clouds: which machines each one can start now, and at what price. + +Each function reads one provider's API and returns `core.Offer` objects. The +controller asks again each time it wants a new worker, so a price drop or new +stock at one cloud moves the next worker there. +""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable + +from .core import Offer + +# Instance types that can run a worker (4 vCPU, 8 GB), and their vCPUs. Arm +# types (c7g, c6g) need an arm64 install that nobody has measured yet: turn +# them on with SUBPLZ_BURST_AWS_TYPES after one test book. +AWS_VCPUS = { + "c6a.xlarge": 4, "c7a.xlarge": 4, "m6a.xlarge": 4, "c6i.xlarge": 4, "c7i.xlarge": 4, + "c7g.xlarge": 4, "c6g.xlarge": 4, "m7g.xlarge": 4, +} +AWS_ARM = {"c7g.xlarge", "c6g.xlarge", "m7g.xlarge"} + + +def aws_offers( + spot_prices: Iterable[dict], + *, + max_price: float, + extra_usd_per_hour: float, + extra_usd_per_book: float, + capacity: int, +) -> list[Offer]: + """Offers from AWS spot price history (DescribeSpotPriceHistory items). + + The newest price of each instance type in each availability zone counts. + A price above `max_price` is left out: the spot request would not start. + `extra_usd_per_hour` is the public IPv4 address and the disk. + """ + newest: dict[tuple[str, str], tuple[str, float]] = {} + for item in spot_prices: + key = (item["InstanceType"], item["AvailabilityZone"]) + stamp = str(item["Timestamp"]) + if key not in newest or stamp > newest[key][0]: + newest[key] = (stamp, float(item["SpotPrice"])) + return [ + Offer( + provider="aws", machine=machine, zone=zone, vcpus=AWS_VCPUS[machine], + usd_per_hour=price + extra_usd_per_hour, per_started_hour=False, + extra_usd_per_book=extra_usd_per_book, capacity=capacity, + ) + for (machine, zone), (_, price) in sorted(newest.items()) + if machine in AWS_VCPUS and price <= max_price + ] + + +def fetch_aws_spot_prices(ec2, machines: list[str], since) -> list[dict]: + """DescribeSpotPriceHistory for Linux, all pages.""" + items: list[dict] = [] + pages = ec2.get_paginator("describe_spot_price_history").paginate( + InstanceTypes=machines, ProductDescriptions=["Linux/UNIX"], StartTime=since, + ) + for page in pages: + items.extend(page["SpotPriceHistory"]) + return items + + +def hetzner_offers( + get_json: Callable[[str], dict], + *, + machines: set[str], + usd_per_eur: float, + ipv4_eur_per_hour: float, + capacity: int, +) -> list[Offer]: + """Offers from the Hetzner Cloud API: the server types for sale in each + location, at the hourly net price, where a datacenter has them in stock.""" + server_types = get_json("/server_types?per_page=50")["server_types"] + in_stock = { + (dc["location"]["name"], type_id) + for dc in get_json("/datacenters?per_page=50")["datacenters"] + for type_id in dc["server_types"]["available"] + } + offers = [] + for st in server_types: + if st["name"] not in machines or st.get("architecture") != "x86" or st.get("deprecation"): + continue + for price in st["prices"]: + if (price["location"], st["id"]) not in in_stock: + continue + eur = float(price["price_hourly"]["net"]) + ipv4_eur_per_hour + offers.append(Offer( + provider="hetzner", machine=st["name"], zone=price["location"], + vcpus=int(st["cores"]), usd_per_hour=eur * usd_per_eur, + per_started_hour=True, capacity=capacity, + )) + return sorted(offers, key=lambda o: (o.machine, o.zone)) diff --git a/tools/burst/revenue.py b/tools/burst/revenue.py new file mode 100644 index 0000000..284e007 --- /dev/null +++ b/tools/burst/revenue.py @@ -0,0 +1,55 @@ +"""The net server revenue of the last 30 days: the base of the burst spend cap. + +Packs count for 30 days after the purchase. A subscription counts while it is +active and paid for the current period. Free credits count for nothing. Stripe +takes 2.9% + 30 cents of each payment, and 0.7% more of a subscription payment +(Stripe Billing). +""" + +from __future__ import annotations + +import math +from datetime import datetime, timedelta + +from sqlalchemy import select +from sqlalchemy.orm import Session + +from backend import pricing +from backend.db import Account, Purchase + + +def net_pack_cents(amount_cents: int) -> int: + return amount_cents - math.ceil(0.029 * amount_cents) - 30 + + +def net_subscription_cents(amount_cents: int) -> int: + return amount_cents - math.ceil(0.036 * amount_cents) - 30 + + +def _aware(when: datetime | None) -> datetime | None: + """SQLite gives back naive times. The app writes UTC.""" + if when is None or when.tzinfo is not None: + return when + from datetime import timezone + return when.replace(tzinfo=timezone.utc) + + +def revenue_30d_usd(session: Session, now: datetime) -> float: + since = now - timedelta(days=30) + total = 0 + for plan_id, amount, created in session.execute( + select(Purchase.plan_id, Purchase.amount_cents, Purchase.created_at) + ): + plan = pricing.get(plan_id, retired=True) + if plan is None or plan.recurring or _aware(created) < since: + continue + total += net_pack_cents(amount) + for plan_id, ends in session.execute( + select(Account.subscription_plan_id, Account.subscription_period_end) + .where(Account.subscription_status == "active") + ): + plan = pricing.get(plan_id or "", retired=True) + if plan is None or not plan.recurring or ends is None or _aware(ends) <= now: + continue + total += net_subscription_cents(plan.price_cents) + return total / 100 diff --git a/tools/burst_controller.py b/tools/burst_controller.py deleted file mode 100644 index 84e453b..0000000 --- a/tools/burst_controller.py +++ /dev/null @@ -1,313 +0,0 @@ -"""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())