feat(zmq): typed structured outputs for the direct backend, with e2e coverage

.github/workflows/e2e-gpu-job.yml

@@ -54,6 +54,11 @@ on:
         type: string
         default: "nixl"
         description: "KV transfer backend for vLLM PD workers: nixl or mooncake"
+      connection_mode:
+        required: false
+        type: string
+        default: ""
+        description: "Wire override for local backends: zmq runs the local cases over ZMQ"
 
 jobs:
   run:
@@ -66,6 +71,7 @@ jobs:
       E2E_RUNTIME: ${{ inputs.engine }}
       E2E_GPU_TIER: ${{ inputs.gpu_tier }}
       E2E_VLLM_KV_BACKEND: ${{ inputs.vllm_kv_backend }}
+      E2E_CONNECTION_MODE: ${{ inputs.connection_mode }}
       ROUTER_LOCAL_MODEL_PATH: /models
     steps:
       - name: Checkout code

.github/workflows/pr-test-rust.yml

@@ -543,6 +543,42 @@ jobs:
       test_dirs: ${{ matrix.test_dirs || 'e2e_test/chat_completions' }}
     secrets: inherit
 
+  e2e-1gpu-chat-zmq:
+    name: e2e-1gpu-chat-zmq (${{ matrix.engine }})
+    needs: [build-wheel, detect-changes]
+    if: >-
+      always()
+      && !cancelled()
+      && needs.build-wheel.result == 'success'
+      && (github.event_name != 'pull_request'
+          || (needs.detect-changes.result == 'success'
+              && (needs.detect-changes.outputs.common == 'true'
+                  || needs.detect-changes.outputs.chat-completions == 'true')))
+    # Same single-worker chat suite as e2e-1gpu-chat, driven over the ZMQ
+    # direct-backend wire. The collection hook deselects the gRPC-only
+    # families (PD, EPD, multi-worker) for a ZMQ lane, so this covers the
+    # local cases only.
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - engine: vllm
+            timeout: 24
+            test_timeout: 18
+          - engine: tokenspeed
+            timeout: 50
+            test_timeout: 18
+    uses: ./.github/workflows/e2e-gpu-job.yml
+    with:
+      engine: ${{ matrix.engine }}
+      gpu_tier: "1"
+      runner: 1-gpu-h100
+      timeout: ${{ matrix.timeout }}
+      test_timeout: ${{ matrix.test_timeout }}
+      test_dirs: e2e_test/chat_completions
+      connection_mode: zmq
+    secrets: inherit
+
   e2e-1gpu-completions:
     name: e2e-1gpu-completions (${{ matrix.engine }})
     needs: [build-wheel, detect-changes]
@@ -1049,7 +1085,7 @@ jobs:
           path: benchmark_go_bindings/
 
   finish:
-    needs: [pre-commit, python-lint, grpc-proto-build-check, build-wheel, python-unit-tests, unit-tests, benchmarks, e2e-1gpu-chat, e2e-1gpu-completions, e2e-1gpu-embeddings, e2e-1gpu-gateway, e2e-1gpu-responses, e2e-2gpu-pd, e2e-4gpu-chat, e2e-4gpu-gateway, e2e-4gpu-epd, e2e-vendor, go-unit-tests, go-bindings-e2e]
+    needs: [pre-commit, python-lint, grpc-proto-build-check, build-wheel, python-unit-tests, unit-tests, benchmarks, e2e-1gpu-chat, e2e-1gpu-chat-zmq, e2e-1gpu-completions, e2e-1gpu-embeddings, e2e-1gpu-gateway, e2e-1gpu-responses, e2e-2gpu-pd, e2e-4gpu-chat, e2e-4gpu-gateway, e2e-4gpu-epd, e2e-vendor, go-unit-tests, go-bindings-e2e]
     if: always()
     runs-on: k8s-runner-cpu
     permissions: {}
@@ -1064,6 +1100,7 @@ jobs:
                 "${{ needs.unit-tests.result }}" == "failure" || \
                 "${{ needs.benchmarks.result }}" == "failure" || \
                 "${{ needs.e2e-1gpu-chat.result }}" == "failure" || \
+                "${{ needs.e2e-1gpu-chat-zmq.result }}" == "failure" || \
                 "${{ needs.e2e-1gpu-completions.result }}" == "failure" || \
                 "${{ needs.e2e-1gpu-embeddings.result }}" == "failure" || \
                 "${{ needs.e2e-1gpu-gateway.result }}" == "failure" || \

e2e_test/fixtures/hooks.py

@@ -14,10 +14,16 @@
 import os
 
 import pytest
-from infra import cleanup_pool, get_runtime
+from infra import ConnectionMode, cleanup_pool, get_connection_mode_override, get_runtime
 
 from .markers import resolve_class_marker
 
+# Local wires a plain (non-PD/EPD) test case can be authored with; these are the
+# only backends a ZMQ lane reuses (mapped onto ZMQ by the ``setup_backend``
+# fixture). Anything else — ``pd_*``/``epd_*``, EPD topology tuples, cloud
+# vendor names — is out of scope for ZMQ.
+_ZMQ_LOCAL_WIRES = frozenset({"grpc", "http"})
+
 # ---------------------------------------------------------------------------
 # Marker registration
 # ---------------------------------------------------------------------------
@@ -117,14 +123,76 @@ def _get_marker(item: pytest.Item, name: str):
     return resolve_class_marker(item, name)
 
 
+def _setup_backend_param(item: pytest.Item):
+    """Return the ``setup_backend`` parametrize value for an item, or None."""
+    callspec = getattr(item, "callspec", None)
+    if callspec is None:
+        return None
+    return getattr(callspec, "params", {}).get("setup_backend")
+
+
+def _is_multi_worker(item: pytest.Item) -> bool:
+    """True when the item's ``workers`` marker asks for a PD/multi-worker topology."""
+    marker = resolve_class_marker(item, "workers")
+    if marker is None:
+        return False
+    kwargs = marker.kwargs
+    if (kwargs.get("count") or 1) > 1:
+        return True
+    return bool(kwargs.get("prefill") or kwargs.get("decode"))
+
+
+def _zmq_dedup_key(item: pytest.Item) -> tuple:
+    """Group key ignoring the ``setup_backend`` value.
+
+    Lets us collapse a case parametrized on both ``grpc`` and ``http`` into a
+    single ZMQ run (they map onto the same wire) while keeping distinct
+    ``api_client`` (or other) parametrizations apart.
+    """
+    callspec = getattr(item, "callspec", None)
+    params = getattr(callspec, "params", {}) or {}
+    others = tuple(sorted((k, repr(v)) for k, v in params.items() if k != "setup_backend"))
+    return (item.nodeid.split("[", 1)[0], others)
+
+
+def _filter_zmq_items(items: list[pytest.Item]) -> tuple[list[pytest.Item], list[pytest.Item]]:
+    """Split items into (kept, deselected) for a ZMQ lane.
+
+    Keeps single-worker local cases (``grpc``/``http`` map onto ZMQ) and drops
+    the gRPC-only families: PD (``pd_*``), EPD (``epd_*`` and topology tuples),
+    and multiple-worker topologies. A case authored for both ``grpc`` and
+    ``http`` is collapsed to one ZMQ run. Items without a ``setup_backend``
+    parametrization are left untouched.
+    """
+    kept: list[pytest.Item] = []
+    deselected: list[pytest.Item] = []
+    groups_with_grpc = {_zmq_dedup_key(it) for it in items if _setup_backend_param(it) == "grpc"}
+    for item in items:
+        param = _setup_backend_param(item)
+        if param is None:
+            kept.append(item)
+            continue
+        if param not in _ZMQ_LOCAL_WIRES or _is_multi_worker(item):
+            deselected.append(item)
+            continue
+        # Collapse the http twin when a grpc one covers the same ZMQ run.
+        if param == "http" and _zmq_dedup_key(item) in groups_with_grpc:
+            deselected.append(item)
+            continue
+        kept.append(item)
+    return kept, deselected
+
+
 def pytest_collection_modifyitems(
     config: pytest.Config,
     items: list[pytest.Item],
 ) -> None:
     """Filter + order collected tests.
 
     Filtering: env vars ``E2E_ENGINE``, ``E2E_VENDOR``, ``E2E_GPU_TIER``
-    select the matching slice when set.
+    select the matching slice when set. When ``E2E_CONNECTION_MODE=zmq`` the
+    lane additionally drops the gRPC-only families (PD, EPD, multi-worker) and
+    collapses ``grpc``/``http`` twins onto a single ZMQ run.
 
     Ordering: items are sorted by ``(backend, model)`` so consecutive
     classes that share a backend cluster together. This is what lets
@@ -157,6 +225,12 @@ def pytest_collection_modifyitems(
             selected.append(item)
         items[:] = selected
 
+    if get_connection_mode_override() == ConnectionMode.ZMQ:
+        kept, deselected = _filter_zmq_items(items)
+        if deselected:
+            config.hook.pytest_deselected(items=deselected)
+            items[:] = kept
+
     items.sort(key=_pool_sort_key)
 
 

e2e_test/fixtures/setup_backend.py

@@ -30,6 +30,7 @@
     ConnectionMode,
     Gateway,
     WorkerType,
+    get_connection_mode_override,
     get_runtime,
     launch_cloud_gateway,
 )
@@ -145,6 +146,11 @@ def setup_backend(request: pytest.FixtureRequest):
     is_pd = backend_name.startswith("pd_")
     protocol = backend_name.replace("epd_", "").replace("pd_", "")
     connection_mode = ConnectionMode(protocol)
+    # A lane can override the local wire (e.g. run grpc/http cases over ZMQ);
+    # PD/EPD keep their own wire since they are excluded from those lanes.
+    mode_override = get_connection_mode_override()
+    if mode_override is not None and not is_pd and not is_epd:
+        connection_mode = mode_override
     engine = get_runtime()
     model_path = get_model_spec(model_id)["model"]
     workers_config = get_marker_kwargs(request, "workers", defaults=_WORKER_DEFAULTS)
@@ -232,18 +238,30 @@ def _setup_local(
         gpus=workers_config.get("gpus"),
         extra_engine_args=workers_config.get("extra_engine_args"),
     )
+    # ZMQ engines dial this gateway's handshake sockets, so they cannot be
+    # reused by a later class's gateway — the pool starts them fresh and the
+    # caller owns their teardown (like the PD path). gRPC/HTTP workers stay
+    # in the pool and outlive the gateway.
+    is_zmq = connection_mode == ConnectionMode.ZMQ
     try:
         _start_gateway(
             gateway,
             gateway_config,
             worker_urls=[w.base_url for w in workers],
             model_path=model_path,
+            backend=engine if is_zmq else None,
         )
         logger.info("%s backend ready at %s", backend_name, gateway.base_url)
         yield backend_name, model_path, _make_openai_client(gateway), gateway
     finally:
-        logger.info("Tearing down %s backend (workers stay in pool)", backend_name)
+        logger.info(
+            "Tearing down %s backend (%s)",
+            backend_name,
+            "stopping ZMQ workers" if is_zmq else "workers stay in pool",
+        )
         gateway.shutdown()
+        if is_zmq:
+            stop_workers(workers)
 
 
 # ---------------------------------------------------------------------------
@@ -463,20 +481,31 @@ def test_router_state(backend_router):
     backend_name = request.param
     model_id = os.environ.get(ENV_MODEL, DEFAULT_MODEL)
     connection_mode = ConnectionMode(backend_name)
+    mode_override = get_connection_mode_override()
+    if mode_override is not None:
+        connection_mode = mode_override
+    engine = get_runtime()
     model_path = get_model_spec(model_id)["model"]
+    is_zmq = connection_mode == ConnectionMode.ZMQ
 
     # Route through the pool so we evict any cached class-scope worker
-    # holding the GPUs we need. The pool retains ownership; we don't stop
-    # the workers ourselves.
+    # holding the GPUs we need. The pool retains ownership of gRPC/HTTP
+    # workers; ZMQ engines are bound to this gateway, so we stop them here.
     workers = get_pool().acquire(
         model_id=model_id,
-        engine=get_runtime(),
+        engine=engine,
         mode=connection_mode,
         count=1,
     )
     gateway = Gateway()
     try:
-        gateway.start(worker_urls=[w.base_url for w in workers], model_path=model_path)
+        gateway.start(
+            worker_urls=[w.base_url for w in workers],
+            model_path=model_path,
+            backend=engine if is_zmq else None,
+        )
         yield gateway
     finally:
         gateway.shutdown()
+        if is_zmq:
+            stop_workers(workers)

e2e_test/infra/__init__.py

@@ -11,6 +11,7 @@
     DEFAULT_RUNTIME,
     DEFAULT_STARTUP_TIMEOUT,
     ENV_BACKENDS,
+    ENV_CONNECTION_MODE,
     ENV_MODEL,
     ENV_MODELS,
     ENV_RUNTIME,
@@ -32,6 +33,7 @@
     ConnectionMode,
     Runtime,
     WorkerType,
+    get_connection_mode_override,
     get_runtime,
     is_mlx,
     is_sglang,
@@ -102,13 +104,15 @@
     "ENV_BACKENDS",
     "ENV_MODEL",
     "ENV_RUNTIME",
+    "ENV_CONNECTION_MODE",
     "ENV_STARTUP_TIMEOUT",
     "ENV_SKIP_MODEL_POOL",
     "ENV_SKIP_BACKEND_SETUP",
     "ENV_SHOW_ROUTER_LOGS",
     "ENV_SHOW_WORKER_LOGS",
     # Runtime helpers
     "get_runtime",
+    "get_connection_mode_override",
     "is_vllm",
     "is_sglang",
     "is_trtllm",

e2e_test/infra/constants.py

@@ -9,6 +9,7 @@ class ConnectionMode(StrEnum):
 
     HTTP = "http"
     GRPC = "grpc"
+    ZMQ = "zmq"
 
 
 class WorkerType(StrEnum):
@@ -35,7 +36,7 @@ class Runtime(StrEnum):
 
 
 # Convenience sets
-LOCAL_MODES = frozenset({ConnectionMode.HTTP, ConnectionMode.GRPC})
+LOCAL_MODES = frozenset({ConnectionMode.HTTP, ConnectionMode.GRPC, ConnectionMode.ZMQ})
 LOCAL_RUNTIMES = frozenset(
     {Runtime.SGLANG, Runtime.VLLM, Runtime.TRTLLM, Runtime.MLX, Runtime.TOKENSPEED}
 )
@@ -59,6 +60,9 @@ class Runtime(StrEnum):
 ENV_RUNTIME = (
     "E2E_RUNTIME"  # Runtime for gRPC tests — one of Runtime.{SGLANG,VLLM,TRTLLM,TOKENSPEED}
 )
+ENV_CONNECTION_MODE = (
+    "E2E_CONNECTION_MODE"  # Per-lane wire override — see get_connection_mode_override
+)
 ENV_STARTUP_TIMEOUT = "E2E_STARTUP_TIMEOUT"
 ENV_SKIP_MODEL_POOL = "SKIP_MODEL_POOL"
 ENV_SKIP_BACKEND_SETUP = "SKIP_BACKEND_SETUP"
@@ -125,6 +129,18 @@ def is_tokenspeed() -> bool:
     return get_runtime() == "tokenspeed"
 
 
+def get_connection_mode_override() -> "ConnectionMode | None":
+    """Per-lane wire-protocol override for local backends.
+
+    Set ``E2E_CONNECTION_MODE`` to run the existing local test cases over a
+    different wire (like ``E2E_RUNTIME`` picks the engine): a ``grpc``/``http``
+    case then runs over that mode without a separate parametrize value. PD/EPD
+    backends keep their own wire. Returns ``None`` when unset.
+    """
+    value = os.environ.get(ENV_CONNECTION_MODE)
+    return ConnectionMode(value.lower()) if value else None
+
+
 ENV_VLLM_KV_BACKEND = "E2E_VLLM_KV_BACKEND"
 
 

e2e_test/infra/gateway.py

@@ -122,6 +122,7 @@ def start(
         igw_mode: bool = False,
         cloud_backend: str | None = None,
         history_backend: str = "memory",
+        backend: str | None = None,
         policy: str = "round_robin",
         timeout: float = DEFAULT_ROUTER_TIMEOUT,
         show_output: bool | None = None,
@@ -229,8 +230,13 @@ def start(
             self.model_path = model_path
             self.pd_mode = False
             self.igw_mode = False
+            mode_args = ["--model-path", model_path, "--worker-urls", *worker_urls]
+            # ZMQ workers share one wire across engine runtimes, so the router
+            # cannot probe the backend from the ipc:// URL — pin it explicitly.
+            if backend is not None:
+                mode_args += ["--router-backend", backend]
             self._launch(
-                mode_args=["--model-path", model_path, "--worker-urls", *worker_urls],
+                mode_args=mode_args,
                 timeout=timeout,
                 show_output=show_output,
                 extra_args=extra_args,