diff --git a/README.md b/README.md index 0a7d1e1..041cda7 100644 --- a/README.md +++ b/README.md @@ -30,6 +30,12 @@ Run a quick local benchmark: python bench/run_bench.py --requests 100 --concurrency 10 ``` +Open the local dashboard: + +```text +http://localhost:8000/dashboard +``` + ## Docker ```bash @@ -38,6 +44,26 @@ docker compose up --build That starts Redis, one gateway, and two workers. +## Benchmarks + +The scenario runner starts local gateway and worker processes, sends load, then writes JSON and markdown results. + +```bash +python bench/run_scenarios.py +``` + +Latest local run: + +| Scenario | RPS | p50 ms | p95 ms | Cache hits | Failures | +| --- | ---: | ---: | ---: | ---: | ---: | +| cache enabled | 211.37 | 5.54 | 312.52 | 71 | 0 | +| cache disabled | 24.92 | 318.85 | 326.38 | 0 | 0 | +| slow worker, round-robin | 25.68 | 67.3 | 1033.24 | 0 | 0 | +| slow worker, load-aware | 128.51 | 4.94 | 425.03 | 0 | 0 | +| dead worker recovery | 698.97 | 11.15 | 28.1 | 0 | 0 | + +Full output is in `bench/results/latest.md` and `bench/results/latest.json`. + ## Training The worker can load a small trained text model artifact from `worker/app/artifacts/tiny_text_model.json`. @@ -57,7 +83,8 @@ If it is using the Mac GPU, the output includes `'device': 'mps'`. - Worker service with a tiny text classifier and optional trained artifact - Redis-backed cache when Redis is available, with an in-memory fallback - Retry on failed worker calls and basic load-aware routing -- Simple benchmark script +- Scenario benchmark runner with saved results +- Local dashboard for gateway and worker state - Unit tests for routing, caching, metrics, and the model This is still intentionally small. The useful part is the system behavior around the model, not the model itself. diff --git a/bench/results/latest.json b/bench/results/latest.json new file mode 100644 index 0000000..900d514 --- /dev/null +++ b/bench/results/latest.json @@ -0,0 +1,212 @@ +{ + "generated_by": "python bench/run_scenarios.py", + "results": [ + { + "name": "cache_enabled", + "description": "One delayed worker, four repeated inputs, cache enabled.", + "router_strategy": "load_aware", + "cache_enabled": true, + "requests": 80, + "concurrency": 8, + "successes": 80, + "failures": 0, + "cache_hits": 71, + "cache_hit_rate": 0.8875, + "total_ms": 378.48, + "requests_per_second": 211.37, + "p50_latency_ms": 5.54, + "p95_latency_ms": 312.52, + "gateway_metrics": { + "requests": 80, + "cache_hits": 71, + "worker_errors": 0, + "cache_hit_rate": 0.8875, + "requests_per_second": 159.35, + "p50_latency_ms": 0.01, + "p95_latency_ms": 294.25 + }, + "workers": [ + { + "url": "http://localhost:8311", + "in_flight": 0, + "requests": 9, + "failures": 0, + "latency_ema_ms": 268.22, + "last_error": null, + "strategy": "load_aware" + } + ] + }, + { + "name": "cache_disabled", + "description": "Same traffic as cache_enabled, but cache disabled.", + "router_strategy": "load_aware", + "cache_enabled": false, + "requests": 80, + "concurrency": 8, + "successes": 80, + "failures": 0, + "cache_hits": 0, + "cache_hit_rate": 0.0, + "total_ms": 3210.55, + "requests_per_second": 24.92, + "p50_latency_ms": 318.85, + "p95_latency_ms": 326.38, + "gateway_metrics": { + "requests": 80, + "cache_hits": 0, + "worker_errors": 0, + "cache_hit_rate": 0.0, + "requests_per_second": 23.93, + "p50_latency_ms": 309.16, + "p95_latency_ms": 316.54 + }, + "workers": [ + { + "url": "http://localhost:8311", + "in_flight": 0, + "requests": 80, + "failures": 0, + "latency_ema_ms": 294.34, + "last_error": null, + "strategy": "load_aware" + } + ] + }, + { + "name": "slow_worker_round_robin", + "description": "One fast worker and one slow worker with round-robin routing.", + "router_strategy": "round_robin", + "cache_enabled": false, + "requests": 80, + "concurrency": 12, + "successes": 80, + "failures": 0, + "cache_hits": 0, + "cache_hit_rate": 0.0, + "total_ms": 3115.31, + "requests_per_second": 25.68, + "p50_latency_ms": 67.3, + "p95_latency_ms": 1033.24, + "gateway_metrics": { + "requests": 80, + "cache_hits": 0, + "worker_errors": 0, + "cache_hit_rate": 0.0, + "requests_per_second": 24.67, + "p50_latency_ms": 108.0, + "p95_latency_ms": 1029.4 + }, + "workers": [ + { + "url": "http://localhost:8311", + "in_flight": 0, + "requests": 57, + "failures": 0, + "latency_ema_ms": 11.48, + "last_error": null, + "strategy": "round_robin" + }, + { + "url": "http://localhost:8312", + "in_flight": 0, + "requests": 40, + "failures": 17, + "latency_ema_ms": 777.5, + "last_error": "ReadTimeout", + "strategy": "round_robin" + } + ] + }, + { + "name": "slow_worker_load_aware", + "description": "Same workers as slow_worker_round_robin, using load-aware routing.", + "router_strategy": "load_aware", + "cache_enabled": false, + "requests": 80, + "concurrency": 12, + "successes": 80, + "failures": 0, + "cache_hits": 0, + "cache_hit_rate": 0.0, + "total_ms": 622.53, + "requests_per_second": 128.51, + "p50_latency_ms": 4.94, + "p95_latency_ms": 425.03, + "gateway_metrics": { + "requests": 80, + "cache_hits": 0, + "worker_errors": 0, + "cache_hit_rate": 0.0, + "requests_per_second": 109.34, + "p50_latency_ms": 3.61, + "p95_latency_ms": 408.65 + }, + "workers": [ + { + "url": "http://localhost:8311", + "in_flight": 0, + "requests": 74, + "failures": 0, + "latency_ema_ms": 3.82, + "last_error": null, + "strategy": "load_aware" + }, + { + "url": "http://localhost:8312", + "in_flight": 0, + "requests": 6, + "failures": 0, + "latency_ema_ms": 365.08, + "last_error": null, + "strategy": "load_aware" + } + ] + }, + { + "name": "dead_worker_recovery", + "description": "One worker always fails, one worker stays healthy.", + "router_strategy": "load_aware", + "cache_enabled": false, + "requests": 60, + "concurrency": 10, + "successes": 60, + "failures": 0, + "cache_hits": 0, + "cache_hit_rate": 0.0, + "total_ms": 85.84, + "requests_per_second": 698.97, + "p50_latency_ms": 11.15, + "p95_latency_ms": 28.1, + "gateway_metrics": { + "requests": 60, + "cache_hits": 0, + "worker_errors": 0, + "cache_hit_rate": 0.0, + "requests_per_second": 275.53, + "p50_latency_ms": 8.53, + "p95_latency_ms": 19.94 + }, + "workers": [ + { + "url": "http://localhost:8311", + "in_flight": 0, + "requests": 5, + "failures": 5, + "latency_ema_ms": 0.0, + "last_error": "HTTPStatusError", + "strategy": "load_aware" + }, + { + "url": "http://localhost:8312", + "in_flight": 0, + "requests": 60, + "failures": 0, + "latency_ema_ms": 8.83, + "last_error": null, + "strategy": "load_aware" + } + ] + } + ] +} diff --git a/bench/results/latest.md b/bench/results/latest.md new file mode 100644 index 0000000..bf46122 --- /dev/null +++ b/bench/results/latest.md @@ -0,0 +1,17 @@ +# Benchmark results + +Generated with `python bench/run_scenarios.py` on the local machine. + +| Scenario | Router | Cache | Requests | RPS | p50 ms | p95 ms | Cache hits | Failures | +| --- | --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | +| cache_enabled | load_aware | True | 80 | 211.37 | 5.54 | 312.52 | 71 | 0 | +| cache_disabled | load_aware | False | 80 | 24.92 | 318.85 | 326.38 | 0 | 0 | +| slow_worker_round_robin | round_robin | False | 80 | 25.68 | 67.3 | 1033.24 | 0 | 0 | +| slow_worker_load_aware | load_aware | False | 80 | 128.51 | 4.94 | 425.03 | 0 | 0 | +| dead_worker_recovery | load_aware | False | 60 | 698.97 | 11.15 | 28.1 | 0 | 0 | + +Notes: + +- Cache scenarios repeat four text inputs to show cache behavior. +- Slow-worker scenarios use one normal worker and one worker with 95 ms artificial delay. +- Dead-worker recovery uses one worker with a 100 percent simulated failure rate. diff --git a/bench/run_bench.py b/bench/run_bench.py index d00acb2..6a76386 100644 --- a/bench/run_bench.py +++ b/bench/run_bench.py @@ -2,6 +2,7 @@ import argparse import asyncio +import json import statistics import time @@ -21,14 +22,21 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--url", default="http://localhost:8000/predict") parser.add_argument("--requests", type=int, default=100) parser.add_argument("--concurrency", type=int, default=10) + parser.add_argument("--unique-requests", action="store_true") return parser.parse_args() -async def call(client: httpx.AsyncClient, url: str, index: int) -> tuple[float, bool]: +async def call(client: httpx.AsyncClient, url: str, index: int, *, unique_requests: bool) -> tuple[float, bool, bool]: started = time.perf_counter() - response = await client.post(url, json={"text": SAMPLES[index % len(SAMPLES)], "request_id": str(index)}) + text = SAMPLES[index % len(SAMPLES)] + if unique_requests: + text = f"{text} {index}" + response = await client.post(url, json={"text": text, "request_id": str(index)}) latency = (time.perf_counter() - started) * 1000 - return latency, response.status_code == 200 + cached = False + if response.status_code == 200: + cached = bool(response.json().get("cached")) + return latency, response.status_code == 200, cached async def main() -> None: @@ -36,27 +44,30 @@ async def main() -> None: semaphore = asyncio.Semaphore(args.concurrency) async with httpx.AsyncClient(timeout=5) as client: - async def limited(index: int) -> tuple[float, bool]: + async def limited(index: int) -> tuple[float, bool, bool]: async with semaphore: - return await call(client, args.url, index) + return await call(client, args.url, index, unique_requests=args.unique_requests) started = time.perf_counter() results = await asyncio.gather(*(limited(index) for index in range(args.requests))) total_ms = (time.perf_counter() - started) * 1000 - latencies = [latency for latency, ok in results if ok] - failures = sum(1 for _, ok in results if not ok) + latencies = [latency for latency, ok, _ in results if ok] + failures = sum(1 for _, ok, _ in results if not ok) + cache_hits = sum(1 for _, ok, cached in results if ok and cached) p95 = statistics.quantiles(latencies, n=20)[18] if len(latencies) >= 20 else max(latencies, default=0) - print({ + print(json.dumps({ "requests": args.requests, "successes": len(latencies), "failures": failures, + "cache_hits": cache_hits, + "cache_hit_rate": round(cache_hits / len(latencies), 4) if latencies else 0, "total_ms": round(total_ms, 2), "requests_per_second": round(args.requests / (total_ms / 1000), 2), "p50_latency_ms": round(statistics.median(latencies), 2) if latencies else 0, "p95_latency_ms": round(p95, 2), - }) + }, indent=2)) if __name__ == "__main__": diff --git a/bench/run_scenarios.py b/bench/run_scenarios.py new file mode 100644 index 0000000..783cbb7 --- /dev/null +++ b/bench/run_scenarios.py @@ -0,0 +1,327 @@ +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import signal +import statistics +import subprocess +import sys +import time +from dataclasses import dataclass +from pathlib import Path +from tempfile import TemporaryDirectory +from typing import Any + +import httpx + + +ROOT = Path(__file__).resolve().parents[1] +RESULTS_DIR = ROOT / "bench" / "results" +BASE_PORT = 8300 + +SAMPLES = [ + "fast stable model service", + "worker timeout and broken response", + "great reliable inference path", + "slow crash error", +] + + +@dataclass(frozen=True) +class WorkerConfig: + name: str + port: int + delay_ms: int = 0 + fail_rate: float = 0.0 + + +@dataclass(frozen=True) +class Scenario: + name: str + description: str + workers: list[WorkerConfig] + router_strategy: str = "load_aware" + cache_enabled: bool = True + requests: int = 80 + concurrency: int = 10 + unique_inputs: bool = False + + +SCENARIOS = [ + Scenario( + name="cache_enabled", + description="One delayed worker, four repeated inputs, cache enabled.", + workers=[WorkerConfig("worker-a", BASE_PORT + 11, delay_ms=35)], + cache_enabled=True, + requests=80, + concurrency=8, + ), + Scenario( + name="cache_disabled", + description="Same traffic as cache_enabled, but cache disabled.", + workers=[WorkerConfig("worker-a", BASE_PORT + 11, delay_ms=35)], + cache_enabled=False, + requests=80, + concurrency=8, + ), + Scenario( + name="slow_worker_round_robin", + description="One fast worker and one slow worker with round-robin routing.", + workers=[ + WorkerConfig("worker-fast", BASE_PORT + 11), + WorkerConfig("worker-slow", BASE_PORT + 12, delay_ms=95), + ], + router_strategy="round_robin", + cache_enabled=False, + requests=80, + concurrency=12, + unique_inputs=True, + ), + Scenario( + name="slow_worker_load_aware", + description="Same workers as slow_worker_round_robin, using load-aware routing.", + workers=[ + WorkerConfig("worker-fast", BASE_PORT + 11), + WorkerConfig("worker-slow", BASE_PORT + 12, delay_ms=95), + ], + router_strategy="load_aware", + cache_enabled=False, + requests=80, + concurrency=12, + unique_inputs=True, + ), + Scenario( + name="dead_worker_recovery", + description="One worker always fails, one worker stays healthy.", + workers=[ + WorkerConfig("worker-dead", BASE_PORT + 11, fail_rate=1.0), + WorkerConfig("worker-healthy", BASE_PORT + 12), + ], + router_strategy="load_aware", + cache_enabled=False, + requests=60, + concurrency=10, + unique_inputs=True, + ), +] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--output-json", type=Path, default=RESULTS_DIR / "latest.json") + parser.add_argument("--output-md", type=Path, default=RESULTS_DIR / "latest.md") + return parser.parse_args() + + +def service_env(extra: dict[str, str]) -> dict[str, str]: + env = os.environ.copy() + env.update(extra) + env["PYTHONPATH"] = str(ROOT) + 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, + ) + + +async def wait_for_health(url: str, timeout_seconds: float = 8.0) -> None: + deadline = time.time() + timeout_seconds + async with httpx.AsyncClient(timeout=1.0) as client: + while time.time() < deadline: + try: + response = await client.get(url) + if response.status_code == 200: + return + except httpx.HTTPError: + await asyncio.sleep(0.1) + raise RuntimeError(f"service did not become healthy: {url}") + + +def stop_processes(processes: list[subprocess.Popen]) -> None: + for process in processes: + if process.poll() is None: + process.send_signal(signal.SIGTERM) + for process in processes: + try: + process.wait(timeout=3) + except subprocess.TimeoutExpired: + process.kill() + + +async def call(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}" + started = time.perf_counter() + response = await client.post(url, json={"text": text, "request_id": str(index)}) + latency_ms = (time.perf_counter() - started) * 1000 + body: dict[str, Any] = {} + if response.headers.get("content-type", "").startswith("application/json"): + body = response.json() + return { + "ok": response.status_code == 200, + "status_code": response.status_code, + "latency_ms": latency_ms, + "cached": bool(body.get("cached")), + "worker_url": body.get("worker_url"), + } + + +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 + + gateway_metrics = (await client.get(url.replace("/predict", "/metrics"))).json() + worker_snapshot = (await client.get(url.replace("/predict", "/workers"))).json() + + 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, + "concurrency": scenario.concurrency, + "successes": len(latencies), + "failures": failures, + "cache_hits": cache_hits, + "cache_hit_rate": round(cache_hits / len(latencies), 4) if latencies else 0, + "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), + "gateway_metrics": gateway_metrics, + "workers": worker_snapshot, + } + + +async def run_scenario(scenario: Scenario, log_dir: Path) -> dict[str, Any]: + processes: list[subprocess.Popen] = [] + gateway_port = BASE_PORT + worker_urls = [] + + 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", + )) + + 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) + return { + "name": scenario.name, + "description": scenario.description, + "router_strategy": scenario.router_strategy, + "cache_enabled": scenario.cache_enabled, + **result, + } + finally: + stop_processes(processes) + + +def write_markdown(path: Path, results: list[dict[str, Any]]) -> None: + lines = [ + "# Benchmark results", + "", + "Generated with `python bench/run_scenarios.py` on the local machine.", + "", + "| Scenario | Router | Cache | Requests | RPS | p50 ms | p95 ms | Cache hits | Failures |", + "| --- | --- | --- | ---: | ---: | ---: | ---: | ---: | ---: |", + ] + for row in results: + lines.append( + f"| {row['name']} | {row['router_strategy']} | {row['cache_enabled']} | " + f"{row['requests']} | {row['requests_per_second']} | {row['p50_latency_ms']} | " + f"{row['p95_latency_ms']} | {row['cache_hits']} | {row['failures']} |" + ) + + lines.extend([ + "", + "Notes:", + "", + "- Cache scenarios repeat four text inputs to show cache behavior.", + "- Slow-worker scenarios use one normal worker and one worker with 95 ms artificial delay.", + "- Dead-worker recovery uses one worker with a 100 percent simulated failure rate.", + "", + ]) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("\n".join(lines)) + + +async def main() -> None: + args = parse_args() + with TemporaryDirectory(prefix="modelmesh-bench-") as directory: + log_dir = Path(directory) + results = [] + for scenario in SCENARIOS: + print(f"running {scenario.name}") + results.append(await run_scenario(scenario, log_dir)) + + payload = { + "generated_by": "python bench/run_scenarios.py", + "results": results, + } + args.output_json.parent.mkdir(parents=True, exist_ok=True) + args.output_json.write_text(json.dumps(payload, indent=2) + "\n") + write_markdown(args.output_md, results) + print(json.dumps(payload, indent=2)) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/docker-compose.yml b/docker-compose.yml index 3d1b1ba..6a79d1d 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -31,6 +31,8 @@ services: WORKER_URLS: http://worker-a:9000,http://worker-b:9000 REDIS_URL: redis://redis:6379/0 REQUEST_TIMEOUT_SECONDS: "1.5" + ROUTER_STRATEGY: load_aware + CACHE_ENABLED: "true" ports: - "8000:8000" depends_on: diff --git a/docs/architecture.md b/docs/architecture.md new file mode 100644 index 0000000..2552661 --- /dev/null +++ b/docs/architecture.md @@ -0,0 +1,37 @@ +# Architecture notes + +ModelMesh is an ML systems project. The model is intentionally small so the repo can focus on serving behavior: routing, caching, retries, observability, and benchmark evidence. + +```mermaid +flowchart LR + client["Client"] --> gateway["FastAPI gateway"] + gateway --> cache["Redis or memory cache"] + gateway --> worker_a["Worker A"] + gateway --> worker_b["Worker B"] + worker_a --> model_a["Text model"] + worker_b --> model_b["Text model"] + gateway --> metrics["Metrics and dashboard"] +``` + +## Request path + +1. The client sends text to `POST /predict`. +2. The gateway checks the prediction cache using the input text as the cache identity. +3. On a miss, the gateway asks the worker pool for candidates. +4. The selected worker returns a prediction with label, score, model version, and worker id. +5. The gateway records latency, cache hits, worker errors, and worker health. + +## Routing + +The gateway supports two routing modes: + +- `round_robin` for a simple baseline. +- `load_aware` for routing based on failures, in-flight requests, and recent latency. + +The benchmark runner uses both modes against the same slow-worker scenario so the difference is measurable. + +## Training + +`training/train_text_model.py` trains the tiny text model with PyTorch. On Apple Silicon it uses `mps`, and it exits instead of silently falling back to CPU unless `--allow-cpu` is passed. + +The trained artifact is plain JSON so the worker can load it without depending on PyTorch at inference time. diff --git a/gateway/app/cache.py b/gateway/app/cache.py index d75eed7..1518a0f 100644 --- a/gateway/app/cache.py +++ b/gateway/app/cache.py @@ -9,12 +9,13 @@ class PredictionCache: - def __init__(self, redis_url: str | None, ttl_seconds: int = 300) -> None: + def __init__(self, redis_url: str | None, ttl_seconds: int = 300, enabled: bool = True) -> None: self.ttl_seconds = ttl_seconds + self.enabled = enabled self.memory: dict[str, dict[str, Any]] = {} self.redis: Redis | None = None - if redis_url: + if enabled and redis_url: try: client = Redis.from_url(redis_url, decode_responses=True) client.ping() @@ -22,11 +23,13 @@ def __init__(self, redis_url: str | None, ttl_seconds: int = 300) -> None: except RedisError: self.redis = None - def key_for(self, payload: dict[str, Any]) -> str: - normalized = json.dumps(payload, sort_keys=True, separators=(",", ":")) + def key_for(self, *, text: str) -> str: + normalized = json.dumps({"text": text}, sort_keys=True, separators=(",", ":")) return "prediction:" + hashlib.sha256(normalized.encode("utf-8")).hexdigest() def get(self, key: str) -> dict[str, Any] | None: + if not self.enabled: + return None if self.redis is not None: try: value = self.redis.get(key) @@ -36,6 +39,8 @@ def get(self, key: str) -> dict[str, Any] | None: return self.memory.get(key) def set(self, key: str, value: dict[str, Any]) -> None: + if not self.enabled: + return if self.redis is not None: try: self.redis.setex(key, self.ttl_seconds, json.dumps(value)) diff --git a/gateway/app/dashboard.py b/gateway/app/dashboard.py new file mode 100644 index 0000000..47db5d1 --- /dev/null +++ b/gateway/app/dashboard.py @@ -0,0 +1,232 @@ +from __future__ import annotations + + +DASHBOARD_HTML = """ + + + + + + ModelMesh + + + +
+
+
+

ModelMesh

+

Gateway metrics and worker state for the local inference mesh.

+
+
loading
+
+ +
+
+
Requests
+
0
+
+
+
Cache hit rate
+
0%
+
+
+
p95 latency
+
0 ms
+
+
+
Worker errors
+
0
+
+
+ +
+

Workers

+ + + + + + + + + + + + +
URLStrategyIn flightRequestsFailuresLatency EMA
+
+
+ + + + +""" diff --git a/gateway/app/main.py b/gateway/app/main.py index 8fb4a95..60f1d0d 100644 --- a/gateway/app/main.py +++ b/gateway/app/main.py @@ -7,13 +7,22 @@ import httpx from fastapi import FastAPI, HTTPException +from fastapi.responses import HTMLResponse from .cache import PredictionCache +from .dashboard import DASHBOARD_HTML from .metrics import MetricsStore from .router import WorkerPool from .schemas import GatewayPrediction, PredictRequest +def env_flag(name: str, default: bool = True) -> bool: + value = os.getenv(name) + if value is None: + return default + return value.lower() not in {"0", "false", "no", "off"} + + def parse_worker_urls(raw: str | None) -> list[str]: if not raw: return ["http://localhost:8011"] @@ -28,15 +37,23 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: app = FastAPI(title="ModelMesh Gateway", version="0.1.0", lifespan=lifespan) -router = WorkerPool(parse_worker_urls(os.getenv("WORKER_URLS"))) -cache = PredictionCache(os.getenv("REDIS_URL")) +router = WorkerPool( + 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)) metrics = MetricsStore() request_timeout = float(os.getenv("REQUEST_TIMEOUT_SECONDS", "1.0")) @app.get("/health") async def health() -> dict[str, str | int]: - return {"status": "ok", "workers": len(router.workers)} + return {"status": "ok", "workers": len(router.workers), "router_strategy": router.strategy} + + +@app.get("/dashboard", response_class=HTMLResponse) +async def dashboard() -> str: + return DASHBOARD_HTML @app.get("/workers") @@ -52,7 +69,7 @@ async def metrics_snapshot() -> dict[str, float | int]: @app.post("/predict", response_model=GatewayPrediction) async def predict(payload: PredictRequest) -> GatewayPrediction: started = time.perf_counter() - cache_key = cache.key_for(payload.model_dump()) + cache_key = cache.key_for(text=payload.text) cached = cache.get(cache_key) if cached is not None: diff --git a/gateway/app/router.py b/gateway/app/router.py index 8068248..f23a683 100644 --- a/gateway/app/router.py +++ b/gateway/app/router.py @@ -16,14 +16,20 @@ class Worker: class WorkerPool: - def __init__(self, worker_urls: list[str]) -> None: + def __init__(self, worker_urls: list[str], strategy: str = "load_aware") -> None: if not worker_urls: raise ValueError("at least one worker URL is required") + if strategy not in {"load_aware", "round_robin"}: + raise ValueError("strategy must be load_aware or round_robin") self.workers = [Worker(url=url.rstrip("/")) for url in worker_urls] + self.strategy = strategy + self._round_robin_index = 0 def candidates(self, attempts: int) -> list[Worker]: count = min(max(attempts, 1), len(self.workers)) + if self.strategy == "round_robin": + return self._round_robin_candidates(count) ranked = sorted(self.workers, key=self._score) return ranked[:count] @@ -56,9 +62,17 @@ def snapshot(self) -> list[dict[str, str | int | float | None]]: "failures": worker.failures, "latency_ema_ms": round(worker.latency_ema_ms, 2), "last_error": worker.last_error, + "strategy": self.strategy, } for worker in self.workers ] def _score(self, worker: Worker) -> tuple[int, int, float, str]: return (worker.failures, worker.in_flight, worker.latency_ema_ms, worker.url) + + def _round_robin_candidates(self, count: int) -> list[Worker]: + selected = [] + for offset in range(count): + selected.append(self.workers[(self._round_robin_index + offset) % len(self.workers)]) + self._round_robin_index = (self._round_robin_index + 1) % len(self.workers) + return selected diff --git a/tests/test_cache.py b/tests/test_cache.py index 0abad5d..210cb27 100644 --- a/tests/test_cache.py +++ b/tests/test_cache.py @@ -3,7 +3,7 @@ def test_memory_cache_round_trip() -> None: cache = PredictionCache(redis_url=None) - key = cache.key_for({"text": "hello", "request_id": None}) + key = cache.key_for(text="hello") value = {"prediction": {"label": "healthy"}, "worker_url": "http://worker"} cache.set(key, value) @@ -11,10 +11,19 @@ def test_memory_cache_round_trip() -> None: assert cache.get(key) == value -def test_cache_key_is_stable_for_field_order() -> None: +def test_cache_key_ignores_request_metadata() -> None: cache = PredictionCache(redis_url=None) - first = cache.key_for({"text": "hello", "request_id": "1"}) - second = cache.key_for({"request_id": "1", "text": "hello"}) + first = cache.key_for(text="hello") + second = cache.key_for(text="hello") assert first == second + + +def test_disabled_cache_is_a_noop() -> None: + cache = PredictionCache(redis_url=None, enabled=False) + key = cache.key_for(text="hello") + + cache.set(key, {"prediction": {"label": "healthy"}}) + + assert cache.get(key) is None diff --git a/tests/test_dashboard.py b/tests/test_dashboard.py new file mode 100644 index 0000000..446f138 --- /dev/null +++ b/tests/test_dashboard.py @@ -0,0 +1,12 @@ +from fastapi.testclient import TestClient + +from gateway.app.main import app + + +def test_dashboard_serves_html() -> None: + with TestClient(app) as client: + response = client.get("/dashboard") + + assert response.status_code == 200 + assert "ModelMesh" in response.text + assert "Gateway metrics" in response.text diff --git a/tests/test_router.py b/tests/test_router.py index 7049bf5..9498619 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -20,6 +20,14 @@ def test_worker_pool_penalizes_failures() -> None: assert router.candidates(1)[0].url == "http://b" assert router.snapshot()[0]["failures"] == 1 + assert router.snapshot()[0]["strategy"] == "load_aware" + + +def test_worker_pool_can_run_round_robin() -> None: + router = WorkerPool(["http://a", "http://b"], strategy="round_robin") + + assert [worker.url for worker in router.candidates(2)] == ["http://a", "http://b"] + assert [worker.url for worker in router.candidates(2)] == ["http://b", "http://a"] def test_worker_pool_requires_worker() -> None: @@ -29,3 +37,12 @@ def test_worker_pool_requires_worker() -> None: assert "worker URL" in str(exc) else: raise AssertionError("expected ValueError") + + +def test_worker_pool_rejects_unknown_strategy() -> None: + try: + WorkerPool(["http://a"], strategy="random") + except ValueError as exc: + assert "strategy" in str(exc) + else: + raise AssertionError("expected ValueError")