diff --git a/docs/cli-options.md b/docs/cli-options.md index 8905f0bf94..41de9b143e 100644 --- a/docs/cli-options.md +++ b/docs/cli-options.md @@ -572,7 +572,7 @@ Token block size for hash-based prompt caching in trace datasets (`mooncake_trac #### `--seq-dist`, `--sequence-distribution` `` -Distribution of (ISL, OSL) pairs with probabilities for mixed workload simulation. Format: `ISL,OSL:prob;ISL,OSL:prob` (semicolons separate pairs, probabilities are percentages 0-100 that must sum to 100). Supports optional stddev: `ISL|stddev,OSL|stddev:prob`. Examples: `128,64:25;512,128:50;1024,256:25` or with variance: `256|10,128|5:40;512|20,256|10:60`. Also supports bracket `[(256,128):40,(512,256):60]` and JSON formats. +Distribution of (ISL, OSL) pairs with probabilities for mixed workload simulation. Format: `ISL,OSL:prob;ISL,OSL:prob` (semicolons separate pairs, probabilities are positive relative weights, normalized across all pairs, that do NOT need to sum to 100). Supports optional stddev: `ISL|stddev,OSL|stddev:prob`. Examples: `128,64:25;512,128:50;1024,256:25` or with variance: `256|10,128|5:40;512|20,256|10:60`. Also supports bracket `[(256,128):40,(512,256):60]` and JSON formats. ### Output Sequence Length (OSL) @@ -1969,7 +1969,7 @@ Token block size for hash-based prompt caching in trace datasets (`mooncake_trac #### `--seq-dist`, `--sequence-distribution` `` -Distribution of (ISL, OSL) pairs with probabilities for mixed workload simulation. Format: `ISL,OSL:prob;ISL,OSL:prob` (semicolons separate pairs, probabilities are percentages 0-100 that must sum to 100). Supports optional stddev: `ISL|stddev,OSL|stddev:prob`. Examples: `128,64:25;512,128:50;1024,256:25` or with variance: `256|10,128|5:40;512|20,256|10:60`. Also supports bracket `[(256,128):40,(512,256):60]` and JSON formats. +Distribution of (ISL, OSL) pairs with probabilities for mixed workload simulation. Format: `ISL,OSL:prob;ISL,OSL:prob` (semicolons separate pairs, probabilities are positive relative weights, normalized across all pairs, that do NOT need to sum to 100). Supports optional stddev: `ISL|stddev,OSL|stddev:prob`. Examples: `128,64:25;512,128:50;1024,256:25` or with variance: `256|10,128|5:40;512|20,256|10:60`. Also supports bracket `[(256,128):40,(512,256):60]` and JSON formats. ### Output Sequence Length (OSL) diff --git a/docs/tutorials/sequence-distributions.md b/docs/tutorials/sequence-distributions.md index 361e319849..33c2f238f5 100644 --- a/docs/tutorials/sequence-distributions.md +++ b/docs/tutorials/sequence-distributions.md @@ -39,6 +39,8 @@ This creates: Values are automatically clamped to be at least 1. +The `:PROB` weights are **relative** — they are normalized at sampling time and do **not** need to sum to 100. The `70;20;10` above happens to total 100, but `7;2;1` produces the identical split, and `50;1` yields ~98%/2%. Each weight must be positive. + ## Supported Formats ### 1. Semicolon Format (Recommended) @@ -80,6 +82,31 @@ Values are automatically clamped to be at least 1. ]} ``` +## YAML config: typed `sequence_distribution` buckets + +The string/CLI formats above only support fixed or `mean|stddev` (normal) ISL/OSL per bucket. In a YAML config, the `prompts.sequence_distribution:` field takes a list of buckets whose `isl` and `osl` accept **any** sampling distribution shape — fixed scalar, `{mean, stddev}` normal, `{mean, median}` log-normal, `{p50, p99[, mean]}` percentile, `{peaks: [...]}` multimodal, or `{points: [...]}` empirical. See [Sampling Distributions in YAML Configs](yaml-distributions.md) for the full set of shapes and the auto-detection rules. + +```yaml +prompts: + sequence_distribution: + # Conversation classes: one bucket drawn per conversation, kept for life. + - first_turn_isl: {p50: 20000, p99: 100000, mean: 30000} # turn 1 seed context + isl: {mean: 300, stddev: 100} # turns 2+ new input + osl: {mean: 250, stddev: 80} # every turn's output cap + probability: 50 + - {isl: {mean: 400000, stddev: 20000}, osl: 64, probability: 1} +``` + +Key behaviors: + +- **One bucket per conversation ("sticky").** Each conversation draws a single bucket at creation (weighted by `probability`) and keeps it for every turn — its ISL and OSL are drawn from that bucket's distributions on each turn. This is how you model correlations like "longer input → shorter output": a huge-context conversation never gets another class's outputs. +- **`first_turn_isl` sizes the seed context.** When a bucket sets `first_turn_isl`, turn 1 draws its ISL from it and turns 2+ draw from `isl` (the per-turn growth). When unset, `isl` applies to all turns of that bucket. +- **Each bucket samples its full shape.** A `{p50, p99, mean}` percentile or `{mean, median}` log-normal bucket is sampled from that distribution on every draw — nothing is flattened to a single mean. +- **`probability` weights are relative.** Normalized across buckets; no need to sum to 100 (each must be positive; `probability: 0` disables a bucket). +- When `sequence_distribution` is set it drives ISL/OSL for every turn: the plain `isl`/`osl` fields are ignored, and combining it with a top-level `first_turn_isl` is a config error — set `first_turn_isl` per bucket instead. + +> **Behavior change:** older releases drew a fresh bucket for every *turn* of a multi-turn conversation. Buckets are now drawn once per *conversation*. Single-turn workloads are unaffected. + ## Examples ### Example Case: Chatbot Workload Simulation diff --git a/docs/tutorials/yaml-distributions.md b/docs/tutorials/yaml-distributions.md index 4341f3b35d..4dea8040c4 100644 --- a/docs/tutorials/yaml-distributions.md +++ b/docs/tutorials/yaml-distributions.md @@ -6,7 +6,7 @@ sidebar-title: Sampling Distributions in YAML Configs # Sampling Distributions in YAML Configs -Several fields in an AIPerf YAML config — input/output token lengths, conversation turn counts, turn delays, image dimensions, audio length, and ranking passage counts — accept a *sampling distribution* instead of a single number. This tutorial covers all five distribution shapes AIPerf supports, the auto-detection rules that pick between them, and the optional `min:`/`max:` clamps that compose with any of them. +Several fields in an AIPerf YAML config — input/output token lengths, conversation turn counts, turn delays, image dimensions, audio length, and ranking passage counts — accept a *sampling distribution* instead of a single number. This tutorial covers all six distribution shapes AIPerf supports, the auto-detection rules that pick between them, and the optional `min:`/`max:` clamps that compose with any of them. If you only ever write `isl: 512`, you've already used a distribution — that scalar is the shorthand for a `FixedDistribution`. Everything below extends from there. @@ -17,26 +17,30 @@ Any field in a YAML config typed as a sampling distribution accepts the full set | Field | Section | What it controls | |---|---|---| | `isl` | `dataset.prompts` (and shorthand at `dataset.isl`) | Input sequence length, in tokens | +| `first_turn_isl` | `dataset.prompts` | First-turn starting-context ISL for growing multi-turn conversations (requires `isl`) | | `osl` | `dataset.prompts` and `dataset.osl` shorthand; also on file datasets | Output sequence length, in tokens | | `turns` | `dataset` | Number of request/response turns per conversation | | `turn_delay` | `dataset` | Delay between turns, in milliseconds | | `width`, `height` | `dataset.images` | Synthetic image dimensions, in pixels | | `length` | `dataset.audio` | Synthetic audio duration, in seconds | +| `shared_system_length` | `dataset.prefix_prompts` | Shared system prompt size, in tokens — sampled once per run (identical across sessions) | +| `user_context_length` | `dataset.prefix_prompts` | Per-session user context size, in tokens — sampled once per session | | `passages`, `passage_tokens`, `query_tokens` | `dataset.rankings` | Rankings/reranking endpoint shapes | Wherever you see `{mean: ..., stddev: ...}` in a template, you can swap in any other shape from this page. -## The five distribution types +## The six distribution types -AIPerf supports five distribution shapes, and **figures out which one you mean from the keys you wrote** — you don't have to add a `type:` key. The discriminator is purely structural: +AIPerf supports six distribution shapes, and **figures out which one you mean from the keys you wrote** — you don't have to add a `type:` key. The discriminator is purely structural, and the keys are checked in a fixed order (`peaks` → `points` → `p50`/`p99` → `median` → `stddev` → `value` → `mean`): | What you wrote | Type | Why | |---|---|---| | `isl: 512` | Fixed | Bare scalar | -| `isl: {mean: 512, stddev: 50}` | Normal | `stddev` present | -| `isl: {mean: 512, median: 400}` | Log-normal | `median` present | | `isl: {peaks: [...]}` | Multimodal | `peaks` present | | `isl: {points: [...]}` | Empirical | `points` present | +| `isl: {p50: 50000, p99: 400000}` | Percentile | `p50`/`p99` present (checked before `median`/`stddev`/`mean`) | +| `isl: {mean: 512, median: 400}` | Log-normal | `median` present | +| `isl: {mean: 512, stddev: 50}` | Normal | `stddev` present | You can override the inference with an explicit `type:` if you'd rather be loud: @@ -44,7 +48,7 @@ You can override the inference with an explicit `type:` if you'd rather be loud: isl: {type: normal, mean: 512, stddev: 50} ``` -`type:` accepts one of `fixed`, `normal`, `lognormal`, `multimodal`, `empirical`. AIPerf strips it after dispatch, so the rest of the dict is parsed normally. +`type:` accepts one of `fixed`, `normal`, `lognormal`, `percentile`, `multimodal`, `empirical`. AIPerf strips it after dispatch, so the rest of the dict is parsed normally. ### Fixed — a constant @@ -93,6 +97,23 @@ Constraints: Use log-normal when modelling sizes that are bounded below by zero and have a long right tail — chat prompt lengths, retrieval-augmented context windows, "most requests are small but some are huge" workloads. +### Percentile — target p50/p99 (and optionally the mean) directly + +Sometimes the workload is specified by its percentiles, not by distribution parameters: "median input is 50k tokens, p99 is 400k, average is 60k." No 2-parameter distribution can hit three targets like that (a normal is symmetric; a log-normal pins p50+p99 but then its mean is decided for you). The `percentile` shape takes the targets and solves the sampling parameters for you at config-validation time: + +```yaml +prompts: + isl: {p50: 50000, p99: 400000} # log-normal fit: p50 and p99 exact + osl: {p50: 200, p99: 2000, mean: 350} # mixture fit: also pins the mean +``` + +- `p50` and `p99` are required; `p99` must be greater than `p50`. +- `mean` is optional. Without it, the mean implied by the log-normal fit applies (for the 50k/400k example that is ~74.6k). With it, AIPerf solves a two-component mixture: a body carrying the median and a small far tail carrying the p99. +- `mean` must be greater than `p50` — percentile models right-skewed shapes only. Infeasible targets (a `mean` below `p50`, or a `mean` too close to `p99`) fail fast in `aiperf config validate` with a message explaining the constraint — not mid-benchmark. +- `min:`/`max:` clamps compose like with every other shape. + +Use percentile when you have production SLO-style numbers and want the synthetic workload to reproduce them without hand-calibrating a multimodal mixture. + ### Multimodal — a mixture of N peaks A weighted mixture of two or more sub-distributions. Each `peak` is itself a distribution, written inline, with an optional `weight`. @@ -168,10 +189,37 @@ A few rules: - Bounds are *inclusive*: `min: 32` means values down to and including 32 are kept; below 32 is clamped up to 32. - `min:` and `max:` must be finite. NaN/inf are rejected at config-validation time so they can't silently disable clamping. - If both are set, `min <= max` is enforced. -- Bounds compose with every shape — Fixed, Normal, Log-normal, Multimodal, and Empirical. +- Bounds compose with every shape — Fixed, Normal, Log-normal, Percentile, Multimodal, and Empirical. For multimodal distributions, a top-level `min`/`max` applies to the *output* of the mixture. If you want different bounds per peak, set `min`/`max` on each peak's sub-distribution instead. +## First-turn context — `first_turn_isl` + +Multi-turn benchmarks with *growing* context (e.g. user-centric mode, where each turn appends to the running conversation) often start with a large seed prompt and then add only small increments per turn. `first_turn_isl` lets you size that opening context separately from the follow-ups: + +```yaml +prompts: + first_turn_isl: {p50: 20000, p99: 100000} # starting context of each conversation + isl: {mean: 300, stddev: 100} # each subsequent turn's new input +``` + +With `sequence_distribution`, put the seed context on the bucket instead — one bucket per conversation, first turn from its `first_turn_isl`, later turns from its `isl`: + +```yaml +prompts: + sequence_distribution: + - first_turn_isl: {p50: 20000, p99: 100000, mean: 30000} + isl: {mean: 300, stddev: 100} + osl: {mean: 250, stddev: 80} + probability: 50 +``` + +- The **first** turn of every conversation draws its ISL from `first_turn_isl`; **every subsequent** turn draws from `isl`. +- Both fields accept any distribution shape on this page — the example pairs a percentile opener with a small normal follow-up. +- `first_turn_isl` **requires `isl`** to be set (it only redefines the first turn; `isl` still sizes the rest). Setting it alone fails at config-validation time. +- When `first_turn_isl` is unset, `isl` applies to all turns. +- `first_turn_isl` cannot be combined with `sequence_distribution` (config error). Buckets carry their own seed context instead — each `sequence_distribution` entry accepts a per-bucket `first_turn_isl`; see [Sequence Length Distributions](sequence-distributions.md). + ## Disambiguation cheat-sheet If AIPerf can't figure out what shape you meant, it errors at config-load time with a message that names the keys it saw. The most common causes: @@ -182,7 +230,9 @@ If AIPerf can't figure out what shape you meant, it errors at config-load time w | `isl: {stddev: 50}` (no `mean`) | Error — Normal requires `mean`. | | `isl: {peaks: [...one entry...]}` | Error — Multimodal requires at least 2 peaks. | | `isl: {value: 512, mean: 600}` | Error — `value` selects Fixed, but `mean` is unknown to Fixed. | -| Passing a string like `"128,64:50;512,128:50"` | Error — that's the legacy `sequence_distribution` string format (semicolon-separated `ISL,OSL:prob` pairs summing to 100), not a sampling distribution. See [Sequence Length Distributions](sequence-distributions.md). | +| `isl: {p50: 50000, p99: 400000, mean: 40000}` | Error — a percentile `mean` must be greater than `p50` (right-skewed shapes only). | +| Passing a string like `"128,64:50;512,128:50"` | Error — that's the legacy `sequence_distribution` string format (semicolon-separated relative-weight `ISL,OSL:prob` pairs), not a sampling distribution. See [Sequence Length Distributions](sequence-distributions.md). | +| `first_turn_isl` next to `sequence_distribution` | Error — seed context lives on each bucket's own `first_turn_isl` when buckets are in play. | When in doubt, run: diff --git a/src/aiperf/common/models/sequence_distribution.py b/src/aiperf/common/models/sequence_distribution.py index cf4d87fd16..768af8cdff 100644 --- a/src/aiperf/common/models/sequence_distribution.py +++ b/src/aiperf/common/models/sequence_distribution.py @@ -11,13 +11,13 @@ different probabilities, enabling simulation of mixed workloads that better represent production traffic patterns. - Supported formats (probabilities must be percentages 0-100): + Supported formats (probabilities are relative weights, normalized at sampling time): - Semicolon: "256,128:40;512,256:60" or "256|10,128|5:40;512|20,256|10:60" - Bracket: "[(256,128):40,(512,256):60]" or "[(256|10,128|5):40,(512|20,256|10):60]" - JSON: '{"pairs": [{"isl": 256, "isl_stddev": 10, "osl": 128, "osl_stddev": 5, "prob": 40}, ...]}' -Note: Probabilities must be specified as percentages (0-100), not fractions (0-1). -This prevents common errors from mixing different probability formats. +Note: Probabilities are positive relative weights and do NOT need to sum to 100. +They are normalized across all pairs at sampling time, so "50;1" yields ~98%/2%. Examples: Basic usage: @@ -32,6 +32,7 @@ from __future__ import annotations +import math import re from dataclasses import dataclass from typing import TYPE_CHECKING @@ -49,9 +50,13 @@ logger = AIPerfLogger(__name__) -def _validate_probability_sum(pairs: list[SequenceLengthPair]) -> None: +def _validate_probability_weights(pairs: list[SequenceLengthPair]) -> None: """ - Validate that probabilities sum to approximately 100.0. + Validate that the pairs carry a positive total relative weight. + + Weights are relative and normalized at sampling time, so they do not need + to sum to 100. The downstream cumulative-probability computation divides by + the summed total, so a non-positive total is the one invalid case. This is a module-level helper used by both SequenceLengthDistribution and DistributionParser to avoid code duplication. @@ -60,14 +65,13 @@ def _validate_probability_sum(pairs: list[SequenceLengthPair]) -> None: pairs: List of SequenceLengthPair objects to validate Raises: - ValueError: If probabilities don't sum to 100.0 (within floating-point tolerance) + ValueError: If the total probability weight is not positive """ total_prob = sum(pair.probability for pair in pairs) - # Allow small floating-point errors - if not np.isclose(total_prob, 100.0, rtol=1e-6, atol=1e-6): + if not math.isfinite(total_prob) or total_prob <= 0: raise ValueError( - f"Probabilities must sum to 100.0, got {total_prob:.6f}. " + f"Probability weights must have a finite positive total, got {total_prob:.6f}. " f"Pairs: {[str(p) for p in pairs]}" ) @@ -92,15 +96,25 @@ def __post_init__(self) -> None: raise ValueError( f"Output sequence length must be positive, got {self.output_seq_len}" ) - if not 0.0 <= self.probability <= 100.0: - raise ValueError(f"Probability must be in [0,100], got {self.probability}") - if self.input_seq_len_stddev < 0.0: + if not math.isfinite(self.probability) or self.probability < 0.0: + raise ValueError( + f"Probability weight must be finite and non-negative, got {self.probability}" + ) + if ( + not math.isfinite(self.input_seq_len_stddev) + or self.input_seq_len_stddev < 0.0 + ): raise ValueError( - f"Input sequence length stddev must be non-negative, got {self.input_seq_len_stddev}" + f"Input sequence length stddev must be finite and non-negative, " + f"got {self.input_seq_len_stddev}" ) - if self.output_seq_len_stddev < 0.0: + if ( + not math.isfinite(self.output_seq_len_stddev) + or self.output_seq_len_stddev < 0.0 + ): raise ValueError( - f"Output sequence length stddev must be non-negative, got {self.output_seq_len_stddev}" + f"Output sequence length stddev must be finite and non-negative, " + f"got {self.output_seq_len_stddev}" ) def __str__(self) -> str: @@ -123,10 +137,11 @@ def __init__(self, pairs: list[SequenceLengthPair]) -> None: Initialize distribution from list of sequence length pairs. Args: - pairs: List of SequenceLengthPair objects. Probabilities must sum to 100.0. + pairs: List of SequenceLengthPair objects. Probabilities are relative + weights normalized at sampling time and need not sum to 100.0. Raises: - ValueError: If pairs is empty or probabilities don't sum to 100.0. + ValueError: If pairs is empty or the total probability weight is not positive. """ if not pairs: raise ValueError( @@ -137,7 +152,7 @@ def __init__(self, pairs: list[SequenceLengthPair]) -> None: # stays pure and works before bootstrap calls rng.init(). self._rng: RandomGenerator | None = None self._pairs = tuple(pairs) # Immutable copy - _validate_probability_sum(list(self._pairs)) + _validate_probability_weights(list(self._pairs)) self._cumulative_probs = self._compute_cumulative_probabilities() logger.debug(f"Created distribution with {len(self._pairs)} pairs: {self}") @@ -150,8 +165,10 @@ def _get_rng(self) -> RandomGenerator: def _compute_cumulative_probabilities(self) -> np.ndarray: """Compute cumulative probability distribution for efficient sampling.""" - # Convert percentages to fractions for internal calculation - probs = [pair.probability / 100.0 for pair in self._pairs] + # Normalize relative weights by their actual total (they need not + # sum to 100); _validate_probability_weights guarantees total > 0. + total = sum(pair.probability for pair in self._pairs) + probs = [pair.probability / total for pair in self._pairs] return np.cumsum(probs, dtype=np.float64) def sample(self) -> tuple[int, int]: @@ -238,17 +255,18 @@ def get_statistics(self) -> dict[str, int | float | list[tuple[int, int, float]] Dictionary with distribution statistics including expected values, variance, and individual pair information. """ - # Expected values (convert percentages to fractions for calculation) - exp_isl = sum(p.input_seq_len * (p.probability / 100.0) for p in self._pairs) - exp_osl = sum(p.output_seq_len * (p.probability / 100.0) for p in self._pairs) + # Normalize relative weights by their actual total (they need not sum + # to 100); _validate_probability_weights guarantees total > 0. + total = sum(p.probability for p in self._pairs) + exp_isl = sum(p.input_seq_len * (p.probability / total) for p in self._pairs) + exp_osl = sum(p.output_seq_len * (p.probability / total) for p in self._pairs) - # Variance calculations var_isl = sum( - (p.probability / 100.0) * (p.input_seq_len - exp_isl) ** 2 + (p.probability / total) * (p.input_seq_len - exp_isl) ** 2 for p in self._pairs ) var_osl = sum( - (p.probability / 100.0) * (p.output_seq_len - exp_osl) ** 2 + (p.probability / total) * (p.output_seq_len - exp_osl) ** 2 for p in self._pairs ) @@ -322,8 +340,8 @@ def validate(cls, dist_str: str) -> list[SequenceLengthPair]: else: pairs = cls._validate_semicolon_format(dist_str) - # Validate probability sum without creating distribution object - _validate_probability_sum(pairs) + # Validate positive total weight without creating distribution object + _validate_probability_weights(pairs) return pairs except Exception as e: diff --git a/src/aiperf/config/dataset/content.py b/src/aiperf/config/dataset/content.py index 5f9366aacd..12eab39baf 100644 --- a/src/aiperf/config/dataset/content.py +++ b/src/aiperf/config/dataset/content.py @@ -77,6 +77,19 @@ class PromptConfig(BaseConfig): ), ] + first_turn_isl: Annotated[ + SamplingDistribution | None, + Field( + default=None, + description="Input sequence length for the FIRST turn of each conversation. " + "In multi-turn benchmarks with growing context (e.g. user-centric mode), " + "this sets the starting-context size while `isl` controls each subsequent " + "turn's new input. When unset, `isl` applies to all turns. " + "Requires `isl` to be set. Cannot be combined with sequence_distribution " + "(set it per bucket there instead).", + ), + ] + osl: Annotated[ SamplingDistribution | None, Field( @@ -120,7 +133,8 @@ class PromptConfig(BaseConfig): default=None, description="Distribution of (ISL, OSL) pairs with probabilities for mixed workload simulation. " "Each entry specifies {isl, osl, probability}. " - "Probabilities are percentages (0-100) and must sum to 100. " + "Probabilities are relative weights (normalized across all entries) and do NOT " + "need to sum to 100 (e.g. probability=50 alongside probability=1 yields ~98%%/2%%). " "When specified, requests are sampled from this distribution instead of using isl/osl fields.", ), ] @@ -134,6 +148,21 @@ def validate_sequence_probabilities( validate_probability_distribution(v) return v + @model_validator(mode="after") + def validate_first_turn_isl_constraints(self) -> PromptConfig: + if self.first_turn_isl is not None and self.sequence_distribution is not None: + raise ValueError( + "first_turn_isl cannot be combined with sequence_distribution; " + "set first_turn_isl on individual sequence_distribution entries " + "instead (each bucket accepts its own first_turn_isl)." + ) + if self.first_turn_isl is not None and self.isl is None: + raise ValueError( + "first_turn_isl requires isl to be set: first_turn_isl sizes the " + "first turn's starting context, isl sizes every subsequent turn." + ) + return self + class PrefixPromptConfig(BaseConfig): """ @@ -175,24 +204,24 @@ class PrefixPromptConfig(BaseConfig): ] shared_system_length: Annotated[ - int | None, + SamplingDistribution | None, Field( - ge=1, default=None, - description="Length of shared system prompt in tokens. " - "This prompt is identical across all sessions and appears as a system message. " + description="Length of shared system prompt in tokens. Accepts a fixed " + "integer or any sampling distribution; sampled ONCE per run so the " + "prompt stays identical across all sessions (its cache-hit purpose). " "First part of a two-part prefix structure with high cache hit rate expected. " "Mutually exclusive with pool_size/length.", ), ] user_context_length: Annotated[ - int | None, + SamplingDistribution | None, Field( - ge=1, default=None, - description="Length of per-session user context prompt in tokens. " - "Each dataset entry gets a unique user context prompt. " + description="Length of per-session user context prompt in tokens. Accepts " + "a fixed integer or any sampling distribution; sampled once PER SESSION " + "so each conversation gets a unique context prompt with its own size. " "Second part of two-part prefix structure with lower cache hit rate expected. " "Mutually exclusive with pool_size/length.", ), @@ -209,6 +238,13 @@ def _validate_prefix_exclusivity(self) -> Self: "pool_size/length and shared_system_length/user_context_length " "are mutually exclusive" ) + for name in ("shared_system_length", "user_context_length"): + dist = getattr(self, name) + if dist is not None and dist.expected_value < 1: + raise ValueError( + f"{name} must have an expected value >= 1, " + f"got {dist.expected_value}" + ) return self diff --git a/src/aiperf/config/distributions.py b/src/aiperf/config/distributions.py index 9fbb778897..da12f41123 100644 --- a/src/aiperf/config/distributions.py +++ b/src/aiperf/config/distributions.py @@ -3,31 +3,42 @@ """AIPerf Configuration - Sampling Distribution Types -5 distribution types, auto-detected from field structure (no ``type:`` key needed): +6 distribution types, auto-detected from field structure (no ``type:`` key needed): isl: 512 # FixedDistribution isl: {mean: 512, stddev: 50} # NormalDistribution isl: {mean: 512, median: 400} # LogNormalDistribution + isl: {p50: 50000, p99: 400000, mean: 60000} # PercentileDistribution isl: {peaks: [{...}, {...}], split: 60} # MultimodalDistribution isl: {points: [{value: 128, weight: 40}, ...]} # EmpiricalDistribution Discriminator rules (checked in order): - scalar int/float -> FixedDistribution - "peaks" in dict -> MultimodalDistribution - "points" in dict -> EmpiricalDistribution - "median" in dict -> LogNormalDistribution - "stddev" in dict -> NormalDistribution - "value" in dict -> FixedDistribution - "mean" alone -> NormalDistribution (stddev defaults to 0) - anything else -> ValueError + scalar int/float -> FixedDistribution + "peaks" in dict -> MultimodalDistribution + "points" in dict -> EmpiricalDistribution + "p50"/"p99" in dict -> PercentileDistribution + "median" in dict -> LogNormalDistribution + "stddev" in dict -> NormalDistribution + "value" in dict -> FixedDistribution + "mean" alone -> NormalDistribution (stddev defaults to 0) + anything else -> ValueError """ from __future__ import annotations import math +import sys +from dataclasses import dataclass from typing import TYPE_CHECKING, Annotated, Any, Self -from pydantic import ConfigDict, Discriminator, Field, Tag, model_validator +from pydantic import ( + ConfigDict, + Discriminator, + Field, + PrivateAttr, + Tag, + model_validator, +) from aiperf.config.base import BaseConfig @@ -105,6 +116,10 @@ def _validate_bounds(self) -> Self: def __getattr__(self, name: str) -> Any: if name == "mean": return self.expected_value + # PrivateAttr reads (e.g. PercentileDistribution._solution) must resolve + # through Pydantic's machinery, which this override would otherwise shadow. + if name in object.__getattribute__(self, "__private_attributes__"): + return super().__getattr__(name) raise AttributeError(f"{type(self).__name__!r} has no attribute {name!r}") def sample(self, rng: RandomGenerator) -> float: @@ -283,6 +298,255 @@ def __repr__(self) -> str: return f"lognormal(mean={self.mean:g}, median={self.median:g})" +_Z99 = 2.3263478740408408 +"""Phi^-1(0.99): the z-score of the 99th percentile of a standard normal.""" + +_SQRT2 = math.sqrt(2.0) + +_LOG_FLOAT_MAX = math.log(sys.float_info.max) +"""Largest argument for which ``math.exp`` returns a finite value (~709.78). + +Used to reject PercentileDistribution targets whose p99/p50 ratio makes the +fitted lognormal's mean or variance factor overflow to a non-finite value. +""" + + +def _phi(x: float) -> float: + """Standard normal CDF via math.erf (no scipy dependency).""" + return 0.5 * (1.0 + math.erf(x / _SQRT2)) + + +def _phi_inv(p: float) -> float: + """Standard normal inverse CDF via bisection. + + Deterministic and dependency-free; 80 bisection rounds over [-10, 10] + give far more precision than the solver's 1e-3 fit tolerances need. + """ + if not 0.0 < p < 1.0: + raise ValueError(f"phi_inv requires 0 < p < 1, got {p}") + lo, hi = -10.0, 10.0 + for _ in range(80): + mid = (lo + hi) / 2.0 + if _phi(mid) < p: + lo = mid + else: + hi = mid + return (lo + hi) / 2.0 + + +@dataclass(frozen=True) +class _PercentileSolution: + """Solved sampling parameters for a PercentileDistribution. + + kind == "lognormal": sample exp(N(mu, sigma)). + kind == "mixture": with probability tail_weight sample + positive-normal(tail_mean, tail_stddev), else + positive-normal(body_mean, body_stddev). + """ + + kind: str + mu: float = 0.0 + sigma: float = 0.0 + body_mean: float = 0.0 + body_stddev: float = 0.0 + tail_mean: float = 0.0 + tail_stddev: float = 0.0 + tail_weight: float = 0.0 + + +def _try_mixture_weight( + p50: float, p99: float, mean: float, body_stddev: float, w: float +) -> _PercentileSolution | None: + """Attempt an exact mixture fit at tail weight ``w``. + + Fixed-point iteration on the two cross-CDF terms; each round solves + body_mean from the p50 equation, tail_mean from the mean equation, and + tail_stddev from the p99 equation. Returns None when this ``w`` cannot + satisfy the targets (caller tries the next weight). + """ + f_tail_at_p50 = 0.0 + f_body_at_p99 = 1.0 + body_mean = tail_mean = tail_stddev = 0.0 + for _ in range(25): + body_cdf_target = (0.5 - w * f_tail_at_p50) / (1.0 - w) + if not 0.0 < body_cdf_target < 1.0: + return None + body_mean = p50 - _phi_inv(body_cdf_target) * body_stddev + tail_mean = (mean - (1.0 - w) * body_mean) / w + tail_cdf_target = (0.99 - (1.0 - w) * f_body_at_p99) / w + if not 0.0 < tail_cdf_target < 1.0: + return None + z = _phi_inv(tail_cdf_target) + if abs(z) < 1e-9: + return None + tail_stddev = (p99 - tail_mean) / z + if tail_stddev <= 0.0: + return None + f_tail_at_p50 = _phi((p50 - tail_mean) / tail_stddev) + f_body_at_p99 = _phi((p99 - body_mean) / body_stddev) + + # Truncation at 0 must be negligible or the sampled stats drift off target. + if body_mean <= 2.0 * body_stddev or tail_mean <= 2.0 * tail_stddev: + return None + achieved_p50 = (1.0 - w) * _phi((p50 - body_mean) / body_stddev) + w * _phi( + (p50 - tail_mean) / tail_stddev + ) + achieved_p99 = (1.0 - w) * _phi((p99 - body_mean) / body_stddev) + w * _phi( + (p99 - tail_mean) / tail_stddev + ) + achieved_mean = (1.0 - w) * body_mean + w * tail_mean + if abs(achieved_p50 - 0.5) > 0.005: + return None + if abs(achieved_p99 - 0.99) > 0.002: + return None + if abs(achieved_mean - mean) > 0.005 * mean: + return None + return _PercentileSolution( + kind="mixture", + body_mean=body_mean, + body_stddev=body_stddev, + tail_mean=tail_mean, + tail_stddev=tail_stddev, + tail_weight=w, + ) + + +def _solve_percentile( + p50: float, p99: float, mean: float | None +) -> _PercentileSolution: + """Solve sampling parameters hitting the percentile targets exactly. + + Without ``mean``: log-normal (2 params, 2 targets, closed form). + With ``mean``: log-normal if the mean already matches the implied one, + otherwise a two-component mixture searched over tail weights and body + spreads. Raises ValueError when no searched shape fits. + """ + mu = math.log(p50) + try: + sigma = math.log(p99 / p50) / _Z99 + implied_mean = p50 * math.exp(sigma**2 / 2.0) + except OverflowError as exc: + raise ValueError( + f"Percentile targets too extreme: p99/p50 ratio {p99 / p50:g} produces a " + f"non-finite mean or variance. Reduce the spread between p50 ({p50:g}) " + f"and p99 ({p99:g})." + ) from exc + # exp(sigma**2) is the lognormal's variance factor; when sigma**2 exceeds the + # largest finite math.exp argument the distribution's second moment is + # non-finite even though the mean may still be representable. Reject so + # downstream `expected_value > 0` gates never see a non-finite value. + if not math.isfinite(implied_mean) or sigma**2 > _LOG_FLOAT_MAX: + raise ValueError( + f"Percentile targets too extreme: p99/p50 ratio {p99 / p50:g} produces a " + f"non-finite mean or variance. Reduce the spread between p50 ({p50:g}) " + f"and p99 ({p99:g})." + ) + if mean is None or abs(mean - implied_mean) <= 0.01 * implied_mean: + return _PercentileSolution(kind="lognormal", mu=mu, sigma=sigma) + + for body_cv in (0.2, 0.1, 0.3): + body_stddev = body_cv * p50 + for w in (0.03, 0.02, 0.05, 0.08, 0.12, 0.18, 0.25, 0.35, 0.45): + solution = _try_mixture_weight(p50, p99, mean, body_stddev, w) + if solution is not None: + return solution + raise ValueError( + f"Percentile targets infeasible: no mixture found for p50={p50:g}, " + f"p99={p99:g}, mean={mean:g}. The mean must lie above p50 and not " + f"approach p99 (a log-normal with these percentiles implies mean~" + f"{implied_mean:g}; omit `mean` to use it). Adjust the targets or " + f"express the shape manually with a multimodal distribution." + ) + + +class PercentileDistribution(Distribution): + """Right-skewed distribution parameterized directly by percentile targets. + + YAML: + isl: {p50: 50000, p99: 400000} # log-normal, exact p50/p99 + isl: {p50: 50000, p99: 400000, mean: 60000} # mixture, also pins the mean + + With only p50 and p99, a log-normal fits both exactly. Adding ``mean`` + covers shapes no 2-parameter distribution can express (e.g. p50=50k, + p99=400k, mean=60k): a body component carries the median while a small + heavy tail carries the p99, solved deterministically at config + validation time so infeasible targets fail fast in `aiperf config + validate` rather than mid-benchmark. + """ + + p50: Annotated[ + float, + Field(gt=0.0, description="Target median (50th percentile) of the samples."), + ] + + p99: Annotated[ + float, + Field( + gt=0.0, + description="Target 99th percentile. Must be greater than p50.", + ), + ] + + mean: Annotated[ + float | None, + Field( + default=None, + gt=0.0, + description=( + "Optional target mean. Must be greater than p50. When omitted, " + "the mean implied by the p50/p99 log-normal fit applies." + ), + ), + ] = None + + _solution: _PercentileSolution | None = PrivateAttr(default=None) + + @model_validator(mode="after") + def validate_and_solve(self) -> Self: + for name, val in (("p50", self.p50), ("p99", self.p99), ("mean", self.mean)): + if val is not None and not math.isfinite(val): + raise ValueError( + f"Percentile target `{name}` must be finite, got {val!r}" + ) + if self.p99 <= self.p50: + raise ValueError(f"p99 ({self.p99}) must be greater than p50 ({self.p50}).") + if self.mean is not None and self.mean <= self.p50: + raise ValueError( + f"mean ({self.mean}) must be greater than p50 ({self.p50}); " + f"percentile distributions model right-skewed shapes. For " + f"left-heavy shapes use a multimodal or empirical distribution." + ) + self._solution = _solve_percentile(self.p50, self.p99, self.mean) + return self + + def _solved(self) -> _PercentileSolution: + if self._solution is None: + self._solution = _solve_percentile(self.p50, self.p99, self.mean) + return self._solution + + def _sample_raw(self, rng: RandomGenerator) -> float: + s = self._solved() + if s.kind == "lognormal": + if s.sigma <= 0: + return self.p50 + return math.exp(rng.sample_normal(s.mu, s.sigma)) + if rng.random() < s.tail_weight: + return rng.sample_positive_normal(s.tail_mean, s.tail_stddev) + return rng.sample_positive_normal(s.body_mean, s.body_stddev) + + @property + def expected_value(self) -> float: + if self.mean is not None: + return self.mean + sigma = math.log(self.p99 / self.p50) / _Z99 + return self.p50 * math.exp(sigma**2 / 2.0) + + def __repr__(self) -> str: + if self.mean is not None: + return f"percentile(p50={self.p50:g}, p99={self.p99:g}, mean={self.mean:g})" + return f"percentile(p50={self.p50:g}, p99={self.p99:g})" + + class PeakEntry(BaseConfig): """A weighted component in a multimodal distribution. @@ -443,11 +707,19 @@ def __repr__(self) -> str: "FixedDistribution": "fixed", "NormalDistribution": "normal", "LogNormalDistribution": "lognormal", + "PercentileDistribution": "percentile", "MultimodalDistribution": "multimodal", "EmpiricalDistribution": "empirical", } -_CANONICAL_TYPES = ("fixed", "normal", "lognormal", "multimodal", "empirical") +_CANONICAL_TYPES = ( + "fixed", + "normal", + "lognormal", + "percentile", + "multimodal", + "empirical", +) def _distribution_discriminator(v: Any) -> str: @@ -458,6 +730,7 @@ def _distribution_discriminator(v: Any) -> str: explicit "type:" -> use it (must be one of _CANONICAL_TYPES) "peaks" in dict -> "multimodal" "points" in dict -> "empirical" + "p50"/"p99" in dict -> "percentile" "median" in dict -> "lognormal" "stddev" in dict -> "normal" "value" in dict -> "fixed" @@ -481,6 +754,8 @@ def _distribution_discriminator(v: Any) -> str: return "multimodal" if "points" in v: return "empirical" + if "p50" in v or "p99" in v: + return "percentile" if "median" in v: return "lognormal" if "stddev" in v: @@ -491,7 +766,7 @@ def _distribution_discriminator(v: Any) -> str: return "normal" raise ValueError( "Cannot determine distribution type from keys. " - "Expected: scalar, {mean+stddev}, {mean+median}, " + "Expected: scalar, {mean+stddev}, {mean+median}, {p50, p99[, mean]}, " "{peaks:[distA, distB]}, or {points:[{value, weight}, ...]}." ) tag = _TAG_MAP.get(type(v).__name__) @@ -504,6 +779,7 @@ def _distribution_discriminator(v: Any) -> str: Annotated[FixedDistribution, Tag("fixed")] | Annotated[NormalDistribution, Tag("normal")] | Annotated[LogNormalDistribution, Tag("lognormal")] + | Annotated[PercentileDistribution, Tag("percentile")] | Annotated[MultimodalDistribution, Tag("multimodal")] | Annotated[EmpiricalDistribution, Tag("empirical")], Discriminator( @@ -511,7 +787,8 @@ def _distribution_discriminator(v: Any) -> str: custom_error_type="invalid_distribution_type", custom_error_message=( "Invalid distribution. Expected: scalar, {mean+stddev}, {mean+median}, " - "{peaks:[{...weight:N}, ...]}, or {points:[{value, weight}, ...]}." + "{p50, p99[, mean]}, {peaks:[{...weight:N}, ...]}, or " + "{points:[{value, weight}, ...]}." ), ), ] @@ -521,6 +798,7 @@ def _distribution_discriminator(v: Any) -> str: 512 -> FixedDistribution {mean: 512, stddev: 50} -> NormalDistribution {mean: 512, median: 400} -> LogNormalDistribution + {p50: 50000, p99: 400000, mean: 60000} -> PercentileDistribution {peaks: [{mean:128, stddev:20, weight:60}, {mean:2048, median:1800, weight:40}]} -> MultimodalDistribution {points: [{value: 128, weight: 40}, ...]} -> EmpiricalDistribution diff --git a/src/aiperf/config/flags/cli_config.py b/src/aiperf/config/flags/cli_config.py index 45afb8b7c8..4b4a31ecc1 100644 --- a/src/aiperf/config/flags/cli_config.py +++ b/src/aiperf/config/flags/cli_config.py @@ -1011,7 +1011,8 @@ def url(self) -> str: Field( default=None, description="Distribution of (ISL, OSL) pairs with probabilities for mixed workload simulation. " - "Format: `ISL,OSL:prob;ISL,OSL:prob` (semicolons separate pairs, probabilities are percentages 0-100 that must sum to 100). " + "Format: `ISL,OSL:prob;ISL,OSL:prob` (semicolons separate pairs, probabilities are positive relative weights, " + "normalized across all pairs, that do NOT need to sum to 100). " "Supports optional stddev: `ISL|stddev,OSL|stddev:prob`. " "Examples: `128,64:25;512,128:50;1024,256:25` or with variance: `256|10,128|5:40;512|20,256|10:60`. " "Also supports bracket `[(256,128):40,(512,256):60]` and JSON formats.", diff --git a/src/aiperf/config/schema/aiperf-config.schema.json b/src/aiperf/config/schema/aiperf-config.schema.json index cbf1475601..4dc2b09c87 100644 --- a/src/aiperf/config/schema/aiperf-config.schema.json +++ b/src/aiperf/config/schema/aiperf-config.schema.json @@ -743,6 +743,9 @@ { "$ref": "#/$defs/LogNormalDistribution" }, + { + "$ref": "#/$defs/PercentileDistribution" + }, { "$ref": "#/$defs/MultimodalDistribution" }, @@ -2505,6 +2508,9 @@ { "$ref": "#/$defs/LogNormalDistribution" }, + { + "$ref": "#/$defs/PercentileDistribution" + }, { "$ref": "#/$defs/MultimodalDistribution" }, @@ -2534,6 +2540,9 @@ { "$ref": "#/$defs/LogNormalDistribution" }, + { + "$ref": "#/$defs/PercentileDistribution" + }, { "$ref": "#/$defs/MultimodalDistribution" }, @@ -2618,6 +2627,9 @@ { "$ref": "#/$defs/LogNormalDistribution" }, + { + "$ref": "#/$defs/PercentileDistribution" + }, { "$ref": "#/$defs/MultimodalDistribution" }, @@ -2647,6 +2659,9 @@ { "$ref": "#/$defs/LogNormalDistribution" }, + { + "$ref": "#/$defs/PercentileDistribution" + }, { "$ref": "#/$defs/MultimodalDistribution" }, @@ -2803,6 +2818,9 @@ { "$ref": "#/$defs/LogNormalDistribution" }, + { + "$ref": "#/$defs/PercentileDistribution" + }, { "$ref": "#/$defs/MultimodalDistribution" }, @@ -7503,6 +7521,9 @@ { "$ref": "#/$defs/LogNormalDistribution" }, + { + "$ref": "#/$defs/PercentileDistribution" + }, { "$ref": "#/$defs/MultimodalDistribution" }, @@ -8535,6 +8556,9 @@ { "$ref": "#/$defs/LogNormalDistribution" }, + { + "$ref": "#/$defs/PercentileDistribution" + }, { "$ref": "#/$defs/MultimodalDistribution" }, @@ -8567,6 +8591,9 @@ { "$ref": "#/$defs/LogNormalDistribution" }, + { + "$ref": "#/$defs/PercentileDistribution" + }, { "$ref": "#/$defs/MultimodalDistribution" }, @@ -9741,6 +9768,9 @@ { "$ref": "#/$defs/LogNormalDistribution" }, + { + "$ref": "#/$defs/PercentileDistribution" + }, { "$ref": "#/$defs/MultimodalDistribution" }, @@ -10059,6 +10089,138 @@ } ] }, + "PercentileDistribution": { + "additionalProperties": false, + "description": "Right-skewed distribution parameterized directly by percentile targets.\n\nYAML:\n isl: {p50: 50000, p99: 400000} # log-normal, exact p50/p99\n isl: {p50: 50000, p99: 400000, mean: 60000} # mixture, also pins the mean\n\nWith only p50 and p99, a log-normal fits both exactly. Adding ``mean``\ncovers shapes no 2-parameter distribution can express (e.g. p50=50k,\np99=400k, mean=60k): a body component carries the median while a small\nheavy tail carries the p99, solved deterministically at config\nvalidation time so infeasible targets fail fast in `aiperf config\nvalidate` rather than mid-benchmark.", + "properties": { + "min": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + }, + { + "type": "string", + "pattern": ".*\\{\\{.*\\}\\}.*", + "description": "Jinja2 template (e.g., '{{ variable }}')." + }, + { + "type": "string", + "pattern": ".*\\$\\{[A-Za-z_][A-Za-z0-9_]*(?::[^}]*)?\\}.*", + "description": "Environment variable (e.g., '${VAR}' or '${VAR:default}')." + } + ], + "default": null, + "description": "Inclusive lower bound; samples below are clamped up. Applies to every distribution type — composes with mean/stddev/median/peaks/points/value.", + "title": "Min", + "x-jinja2-supported": true + }, + "max": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + }, + { + "type": "string", + "pattern": ".*\\{\\{.*\\}\\}.*", + "description": "Jinja2 template (e.g., '{{ variable }}')." + }, + { + "type": "string", + "pattern": ".*\\$\\{[A-Za-z_][A-Za-z0-9_]*(?::[^}]*)?\\}.*", + "description": "Environment variable (e.g., '${VAR}' or '${VAR:default}')." + } + ], + "default": null, + "description": "Inclusive upper bound; samples above are clamped down.", + "title": "Max", + "x-jinja2-supported": true + }, + "p50": { + "description": "Target median (50th percentile) of the samples.", + "title": "P50", + "oneOf": [ + { + "description": "Target median (50th percentile) of the samples.", + "exclusiveMinimum": 0.0, + "title": "P50", + "type": "number" + }, + { + "type": "string", + "pattern": ".*\\{\\{.*\\}\\}.*", + "description": "Jinja2 template (e.g., '{{ variable }}')." + }, + { + "type": "string", + "pattern": ".*\\$\\{[A-Za-z_][A-Za-z0-9_]*(?::[^}]*)?\\}.*", + "description": "Environment variable (e.g., '${VAR}' or '${VAR:default}')." + } + ], + "x-jinja2-supported": true + }, + "p99": { + "description": "Target 99th percentile. Must be greater than p50.", + "title": "P99", + "oneOf": [ + { + "description": "Target 99th percentile. Must be greater than p50.", + "exclusiveMinimum": 0.0, + "title": "P99", + "type": "number" + }, + { + "type": "string", + "pattern": ".*\\{\\{.*\\}\\}.*", + "description": "Jinja2 template (e.g., '{{ variable }}')." + }, + { + "type": "string", + "pattern": ".*\\$\\{[A-Za-z_][A-Za-z0-9_]*(?::[^}]*)?\\}.*", + "description": "Environment variable (e.g., '${VAR}' or '${VAR:default}')." + } + ], + "x-jinja2-supported": true + }, + "mean": { + "anyOf": [ + { + "exclusiveMinimum": 0.0, + "type": "number" + }, + { + "type": "null" + }, + { + "type": "string", + "pattern": ".*\\{\\{.*\\}\\}.*", + "description": "Jinja2 template (e.g., '{{ variable }}')." + }, + { + "type": "string", + "pattern": ".*\\$\\{[A-Za-z_][A-Za-z0-9_]*(?::[^}]*)?\\}.*", + "description": "Environment variable (e.g., '${VAR}' or '${VAR:default}')." + } + ], + "default": null, + "description": "Optional target mean. Must be greater than p50. When omitted, the mean implied by the p50/p99 log-normal fit applies.", + "title": "Mean", + "x-jinja2-supported": true + } + }, + "required": [ + "p50", + "p99" + ], + "title": "PercentileDistribution", + "type": "object", + "x-kubernetes-preserve-unknown-fields": true + }, "PoissonPhase": { "additionalProperties": false, "description": "Poisson-distributed request arrivals at the target rate.", @@ -10473,52 +10635,66 @@ "sharedSystemLength": { "anyOf": [ { - "minimum": 1, - "type": "integer" + "oneOf": [ + { + "$ref": "#/$defs/FixedDistribution" + }, + { + "$ref": "#/$defs/NormalDistribution" + }, + { + "$ref": "#/$defs/LogNormalDistribution" + }, + { + "$ref": "#/$defs/PercentileDistribution" + }, + { + "$ref": "#/$defs/MultimodalDistribution" + }, + { + "$ref": "#/$defs/EmpiricalDistribution" + } + ] }, { "type": "null" - }, - { - "type": "string", - "pattern": ".*\\{\\{.*\\}\\}.*", - "description": "Jinja2 template (e.g., '{{ variable }}')." - }, - { - "type": "string", - "pattern": ".*\\$\\{[A-Za-z_][A-Za-z0-9_]*(?::[^}]*)?\\}.*", - "description": "Environment variable (e.g., '${VAR}' or '${VAR:default}')." } ], "default": null, - "description": "Length of shared system prompt in tokens. This prompt is identical across all sessions and appears as a system message. First part of a two-part prefix structure with high cache hit rate expected. Mutually exclusive with pool_size/length.", - "title": "Sharedsystemlength", - "x-jinja2-supported": true + "description": "Length of shared system prompt in tokens. Accepts a fixed integer or any sampling distribution; sampled ONCE per run so the prompt stays identical across all sessions (its cache-hit purpose). First part of a two-part prefix structure with high cache hit rate expected. Mutually exclusive with pool_size/length.", + "title": "Sharedsystemlength" }, "userContextLength": { "anyOf": [ { - "minimum": 1, - "type": "integer" + "oneOf": [ + { + "$ref": "#/$defs/FixedDistribution" + }, + { + "$ref": "#/$defs/NormalDistribution" + }, + { + "$ref": "#/$defs/LogNormalDistribution" + }, + { + "$ref": "#/$defs/PercentileDistribution" + }, + { + "$ref": "#/$defs/MultimodalDistribution" + }, + { + "$ref": "#/$defs/EmpiricalDistribution" + } + ] }, { "type": "null" - }, - { - "type": "string", - "pattern": ".*\\{\\{.*\\}\\}.*", - "description": "Jinja2 template (e.g., '{{ variable }}')." - }, - { - "type": "string", - "pattern": ".*\\$\\{[A-Za-z_][A-Za-z0-9_]*(?::[^}]*)?\\}.*", - "description": "Environment variable (e.g., '${VAR}' or '${VAR:default}')." } ], "default": null, - "description": "Length of per-session user context prompt in tokens. Each dataset entry gets a unique user context prompt. Second part of two-part prefix structure with lower cache hit rate expected. Mutually exclusive with pool_size/length.", - "title": "Usercontextlength", - "x-jinja2-supported": true + "description": "Length of per-session user context prompt in tokens. Accepts a fixed integer or any sampling distribution; sampled once PER SESSION so each conversation gets a unique context prompt with its own size. Second part of two-part prefix structure with lower cache hit rate expected. Mutually exclusive with pool_size/length.", + "title": "Usercontextlength" } }, "title": "PrefixPromptConfig", @@ -10541,6 +10717,9 @@ { "$ref": "#/$defs/LogNormalDistribution" }, + { + "$ref": "#/$defs/PercentileDistribution" + }, { "$ref": "#/$defs/MultimodalDistribution" }, @@ -10557,6 +10736,38 @@ "description": "Input sequence length in tokens. Can be a fixed integer (e.g., 512) or distribution {mean: 512, stddev: 50}. AIPerf generates prompts with lengths following a normal distribution around the mean (±stddev). Ignored when sequence_distribution is specified.", "title": "Isl" }, + "firstTurnIsl": { + "anyOf": [ + { + "oneOf": [ + { + "$ref": "#/$defs/FixedDistribution" + }, + { + "$ref": "#/$defs/NormalDistribution" + }, + { + "$ref": "#/$defs/LogNormalDistribution" + }, + { + "$ref": "#/$defs/PercentileDistribution" + }, + { + "$ref": "#/$defs/MultimodalDistribution" + }, + { + "$ref": "#/$defs/EmpiricalDistribution" + } + ] + }, + { + "type": "null" + } + ], + "default": null, + "description": "Input sequence length for the FIRST turn of each conversation. In multi-turn benchmarks with growing context (e.g. user-centric mode), this sets the starting-context size while `isl` controls each subsequent turn's new input. When unset, `isl` applies to all turns. Requires `isl` to be set. Cannot be combined with sequence_distribution (set it per bucket there instead).", + "title": "Firstturnisl" + }, "osl": { "anyOf": [ { @@ -10570,6 +10781,9 @@ { "$ref": "#/$defs/LogNormalDistribution" }, + { + "$ref": "#/$defs/PercentileDistribution" + }, { "$ref": "#/$defs/MultimodalDistribution" }, @@ -10649,7 +10863,7 @@ } ], "default": null, - "description": "Distribution of (ISL, OSL) pairs with probabilities for mixed workload simulation. Each entry specifies {isl, osl, probability}. Probabilities are percentages (0-100) and must sum to 100. When specified, requests are sampled from this distribution instead of using isl/osl fields.", + "description": "Distribution of (ISL, OSL) pairs with probabilities for mixed workload simulation. Each entry specifies {isl, osl, probability}. Probabilities are relative weights (normalized across all entries) and do NOT need to sum to 100 (e.g. probability=50 alongside probability=1 yields ~98%%/2%%). When specified, requests are sampled from this distribution instead of using isl/osl fields.", "title": "Sequencedistribution" } }, @@ -10853,6 +11067,9 @@ { "$ref": "#/$defs/LogNormalDistribution" }, + { + "$ref": "#/$defs/PercentileDistribution" + }, { "$ref": "#/$defs/MultimodalDistribution" }, @@ -10885,6 +11102,9 @@ { "$ref": "#/$defs/LogNormalDistribution" }, + { + "$ref": "#/$defs/PercentileDistribution" + }, { "$ref": "#/$defs/MultimodalDistribution" }, @@ -10917,6 +11137,9 @@ { "$ref": "#/$defs/LogNormalDistribution" }, + { + "$ref": "#/$defs/PercentileDistribution" + }, { "$ref": "#/$defs/MultimodalDistribution" }, @@ -11589,7 +11812,7 @@ }, "SequenceDistributionEntry": { "additionalProperties": false, - "description": "Defines a single entry in an ISL/OSL probability distribution.\n\nAIPerf supports multi-modal token length distributions, allowing\ndifferent ISL/OSL combinations with relative frequencies for\nrealistic workload modeling.\n\nYAML Representation:\n sequence_distribution:\n - {isl: 128, osl: 64, probability: 40}\n - {isl: {mean: 512, stddev: 50}, osl: 256, probability: 35}\n - {isl: {mean: 2048, median: 1800}, osl: 512, probability: 25}", + "description": "Defines a single entry in an ISL/OSL probability distribution.\n\nAIPerf supports multi-modal token length distributions, allowing\ndifferent ISL/OSL combinations with relative frequencies for\nrealistic workload modeling.\n\nYAML Representation:\n sequence_distribution:\n - {isl: 128, osl: 64, probability: 40}\n - {isl: {mean: 512, stddev: 50}, osl: 256, probability: 35}\n - {first_turn_isl: {p50: 20000, p99: 100000}, isl: {mean: 300, stddev: 100},\n osl: 512, probability: 25}", "examples": [ { "isl": 128, @@ -11617,6 +11840,9 @@ { "$ref": "#/$defs/LogNormalDistribution" }, + { + "$ref": "#/$defs/PercentileDistribution" + }, { "$ref": "#/$defs/MultimodalDistribution" }, @@ -11662,6 +11888,38 @@ "title": "Islstddev", "x-jinja2-supported": true }, + "firstTurnIsl": { + "anyOf": [ + { + "oneOf": [ + { + "$ref": "#/$defs/FixedDistribution" + }, + { + "$ref": "#/$defs/NormalDistribution" + }, + { + "$ref": "#/$defs/LogNormalDistribution" + }, + { + "$ref": "#/$defs/PercentileDistribution" + }, + { + "$ref": "#/$defs/MultimodalDistribution" + }, + { + "$ref": "#/$defs/EmpiricalDistribution" + } + ] + }, + { + "type": "null" + } + ], + "default": null, + "description": "Input sequence length for the FIRST turn of conversations drawn from this bucket. Sets the starting-context size of the conversation class while `isl` sizes each subsequent turn's new input. When unset, `isl` applies to all turns. Accepts a fixed integer or any distribution shape, including percentile {p50, p99, mean}.", + "title": "Firstturnisl" + }, "osl": { "description": "Output sequence length (tokens). Can be a fixed integer, a {mean, stddev} distribution, or any typed distribution.", "oneOf": [ @@ -11674,6 +11932,9 @@ { "$ref": "#/$defs/LogNormalDistribution" }, + { + "$ref": "#/$defs/PercentileDistribution" + }, { "$ref": "#/$defs/MultimodalDistribution" }, @@ -11720,12 +11981,11 @@ "x-jinja2-supported": true }, "probability": { - "description": "Relative probability weight for this distribution bucket (0-100). Weights are normalized across all entries. Example: probability=40 means 40%% of requests use this ISL/OSL.", + "description": "Relative probability weight for this distribution bucket. Zero disables the bucket (it is never sampled), which supports sweep-templated configs. Weights are normalized across all entries and do NOT need to sum to 100 (probability=50 alongside probability=1 yields ~98%%/2%%).", "title": "Probability", "oneOf": [ { - "description": "Relative probability weight for this distribution bucket (0-100). Weights are normalized across all entries. Example: probability=40 means 40%% of requests use this ISL/OSL.", - "maximum": 100.0, + "description": "Relative probability weight for this distribution bucket. Zero disables the bucket (it is never sampled), which supports sweep-templated configs. Weights are normalized across all entries and do NOT need to sum to 100 (probability=50 alongside probability=1 yields ~98%%/2%%).", "minimum": 0.0, "title": "Probability", "type": "number" @@ -12341,6 +12601,9 @@ { "$ref": "#/$defs/LogNormalDistribution" }, + { + "$ref": "#/$defs/PercentileDistribution" + }, { "$ref": "#/$defs/MultimodalDistribution" }, @@ -12370,6 +12633,9 @@ { "$ref": "#/$defs/LogNormalDistribution" }, + { + "$ref": "#/$defs/PercentileDistribution" + }, { "$ref": "#/$defs/MultimodalDistribution" }, @@ -12472,6 +12738,9 @@ { "$ref": "#/$defs/LogNormalDistribution" }, + { + "$ref": "#/$defs/PercentileDistribution" + }, { "$ref": "#/$defs/MultimodalDistribution" }, @@ -12501,6 +12770,9 @@ { "$ref": "#/$defs/LogNormalDistribution" }, + { + "$ref": "#/$defs/PercentileDistribution" + }, { "$ref": "#/$defs/MultimodalDistribution" }, diff --git a/src/aiperf/config/types.py b/src/aiperf/config/types.py index bc42094482..04c631a6a8 100644 --- a/src/aiperf/config/types.py +++ b/src/aiperf/config/types.py @@ -11,6 +11,7 @@ from __future__ import annotations +import math from typing import Annotated, Any from pydantic import ConfigDict, Field, model_validator @@ -54,7 +55,8 @@ class SequenceDistributionEntry(BaseConfig): sequence_distribution: - {isl: 128, osl: 64, probability: 40} - {isl: {mean: 512, stddev: 50}, osl: 256, probability: 35} - - {isl: {mean: 2048, median: 1800}, osl: 512, probability: 25} + - {first_turn_isl: {p50: 20000, p99: 100000}, isl: {mean: 300, stddev: 100}, + osl: 512, probability: 25} """ isl: Annotated[ @@ -77,6 +79,18 @@ class SequenceDistributionEntry(BaseConfig): ), ] + first_turn_isl: Annotated[ + SamplingDistribution | None, + Field( + default=None, + description="Input sequence length for the FIRST turn of conversations " + "drawn from this bucket. Sets the starting-context size of the " + "conversation class while `isl` sizes each subsequent turn's new input. " + "When unset, `isl` applies to all turns. Accepts a fixed integer or any " + "distribution shape, including percentile {p50, p99, mean}.", + ), + ] + osl: Annotated[ SamplingDistribution, Field( @@ -128,10 +142,12 @@ def merge_stddev_shorthand(cls, data: Any) -> Any: float, Field( ge=0.0, - le=100.0, - description="Relative probability weight for this distribution bucket (0-100). " - "Weights are normalized across all entries. " - "Example: probability=40 means 40%% of requests use this ISL/OSL.", + allow_inf_nan=False, + description="Relative probability weight for this distribution bucket. " + "Zero disables the bucket (it is never sampled), which supports " + "sweep-templated configs. Weights are normalized across all entries and " + "do NOT need to sum to 100 (probability=50 alongside probability=1 yields " + "~98%%/2%%).", ), ] @@ -155,11 +171,19 @@ def merge_stddev_shorthand(cls, data: Any) -> Any: def validate_probability_distribution( entries: list[SequenceDistributionEntry], ) -> list[SequenceDistributionEntry]: - """Validate that a probability distribution sums to approximately 100.""" + """Validate a relative-weight probability distribution. + + Weights are relative and normalized at sampling time, so they do not need + to sum to 100. The only requirement is at least one entry with a positive + total weight: the downstream sampler normalizes by that total, so a + non-positive total would divide by zero. + """ + if not entries: + raise ValueError("Sequence distribution must contain at least one entry.") total = sum(entry.probability for entry in entries) - if not (99.0 <= total <= 101.0): + if not math.isfinite(total) or total <= 0: raise ValueError( - f"Sequence distribution probabilities must sum to ~100, got {total}. " + f"Sequence distribution weights must have a finite positive total, got {total}. " f"Individual probabilities: {[e.probability for e in entries]}" ) return entries diff --git a/src/aiperf/dataset/composer/base.py b/src/aiperf/dataset/composer/base.py index a7b08f0640..f14ed8fb31 100644 --- a/src/aiperf/dataset/composer/base.py +++ b/src/aiperf/dataset/composer/base.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations +import bisect import inspect from abc import ABC, abstractmethod from typing import TYPE_CHECKING @@ -10,10 +11,6 @@ from aiperf.common.enums import ConversationContextMode, ModelSelectionStrategy from aiperf.common.mixins import AIPerfLoggerMixin from aiperf.common.models import Conversation, Turn -from aiperf.common.models.sequence_distribution import ( - SequenceLengthDistribution, - SequenceLengthPair, -) from aiperf.common.tokenizer import Tokenizer from aiperf.config.dataset import FileDataset, SyntheticDataset from aiperf.dataset.generator.audio import AudioGenerator @@ -22,6 +19,7 @@ from aiperf.dataset.generator.video import VideoGenerator if TYPE_CHECKING: + from aiperf.common.random_generator import RandomGenerator from aiperf.config.dataset import VideoConfig from aiperf.config.dataset.content import ( AudioConfig, @@ -31,6 +29,59 @@ ) from aiperf.config.distributions import SamplingDistribution from aiperf.config.resolution.plan import BenchmarkRun + from aiperf.config.types import SequenceDistributionEntry + + +class _TypedSequenceDistribution: + """Weighted ISL/OSL buckets whose per-bucket typed distributions sample + their full shape (the legacy SequenceLengthDistribution only supported + fixed-or-normal buckets). + + Weights are relative and normalized by their total here; config-level + validation only guarantees a positive total (so this division is safe). + """ + + def __init__( + self, entries: list[SequenceDistributionEntry], rng_instance: RandomGenerator + ) -> None: + # Zero-weight buckets are valid config (disabled buckets in sweep + # templates); drop them so they never sample and never widen a + # cumulative interval. Config-level validation normally guarantees a + # positive-weight entry survives; the guard below is defensive. + self._entries = [e for e in entries if e.probability > 0] + if not self._entries: + raise ValueError( + "sequence_distribution requires at least one positive-weight entry" + ) + self._rng = rng_instance + total = sum(e.probability for e in self._entries) + cumulative = 0.0 + self._cumulative: list[float] = [] + for entry in self._entries: + cumulative += entry.probability / total + self._cumulative.append(cumulative) + + def sample_bucket(self) -> SequenceDistributionEntry: + """Draw one weighted bucket. Called once per conversation: the bucket + is the conversation's workload class and stays fixed for its life.""" + r = self._rng.random() + idx = bisect.bisect_right(self._cumulative, r) + return self._entries[min(idx, len(self._entries) - 1)] + + def sample_lengths( + self, bucket: SequenceDistributionEntry, *, is_first: bool = False + ) -> tuple[int, int]: + """Draw one (ISL, OSL) pair from the held bucket. The first turn's ISL + comes from the bucket's first_turn_isl (seed context) when set.""" + isl_dist = ( + bucket.first_turn_isl + if is_first and bucket.first_turn_isl is not None + else bucket.isl + ) + return ( + isl_dist.sample_int(self._rng), + bucket.osl.sample_int(self._rng), + ) class BaseDatasetComposer(AIPerfLoggerMixin, ABC): @@ -81,17 +132,20 @@ def __init__( self.video_generator = VideoGenerator(self._synthetic_video) self._model_selector_rng = rng.derive("composer.turn.model_selection") - self._max_tokens_rng = rng.derive("composer.turn.max_tokens") + self._seq_len_rng = rng.derive("composer.turn.sequence_lengths") self.turn_count = 0 # ``PromptConfig.sequence_distribution`` is a # ``list[SequenceDistributionEntry]`` of typed ``SamplingDistribution`` - # objects. Convert each entry directly to a ``SequenceLengthPair`` - # (extracting mean + stddev from the underlying distribution) and - # build the runtime distribution without re-serializing through - # ``DistributionParser.parse``, which only accepts strings. - self._seq_distribution = self._build_sequence_distribution() + # objects. Each conversation draws ONE bucket at creation and keeps it + # for every turn (sticky per-conversation workload class); the runtime + # sampler draws that bucket's ISL/OSL from their full distribution shape + # (lognormal/multimodal/empirical/percentile), not a flattened + # mean+stddev. + self._seq_distribution: _TypedSequenceDistribution | None = ( + self._build_sequence_distribution() + ) # Cache for turn-level sequence lengths to ensure ISL/OSL pairing consistency self._turn_sequence_cache: dict[int, tuple[int, int]] = {} @@ -106,33 +160,23 @@ def create_dataset(self) -> list[Conversation]: """ ... - def _build_sequence_distribution(self) -> SequenceLengthDistribution | None: - """Build a runtime sequence-length distribution from config entries. + def _build_sequence_distribution(self) -> _TypedSequenceDistribution | None: + """Build the runtime sequence-length sampler from config entries. ``PromptConfig.sequence_distribution`` is a list of ``SequenceDistributionEntry`` carrying typed ``SamplingDistribution`` - ISL/OSL fields (Fixed/Normal/LogNormal/...). Pull the mean and the - normal-distribution stddev (0 for non-normal types) off each entry to - construct ``SequenceLengthPair`` directly. ``DistributionParser.parse`` - only accepts strings and would reject this list shape. + ISL/OSL fields. Each bucket samples its full distribution shape — + lognormal/multimodal/empirical/percentile buckets are NOT flattened + to mean+stddev. """ if self._synthetic_prompts is None: return None entries = self._synthetic_prompts.sequence_distribution if not entries: return None - - pairs = [ - SequenceLengthPair( - input_seq_len=int(entry.isl.expected_value), - output_seq_len=int(entry.osl.expected_value), - probability=float(entry.probability), - input_seq_len_stddev=float(getattr(entry.isl, "stddev", 0.0) or 0.0), - output_seq_len_stddev=float(getattr(entry.osl, "stddev", 0.0) or 0.0), - ) - for entry in entries - ] - return SequenceLengthDistribution(pairs) + return _TypedSequenceDistribution( + entries, rng.derive("composer.sequence.distribution") + ) def _osl_distribution(self) -> SamplingDistribution | None: """Resolve the OSL distribution to use as a fallback for max_tokens. @@ -170,14 +214,30 @@ def _select_model_name(self) -> str: else: raise ValueError(f"Invalid model selection strategy: {strategy}.") - def _get_turn_sequence_lengths(self, turn_id: int) -> tuple[int, int]: - """Get or sample ISL/OSL pair for a specific turn, ensuring consistency. + def _get_turn_sequence_lengths( + self, + turn_id: int, + *, + is_first: bool = False, + bucket: SequenceDistributionEntry | None = None, + ) -> tuple[int, int]: + """Sample (or return the cached) ISL/OSL pair for a specific turn. - This method caches the sequence lengths per turn to ensure that the same - ISL/OSL pair is used for both prompt generation and max_tokens setting. + Both lengths are drawn from their full typed distributions + (Fixed/Normal/LogNormal/Multimodal/Empirical/Percentile) exactly once + per turn; the cache guarantees prompt generation and max_tokens see + the same pair. Args: turn_id: Unique identifier for the turn + is_first: When True the ISL is drawn from the first-turn starting- + context distribution instead of ``isl`` (``prompts.first_turn_isl`` + on the plain path, ``bucket.first_turn_isl`` on the + sequence_distribution path). + bucket: The conversation's sticky sequence_distribution bucket. + Required on the sequence_distribution path when the cache is + cold; None falls back to a fresh bucket draw (non-synthetic + composers that never threaded one, e.g. rankings max_tokens). Returns: Tuple of (input_seq_len, output_seq_len) @@ -186,24 +246,32 @@ def _get_turn_sequence_lengths(self, turn_id: int) -> tuple[int, int]: return self._turn_sequence_cache[turn_id] if self._seq_distribution is None: - isl_mean = ( - int(self._synthetic_prompts.isl.expected_value) - if self._synthetic_prompts is not None - and self._synthetic_prompts.isl is not None + prompts = self._synthetic_prompts + isl_dist = None + if prompts is not None: + isl_dist = ( + prompts.first_turn_isl + if is_first and prompts.first_turn_isl is not None + else prompts.isl + ) + isl = ( + isl_dist.sample_int(self._seq_len_rng) + if isl_dist is not None and isl_dist.expected_value > 0 else 0 ) - osl_mean = ( - int(self._synthetic_prompts.osl.expected_value) - if self._synthetic_prompts is not None - and self._synthetic_prompts.osl is not None + osl_dist = self._osl_distribution() + osl = ( + osl_dist.sample_int(self._seq_len_rng) + if osl_dist is not None and osl_dist.expected_value > 0 else None ) - seq_lengths = ( - isl_mean, - osl_mean or max(128, isl_mean // 2), - ) + seq_lengths = (isl, osl if osl is not None else max(128, isl // 2)) else: - seq_lengths = self._seq_distribution.sample() + if bucket is None: + bucket = self._seq_distribution.sample_bucket() + seq_lengths = self._seq_distribution.sample_lengths( + bucket, is_first=is_first + ) self._turn_sequence_cache[turn_id] = seq_lengths return seq_lengths @@ -216,7 +284,9 @@ def _clear_turn_cache(self, turn_id: int) -> None: """ self._turn_sequence_cache.pop(turn_id, None) - def _set_max_tokens(self, turn: Turn) -> None: + def _set_max_tokens( + self, turn: Turn, bucket: SequenceDistributionEntry | None = None + ) -> None: """Set max_tokens for the turn based on the sequence distribution or output configuration. If the turn already has max_tokens set (e.g., from per-line input data), @@ -225,6 +295,8 @@ def _set_max_tokens(self, turn: Turn) -> None: Args: turn: The turn object to finalize. + bucket: The conversation's sticky sequence_distribution bucket, so the + cached OSL comes from the same class the turn's ISL was drawn from. """ if turn.max_tokens is not None: return @@ -232,21 +304,18 @@ def _set_max_tokens(self, turn: Turn) -> None: if self._seq_distribution is not None: # Use cached sequence distribution to get OSL (ensures ISL/OSL pairing consistency) turn_id = id(turn) - _, osl = self._get_turn_sequence_lengths(turn_id) + _, osl = self._get_turn_sequence_lengths(turn_id, bucket=bucket) if osl > 0: turn.max_tokens = osl else: osl_dist = self._osl_distribution() - if osl_dist is not None: - osl_mean = int(osl_dist.expected_value) - if osl_mean <= 0: - return - osl_stddev = int(getattr(osl_dist, "stddev", 0.0) or 0.0) - turn.max_tokens = self._max_tokens_rng.sample_positive_normal_integer( - osl_mean, osl_stddev - ) + if osl_dist is not None and osl_dist.expected_value > 0: + _, osl = self._get_turn_sequence_lengths(id(turn)) + turn.max_tokens = osl - def _finalize_turn(self, turn: Turn) -> None: + def _finalize_turn( + self, turn: Turn, bucket: SequenceDistributionEntry | None = None + ) -> None: """Finalize a turn by populating all required metadata fields. This method handles: @@ -256,10 +325,12 @@ def _finalize_turn(self, turn: Turn) -> None: Args: turn: The turn object to finalize. + bucket: The conversation's sticky sequence_distribution bucket, + forwarded so max_tokens is drawn from the same class. """ if turn.model is None: turn.model = self._select_model_name() - self._set_max_tokens(turn) + self._set_max_tokens(turn, bucket) # Clear cached sequence lengths for this turn to free memory turn_id = id(turn) diff --git a/src/aiperf/dataset/composer/synthetic.py b/src/aiperf/dataset/composer/synthetic.py index 83fb6386a7..88a178d05a 100644 --- a/src/aiperf/dataset/composer/synthetic.py +++ b/src/aiperf/dataset/composer/synthetic.py @@ -13,6 +13,7 @@ if TYPE_CHECKING: from aiperf.config.resolution.plan import BenchmarkRun + from aiperf.config.types import SequenceDistributionEntry def _expected(distribution: object | None) -> float: @@ -22,11 +23,6 @@ def _expected(distribution: object | None) -> float: return float(getattr(distribution, "expected_value", 0.0)) -def _stddev_int(distribution: object | None) -> int: - """Integer stddev of a distribution (Normal only) or 0.""" - return int(getattr(distribution, "stddev", 0) or 0) - - class SyntheticDatasetComposer(BaseDatasetComposer): def __init__(self, *, run: BenchmarkRun, tokenizer: Tokenizer | None, **kwargs): super().__init__(run=run, tokenizer=tokenizer, **kwargs) @@ -49,10 +45,8 @@ def __init__(self, *, run: BenchmarkRun, tokenizer: Tokenizer | None, **kwargs): # user didn't configure them — apply the canonical defaults # (turn=1, batch=1, delay=0) explicitly here. self._num_entries = dataset.entries - self._turn_mean = max(1, int(_expected(dataset.turns))) - self._turn_stddev = _stddev_int(dataset.turns) - self._turn_delay_mean = _expected(dataset.turn_delay) - self._turn_delay_stddev = _stddev_int(dataset.turn_delay) + self._turns_dist = dataset.turns + self._turn_delay_dist = dataset.turn_delay self._turn_delay_ratio = dataset.turn_delay_ratio self._prompt_batch_size = ( dataset.prompts.batch_size if dataset.prompts is not None else 1 @@ -66,13 +60,13 @@ def __init__(self, *, run: BenchmarkRun, tokenizer: Tokenizer | None, **kwargs): self._video_batch_size = ( dataset.video.batch_size if dataset.video is not None else 1 ) - self._isl_stddev = _stddev_int( - dataset.prompts.isl if dataset.prompts is not None else None - ) # Inclusion flags (computed once at init). - self._include_prompt = ( - _expected(dataset.prompts.isl if dataset.prompts is not None else None) > 0 + # A sequence_distribution drives ISL/OSL directly, so text is included + # even when prompts.isl is left unset (the isl field is ignored then). + self._include_prompt = dataset.prompts is not None and ( + _expected(dataset.prompts.isl) > 0 + or bool(dataset.prompts.sequence_distribution) ) self._include_image = ( dataset.images is not None @@ -111,14 +105,22 @@ def create_dataset(self) -> list[Conversation]: for _ in range(self._num_entries): conversation = Conversation(session_id=self.session_id_generator.next()) - num_turns = self._turn_sampler_rng.sample_positive_normal_integer( - self._turn_mean, - self._turn_stddev, + # The conversation's workload class: drawn once, kept for all turns. + bucket = ( + self._seq_distribution.sample_bucket() + if self._seq_distribution is not None + else None + ) + + num_turns = ( + self._turns_dist.sample_int(self._turn_sampler_rng) + if self._turns_dist is not None + else 1 ) self.logger.debug("Creating conversation with %d turns", num_turns) for turn_idx in range(num_turns): - turn = self._create_turn(is_first=(turn_idx == 0)) + turn = self._create_turn(is_first=(turn_idx == 0), bucket=bucket) conversation.turns.append(turn) conversations.append(conversation) @@ -126,7 +128,9 @@ def create_dataset(self) -> list[Conversation]: self._finalize_conversations(conversations) return conversations - def _create_turn(self, is_first: bool) -> Turn: + def _create_turn( + self, is_first: bool, bucket: SequenceDistributionEntry | None = None + ) -> Turn: """Create a turn object that contains synthetic payloads to send. It generates multi-modal data (e.g. text, image, audio) using synthetic @@ -134,6 +138,8 @@ def _create_turn(self, is_first: bool) -> Turn: Args: is_first: Whether the turn is the first turn in the conversation. + bucket: The conversation's sticky sequence_distribution bucket + (None when no sequence_distribution is configured). Returns: Turn: A dataset representation of a single turn. @@ -141,7 +147,7 @@ def _create_turn(self, is_first: bool) -> Turn: turn = Turn() if self.include_prompt: - turn.texts.append(self._generate_text_payloads(turn, is_first)) + turn.texts.append(self._generate_text_payloads(turn, is_first, bucket)) if self.include_image: turn.images.append(self._generate_image_payloads()) if self.include_audio: @@ -149,12 +155,13 @@ def _create_turn(self, is_first: bool) -> Turn: if self.include_video: turn.videos.append(self._generate_video_payloads()) - if not is_first and self._turn_delay_mean > 0: - delay = self._delay_sampler_rng.sample_positive_normal_integer( - int(self._turn_delay_mean), - self._turn_delay_stddev, - ) - turn.delay = delay * self._turn_delay_ratio + if ( + not is_first + and self._turn_delay_dist is not None + and self._turn_delay_dist.expected_value > 0 + ): + delay = self._turn_delay_dist.sample(self._delay_sampler_rng) + turn.delay = max(0.0, delay) * self._turn_delay_ratio if not turn.texts and not turn.images and not turn.audios and not turn.videos: self.logger.warning( @@ -163,16 +170,23 @@ def _create_turn(self, is_first: bool) -> Turn: "setting the mean to a positive value." ) - self._finalize_turn(turn) + self._finalize_turn(turn, bucket) return turn - def _generate_text_payloads(self, turn: Turn, is_first: bool) -> Text: + def _generate_text_payloads( + self, + turn: Turn, + is_first: bool, + bucket: SequenceDistributionEntry | None = None, + ) -> Text: """Generate text payloads for a single turn. Args: turn: The turn object (used for caching sequence lengths) is_first: Whether the turn is the first turn in the conversation. + bucket: The conversation's sticky sequence_distribution bucket + (None when no sequence_distribution is configured). Returns: Text: A text payload object. @@ -190,14 +204,18 @@ def _generate_text_payloads(self, turn: Turn, is_first: bool) -> Text: # Sample ISL/OSL pair for this request (cached for consistency) turn_id = id(turn) - isl, _ = self._get_turn_sequence_lengths(turn_id) - - # Preserve original variance unless sequence distribution is active - stddev = 0 if self._seq_distribution is not None else self._isl_stddev + isl, _ = self._get_turn_sequence_lengths( + turn_id, is_first=is_first, bucket=bucket + ) + # `isl` was already drawn from the full typed distribution inside + # `_get_turn_sequence_lengths` via `sample_int(rng)`. Passing a + # stddev here would apply a SECOND normal sample around that draw, + # compounding the variance to stddev*sqrt(2). Per-turn variance + # belongs to the distribution; per-prompt generation is + # deterministic at the turn's target. for _ in range(self._prompt_batch_size): - # Generate prompt content using the sampled input sequence length - content = self.prompt_generator.generate(mean=isl, stddev=stddev) + content = self.prompt_generator.generate(mean=isl, stddev=0) # Add prefix prompt if this is the first turn and prefix is enabled if is_first and self.prefix_prompt_enabled: diff --git a/src/aiperf/dataset/composer/synthetic_rankings.py b/src/aiperf/dataset/composer/synthetic_rankings.py index 459ce54dce..f5e888a34a 100644 --- a/src/aiperf/dataset/composer/synthetic_rankings.py +++ b/src/aiperf/dataset/composer/synthetic_rankings.py @@ -16,11 +16,6 @@ from aiperf.config.resolution.plan import BenchmarkRun -def _stddev(distribution: object | None) -> int: - """Return integer stddev of a distribution (Normal only) or 0 otherwise.""" - return int(getattr(distribution, "stddev", 0) or 0) - - class SyntheticRankingsDatasetComposer(BaseDatasetComposer): """Composer that generates synthetic data for the Rankings endpoint. @@ -56,13 +51,9 @@ def create_dataset(self) -> list[Conversation]: Each conversation contains one turn with one query and multiple passages. """ conversations: list[Conversation] = [] - num_passages_mean = int(self._rankings.passages.expected_value) - num_passages_std = _stddev(self._rankings.passages) for _ in range(self._num_entries): - num_passages = self._passages_rng.sample_positive_normal_integer( - num_passages_mean, num_passages_std - ) + num_passages = self._rankings.passages.sample_int(self._passages_rng) conversation = Conversation(session_id=self.session_id_generator.next()) turn = self._create_turn(num_passages=num_passages) conversation.turns.append(turn) @@ -85,23 +76,15 @@ def _create_turn(self, num_passages: int) -> Turn: turn = Turn() query_text = self.prompt_generator.generate_prompt( - self.prompt_generator.calculate_num_tokens( - int(self._rankings.query_tokens.expected_value), - _stddev(self._rankings.query_tokens), - ) + self._rankings.query_tokens.sample_int(self._query_token_rng) ) query = Text(name="query", contents=[query_text]) # Generate passages with rankings-specific token counts (per passage) passages = Text(name="passages") - passage_token_mean = int(self._rankings.passage_tokens.expected_value) - passage_token_stddev = _stddev(self._rankings.passage_tokens) for _ in range(num_passages): passage_text = self.prompt_generator.generate_prompt( - self.prompt_generator.calculate_num_tokens( - passage_token_mean, - passage_token_stddev, - ) + self._rankings.passage_tokens.sample_int(self._passages_token_rng) ) passages.contents.append(passage_text) diff --git a/src/aiperf/dataset/generator/audio.py b/src/aiperf/dataset/generator/audio.py index 8586f5c51f..cbb393a29f 100644 --- a/src/aiperf/dataset/generator/audio.py +++ b/src/aiperf/dataset/generator/audio.py @@ -135,15 +135,11 @@ def generate(self, *args, **kwargs) -> str: - bit depth is not supported (must be 8, 16, 24, or 32) """ length_dist = self.config.length - length_mean = float(length_dist.expected_value) - length_stddev = float(getattr(length_dist, "stddev", 0) or 0) - if length_mean < 0.01: + if float(length_dist.expected_value) < 0.01: raise ConfigurationError("Audio length must be greater than 0.01 seconds") - # Sample audio length (in seconds) using rejection sampling - audio_length = self._duration_rng.sample_normal( - length_mean, length_stddev, lower=0.01 - ) + # Sample audio length (in seconds) from the full typed distribution + audio_length = max(0.01, length_dist.sample(self._duration_rng)) # Randomly select sampling rate and bit depth sampling_rate_hz = int( diff --git a/src/aiperf/dataset/generator/image.py b/src/aiperf/dataset/generator/image.py index 40b2115a19..84a833ed87 100644 --- a/src/aiperf/dataset/generator/image.py +++ b/src/aiperf/dataset/generator/image.py @@ -107,12 +107,8 @@ def generate(self, *args, **kwargs) -> str: width_dist = self.config.width height_dist = self.config.height - width = self._dimensions_rng.sample_positive_normal_integer( - int(width_dist.expected_value), int(getattr(width_dist, "stddev", 0) or 0) - ) - height = self._dimensions_rng.sample_positive_normal_integer( - int(height_dist.expected_value), int(getattr(height_dist, "stddev", 0) or 0) - ) + width = width_dist.sample_int(self._dimensions_rng) + height = height_dist.sample_int(self._dimensions_rng) image = self._create_source_image(width, height) self.debug( diff --git a/src/aiperf/dataset/generator/prompt.py b/src/aiperf/dataset/generator/prompt.py index f375817bf5..763e135db8 100644 --- a/src/aiperf/dataset/generator/prompt.py +++ b/src/aiperf/dataset/generator/prompt.py @@ -54,6 +54,7 @@ def __init__( self._length_rng = rng.derive("dataset.prompt.length") self._corpus_rng = rng.derive("dataset.prompt.corpus") self._prefix_rng = rng.derive("dataset.prompt.prefix") + self._context_len_rng = rng.derive("dataset.prompt.context_length") super().__init__(tokenizer=tokenizer, **kwargs) @@ -371,14 +372,15 @@ def _generate_shared_system_prompt(self) -> None: if self._tokenized_corpus is None: raise NotInitializedError("Tokenized corpus is not initialized.") - length = ( + dist = ( self.prefix_prompts.shared_system_length if self.prefix_prompts is not None else None ) - if length is None: + if dist is None: return + length = dist.sample_int(self._context_len_rng) self._shared_system_prompt = self.generate_prompt(length) self.debug(lambda: f"Generated shared system prompt with {length} tokens") @@ -417,12 +419,12 @@ def generate_user_context_prompt(self, session_index: int) -> str: if self._tokenized_corpus is None: raise NotInitializedError("Tokenized corpus is not initialized.") - length = ( + dist = ( self.prefix_prompts.user_context_length if self.prefix_prompts is not None else None ) - if length is None: + if dist is None: raise InvalidStateError( "User context prompt length is not configured. " "Ensure --user-context-prompt-length is specified." @@ -430,6 +432,7 @@ def generate_user_context_prompt(self, session_index: int) -> str: # Generate new prompts on-demand as needed while session_index >= len(self._user_context_prompts): + length = dist.sample_int(self._context_len_rng) new_prompt = self.generate_prompt(length) self._user_context_prompts.append(new_prompt) self.debug( diff --git a/src/aiperf/timing/strategies/user_centric_rate.py b/src/aiperf/timing/strategies/user_centric_rate.py index 360e3770c2..dcd063507b 100644 --- a/src/aiperf/timing/strategies/user_centric_rate.py +++ b/src/aiperf/timing/strategies/user_centric_rate.py @@ -178,11 +178,18 @@ def _generate_next_user( self._next_user_id += 1 sampled = self._conversation_source.next(x_correlation_id=str(user_id)) + # Virtual-history turn budgets derive from the dataset AVERAGE turn + # count; the drawn conversation may be shorter (per-conversation + # `turns` distributions), and the worker rejects num_turns beyond + # the conversation's real length. + conversation_turns = len(sampled.metadata.turns) user = User( user_id=user_id, sampled=sampled, next_send_time=target_perf_sec or 0.0, - max_turns=max_turns or len(sampled.metadata.turns), + max_turns=min(max_turns, conversation_turns) + if max_turns + else conversation_turns, order=order or 0, ) self._session_to_user[user.x_correlation_id] = user diff --git a/tests/unit/common/models/test_sequence_distribution.py b/tests/unit/common/models/test_sequence_distribution.py index 5f0a2dec01..cd7538286a 100644 --- a/tests/unit/common/models/test_sequence_distribution.py +++ b/tests/unit/common/models/test_sequence_distribution.py @@ -14,6 +14,7 @@ import numpy as np import pytest +from pytest import param from aiperf.common import random_generator as rng from aiperf.common.models.sequence_distribution import ( @@ -61,32 +62,64 @@ def test_invalid_output_length(self): SequenceLengthPair(256, -1, 50.0) def test_invalid_probability(self): - """Test validation of probability values.""" - with pytest.raises(ValueError, match="Probability must be in \\[0,100\\]"): + """Negative weights are rejected; zero is a valid disabled-bucket weight.""" + with pytest.raises( + ValueError, match="Probability weight must be finite and non-negative" + ): SequenceLengthPair(256, 128, -10.0) - with pytest.raises(ValueError): - SequenceLengthPair(256, 128, 110.0) + # Zero is a valid relative weight (disabled bucket); it must not raise. + SequenceLengthPair(256, 128, 0.0) + + @pytest.mark.parametrize( + "bad", + [ + param(float("nan"), id="nan"), + param(float("inf"), id="inf"), + param(float("-inf"), id="neg-inf"), + ], + ) # fmt: skip + def test_non_finite_probability_rejected(self, bad): + """NaN/inf weights poison cumulative normalization; reject at construction.""" + with pytest.raises(ValueError, match="Probability weight must be finite"): + SequenceLengthPair(256, 128, bad) + + @pytest.mark.parametrize( + "bad", + [ + param(float("nan"), id="nan"), + param(float("inf"), id="inf"), + ], + ) # fmt: skip + def test_non_finite_stddev_rejected(self, bad): + """NaN/inf stddevs poison per-request length sampling; reject at construction.""" + with pytest.raises(ValueError, match="stddev must be finite"): + SequenceLengthPair(256, 128, 50.0, input_seq_len_stddev=bad) + with pytest.raises(ValueError, match="stddev must be finite"): + SequenceLengthPair(256, 128, 50.0, output_seq_len_stddev=bad) def test_invalid_input_stddev(self): """Test validation of negative input sequence length standard deviation.""" with pytest.raises( - ValueError, match="Input sequence length stddev must be non-negative" + ValueError, + match="Input sequence length stddev must be finite and non-negative", ): SequenceLengthPair(256, 128, 50.0, input_seq_len_stddev=-1.0) def test_invalid_output_stddev(self): """Test validation of negative output sequence length standard deviation.""" with pytest.raises( - ValueError, match="Output sequence length stddev must be non-negative" + ValueError, + match="Output sequence length stddev must be finite and non-negative", ): SequenceLengthPair(256, 128, 50.0, output_seq_len_stddev=-2.0) def test_boundary_probabilities(self): - """Test boundary probability values.""" - # Should work - SequenceLengthPair(256, 128, 0.0) + """Test boundary probability values (positive relative weights).""" + # Weights are relative: small and large positive values are both valid. + SequenceLengthPair(256, 128, 0.001) SequenceLengthPair(256, 128, 100.0) + SequenceLengthPair(256, 128, 250.0) def test_immutability(self): """Test that pairs are immutable.""" @@ -225,16 +258,17 @@ def test_empty_pairs_validation(self): with pytest.raises(ValueError, match="at least one sequence length pair"): SequenceLengthDistribution([]) - def test_probability_sum_validation(self): - """Test validation of probability sum.""" - # Probabilities don't sum to 100.0 - invalid_pairs = [ + def test_relative_weights_do_not_require_sum_100(self): + """Probabilities are relative weights; a non-100 sum is accepted.""" + # Sum = 70.0 -> normalized to 30/70 and 40/70 at sampling time. + pairs = [ SequenceLengthPair(256, 128, 30.0), - SequenceLengthPair(512, 256, 40.0), # Sum = 70.0 + SequenceLengthPair(512, 256, 40.0), ] - with pytest.raises(ValueError, match="must sum to 100.0"): - SequenceLengthDistribution(invalid_pairs) + dist = SequenceLengthDistribution(pairs) + for isl, osl in dist.sample_batch(100): + assert (isl, osl) in [(256, 128), (512, 256)] def test_probability_sum_tolerance(self): """Test that small floating-point errors are tolerated.""" @@ -248,6 +282,17 @@ def test_probability_sum_tolerance(self): dist = SequenceLengthDistribution(pairs) assert dist is not None + def test_zero_weight_pair_never_sampled(self): + """A zero-weight pair forms a zero-width cumulative interval and is never + selected by np.searchsorted(side='right').""" + pairs = [ + SequenceLengthPair(100, 10, 0.0), + SequenceLengthPair(4000, 200, 50.0), + ] + dist = SequenceLengthDistribution(pairs) + samples = dist.sample_batch(500) + assert all(isl == 4000 and osl == 200 for isl, osl in samples) + def test_statistics_calculation(self): """Test distribution statistics calculation.""" dist = SequenceLengthDistribution(self.multi_pair) @@ -262,6 +307,22 @@ def test_statistics_calculation(self): assert stats["num_pairs"] == 2 assert abs(stats["total_probability"] - 100.0) < 1e-10 + def test_statistics_normalizes_by_actual_total_not_100(self): + """Relative weights that do not sum to 100 normalize by their true total. + + Weights 50 and 1 -> expected values weighted by 50/51 and 1/51, not /100. + """ + pairs = [ + SequenceLengthPair(1000, 200, 50.0), + SequenceLengthPair(2000, 400, 1.0), + ] + dist = SequenceLengthDistribution(pairs) + stats = dist.get_statistics() + + assert stats["expected_isl"] == pytest.approx(1000 * 50 / 51 + 2000 * 1 / 51) + assert stats["expected_osl"] == pytest.approx(200 * 50 / 51 + 400 * 1 / 51) + assert stats["total_probability"] == pytest.approx(51.0) + def test_string_representation(self): """Test string representation.""" dist = SequenceLengthDistribution(self.multi_pair) @@ -359,11 +420,15 @@ def test_semicolon_format_mixed_stddev(self): assert dist.pairs[0] == SequenceLengthPair(256, 128, 60.0, 10.0, 0.0) assert dist.pairs[1] == SequenceLengthPair(512, 256, 40.0, 0.0, 15.0) - def test_semicolon_format_invalid_fractions(self): - """Test that fractions are properly rejected (percentage-only enforcement).""" + def test_semicolon_format_relative_weight_fractions(self): + """Fractional weights are accepted as relative weights (0.6:0.4 -> 60/40).""" dist_str = "256,128:0.6;512,256:0.4" - with pytest.raises(ValueError, match="Probabilities must sum to 100.0"): - DistributionParser.parse(dist_str) + dist = DistributionParser.parse(dist_str) + + assert len(dist.pairs) == 2 + samples = dist.sample_batch(2000) + small_frac = sum(1 for isl, _ in samples if isl == 256) / len(samples) + assert 0.5 < small_frac < 0.7 def test_bracket_format_parsing(self): """Test parsing bracket format with percentages.""" @@ -401,6 +466,15 @@ def test_json_format_with_stddev(self): assert dist.pairs[0] == SequenceLengthPair(256, 128, 60.0, 10.0, 5.0) assert dist.pairs[1] == SequenceLengthPair(512, 256, 40.0, 20.0, 15.0) + def test_zero_weight_pair_parses_and_never_sampled(self): + """A legacy string with a zero-weight pair is valid; the pair never samples.""" + dist = DistributionParser.parse("100,10:0;4000,10:50") + + assert len(dist.pairs) == 2 + assert dist.pairs[0].probability == 0.0 + samples = dist.sample_batch(500) + assert all(isl == 4000 for isl, _ in samples) + def test_single_pair_parsing(self): """Test parsing single pair.""" dist_str = "1024,512:100" @@ -424,9 +498,9 @@ def test_invalid_format_parsing(self): "256,128", # Missing probability "256:50", # Missing OSL "invalid", - "256,128:110", # Invalid probability (>100) - "256,128:-10", # Invalid probability (<0) - "256,128:0.6", # Fraction not allowed (percentage-only) + "256,128:-10", # Invalid probability weight (must be positive) + "256,128:inf", # Non-finite probability weight + "256,128:nan", # Non-finite probability weight '{"invalid": "json"}', # Invalid JSON structure "256|,128:100", # Empty stddev "256|-5,128:100", # Negative stddev @@ -671,13 +745,24 @@ class MockComposer(BaseDatasetComposer): def create_dataset(self): return [] - # Set up composer with distribution + # Set up composer with distribution. The runtime _seq_distribution is a + # _TypedSequenceDistribution (fixed-scalar buckets here for determinism). + from aiperf.config.types import SequenceDistributionEntry + from aiperf.dataset.composer.base import _TypedSequenceDistribution + composer = MockComposer.__new__(MockComposer) - dist = DistributionParser.parse("128,64:50;256,128:50") - composer._seq_distribution = dist + entries = [ + SequenceDistributionEntry.model_validate( + {"isl": 128, "osl": 64, "probability": 50} + ), + SequenceDistributionEntry.model_validate( + {"isl": 256, "osl": 128, "probability": 50} + ), + ] + composer._seq_distribution = _TypedSequenceDistribution( + entries, rng.derive("test_composer") + ) composer._turn_sequence_cache = {} - # Use the global RNG instead of _seq_rng - composer._rng = rng.derive("test_composer") # Create a turn and get its ID turn = Turn() @@ -714,15 +799,20 @@ class MockComposer(BaseDatasetComposer): def create_dataset(self): return [] - # Set up composer with distribution + # Set up composer with a single fixed bucket for predictable results. + from aiperf.config.types import SequenceDistributionEntry + from aiperf.dataset.composer.base import _TypedSequenceDistribution + composer = MockComposer.__new__(MockComposer) - dist = DistributionParser.parse( - "100,50:100" - ) # Single pair for predictable results - composer._seq_distribution = dist + composer._seq_distribution = _TypedSequenceDistribution( + [ + SequenceDistributionEntry.model_validate( + {"isl": 100, "osl": 50, "probability": 100} + ) + ], + rng.derive("test_composer2"), + ) composer._turn_sequence_cache = {} - # Use the global RNG instead of _seq_rng - composer._rng = rng.derive("test_composer2") # Create two different turns turn1 = Turn() diff --git a/tests/unit/config/test_camel_case_config.py b/tests/unit/config/test_camel_case_config.py index 13bf56b49e..e46d788a8d 100644 --- a/tests/unit/config/test_camel_case_config.py +++ b/tests/unit/config/test_camel_case_config.py @@ -477,8 +477,8 @@ def test_dataset_camel_case_fields(self) -> None: ds = config.benchmark.get_dataset("main") assert ds.prompts.batch_size == 4 assert ds.prompts.block_size == 128 - assert ds.prefix_prompts.shared_system_length == 100 - assert ds.prefix_prompts.user_context_length == 50 + assert ds.prefix_prompts.shared_system_length.expected_value == 100 + assert ds.prefix_prompts.user_context_length.expected_value == 50 assert ds.turn_delay.expected_value == 200.0 assert ds.turn_delay_ratio == 2.0 diff --git a/tests/unit/config/test_distributions.py b/tests/unit/config/test_distributions.py index 852b5444d1..d4d6205f7c 100644 --- a/tests/unit/config/test_distributions.py +++ b/tests/unit/config/test_distributions.py @@ -19,6 +19,7 @@ from pytest import param from aiperf.common import random_generator as rng +from aiperf.config.dataset.content import PrefixPromptConfig, PromptConfig from aiperf.config.distributions import ( Distribution, EmpiricalDistribution, @@ -30,7 +31,10 @@ PeakEntry, SamplingDistribution, ) -from aiperf.config.types import SequenceDistributionEntry +from aiperf.config.types import ( + SequenceDistributionEntry, + validate_probability_distribution, +) # TypeAdapter for the discriminated union _TA = TypeAdapter(SamplingDistribution) @@ -538,3 +542,168 @@ def test_osl_stddev_shorthand_creates_normal(self) -> None: assert isinstance(entry.osl, NormalDistribution) assert entry.osl.mean == 256.0 assert entry.osl.stddev == 25.0 + + +# ============================================================ +# 9. first_turn_isl + relative bucket weights +# ============================================================ + + +def test_first_turn_isl_without_isl_raises(): + with pytest.raises(ValidationError, match="first_turn_isl"): + PromptConfig(first_turn_isl={"mean": 20000, "stddev": 10}) + + +def test_sequence_distribution_relative_weights_accepted(): + cfg = PromptConfig( + sequence_distribution=[ + {"isl": 60000, "osl": 100, "probability": 50}, + {"isl": 400000, "osl": 10, "probability": 1}, + ] + ) + assert len(cfg.sequence_distribution) == 2 + + +def test_sequence_distribution_zero_total_weight_rejected(): + with pytest.raises(ValidationError): + PromptConfig( + sequence_distribution=[ + {"isl": 100, "osl": 10, "probability": 0}, + ] + ) + + +def test_legacy_string_parser_accepts_relative_weights(): + from aiperf.common.models.sequence_distribution import DistributionParser + + dist = DistributionParser.parse("100,10:50;4000,10:1") + samples = [dist.sample() for _ in range(2000)] + small = sum(1 for isl, _ in samples if isl == 100) / len(samples) + assert small > 0.94 + + +# ============================================================ +# 10. Per-bucket first_turn_isl + top-level exclusivity +# ============================================================ + + +class TestSequenceDistributionEntryFiniteWeights: + @pytest.mark.parametrize( + "bad", + [ + param(float("nan"), id="nan"), + param(float("inf"), id="inf"), + param(float("-inf"), id="neg-inf"), + ], + ) # fmt: skip + def test_non_finite_probability_rejected(self, bad): + """inf/inf normalizes to NaN in _TypedSequenceDistribution; reject at config.""" + with pytest.raises(ValidationError): + SequenceDistributionEntry.model_validate( + {"isl": 128, "osl": 64, "probability": bad} + ) + + def test_overflowing_total_rejected(self): + """Per-entry weights can be finite while their sum overflows to inf.""" + entries = [ + SequenceDistributionEntry.model_validate( + {"isl": 128, "osl": 64, "probability": 1e308} + ) + for _ in range(2) + ] + with pytest.raises(ValueError, match="finite positive total"): + validate_probability_distribution(entries) + + +class TestSequenceDistributionEntryFirstTurnIsl: + def test_entry_accepts_first_turn_isl_distribution(self): + entry = SequenceDistributionEntry.model_validate( + { + "first_turn_isl": {"p50": 20000, "p99": 100000, "mean": 30000}, + "isl": {"mean": 300, "stddev": 100}, + "osl": {"mean": 250, "stddev": 80}, + "probability": 50, + } + ) + assert entry.first_turn_isl is not None + assert entry.first_turn_isl.expected_value == pytest.approx(30000, rel=0.01) + + def test_entry_first_turn_isl_defaults_to_none(self): + entry = SequenceDistributionEntry.model_validate( + {"isl": 128, "osl": 64, "probability": 100} + ) + assert entry.first_turn_isl is None + + def test_entry_first_turn_isl_scalar_coerces_to_fixed(self): + entry = SequenceDistributionEntry.model_validate( + {"first_turn_isl": 40000, "isl": 128, "osl": 64, "probability": 100} + ) + assert entry.first_turn_isl.expected_value == 40000 + + +class TestPromptConfigFirstTurnIslSeqDistConflict: + def test_top_level_first_turn_isl_with_sequence_distribution_rejected(self): + with pytest.raises( + ValueError, match="cannot be combined with sequence_distribution" + ): + PromptConfig.model_validate( + { + "first_turn_isl": {"mean": 20000, "stddev": 100}, + "isl": {"mean": 300, "stddev": 100}, + "sequence_distribution": [ + {"isl": 128, "osl": 64, "probability": 100} + ], + } + ) + + def test_per_bucket_first_turn_isl_with_sequence_distribution_accepted(self): + cfg = PromptConfig.model_validate( + { + "sequence_distribution": [ + { + "first_turn_isl": {"mean": 20000, "stddev": 100}, + "isl": 128, + "osl": 64, + "probability": 100, + } + ] + } + ) + assert cfg.sequence_distribution[0].first_turn_isl is not None + + def test_first_turn_isl_without_isl_still_rejected(self): + with pytest.raises(ValueError, match="first_turn_isl requires isl"): + PromptConfig.model_validate( + {"first_turn_isl": {"mean": 20000, "stddev": 100}} + ) + + +class TestPrefixPromptLengthDistributions: + def test_scalar_lengths_still_parse_as_fixed(self): + cfg = PrefixPromptConfig.model_validate( + {"shared_system_length": 2048, "user_context_length": 512} + ) + assert cfg.shared_system_length.expected_value == 2048 + assert cfg.user_context_length.expected_value == 512 + + def test_distribution_lengths_accepted(self): + cfg = PrefixPromptConfig.model_validate( + { + "shared_system_length": {"mean": 2048, "stddev": 512, "min": 512}, + "user_context_length": {"p50": 4000, "p99": 32000, "mean": 6000}, + } + ) + assert cfg.shared_system_length.expected_value == pytest.approx(2048, rel=0.01) + assert cfg.user_context_length.expected_value == pytest.approx(6000, rel=0.01) + + def test_sub_one_expected_value_rejected(self): + with pytest.raises(ValueError, match="expected value >= 1"): + PrefixPromptConfig.model_validate( + {"user_context_length": {"mean": 0.2, "stddev": 0}} + ) + + def test_pool_fields_remain_ints(self): + with pytest.raises(ValueError): + PrefixPromptConfig.model_validate( + {"pool_size": 4, "length": {"mean": 100, "stddev": 10}} + ) diff --git a/tests/unit/config/test_percentile_distribution.py b/tests/unit/config/test_percentile_distribution.py new file mode 100644 index 0000000000..26b5017370 --- /dev/null +++ b/tests/unit/config/test_percentile_distribution.py @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import math + +import pytest +from pydantic import TypeAdapter +from pytest import param + +from aiperf.common import random_generator as rng +from aiperf.config.distributions import ( + PercentileDistribution, + SamplingDistribution, +) + +ADAPTER = TypeAdapter(SamplingDistribution) + + +def _percentile(sorted_vals: list[float], q: float) -> float: + idx = min(int(q * len(sorted_vals)), len(sorted_vals) - 1) + return sorted_vals[idx] + + +class TestDiscriminator: + @pytest.mark.parametrize( + "payload", + [ + param({"p50": 50000, "p99": 400000}, id="p50_p99"), + param({"p50": 50000, "p99": 400000, "mean": 60000}, id="with_mean"), + param({"type": "percentile", "p50": 100, "p99": 500}, id="explicit_type"), + ], + ) # fmt: skip + def test_discriminator_p50_key_selects_percentile(self, payload: dict) -> None: + dist = ADAPTER.validate_python(payload) + assert isinstance(dist, PercentileDistribution) + + def test_discriminator_mean_with_p50_not_normal(self) -> None: + # "mean" alone maps to Normal; p50 presence must win over mean. + dist = ADAPTER.validate_python({"mean": 60000, "p50": 50000, "p99": 400000}) + assert isinstance(dist, PercentileDistribution) + + +class TestValidation: + @pytest.mark.parametrize( + "payload,match", + [ + param({"p50": 400000, "p99": 50000}, "p99", id="p99_below_p50"), + param({"p50": 100, "p99": 100}, "p99", id="p99_equal_p50"), + param({"p50": 50000, "p99": 400000, "mean": 40000}, "mean", id="mean_below_p50"), + param({"p50": 50000, "p99": 400000, "mean": 399999}, "infeasible|mean", id="mean_infeasible_high"), + param({"p50": 0, "p99": 100}, "greater than 0", id="p50_zero"), + ], + ) # fmt: skip + def test_invalid_targets_raise(self, payload: dict, match: str) -> None: + with pytest.raises(ValueError, match=match): + ADAPTER.validate_python(payload) + + def test_non_finite_rejected(self) -> None: + with pytest.raises(ValueError): + ADAPTER.validate_python({"p50": 50000, "p99": float("inf")}) + + @pytest.mark.parametrize( + "p50,p99", + [ + param(1, 1e40, id="overflow_ratio"), + param(100, 1e35, id="inf_implied_mean"), + ], + ) # fmt: skip + def test_extreme_ratio_rejected(self, p50: float, p99: float) -> None: + with pytest.raises(ValueError, match="too extreme|finite"): + PercentileDistribution(p50=p50, p99=p99) + + +class TestLogNormalMode: + """p50 + p99 without mean fits a log-normal exactly.""" + + def test_expected_value_is_lognormal_implied_mean(self) -> None: + dist = PercentileDistribution(p50=50000, p99=400000) + sigma = math.log(400000 / 50000) / 2.3263478740408408 + assert dist.expected_value == pytest.approx( + 50000 * math.exp(sigma**2 / 2), rel=1e-9 + ) + + def test_sampled_percentiles_match_targets(self) -> None: + dist = PercentileDistribution(p50=50000, p99=400000) + r = rng.derive("test.percentile.lognormal") + vals = sorted(dist.sample(r) for _ in range(100_000)) + assert _percentile(vals, 0.50) == pytest.approx(50000, rel=0.03) + assert _percentile(vals, 0.99) == pytest.approx(400000, rel=0.05) + + +class TestMixtureMode: + """p50 + p99 + mean solves a two-component mixture (the headline use-case).""" + + def test_headline_use_case_hits_all_three_targets(self) -> None: + dist = PercentileDistribution(p50=50000, p99=400000, mean=60000) + r = rng.derive("test.percentile.mixture") + vals = sorted(dist.sample(r) for _ in range(200_000)) + assert _percentile(vals, 0.50) == pytest.approx(50000, rel=0.03) + assert _percentile(vals, 0.99) == pytest.approx(400000, rel=0.05) + assert sum(vals) / len(vals) == pytest.approx(60000, rel=0.03) + + def test_expected_value_returns_configured_mean(self) -> None: + dist = PercentileDistribution(p50=50000, p99=400000, mean=60000) + assert dist.expected_value == pytest.approx(60000) + + def test_mean_matching_lognormal_uses_lognormal(self) -> None: + sigma = math.log(8.0) / 2.3263478740408408 + implied = 50000 * math.exp(sigma**2 / 2) + dist = PercentileDistribution(p50=50000, p99=400000, mean=implied) + r = rng.derive("test.percentile.implied") + vals = sorted(dist.sample(r) for _ in range(100_000)) + assert _percentile(vals, 0.50) == pytest.approx(50000, rel=0.03) + + @pytest.mark.parametrize( + "p50,p99,mean", + [ + param(50000, 400000, 60000, id="headline"), + param(50000, 400000, 100000, id="mid_mean"), + param(128, 8192, 512, id="small_tokens"), + param(1000, 2000, 1100, id="tight_spread"), + ], + ) # fmt: skip + def test_solver_various_targets_all_hit( + self, p50: float, p99: float, mean: float + ) -> None: + dist = PercentileDistribution(p50=p50, p99=p99, mean=mean) + r = rng.derive(f"test.percentile.{p50}.{mean}") + vals = sorted(dist.sample(r) for _ in range(200_000)) + assert _percentile(vals, 0.50) == pytest.approx(p50, rel=0.05) + assert _percentile(vals, 0.99) == pytest.approx(p99, rel=0.06) + assert sum(vals) / len(vals) == pytest.approx(mean, rel=0.04) + + +class TestComposition: + def test_min_max_clamps_compose(self) -> None: + dist = ADAPTER.validate_python( + {"p50": 50000, "p99": 400000, "mean": 60000, "max": 600000, "min": 1000} + ) + r = rng.derive("test.percentile.clamp") + vals = [dist.sample(r) for _ in range(50_000)] + assert max(vals) <= 600000 + assert min(vals) >= 1000 + + def test_sample_int_returns_positive_ints(self) -> None: + dist = PercentileDistribution(p50=50, p99=400) + r = rng.derive("test.percentile.int") + vals = [dist.sample_int(r) for _ in range(1000)] + assert all(isinstance(v, int) and v >= 1 for v in vals) + + def test_serialization_round_trip(self) -> None: + dist = PercentileDistribution(p50=50000, p99=400000, mean=60000) + again = ADAPTER.validate_python(dist.model_dump()) + r = rng.derive("test.percentile.roundtrip") + assert isinstance(again, PercentileDistribution) + assert again.sample(r) > 0 + + def test_repr_shows_targets(self) -> None: + assert "percentile" in repr(PercentileDistribution(p50=100, p99=500)) diff --git a/tests/unit/dataset/composer/test_base_composer.py b/tests/unit/dataset/composer/test_base_composer.py index c97d2444c7..43449c7905 100644 --- a/tests/unit/dataset/composer/test_base_composer.py +++ b/tests/unit/dataset/composer/test_base_composer.py @@ -71,17 +71,15 @@ def test_initialization_with_sequence_distribution( run=make_run(sequence_dist_config), tokenizer=mock_tokenizer ) - # Distribution was built directly from the v2 entries (not via - # DistributionParser.parse, which only accepts strings). - assert composer._seq_distribution is not None - pairs = composer._seq_distribution.pairs - assert len(pairs) == 2 - assert pairs[0].input_seq_len == 100 - assert pairs[0].output_seq_len == 25 - assert pairs[0].probability == 50.0 - assert pairs[1].input_seq_len == 200 - assert pairs[1].output_seq_len == 50 - assert pairs[1].probability == 50.0 + # The runtime sampler is built directly from the typed v2 entries (not + # via DistributionParser.parse, which only accepts strings). Each bucket + # has stddev 0, so sampling is deterministic per bucket: every draw must + # be one of the two configured (isl, osl) pairs and both must appear. + # Buckets are drawn per conversation (sticky), then lengths per turn. + dist = composer._seq_distribution + assert dist is not None + samples = {dist.sample_lengths(dist.sample_bucket()) for _ in range(100)} + assert samples == {(100, 25), (200, 50)} def test_model_selection_round_robin(self, base_config, mock_tokenizer): """Test round robin model selection.""" @@ -136,24 +134,23 @@ def test_get_turn_sequence_lengths_with_distribution( def test_get_turn_sequence_lengths_without_distribution( self, base_config, mock_tokenizer ): - """Test getting sequence lengths without distribution (fallback).""" + """Without a seq-dist, ISL/OSL are drawn per turn from the typed + distributions (not flattened to the distribution mean).""" composer = ConcreteBaseComposer( run=make_run(base_config), tokenizer=mock_tokenizer ) turn_id = 12345 - result = composer._get_turn_sequence_lengths(turn_id) + isl, osl = composer._get_turn_sequence_lengths(turn_id) - # Should use fallback values from config - expected = ( - base_config.prompt_input_tokens_mean, - base_config.prompt_output_tokens_mean, - ) - assert result == expected + # ISL ~ Normal(mean=100, stddev=10); OSL ~ Normal(mean=50, stddev=5). + # Sampled per turn, so they land near (not exactly at) the means. + assert 60 < isl < 140 + assert 30 < osl < 70 - # Should be cached + # Cached for reuse within the same turn (ISL/OSL pairing). assert turn_id in composer._turn_sequence_cache - assert composer._turn_sequence_cache[turn_id] == expected + assert composer._turn_sequence_cache[turn_id] == (isl, osl) def test_clear_turn_cache(self, sequence_dist_config, mock_tokenizer): """Test clearing turn cache.""" diff --git a/tests/unit/dataset/composer/test_distribution_wiring.py b/tests/unit/dataset/composer/test_distribution_wiring.py new file mode 100644 index 0000000000..1d54446725 --- /dev/null +++ b/tests/unit/dataset/composer/test_distribution_wiring.py @@ -0,0 +1,338 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Typed SamplingDistribution shapes must actually sample at runtime. + +Regression tests for the flattening bug where lognormal/multimodal/ +empirical/percentile ISL/OSL configs silently collapsed to a constant at +the distribution mean (only Normal's stddev survived). +""" + +import statistics +from typing import Any + +import pytest + +from aiperf.dataset.composer.synthetic import SyntheticDatasetComposer +from tests.harness.fake_tokenizer import FakeTokenizer +from tests.unit.conftest import make_benchmark_run + + +def _make_composer( + isl: Any = None, + osl: Any = None, + entries: int = 1000, + sequence_distribution: Any = None, + first_turn_isl: Any = None, + **dataset_overrides: Any, +) -> SyntheticDatasetComposer: + """Build a SyntheticDatasetComposer whose prompts.isl/osl are the given + distribution dicts (or fixed scalars). + + Uses the native BenchmarkConfig path (rather than CLIConfig, which is flat + and cannot express typed distributions): the top-level ``isl``/``osl`` + shortcuts hoist into ``prompts.{isl,osl}``. A ``sequence_distribution`` list + and ``first_turn_isl`` are nested under ``prompts`` (they have no top-level + shorthand). A FakeTokenizer keeps prompt generation cheap. + + The composer clears its per-turn sequence cache inside ``_finalize_turn`` + to free memory; the tests need to observe every per-turn sample, so the + clear is disabled here. Sampling still happens exactly once per turn (the + cache lookup guarantees it), so this does not change what is drawn. + """ + dataset: dict[str, Any] = { + "name": "default", + "type": "synthetic", + "entries": entries, + **dataset_overrides, + } + if isl is not None: + dataset["isl"] = isl + if osl is not None: + dataset["osl"] = osl + if sequence_distribution is not None: + dataset.setdefault("prompts", {})["sequence_distribution"] = ( + sequence_distribution + ) + if first_turn_isl is not None: + dataset.setdefault("prompts", {})["first_turn_isl"] = first_turn_isl + + run = make_benchmark_run(extra={"datasets": [dataset]}) + composer = SyntheticDatasetComposer(run=run, tokenizer=FakeTokenizer()) + composer._clear_turn_cache = lambda turn_id: None + return composer + + +class TestPlainIslOslSampling: + def test_lognormal_isl_varies_per_turn(self): + composer = _make_composer(isl={"mean": 1000, "median": 600}) + composer.create_dataset() + isls = [pair[0] for pair in composer._turn_sequence_cache.values()] + assert len(set(isls)) > 10 # was: constant 1000 for every turn + assert statistics.median(isls) < statistics.fmean(isls) # right skew + + def test_multimodal_isl_produces_both_modes(self): + composer = _make_composer( + isl={ + "peaks": [ + {"mean": 100, "stddev": 5, "weight": 50}, + {"mean": 10000, "stddev": 50, "weight": 50}, + ] + } + ) + composer.create_dataset() + isls = [pair[0] for pair in composer._turn_sequence_cache.values()] + low = [v for v in isls if v < 1000] + high = [v for v in isls if v > 5000] + assert len(low) + len(high) == len(isls) + assert 0.3 < len(low) / len(isls) < 0.7 + + def test_empirical_isl_only_configured_values(self): + composer = _make_composer( + isl={ + "points": [{"value": 128, "weight": 50}, {"value": 4096, "weight": 50}] + } + ) + composer.create_dataset() + isls = {pair[0] for pair in composer._turn_sequence_cache.values()} + assert isls <= {128, 4096} + assert isls == {128, 4096} + + def test_percentile_isl_hits_targets(self): + composer = _make_composer( + isl={"p50": 5000, "p99": 40000, "mean": 6000}, entries=4000 + ) + composer.create_dataset() + isls = sorted(pair[0] for pair in composer._turn_sequence_cache.values()) + assert isls[len(isls) // 2] == pytest.approx(5000, rel=0.10) + assert statistics.fmean(isls) == pytest.approx(6000, rel=0.08) + + +class TestNoDoubleSampling: + def test_normal_isl_observed_stddev_matches_configured(self, monkeypatch): + """Configured stddev 50 must yield observed ~50, not ~71 (sqrt(2)x). + + Capture the mean/stddev PromptGenerator.generate receives: variance + must come from the per-turn sample (stddev arg == 0).""" + captured = [] + composer = _make_composer(isl={"mean": 550, "stddev": 50}, entries=500) + + original = composer.prompt_generator.generate + + def spy(*, mean, stddev): + captured.append((mean, stddev)) + return original(mean=mean, stddev=stddev) + + monkeypatch.setattr(composer.prompt_generator, "generate", spy) + composer.create_dataset() + assert all(s == 0 for _, s in captured) + means = [m for m, _ in captured] + assert 35 < statistics.stdev(means) < 65 + + +class TestMaxTokensPairing: + def test_osl_distribution_sets_varied_max_tokens(self): + composer = _make_composer(isl=512, osl={"mean": 256, "stddev": 64}) + conversations = composer.create_dataset() + max_tokens = [t.max_tokens for c in conversations for t in c.turns] + assert len(set(max_tokens)) > 5 + assert statistics.fmean(max_tokens) == pytest.approx(256, rel=0.10) + + def test_osl_and_isl_come_from_same_turn_sample(self): + """max_tokens must equal the OSL cached for that turn (single draw).""" + composer = _make_composer(isl=512, osl={"mean": 256, "stddev": 64}) + conversations = composer.create_dataset() + # every max_tokens value must appear as an OSL in the cache + cached_osls = sorted(p[1] for p in composer._turn_sequence_cache.values()) + turn_max = sorted(t.max_tokens for c in conversations for t in c.turns) + assert turn_max == cached_osls + + def test_osl_zero_mean_disables_max_tokens(self): + composer = _make_composer(isl=512, osl={"mean": 0}) + conversations = composer.create_dataset() + assert all(t.max_tokens is None for c in conversations for t in c.turns) + + +class TestTurnsAndDelayWiring: + def test_lognormal_turns_varies(self): + composer = _make_composer(isl=128, turns={"mean": 6, "median": 4}, entries=300) + conversations = composer.create_dataset() + counts = [len(c.turns) for c in conversations] + assert len(set(counts)) > 3 # was: constant 6 + assert statistics.median(counts) <= statistics.fmean(counts) # right skew + + def test_empirical_turn_delay_only_configured_values(self): + composer = _make_composer( + isl=128, + turns=3, + turn_delay={ + "points": [{"value": 100, "weight": 50}, {"value": 5000, "weight": 50}] + }, + entries=100, + ) + conversations = composer.create_dataset() + delays = { + t.delay for c in conversations for t in c.turns if t.delay is not None + } + assert delays <= {100, 5000, 100.0, 5000.0} + assert len(delays) == 2 # was: constant 2550 (empirical mean) + + +class TestTypedSequenceDistributionBuckets: + def test_lognormal_bucket_actually_skews(self): + composer = _make_composer( + sequence_distribution=[ + {"isl": {"mean": 2000, "median": 1000}, "osl": 100, "probability": 100}, + ] + ) + composer.create_dataset() + isls = [p[0] for p in composer._turn_sequence_cache.values()] + assert len(set(isls)) > 10 # was: constant 2000 (lognormal flattened) + assert statistics.median(isls) < statistics.fmean(isls) + + def test_bucket_pairing_never_crosses(self): + composer = _make_composer( + sequence_distribution=[ + { + "isl": {"mean": 100, "stddev": 5}, + "osl": {"mean": 10, "stddev": 1}, + "probability": 50, + }, + { + "isl": {"mean": 10000, "stddev": 50}, + "osl": {"mean": 1000, "stddev": 10}, + "probability": 50, + }, + ] + ) + composer.create_dataset() + for isl, osl in composer._turn_sequence_cache.values(): + assert (isl < 1000) == (osl < 100) # small isl with small osl only + + def test_bucket_weights_respected(self): + composer = _make_composer( + entries=2000, + sequence_distribution=[ + {"isl": 100, "osl": 10, "probability": 80}, + {"isl": 10000, "osl": 1000, "probability": 20}, + ], + ) + composer.create_dataset() + isls = [p[0] for p in composer._turn_sequence_cache.values()] + small_frac = sum(1 for v in isls if v == 100) / len(isls) + assert 0.72 < small_frac < 0.88 + + def test_percentile_inside_bucket(self): + composer = _make_composer( + entries=4000, + sequence_distribution=[ + { + "isl": {"p50": 5000, "p99": 40000, "mean": 6000}, + "osl": 100, + "probability": 100, + }, + ], + ) + composer.create_dataset() + isls = sorted(p[0] for p in composer._turn_sequence_cache.values()) + assert isls[len(isls) // 2] == pytest.approx(5000, rel=0.10) + + +class TestFirstTurnIsl: + def test_first_turn_uses_starting_distribution_subsequent_use_isl(self): + composer = _make_composer( + isl={"mean": 200, "stddev": 5}, + first_turn_isl={"mean": 20000, "stddev": 10}, + turns=4, + entries=50, + ) + conversations = composer.create_dataset() + cache = composer._turn_sequence_cache + for conv in conversations: + assert cache[id(conv.turns[0])][0] > 10000 # starting context size + for turn in conv.turns[1:]: + assert cache[id(turn)][0] < 1000 # per-turn new input + + def test_first_turn_isl_unset_isl_applies_to_all_turns(self): + composer = _make_composer(isl={"mean": 200, "stddev": 5}, turns=3, entries=30) + conversations = composer.create_dataset() + cache = composer._turn_sequence_cache + for conv in conversations: + for turn in conv.turns: + assert cache[id(turn)][0] < 1000 + + +class TestRelativeBucketWeights: + def test_weights_not_summing_to_100_normalized(self): + composer = _make_composer( + entries=2000, + sequence_distribution=[ + {"isl": 100, "osl": 10, "probability": 50}, + {"isl": 10000, "osl": 1000, "probability": 1}, + ], + ) + composer.create_dataset() + isls = [p[0] for p in composer._turn_sequence_cache.values()] + small_frac = sum(1 for v in isls if v == 100) / len(isls) + assert small_frac > 0.94 # 50:1 ~ 98% + + +class TestZeroWeightBuckets: + def test_zero_weight_bucket_validates_and_never_samples(self): + """A probability: 0 bucket is valid config and is never sampled.""" + from aiperf.common import random_generator as rng + from aiperf.config.types import SequenceDistributionEntry + from aiperf.dataset.composer.base import _TypedSequenceDistribution + + entries = [ + SequenceDistributionEntry.model_validate( + {"isl": 100, "osl": 10, "probability": 0} + ), + SequenceDistributionEntry.model_validate( + {"isl": 5000, "osl": 500, "probability": 100} + ), + ] + dist = _TypedSequenceDistribution(entries, rng.derive("test.zero.bucket")) + draws = [dist.sample_lengths(dist.sample_bucket()) for _ in range(500)] + assert all(isl == 5000 and osl == 500 for isl, osl in draws) + + def test_zero_weight_bucket_end_to_end_via_composer(self): + composer = _make_composer( + entries=500, + sequence_distribution=[ + {"isl": 100, "osl": 10, "probability": 0}, + {"isl": 5000, "osl": 500, "probability": 100}, + ], + ) + composer.create_dataset() + isls = {p[0] for p in composer._turn_sequence_cache.values()} + assert isls == {5000} # zero-weight bucket never contributes + + def test_all_zero_list_rejected_by_validator(self): + from aiperf.config.types import ( + SequenceDistributionEntry, + validate_probability_distribution, + ) + + entries = [ + SequenceDistributionEntry.model_validate( + {"isl": 100, "osl": 10, "probability": 0} + ), + SequenceDistributionEntry.model_validate( + {"isl": 5000, "osl": 500, "probability": 0} + ), + ] + with pytest.raises(ValueError, match="positive total"): + validate_probability_distribution(entries) + + def test_all_zero_list_rejected_by_typed_distribution(self): + from aiperf.common import random_generator as rng + from aiperf.config.types import SequenceDistributionEntry + from aiperf.dataset.composer.base import _TypedSequenceDistribution + + entries = [ + SequenceDistributionEntry.model_validate( + {"isl": 100, "osl": 10, "probability": 0} + ), + ] + with pytest.raises(ValueError, match="positive-weight entry"): + _TypedSequenceDistribution(entries, rng.derive("test.all.zero")) diff --git a/tests/unit/dataset/composer/test_percentile_e2e.py b/tests/unit/dataset/composer/test_percentile_e2e.py new file mode 100644 index 0000000000..f2f3341a89 --- /dev/null +++ b/tests/unit/dataset/composer/test_percentile_e2e.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""End-to-end: a YAML percentile ISL config produces a dataset whose sampled +ISLs hit the configured p50/p99/mean — the headline user scenario.""" + +import statistics + +import pytest + +# Reuse the composer construction helper (native BenchmarkConfig path that can +# express typed distributions) from the wiring regression suite. +from tests.unit.dataset.composer.test_distribution_wiring import _make_composer + + +def test_percentile_isl_end_to_end_hits_p50_p99_mean(): + composer = _make_composer( + isl={"p50": 5000, "p99": 40000, "mean": 6000}, + osl=128, + entries=6000, + ) + composer.create_dataset() + isls = sorted(p[0] for p in composer._turn_sequence_cache.values()) + n = len(isls) + assert isls[n // 2] == pytest.approx(5000, rel=0.08) + assert isls[int(n * 0.99)] == pytest.approx(40000, rel=0.12) + assert statistics.fmean(isls) == pytest.approx(6000, rel=0.06) + + +def test_sticky_bucket_first_turn_percentile_end_to_end(): + """Headline combined scenario: a conversation class whose SEED context + hits {p50, p99, mean} percentile targets while later turns grow by the + per-turn isl — all through one sticky sequence_distribution bucket.""" + composer = _make_composer( + sequence_distribution=[ + { + "first_turn_isl": {"p50": 5000, "p99": 40000, "mean": 6000}, + "isl": {"mean": 300, "stddev": 100}, + "osl": 128, + "probability": 100, + } + ], + entries=4000, + turns=2, + ) + conversations = composer.create_dataset() + first = sorted( + composer._turn_sequence_cache[id(c.turns[0])][0] for c in conversations + ) + later = [ + composer._turn_sequence_cache[id(t)][0] + for c in conversations + for t in c.turns[1:] + ] + n = len(first) + assert first[n // 2] == pytest.approx(5000, rel=0.08) + assert first[int(n * 0.99)] == pytest.approx(40000, rel=0.12) + assert statistics.fmean(first) == pytest.approx(6000, rel=0.06) + assert statistics.fmean(later) == pytest.approx(300, rel=0.10) diff --git a/tests/unit/dataset/composer/test_sticky_sequence_buckets.py b/tests/unit/dataset/composer/test_sticky_sequence_buckets.py new file mode 100644 index 0000000000..ab1dcefa6c --- /dev/null +++ b/tests/unit/dataset/composer/test_sticky_sequence_buckets.py @@ -0,0 +1,124 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Sticky sequence_distribution buckets: each conversation draws ONE bucket +at creation and keeps it for every turn; per-bucket first_turn_isl sizes the +seed context while isl sizes subsequent turns.""" + +import bisect + +from aiperf.common import random_generator as rng +from aiperf.config.dataset.content import PromptConfig +from tests.unit.dataset.composer.test_distribution_wiring import _make_composer + +# Disjoint ranges so a single sample proves bucket membership. +SMALL_BUCKET = { + "isl": {"mean": 150, "stddev": 10}, + "osl": {"mean": 50, "stddev": 5}, + "probability": 50, +} +LARGE_BUCKET = { + "isl": {"mean": 150000, "stddev": 1000}, + "osl": {"mean": 5000, "stddev": 100}, + "probability": 50, +} + + +class TestStickyBuckets: + def test_conversations_never_mix_buckets(self): + composer = _make_composer( + sequence_distribution=[SMALL_BUCKET, LARGE_BUCKET], + entries=300, + turns=4, + ) + conversations = composer.create_dataset() + small = large = 0 + for conv in conversations: + pairs = [composer._turn_sequence_cache[id(t)] for t in conv.turns] + if pairs[0][0] < 1000: + small += 1 + assert all(isl < 1000 and osl < 1000 for isl, osl in pairs) + else: + large += 1 + assert all(isl > 100000 and osl > 1000 for isl, osl in pairs) + # Both classes must actually appear (weights are 50/50). + assert small > 60 and large > 60 + + def test_single_turn_workloads_still_sample_both_buckets(self): + composer = _make_composer( + sequence_distribution=[SMALL_BUCKET, LARGE_BUCKET], entries=400 + ) + composer.create_dataset() + isls = [p[0] for p in composer._turn_sequence_cache.values()] + low = sum(1 for v in isls if v < 1000) + assert 0.3 < low / len(isls) < 0.7 + + +class TestFirstTurnIslRouting: + def test_first_turn_uses_seed_context_later_turns_use_isl(self): + composer = _make_composer( + sequence_distribution=[ + { + "first_turn_isl": {"mean": 20000, "stddev": 500}, + "isl": {"mean": 300, "stddev": 50}, + "osl": 128, + "probability": 100, + } + ], + entries=100, + turns=3, + ) + conversations = composer.create_dataset() + for conv in conversations: + isls = [composer._turn_sequence_cache[id(t)][0] for t in conv.turns] + assert isls[0] > 10000 # seed context + assert all(v < 5000 for v in isls[1:]) # per-turn growth + + def test_bucket_without_first_turn_isl_falls_back_to_isl(self): + composer = _make_composer( + sequence_distribution=[ + {"isl": {"mean": 300, "stddev": 50}, "osl": 128, "probability": 100} + ], + entries=100, + turns=3, + ) + conversations = composer.create_dataset() + for conv in conversations: + isls = [composer._turn_sequence_cache[id(t)][0] for t in conv.turns] + assert all(v < 5000 for v in isls) + + +class TestSingleTurnDrawStreamCompat: + def test_single_turn_draws_match_legacy_ordering(self): + """Per single-turn conversation the RNG stream must see exactly the + legacy draw order: bucket pick (one random()), then ISL, then OSL. + Replays the pre-sticky algorithm on an identically-derived stream. + """ + bucket_dicts = [SMALL_BUCKET, LARGE_BUCKET] + # Sanity: derive() must be reproducible for the replay to be a valid oracle. + assert ( + rng.derive("compat.check").random() == rng.derive("compat.check").random() + ) + + composer = _make_composer(sequence_distribution=bucket_dicts, entries=50) + composer.create_dataset() + actual = list(composer._turn_sequence_cache.values()) + + parsed = PromptConfig.model_validate( + {"sequence_distribution": bucket_dicts} + ).sequence_distribution + replay_rng = rng.derive("composer.sequence.distribution") + total = sum(e.probability for e in parsed) + cumulative: list[float] = [] + acc = 0.0 + for e in parsed: + acc += e.probability / total + cumulative.append(acc) + expected = [] + for _ in range(50): + r = replay_rng.random() + idx = min(bisect.bisect_right(cumulative, r), len(parsed) - 1) + entry = parsed[idx] + expected.append( + (entry.isl.sample_int(replay_rng), entry.osl.sample_int(replay_rng)) + ) + assert actual == expected diff --git a/tests/unit/dataset/composer/test_synthetic_rankings_composer.py b/tests/unit/dataset/composer/test_synthetic_rankings_composer.py index 3f1943b014..cba571ecb9 100644 --- a/tests/unit/dataset/composer/test_synthetic_rankings_composer.py +++ b/tests/unit/dataset/composer/test_synthetic_rankings_composer.py @@ -1,9 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 +import statistics from aiperf.common.models import Conversation, Turn from aiperf.dataset.composer.synthetic_rankings import SyntheticRankingsDatasetComposer +from tests.harness.fake_tokenizer import FakeTokenizer +from tests.unit.conftest import make_benchmark_run from tests.unit.dataset.composer.conftest import make_run @@ -79,6 +82,33 @@ def test_reproducibility_fixed_seed(synthetic_config, mock_tokenizer): assert t1.texts[1].contents == t2.texts[1].contents +def _make_rankings_composer(entries: int, **rankings): + """Build a rankings composer whose rankings.* are typed distributions. + + Uses the native BenchmarkConfig path (CLIConfig is flat and cannot express + lognormal/empirical distributions) with a FakeTokenizer. + """ + dataset = { + "name": "default", + "type": "synthetic", + "entries": entries, + "rankings": rankings, + } + run = make_benchmark_run(extra={"datasets": [dataset]}) + return SyntheticRankingsDatasetComposer(run=run, tokenizer=FakeTokenizer()) + + +class TestRankingsWiring: + def test_lognormal_passages_varies(self): + composer = _make_rankings_composer( + entries=300, passages={"mean": 8, "median": 5} + ) + conversations = composer.create_dataset() + counts = [len(c.turns[0].texts[1].contents) for c in conversations] + assert len(set(counts)) > 3 # was: constant 8 (lognormal flattened) + assert statistics.median(counts) <= statistics.fmean(counts) # right skew + + def test_rankings_specific_token_options(synthetic_config, mock_tokenizer): """Test that rankings-specific token options are used for query and passages.""" synthetic_config.rankings_passages_mean = 3 diff --git a/tests/unit/dataset/generator/test_audio_generator.py b/tests/unit/dataset/generator/test_audio_generator.py index 732dbc7969..46851b94f5 100644 --- a/tests/unit/dataset/generator/test_audio_generator.py +++ b/tests/unit/dataset/generator/test_audio_generator.py @@ -3,17 +3,20 @@ import base64 import io +import statistics import sys +from unittest.mock import patch import numpy as np import pytest import soundfile as sf +import aiperf.dataset.generator.audio as audio_module from aiperf.common import random_generator as rng from aiperf.common.enums import AudioFormat from aiperf.common.exceptions import ConfigurationError from aiperf.config.dataset.content import AudioConfig -from aiperf.config.distributions import NormalDistribution +from aiperf.config.distributions import LogNormalDistribution, NormalDistribution from aiperf.dataset.generator import ( AudioGenerator, ) @@ -199,6 +202,35 @@ def test_audio_below_min_length_raises(): audio_generator.generate() +def test_lognormal_audio_length_varies(): + """Audio length samples its full lognormal shape, not the flattened mean. + + A lognormal has no ``stddev`` attribute, so the old code collapsed it to a + constant at the mean; num_samples (length * fixed 8kHz rate) is captured to + observe the drawn length without decoding. + """ + config = AudioConfig( + length=LogNormalDistribution(mean=3.0, median=2.0), + sample_rates=[8.0], + depths=[16], + format=AudioFormat.WAV, + channels=1, + ) + generator = AudioGenerator(config) + + with patch.object( + audio_module, + "generate_noise_signal", + wraps=audio_module.generate_noise_signal, + ) as spy: + for _ in range(100): + generator.generate() + + num_samples = [call.args[1] for call in spy.call_args_list] + assert len(set(num_samples)) > 3 # was: constant (lognormal flattened to mean) + assert statistics.median(num_samples) < statistics.fmean(num_samples) # right skew + + class TestAudioBitDepth: """Test suite for audio bit depth support, including 8-bit unsigned WAV.""" diff --git a/tests/unit/dataset/generator/test_image_generator.py b/tests/unit/dataset/generator/test_image_generator.py index 3980809459..583fe070cf 100644 --- a/tests/unit/dataset/generator/test_image_generator.py +++ b/tests/unit/dataset/generator/test_image_generator.py @@ -15,7 +15,11 @@ ImageSourceSamplingStrategy, ) from aiperf.config.dataset.content import ImageConfig -from aiperf.config.distributions import NormalDistribution +from aiperf.config.distributions import ( + EmpiricalDistribution, + EmpiricalPoint, + NormalDistribution, +) from aiperf.dataset.generator import ImageGenerator # v1 had separate ImageWidthConfig / ImageHeightConfig dataclasses; v2 uses @@ -365,6 +369,49 @@ def generate_with_seed(seed): assert generate_with_seed(12345) == generate_with_seed(12345) +class TestImageDimensionsWiring: + """Width/height sample from their full typed distribution, not a flattened mean.""" + + def test_empirical_image_dimensions_exact_values(self): + # Empirical width/height flatten to expected_value (640) under the old + # code; the typed path must draw only the configured discrete values. + config = ImageConfig( + batch_size=1, + width=EmpiricalDistribution( + points=[ + EmpiricalPoint(value=256, weight=50), + EmpiricalPoint(value=1024, weight=50), + ] + ), + height=EmpiricalDistribution( + points=[ + EmpiricalPoint(value=256, weight=50), + EmpiricalPoint(value=1024, weight=50), + ] + ), + format=ImageFormat.PNG, + source=ImageSource.NOISE, + ) + generator = ImageGenerator(config) + + captured: list[tuple[int, int]] = [] + original = generator._create_source_image + + def spy(width: int, height: int): + captured.append((width, height)) + return original(width, height) + + generator._create_source_image = spy + for _ in range(50): + generator.generate() + + widths = {w for w, _ in captured} + heights = {h for _, h in captured} + assert widths <= {256, 1024} + assert heights <= {256, 1024} + assert widths == {256, 1024} # both modes appear (was: constant 640) + + class TestImageGeneratorNoiseMode: """Tests for noise source mode.""" diff --git a/tests/unit/dataset/generator/test_prompt_generator.py b/tests/unit/dataset/generator/test_prompt_generator.py index 3b216ad23c..85dc6c3b7d 100644 --- a/tests/unit/dataset/generator/test_prompt_generator.py +++ b/tests/unit/dataset/generator/test_prompt_generator.py @@ -840,3 +840,35 @@ def test_build_token_sequence_same_validation_as_generate_cached( # This should raise same error as _generate_cached_prompt with pytest.raises(ConfigurationError): generator._build_token_sequence(10, [1, 2, 3], 5) # final_block_size = 0 + + +class TestPrefixLengthDistributionSampling: + def _make_generator(self, prefix_prompts: PrefixPromptConfig): + from aiperf.dataset.generator.prompt import PromptGenerator + from tests.harness.fake_tokenizer import FakeTokenizer + + return PromptGenerator( + prompts=None, prefix_prompts=prefix_prompts, tokenizer=FakeTokenizer() + ) + + def test_shared_system_prompt_sampled_once_and_stable(self): + gen = self._make_generator( + PrefixPromptConfig.model_validate( + {"shared_system_length": {"mean": 500, "stddev": 100, "min": 50}} + ) + ) + first = gen.get_shared_system_prompt() + assert first == gen.get_shared_system_prompt() # identical across calls + n_tokens = len(gen.tokenizer.encode(first)) + assert 50 <= n_tokens <= 1200 # within the clamped distribution's support + + def test_user_context_lengths_vary_across_sessions_and_are_stable_per_session(self): + gen = self._make_generator( + PrefixPromptConfig.model_validate( + {"user_context_length": {"mean": 400, "stddev": 150, "min": 20}} + ) + ) + prompts = [gen.generate_user_context_prompt(i) for i in range(30)] + lengths = [len(gen.tokenizer.encode(p)) for p in prompts] + assert len(set(lengths)) > 5 # sizes actually vary per session + assert gen.generate_user_context_prompt(7) == prompts[7] # stable per index diff --git a/tests/unit/property/_numeric_bounds_baseline.txt b/tests/unit/property/_numeric_bounds_baseline.txt index 0b607afe32..08e543923e 100644 --- a/tests/unit/property/_numeric_bounds_baseline.txt +++ b/tests/unit/property/_numeric_bounds_baseline.txt @@ -269,6 +269,8 @@ PlotSpec.ci_level: numeric field has no ge/gt/le/lt bound and is not FiniteFloat PluginEntry.priority: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. PluginSpec.priority: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. PowerMetricUnitInfo.watts: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +PrefixPromptConfig.shared_system_length: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +PrefixPromptConfig.user_context_length: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. ProcessingStats.errors: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. ProcessingStats.processed: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. ProcessRecordsCommand.request_ns: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. @@ -286,6 +288,7 @@ ProfileResults.start_ns: numeric field has no ge/gt/le/lt bound and is not Finit ProfileResults.timeslice_metric_results: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. ProfileResults.total_expected: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. ProfileStartCommand.request_ns: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +PromptConfig.first_turn_isl: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. PromptConfig.isl: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. PromptConfig.osl: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. PublicDataset.random_seed: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. @@ -327,6 +330,7 @@ SageMakerDataCaptureTrace.timestamp: numeric field has no ge/gt/le/lt bound and _SearchPlannerSettings.SLA_PRECISION_REQUESTS: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. SearchRecipeContext.sla_targets: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. SearchRecipeOutput.slos: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. +SequenceDistributionEntry.first_turn_isl: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. ServerMetricsEndpointSummary.info: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. ServerMetricsRecord.endpoint_latency_ns: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. ServerMetricsRecord.first_byte_ns: numeric field has no ge/gt/le/lt bound and is not FiniteFloat. Add a Pydantic numeric constraint, annotate as FiniteFloat, or add to NUMERIC_BOUNDS_WHITELIST. diff --git a/tests/unit/property/_strategies.py b/tests/unit/property/_strategies.py index 54259c0804..f025319ed8 100644 --- a/tests/unit/property/_strategies.py +++ b/tests/unit/property/_strategies.py @@ -370,6 +370,27 @@ def lognormal_distribution_inputs() -> st.SearchStrategy[dict]: ) +def percentile_distribution_inputs() -> st.SearchStrategy[dict]: + """Percentile dicts: p50 < p99, optional feasible mean between them.""" + + def _build(p50: float, ratio: float, mean_frac: float | None) -> dict: + p99 = p50 * ratio + d = {"p50": p50, "p99": p99} + if mean_frac is not None: + d["mean"] = p50 + mean_frac * (p99 - p50) + return d + + return st.builds( + _build, + st.floats(min_value=10.0, max_value=1e6, allow_nan=False, allow_infinity=False), + # Ratio max is intentionally extreme (1e30) so the fuzzer drives the + # finite-mean/variance guard in _solve_percentile; the resulting + # ValueError is in the property test's ALLOWED set. + st.floats(min_value=1.5, max_value=1e30), + st.one_of(st.none(), st.floats(min_value=0.05, max_value=0.3)), + ) + + def multimodal_distribution_inputs() -> st.SearchStrategy[dict]: peak = st.fixed_dictionaries( {"mean": adversarial_floats(), "stddev": adversarial_floats()}, diff --git a/tests/unit/property/test_finite_invariants.py b/tests/unit/property/test_finite_invariants.py index 3a141c0caa..cd554244d5 100644 --- a/tests/unit/property/test_finite_invariants.py +++ b/tests/unit/property/test_finite_invariants.py @@ -194,6 +194,8 @@ def _iter_pydantic_models() -> list[type]: "NormalDistribution.max": "inherited; validated by Distribution._validate_bounds", "LogNormalDistribution.min": "inherited; validated by Distribution._validate_bounds", "LogNormalDistribution.max": "inherited; validated by Distribution._validate_bounds", + "PercentileDistribution.min": "inherited; validated by Distribution._validate_bounds", + "PercentileDistribution.max": "inherited; validated by Distribution._validate_bounds", "MultimodalDistribution.min": "inherited; validated by Distribution._validate_bounds", "MultimodalDistribution.max": "inherited; validated by Distribution._validate_bounds", "EmpiricalDistribution.min": "inherited; validated by Distribution._validate_bounds", diff --git a/tests/unit/property/test_pydantic_field_fuzz.py b/tests/unit/property/test_pydantic_field_fuzz.py index bc306613dd..ead67e4fdc 100644 --- a/tests/unit/property/test_pydantic_field_fuzz.py +++ b/tests/unit/property/test_pydantic_field_fuzz.py @@ -31,6 +31,7 @@ LogNormalDistribution, MultimodalDistribution, NormalDistribution, + PercentileDistribution, SamplingDistribution, ) from aiperf.config.flags.cli_config import CLIConfig @@ -56,6 +57,7 @@ multi_run_inputs, multimodal_distribution_inputs, normal_distribution_inputs, + percentile_distribution_inputs, sampling_dimension_inputs, scenario_sweep_inputs, search_space_dimension_inputs, @@ -197,6 +199,14 @@ def test_lognormal_distribution_never_unhandled(data: dict) -> None: _check_no_unhandled(LogNormalDistribution, data) +@PROFILE +@given(data=percentile_distribution_inputs()) +def test_percentile_distribution_never_unhandled(data: dict) -> None: + # Infeasible (p50, p99, mean) triples raise ValueError from the solver, + # which is in ALLOWED -- _check_no_unhandled treats it as a clean rejection. + _check_no_unhandled(PercentileDistribution, data) + + @PROFILE @given(data=multimodal_distribution_inputs()) def test_multimodal_distribution_never_unhandled(data: dict) -> None: @@ -217,6 +227,7 @@ def test_empirical_distribution_never_unhandled(data: dict) -> None: data=fixed_distribution_inputs() | normal_distribution_inputs() | lognormal_distribution_inputs() + | percentile_distribution_inputs() | multimodal_distribution_inputs() | empirical_distribution_inputs() ) diff --git a/tests/unit/timing/strategies/test_user_centric_rate.py b/tests/unit/timing/strategies/test_user_centric_rate.py index e0b86d3eb2..8142195918 100644 --- a/tests/unit/timing/strategies/test_user_centric_rate.py +++ b/tests/unit/timing/strategies/test_user_centric_rate.py @@ -12,6 +12,7 @@ TWO_TURN = [("c1", 2), ("c2", 2), ("c3", 2), ("c4", 2), ("c5", 2)] MULTI_TURN = [("c1", 3), ("c2", 3), ("c3", 3), ("c4", 3)] +VARYING_TURNS = [("c1", 8), ("c2", 2), ("c3", 8), ("c4", 2)] class TestUserCentricInit: @@ -83,6 +84,33 @@ async def test_issues_expected_credits( assert len(h.sent_credits) >= expected +@pytest.mark.asyncio +class TestVaryingTurnCounts: + async def test_max_turns_clamped_to_sampled_conversation_length( + self, create_orchestrator_harness + ) -> None: + """Virtual-history warm start derives turns_to_send from the dataset + AVERAGE turn count; a user who draws a conversation shorter than that + average must clamp to the conversation's real length. The worker-side + session manager rejects num_turns > len(conversation.turns), so an + unclamped credit errors every such session at benchmark start.""" + h: OrchestratorHarness = create_orchestrator_harness( + conversations=VARYING_TURNS * 5, + user_centric_rate=20.0, + num_users=10, + request_count=40, + ) + lengths = dict(VARYING_TURNS) + await h.run_with_auto_return() + assert h.sent_credits + for credit in h.sent_credits: + assert credit.num_turns <= lengths[credit.conversation_id], ( + f"credit for {credit.conversation_id} carries " + f"num_turns={credit.num_turns} > conversation length " + f"{lengths[credit.conversation_id]}" + ) + + @pytest.mark.asyncio class TestSessionTracking: async def test_unique_correlation_ids(self, create_orchestrator_harness) -> None: