From af2eb019b534026738953ba5e81c40bbcee3b2c4 Mon Sep 17 00:00:00 2001 From: xianlubird Date: Mon, 31 Aug 2026 13:42:34 +0800 Subject: [PATCH 1/4] feat(sglang): expose cross-encoder reranking through /v1/rerank Signed-off-by: xianlubird --- .../src/dynamo/sglang/init_embedding.py | 2 +- components/src/dynamo/sglang/protocol.py | 16 +- components/src/dynamo/sglang/register.py | 2 +- .../sglang/request_handlers/__init__.py | 2 + .../embedding/embedding_handler.py | 9 + .../request_handlers/rerank/__init__.py | 6 + .../request_handlers/rerank/rerank_handler.py | 124 +++++++++++++ .../tests/test_sglang_rerank_handler.py | 175 ++++++++++++++++++ lib/bindings/python/rust/lib.rs | 9 + lib/llm/src/discovery/model.rs | 14 +- lib/llm/src/discovery/model_manager.rs | 54 +++++- lib/llm/src/discovery/watcher.rs | 27 ++- lib/llm/src/discovery/worker_set.rs | 19 +- lib/llm/src/endpoint_type.rs | 4 + lib/llm/src/http/service/metrics.rs | 5 + lib/llm/src/http/service/openai.rs | 108 ++++++++++- lib/llm/src/http/service/service_v2.rs | 11 ++ lib/llm/src/model_type.rs | 26 +++ lib/llm/src/protocols/openai.rs | 1 + lib/llm/src/protocols/openai/rerank.rs | 156 ++++++++++++++++ .../src/protocols/openai/rerank/aggregator.rs | 63 +++++++ lib/llm/src/types.rs | 12 ++ lib/llm/tests/http-service.rs | 93 +++++++++- 23 files changed, 922 insertions(+), 16 deletions(-) create mode 100644 components/src/dynamo/sglang/request_handlers/rerank/__init__.py create mode 100644 components/src/dynamo/sglang/request_handlers/rerank/rerank_handler.py create mode 100644 components/src/dynamo/sglang/tests/test_sglang_rerank_handler.py create mode 100644 lib/llm/src/protocols/openai/rerank.rs create mode 100644 lib/llm/src/protocols/openai/rerank/aggregator.rs diff --git a/components/src/dynamo/sglang/init_embedding.py b/components/src/dynamo/sglang/init_embedding.py index ebbe0cb3df9e..fb81f20526f1 100644 --- a/components/src/dynamo/sglang/init_embedding.py +++ b/components/src/dynamo/sglang/init_embedding.py @@ -87,7 +87,7 @@ async def init_embedding( server_args, dynamo_args, input_type=ModelInput.Text, - output_type=ModelType.Embedding, + output_type=ModelType.Embedding | ModelType.Rerank, readiness_gate=ready_event, worker_type=WorkerType.Aggregated, needs=[], diff --git a/components/src/dynamo/sglang/protocol.py b/components/src/dynamo/sglang/protocol.py index 54a7a393231e..5070cf48c008 100644 --- a/components/src/dynamo/sglang/protocol.py +++ b/components/src/dynamo/sglang/protocol.py @@ -57,12 +57,22 @@ class EmbeddingRequest(BaseModel): model: str input: EmbeddingInput user: Optional[str] = None - dimensions: Optional[ - int - ] = None # only supported in text-embedding-3 and later models from OpenAI + dimensions: Optional[int] = ( + None # only supported in text-embedding-3 and later models from OpenAI + ) encoding_format: Literal["float", "base64"] = "float" +class RerankRequest(BaseModel): + """SGLang-compatible text-only cross-encoder rerank request.""" + + model: str + query: str + documents: List[str] + top_n: Optional[int] = None + return_documents: bool = True + + class DisaggPreprocessedRequest(BaseModel): request: Union[PreprocessedRequest, ChatCompletionRequest] sampling_params: dict diff --git a/components/src/dynamo/sglang/register.py b/components/src/dynamo/sglang/register.py index cd85f81e3aaa..efc4ccf75489 100644 --- a/components/src/dynamo/sglang/register.py +++ b/components/src/dynamo/sglang/register.py @@ -151,7 +151,7 @@ async def _register_model_with_runtime_config( ) input_type = ModelInput.Text # Only override output_type for chat models, not for embeddings - if output_type != ModelType.Embedding: + if not output_type.supports_embedding(): output_type = ModelType.Chat if runtime_config is not None and _supports_engine_generate( diff --git a/components/src/dynamo/sglang/request_handlers/__init__.py b/components/src/dynamo/sglang/request_handlers/__init__.py index 0c1fde3e2bc8..bae5130cc002 100644 --- a/components/src/dynamo/sglang/request_handlers/__init__.py +++ b/components/src/dynamo/sglang/request_handlers/__init__.py @@ -14,6 +14,7 @@ "PrefillWorkerHandler": ".llm", # Embedding handlers "EmbeddingWorkerHandler": ".embedding", + "RerankWorkerHandler": ".rerank", # Image diffusion handlers "ImageDiffusionWorkerHandler": ".image_diffusion", # Video generation handlers @@ -45,6 +46,7 @@ def __getattr__(name: str) -> Any: "PrefillWorkerHandler", # Embedding handlers "EmbeddingWorkerHandler", + "RerankWorkerHandler", # Image diffusion handlers "ImageDiffusionWorkerHandler", # Video generation handlers diff --git a/components/src/dynamo/sglang/request_handlers/embedding/embedding_handler.py b/components/src/dynamo/sglang/request_handlers/embedding/embedding_handler.py index 17636c74c37e..fe0fb1ddc83f 100644 --- a/components/src/dynamo/sglang/request_handlers/embedding/embedding_handler.py +++ b/components/src/dynamo/sglang/request_handlers/embedding/embedding_handler.py @@ -20,6 +20,7 @@ observe_embedding_input_tokens, ) from dynamo.sglang.request_handlers.handler_base import BaseWorkerHandler +from dynamo.sglang.request_handlers.rerank import RerankWorkerHandler def _encode_floats_to_base64(floats: List[float]) -> str: @@ -44,6 +45,9 @@ def __init__( shutdown_event: Optional[asyncio.Event] = None, ): super().__init__(engine, config, publisher, None, shutdown_event) + self.rerank_handler = RerankWorkerHandler( + engine, enable_trace=self.enable_trace + ) logging.info("Embedding worker handler initialized") def cleanup(self) -> None: @@ -97,6 +101,11 @@ async def generate( request: Embedding request dictionary. context: Context object for cancellation handling. """ + if "query" in request or "documents" in request: + async for response in self.rerank_handler.generate(request, context): + yield response + return + embedding_input = request.get("input") if isinstance(embedding_input, str): input_type = "str" diff --git a/components/src/dynamo/sglang/request_handlers/rerank/__init__.py b/components/src/dynamo/sglang/request_handlers/rerank/__init__.py new file mode 100644 index 000000000000..d379f644f529 --- /dev/null +++ b/components/src/dynamo/sglang/request_handlers/rerank/__init__.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dynamo.sglang.request_handlers.rerank.rerank_handler import RerankWorkerHandler + +__all__ = ["RerankWorkerHandler"] diff --git a/components/src/dynamo/sglang/request_handlers/rerank/rerank_handler.py b/components/src/dynamo/sglang/request_handlers/rerank/rerank_handler.py new file mode 100644 index 000000000000..64369bf494c4 --- /dev/null +++ b/components/src/dynamo/sglang/request_handlers/rerank/rerank_handler.py @@ -0,0 +1,124 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import heapq +from collections.abc import AsyncGenerator +from typing import Any + +import sglang as sgl +from sglang.srt.managers.io_struct import EmbeddingReqInput + +from dynamo._core import Context +from dynamo.sglang.protocol import RerankRequest + + +def _looks_like_decoder_reranker(tokenizer_manager: Any) -> bool: + """Detect decoder-only Qwen3 rerankers that need SGLang's generation path.""" + tokenizer = getattr(tokenizer_manager, "tokenizer", None) + chat_template = getattr(tokenizer, "chat_template", "") + template = chat_template.lower() if isinstance(chat_template, str) else "" + model_config = getattr(tokenizer_manager, "model_config", None) + model_path = str(getattr(model_config, "model_path", "")).lower() + has_yes_no_template = ( + "answer can only be" in template and "yes" in template and "no" in template + ) + return ( + has_yes_no_template + or "qwen3-reranker" in model_path + or "qwen3-vl-reranker" in model_path + ) + + +class RerankWorkerHandler: + """Runs stable text-only cross-encoder reranking through SGLang pooling.""" + + def __init__(self, engine: sgl.Engine, *, enable_trace: bool) -> None: + self.engine = engine + self.enable_trace = enable_trace + + @staticmethod + def _validate(request: RerankRequest) -> None: + if not request.query.strip(): + raise ValueError("Query cannot be empty or whitespace only") + if not request.documents: + raise ValueError("Documents cannot be empty") + if any(not document.strip() for document in request.documents): + raise ValueError("Each document cannot be empty or whitespace only") + if request.top_n is not None and request.top_n < 1: + raise ValueError("Parameter 'top_n' must be larger than 0") + + @staticmethod + def _score(item: Any, index: int) -> float: + if not isinstance(item, dict) or "embedding" not in item: + raise ValueError(f"Missing embedding score for rerank at index {index}") + score = item["embedding"] + if isinstance(score, list): + if not score or not isinstance(score[0], (int, float)): + raise ValueError( + f"Invalid embedding score for rerank at index {index}: {score!r}" + ) + score = score[0] + if not isinstance(score, (int, float)): + raise ValueError( + f"Invalid embedding score for rerank at index {index}: {score!r}" + ) + return float(score) + + async def generate( + self, request: dict[str, Any], context: Context + ) -> AsyncGenerator[list[dict[str, Any]], None]: + rerank_request = RerankRequest(**request) + self._validate(rerank_request) + + tokenizer_manager = self.engine.tokenizer_manager + if _looks_like_decoder_reranker(tokenizer_manager): + raise ValueError( + "Dynamo's SGLang /v1/rerank integration currently supports " + "cross-encoder rerankers only; decoder-only Qwen3 and Qwen3-VL " + "rerankers are not supported" + ) + + pairs = [ + [rerank_request.query, document] for document in rerank_request.documents + ] + trace_id = context.trace_id + request_ids = ( + [f"{trace_id}-{index}" for index in range(len(pairs))] + if trace_id is not None + else None + ) + internal_request = EmbeddingReqInput( + text=pairs, + is_cross_encoder_request=True, + external_trace_header=( + context.trace_headers() if self.enable_trace else None + ), + rid=request_ids, + ) + responses = tokenizer_manager.generate_request(internal_request, None) + result = await anext(responses) + if not isinstance(result, list): + result = [result] + if len(result) != len(rerank_request.documents): + raise ValueError( + "SGLang returned an unexpected number of rerank scores: " + f"expected {len(rerank_request.documents)}, got {len(result)}" + ) + + ranked = [] + for index, item in enumerate(result): + response = {"score": self._score(item, index), "index": index} + if rerank_request.return_documents: + response["document"] = rerank_request.documents[index] + meta_info = item.get("meta_info") if isinstance(item, dict) else None + if meta_info is not None: + response["meta_info"] = meta_info + ranked.append(response) + + if rerank_request.top_n is not None: + ranked = heapq.nlargest( + rerank_request.top_n, ranked, key=lambda item: item["score"] + ) + else: + ranked.sort(key=lambda item: item["score"], reverse=True) + yield ranked diff --git a/components/src/dynamo/sglang/tests/test_sglang_rerank_handler.py b/components/src/dynamo/sglang/tests/test_sglang_rerank_handler.py new file mode 100644 index 000000000000..cddd636af892 --- /dev/null +++ b/components/src/dynamo/sglang/tests/test_sglang_rerank_handler.py @@ -0,0 +1,175 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for SGLang cross-encoder reranking.""" + +from types import SimpleNamespace +from typing import Any + +import pytest + +pytest.importorskip( + "sglang.srt.managers.io_struct", reason="sglang not installed in this container" +) + +from dynamo.sglang.request_handlers.embedding import ( # noqa: E402 + EmbeddingWorkerHandler, +) +from dynamo.sglang.request_handlers.rerank import RerankWorkerHandler # noqa: E402 + +pytestmark = [ + pytest.mark.unit, + pytest.mark.sglang, + pytest.mark.core, + pytest.mark.gpu_0, + pytest.mark.profiled_vram_gib(0), + pytest.mark.pre_merge, +] + + +class _TokenizerManager: + def __init__( + self, + scores: list[float], + *, + chat_template: str = "", + model_path: str = "BAAI/bge-reranker-v2-m3", + ) -> None: + self.scores = scores + self.requests: list[tuple[Any, Any]] = [] + self.tokenizer = SimpleNamespace(chat_template=chat_template) + self.model_config = SimpleNamespace(model_path=model_path) + + async def generate_request(self, request: Any, raw_request: Any): + request.normalize_batch_and_arguments() + self.requests.append((request, raw_request)) + yield [ + { + "embedding": [score], + "meta_info": {"prompt_tokens": index + 1}, + } + for index, score in enumerate(self.scores) + ] + + +class _Engine: + def __init__(self, tokenizer_manager: _TokenizerManager) -> None: + self.tokenizer_manager = tokenizer_manager + + +class _Context: + trace_id = "rerank-trace" + + def trace_headers(self) -> dict[str, str]: + return {"traceparent": "00-test"} + + +def _handler( + scores: list[float], + *, + chat_template: str = "", + model_path: str = "BAAI/bge-reranker-v2-m3", +) -> tuple[RerankWorkerHandler, _TokenizerManager]: + manager = _TokenizerManager( + scores, chat_template=chat_template, model_path=model_path + ) + return RerankWorkerHandler(_Engine(manager), enable_trace=True), manager + + +@pytest.mark.asyncio +async def test_builds_pairs_sorts_scores_and_applies_top_n(): + handler, manager = _handler([0.2, 0.9, 0.5]) + outputs = [ + output + async for output in handler.generate( + { + "model": "reranker", + "query": "query", + "documents": ["zero", "one", "two"], + "top_n": 2, + "return_documents": True, + }, + _Context(), + ) + ] + + assert [item["index"] for item in outputs[0]] == [1, 2] + assert [item["document"] for item in outputs[0]] == ["one", "two"] + [(request, raw_request)] = manager.requests + assert raw_request is None + assert request.text == [["query", "zero"], ["query", "one"], ["query", "two"]] + assert request.is_cross_encoder_request is True + assert request.rid == ["rerank-trace-0", "rerank-trace-1", "rerank-trace-2"] + assert request.external_trace_header == {"traceparent": "00-test"} + + +@pytest.mark.asyncio +async def test_omits_documents_when_not_requested(): + handler, _ = _handler([0.1, 0.8]) + [output] = [ + output + async for output in handler.generate( + { + "model": "reranker", + "query": "query", + "documents": ["a", "b"], + "return_documents": False, + }, + _Context(), + ) + ] + assert [item["index"] for item in output] == [1, 0] + assert all("document" not in item for item in output) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request", + [ + {"model": "m", "query": " ", "documents": ["doc"]}, + {"model": "m", "query": "query", "documents": []}, + {"model": "m", "query": "query", "documents": [""]}, + {"model": "m", "query": "query", "documents": ["doc"], "top_n": 0}, + ], +) +async def test_rejects_invalid_requests_before_inference(request): + handler, manager = _handler([0.5]) + with pytest.raises(ValueError): + _ = [output async for output in handler.generate(request, _Context())] + assert manager.requests == [] + + +@pytest.mark.asyncio +async def test_rejects_decoder_only_qwen3_reranker(): + handler, manager = _handler( + [0.5], + chat_template='The answer can only be "yes" or "no"', + model_path="Qwen/Qwen3-Reranker-0.6B", + ) + with pytest.raises(ValueError, match="cross-encoder rerankers only"): + _ = [ + output + async for output in handler.generate( + {"model": "m", "query": "query", "documents": ["doc"]}, + _Context(), + ) + ] + assert manager.requests == [] + + +@pytest.mark.asyncio +async def test_embedding_endpoint_dispatches_rerank_shape(): + class _RerankHandler: + async def generate(self, request, context): + yield [{"score": 0.7, "index": 0}] + + handler = EmbeddingWorkerHandler.__new__(EmbeddingWorkerHandler) + handler.rerank_handler = _RerankHandler() + output = [ + item + async for item in handler.generate( + {"model": "m", "query": "query", "documents": ["doc"]}, + _Context(), + ) + ] + assert output == [[{"score": 0.7, "index": 0}]] diff --git a/lib/bindings/python/rust/lib.rs b/lib/bindings/python/rust/lib.rs index 8c2fe813138c..b3d736757246 100644 --- a/lib/bindings/python/rust/lib.rs +++ b/lib/bindings/python/rust/lib.rs @@ -1052,6 +1052,11 @@ impl ModelType { inner: llm_rs::model_type::ModelType::Pooling, }; + #[classattr] + const Rerank: Self = ModelType { + inner: llm_rs::model_type::ModelType::Rerank, + }; + fn supports_chat(&self) -> bool { self.inner.supports_chat() } @@ -1068,6 +1073,10 @@ impl ModelType { self.inner.supports_pooling() } + fn supports_rerank(&self) -> bool { + self.inner.supports_rerank() + } + fn __or__(&self, other: &Self) -> Self { ModelType { inner: self.inner | other.inner, diff --git a/lib/llm/src/discovery/model.rs b/lib/llm/src/discovery/model.rs index bfc37dbcc178..7b687eefadea 100644 --- a/lib/llm/src/discovery/model.rs +++ b/lib/llm/src/discovery/model.rs @@ -27,7 +27,7 @@ use crate::types::{ classify::OpenAIClassifyStreamingEngine, completions::OpenAICompletionsStreamingEngine, embeddings::OpenAIEmbeddingsStreamingEngine, generate::GenerateStreamingEngine, images::OpenAIImagesStreamingEngine, pooling::OpenAIPoolingStreamingEngine, - videos::OpenAIVideosStreamingEngine, + rerank::OpenAIRerankStreamingEngine, videos::OpenAIVideosStreamingEngine, }, }; @@ -240,6 +240,13 @@ impl Model { .any(|entry| entry.value().has_pooling_engine()) } + /// Check if any WorkerSet has a rerank engine. + pub fn has_rerank_engine(&self) -> bool { + self.worker_sets + .iter() + .any(|entry| entry.value().has_rerank_engine()) + } + /// Check if any WorkerSet has a tensor engine. pub fn has_tensor_engine(&self) -> bool { self.worker_sets @@ -578,6 +585,11 @@ impl Model { .ok_or_else(|| self.engine_error(self.has_pooling_engine())) } + pub fn get_rerank_engine(&self) -> Result { + self.select_worker_set_with(|ws| ws.rerank_engine.clone()) + .ok_or_else(|| self.engine_error(self.has_rerank_engine())) + } + pub fn get_images_engine(&self) -> Result { self.select_worker_set_with(|ws| ws.images_engine.clone()) .ok_or_else(|| self.engine_error(self.has_images_engine())) diff --git a/lib/llm/src/discovery/model_manager.rs b/lib/llm/src/discovery/model_manager.rs index cbacfaae72e6..6f7397dd01f5 100644 --- a/lib/llm/src/discovery/model_manager.rs +++ b/lib/llm/src/discovery/model_manager.rs @@ -54,7 +54,7 @@ use crate::{ classify::OpenAIClassifyStreamingEngine, completions::OpenAICompletionsStreamingEngine, embeddings::OpenAIEmbeddingsStreamingEngine, generate::GenerateStreamingEngine, images::OpenAIImagesStreamingEngine, pooling::OpenAIPoolingStreamingEngine, - videos::OpenAIVideosStreamingEngine, + rerank::OpenAIRerankStreamingEngine, videos::OpenAIVideosStreamingEngine, }, }, worker_type::WorkerType, @@ -1117,6 +1117,16 @@ impl ModelManager { .collect() } + pub fn list_rerank_models(&self) -> Vec { + self.catalog + .load() + .models + .iter() + .filter(|(_, model)| model.has_rerank_engine()) + .map(|(name, _)| name.clone()) + .collect() + } + pub fn list_tensor_models(&self) -> Vec { self.catalog .load() @@ -1234,6 +1244,18 @@ impl ModelManager { .get_pooling_engine() } + pub fn get_rerank_engine( + &self, + model: &str, + ) -> Result { + self.catalog + .load() + .models + .get(model) + .ok_or_else(|| ModelManagerError::ModelNotFound(model.to_string()))? + .get_rerank_engine() + } + pub fn get_completions_engine( &self, model: &str, @@ -1547,6 +1569,29 @@ impl ModelManager { Ok(()) } + pub fn add_rerank_model( + &self, + model: &str, + card_checksum: &str, + engine: OpenAIRerankStreamingEngine, + ) -> Result<(), ModelManagerError> { + let _reservation = self.reservation_lock.lock(); + let model_entry = self.get_or_create_model(model); + if model_entry.has_rerank_engine() { + return Err(ModelManagerError::ModelAlreadyExists(model.to_string())); + } + let namespace = format!("__local_rerank_{}", model); + let mut ws = WorkerSet::new( + namespace.clone(), + card_checksum.to_string(), + Self::aggregated_local_card(), + ); + ws.rerank_engine = Some(engine); + model_entry.add_worker_set(namespace, Arc::new(ws)); + self.publish_catalog_locked(); + Ok(()) + } + pub fn add_tensor_model( &self, model: &str, @@ -1764,6 +1809,13 @@ impl ModelManager { .ok_or_else(|| ModelManagerError::ModelNotFound(model.to_string())) } + pub fn remove_rerank_model(&self, model: &str) -> Result<(), ModelManagerError> { + let namespace = format!("__local_rerank_{}", model); + self.remove_worker_set(model, &namespace) + .map(|_| ()) + .ok_or_else(|| ModelManagerError::ModelNotFound(model.to_string())) + } + pub fn remove_images_model(&self, model: &str) -> Result<(), ModelManagerError> { let namespace = format!("__local_images_{}", model); self.remove_worker_set(model, &namespace) diff --git a/lib/llm/src/discovery/watcher.rs b/lib/llm/src/discovery/watcher.rs index f3430f4503d5..2c71950ba09f 100644 --- a/lib/llm/src/discovery/watcher.rs +++ b/lib/llm/src/discovery/watcher.rs @@ -54,6 +54,7 @@ use crate::{ embeddings::{NvCreateEmbeddingRequest, NvCreateEmbeddingResponse}, images::{NvCreateImageRequest, NvImagesResponse}, pooling::{NvCreatePoolingRequest, NvCreatePoolingResponse}, + rerank::{NvCreateRerankRequest, NvCreateRerankResponse}, videos::{NvCreateVideoRequest, NvVideosResponse}, }, tensor::{NvCreateTensorRequest, NvCreateTensorResponse}, @@ -241,6 +242,7 @@ const ALL_MODEL_TYPES: &[ModelType] = &[ ModelType::Realtime, ModelType::Classify, ModelType::Pooling, + ModelType::Rerank, ]; /// Returns true if no models in the manager support the given model type. @@ -265,6 +267,8 @@ fn is_model_type_list_empty(manager: &ModelManager, model_type: ModelType) -> bo manager.list_classify_models().is_empty() } else if model_type == ModelType::Pooling { manager.list_pooling_models().is_empty() + } else if model_type == ModelType::Rerank { + manager.list_rerank_models().is_empty() } else { true } @@ -874,6 +878,17 @@ where worker_set.pooling_engine = Some(Arc::new(push_router)); } + if card.model_type.supports_rerank() { + let push_router = PushRouter::< + NvCreateRerankRequest, + Annotated, + >::from_client_with_monitor( + client.clone(), router_config.router_mode, None + ) + .await?; + worker_set.rerank_engine = Some(Arc::new(push_router)); + } + if card.model_type.supports_chat() { let chat_router = PushRouter::< NvCreateChatCompletionRequest, @@ -1013,7 +1028,7 @@ where // prefill is routed off `worker_type`.) anyhow::bail!( "Unsupported model configuration: {} with {} input. Supported combinations: \ - Tokens+(Chat|Completions), Text+(Chat|Completions|Images|Audios|Videos|Embeddings|Classify|Pooling|Realtime), \ + Tokens+(Chat|Completions), Text+(Chat|Completions|Images|Audios|Videos|Embeddings|Classify|Pooling|Rerank|Realtime), \ Tokens+Embeddings, Tensor+TensorBased", card.model_type, card.model_input.as_str() @@ -1820,16 +1835,17 @@ mod tests { assert!(is_model_type_list_empty(&mm, ModelType::Realtime)); assert!(is_model_type_list_empty(&mm, ModelType::Classify)); assert!(is_model_type_list_empty(&mm, ModelType::Pooling)); + assert!(is_model_type_list_empty(&mm, ModelType::Rerank)); } #[test] fn removal_cards_contain_only_the_empty_model_type() { let mm = ModelManager::new(); let mut card = ModelDeploymentCard::with_name_only("model"); - card.model_type = ModelType::Classify | ModelType::Pooling; + card.model_type = ModelType::Classify | ModelType::Pooling | ModelType::Rerank; let removed_cards = removed_model_cards(&mm, &card); - assert_eq!(removed_cards.len(), 2); + assert_eq!(removed_cards.len(), 3); assert!( removed_cards .iter() @@ -1840,6 +1856,11 @@ mod tests { .iter() .any(|card| card.model_type == ModelType::Pooling) ); + assert!( + removed_cards + .iter() + .any(|card| card.model_type == ModelType::Rerank) + ); assert!( removed_cards .iter() diff --git a/lib/llm/src/discovery/worker_set.rs b/lib/llm/src/discovery/worker_set.rs index 47345c7e013a..a31dc4c2ce1f 100644 --- a/lib/llm/src/discovery/worker_set.rs +++ b/lib/llm/src/discovery/worker_set.rs @@ -27,7 +27,7 @@ use crate::{ classify::OpenAIClassifyStreamingEngine, completions::OpenAICompletionsStreamingEngine, embeddings::OpenAIEmbeddingsStreamingEngine, generate::GenerateStreamingEngine, images::OpenAIImagesStreamingEngine, pooling::OpenAIPoolingStreamingEngine, - videos::OpenAIVideosStreamingEngine, + rerank::OpenAIRerankStreamingEngine, videos::OpenAIVideosStreamingEngine, }, }, }; @@ -152,6 +152,7 @@ pub struct WorkerSet { pub(crate) embeddings_engine: Option, pub(crate) classify_engine: Option, pub(crate) pooling_engine: Option, + pub(crate) rerank_engine: Option, pub(crate) images_engine: Option, pub(crate) videos_engine: Option, pub(crate) audios_engine: Option, @@ -198,6 +199,7 @@ impl WorkerSet { embeddings_engine: None, classify_engine: None, pooling_engine: None, + rerank_engine: None, images_engine: None, videos_engine: None, audios_engine: None, @@ -269,6 +271,10 @@ impl WorkerSet { self.pooling_engine.is_some() } + pub fn has_rerank_engine(&self) -> bool { + self.rerank_engine.is_some() + } + pub fn has_images_engine(&self) -> bool { self.images_engine.is_some() } @@ -315,6 +321,7 @@ impl WorkerSet { || self.has_embeddings_engine() || self.has_classify_engine() || self.has_pooling_engine() + || self.has_rerank_engine() || self.has_images_engine() || self.has_tensor_engine() || self.has_videos_engine() @@ -407,6 +414,7 @@ impl WorkerSet { retain_for_requests!(embeddings_engine); retain_for_requests!(classify_engine); retain_for_requests!(pooling_engine); + retain_for_requests!(rerank_engine); retain_for_requests!(images_engine); retain_for_requests!(videos_engine); retain_for_requests!(audios_engine); @@ -440,6 +448,7 @@ impl WorkerSet { embeddings_engine: lora_context_engine(&self.embeddings_engine, &lora_name), classify_engine: lora_context_engine(&self.classify_engine, &lora_name), pooling_engine: lora_context_engine(&self.pooling_engine, &lora_name), + rerank_engine: lora_context_engine(&self.rerank_engine, &lora_name), images_engine: lora_context_engine(&self.images_engine, &lora_name), videos_engine: lora_context_engine(&self.videos_engine, &lora_name), audios_engine: lora_context_engine(&self.audios_engine, &lora_name), @@ -491,6 +500,7 @@ mod tests { use crate::types::openai::embeddings::{NvCreateEmbeddingRequest, NvCreateEmbeddingResponse}; use crate::types::openai::images::{NvCreateImageRequest, NvImagesResponse}; use crate::types::openai::pooling::{NvCreatePoolingRequest, NvCreatePoolingResponse}; + use crate::types::openai::rerank::{NvCreateRerankRequest, NvCreateRerankResponse}; use crate::types::openai::videos::{NvCreateVideoRequest, NvVideosResponse}; use async_trait::async_trait; use dynamo_runtime::engine::AsyncEngine; @@ -616,6 +626,7 @@ mod tests { assert!(!ws.has_embeddings_engine()); assert!(!ws.has_classify_engine()); assert!(!ws.has_pooling_engine()); + assert!(!ws.has_rerank_engine()); assert!(!ws.has_images_engine()); assert!(!ws.has_videos_engine()); assert!(!ws.has_audios_engine()); @@ -676,6 +687,12 @@ mod tests { StubEngine::::new(), "pooling" ); + check!( + rerank_engine, + has_rerank_engine, + StubEngine::::new(), + "rerank" + ); check!( images_engine, has_images_engine, diff --git a/lib/llm/src/endpoint_type.rs b/lib/llm/src/endpoint_type.rs index 21ed99e311c2..8cf450b4548b 100644 --- a/lib/llm/src/endpoint_type.rs +++ b/lib/llm/src/endpoint_type.rs @@ -24,6 +24,8 @@ pub enum EndpointType { Classify, /// Pooling API (raw pooler output: token embeddings, logits, rewards) Pooling, + /// Cross-encoder reranking API + Rerank, /// Responses API Responses, /// Anthropic Messages API @@ -46,6 +48,7 @@ impl EndpointType { Self::Realtime => "realtime", Self::Classify => "classify", Self::Pooling => "pooling", + Self::Rerank => "rerank", Self::Responses => "responses", Self::AnthropicMessages => "anthropic_messages", Self::Generate => "generate", @@ -64,6 +67,7 @@ impl EndpointType { Self::Realtime, Self::Classify, Self::Pooling, + Self::Rerank, Self::Responses, Self::AnthropicMessages, Self::Generate, diff --git a/lib/llm/src/http/service/metrics.rs b/lib/llm/src/http/service/metrics.rs index 4eaea51efdae..a51c87bce022 100644 --- a/lib/llm/src/http/service/metrics.rs +++ b/lib/llm/src/http/service/metrics.rs @@ -468,6 +468,9 @@ pub enum Endpoint { /// Classification (sequence classification / cross-encoder pooling) Classify, + /// Cross-encoder reranking + Rerank, + /// Pooling (raw pooler output) Pooling, @@ -1555,6 +1558,7 @@ impl std::fmt::Display for Endpoint { Endpoint::ChatCompletions => write!(f, "chat_completions"), Endpoint::Embeddings => write!(f, "embeddings"), Endpoint::Classify => write!(f, "classify"), + Endpoint::Rerank => write!(f, "rerank"), Endpoint::Pooling => write!(f, "pooling"), Endpoint::Images => write!(f, "images"), Endpoint::Videos => write!(f, "videos"), @@ -1574,6 +1578,7 @@ impl Endpoint { Endpoint::ChatCompletions => "chat_completions", Endpoint::Embeddings => "embeddings", Endpoint::Classify => "classify", + Endpoint::Rerank => "rerank", Endpoint::Pooling => "pooling", Endpoint::Images => "images", Endpoint::Videos => "videos", diff --git a/lib/llm/src/http/service/openai.rs b/lib/llm/src/http/service/openai.rs index 15cb1e6c5e5e..20d05bc65a26 100644 --- a/lib/llm/src/http/service/openai.rs +++ b/lib/llm/src/http/service/openai.rs @@ -76,6 +76,7 @@ use crate::protocols::openai::{ NvCreatePoolingRequest, NvCreatePoolingResponse, PoolingEmbedDType, PoolingEncodingFormat, PoolingEndianness, PoolingOutput, }, + rerank::{NvCreateRerankRequest, NvCreateRerankResponse}, responses::{ NvCreateResponse, NvResponse, ResponseParams, ResponsesConversionError, chat_completion_to_response, @@ -1636,7 +1637,92 @@ async fn classify( Ok(Json(response).into_response()) } -fn pooling_or_classify_bad_request(message: String) -> ErrorResponse { +#[tracing::instrument(skip_all)] +async fn rerank( + State(state): State>, + headers: HeaderMap, + Json(mut request): Json, +) -> Result { + check_ready(&state)?; + check_model_serving_ready(&state, &request.model)?; + + if !state.nvext_enabled() { + warn_nvext_disabled( + "rerank", + request + .nvext + .as_ref() + .is_some_and(|nvext| nvext.annotations.is_some()), + ); + request.nvext = None; + } + + let canonical = state.manager().resolve_canonical_name(&request.model); + if canonical != request.model { + request.model = canonical; + } + let request_id = get_or_create_request_id(&headers); + let request = context_from_headers(request, request_id, &headers)?; + let request_id = request.id().to_string(); + let model = &request.model; + let metric_model = state.manager().metric_model_for(model).to_string(); + let mut inflight = state.metrics_clone().create_inflight_guard( + &metric_model, + Endpoint::Rerank, + false, + &request_id, + ); + + if let Err(message) = request.validate_semantics() { + inflight.mark_error(ErrorType::Validation); + return Err(pooling_family_bad_request(message.to_string())); + } + + let http_queue_guard = state.metrics_clone().create_http_queue_guard(&metric_model); + let engine = state.manager().get_rerank_engine(model).map_err(|error| { + let response = ErrorMessage::from_model_error(&error); + inflight.mark_error(extract_error_type_from_response(&response)); + response + })?; + let mut response_collector = state + .metrics_clone() + .create_response_collector(&metric_model); + let model_name = model.to_string(); + let stream = engine.generate(request).await.map_err(|error| { + if super::metrics::request_was_rejected(error.as_ref()) { + state + .metrics_clone() + .inc_rejection(&model_name, Endpoint::Rerank); + } + let response = ErrorMessage::from_anyhow(error, "Failed to generate reranking"); + inflight.mark_error(extract_error_type_from_response(&response)); + response + })?; + + let mut http_queue_guard = Some(http_queue_guard); + let stream = stream.inspect(move |response| { + process_response_and_observe_metrics( + response, + &mut response_collector, + &mut http_queue_guard, + ); + }); + let response = NvCreateRerankResponse::from_annotated_stream(stream) + .await + .map_err(|error| { + let response = ErrorMessage::from_anyhow( + anyhow::Error::new(error), + "Failed to fold rerank stream", + ); + inflight.mark_error(extract_error_type_from_response(&response)); + response + })?; + + inflight.mark_ok(); + Ok(Json(response).into_response()) +} + +fn pooling_family_bad_request(message: String) -> ErrorResponse { let code = StatusCode::BAD_REQUEST; ( code, @@ -1652,7 +1738,7 @@ fn pooling_or_classify_bad_request(message: String) -> ErrorResponse { fn validate_pooling_cache_salt(cache_salt: Option<&str>) -> Result<(), ErrorResponse> { if cache_salt == Some("") { - return Err(pooling_or_classify_bad_request( + return Err(pooling_family_bad_request( "Parameter 'cache_salt' must be a non-empty string if provided.".to_string(), )); } @@ -1866,7 +1952,7 @@ async fn pooling( // vLLM currently rejects dimensionality reduction on `/pooling`. if request.dimensions.is_some() { inflight.mark_error(ErrorType::Validation); - return Err(pooling_or_classify_bad_request( + return Err(pooling_family_bad_request( "dimensions is currently not supported".to_string(), )); } @@ -4011,6 +4097,22 @@ pub fn classify_router( (vec![doc], router) } +/// Create an Axum [`Router`] for SGLang-compatible cross-encoder reranking. +/// If no path is provided, the default path is `/v1/rerank`. +pub fn rerank_router( + state: Arc, + path: Option, +) -> (Vec, Router) { + let path = path.unwrap_or("/v1/rerank".to_string()); + let doc = RouteDoc::new(axum::http::Method::POST, &path); + let router = Router::new() + .route(&path, post(rerank)) + .layer(middleware::from_fn(smart_json_error_middleware)) + .layer(axum::extract::DefaultBodyLimit::max(get_body_limit())) + .with_state(state); + (vec![doc], router) +} + /// Create an Axum [`Router`] for the `/v1/pooling` endpoint (raw pooler output /// from pooling-runner models). If no path is provided, the default path is /// `/v1/pooling`. Deployments migrating clients from native `vllm-serve` diff --git a/lib/llm/src/http/service/service_v2.rs b/lib/llm/src/http/service/service_v2.rs index dee46c0e9139..79bb520da9a2 100644 --- a/lib/llm/src/http/service/service_v2.rs +++ b/lib/llm/src/http/service/service_v2.rs @@ -379,6 +379,7 @@ struct StateFlags { cmpl_endpoints_enabled: AtomicBool, embeddings_endpoints_enabled: AtomicBool, classify_endpoints_enabled: AtomicBool, + rerank_endpoints_enabled: AtomicBool, pooling_endpoints_enabled: AtomicBool, images_endpoints_enabled: AtomicBool, videos_endpoints_enabled: AtomicBool, @@ -397,6 +398,7 @@ impl StateFlags { EndpointType::Completion => self.cmpl_endpoints_enabled.load(Ordering::Relaxed), EndpointType::Embedding => self.embeddings_endpoints_enabled.load(Ordering::Relaxed), EndpointType::Classify => self.classify_endpoints_enabled.load(Ordering::Relaxed), + EndpointType::Rerank => self.rerank_endpoints_enabled.load(Ordering::Relaxed), EndpointType::Pooling => self.pooling_endpoints_enabled.load(Ordering::Relaxed), EndpointType::Images => self.images_endpoints_enabled.load(Ordering::Relaxed), EndpointType::Videos => self.videos_endpoints_enabled.load(Ordering::Relaxed), @@ -425,6 +427,9 @@ impl StateFlags { EndpointType::Classify => self .classify_endpoints_enabled .store(enabled, Ordering::Relaxed), + EndpointType::Rerank => self + .rerank_endpoints_enabled + .store(enabled, Ordering::Relaxed), EndpointType::Pooling => self .pooling_endpoints_enabled .store(enabled, Ordering::Relaxed), @@ -474,6 +479,7 @@ impl State { cmpl_endpoints_enabled: AtomicBool::new(false), embeddings_endpoints_enabled: AtomicBool::new(false), classify_endpoints_enabled: AtomicBool::new(false), + rerank_endpoints_enabled: AtomicBool::new(false), pooling_endpoints_enabled: AtomicBool::new(false), images_endpoints_enabled: AtomicBool::new(false), videos_endpoints_enabled: AtomicBool::new(false), @@ -1038,6 +1044,8 @@ static HTTP_SVC_CMP_PATH_ENV: &str = "DYN_HTTP_SVC_CMP_PATH"; static HTTP_SVC_EMB_PATH_ENV: &str = "DYN_HTTP_SVC_EMB_PATH"; /// Environment variable to set the classify endpoint path (default: `/v1/classify`) static HTTP_SVC_CLASSIFY_PATH_ENV: &str = "DYN_HTTP_SVC_CLASSIFY_PATH"; +/// Environment variable to set the rerank endpoint path (default: `/v1/rerank`) +static HTTP_SVC_RERANK_PATH_ENV: &str = "DYN_HTTP_SVC_RERANK_PATH"; /// Environment variable to set the pooling endpoint path (default: `/v1/pooling`) static HTTP_SVC_POOLING_PATH_ENV: &str = "DYN_HTTP_SVC_POOLING_PATH"; /// Environment variable to set the responses endpoint path (default: `/v1/responses`) @@ -1445,6 +1453,8 @@ impl HttpServiceConfigBuilder { super::openai::embeddings_router(state.clone(), var(HTTP_SVC_EMB_PATH_ENV).ok()); let (classify_docs, classify_route) = super::openai::classify_router(state.clone(), var(HTTP_SVC_CLASSIFY_PATH_ENV).ok()); + let (rerank_docs, rerank_route) = + super::openai::rerank_router(state.clone(), var(HTTP_SVC_RERANK_PATH_ENV).ok()); let (pooling_docs, pooling_route) = super::openai::pooling_router(state.clone(), var(HTTP_SVC_POOLING_PATH_ENV).ok()); let (images_docs, images_route) = super::openai::images_router(state.clone(), None); @@ -1461,6 +1471,7 @@ impl HttpServiceConfigBuilder { endpoint_routes.insert(EndpointType::Completion, (cmpl_docs, cmpl_route)); endpoint_routes.insert(EndpointType::Embedding, (embed_docs, embed_route)); endpoint_routes.insert(EndpointType::Classify, (classify_docs, classify_route)); + endpoint_routes.insert(EndpointType::Rerank, (rerank_docs, rerank_route)); endpoint_routes.insert(EndpointType::Pooling, (pooling_docs, pooling_route)); endpoint_routes.insert(EndpointType::Images, (images_docs, images_route)); endpoint_routes.insert(EndpointType::Videos, (videos_docs, videos_route)); diff --git a/lib/llm/src/model_type.rs b/lib/llm/src/model_type.rs index fa3004b808ee..721dbca067b2 100644 --- a/lib/llm/src/model_type.rs +++ b/lib/llm/src/model_type.rs @@ -62,6 +62,8 @@ bitflags! { /// mounts its `/pooling` alongside those surfaces for every /// pooling-runner model. const Pooling = 1 << 10; + /// Cross-encoder relevance scoring served on `/v1/rerank`. + const Rerank = 1 << 11; } } @@ -107,6 +109,9 @@ impl ModelType { pub fn supports_pooling(&self) -> bool { self.contains(ModelType::Pooling) } + pub fn supports_rerank(&self) -> bool { + self.contains(ModelType::Rerank) + } pub fn as_vec(&self) -> Vec<&'static str> { let mut result = Vec::new(); @@ -143,6 +148,9 @@ impl ModelType { if self.supports_pooling() { result.push("pooling"); } + if self.supports_rerank() { + result.push("rerank"); + } result } @@ -183,6 +191,9 @@ impl ModelType { if self.supports_pooling() { result.push(ModelType::Pooling); } + if self.supports_rerank() { + result.push(ModelType::Rerank); + } result } @@ -233,6 +244,9 @@ impl ModelType { if self.contains(Self::Pooling) { endpoint_types.push(crate::endpoint_type::EndpointType::Pooling); } + if self.contains(Self::Rerank) { + endpoint_types.push(crate::endpoint_type::EndpointType::Rerank); + } // [gluo NOTE] ModelType::Tensor doesn't map to any endpoint type, // current use of endpoint type is LLM specific and so does the HTTP // server that uses it. @@ -414,6 +428,18 @@ mod tests { ); } + #[test] + fn rerank_capability_maps_to_endpoint() { + assert_eq!(ModelType::Rerank.bits(), 1 << 11); + assert!(ModelType::Rerank.supports_rerank()); + assert_eq!(ModelType::Rerank.as_vec(), vec!["rerank"]); + assert_eq!(ModelType::Rerank.units(), vec![ModelType::Rerank]); + assert_eq!( + ModelType::Rerank.as_endpoint_types(), + vec![EndpointType::Rerank] + ); + } + #[test] fn token_generating_models_do_not_imply_vllm_generate_support() { assert!( diff --git a/lib/llm/src/protocols/openai.rs b/lib/llm/src/protocols/openai.rs index a594a3c595b1..4029712c7ba3 100644 --- a/lib/llm/src/protocols/openai.rs +++ b/lib/llm/src/protocols/openai.rs @@ -26,6 +26,7 @@ pub mod generate; pub mod images; pub mod models; pub mod pooling; +pub mod rerank; pub mod responses; pub mod stream_aggregator; pub mod tools; diff --git a/lib/llm/src/protocols/openai/rerank.rs b/lib/llm/src/protocols/openai/rerank.rs new file mode 100644 index 000000000000..cc22857699f0 --- /dev/null +++ b/lib/llm/src/protocols/openai/rerank.rs @@ -0,0 +1,156 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Protocol types for SGLang-compatible cross-encoder reranking at +//! `POST /v1/rerank`. + +use dynamo_runtime::protocols::annotated::AnnotationsProvider; +use serde::{Deserialize, Serialize}; +use utoipa::ToSchema; +use validator::Validate; + +mod aggregator; + +pub use super::embeddings::{NvExt, NvExtProvider}; + +/// Request for the `/v1/rerank` endpoint. +#[derive(ToSchema, Serialize, Deserialize, Validate, Debug, Clone)] +pub struct NvCreateRerankRequest { + /// The model to use for reranking. + pub model: String, + + /// Query to compare with each document. + pub query: String, + + /// Candidate documents to rank. + pub documents: Vec, + + /// Maximum number of ranked documents to return. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub top_n: Option, + + /// Include the source document in each result. + #[serde(default = "default_return_documents")] + pub return_documents: bool, + + #[serde(skip_serializing_if = "Option::is_none")] + pub nvext: Option, +} + +fn default_return_documents() -> bool { + true +} + +impl NvCreateRerankRequest { + /// Validate the cross-encoder request before dispatching it to a worker. + pub fn validate_semantics(&self) -> Result<(), &'static str> { + if self.query.trim().is_empty() { + return Err("Query cannot be empty or whitespace only"); + } + if self.documents.is_empty() { + return Err("Documents cannot be empty"); + } + if self + .documents + .iter() + .any(|document| document.trim().is_empty()) + { + return Err("Each document cannot be empty or whitespace only"); + } + if self.top_n == Some(0) { + return Err("Parameter 'top_n' must be larger than 0"); + } + Ok(()) + } +} + +/// One result in the SGLang-compatible rerank response array. +#[derive(ToSchema, Serialize, Deserialize, Debug, Clone, PartialEq)] +pub struct RerankResult { + pub score: f64, + pub index: usize, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub document: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub meta_info: Option, +} + +/// SGLang returns reranking results as a bare JSON array. +#[derive(ToSchema, Serialize, Deserialize, Debug, Clone, Default, PartialEq)] +#[serde(transparent)] +pub struct NvCreateRerankResponse(pub Vec); + +impl NvExtProvider for NvCreateRerankRequest { + fn nvext(&self) -> Option<&NvExt> { + self.nvext.as_ref() + } +} + +impl AnnotationsProvider for NvCreateRerankRequest { + fn annotations(&self) -> Option> { + self.nvext + .as_ref() + .and_then(|nvext| nvext.annotations.clone()) + } + + fn has_annotation(&self, annotation: &str) -> bool { + self.nvext + .as_ref() + .and_then(|nvext| nvext.annotations.as_ref()) + .is_some_and(|annotations| annotations.iter().any(|item| item == annotation)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn request_defaults_match_sglang() { + let request: NvCreateRerankRequest = serde_json::from_value(json!({ + "model": "BAAI/bge-reranker-v2-m3", + "query": "What is Dynamo?", + "documents": ["A", "B"] + })) + .unwrap(); + + assert_eq!(request.top_n, None); + assert!(request.return_documents); + assert!(request.validate_semantics().is_ok()); + } + + #[test] + fn invalid_requests_are_rejected() { + for (query, documents, top_n) in [ + ("", vec!["doc"], None), + ("query", vec![], None), + ("query", vec![" "], None), + ("query", vec!["doc"], Some(0)), + ] { + let request = NvCreateRerankRequest { + model: "model".to_string(), + query: query.to_string(), + documents: documents.into_iter().map(str::to_string).collect(), + top_n, + return_documents: true, + nvext: None, + }; + assert!(request.validate_semantics().is_err()); + } + } + + #[test] + fn response_serializes_as_bare_array_and_omits_documents() { + let response = NvCreateRerankResponse(vec![RerankResult { + score: 0.9, + index: 1, + document: None, + meta_info: None, + }]); + assert_eq!( + serde_json::to_value(response).unwrap(), + json!([{"score": 0.9, "index": 1}]) + ); + } +} diff --git a/lib/llm/src/protocols/openai/rerank/aggregator.rs b/lib/llm/src/protocols/openai/rerank/aggregator.rs new file mode 100644 index 000000000000..799aa6641494 --- /dev/null +++ b/lib/llm/src/protocols/openai/rerank/aggregator.rs @@ -0,0 +1,63 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::NvCreateRerankResponse; +use crate::protocols::{ + Annotated, + codec::{Message, SseCodecError}, + convert_sse_stream, + openai::stream_aggregator::{StreamAggregable, aggregate_stream}, +}; +use dynamo_runtime::{engine::DataStream, error::DynamoError}; +use futures::Stream; + +impl StreamAggregable for NvCreateRerankResponse { + fn empty() -> Self { + Self::default() + } + + fn merge(&mut self, next: Self) { + self.0.extend(next.0); + } +} + +impl NvCreateRerankResponse { + pub async fn from_sse_stream( + stream: DataStream>, + ) -> Result { + Self::from_annotated_stream(convert_sse_stream::(stream)).await + } + + pub async fn from_annotated_stream( + stream: impl Stream>, + ) -> Result { + aggregate_stream(stream).await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::protocols::openai::rerank::RerankResult; + use futures::stream; + + #[tokio::test] + async fn aggregates_worker_chunks() { + let item = |index, score| { + Annotated::from_data(NvCreateRerankResponse(vec![RerankResult { + score, + index, + document: None, + meta_info: None, + }])) + }; + let response = NvCreateRerankResponse::from_annotated_stream(stream::iter(vec![ + item(1, 0.9), + item(0, 0.2), + ])) + .await + .unwrap(); + assert_eq!(response.0.len(), 2); + assert_eq!(response.0[0].index, 1); + } +} diff --git a/lib/llm/src/types.rs b/lib/llm/src/types.rs index 7a21abd4140a..fcddbb7862a3 100644 --- a/lib/llm/src/types.rs +++ b/lib/llm/src/types.rs @@ -106,6 +106,18 @@ pub mod openai { ServerStreamingEngine>; } + pub mod rerank { + use super::*; + + pub use protocols::openai::rerank::{NvCreateRerankRequest, NvCreateRerankResponse}; + + pub type OpenAIRerankUnaryEngine = + UnaryEngine; + + pub type OpenAIRerankStreamingEngine = + ServerStreamingEngine>; + } + pub mod images { use super::*; diff --git a/lib/llm/tests/http-service.rs b/lib/llm/tests/http-service.rs index eda6304155b6..6a9e3bc4855d 100644 --- a/lib/llm/tests/http-service.rs +++ b/lib/llm/tests/http-service.rs @@ -569,6 +569,7 @@ fn compute_index(endpoint: &Endpoint, request_type: &RequestType, status: &Statu Endpoint::ChatCompletions => 1, Endpoint::Embeddings => todo!(), Endpoint::Classify => todo!(), + Endpoint::Rerank => todo!(), Endpoint::Pooling => todo!(), Endpoint::Responses => todo!(), Endpoint::AnthropicMessages => todo!(), @@ -2219,12 +2220,51 @@ impl } } +#[async_trait] +impl + AsyncEngine< + SingleIn, + ManyOut>, + Error, + > for UncalledPoolingFamilyEngine +{ + async fn generate( + &self, + request: SingleIn, + ) -> Result< + ManyOut>, + Error, + > { + use dynamo_llm::protocols::openai::rerank::{NvCreateRerankResponse, RerankResult}; + + let (_request, context) = request.transfer(()); + let response_ctx = context.context(); + let stream = stream! { + yield Annotated::from_data(NvCreateRerankResponse(vec![ + RerankResult { + score: 0.9, + index: 1, + document: Some("second".to_string()), + meta_info: None, + }, + RerankResult { + score: 0.2, + index: 0, + document: Some("first".to_string()), + meta_info: None, + }, + ])); + }; + Ok(ResponseStream::new(Box::pin(stream), response_ctx)) + } +} + /// A request rejected by handler-local validation must still be counted in /// `requests_total` with `error_type=validation`, like `chat_completions`. /// Validating before the inflight guard would drop these 400s from metrics /// (and from the "request completed" log the guard emits on drop). #[tokio::test] -async fn test_classify_and_pooling_validation_errors_are_metered() { +async fn test_pooling_family_validation_errors_are_metered() { let (listener, port) = bind_random_port().await; let service = HttpService::builder().port(port).build().unwrap(); service @@ -2233,6 +2273,9 @@ async fn test_classify_and_pooling_validation_errors_are_metered() { service .enable_model_endpoint(dynamo_llm::endpoint_type::EndpointType::Pooling, true) .unwrap(); + service + .enable_model_endpoint(dynamo_llm::endpoint_type::EndpointType::Rerank, true) + .unwrap(); let state = service.state_clone(); let manager = state.manager(); @@ -2250,7 +2293,10 @@ async fn test_classify_and_pooling_validation_errors_are_metered() { .add_classify_model("foo", card.mdcsum(), engine.clone()) .unwrap(); manager - .add_pooling_model("foo", card.mdcsum(), engine) + .add_pooling_model("foo", card.mdcsum(), engine.clone()) + .unwrap(); + manager + .add_rerank_model("foo", card.mdcsum(), engine) .unwrap(); let metrics = state.metrics_clone(); @@ -2276,6 +2322,26 @@ async fn test_classify_and_pooling_validation_errors_are_metered() { 1, ); + // ==== /v1/rerank: successful responses preserve SGLang's bare-array shape ==== + let response = client + .post(format!("http://localhost:{port}/v1/rerank")) + .json(&serde_json::json!({ + "model": "foo", + "query": "query", + "documents": ["first", "second"] + })) + .send() + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + response.json::().await.unwrap(), + serde_json::json!([ + {"score": 0.9, "index": 1, "document": "second"}, + {"score": 0.2, "index": 0, "document": "first"} + ]) + ); + // ==== /v1/pooling: empty cache_salt ==== let response = client .post(format!("http://localhost:{port}/v1/pooling")) @@ -2312,6 +2378,29 @@ async fn test_classify_and_pooling_validation_errors_are_metered() { 2, ); + // ==== /v1/rerank: top_n must be positive ==== + let response = client + .post(format!("http://localhost:{port}/v1/rerank")) + .json(&serde_json::json!({ + "model": "foo", + "query": "query", + "documents": ["document"], + "top_n": 0 + })) + .send() + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + compare_counter( + &metrics, + "foo", + &Endpoint::Rerank, + &RequestType::Unary, + &Status::Error, + &ErrorType::Validation, + 1, + ); + cancel_token.cancel(); task.await.unwrap().unwrap(); } From 92b6abc32430e0e6e33687f9d3735ab21a508d0e Mon Sep 17 00:00:00 2001 From: xianlubird Date: Mon, 31 Aug 2026 14:19:21 +0800 Subject: [PATCH 2/4] fix(sglang): address rerank CI failures Signed-off-by: xianlubird --- components/src/dynamo/sglang/protocol.py | 6 +++--- .../src/dynamo/sglang/request_handlers/__init__.py | 12 +++--------- .../request_handlers/embedding/embedding_handler.py | 6 +++--- .../sglang/request_handlers/rerank/rerank_handler.py | 5 ++++- .../sglang/tests/test_sglang_rerank_handler.py | 6 +++--- lib/bindings/python/src/dynamo/_core.pyi | 10 ++++++++-- 6 files changed, 24 insertions(+), 21 deletions(-) diff --git a/components/src/dynamo/sglang/protocol.py b/components/src/dynamo/sglang/protocol.py index 5070cf48c008..0ddffb92d085 100644 --- a/components/src/dynamo/sglang/protocol.py +++ b/components/src/dynamo/sglang/protocol.py @@ -57,9 +57,9 @@ class EmbeddingRequest(BaseModel): model: str input: EmbeddingInput user: Optional[str] = None - dimensions: Optional[int] = ( - None # only supported in text-embedding-3 and later models from OpenAI - ) + dimensions: Optional[ + int + ] = None # only supported in text-embedding-3 and later models from OpenAI encoding_format: Literal["float", "base64"] = "float" diff --git a/components/src/dynamo/sglang/request_handlers/__init__.py b/components/src/dynamo/sglang/request_handlers/__init__.py index bae5130cc002..3799365ca0a9 100644 --- a/components/src/dynamo/sglang/request_handlers/__init__.py +++ b/components/src/dynamo/sglang/request_handlers/__init__.py @@ -37,22 +37,16 @@ def __getattr__(name: str) -> Any: __all__ = [ - # Base handlers "BaseGenerativeHandler", "BaseWorkerHandler", - # LLM handlers "DecodeWorkerHandler", "DiffusionWorkerHandler", - "PrefillWorkerHandler", - # Embedding handlers "EmbeddingWorkerHandler", - "RerankWorkerHandler", - # Image diffusion handlers "ImageDiffusionWorkerHandler", - # Video generation handlers - "VideoGenerationWorkerHandler", - # Multimodal handlers "MultimodalEncodeWorkerHandler", "MultimodalPrefillWorkerHandler", "MultimodalWorkerHandler", + "PrefillWorkerHandler", + "RerankWorkerHandler", + "VideoGenerationWorkerHandler", ] diff --git a/components/src/dynamo/sglang/request_handlers/embedding/embedding_handler.py b/components/src/dynamo/sglang/request_handlers/embedding/embedding_handler.py index fe0fb1ddc83f..8879f5a1aebb 100644 --- a/components/src/dynamo/sglang/request_handlers/embedding/embedding_handler.py +++ b/components/src/dynamo/sglang/request_handlers/embedding/embedding_handler.py @@ -93,7 +93,7 @@ async def _encode_input( async def generate( self, request: dict, context: Context - ) -> AsyncGenerator[Dict[str, Any], None]: + ) -> AsyncGenerator[Dict[str, Any] | List[Dict[str, Any]], None]: """ Generate embeddings for the given input. @@ -102,8 +102,8 @@ async def generate( context: Context object for cancellation handling. """ if "query" in request or "documents" in request: - async for response in self.rerank_handler.generate(request, context): - yield response + async for rerank_response in self.rerank_handler.generate(request, context): + yield rerank_response return embedding_input = request.get("input") diff --git a/components/src/dynamo/sglang/request_handlers/rerank/rerank_handler.py b/components/src/dynamo/sglang/request_handlers/rerank/rerank_handler.py index 64369bf494c4..8eb81ad9a9e4 100644 --- a/components/src/dynamo/sglang/request_handlers/rerank/rerank_handler.py +++ b/components/src/dynamo/sglang/request_handlers/rerank/rerank_handler.py @@ -107,7 +107,10 @@ async def generate( ranked = [] for index, item in enumerate(result): - response = {"score": self._score(item, index), "index": index} + response: dict[str, Any] = { + "score": self._score(item, index), + "index": index, + } if rerank_request.return_documents: response["document"] = rerank_request.documents[index] meta_info = item.get("meta_info") if isinstance(item, dict) else None diff --git a/components/src/dynamo/sglang/tests/test_sglang_rerank_handler.py b/components/src/dynamo/sglang/tests/test_sglang_rerank_handler.py index cddd636af892..2c1573f47b51 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_rerank_handler.py +++ b/components/src/dynamo/sglang/tests/test_sglang_rerank_handler.py @@ -124,7 +124,7 @@ async def test_omits_documents_when_not_requested(): @pytest.mark.asyncio @pytest.mark.parametrize( - "request", + "bad_request", [ {"model": "m", "query": " ", "documents": ["doc"]}, {"model": "m", "query": "query", "documents": []}, @@ -132,10 +132,10 @@ async def test_omits_documents_when_not_requested(): {"model": "m", "query": "query", "documents": ["doc"], "top_n": 0}, ], ) -async def test_rejects_invalid_requests_before_inference(request): +async def test_rejects_invalid_requests_before_inference(bad_request): handler, manager = _handler([0.5]) with pytest.raises(ValueError): - _ = [output async for output in handler.generate(request, _Context())] + _ = [output async for output in handler.generate(bad_request, _Context())] assert manager.requests == [] diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index ca25958f2224..f7c779ac195f 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -1645,8 +1645,8 @@ class ModelInput: class ModelType: """OpenAI-style surfaces supported by a model. - Values are Chat, Completions, Embedding, Classify, Pooling, TensorBased, - Images, Audios, Videos, Realtime, and Empty (no OpenAI surface). + Values are Chat, Completions, Embedding, Classify, Pooling, Rerank, + TensorBased, Images, Audios, Videos, Realtime, and Empty (no OpenAI surface). """ # No OpenAI surface — used by prefill / encode workers whose role is # carried by WorkerType. Symmetric with the other ModelType.Foo members. @@ -1668,6 +1668,8 @@ class ModelType: # Raw pooler output served on /v1/pooling (token embeddings, logits, rewards). # Usually combined with Classify or Embedding: ModelType.Classify | ModelType.Pooling. Pooling: ModelType + # Cross-encoder relevance scoring served on /v1/rerank. + Rerank: ModelType def __or__(self, other: ModelType) -> ModelType: ... @@ -1688,6 +1690,10 @@ class ModelType: """Return True if this model type supports /v1/pooling.""" ... + def supports_rerank(self) -> bool: + """Return True if this model type supports /v1/rerank.""" + ... + class RouterMode: """Router mode for load balancing requests across workers""" RoundRobin: "RouterMode" From a0ebce919ba91c29314db97f1175c39245db7f8d Mon Sep 17 00:00:00 2001 From: xianlubird Date: Wed, 9 Sep 2026 14:12:46 +0800 Subject: [PATCH 3/4] fix(sglang): isolate reranking in a dedicated worker Signed-off-by: xianlubird --- components/src/dynamo/sglang/args.py | 15 +- components/src/dynamo/sglang/backend_args.py | 20 +++ components/src/dynamo/sglang/health_check.py | 30 ++++ .../src/dynamo/sglang/init_embedding.py | 93 +++++++---- components/src/dynamo/sglang/init_rerank.py | 27 ++++ components/src/dynamo/sglang/main.py | 5 + components/src/dynamo/sglang/publisher.py | 8 +- components/src/dynamo/sglang/register.py | 8 +- .../embedding/embedding_handler.py | 11 +- .../request_handlers/rerank/rerank_handler.py | 38 +++-- .../sglang/tests/test_sglang_publisher.py | 7 +- .../tests/test_sglang_rerank_handler.py | 152 ++++++++++++++---- .../dynamo/sglang/tests/test_sglang_unit.py | 1 + examples/backends/sglang/launch/agg_rerank.sh | 75 +++++++++ lib/llm/src/model_type.rs | 29 ++++ 15 files changed, 429 insertions(+), 90 deletions(-) create mode 100644 components/src/dynamo/sglang/init_rerank.py create mode 100755 examples/backends/sglang/launch/agg_rerank.sh diff --git a/components/src/dynamo/sglang/args.py b/components/src/dynamo/sglang/args.py index 429eaf35b5d0..e6ebfd2aebe9 100644 --- a/components/src/dynamo/sglang/args.py +++ b/components/src/dynamo/sglang/args.py @@ -103,6 +103,8 @@ def _unsupported_fpm_trace_role(dynamo_config: DynamoConfig) -> Optional[str]: """Return the worker role when the selected path does not create an FPM relay.""" if is_snapshot_enabled(): return "snapshot" + if dynamo_config.rerank_worker: + return "rerank" if dynamo_config.embedding_worker: return "embedding" if ( @@ -441,9 +443,14 @@ async def parse_args(args: list[str]) -> Config: if dynamo_config.enable_multimodal: parsed_args.enable_multimodal = True - # If --embedding-worker is set, also set SGLang's --is-embedding flag - if dynamo_config.embedding_worker: + # Both dedicated pooling modes use SGLang's embedding engine. + if dynamo_config.embedding_worker or dynamo_config.rerank_worker: parsed_args.is_embedding = True + if dynamo_config.rerank_worker and ( + parsed_args.disaggregation_mode != "null" + or getattr(parsed_args, "dllm_algorithm", None) + ): + raise ValueError("--rerank-worker requires aggregated cross-encoder serving") # Enable encoder_only mode for multimodal encode workers to load only vision encoder # This significantly reduces memory usage by avoiding loading the full LLM weights @@ -452,7 +459,9 @@ async def parse_args(args: list[str]) -> Config: endpoint = dynamo_config.endpoint if endpoint is None: - if dynamo_config.embedding_worker: + if dynamo_config.rerank_worker: + endpoint = f"dyn://{namespace}.rerank.generate" + elif dynamo_config.embedding_worker: endpoint = f"dyn://{namespace}.backend.generate" elif dynamo_config.image_diffusion_worker: endpoint = f"dyn://{namespace}.backend.generate" diff --git a/components/src/dynamo/sglang/backend_args.py b/components/src/dynamo/sglang/backend_args.py index 8c7c051dcda6..ee37ee7ade31 100644 --- a/components/src/dynamo/sglang/backend_args.py +++ b/components/src/dynamo/sglang/backend_args.py @@ -109,6 +109,14 @@ def add_arguments(self, parser: argparse.ArgumentParser) -> None: help="Run as embedding worker component (Dynamo flag, also sets SGLang's --is-embedding).", ) + add_negatable_bool_argument( + g, + flag_name="--rerank-worker", + env_var="DYN_SGL_RERANK_WORKER", + default=False, + help="Run a dedicated text cross-encoder rerank worker. Requires a rerank-capable frontend; also sets SGLang's --is-embedding.", + ) + add_negatable_bool_argument( g, flag_name="--image-diffusion-worker", @@ -186,6 +194,7 @@ class DynamoSGLangConfig(ConfigBase): dedicated_mm_encoder: bool = False embedding_transfer_mode: EmbeddingTransferMode embedding_worker: bool + rerank_worker: bool = False image_diffusion_worker: bool disagg_config: Optional[str] = None @@ -203,6 +212,17 @@ class DynamoSGLangConfig(ConfigBase): def validate(self) -> None: _reject_removed_multimodal_env_vars() + if self.rerank_worker and ( + self.embedding_worker + or self.image_diffusion_worker + or self.video_generation_worker + or self.enable_multimodal + or self.dedicated_mm_encoder + or self.frontend_decoding + ): + raise ValueError( + "--rerank-worker cannot be combined with other worker modes" + ) if not isinstance(self.embedding_transfer_mode, EmbeddingTransferMode): self.embedding_transfer_mode = EmbeddingTransferMode( diff --git a/components/src/dynamo/sglang/health_check.py b/components/src/dynamo/sglang/health_check.py index 1b1296dd2095..df1ad3717a19 100644 --- a/components/src/dynamo/sglang/health_check.py +++ b/components/src/dynamo/sglang/health_check.py @@ -87,6 +87,36 @@ def __init__( super().__init__() +class SglangEmbeddingHealthCheckPayload(HealthCheckPayload): + """Send an embedding request through the worker's normal encode path.""" + + def __init__( + self, + model_name: str, + engine: Optional[sgl.Engine] = None, + use_text_input: bool = False, + ) -> None: + self.default_payload = { + "model": model_name, + "input": ( + "Test" if use_text_input else [_get_bos_token_id_from_engine(engine)] + ), + } + super().__init__() + + +class SglangRerankHealthCheckPayload(HealthCheckPayload): + """Probe the cross-encoder path with a single query/document pair.""" + + def __init__(self, model: str) -> None: + self.default_payload = { + "model": model, + "query": "health check", + "documents": ["health check"], + } + super().__init__() + + class SglangDisaggHealthCheckPayload(HealthCheckPayload): """SGLang-specific health check payload for PD-disaggregated mode. diff --git a/components/src/dynamo/sglang/init_embedding.py b/components/src/dynamo/sglang/init_embedding.py index fb81f20526f1..107295e6825d 100644 --- a/components/src/dynamo/sglang/init_embedding.py +++ b/components/src/dynamo/sglang/init_embedding.py @@ -12,13 +12,16 @@ from dynamo.llm import ModelInput, ModelType, WorkerType from dynamo.runtime import DistributedRuntime from dynamo.sglang.args import Config -from dynamo.sglang.health_check import SglangHealthCheckPayload +from dynamo.sglang.health_check import ( + SglangEmbeddingHealthCheckPayload, + SglangRerankHealthCheckPayload, +) from dynamo.sglang.publisher import ( set_forward_pass_metrics_worker_id, setup_sgl_metrics, ) from dynamo.sglang.register import register_model_with_readiness_gate -from dynamo.sglang.request_handlers import EmbeddingWorkerHandler +from dynamo.sglang.request_handlers import EmbeddingWorkerHandler, RerankWorkerHandler from dynamo.sglang.request_handlers.embedding.metrics import init_embedding_metrics @@ -30,6 +33,25 @@ async def init_embedding( run_deferred_handlers: Callable[[], Awaitable[None]] | None = None, ) -> None: """Initialize embedding worker component""" + await _init_pooling( + runtime, + config, + shutdown_event, + shutdown_endpoints, + run_deferred_handlers, + rerank=False, + ) + + +async def _init_pooling( + runtime: DistributedRuntime, + config: Config, + shutdown_event: asyncio.Event, + shutdown_endpoints: list, + run_deferred_handlers: Callable[[], Awaitable[None]] | None = None, + *, + rerank: bool, +) -> None: server_args, dynamo_args = config.server_args, config.dynamo_args generate_endpoint = runtime.endpoint( @@ -45,35 +67,43 @@ async def init_embedding( engine, config, generate_endpoint ) - # Wire ``dynamo_embedding_*`` histograms to the worker's /metrics - # endpoint via a dedicated CollectorRegistry — the default Prometheus - # global registry is not exposed by the Dynamo SGLang publisher. - # Done AFTER engine init (and after setup_sgl_metrics) so the lazy - # ``prometheus_client`` import in metrics.py doesn't interfere with - # SGLang's multiprocess Prometheus setup. - from prometheus_client import CollectorRegistry + if not rerank: + # Wire ``dynamo_embedding_*`` histograms to the worker's /metrics + # endpoint via a dedicated CollectorRegistry — the default Prometheus + # global registry is not exposed by the Dynamo SGLang publisher. + # Done AFTER engine init (and after setup_sgl_metrics) so the lazy + # ``prometheus_client`` import in metrics.py doesn't interfere with + # SGLang's multiprocess Prometheus setup. + from prometheus_client import CollectorRegistry - embedding_metrics_registry = CollectorRegistry() - register_engine_metrics_callback( - endpoint=generate_endpoint, - registry=embedding_metrics_registry, - metric_prefix_filters=["dynamo_embedding_"], - namespace_name=dynamo_args.namespace, - component_name=dynamo_args.component, - endpoint_name=dynamo_args.endpoint, - model_name=server_args.served_model_name, - ) - init_embedding_metrics(embedding_metrics_registry) + embedding_metrics_registry = CollectorRegistry() + register_engine_metrics_callback( + endpoint=generate_endpoint, + registry=embedding_metrics_registry, + metric_prefix_filters=["dynamo_embedding_"], + namespace_name=dynamo_args.namespace, + component_name=dynamo_args.component, + endpoint_name=dynamo_args.endpoint, + model_name=server_args.served_model_name, + ) + init_embedding_metrics(embedding_metrics_registry) ready_event = asyncio.Event() - - handler = EmbeddingWorkerHandler(engine, config, publisher, shutdown_event) - health_check_payload = SglangHealthCheckPayload( - engine, use_text_input=dynamo_args.use_sglang_tokenizer - ).to_dict() - - register_model_taint_route(runtime, generate_endpoint) + handler: EmbeddingWorkerHandler | RerankWorkerHandler | None = None try: + if rerank: + handler = RerankWorkerHandler(engine, config, publisher, shutdown_event) + health_check_payload = SglangRerankHealthCheckPayload( + server_args.served_model_name or server_args.model_path + ).to_dict() + else: + handler = EmbeddingWorkerHandler(engine, config, publisher, shutdown_event) + health_check_payload = SglangEmbeddingHealthCheckPayload( + server_args.served_model_name or server_args.model_path, + engine, + use_text_input=dynamo_args.use_sglang_tokenizer, + ).to_dict() + register_model_taint_route(runtime, generate_endpoint) await asyncio.gather( generate_endpoint.serve_endpoint( handler.generate, @@ -87,14 +117,14 @@ async def init_embedding( server_args, dynamo_args, input_type=ModelInput.Text, - output_type=ModelType.Embedding | ModelType.Rerank, + output_type=ModelType.Rerank if rerank else ModelType.Embedding, readiness_gate=ready_event, worker_type=WorkerType.Aggregated, needs=[], ), ) except Exception as e: - logging.error(f"Failed to serve embedding endpoints: {e}") + logging.error("Failed to serve pooling endpoint: %s", e) raise finally: metrics_task.cancel() @@ -103,7 +133,10 @@ async def init_embedding( except asyncio.CancelledError: logging.info("Metrics task successfully cancelled") pass - handler.cleanup() + if handler is not None: + handler.cleanup() + else: + engine.shutdown() if run_deferred_handlers is not None: logging.info("Running deferred handlers") await run_deferred_handlers() diff --git a/components/src/dynamo/sglang/init_rerank.py b/components/src/dynamo/sglang/init_rerank.py new file mode 100644 index 000000000000..347ab5e691e1 --- /dev/null +++ b/components/src/dynamo/sglang/init_rerank.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +from typing import Awaitable, Callable + +from dynamo.runtime import DistributedRuntime +from dynamo.sglang.args import Config +from dynamo.sglang.init_embedding import _init_pooling + + +async def init_rerank( + runtime: DistributedRuntime, + config: Config, + shutdown_event: asyncio.Event, + shutdown_endpoints: list, + run_deferred_handlers: Callable[[], Awaitable[None]] | None = None, +) -> None: + """Serve and advertise a dedicated cross-encoder worker.""" + await _init_pooling( + runtime, + config, + shutdown_event, + shutdown_endpoints, + run_deferred_handlers, + rerank=True, + ) diff --git a/components/src/dynamo/sglang/main.py b/components/src/dynamo/sglang/main.py index 91d02e166a95..54ce0e4da2ad 100644 --- a/components/src/dynamo/sglang/main.py +++ b/components/src/dynamo/sglang/main.py @@ -29,6 +29,7 @@ init_multimodal_prefill_worker, init_multimodal_worker, ) +from dynamo.sglang.init_rerank import init_rerank from dynamo.sglang.shutdown import install_graceful_shutdown from dynamo.sglang.snapshot import prepare_snapshot_engine @@ -88,6 +89,10 @@ async def worker(argv: list[str] | None = None): await init_video_diffusion( runtime, config, shutdown_endpoints, run_deferred_handlers ) + elif config.dynamo_args.rerank_worker: + await init_rerank( + runtime, config, shutdown_event, shutdown_endpoints, run_deferred_handlers + ) elif config.dynamo_args.embedding_worker: await init_embedding( runtime, diff --git a/components/src/dynamo/sglang/publisher.py b/components/src/dynamo/sglang/publisher.py index 54892c47f468..17246bdf1e8b 100644 --- a/components/src/dynamo/sglang/publisher.py +++ b/components/src/dynamo/sglang/publisher.py @@ -479,7 +479,7 @@ async def setup_sgl_metrics( and starts a ``DynamoSglangPublisher`` that pulls scheduler metrics over ZMQ and (optionally) forwards KV events / FPM stats. - For **embedding workers** (``config.dynamo_args.embedding_worker``), + For **embedding and rerank workers**, the chat-shaped pipeline is **skipped entirely**: pooling engines have no KV cache, no prefill/decode phase, and no scheduler metrics worth collecting, so every metric in that pipeline would emit zeros @@ -500,9 +500,11 @@ async def setup_sgl_metrics( """ metrics_labels = [("model", engine.server_args.served_model_name)] - if getattr(config.dynamo_args, "embedding_worker", False): + if getattr(config.dynamo_args, "embedding_worker", False) or getattr( + config.dynamo_args, "rerank_worker", False + ): logging.info( - "Embedding worker: skipping chat-shaped Prometheus + KV-event " + "Pooling worker: skipping chat-shaped Prometheus + KV-event " "wiring (no KV cache, no prefill/decode, no scheduler metrics). " "Embedding-shaped metrics are registered separately." ) diff --git a/components/src/dynamo/sglang/register.py b/components/src/dynamo/sglang/register.py index efc4ccf75489..d55fd50c0eb3 100644 --- a/components/src/dynamo/sglang/register.py +++ b/components/src/dynamo/sglang/register.py @@ -145,14 +145,14 @@ async def _register_model_with_runtime_config( """ runtime_config = await get_runtime_config(engine, server_args, dynamo_args) - if dynamo_args.use_sglang_tokenizer: + if dynamo_args.use_sglang_tokenizer and not ( + output_type.supports_embedding() or output_type.supports_rerank() + ): logging.warning( "Using the sglang tokenizer/detokenizer instead. The dynamo tokenizer/detokenizer will not be used and only v1/chat/completions will be available" ) input_type = ModelInput.Text - # Only override output_type for chat models, not for embeddings - if not output_type.supports_embedding(): - output_type = ModelType.Chat + output_type = ModelType.Chat if runtime_config is not None and _supports_engine_generate( input_type, output_type, worker_type diff --git a/components/src/dynamo/sglang/request_handlers/embedding/embedding_handler.py b/components/src/dynamo/sglang/request_handlers/embedding/embedding_handler.py index 8879f5a1aebb..17636c74c37e 100644 --- a/components/src/dynamo/sglang/request_handlers/embedding/embedding_handler.py +++ b/components/src/dynamo/sglang/request_handlers/embedding/embedding_handler.py @@ -20,7 +20,6 @@ observe_embedding_input_tokens, ) from dynamo.sglang.request_handlers.handler_base import BaseWorkerHandler -from dynamo.sglang.request_handlers.rerank import RerankWorkerHandler def _encode_floats_to_base64(floats: List[float]) -> str: @@ -45,9 +44,6 @@ def __init__( shutdown_event: Optional[asyncio.Event] = None, ): super().__init__(engine, config, publisher, None, shutdown_event) - self.rerank_handler = RerankWorkerHandler( - engine, enable_trace=self.enable_trace - ) logging.info("Embedding worker handler initialized") def cleanup(self) -> None: @@ -93,7 +89,7 @@ async def _encode_input( async def generate( self, request: dict, context: Context - ) -> AsyncGenerator[Dict[str, Any] | List[Dict[str, Any]], None]: + ) -> AsyncGenerator[Dict[str, Any], None]: """ Generate embeddings for the given input. @@ -101,11 +97,6 @@ async def generate( request: Embedding request dictionary. context: Context object for cancellation handling. """ - if "query" in request or "documents" in request: - async for rerank_response in self.rerank_handler.generate(request, context): - yield rerank_response - return - embedding_input = request.get("input") if isinstance(embedding_input, str): input_type = "str" diff --git a/components/src/dynamo/sglang/request_handlers/rerank/rerank_handler.py b/components/src/dynamo/sglang/request_handlers/rerank/rerank_handler.py index 8eb81ad9a9e4..0b2f8165b189 100644 --- a/components/src/dynamo/sglang/request_handlers/rerank/rerank_handler.py +++ b/components/src/dynamo/sglang/request_handlers/rerank/rerank_handler.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import asyncio import heapq from collections.abc import AsyncGenerator from typing import Any @@ -9,7 +10,10 @@ from sglang.srt.managers.io_struct import EmbeddingReqInput from dynamo._core import Context +from dynamo.sglang.args import Config from dynamo.sglang.protocol import RerankRequest +from dynamo.sglang.publisher import DynamoSglangPublisher +from dynamo.sglang.request_handlers.handler_base import BaseWorkerHandler def _looks_like_decoder_reranker(tokenizer_manager: Any) -> bool: @@ -29,12 +33,31 @@ def _looks_like_decoder_reranker(tokenizer_manager: Any) -> bool: ) -class RerankWorkerHandler: +class RerankWorkerHandler(BaseWorkerHandler): """Runs stable text-only cross-encoder reranking through SGLang pooling.""" - def __init__(self, engine: sgl.Engine, *, enable_trace: bool) -> None: - self.engine = engine - self.enable_trace = enable_trace + def __init__( + self, + engine: sgl.Engine, + config: Config, + publisher: DynamoSglangPublisher | None = None, + shutdown_event: asyncio.Event | None = None, + ) -> None: + self.validate_engine(engine) + super().__init__(engine, config, publisher, None, shutdown_event) + + @staticmethod + def validate_engine(engine: sgl.Engine) -> None: + if _looks_like_decoder_reranker(engine.tokenizer_manager): + raise ValueError( + "Dynamo's SGLang /v1/rerank integration currently supports " + "cross-encoder rerankers only; decoder-only Qwen3 and Qwen3-VL " + "rerankers are not supported" + ) + + def cleanup(self) -> None: + super().cleanup() + self.engine.shutdown() @staticmethod def _validate(request: RerankRequest) -> None: @@ -71,13 +94,6 @@ async def generate( self._validate(rerank_request) tokenizer_manager = self.engine.tokenizer_manager - if _looks_like_decoder_reranker(tokenizer_manager): - raise ValueError( - "Dynamo's SGLang /v1/rerank integration currently supports " - "cross-encoder rerankers only; decoder-only Qwen3 and Qwen3-VL " - "rerankers are not supported" - ) - pairs = [ [rerank_request.query, document] for document in rerank_request.documents ] diff --git a/components/src/dynamo/sglang/tests/test_sglang_publisher.py b/components/src/dynamo/sglang/tests/test_sglang_publisher.py index 0aea2168f9ee..f4c5d695bcfd 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_publisher.py +++ b/components/src/dynamo/sglang/tests/test_sglang_publisher.py @@ -581,7 +581,10 @@ def shutdown(self): @pytest.mark.asyncio -async def test_setup_sgl_metrics_skips_chat_pipeline_for_embedding_worker(monkeypatch): +@pytest.mark.parametrize("role", ["embedding_worker", "rerank_worker"]) +async def test_setup_sgl_metrics_skips_chat_pipeline_for_pooling_worker( + monkeypatch, role +): """``setup_sgl_metrics`` short-circuits for embedding workers. Chat-shaped collectors (``sglang:*`` multiproc metrics, the Dynamo @@ -624,7 +627,7 @@ def _wrapped(*_a, **_kw): ) ) config = SimpleNamespace( - dynamo_args=SimpleNamespace(embedding_worker=True), + dynamo_args=SimpleNamespace(**{role: True}), server_args=engine.server_args, ) generate_endpoint = SimpleNamespace() diff --git a/components/src/dynamo/sglang/tests/test_sglang_rerank_handler.py b/components/src/dynamo/sglang/tests/test_sglang_rerank_handler.py index 2c1573f47b51..1c6b0f53452e 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_rerank_handler.py +++ b/components/src/dynamo/sglang/tests/test_sglang_rerank_handler.py @@ -3,8 +3,11 @@ """Unit tests for SGLang cross-encoder reranking.""" +import argparse +import asyncio from types import SimpleNamespace from typing import Any +from unittest.mock import AsyncMock, Mock import pytest @@ -12,6 +15,18 @@ "sglang.srt.managers.io_struct", reason="sglang not installed in this container" ) +pytest.importorskip( + "sglang.srt.speculative.spec_info", reason="full SGLang runtime is required" +) + +from dynamo.llm import ModelType # noqa: E402 +from dynamo.sglang import init_embedding as pooling_init # noqa: E402 +from dynamo.sglang.backend_args import ( # noqa: E402 + DynamoSGLangArgGroup, + DynamoSGLangConfig, +) +from dynamo.sglang.health_check import SglangRerankHealthCheckPayload # noqa: E402 +from dynamo.sglang.init_rerank import init_rerank # noqa: E402 from dynamo.sglang.request_handlers.embedding import ( # noqa: E402 EmbeddingWorkerHandler, ) @@ -73,7 +88,10 @@ def _handler( manager = _TokenizerManager( scores, chat_template=chat_template, model_path=model_path ) - return RerankWorkerHandler(_Engine(manager), enable_trace=True), manager + handler = RerankWorkerHandler.__new__(RerankWorkerHandler) + handler.engine = _Engine(manager) + handler.enable_trace = True + return handler, manager @pytest.mark.asyncio @@ -139,37 +157,117 @@ async def test_rejects_invalid_requests_before_inference(bad_request): assert manager.requests == [] -@pytest.mark.asyncio -async def test_rejects_decoder_only_qwen3_reranker(): - handler, manager = _handler( - [0.5], - chat_template='The answer can only be "yes" or "no"', - model_path="Qwen/Qwen3-Reranker-0.6B", - ) +def test_rejects_decoder_only_qwen3_reranker_at_startup(): + manager = _TokenizerManager([0.5], model_path="Qwen/Qwen3-Reranker-0.6B") with pytest.raises(ValueError, match="cross-encoder rerankers only"): + RerankWorkerHandler(_Engine(manager), SimpleNamespace()) + assert manager.requests == [] + + +@pytest.mark.asyncio +async def test_embedding_endpoint_does_not_dispatch_rerank(): + handler = EmbeddingWorkerHandler.__new__(EmbeddingWorkerHandler) + with pytest.raises(ValueError): _ = [ - output - async for output in handler.generate( - {"model": "m", "query": "query", "documents": ["doc"]}, - _Context(), + item + async for item in handler.generate( + {"model": "m", "query": "query", "documents": ["doc"]}, _Context() ) ] - assert manager.requests == [] @pytest.mark.asyncio -async def test_embedding_endpoint_dispatches_rerank_shape(): - class _RerankHandler: - async def generate(self, request, context): - yield [{"score": 0.7, "index": 0}] +@pytest.mark.parametrize("rerank", [False, True]) +@pytest.mark.parametrize("fail_startup", [False, True]) +async def test_dedicated_pooling_registration_health_and_cleanup( + monkeypatch, rerank, fail_startup +): + module = pooling_init + monkeypatch.delenv("DYN_HEALTH_CHECK_PAYLOAD", raising=False) - handler = EmbeddingWorkerHandler.__new__(EmbeddingWorkerHandler) - handler.rerank_handler = _RerankHandler() - output = [ - item - async for item in handler.generate( - {"model": "m", "query": "query", "documents": ["doc"]}, - _Context(), - ) - ] - assert output == [[{"score": 0.7, "index": 0}]] + engine = Mock() + engine.tokenizer_manager.tokenizer.bos_token_id = 1 + endpoint = Mock() + endpoint.serve_endpoint = AsyncMock() + runtime = Mock() + runtime.endpoint.return_value = endpoint + config = SimpleNamespace( + server_args=SimpleNamespace(served_model_name="pooler", model_path="pooler"), + dynamo_args=SimpleNamespace( + namespace="test", + component="rerank" if rerank else "backend", + endpoint="generate", + use_sglang_tokenizer=False, + ), + ) + handler = Mock() + handlers = [Mock(return_value=handler), Mock(return_value=handler)] + monkeypatch.setattr(module, "EmbeddingWorkerHandler", handlers[0]) + monkeypatch.setattr(module, "RerankWorkerHandler", handlers[1]) + monkeypatch.setattr(module.sgl, "Engine", Mock(return_value=engine)) + metrics_task = asyncio.create_task(asyncio.Event().wait()) + monkeypatch.setattr( + module, "setup_sgl_metrics", AsyncMock(return_value=(None, metrics_task, [])) + ) + register = AsyncMock() + monkeypatch.setattr(module, "register_model_with_readiness_gate", register) + monkeypatch.setattr(module, "register_model_taint_route", Mock()) + monkeypatch.setattr(module, "register_engine_metrics_callback", Mock()) + monkeypatch.setattr(module, "init_embedding_metrics", Mock()) + deferred = AsyncMock() + shutdown_endpoints = [] + + init = init_rerank if rerank else module.init_embedding + if fail_startup: + handlers[int(rerank)].side_effect = ValueError("unsupported model") + with pytest.raises(ValueError, match="unsupported model"): + await init(runtime, config, asyncio.Event(), shutdown_endpoints, deferred) + register.assert_not_called() + endpoint.serve_endpoint.assert_not_called() + engine.shutdown.assert_called_once() + assert metrics_task.cancelled() + deferred.assert_awaited_once() + return + await init(runtime, config, asyncio.Event(), shutdown_endpoints, deferred) + + assert register.call_args.kwargs["output_type"] == ( + ModelType.Rerank if rerank else ModelType.Embedding + ) + assert shutdown_endpoints == [endpoint] + handlers[int(rerank)].assert_called_once() + handlers[int(not rerank)].assert_not_called() + assert endpoint.serve_endpoint.call_args.args[0] == handler.generate + payload = endpoint.serve_endpoint.call_args.kwargs["health_check_payload"] + if rerank: + assert payload["query"] and payload["documents"] + assert "input" not in payload + else: + assert payload["input"] == [1] + assert "query" not in payload + handler.cleanup.assert_called_once() + assert metrics_task.cancelled() + deferred.assert_awaited_once() + + +def test_rerank_flag_defaults_and_rejects_embedding_combination(monkeypatch): + monkeypatch.delenv("DYN_SGL_RERANK_WORKER", raising=False) + parser = argparse.ArgumentParser() + DynamoSGLangArgGroup().add_arguments(parser) + assert not parser.parse_args([]).rerank_worker + args = parser.parse_args(["--rerank-worker"]) + config = DynamoSGLangConfig.from_cli_args(args) + config.validate() + assert config.rerank_worker and not config.embedding_worker + args = parser.parse_args(["--rerank-worker", "--embedding-worker"]) + with pytest.raises(ValueError, match="cannot be combined"): + DynamoSGLangConfig.from_cli_args(args).validate() + + +@pytest.mark.asyncio +async def test_rerank_health_check_runs_cross_encoder(monkeypatch): + monkeypatch.delenv("DYN_HEALTH_CHECK_PAYLOAD", raising=False) + handler, manager = _handler([0.5]) + payload = SglangRerankHealthCheckPayload("pooler").to_dict() + outputs = [item async for item in handler.generate(payload, _Context())] + assert outputs[0][0]["score"] == 0.5 + assert manager.requests[0][0].is_cross_encoder_request diff --git a/components/src/dynamo/sglang/tests/test_sglang_unit.py b/components/src/dynamo/sglang/tests/test_sglang_unit.py index df37ed92e866..8f43e97877b7 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_unit.py +++ b/components/src/dynamo/sglang/tests/test_sglang_unit.py @@ -1010,6 +1010,7 @@ async def test_invalid_fpm_trace_is_disabled_by_arg_parser( ("overrides", "role"), [ ({"embedding_worker": True}, "embedding"), + ({"rerank_worker": True}, "rerank"), ({"multimodal_encode_worker": True}, "dedicated multimodal"), ({"multimodal_worker": True}, "dedicated multimodal"), ({"image_diffusion_worker": True}, "image diffusion"), diff --git a/examples/backends/sglang/launch/agg_rerank.sh b/examples/backends/sglang/launch/agg_rerank.sh new file mode 100755 index 000000000000..beedefec4c08 --- /dev/null +++ b/examples/backends/sglang/launch/agg_rerank.sh @@ -0,0 +1,75 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Dedicated text cross-encoder reranking. Upgrade the frontend before starting this worker. +# 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_sglang_gpu_mem_args +source "$SCRIPT_DIR/../../../common/launch_utils.sh" # print_launch_banner, wait_any_exit + +# Default values +MODEL="BAAI/bge-reranker-v2-m3" + +# Parse command line arguments +EXTRA_ARGS=() +while [[ $# -gt 0 ]]; do + case $1 in + --model-path) + MODEL="$2" + shift 2 + ;; + -h|--help) + echo "Usage: $0 [OPTIONS]" + echo "Options:" + echo " --model-path Specify model (default: $MODEL)" + echo " -h, --help Show this help message" + echo "" + echo "Note: System metrics are enabled by default on port 8081 (worker)" + exit 0 + ;; + *) + EXTRA_ARGS+=("$1") + shift + ;; + esac +done + +GPU_MEM_ARGS=$(build_sglang_gpu_mem_args) + +HTTP_PORT="${DYN_HTTP_PORT:-8000}" +print_launch_banner --no-curl "Launching Rerank Worker" "$MODEL" "$HTTP_PORT" + +print_curl_footer <(rerank_wire).is_err()); + assert!(serde_json::from_value::(wire).is_ok()); + + card.model_type = ModelType::Embedding | ModelType::Rerank; + assert!(serde_json::from_value::(serde_json::to_value(card).unwrap()).is_err()); + } + #[test] fn realtime_bit_position() { assert_eq!(ModelType::Realtime.bits(), 1 << 8); From 15aa12d3b22fdbea1e5cafca75b7e552eccc75fb Mon Sep 17 00:00:00 2001 From: xianlubird Date: Thu, 10 Sep 2026 10:26:23 +0800 Subject: [PATCH 4/4] test(sglang): adapt pooling mocks to resolved server args Signed-off-by: xianlubird --- .../src/dynamo/sglang/tests/test_sglang_rerank_handler.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/components/src/dynamo/sglang/tests/test_sglang_rerank_handler.py b/components/src/dynamo/sglang/tests/test_sglang_rerank_handler.py index 1c6b0f53452e..2d9c576ab28f 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_rerank_handler.py +++ b/components/src/dynamo/sglang/tests/test_sglang_rerank_handler.py @@ -186,6 +186,9 @@ async def test_dedicated_pooling_registration_health_and_cleanup( monkeypatch.delenv("DYN_HEALTH_CHECK_PAYLOAD", raising=False) engine = Mock() + engine.server_args = SimpleNamespace( + served_model_name="resolved-pooler", model_path="pooler" + ) engine.tokenizer_manager.tokenizer.bos_token_id = 1 endpoint = Mock() endpoint.serve_endpoint = AsyncMock() @@ -193,6 +196,7 @@ async def test_dedicated_pooling_registration_health_and_cleanup( runtime.endpoint.return_value = endpoint config = SimpleNamespace( server_args=SimpleNamespace(served_model_name="pooler", model_path="pooler"), + use_resolved_server_args=Mock(return_value=engine.server_args), dynamo_args=SimpleNamespace( namespace="test", component="rerank" if rerank else "backend", @@ -230,6 +234,8 @@ async def test_dedicated_pooling_registration_health_and_cleanup( return await init(runtime, config, asyncio.Event(), shutdown_endpoints, deferred) + config.use_resolved_server_args.assert_called_once_with(engine.server_args) + assert register.call_args.args[2] is engine.server_args assert register.call_args.kwargs["output_type"] == ( ModelType.Rerank if rerank else ModelType.Embedding ) @@ -238,6 +244,7 @@ async def test_dedicated_pooling_registration_health_and_cleanup( handlers[int(not rerank)].assert_not_called() assert endpoint.serve_endpoint.call_args.args[0] == handler.generate payload = endpoint.serve_endpoint.call_args.kwargs["health_check_payload"] + assert payload["model"] == "resolved-pooler" if rerank: assert payload["query"] and payload["documents"] assert "input" not in payload