tests/test_burst.py (19709 bytes)
1 """The burst controller (tools/burst): which workers it starts, where, and when it stops.""" 2 3 from __future__ import annotations 4 5 import json 6 import math 7 import sys 8 import time 9 from datetime import datetime, timedelta, timezone 10 from pathlib import Path 11 12 import fakeredis 13 import pytest 14 from hypothesis import given, settings 15 from hypothesis import strategies as st 16 from rq import Queue, Worker 17 18 from backend.db import Account, Purchase, SessionLocal 19 from tools.burst import __main__ as run_mod 20 from tools.burst import offers as offers_mod 21 from tools.burst.core import ( 22 NAME_PREFIX, Config, Memory, Offer, decide, machine_cost, spend_30d, spend_cap_usd, 23 ) 24 from tools.burst.revenue import net_pack_cents, net_subscription_cents, revenue_30d_usd 25 26 REF = "new" 27 NOW = 1_000_000.0 28 MIN = 60.0 29 AWS = Offer("aws", "c6a.xlarge", "us-west-2a", 4, 0.07) 30 HETZNER = Offer("hetzner", "cx33", "fsn1", 4, 0.0164, per_started_hour=True) 31 32 33 def worker(created_ago: float, ref: str = REF, offer: Offer = AWS) -> dict: 34 return {"created": NOW - created_ago, "ref": ref, "provider": offer.provider, 35 "machine": offer.machine, "zone": offer.zone, "usd_per_hour": offer.usd_per_hour, 36 "per_started_hour": offer.per_started_hour} 37 38 39 # --- invariants over any state ------------------------------------------------- 40 41 names = st.sampled_from([f"{NAME_PREFIX}{i}" for i in range(8)]) 42 offer_st = st.builds( 43 Offer, provider=st.sampled_from(["aws", "hetzner"]), machine=st.sampled_from(["a", "b"]), 44 zone=st.sampled_from(["z1", "z2"]), vcpus=st.sampled_from([2, 4, 8]), 45 usd_per_hour=st.floats(0.005, 0.2), per_started_hour=st.booleans(), 46 capacity=st.integers(0, 3), 47 ) 48 49 50 @st.composite 51 def states(draw): 52 pool = draw(st.sets(names, max_size=8)) 53 workers = {n: worker(draw(st.floats(0, 60 * MIN)), draw(st.sampled_from(["old", REF])), 54 draw(st.sampled_from([AWS, HETZNER]))) for n in pool} 55 live = {n: draw(st.sampled_from(["busy", "idle"])) for n in pool if draw(st.booleans())} 56 registered = {n for n in pool if n in live or draw(st.booleans())} 57 draining = {n for n in registered if draw(st.booleans())} 58 idle_since = {n: NOW - draw(st.floats(0, 30 * MIN)) for n in live if live[n] == "idle"} 59 memory = Memory(workers=workers, registered=registered, idle_since=idle_since, 60 draining=draining, cooldown=draw(st.sampled_from([{}, {"aws": NOW + 60}]))) 61 offers = draw(st.lists(offer_st, max_size=4, unique=True)) 62 cap = draw(st.sampled_from([math.inf, 0.0, 0.5, 5.0])) 63 return (draw(st.integers(0, 12)), live, memory, Config(max_workers=draw(st.integers(0, 6))), 64 offers, cap) 65 66 67 @settings(max_examples=500) 68 @given(states()) 69 def test_invariants(state): 70 waiting, live, memory, config, offers, cap = state 71 d, after = decide(NOW, waiting, live, memory, config, REF, offers, cap) 72 added = [n for n, _ in d.add] 73 74 # A worker that rq knows is never destroyed: it must stop by itself first. 75 assert not d.remove & set(live) 76 assert set(after.workers) == (set(memory.workers) - d.remove) | set(added) 77 assert all(n.startswith(NAME_PREFIX) and n not in memory.workers and n not in live for n in added) 78 # The cap on workers holds, and no more new workers than waiting jobs. 79 assert len([n for n in after.workers if n not in after.draining]) <= config.max_workers 80 assert len(added) <= waiting 81 # Only a worker that rq knows gets the shutdown command. 82 assert d.drain <= set(live) 83 # New workers come from offers that have room, from a cloud that is not resting. 84 for offer in {o for _, o in d.add}: 85 assert offer in offers and memory.cooldown.get(offer.provider, 0) <= NOW 86 assert sum(1 for _, o in d.add if o == offer) <= offer.capacity 87 # Each machine that ends goes to the ledger, once. 88 assert sorted(e["name"] for e in after.ledger if e["end"] == NOW) == sorted(d.remove) 89 90 91 @settings(max_examples=500) 92 @given(states()) 93 def test_a_new_worker_goes_to_the_cheapest_offer_that_fits(state): 94 waiting, live, memory, config, offers, cap = state 95 d, _ = decide(NOW, waiting, live, memory, config, REF, offers, cap) 96 if not d.add: 97 return 98 first = d.add[0][1] 99 fitting = [o for o in offers if o.capacity > 0 and memory.cooldown.get(o.provider, 0) <= NOW] 100 cheaper = [o for o in fitting if o.usd_per_book(config) < first.usd_per_book(config)] 101 # A cheaper offer is left out only when the spend cap has no room for its 102 # reserve, which is its hourly price: so it must cost more by the hour. 103 # (A counterexample of 2026-09-22: 8 vCPU at $0.30/h is cheaper by the book 104 # than 2 vCPU at $0.125/h, and a small cap takes only the second.) 105 for o in cheaper: 106 assert o.usd_per_hour > first.usd_per_hour 107 if cap == math.inf: 108 assert not cheaper 109 110 111 @settings(max_examples=300) 112 @given(states()) 113 def test_the_spend_cap_holds(state): 114 waiting, live, memory, config, offers, cap = state 115 d, after = decide(NOW, waiting, live, memory, config, REF, offers, cap) 116 if d.add: 117 reserve = sum(config.reserve_hours * w["usd_per_hour"] for n, w in after.workers.items() 118 if n not in after.draining) 119 assert spend_30d(NOW, after) + reserve <= cap + 1e-9 120 121 122 # --- scenarios ------------------------------------------------------------------- 123 124 def test_a_backlog_adds_workers_on_the_cheapest_cloud_up_to_the_cap(): 125 d, after = decide(NOW, 5, {}, Memory(), Config(max_workers=3), REF, [AWS, HETZNER]) 126 assert [o for _, o in d.add] == [HETZNER, HETZNER, HETZNER] 127 assert all(after.workers[n]["provider"] == "hetzner" for n, _ in d.add) 128 129 130 def test_when_the_cheapest_cloud_is_full_the_next_one_gets_the_worker(): 131 small = Offer("hetzner", "cx33", "fsn1", 4, 0.0164, per_started_hour=True, capacity=1) 132 d, _ = decide(NOW, 3, {}, Memory(), Config(max_workers=3), REF, [AWS, small]) 133 assert [o.provider for _, o in d.add] == ["hetzner", "aws", "aws"] 134 135 136 def test_a_resting_cloud_gets_no_worker(): 137 memory = Memory(cooldown={"hetzner": NOW + 10 * MIN}) 138 d, _ = decide(NOW, 1, {}, memory, Config(max_workers=3), REF, [AWS, HETZNER]) 139 assert [o.provider for _, o in d.add] == ["aws"] 140 141 142 def test_a_small_cap_takes_the_offer_with_the_lower_hourly_price(): 143 # Found by the property test: the big machine is cheaper by the book, but 144 # the cap has room only for the reserve of the small one. 145 small = Offer("aws", "small", "z1", 2, 0.125) # $0.52 a book, $0.31 reserve 146 big = Offer("hetzner", "big", "z1", 8, 0.35) # $0.40 a book, $0.88 reserve 147 d, _ = decide(NOW, 1, {}, Memory(), Config(max_workers=1), REF, [small, big], cap_usd=0.5) 148 assert [o for _, o in d.add] == [small] 149 d, _ = decide(NOW, 1, {}, Memory(), Config(max_workers=1), REF, [small, big]) 150 assert [o for _, o in d.add] == [big] 151 152 153 def test_egress_can_make_the_cheaper_machine_the_dearer_book(): 154 aws_with_egress = Offer("aws", "c6a.xlarge", "us-west-2a", 4, 0.03, extra_usd_per_book=0.2) 155 d, _ = decide(NOW, 1, {}, Memory(), Config(max_workers=1), REF, [aws_with_egress, HETZNER]) 156 assert d.add[0][1] == HETZNER 157 158 159 def test_the_spend_cap_stops_new_workers_but_not_running_ones(): 160 memory = Memory(workers={"subplz-burst-a": worker(60 * MIN)}, registered={"subplz-burst-a"}, 161 ledger=[{"name": "old", "provider": "aws", "start": NOW - 50 * 3600, 162 "end": NOW - 10 * 3600, "usd_per_hour": 0.07, "per_started_hour": False}]) 163 # 40 h at $0.07 = $2.80 spent, plus 1 h of the running worker, and its reserve. 164 d, after = decide(NOW, 3, {"subplz-burst-a": "busy"}, memory, Config(max_workers=3), REF, 165 [AWS], cap_usd=3.2) 166 assert d.add == () and not d.drain and "subplz-burst-a" in after.workers 167 d, _ = decide(NOW, 3, {"subplz-burst-a": "busy"}, memory, Config(max_workers=3), REF, 168 [AWS], cap_usd=10.0) 169 assert len(d.add) == 2 170 171 172 def test_an_hourly_machine_stops_late_in_its_paid_hour(): 173 name = "subplz-burst-h" 174 for age, stops in ((70 * MIN, False), (115 * MIN, True)): 175 memory = Memory(workers={name: worker(age, offer=HETZNER)}, registered={name}, 176 idle_since={name: NOW - 20 * MIN}) 177 d, _ = decide(NOW, 0, {name: "idle"}, memory, Config(max_workers=3), REF) 178 assert (name in d.drain) is stops, age 179 180 181 def test_an_idle_worker_stops_after_the_grace_time_then_its_machine_goes(): 182 name = "subplz-burst-a" 183 memory = Memory(workers={name: worker(60 * MIN)}, registered={name}, 184 idle_since={name: NOW - 11 * MIN}) 185 d, after = decide(NOW, 0, {name: "idle"}, memory, Config(max_workers=3), REF) 186 assert d.drain == {name} and not d.remove 187 d, after = decide(NOW + MIN, 0, {}, after, Config(max_workers=3), REF) 188 assert d.remove == {name} and after.workers == {} 189 assert after.ledger[-1]["name"] == name and after.ledger[-1]["end"] == NOW + MIN 190 191 192 def test_a_machine_that_never_starts_its_worker_is_destroyed_and_replaced(): 193 memory = Memory(workers={"subplz-burst-a": worker(21 * MIN)}) 194 d, _ = decide(NOW, 1, {}, memory, Config(max_workers=1), REF, [AWS]) 195 assert d.remove == {"subplz-burst-a"} and len(d.add) == 1 196 197 198 def test_after_a_deploy_the_workers_of_the_old_commit_are_replaced(): 199 memory = Memory( 200 workers={"subplz-burst-idle": worker(60 * MIN, "old"), "subplz-burst-busy": worker(60 * MIN, "old"), 201 "subplz-burst-boot": worker(2 * MIN, "old")}, 202 registered={"subplz-burst-idle", "subplz-burst-busy"}, 203 ) 204 live = {"subplz-burst-idle": "idle", "subplz-burst-busy": "busy"} 205 d, after = decide(NOW, 2, live, memory, Config(max_workers=3), REF, [AWS]) 206 assert d.drain == {"subplz-burst-idle", "subplz-burst-busy"} 207 assert d.remove == {"subplz-burst-boot"} 208 assert len(d.add) == 2 and all(after.workers[n]["ref"] == REF for n, _ in d.add) 209 210 211 def test_a_cap_of_zero_stops_every_worker_and_starts_none(): 212 memory = Memory( 213 workers={"subplz-burst-boot": worker(MIN), "subplz-burst-idle": worker(60 * MIN), 214 "subplz-burst-busy": worker(60 * MIN)}, 215 registered={"subplz-burst-idle", "subplz-burst-busy"}, 216 ) 217 live = {"subplz-burst-idle": "idle", "subplz-burst-busy": "busy"} 218 d, _ = decide(NOW, 4, live, memory, Config(max_workers=0), REF, [AWS]) 219 assert d.remove == {"subplz-burst-boot"} 220 assert d.drain == {"subplz-burst-idle", "subplz-burst-busy"} and d.add == () 221 222 223 def test_machine_cost_and_spend(): 224 assert machine_cost(0, 5400, 0.1, False) == pytest.approx(0.15) 225 assert machine_cost(0, 5400, 0.1, True) == pytest.approx(0.2) # 2 started hours 226 memory = Memory(workers={"w": worker(3600)}, 227 ledger=[{"name": "x", "provider": "aws", "start": 0, "end": NOW - 31 * 86400, 228 "usd_per_hour": 9.0, "per_started_hour": False}]) 229 assert spend_30d(NOW, memory) == pytest.approx(0.07) # the old machine is out of the window 230 231 232 def test_spend_cap(): 233 assert spend_cap_usd(0.0) == 5.0 # the floor serves the first buyer 234 assert spend_cap_usd(115.85) == pytest.approx(34.755) 235 assert spend_cap_usd(1000.0) == 100.0 236 237 238 # --- live offers -------------------------------------------------------------------- 239 240 def test_aws_offers_take_the_newest_price_and_skip_a_price_above_the_maximum(): 241 items = [ 242 {"InstanceType": "c6a.xlarge", "AvailabilityZone": "us-west-2a", "SpotPrice": "0.090", 243 "Timestamp": "2026-09-22T10:00:00Z"}, 244 {"InstanceType": "c6a.xlarge", "AvailabilityZone": "us-west-2a", "SpotPrice": "0.050", 245 "Timestamp": "2026-09-22T11:00:00Z"}, 246 {"InstanceType": "c7a.xlarge", "AvailabilityZone": "us-west-2b", "SpotPrice": "0.080", 247 "Timestamp": "2026-09-22T11:00:00Z"}, 248 {"InstanceType": "x9.huge", "AvailabilityZone": "us-west-2b", "SpotPrice": "0.001", 249 "Timestamp": "2026-09-22T11:00:00Z"}, 250 ] 251 found = offers_mod.aws_offers(items, max_price=0.0625, extra_usd_per_hour=0.0072, 252 extra_usd_per_book=0.0, capacity=1) 253 assert found == [Offer("aws", "c6a.xlarge", "us-west-2a", 4, pytest.approx(0.0572), False, 0.0, 1)] 254 255 256 def test_hetzner_offers_come_from_stock_at_the_net_price(): 257 api = { 258 "/server_types?per_page=50": {"server_types": [ 259 {"id": 1, "name": "cx33", "cores": 4, "architecture": "x86", "deprecation": None, 260 "prices": [{"location": "fsn1", "price_hourly": {"net": "0.0136", "gross": "0.0162"}}, 261 {"location": "nbg1", "price_hourly": {"net": "0.0136", "gross": "0.0162"}}]}, 262 {"id": 2, "name": "cax21", "cores": 4, "architecture": "arm", "deprecation": None, 263 "prices": [{"location": "fsn1", "price_hourly": {"net": "0.0100", "gross": "0.0119"}}]}, 264 {"id": 3, "name": "cpx31", "cores": 4, "architecture": "x86", "deprecation": None, 265 "prices": [{"location": "hil", "price_hourly": {"net": "0.0997", "gross": "0.0997"}}]}, 266 ]}, 267 "/datacenters?per_page=50": {"datacenters": [ 268 {"location": {"name": "fsn1"}, "server_types": {"available": [1, 2]}}, 269 {"location": {"name": "nbg1"}, "server_types": {"available": []}}, # sold out 270 {"location": {"name": "hil"}, "server_types": {"available": [3]}}, 271 ]}, 272 } 273 found = offers_mod.hetzner_offers(api.__getitem__, machines={"cx33", "cax21", "cpx31"}, 274 usd_per_eur=1.15, ipv4_eur_per_hour=0.0007, capacity=2) 275 assert [(o.machine, o.zone) for o in found] == [("cpx31", "hil"), ("cx33", "fsn1")] 276 assert found[1].usd_per_hour == pytest.approx((0.0136 + 0.0007) * 1.15) 277 assert all(o.per_started_hour and o.capacity == 2 for o in found) 278 279 280 # --- revenue ------------------------------------------------------------------------- 281 282 def test_net_after_stripe(): 283 assert [net_pack_cents(a) for a in (499, 3999, 17499)] == [454, 3853, 16961] 284 assert [net_subscription_cents(a) for a in (499, 999, 5000)] == [451, 933, 4790] 285 286 287 def test_revenue_counts_recent_packs_and_active_subscriptions(): 288 now = datetime(2026, 9, 22, tzinfo=timezone.utc) 289 tag = str(time.time_ns()) 290 with SessionLocal() as s: 291 base = revenue_30d_usd(s, now) 292 s.add(Account(id=f"acct_a{tag}", device_token=f"a{tag}")) 293 s.add(Account(id=f"acct_b{tag}", device_token=f"b{tag}", subscription_plan_id="month30", 294 subscription_status="active", subscription_period_end=now + timedelta(days=9))) 295 s.add(Account(id=f"acct_c{tag}", device_token=f"c{tag}", subscription_plan_id="unlimited", 296 subscription_status="canceled", subscription_period_end=now + timedelta(days=9))) 297 s.add(Purchase(account_id=f"acct_a{tag}", plan_id="pack100", credits=100, amount_cents=3999, 298 stripe_session_id=f"cs_new{tag}", created_at=now - timedelta(days=3))) 299 s.add(Purchase(account_id=f"acct_a{tag}", plan_id="pack10", credits=10, amount_cents=499, 300 stripe_session_id=f"cs_old{tag}", created_at=now - timedelta(days=40))) 301 s.commit() 302 assert revenue_30d_usd(s, now) - base == pytest.approx((3853 + 933) / 100) 303 304 305 # --- one run, with fake clouds and a fake Terraform ------------------------------------ 306 307 def register(conn, rq_name: str, host: str, state: str) -> None: 308 w = Worker([Queue(run_mod.QUEUE_NAME, connection=conn)], connection=conn, name=rq_name) 309 w.hostname = host 310 w.register_birth() 311 w.set_state(state) 312 313 314 def test_the_controller_counts_the_paid_queue_of_the_app(): 315 from backend.queue import PAID_QUEUE 316 assert run_mod.QUEUE_NAME == PAID_QUEUE 317 318 319 def test_read_queue_reports_burst_workers_only(): 320 conn = fakeredis.FakeRedis() 321 register(conn, "rq-1", "subplz-burst-1-0", "busy") 322 register(conn, "rq-web", "racknerd-655a1e9", "busy") # the web server's own worker 323 Queue(run_mod.QUEUE_NAME, connection=conn).enqueue("os.getcwd") 324 waiting, live, rq_names = run_mod.read_queue(conn) 325 assert (waiting, live, rq_names) == (1, {"subplz-burst-1-0": "busy"}, {"subplz-burst-1-0": "rq-1"}) 326 327 328 def test_terraform_gets_the_burst_keys_not_the_storage_keys(monkeypatch): 329 monkeypatch.setenv("AWS_ACCESS_KEY_ID", "r2-key") 330 monkeypatch.setenv("AWS_ENDPOINT_URL_S3", "https://r2.example") 331 monkeypatch.setenv("SUBPLZ_WEB_DATABASE_URL", "postgresql://secret") 332 monkeypatch.setenv("SUBPLZ_BURST_AWS_ACCESS_KEY_ID", "ec2-key") 333 monkeypatch.setenv("HCLOUD_TOKEN", "hz") 334 aws = run_mod.terraform_env("aws") 335 assert aws["AWS_ACCESS_KEY_ID"] == "ec2-key" and "AWS_ENDPOINT_URL_S3" not in aws 336 assert "SUBPLZ_WEB_DATABASE_URL" not in aws 337 hetzner = run_mod.terraform_env("hetzner") 338 assert hetzner["HCLOUD_TOKEN"] == "hz" and "AWS_ACCESS_KEY_ID" not in hetzner 339 340 341 @pytest.fixture 342 def world(tmp_path, monkeypatch): 343 conn = fakeredis.FakeRedis() 344 monkeypatch.setattr("redis.Redis.from_url", lambda url: conn) 345 released = [] 346 monkeypatch.setattr(run_mod, "ensure_release", lambda app_dir, ref, bucket: released.append(ref)) 347 monkeypatch.setattr(run_mod, "revenue_30d", lambda: 0.0) # the floor: $5 348 offers = [AWS, HETZNER] 349 monkeypatch.setattr(run_mod, "gather_offers", lambda memory: list(offers)) 350 calls = tmp_path / "calls.txt" 351 fake = tmp_path / "fake_terraform.py" 352 fake.write_text( 353 "import sys, pathlib\n" 354 f"pathlib.Path({str(calls)!r}).open('a').write(' '.join(sys.argv[1:]) + '\\n')\n" 355 f"bad = pathlib.Path({str(tmp_path / 'fail')!r}).read_text()\n" 356 "sys.exit(1 if bad and bad in ' '.join(sys.argv) else 0)\n" 357 ) 358 (tmp_path / "fail").write_text("") 359 360 def run(extra=()): 361 argv = ["--redis-url", "redis://x", "--bucket", "b", "--max-workers", "2", 362 "--app-dir", str(Path(__file__).resolve().parents[1]), 363 "--modules-dir", str(tmp_path / "modules"), "--state-dir", str(tmp_path / "state"), 364 "--terraform", sys.executable, str(fake), *extra] 365 return run_mod.main(argv) 366 367 return conn, run, tmp_path, calls, released 368 369 370 def enqueue(conn, n: int) -> None: 371 for _ in range(n): 372 Queue(run_mod.QUEUE_NAME, connection=conn).enqueue("os.getcwd") 373 374 375 def test_a_run_starts_workers_on_the_cheapest_cloud(world): 376 conn, run, tmp, calls, released = world 377 enqueue(conn, 3) 378 assert run() == 0 379 hetzner = json.loads((tmp / "state" / "workers-hetzner.tfvars.json").read_text())["workers"] 380 assert len(hetzner) == 2 and all(w["machine"] == "cx33" for w in hetzner.values()) 381 assert "modules" in calls.read_text() and "hetzner" in calls.read_text() 382 assert not (tmp / "state" / "workers-aws.tfvars.json").exists() # nothing changed there 383 assert len(released) == 1 384 385 386 def test_a_failed_apply_rests_the_cloud_and_the_next_run_uses_another(world): 387 conn, run, tmp, calls, released = world 388 enqueue(conn, 1) 389 (tmp / "fail").write_text("hetzner") 390 assert run() == 1 391 state = json.loads((tmp / "state" / "controller.json").read_text()) 392 assert state["workers"] == {} and state["cooldown"]["hetzner"] > time.time() 393 assert json.loads((tmp / "state" / "workers-hetzner.tfvars.json").read_text()) == {"workers": {}} 394 assert run() == 0 # Hetzner rests: AWS gets the worker 395 aws = json.loads((tmp / "state" / "workers-aws.tfvars.json").read_text())["workers"] 396 assert [w["machine"] for w in aws.values()] == ["c6a.xlarge"] 397 398 399 def test_a_dry_run_changes_nothing(world): 400 conn, run, tmp, calls, released = world 401 enqueue(conn, 1) 402 assert run(["--dry-run"]) == 0 403 assert not calls.exists() and not (tmp / "state").exists() and not released