feat(vllm): add aggregated text-embedding worker (DIS-2092)

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."
+            )

components/src/dynamo/vllm/handlers.py

@@ -2748,3 +2748,304 @@ 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``).
+        """
+        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()]
+            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
+        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):
+        """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"
+    )