tools/burst/core.py (10047 bytes)
1 """The decisions of the burst controller. Pure: no I/O, no clock, no network. 2 3 `decide` takes what rq reports about the queue and the workers, the live 4 offers of the clouds, and the spend cap. It returns which workers to start (on 5 the cheapest offer), which to stop, and which machines to destroy. 6 """ 7 8 from __future__ import annotations 9 10 import math 11 from dataclasses import dataclass, field 12 13 NAME_PREFIX = "subplz-burst-" 14 HOUR = 3600.0 15 WINDOW = 30 * 86400.0 # the spend cap and the revenue both look back 30 days 16 17 18 @dataclass(frozen=True) 19 class Config: 20 max_workers: int 21 # A new machine installs the speech stack before its worker starts. 22 boot_timeout_s: float = 20 * 60 23 # An idle worker waits this long for a new job before it stops. 24 idle_grace_s: float = 10 * 60 25 # A 10-hour book on 4 vCPU: the web server does it at about 5x real time. 26 book_hours_4vcpu: float = 2.0 27 boot_hours: float = 0.15 28 # The cap keeps this much work of each worker in reserve. 29 reserve_hours: float = 2.5 30 # A cloud that bills each started hour gets its worker stopped only late in 31 # the hour that is paid already. 32 hour_drain_after_s: float = 50 * 60 33 34 35 @dataclass(frozen=True) 36 class Offer: 37 """A machine that one cloud can start now, at its live price.""" 38 39 provider: str # "aws" or "hetzner": the Terraform module infra/burst/workers/<provider> 40 machine: str # instance type or server type 41 zone: str # AWS availability zone or Hetzner location 42 vcpus: int 43 usd_per_hour: float # all that the machine costs per hour: price, IPv4 address, disk 44 per_started_hour: bool = False # Hetzner bills each started hour in full 45 extra_usd_per_book: float = 0.0 # for example the transfer fee of the results 46 capacity: int = 1000 # machines of this offer that the account can start now 47 48 def usd_per_book(self, config: Config) -> float: 49 hours = config.book_hours_4vcpu * 4 / self.vcpus + config.boot_hours 50 return self.usd_per_hour * hours + self.extra_usd_per_book 51 52 53 @dataclass 54 class Memory: 55 """What the controller keeps between runs, in a JSON file.""" 56 57 # name -> {"created", "ref", "provider", "machine", "zone", "usd_per_hour", "per_started_hour"} 58 workers: dict[str, dict] = field(default_factory=dict) 59 # Names that had an rq worker at least once. 60 registered: set[str] = field(default_factory=set) 61 # name -> time from which the worker was idle without a break. 62 idle_since: dict[str, float] = field(default_factory=dict) 63 # Names that received the rq shutdown command. 64 draining: set[str] = field(default_factory=set) 65 # Machines that ended in the last 30 days: 66 # {"name", "provider", "start", "end", "usd_per_hour", "per_started_hour"}. 67 ledger: list[dict] = field(default_factory=list) 68 # provider -> time until which the controller does not use it (an apply failed). 69 cooldown: dict[str, float] = field(default_factory=dict) 70 71 def to_json(self) -> dict: 72 return { 73 "workers": self.workers, 74 "registered": sorted(self.registered), 75 "idle_since": self.idle_since, 76 "draining": sorted(self.draining), 77 "ledger": self.ledger, 78 "cooldown": self.cooldown, 79 } 80 81 @classmethod 82 def from_json(cls, data: dict) -> Memory: 83 return cls( 84 workers=dict(data.get("workers", {})), 85 registered=set(data.get("registered", [])), 86 idle_since=dict(data.get("idle_since", {})), 87 draining=set(data.get("draining", [])), 88 ledger=list(data.get("ledger", [])), 89 cooldown=dict(data.get("cooldown", {})), 90 ) 91 92 93 @dataclass(frozen=True) 94 class Decision: 95 add: tuple[tuple[str, Offer], ...] # new machines: (name, offer) 96 drain: frozenset[str] # send these workers the rq shutdown command 97 remove: frozenset[str] # destroy these machines 98 99 100 def machine_cost(start: float, end: float, usd_per_hour: float, per_started_hour: bool) -> float: 101 hours = max(0.0, end - start) / HOUR 102 if per_started_hour: 103 hours = math.ceil(hours) 104 return hours * usd_per_hour 105 106 107 def spend_30d(now: float, memory: Memory) -> float: 108 """What the machines of the last 30 days cost: the ended ones and the running ones.""" 109 total = 0.0 110 for entry in memory.ledger: 111 if entry["end"] >= now - WINDOW: 112 total += machine_cost(entry["start"], entry["end"], entry["usd_per_hour"], 113 entry["per_started_hour"]) 114 for info in memory.workers.values(): 115 total += machine_cost(info["created"], now, info["usd_per_hour"], info["per_started_hour"]) 116 return total 117 118 119 def spend_cap_usd(revenue_30d_usd: float, k: float = 0.30, floor: float = 5.0, 120 ceiling: float = 100.0) -> float: 121 """The burst budget of the last 30 days: a share of the net revenue, with a floor 122 so that the first buyer gets a burst worker, and a ceiling.""" 123 return min(ceiling, max(floor, k * revenue_30d_usd)) 124 125 126 def decide( 127 now: float, 128 waiting: int, 129 live: dict[str, str], 130 memory: Memory, 131 config: Config, 132 ref: str, 133 offers: list[Offer] = (), 134 cap_usd: float = math.inf, 135 ) -> tuple[Decision, Memory]: 136 """The next set of burst workers. 137 138 `waiting` is the number of paid jobs that no worker has taken. `live` maps 139 the host name of each burst worker that rq knows to "busy" or "idle". `ref` 140 is the app commit for new workers. `offers` are the machines that the clouds 141 can start now. `cap_usd` is the spend cap of the last 30 days. 142 """ 143 keep: set[str] = set() 144 remove: set[str] = set() 145 drain: set[str] = set() 146 registered = set(memory.registered) 147 idle_since = dict(memory.idle_since) 148 draining = set(memory.draining) 149 booting: list[str] = [] 150 idle_ready: list[str] = [] 151 busy: list[str] = [] 152 153 for name, info in memory.workers.items(): 154 state = live.get(name) 155 stale = info.get("ref") != ref 156 if state is None: 157 if name in registered or stale: 158 remove.add(name) # it stopped, or it would run an old commit 159 elif now - info["created"] >= config.boot_timeout_s: 160 remove.add(name) # it never started 161 else: 162 keep.add(name) 163 booting.append(name) 164 continue 165 166 keep.add(name) 167 registered.add(name) 168 if state == "busy": 169 idle_since.pop(name, None) 170 if stale and name not in draining: 171 drain.add(name) # it finishes the job, then stops 172 elif name not in draining: 173 busy.append(name) 174 continue 175 176 idle_since.setdefault(name, now) 177 if name in draining: 178 continue # it stops soon 179 late_in_paid_hour = (not info.get("per_started_hour") 180 or (now - info["created"]) % HOUR >= config.hour_drain_after_s) 181 if stale or (waiting == 0 and now - idle_since[name] >= config.idle_grace_s 182 and late_in_paid_hour): 183 drain.add(name) 184 else: 185 idle_ready.append(name) 186 187 # Above the cap on workers (the operator lowered it): stop the cheapest first. 188 active = len(booting) + len(idle_ready) + len(busy) 189 excess = active - config.max_workers 190 while excess > 0 and booting: 191 name = booting.pop() 192 keep.discard(name) 193 remove.add(name) 194 excess -= 1 195 for group in (idle_ready, busy): 196 while excess > 0 and group: 197 drain.add(group.pop()) 198 excess -= 1 199 200 # The machines that end now go to the ledger of the spend cap. 201 ledger = [e for e in memory.ledger if e["end"] >= now - WINDOW] 202 for name in remove: 203 info = memory.workers[name] 204 ledger.append({"name": name, "provider": info["provider"], "start": info["created"], 205 "end": now, "usd_per_hour": info["usd_per_hour"], 206 "per_started_hour": info["per_started_hour"]}) 207 208 kept = {n: memory.workers[n] for n in keep} 209 budget = cap_usd - spend_30d(now, Memory(workers=kept, ledger=ledger)) 210 reserved = sum(config.reserve_hours * info["usd_per_hour"] 211 for n, info in kept.items() if n not in draining | drain) 212 213 active = len(booting) + len(idle_ready) + len(busy) 214 need = waiting - len(idle_ready) - len(booting) 215 room = config.max_workers - active 216 usable = sorted( 217 (o for o in offers if memory.cooldown.get(o.provider, 0.0) <= now and o.capacity > 0), 218 key=lambda o: (o.usd_per_book(config), o.provider, o.machine, o.zone), 219 ) 220 left = {o: o.capacity for o in usable} 221 taken = set(memory.workers) | set(live) 222 add: list[tuple[str, Offer]] = [] 223 i = 0 224 while len(add) < min(need, room): 225 offer = next((o for o in usable if left[o] > 0 226 and reserved + config.reserve_hours * o.usd_per_hour <= budget), None) 227 if offer is None: 228 break # no cloud can start one, or the spend cap is reached 229 name = f"{NAME_PREFIX}{int(now)}-{i}" 230 i += 1 231 if name in taken: 232 continue 233 left[offer] -= 1 234 reserved += config.reserve_hours * offer.usd_per_hour 235 add.append((name, offer)) 236 237 workers = dict(kept) 238 for name, offer in add: 239 workers[name] = { 240 "created": now, "ref": ref, "provider": offer.provider, "machine": offer.machine, 241 "zone": offer.zone, "usd_per_hour": offer.usd_per_hour, 242 "per_started_hour": offer.per_started_hour, 243 } 244 new_memory = Memory( 245 workers=workers, 246 registered=registered & keep, 247 idle_since={n: t for n, t in idle_since.items() if n in keep and live.get(n) == "idle"}, 248 draining=(draining | drain) & keep, 249 ledger=ledger, 250 cooldown={p: t for p, t in memory.cooldown.items() if t > now}, 251 ) 252 return Decision(add=tuple(add), drain=frozenset(drain), remove=frozenset(remove)), new_memory