sgl-router: experimental Rust HTTP router for SGLang worker pools (#25851)
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Kangyan-Zhou
co-authored by
Claude Opus 4.7
Claude Opus 4.7
parent
aae04b1241
commit
6e8fe176be
@@ -0,0 +1,266 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""Content-based cross-router routing test for cache-aware-zmq.
|
||||
|
||||
Two routers + two SGLang workers + one shared model. Each router runs an
|
||||
independent ``cache_aware_zmq`` policy whose ``KvEventIndex`` subscribes
|
||||
to **both** workers' KV publishers.
|
||||
|
||||
The test warms each worker with a DIFFERENT prefix DIRECTLY (bypassing
|
||||
both routers), then sends those prefixes through each router and
|
||||
asserts that routing follows the prefix CONTENT: ``PREFIX_X`` lands on
|
||||
the worker holding X, ``PREFIX_Y`` lands on the worker holding Y, on
|
||||
both routers.
|
||||
|
||||
# Why content-based, not convergence
|
||||
|
||||
An earlier version of this test asserted that both routers converged on
|
||||
the *same dominant worker* after a one-prefix warmup. That property
|
||||
sounds like it pins the ZMQ-fan-out contract, but it doesn't: when the
|
||||
KV-event path is broken (subscribers never opened, e.g. a worker's
|
||||
``/server_info`` lacks the ``kv_events`` block), ``cache_aware_zmq``
|
||||
silently degrades to **min-load** — which, with sequential requests
|
||||
holding ``active_load`` at zero, picks the same worker deterministically
|
||||
on every call within a router. Both routers' min-load picks happened to
|
||||
agree often enough (about half the time, modulo HashSet seed) to make
|
||||
the convergence assertion pass even when no event ever flowed.
|
||||
|
||||
Content-based routing is uniquely sensitive to the KV-event path. Two
|
||||
disjoint prefixes warmed on two different workers can only be routed
|
||||
correctly if the router knows *which worker holds which content* — the
|
||||
only mechanism that supplies that information is the ``BlockStored``
|
||||
event stream. Under min-load fallback, both prefixes route to the same
|
||||
default worker on each router, so the ``PREFIX_Y → worker_y`` assertion
|
||||
fails regardless of which worker min-load defaults to.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import time
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from infra.gateway import Gateway
|
||||
from infra.model_pool import PASSTHROUGH_CHAT_TEMPLATE_PATH, spawn_worker
|
||||
from infra.model_specs import get_model_spec
|
||||
|
||||
# Disjoint prefixes — share no common opening text, so block 0 hashes
|
||||
# differ from the first block onward and each worker's HashTree
|
||||
# contribution is uniquely identifying.
|
||||
#
|
||||
# Length matters: each prefix must span ≥2 SGLang blocks at the default
|
||||
# block_size of 64 tokens so the worker actually emits BlockStored
|
||||
# events. Below that, the publisher stays quiet and we'd be testing
|
||||
# min-load by accident — the exact failure mode this test exists to
|
||||
# rule out.
|
||||
_PREFIX_X_BODY = (
|
||||
"Apricot bouquet cinnamon dewdrop elderflower fennel garlic "
|
||||
"hibiscus indigo jasmine kumquat lavender mint nutmeg oregano "
|
||||
"paprika quince rosemary saffron tarragon. "
|
||||
)
|
||||
PREFIX_X = (_PREFIX_X_BODY * 8).strip()
|
||||
|
||||
_PREFIX_Y_BODY = (
|
||||
"Zephyr yellow xylophone wombat vortex umbrella thistle saffron "
|
||||
"quartz peppermint orchid nightshade marigold lemongrass kale "
|
||||
"juniper iris hyacinth gardenia foxglove. "
|
||||
)
|
||||
PREFIX_Y = (_PREFIX_Y_BODY * 8).strip()
|
||||
|
||||
|
||||
_REQ_TOTAL_RE = re.compile(
|
||||
r"^sgl_router_requests_total\{([^}]*)\}\s+(\d+(?:\.\d+)?)\s*$"
|
||||
)
|
||||
_LABEL_RE = re.compile(r'(\w+)="([^"]*)"')
|
||||
|
||||
|
||||
def _success_counts_by_worker(router_url: str) -> dict[str, int]:
|
||||
"""Scrape ``/metrics`` and return ``{worker_url: success_count}``."""
|
||||
r = httpx.get(f"{router_url}/metrics", timeout=5.0)
|
||||
r.raise_for_status()
|
||||
counts: dict[str, int] = {}
|
||||
for line in r.text.splitlines():
|
||||
m = _REQ_TOTAL_RE.match(line)
|
||||
if not m:
|
||||
continue
|
||||
labels = dict(_LABEL_RE.findall(m.group(1)))
|
||||
if labels.get("outcome") != "success":
|
||||
continue
|
||||
worker = labels.get("worker_url")
|
||||
if not worker:
|
||||
continue
|
||||
try:
|
||||
counts[worker] = counts.get(worker, 0) + int(float(m.group(2)))
|
||||
except ValueError:
|
||||
continue
|
||||
return counts
|
||||
|
||||
|
||||
def _send_chat(url: str, model_id: str, prompt: str) -> int:
|
||||
"""POST one chat completion; return the HTTP status."""
|
||||
r = httpx.post(
|
||||
f"{url}/v1/chat/completions",
|
||||
json={
|
||||
"model": model_id,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"max_tokens": 4,
|
||||
"stream": False,
|
||||
},
|
||||
timeout=60.0,
|
||||
)
|
||||
return r.status_code
|
||||
|
||||
|
||||
def _direct_warm(worker_url: str, model_id: str, prefix: str) -> None:
|
||||
"""Send one ``/v1/chat/completions`` request with ``prefix`` DIRECTLY to a worker.
|
||||
|
||||
The KV-event publisher emits ``BlockStored`` as the request's
|
||||
prompt blocks commit to that worker's cache; routers subscribed to
|
||||
the publisher receive the event and add ``(block_hash → worker)``
|
||||
entries to their ``HashTree``. The test then exercises those
|
||||
entries by routing through the router.
|
||||
|
||||
Direct-warming (rather than going through a router) is the load-
|
||||
bearing detail: routing through a router would itself choose which
|
||||
worker to populate, so the two workers' HashTree state would no
|
||||
longer be uniquely identifying.
|
||||
|
||||
Token alignment with the router — ``cache_aware_zmq`` hashes
|
||||
``messages[*].content`` RAW (``cache_aware_zmq.rs::extract_prompt_text``)
|
||||
using ``add_special_tokens=false``. By default SGLang's chat
|
||||
endpoint would wrap ``prefix`` in the model's chat template before
|
||||
tokenizing — adding role tags, end-of-turn markers, and a
|
||||
generation prompt — and the resulting block hashes would never
|
||||
match what the router computes from raw content.
|
||||
|
||||
The test launches each worker with ``--chat-template
|
||||
<PASSTHROUGH_CHAT_TEMPLATE_PATH>``: a Jinja template that emits
|
||||
only ``messages[*].content`` (the same shape the router extracts),
|
||||
and which combines with Transformers' ``apply_chat_template(
|
||||
tokenize=True, add_special_tokens=False)`` to produce the same
|
||||
token stream the router will compute. So warm and route hash the
|
||||
same blocks via the same endpoint.
|
||||
"""
|
||||
r = httpx.post(
|
||||
f"{worker_url}/v1/chat/completions",
|
||||
json={
|
||||
"model": model_id,
|
||||
"messages": [{"role": "user", "content": prefix}],
|
||||
"max_tokens": 4,
|
||||
"stream": False,
|
||||
},
|
||||
timeout=60.0,
|
||||
)
|
||||
assert (
|
||||
r.status_code == 200
|
||||
), f"direct warm to {worker_url} failed: HTTP {r.status_code} {r.text!r}"
|
||||
|
||||
|
||||
def _route_through(router_url: str, model_id: str, prompt: str) -> str:
|
||||
"""Send one request through ``router_url``; return which worker handled it.
|
||||
|
||||
Computed by diffing the per-worker success-counter on ``/metrics``
|
||||
around the call. Asserts exactly one worker absorbed the request
|
||||
(no partial counts, no cancellation race).
|
||||
"""
|
||||
before = _success_counts_by_worker(router_url)
|
||||
code = _send_chat(router_url, model_id, prompt)
|
||||
assert code == 200, f"request to {router_url} failed: HTTP {code}"
|
||||
after = _success_counts_by_worker(router_url)
|
||||
deltas = {w: after.get(w, 0) - before.get(w, 0) for w in set(after) | set(before)}
|
||||
winners = [w for w, d in deltas.items() if d > 0]
|
||||
assert (
|
||||
len(winners) == 1
|
||||
), f"expected exactly one worker delta on {router_url}, got {deltas}"
|
||||
return winners[0]
|
||||
|
||||
|
||||
@pytest.mark.real_gpu
|
||||
@pytest.mark.slow
|
||||
def test_two_routers_route_by_prefix_content(
|
||||
router_binary, # noqa: ARG001 — fixture forces release-binary presence
|
||||
gpu_allocator,
|
||||
):
|
||||
"""Each router must route by prefix CONTENT, agreeing across routers.
|
||||
|
||||
With each worker direct-warmed by a different disjoint prefix, the
|
||||
only way a router can route ``PREFIX_X → worker_x`` AND
|
||||
``PREFIX_Y → worker_y`` is by consulting a HashTree populated from
|
||||
the BlockStored events the workers emit. Min-load fallback (the
|
||||
failure mode when no SUB socket opened) is content-blind and would
|
||||
route both prefixes to whichever worker its tiebreaker prefers.
|
||||
"""
|
||||
spec = get_model_spec("qwen3-0.6b")
|
||||
gpus = gpu_allocator.acquire(2)
|
||||
# Passthrough chat template — see _direct_warm for the rationale. Both
|
||||
# workers must run with the same template; otherwise their KV blocks
|
||||
# would hash template-wrapped tokens while the router hashes raw
|
||||
# content, and every lookup would miss the tree.
|
||||
worker_chat_template_args = ["--chat-template", PASSTHROUGH_CHAT_TEMPLATE_PATH]
|
||||
try:
|
||||
with (
|
||||
spawn_worker(
|
||||
"qwen3-0.6b",
|
||||
gpu_ids=[gpus[0]],
|
||||
enable_kv_events=True,
|
||||
extra_args=worker_chat_template_args,
|
||||
) as worker_x,
|
||||
spawn_worker(
|
||||
"qwen3-0.6b",
|
||||
gpu_ids=[gpus[1]],
|
||||
enable_kv_events=True,
|
||||
extra_args=worker_chat_template_args,
|
||||
) as worker_y,
|
||||
Gateway() as router_a,
|
||||
Gateway() as router_b,
|
||||
):
|
||||
worker_urls = [worker_x.url, worker_y.url]
|
||||
for gw in (router_a, router_b):
|
||||
gw.start_regular(
|
||||
model_id=spec["model"],
|
||||
tokenizer_path=spec["model"],
|
||||
worker_urls=worker_urls,
|
||||
policy="cache_aware_zmq",
|
||||
timeout=120.0,
|
||||
)
|
||||
|
||||
# 1. Direct-warm each worker with its own prefix. Must happen
|
||||
# AFTER both routers have started — ZMQ PUB/SUB doesn't
|
||||
# replay messages emitted before SUB attaches, so any
|
||||
# BlockStored event predating subscription is lost and
|
||||
# the HashTree never sees it.
|
||||
_direct_warm(worker_x.url, spec["model"], PREFIX_X)
|
||||
_direct_warm(worker_y.url, spec["model"], PREFIX_Y)
|
||||
|
||||
# 2. Drain the SUB mpsc + pump-apply path. Sub-second under
|
||||
# loopback ZMQ; 2 s leaves comfortable headroom.
|
||||
time.sleep(2.0)
|
||||
|
||||
# 3. Content-routing assertion (×4): each prefix must land
|
||||
# on the worker that holds it, on either router.
|
||||
#
|
||||
# The four assertions below are independently strong:
|
||||
# min-load fallback routes both prefixes on a given
|
||||
# router to a single default worker, so for ANY broken-
|
||||
# fan-out scenario at least one of the four fails.
|
||||
for router, label in ((router_a, "A"), (router_b, "B")):
|
||||
landed = _route_through(router.base_url, spec["model"], PREFIX_X)
|
||||
assert landed == worker_x.url, (
|
||||
f"router {label}: PREFIX_X must route to worker_x "
|
||||
f"({worker_x.url}); landed on {landed}. "
|
||||
f"Likely cause: HashTree is empty — KV-event "
|
||||
f"subscriber never opened, or BlockStored events "
|
||||
f"never reached the pump."
|
||||
)
|
||||
landed = _route_through(router.base_url, spec["model"], PREFIX_Y)
|
||||
assert landed == worker_y.url, (
|
||||
f"router {label}: PREFIX_Y must route to worker_y "
|
||||
f"({worker_y.url}); landed on {landed}. "
|
||||
f"Likely cause: HashTree is empty — KV-event "
|
||||
f"subscriber never opened, or BlockStored events "
|
||||
f"never reached the pump."
|
||||
)
|
||||
finally:
|
||||
gpu_allocator.release(gpus)
|
||||
@@ -0,0 +1,98 @@
|
||||
"""Basic chat-completions correctness — ported from SMG's
|
||||
``e2e_test/chat_completions/test_validation.py``, narrowed to the
|
||||
subset that exercises sgl-router (not SMG's per-message validators).
|
||||
|
||||
The shape:
|
||||
- single-worker regular-mode router
|
||||
- non-streaming + streaming chat completion
|
||||
- assistant message non-empty, role correct, finish_reason set
|
||||
|
||||
These are the smoke tests that run first; if they pass, the heavier
|
||||
multi-worker acceptance tests are worth running.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from infra.gateway import Gateway
|
||||
from infra.model_pool import spawn_worker
|
||||
from infra.model_specs import get_model_spec
|
||||
|
||||
|
||||
@pytest.mark.real_gpu
|
||||
def test_chat_non_streaming_returns_assistant_message(
|
||||
router_binary, # noqa: ARG001
|
||||
gpu_allocator,
|
||||
):
|
||||
gpu = gpu_allocator.acquire(1)
|
||||
try:
|
||||
with spawn_worker("qwen3-0.6b", gpu_ids=gpu) as worker:
|
||||
spec = get_model_spec("qwen3-0.6b")
|
||||
with Gateway() as gw:
|
||||
gw.start_regular(
|
||||
model_id=spec["model"],
|
||||
tokenizer_path=spec["model"],
|
||||
worker_urls=[worker.url],
|
||||
timeout=120.0,
|
||||
)
|
||||
resp = httpx.post(
|
||||
f"{gw.base_url}/v1/chat/completions",
|
||||
json={
|
||||
"model": spec["model"],
|
||||
"messages": [{"role": "user", "content": "Say hi."}],
|
||||
"max_tokens": 16,
|
||||
"stream": False,
|
||||
},
|
||||
timeout=60.0,
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
body = resp.json()
|
||||
choice = body["choices"][0]
|
||||
assert choice["message"]["role"] == "assistant"
|
||||
assert choice["message"][
|
||||
"content"
|
||||
], f"empty assistant content: {choice!r}"
|
||||
assert choice.get("finish_reason"), choice
|
||||
finally:
|
||||
gpu_allocator.release(gpu)
|
||||
|
||||
|
||||
@pytest.mark.real_gpu
|
||||
def test_chat_streaming_emits_sse_chunks_with_done(
|
||||
router_binary, # noqa: ARG001
|
||||
gpu_allocator,
|
||||
):
|
||||
gpu = gpu_allocator.acquire(1)
|
||||
try:
|
||||
with spawn_worker("qwen3-0.6b", gpu_ids=gpu) as worker:
|
||||
spec = get_model_spec("qwen3-0.6b")
|
||||
with Gateway() as gw:
|
||||
gw.start_regular(
|
||||
model_id=spec["model"],
|
||||
tokenizer_path=spec["model"],
|
||||
worker_urls=[worker.url],
|
||||
timeout=120.0,
|
||||
)
|
||||
chunks: list[str] = []
|
||||
with httpx.stream(
|
||||
"POST",
|
||||
f"{gw.base_url}/v1/chat/completions",
|
||||
json={
|
||||
"model": spec["model"],
|
||||
"messages": [{"role": "user", "content": "Say hi."}],
|
||||
"max_tokens": 16,
|
||||
"stream": True,
|
||||
},
|
||||
timeout=60.0,
|
||||
) as resp:
|
||||
assert resp.status_code == 200, resp.read().decode()
|
||||
for line in resp.iter_lines():
|
||||
if line.startswith("data:"):
|
||||
chunks.append(line.strip())
|
||||
assert len(chunks) >= 2, f"expected >=2 SSE chunks, got: {chunks}"
|
||||
assert any(
|
||||
"[DONE]" in c for c in chunks
|
||||
), f"no [DONE] terminator in stream: {chunks}"
|
||||
finally:
|
||||
gpu_allocator.release(gpu)
|
||||
@@ -0,0 +1,327 @@
|
||||
"""Pytest fixtures for ``experimental/sgl-router/tests/e2e/``.
|
||||
|
||||
Two flavors of fixtures coexist here:
|
||||
|
||||
1. **Session-scoped smoke fixtures** (``sglang_server`` + ``router``) —
|
||||
launch ONE SGLang worker + ONE router on fixed ports for the whole
|
||||
test session. Used by the lightweight ``test_chat_smoke.py`` /
|
||||
``test_tokenize_smoke.py`` files. These are the cheap "did the
|
||||
binary start at all" sanity tests.
|
||||
|
||||
2. **Per-test multi-worker fixtures** (``router_binary`` +
|
||||
``gpu_allocator``) — just enough infra for the acceptance tests in
|
||||
``chat_completions/`` to bring up their own multi-worker
|
||||
topologies. Backed by the ``infra.gateway.Gateway`` and
|
||||
``infra.model_pool.spawn_worker`` helpers.
|
||||
|
||||
Both sets share the same release binary; ``SGL_ROUTER_BINARY`` env var
|
||||
overrides the path for both.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Iterator
|
||||
from pathlib import Path
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Make `from infra import gateway, model_pool, model_specs` resolve from
|
||||
# tests under tests/e2e/ without requiring a sibling `__init__.py` chain.
|
||||
# Mirrors SMG's e2e_test/conftest.py sys.path setup.
|
||||
_E2E_DIR = Path(__file__).resolve().parent
|
||||
if str(_E2E_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(_E2E_DIR))
|
||||
|
||||
MODEL = "Qwen/Qwen3-0.6B"
|
||||
SGLANG_PORT = 30000
|
||||
ROUTER_PORT = 8090
|
||||
|
||||
# Path to the release binary. This file lives at
|
||||
# `experimental/sgl-router/tests/e2e/conftest.py`, so:
|
||||
# parent = tests/e2e/
|
||||
# parent.parent = tests/
|
||||
# parent.parent.parent = experimental/sgl-router/ ← cargo workspace root
|
||||
# A previous version used `parent.parent / "target"`, which pointed at
|
||||
# `experimental/sgl-router/tests/target/` and silently broke every
|
||||
# fixture that tries to launch the router binary (CI's
|
||||
# `cargo build --release` lands the artifact at
|
||||
# `experimental/sgl-router/target/release/sgl-router`, not under
|
||||
# `tests/`).
|
||||
_SGL_ROUTER_ROOT = Path(__file__).parent.parent.parent
|
||||
_BINARY = (
|
||||
Path(os.environ.get("CARGO_TARGET_DIR", str(_SGL_ROUTER_ROOT / "target")))
|
||||
/ "release"
|
||||
/ "sgl-router"
|
||||
)
|
||||
|
||||
|
||||
def _wait_http(url: str, timeout: int = 120) -> None:
|
||||
"""Poll *url* until it returns 2xx or raises RuntimeError on timeout."""
|
||||
deadline = time.time() + timeout
|
||||
last_exc: Exception | None = None
|
||||
while time.time() < deadline:
|
||||
try:
|
||||
resp = httpx.get(url, timeout=5)
|
||||
if resp.status_code < 300:
|
||||
return
|
||||
except Exception as exc: # noqa: BLE001
|
||||
last_exc = exc
|
||||
time.sleep(5)
|
||||
raise RuntimeError(
|
||||
f"Timed out waiting for {url} after {timeout}s (last error: {last_exc})"
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def sglang_server():
|
||||
"""Launch a real SGLang server on port 30000 and wait until healthy."""
|
||||
# Stream the server's stdout/stderr to a file rather than capturing
|
||||
# to subprocess.PIPE. The launch_server startup log is verbose (model
|
||||
# download, JIT warmup, NCCL init); once a PIPE'd output fills its
|
||||
# ~64 KB OS buffer with nothing reading it, the SGLang process
|
||||
# blocks on stdout write and never reaches "Server started" — the
|
||||
# health probe then times out at 300 s and we have no visibility
|
||||
# into *why*. A real log file fixes both (no buffer pressure, and
|
||||
# the file is dumped on failure for triage).
|
||||
log_path = Path(tempfile.gettempdir()) / f"sglang-server-{SGLANG_PORT}.log"
|
||||
log_handle = open(log_path, "w", buffering=1) # line-buffered
|
||||
proc = subprocess.Popen(
|
||||
[
|
||||
"python3",
|
||||
"-m",
|
||||
"sglang.launch_server",
|
||||
"--model-path",
|
||||
MODEL,
|
||||
"--port",
|
||||
str(SGLANG_PORT),
|
||||
"--tp",
|
||||
"1",
|
||||
],
|
||||
stdout=log_handle,
|
||||
stderr=subprocess.STDOUT,
|
||||
)
|
||||
|
||||
try:
|
||||
_wait_http(f"http://localhost:{SGLANG_PORT}/health", timeout=300)
|
||||
except Exception:
|
||||
# Dump the server log so the operator can see why startup failed
|
||||
# (model download error, port conflict, OOM, JIT crash, etc.).
|
||||
proc.send_signal(signal.SIGTERM)
|
||||
try:
|
||||
proc.wait(timeout=30)
|
||||
except subprocess.TimeoutExpired:
|
||||
proc.kill()
|
||||
proc.wait()
|
||||
log_handle.flush()
|
||||
log_handle.close()
|
||||
try:
|
||||
tail = log_path.read_text(errors="replace").splitlines()[-200:]
|
||||
except OSError:
|
||||
tail = ["(server log unreadable)"]
|
||||
logger.error(
|
||||
"sglang_server fixture failed; last 200 log lines from %s:\n%s",
|
||||
log_path,
|
||||
"\n".join(tail),
|
||||
)
|
||||
raise
|
||||
|
||||
yield f"http://localhost:{SGLANG_PORT}"
|
||||
|
||||
proc.send_signal(signal.SIGTERM)
|
||||
try:
|
||||
proc.wait(timeout=30)
|
||||
except subprocess.TimeoutExpired:
|
||||
proc.kill()
|
||||
proc.wait()
|
||||
log_handle.flush()
|
||||
log_handle.close()
|
||||
|
||||
|
||||
def _find_tokenizer_path(model: str) -> str:
|
||||
"""Locate the tokenizer.json for *model* from the local HF Hub cache.
|
||||
|
||||
Falls back to the model string itself (a valid HF Hub repo identifier
|
||||
that dynamo-tokenizers can resolve at runtime) when the cache is absent.
|
||||
"""
|
||||
try:
|
||||
from huggingface_hub import try_to_load_from_cache # type: ignore[import]
|
||||
|
||||
path = try_to_load_from_cache(model, "tokenizer.json")
|
||||
if path and Path(path).is_file():
|
||||
return str(path)
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
# Let dynamo-tokenizers resolve the repo identifier directly.
|
||||
return model
|
||||
|
||||
|
||||
def build_smoke_router_config(
|
||||
*,
|
||||
host: str,
|
||||
port: int,
|
||||
model: str,
|
||||
tokenizer_path: str,
|
||||
sglang_url: str,
|
||||
) -> str:
|
||||
"""Build the TOML the smoke `router` fixture writes to disk.
|
||||
|
||||
Returns ``main_config_text`` carrying ``[server]``, ``[[models]]``,
|
||||
and ``[discovery] backend = "static_urls"`` with the worker URL
|
||||
inline. The Rust ``Config`` struct requires a ``[discovery]``
|
||||
section (``DiscoveryConfig`` has no ``#[serde(default)]``) and has
|
||||
no top-level ``workers`` field. The previous ``static_file``
|
||||
backend was replaced by ``static_urls`` (which holds the URL list
|
||||
inline rather than via a side-car file).
|
||||
"""
|
||||
return f"""\
|
||||
[server]
|
||||
host = "{host}"
|
||||
port = {port}
|
||||
|
||||
[[models]]
|
||||
id = "{model}"
|
||||
tokenizer_path = "{tokenizer_path}"
|
||||
|
||||
[discovery]
|
||||
backend = "static_urls"
|
||||
|
||||
[discovery.static_urls]
|
||||
urls = ["{sglang_url}"]
|
||||
"""
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def router(sglang_server): # noqa: ARG001 (sglang_server must start first)
|
||||
"""Launch sgl-router on port 8090 pointed at the SGLang worker."""
|
||||
tok_path = _find_tokenizer_path(MODEL)
|
||||
cfg_handle = tempfile.NamedTemporaryFile(mode="w", suffix=".toml", delete=False)
|
||||
cfg_path = Path(cfg_handle.name)
|
||||
main_text = build_smoke_router_config(
|
||||
host="0.0.0.0",
|
||||
port=ROUTER_PORT,
|
||||
model=MODEL,
|
||||
tokenizer_path=tok_path,
|
||||
sglang_url=f"http://localhost:{SGLANG_PORT}",
|
||||
)
|
||||
cfg_handle.write(main_text)
|
||||
cfg_handle.close()
|
||||
|
||||
try:
|
||||
proc = subprocess.Popen(
|
||||
[str(_BINARY), "--config", str(cfg_path)],
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
)
|
||||
|
||||
try:
|
||||
_wait_http(f"http://localhost:{ROUTER_PORT}/readyz", timeout=60)
|
||||
except Exception:
|
||||
proc.send_signal(signal.SIGTERM)
|
||||
proc.wait(timeout=30)
|
||||
raise
|
||||
|
||||
yield f"http://localhost:{ROUTER_PORT}"
|
||||
|
||||
proc.send_signal(signal.SIGTERM)
|
||||
try:
|
||||
proc.wait(timeout=30)
|
||||
except subprocess.TimeoutExpired:
|
||||
proc.kill()
|
||||
proc.wait()
|
||||
finally:
|
||||
cfg_path.unlink(missing_ok=True)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Per-test multi-worker acceptance fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _detect_gpu_count() -> int:
|
||||
"""Count visible GPUs via ``nvidia-smi``. Returns 0 when no NVIDIA GPU
|
||||
is available (CI on CPU-only runners, dev laptops, etc.).
|
||||
"""
|
||||
try:
|
||||
out = subprocess.check_output(
|
||||
["nvidia-smi", "--query-gpu=index", "--format=csv,noheader"],
|
||||
stderr=subprocess.DEVNULL,
|
||||
timeout=5.0,
|
||||
)
|
||||
except (FileNotFoundError, subprocess.SubprocessError):
|
||||
return 0
|
||||
return len([ln for ln in out.decode().splitlines() if ln.strip()])
|
||||
|
||||
|
||||
class GPUAllocator:
|
||||
"""Single-process GPU index allocator. Test-scoped; not safe for
|
||||
cross-process use (pytest-xdist) — each worker would race over the
|
||||
full GPU set. Acceptance tests run serially, so this is fine.
|
||||
"""
|
||||
|
||||
def __init__(self, total: int):
|
||||
self.total = total
|
||||
self._free: list[int] = list(range(total))
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def acquire(self, n: int = 1) -> list[int]:
|
||||
with self._lock:
|
||||
if n > len(self._free):
|
||||
raise pytest.skip.Exception(
|
||||
f"requested {n} GPUs, only {len(self._free)}/{self.total} free"
|
||||
)
|
||||
picked = self._free[:n]
|
||||
self._free = self._free[n:]
|
||||
return picked
|
||||
|
||||
def release(self, ids: list[int]) -> None:
|
||||
with self._lock:
|
||||
self._free.extend(ids)
|
||||
self._free.sort()
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def router_binary() -> Path:
|
||||
"""Locate the release ``sgl-router`` binary or skip the session.
|
||||
|
||||
Used by the multi-worker acceptance tests (which spawn their own
|
||||
Gateway per test instead of using the session-scoped ``router``
|
||||
fixture).
|
||||
"""
|
||||
env_path = os.environ.get("SGL_ROUTER_BINARY")
|
||||
candidates: list[Path] = []
|
||||
if env_path:
|
||||
candidates.append(Path(env_path))
|
||||
candidates.append(_BINARY)
|
||||
for c in candidates:
|
||||
if c.exists():
|
||||
return c
|
||||
pytest.skip(
|
||||
"sgl-router release binary not found at any of: "
|
||||
+ ", ".join(str(c) for c in candidates)
|
||||
+ ". Build with `cargo build --release` in experimental/sgl-router/."
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def gpu_allocator() -> Iterator[GPUAllocator]:
|
||||
"""Session-scoped GPU index allocator. Skips the entire session when
|
||||
no GPUs are visible — acceptance tests under chat_completions/ are
|
||||
real-GPU.
|
||||
"""
|
||||
n = _detect_gpu_count()
|
||||
if n == 0:
|
||||
pytest.skip(
|
||||
"no NVIDIA GPUs visible to nvidia-smi; acceptance tests are GPU-only"
|
||||
)
|
||||
yield GPUAllocator(n)
|
||||
@@ -0,0 +1,404 @@
|
||||
"""Minimal sgl-router Gateway class — adapted from SMG's e2e_test/infra/gateway.py.
|
||||
|
||||
Differences from SMG:
|
||||
- SMG drives a Python launcher (`python3 -m sglang_router.launch_router`)
|
||||
with worker URLs on the CLI.
|
||||
- sgl-router uses a Rust binary (`experimental/sgl-router/target/release/sgl-router`)
|
||||
with a TOML config file. Worker discovery is config-file-based; this
|
||||
Gateway writes a TOML to a tempfile and execs the binary with
|
||||
`--config <tempfile>`.
|
||||
|
||||
Supported lifecycles:
|
||||
- Regular mode: one model, N worker URLs, single policy.
|
||||
- PD mode: one model, prefill_workers + decode_workers (lists of URLs),
|
||||
discovery emits separate `WorkerMode::Prefill` / `WorkerMode::Decode`
|
||||
entries. The router resolves PD pool isolation at request time.
|
||||
|
||||
Use as a context manager:
|
||||
|
||||
with Gateway() as gw:
|
||||
gw.start_regular(model_path="...", worker_urls=[...])
|
||||
resp = httpx.post(f"{gw.base_url}/v1/chat/completions", json=...)
|
||||
|
||||
or pytest fixture style (see e2e_test/conftest.py).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import socket
|
||||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Repo-relative path to the release binary. Set ``SGL_ROUTER_BINARY`` to
|
||||
# override (e.g. a debug build, or a non-default ``CARGO_TARGET_DIR``).
|
||||
# This file is at `experimental/sgl-router/tests/e2e/infra/gateway.py`,
|
||||
# so four `.parent` hops to reach the sgl-router workspace root
|
||||
# (infra → e2e → tests → sgl-router). Cargo lands the binary at
|
||||
# `experimental/sgl-router/target/release/sgl-router`. A previous
|
||||
# version used three hops and pointed at `tests/target/`, which
|
||||
# would have broken any test that actually launches the router via
|
||||
# this helper.
|
||||
DEFAULT_BINARY = (
|
||||
Path(__file__).resolve().parent.parent.parent.parent
|
||||
/ "target"
|
||||
/ "release"
|
||||
/ "sgl-router"
|
||||
)
|
||||
|
||||
|
||||
def _get_open_port() -> int:
|
||||
"""Reserve an ephemeral TCP port in [20000, 55535].
|
||||
|
||||
The router itself doesn't have the ``port + 10000`` gRPC-derivation
|
||||
constraint that SGLang's launch_server does, but we cap the range
|
||||
anyway so the e2e helpers behave consistently across components.
|
||||
"""
|
||||
for _ in range(50):
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
port = s.getsockname()[1]
|
||||
if 20000 <= port <= 55535:
|
||||
return port
|
||||
raise RuntimeError(
|
||||
"could not allocate an ephemeral port in [20000, 55535] after 50 tries"
|
||||
)
|
||||
|
||||
|
||||
def _resolve_tokenizer_path(tokenizer_path: str) -> str:
|
||||
"""Resolve a HuggingFace repo ID to a local ``tokenizer.json`` path.
|
||||
|
||||
sgl-router's tokenizer loader treats the input as a filesystem path and
|
||||
inspects its extension; a bare HF id like ``Qwen/Qwen3-0.6B`` looks
|
||||
like a file with extension ``.6B`` and is rejected. When the HF Hub
|
||||
cache already has the tokenizer, point the loader at the on-disk
|
||||
``tokenizer.json`` directly. Pass paths/URLs through unchanged.
|
||||
"""
|
||||
p = Path(tokenizer_path)
|
||||
if p.exists():
|
||||
return str(p)
|
||||
try:
|
||||
from huggingface_hub import try_to_load_from_cache # type: ignore[import]
|
||||
|
||||
cached = try_to_load_from_cache(tokenizer_path, "tokenizer.json")
|
||||
if cached and Path(cached).is_file():
|
||||
return str(cached)
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
return tokenizer_path
|
||||
|
||||
|
||||
@dataclass
|
||||
class WorkerInfo:
|
||||
"""Worker visible to the gateway via ``/v1/models``-style introspection.
|
||||
|
||||
Mirrors SMG's WorkerInfo shape so test code reads the same. sgl-router
|
||||
does not currently surface a `/v1/workers` admin API — this is a
|
||||
placeholder for a future admin surface; current tests scrape
|
||||
`/metrics` for per-worker observability instead.
|
||||
"""
|
||||
|
||||
id: str
|
||||
url: str
|
||||
model: str | None = None
|
||||
status: str = "unknown"
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
class Gateway:
|
||||
"""Lifecycle-managed sgl-router instance for e2e tests.
|
||||
|
||||
Not thread-safe; assume one Gateway per test (or per fixture scope).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str = "127.0.0.1",
|
||||
port: int | None = None,
|
||||
binary: Path | None = None,
|
||||
proxy_request_timeout_secs: int | None = None,
|
||||
stale_request_timeout_secs: int | None = None,
|
||||
):
|
||||
self.host = host
|
||||
self.port = port or _get_open_port()
|
||||
self.base_url = f"http://{self.host}:{self.port}"
|
||||
# Resolve binary from env override, explicit arg, or repo default.
|
||||
env_binary = os.environ.get("SGL_ROUTER_BINARY")
|
||||
if binary is not None:
|
||||
self.binary = Path(binary)
|
||||
elif env_binary:
|
||||
self.binary = Path(env_binary)
|
||||
else:
|
||||
self.binary = DEFAULT_BINARY
|
||||
|
||||
# Test-side overrides for the router's tunables. Both default to
|
||||
# `None`, in which case the router uses its production defaults
|
||||
# (60 s proxy timeout, 300 s stale-request timeout). Tests set
|
||||
# these short so per-request failures and stale-request expiry
|
||||
# surface within the test's wall-time budget.
|
||||
self.proxy_request_timeout_secs = proxy_request_timeout_secs
|
||||
self.stale_request_timeout_secs = stale_request_timeout_secs
|
||||
|
||||
self.process: subprocess.Popen | None = None
|
||||
self._config_path: Path | None = None
|
||||
self._started: bool = False
|
||||
# Track child workers we spawned so __exit__ can tear them down.
|
||||
self._owned_workers: list[subprocess.Popen] = []
|
||||
|
||||
# ----- context manager -------------------------------------------------
|
||||
|
||||
def __enter__(self) -> "Gateway":
|
||||
return self
|
||||
|
||||
def __exit__(self, *exc) -> None:
|
||||
self.shutdown()
|
||||
|
||||
# ----- start ----------------------------------------------------------
|
||||
|
||||
def start_regular(
|
||||
self,
|
||||
*,
|
||||
model_id: str,
|
||||
tokenizer_path: str,
|
||||
worker_urls: list[str],
|
||||
policy: str = "round_robin",
|
||||
extra_models: list[dict] | None = None,
|
||||
timeout: float = 60.0,
|
||||
) -> None:
|
||||
"""Start the router in regular (non-PD) mode.
|
||||
|
||||
Args:
|
||||
model_id: The model identifier the router will dispatch under.
|
||||
tokenizer_path: Path or HF ID for the tokenizer the router uses
|
||||
for cache-aware tokenization.
|
||||
worker_urls: URLs of already-running ``sglang.launch_server``
|
||||
instances. The router uses ``static_urls`` discovery;
|
||||
each worker's mode (plain) and any disaggregation
|
||||
metadata are learned from ``/server_info``.
|
||||
policy: Policy kind — ``round_robin``, ``random``, ``power_of_two``,
|
||||
or ``cache_aware_zmq``.
|
||||
timeout: How long to wait for ``/readyz`` before giving up.
|
||||
"""
|
||||
self._launch(
|
||||
self._build_config(
|
||||
model_id=model_id,
|
||||
tokenizer_path=tokenizer_path,
|
||||
urls=list(worker_urls),
|
||||
policy=policy,
|
||||
extra_models=extra_models or [],
|
||||
),
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
def start_pd(
|
||||
self,
|
||||
*,
|
||||
model_id: str,
|
||||
tokenizer_path: str,
|
||||
prefill_urls: list[str],
|
||||
decode_urls: list[str],
|
||||
policy: str = "round_robin",
|
||||
timeout: float = 60.0,
|
||||
) -> None:
|
||||
"""Start the router in PD-disaggregated mode.
|
||||
|
||||
All prefill + decode URLs go into one ``static_urls`` list. The
|
||||
router seeds each worker as ``WorkerMode::Plain`` and the
|
||||
manager's ``/server_info`` introspect step overrides mode +
|
||||
``bootstrap_port`` from the worker's self-disclosure. Workers
|
||||
must have been launched with ``--disaggregation-mode`` and
|
||||
``--disaggregation-bootstrap-port`` for the PD role to be
|
||||
picked up (see ``model_pool.spawn_worker``); modern SGLang is
|
||||
assumed.
|
||||
"""
|
||||
self._launch(
|
||||
self._build_config(
|
||||
model_id=model_id,
|
||||
tokenizer_path=tokenizer_path,
|
||||
urls=list(prefill_urls) + list(decode_urls),
|
||||
policy=policy,
|
||||
extra_models=[],
|
||||
),
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
# ----- shutdown --------------------------------------------------------
|
||||
|
||||
def shutdown(self) -> None:
|
||||
"""SIGTERM the router; SIGKILL after 30s. Idempotent."""
|
||||
if self.process is not None and self.process.poll() is None:
|
||||
try:
|
||||
self.process.send_signal(signal.SIGTERM)
|
||||
try:
|
||||
self.process.wait(timeout=30)
|
||||
except subprocess.TimeoutExpired:
|
||||
self.process.kill()
|
||||
self.process.wait()
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
self.process = None
|
||||
if self._config_path and self._config_path.exists():
|
||||
self._config_path.unlink(missing_ok=True)
|
||||
self._config_path = None
|
||||
self._started = False
|
||||
# Tear down any owned upstream workers.
|
||||
for w in self._owned_workers:
|
||||
if w.poll() is None:
|
||||
try:
|
||||
w.send_signal(signal.SIGTERM)
|
||||
try:
|
||||
w.wait(timeout=30)
|
||||
except subprocess.TimeoutExpired:
|
||||
w.kill()
|
||||
w.wait()
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
self._owned_workers.clear()
|
||||
|
||||
# ----- HTTP introspection helpers -------------------------------------
|
||||
|
||||
def healthy(self, timeout: float = 5.0) -> bool:
|
||||
try:
|
||||
resp = httpx.get(f"{self.base_url}/healthz", timeout=timeout)
|
||||
return resp.status_code == 200
|
||||
except (httpx.RequestError, httpx.TimeoutException):
|
||||
return False
|
||||
|
||||
def ready(self, timeout: float = 5.0) -> bool:
|
||||
try:
|
||||
resp = httpx.get(f"{self.base_url}/readyz", timeout=timeout)
|
||||
return resp.status_code == 200
|
||||
except (httpx.RequestError, httpx.TimeoutException):
|
||||
return False
|
||||
|
||||
def metrics_text(self, timeout: float = 5.0) -> str | None:
|
||||
try:
|
||||
resp = httpx.get(f"{self.base_url}/metrics", timeout=timeout)
|
||||
if resp.status_code == 200:
|
||||
return resp.text
|
||||
return None
|
||||
except (httpx.RequestError, httpx.TimeoutException):
|
||||
return None
|
||||
|
||||
# ----- internals ------------------------------------------------------
|
||||
|
||||
def _build_config(
|
||||
self,
|
||||
*,
|
||||
model_id: str,
|
||||
tokenizer_path: str,
|
||||
urls: list[str],
|
||||
policy: str,
|
||||
extra_models: list[dict],
|
||||
) -> str:
|
||||
resolved_tokenizer = _resolve_tokenizer_path(tokenizer_path)
|
||||
|
||||
extra_model_toml = ""
|
||||
for em in extra_models:
|
||||
extra_model_toml += (
|
||||
f'\n[[models]]\nid = "{em["id"]}"\n'
|
||||
f'tokenizer_path = "{_resolve_tokenizer_path(em["tokenizer_path"])}"\n'
|
||||
f'policy = "{em.get("policy", policy)}"\n'
|
||||
)
|
||||
|
||||
# Optional tunables — only emit the [proxy] and [active_load]
|
||||
# sections if a test has overridden them, so production defaults
|
||||
# apply otherwise.
|
||||
proxy_section = ""
|
||||
if self.proxy_request_timeout_secs is not None:
|
||||
proxy_section = (
|
||||
f"\n[proxy]\nrequest_timeout_secs = {self.proxy_request_timeout_secs}\n"
|
||||
)
|
||||
active_load_section = ""
|
||||
if self.stale_request_timeout_secs is not None:
|
||||
active_load_section = (
|
||||
f"\n[active_load]\nstale_request_timeout_secs = "
|
||||
f"{self.stale_request_timeout_secs}\n"
|
||||
)
|
||||
|
||||
urls_toml = ", ".join(f'"{u}"' for u in urls)
|
||||
|
||||
return f"""\
|
||||
[server]
|
||||
host = "{self.host}"
|
||||
port = {self.port}
|
||||
|
||||
[[models]]
|
||||
id = "{model_id}"
|
||||
tokenizer_path = "{resolved_tokenizer}"
|
||||
policy = "{policy}"
|
||||
{extra_model_toml}
|
||||
|
||||
[discovery]
|
||||
backend = "static_urls"
|
||||
|
||||
[discovery.static_urls]
|
||||
urls = [{urls_toml}]
|
||||
{proxy_section}{active_load_section}"""
|
||||
|
||||
def _launch(self, config_text: str, *, timeout: float) -> None:
|
||||
if not self.binary.exists():
|
||||
raise RuntimeError(
|
||||
f"sgl-router binary not found at {self.binary}. "
|
||||
"Build it first: `cd experimental/sgl-router && cargo build --release` "
|
||||
"or set SGL_ROUTER_BINARY to the binary path."
|
||||
)
|
||||
# Write the main config.
|
||||
fd, path = tempfile.mkstemp(suffix=".toml", prefix="sgl-router-")
|
||||
os.close(fd)
|
||||
self._config_path = Path(path)
|
||||
self._config_path.write_text(config_text, encoding="utf-8")
|
||||
logger.info("sgl-router config: %s", self._config_path)
|
||||
logger.debug("sgl-router config text:\n%s", config_text)
|
||||
|
||||
self.process = subprocess.Popen(
|
||||
[str(self.binary), "--config", str(self._config_path)],
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
start_new_session=True,
|
||||
)
|
||||
|
||||
try:
|
||||
self._wait_ready(timeout=timeout)
|
||||
except Exception:
|
||||
self.shutdown()
|
||||
raise
|
||||
self._started = True
|
||||
|
||||
def _wait_ready(self, *, timeout: float) -> None:
|
||||
deadline = time.time() + timeout
|
||||
last_exc: Exception | None = None
|
||||
while time.time() < deadline:
|
||||
if self.process is not None and self.process.poll() is not None:
|
||||
# Process exited early — surface stdout/stderr.
|
||||
out = b""
|
||||
try:
|
||||
if self.process.stdout is not None:
|
||||
out = self.process.stdout.read() or b""
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
raise RuntimeError(
|
||||
f"sgl-router exited during startup with code "
|
||||
f"{self.process.returncode}. output:\n{out.decode(errors='replace')}",
|
||||
)
|
||||
try:
|
||||
resp = httpx.get(f"{self.base_url}/readyz", timeout=2.0)
|
||||
if resp.status_code == 200:
|
||||
return
|
||||
except (httpx.RequestError, httpx.TimeoutException) as exc:
|
||||
last_exc = exc
|
||||
time.sleep(0.5)
|
||||
raise TimeoutError(
|
||||
f"sgl-router did not become ready at {self.base_url} within {timeout}s "
|
||||
f"(last error: {last_exc})"
|
||||
)
|
||||
@@ -0,0 +1,228 @@
|
||||
"""Minimal SGLang worker spawner for sgl-router e2e tests.
|
||||
|
||||
Adapted from SMG's e2e_test/infra/model_pool.py — the 1200-line original
|
||||
manages a pool of long-lived workers across many tests; here we only
|
||||
need a thin wrapper around ``sglang.launch_server`` that:
|
||||
|
||||
- allocates GPU(s) for the worker (via ``CUDA_VISIBLE_DEVICES``),
|
||||
- spawns ``python3 -m sglang.launch_server`` with the right args,
|
||||
- waits for ``/health`` to come up,
|
||||
- optionally injects ``--kv-events-config`` so the worker exposes
|
||||
the ``kv_events`` block on ``/server_info``.
|
||||
|
||||
A test owns a ``ModelInstance`` for its duration; teardown shuts the
|
||||
worker down. No cross-test pooling — the acceptance tests are slow
|
||||
enough already (model load dominates) that pooling complexity wasn't
|
||||
worth porting.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import socket
|
||||
import subprocess
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
import httpx
|
||||
|
||||
from .model_specs import get_model_spec
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Passthrough Jinja chat template that emits ONLY `messages[*].content`
|
||||
# joined with `\n` — matching the router's cache_aware_zmq prompt
|
||||
# extraction. A worker launched with
|
||||
# ``--chat-template <PASSTHROUGH_CHAT_TEMPLATE_PATH>`` tokenizes the
|
||||
# raw content string, so its KV-block hashes align with what the
|
||||
# router computes from the same chat-completions request. Test-only.
|
||||
PASSTHROUGH_CHAT_TEMPLATE_PATH = str(
|
||||
Path(__file__).parent / "passthrough_chat_template.jinja"
|
||||
)
|
||||
|
||||
|
||||
def _get_open_port() -> int:
|
||||
"""Allocate an ephemeral TCP port in the range [20000, 55535].
|
||||
|
||||
SGLang derives its internal gRPC port as ``http_port + 10000``; if the
|
||||
kernel hands us an ephemeral port above 55535, that derivation overflows
|
||||
65535 and ``ServerArgs.__post_init__`` rejects it. Retrying a bounded
|
||||
number of times keeps us safely below the ceiling without hand-rolling
|
||||
a port registry.
|
||||
"""
|
||||
for _ in range(50):
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
port = s.getsockname()[1]
|
||||
if 20000 <= port <= 55535:
|
||||
return port
|
||||
raise RuntimeError(
|
||||
"could not allocate an ephemeral port in [20000, 55535] after 50 tries; "
|
||||
"SGLang derives its internal gRPC port as http_port + 10000 and "
|
||||
"rejects values above 65535"
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelInstance:
|
||||
"""A running ``sglang.launch_server`` process.
|
||||
|
||||
Use as a context manager:
|
||||
|
||||
with spawn_worker("qwen3-0.6b", gpu_ids=[0]) as inst:
|
||||
httpx.post(f"{inst.url}/generate", ...)
|
||||
"""
|
||||
|
||||
url: str
|
||||
port: int
|
||||
process: subprocess.Popen
|
||||
model_id: str
|
||||
gpu_ids: list[int] = field(default_factory=list)
|
||||
kv_events_endpoint: str | None = None
|
||||
|
||||
def __enter__(self) -> "ModelInstance":
|
||||
return self
|
||||
|
||||
def __exit__(self, *exc) -> None:
|
||||
self.shutdown()
|
||||
|
||||
def shutdown(self) -> None:
|
||||
if self.process is not None and self.process.poll() is None:
|
||||
try:
|
||||
self.process.send_signal(signal.SIGTERM)
|
||||
try:
|
||||
self.process.wait(timeout=60)
|
||||
except subprocess.TimeoutExpired:
|
||||
self.process.kill()
|
||||
self.process.wait()
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
|
||||
|
||||
def spawn_worker(
|
||||
model_id: str,
|
||||
*,
|
||||
gpu_ids: list[int],
|
||||
port: int | None = None,
|
||||
enable_kv_events: bool = False,
|
||||
kv_events_port: int | None = None,
|
||||
disagg_mode: str | None = None,
|
||||
bootstrap_port: int | None = None,
|
||||
extra_args: list[str] | None = None,
|
||||
timeout: float = 600.0,
|
||||
) -> ModelInstance:
|
||||
"""Spawn a single ``sglang.launch_server`` and wait for ``/health``.
|
||||
|
||||
Args:
|
||||
model_id: Key into :data:`model_specs.MODEL_SPECS`.
|
||||
gpu_ids: Concrete GPU indices to bind via ``CUDA_VISIBLE_DEVICES``.
|
||||
port: HTTP port; auto-assigned if None.
|
||||
enable_kv_events: If True, inject ``--kv-events-config`` with a
|
||||
ZMQ publisher so the router's introspection picks up the
|
||||
kv_events block from ``/server_info`` (Patch 1).
|
||||
kv_events_port: ZMQ publisher port. Auto-assigned if None and
|
||||
``enable_kv_events`` is True.
|
||||
disagg_mode: "prefill" or "decode" for PD-disagg launches; passed
|
||||
through as ``--disaggregation-mode``.
|
||||
bootstrap_port: PD-disagg bootstrap port (prefill side only).
|
||||
extra_args: Additional CLI args appended verbatim.
|
||||
timeout: Health-check timeout. Cold-start on a fresh GPU can be
|
||||
slow; default is 10 minutes.
|
||||
"""
|
||||
spec = get_model_spec(model_id)
|
||||
port = port or _get_open_port()
|
||||
base_url = f"http://127.0.0.1:{port}"
|
||||
|
||||
cmd = [
|
||||
"python3",
|
||||
"-m",
|
||||
"sglang.launch_server",
|
||||
"--model-path",
|
||||
spec["model"],
|
||||
"--port",
|
||||
str(port),
|
||||
"--host",
|
||||
"127.0.0.1",
|
||||
"--tp",
|
||||
str(spec.get("tp", 1)),
|
||||
]
|
||||
cmd.extend(spec.get("worker_args", []) or [])
|
||||
|
||||
kv_events_endpoint: str | None = None
|
||||
if enable_kv_events:
|
||||
kv_port = kv_events_port or _get_open_port()
|
||||
kv_events_endpoint = f"tcp://*:{kv_port}"
|
||||
kv_cfg = {
|
||||
"publisher": "zmq",
|
||||
"endpoint": kv_events_endpoint,
|
||||
"topic": "kv",
|
||||
}
|
||||
cmd.extend(["--kv-events-config", json.dumps(kv_cfg)])
|
||||
|
||||
if disagg_mode is not None:
|
||||
cmd.extend(["--disaggregation-mode", disagg_mode])
|
||||
if bootstrap_port is not None:
|
||||
cmd.extend(["--disaggregation-bootstrap-port", str(bootstrap_port)])
|
||||
|
||||
if extra_args:
|
||||
cmd.extend(extra_args)
|
||||
|
||||
env = os.environ.copy()
|
||||
env["CUDA_VISIBLE_DEVICES"] = ",".join(str(g) for g in gpu_ids)
|
||||
logger.info(
|
||||
"spawning sglang worker: model=%s port=%d gpus=%s disagg=%s",
|
||||
model_id,
|
||||
port,
|
||||
gpu_ids,
|
||||
disagg_mode,
|
||||
)
|
||||
|
||||
proc = subprocess.Popen(
|
||||
cmd,
|
||||
env=env,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
start_new_session=True,
|
||||
)
|
||||
|
||||
inst = ModelInstance(
|
||||
url=base_url,
|
||||
port=port,
|
||||
process=proc,
|
||||
model_id=model_id,
|
||||
gpu_ids=list(gpu_ids),
|
||||
kv_events_endpoint=kv_events_endpoint,
|
||||
)
|
||||
|
||||
# Wait for /health. Cold-start on H200 with weights uncached can take
|
||||
# ~5 minutes; CI configurations should pre-warm.
|
||||
deadline = time.time() + timeout
|
||||
while time.time() < deadline:
|
||||
if proc.poll() is not None:
|
||||
out = b""
|
||||
try:
|
||||
if proc.stdout is not None:
|
||||
out = proc.stdout.read() or b""
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
raise RuntimeError(
|
||||
f"sglang worker exited during startup with code {proc.returncode}; "
|
||||
f"cmd: {' '.join(cmd)}\noutput:\n{out.decode(errors='replace')}",
|
||||
)
|
||||
try:
|
||||
resp = httpx.get(f"{base_url}/health", timeout=2.0)
|
||||
if resp.status_code == 200:
|
||||
logger.info("sglang worker ready at %s", base_url)
|
||||
return inst
|
||||
except (httpx.RequestError, httpx.TimeoutException):
|
||||
pass
|
||||
time.sleep(2.0)
|
||||
|
||||
inst.shutdown()
|
||||
raise TimeoutError(
|
||||
f"sglang worker did not become healthy at {base_url} within {timeout}s",
|
||||
)
|
||||
@@ -0,0 +1,77 @@
|
||||
"""Model specifications for sgl-router e2e tests.
|
||||
|
||||
Adapted from SMG's e2e_test/infra/model_specs.py. The same dict-of-dicts
|
||||
shape (so test code reads the same) but the entries are narrower —
|
||||
sgl-router tests today target small/medium models only; the larger
|
||||
function-calling / reasoning models from SMG are out of scope.
|
||||
|
||||
Each entry:
|
||||
- model: HuggingFace path or local path (env-resolved)
|
||||
- memory_gb: estimated single-GPU footprint
|
||||
- tp: tensor-parallel size (= GPUs needed)
|
||||
- features: feature tags for filtering
|
||||
- worker_args: optional extra `sglang.launch_server` flags
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
# Local-cache root for CI / cluster nodes that pre-download HF weights.
|
||||
# Mirrors the SMG `ROUTER_LOCAL_MODEL_PATH` env var.
|
||||
ROUTER_LOCAL_MODEL_PATH = os.environ.get("ROUTER_LOCAL_MODEL_PATH", "")
|
||||
|
||||
|
||||
def _resolve_model_path(hf_path: str) -> str:
|
||||
"""Prefer a local copy of the model when one exists under
|
||||
``ROUTER_LOCAL_MODEL_PATH``; otherwise fall back to the HuggingFace ID.
|
||||
"""
|
||||
if ROUTER_LOCAL_MODEL_PATH:
|
||||
local_path = os.path.join(ROUTER_LOCAL_MODEL_PATH, hf_path)
|
||||
if os.path.exists(local_path):
|
||||
return local_path
|
||||
return hf_path
|
||||
|
||||
|
||||
MODEL_SPECS: dict[str, dict] = {
|
||||
# Fast-start tiny model for convergence / decode-affinity / stale-request
|
||||
# tests. Single GPU, ~2 GB weights, sub-30s start on a warm cache.
|
||||
"qwen3-0.6b": {
|
||||
"model": _resolve_model_path("Qwen/Qwen3-0.6B"),
|
||||
"memory_gb": 4,
|
||||
"tp": 1,
|
||||
"features": ["chat", "streaming"],
|
||||
},
|
||||
# Standard small chat model — matches SMG's `llama-1b` entry.
|
||||
"llama-1b": {
|
||||
"model": _resolve_model_path("meta-llama/Llama-3.2-1B-Instruct"),
|
||||
"memory_gb": 4,
|
||||
"tp": 1,
|
||||
"features": ["chat", "streaming"],
|
||||
},
|
||||
# Primary 8B chat model — matches SMG's `llama-8b`.
|
||||
"llama-8b": {
|
||||
"model": _resolve_model_path("meta-llama/Llama-3.1-8B-Instruct"),
|
||||
"memory_gb": 16,
|
||||
"tp": 1,
|
||||
"features": ["chat", "streaming"],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def get_model_spec(model_id: str) -> dict:
|
||||
"""Return the spec dict for ``model_id``; KeyError if absent."""
|
||||
if model_id not in MODEL_SPECS:
|
||||
raise KeyError(
|
||||
f"Unknown model: {model_id}. Available: {list(MODEL_SPECS.keys())}"
|
||||
)
|
||||
return MODEL_SPECS[model_id]
|
||||
|
||||
|
||||
def get_models_with_feature(feature: str) -> list[str]:
|
||||
"""Filter model IDs by feature tag (e.g. ``streaming``, ``chat``)."""
|
||||
return [
|
||||
model_id
|
||||
for model_id, spec in MODEL_SPECS.items()
|
||||
if feature in spec.get("features", [])
|
||||
]
|
||||
@@ -0,0 +1,13 @@
|
||||
{#-
|
||||
Passthrough chat template for cache-aware-zmq e2e tests.
|
||||
|
||||
Emits ONLY `messages[*].content` joined with `\n` — no role markers,
|
||||
no special tokens, no generation prompt. This is the SAME shape the
|
||||
router's cache_aware_zmq policy produces in `extract_prompt_text`,
|
||||
so a worker launched with `--chat-template <this file>` tokenizes the
|
||||
same string the router will tokenize for routing — making block
|
||||
hashes align across worker KV cache and router HashTree.
|
||||
|
||||
Use only for tests; not appropriate for any real chat workload.
|
||||
-#}
|
||||
{{- messages | map(attribute='content') | join('\n') -}}
|
||||
@@ -0,0 +1,6 @@
|
||||
FROM python:3.12-slim
|
||||
WORKDIR /app
|
||||
RUN pip install --no-cache-dir fastapi uvicorn
|
||||
COPY fake_worker.py .
|
||||
EXPOSE 30000
|
||||
CMD ["python", "fake_worker.py"]
|
||||
@@ -0,0 +1,39 @@
|
||||
# syntax=docker/dockerfile:1.6
|
||||
# Build sgl-router binary for k8s integration E2E.
|
||||
# Context root: repo root (one level above experimental/sgl-router/).
|
||||
|
||||
# Matches rust-toolchain.toml's pinned channel, avoiding an in-build rustup channel-sync.
|
||||
FROM rust:1.90-bookworm AS builder
|
||||
|
||||
# Pin to the exact toolchain pre-installed in the base image so rustup
|
||||
# doesn't try to sync the channel manifest when it sees rust-toolchain.toml's
|
||||
# `channel = "1.90"`.
|
||||
ENV RUSTUP_TOOLCHAIN=1.90.0
|
||||
|
||||
# libssl-dev + pkg-config ship with rust:1.90-bookworm already; no apt-get needed.
|
||||
|
||||
WORKDIR /build
|
||||
|
||||
# Copy just the sgl-router crate (context is the repo root)
|
||||
COPY experimental/sgl-router /build/experimental/sgl-router
|
||||
|
||||
RUN --mount=type=cache,target=/usr/local/cargo/registry \
|
||||
--mount=type=cache,target=/usr/local/cargo/git \
|
||||
--mount=type=cache,target=/build/experimental/sgl-router/target \
|
||||
cd /build/experimental/sgl-router \
|
||||
&& cargo build --release --bin sgl-router \
|
||||
&& cp target/release/sgl-router /usr/local/bin/sgl-router
|
||||
|
||||
FROM debian:bookworm-slim
|
||||
|
||||
RUN apt-get update && apt-get install -y ca-certificates \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY --from=builder /usr/local/bin/sgl-router /usr/local/bin/sgl-router
|
||||
|
||||
# Tiny tokenizer fixture used by the E2E config
|
||||
COPY experimental/sgl-router/tests/fixtures/tiny_tokenizer.json /etc/tokenizer/tiny.json
|
||||
|
||||
EXPOSE 8090
|
||||
|
||||
ENTRYPOINT ["sgl-router"]
|
||||
@@ -0,0 +1,251 @@
|
||||
"""Pytest configuration for sgl-router K8s integration tests.
|
||||
|
||||
These tests require:
|
||||
- A kind cluster named 'sgl-router-kind'
|
||||
- The sgl-router:e2e and sgl-router-fake-worker:e2e images loaded into kind
|
||||
- kubectl configured to use the kind-sgl-router-kind context
|
||||
|
||||
Setup: ./tests/e2e/k8s_integration/setup.sh
|
||||
Teardown: ./tests/e2e/k8s_integration/setup.sh teardown
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import socket
|
||||
import subprocess
|
||||
import time
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
NAMESPACE = "sgl-router-test"
|
||||
CLUSTER_NAME = "sgl-router-kind"
|
||||
KUBECTL_CONTEXT = f"kind-{CLUSTER_NAME}"
|
||||
|
||||
# sgl-router discovery reconciliation: if the watcher misses an event the
|
||||
# reconciler fires within ~60s. Tests that exercise removal wait up to 90s.
|
||||
RECONCILIATION_WAIT_SECS = 90
|
||||
|
||||
# Errors safe to retry while polling (transport-level only — HTTP 4xx/5xx
|
||||
# are intentionally NOT included so real regressions surface immediately).
|
||||
_TRANSIENT_ERRORS = (
|
||||
httpx.TransportError,
|
||||
httpx.TimeoutException,
|
||||
ConnectionError,
|
||||
OSError,
|
||||
)
|
||||
|
||||
|
||||
def pytest_configure(config):
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"slow: marks tests that wait for multiple reconciliation cycles "
|
||||
"(deselect with '-m \"not slow\"')",
|
||||
)
|
||||
|
||||
|
||||
def _kubectl(
|
||||
*args: str,
|
||||
check: bool = True,
|
||||
capture: bool = True,
|
||||
) -> subprocess.CompletedProcess:
|
||||
cmd = ["kubectl", "--context", KUBECTL_CONTEXT, *args]
|
||||
logger.debug("Running: %s", " ".join(cmd))
|
||||
return subprocess.run(cmd, capture_output=capture, text=True, check=check)
|
||||
|
||||
|
||||
def _apply_from_stdin(yaml_content: str) -> subprocess.CompletedProcess:
|
||||
return subprocess.run(
|
||||
["kubectl", "--context", KUBECTL_CONTEXT, "apply", "-f", "-"],
|
||||
input=yaml_content,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
|
||||
|
||||
def _wait_for_deployment_ready(
|
||||
name: str,
|
||||
namespace: str = NAMESPACE,
|
||||
timeout: int = 180,
|
||||
) -> None:
|
||||
_kubectl(
|
||||
"rollout",
|
||||
"status",
|
||||
f"deployment/{name}",
|
||||
"-n",
|
||||
namespace,
|
||||
f"--timeout={timeout}s",
|
||||
)
|
||||
|
||||
|
||||
def _wait_for_pod_ready(
|
||||
name: str,
|
||||
namespace: str = NAMESPACE,
|
||||
timeout: int = 120,
|
||||
) -> None:
|
||||
_kubectl(
|
||||
"wait",
|
||||
"--for=condition=Ready",
|
||||
f"pod/{name}",
|
||||
"-n",
|
||||
namespace,
|
||||
f"--timeout={timeout}s",
|
||||
)
|
||||
|
||||
|
||||
def _wait_for_port(port: int, proc: subprocess.Popen, timeout: int = 15) -> None:
|
||||
"""Poll until a TCP connection to localhost:port succeeds."""
|
||||
deadline = time.time() + timeout
|
||||
while time.time() < deadline:
|
||||
if proc.poll() is not None:
|
||||
stderr = proc.stderr.read().decode() if proc.stderr else ""
|
||||
raise RuntimeError(f"port-forward process exited early: {stderr}")
|
||||
try:
|
||||
with socket.create_connection(("127.0.0.1", port), timeout=1):
|
||||
return
|
||||
except OSError:
|
||||
time.sleep(0.5)
|
||||
raise TimeoutError(f"Port {port} not ready after {timeout}s")
|
||||
|
||||
|
||||
def _port_forward_start(
|
||||
namespace: str,
|
||||
service: str,
|
||||
local_port: int,
|
||||
remote_port: int,
|
||||
) -> subprocess.Popen:
|
||||
"""Start kubectl port-forward and wait until the port is reachable."""
|
||||
cmd = [
|
||||
"kubectl",
|
||||
"--context",
|
||||
KUBECTL_CONTEXT,
|
||||
"port-forward",
|
||||
f"svc/{service}",
|
||||
f"{local_port}:{remote_port}",
|
||||
"-n",
|
||||
namespace,
|
||||
]
|
||||
logger.info("Starting port-forward: %s", " ".join(cmd))
|
||||
proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
||||
_wait_for_port(local_port, proc)
|
||||
return proc
|
||||
|
||||
|
||||
def _cleanup_port_forward(name: str, pf: subprocess.Popen) -> None:
|
||||
try:
|
||||
pf.terminate()
|
||||
pf.wait(timeout=10)
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.warning(
|
||||
"Port-forward %s did not exit on SIGTERM after 10s; killing", name
|
||||
)
|
||||
pf.kill()
|
||||
try:
|
||||
pf.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.warning("Port-forward %s still running after SIGKILL", name)
|
||||
except Exception as exc:
|
||||
logger.warning("Error cleaning up %s port-forward: %s", name, exc)
|
||||
|
||||
rc = pf.returncode
|
||||
stderr = pf.stderr.read().decode() if pf.stderr else ""
|
||||
if rc != -15:
|
||||
suffix = f": {stderr.strip()}" if stderr.strip() else ""
|
||||
logger.warning("Port-forward %s exited rc=%s%s", name, rc, suffix)
|
||||
else:
|
||||
logger.debug("Port-forward %s exited cleanly (rc=%s)", name, rc)
|
||||
|
||||
|
||||
def _poll_until(
|
||||
predicate,
|
||||
description: str,
|
||||
timeout: int,
|
||||
interval: float = 5,
|
||||
) -> bool:
|
||||
"""Poll predicate until True, or raise TimeoutError.
|
||||
|
||||
Only transient network errors are retried; HTTP status errors and
|
||||
programming errors propagate immediately.
|
||||
"""
|
||||
deadline = time.time() + timeout
|
||||
last_error = None
|
||||
attempts = 0
|
||||
while time.time() < deadline:
|
||||
try:
|
||||
attempts += 1
|
||||
if predicate():
|
||||
logger.info(
|
||||
"Condition met: %s (after %d attempts)", description, attempts
|
||||
)
|
||||
return True
|
||||
except _TRANSIENT_ERRORS as exc:
|
||||
last_error = exc
|
||||
logger.debug("Transient error on attempt %d: %s", attempts, exc)
|
||||
time.sleep(interval)
|
||||
msg = f"Timeout waiting for: {description} (after {timeout}s, {attempts} attempts)"
|
||||
if last_error:
|
||||
msg += f" — last error: {last_error}"
|
||||
raise TimeoutError(msg)
|
||||
|
||||
|
||||
def _get_router_url(router_base: str) -> str:
|
||||
return router_base
|
||||
|
||||
|
||||
def _router_is_healthy(router_base: str) -> bool:
|
||||
try:
|
||||
r = httpx.get(f"{router_base}/healthz", timeout=3.0)
|
||||
return r.status_code == 200
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def k8s_cluster():
|
||||
"""Assert the kind cluster exists and kubectl context is reachable."""
|
||||
result = subprocess.run(
|
||||
["kind", "get", "clusters"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
if CLUSTER_NAME not in result.stdout.splitlines():
|
||||
pytest.skip(
|
||||
f"kind cluster '{CLUSTER_NAME}' not found — run "
|
||||
f"./tests/e2e/k8s_integration/setup.sh first"
|
||||
)
|
||||
_kubectl("cluster-info")
|
||||
return True
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def router_port_forward(k8s_cluster):
|
||||
"""Per-test port-forward to sgl-router service.
|
||||
|
||||
Function-scoped because some tests (notably
|
||||
test_lifecycle.TestRouterRestart) force-delete the router pod;
|
||||
a session-scoped port-forward would be bound to the deleted pod's
|
||||
network namespace and stay dead for all subsequent tests in the
|
||||
suite. Per-test setup costs ~1-2s.
|
||||
"""
|
||||
_wait_for_deployment_ready("sgl-router")
|
||||
pf = _port_forward_start(NAMESPACE, "sgl-router", 8090, 8090)
|
||||
try:
|
||||
_poll_until(
|
||||
lambda: _router_is_healthy("http://127.0.0.1:8090"),
|
||||
"sgl-router /healthz returns 200",
|
||||
timeout=30,
|
||||
interval=1,
|
||||
)
|
||||
yield "http://127.0.0.1:8090"
|
||||
finally:
|
||||
_cleanup_port_forward("sgl-router", pf)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def router_url(router_port_forward):
|
||||
return router_port_forward
|
||||
@@ -0,0 +1,73 @@
|
||||
"""Minimal fake SGLang worker for kind E2E integration testing.
|
||||
|
||||
Responds to:
|
||||
GET /health -> {"status": "ok"}
|
||||
GET /server_info -> {"served_model_name": MODEL_ID}
|
||||
GET /v1/models -> list with a single MODEL_ID model entry
|
||||
POST /v1/chat/completions -> echoes the last user message back
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
import uvicorn
|
||||
from fastapi import FastAPI, Request
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
MODEL_ID = os.environ.get("MODEL_ID", "tiny")
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
async def health():
|
||||
return {"status": "ok"}
|
||||
|
||||
|
||||
@app.get("/server_info")
|
||||
async def server_info():
|
||||
# The sgl-router worker manager fetches this on every Added event and
|
||||
# uses `served_model_name` to populate the registry's model index.
|
||||
return {"served_model_name": MODEL_ID}
|
||||
|
||||
|
||||
@app.get("/v1/models")
|
||||
async def models():
|
||||
return {
|
||||
"object": "list",
|
||||
"data": [
|
||||
{
|
||||
"id": MODEL_ID,
|
||||
"object": "model",
|
||||
"created": 0,
|
||||
"owned_by": "sglang",
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
@app.post("/v1/chat/completions")
|
||||
async def chat_completions(request: Request):
|
||||
payload = await request.json()
|
||||
messages = payload.get("messages", [])
|
||||
last_content = messages[-1]["content"] if messages else ""
|
||||
return {
|
||||
"id": "chatcmpl-mock",
|
||||
"object": "chat.completion",
|
||||
"model": payload.get("model", MODEL_ID),
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": f"echo: {last_content}",
|
||||
},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
uvicorn.run(app, host="0.0.0.0", port=30000)
|
||||
@@ -0,0 +1,4 @@
|
||||
apiVersion: v1
|
||||
kind: Namespace
|
||||
metadata:
|
||||
name: sgl-router-test
|
||||
@@ -0,0 +1,33 @@
|
||||
# Cluster-wide RBAC for the cross-namespace discovery test.
|
||||
# Distinct ServiceAccount/ClusterRole names to avoid collision with
|
||||
# the namespace-scoped Role in rbac.yaml used by the default router.
|
||||
apiVersion: v1
|
||||
kind: ServiceAccount
|
||||
metadata:
|
||||
name: sgl-router-cluster
|
||||
namespace: sgl-router-test
|
||||
---
|
||||
apiVersion: rbac.authorization.k8s.io/v1
|
||||
kind: ClusterRole
|
||||
metadata:
|
||||
name: sgl-router-cluster
|
||||
rules:
|
||||
- apiGroups: ["discovery.k8s.io"]
|
||||
resources: ["endpointslices"]
|
||||
verbs: ["get", "list", "watch"]
|
||||
- apiGroups: [""]
|
||||
resources: ["services", "pods"]
|
||||
verbs: ["get", "list", "watch"]
|
||||
---
|
||||
apiVersion: rbac.authorization.k8s.io/v1
|
||||
kind: ClusterRoleBinding
|
||||
metadata:
|
||||
name: sgl-router-cluster
|
||||
roleRef:
|
||||
apiGroup: rbac.authorization.k8s.io
|
||||
kind: ClusterRole
|
||||
name: sgl-router-cluster
|
||||
subjects:
|
||||
- kind: ServiceAccount
|
||||
name: sgl-router-cluster
|
||||
namespace: sgl-router-test
|
||||
@@ -0,0 +1,34 @@
|
||||
apiVersion: v1
|
||||
kind: ServiceAccount
|
||||
metadata:
|
||||
name: sgl-router
|
||||
namespace: sgl-router-test
|
||||
---
|
||||
apiVersion: rbac.authorization.k8s.io/v1
|
||||
kind: Role
|
||||
metadata:
|
||||
name: sgl-router
|
||||
namespace: sgl-router-test
|
||||
rules:
|
||||
# EndpointSlice watch (k8s discovery backend)
|
||||
- apiGroups: ["discovery.k8s.io"]
|
||||
resources: ["endpointslices"]
|
||||
verbs: ["get", "list", "watch"]
|
||||
# Service list/watch (needed to resolve EndpointSlice owner)
|
||||
- apiGroups: [""]
|
||||
resources: ["services", "pods"]
|
||||
verbs: ["get", "list", "watch"]
|
||||
---
|
||||
apiVersion: rbac.authorization.k8s.io/v1
|
||||
kind: RoleBinding
|
||||
metadata:
|
||||
name: sgl-router
|
||||
namespace: sgl-router-test
|
||||
roleRef:
|
||||
apiGroup: rbac.authorization.k8s.io
|
||||
kind: Role
|
||||
name: sgl-router
|
||||
subjects:
|
||||
- kind: ServiceAccount
|
||||
name: sgl-router
|
||||
namespace: sgl-router-test
|
||||
+60
@@ -0,0 +1,60 @@
|
||||
# sgl-router deployment with ClusterRole for cross-namespace discovery test.
|
||||
# Watches workers in ALL namespaces via cluster-scoped EndpointSlice access.
|
||||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
metadata:
|
||||
name: sgl-router-cluster
|
||||
namespace: sgl-router-test
|
||||
spec:
|
||||
replicas: 1
|
||||
selector:
|
||||
matchLabels:
|
||||
app: sgl-router-cluster
|
||||
template:
|
||||
metadata:
|
||||
labels:
|
||||
app: sgl-router-cluster
|
||||
spec:
|
||||
serviceAccountName: sgl-router-cluster
|
||||
containers:
|
||||
- name: router
|
||||
image: sgl-router:e2e
|
||||
imagePullPolicy: Never
|
||||
args:
|
||||
- "--config"
|
||||
- "/etc/config/router-cluster.toml"
|
||||
ports:
|
||||
- containerPort: 8091
|
||||
name: http
|
||||
readinessProbe:
|
||||
httpGet:
|
||||
path: /readyz
|
||||
port: 8091
|
||||
initialDelaySeconds: 3
|
||||
periodSeconds: 3
|
||||
livenessProbe:
|
||||
httpGet:
|
||||
path: /healthz
|
||||
port: 8091
|
||||
initialDelaySeconds: 5
|
||||
periodSeconds: 10
|
||||
volumeMounts:
|
||||
- name: config
|
||||
mountPath: /etc/config
|
||||
volumes:
|
||||
- name: config
|
||||
configMap:
|
||||
name: sgl-router-cluster-config
|
||||
---
|
||||
apiVersion: v1
|
||||
kind: Service
|
||||
metadata:
|
||||
name: sgl-router-cluster
|
||||
namespace: sgl-router-test
|
||||
spec:
|
||||
selector:
|
||||
app: sgl-router-cluster
|
||||
ports:
|
||||
- name: http
|
||||
port: 8091
|
||||
targetPort: 8091
|
||||
@@ -0,0 +1,58 @@
|
||||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
metadata:
|
||||
name: sgl-router
|
||||
namespace: sgl-router-test
|
||||
spec:
|
||||
replicas: 1
|
||||
selector:
|
||||
matchLabels:
|
||||
app: sgl-router
|
||||
template:
|
||||
metadata:
|
||||
labels:
|
||||
app: sgl-router
|
||||
spec:
|
||||
serviceAccountName: sgl-router
|
||||
containers:
|
||||
- name: router
|
||||
image: sgl-router:e2e
|
||||
imagePullPolicy: Never
|
||||
args:
|
||||
- "--config"
|
||||
- "/etc/config/router.toml"
|
||||
ports:
|
||||
- containerPort: 8090
|
||||
name: http
|
||||
readinessProbe:
|
||||
httpGet:
|
||||
path: /readyz
|
||||
port: 8090
|
||||
initialDelaySeconds: 3
|
||||
periodSeconds: 3
|
||||
livenessProbe:
|
||||
httpGet:
|
||||
path: /healthz
|
||||
port: 8090
|
||||
initialDelaySeconds: 5
|
||||
periodSeconds: 10
|
||||
volumeMounts:
|
||||
- name: config
|
||||
mountPath: /etc/config
|
||||
volumes:
|
||||
- name: config
|
||||
configMap:
|
||||
name: sgl-router-config
|
||||
---
|
||||
apiVersion: v1
|
||||
kind: Service
|
||||
metadata:
|
||||
name: sgl-router
|
||||
namespace: sgl-router-test
|
||||
spec:
|
||||
selector:
|
||||
app: sgl-router
|
||||
ports:
|
||||
- name: http
|
||||
port: 8090
|
||||
targetPort: 8090
|
||||
@@ -0,0 +1,2 @@
|
||||
httpx==0.27.2
|
||||
pytest==8.3.3
|
||||
+191
@@ -0,0 +1,191 @@
|
||||
#!/usr/bin/env bash
|
||||
# Bootstrap a kind cluster for sgl-router K8s integration E2E tests.
|
||||
#
|
||||
# Prerequisites: Docker, kind, kubectl
|
||||
#
|
||||
# Usage:
|
||||
# ./tests/e2e/k8s_integration/setup.sh # full setup
|
||||
# ./tests/e2e/k8s_integration/setup.sh teardown # delete the cluster
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
|
||||
REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../../.." && pwd)" # repo root (above experimental/)
|
||||
SGL_ROUTER_DIR="${REPO_ROOT}/experimental/sgl-router"
|
||||
CLUSTER_NAME="${CLUSTER:-sgl-router-kind}"
|
||||
NAMESPACE="${NAMESPACE:-sgl-router-test}"
|
||||
CONTEXT="kind-${CLUSTER_NAME}"
|
||||
MANIFESTS_DIR="${SCRIPT_DIR}/manifests"
|
||||
|
||||
log() { echo "==> $*"; }
|
||||
|
||||
teardown() {
|
||||
log "Tearing down cluster '${CLUSTER_NAME}'..."
|
||||
if kind get clusters 2>/dev/null | grep -q "^${CLUSTER_NAME}$"; then
|
||||
kind delete cluster --name "${CLUSTER_NAME}"
|
||||
else
|
||||
log "Cluster '${CLUSTER_NAME}' not found, nothing to tear down."
|
||||
fi
|
||||
log "Done."
|
||||
}
|
||||
|
||||
if [[ "${1:-}" == "teardown" ]]; then
|
||||
teardown
|
||||
exit 0
|
||||
fi
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Step 1: Create kind cluster (idempotent)
|
||||
# ---------------------------------------------------------------------------
|
||||
if kind get clusters 2>/dev/null | grep -q "^${CLUSTER_NAME}$"; then
|
||||
log "Kind cluster '${CLUSTER_NAME}' already exists — reusing."
|
||||
else
|
||||
log "Creating kind cluster '${CLUSTER_NAME}'..."
|
||||
kind create cluster --name "${CLUSTER_NAME}" --wait 60s
|
||||
fi
|
||||
|
||||
kubectl config use-context "${CONTEXT}"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Step 2: Build Docker images (unless SKIP_DOCKER_BUILD=1)
|
||||
# ---------------------------------------------------------------------------
|
||||
if [[ "${SKIP_DOCKER_BUILD:-}" == "1" ]]; then
|
||||
log "SKIP_DOCKER_BUILD=1 — skipping docker build; expecting images to exist locally."
|
||||
for img in sgl-router:e2e sgl-router-fake-worker:e2e; do
|
||||
if ! docker image inspect "${img}" >/dev/null 2>&1; then
|
||||
log "ERROR: ${img} not found locally; cannot continue without building."
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
else
|
||||
log "Building sgl-router:e2e from ${REPO_ROOT} ..."
|
||||
docker build \
|
||||
-f "${SCRIPT_DIR}/Dockerfile.router" \
|
||||
-t sgl-router:e2e \
|
||||
"${REPO_ROOT}"
|
||||
|
||||
log "Building sgl-router-fake-worker:e2e ..."
|
||||
docker build \
|
||||
-f "${SCRIPT_DIR}/Dockerfile.fake_worker" \
|
||||
-t sgl-router-fake-worker:e2e \
|
||||
"${SCRIPT_DIR}"
|
||||
fi
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Step 3: Load images into kind
|
||||
# ---------------------------------------------------------------------------
|
||||
log "Loading images into kind cluster '${CLUSTER_NAME}'..."
|
||||
kind load docker-image sgl-router:e2e --name "${CLUSTER_NAME}"
|
||||
kind load docker-image sgl-router-fake-worker:e2e --name "${CLUSTER_NAME}"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Step 4: Apply namespace and RBAC
|
||||
# ---------------------------------------------------------------------------
|
||||
log "Applying namespace and RBAC..."
|
||||
kubectl --context "${CONTEXT}" apply -f "${MANIFESTS_DIR}/namespace.yaml"
|
||||
kubectl --context "${CONTEXT}" apply -f "${MANIFESTS_DIR}/rbac.yaml"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Step 5: Deploy 3 fake-worker replicas behind a Service
|
||||
# The Service causes K8s to auto-create an EndpointSlice, which
|
||||
# the sgl-router K8s discovery backend watches.
|
||||
# ---------------------------------------------------------------------------
|
||||
log "Deploying fake-worker Deployment + Service (3 replicas, app=sglang)..."
|
||||
kubectl --context "${CONTEXT}" -n "${NAMESPACE}" apply -f - <<EOF
|
||||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
metadata:
|
||||
name: fake-worker
|
||||
namespace: ${NAMESPACE}
|
||||
labels:
|
||||
app: sglang
|
||||
spec:
|
||||
replicas: 3
|
||||
selector:
|
||||
matchLabels:
|
||||
app: sglang
|
||||
template:
|
||||
metadata:
|
||||
labels:
|
||||
app: sglang
|
||||
spec:
|
||||
containers:
|
||||
- name: worker
|
||||
image: sgl-router-fake-worker:e2e
|
||||
imagePullPolicy: Never
|
||||
ports:
|
||||
- containerPort: 30000
|
||||
readinessProbe:
|
||||
httpGet:
|
||||
path: /health
|
||||
port: 30000
|
||||
initialDelaySeconds: 2
|
||||
periodSeconds: 3
|
||||
---
|
||||
apiVersion: v1
|
||||
kind: Service
|
||||
metadata:
|
||||
name: fake-worker
|
||||
namespace: ${NAMESPACE}
|
||||
labels:
|
||||
app: sglang
|
||||
spec:
|
||||
selector:
|
||||
app: sglang
|
||||
ports:
|
||||
- port: 30000
|
||||
targetPort: 30000
|
||||
EOF
|
||||
|
||||
log "Waiting for fake-worker rollout..."
|
||||
kubectl --context "${CONTEXT}" -n "${NAMESPACE}" rollout status deployment/fake-worker --timeout=120s
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Step 6: Create sgl-router ConfigMap with k8s discovery pointing at the
|
||||
# namespace where fake-worker pods live.
|
||||
# ---------------------------------------------------------------------------
|
||||
log "Creating sgl-router-config ConfigMap..."
|
||||
ROUTER_CONFIG="[server]
|
||||
host = \"0.0.0.0\"
|
||||
port = 8090
|
||||
|
||||
[[models]]
|
||||
id = \"tiny\"
|
||||
tokenizer_path = \"/etc/tokenizer/tiny.json\"
|
||||
policy = \"round_robin\"
|
||||
# Aggressive breaker so a terminating pod's connection-refused
|
||||
# immediately excludes it from the next request's candidate set —
|
||||
# the reconciliation tests scale workers rapidly and depend on
|
||||
# fast worker eviction to absorb the churn.
|
||||
circuit_breaker = { threshold = 1, cool_down_secs = 5 }
|
||||
|
||||
[discovery]
|
||||
backend = \"k8s\"
|
||||
|
||||
[discovery.k8s]
|
||||
namespace = \"${NAMESPACE}\"
|
||||
label_selector = \"app=sglang\""
|
||||
|
||||
kubectl --context "${CONTEXT}" -n "${NAMESPACE}" create configmap sgl-router-config \
|
||||
--from-literal=router.toml="${ROUTER_CONFIG}" \
|
||||
--dry-run=client -o yaml \
|
||||
| kubectl --context "${CONTEXT}" apply -f -
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Step 7: Deploy sgl-router
|
||||
# ---------------------------------------------------------------------------
|
||||
log "Deploying sgl-router..."
|
||||
kubectl --context "${CONTEXT}" apply -f "${MANIFESTS_DIR}/router.yaml"
|
||||
|
||||
log "Waiting for sgl-router rollout..."
|
||||
kubectl --context "${CONTEXT}" -n "${NAMESPACE}" rollout status deployment/sgl-router --timeout=300s
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Done
|
||||
# ---------------------------------------------------------------------------
|
||||
log ""
|
||||
log "Setup complete! Run the integration tests with:"
|
||||
log " pytest tests/e2e/k8s_integration/ -v -s"
|
||||
log ""
|
||||
log "To tear down:"
|
||||
log " ./tests/e2e/k8s_integration/setup.sh teardown"
|
||||
@@ -0,0 +1,265 @@
|
||||
"""Cross-namespace service discovery integration test.
|
||||
|
||||
Validates that a sgl-router instance with cluster-wide RBAC and no namespace
|
||||
filter in its k8s discovery config watches EndpointSlices in all namespaces.
|
||||
Workers deployed in a second namespace (sgl-router-test-extra) must be
|
||||
discovered alongside those in the primary namespace.
|
||||
|
||||
This test deploys a separate router Deployment (sgl-router-cluster) with a
|
||||
ClusterRole that grants EndpointSlice access across all namespaces.
|
||||
|
||||
Run with:
|
||||
pytest tests/e2e/k8s_integration/test_cross_namespace.py -v -s
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from conftest import (
|
||||
KUBECTL_CONTEXT,
|
||||
NAMESPACE,
|
||||
_apply_from_stdin,
|
||||
_cleanup_port_forward,
|
||||
_kubectl,
|
||||
_poll_until,
|
||||
_port_forward_start,
|
||||
_wait_for_deployment_ready,
|
||||
logger,
|
||||
)
|
||||
|
||||
MANIFESTS_DIR = Path(__file__).parent / "manifests"
|
||||
EXTRA_NAMESPACE = "sgl-router-test-extra"
|
||||
CLUSTER_ROUTER_PORT = 8093
|
||||
|
||||
|
||||
def _deploy_fake_worker_in_ns(name: str, namespace: str) -> None:
|
||||
"""Deploy a fake-worker pod with imagePullPolicy=Never in the given namespace."""
|
||||
pod_manifest = {
|
||||
"apiVersion": "v1",
|
||||
"kind": "Pod",
|
||||
"metadata": {
|
||||
"name": name,
|
||||
"namespace": namespace,
|
||||
"labels": {"app": "sglang", "cross-ns-test": "true"},
|
||||
},
|
||||
"spec": {
|
||||
"containers": [
|
||||
{
|
||||
"name": "worker",
|
||||
"image": "sgl-router-fake-worker:e2e",
|
||||
"imagePullPolicy": "Never",
|
||||
"ports": [{"containerPort": 30000}],
|
||||
"readinessProbe": {
|
||||
"httpGet": {"path": "/health", "port": 30000},
|
||||
"initialDelaySeconds": 2,
|
||||
"periodSeconds": 3,
|
||||
},
|
||||
}
|
||||
]
|
||||
},
|
||||
}
|
||||
proc = subprocess.run(
|
||||
["kubectl", "--context", KUBECTL_CONTEXT, "apply", "-f", "-"],
|
||||
input=json.dumps(pod_manifest),
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
if proc.returncode != 0:
|
||||
raise RuntimeError(
|
||||
f"Failed to deploy pod {name} in namespace {namespace} "
|
||||
f"(rc={proc.returncode}): {proc.stderr.strip()!r}"
|
||||
)
|
||||
logger.info("Deployed worker %s in namespace %s", name, namespace)
|
||||
|
||||
|
||||
def _safe_delete_pod(name: str, namespace: str) -> None:
|
||||
try:
|
||||
_kubectl(
|
||||
"delete",
|
||||
"pod",
|
||||
name,
|
||||
"-n",
|
||||
namespace,
|
||||
"--ignore-not-found",
|
||||
"--force",
|
||||
"--grace-period=0",
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning("Cleanup failed for pod %s in ns %s: %s", name, namespace, exc)
|
||||
|
||||
|
||||
def _ensure_namespace(name: str) -> None:
|
||||
manifest = {"apiVersion": "v1", "kind": "Namespace", "metadata": {"name": name}}
|
||||
_apply_from_stdin(json.dumps(manifest))
|
||||
|
||||
|
||||
def _ensure_service_in_ns(namespace: str, selector: str = "app=sglang") -> None:
|
||||
"""Create a Service so K8s auto-creates an EndpointSlice for cross-ns workers.
|
||||
|
||||
Service `metadata.labels` propagates to the auto-created EndpointSlice's
|
||||
labels — and the cluster-scoped router filters slices server-side by
|
||||
`app=sglang,cross-ns-test=true`. Without those labels on the Service,
|
||||
its EndpointSlice gets filtered out and the cross-ns worker is invisible.
|
||||
"""
|
||||
svc_manifest = {
|
||||
"apiVersion": "v1",
|
||||
"kind": "Service",
|
||||
"metadata": {
|
||||
"name": "fake-worker",
|
||||
"namespace": namespace,
|
||||
"labels": {"app": "sglang", "cross-ns-test": "true"},
|
||||
},
|
||||
"spec": {
|
||||
"selector": {"app": "sglang", "cross-ns-test": "true"},
|
||||
"ports": [{"port": 30000, "targetPort": 30000}],
|
||||
},
|
||||
}
|
||||
_apply_from_stdin(json.dumps(svc_manifest))
|
||||
|
||||
|
||||
def _can_route(router_url: str) -> bool:
|
||||
try:
|
||||
r = httpx.post(
|
||||
f"{router_url}/v1/chat/completions",
|
||||
json={
|
||||
"model": "tiny",
|
||||
"messages": [{"role": "user", "content": "cross-ns"}],
|
||||
},
|
||||
timeout=8.0,
|
||||
)
|
||||
return r.status_code == 200
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def cluster_scoped_router(k8s_cluster):
|
||||
"""Deploy the cluster-scoped RBAC + router, plus a second namespace."""
|
||||
rbac_manifest = MANIFESTS_DIR / "rbac-cluster-scoped.yaml"
|
||||
router_manifest = MANIFESTS_DIR / "router-cluster-scoped.yaml"
|
||||
|
||||
_kubectl("apply", "-f", str(rbac_manifest))
|
||||
_ensure_namespace(EXTRA_NAMESPACE)
|
||||
_ensure_service_in_ns(EXTRA_NAMESPACE)
|
||||
|
||||
# ConfigMap for the cluster-scoped router: empty namespace = watch all
|
||||
cluster_config = """[server]
|
||||
host = "0.0.0.0"
|
||||
port = 8091
|
||||
|
||||
[[models]]
|
||||
id = "tiny"
|
||||
tokenizer_path = "/etc/tokenizer/tiny.json"
|
||||
policy = "round_robin"
|
||||
|
||||
[discovery]
|
||||
backend = "k8s"
|
||||
|
||||
[discovery.k8s]
|
||||
namespace = ""
|
||||
label_selector = "app=sglang,cross-ns-test=true"
|
||||
"""
|
||||
_kubectl(
|
||||
"create",
|
||||
"configmap",
|
||||
"sgl-router-cluster-config",
|
||||
f"--from-literal=router-cluster.toml={cluster_config}",
|
||||
"-n",
|
||||
NAMESPACE,
|
||||
"--dry-run=client",
|
||||
"-o",
|
||||
"yaml",
|
||||
check=True,
|
||||
)
|
||||
# pipe through apply
|
||||
proc = _kubectl(
|
||||
"create",
|
||||
"configmap",
|
||||
"sgl-router-cluster-config",
|
||||
f"--from-literal=router-cluster.toml={cluster_config}",
|
||||
"-n",
|
||||
NAMESPACE,
|
||||
"--dry-run=client",
|
||||
"-o",
|
||||
"yaml",
|
||||
)
|
||||
_apply_from_stdin(proc.stdout)
|
||||
|
||||
_kubectl("apply", "-f", str(router_manifest))
|
||||
|
||||
# The cluster-scoped router's /readyz blocks on registry-not-empty, so
|
||||
# without at least one matching worker the rollout-status check below
|
||||
# would hang for 180s. Deploy a "bootstrap" worker in EXTRA_NAMESPACE
|
||||
# with the label_selector match (app=sglang,cross-ns-test=true) so the
|
||||
# router's k8s discovery picks it up before the readiness probe runs.
|
||||
# The test body adds a SECOND worker later to verify dynamic discovery.
|
||||
bootstrap_worker = "cross-ns-worker-bootstrap"
|
||||
_deploy_fake_worker_in_ns(bootstrap_worker, EXTRA_NAMESPACE)
|
||||
|
||||
pf = None
|
||||
try:
|
||||
_wait_for_deployment_ready("sgl-router-cluster")
|
||||
pf = _port_forward_start(
|
||||
NAMESPACE, "sgl-router-cluster", CLUSTER_ROUTER_PORT, 8091
|
||||
)
|
||||
yield f"http://127.0.0.1:{CLUSTER_ROUTER_PORT}"
|
||||
finally:
|
||||
if pf is not None:
|
||||
_cleanup_port_forward("cluster_router", pf)
|
||||
_safe_delete_pod(bootstrap_worker, EXTRA_NAMESPACE)
|
||||
_kubectl(
|
||||
"delete", "-f", str(router_manifest), "--ignore-not-found", check=False
|
||||
)
|
||||
_kubectl("delete", "-f", str(rbac_manifest), "--ignore-not-found", check=False)
|
||||
_kubectl(
|
||||
"delete",
|
||||
"namespace",
|
||||
EXTRA_NAMESPACE,
|
||||
"--ignore-not-found",
|
||||
"--wait=true",
|
||||
"--timeout=60s",
|
||||
check=False,
|
||||
)
|
||||
|
||||
|
||||
class TestClusterWideDiscovery:
|
||||
"""Router with ClusterRole and no namespace filter sees workers in every namespace."""
|
||||
|
||||
def test_router_routes_to_worker_in_extra_namespace(self, cluster_scoped_router):
|
||||
"""Deploy one fake-worker pod in the extra namespace behind a Service;
|
||||
the cluster-scoped router must discover it (via its EndpointSlice) and
|
||||
successfully route a chat completion to it."""
|
||||
router_url = cluster_scoped_router
|
||||
worker_name = "cross-ns-worker-extra"
|
||||
|
||||
try:
|
||||
_deploy_fake_worker_in_ns(worker_name, EXTRA_NAMESPACE)
|
||||
|
||||
_poll_until(
|
||||
lambda: _can_route(router_url),
|
||||
"cluster-scoped router routes to worker in extra namespace",
|
||||
timeout=60,
|
||||
interval=3,
|
||||
)
|
||||
|
||||
r = httpx.post(
|
||||
f"{router_url}/v1/chat/completions",
|
||||
json={
|
||||
"model": "tiny",
|
||||
"messages": [
|
||||
{"role": "user", "content": "cross-namespace routing"}
|
||||
],
|
||||
},
|
||||
timeout=15.0,
|
||||
)
|
||||
assert r.status_code == 200, f"expected 200, got {r.status_code}: {r.text}"
|
||||
assert "echo:" in r.json()["choices"][0]["message"]["content"]
|
||||
finally:
|
||||
_safe_delete_pod(worker_name, EXTRA_NAMESPACE)
|
||||
@@ -0,0 +1,84 @@
|
||||
"""E2E: sgl-router K8s discovery — basic routing.
|
||||
|
||||
Verifies that sgl-router, configured with the k8s EndpointSlice backend,
|
||||
discovers the 3 fake-worker replicas deployed by setup.sh and successfully
|
||||
routes chat-completion requests to them.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from conftest import NAMESPACE, _kubectl, _poll_until, logger
|
||||
|
||||
|
||||
def _scale_fake_worker(replicas: int) -> None:
|
||||
_kubectl(
|
||||
"scale",
|
||||
"deployment/fake-worker",
|
||||
f"--replicas={replicas}",
|
||||
"-n",
|
||||
NAMESPACE,
|
||||
)
|
||||
|
||||
|
||||
def test_router_routes_chat_to_a_worker(router_url):
|
||||
"""A /v1/chat/completions request through the router returns 200 with the
|
||||
fake-worker echo payload, proving end-to-end routing works."""
|
||||
r = httpx.post(
|
||||
f"{router_url}/v1/chat/completions",
|
||||
json={
|
||||
"model": "tiny",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"stream": False,
|
||||
},
|
||||
timeout=15.0,
|
||||
)
|
||||
assert r.status_code == 200, f"expected 200, got {r.status_code}: {r.text}"
|
||||
body = r.json()
|
||||
assert "echo:" in body["choices"][0]["message"]["content"]
|
||||
|
||||
|
||||
def test_router_lists_model(router_url):
|
||||
"""GET /v1/models returns the 'tiny' model entry from the router config."""
|
||||
r = httpx.get(f"{router_url}/v1/models", timeout=10.0)
|
||||
assert r.status_code == 200, f"expected 200, got {r.status_code}: {r.text}"
|
||||
body = r.json()
|
||||
ids = [m["id"] for m in body["data"]]
|
||||
assert "tiny" in ids, f"expected 'tiny' in model list, got {ids}"
|
||||
|
||||
|
||||
def test_router_discovers_multiple_workers(router_url):
|
||||
"""Scale down from 3 to 1 and back to 3 replicas; router must continue
|
||||
routing successfully after each transition (EndpointSlice watch reflects
|
||||
the change)."""
|
||||
# First confirm baseline routing
|
||||
r = httpx.post(
|
||||
f"{router_url}/v1/chat/completions",
|
||||
json={
|
||||
"model": "tiny",
|
||||
"messages": [{"role": "user", "content": "scale-test"}],
|
||||
},
|
||||
timeout=15.0,
|
||||
)
|
||||
assert r.status_code == 200
|
||||
|
||||
# Scale down to 1 — router should still route after reconverging
|
||||
_scale_fake_worker(1)
|
||||
_poll_until(
|
||||
lambda: httpx.post(
|
||||
f"{router_url}/v1/chat/completions",
|
||||
json={
|
||||
"model": "tiny",
|
||||
"messages": [{"role": "user", "content": "post-scale-down"}],
|
||||
},
|
||||
timeout=10.0,
|
||||
).status_code
|
||||
== 200,
|
||||
"router routes after scale-down to 1",
|
||||
timeout=60,
|
||||
interval=3,
|
||||
)
|
||||
|
||||
# Restore to 3
|
||||
_scale_fake_worker(3)
|
||||
@@ -0,0 +1,150 @@
|
||||
"""Worker lifecycle integration tests.
|
||||
|
||||
Covers:
|
||||
1. Scaling replicas up — new EndpointSlice entries are discovered.
|
||||
2. Scaling replicas down — removed endpoints are deregistered.
|
||||
3. Router restart — after the router pod is killed, the Deployment restarts
|
||||
it and it re-lists the existing EndpointSlice entries without duplicates.
|
||||
|
||||
These tests DO NOT use a /workers admin API (sgl-router does not expose
|
||||
one). They verify behaviour through /v1/chat/completions responses and
|
||||
by driving the deployment scale.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from conftest import (
|
||||
NAMESPACE,
|
||||
_cleanup_port_forward,
|
||||
_kubectl,
|
||||
_poll_until,
|
||||
_port_forward_start,
|
||||
_wait_for_deployment_ready,
|
||||
logger,
|
||||
)
|
||||
|
||||
ROUTER_RESTART_PORT = 8092
|
||||
|
||||
|
||||
def _scale(deployment: str, replicas: int) -> None:
|
||||
_kubectl(
|
||||
"scale", f"deployment/{deployment}", f"--replicas={replicas}", "-n", NAMESPACE
|
||||
)
|
||||
|
||||
|
||||
def _can_route(router_url: str) -> bool:
|
||||
try:
|
||||
r = httpx.post(
|
||||
f"{router_url}/v1/chat/completions",
|
||||
json={"model": "tiny", "messages": [{"role": "user", "content": "ping"}]},
|
||||
timeout=8.0,
|
||||
)
|
||||
return r.status_code == 200
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
class TestScaleUp:
|
||||
"""Scaling fake-worker replicas up must not break routing."""
|
||||
|
||||
def test_router_routes_after_scale_up(self, router_url):
|
||||
"""Restore 3 replicas (in case a prior test left 1), verify routing."""
|
||||
_scale("fake-worker", 3)
|
||||
_poll_until(
|
||||
lambda: _can_route(router_url),
|
||||
"router routes after scale-up to 3",
|
||||
timeout=60,
|
||||
interval=3,
|
||||
)
|
||||
|
||||
|
||||
class TestScaleDown:
|
||||
"""Scaling to 0 then back up must restore routing."""
|
||||
|
||||
def test_router_recovers_after_scale_to_zero_and_back(self, router_url):
|
||||
try:
|
||||
_scale("fake-worker", 0)
|
||||
# After scale-to-0 the router may return 503 (no healthy workers)
|
||||
# That is expected behaviour — assert it transitions back on scale-up.
|
||||
_scale("fake-worker", 2)
|
||||
_poll_until(
|
||||
lambda: _can_route(router_url),
|
||||
"router routes again after scale-up from 0",
|
||||
timeout=90,
|
||||
interval=3,
|
||||
)
|
||||
finally:
|
||||
_scale("fake-worker", 3)
|
||||
|
||||
|
||||
class TestRouterRestart:
|
||||
"""Killing the router pod forces a Deployment restart; the new pod must
|
||||
re-discover workers via the EndpointSlice watch without duplicates."""
|
||||
|
||||
def test_router_rediscovers_workers_after_restart(self, k8s_cluster):
|
||||
# Use a dedicated port to avoid clashing with the session fixture
|
||||
pf_holder: list = [None]
|
||||
try:
|
||||
_wait_for_deployment_ready("sgl-router")
|
||||
pf_holder[0] = _port_forward_start(
|
||||
NAMESPACE, "sgl-router", ROUTER_RESTART_PORT, 8090
|
||||
)
|
||||
restart_url = f"http://127.0.0.1:{ROUTER_RESTART_PORT}"
|
||||
|
||||
# Baseline: routing works pre-restart
|
||||
_poll_until(
|
||||
lambda: _can_route(restart_url),
|
||||
"baseline routing works pre-restart",
|
||||
timeout=30,
|
||||
interval=2,
|
||||
)
|
||||
|
||||
# Kill the router pod — the Deployment ReplicaSet will restart it
|
||||
res = _kubectl(
|
||||
"get",
|
||||
"pod",
|
||||
"-n",
|
||||
NAMESPACE,
|
||||
"-l",
|
||||
"app=sgl-router",
|
||||
"-o",
|
||||
"jsonpath={.items[0].metadata.name}",
|
||||
check=False,
|
||||
)
|
||||
old_pod = res.stdout.strip()
|
||||
if old_pod:
|
||||
_kubectl(
|
||||
"delete",
|
||||
"pod",
|
||||
old_pod,
|
||||
"-n",
|
||||
NAMESPACE,
|
||||
"--force",
|
||||
"--grace-period=0",
|
||||
)
|
||||
|
||||
# Tear down the old port-forward before waiting for the new pod
|
||||
if pf_holder[0] is not None:
|
||||
_cleanup_port_forward("router-restart-pre-kill", pf_holder[0])
|
||||
pf_holder[0] = None
|
||||
|
||||
_wait_for_deployment_ready("sgl-router")
|
||||
|
||||
pf_holder[0] = _port_forward_start(
|
||||
NAMESPACE, "sgl-router", ROUTER_RESTART_PORT, 8090
|
||||
)
|
||||
|
||||
# After restart, routing must come back (EndpointSlice re-watch)
|
||||
_poll_until(
|
||||
lambda: _can_route(restart_url),
|
||||
"routing restored after router restart",
|
||||
timeout=60,
|
||||
interval=3,
|
||||
)
|
||||
finally:
|
||||
if pf_holder[0] is not None:
|
||||
_cleanup_port_forward("router-restart", pf_holder[0])
|
||||
@@ -0,0 +1,163 @@
|
||||
"""K8s discovery reconciliation integration tests.
|
||||
|
||||
Tests verify that:
|
||||
1. The K8s EndpointSlice watcher correctly discovers new workers as Services
|
||||
and backing Deployments are updated.
|
||||
2. Workers are removed from the router's registry after the backing EndpointSlice
|
||||
entries disappear (pod deleted / deployment scaled to 0).
|
||||
3. After a simulated watch-connection interruption (router restarted), the
|
||||
registry converges back to the correct worker set.
|
||||
|
||||
Note: sgl-router does not currently expose a Prometheus /metrics endpoint,
|
||||
so the SMG-style metric assertions are not used here. Disconnect/reconnect
|
||||
coverage is provided by test_lifecycle.TestRouterRestart.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from conftest import (
|
||||
NAMESPACE,
|
||||
RECONCILIATION_WAIT_SECS,
|
||||
_kubectl,
|
||||
_poll_until,
|
||||
logger,
|
||||
)
|
||||
|
||||
|
||||
def _scale_fake_worker(replicas: int) -> None:
|
||||
_kubectl(
|
||||
"scale", "deployment/fake-worker", f"--replicas={replicas}", "-n", NAMESPACE
|
||||
)
|
||||
|
||||
|
||||
def _can_route(router_url: str) -> bool:
|
||||
try:
|
||||
r = httpx.post(
|
||||
f"{router_url}/v1/chat/completions",
|
||||
json={
|
||||
"model": "tiny",
|
||||
"messages": [{"role": "user", "content": "reconcile"}],
|
||||
},
|
||||
timeout=8.0,
|
||||
)
|
||||
return r.status_code == 200
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
class TestWatcherDiscovery:
|
||||
"""The EndpointSlice watcher discovers new endpoints on Deployment scale-up."""
|
||||
|
||||
def test_watcher_discovers_new_endpoints_on_scale_up(self, router_url):
|
||||
"""Scale from 1 to 3 replicas; router must continue routing successfully."""
|
||||
_scale_fake_worker(1)
|
||||
# Wait for scale-down to propagate and routing to stabilise
|
||||
_poll_until(
|
||||
lambda: _can_route(router_url),
|
||||
"router routes with 1 replica",
|
||||
timeout=60,
|
||||
interval=3,
|
||||
)
|
||||
|
||||
_scale_fake_worker(3)
|
||||
_poll_until(
|
||||
lambda: _can_route(router_url),
|
||||
"router routes with 3 replicas (after scale-up)",
|
||||
timeout=60,
|
||||
interval=3,
|
||||
)
|
||||
|
||||
|
||||
class TestStaleEndpointRemoval:
|
||||
"""When fake-worker replicas drop, the router must stop routing to the
|
||||
removed endpoints.
|
||||
|
||||
Because sgl-router has no /workers admin API, we verify removal
|
||||
indirectly: scale to 0, assert the router returns non-200 (or at least
|
||||
that scaling back to 2 restores routing), then restore.
|
||||
"""
|
||||
|
||||
def test_routing_restores_after_scale_down_and_back_up(self, router_url):
|
||||
"""Scale to 0 (no workers → expect non-200), then restore to 2.
|
||||
After restore the router must route again within the reconciliation window.
|
||||
"""
|
||||
try:
|
||||
_scale_fake_worker(0)
|
||||
|
||||
# Expect routing to fail eventually (503 or connection error)
|
||||
deadline = time.time() + RECONCILIATION_WAIT_SECS
|
||||
routing_failed = False
|
||||
while time.time() < deadline:
|
||||
try:
|
||||
r = httpx.post(
|
||||
f"{router_url}/v1/chat/completions",
|
||||
json={
|
||||
"model": "tiny",
|
||||
"messages": [{"role": "user", "content": "no-workers"}],
|
||||
},
|
||||
timeout=5.0,
|
||||
)
|
||||
if r.status_code != 200:
|
||||
routing_failed = True
|
||||
break
|
||||
except Exception:
|
||||
routing_failed = True
|
||||
break
|
||||
time.sleep(3)
|
||||
|
||||
# If after RECONCILIATION_WAIT_SECS the router is still routing,
|
||||
# that means old endpoints are cached — not necessarily wrong for
|
||||
# a watcher that hasn't ticked yet, but log a warning.
|
||||
if not routing_failed:
|
||||
logger.warning(
|
||||
"Router still returning 200 after scale-to-0; "
|
||||
"EndpointSlice event may be delayed — continuing test."
|
||||
)
|
||||
|
||||
# Restore workers and verify routing comes back
|
||||
_scale_fake_worker(2)
|
||||
_poll_until(
|
||||
lambda: _can_route(router_url),
|
||||
"routing restored after scale back up to 2",
|
||||
timeout=RECONCILIATION_WAIT_SECS,
|
||||
interval=3,
|
||||
)
|
||||
finally:
|
||||
_scale_fake_worker(3)
|
||||
|
||||
|
||||
class TestReconciliationConsistency:
|
||||
"""Routing remains stable over multiple reconciliation windows with steady
|
||||
worker state — no spurious deregistrations or duplicate registrations."""
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_routing_stable_over_multiple_reconciliation_cycles(self, router_url):
|
||||
"""Deploy 3 workers, sample routing success over ~150s (2 reconciliation
|
||||
cycles + margin), assert no interruptions."""
|
||||
_scale_fake_worker(3)
|
||||
_poll_until(
|
||||
lambda: _can_route(router_url),
|
||||
"baseline routing with 3 workers",
|
||||
timeout=30,
|
||||
interval=2,
|
||||
)
|
||||
|
||||
# Sample every 15s for 150s
|
||||
wait_secs = RECONCILIATION_WAIT_SECS + 60
|
||||
end_time = time.time() + wait_secs
|
||||
failures = []
|
||||
while time.time() < end_time:
|
||||
ok = _can_route(router_url)
|
||||
if not ok:
|
||||
failures.append(time.time())
|
||||
time.sleep(15)
|
||||
|
||||
assert not failures, (
|
||||
f"Routing failed at {len(failures)} sample(s) during stability window; "
|
||||
f"timestamps: {failures}"
|
||||
)
|
||||
@@ -0,0 +1,26 @@
|
||||
# Pytest configuration for sgl-router tests/e2e/.
|
||||
# Lives next to conftest.py so `pytest experimental/sgl-router/tests/e2e/`
|
||||
# picks it up automatically.
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
minversion = "8.0"
|
||||
# Default discovery: smoke tests (top-level test_*.py) and the
|
||||
# multi-worker chat_completions suite. k8s_integration is intentionally
|
||||
# not in the default set — it requires a kind/k8s cluster and is
|
||||
# invoked explicitly.
|
||||
testpaths = [
|
||||
".",
|
||||
"chat_completions",
|
||||
]
|
||||
python_files = ["test_*.py"]
|
||||
python_classes = ["Test*"]
|
||||
python_functions = ["test_*"]
|
||||
markers = [
|
||||
"real_gpu: requires at least one NVIDIA GPU (skipped on CPU-only hosts)",
|
||||
"pd_mode: requires the router started in PD-disaggregation mode",
|
||||
"slow: takes >30s (model load, multi-request convergence checks)",
|
||||
]
|
||||
log_cli = true
|
||||
log_cli_level = "INFO"
|
||||
log_cli_format = "%(asctime)s [%(levelname)s] %(name)s: %(message)s"
|
||||
log_cli_date_format = "%H:%M:%S"
|
||||
@@ -0,0 +1,15 @@
|
||||
httpx==0.27.2
|
||||
pytest==8.3.3
|
||||
pytest-asyncio==0.24.0
|
||||
# huggingface_hub is intentionally NOT pinned here. SGLang's
|
||||
# `scripts/ci/cuda/ci_install_dependency.sh` already installs a
|
||||
# version compatible with the rest of its transitive deps
|
||||
# (transformers / diffusers / kernels, which require
|
||||
# huggingface_hub >= 1.5 / >= 0.34 / >= 1.3 respectively). An earlier
|
||||
# pin of `huggingface_hub==0.26.2` here got installed AFTER the SGLang
|
||||
# deps and downgraded huggingface_hub past `is_offline_mode`'s top-
|
||||
# level export, which broke `from sglang.srt.server_args import …`
|
||||
# at module import time and turned every smoke test into a 5-minute
|
||||
# `/health` timeout with no actionable signal.
|
||||
# The e2e suite only uses huggingface_hub's `try_to_load_from_cache`,
|
||||
# which is available in every release SGLang would install.
|
||||
@@ -0,0 +1,64 @@
|
||||
"""
|
||||
Smoke tests for /v1/models and /v1/chat/completions (streaming + non-streaming).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
MODEL = "Qwen/Qwen3-0.6B"
|
||||
|
||||
|
||||
def test_models(router: str) -> None:
|
||||
"""GET /v1/models must list the configured model."""
|
||||
resp = httpx.get(f"{router}/v1/models", timeout=30)
|
||||
assert resp.status_code == 200, resp.text
|
||||
data = resp.json()
|
||||
ids = [m["id"] for m in data.get("data", [])]
|
||||
assert any(
|
||||
MODEL in mid for mid in ids
|
||||
), f"Model {MODEL!r} not found in /v1/models response: {ids}"
|
||||
|
||||
|
||||
def test_chat_non_streaming(router: str) -> None:
|
||||
"""POST /v1/chat/completions (stream=False) returns an assistant message."""
|
||||
payload = {
|
||||
"model": MODEL,
|
||||
"messages": [{"role": "user", "content": "Say hi."}],
|
||||
"max_tokens": 10,
|
||||
"stream": False,
|
||||
}
|
||||
resp = httpx.post(f"{router}/v1/chat/completions", json=payload, timeout=60)
|
||||
assert resp.status_code == 200, resp.text
|
||||
body = resp.json()
|
||||
choice = body["choices"][0]
|
||||
assert choice["message"]["role"] == "assistant"
|
||||
assert choice["message"]["content"], "Expected non-empty assistant content"
|
||||
|
||||
|
||||
def test_chat_streaming(router: str) -> None:
|
||||
"""POST /v1/chat/completions (stream=True) returns >=2 SSE chunks incl. [DONE]."""
|
||||
payload = {
|
||||
"model": MODEL,
|
||||
"messages": [{"role": "user", "content": "Say hi."}],
|
||||
"max_tokens": 10,
|
||||
"stream": True,
|
||||
}
|
||||
chunks: list[str] = []
|
||||
with httpx.stream(
|
||||
"POST",
|
||||
f"{router}/v1/chat/completions",
|
||||
json=payload,
|
||||
timeout=60,
|
||||
) as resp:
|
||||
assert resp.status_code == 200, resp.read().decode()
|
||||
for line in resp.iter_lines():
|
||||
line = line.strip()
|
||||
if line.startswith("data:"):
|
||||
chunks.append(line)
|
||||
|
||||
assert len(chunks) >= 2, f"Expected >=2 SSE chunks, got {len(chunks)}: {chunks}"
|
||||
assert any(
|
||||
"[DONE]" in c for c in chunks
|
||||
), f"No [DONE] chunk found in SSE stream: {chunks}"
|
||||
@@ -0,0 +1,37 @@
|
||||
"""
|
||||
Smoke test for /v1/tokenize and /v1/detokenize round-trip.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
|
||||
MODEL = "Qwen/Qwen3-0.6B"
|
||||
TEXT = "Hello, world!"
|
||||
|
||||
|
||||
def test_tokenize_round_trip(router: str) -> None:
|
||||
"""POST /v1/tokenize then /v1/detokenize must recover the original text."""
|
||||
# Tokenize
|
||||
tok_resp = httpx.post(
|
||||
f"{router}/v1/tokenize",
|
||||
json={"model": MODEL, "prompt": TEXT},
|
||||
timeout=30,
|
||||
)
|
||||
assert tok_resp.status_code == 200, tok_resp.text
|
||||
tokens = tok_resp.json()["tokens"]
|
||||
assert (
|
||||
isinstance(tokens, list) and len(tokens) > 0
|
||||
), f"Expected non-empty token list, got: {tokens}"
|
||||
|
||||
# Detokenize
|
||||
detok_resp = httpx.post(
|
||||
f"{router}/v1/detokenize",
|
||||
json={"model": MODEL, "tokens": tokens},
|
||||
timeout=30,
|
||||
)
|
||||
assert detok_resp.status_code == 200, detok_resp.text
|
||||
recovered = detok_resp.json()["text"]
|
||||
assert (
|
||||
TEXT in recovered or recovered in TEXT
|
||||
), f"Round-trip mismatch: original={TEXT!r}, recovered={recovered!r}"
|
||||
Reference in New Issue
Block a user