test(zmq): direct-backend e2e tests and CI lanes - #2060
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe PR adds optional connection-mode workflow wiring, separates diagnostic artifacts by mode, adds ZMQ validation and command-builder tests, filters redundant ZMQ fixture cases, skips unsupported TokenSpeed tests, and runs vLLM and TokenSpeed ZMQ chat tests in CI. ChangesZMQ E2E coverage
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant CIWorkflow
participant E2EGPUWorkflow
participant E2EWorker
CIWorkflow->>E2EGPUWorkflow: pass connection_mode
E2EGPUWorkflow->>E2EWorker: export E2E_CONNECTION_MODE
E2EWorker->>E2EGPUWorkflow: write mode-specific diagnostic artifact
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Clean PR — thorough unit tests for the ZMQ harness (connection-mode parsing, engine validation, headless command builders, wire-family dedup/deselect hooks) and well-structured CI integration. All test assertions verified against source implementations. No issues found.
0 🔴 Important · 0 🟡 Nit · 0 🟣 Pre-existing
6161ff8 to
db11d98
Compare
873cdd8 to
6a9a08a
Compare
db11d98 to
55e1108
Compare
6a9a08a to
410a8b1
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
e2e_test/fixtures/test_hooks_zmq_filter.py (1)
73-78: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win🟡 Nit — Test gRPC preference in both collection orders.
The test only checks
[grpc, http]. Add[http, grpc]and still require that gRPC is kept. This detects an order-dependent filter that keeps HTTP when collection order changes.As per coding guidelines, use the pr-test-analyzer agent to verify tests adequately cover new or changed functionality.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@e2e_test/fixtures/test_hooks_zmq_filter.py` around lines 73 - 78, Extend test_grpc_http_twins_collapse_to_grpc to exercise both collection orders, including [http, grpc], and assert that grpc remains kept while http is deselected in each case. Use the pr-test-analyzer agent to verify the tests adequately cover this filtering behavior.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@e2e_test/fixtures/test_connection_mode_validation.py`:
- Around line 27-34: Expand the engine parameter list in
test_non_zmq_modes_accept_any_engine to include Runtime.TRTLLM.value and
Runtime.MLX.value alongside the existing runtimes, ensuring both GRPC and HTTP
modes are validated against every supported engine.
- Line 9: Remove the test-collection dependency on setup_backend in
test_connection_mode_validation.py so missing anthropic/openai SDKs cannot
silently skip _validate_connection_mode coverage. Move or reuse
_validate_connection_mode from a dependency-free location such as
infra.constants, and keep the validation tests runnable through the Python
unit-test path without cloud SDK imports.
---
Nitpick comments:
In `@e2e_test/fixtures/test_hooks_zmq_filter.py`:
- Around line 73-78: Extend test_grpc_http_twins_collapse_to_grpc to exercise
both collection orders, including [http, grpc], and assert that grpc remains
kept while http is deselected in each case. Use the pr-test-analyzer agent to
verify the tests adequately cover this filtering behavior.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 8871e97b-dd96-4223-9e5f-8028a51e913e
📒 Files selected for processing (7)
.github/workflows/e2e-gpu-job.yml.github/workflows/pr-test-rust.ymle2e_test/fixtures/test_connection_mode_validation.pye2e_test/fixtures/test_hooks_zmq_filter.pye2e_test/infra/test_connection_mode.pye2e_test/infra/test_zmq_cmd_builders.pyscripts/ci_install_tokenspeed.sh
| from infra.constants import ConnectionMode, Runtime | ||
|
|
||
| # setup_backend pulls in the cloud SDKs; skip if the env lacks them. | ||
| setup_backend = pytest.importorskip("fixtures.setup_backend") |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate the jobs that collect this test and the dependency-install steps.
rg -n -C 4 \
'test_connection_mode_validation|python-unit-tests|ci_install_e2e_deps|fixtures\.setup_backend' \
.github scripts e2e_test 2>/dev/null || trueRepository: smg-project/smg
Length of output: 9462
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== candidate workflow names =="
git ls-files .github/workflows scripts e2e_test | sort
echo
echo "== setup_backend contents =="
fd -a 'setup_backend\.py$|setup_backend/py$' e2e_test 2>/dev/null | sed 's#^\./##'
setup_files=$(fd 'setup_backend\.py$' e2e_test 2>/dev/null || true)
for f in $setup_files; do
echo "--- $f ---"
wc -l "$f"
sed -n '1,220p' "$f"
done
echo
echo "== test_connection_mode_validation outline and contents =="
wc -l e2e_test/fixtures/test_connection_mode_validation.py
sed -n '1,220p' e2e_test/fixtures/test_connection_mode_validation.py
echo
echo "== dependency install script =="
sed -n '1,220p' scripts/ci_install_e2e_deps.sh
echo
echo "== pytest config and invocation references =="
rg -n -C 3 'pytest|e2e_test/fixtures/test_connection_mode_validation|fixtures/test_connection_mode_validation|test_connection_mode_validation|setup_backend|python-unit-tests|ci_install_e2e_deps' pyproject.toml pytest.ini setup.cfg tox.ini .github scripts pyproject*.toml 2>/dev/null || trueRepository: smg-project/smg
Length of output: 40171
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== setup_backend imports section =="
sed -n '1,70p' e2e_test/fixtures/setup_backend.py
echo
echo "== cloud SDK / setup_backend dependent imports =="
python3 - <<'PY'
import ast
from pathlib import Path
path = Path("e2e_test/fixtures/setup_backend.py")
tree = ast.parse(path.read_text(), path=path)
imports = []
functions = {}
classes = {}
for node in ast.walk(tree):
if isinstance(node, ast.FunctionDef):
functions[node.name] = (node.first_arg_line := node.lineno)
elif isinstance(node, ast.ClassDef):
classes[node.name] = node.lineno
for node in tree.body:
if isinstance(node, ast.Import):
for alias in node.names:
imports.append(alias.name)
elif isinstance(node, ast.ImportFrom):
imports.append((node.module, {alias.name for alias in node.names}))
for name in ("_setup_cloud", "_make_client", "_start_worker", "_start_gateway", "setup_backend"):
print(f"{name}: imports={functions.get(name, [])}")
print("top-level imports:", imports)
print("cloud packages in imports:", [x for x in imports if isinstance(x, str) and x in {"anthropic", "openai"}])
PY
echo
echo "== e2e_test metadata =="
for f in pyproject.toml setup.py setup.cfg poetry.lock uv.lock requirements*.txt; do
[ -f "e2e_test/$f" ] && { echo "--- e2e_test/$f ---"; sed -n '1,220p' "e2e_test/$f"; }
done
for f in e2e_test/pyproject.toml setup.py setup.cfg; do [ -f "$f" ] && { echo "--- $f ---"; sed -n '1,220p' "$f"; }; done
echo
echo "== changed files hint =="
git status --short 2>/dev/null | sed -n '1,120p' || trueRepository: smg-project/smg
Length of output: 2456
🔴 Do not make connection-mode validation dependent on fixtures.setup_backend.
test_connection_mode_validation.py is collected with the Python unit-test path, but importing fixtures.setup_backend imports anthropic and openai. A missing cloud SDK makes these unit tests successful skips instead of running the _validate_connection_mode coverage. Move the validator to infra.constants, a dependency-free test helper, or make the skip fail/noise loudly instead of silencing the suite.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@e2e_test/fixtures/test_connection_mode_validation.py` at line 9, Remove the
test-collection dependency on setup_backend in
test_connection_mode_validation.py so missing anthropic/openai SDKs cannot
silently skip _validate_connection_mode coverage. Move or reuse
_validate_connection_mode from a dependency-free location such as
infra.constants, and keep the validation tests runnable through the Python
unit-test path without cloud SDK imports.
Source: Coding guidelines
| @pytest.mark.parametrize("mode", [ConnectionMode.GRPC, ConnectionMode.HTTP]) | ||
| @pytest.mark.parametrize( | ||
| "engine", | ||
| [Runtime.SGLANG.value, Runtime.VLLM.value, Runtime.TOKENSPEED.value], | ||
| ) | ||
| def test_non_zmq_modes_accept_any_engine(mode, engine): | ||
| # Non-ZMQ wires impose no engine restriction. | ||
| setup_backend._validate_connection_mode(mode, engine) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🟡 Nit — Cover all engines for non-ZMQ modes.
test_non_zmq_modes_accept_any_engine omits Runtime.TRTLLM and Runtime.MLX, although lines 20-24 establish that both are valid engine inputs for the same validator. Add them to prevent an accidental gRPC or HTTP restriction from passing untested.
Proposed test expansion
"engine",
- [Runtime.SGLANG.value, Runtime.VLLM.value, Runtime.TOKENSPEED.value],
+ [
+ Runtime.SGLANG.value,
+ Runtime.VLLM.value,
+ Runtime.TOKENSPEED.value,
+ Runtime.TRTLLM.value,
+ Runtime.MLX.value,
+ ],
)As per coding guidelines, use the pr-test-analyzer agent to verify tests adequately cover new or changed functionality.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| @pytest.mark.parametrize("mode", [ConnectionMode.GRPC, ConnectionMode.HTTP]) | |
| @pytest.mark.parametrize( | |
| "engine", | |
| [Runtime.SGLANG.value, Runtime.VLLM.value, Runtime.TOKENSPEED.value], | |
| ) | |
| def test_non_zmq_modes_accept_any_engine(mode, engine): | |
| # Non-ZMQ wires impose no engine restriction. | |
| setup_backend._validate_connection_mode(mode, engine) | |
| `@pytest.mark.parametrize`("mode", [ConnectionMode.GRPC, ConnectionMode.HTTP]) | |
| `@pytest.mark.parametrize`( | |
| "engine", | |
| [ | |
| Runtime.SGLANG.value, | |
| Runtime.VLLM.value, | |
| Runtime.TOKENSPEED.value, | |
| Runtime.TRTLLM.value, | |
| Runtime.MLX.value, | |
| ], | |
| ) | |
| def test_non_zmq_modes_accept_any_engine(mode, engine): | |
| # Non-ZMQ wires impose no engine restriction. | |
| setup_backend._validate_connection_mode(mode, engine) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@e2e_test/fixtures/test_connection_mode_validation.py` around lines 27 - 34,
Expand the engine parameter list in test_non_zmq_modes_accept_any_engine to
include Runtime.TRTLLM.value and Runtime.MLX.value alongside the existing
runtimes, ensuring both GRPC and HTTP modes are validated against every
supported engine.
Source: Coding guidelines
410a8b1 to
d4983d9
Compare
27281ab to
a0d8519
Compare
e572c3c to
fb8136e
Compare
- Add unit tests for the ZMQ test harness: connection-mode parsing and engine-capability validation, headless ZMQ command builders, and the pytest wire-family dedup/deselect hooks. - Wire the ZMQ lanes into CI: run the Rust ZMQ tests on the PR lane (installing TokenSpeed) and add the direct-ZMQ e2e job to the GPU lane. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> (cherry picked from commit 6a9a08a)
fb8136e to
2289ae8
Compare
E2E_ZMQ_ENGINE_COUNT runs a ZMQ lane with grouped workers, riding the lanes #2060 landed: the vLLM worker builder appends the engine-level --data-parallel-size (flowing through the same smg serve launcher as production), the gateway gains --zmq-engine-count so its handshake awaits every engine, and start_workers sizes the worker's GPU slice as tp x engine count. e2e-2gpu-chat-zmq-dp runs the chat suite with dp=2 vLLM groups on the 2-GPU runner, exercising the grouped handshake, the connector's least-loaded selection, and the wave protocol live. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
E2E_ZMQ_ENGINE_COUNT runs a ZMQ lane with grouped workers, riding the lanes #2060 landed: the vLLM worker builder appends the engine-level --data-parallel-size (flowing through the same smg serve launcher as production), the gateway gains --zmq-engine-count so its handshake awaits every engine, and start_workers sizes the worker's GPU slice as tp x engine count. e2e-2gpu-chat-zmq-dp runs the chat suite with dp=2 vLLM groups on the 2-GPU runner, exercising the grouped handshake, the connector's least-loaded selection, and the wave protocol live. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
Description
Problem
The direct-ZMQ backend and its e2e harness are not covered by tests or run in CI,
so regressions in ZMQ dispatch, capability validation, or the headless command
builders would go undetected. Nothing has ever exercised the ZMQ wire end to end
on a GPU lane.
Solution
Add unit tests for the ZMQ harness and wire the ZMQ lanes into CI: a
connection_modeinput on the reusable GPU e2e job, and an
e2e-1gpu-chat-zmqlane that runs theexisting single-worker chat suite over the ZMQ wire for both vLLM and TokenSpeed.
Changes
e2e_test/infra/{test_connection_mode.py,test_zmq_cmd_builders.py}—connection-mode parsing, engine-capability validation, headless command builders.
e2e_test/fixtures/{test_connection_mode_validation.py,test_hooks_zmq_filter.py}— wire-family dedup/deselect hook coverage.
.github/workflows/e2e-gpu-job.yml—connection_modeinput, folded into thelog-artifact name so the gRPC and ZMQ legs don't collide.
.github/workflows/pr-test-rust.yml— thee2e-1gpu-chat-zmqlane, gated on thesame change filters as the gRPC chat lane and added to
finish.scripts/ci_install_tokenspeed.sh— pin bump to a TokenSpeed revision with theZMQ entrypoint.
Test Plan
pytest e2e_test/infra/test_connection_mode.py e2e_test/infra/test_zmq_cmd_builders.py e2e_test/fixtures/test_connection_mode_validation.py e2e_test/fixtures/test_hooks_zmq_filter.py→ 33 passed, 3 skipped (the 3 skips need the
smgwheel installed, which CI has).Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesRebuilt on current
main. The rest of the original 7-PR stack (core dispatch,structured outputs, multimodal, EOS forwarding, harmony stop matcher, e2e infra)
has already landed in reworked form, so this is now a standalone tests+CI change.