feat(vllm): OpenAI embeddings dimensions truncation (DIS-2093 vLLM side)

[tzulingk]

Overview:

Adds OpenAI /v1/embeddings dimensions (Matryoshka truncation) support to the vLLM embedding worker. Mirror of PR #9722 which did the same for SGLang.

Scope: Only changes the vLLM EmbeddingWorkerHandler from #9713. Once #9713 and #9722 both land, this PR closes out the vLLM half of DIS-2093 (the dimensions half).

Details:

In EmbeddingWorkerHandler.generate, the handler now reads request.get("dimensions") once before the per-input loop, validates it (must be a positive int), then per embedding asserts dimensions <= len(embedding) and slices to embedding[:dimensions]. Out-of-range values raise ValueError which the Rust frontend surfaces as HTTP 400.

Test: embedding_agg config in tests/serve/test_vllm.py gains a fourth payload — dimensions=128 against Qwen3-Embedding-0.6B (hidden dim 1024). Built inline using EmbeddingPayload(...) directly because the embedding_payload(..., extra_body=...) helper from PR #9722 isn't on this branch's base.

Not included — encoding_format=base64

Same as the SGLang half of DIS-2093: the Rust frontend's response type is Vec<f32> (inherited from the upstream async_openai::types::embeddings::* re-export) and rejects base64 strings. End-to-end support requires owning the embedding response type in lib/protocols/. Tracked in DIS-2099.

Stacking

Based on feat/vllm-embedding-worker-mvp (PR #9713). Will rebase to main once #9713 lands.

Reviewer focus

1. components/src/dynamo/vllm/handlers.py::EmbeddingWorkerHandler.generate — the new validation + slicing.
2. tests/serve/test_vllm.py — one new dimensions=128 payload.

Refs DIS-2093.

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features

  • Added embedding worker mode to serve embedding models with new --embedding-worker configuration flag
  • Implemented OpenAI-compatible /v1/embeddings endpoint supporting text, token, and batched input formats
  • Added dimension truncation support for embeddings
  • New aggregated embedding model launch script provided

• Tests

  • Added comprehensive test coverage for embedding worker configuration, validation, and API behavior

[coderabbitai]

Walkthrough

This PR adds support for dedicated embedding worker mode in the Dynamo vLLM wrapper, enabling OpenAI-compatible embedding serving in aggregated disaggregation mode. It introduces CLI configuration with strict validation, a specialized handler for embedding requests, factory-level routing, and end-to-end launch and integration testing.

Changes

Embedding Worker Feature

Layer / File(s)	Summary
Config flag and validation rules components/src/dynamo/vllm/backend_args.py, components/src/dynamo/vllm/tests/test_backend_args.py, components/src/dynamo/vllm/tests/test_vllm_unit.py	Added --embedding-worker CLI flag and embedding_worker config field to DynamoVllmConfig. Validation enforces embedding worker is only allowed in agg disaggregation mode, mutually exclusive with multimodal roles and enable_multimodal, and incompatible with benchmark_mode.
Embedding input classification components/src/dynamo/vllm/tests/test_vllm_worker_handler.py	Test suite validates _classify_embedding_input for OpenAI-spec embedding input shapes: single string, list of strings, token-id lists, and batched prompts, with strict rejection of mixed-type and invalid inputs.
Worker factory embedding routing components/src/dynamo/vllm/worker_factory.py, components/src/dynamo/vllm/tests/test_vllm_worker_factory.py	Factory imports EmbeddingWorkerHandler, checks config.embedding_worker, and routes to _create_embedding_worker coroutine to initialize vLLM engine, construct handler, serve endpoint, and register embedding model with guaranteed cleanup.
Launch script and integration test examples/backends/vllm/launch/agg_embed.sh, tests/serve/test_vllm.py	agg_embed.sh launches aggregated embedding serving with frontend and vLLM worker processes; integration test adds embedding_agg configuration covering default, single, batch, and dimension-truncation embedding payload variants.

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 40.74% 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 clearly and specifically identifies the main change: adding OpenAI embeddings dimensions truncation support to vLLM, with precise issue reference (DIS-2093).
Description check	✅ Passed	The PR description follows the template structure with clear Overview, Details, and Related Issues sections; it comprehensively explains the change, scope, implementation, limitations, and stacking context.
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.

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

✨ Finishing Touches

⚔️ Resolve merge conflicts

• Resolve merge conflict in branch feat/vllm-embedding-dimensions

---

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]

🧹 Nitpick comments (1)

components/src/dynamo/vllm/handlers.py (1)

2887-2889: ⚡ Quick win

Move PoolingParams to module scope.

This function-local import violates the repo Python rule and adds import work on every embeddings request. If the lazy-load concern is real, the embedding handler should live in its own module instead of importing inside generate().

♻️ Proposed fix

+from vllm import PoolingParams
 from vllm.config import ModelConfig, VllmConfig
 from vllm.inputs import EmbedsPrompt, TextPrompt, TokensPrompt
@@
-        # Lazy import to avoid pulling PoolingParams into handlers.py at module
-        # load time for non-embedding workers.
-        from vllm import PoolingParams
-
         model_name = request.get("model") or self.config.served_model_name or ""

As per coding guidelines, "Keep all imports at module scope (flag any import inside functions/classes)."

🤖 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 `@components/src/dynamo/vllm/handlers.py` around lines 2887 - 2889, The local
import of PoolingParams inside generate() should be moved to module scope to
comply with repo rules and avoid per-request import overhead: remove the
in-function "from vllm import PoolingParams" and add "from vllm import
PoolingParams" at the top of this handlers.py module; if the original lazy-load
rationale is still required, extract the embedding handler into a dedicated
module (e.g., vllm_embedding_handler) and place the PoolingParams import at that
module's top so generate() no longer contains inline imports.

🤖 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.

Nitpick comments:
In `@components/src/dynamo/vllm/handlers.py`:
- Around line 2887-2889: The local import of PoolingParams inside generate()
should be moved to module scope to comply with repo rules and avoid per-request
import overhead: remove the in-function "from vllm import PoolingParams" and add
"from vllm import PoolingParams" at the top of this handlers.py module; if the
original lazy-load rationale is still required, extract the embedding handler
into a dedicated module (e.g., vllm_embedding_handler) and place the
PoolingParams import at that module's top so generate() no longer contains
inline imports.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 3e14961f-ffff-475d-bc9e-db07488f7023

📥 Commits

Reviewing files that changed from the base of the PR and between 066a7a5 and 2a231ed.

📒 Files selected for processing (9)

• components/src/dynamo/vllm/backend_args.py
• components/src/dynamo/vllm/handlers.py
• components/src/dynamo/vllm/tests/test_backend_args.py
• components/src/dynamo/vllm/tests/test_vllm_unit.py
• components/src/dynamo/vllm/tests/test_vllm_worker_factory.py
• components/src/dynamo/vllm/tests/test_vllm_worker_handler.py
• components/src/dynamo/vllm/worker_factory.py
• examples/backends/vllm/launch/agg_embed.sh
• tests/serve/test_vllm.py