Repository navigation
feat(e2e): EPD multimodal smoke CI on 4-gpu-h100 - #1924
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughAdds TokenSpeed EPD worker and gateway support, multimodal routing tests, a Qwen3.5-9B model specification, and CI workflow and CUDA installation updates for dedicated four-GPU EPD testing. ChangesTokenSpeed EPD multimodal execution
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant E2E_Test
participant Gateway
participant EncodeWorker
participant PrefillWorker
participant DecodeWorker
E2E_Test->>Gateway: Send multimodal chat completion
Gateway->>EncodeWorker: Dispatch image encode request
EncodeWorker->>PrefillWorker: Transfer encoded data
PrefillWorker->>DecodeWorker: Start token generation
DecodeWorker-->>E2E_Test: Return completion response
Possibly related PRs
Suggested labels: Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
github-actions
Bot
added
ci
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces EPD (Encode-Prefill-Decode) multimodal Chat Completions end-to-end tests and adds support for EPD disaggregation topologies. It includes updates to backend setup, gateway argument building, and worker management to support separate encode, prefill, and decode workers using TokenSpeed. Additionally, unit tests for the EPD command-builder logic are added, and the Qwen3.5-9B model is integrated. Feedback is provided regarding a potential security improvement: using tempfile.mkdtemp instead of a predictable path in the shared temporary directory to avoid permission conflicts or hijacking vulnerabilities.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| FIXTURES_DIR = Path(__file__).parent.parent / "fixtures" / "images" | ||
| DOG_IMAGE_PATH = FIXTURES_DIR / "dog.jpg" # Black labrador puppy (checked in) | ||
|
|
||
| _LOG_DIR = Path(tempfile.gettempdir()) / f"smg-e2e-epd-{os.getpid()}" |
There was a problem hiding this comment.
Using a predictable path in the shared temporary directory (tempfile.gettempdir()) can lead to permission conflicts or hijacking vulnerabilities on shared systems or multi-user CI environments. It is safer and more robust to use tempfile.mkdtemp to create a uniquely named, securely permissioned directory at import time.
| _LOG_DIR = Path(tempfile.gettempdir()) / f"smg-e2e-epd-{os.getpid()}" | |
| _LOG_DIR = Path(tempfile.mkdtemp(prefix="smg-e2e-epd-")) |
Signed-off-by: key4ng <rukeyang@gmail.com>
…kers Signed-off-by: key4ng <rukeyang@gmail.com>
Signed-off-by: key4ng <rukeyang@gmail.com>
Signed-off-by: key4ng <rukeyang@gmail.com>
Signed-off-by: key4ng <rukeyang@gmail.com>
Signed-off-by: key4ng <rukeyang@gmail.com>
Signed-off-by: key4ng <rukeyang@gmail.com>
…alse pass) Signed-off-by: key4ng <rukeyang@gmail.com>
Signed-off-by: key4ng <rukeyang@gmail.com>
…d fixes Signed-off-by: key4ng <rukeyang@gmail.com>
…pts reasoning_content Signed-off-by: key4ng <rukeyang@gmail.com>
57608c2 to
c8ab2e0
Compare
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 57608c2d02
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| e2e-4gpu-epd: | ||
| name: e2e-4gpu-epd (tokenspeed) | ||
| needs: [e2e-1gpu-chat] |
There was a problem hiding this comment.
Adding this workflow job does not make its result part of the aggregate finish check: I inspected the finish job in this same workflow, and its needs list and failure condition still omit e2e-4gpu-epd. If branch protection relies on finish, this new smoke can fail or still be running while finish reports success, so the EPD coverage added here would not actually block merges.
Useful? React with 👍 / 👎.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| if pf.bootstrap_port is not None: | ||
| args.append(str(pf.bootstrap_port)) | ||
| for dc in decode_workers: | ||
| args += ["--decode", dc.worker_url] |
There was a problem hiding this comment.
🟡 Nit: This line uses dc.worker_url while the encode/prefill loops above use en.base_url / pf.base_url. Since worker_url is just an alias for base_url (worker.py:62-64), this is functionally identical — but mixing both names in the same 10-line function reads like they might be different properties. Using dc.base_url here would make the intent clearer.
| args += ["--decode", dc.worker_url] | |
| args += ["--decode", dc.base_url] |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
e2e_test/infra/worker.py (1)
579-600: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winRelease reserved ports when worker startup fails.
worker.start()can fail before establishing a live process, butstop_workers()callsWorker.stop(), which returns early whenprocessisNoneor already exited. The newly reserveddist_init_port, bootstrap port, and service port therefore leak after failed starts.Move reservation cleanup outside the process-liveness early return, or explicitly release these ports in the startup exception path.
🤖 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/infra/worker.py` around lines 579 - 600, Ensure reserved service, bootstrap, and dist_init ports are released when Worker.start() fails before a live process exists. Update Worker.stop() to perform port cleanup before its process-liveness early return, or invoke equivalent cleanup from the startup exception path, while preserving normal process termination behavior.
🤖 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 @.github/workflows/e2e-gpu-job.yml:
- Around line 134-136: Update the “Download extra models” workflow step so
inputs.extra_models is passed through the step’s env configuration rather than
interpolated into run. In the shell command, safely split the environment value
into arguments before invoking scripts/ci_download_model.sh, preserving support
for multiple model names without allowing input text to become shell syntax.
In `@e2e_test/chat_completions/test_epd_multimodal.py`:
- Around line 120-124: Replace the broad substring assertions in the multimodal
image tests around the answer validation with normalized exact-answer checks and
disjoint expected classifications for each image. Require the pug case to return
“pug” and ensure color cases accept only their intended normalized color,
preventing generic or negated text from passing.
In `@grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py`:
- Around line 201-205: Make the acceptance event in
grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py lines 201-205 use
a logger and level preserved in worker logs. Update the assertion in
e2e_test/chat_completions/test_epd_multimodal.py lines 76-78 to match the
current request ID rather than counting generic router dispatch markers, so
delayed or unrelated requests cannot satisfy it.
---
Outside diff comments:
In `@e2e_test/infra/worker.py`:
- Around line 579-600: Ensure reserved service, bootstrap, and dist_init ports
are released when Worker.start() fails before a live process exists. Update
Worker.stop() to perform port cleanup before its process-liveness early return,
or invoke equivalent cleanup from the startup exception path, while preserving
normal process termination behavior.
🪄 Autofix (Beta)
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: ASSERTIVE
Plan: Pro
Run ID: 532489ff-9008-4809-b1c0-3b0ef7daa1ea
📒 Files selected for processing (12)
.github/workflows/e2e-gpu-job.yml.github/workflows/pr-test-rust.ymle2e_test/chat_completions/test_epd_multimodal.pye2e_test/fixtures/setup_backend.pye2e_test/infra/constants.pye2e_test/infra/gateway.pye2e_test/infra/model_specs.pye2e_test/infra/test_epd_cmd_builders.pye2e_test/infra/worker.pygrpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.pyscripts/ci_download_model.shscripts/ci_install_tokenspeed.sh
| # (2) The answer is correct about the image — only possible if the encoder's | ||
| # embeddings reached prefill+decode (EPD-only gateway; no fallback path). | ||
| assert any(k in text.lower() for k in keywords), ( | ||
| f"expected one of {keywords} in the answer, got: {text!r}" | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Require answers that distinguish each image.
Substring matching allows false passes: both animal cases accept "dog"/"puppy", and generic or negated text containing a color also passes. The test can therefore succeed without proving that different images produced different classifications.
Use normalized exact answers for colors and disjoint expectations—such as requiring "pug" for the pug image—or use images from different species.
Also applies to: 136-152
🤖 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/chat_completions/test_epd_multimodal.py` around lines 120 - 124,
Replace the broad substring assertions in the multimodal image tests around the
answer validation with normalized exact-answer checks and disjoint expected
classifications for each image. Require the pug case to return “pug” and ensure
color cases accept only their intended normalized color, preventing generic or
negated text from passing.
| logger.info( | ||
| "EPD encode: accepted request_id=%s room=%s", | ||
| request.request_id, | ||
| bootstrap_room, | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Make the acceptance assertion request-scoped and observable.
The new acceptance event is logged at INFO, while the E2E test states that TokenSpeed logging suppresses it and instead counts an unscoped router marker. An unrelated or delayed dispatch can therefore satisfy the assertion.
grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py#L201-L205: emit the acceptance event through a logger/level preserved in worker logs.e2e_test/chat_completions/test_epd_multimodal.py#L76-L78: correlate the assertion with the current request ID instead of counting generic dispatch lines.
📍 Affects 2 files
grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py#L201-L205(this comment)e2e_test/chat_completions/test_epd_multimodal.py#L76-L78
🤖 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 `@grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py` around lines
201 - 205, Make the acceptance event in
grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py lines 201-205 use
a logger and level preserved in worker logs. Update the assertion in
e2e_test/chat_completions/test_epd_multimodal.py lines 76-78 to match the
current request ID rather than counting generic router dispatch markers, so
delayed or unrelated requests cannot satisfy it.
…args, mkdtemp log dir Signed-off-by: key4ng <rukeyang@gmail.com>
|
Thanks for the reviews! Fixed
Intentionally deferred
Kept as-is (with rationale)
Note: the worker-side |
|
You have reached your Codex usage limits. You can see your limits in the Codex usage dashboard. |
|
Great to hear that the |
|
Tip For best results, initiate chat on the files or code changes.
The fixes all sound solid:
On the deferred items:
Good call also documenting that the router-side (ノ◕‿◕)ノ*:・゚✧ ✏️ Learnings added
|
Signed-off-by: key4ng <rukeyang@gmail.com>
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
e2e_test/infra/gateway.py (1)
136-148: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winReject incomplete EPD/PD topologies before launch.
encode_workers=[]or a missingprefill_workers/decode_workersside is treated as a valid mode, converted to empty lists, and started with only a health check. That can leave the gateway reporting healthy until the first request. Require all worker lists for EPD, and both lists for PD, or fail fast on partial configs.🤖 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/infra/gateway.py` around lines 136 - 148, Validate the worker-list contents before mode selection in the gateway initialization flow: EPD requires non-empty encode_workers, prefill_workers, and decode_workers, while PD requires non-empty prefill_workers and decode_workers. Reject partial or empty configurations with ValueError before launch, and ensure only complete topologies contribute to modes_specified.
🤖 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 @.github/workflows/e2e-gpu-job.yml:
- Around line 134-141: Update the “Download extra models” step to parse
EXTRA_MODELS into a shell array using space-separated words, then invoke
scripts/ci_download_model.sh with the array elements quoted via "${models[@]}".
Preserve the existing environment-variable-based input handling and conditional
execution while preventing pathname expansion.
In `@e2e_test/chat_completions/test_epd_multimodal.py`:
- Line 48: Replace the module-level tempfile.mkdtemp call used for _LOG_DIR with
pytest-managed temporary storage, or add explicit teardown that removes the
created directory after the E2E tests complete. Preserve _LOG_DIR’s existing
Path-based usage while ensuring every smg-e2e-epd-* directory is cleaned up.
---
Outside diff comments:
In `@e2e_test/infra/gateway.py`:
- Around line 136-148: Validate the worker-list contents before mode selection
in the gateway initialization flow: EPD requires non-empty encode_workers,
prefill_workers, and decode_workers, while PD requires non-empty prefill_workers
and decode_workers. Reject partial or empty configurations with ValueError
before launch, and ensure only complete topologies contribute to
modes_specified.
🪄 Autofix (Beta)
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: ASSERTIVE
Plan: Pro
Run ID: 85443d0a-39e2-45e6-a5b5-5a1a370d70ce
📒 Files selected for processing (3)
.github/workflows/e2e-gpu-job.ymle2e_test/chat_completions/test_epd_multimodal.pye2e_test/infra/gateway.py
2 participants
Description
Problem
TokenSpeed EPD (encode-prefill-decode) disaggregation is live in the gateway (
--epd-disaggregation,--encode,RoutingMode::EncodePrefillDecode) and the TokenSpeed gRPC encode servicer, but there is no automated coverage — nothing exercises the encode→prefill→decode multimodal path on change.Solution
Extend the e2e harness with EPD support and add a
4-gpu-h100smoke that runs a TokenSpeed EPD vision model across four worker-count topologies (1e1p1d,1e2p1d,2e1p1d,1e1p2d, all tp=1). The test asserts real EPD participation (the encode worker's own per-request accept log), not just that a plausible answer returned — a single-worker fallback fails it.Changes
WorkerType.ENCODE+ a per-requestEPD encode: acceptedmarker in the TokenSpeed encode servicer.worker.py: TokenSpeed disaggregation launch flags (--disaggregation-mode {encode|prefill|decode}, bootstrap port for encode+prefill,--disaggregation-transfer-backend mooncake, unique--dist-init-addrper worker, prefix-caching/enforce-eager per role,--skip-server-warmup) + NVLink/warmup env.gateway.py:build_epd_mode_args()+ an--epd-disaggregationlaunch mode (--encode/--prefill/--decode,--encode-policy consistent_hashing,--multimodal-tensor-transport inline).setup_backend.py: anepd_grpcfixture mode (_setup_epd) launching N encode + N prefill + N decode workers with sequential GPU offsets.model_specs.py:Qwen/Qwen3.5-9B(multimodal, tp=1, FA3,skip_tier_download).test_epd_multimodal.py: the smoke test over the four topologies (request-scoped encode-acceptance assertion).e2e-4gpu-epdjob inpr-test-rust.yml;skip_tier_download+extra_modelsmodel provisioning ine2e-gpu-job.yml/ci_download_model.sh.Note:
build_epd_mode_argsintentionally leaves prefill/decode routing at the gateway defaults and only sets--encode-policy.Test Plan
pytest e2e_test/infra/test_epd_cmd_builders.py -v— 10 tests cover the disaggregation flags, EPD gateway args, and the model spec.4-gpu-h100runner viae2e-4gpu-epdacross all four topologies. First-run watch-items: Mooncake transport under non-privileged CI (de-risk1e1p1dfirst) and encode-marker visibility under the worker--log-level warning.Checklist
cargo +nightly fmtpasses (N/A — no Rust changes beyond a one-line log)cargo clippy --all-targets --all-features -- -D warningspasses (N/A)🤖 Generated with Claude Code
Summary by CodeRabbit
extra_modelsinput.