feat: add nim image retrieval endpoint support

Author: ajcasagrande

src/aiperf/common/enums/metric_enums.py

@@ -188,6 +188,10 @@ class GenericMetricUnit(BaseMetricUnit):
     RATIO = _unit("ratio")
     USER = _unit("user")
     PERCENT = _unit("%")
+    IMAGE = _unit("image")
+    IMAGES = _unit("images")
+    VIDEO = _unit("video")
+    VIDEOS = _unit("videos")
 
 
 class PowerMetricUnitInfo(BaseMetricUnitInfo):
@@ -289,7 +293,11 @@ class MetricOverTimeUnitInfo(BaseMetricUnitInfo):
     @model_validator(mode="after")
     def _set_tag(self: Self) -> Self:
         """Set the tag based on the existing units. ie. requests/sec, tokens/sec, etc."""
-        self.tag = f"{self.primary_unit}/{self.time_unit}"
+        self.tag = (
+            f"{self.primary_unit}/{self.time_unit}"
+            if not self.inverted
+            else f"{self.time_unit}/{self.primary_unit}"
+        )
         if self.third_unit:
             # If there is a third unit, add it to the tag. ie. tokens/sec/user
             self.tag += f"/{self.third_unit}"
@@ -302,6 +310,7 @@ def _set_tag(self: Self) -> Self:
     primary_unit: "MetricUnitT"
     time_unit: MetricTimeUnit | MetricTimeUnitInfo
     third_unit: "MetricUnitT | None" = None
+    inverted: bool = False
 
     def convert_to(self, other_unit: "MetricUnitT", value: int | float) -> float:
         """Convert a value from this unit to another unit."""
@@ -342,6 +351,24 @@ class MetricOverTimeUnit(BaseMetricUnit):
         time_unit=MetricTimeUnit.SECONDS,
         third_unit=GenericMetricUnit.USER,
     )
+    IMAGES_PER_SECOND = MetricOverTimeUnitInfo(
+        primary_unit=GenericMetricUnit.IMAGES,
+        time_unit=MetricTimeUnit.SECONDS,
+    )
+    MS_PER_IMAGE = MetricOverTimeUnitInfo(
+        time_unit=MetricTimeUnit.MILLISECONDS,
+        primary_unit=GenericMetricUnit.IMAGE,
+        inverted=True,
+    )
+    VIDEOS_PER_SECOND = MetricOverTimeUnitInfo(
+        primary_unit=GenericMetricUnit.VIDEOS,
+        time_unit=MetricTimeUnit.SECONDS,
+    )
+    MS_PER_VIDEO = MetricOverTimeUnitInfo(
+        time_unit=MetricTimeUnit.MILLISECONDS,
+        primary_unit=GenericMetricUnit.VIDEO,
+        inverted=True,
+    )
 
     @cached_property
     def info(self) -> MetricOverTimeUnitInfo:
@@ -363,6 +390,11 @@ def third_unit(self) -> "MetricUnitT | None":
         """Get the third unit (if applicable)."""
         return self.info.third_unit
 
+    @cached_property
+    def inverted(self) -> bool:
+        """Whether the metric is inverted (e.g. time / metric)."""
+        return self.info.inverted
+
 
 class MetricType(CaseInsensitiveStrEnum):
     """Defines the possible types of metrics."""
@@ -643,6 +675,9 @@ class MetricFlags(Flag):
     TOKENIZES_INPUT_ONLY = 1 << 12
     """Metrics that are only applicable when the endpoint tokenizes input text."""
 
+    SUPPORTS_VIDEO_ONLY = 1 << 13
+    """Metrics that are only applicable to video-based endpoints."""
+
     def has_flags(self, flags: "MetricFlags") -> bool:
         """Return True if the metric has ALL of the given flag(s) (regardless of other flags)."""
         # Bitwise AND will return the input flags only if all of the given flags are present.

src/aiperf/common/enums/plugin_enums.py

@@ -32,6 +32,7 @@ class EndpointType(CaseInsensitiveStrEnum):
     NIM_RANKINGS = "nim_rankings"
     SOLIDO_RAG = "solido_rag"
     TEMPLATE = "template"
+    IMAGE_RETRIEVAL = "image_retrieval"
 
 
 class TransportType(CaseInsensitiveStrEnum):

src/aiperf/common/messages/base_messages.py

@@ -1,39 +1,28 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
-import json
 import time
 from typing import ClassVar
 
 from pydantic import Field
 
 from aiperf.common.aiperf_logger import AIPerfLogger
 from aiperf.common.enums.message_enums import MessageType
-from aiperf.common.models.base_models import AIPerfBaseModel, exclude_if_none
+from aiperf.common.models.base_models import AIPerfBaseModel
 from aiperf.common.models.error_models import ErrorDetails
 from aiperf.common.types import MessageTypeT
+from aiperf.common.utils import load_json_str
 
 _logger = AIPerfLogger(__name__)
 
 
-@exclude_if_none("request_ns", "request_id")
 class Message(AIPerfBaseModel):
-    """Base message class for optimized message handling. Based on the AIPerfBaseModel class,
-    so it supports @exclude_if_none decorator. see :class:`AIPerfBaseModel` for more details.
+    """Base message class for optimized message handling.
 
     This class provides a base for all messages, including common fields like message_type,
-    request_ns, and request_id. It also supports optional field exclusion based on the
-    @exclude_if_none decorator.
+    request_ns, and request_id.
 
     Each message model should inherit from this class, set the message_type field,
     and define its own additional fields.
-
-    Example:
-    ```python
-    @exclude_if_none("some_field")
-    class ExampleMessage(Message):
-        some_field: int | None = Field(default=None)
-        other_field: int = Field(default=1)
-    ```
     """
 
     _message_type_lookup: ClassVar[dict[MessageTypeT, type["Message"]]] = {}
@@ -71,7 +60,7 @@ def __get_validators__(cls):
     def from_json(cls, json_str: str | bytes | bytearray) -> "Message":
         """Deserialize a message from a JSON string, attempting to auto-detect the message type.
         NOTE: If you already know the message type, use the more performant :meth:`from_json_with_type` instead."""
-        data = json.loads(json_str)
+        data = load_json_str(json_str)
         message_type = data.get("message_type")
         if not message_type:
             raise ValueError(f"Missing message_type: {json_str}")
@@ -97,7 +86,7 @@ def from_json_with_type(
         return message_class.model_validate_json(json_str)
 
     def __str__(self) -> str:
-        return self.model_dump_json()
+        return self.model_dump_json(exclude_none=True)
 
 
 class RequiresRequestNSMixin(Message):

src/aiperf/common/messages/command_messages.py

@@ -1,6 +1,5 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
-import json
 import uuid
 from typing import Any, ClassVar
 
@@ -18,14 +17,13 @@
 from aiperf.common.models import (
     ErrorDetails,
     ProcessRecordsResult,
-    exclude_if_none,
 )
 from aiperf.common.types import CommandTypeT, MessageTypeT, ServiceTypeT
+from aiperf.common.utils import load_json_str
 
 _logger = AIPerfLogger(__name__)
 
 
-@exclude_if_none("target_service_id", "target_service_type")
 class TargetedServiceMessage(BaseServiceMessage):
     """Message that can be targeted to a specific service by id or type.
     If both `target_service_type` and `target_service_id` are None, the message is
@@ -80,7 +78,7 @@ def __init_subclass__(cls, **kwargs):
     @classmethod
     def from_json(cls, json_str: str | bytes | bytearray) -> "CommandMessage":
         """Deserialize a command message from a JSON string, attempting to auto-detect the command type."""
-        data = json.loads(json_str)
+        data = load_json_str(json_str)
         command_type = data.get("command")
         if not command_type:
             raise ValueError(f"Missing command: {json_str}")
@@ -139,7 +137,7 @@ def __init_subclass__(cls, **kwargs):
     @classmethod
     def from_json(cls, json_str: str | bytes | bytearray) -> "CommandResponse":
         """Deserialize a command response message from a JSON string, attempting to auto-detect the command response type."""
-        data = json.loads(json_str)
+        data = load_json_str(json_str)
         status = data.get("status")
         if not status:
             raise ValueError(f"Missing command response status: {json_str}")
@@ -259,7 +257,6 @@ class SpawnWorkersCommand(CommandMessage):
     num_workers: int = Field(..., description="Number of workers to spawn")
 
 
-@exclude_if_none("worker_ids", "num_workers")
 class ShutdownWorkersCommand(CommandMessage):
     command: CommandTypeT = CommandType.SHUTDOWN_WORKERS
 

src/aiperf/common/models/__init__.py

@@ -10,7 +10,6 @@
 ########################################################################
 from aiperf.common.models.base_models import (
     AIPerfBaseModel,
-    exclude_if_none,
 )
 from aiperf.common.models.credit_models import (
     CreditPhaseConfig,
@@ -69,6 +68,7 @@
     BaseInferenceServerResponse,
     BaseResponseData,
     EmbeddingResponseData,
+    ImageRetrievalResponseData,
     MetricRecordInfo,
     MetricRecordMetadata,
     MetricResult,
@@ -144,6 +144,7 @@
     "GpuTelemetrySnapshot",
     "IOCounters",
     "Image",
+    "ImageRetrievalResponseData",
     "InputsFile",
     "JsonExportData",
     "JsonMetricResult",
@@ -194,6 +195,5 @@
     "WorkerTaskStats",
     "create_balanced_distribution",
     "create_uniform_distribution",
-    "exclude_if_none",
     "logger",
 ]

src/aiperf/common/models/base_models.py

@@ -1,54 +1,14 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
-from typing import Any, ClassVar
-
-from pydantic import BaseModel, ConfigDict, model_serializer
-
-from aiperf.common.types import AIPerfBaseModelT
-
-
-def exclude_if_none(*field_names: str):
-    """Decorator to set the _exclude_if_none_fields class attribute to the set of
-    field names that should be excluded if they are None.
-    """
-
-    def decorator(model: type[AIPerfBaseModelT]) -> type[AIPerfBaseModelT]:
-        # This attribute is defined by the AIPerfBaseModel class.
-        if not hasattr(model, "_exclude_if_none_fields"):
-            model._exclude_if_none_fields = set()
-        model._exclude_if_none_fields.update(set(field_names))
-        return model
-
-    return decorator
+from pydantic import BaseModel, ConfigDict
 
 
 class AIPerfBaseModel(BaseModel):
-    """Base model for all AIPerf Pydantic models. This class is configured to allow
-    arbitrary types to be used as fields as to allow for more flexible model definitions
-    by end users without breaking the existing code.
-
-    The @exclude_if_none decorator can also be used to specify which fields
-    should be excluded from the serialized model if they are None. This is a workaround
-    for the fact that pydantic does not support specifying exclude_none on a per-field basis.
-    """
+    """Base model for all AIPerf Pydantic models.
 
-    _exclude_if_none_fields: ClassVar[set[str]] = set()
-    """Set of field names that should be excluded from the serialized model if they
-    are None. This is set by the @exclude_if_none decorator.
+    This class is configured to allow arbitrary types to be used as fields
+    to allow for more flexible model definitions by end users without breaking
+    existing code.
     """
 
-    # Allow extras by default to be more flexible for end users
     model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
-
-    @model_serializer
-    def _serialize_model(self) -> dict[str, Any]:
-        """Serialize the model to a dictionary.
-
-        This method overrides the default serializer to exclude fields that with a
-        value of None and were marked with the @exclude_if_none decorator.
-        """
-        return {
-            k: v
-            for k, v in self
-            if not (k in self._exclude_if_none_fields and v is None)
-        }

src/aiperf/common/models/dataset_models.py

@@ -6,7 +6,7 @@
 from pydantic import Field
 
 from aiperf.common.enums import MediaType
-from aiperf.common.models.base_models import AIPerfBaseModel, exclude_if_none
+from aiperf.common.models.base_models import AIPerfBaseModel
 from aiperf.common.types import MediaTypeT
 
 
@@ -45,7 +45,6 @@ class Video(Media):
     media_type: ClassVar[MediaTypeT] = MediaType.VIDEO
 
 
-@exclude_if_none("role")
 class Turn(AIPerfBaseModel):
     """A dataset representation of a single turn within a conversation.
 

src/aiperf/common/models/export_models.py

@@ -3,17 +3,14 @@
 
 from datetime import datetime
 
-from pydantic import ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field
 
 from aiperf.common.config import UserConfig
 from aiperf.common.models import ErrorDetailsCount
-from aiperf.common.models.base_models import AIPerfBaseModel, exclude_if_none
+from aiperf.common.models.base_models import AIPerfBaseModel
 
 
-@exclude_if_none(
-    "min", "max", "p1", "p5", "p10", "p25", "p50", "p75", "p90", "p95", "p99", "std"
-)
-class JsonMetricResult(AIPerfBaseModel):
+class JsonMetricResult(BaseModel):
     """The result values of a single metric for JSON export.
 
     NOTE:
@@ -70,7 +67,7 @@ class TelemetryExportData(AIPerfBaseModel):
     endpoints: dict[str, EndpointData]
 
 
-class JsonExportData(AIPerfBaseModel):
+class JsonExportData(BaseModel):
     """Summary data to be exported to a JSON file.
 
     NOTE:

src/aiperf/common/models/record_models.py

@@ -622,6 +622,18 @@ class RankingsResponseData(BaseResponseData):
     )
 
 
+class ImageRetrievalResponseData(BaseResponseData):
+    """Parsed image retrieval response data."""
+
+    data: list[dict[str, Any]] = Field(
+        ..., description="The image retrieval data from the response."
+    )
+
+    def get_text(self) -> str:
+        """Get the text of the response (empty for image retrieval)."""
+        return ""
+
+
 class ParsedResponse(AIPerfBaseModel):
     """Parsed response from a inference client."""
 
@@ -633,6 +645,7 @@ class ParsedResponse(AIPerfBaseModel):
         | TextResponseData
         | EmbeddingResponseData
         | RankingsResponseData
+        | ImageRetrievalResponseData
         | BaseResponseData
         | None
     ] = Field(

src/aiperf/common/models/sequence_distribution.py

@@ -32,12 +32,14 @@
 
 from __future__ import annotations
 
-import json
 import logging
 import re
 from dataclasses import dataclass
 
 import numpy as np
+import orjson
+
+from aiperf.common.utils import load_json_str
 
 logger = logging.getLogger(__name__)
 
@@ -337,9 +339,9 @@ def parse(cls, dist_str: str) -> SequenceLengthDistribution:
     def _parse_json_format(cls, json_str: str) -> SequenceLengthDistribution:
         """Parse JSON format: {"pairs": [{"isl": 256, "isl_stddev": 10, "osl": 128, "osl_stddev": 5, "prob": 40}, ...]}"""
         try:
-            data = json.loads(json_str)
-        except json.JSONDecodeError as e:
-            raise ValueError(f"Invalid JSON format: {e}") from None
+            data = load_json_str(json_str)
+        except orjson.JSONDecodeError as e:
+            raise ValueError(f"Invalid JSON format: {e}") from e
 
         # Validate structure outside the JSON parsing try-catch
         if "pairs" not in data: