feat(server): serve text reranking with AIPerf

CMakeLists.txt

@@ -529,6 +529,13 @@ if(TRTMC_BUILD_TESTS)
     target_compile_options(test_server_embedding_worker PRIVATE -Wall -Wextra -Wpedantic)
     add_test(NAME server_embedding_worker COMMAND test_server_embedding_worker)
     set_tests_properties(server_embedding_worker PROPERTIES LABELS cpu)
+    add_executable(test_server_ranking_worker apps/server/tests/test_ranking_worker.cpp)
+    target_include_directories(test_server_ranking_worker PRIVATE
+      ${PROJECT_SOURCE_DIR}/apps ${PROJECT_SOURCE_DIR}/core/runtime/include)
+    target_link_libraries(test_server_ranking_worker PRIVATE trtmc_server nlohmann_json::nlohmann_json)
+    target_compile_options(test_server_ranking_worker PRIVATE -Wall -Wextra -Wpedantic)
+    add_test(NAME server_ranking_worker COMMAND test_server_ranking_worker)
+    set_tests_properties(server_ranking_worker PROPERTIES LABELS cpu)
   endif()
 
   add_library(trtmc_test_cli_fixture SHARED apps/cli/tests/test_family_cli.cpp)

apps/aiperf_qual/README.md

@@ -356,16 +356,22 @@ modes require their own evidence.
 | Profile | Endpoint / AIPerf type | Qualified workload | Evidence |
 | --- | --- | --- | --- |
 | `bge-small-en-v1.5` | `/v1/embeddings` / `embeddings` | One sentence per request, FP16, context 512, concurrency 1; five fixed example sentences, 5 warmups and 50 measured requests | [Stage 1](https://github.com/NVIDIA/TensorRT-Model-Connect/issues/1668) |
+| `nemotron-rerank-vl-1b-v2` | `/v1/ranking` / `nim_rankings` | One query plus two fixed text passages, FP16, context 685, concurrency 1; 5 warmups and 50 measured requests | [Stage 2](https://github.com/NVIDIA/TensorRT-Model-Connect/issues/1669) |
 
 The embedding pilot passed the unchanged bert family correctness checks and a
 real published AIPerf 0.13.0 run on one GB300 with TensorRT 11.1.0.106. All
-55 responses contained complete 384-value vectors; client/server request IDs
-and timing exports were matched. Additional epic stages remain planned until
-their native correctness and actual client runs are complete.
+55 responses contained complete 384-value vectors.
+
+The reranking pilot passed the unchanged eagle_vlm family correctness checks
+and the same published client run configuration. All 55 responses retained
+real finite scores and original passage indices. Both pilots matched all
+client/server request IDs and timing exports. Additional epic stages remain
+planned until their native correctness and actual client runs are complete.
 
 Use [Profile Embeddings with AIPerf](../../website/docs/user-guides/profile-embeddings-with-aiperf.md)
-for the short recipe. Native Task-call wall time stays separate from client HTTP
-latency; this small run is an integration example, not a performance gate.
+or [Profile Reranking with AIPerf](../../website/docs/user-guides/profile-reranking-with-aiperf.md)
+for the short recipes. Native Task-call wall time stays separate from client
+HTTP latency; these small runs are integration examples, not performance gates.
 
 ## Text profiling with the persistent server
 

apps/aiperf_qual/examples/ranking-inputs.jsonl

@@ -0,0 +1,3 @@
+{"texts":[{"name":"query","contents":["What is the capital of France?"]},{"name":"passages","contents":["Paris is the capital of France.","Bananas are yellow fruit."]}]}
+{"texts":[{"name":"query","contents":["What does a sentence embedding represent?"]},{"name":"passages","contents":["A sentence embedding is a numeric representation of text.","An oven is used to bake bread."]}]}
+{"texts":[{"name":"query","contents":["What does AIPerf measure?"]},{"name":"passages","contents":["AIPerf measures latency and throughput of inference requests.","Rain falls from clouds."]}]}

apps/perf_serving/tests/test_protocol.py

@@ -82,6 +82,16 @@ def test_stop_sequences_truncate_text():
         openai.stop_sequences({"stop": [""]})
 
 
+@pytest.mark.parametrize("body", [
+    {"query": {"text": "q"}, "passages": []},
+    {"query": {"text": "q"}, "passages": [{"text": "a"}], "truncate": "END"},
+    {"query": {"text": 1}, "passages": [{"text": "a"}]},
+])
+def test_ranking_rejects_unsupported_truncation_and_non_text_inputs(body):
+    with pytest.raises(openai.RequestError):
+        openai.ranking(body, {})
+
+
 def test_ranking_and_image_generation_mapping():
     ranking = openai.ranking({"query": {"text": "q"}, "passages": [{"text": "a"}, {"text": "b"}]}, {})
     assert ranking == {"query": "q", "documents": ["a", "b"]}

apps/perf_serving/tests/test_trtmc_backend.py

@@ -109,3 +109,13 @@ def test_isolate_requests_uses_a_fresh_worker_per_request(tmp_path, worker):
         assert backend.describe()["isolate_requests"] is True
     finally:
         backend.close()
+
+
+def test_ranking_transport_is_complete_without_full_logging(tmp_path, worker):
+    backend = TrtmcWorkerBackend(worker=worker, session={**session(), "operation": "rerank"},
+                                scratch=tmp_path, full_observations=False)
+    try:
+        result = backend.invoke({"query": "q", "documents": ["a", "b"]}, tmp_path / "out")
+        assert result.observation["full_observation"] is True
+    finally:
+        backend.close()

apps/perf_serving/trtmc_perf_serving/backends/trtmc.py

@@ -87,8 +87,8 @@ def invoke(self, request: Mapping[str, Any], artifact_base: Path) -> Invocation:
         self._served += 1
         message_id = uuid.uuid4().hex
         self._write({"id": message_id, "request": dict(request), "artifact_base": str(artifact_base),
-                     # API vectors must survive transport regardless of log compaction.
-                     "full_observation": self._full_observations or self.operation in {"encode", "embed"}})
+                     # API vectors and scores must survive transport regardless of log compaction.
+                     "full_observation": self._full_observations or self.operation in {"encode", "embed", "rerank"}})
         reply = self._read(self._request_timeout_s, "worker request")
         if reply.get("id") != message_id:
             self._retire()

apps/perf_serving/trtmc_perf_serving/openai.py

@@ -140,11 +140,16 @@ def embeddings(body: Mapping[str, Any], base: Mapping[str, Any]) -> dict[str, An
 def ranking(body: Mapping[str, Any], base: Mapping[str, Any]) -> dict[str, Any]:
     """NIM ``/v1/ranking``: {"query": {"text"}, "passages": [{"text"}]}."""
     _check_fields(body, {"query", "passages", "truncate"})
+    if body.get("truncate", "NONE") != "NONE":
+        raise RequestError("only truncate=NONE is supported")
     try:
         query = body["query"]["text"]
         documents = [passage["text"] for passage in body["passages"]]
     except (KeyError, TypeError) as error:
         raise RequestError("ranking requires query.text and passages[].text") from error
+    if (not isinstance(query, str) or not query or not documents or
+            any(not isinstance(document, str) or not document for document in documents)):
+        raise RequestError("ranking requires non-empty query and passage strings")
     return {**base, "query": query, "documents": documents}
 
 

apps/server/native_worker.cpp

@@ -318,6 +318,40 @@ int run_embedding_protocol(Invoke&& invoke, Values&& extract_values, std::istrea
         input, output);
 }
 
+template <class Invoke, class Scores>
+int run_ranking_protocol(Invoke&& invoke, Scores&& extract_scores, std::istream& input,
+                         std::ostream& output) {
+    const Json ready{
+        {"event", "ready"}, {"protocol_version", 2}, {"capabilities", Json::array({"reranking"})}};
+    return run_protocol(
+        ready,
+        {{"rerank",
+          [&](const Json& request) {
+              const auto query = string_field(request, "query");
+              const auto found = request.find("documents");
+              if (found == request.end() || !found->is_array() || found->empty() ||
+                  !std::all_of(found->begin(), found->end(), [](const Json& value) {
+                      return value.is_string() && !value.get_ref<const std::string&>().empty();
+                  }))
+                  throw ProtocolError("documents must be a non-empty list of non-empty strings");
+              const auto documents = found->get<std::vector<std::string>>();
+              const auto started = std::chrono::steady_clock::now();
+              auto result = invoke(query, documents);
+              const auto wall_ms = std::chrono::duration<double, std::milli>(
+                                       std::chrono::steady_clock::now() - started)
+                                       .count();
+              const auto scores = extract_scores(result);
+              if (scores.size() != documents.size() ||
+                  !std::all_of(scores.begin(), scores.end(),
+                               [](float score) { return std::isfinite(score); }))
+                  throw std::runtime_error("ranking Task returned invalid scores");
+              return Json{{"scores", scores},
+                          {"model_call_ms", wall_ms},
+                          {"timing_scope", "public_task_call_wall"}};
+          }}},
+        input, output);
+}
+
 int run_existing_worker(ITask& task, std::istream& input, std::ostream& output) {
     // Legacy bundle modes select the declared public Task before invocation.
     // Several families implement more C++ interfaces than a given bundle enables.
@@ -337,6 +371,17 @@ int run_existing_worker(ITask& task, std::istream& input, std::ostream& output)
             [&](const std::string& text) { return embedding->embed(text); },
             checked_embedding_values, input, output);
     }
+
+    if (std::string_view(task.task()) == IReranking::kTask) {
+        auto* ranking = dynamic_cast<IReranking*>(&task);
+        if (!ranking)
+            throw std::invalid_argument("bundle does not implement reranking");
+        return run_ranking_protocol(
+            [&](const std::string& query, const std::vector<std::string>& documents) {
+                return ranking->rerank_batch(query, documents);
+            },
+            [](const std::vector<float>& scores) { return scores; }, input, output);
+    }
     return run_text_worker(task, input, output);
 }
 
@@ -359,6 +404,20 @@ int run_model_worker(const Model& model, std::istream& input, std::ostream& outp
                                       },
                                       input, output);
     }
+
+    if (model.supports<TextQueryDocumentsToRelevance>()) {
+        const auto task = model.task<TextQueryDocumentsToRelevance>();
+        return run_ranking_protocol(
+            [&](const std::string& query, const std::vector<std::string>& documents) {
+                return task.run({query, documents});
+            },
+            [](const DocumentRelevanceResult& result) {
+                if (result.kind() != TRTMC_SCORE_LOGIT)
+                    throw std::runtime_error("NIM ranking requires raw logit scores");
+                return std::vector<float>(result.scores().begin(), result.scores().end());
+            },
+            input, output);
+    }
     return run_text_worker(model, input, output);
 }
 

apps/server/python/tests/test_embeddings.py

@@ -121,6 +121,7 @@ def test_capability_admission_limits_and_rejection_records(tmp_path):
 
 @pytest.mark.parametrize("updates", [
     {"values": {"$array_length": 384}}, {"dim": 383}, {"values": [True], "dim": 1},
+    {"values": [10**400], "dim": 1}, {"model_call_ms": 10**400},
     {"model_call_ms": -1}, {"timing_scope": "unknown"},
 ])
 def test_invalid_native_output_is_rejected_and_retires_lane(updates):

apps/server/python/tests/test_rankings.py

@@ -0,0 +1,146 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from concurrent.futures import Future
+import json
+
+from fastapi.testclient import TestClient
+import pytest
+
+from trtmc_server.app import ServerConfig, create_app
+from trtmc_server.errors import ModelNotFoundError, WorkerProtocolError, WorkerSaturatedError
+from trtmc_server.registry import ModelRegistry
+
+
+class Session:
+    replica_name = "ranker-1"
+
+    def __init__(self, registry):
+        self.registry = registry
+
+    def submit(self, operation, payload):
+        self.registry.calls.append((operation, payload))
+        future = Future()
+        future.set_result(self.registry.result)
+        return future
+
+    def retire(self, error):
+        self.registry.retired = True
+        future = Future()
+        future.set_result(None)
+        return future
+
+    def close(self):
+        self.registry.closed += 1
+
+
+class Registry:
+    def __init__(self):
+        self.calls = []
+        self.result = {"scores": [-2.0, 3.0], "model_call_ms": 1.5, "timing_scope": "public_task_call_wall"}
+        self.retired = False
+        self.closed = 0
+        self.saturated = False
+
+    def start(self):
+        pass
+
+    def close(self):
+        pass
+
+    def supports(self, model, capability):
+        if model == "missing":
+            raise ModelNotFoundError("model is not registered")
+        return model == "ranker" and capability == "reranking"
+
+    def acquire(self, model):
+        if self.saturated:
+            raise WorkerSaturatedError("busy")
+        return Session(self)
+
+
+BODY = {"model": "ranker", "query": {"text": "Exact query: café"},
+        "passages": [{"text": "first"}, {"text": "second"}]}
+
+
+def client_for(registry, records=None):
+    return TestClient(create_app(registry, ServerConfig(max_body_bytes=1024, max_prompt_bytes=128,
+                                                       max_generation_tokens=16, records=records)))
+
+
+def test_real_scores_sorted_with_original_indices_and_timing(tmp_path):
+    registry = Registry()
+    path = tmp_path / "records.jsonl"
+    with client_for(registry, path) as client:
+        response = client.post("/v1/ranking", json=BODY, headers={"X-Request-ID": "ranking-1"})
+    assert response.status_code == 200 and response.headers["X-Request-ID"] == "ranking-1"
+    assert response.json() == {"rankings": [{"index": 1, "logit": 3.0}, {"index": 0, "logit": -2.0}]}
+    assert registry.calls == [("rerank", {"query": "Exact query: café", "documents": ["first", "second"]})]
+    assert registry.closed == 1
+    record = json.loads(path.read_text())
+    assert record["request_id"] == "ranking-1" and record["model_call_ms"] == 1.5
+    assert record["timing_scope"] == "public_task_call_wall" and record["passages"] == 2
+    assert record["passage_bytes"] == [5, 6] and record["completion_tokens"] is None
+    assert "Exact query" not in path.read_text()
+
+
+def test_ties_preserve_original_passage_order():
+    registry = Registry()
+    registry.result["scores"] = [3.0, 3.0]
+    with client_for(registry) as client:
+        response = client.post("/v1/ranking", json=BODY)
+    assert [item["index"] for item in response.json()["rankings"]] == [0, 1]
+
+
+@pytest.mark.parametrize("extra", [
+    {"truncate": "END"}, {"stream": True}, {"top_n": 1}, {"passages": []},
+    {"query": {"text": ""}}, {"passages": [{"text": 1}]}, {"passages": [{"image": "x"}]},
+])
+def test_unsupported_inputs_are_rejected_before_execution(extra):
+    registry = Registry()
+    with client_for(registry) as client:
+        response = client.post("/v1/ranking", json={**BODY, **extra})
+    assert response.status_code == 400 and not registry.calls
+
+
+def test_capability_admission_limits_and_rejection_records(tmp_path):
+    registry = Registry()
+    path = tmp_path / "records.jsonl"
+    with client_for(registry, path) as client:
+        for model, status in [("missing", 404), ("text-model", 400)]:
+            assert client.post("/v1/ranking", json={**BODY, "model": model}).status_code == status
+        registry.saturated = True
+        busy = client.post("/v1/ranking", json=BODY)
+        assert busy.status_code == 429 and busy.headers["Retry-After"] == "1"
+        registry.saturated = False
+        assert client.post("/v1/ranking", json={**BODY, "query": {"text": "x" * 129}}).status_code == 413
+        assert client.post("/v1/ranking", content=b"x" * 1025).status_code == 413
+    assert not registry.calls
+    records = [json.loads(line) for line in path.read_text().splitlines()]
+    assert [row["status"] for row in records] == [404, 400, 429, 413, 413]
+    assert [row["query_bytes"] for row in records[:4]] == [18, 18, 18, 129]
+    for row in records[:4]:
+        assert row["passages"] == row["input_items"] == 2 and row["passage_bytes"] == [5, 6]
+
+
+@pytest.mark.parametrize("updates", [
+    {"scores": []}, {"scores": [True, 2.0]}, {"scores": {"$array_length": 2}},
+    {"scores": [10**400, 2.0]}, {"model_call_ms": 10**400},
+    {"model_call_ms": -1}, {"timing_scope": "unknown"},
+])
+def test_invalid_native_output_is_rejected_and_retires_lane(updates, tmp_path):
+    registry = Registry()
+    registry.result.update(updates)
+    path = tmp_path / "records.jsonl"
+    with client_for(registry, path) as client:
+        response = client.post("/v1/ranking", json=BODY)
+    assert response.status_code == 502 and registry.retired and registry.closed == 1
+    record = json.loads(path.read_text())
+    assert record["passages"] == 2 and record["query_bytes"] == 18
+    assert record["passage_bytes"] == [5, 6]
+
+
+def test_registry_accepts_ranking_capability_and_rejects_unknown_only():
+    ModelRegistry._validate_ready({"event": "ready", "protocol_version": 2, "capabilities": ["reranking"]})
+    with pytest.raises(WorkerProtocolError):
+        ModelRegistry._validate_ready({"event": "ready", "protocol_version": 2, "capabilities": ["unknown"]})

apps/server/python/trtmc_server/app.py

@@ -31,10 +31,10 @@
     WorkerTimeoutError,
 )
 from .metrics import Metrics
-from .protocol import chat_prompt, extract_embedding, extract_result, generation_config, public_worker_error
+from .protocol import chat_prompt, extract_embedding, extract_ranking, extract_result, generation_config, public_worker_error
 from .records import RequestPolicy, RequestRecords
 from .registry import ModelRegistry
-from .schemas import ChatCompletionRequest, CompletionRequest, EmbeddingRequest, GenerationRequest
+from .schemas import ChatCompletionRequest, CompletionRequest, EmbeddingRequest, GenerationRequest, RankingRequest
 from .worker import WorkerSession
 
 
@@ -66,6 +66,7 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
             "/v1/completions",
             "/v1/chat/completions",
             "/v1/embeddings",
+            "/v1/ranking",
         }:
             await self.app(scope, receive, send)
             return
@@ -568,6 +569,28 @@ async def embeddings(http_request: Request, request: EmbeddingRequest) -> Any:
         return {"object": "list", "model": request.model,
                 "data": [{"object": "embedding", "index": 0, "embedding": vector}]}
 
+    @app.post("/v1/ranking")
+    async def ranking(http_request: Request, request: RankingRequest) -> Any:
+        route = http_request.url.path
+        context = http_request.state.timing
+        context["model"] = request.model
+        documents = [passage.text for passage in request.passages]
+        context.update(input_items=len(documents), passages=len(documents),
+                       query_bytes=len(request.query.text.encode("utf-8")),
+                       passage_bytes=[len(text.encode("utf-8")) for text in documents])
+        size = context["query_bytes"] + sum(context["passage_bytes"])
+        if size > config.max_prompt_bytes:
+            metrics.reject(route, 413)
+            return error_response(413, "content_too_large", "query and passages exceed the server limit")
+        result = await execute_task(http_request, request.model, "reranking", "rerank",
+                                    {"query": request.query.text, "documents": documents},
+                                    lambda value: extract_ranking(value, len(documents)))
+        if isinstance(result, JSONResponse):
+            return result
+        scores, _ = extract_ranking(result, len(documents))
+        order = sorted(range(len(scores)), key=lambda index: -scores[index])
+        return {"rankings": [{"index": index, "logit": scores[index]} for index in order]}
+
     @app.post("/v1/chat/completions")
     async def chat_completions(http_request: Request, request: ChatCompletionRequest) -> Any:
         try: