From cc838b69531c2173a93ae1a193b797db2012d129 Mon Sep 17 00:00:00 2001 From: Alex McC <319643551+hello-alexmcc@users.noreply.github.com> Date: Mon, 7 Sep 2026 16:49:15 -0700 Subject: [PATCH 1/4] test(e2e): cover overload sheds, routing-key pinning, and worker restart Three merged router fixes changed behaviour a client can observe without leaving an end-to-end check behind: - #2417: an overload shed must answer 503 with its own worker_overload_protection_shed code and a Retry-After, not the generic no_available_workers. The engine is pinned to one running request so a burst piles up in its queue; a probe sent while the queue is deep must be shed, and the worker must serve again once the burst drains. - #2355: with --routing-key-override a body rid must outrank the routing key header. Two workers under the manual policy: a header key stays sticky, and a request whose header names one lineage and whose body rid names another lands on the rid's worker. The serving worker is observed through the gateway's in-flight load during a long generation. - #2427: a gRPC worker stopped under a live gateway must be reported unhealthy, requests must fail fast rather than hang, and after the worker restarts on the same port it must return to rotation. The test borrows the session pool's worker so it does not race a cached worker for the GPU. Signed-off-by: Alex McC <319643551+hello-alexmcc@users.noreply.github.com> --- e2e_test/router/test_overload.py | 174 +++++++++++++++++++++++++ e2e_test/router/test_routing_key.py | 119 +++++++++++++++++ e2e_test/router/test_worker_restart.py | 132 +++++++++++++++++++ 3 files changed, 425 insertions(+) create mode 100644 e2e_test/router/test_overload.py create mode 100644 e2e_test/router/test_routing_key.py create mode 100644 e2e_test/router/test_worker_restart.py diff --git a/e2e_test/router/test_overload.py b/e2e_test/router/test_overload.py new file mode 100644 index 0000000000..0db81bf4c8 --- /dev/null +++ b/e2e_test/router/test_overload.py @@ -0,0 +1,174 @@ +"""Overload protection sheds with its own error code and recovers. + +``--worker-overload-protection`` takes a worker whose waiting queue crosses +``--worker-overload-waiting-requests`` out of routing; when every worker for +the model is overloaded the gateway sheds immediately with 503 instead of +queueing. The shed used to reuse the generic ``no_available_workers`` code, so +a downstream gateway could not tell an overload from a dead fleet (#2417). + +The engine is pinned to one running request so a burst piles up in its +waiting queue, the gateway polls loads every second, and a probe sent while +the queue is deep must be shed with ``worker_overload_protection_shed`` and a +``Retry-After``. Once the burst drains the worker must serve again. + +Usage: + E2E_RUNTIME=sglang pytest e2e_test/router/test_overload.py -v +""" + +from __future__ import annotations + +import logging +import threading +import time + +import httpx +import pytest + +logger = logging.getLogger(__name__) + +SHED_CODE = "worker_overload_protection_shed" +WAITING_THRESHOLD = 2 +BURST_SIZE = 12 +_MODEL = "meta-llama/Llama-3.2-1B-Instruct" +_GATEWAY_ARGS = [ + "--worker-overload-protection", + "--worker-overload-waiting-requests", + str(WAITING_THRESHOLD), + "--load-monitor-interval", + "1", +] + + +def _chat_body(model: str, max_tokens: int, *, ignore_eos: bool) -> dict: + return { + "model": model, + "messages": [{"role": "user", "content": "Write a long story about the sea."}], + "max_tokens": max_tokens, + "temperature": 0.0, + "ignore_eos": ignore_eos, + } + + +def _post(gateway, body: dict, timeout: float) -> httpx.Response: + return httpx.post(f"{gateway.base_url}/v1/chat/completions", json=body, timeout=timeout) + + +def _error_code(resp: httpx.Response) -> str | None: + try: + body = resp.json() + except ValueError: + return None + error = body.get("error", body) + return error.get("code") if isinstance(error, dict) else None + + +def _wait_until_served(gateway, model: str, timeout: float) -> None: + """Retry a tiny request until the worker answers 200 again.""" + deadline = time.monotonic() + timeout + last = "no attempt" + while time.monotonic() < deadline: + try: + resp = _post(gateway, _chat_body(model, 4, ignore_eos=False), timeout=60.0) + except httpx.HTTPError as exc: + last = repr(exc) + else: + if resp.status_code == 200: + return + last = f"{resp.status_code} {resp.text[:200]}" + time.sleep(1.0) + pytest.fail(f"worker did not serve within {timeout:.0f}s; last: {last}") + + +def _wait_for_deep_queue(gateway, min_waiting: int, timeout: float) -> dict: + """Poll /loads until a report shows at least ``min_waiting`` queued requests.""" + deadline = time.monotonic() + timeout + last: list[dict] = [] + while time.monotonic() < deadline: + resp = httpx.get(f"{gateway.base_url}/loads", timeout=5.0) + assert resp.status_code == 200, resp.text + last = resp.json().get("loads", []) + for entry in last: + if entry.get("num_waiting_reqs", 0) >= min_waiting: + return entry + time.sleep(0.2) + pytest.fail(f"engine queue never reached {min_waiting} waiting requests; last reports: {last}") + + +class _OverloadShedBase: + """Shared body; subclasses pin the engine and its single-slot flag.""" + + def test_shed_carries_its_own_code_and_recovers(self, setup_backend): + _, model, _, gateway = setup_backend + _wait_until_served(gateway, model, timeout=60.0) + + results: list[httpx.Response | BaseException] = [] + + def _long_request() -> None: + try: + results.append(_post(gateway, _chat_body(model, 256, ignore_eos=True), 240.0)) + except BaseException as exc: # reported below + results.append(exc) + + threads = [threading.Thread(target=_long_request, daemon=True) for _ in range(BURST_SIZE)] + for t in threads: + t.start() + + try: + queue = _wait_for_deep_queue(gateway, WAITING_THRESHOLD, timeout=20.0) + logger.info("queue is deep: %s", queue) + + probe = _post(gateway, _chat_body(model, 4, ignore_eos=False), timeout=240.0) + logger.info( + "probe: status=%s retry-after=%s code=%s", + probe.status_code, + probe.headers.get("retry-after"), + _error_code(probe), + ) + assert probe.status_code == 503, ( + f"probe was not shed while the queue was deep: {probe.status_code} {probe.text[:200]}" + ) + assert _error_code(probe) == SHED_CODE, ( + f"shed carried the wrong code: {probe.text[:300]}" + ) + retry_after = probe.headers.get("retry-after") + assert retry_after is not None and int(retry_after) >= 1, ( + f"shed without a usable Retry-After: {retry_after!r}" + ) + finally: + for t in threads: + t.join(timeout=300) + + transport_errors = [r for r in results if isinstance(r, BaseException)] + assert not transport_errors, ( + f"burst requests failed at the transport: {transport_errors[:3]}" + ) + responses = [r for r in results if isinstance(r, httpx.Response)] + for resp in responses: + assert resp.status_code == 200 or ( + resp.status_code == 503 and _error_code(resp) == SHED_CODE + ), f"burst request neither served nor shed: {resp.status_code} {resp.text[:200]}" + assert any(r.status_code == 200 for r in responses), "no burst request was served" + + _wait_until_served(gateway, model, timeout=90.0) + + +@pytest.mark.engine("sglang") +@pytest.mark.gpu(1) +@pytest.mark.e2e +@pytest.mark.model(_MODEL) +@pytest.mark.workers(count=1, extra_engine_args=["--max-running-requests", "1"]) +@pytest.mark.gateway(extra_args=_GATEWAY_ARGS) +@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True) +class TestOverloadShedSglang(_OverloadShedBase): + """SGLang runs one request at a time; the rest queue.""" + + +@pytest.mark.engine("vllm") +@pytest.mark.gpu(1) +@pytest.mark.e2e +@pytest.mark.model(_MODEL) +@pytest.mark.workers(count=1, extra_engine_args=["--max-num-seqs", "1"]) +@pytest.mark.gateway(extra_args=_GATEWAY_ARGS) +@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True) +class TestOverloadShedVllm(_OverloadShedBase): + """vLLM runs one sequence at a time; the rest queue.""" diff --git a/e2e_test/router/test_routing_key.py b/e2e_test/router/test_routing_key.py new file mode 100644 index 0000000000..8520d10f6f --- /dev/null +++ b/e2e_test/router/test_routing_key.py @@ -0,0 +1,119 @@ +"""Routing keys stay sticky and a body ``rid`` outranks the header. + +With ``--routing-key-override`` every request must take the buffered path so +the body's ``rid`` can win over the routing-key header. A streamed +pass-through had silently dropped that precedence, so two requests from one +``rid`` lineage could land on different workers (#2355). Which worker served +a request is observed through the gateway's in-flight load while a long +generation is running. + +Usage: + E2E_RUNTIME=sglang pytest e2e_test/router/test_routing_key.py -v +""" + +from __future__ import annotations + +import logging +import threading +import time +from collections.abc import Callable + +import httpx +import pytest + +logger = logging.getLogger(__name__) + +KEY_HEADER = "x-smg-routing-key" +_GATEWAY_ARGS = ["--routing-key-override"] + + +def _request(gateway, model: str, *, key: str, rid: str | None) -> Callable[[], None]: + """Build a sender for one long generation carrying a header key and maybe a body rid.""" + body: dict = { + "model": model, + "messages": [{"role": "user", "content": "Write a long story about a river."}], + "max_tokens": 200, + "temperature": 0.0, + "ignore_eos": True, + } + if rid is not None: + body["rid"] = rid + + def _send() -> None: + resp = httpx.post( + f"{gateway.base_url}/v1/chat/completions", + headers={KEY_HEADER: key}, + json=body, + timeout=120.0, + ) + assert resp.status_code == 200, f"{resp.status_code} {resp.text[:200]}" + + return _send + + +def _serving_worker(gateway, send: Callable[[], None]) -> str: + """Run ``send`` and return the URL of the worker that carried it in flight.""" + failures: list[BaseException] = [] + + def _run() -> None: + try: + send() + except BaseException as exc: # surfaced after the join + failures.append(exc) + + thread = threading.Thread(target=_run, daemon=True) + thread.start() + busy: set[str] = set() + deadline = time.monotonic() + 30.0 + while time.monotonic() < deadline and thread.is_alive(): + busy.update(w.url for w in gateway.list_workers(strict=True) if w.pending_requests > 0) + if busy: + break + time.sleep(0.02) + thread.join(timeout=120) + assert not failures, f"request failed: {failures[0]!r}" + assert len(busy) == 1, f"expected exactly one worker to carry the request, saw {sorted(busy)}" + return busy.pop() + + +@pytest.mark.engine("sglang", "vllm") +@pytest.mark.gpu(2) +@pytest.mark.e2e +@pytest.mark.model("meta-llama/Llama-3.2-1B-Instruct") +@pytest.mark.workers(count=2) +@pytest.mark.gateway(policy="manual", extra_args=_GATEWAY_ARGS) +@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True) +class TestRoutingKeyPinning: + """Two workers, manual policy: keys pin, and the body rid decides the key.""" + + def test_header_key_is_sticky(self, setup_backend): + _, model, _, gateway = setup_backend + assert len(gateway.list_workers(strict=True)) == 2 + + served = { + _serving_worker(gateway, _request(gateway, model, key="session-sticky", rid=None)) + for _ in range(3) + } + + assert len(served) == 1, f"one key reached more than one worker: {sorted(served)}" + + def test_body_rid_outranks_header_key(self, setup_backend): + _, model, _, gateway = setup_backend + + # Find two keys the policy assigned to different workers. + homes: dict[str, str] = {} + for i in range(8): + key = f"lineage-{i}" + homes[key] = _serving_worker(gateway, _request(gateway, model, key=key, rid=None)) + if len(set(homes.values())) == 2: + break + assert len(set(homes.values())) == 2, f"every key landed on one worker: {homes}" + key_a, key_b = list(homes)[-2:] + worker_a, worker_b = homes[key_a], homes[key_b] + if worker_a == worker_b: # the last two keys share a home; pick a differing pair + key_a = next(k for k, w in homes.items() if w != worker_b) + worker_a = homes[key_a] + logger.info("homes: %s -> %s, %s -> %s", key_a, worker_a, key_b, worker_b) + + assert _serving_worker(gateway, _request(gateway, model, key=key_a, rid=key_b)) == worker_b + assert _serving_worker(gateway, _request(gateway, model, key=key_b, rid=key_a)) == worker_a diff --git a/e2e_test/router/test_worker_restart.py b/e2e_test/router/test_worker_restart.py new file mode 100644 index 0000000000..3842f836c3 --- /dev/null +++ b/e2e_test/router/test_worker_restart.py @@ -0,0 +1,132 @@ +"""A worker that dies and returns on the same port rejoins rotation. + +The gateway must notice a dead gRPC worker through its health probes, fail +fast while it is gone rather than hang, and put the worker back into rotation +once it is healthy again. The SGLang gRPC server used to race its own IPC +binding during startup (#2427); restarting the worker under a live gateway is +the closest end-to-end reproduction of that path. + +Usage: + E2E_RUNTIME=sglang pytest e2e_test/router/test_worker_restart.py -v +""" + +from __future__ import annotations + +import logging +import os +import time + +import httpx +import pytest +from infra import ConnectionMode, Gateway, get_pool +from infra.constants import get_runtime +from infra.model_specs import get_model_spec + +logger = logging.getLogger(__name__) + +MODEL = "meta-llama/Llama-3.2-1B-Instruct" +_HEALTH_ARGS = [ + "--health-check-interval-secs", + "1", + "--health-check-timeout-secs", + "2", + "--health-failure-threshold", + "1", + "--health-success-threshold", + "1", +] + + +def _status(gateway: Gateway, url: str) -> str | None: + for worker in gateway.list_workers(strict=True): + if worker.url == url: + return worker.status + return None + + +def _wait_for_status(gateway: Gateway, url: str, wanted: str, timeout: float) -> None: + deadline = time.monotonic() + timeout + seen: str | None = None + while time.monotonic() < deadline: + seen = _status(gateway, url) + if seen == wanted: + return + time.sleep(0.5) + pytest.fail(f"worker {url} never became {wanted} within {timeout:.0f}s (last: {seen})") + + +def _chat(gateway: Gateway, model_path: str, timeout: float) -> httpx.Response: + return httpx.post( + f"{gateway.base_url}/v1/chat/completions", + json={ + "model": model_path, + "messages": [{"role": "user", "content": "Say hello."}], + "max_tokens": 4, + }, + timeout=timeout, + ) + + +def _wait_until_served(gateway: Gateway, model_path: str, timeout: float) -> None: + deadline = time.monotonic() + timeout + last = "no attempt" + while time.monotonic() < deadline: + try: + resp = _chat(gateway, model_path, timeout=30.0) + except httpx.HTTPError as exc: + last = repr(exc) + else: + if resp.status_code == 200: + return + last = f"{resp.status_code} {resp.text[:200]}" + time.sleep(1.0) + pytest.fail(f"gateway did not serve within {timeout:.0f}s; last: {last}") + + +@pytest.mark.engine("sglang", "vllm") +@pytest.mark.gpu(1) +@pytest.mark.e2e +class TestWorkerRestart: + """Stop the pooled gRPC worker under a live gateway, then bring it back.""" + + def test_worker_returns_to_rotation_after_restart(self): + engine = get_runtime() + model_path = get_model_spec(MODEL)["model"] + # The session pool owns the GPU; borrowing its worker avoids racing a + # cached worker for memory and leaves a healthy worker behind for the + # next class. + worker = get_pool().acquire( + model_id=MODEL, + engine=engine, + mode=ConnectionMode.GRPC, + count=1, + log_dir=os.environ.get("E2E_LOG_DIR"), + )[0] + gateway = Gateway() + try: + gateway.start( + worker_urls=[worker.base_url], model_path=model_path, extra_args=_HEALTH_ARGS + ) + _wait_until_served(gateway, model_path, timeout=120.0) + + worker.stop() + _wait_for_status(gateway, worker.base_url, "unhealthy", timeout=30.0) + started = time.monotonic() + resp = _chat(gateway, model_path, timeout=30.0) + elapsed = time.monotonic() - started + logger.info( + "while down: status=%s after %.1fs body=%s", + resp.status_code, + elapsed, + resp.text[:200], + ) + assert 500 <= resp.status_code < 600, ( + f"request during the outage should fail, got {resp.status_code}: {resp.text[:200]}" + ) + assert elapsed < 20.0, f"request during the outage hung for {elapsed:.1f}s" + + worker.start() # same port, same URL + _wait_for_status(gateway, worker.base_url, "healthy", timeout=240.0) + _wait_until_served(gateway, model_path, timeout=60.0) + finally: + gateway.shutdown() From 5abe7fdecb2fecafd983ef286127b2cc43624ce1 Mon Sep 17 00:00:00 2001 From: Alex McC <319643551+hello-alexmcc@users.noreply.github.com> Date: Mon, 7 Sep 2026 17:10:42 -0700 Subject: [PATCH 2/4] test(e2e): run the routing-key pinning check on TokenSpeed too Signed-off-by: Alex McC <319643551+hello-alexmcc@users.noreply.github.com> --- e2e_test/router/test_routing_key.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/e2e_test/router/test_routing_key.py b/e2e_test/router/test_routing_key.py index 8520d10f6f..ebb216e372 100644 --- a/e2e_test/router/test_routing_key.py +++ b/e2e_test/router/test_routing_key.py @@ -76,7 +76,7 @@ def _run() -> None: return busy.pop() -@pytest.mark.engine("sglang", "vllm") +@pytest.mark.engine("sglang", "vllm", "tokenspeed") @pytest.mark.gpu(2) @pytest.mark.e2e @pytest.mark.model("meta-llama/Llama-3.2-1B-Instruct") From 8511d8d8462f885e6f00ef3ad93c47407dae4553 Mon Sep 17 00:00:00 2001 From: Alex McC <319643551+hello-alexmcc@users.noreply.github.com> Date: Mon, 7 Sep 2026 17:46:27 -0700 Subject: [PATCH 3/4] test(e2e): expect the 503 for a worker outage and satisfy CI's mypy The first CI run of the worker-restart test caught the gateway answering 404 model_not_found while the only worker was restarting, on SGLang and on vLLM. That answer is fixed in #2465; the test now asserts the 503 with no_available_workers instead of any 5xx. The queue-depth helper raises instead of calling pytest.fail so mypy sees the missing return without pytest stubs, which is how the lint job runs it. Signed-off-by: Alex McC <319643551+hello-alexmcc@users.noreply.github.com> --- e2e_test/router/test_overload.py | 4 +++- e2e_test/router/test_worker_restart.py | 17 +++++++++++++++-- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/e2e_test/router/test_overload.py b/e2e_test/router/test_overload.py index 0db81bf4c8..e22e208e6b 100644 --- a/e2e_test/router/test_overload.py +++ b/e2e_test/router/test_overload.py @@ -91,7 +91,9 @@ def _wait_for_deep_queue(gateway, min_waiting: int, timeout: float) -> dict: if entry.get("num_waiting_reqs", 0) >= min_waiting: return entry time.sleep(0.2) - pytest.fail(f"engine queue never reached {min_waiting} waiting requests; last reports: {last}") + raise AssertionError( + f"engine queue never reached {min_waiting} waiting requests; last reports: {last}" + ) class _OverloadShedBase: diff --git a/e2e_test/router/test_worker_restart.py b/e2e_test/router/test_worker_restart.py index 3842f836c3..2de588da00 100644 --- a/e2e_test/router/test_worker_restart.py +++ b/e2e_test/router/test_worker_restart.py @@ -67,6 +67,15 @@ def _chat(gateway: Gateway, model_path: str, timeout: float) -> httpx.Response: ) +def _error_code(resp: httpx.Response) -> str | None: + try: + body = resp.json() + except ValueError: + return None + error = body.get("error", body) + return error.get("code") if isinstance(error, dict) else None + + def _wait_until_served(gateway: Gateway, model_path: str, timeout: float) -> None: deadline = time.monotonic() + timeout last = "no attempt" @@ -120,9 +129,13 @@ def test_worker_returns_to_rotation_after_restart(self): elapsed, resp.text[:200], ) - assert 500 <= resp.status_code < 600, ( - f"request during the outage should fail, got {resp.status_code}: {resp.text[:200]}" + # The model exists and its worker is merely down: that is a 503 + # no_available_workers (#2465), not a 404 that tells the client the + # model is gone. + assert resp.status_code == 503, ( + f"request during the outage should be a 503, got {resp.status_code}: {resp.text[:200]}" ) + assert _error_code(resp) == "no_available_workers", resp.text[:200] assert elapsed < 20.0, f"request during the outage hung for {elapsed:.1f}s" worker.start() # same port, same URL From 78c7d76edf06f04dd70c0ca38ca6df6c0f7376e0 Mon Sep 17 00:00:00 2001 From: Alex McC <319643551+hello-alexmcc@users.noreply.github.com> Date: Tue, 8 Sep 2026 21:48:22 -0700 Subject: [PATCH 4/4] test(e2e): run the routing-key pinning checks on a policy the sticky override wraps The body-rid-outranks-header guarantee comes from the registry's sticky override, which only wraps policies that do not key on the header themselves; under manual the policy reads the header alone and ignores the body rid, so the check failed on every attempt. Run it on round_robin, where the override applies. Signed-off-by: Alex McC <319643551+hello-alexmcc@users.noreply.github.com> --- e2e_test/router/test_routing_key.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/e2e_test/router/test_routing_key.py b/e2e_test/router/test_routing_key.py index ebb216e372..abfa68ea9f 100644 --- a/e2e_test/router/test_routing_key.py +++ b/e2e_test/router/test_routing_key.py @@ -81,10 +81,14 @@ def _run() -> None: @pytest.mark.e2e @pytest.mark.model("meta-llama/Llama-3.2-1B-Instruct") @pytest.mark.workers(count=2) -@pytest.mark.gateway(policy="manual", extra_args=_GATEWAY_ARGS) +# The sticky override is what lets a body rid outrank the header, and it only +# wraps policies that do not key on the header themselves: under ``manual`` +# (and ``consistent_hashing``) the policy reads the header alone and a body +# rid is ignored. +@pytest.mark.gateway(policy="round_robin", extra_args=_GATEWAY_ARGS) @pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True) class TestRoutingKeyPinning: - """Two workers, manual policy: keys pin, and the body rid decides the key.""" + """Two workers, sticky override on round_robin: keys pin, and the body rid decides the key.""" def test_header_key_is_sticky(self, setup_backend): _, model, _, gateway = setup_backend