[model-gateway] rename py_test to e2e_test (#16454)

2026-01-04 22:24:34 -08:00
parent f02d82211a
commit e6fe092dcc
38 changed files with 9 additions and 9 deletions
@@ -0,0 +1,238 @@
+import threading
+import time
+from types import SimpleNamespace
+
+import pytest
+import requests
+
+from sglang.test.run_eval import run_eval
+
+
+def _wait_for_workers(
+    base_url: str, expected_count: int, timeout: float = 60.0, headers: dict = None
+) -> None:
+    """Poll /workers endpoint until expected number of workers are registered."""
+    start = time.perf_counter()
+    with requests.Session() as session:
+        while time.perf_counter() - start < timeout:
+            try:
+                r = session.get(f"{base_url}/workers", headers=headers, timeout=5)
+                if r.status_code == 200:
+                    workers = r.json().get("workers", [])
+                    if len(workers) >= expected_count:
+                        return
+            except requests.RequestException:
+                pass
+            time.sleep(0.5)
+    raise TimeoutError(
+        f"Expected {expected_count} workers at {base_url}, timed out after {timeout}s"
+    )
+
+
+def _get_worker_id_by_url(base_url: str, worker_url: str, headers: dict = None) -> str:
+    r = requests.get(f"{base_url}/workers", headers=headers, timeout=10)
+    r.raise_for_status()
+    workers = r.json().get("workers", [])
+    worker_id = next((w.get("id") for w in workers if w.get("url") == worker_url), None)
+    if not worker_id:
+        raise RuntimeError(
+            f"Could not find worker_id for url={worker_url}. workers={workers}"
+        )
+    return worker_id
+
+
+@pytest.mark.e2e
+def test_mmlu(e2e_router_only_rr, e2e_two_workers_dp2, e2e_model):
+    # Attach two dp=2 workers (total 4 GPUs) to a fresh router-only instance
+    base = e2e_router_only_rr.url
+    for w in e2e_two_workers_dp2:
+        r = requests.post(f"{base}/workers", json={"url": w.url}, timeout=180)
+        assert (
+            r.status_code == 202
+        ), f"Expected 202 ACCEPTED, got {r.status_code}: {r.text}"
+
+    # Wait for workers to be registered
+    _wait_for_workers(base, expected_count=2, timeout=60.0)
+
+    args = SimpleNamespace(
+        base_url=base,
+        model=e2e_model,
+        eval_name="mmlu",
+        num_examples=64,
+        num_threads=32,
+        temperature=0.1,
+    )
+    metrics = run_eval(args)
+    assert metrics["score"] >= 0.65
+
+
+@pytest.mark.e2e
+def test_genai_bench(
+    e2e_router_only_rr, e2e_two_workers_dp2, e2e_model, genai_bench_runner
+):
+    """Attach a worker to the regular router and run a short genai-bench."""
+    base = e2e_router_only_rr.url
+    for w in e2e_two_workers_dp2:
+        r = requests.post(f"{base}/workers", json={"url": w.url}, timeout=180)
+        assert (
+            r.status_code == 202
+        ), f"Expected 202 ACCEPTED, got {r.status_code}: {r.text}"
+
+    # Wait for workers to be registered
+    _wait_for_workers(base, expected_count=2, timeout=60.0)
+
+    genai_bench_runner(
+        router_url=base,
+        model_path=e2e_model,
+        experiment_folder="benchmark_round_robin_regular",
+        thresholds={
+            "ttft_mean_max": 6,
+            "e2e_latency_mean_max": 14,
+            "input_throughput_mean_min": 800,  # temp relax from 1000 to 800 for now
+            "output_throughput_mean_min": 12,
+            # Enforce GPU utilization p50 >= 99% during the run.
+            "gpu_util_p50_min": 99,
+        },
+        kill_procs=e2e_two_workers_dp2,
+    )
+
+
+@pytest.mark.e2e
+def test_add_and_remove_worker_live(e2e_router_only_rr, e2e_primary_worker, e2e_model):
+    base = e2e_router_only_rr.url
+    worker_url = e2e_primary_worker.url
+
+    r = requests.post(f"{base}/workers", json={"url": worker_url}, timeout=180)
+    assert r.status_code == 202, f"Expected 202 ACCEPTED, got {r.status_code}: {r.text}"
+
+    # Wait for worker to be registered
+    _wait_for_workers(base, expected_count=1, timeout=60.0)
+
+    with requests.Session() as s:
+        for i in range(8):
+            r = s.post(
+                f"{base}/v1/completions",
+                json={
+                    "model": e2e_model,
+                    "prompt": f"x{i}",
+                    "max_tokens": 1,
+                    "stream": False,
+                },
+                timeout=120,
+            )
+            r.raise_for_status()
+
+    # Remove the worker
+    worker_id = _get_worker_id_by_url(base, worker_url)
+    r = requests.delete(f"{base}/workers/{worker_id}", timeout=60)
+    assert r.status_code == 202, f"Expected 202 ACCEPTED, got {r.status_code}: {r.text}"
+
+
+@pytest.mark.e2e
+def test_lazy_fault_tolerance_live(e2e_router_only_rr, e2e_primary_worker, e2e_model):
+    base = e2e_router_only_rr.url
+    worker = e2e_primary_worker
+
+    r = requests.post(f"{base}/workers", json={"url": worker.url}, timeout=180)
+    assert r.status_code == 202, f"Expected 202 ACCEPTED, got {r.status_code}: {r.text}"
+
+    # Wait for worker to be registered
+    _wait_for_workers(base, expected_count=1, timeout=60.0)
+
+    def killer():
+        time.sleep(10)
+        try:
+            worker.proc.terminate()
+        except Exception:
+            pass
+
+    t = threading.Thread(target=killer, daemon=True)
+    t.start()
+
+    args = SimpleNamespace(
+        base_url=base,
+        model=e2e_model,
+        eval_name="mmlu",
+        num_examples=32,
+        num_threads=16,
+        temperature=0.0,
+    )
+    metrics = run_eval(args)
+    assert 0.0 <= metrics["score"] <= 1.0
+
+
+@pytest.mark.e2e
+def test_dp_aware_worker_expansion_and_api_key(
+    e2e_model,
+    e2e_router_only_rr_dp_aware_api,
+    e2e_worker_dp2_api,
+):
+    """
+    Launch a router-only instance in dp_aware mode and a single worker with dp_size=2
+    and API key protection. Verify expansion, auth enforcement, and basic eval.
+    """
+    import os
+
+    router_url = e2e_router_only_rr_dp_aware_api.url
+    worker_url = e2e_worker_dp2_api.url
+    api_key = e2e_router_only_rr_dp_aware_api.api_key
+
+    # Attach worker; router should expand to dp_size logical workers
+    r = requests.post(
+        f"{router_url}/workers",
+        json={"url": worker_url, "api_key": api_key},
+        headers={"Authorization": f"Bearer {api_key}"},
+        timeout=180,
+    )
+    assert r.status_code == 202, f"Expected 202 ACCEPTED, got {r.status_code}: {r.text}"
+
+    # Wait for workers to be registered and expanded
+    _wait_for_workers(
+        router_url,
+        expected_count=2,
+        timeout=60.0,
+        headers={"Authorization": f"Bearer {api_key}"},
+    )
+
+    # Verify the expanded workers have correct URLs
+    r = requests.get(
+        f"{router_url}/workers",
+        headers={"Authorization": f"Bearer {api_key}"},
+        timeout=30,
+    )
+    r.raise_for_status()
+    workers = r.json().get("workers", [])
+    urls = [w["url"] for w in workers]
+    assert len(urls) == 2
+    assert set(urls) == {f"{worker_url}@0", f"{worker_url}@1"}
+
+    # Verify API key enforcement
+    # 1) Without Authorization -> Should get 401 Unauthorized
+    r = requests.post(
+        f"{router_url}/v1/completions",
+        json={"model": e2e_model, "prompt": "hi", "max_tokens": 1},
+        timeout=60,
+    )
+    assert r.status_code == 401
+
+    # 2) With correct Authorization -> 200
+    r = requests.post(
+        f"{router_url}/v1/completions",
+        json={"model": e2e_model, "prompt": "hi", "max_tokens": 1},
+        headers={"Authorization": f"Bearer {api_key}"},
+        timeout=60,
+    )
+    assert r.status_code == 200
+
+    # Finally, run MMLU eval through the router with auth
+    os.environ["OPENAI_API_KEY"] = api_key
+    args = SimpleNamespace(
+        base_url=router_url,
+        model=e2e_model,
+        eval_name="mmlu",
+        num_examples=64,
+        num_threads=32,
+        temperature=0.1,
+    )
+    metrics = run_eval(args)
+    assert metrics["score"] >= 0.65