Repository navigation
fix(e2e): respect workers_config in vLLM/TRT-LLM gRPC multi-worker setup - #525
Conversation
_setup_grpc_backend() accepted workers_config but never read it — it always called model_pool.get_grpc_worker() which returns exactly one worker. When nightly benchmarks ran multi-worker tests (count=4), vLLM gRPC got 1 worker while HTTP got 4, causing a 2-24× performance gap that made protocol comparison results meaningless. What changed: - e2e_test/fixtures/setup_backend.py: rewrite _setup_grpc_backend() to read workers_config["count"] and, when count > 1, use get_workers_by_type() + launch_workers() to acquire/launch the requested number of gRPC workers — the same pattern used by _setup_local_backend() for SGLang. Why: - Nightly benchmark run showed HTTP winning every metric by 4-27% aggregate, with vLLM multi-worker gRPC TTFT p99 up to 24× worse. Investigation of gateway logs revealed vLLM multi-worker tests had 4 HTTP workers but only 1 gRPC worker behind the gateway. How: - Single-worker path (count=1) is unchanged — still uses get_grpc_worker() for backward compatibility. - Multi-worker path mirrors _setup_local_backend(): finds existing GRPC REGULAR workers, releases wrong-mode workers, launches missing ones via launch_workers() with WorkerIdentity, and collects all worker URLs for the gateway. - Cleanup paths iterate over all instances instead of a single one. - Both vLLM and TRT-LLM are fixed since both route through _setup_grpc_backend() at line 125. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
Summary of ChangesHello @slin1237, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request addresses a critical issue in the e2e testing infrastructure where vLLM and TRT-LLM gRPC multi-worker configurations were not being correctly provisioned. Previously, only a single gRPC worker was launched regardless of the Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
📝 WalkthroughWalkthroughRefactors test fixture backend setup to add multi-worker support across gRPC, local, and PD backends: discovers/releases/reuses/launches worker instances, aggregates worker URLs and model_path, and updates cleanup and logging to handle multiple acquired workers. Changes
Sequence Diagram(s)sequenceDiagram
participant Setup as Setup Function
participant Pool as Model Pool
participant Workers as Worker Instances
participant Gateway as Gateway
rect rgba(100, 150, 200, 0.5)
note over Setup,Workers: Multi-worker setup (num_workers > 1)
Setup->>Pool: get_workers_by_type(model_id)
Pool-->>Workers: list discovered instances
Setup->>Workers: filter by ConnectionMode / release non-matching
alt enough existing workers
Setup->>Workers: reuse required instances
else need more workers
Setup->>Workers: launch missing workers
Workers-->>Pool: register/acquire instances
end
Setup->>Setup: collect worker_urls & model_path
Setup->>Gateway: start(prefill_workers, decode_workers, worker_urls, model_path)
Gateway-->>Setup: ready
end
rect rgba(150, 100, 200, 0.5)
note over Setup,Workers: Cleanup on failure
Setup->>Workers: release all acquired instances
Workers-->>Setup: released
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Comment |
![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
The pull request correctly addresses the issue where _setup_grpc_backend ignored the worker count for vLLM and TRT-LLM gRPC setups. By adopting the multi-worker pattern from _setup_local_backend, it now supports launching and managing multiple gRPC workers. I've suggested improvements to ensure robust resource cleanup in case of failures, more precise validation of the launched worker count, and a review of the timeout mechanism for launching multiple workers. Additionally, I noted that this logic is now duplicated with _setup_local_backend, which presents a refactoring opportunity.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@e2e_test/fixtures/setup_backend.py`:
- Around line 509-535: The else branch leaks already-acquired existing_grpc if
model_pool.launch_workers raises and silently proceeds when launch_workers
returns an empty list; fix by (1) building an acquired_workers list (include
existing_grpc) before calling model_pool.launch_workers so cleanup will release
them on exceptions (same pattern as acquired_prefills/decodes), (2) after
calling model_pool.launch_workers, immediately check if new_instances is empty
and call pytest.fail with a clear message if so, and (3) only set instances =
existing_grpc + new_instances after the success checks so the later exception
handler will see the correct acquired list; refer to existing_grpc,
new_instances, instances, and model_pool.launch_workers to locate the changes.
- Around line 579-584: The log currently prints the configured worker target via
num_workers which can be misleading; update the log call that formats "Setup %s
gRPC backend: model=%s, workers=%d, gateway=%s, policy=%s" to pass
len(instances) (the actual allocated worker count) instead of num_workers so the
message reflects real allocation (keep the other fields runtime_label, model_id,
gateway.base_url, gateway_config["policy"] unchanged).
…worker Address PR review feedback for two bugs: 1. Resource leak: existing_grpc workers (already acquired by get_workers_by_type) were not tracked in `instances` until after launch_workers completed. If launch_workers raised, those workers were never released — blocking GPU allocation for subsequent tests. Fix: assign `instances = list(existing_grpc)` immediately and use `instances.append()` for new workers. 2. Silent under-provisioning: when launch_workers returned [] due to insufficient GPUs, `instances = existing_grpc + []` could be non-empty (e.g. 1 of 4 needed), passing the `if not instances` check and running the test with fewer workers than configured. Fix: fail explicitly when launch_workers returns [] and check `len(instances) < num_workers` instead of `not instances`. Also fix the setup log to report len(instances) (actual) instead of num_workers (configured target). Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
e2e_test/fixtures/setup_backend.py (2)
704-711:⚠️ Potential issue | 🟡 MinorLog still reports configured
num_workers, not actuallen(instances).Line 708 logs
num_workers(the configured target). The identical issue was fixed in_setup_grpc_backend(line 594 now useslen(instances)). Apply the same fix here for consistency and observability.🔧 Proposed fix
logger.info( "Setup %s backend: model=%s, workers=%d, gateway=%s, policy=%s", backend_name, model_id, - num_workers, + len(instances), gateway.base_url, gateway_config["policy"], )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@e2e_test/fixtures/setup_backend.py` around lines 704 - 711, The logger call uses the configured num_workers rather than the actual launched instance count; update the logger.info invocation that currently prints backend_name, model_id, num_workers, gateway.base_url, gateway_config["policy"] to instead pass len(instances) for the workers value (mirror the fix made in _setup_grpc_backend), so the message reports the actual number of launched instances.
640-658:⚠️ Potential issue | 🟠 Major
_setup_local_backendelse branch has the same two bugs that were just fixed in_setup_grpc_backend— resource leak + silent under-provisioning.Bug 1 —
existing_for_modeworkers leak whenlaunch_workersraises.
instancesis[]whenlaunch_workersexecutes (line 651). If it raises, theexceptat line 667 iterates over an empty list and all workers already acquired viaget_workers_by_typeare silently leaked.Bug 2 — Partial allocation silently proceeds when
launch_workersreturns[].
There is noif not new_instances: pytest.fail(...)guard. If GPU capacity is insufficient,new_instances = [], soinstances = existing_for_mode + []. Whenexisting_for_modeis non-empty (e.g., 1 of 3 needed workers already exists),if not instances:at line 657 isFalseand the test runs under-provisioned — exactly the bug this PR was written to fix in gRPC.Apply the same three-part fix used in
_setup_grpc_backend(lines 502, 533–537, 543–547):🐛 Proposed fix mirroring _setup_grpc_backend pattern
else: missing = num_workers - len(existing_for_mode) workers_to_launch = [ WorkerIdentity( model_id, connection_mode, WorkerType.REGULAR, len(existing_for_mode) + i, ) for i in range(missing) ] + instances = list(existing_for_mode) # track before launch so cleanup fires on failure new_instances = model_pool.launch_workers(workers_to_launch, startup_timeout=300) + if not new_instances: + pytest.fail( + f"Failed to launch {missing} workers for {model_id}: " + f"GPU allocation failed" + ) # Acquire newly launched instances for inst in new_instances: inst.acquire() - instances = existing_for_mode + new_instances + instances.append(inst) - if not instances: - pytest.fail(f"Failed to get {num_workers} workers for {model_id}") + if len(instances) < num_workers: + pytest.fail( + f"Only got {len(instances)}/{num_workers} workers for {model_id}" + )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@e2e_test/fixtures/setup_backend.py` around lines 640 - 658, In _setup_local_backend, fix two issues: (1) avoid leaking already-acquired workers in existing_for_mode if model_pool.launch_workers raises by wrapping the launch call in try/except and releasing (calling .release() or the same cleanup used elsewhere) all entries in existing_for_mode on exception before re-raising or failing; (2) prevent silent under-provisioning by checking new_instances after the launch and calling pytest.fail (with a clear message) if new_instances is empty (mirroring the _setup_grpc_backend pattern); ensure you still acquire newly launched instances (inst.acquire()) and then set instances = existing_for_mode + new_instances and finally assert/pytest.fail if instances is empty. Reference symbols: _setup_local_backend, existing_for_mode, model_pool.launch_workers, new_instances, inst.acquire(), instances, pytest.fail.
♻️ Duplicate comments (1)
e2e_test/fixtures/setup_backend.py (1)
492-608:_setup_grpc_backendmulti-worker implementation looks correct — fixes from past review are properly applied.All three issues flagged in the previous review cycle are confirmed resolved:
- Resource leak:
instances = list(existing_grpc)at line 502 pre-seeds cleanup tracking beforelaunch_workers, matching the_setup_pd_backend_commonpattern.- Silent under-provisioning:
if not new_instances: pytest.fail(...)at lines 533–537 explicitly fails fast; the guard at line 543 useslen(instances) < num_workersinstead of the weakerif not instances:.- Log accuracy: line 594 now emits
len(instances)(actual acquired) instead ofnum_workers(configured target).🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@e2e_test/fixtures/setup_backend.py` around lines 492 - 608, The implementation of _setup_grpc_backend looks good and the previous issues are resolved; remove the stray duplicate review tag by deleting the redundant "[duplicate_comment]" marker from the PR review/comment text so the approval is unambiguous (no code changes required to functions like _setup_grpc_backend, instances tracking, or Gateway usage).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@e2e_test/fixtures/setup_backend.py`:
- Around line 704-711: The logger call uses the configured num_workers rather
than the actual launched instance count; update the logger.info invocation that
currently prints backend_name, model_id, num_workers, gateway.base_url,
gateway_config["policy"] to instead pass len(instances) for the workers value
(mirror the fix made in _setup_grpc_backend), so the message reports the actual
number of launched instances.
- Around line 640-658: In _setup_local_backend, fix two issues: (1) avoid
leaking already-acquired workers in existing_for_mode if
model_pool.launch_workers raises by wrapping the launch call in try/except and
releasing (calling .release() or the same cleanup used elsewhere) all entries in
existing_for_mode on exception before re-raising or failing; (2) prevent silent
under-provisioning by checking new_instances after the launch and calling
pytest.fail (with a clear message) if new_instances is empty (mirroring the
_setup_grpc_backend pattern); ensure you still acquire newly launched instances
(inst.acquire()) and then set instances = existing_for_mode + new_instances and
finally assert/pytest.fail if instances is empty. Reference symbols:
_setup_local_backend, existing_for_mode, model_pool.launch_workers,
new_instances, inst.acquire(), instances, pytest.fail.
---
Duplicate comments:
In `@e2e_test/fixtures/setup_backend.py`:
- Around line 492-608: The implementation of _setup_grpc_backend looks good and
the previous issues are resolved; remove the stray duplicate review tag by
deleting the redundant "[duplicate_comment]" marker from the PR review/comment
text so the approval is unambiguous (no code changes required to functions like
_setup_grpc_backend, instances tracking, or Gateway usage).
1 participant
Summary
_setup_grpc_backend()to actually readworkers_config["count"]and launch the requested number of gRPC workers for vLLM and TRT-LLM multi-worker testsWhat changed
e2e_test/fixtures/setup_backend.py: Rewrote_setup_grpc_backend()to support multi-worker mode using the same pattern as_setup_local_backend()—get_workers_by_type()to find existing gRPC REGULAR workers,launch_workers()withWorkerIdentityfor missing ones, proper cleanup of all instancesWhy
Nightly benchmark run showed HTTP winning every metric by 4-27% aggregate, with vLLM multi-worker gRPC TTFT p99 up to 24× worse than HTTP. Investigation of gateway logs revealed the root cause: vLLM multi-worker tests had 4 HTTP workers but only 1 gRPC worker behind the gateway.
_setup_grpc_backend()acceptedworkers_configas a parameter but never read it — it always calledmodel_pool.get_grpc_worker()which returns exactly one worker. SGLang was unaffected because SGLang gRPC falls through to_setup_local_backend()which correctly handles multi-worker.How
count > 1: usesget_workers_by_type(model_id, WorkerType.REGULAR), filters byConnectionMode.GRPC, launches missing workers vialaunch_workers(), collects all worker URLs for the gatewaycount == 1: preserves existingget_grpc_worker()path (no behavior change)_setup_grpc_backend()(line 125:if is_vllm() or is_trtllm())Test plan
_setup_local_backend)Summary by CodeRabbit