Canary the Free Lane Before You Commit 2,000 Jobs

#!/usr/bin/env python3
"""probe_free_lane.py - quote free-lane capacity before committing a batch."""
import argparse, json, os, time, urllib.error, urllib.request
from concurrent.futures import ThreadPoolExecutor

def pct(xs, p):
    if not xs:
        return None
    xs = sorted(xs)
    return xs[max(0, min(len(xs) - 1, int(round(p / 100 * (len(xs) - 1)))))]

def probe(args, i):
    body = json.dumps({
        "model": args.model,
        "messages": [{"role": "user", "content": "Reply with the single word OK."}],
        "max_tokens": 8,
        "temperature": 0,
    }).encode()
    req = urllib.request.Request(
        args.endpoint.rstrip("/") + "/chat/completions",
        data=body,
        headers={
            "Content-Type": "application/json",
            "Authorization": "Bearer " + os.environ.get(args.key_env, ""),
            "X-Canary-Id": args.canary_id,
        },
    )
    t0, status, tokens, note = time.monotonic(), 0, 0, ""
    try:
        with urllib.request.urlopen(req, timeout=args.timeout) as r:
            status = r.status
            payload = json.loads(r.read() or b"{}")
            tokens = int(payload.get("usage", {}).get("total_tokens", 0))
    except urllib.error.HTTPError as e:
        status, note = e.code, (e.headers.get("Retry-After") or "")[:8]
    except Exception as e:
        status, note = -1, type(e).__name__
    rec = {"i": i, "status": status, "latency_s": round(time.monotonic() - t0, 3),
           "tokens": tokens, "retry_after": note, "canary_id": args.canary_id}
    print(json.dumps(rec), flush=True)
    return rec

if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--endpoint", required=True)
    p.add_argument("--model", required=True)
    p.add_argument("--key-env", default="FREE_LANE_KEY")
    p.add_argument("--canary-id", default=time.strftime("canary-%Y%m%dT%H%M"))
    p.add_argument("--size", type=int, default=12)
    p.add_argument("--concurrency", type=int, default=3)
    p.add_argument("--jobs", type=int, default=2000)
    p.add_argument("--slack-min", type=float, default=240.0)
    p.add_argument("--timeout", type=float, default=30.0)
    a = p.parse_args()

    with ThreadPoolExecutor(max_workers=a.concurrency) as ex:
        recs = list(ex.map(lambda i: probe(a, i), range(a.size)))

    ok = [r for r in recs if r["status"] == 200]
    lat = [r["latency_s"] for r in ok]
    accept = len(ok) / len(recs)
    p95 = pct(lat, 95) or a.timeout
    waves = a.jobs / max(1, a.concurrency)
    projected_min = waves * p95 / 60
    tokens = sum(r["tokens"] for r in ok)
    burn = tokens / max(1, len(ok)) * a.jobs

    if accept < 0.80 or projected_min > a.slack_min:
        action = "ESCALATE_OR_DEFER"
    elif accept < 0.95 or projected_min > 0.70 * a.slack_min:
        action = "SPLIT"
    else:
        action = "RUN_FREE"

    print(json.dumps({
        "action": action, "accept_rate": round(accept, 3),
        "p95_s": p95, "projected_min_optimistic": round(projected_min, 1),
        "slack_min": a.slack_min, "projected_tokens": int(burn),
        "short_circuit": all(r["status"] == 429 for r in recs[:4]),
    }, indent=2))

export FREE_LANE_KEY=...
python3 probe_free_lane.py \
  --endpoint https://your-endpoint/v1 \
  --model your-free-model \
  --size 12 --concurrency 3 --jobs 2000 --slack-min 240

{
  "action": "RUN_FREE",
  "accept_rate": 1.0,
  "p95_s": 2.4,
  "projected_min_optimistic": 26.7,
  "slack_min": 240.0,
  "projected_tokens": 14800,
  "short_circuit": false
}

# fake_lane.py - local 429 injector for probe rehearsal
import json, random, time
from http.server import BaseHTTPRequestHandler, HTTPServer

FAIL_RATE, LATENCY = 0.45, 0.08

class H(BaseHTTPRequestHandler):
    def do_POST(self):
        self.rfile.read(int(self.headers.get("Content-Length", 0)))
        time.sleep(LATENCY)
        if random.random() < FAIL_RATE:
            self.send_response(429)
            self.send_header("Retry-After", "30")
            self.end_headers()
            self.wfile.write(b'{"error":"rate_limited"}')
            return
        body = json.dumps({"choices": [{"message": {"content": "OK"}}],
                           "usage": {"total_tokens": 26}}).encode()
        self.send_response(200)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

HTTPServer(("127.0.0.1", 8099), H).serve_forever()