test(vllm): multi-worker embedding routing tests (DIS-2098)

[tzulingk]

Overview:

Adds end-to-end + unit coverage for the two routing properties the embedding worker pool exposes today: load-balancing across multiple workers serving the same model, and name-keyed dispatch across workers serving different models. Both paths run through the same model_manager → select_worker_set_with() machinery that already serves chat traffic; the work here is explicit embedding-side exercise of that machinery against /v1/embeddings.

Linear: https://linear.app/nvidia/issue/DIS-2098

Architecture

The frontend dispatches each /v1/embeddings request through one lookup:

POST /v1/embeddings  {"model": "M", "input": ...}
        │
        ▼
model_manager.get_embeddings_engine(M)   # name-keyed DashMap lookup
        │
        ▼
select_worker_set_with(eligible)         # weighted-random across matching workers
        │
        ▼
Worker on its own GPU / DYN_SYSTEM_PORT / endpoint path

A worker registers itself with the runtime under --endpoint <namespace>.<component>.<endpoint>. The runtime indexes the (worker, model) pair by model name, so the frontend never needs to know endpoint paths — it asks "who serves model M?" and gets back the set of workers that registered M, then picks one.

Two consequences shape the test setup:

• One model per endpoint path. The runtime rejects a second worker that tries to register a different model on the same namespace.component.endpoint. So the multi-worker launch script gives each worker its own unique --endpoint (here: dynamo.embed-worker-1.generate and dynamo.embed-worker-2.generate). They land in the same model registry from the frontend's perspective, but the runtime treats them as independent worker instances.

• Verification is structural, not behavioral. Each worker exposes /metrics on its own DYN_SYSTEM_PORT. The test scrapes dynamo_component_requests_total per worker — that counter increments on the worker that actually served the request, so the dispatch decision is observable from outside the frontend.

What's checked

1. Same-model load balancing — two workers register Qwen3-Embedding-0.6B (under different endpoint paths). The frontend's worker-set lookup returns both; select_worker_set_with() weighted-randomly picks one per request. The test sends a 20-request burst and asserts both workers' dynamo_component_requests_total counters moved.

2. Multi-model dispatch — worker A registers Qwen3-Embedding-0.6B, worker B registers BAAI/bge-small-en-v1.5. The frontend's name-keyed get_embeddings_engine(model) lookup must return only the matching worker. The test sends a burst with model="Qwen3...", snapshots per-worker counters, sends another burst with model="BAAI...", snapshots again. Burst A must increment worker A's counter only; burst B must increment worker B's counter only.

Why per-burst deltas, not absolute counts

The multi-model test sends two bursts in sequence:

Burst A:  10 requests with model="Qwen3..."   → should only reach worker A
Burst B:  10 requests with model="BAAI..."    → should only reach worker B

Each worker's dynamo_component_requests_total counter increments on every request it serves. The interesting question to verify is: did burst B traffic stay out of worker A?

A naive check on absolute totals cannot see a dispatch leak. Suppose dispatch were broken and burst B leaked 3 requests onto worker A:

                       Worker A     Worker B
After burst A          10           0
After burst B          13           10        ← B got its 10, A got +3 leak

An "is worker B's counter > 0?" assertion still passes (10 > 0), and an "is worker A's counter == 0?" assertion fails for the wrong reason — worker A is at 13 because of burst A, not because of the burst B leak. State from burst A dominates anything you can read from absolute totals.

The fix is to snapshot each worker's counter at the start of each burst and assert on the delta through the end:

Burst B start:   snapshot {A: 10, B: 0}
Burst B end:     read     {A: 13, B: 10}
Delta:                    {A: +3, B: +10}    ← leak now visible

Now "during burst B, only worker B's delta is positive" actually has bite: delta[A] = 3 is a real signal of dispatch breakage, not a remnant of burst A. "Wrong-model traffic stayed out of worker A during burst B" is a claim about a specific time window, so it needs two snapshots (start and end) to be testable at all.

Why indices, not absolute ports

The harness assigns dynamic ports to each test run, so we can't hard-code DYN_SYSTEM_PORT1 = 8081 in expectations. The payload accepts expected_worker_indices_with_delta: set[int] where 0 = first worker in system_ports list, 1 = second. Port mapping (DefaultPort.SYSTEM{1,2} → actual dynamic port) happens inside the harness before the payload runs.

Pieces

• examples/backends/vllm/launch/agg_embed_multiworker.sh — 1 frontend + 2 embedding workers (one per GPU). Each worker gets its own DYN_SYSTEM_PORT and its own unique --endpoint <namespace>.<component>.<endpoint> so registration doesn't collide. Pooler config mirrors agg_embed.sh (--runner pooling, MEAN pooler, fp32).
• tests/utils/payloads.py — new EmbeddingMultiWorkerDispatchPayload. Takes a list of system ports + a set of worker indices that should observe non-zero delta. Snapshots /metrics on the first request of a burst, deltas against the snapshot on the last.
• tests/serve/test_vllm.py — two pytest.mark.gpu_2 / pre_merge test functions: test_embedding_multi_worker_same_model_load_balance, test_embedding_multi_worker_multi_model_dispatch.
• tests/utils/test_embedding_dispatch_payload.py — six unit tests for the baseline+delta logic. Monkeypatches requests.get with a canned per-port /metrics sequence so the payload runs without a GPU.

Where should the reviewer start?

1. tests/utils/payloads.py EmbeddingMultiWorkerDispatchPayload — the diff semantics + index-based assertion are the load-bearing pieces of the test surface.
2. tests/utils/test_embedding_dispatch_payload.py — six small tests that exercise the payload in isolation, no GPU.
3. examples/backends/vllm/launch/agg_embed_multiworker.sh — confirm the per-worker --endpoint + DYN_SYSTEM_PORT wiring matches what the harness expects.
4. tests/serve/test_vllm.py — the two new test functions, modeled on test_lora_aggregated_router which already uses the 2-system-ports + 2-GPU pattern.

Summary by CodeRabbit

• New Features

  • Added support for multi-worker embedding serving with automatic load balancing across multiple GPU workers.
  • Enabled flexible assignment of embedding models to specific workers for optimized resource utilization.

• Tests

  • Added comprehensive end-to-end test coverage for multi-worker embedding deployments, validating load distribution and model-based worker routing.

[coderabbitai]

Walkthrough

This PR adds complete multi-worker embedding infrastructure for vLLM, including a launcher script that orchestrates a routing frontend and two GPU-pinned workers, a new test payload that validates multi-worker dispatch via Prometheus metrics, unit tests for the payload logic, and end-to-end tests for load balancing and model-name-based routing.

Changes

Multi-Worker Embedding Test Infrastructure

Layer / File(s)	Summary
Embedding dispatch payload contract tests/utils/payloads.py	EmbeddingMultiWorkerDispatchPayload captures Prometheus baseline counters on first request and validates final-request deltas match expected worker indices, with optional minimum total-delta enforcement.
Unit tests for dispatch payload tests/utils/test_embedding_dispatch_payload.py	Mocked Prometheus metrics server and fixture validate payload baseline capture, delta computation, per-worker routing assertion, multi-model dispatch accuracy, and bounds enforcement across six test cases.
Multi-worker launcher script examples/backends/vllm/launch/agg_embed_multiworker.sh	Bash script accepts two models, starts routing frontend, spawns two vLLM embedding workers on separate GPUs with unique system metric ports, common vLLM config, and exits on first worker failure.
End-to-end tests for multi-worker embedding tests/serve/test_vllm.py	E2e tests deploy the launcher with same or different models and validate routing: same-model test confirms load-balance delta distribution; multi-model test confirms model-name dispatch and MAX_MODEL_LEN configuration.

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 64.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately summarizes the main change: adding multi-worker embedding routing tests for vLLM, covering load-balancing and name-keyed dispatch scenarios.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description check	✅ Passed	The PR description is comprehensive, well-structured, and addresses all required template sections with detailed explanations.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🤖 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 `@examples/backends/vllm/launch/agg_embed_multiworker.sh`:
- Around line 54-56: The script uses SYSTEM_PORT1="${DYN_SYSTEM_PORT1:-8081}"
which ignores a caller-supplied DYN_SYSTEM_PORT; change the assignment for
SYSTEM_PORT1 to prefer DYN_SYSTEM_PORT and fall back to DYN_SYSTEM_PORT1 and
then 8081 so callers can set a single DYN_SYSTEM_PORT for worker-1 while
preserving existing DYN_SYSTEM_PORT1 behavior; leave SYSTEM_PORT2 as-is
(DYN_SYSTEM_PORT2) so multi-worker distinct ports still work.

🪄 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: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: fc557b45-1e6f-4262-8a2a-04d75baec0a5

📥 Commits

Reviewing files that changed from the base of the PR and between 6e37679 and 5fc95bf.

📒 Files selected for processing (4)

• examples/backends/vllm/launch/agg_embed_multiworker.sh
• tests/serve/test_vllm.py
• tests/utils/payloads.py
• tests/utils/test_embedding_dispatch_payload.py

[biswapanda]

lgtm with a minor comment

[biswapanda]

we'd add a k8s based deploy test for a known embedding model + deterministic/greedy sampling

[biswapanda]

Let's add a k8s based end-to-end deploy test for a known embedding model + deterministic/greedy sampling.
It can be this same PR or a follow up PR.

[tzulingk]

@biswapanda https://linear.app/nvidia/issue/DIS-2138/e2e-validation-embedding-worker-stack-on-vllm-combined-prs-9765-9830 did the E2E k8 tests

[tzulingk]

@biswapanda thanks for the review.

The k8s e2e is actually already done — see DIS-2138, where I ran the 9-test matrix against a real Qwen3-Embedding-0.6B deployed on dynamo-aks-dev (arm64 GB200) and posted full results (8/9 pass; the 9th is a pre-existing Rust frontend HTTP-500-vs-400 mapping issue unrelated to this PR).

Taking your wording "this same PR or a follow up PR" — turning that one-shot manual run into a recurring CI job is real infra work (new K8s test-runner harness, namespace lifecycle, image-build coordination). I've filed DIS-2153 as a follow-up so this PR stays scoped to its title and we don't bundle CI infra here. Happy to discuss whether DIS-2153 should be done before merging this — but my reading is the in-process embedding_multi_worker_* tests already in this PR cover the routing + canary behavior in CI, and the K8s-deploy layer is orthogonal (operator + service discovery, both already exercised by the merged #9886 / #9830 / #9887 in chat-worker CI elsewhere).

[biswapanda]

@biswapanda thanks for the review.

The k8s e2e is actually already done — see DIS-2138, where I ran the 9-test matrix against a real Qwen3-Embedding-0.6B deployed on dynamo-aks-dev (arm64 GB200) and posted full results (8/9 pass; the 9th is a pre-existing Rust frontend HTTP-500-vs-400 mapping issue unrelated to this PR).

Taking your wording "this same PR or a follow up PR" — turning that one-shot manual run into a recurring CI job is real infra work (new K8s test-runner harness, namespace lifecycle, image-build coordination). I've filed DIS-2153 as a follow-up so this PR stays scoped to its title and we don't bundle CI infra here. Happy to discuss whether DIS-2153 should be done before merging this — but my reading is the in-process embedding_multi_worker_* tests already in this PR cover the routing + canary behavior in CI, and the K8s-deploy layer is orthogonal (operator + service discovery, both already exercised by the merged #9886 / #9830 / #9887 in chat-worker CI elsewhere).

Thanks @tzulingk. yes k8s tests are non-blocker and can be a follow up.