From 18fbbcb074acd05188f683be94512f85481296bb Mon Sep 17 00:00:00 2001 From: Tzu-Ling Date: Wed, 20 May 2026 10:22:52 -0500 Subject: [PATCH 1/2] feat(vllm): add aggregated text-embedding worker MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds an aggregated text-embedding worker shape to Dynamo's vLLM backend. Users can move an existing pooling-model vLLM deployment (e.g. Qwen3-Embedding-0.6B) onto Dynamo+vLLM without changing engine args. The Rust frontend's `/v1/embeddings` route and `ModelType::Embedding` plumbing have existed since the SGLang embedding backend shipped; this commit registers a vLLM worker on the same endpoint. Architecture ============ A new `--embedding-worker` flag in `backend_args.py` selects a different worker shape: pooling-mode `AsyncLLM` fronted by a new `EmbeddingWorkerHandler` that calls `AsyncLLM.encode(prompt, pooling_params, request_id)` and converts the `PoolingRequestOutput` to an OpenAI-shaped response. Dispatch happens in `WorkerFactory.create()` — the embedding branch is checked first, before the existing `disaggregation_mode` switch, because it's a different worker shape rather than a variant of decode. `EmbeddingWorkerHandler` is a standalone class that does NOT inherit `BaseWorkerHandler`. The base does generation-only init (media loaders, KV-block lookup, embedding cache manager) that would either fail or be meaningless on a pooling engine. Pooling inference is a single forward pass with no KV cache, no multimodal, no streamed decode — a separate class is clearer than overriding most of the base's behavior. Intentionally skipped on the embedding path ============================================ - KV-events publisher: no KV cache, nothing to publish. - Forward-pass-metrics relay: relays decode-phase ZMQ metrics; no decode here. - InstrumentedScheduler: hard-codes pooling_params=None (would silently disable the pooling pass) and emits decode-shaped metrics that don't apply. Only installed when --benchmark-mode is set, which is rejected for embedding workers via `_validate_embedding_worker_exclusivity()`. - P/D disaggregation: rejected at parse time. `--embedding-worker` combined with any non-`agg` `--disaggregation-mode` raises `ValueError`. Request handling ================ - Inputs accepted per the OpenAI /v1/embeddings spec: `str`, `list[str]`, `list[int]`, `list[list[int]]`. Mixed lists are rejected with a clear `TypeError`. Token-id forms are passed as `vllm.inputs.TokensPrompt` so the engine skips its own tokenizer. - Per-request `request_id` is derived from `context.id()` (matches the chat/completion paths) so concurrent embeddings never collide inside `AsyncLLM`. - Each encode call is wrapped in `_abort_monitor` so client cancellation or `shutdown_event` calls `engine_client.abort()` and propagates `EngineShutdown` rather than leaving GPU work running. - `VllmEngineMonitor` is wired in `__init__` so a dead pooling engine triggers `shutdown_event` instead of leaving the endpoint registered serving failures. Validation ========== `DynamoVllmConfig._validate_embedding_worker_exclusivity()` rejects `--embedding-worker` combined with: any non-`agg` disaggregation, multimodal flags, `--enable-multimodal`, or `--benchmark-mode`. Launch script ============= `examples/backends/vllm/launch/agg_embed.sh` defaults to `Qwen/Qwen3-Embedding-0.6B` with the standard pooling args: `--runner pooling --dtype float32 --pooler-config '{"pooling_type":"MEAN","use_activation":false}'`. Tests ===== - `test_vllm_unit.py::TestEmbeddingWorkerFlag`: parse-time flag acceptance + exclusion combinations. - `test_vllm_worker_factory.py::test_embedding_worker_takes_priority`: dispatch precedence over decode/prefill/encode paths. - `test_vllm_worker_handler.py::TestClassifyEmbeddingInput`: all four input shapes plus mixed-list / bool / empty rejection. - `test_backend_args.py::TestEmbeddingWorkerExclusivity`: covers the benchmark-mode and multimodal rejections. - `test_vllm.py` adds an `embedding_agg` config in the `pytest.mark.core` lane that exercises the launch script + a range of input shapes. Signed-off-by: Tzu-Ling --- components/src/dynamo/vllm/backend_args.py | 35 +++ components/src/dynamo/vllm/handlers.py | 288 ++++++++++++++++++ .../dynamo/vllm/tests/test_backend_args.py | 58 ++++ .../src/dynamo/vllm/tests/test_vllm_unit.py | 67 ++++ .../vllm/tests/test_vllm_worker_factory.py | 16 + .../vllm/tests/test_vllm_worker_handler.py | 60 ++++ components/src/dynamo/vllm/worker_factory.py | 99 ++++++ examples/backends/vllm/launch/agg_embed.sh | 91 ++++++ tests/serve/test_vllm.py | 47 +++ 9 files changed, 761 insertions(+) create mode 100755 examples/backends/vllm/launch/agg_embed.sh diff --git a/components/src/dynamo/vllm/backend_args.py b/components/src/dynamo/vllm/backend_args.py index b44d2adeef79..441e2e291c2f 100644 --- a/components/src/dynamo/vllm/backend_args.py +++ b/components/src/dynamo/vllm/backend_args.py @@ -129,6 +129,16 @@ def add_arguments(self, parser) -> None: choices=[m.value for m in EmbeddingTransferMode], ) + add_negatable_bool_argument( + g, + flag_name="--embedding-worker", + env_var="DYN_VLLM_EMBEDDING_WORKER", + default=False, + help="Run as a text-embedding worker. Engine must be started with " + "vLLM's --runner pooling. Skips KV-events, KV router registration, " + "and InstrumentedScheduler injection (none apply to pooling models).", + ) + # Headless mode for multi-node TP/PP add_negatable_bool_argument( g, @@ -259,6 +269,7 @@ class DynamoVllmConfig(ConfigBase): embedding_transfer_mode: Union[ str, EmbeddingTransferMode ] # resolved to enum in validate() + embedding_worker: bool = False # Headless mode for multi-node TP/PP headless: bool = False @@ -284,6 +295,7 @@ def validate(self) -> None: self._resolve_embedding_transfer_mode() self._validate_multimodal_role_exclusivity() self._validate_multimodal_requires_flag() + self._validate_embedding_worker_exclusivity() def _resolve_embedding_transfer_mode(self) -> None: """Resolve embedding_transfer_mode from string to enum.""" @@ -441,3 +453,26 @@ def _validate_multimodal_requires_flag(self) -> None: raise ValueError( "Use --enable-multimodal when enabling any multimodal component" ) + + def _validate_embedding_worker_exclusivity(self) -> None: + """Embedding worker is aggregated-only and exclusive of multimodal roles.""" + if not self.embedding_worker: + return + if self.disaggregation_mode != DisaggregationMode.AGGREGATED: + raise ValueError( + "--embedding-worker is only valid with --disaggregation-mode=agg " + f"(got {self.disaggregation_mode.value if isinstance(self.disaggregation_mode, DisaggregationMode) else self.disaggregation_mode}). " + "Pooling models do not have prefill/decode phases." + ) + if self._count_multimodal_roles() > 0 or self.enable_multimodal: + raise ValueError( + "--embedding-worker cannot be combined with multimodal flags." + ) + if self.benchmark_mode is not None: + raise ValueError( + "--embedding-worker cannot be combined with --benchmark-mode. " + "Benchmark mode injects InstrumentedScheduler, which is a " + "generation scheduler and not compatible with pooling engines. " + "Embedding workers do not run generation, so prefill/decode " + "benchmark sweeps are not meaningful." + ) diff --git a/components/src/dynamo/vllm/handlers.py b/components/src/dynamo/vllm/handlers.py index b5053d99c1d6..d925fd3fca18 100644 --- a/components/src/dynamo/vllm/handlers.py +++ b/components/src/dynamo/vllm/handlers.py @@ -2748,3 +2748,291 @@ def _build_embedding_params( # as request input. return build_qwen_embedding_params(multi_modal_data, self._qwen_grid_params) return None + + +class EmbeddingWorkerHandler: + """Standalone handler for OpenAI /v1/embeddings requests on vLLM. + + Does NOT inherit BaseWorkerHandler. The base class does generation-only + init (media loaders, KV-block lookup via get_dp_range_for_worker, embedding + cache manager) that would either fail or be meaningless on a pooling + engine. Embedding inference is a single forward pass with no KV cache, no + multimodal data, and no streamed decode. + """ + + def __init__( + self, + runtime, + engine: Any, + config: Config, + shutdown_event: Optional[asyncio.Event] = None, + ) -> None: + self.runtime = runtime + self.engine_client = engine + self.config = config + self.shutdown_event = shutdown_event + # Dead-engine detection: VllmEngineMonitor polls AsyncLLM and triggers + # shutdown_event + process exit on EngineDeadError. Without this, a + # crashed pooling engine leaves the endpoint registered and serves + # failures. + self.engine_monitor = VllmEngineMonitor(runtime, engine, shutdown_event) + logger.info("Embedding worker handler initialized") + + def cleanup(self) -> None: + """Release resources owned by this handler. + + AsyncLLM lifecycle is owned by the worker factory / runtime; the + engine monitor cancels its background tasks via ``__del__``. + """ + return None + + async def _monitor_abort(self, context: Context, request_id: str) -> None: + """Background task: abort the encode if context is cancelled or + shutdown_event fires. Raises EngineShutdown on shutdown so the + ``_abort_monitor`` context manager can propagate it. + + Mirrors ``BaseWorkerHandler._monitor_abort`` but trimmed for the + embedding path (no ``is_prefill``, no ``abort_guard``). + """ + try: + # `list[Any]` mirrors BaseWorkerHandler._monitor_abort: the + # iterable mixes the Future from async_killed_or_stopped() with + # the Task from shutdown_event.wait(). + wait_for: list[Any] = [context.async_killed_or_stopped()] + shutdown_task = None + if self.shutdown_event is not None: + shutdown_task = asyncio.create_task(self.shutdown_event.wait()) + wait_for.append(shutdown_task) + + done, pending = await asyncio.wait( + wait_for, return_when=asyncio.FIRST_COMPLETED + ) + + for task in pending: + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + logger.debug(f"Aborting embedding request ID: {request_id}") + try: + await asyncio.shield(self.engine_client.abort(request_id)) + except asyncio.CancelledError: + logger.debug( + f"Abort shielded from cancellation for embedding request " + f"{request_id}, continuing in background" + ) + + if shutdown_task is not None and shutdown_task in done: + raise EngineShutdown("Engine was shut down during embedding.") + except asyncio.CancelledError: + pass + except EngineShutdown: + raise + except Exception as e: + # Unexpected failure in the monitor task — log and propagate so + # `_abort_monitor.__aexit__` surfaces it via ``task.result()`` + # rather than silently leaving the encode unmanaged. + logger.error( + f"Error in embedding abort monitor for request {request_id}: {e}" + ) + raise + + @asynccontextmanager + async def _abort_monitor(self, context: Context, request_id: str): + """Create + tear down an abort monitor task around one encode call. + + On exit, re-raises EngineShutdown if the monitor caught a shutdown. + """ + task = asyncio.create_task(self._monitor_abort(context, request_id)) + try: + yield task + finally: + if not task.done(): + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + else: + # Re-raise EngineShutdown if the monitor task raised it. + task.result() + + async def generate( + self, request: dict, context: Context + ) -> AsyncIterator[Dict[str, Any]]: + """Handle one OpenAI /v1/embeddings request. + + The Rust frontend forwards the request dict directly. Expected keys: + ``model: str``, ``input: str | list[str] | list[int] | list[list[int]]``. + Optional ``dimensions`` and ``encoding_format`` fields are currently + ignored. + """ + # 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 "" + input_field = request.get("input") + if input_field is None: + raise ValueError("Embedding request missing required 'input' field") + + # Per OpenAI spec, `input` can be: + # - str : single text prompt + # - list[str] : batch of text prompts + # - list[int] : single pre-tokenized prompt (token IDs) + # - list[list[int]]: batch of pre-tokenized prompts + # Token-id forms must be passed to vLLM as TokensPrompt so the engine + # skips its own tokenizer; the previous str()-coercion path turned + # `[1, 2, 3]` into three text prompts ("1", "2", "3") instead of one. + prompts: list[Any] = _classify_embedding_input(input_field) + + pooling_params = PoolingParams() + # Use the per-request context id (same as the chat/completion paths + # in this file) so concurrent embeddings never collide inside + # ``AsyncLLM``. ``context.trace_id`` is a distributed-trace id + # shared by every request in a trace and ``id(context)`` can be + # reused across short-lived ``Context`` objects, so neither is + # unique enough to scope a vLLM ``request_id``. + base_request_id = context.id() + + embedding_objects: list[Dict[str, Any]] = [] + prompt_tokens = 0 + + for idx, prompt in enumerate(prompts): + request_id = f"{base_request_id}-{idx}" + encode_arg: Any = ( + prompt + if isinstance(prompt, str) + else TokensPrompt(prompt_token_ids=prompt) + ) + final_output = None + async with self._abort_monitor(context, request_id): + async for out in self.engine_client.encode( + prompt=encode_arg, + pooling_params=pooling_params, + request_id=request_id, + ): + final_output = out + + if final_output is None: + raise RuntimeError( + f"vLLM engine.encode produced no output for input index {idx}" + ) + + embedding_objects.append( + { + "object": "embedding", + "embedding": _pooling_output_to_list(final_output.outputs.data), + "index": idx, + } + ) + token_ids = getattr(final_output, "prompt_token_ids", None) or [] + prompt_tokens += len(token_ids) + + yield { + "object": "list", + "data": embedding_objects, + "model": model_name, + "usage": { + "prompt_tokens": prompt_tokens, + "total_tokens": prompt_tokens, + }, + } + + +def _is_token_id(x: Any) -> bool: + """True iff ``x`` is an int that could be a vLLM token id. + + Filters out ``bool`` (subclass of int) so ``[True, False]`` is not + accepted as a tokenized prompt. + """ + return isinstance(x, int) and not isinstance(x, bool) + + +def _classify_embedding_input(input_field: Any) -> list[Any]: + """Map an OpenAI ``input`` payload to a list of vLLM-ready prompts. + + Returns a list whose elements are either: + - ``str`` — passed straight to ``engine.encode`` as text, or + - ``list[int]`` — wrapped in ``TokensPrompt`` by the caller. + + Rejects mixed lists (e.g. ``["foo", 42]`` or ``[[1, 2], "bar"]``) with + a clear ``TypeError`` rather than silently coercing. + """ + if isinstance(input_field, str): + return [input_field] + if not isinstance(input_field, list): + raise TypeError( + f"Invalid 'input' type {type(input_field).__name__}; " + "expected str, list[str], list[int], or list[list[int]]" + ) + if not input_field: + raise ValueError("Embedding request 'input' must be non-empty") + + first = input_field[0] + if isinstance(first, str): + texts: list[str] = [] + for item in input_field: + if not isinstance(item, str): + raise TypeError( + "'input' list mixes str and non-str entries; pass either " + "all strings or all token-id arrays" + ) + texts.append(item) + return texts + if _is_token_id(first): + token_ids: list[int] = [] + for item in input_field: + if not _is_token_id(item): + raise TypeError( + "'input' list mixes int and non-int entries; for tokenized " + "input pass all integers (single prompt) or list[list[int]]" + ) + token_ids.append(item) + # Single tokenized prompt. + return [token_ids] + if isinstance(first, list): + prompts: list[list[int]] = [] + for i, item in enumerate(input_field): + if not isinstance(item, list): + raise TypeError( + f"'input' list element at index {i} must be a list of " + "ints (token IDs); mixed batches are not supported" + ) + inner: list[int] = [] + for x in item: + if not _is_token_id(x): + raise TypeError( + f"'input' list element at index {i} must be a list of " + "ints (token IDs); mixed batches are not supported" + ) + inner.append(x) + prompts.append(inner) + return prompts + raise TypeError( + f"Unsupported 'input' element type {type(first).__name__}; " + "expected str, int, or list[int]" + ) + + +def _pooling_output_to_list(data: Any) -> list[float]: + """Convert a vLLM PoolingOutput.data tensor (or list) to a flat list[float]. + + vLLM's pooling pipeline can return a tensor with a singleton batch dim + (shape ``(1, hidden_dim)``) instead of a 1D vector (shape ``(hidden_dim,)``). + The OpenAI ``/v1/embeddings`` response expects ``data[].embedding`` to be a + flat array of floats, so we flatten unconditionally. + """ + if isinstance(data, torch.Tensor): + return data.detach().cpu().flatten().tolist() + if isinstance(data, (list, tuple)): + # Already a list — flatten one level if it's a list-of-lists. + if data and isinstance(data[0], (list, tuple)): + return [float(x) for row in data for x in row] + return [float(x) for x in data] + raise TypeError( + f"Unsupported PoolingOutput.data type {type(data).__name__}; " + "expected torch.Tensor or list" + ) diff --git a/components/src/dynamo/vllm/tests/test_backend_args.py b/components/src/dynamo/vllm/tests/test_backend_args.py index 41c54e4d3122..117e1d97a4d6 100644 --- a/components/src/dynamo/vllm/tests/test_backend_args.py +++ b/components/src/dynamo/vllm/tests/test_backend_args.py @@ -37,6 +37,9 @@ def create_config() -> DynamoVllmConfig: config.multimodal_worker = False config.multimodal_encode_worker = False config.multimodal_decode_worker = False + config.enable_multimodal = False + config.embedding_worker = False + config.benchmark_mode = None return config @@ -119,3 +122,58 @@ def test_decode_worker(self, mode): else: with pytest.raises(ValueError): config._resolve_disaggregation_model_from_legacy_multimodal_flags() + + +class TestEmbeddingWorkerExclusivity: + """--embedding-worker rejects combinations that don't make sense for a + pooling engine (non-aggregated disagg, multimodal, benchmark-mode). + """ + + def test_baseline_aggregated_is_accepted(self): + config = create_config() + config.embedding_worker = True + config.disaggregation_mode = DisaggregationMode.AGGREGATED + # Must not raise. + config._validate_embedding_worker_exclusivity() + + @pytest.mark.parametrize( + "mode", + [ + DisaggregationMode.PREFILL, + DisaggregationMode.DECODE, + DisaggregationMode.ENCODE, + ], + ) + def test_non_aggregated_disagg_rejected(self, mode): + config = create_config() + config.embedding_worker = True + config.disaggregation_mode = mode + with pytest.raises(ValueError, match="disaggregation-mode=agg"): + config._validate_embedding_worker_exclusivity() + + def test_multimodal_combination_rejected(self): + config = create_config() + config.embedding_worker = True + config.disaggregation_mode = DisaggregationMode.AGGREGATED + config.enable_multimodal = True + with pytest.raises(ValueError, match="multimodal"): + config._validate_embedding_worker_exclusivity() + + def test_benchmark_mode_rejected(self): + # The bug surfaced by review: --embedding-worker + --benchmark-mode + # silently injected InstrumentedScheduler (a generation scheduler) on + # the pooling engine. Validation must reject the combination upfront. + config = create_config() + config.embedding_worker = True + config.disaggregation_mode = DisaggregationMode.AGGREGATED + config.benchmark_mode = "agg" + with pytest.raises(ValueError, match="benchmark-mode"): + config._validate_embedding_worker_exclusivity() + + def test_no_op_when_embedding_worker_disabled(self): + # Validator must not punish callers that have benchmark_mode set + # but are not running an embedding worker. + config = create_config() + config.embedding_worker = False + config.benchmark_mode = "agg" + config._validate_embedding_worker_exclusivity() diff --git a/components/src/dynamo/vllm/tests/test_vllm_unit.py b/components/src/dynamo/vllm/tests/test_vllm_unit.py index ce4b0065468a..deb6c66576d8 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_unit.py +++ b/components/src/dynamo/vllm/tests/test_vllm_unit.py @@ -819,3 +819,70 @@ def test_no_runner_attr_skipped_gracefully(self): update_engine_config_with_dynamo(dynamo_cfg, engine_cfg) assert not hasattr(engine_cfg, "runner") + + +class TestEmbeddingWorkerFlag: + """Parsing + validation for --embedding-worker.""" + + def test_default_false(self, mock_vllm_cli): + """Without --embedding-worker, the flag is False.""" + mock_vllm_cli("--model", "Qwen/Qwen3-0.6B") + config = parse_args() + assert config.embedding_worker is False + + def test_flag_sets_true(self, mock_vllm_cli): + """--embedding-worker on its own with default agg mode parses cleanly.""" + mock_vllm_cli( + "--model", + "Qwen/Qwen3-0.6B", + "--embedding-worker", + "--runner", + "pooling", + ) + config = parse_args() + assert config.embedding_worker is True + + def test_rejects_prefill_disagg(self, mock_vllm_cli): + """--embedding-worker combined with --disaggregation-mode prefill is rejected.""" + mock_vllm_cli( + "--model", + "Qwen/Qwen3-0.6B", + "--embedding-worker", + "--runner", + "pooling", + "--disaggregation-mode", + "prefill", + "--kv-transfer-config", + '{"kv_connector":"NixlConnector","kv_role":"kv_both"}', + ) + with pytest.raises(ValueError, match="--embedding-worker is only valid"): + parse_args() + + def test_rejects_decode_disagg(self, mock_vllm_cli): + """--embedding-worker combined with --disaggregation-mode decode is rejected.""" + mock_vllm_cli( + "--model", + "Qwen/Qwen3-0.6B", + "--embedding-worker", + "--runner", + "pooling", + "--disaggregation-mode", + "decode", + ) + with pytest.raises(ValueError, match="--embedding-worker is only valid"): + parse_args() + + def test_rejects_multimodal_combo(self, mock_vllm_cli): + """--embedding-worker combined with multimodal flags is rejected.""" + mock_vllm_cli( + "--model", + "Qwen/Qwen3-0.6B", + "--embedding-worker", + "--runner", + "pooling", + "--enable-multimodal", + ) + with pytest.raises( + ValueError, match="--embedding-worker cannot be combined with multimodal" + ): + parse_args() diff --git a/components/src/dynamo/vllm/tests/test_vllm_worker_factory.py b/components/src/dynamo/vllm/tests/test_vllm_worker_factory.py index 6fee3d4d99a3..92f98bd6aee9 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_worker_factory.py +++ b/components/src/dynamo/vllm/tests/test_vllm_worker_factory.py @@ -32,6 +32,7 @@ def _make_config(**overrides) -> Mock: "omni": False, "route_to_encoder": False, "disaggregation_mode": DisaggregationMode.AGGREGATED, + "embedding_worker": False, } defaults.update(overrides) return Mock(**defaults) @@ -54,6 +55,7 @@ def factory(self) -> WorkerFactory: factory._create_multimodal_worker = AsyncMock() # type: ignore[assignment] factory._create_prefill_worker = AsyncMock() # type: ignore[assignment] factory._create_decode_worker = AsyncMock() # type: ignore[assignment] + factory._create_embedding_worker = AsyncMock() # type: ignore[assignment] return factory # Tests for non-legacy worker config, 'route_to_encode' is worker internal config @@ -105,6 +107,20 @@ async def test_encode(self, factory: WorkerFactory, route_to_encode: bool) -> No factory._create_multimodal_encode_worker.assert_called_once() # type: ignore[union-attr] + async def test_embedding_worker_takes_priority( + self, factory: WorkerFactory + ) -> None: + """--embedding-worker is checked first; disaggregation_mode is ignored.""" + config = _make_config(embedding_worker=True) + shutdown_event = asyncio.Event() + + await factory.create(Mock(), config, shutdown_event, []) + + factory._create_embedding_worker.assert_called_once() # type: ignore[union-attr] + factory._create_decode_worker.assert_not_called() # type: ignore[union-attr] + factory._create_prefill_worker.assert_not_called() # type: ignore[union-attr] + factory._create_multimodal_encode_worker.assert_not_called() # type: ignore[union-attr] + async def test_passes_snapshot_engine(self, factory: WorkerFactory) -> None: config = _make_config(multimodal_worker=True) runtime = Mock() diff --git a/components/src/dynamo/vllm/tests/test_vllm_worker_handler.py b/components/src/dynamo/vllm/tests/test_vllm_worker_handler.py index 8509210e96ac..007cc38122b0 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_worker_handler.py +++ b/components/src/dynamo/vllm/tests/test_vllm_worker_handler.py @@ -1028,3 +1028,63 @@ async def _empty_gen(*args, **kwargs): handler.engine_client.abort.assert_not_called() if guard._abort_task is not None: assert guard._abort_task.done() + + +class TestClassifyEmbeddingInput: + """Unit tests for the embedding input classifier. + + Covers the four OpenAI-spec input shapes (str / list[str] / list[int] / + list[list[int]]) and the previous bug where `[1, 2, 3]` was silently + coerced to three text prompts via `str(item)`. Pure-function logic, no + async / vLLM engine needed. + """ + + def test_single_string(self): + assert mod._classify_embedding_input("hello") == ["hello"] + + def test_list_of_strings(self): + result = mod._classify_embedding_input(["a", "b", "c"]) + assert result == ["a", "b", "c"] + + def test_list_of_ints_is_one_tokenized_prompt(self): + # The bug: previously this returned ["1", "2", "3"] (three text + # prompts). Correct behavior is one tokenized prompt. + result = mod._classify_embedding_input([1, 2, 3]) + assert result == [[1, 2, 3]] + + def test_list_of_list_of_ints_is_batch_of_tokenized_prompts(self): + result = mod._classify_embedding_input([[1, 2], [3, 4, 5]]) + assert result == [[1, 2], [3, 4, 5]] + + def test_mixed_str_and_int_rejected(self): + with pytest.raises(TypeError, match="mixes str and non-str"): + mod._classify_embedding_input(["hello", 42]) + + def test_mixed_int_and_str_rejected(self): + with pytest.raises(TypeError, match="mixes int and non-int"): + mod._classify_embedding_input([1, "two"]) + + def test_mixed_list_of_lists_with_str_rejected(self): + with pytest.raises(TypeError, match="must be a list of"): + mod._classify_embedding_input([[1, 2], "three"]) + + def test_inner_list_with_non_int_rejected(self): + with pytest.raises(TypeError, match="must be a list of"): + mod._classify_embedding_input([[1, 2], [3.5, 4]]) + + def test_bool_is_not_treated_as_int(self): + # `bool` is a subclass of `int`; token ids must be real ints. + with pytest.raises(TypeError): + mod._classify_embedding_input([True, False]) + + def test_empty_list_rejected(self): + with pytest.raises(ValueError, match="must be non-empty"): + mod._classify_embedding_input([]) + + def test_unsupported_top_level_type_rejected(self): + with pytest.raises(TypeError, match="Invalid 'input' type"): + mod._classify_embedding_input({"text": "hi"}) + + def test_unsupported_element_type_rejected(self): + with pytest.raises(TypeError, match="Unsupported 'input' element"): + mod._classify_embedding_input([3.14, 2.71]) diff --git a/components/src/dynamo/vllm/worker_factory.py b/components/src/dynamo/vllm/worker_factory.py index 5ff59c718e65..9860cf7ac5f5 100644 --- a/components/src/dynamo/vllm/worker_factory.py +++ b/components/src/dynamo/vllm/worker_factory.py @@ -30,6 +30,7 @@ from .handlers import ( BaseWorkerHandler, DecodeWorkerHandler, + EmbeddingWorkerHandler, PrefillWorkerHandler, get_dp_range_for_worker, ) @@ -136,6 +137,16 @@ async def create( ) -> None: """Create the appropriate multimodal worker based on config flags.""" + # Embedding worker is selected first because it crosses worker shapes + # (pooling AsyncLLM, ModelType.Embedding) rather than being a variant + # of decode. Aggregated-only — exclusivity with disagg modes is + # enforced earlier in DynamoVllmConfig._validate_embedding_worker_exclusivity. + if config.embedding_worker: + await self._create_embedding_worker( + runtime, config, shutdown_event, shutdown_endpoints + ) + return + # NOTE: --benchmark-mode is only supported for prefill/decode workers. # The encode worker path does not wire benchmark waiting or # the get_perf_metrics endpoint. @@ -193,6 +204,94 @@ async def _create_multimodal_encode_worker( finally: handler.cleanup() + async def _create_embedding_worker( + self, + runtime: DistributedRuntime, + config: Config, + shutdown_event: asyncio.Event, + shutdown_endpoints: list, # mutated in place + ) -> None: + """Initialize an aggregated text-embedding worker. + + Pooling models have no KV cache, no decode phase, and no streamed + output, so several pieces of the decode-worker setup are intentionally + skipped here: + + - KV-events publisher: no KV cache → nothing to publish. + - Forward-pass-metrics relay: relays decode-phase ZMQ metrics; no + decode here. + - StatLoggerFactory wiring: built around per-batch sampling/decoding + stats which the pooling engine does not emit. + - InstrumentedScheduler: hard-codes ``pooling_params=None`` (see + components/src/dynamo/vllm/instrumented_scheduler.py), which would + silently disable the pooling pass. ``setup_vllm_engine`` only + installs it when ``--benchmark-mode`` is set, which is rejected + for embedding workers via config validation. + + We are deliberately not extending ``--benchmark-mode`` with an + ``embed`` choice. That flag exists primarily to expose a worker's + capability curve (RPS / p99 vs. concurrency, throughput knee) at + startup for capacity planning, engine-arg tuning, and as input to + the Dynamo planner's auto-scaling decisions. Decode workloads + benefit because they have many interacting knobs (max-num-seqs, + chunked prefill, prefill/decode mix). Embedding workloads are + essentially ``(batch_size × ISL → latency)`` -- a clean two-axis + function -- so the value of in-process self-profiling is much + lower than external HTTP load testing, which is what every other + embedding-serving stack uses anyway. The single remaining wedge + is planner integration: if/when the Dynamo planner needs + in-process embedding capability curves to auto-scale embedding + fleets, add ``--benchmark-mode embed`` at that point together + with the planner's embedding-capability model. + + The engine itself is the standard ``AsyncLLM`` constructed by + ``setup_vllm_engine``; pooling vs. generation is selected by the + user's ``--runner pooling`` argument flowing through ``engine_args``. + """ + generate_endpoint = runtime.endpoint( + f"{config.namespace}.{config.component}.{config.endpoint}" + ) + shutdown_endpoints[:] = [generate_endpoint] + + fpm_worker_id = str(generate_endpoint.connection_id()) + factory = StatLoggerFactory(endpoint=generate_endpoint) + ( + engine_client, + vllm_config, + _default_sampling_params, + _prometheus_temp_dir, + _component_gauges, + ) = self.setup_vllm_engine(config, factory, fpm_worker_id=fpm_worker_id) + + handler = EmbeddingWorkerHandler( + runtime=runtime, + engine=engine_client, + config=config, + shutdown_event=shutdown_event, + ) + + logger.info("Starting to serve the embedding worker endpoint...") + try: + await asyncio.gather( + generate_endpoint.serve_endpoint( + handler.generate, + metrics_labels=[("model", config.model)], + ), + self.register_vllm_model( + ModelInput.Text, + ModelType.Embedding, + generate_endpoint, + config, + engine_client, + vllm_config, + ), + ) + except Exception as e: + logger.error(f"Failed to serve embedding worker endpoint: {e}") + raise + finally: + handler.cleanup() + async def _maybe_wait_for_failover_lock( self, handler, diff --git a/examples/backends/vllm/launch/agg_embed.sh b/examples/backends/vllm/launch/agg_embed.sh new file mode 100755 index 000000000000..5224a2ec9ae3 --- /dev/null +++ b/examples/backends/vllm/launch/agg_embed.sh @@ -0,0 +1,91 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Aggregated embedding model serving. +# GPUs: 1 + +set -e +trap 'echo Cleaning up...; kill 0' EXIT + +SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" +source "$SCRIPT_DIR/../../../common/gpu_utils.sh" # build_vllm_gpu_mem_args +source "$SCRIPT_DIR/../../../common/launch_utils.sh" # print_launch_banner, wait_any_exit + +# Default embedding model. Smaller alternatives: +# `BAAI/bge-small-en-v1.5` (CPU-friendly), `intfloat/e5-small-v2`. +MODEL="Qwen/Qwen3-Embedding-0.6B" + +# Parse command line arguments +EXTRA_ARGS=() +while [[ $# -gt 0 ]]; do + case $1 in + --model) + MODEL="$2" + shift 2 + ;; + -h|--help) + echo "Usage: $0 [OPTIONS]" + echo "Options:" + echo " --model Specify embedding model (default: $MODEL)" + echo " -h, --help Show this help message" + echo "" + echo "Any additional options are passed through to dynamo.vllm." + echo "Note: --runner pooling, --dtype float32, and --pooler-config" + echo "are set here. Override via EXTRA_ARGS if your model requires" + echo "different pooling." + exit 0 + ;; + *) + EXTRA_ARGS+=("$1") + shift + ;; + esac +done + +GPU_MEM_ARGS=$(build_vllm_gpu_mem_args) + +HTTP_PORT="${DYN_HTTP_PORT:-8000}" +print_launch_banner --no-curl "Launching Embedding Worker (1 GPU)" "$MODEL" "$HTTP_PORT" + +print_curl_footer < Date: Thu, 21 May 2026 09:48:23 -0500 Subject: [PATCH 2/2] fix(vllm/embedding): cancel shutdown_event.wait() task on _monitor_abort exit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit On the success path of an embedding request, ``_abort_monitor.__aexit__`` cancels the running ``_monitor_abort`` task while it's blocked in ``asyncio.wait``. The CancelledError propagates out of ``asyncio.wait`` and jumps directly to ``except asyncio.CancelledError: pass``, short-circuiting past the ``for task in pending: task.cancel()`` cleanup loop. The result: the ``shutdown_event.wait()`` task created inside ``_monitor_abort`` is never cancelled — one leaked task per embedding request, accumulating for the lifetime of the worker. Fix: add a ``finally`` block that cancels ``shutdown_task`` on every exit path (normal completion, cancellation, EngineShutdown, unexpected exception). Addresses dynamo-ops review on PR #9713. Signed-off-by: Tzu-Ling --- components/src/dynamo/vllm/handlers.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/components/src/dynamo/vllm/handlers.py b/components/src/dynamo/vllm/handlers.py index d925fd3fca18..f0403f3b2348 100644 --- a/components/src/dynamo/vllm/handlers.py +++ b/components/src/dynamo/vllm/handlers.py @@ -2794,12 +2794,12 @@ async def _monitor_abort(self, context: Context, request_id: str) -> None: Mirrors ``BaseWorkerHandler._monitor_abort`` but trimmed for the embedding path (no ``is_prefill``, no ``abort_guard``). """ + shutdown_task: Optional[asyncio.Task] = None try: # `list[Any]` mirrors BaseWorkerHandler._monitor_abort: the # iterable mixes the Future from async_killed_or_stopped() with # the Task from shutdown_event.wait(). wait_for: list[Any] = [context.async_killed_or_stopped()] - shutdown_task = None if self.shutdown_event is not None: shutdown_task = asyncio.create_task(self.shutdown_event.wait()) wait_for.append(shutdown_task) @@ -2838,6 +2838,19 @@ async def _monitor_abort(self, context: Context, request_id: str) -> None: f"Error in embedding abort monitor for request {request_id}: {e}" ) raise + finally: + # On the success path the wrapping ``_abort_monitor`` cancels + # this coroutine while it's blocked in ``asyncio.wait``, which + # short-circuits past the pending-task cleanup loop above and + # leaves ``shutdown_task`` (the ``shutdown_event.wait()`` task) + # pending forever — one leaked task per embedding request. + # Cancel it here on every exit path. + if shutdown_task is not None and not shutdown_task.done(): + shutdown_task.cancel() + try: + await shutdown_task + except asyncio.CancelledError: + pass @asynccontextmanager async def _abort_monitor(self, context: Context, request_id: str):