From 70ef37bca5537f965e5dac26f922d4ad5fa171be Mon Sep 17 00:00:00 2001 From: Prakhar Verma Date: Sun, 10 May 2026 20:06:29 -0400 Subject: [PATCH 1/2] Clean up gateway handlers --- gateway/app/main.py | 168 +++++++++++++++++++++++++++++--------------- 1 file changed, 112 insertions(+), 56 deletions(-) diff --git a/gateway/app/main.py b/gateway/app/main.py index 5954fd6..f94ae1f 100644 --- a/gateway/app/main.py +++ b/gateway/app/main.py @@ -3,7 +3,7 @@ import os import time from contextlib import asynccontextmanager -from typing import AsyncIterator +from typing import Any, AsyncIterator import httpx from fastapi import FastAPI, HTTPException @@ -12,9 +12,14 @@ from .cache import PredictionCache from .dashboard import DASHBOARD_HTML from .metrics import MetricsStore -from .router import WorkerPool +from .router import Worker, WorkerPool from .schemas import GatewayPrediction, PredictRequest, WorkerControlRequest +DEFAULT_WORKER_URL = "http://localhost:8011" +MAX_WORKER_ATTEMPTS = 2 +JsonObject = dict[str, Any] +WorkerRow = dict[str, str | int | float | None] + def env_flag(name: str, default: bool = True) -> bool: value = os.getenv(name) @@ -36,7 +41,7 @@ def parse_worker_hostports(env: dict[str, str] | None = None) -> list[str]: def parse_worker_urls(raw: str | None) -> list[str]: if not raw: - return parse_worker_hostports() or ["http://localhost:8011"] + return parse_worker_hostports() or [DEFAULT_WORKER_URL] return [item.strip() for item in raw.split(",") if item.strip()] @@ -52,14 +57,96 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: parse_worker_urls(os.getenv("WORKER_URLS")), strategy=os.getenv("ROUTER_STRATEGY", "load_aware"), ) -cache = PredictionCache(os.getenv("REDIS_URL"), enabled=env_flag("CACHE_ENABLED", default=True)) +cache = PredictionCache( + os.getenv("REDIS_URL"), + enabled=env_flag("CACHE_ENABLED", default=True), +) metrics = MetricsStore() request_timeout = float(os.getenv("REQUEST_TIMEOUT_SECONDS", "1.0")) +def elapsed_ms(started: float) -> float: + return (time.perf_counter() - started) * 1000 + + +def gateway_client() -> httpx.AsyncClient: + return app.state.client + + +def worker_for_index(worker_index: int) -> Worker: + if worker_index < 0 or worker_index >= len(router.workers): + raise HTTPException(status_code=404, detail="worker not found") + return router.workers[worker_index] + + +async def read_worker_config(worker_url: str) -> JsonObject | None: + try: + response = await gateway_client().get(f"{worker_url}/config", timeout=0.4) + response.raise_for_status() + except httpx.HTTPError: + return None + return response.json() + + +def add_worker_config( + row: WorkerRow, + config: JsonObject | None, +) -> None: + if config is None: + row["control_status"] = "offline" + return + + row["worker_id"] = config.get("worker_id") + row["model_version"] = config.get("model_version") + row["delay_ms"] = config.get("delay_ms") + row["fail_rate"] = config.get("fail_rate") + row["control_status"] = "ok" + + +def cached_prediction(cached: JsonObject, latency_ms: float) -> GatewayPrediction: + return GatewayPrediction( + prediction=cached["prediction"], + cached=True, + worker_url=cached["worker_url"], + attempts=0, + latency_ms=round(latency_ms, 2), + ) + + +async def control_remote_worker( + worker: Worker, + update: WorkerControlRequest, +) -> JsonObject: + try: + response = await gateway_client().post( + f"{worker.url}/control", + json=update.model_dump(exclude_none=True), + timeout=request_timeout, + ) + response.raise_for_status() + except httpx.HTTPError as exc: + detail = f"worker control failed: {exc.__class__.__name__}" + raise HTTPException(status_code=502, detail=detail) from exc + return response.json() + + +async def request_prediction(worker: Worker, payload: PredictRequest) -> JsonObject: + response = await gateway_client().post( + f"{worker.url}/predict", + json=payload.model_dump(), + timeout=request_timeout, + ) + response.raise_for_status() + return response.json() + + @app.get("/health") async def health() -> dict[str, str | int]: - return {"status": "ok", "workers": len(router.workers), "router_strategy": router.strategy} + return { + "status": "ok", + "workers": len(router.workers), + "router_strategy": router.strategy, + } @app.get("/dashboard", response_class=HTMLResponse) @@ -68,43 +155,22 @@ async def dashboard() -> str: @app.get("/workers") -async def workers() -> list[dict[str, str | int | float | None]]: - client: httpx.AsyncClient = app.state.client +async def workers() -> list[WorkerRow]: rows = router.snapshot() for index, row in enumerate(rows): row["index"] = index - try: - response = await client.get(f"{row['url']}/config", timeout=0.4) - response.raise_for_status() - config = response.json() - row["worker_id"] = config.get("worker_id") - row["model_version"] = config.get("model_version") - row["delay_ms"] = config.get("delay_ms") - row["fail_rate"] = config.get("fail_rate") - row["control_status"] = "ok" - except httpx.HTTPError: - row["control_status"] = "offline" + add_worker_config(row, await read_worker_config(str(row["url"]))) return rows @app.post("/workers/{worker_index}/control") -async def control_worker(worker_index: int, update: WorkerControlRequest) -> dict[str, object]: - if worker_index < 0 or worker_index >= len(router.workers): - raise HTTPException(status_code=404, detail="worker not found") - - worker = router.workers[worker_index] - client: httpx.AsyncClient = app.state.client - try: - response = await client.post( - f"{worker.url}/control", - json=update.model_dump(exclude_none=True), - timeout=request_timeout, - ) - response.raise_for_status() - except httpx.HTTPError as exc: - raise HTTPException(status_code=502, detail=f"worker control failed: {exc.__class__.__name__}") from exc - - return {"worker_index": worker_index, "worker_url": worker.url, "config": response.json()} +async def control_worker( + worker_index: int, + update: WorkerControlRequest, +) -> dict[str, object]: + worker = worker_for_index(worker_index) + config = await control_remote_worker(worker, update) + return {"worker_index": worker_index, "worker_url": worker.url, "config": config} @app.get("/metrics") @@ -119,35 +185,22 @@ async def predict(payload: PredictRequest) -> GatewayPrediction: cached = cache.get(cache_key) if cached is not None: - latency_ms = (time.perf_counter() - started) * 1000 + latency_ms = elapsed_ms(started) metrics.record_request(latency_ms, cached=True) - return GatewayPrediction( - prediction=cached["prediction"], - cached=True, - worker_url=cached["worker_url"], - attempts=0, - latency_ms=round(latency_ms, 2), - ) + return cached_prediction(cached, latency_ms) errors: list[str] = [] attempts = 0 - client: httpx.AsyncClient = app.state.client - for worker in router.candidates(attempts=2): + for worker in router.candidates(attempts=MAX_WORKER_ATTEMPTS): attempts += 1 router.mark_start(worker) worker_started = time.perf_counter() try: - response = await client.post( - f"{worker.url}/predict", - json=payload.model_dump(), - timeout=request_timeout, - ) - response.raise_for_status() - prediction = response.json() - worker_latency_ms = (time.perf_counter() - worker_started) * 1000 + prediction = await request_prediction(worker, payload) + worker_latency_ms = elapsed_ms(worker_started) router.mark_success(worker, worker_latency_ms) - latency_ms = (time.perf_counter() - started) * 1000 + latency_ms = elapsed_ms(started) result = { "prediction": prediction, "worker_url": worker.url, @@ -162,6 +215,9 @@ async def predict(payload: PredictRequest) -> GatewayPrediction: router.mark_failure(worker, error) errors.append(f"{worker.url}: {error}") - latency_ms = (time.perf_counter() - started) * 1000 + latency_ms = elapsed_ms(started) metrics.record_request(latency_ms, cached=False, worker_error=True) - raise HTTPException(status_code=503, detail={"message": "all workers failed", "errors": errors}) + raise HTTPException( + status_code=503, + detail={"message": "all workers failed", "errors": errors}, + ) From 0d265130c6a3faa4e52544aa0620e742b4565abb Mon Sep 17 00:00:00 2001 From: Prakhar Verma Date: Sun, 10 May 2026 20:06:33 -0400 Subject: [PATCH 2/2] Clean up benchmark and eval scripts --- bench/run_scenarios.py | 172 +++++++++++++++++++------------- training/evaluate_text_model.py | 130 ++++++++++++++---------- 2 files changed, 182 insertions(+), 120 deletions(-) diff --git a/bench/run_scenarios.py b/bench/run_scenarios.py index 783cbb7..fe931b8 100644 --- a/bench/run_scenarios.py +++ b/bench/run_scenarios.py @@ -122,15 +122,65 @@ def service_env(extra: dict[str, str]) -> dict[str, str]: return env -def start_process(args: list[str], env: dict[str, str], log_path: Path) -> subprocess.Popen: - log_file = log_path.open("w") - return subprocess.Popen( - args, - cwd=ROOT, - env=env, - stdout=log_file, - stderr=subprocess.STDOUT, - text=True, +def service_url(port: int) -> str: + return f"http://localhost:{port}" + + +def uvicorn_args(module: str, port: int) -> list[str]: + return [ + sys.executable, + "-m", + "uvicorn", + module, + "--host", + "127.0.0.1", + "--port", + str(port), + ] + + +def start_process( + args: list[str], + env: dict[str, str], + log_path: Path, +) -> subprocess.Popen: + with log_path.open("w") as log_file: + return subprocess.Popen( + args, + cwd=ROOT, + env=env, + stdout=log_file, + stderr=subprocess.STDOUT, + text=True, + ) + + +def start_worker(worker: WorkerConfig, scenario: Scenario, log_dir: Path) -> subprocess.Popen: + return start_process( + uvicorn_args("worker.app.main:app", worker.port), + service_env({ + "WORKER_ID": worker.name, + "SIMULATE_DELAY_MS": str(worker.delay_ms), + "FAIL_RATE": str(worker.fail_rate), + }), + log_dir / f"{scenario.name}-{worker.name}.log", + ) + + +def start_gateway( + scenario: Scenario, + worker_urls: list[str], + log_dir: Path, +) -> subprocess.Popen: + return start_process( + uvicorn_args("gateway.app.main:app", BASE_PORT), + service_env({ + "WORKER_URLS": ",".join(worker_urls), + "ROUTER_STRATEGY": scenario.router_strategy, + "CACHE_ENABLED": str(scenario.cache_enabled).lower(), + "REQUEST_TIMEOUT_SECONDS": "1.0", + }), + log_dir / f"{scenario.name}-gateway.log", ) @@ -143,7 +193,8 @@ async def wait_for_health(url: str, timeout_seconds: float = 8.0) -> None: if response.status_code == 200: return except httpx.HTTPError: - await asyncio.sleep(0.1) + pass + await asyncio.sleep(0.1) raise RuntimeError(f"service did not become healthy: {url}") @@ -158,7 +209,12 @@ def stop_processes(processes: list[subprocess.Popen]) -> None: process.kill() -async def call(client: httpx.AsyncClient, url: str, index: int, unique_inputs: bool) -> dict[str, Any]: +async def send_prediction( + client: httpx.AsyncClient, + url: str, + index: int, + unique_inputs: bool, +) -> dict[str, Any]: text = SAMPLES[index % len(SAMPLES)] if unique_inputs: text = f"{text} sample {index}" @@ -177,24 +233,24 @@ async def call(client: httpx.AsyncClient, url: str, index: int, unique_inputs: b } -async def run_load(url: str, scenario: Scenario) -> dict[str, Any]: - semaphore = asyncio.Semaphore(scenario.concurrency) - async with httpx.AsyncClient(timeout=5.0) as client: - async def limited(index: int) -> dict[str, Any]: - async with semaphore: - return await call(client, url, index, scenario.unique_inputs) - - started = time.perf_counter() - rows = await asyncio.gather(*(limited(index) for index in range(scenario.requests))) - total_ms = (time.perf_counter() - started) * 1000 +def p95_latency(latencies: list[float]) -> float: + if not latencies: + return 0.0 + if len(latencies) < 20: + return max(latencies) + return statistics.quantiles(latencies, n=20)[18] - gateway_metrics = (await client.get(url.replace("/predict", "/metrics"))).json() - worker_snapshot = (await client.get(url.replace("/predict", "/workers"))).json() +def summarize_load( + scenario: Scenario, + rows: list[dict[str, Any]], + total_ms: float, + gateway_metrics: dict[str, Any], + worker_snapshot: list[dict[str, Any]], +) -> dict[str, Any]: latencies = [row["latency_ms"] for row in rows if row["ok"]] failures = sum(1 for row in rows if not row["ok"]) cache_hits = sum(1 for row in rows if row["ok"] and row["cached"]) - p95 = statistics.quantiles(latencies, n=20)[18] if len(latencies) >= 20 else max(latencies, default=0) return { "requests": scenario.requests, @@ -206,64 +262,44 @@ async def limited(index: int) -> dict[str, Any]: "total_ms": round(total_ms, 2), "requests_per_second": round(scenario.requests / (total_ms / 1000), 2), "p50_latency_ms": round(statistics.median(latencies), 2) if latencies else 0, - "p95_latency_ms": round(p95, 2), + "p95_latency_ms": round(p95_latency(latencies), 2), "gateway_metrics": gateway_metrics, "workers": worker_snapshot, } +async def run_load(url: str, scenario: Scenario) -> dict[str, Any]: + semaphore = asyncio.Semaphore(scenario.concurrency) + async with httpx.AsyncClient(timeout=5.0) as client: + async def limited(index: int) -> dict[str, Any]: + async with semaphore: + return await send_prediction(client, url, index, scenario.unique_inputs) + + started = time.perf_counter() + rows = await asyncio.gather(*(limited(index) for index in range(scenario.requests))) + total_ms = (time.perf_counter() - started) * 1000 + + gateway_metrics = (await client.get(url.replace("/predict", "/metrics"))).json() + worker_snapshot = (await client.get(url.replace("/predict", "/workers"))).json() + + return summarize_load(scenario, rows, total_ms, gateway_metrics, worker_snapshot) + + async def run_scenario(scenario: Scenario, log_dir: Path) -> dict[str, Any]: processes: list[subprocess.Popen] = [] - gateway_port = BASE_PORT - worker_urls = [] + worker_urls = [service_url(worker.port) for worker in scenario.workers] try: for worker in scenario.workers: - worker_urls.append(f"http://localhost:{worker.port}") - processes.append(start_process( - [ - sys.executable, - "-m", - "uvicorn", - "worker.app.main:app", - "--host", - "127.0.0.1", - "--port", - str(worker.port), - ], - service_env({ - "WORKER_ID": worker.name, - "SIMULATE_DELAY_MS": str(worker.delay_ms), - "FAIL_RATE": str(worker.fail_rate), - }), - log_dir / f"{scenario.name}-{worker.name}.log", - )) + processes.append(start_worker(worker, scenario, log_dir)) for url in worker_urls: await wait_for_health(f"{url}/health") - processes.append(start_process( - [ - sys.executable, - "-m", - "uvicorn", - "gateway.app.main:app", - "--host", - "127.0.0.1", - "--port", - str(gateway_port), - ], - service_env({ - "WORKER_URLS": ",".join(worker_urls), - "ROUTER_STRATEGY": scenario.router_strategy, - "CACHE_ENABLED": "true" if scenario.cache_enabled else "false", - "REQUEST_TIMEOUT_SECONDS": "1.0", - }), - log_dir / f"{scenario.name}-gateway.log", - )) - await wait_for_health(f"http://localhost:{gateway_port}/health") - - result = await run_load(f"http://localhost:{gateway_port}/predict", scenario) + processes.append(start_gateway(scenario, worker_urls, log_dir)) + await wait_for_health(f"{service_url(BASE_PORT)}/health") + + result = await run_load(f"{service_url(BASE_PORT)}/predict", scenario) return { "name": scenario.name, "description": scenario.description, diff --git a/training/evaluate_text_model.py b/training/evaluate_text_model.py index af1cfb6..c038edc 100644 --- a/training/evaluate_text_model.py +++ b/training/evaluate_text_model.py @@ -3,32 +3,43 @@ import argparse import json import sys +from dataclasses import dataclass from pathlib import Path +from typing import Any ROOT = Path(__file__).resolve().parents[1] -sys.path.insert(0, str(ROOT)) +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) OUTPUT_JSON = ROOT / "training" / "evaluation" / "latest.json" OUTPUT_MD = ROOT / "training" / "evaluation" / "latest.md" +LABELS = ("healthy", "risky") + + +@dataclass(frozen=True) +class ChallengeExample: + expected: str + text: str + CHALLENGE_SET = [ - ("healthy", "fast stable reliable service"), - ("risky", "slow broken timeout error"), - ("healthy", "worker recovered after timeout and is stable"), - ("risky", "cache hit but worker failed with timeout"), - ("risky", "great latency but retry caused timeout"), - ("healthy", "responsive route completed request"), - ("risky", "unhealthy worker returned fast response"), - ("healthy", "stable gateway handled burst traffic"), - ("risky", "prediction failed after worker crash"), - ("risky", "service is healthy but slow"), - ("healthy", "reliable cache and responsive worker"), - ("risky", "broken route recovered with retry"), - ("healthy", "cache recovered and request completed"), - ("risky", "gateway retry loop after timeout"), - ("healthy", "responsive worker handled traffic burst"), - ("risky", "mesh returned broken prediction error"), + ChallengeExample("healthy", "fast stable reliable service"), + ChallengeExample("risky", "slow broken timeout error"), + ChallengeExample("healthy", "worker recovered after timeout and is stable"), + ChallengeExample("risky", "cache hit but worker failed with timeout"), + ChallengeExample("risky", "great latency but retry caused timeout"), + ChallengeExample("healthy", "responsive route completed request"), + ChallengeExample("risky", "unhealthy worker returned fast response"), + ChallengeExample("healthy", "stable gateway handled burst traffic"), + ChallengeExample("risky", "prediction failed after worker crash"), + ChallengeExample("risky", "service is healthy but slow"), + ChallengeExample("healthy", "reliable cache and responsive worker"), + ChallengeExample("risky", "broken route recovered with retry"), + ChallengeExample("healthy", "cache recovered and request completed"), + ChallengeExample("risky", "gateway retry loop after timeout"), + ChallengeExample("healthy", "responsive worker handled traffic burst"), + ChallengeExample("risky", "mesh returned broken prediction error"), ] @@ -40,48 +51,60 @@ def parse_args() -> argparse.Namespace: return parser.parse_args() -def evaluate(model_artifact: Path | None = None) -> dict: - from worker.app.model import load_model +def empty_confusion_matrix() -> dict[str, dict[str, int]]: + return {expected: {predicted: 0 for predicted in LABELS} for expected in LABELS} - model = load_model(str(model_artifact) if model_artifact else None) - rows = [] - labels = ["healthy", "risky"] - confusion = { - "healthy": {"healthy": 0, "risky": 0}, - "risky": {"healthy": 0, "risky": 0}, + +def score_example(model: Any, example: ChallengeExample) -> dict[str, Any]: + prediction = model.predict(example.text) + predicted = str(prediction["label"]) + return { + "text": example.text, + "expected": example.expected, + "predicted": predicted, + "score": prediction["score"], + "correct": example.expected == predicted, } - for expected, text in CHALLENGE_SET: - prediction = model.predict(text) - predicted = str(prediction["label"]) - confusion[expected][predicted] += 1 - rows.append({ - "text": text, - "expected": expected, - "predicted": predicted, - "score": prediction["score"], - "correct": expected == predicted, - }) - total = len(rows) - correct = sum(1 for row in rows if row["correct"]) - per_label = {} - for label in labels: +def build_confusion_matrix(rows: list[dict[str, Any]]) -> dict[str, dict[str, int]]: + confusion = empty_confusion_matrix() + for row in rows: + confusion[row["expected"]][row["predicted"]] += 1 + return confusion + + +def safe_rate(numerator: int, denominator: int) -> float: + return round(numerator / denominator, 4) if denominator else 0 + + +def label_metrics(confusion: dict[str, dict[str, int]]) -> dict[str, dict[str, float]]: + per_label: dict[str, dict[str, float]] = {} + for label in LABELS: true_positive = confusion[label][label] false_positive = sum( - confusion[other][label] for other in labels if other != label + confusion[other][label] for other in LABELS if other != label ) false_negative = sum( - confusion[label][other] for other in labels if other != label + confusion[label][other] for other in LABELS if other != label ) precision_base = true_positive + false_positive recall_base = true_positive + false_negative - precision = true_positive / precision_base if precision_base else 0 - recall = true_positive / recall_base if recall_base else 0 per_label[label] = { - "precision": round(precision, 4), - "recall": round(recall, 4), + "precision": safe_rate(true_positive, precision_base), + "recall": safe_rate(true_positive, recall_base), } + return per_label + + +def evaluate(model_artifact: Path | None = None) -> dict[str, Any]: + from worker.app.model import load_model + + model = load_model(str(model_artifact) if model_artifact else None) + rows = [score_example(model, example) for example in CHALLENGE_SET] + confusion = build_confusion_matrix(rows) + total = len(rows) + correct = sum(1 for row in rows if row["correct"]) return { "model_version": model.version, @@ -90,15 +113,12 @@ def evaluate(model_artifact: Path | None = None) -> dict: "accuracy": round(correct / total, 4), "correct": correct, "confusion_matrix": confusion, - "per_label": per_label, + "per_label": label_metrics(confusion), "rows": rows, } -def write_outputs(result: dict, json_path: Path, md_path: Path) -> None: - json_path.parent.mkdir(parents=True, exist_ok=True) - json_path.write_text(json.dumps(result, indent=2) + "\n") - +def render_markdown(result: dict[str, Any]) -> str: lines = [ "# Model evaluation", "", @@ -149,7 +169,13 @@ def write_outputs(result: dict, json_path: Path, md_path: Path) -> None: ) lines.append("") - md_path.write_text("\n".join(lines)) + return "\n".join(lines) + + +def write_outputs(result: dict[str, Any], json_path: Path, md_path: Path) -> None: + json_path.parent.mkdir(parents=True, exist_ok=True) + json_path.write_text(json.dumps(result, indent=2) + "\n") + md_path.write_text(render_markdown(result)) def main() -> None: