test(e2e): cover overload sheds, routing-key pinning, and worker restart

e2e_test/router/test_overload.py

@@ -0,0 +1,176 @@
+"""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)
+    raise AssertionError(
+        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."""

e2e_test/router/test_routing_key.py

@@ -0,0 +1,123 @@
+"""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", "tokenspeed")
+@pytest.mark.gpu(2)
+@pytest.mark.e2e
+@pytest.mark.model("meta-llama/Llama-3.2-1B-Instruct")
+@pytest.mark.workers(count=2)
+# 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, 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
+        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