From acdf2751a389e0ccc0e7e972fdb1b9565ff836e2 Mon Sep 17 00:00:00 2001 From: Alex Cheema Date: Thu, 12 Feb 2026 10:58:54 -0800 Subject: [PATCH] feat: add deterministic inference snapshot test Launch mlx-community/Qwen3-0.6B-4bit on the cluster, send a chat completion with seed=42 and temperature=0, and verify the output matches a committed snapshot. Tests inference determinism end-to-end. Co-Authored-By: Claude Opus 4.6 --- e2e/conftest.py | 47 +++++++++++++++++++++- e2e/snapshots/inference.json | 8 ++++ e2e/test_inference_snapshot.py | 71 ++++++++++++++++++++++++++++++++++ 3 files changed, 125 insertions(+), 1 deletion(-) create mode 100644 e2e/snapshots/inference.json create mode 100644 e2e/test_inference_snapshot.py diff --git a/e2e/conftest.py b/e2e/conftest.py index 630c28ce..3b6bd5b6 100644 --- a/e2e/conftest.py +++ b/e2e/conftest.py @@ -1,10 +1,11 @@ """Shared E2E test infrastructure for exo cluster tests.""" import asyncio +import json import os import sys from pathlib import Path -from urllib.request import urlopen +from urllib.request import urlopen, Request from urllib.error import URLError E2E_DIR = Path(__file__).parent.resolve() @@ -112,3 +113,47 @@ class Cluster: await self.wait_for("Nodes discovered each other", nodes_discovered) await self.wait_for("Master election resolved", master_elected) await self.wait_for("API responding", api_responding) + + async def _api(self, method: str, path: str, body: dict | None = None, timeout: int = 30) -> dict: + """Make an API request to the cluster. Returns parsed JSON.""" + url = f"http://localhost:52415{path}" + data = json.dumps(body).encode() if body else None + req = Request(url, data=data, headers={"Content-Type": "application/json"}, method=method) + loop = asyncio.get_event_loop() + resp_bytes = await loop.run_in_executor(None, lambda: urlopen(req, timeout=timeout).read()) + return json.loads(resp_bytes) + + async def place_model(self, model: str, timeout: int = 600): + """Place a model instance on the cluster (triggers download) and wait until it's ready.""" + await self._api("POST", "/place_instance", {"model_id": model}) + + async def model_ready(): + try: + resp = await self._api("GET", "/v1/models") + return any(m.get("id") == model for m in resp.get("data", [])) + except Exception: + return False + + await self.wait_for(f"Model {model} ready", model_ready, timeout=timeout) + + async def chat(self, model: str, messages: list[dict], timeout: int = 600, **kwargs) -> dict: + """Send a chat completion request. Retries until model is downloaded and inference completes.""" + body = json.dumps({"model": model, "messages": messages, **kwargs}).encode() + deadline = asyncio.get_event_loop().time() + timeout + last_error = None + + while asyncio.get_event_loop().time() < deadline: + try: + req = Request( + "http://localhost:52415/v1/chat/completions", + data=body, + headers={"Content-Type": "application/json"}, + ) + loop = asyncio.get_event_loop() + resp_bytes = await loop.run_in_executor(None, lambda: urlopen(req, timeout=300).read()) + return json.loads(resp_bytes) + except Exception as e: + last_error = e + await asyncio.sleep(5) + + raise TimeoutError(f"Chat request failed after {timeout}s: {last_error}") diff --git a/e2e/snapshots/inference.json b/e2e/snapshots/inference.json new file mode 100644 index 00000000..d2721114 --- /dev/null +++ b/e2e/snapshots/inference.json @@ -0,0 +1,8 @@ +{ + "model": "mlx-community/Qwen3-0.6B-4bit", + "seed": 42, + "temperature": 0, + "prompt": "What is 2+2? Reply with just the number.", + "max_tokens": 32, + "content": "\nOkay, so I need to figure out what 2+2 is. Let me think. Well, if you add 2 and 2 together" +} diff --git a/e2e/test_inference_snapshot.py b/e2e/test_inference_snapshot.py new file mode 100644 index 00000000..8e3e06d1 --- /dev/null +++ b/e2e/test_inference_snapshot.py @@ -0,0 +1,71 @@ +"""Test: Deterministic inference output (snapshot test). + +Sends a chat completion request with a fixed seed and temperature=0, +then verifies the output matches a known-good snapshot. This ensures +inference produces consistent results across runs. +""" + +import asyncio +import json +import sys +from pathlib import Path +sys.path.insert(0, str(Path(__file__).parent)) +from conftest import Cluster + +MODEL = "mlx-community/Qwen3-0.6B-4bit" +SEED = 42 +PROMPT = "What is 2+2? Reply with just the number." +MAX_TOKENS = 32 +SNAPSHOT_FILE = Path(__file__).parent / "snapshots" / "inference.json" + + +async def main(): + async with Cluster("inference_snapshot") as cluster: + await cluster.build() + await cluster.start() + await cluster.assert_healthy() + + # Launch the model instance (triggers download + placement) + print(f" Launching model {MODEL}...") + await cluster.place_model(MODEL) + + print(f" Sending chat completion (seed={SEED}, temperature=0)...") + resp = await cluster.chat( + model=MODEL, + messages=[{"role": "user", "content": PROMPT}], + seed=SEED, + temperature=0, + max_tokens=MAX_TOKENS, + ) + + content = resp["choices"][0]["message"]["content"] + print(f" Response: {content!r}") + + # Load or create snapshot + if SNAPSHOT_FILE.exists(): + snapshot = json.loads(SNAPSHOT_FILE.read_text()) + expected = snapshot["content"] + assert content == expected, ( + f"Snapshot mismatch!\n" + f" Expected: {expected!r}\n" + f" Got: {content!r}\n" + f" Delete {SNAPSHOT_FILE} to regenerate." + ) + print(f" Output matches snapshot") + else: + SNAPSHOT_FILE.parent.mkdir(parents=True, exist_ok=True) + SNAPSHOT_FILE.write_text(json.dumps({ + "model": MODEL, + "seed": SEED, + "temperature": 0, + "prompt": PROMPT, + "max_tokens": MAX_TOKENS, + "content": content, + }, indent=2) + "\n") + print(f" Snapshot created: {SNAPSHOT_FILE}") + + print("PASSED: inference_snapshot") + + +if __name__ == "__main__": + asyncio.run(main())