Rerank support for vLLM in HuggingFace Serving Runtime (kserve#4376)

Signed-off-by: ayush <ayush.sawant@nutanix.com>
Signed-off-by: tmbochenski <31276485+tmbochenski@users.noreply.github.com>
Co-authored-by: tmbochenski <31276485+tmbochenski@users.noreply.github.com>

Author: AyushSawant18588

python/huggingfaceserver/huggingfaceserver/encoder_model.py

@@ -60,6 +60,8 @@
     EmbeddingRequest,
     EmbeddingResponseData,
     ErrorResponse,
+    Rerank,
+    RerankRequest,
     UsageInfo,
 )
 
@@ -439,3 +441,13 @@ async def create_embedding(
 
         except Exception as e:
             raise OpenAIError(f"Error during embedding creation: {e}") from e
+
+    async def create_rerank(
+        self,
+        request: RerankRequest,
+        raw_request: Optional[Request] = None,
+        context: Optional[Dict[str, Any]] = None,
+    ) -> Union[AsyncGenerator[str, None], Rerank, ErrorResponse]:
+        raise OpenAIError(
+            "Rerank is not implemented for Encoder model with huggingface backend"
+        )

python/huggingfaceserver/huggingfaceserver/vllm/vllm_model.py

@@ -12,37 +12,43 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional, Union, AsyncGenerator
-import torch
 from argparse import Namespace
-from fastapi import Request
+from typing import Any, Dict, Optional, Union, AsyncGenerator
 from http import HTTPStatus
 
-from kserve.protocol.rest.openai.errors import create_error_response
-from kserve.protocol.rest.openai import OpenAIEncoderModel, OpenAIGenerativeModel
-from kserve.protocol.rest.openai.types import (
-    Completion,
-    ChatCompletion,
-    CompletionRequest,
-    ChatCompletionRequest,
-    EmbeddingRequest,
-    Embedding,
-    ErrorResponse,
-)
-
-import vllm.envs as envs
+import torch
+from fastapi import Request
 from vllm import AsyncEngineArgs
+import vllm.envs as envs
 from vllm.entrypoints.logger import RequestLogger
 from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.openai.serving_completion import OpenAIServingCompletion
 from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
 from vllm.entrypoints.openai.serving_embedding import OpenAIServingEmbedding
+from vllm.entrypoints.openai.serving_score import ServingScores
 from vllm.entrypoints.openai.tool_parsers import ToolParserManager
 from vllm.entrypoints.openai.serving_models import BaseModelPath, OpenAIServingModels
 from vllm.entrypoints.openai.cli_args import validate_parsed_serve_args
 from vllm.entrypoints.chat_utils import load_chat_template
 from vllm.entrypoints.openai.protocol import ErrorResponse as engineError
 from vllm.entrypoints.openai.reasoning_parsers import ReasoningParserManager
+
+from kserve.protocol.rest.openai.errors import create_error_response
+from kserve.protocol.rest.openai import (
+    OpenAIEncoderModel,
+    OpenAIGenerativeModel,
+)
+from kserve.protocol.rest.openai.types import (
+    Completion,
+    ChatCompletion,
+    CompletionRequest,
+    ChatCompletionRequest,
+    EmbeddingRequest,
+    Embedding,
+    ErrorResponse,
+    RerankRequest,
+    Rerank,
+)
 from .utils import build_async_engine_client_from_engine_args, build_vllm_engine_args
 
 
@@ -53,7 +59,11 @@ class VLLMModel(
     vllm_engine_args: AsyncEngineArgs = None
     args: Namespace = None
     ready: bool = False
+    openai_serving_models: Optional[OpenAIServingModels] = None
     openai_serving_completion: Optional[OpenAIServingCompletion] = None
+    openai_serving_chat: Optional[OpenAIServingChat] = None
+    openai_serving_embedding: Optional[OpenAIServingEmbedding] = None
+    serving_reranking: Optional[ServingScores] = None
 
     def __init__(
         self,
@@ -68,6 +78,9 @@ def __init__(
         self.vllm_engine_args = engine_args
         self.request_logger = request_logger
         self.model_name = model_name
+        self.base_model_paths = []
+        self.log_stats = True
+        self.model_config = None
 
     async def start_engine(self):
         if self.args.tool_parser_plugin and len(self.args.tool_parser_plugin) > 3:
@@ -169,6 +182,17 @@ async def start_engine(self):
                 else None
             )
 
+            self.serving_reranking = (
+                ServingScores(
+                    self.engine_client,
+                    self.model_config,
+                    self.openai_serving_models,
+                    request_logger=self.request_logger,
+                )
+                if self.model_config.task == "score"
+                else None
+            )
+
         self.ready = True
         return self.ready
 
@@ -201,6 +225,11 @@ async def create_completion(
         raw_request: Optional[Request] = None,
         context: Optional[Dict[str, Any]] = None,
     ) -> Union[AsyncGenerator[str, None], Completion, ErrorResponse]:
+        if self.openai_serving_completion is None:
+            return create_error_response(
+                message="The model does not support Completions API",
+                status_code=HTTPStatus.BAD_REQUEST,
+            )
         response = await self.openai_serving_completion.create_completion(
             request, raw_request
         )
@@ -221,6 +250,11 @@ async def create_chat_completion(
         raw_request: Optional[Request] = None,
         context: Optional[Dict[str, Any]] = None,
     ) -> Union[AsyncGenerator[str, None], ChatCompletion, ErrorResponse]:
+        if self.openai_serving_chat is None:
+            return create_error_response(
+                message="The model does not support Chat Completions API",
+                status_code=HTTPStatus.BAD_REQUEST,
+            )
         response = await self.openai_serving_chat.create_chat_completion(
             request, raw_request
         )
@@ -241,6 +275,11 @@ async def create_embedding(
         raw_request: Optional[Request] = None,
         context: Optional[Dict[str, Any]] = None,
     ) -> Union[AsyncGenerator[str, None], Embedding, ErrorResponse]:
+        if self.openai_serving_embedding is None:
+            return create_error_response(
+                message="The model does not support Embeddings API",
+                status_code=HTTPStatus.BAD_REQUEST,
+            )
         response = await self.openai_serving_embedding.create_embedding(
             request, raw_request
         )
@@ -254,3 +293,26 @@ async def create_embedding(
             )
 
         return response
+
+    async def create_rerank(
+        self,
+        request: RerankRequest,
+        raw_request: Optional[Request] = None,
+        context: Optional[Dict[str, Any]] = None,
+    ) -> Union[AsyncGenerator[str, None], Rerank, ErrorResponse]:
+        if self.serving_reranking is None:
+            return create_error_response(
+                message="The model does not support Rerank API",
+                status_code=HTTPStatus.BAD_REQUEST,
+            )
+        response = await self.serving_reranking.do_rerank(request, raw_request)
+
+        if isinstance(response, engineError):
+            return create_error_response(
+                message=response.message,
+                err_type=response.type,
+                param=response.param,
+                status_code=HTTPStatus(response.code),
+            )
+
+        return response

python/huggingfaceserver/tests/test_vllm_rerank.py

@@ -0,0 +1,117 @@
+# Copyright 2025 The KServe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import pytest
+import requests
+
+from kserve.protocol.rest.openai.types import Rerank
+from server import RemoteOpenAIServer
+
+
+MODEL = "BAAI/bge-reranker-base"
+MODEL_NAME = "test-model"
+
+
+@pytest.fixture(scope="module")
+def server():  # noqa: F811
+    args = [
+        # use half precision for speed and memory savings in CI environment
+        "--dtype",
+        "bfloat16",
+        "--max-model-len",
+        "100",
+        "--enforce-eager",
+    ]
+
+    with RemoteOpenAIServer(MODEL, MODEL_NAME, args) as remote_server:
+        yield remote_server
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "model_name",
+    [MODEL_NAME],
+)
+async def test_rerank_texts(server: RemoteOpenAIServer, model_name: str):
+    query = "What is the capital of France?"
+    documents = [
+        "The capital of Brazil is Brasilia.",
+        "The capital of France is Paris.",
+    ]
+
+    rerank_response = requests.post(
+        server.url_for("openai/v1", "rerank"),
+        json={
+            "model": model_name,
+            "query": query,
+            "documents": documents,
+        },
+    )
+    rerank_response.raise_for_status()
+    rerank = Rerank.model_validate(rerank_response.json())
+
+    assert rerank.id is not None
+    assert rerank.results is not None
+    assert len(rerank.results) == 2
+    assert rerank.results[0].relevance_score >= 0.9
+    assert rerank.results[1].relevance_score <= 0.01
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "model_name",
+    [MODEL_NAME],
+)
+async def test_top_n(server: RemoteOpenAIServer, model_name: str):
+    query = "What is the capital of France?"
+    documents = [
+        "The capital of Brazil is Brasilia.",
+        "The capital of France is Paris.",
+        "Cross-encoder models are neat",
+    ]
+
+    rerank_response = requests.post(
+        server.url_for("openai/v1", "rerank"),
+        json={"model": model_name, "query": query, "documents": documents, "top_n": 2},
+    )
+    rerank_response.raise_for_status()
+    rerank = Rerank.model_validate(rerank_response.json())
+
+    assert rerank.id is not None
+    assert rerank.results is not None
+    assert len(rerank.results) == 2
+    assert rerank.results[0].relevance_score >= 0.9
+    assert rerank.results[1].relevance_score <= 0.01
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "model_name",
+    [MODEL_NAME],
+)
+async def test_rerank_max_model_len(server: RemoteOpenAIServer, model_name: str):
+    query = "What is the capital of France?" * 100
+    documents = [
+        "The capital of Brazil is Brasilia.",
+        "The capital of France is Paris.",
+    ]
+
+    rerank_response = requests.post(
+        server.url_for("openai/v1", "rerank"),
+        json={"model": model_name, "query": query, "documents": documents},
+    )
+    assert rerank_response.status_code == 400
+    # Assert just a small fragments of the response
+    assert "Please reduce the length of the input." in rerank_response.text