feat: add VLM/image endpoints, harden KServe V2 transport, and expand test coverage

Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>

Author: ajcasagrande

docs/dev/adding-grpc-endpoints.md

@@ -11,7 +11,7 @@ This guide explains how to add support for new gRPC-based inference protocols in
 
 AIPerf separates **endpoints** (payload formatting and response parsing), **transports** (wire protocol and connection management), and **serializers** (proto-specific byte conversion):
 
-```
+```text
 InferenceClient
   |
   |-- Endpoint (format_payload / parse_response)
@@ -38,7 +38,7 @@ InferenceClient
 |---|---|---|
 | Proto definitions (`.proto`) | Yes | - |
 | Serializer class (dict <-> protobuf bytes) | Yes | - |
-| `payload_converter.py` (dict <-> protobuf objects) | Yes | - |
+| Payload converter (dict <-> protobuf objects) | Yes | - |
 | Endpoint class (`format_payload` / `parse_response`) | Yes | - |
 | `GrpcTransport` (timing, tracing, cancellation) | - | Yes |
 | `GenericGrpcClient` (raw bytes over gRPC) | - | Yes |
@@ -159,9 +159,9 @@ Usage: `aiperf profile --endpoint-type my_v2_endpoint --url grpc://triton:8001 .
 
 Create your `.proto` file and generate stubs:
 
-```
+```text
 src/aiperf/transports/grpc/proto/my_service.proto
-src/aiperf/transports/grpc/proto/my_service_pb2.py      (generated)
+src/aiperf/transports/grpc/proto/my_service_pb2.py       (generated)
 src/aiperf/transports/grpc/proto/my_service_pb2_grpc.py  (generated)
 ```
 
@@ -460,5 +460,5 @@ When adding a new endpoint reusing an existing protocol (Strategy A):
 - [Source: GrpcTransport](../../src/aiperf/transports/grpc/grpc_transport.py) -- Generic transport implementation
 - [Source: GenericGrpcClient](../../src/aiperf/transports/grpc/grpc_client.py) -- Proto-free gRPC client
 - [Source: KServeV2GrpcSerializer](../../src/aiperf/transports/grpc/kserve_v2_serializers.py) -- Reference serializer implementation
-- [Source: payload_converter](../../src/aiperf/transports/grpc/payload_converter.py) -- V2 dict/protobuf conversion
+- [Source: KServeV2GrpcSerializer](../../src/aiperf/transports/grpc/kserve_v2_serializers.py) -- V2 dict/protobuf conversion
 - [Source: InferenceClient](../../src/aiperf/workers/inference_client.py) -- Transport/endpoint wiring

docs/tutorials/grpc-transport.md

@@ -79,7 +79,7 @@ GenericGrpcClient                   -- proto-free, sends/receives raw bytes via
 Triton / TRT-LLM Server
 ```
 
-The endpoint never knows it's running over gRPC. The serializer (e.g., `KServeV2GrpcSerializer`) converts the endpoint's dict payload to protobuf bytes on the way out, and converts response bytes back to a JSON-serialized `TextResponse` on the way in. This means all existing `--extra-inputs` options (like `v2_input_name`, `v2_output_name`) work identically over gRPC.
+The endpoint never knows it's running over gRPC. The serializer (e.g., `KServeV2GrpcSerializer`) converts the endpoint's dict payload to protobuf bytes on the way out, and converts response bytes back to a V2 JSON-format dict on the way in. The transport layer then wraps this dict as a `TextResponse`. This means all existing `--extra-inputs` options (like `v2_input_name`, `v2_output_name`) work identically over gRPC.
 
 The serializer class and gRPC method paths are declared in `plugins.yaml` endpoint metadata, so adding support for a new gRPC protocol requires only a new serializer — no transport changes.
 
@@ -326,5 +326,5 @@ When choosing between HTTP and gRPC for V2 inference:
 - [Source: grpc_transport.py](../../src/aiperf/transports/grpc/grpc_transport.py) - Generic transport implementation
 - [Source: grpc_client.py](../../src/aiperf/transports/grpc/grpc_client.py) - Proto-free gRPC client
 - [Source: kserve_v2_serializers.py](../../src/aiperf/transports/grpc/kserve_v2_serializers.py) - KServe V2 serializer
-- [Source: payload_converter.py](../../src/aiperf/transports/grpc/payload_converter.py) - Dict/protobuf conversion
+- [Source: kserve_v2_serializers.py](../../src/aiperf/transports/grpc/kserve_v2_serializers.py) - V2 dict/protobuf conversion
 - [Source: status_mapping.py](../../src/aiperf/transports/grpc/status_mapping.py) - gRPC to HTTP status mapping

docs/tutorials/kserve.md

@@ -9,14 +9,16 @@ AIPerf provides first-class support for benchmarking [KServe](https://kserve.git
 
 ## Endpoint Types
 
-AIPerf provides five KServe-specific endpoint types:
+AIPerf provides seven KServe-specific endpoint types:
 
 | Endpoint Type | Protocol | URL Path | Streaming | Token Metrics | Use Case |
 |---|---|---|---|---|---|
 | `kserve_chat` | OpenAI-compatible | `/openai/v1/chat/completions` | Yes | Yes | LLMs via vLLM/TRT-LLM on KServe |
 | `kserve_completions` | OpenAI-compatible | `/openai/v1/completions` | Yes | Yes | Text completions via vLLM/TRT-LLM on KServe |
 | `kserve_embeddings` | OpenAI-compatible | `/openai/v1/embeddings` | No | No | Embedding models on KServe |
 | `kserve_v2_infer` | V2 Open Inference Protocol | `/v2/models/{model_name}/infer` | Yes (gRPC) | Yes | Triton/TRT-LLM tensor inference |
+| `kserve_v2_embeddings` | V2 Open Inference Protocol | `/v2/models/{model_name}/infer` | No | No | Triton/TRT-LLM embedding models |
+| `kserve_v2_rankings` | V2 Open Inference Protocol | `/v2/models/{model_name}/infer` | No | No | Triton/TRT-LLM reranking models |
 | `kserve_v1_predict` | V1 TensorFlow Serving | `/v1/models/{model_name}:predict` | No | No | Legacy TF Serving-style models |
 
 **Token Metrics**: When "Yes", AIPerf computes token-based metrics (input/output token counts, tokens per second). When "No", only request-level metrics (latency, throughput) are available.
@@ -353,6 +355,8 @@ Each KServe endpoint type includes a `health_path` in its metadata for pre-fligh
 | `kserve_completions` | `/openai/v1/models` | Lists available OpenAI-compatible models |
 | `kserve_embeddings` | `/openai/v1/models` | Lists available OpenAI-compatible models |
 | `kserve_v2_infer` | `/v2/models/{model_name}/ready` | V2 model readiness check |
+| `kserve_v2_embeddings` | `/v2/models/{model_name}/ready` | V2 model readiness check |
+| `kserve_v2_rankings` | `/v2/models/{model_name}/ready` | V2 model readiness check |
 | `kserve_v1_predict` | `/v1/models/{model_name}` | V1 model metadata/status |
 
 Health paths that contain `{model_name}` are resolved using the same template substitution as endpoint paths.
@@ -368,6 +372,8 @@ Health paths that contain `{model_name}` are resolved using the same template su
 | KServe + vLLM (embeddings) | `kserve_embeddings` | Vector embeddings |
 | KServe + Triton (text) | `kserve_v2_infer` | Wraps text as BYTES tensors |
 | KServe + TRT-LLM via Triton | `kserve_v2_infer` | Standard Triton text pipeline |
+| KServe + Triton (embeddings) | `kserve_v2_embeddings` | V2 BYTES tensor embedding models |
+| KServe + Triton (reranking) | `kserve_v2_rankings` | V2 BYTES tensor reranking models |
 | KServe + TF Serving model | `kserve_v1_predict` | Legacy instance-based format |
 | KServe + custom model server | `kserve_v1_predict` or `template` | Depends on API format |
 | Non-KServe vLLM/TRT-LLM | `chat` or `completions` | Use standard endpoints for direct deployments |

src/aiperf/endpoints/base_rankings_endpoint.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
 
 from abc import abstractmethod
 from typing import Any
@@ -98,7 +99,7 @@ def format_payload(self, request_info: RequestInfo) -> dict[str, Any]:
         turn = request_info.turns[0]
         model_endpoint = request_info.model_endpoint
 
-        if turn.max_tokens:
+        if turn.max_tokens is not None:
             self.warning("Max_tokens is provided but is not supported for rankings.")
 
         query_text, passage_texts = self._extract_query_and_passages(turn)

src/aiperf/endpoints/kserve_v2_embeddings.py

@@ -93,14 +93,10 @@ def parse_response(
         if not outputs:
             return None
 
-        # Find output tensor by name, fallback to first
-        output = None
-        for o in outputs:
-            if o.get("name") == self._output_name:
-                output = o
-                break
-        if output is None:
-            output = outputs[0]
+        # Find the output tensor with the matching name
+        output = next(
+            (o for o in outputs if o.get("name") == self._output_name), outputs[0]
+        )
 
         data = output.get("data")
         if not data:

src/aiperf/endpoints/kserve_v2_images.py

@@ -124,14 +124,10 @@ def parse_response(
         if not outputs:
             return None
 
-        # Find output tensor by name, fallback to first
-        output = None
-        for o in outputs:
-            if o.get("name") == self._output_name:
-                output = o
-                break
-        if output is None:
-            output = outputs[0]
+        # Find the output tensor with the matching name
+        output = next(
+            (o for o in outputs if o.get("name") == self._output_name), outputs[0]
+        )
 
         data = output.get("data")
         if not isinstance(data, list) or not data:

src/aiperf/endpoints/kserve_v2_infer.py

@@ -13,6 +13,67 @@
 from aiperf.endpoints.base_endpoint import BaseEndpoint
 
 
+def _extract_v2_text(output: dict[str, Any]) -> str | None:
+    """Extract text from a V2 BYTES output tensor.
+
+    Args:
+        output: V2 output tensor dict with ``data`` key.
+
+    Returns:
+        First data element as string, or None if empty.
+    """
+    data = output.get("data")
+    if isinstance(data, list) and len(data) > 0 and data[0] is not None:
+        return str(data[0])
+    return None
+
+
+def parse_v2_text_response(
+    endpoint: BaseEndpoint,
+    response: InferenceServerResponse,
+    output_name: str,
+) -> ParsedResponse | None:
+    """Parse V2 inference response, extracting text from BYTES output tensor.
+
+    Shared by KServeV2InferEndpoint and KServeV2VLMEndpoint since both
+    produce text output in the same tensor format.
+
+    Args:
+        endpoint: Endpoint instance (for make_text_response_data).
+        response: Raw response from inference server.
+        output_name: Expected output tensor name.
+
+    Returns:
+        Parsed response with extracted text content, or None if no content.
+    """
+    json_obj = response.get_json()
+    if not json_obj:
+        return None
+
+    outputs = json_obj.get("outputs")
+    if not outputs:
+        return None
+
+    for output in outputs:
+        if output.get("name") == output_name:
+            text = _extract_v2_text(output)
+            if text is not None:
+                return ParsedResponse(
+                    perf_ns=response.perf_ns,
+                    data=endpoint.make_text_response_data(text),
+                )
+
+    for output in outputs:
+        text = _extract_v2_text(output)
+        if text is not None:
+            return ParsedResponse(
+                perf_ns=response.perf_ns,
+                data=endpoint.make_text_response_data(text),
+            )
+
+    return None
+
+
 class KServeV2InferEndpoint(BaseEndpoint):
     """KServe V2 Open Inference Protocol endpoint for Triton/TRT-LLM.
 
@@ -91,32 +152,4 @@ def parse_response(
         Returns:
             Parsed response with extracted text content, or None if no content
         """
-        json_obj = response.get_json()
-        if not json_obj:
-            return None
-
-        outputs = json_obj.get("outputs")
-        if not outputs:
-            return None
-
-        for output in outputs:
-            if output.get("name") == self._output_name:
-                data = output.get("data")
-                if isinstance(data, list) and len(data) > 0 and data[0] is not None:
-                    text = str(data[0])
-                    return ParsedResponse(
-                        perf_ns=response.perf_ns,
-                        data=self.make_text_response_data(text),
-                    )
-
-        # Fallback: try first output with data
-        for output in outputs:
-            data = output.get("data")
-            if isinstance(data, list) and len(data) > 0 and data[0] is not None:
-                text = str(data[0])
-                return ParsedResponse(
-                    perf_ns=response.perf_ns,
-                    data=self.make_text_response_data(text),
-                )
-
-        return None
+        return parse_v2_text_response(self, response, self._output_name)

src/aiperf/endpoints/kserve_v2_rankings.py

@@ -114,16 +114,19 @@ def extract_rankings(self, json_obj: dict[str, Any]) -> list[dict[str, Any]]:
         if not outputs:
             return []
 
-        output = None
-        for o in outputs:
-            if o.get("name") == self._output_name:
-                output = o
-                break
-        if output is None:
-            output = outputs[0]
+        # Find the output tensor with the matching name
+        output = next(
+            (o for o in outputs if o.get("name") == self._output_name), outputs[0]
+        )
 
         data = output.get("data")
         if not data:
             return []
 
-        return [{"index": i, "score": float(s)} for i, s in enumerate(data)]
+        results = []
+        for i, s in enumerate(data):
+            try:
+                results.append({"index": i, "score": float(s)})
+            except (ValueError, TypeError):
+                self.warning(f"Skipping non-numeric score at index {i}: {s!r}")
+        return results

src/aiperf/endpoints/kserve_v2_vlm.py

@@ -11,6 +11,7 @@
     RequestInfo,
 )
 from aiperf.endpoints.base_endpoint import BaseEndpoint
+from aiperf.endpoints.kserve_v2_infer import parse_v2_text_response
 
 
 class KServeV2VLMEndpoint(BaseEndpoint):
@@ -108,32 +109,4 @@ def parse_response(
         Returns:
             Parsed response with extracted text content, or None if no content
         """
-        json_obj = response.get_json()
-        if not json_obj:
-            return None
-
-        outputs = json_obj.get("outputs")
-        if not outputs:
-            return None
-
-        for output in outputs:
-            if output.get("name") == self._output_name:
-                data = output.get("data")
-                if isinstance(data, list) and len(data) > 0 and data[0] is not None:
-                    text = str(data[0])
-                    return ParsedResponse(
-                        perf_ns=response.perf_ns,
-                        data=self.make_text_response_data(text),
-                    )
-
-        # Fallback: try first output with data
-        for output in outputs:
-            data = output.get("data")
-            if isinstance(data, list) and len(data) > 0 and data[0] is not None:
-                text = str(data[0])
-                return ParsedResponse(
-                    perf_ns=response.perf_ns,
-                    data=self.make_text_response_data(text),
-                )
-
-        return None
+        return parse_v2_text_response(self, response, self._output_name)

src/aiperf/transports/base_transports.py

@@ -6,7 +6,7 @@
 import importlib.metadata as importlib_metadata
 from abc import ABC, abstractmethod
 from collections.abc import Awaitable, Callable
-from typing import Protocol, runtime_checkable
+from typing import Any, Protocol, runtime_checkable
 from urllib.parse import parse_qs, urlencode, urlparse, urlunparse
 
 from aiperf.common.mixins import AIPerfLifecycleMixin
@@ -39,7 +39,7 @@
 class TransportProtocol(AIPerfLifecycleProtocol, Protocol):
     """Protocol for a transport that sends requests to an inference server."""
 
-    def __init__(self, **kwargs) -> None: ...
+    def __init__(self, **kwargs: Any) -> None: ...
 
     def get_transport_headers(self, request_info: RequestInfo) -> dict[str, str]: ...
 
@@ -60,7 +60,7 @@ class BaseTransport(AIPerfLifecycleMixin, ABC):
     Transports handle the protocol layer (HTTP, gRPC, etc.).
     """
 
-    def __init__(self, model_endpoint: ModelEndpointInfo, **kwargs) -> None:
+    def __init__(self, model_endpoint: ModelEndpointInfo, **kwargs: Any) -> None:
         super().__init__(**kwargs)
         self.model_endpoint: ModelEndpointInfo = model_endpoint
         self.user_agent: str = f"aiperf/{importlib_metadata.version('aiperf')}"
@@ -162,7 +162,7 @@ async def send_request(
         Args:
             request_info: Request context and metadata
             payload: Request payload (format depends on transport)
-            first_token_callback: Optional callback fired on first SSE message with ttft_ns
+            first_token_callback: Optional callback fired on first response with ttft_ns
 
         Returns:
             Record containing responses, timing, and any errors