| """Fail-closed local cost reservations for serialized Modal reference jobs.""" |
|
|
| import fcntl |
| import json |
| import math |
| import os |
| import time |
| import uuid |
| from contextlib import contextmanager |
| from pathlib import Path |
|
|
|
|
| class Budget: |
| def __init__(self, path, cap=250.0): |
| self.path, self.cap = Path(path), cap |
| if not math.isfinite(cap) or not 0 < cap <= 250: |
| raise ValueError("The authorized cap is at most US$250") |
|
|
| @contextmanager |
| def reserve(self, amount, label): |
| if not math.isfinite(amount) or amount <= 0: |
| raise ValueError("Positive finite reservation required") |
| self.path.parent.mkdir(parents=True, exist_ok=True) |
| with self.path.with_suffix(".lock").open("a") as lock: |
| |
| try: |
| fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) |
| except BlockingIOError: |
| raise RuntimeError("Another reference job already owns the GPU budget") from None |
| ledger = ( |
| json.loads(self.path.read_text()) if self.path.exists() else {"cap": self.cap, "entries": []} |
| ) |
| if ledger["cap"] != self.cap: |
| raise ValueError("Budget cap changed") |
| if any(e["status"] == "reserved" for e in ledger["entries"]): |
| raise RuntimeError( |
| "Unresolved prior dispatch: reconcile its remote status before another GPU job" |
| ) |
| if any( |
| type(e.get("reserved_usd")) not in (int, float) |
| or not math.isfinite(e["reserved_usd"]) |
| or e["reserved_usd"] <= 0 |
| or e.get("status") not in ("reserved", "completed_conservative_charge") |
| for e in ledger["entries"] |
| ): |
| raise ValueError("Corrupt budget ledger") |
| used = sum(e["reserved_usd"] for e in ledger["entries"]) |
| if used + amount > self.cap: |
| raise RuntimeError("Modal budget exhausted before dispatch") |
| entry = { |
| "id": uuid.uuid4().hex, |
| "label": label, |
| "reserved_usd": amount, |
| "status": "reserved", |
| "time": time.time(), |
| } |
| ledger["entries"].append(entry) |
| self._write(ledger) |
| |
| yield entry |
| entry["status"] = "completed_conservative_charge" |
| entry["completed"] = time.time() |
| self._write(ledger) |
|
|
| def _write(self, value): |
| tmp = self.path.with_suffix(".tmp") |
| with tmp.open("w") as file: |
| file.write(json.dumps(value, indent=2)) |
| file.flush() |
| os.fsync(file.fileno()) |
| tmp.replace(self.path) |
|
|
