test(gms): add reproducible vLLM and SGLang E2E coverage

Signed-off-by: mkhadkevich <mkhadkevich@nvidia.com>

Author: hutm

tests/router/common.py

@@ -2149,6 +2149,7 @@ def _test_router_decisions_disagg(
     durable_kv_events: bool = False,
     router_aic_config: Optional[dict[str, Any]] = None,
     enable_bootstrap: bool = False,
+    strict_timing: bool = True,
 ):
     """Validate KV cache prefix reuse in disaggregated prefill-decode setup via HTTP frontend.
 
@@ -2172,6 +2173,7 @@ def _test_router_decisions_disagg(
         store_backend: Storage backend to use ("etcd" or "file"). Defaults to "etcd".
         durable_kv_events: If True, use durable KV events (JetStream). Defaults to False.
         router_aic_config: Optional AIC router perf-model config for frontend KV routing.
+        strict_timing: If False, allow backend responses that omit optional timing fields.
 
     Raises:
         AssertionError: If prefill_worker_ids differ across requests (prefix reuse failure)
@@ -2216,7 +2218,7 @@ def _test_router_decisions_disagg(
         )
 
         async def send_progressive_requests():
-            """Send 4 progressive requests with overlapping prefixes and collect worker IDs."""
+            """Send progressive requests with overlapping prefixes and collect worker IDs."""
             prefill_worker_ids = []
             decode_worker_ids = []
 
@@ -2252,10 +2254,13 @@ async def send_progressive_requests():
                             response.status == 200
                         ), f"Request {i + 1} failed with status {response.status}"
 
-                        # Collect all chunks and look for nvext with worker_id and timing
+                        # Collect all chunks and look for nvext with worker_id, timing,
+                        # and actual generated text. A backend can otherwise surface a
+                        # transfer failure as a 200 stream with worker IDs but no tokens.
                         prefill_wid = None
                         decode_wid = None
                         timing_info = None
+                        generated_parts: list[str] = []
 
                         async for line in response.content:
                             if not line:
@@ -2271,6 +2276,12 @@ async def send_progressive_requests():
 
                             try:
                                 data = json.loads(data_str)
+                                for choice in data.get("choices", []):
+                                    delta = choice.get("delta") or {}
+                                    content = delta.get("content")
+                                    if content:
+                                        generated_parts.append(content)
+
                                 # Check for nvext in the response
                                 nvext = data.get("nvext", {})
                                 if nvext:
@@ -2291,9 +2302,15 @@ async def send_progressive_requests():
                             except json.JSONDecodeError:
                                 continue
 
+                        generated_text = "".join(generated_parts)
                         logger.info(
                             f"Request {i + 1}: prefill_worker_id={prefill_wid}, "
-                            f"decode_worker_id={decode_wid}, timing={timing_info}"
+                            f"decode_worker_id={decode_wid}, timing={timing_info}, "
+                            f"generated_chars={len(generated_text)}"
+                        )
+                        assert generated_text.strip(), (
+                            f"Request {i + 1}: Expected generated content in the response stream, "
+                            "but the backend returned no text tokens"
                         )
 
                         if prefill_wid is not None:
@@ -2307,7 +2324,12 @@ async def send_progressive_requests():
                         assert (
                             timing_info is not None
                         ), f"Request {i + 1}: Expected timing info in final chunk, got None"
-                        verify_response_timing(timing_info, disagg=not enable_bootstrap)
+                        verify_response_timing(
+                            timing_info,
+                            disagg=not enable_bootstrap,
+                            require_ttft=strict_timing,
+                            require_kv_transfer_latency=strict_timing,
+                        )
 
                     # Small delay between requests
                     await asyncio.sleep(1)

tests/router/e2e_harness.py

@@ -4,6 +4,7 @@
 import logging
 import os
 import time
+from contextlib import ExitStack, contextmanager
 from typing import Any, Callable, ContextManager
 
 from tests.router.common import (
@@ -19,6 +20,71 @@
 
 logger = logging.getLogger(__name__)
 
+
+def env_int(name: str, default: int) -> int:
+    return int(os.environ.get(name, default))
+
+
+def env_optional_int(name: str) -> int | None:
+    raw = os.environ.get(name)
+    return None if raw is None else int(raw)
+
+
+def env_float(name: str, default: float) -> float:
+    return float(os.environ.get(name, default))
+
+
+def env_bool(name: str, default: bool) -> bool:
+    raw = os.environ.get(name)
+    return default if raw is None else raw.lower() not in {"0", "false", "no", "off"}
+
+
+def resolve_router_gpu_start_index(gpu_start_index: int) -> int:
+    override = os.environ.get("DYNAMO_ROUTER_E2E_GPU_START_INDEX")
+    if override is None:
+        return gpu_start_index
+    try:
+        base_index = int(override)
+    except ValueError as exc:
+        raise ValueError(
+            "DYNAMO_ROUTER_E2E_GPU_START_INDEX must be an integer"
+        ) from exc
+    return base_index + gpu_start_index
+
+
+def _resolve_router_e2e_num_requests(default: int = 10) -> int:
+    override = os.environ.get("DYNAMO_ROUTER_E2E_NUM_REQUESTS")
+    if override is None:
+        return default
+    try:
+        value = int(override)
+    except ValueError as exc:
+        raise ValueError("DYNAMO_ROUTER_E2E_NUM_REQUESTS must be an integer") from exc
+    if value <= 0:
+        raise ValueError("DYNAMO_ROUTER_E2E_NUM_REQUESTS must be positive")
+    return value
+
+
+@contextmanager
+def maybe_router_gms_servers():
+    if os.environ.get("DYNAMO_ROUTER_E2E_ENABLE_GMS") != "1":
+        yield
+        return
+
+    from tests.gpu_memory_service.common.gms import GMSServer
+
+    devices = os.environ.get("DYNAMO_ROUTER_E2E_GMS_DEVICES", "0,1")
+    tags = os.environ.get("DYNAMO_ROUTER_E2E_GMS_TAGS", "weights,kv_cache")
+    device_ids = [int(part.strip()) for part in devices.split(",") if part.strip()]
+    tag_names = [part.strip() for part in tags.split(",") if part.strip()]
+
+    with ExitStack() as stack:
+        for device in device_ids:
+            for tag in tag_names:
+                stack.enter_context(GMSServer(device=device, tag=tag))
+        yield
+
+
 TEST_PROMPT = (
     "In a quiet meadow tucked between rolling hills, a plump gray rabbit nibbled on "
     "clover beneath the shade of a gnarled oak tree. Its ears twitched at the faint "
@@ -214,15 +280,17 @@ def run_basic_router_test(
         },
         engine_process_kwargs=engine_process_kwargs,
     )
-    with process as engine_workers:
+    with ExitStack() as stack:
+        stack.enter_context(maybe_router_gms_servers())
+        engine_workers = stack.enter_context(process)
         frontend_port = allocate_frontend_ports(request, 1)[0]
         _test_router_basic(
             engine_workers=engine_workers,
             block_size=block_size,
             request=request,
             frontend_port=frontend_port,
             test_payload=test_payload or build_test_payload(model_name),
-            num_requests=num_requests,
+            num_requests=_resolve_router_e2e_num_requests(num_requests),
             frontend_timeout=frontend_timeout,
             store_backend="etcd",
             request_plane=request_plane,
@@ -263,7 +331,9 @@ def run_router_decisions_test(
         default_process_kwargs=default_process_kwargs,
         engine_process_kwargs=engine_process_kwargs,
     )
-    with process as engine_workers:
+    with ExitStack() as stack:
+        stack.enter_context(maybe_router_gms_servers())
+        engine_workers = stack.enter_context(process)
         endpoint = get_engine_endpoint(engine_workers, request_plane, component_name)
         scenario_kwargs = dict(test_kwargs or {})
         for argument, attribute in (
@@ -302,6 +372,7 @@ def run_disagg_router_decisions_test(
     | None = None,
     test_payload: dict[str, Any] | None = None,
     test_kwargs: dict[str, Any] | None = None,
+    strict_timing: bool = True,
 ):
     shared_namespace = f"test-namespace-{generate_random_suffix()}"
     frontend_port = allocate_frontend_ports(request, 1)[0]
@@ -316,6 +387,7 @@ def run_disagg_router_decisions_test(
     }
 
     def run_test(prefill_workers, decode_workers):
+        scenario_kwargs = {"strict_timing": strict_timing, **(test_kwargs or {})}
         _test_router_decisions_disagg(
             prefill_workers=prefill_workers,
             decode_workers=decode_workers,
@@ -324,29 +396,35 @@ def run_test(prefill_workers, decode_workers):
             frontend_port=frontend_port,
             test_payload=test_payload or build_test_payload(model_name),
             request_plane=request_plane,
-            **(test_kwargs or {}),
+            **scenario_kwargs,
         )
 
-    if worker_context_factory is not None:
-        with worker_context_factory(shared_namespace) as workers:
+    with ExitStack() as stack:
+        stack.enter_context(maybe_router_gms_servers())
+        if worker_context_factory is not None:
+            workers = stack.enter_context(worker_context_factory(shared_namespace))
             run_test(*workers)
-        return
-
-    with engine_process_cls(
-        request,
-        num_workers=num_prefill_workers,
-        request_plane=request_plane,
-        **{engine_args_name: engine_args},
-        **prefill_kwargs,
-    ) as prefill_workers:
-        with engine_process_cls(
-            request,
-            num_workers=num_decode_workers,
-            request_plane=request_plane,
-            **{engine_args_name: engine_args},
-            **decode_kwargs,
-        ) as decode_workers:
-            run_test(prefill_workers, decode_workers)
+            return
+
+        prefill_workers = stack.enter_context(
+            engine_process_cls(
+                request,
+                num_workers=num_prefill_workers,
+                request_plane=request_plane,
+                **{engine_args_name: engine_args},
+                **prefill_kwargs,
+            )
+        )
+        decode_workers = stack.enter_context(
+            engine_process_cls(
+                request,
+                num_workers=num_decode_workers,
+                request_plane=request_plane,
+                **{engine_args_name: engine_args},
+                **decode_kwargs,
+            )
+        )
+        run_test(prefill_workers, decode_workers)
 
 
 def run_indexers_sync_test(
@@ -383,7 +461,9 @@ def run_indexers_sync_test(
         },
         engine_process_kwargs=engine_process_kwargs,
     )
-    with process as engine_workers:
+    with ExitStack() as stack:
+        stack.enter_context(maybe_router_gms_servers())
+        engine_workers = stack.enter_context(process)
         _test_router_indexers_sync(
             engine_workers=engine_workers,
             block_size=block_size,