feat: add rankings prompt and query token options (#498)

Author: debermudez

docs/cli_options.md

@@ -124,6 +124,26 @@ Mean number of passages per rankings entry (per query)(default 1).
 Stddev for passages per rankings entry (default 0).
 <br>_Default: `0`_
 
+#### `--rankings-passages-prompt-token-mean` `<int>`
+
+Mean number of tokens in a passage entry for rankings (default 550).
+<br>_Default: `550`_
+
+#### `--rankings-passages-prompt-token-stddev` `<int>`
+
+Stddev for number of tokens in a passage entry for rankings (default 0).
+<br>_Default: `0`_
+
+#### `--rankings-query-prompt-token-mean` `<int>`
+
+Mean number of tokens in a query entry for rankings (default 550).
+<br>_Default: `550`_
+
+#### `--rankings-query-prompt-token-stddev` `<int>`
+
+Stddev for number of tokens in a query entry for rankings (default 0).
+<br>_Default: `0`_
+
 ## Audio Input Options
 
 #### `--audio-batch-size`, `--batch-size-audio` `<int>`

docs/tutorials/rankings.md

@@ -43,10 +43,14 @@ aiperf profile \
     --request-count 10 \
     --rankings-passages-mean 5 \
     --rankings-passages-stddev 1 \
-    --prompt-input-tokens-mean 32 \
-    --prompt-input-tokens-stddev 8
+    --rankings-passages-prompt-token-mean 32 \
+    --rankings-passages-prompt-token-stddev 8 \
+    --rankings-query-prompt-token-mean 16 \
+    --rankings-query-prompt-token-stddev 4
 ```
 
+> **Note:** The rankings-specific token options cannot be used together with `--prompt-input-tokens-mean` or `--prompt-input-tokens-stddev`. Use the rankings-specific options for controlling token counts in rankings queries and passages.
+
 ### Profile using Custom Inputs
 
 Create a file named rankings.jsonl where each line represents a ranking request with a query and one or more passages.

src/aiperf/common/config/config_defaults.py

@@ -59,6 +59,10 @@ class InputDefaults:
     NUM_DATASET_ENTRIES = 100
     RANKINGS_PASSAGES_MEAN = 1
     RANKINGS_PASSAGES_STDDEV = 0
+    RANKINGS_PASSAGES_PROMPT_TOKEN_MEAN = 550
+    RANKINGS_PASSAGES_PROMPT_TOKEN_STDDEV = 0
+    RANKINGS_QUERY_PROMPT_TOKEN_MEAN = 550
+    RANKINGS_QUERY_PROMPT_TOKEN_STDDEV = 0
 
 
 @dataclass(frozen=True)

src/aiperf/common/config/input_config.py

@@ -331,6 +331,62 @@ def validate_goodput(self) -> Self:
         ),
     ] = InputDefaults.RANKINGS_PASSAGES_STDDEV
 
+    rankings_passages_prompt_token_mean: Annotated[
+        int,
+        Field(
+            ge=1,
+            description=(
+                "Mean number of tokens in a passage entry for rankings (default 550)."
+            ),
+        ),
+        CLIParameter(
+            name=("--rankings-passages-prompt-token-mean",),
+            group=_CLI_GROUP,
+        ),
+    ] = InputDefaults.RANKINGS_PASSAGES_PROMPT_TOKEN_MEAN
+
+    rankings_passages_prompt_token_stddev: Annotated[
+        int,
+        Field(
+            ge=0,
+            description=(
+                "Stddev for number of tokens in a passage entry for rankings (default 0)."
+            ),
+        ),
+        CLIParameter(
+            name=("--rankings-passages-prompt-token-stddev",),
+            group=_CLI_GROUP,
+        ),
+    ] = InputDefaults.RANKINGS_PASSAGES_PROMPT_TOKEN_STDDEV
+
+    rankings_query_prompt_token_mean: Annotated[
+        int,
+        Field(
+            ge=1,
+            description=(
+                "Mean number of tokens in a query entry for rankings (default 550)."
+            ),
+        ),
+        CLIParameter(
+            name=("--rankings-query-prompt-token-mean",),
+            group=_CLI_GROUP,
+        ),
+    ] = InputDefaults.RANKINGS_QUERY_PROMPT_TOKEN_MEAN
+
+    rankings_query_prompt_token_stddev: Annotated[
+        int,
+        Field(
+            ge=0,
+            description=(
+                "Stddev for number of tokens in a query entry for rankings (default 0)."
+            ),
+        ),
+        CLIParameter(
+            name=("--rankings-query-prompt-token-stddev",),
+            group=_CLI_GROUP,
+        ),
+    ] = InputDefaults.RANKINGS_QUERY_PROMPT_TOKEN_STDDEV
+
     audio: AudioConfig = AudioConfig()
     image: ImageConfig = ImageConfig()
     video: VideoConfig = VideoConfig()

src/aiperf/common/config/user_config.py

@@ -20,6 +20,7 @@
 from aiperf.common.config.output_config import OutputConfig
 from aiperf.common.config.tokenizer_config import TokenizerConfig
 from aiperf.common.enums import CustomDatasetType, GPUTelemetryMode
+from aiperf.common.enums.plugin_enums import EndpointType
 from aiperf.common.enums.timing_enums import RequestRateMode, TimingMode
 from aiperf.common.utils import load_json_str
 
@@ -398,3 +399,68 @@ def validate_concurrency_limits(self) -> Self:
             )
 
         return self
+
+    @model_validator(mode="after")
+    def validate_rankings_token_options(self) -> Self:
+        """Validate rankings token options usage."""
+
+        # Check if prompt input tokens have been changed from defaults
+        prompt_tokens_modified = any(
+            field in self.input.prompt.input_tokens.model_fields_set
+            for field in ["mean", "stddev"]
+        )
+
+        # Check if any rankings-specific token options have been changed from defaults
+        rankings_token_fields = [
+            "rankings_passages_prompt_token_mean",
+            "rankings_passages_prompt_token_stddev",
+            "rankings_query_prompt_token_mean",
+            "rankings_query_prompt_token_stddev",
+        ]
+        rankings_tokens_modified = any(
+            field in self.input.model_fields_set for field in rankings_token_fields
+        )
+
+        # Check if any rankings-specific passage options have been changed from defaults
+        rankings_passages_fields = [
+            "rankings_passages_mean",
+            "rankings_passages_stddev",
+        ]
+        rankings_passages_modified = any(
+            field in self.input.model_fields_set for field in rankings_passages_fields
+        )
+
+        rankings_options_modified = (
+            rankings_tokens_modified or rankings_passages_modified
+        )
+
+        endpoint_type_is_rankings = "rankings" in self.endpoint.type.lower()
+
+        # Validate that rankings options are only used with rankings endpoints
+        rankings_endpoints = [
+            endpoint_type
+            for endpoint_type in EndpointType
+            if "rankings" in endpoint_type.lower()
+        ]
+        if rankings_options_modified and not endpoint_type_is_rankings:
+            raise ValueError(
+                f"Rankings-specific options (--rankings-passages-mean, --rankings-passages-stddev, "
+                "--rankings-passages-prompt-token-mean, --rankings-passages-prompt-token-stddev, "
+                "--rankings-query-prompt-token-mean, --rankings-query-prompt-token-stddev) "
+                "can only be used with rankings endpoint types "
+                f"Rankings endpoints: ({', '.join(rankings_endpoints)})."
+            )
+
+        # Validate that prompt tokens and rankings tokens are not both set
+        if prompt_tokens_modified and (
+            rankings_tokens_modified or endpoint_type_is_rankings
+        ):
+            raise ValueError(
+                "The --prompt-input-tokens-mean/--prompt-input-tokens-stddev options "
+                "cannot be used together with rankings-specific token options or the rankings endpoints"
+                "Ranking options: (--rankings-passages-prompt-token-mean, --rankings-passages-prompt-token-stddev, "
+                "--rankings-query-prompt-token-mean, --rankings-query-prompt-token-stddev, ). "
+                f"Rankings endpoints: ({', '.join(rankings_endpoints)})."
+                "Please use only one set of options."
+            )
+        return self

src/aiperf/common/enums/__init__.py

@@ -56,6 +56,7 @@
     FrequencyMetricUnit,
     FrequencyMetricUnitInfo,
     GenericMetricUnit,
+    MetricDictValueTypeT,
     MetricFlags,
     MetricOverTimeUnit,
     MetricOverTimeUnitInfo,
@@ -64,8 +65,10 @@
     MetricTimeUnit,
     MetricTimeUnitInfo,
     MetricType,
+    MetricUnitT,
     MetricValueType,
     MetricValueTypeInfo,
+    MetricValueTypeT,
     MetricValueTypeVarT,
     PowerMetricUnit,
     PowerMetricUnitInfo,
@@ -141,6 +144,7 @@
     "LifecycleState",
     "MediaType",
     "MessageType",
+    "MetricDictValueTypeT",
     "MetricFlags",
     "MetricOverTimeUnit",
     "MetricOverTimeUnitInfo",
@@ -149,8 +153,10 @@
     "MetricTimeUnit",
     "MetricTimeUnitInfo",
     "MetricType",
+    "MetricUnitT",
     "MetricValueType",
     "MetricValueTypeInfo",
+    "MetricValueTypeT",
     "MetricValueTypeVarT",
     "ModelSelectionStrategy",
     "PowerMetricUnit",

src/aiperf/dataset/__init__.py

@@ -40,6 +40,7 @@
     BaseLoader,
     BasePublicDatasetLoader,
     CustomDatasetT,
+    Filename,
     MediaConversionMixin,
     MooncakeTrace,
     MooncakeTraceDatasetLoader,
@@ -70,6 +71,7 @@
     "CustomDatasetT",
     "DEFAULT_CORPUS_FILE",
     "DatasetManager",
+    "Filename",
     "ImageGenerator",
     "MP3_SUPPORTED_SAMPLE_RATES",
     "MediaConversionMixin",

src/aiperf/dataset/composer/synthetic_rankings.py

@@ -23,6 +23,8 @@ def __init__(self, config: UserConfig, tokenizer: Tokenizer):
 
         self.session_id_generator = SessionIDGenerator(seed=config.input.random_seed)
         self._passages_rng = rng.derive("dataset.rankings.passages")
+        self._passages_token_rng = rng.derive("dataset.rankings.passages.tokens")
+        self._query_token_rng = rng.derive("dataset.rankings.query.tokens")
 
         # Set default sampling strategy for synthetic rankings dataset if not explicitly set
         if self.config.input.dataset_sampling_strategy is None:
@@ -33,12 +35,6 @@ def __init__(self, config: UserConfig, tokenizer: Tokenizer):
                 f"Using default sampling strategy for synthetic rankings dataset: {InputDefaults.DATASET_SAMPLING_STRATEGY}"
             )
 
-        if self.config.input.prompt.input_tokens.mean <= 0:
-            raise ValueError(
-                "Synthetic rankings data generation requires text prompts to be enabled. "
-                "Please set --prompt-input-tokens-mean > 0."
-            )
-
     def create_dataset(self) -> list[Conversation]:
         """Generate synthetic dataset for the rankings endpoint.
 
@@ -64,17 +60,22 @@ def _create_turn(self, num_passages: int) -> Turn:
         """Create a single ranking turn with one synthetic query and multiple synthetic passages."""
         turn = Turn()
 
-        query_text = self.prompt_generator.generate(
-            mean=self.config.input.prompt.input_tokens.mean,
-            stddev=self.config.input.prompt.input_tokens.stddev,
+        query_text = self.prompt_generator.generate_prompt(
+            self.prompt_generator.calculate_num_tokens(
+                self.config.input.rankings_query_prompt_token_mean,
+                self.config.input.rankings_query_prompt_token_stddev,
+            )
         )
         query = Text(name="query", contents=[query_text])
 
+        # Generate passages with rankings-specific token counts (per passage)
         passages = Text(name="passages")
         for _ in range(num_passages):
-            passage_text = self.prompt_generator.generate(
-                mean=self.config.input.prompt.input_tokens.mean,
-                stddev=self.config.input.prompt.input_tokens.stddev,
+            passage_text = self.prompt_generator.generate_prompt(
+                self.prompt_generator.calculate_num_tokens(
+                    self.config.input.rankings_passages_prompt_token_mean,
+                    self.config.input.rankings_passages_prompt_token_stddev,
+                )
             )
             passages.contents.append(passage_text)
 

src/aiperf/dataset/generator/prompt.py

@@ -123,7 +123,7 @@ def _create_prefix_prompt_pool(self) -> None:
             raise NotInitializedError("Tokenized corpus is not initialized.")
 
         self._prefix_prompts = [
-            self._generate_prompt(self.config.prefix_prompt.length)
+            self.generate_prompt(self.config.prefix_prompt.length)
             for _ in range(self.config.prefix_prompt.pool_size)
         ]
         self.debug(
@@ -137,6 +137,7 @@ def generate(
         hash_ids: list[int] | None = None,
     ) -> str:
         """Generate a synthetic prompt with the configuration parameters.
+        Serves as a wrapper around other internal methods to provide a unified interface.
 
         Args:
             mean: The mean of the normal distribution.
@@ -151,10 +152,24 @@ def generate(
                 mean, hash_ids, self.config.input_tokens.block_size
             )
 
-        num_tokens = self._length_rng.sample_positive_normal_integer(mean, stddev)
-        return self._generate_prompt(num_tokens)
+        num_tokens = self.calculate_num_tokens(mean, stddev)
+        return self.generate_prompt(num_tokens)
 
-    def _generate_prompt(self, num_tokens: int) -> str:
+    def calculate_num_tokens(
+        self,
+        mean: int | None = None,
+        stddev: int | None = None,
+    ) -> int:
+        """Calculate the number of tokens for a prompt based on a normal distribution.
+
+        Args:
+            mean: The mean of the normal distribution.
+            stddev: The standard deviation of the normal distribution.
+        """
+
+        return self._length_rng.sample_positive_normal_integer(mean, stddev)
+
+    def generate_prompt(self, num_tokens: int) -> str:
         """Generate a prompt containing exactly `num_tokens` number of tokens.
 
         Args:
@@ -175,7 +190,7 @@ def _generate_cached_prompt(
         Generate a prompt containing exactly `num_tokens` by reusing previously generated prompts
         stored in `_cache`. Each hash index in `hash_ids` corresponds to a block of
         `block_size` tokens. If a hash index is found in `_cache`, its stored prompt is reused.
-        Otherwise, a new prompt is generated using `_generate_prompt()` and stored in `_cache`.
+        Otherwise, a new prompt is generated using `generate_prompt()` and stored in `_cache`.
 
         Args:
             num_tokens: The number of tokens required in the prompt.

src/aiperf/dataset/loader/__init__.py

@@ -33,6 +33,7 @@
     MultiTurnDatasetLoader,
 )
 from aiperf.dataset.loader.random_pool import (
+    Filename,
     RandomPoolDatasetLoader,
 )
 from aiperf.dataset.loader.sharegpt import (
@@ -48,6 +49,7 @@
     "BaseLoader",
     "BasePublicDatasetLoader",
     "CustomDatasetT",
+    "Filename",
     "MediaConversionMixin",
     "MooncakeTrace",
     "MooncakeTraceDatasetLoader",