feat(dataset): typed distribution sampling, percentile targets, sticky sequence buckets, first-turn context#1130
feat(dataset): typed distribution sampling, percentile targets, sticky sequence buckets, first-turn context#1130ajcasagrande wants to merge 14 commits into
Conversation
Add a sixth sampling distribution parameterized by percentile targets:
{p50, p99} fits a log-normal exactly; {p50, p99, mean} solves a
two-component truncated-normal mixture deterministically at validation
time (stdlib math only, no scipy). Register in the discriminated union
via a p50/p99 discriminator rule ahead of the median check. Fix the base
Distribution.__getattr__ to delegate PrivateAttr reads to Pydantic so the
solved-parameters cache resolves.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
…ning, no double-sampling) Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
Replace the mean+stddev flattening in _build_sequence_distribution with a private _TypedSequenceDistribution runtime class: pick a bucket by weight, then draw entry.isl.sample_int(rng) / entry.osl.sample_int(rng) so lognormal/multimodal/empirical/percentile buckets keep their full shape. Also fix SyntheticDatasetComposer._include_prompt to enable text when a sequence_distribution is configured (native YAML seq-dist without isl previously raised 'All synthetic data are disabled'; the CLI path masked this by defaulting isl.mean=550). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
…dio, rankings Rewire the remaining SamplingDistribution consumers to draw from their full typed shape (lognormal/multimodal/empirical/percentile) instead of flattening to (expected_value, stddev-if-Normal): - synthetic.py: turns via sample_int (floors at 1); turn_delay via sample() gated on expected_value > 0 (preserves mean-0 disables delay), now float. - image.py: width/height via sample_int. - audio.py: length via sample() clamped to max(0.01, ...); keeps < 0.01 config error on expected_value. - synthetic_rankings.py: passages/query_tokens/passage_tokens via sample_int, removing the calculate_num_tokens double-sample. Delete now-dead _stddev_int (synthetic) and _stddev (rankings) helpers; keep _expected for the config-time _include_* gates. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
…ights Add PromptConfig.first_turn_isl to size the first turn's starting context in growing multi-turn benchmarks while isl sizes each subsequent turn's new input; wired through _get_turn_sequence_lengths(is_first=...). Relax ISL/OSL sequence-distribution bucket weights to positive relative weights (normalized at sampling time); they no longer must sum to 100. Applies to both the typed sequence_distribution path and the legacy CLI string parser. Non-positive totals are still rejected to guard the normalization divide. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
…n buckets Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
…t buckets, get_statistics normalization Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
Top-level first_turn_isl now conflicts with sequence_distribution (ValueError) instead of being silently ignored; each bucket carries its own optional first_turn_isl for conversation-class seed context. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
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 per-turn growth. Single-turn draw ordering is byte-identical to the previous per-turn sampling. Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
shared_system_length samples once per run (prompt identical across sessions); user_context_length samples once per session. Scalar ints still coerce to FixedDistribution; CLI flags unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
…ation Plus first-turn docs coherence: per-bucket first_turn_isl examples and the sequence_distribution conflict error documented in the tutorial. Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
ajcasagrande
requested review from
FrankD412,
debermudez,
ganeshku1,
ilana-n,
lkomali and
matthewkotila
as code owners
Try out this PRQuick install: pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@f5fa59348e472514ae1bd6d7090ecc42ff3dd23aRecommended with virtual environment (using uv): uv venv --python 3.12 && source .venv/bin/activate
uv pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@f5fa59348e472514ae1bd6d7090ecc42ff3dd23aLast updated for commit: |
|
…n length Virtual-history users derive turns_to_send from the dataset AVERAGE turn count; with per-conversation turns distributions actually varying (fixed on this branch), a user drawing a shorter conversation carried num_turns beyond the conversation's real length and the worker errored every such session at benchmark start. Latent on main, where the turns distribution flattened to a constant so average == every length. Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
![dynamo-review-agent[bot]](https://avatars.githubusercontent.com/u/201626793?s=60&v=4){kind=small}
| ) | ||
| if not 0.0 <= self.probability <= 100.0: | ||
| raise ValueError(f"Probability must be in [0,100], got {self.probability}") | ||
| if self.probability < 0.0: |
There was a problem hiding this comment.
Legacy sequence-distribution parsing now accepts NaN/Inf weights because the new check only rejects negative values, which can poison cumulative probability normalization with non-finite values. Fix: reject non-finite weights and require the summed weight to be finite and positive.
🤖 AI Fix
Update src/aiperf/common/models/sequence_distribution.py so SequenceLengthPair.__post_init__ rejects not math.isfinite(self.probability) as well as negatives, and _validate_probability_weights rejects not math.isfinite(total_prob) before allowing normalization.
There was a problem hiding this comment.
Fixed in f5fa593. SequenceLengthPair.__post_init__ now rejects non-finite probability and both stddevs, and _validate_probability_weights requires a finite positive total. Parser-level coverage added (256,128:inf / :nan rejected) plus construction-level nan/inf/-inf cases.
| 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. " |
There was a problem hiding this comment.
Typed YAML sequence_distribution buckets can accept infinite probability values, causing _TypedSequenceDistribution to compute inf / inf as NaN during sampling. Fix: disallow non-finite probabilities and validate that the total weight is finite and positive.
🤖 AI Fix
Update src/aiperf/config/types.py in SequenceDistributionEntry and validate_probability_distribution to reject non-finite probability values, for example by adding allow_inf_nan=False to the Field and checking math.isfinite(total) before total > 0.
There was a problem hiding this comment.
Fixed in f5fa593. SequenceDistributionEntry.probability now carries allow_inf_nan=False (inf was indeed accepted by ge=0.0; NaN was already rejected), and validate_probability_distribution rejects a total that overflows to inf even when every per-entry weight is finite (e.g. two 1e308 weights). Regression tests for both layers.
|
Follow-up commit a9173ef: running the headline config end-to-end against the mock server surfaced a latent user-centric bug this branch exposes — virtual-history warm-start users derive their remaining-turn budget from the dataset average turn count, and with 🤖 Generated with Claude Code |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (4)
🚧 Files skipped from review as they are similar to previous changes (4)
WalkthroughThis PR changes sequence-distribution probabilities to positive relative weights, adds ChangesRelative weights, percentile distribution, and typed sequence sampling
Estimated code review effort: 4 (Complex) | ~75 minutes Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (3)
src/aiperf/common/models/sequence_distribution.py (1)
14-20: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMinor: doc example simplifies the actual grammar.
Note: Probabilities are positive relative weights and do NOT need to sum to 100. The
"50;1"shorthand in the example doesn't match the actualISL,OSL:PROBgrammar shown just above it, which could confuse readers trying to map the example to real syntax.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/common/models/sequence_distribution.py` around lines 14 - 20, The docstring example in SequenceDistribution currently shows a shorthand that does not match the actual semicolon grammar, which can mislead readers. Update the examples in the sequence distribution documentation to use the real ISL,OSL:PROB syntax consistently, and if you keep a normalization example, make sure it is expressed with valid pair syntax in the relevant class or parser docs so the examples align with the supported formats.tests/unit/config/test_distributions.py (1)
573-579: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMove local import to top of file.
DistributionParseris imported inside the test function; move it to the module-level imports.♻️ Proposed fix
+from aiperf.common.models.sequence_distribution import DistributionParser + 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")As per path instructions,
tests/**/*.{py,pyi}: "Keep imports at the top of the file, use fixtures for setup, and keep one focus per test."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/config/test_distributions.py` around lines 573 - 579, The test `test_legacy_string_parser_accepts_relative_weights` imports `DistributionParser` inside the function, but the module-level imports should be used instead. Move the `DistributionParser` import to the top of the file with the other imports, and keep the test body focused on parsing and sampling only.Source: Path instructions
tests/unit/dataset/composer/test_distribution_wiring.py (1)
280-338: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMove local imports to module top.
random_generator,SequenceDistributionEntry,_TypedSequenceDistribution, andvalidate_probability_distributionare imported inside three separate test methods inTestZeroWeightBuckets. Consolidate these into the top-level imports.♻️ Proposed fix
+from aiperf.common import random_generator as rng +from aiperf.config.types import ( + SequenceDistributionEntry, + validate_probability_distribution, +) +from aiperf.dataset.composer.base import _TypedSequenceDistribution + 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 = [(similar cleanup applies to the other two methods)
As per path instructions,
tests/**/*.{py,pyi}: "Keep imports at the top of the file, use fixtures for setup, and keep one focus per test."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/dataset/composer/test_distribution_wiring.py` around lines 280 - 338, The TestZeroWeightBuckets methods are doing repeated local imports instead of using the module-level imports expected in tests. Move random_generator, SequenceDistributionEntry, _TypedSequenceDistribution, and validate_probability_distribution to the top of the test module, then remove the inline imports from test_zero_weight_bucket_validates_and_never_samples, test_zero_weight_bucket_end_to_end_via_composer, test_all_zero_list_rejected_by_validator, and test_all_zero_list_rejected_by_typed_distribution so the tests only focus on behavior.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/aiperf/config/schema/aiperf-config.schema.json`:
- Around line 9771-9773: The PeakEntry schema currently only allows the wrapped
PercentileDistribution shape, but the parser also accepts an inline {p50, p99,
weight} variant. Update the PeakEntry union in aiperf-config.schema.json to add
a sibling inline PercentileDistribution branch alongside the existing $ref,
using the same PercentileDistribution symbol so schema validation matches the
Python parser behavior.
In `@src/aiperf/dataset/composer/base.py`:
- Around line 71-84: The sample_lengths method in SequenceDistributionComposer
currently samples bucket.osl unconditionally, so a zero-valued OSL still becomes
at least 1 instead of disabling output. Update the bucketed path to mirror the
plain-path behavior by checking the OSL distribution’s expected value before
sampling, and return 0 for the OSL when it is disabled. Keep the fix localized
to sample_lengths and use the existing bucket.osl and isl_dist logic to preserve
first-turn behavior.
In `@tests/unit/config/test_percentile_distribution.py`:
- Line 69: The pytest.raises regex patterns in the percentile distribution tests
use normal strings, which triggers Ruff RUF043; update the match= arguments in
the pytest.raises blocks to raw string literals in the relevant test cases so
the regex alternation is preserved. Apply this to the assertions around the
infeasible/mean case and the too extreme/finite case in the percentile
distribution test module.
---
Nitpick comments:
In `@src/aiperf/common/models/sequence_distribution.py`:
- Around line 14-20: The docstring example in SequenceDistribution currently
shows a shorthand that does not match the actual semicolon grammar, which can
mislead readers. Update the examples in the sequence distribution documentation
to use the real ISL,OSL:PROB syntax consistently, and if you keep a
normalization example, make sure it is expressed with valid pair syntax in the
relevant class or parser docs so the examples align with the supported formats.
In `@tests/unit/config/test_distributions.py`:
- Around line 573-579: The test
`test_legacy_string_parser_accepts_relative_weights` imports
`DistributionParser` inside the function, but the module-level imports should be
used instead. Move the `DistributionParser` import to the top of the file with
the other imports, and keep the test body focused on parsing and sampling only.
In `@tests/unit/dataset/composer/test_distribution_wiring.py`:
- Around line 280-338: The TestZeroWeightBuckets methods are doing repeated
local imports instead of using the module-level imports expected in tests. Move
random_generator, SequenceDistributionEntry, _TypedSequenceDistribution, and
validate_probability_distribution to the top of the test module, then remove the
inline imports from test_zero_weight_bucket_validates_and_never_samples,
test_zero_weight_bucket_end_to_end_via_composer,
test_all_zero_list_rejected_by_validator, and
test_all_zero_list_rejected_by_typed_distribution so the tests only focus on
behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 6c273293-9e9d-4ccc-912a-a87cd11a8347
📒 Files selected for processing (31)
docs/cli-options.mddocs/tutorials/sequence-distributions.mddocs/tutorials/yaml-distributions.mdsrc/aiperf/common/models/sequence_distribution.pysrc/aiperf/config/dataset/content.pysrc/aiperf/config/distributions.pysrc/aiperf/config/flags/cli_config.pysrc/aiperf/config/schema/aiperf-config.schema.jsonsrc/aiperf/config/types.pysrc/aiperf/dataset/composer/base.pysrc/aiperf/dataset/composer/synthetic.pysrc/aiperf/dataset/composer/synthetic_rankings.pysrc/aiperf/dataset/generator/audio.pysrc/aiperf/dataset/generator/image.pysrc/aiperf/dataset/generator/prompt.pytests/unit/common/models/test_sequence_distribution.pytests/unit/config/test_camel_case_config.pytests/unit/config/test_distributions.pytests/unit/config/test_percentile_distribution.pytests/unit/dataset/composer/test_base_composer.pytests/unit/dataset/composer/test_distribution_wiring.pytests/unit/dataset/composer/test_percentile_e2e.pytests/unit/dataset/composer/test_sticky_sequence_buckets.pytests/unit/dataset/composer/test_synthetic_rankings_composer.pytests/unit/dataset/generator/test_audio_generator.pytests/unit/dataset/generator/test_image_generator.pytests/unit/dataset/generator/test_prompt_generator.pytests/unit/property/_numeric_bounds_baseline.txttests/unit/property/_strategies.pytests/unit/property/test_finite_invariants.pytests/unit/property/test_pydantic_field_fuzz.py
| { | ||
| "$ref": "#/$defs/PercentileDistribution" | ||
| }, |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Inspect PeakEntry parsing to confirm inline percentile shorthand is/ isn't supported.
rg -nP -C6 '\bclass PeakEntry\b|weight.*pop|extract.*weight' src/aiperf/common/models src/aiperf/configRepository: ai-dynamo/aiperf
Length of output: 2428
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the PeakEntry model and the relevant schema definitions.
sed -n '540,610p' src/aiperf/config/distributions.py
printf '\n--- schema excerpt ---\n'
sed -n '9680,9810p' src/aiperf/config/schema/aiperf-config.schema.jsonRepository: ai-dynamo/aiperf
Length of output: 7527
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Check whether PeakEntry schema already has an inline PercentileDistribution variant.
rg -n -C3 'PercentileDistribution peak|Inline Percentile|PeakEntry' src/aiperf/config/schema/aiperf-config.schema.jsonRepository: ai-dynamo/aiperf
Length of output: 1383
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read the PeakEntry variants to see whether an inline percentile shorthand exists.
sed -n '9751,9865p' src/aiperf/config/schema/aiperf-config.schema.jsonRepository: ai-dynamo/aiperf
Length of output: 4428
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Find the remaining PeakEntry anyOf variants and look for a percentile shorthand.
sed -n '9865,9965p' src/aiperf/config/schema/aiperf-config.schema.jsonRepository: ai-dynamo/aiperf
Length of output: 3794
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the rest of PeakEntry's inline variants for a percentile shorthand.
sed -n '9965,10080p' src/aiperf/config/schema/aiperf-config.schema.jsonRepository: ai-dynamo/aiperf
Length of output: 3938
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect SamplingDistribution and PercentileDistribution definitions.
rg -n -C4 'class PercentileDistribution|SamplingDistribution\s*=|class PeakEntry' src/aiperf/config/distributions.pyRepository: ai-dynamo/aiperf
Length of output: 1265
Add an inline PercentileDistribution peak variant
PeakEntry accepts inline {p50, p99, weight} in Python, but the schema only exposes the wrapper form. Add a sibling inline PercentileDistribution branch so schema validation matches the parser.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/aiperf/config/schema/aiperf-config.schema.json` around lines 9771 - 9773,
The PeakEntry schema currently only allows the wrapped PercentileDistribution
shape, but the parser also accepts an inline {p50, p99, weight} variant. Update
the PeakEntry union in aiperf-config.schema.json to add a sibling inline
PercentileDistribution branch alongside the existing $ref, using the same
PercentileDistribution symbol so schema validation matches the Python parser
behavior.
| 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), | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm sample_int()'s floor-to-1 behavior is shared by all SamplingDistribution
# subtypes (Fixed/Normal/LogNormal/Multimodal/Empirical/Percentile) and not
# overridden anywhere to special-case zero.
rg -n "def sample_int" src/aiperf/config/distributions.py -A3Repository: ai-dynamo/aiperf
Length of output: 291
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the relevant implementation and nearby helpers.
sed -n '1,220p' src/aiperf/dataset/composer/base.py
printf '\n--- distributions ---\n'
sed -n '1,260p' src/aiperf/config/distributions.pyRepository: ai-dynamo/aiperf
Length of output: 18690
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate any existing handling of zero-mean OSL or max_tokens disabling.
rg -n "expected_value > 0|osl: \{mean: 0\}|max_tokens=0|sample_lengths\(" src/aiperf -A3 -B3Repository: ai-dynamo/aiperf
Length of output: 4735
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read the turn-length selection path that uses the bucket sampler.
sed -n '228,320p' src/aiperf/dataset/composer/base.pyRepository: ai-dynamo/aiperf
Length of output: 3967
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the sequence-distribution config schema and any docs for OSL=0 semantics.
rg -n "SequenceDistributionEntry|first_turn_isl|osl:.*mean: 0|disables output|max_tokens" src/aiperf/config src/aiperf/dataset -A4 -B4Repository: ai-dynamo/aiperf
Length of output: 50372
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read the exact turn-length logic around the bucket and plain paths.
sed -n '240,320p' src/aiperf/dataset/composer/base.pyRepository: ai-dynamo/aiperf
Length of output: 3296
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the sequence distribution entry schema.
rg -n "class SequenceDistributionEntry|sequence_distribution" src/aiperf/config -A40 -B10Repository: ai-dynamo/aiperf
Length of output: 50372
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the exact distribution contract for zero-valued OSL and sampling semantics.
sed -n '1,220p' src/aiperf/config/types.pyRepository: ai-dynamo/aiperf
Length of output: 6285
src/aiperf/dataset/composer/base.py: guard bucketed OSL before sampling
bucket.osl.sample_int(self._rng) always floors to at least 1, so a bucket with osl: {mean: 0} still produces max_tokens=1 instead of disabling output like the plain path does. Mirror the plain-path expected_value > 0 guard here and return 0 when the bucket OSL is disabled.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/aiperf/dataset/composer/base.py` around lines 71 - 84, The sample_lengths
method in SequenceDistributionComposer currently samples bucket.osl
unconditionally, so a zero-valued OSL still becomes at least 1 instead of
disabling output. Update the bucketed path to mirror the plain-path behavior by
checking the OSL distribution’s expected value before sampling, and return 0 for
the OSL when it is disabled. Keep the fix localized to sample_lengths and use
the existing bucket.osl and isl_dist logic to preserve first-turn behavior.
| ], | ||
| ) # fmt: skip | ||
| def test_extreme_ratio_rejected(self, p50: float, p99: float) -> None: | ||
| with pytest.raises(ValueError, match="too extreme|finite"): |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Use a raw string for the regex match= pattern.
The | alternation is intended regex, but the non-raw string trips Ruff RUF043. Same applies to match="infeasible|mean" on Line 49.
🧹 Proposed fix
- with pytest.raises(ValueError, match="too extreme|finite"):
+ with pytest.raises(ValueError, match=r"too extreme|finite"):📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| with pytest.raises(ValueError, match="too extreme|finite"): | |
| with pytest.raises(ValueError, match=r"too extreme|finite"): |
🧰 Tools
🪛 Ruff (0.15.20)
[warning] 69-69: Pattern passed to match= contains metacharacters but is neither escaped nor raw
(RUF043)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unit/config/test_percentile_distribution.py` at line 69, The
pytest.raises regex patterns in the percentile distribution tests use normal
strings, which triggers Ruff RUF043; update the match= arguments in the
pytest.raises blocks to raw string literals in the relevant test cases so the
regex alternation is preserved. Apply this to the assertions around the
infeasible/mean case and the too extreme/finite case in the percentile
distribution test module.
Source: Linters/SAST tools
There was a problem hiding this comment.
🧹 Nitpick comments (1)
src/aiperf/timing/strategies/user_centric_rate.py (1)
181-196: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueClamp logic correctly enforces the conversation-length invariant.
conversation_turns = len(sampled.metadata.turns)and the resultingmin(max_turns, conversation_turns)correctly preventsUser.max_turnsfrom exceeding the sampled conversation's actual length, matching theUserSessionManagercontract that rejectsnum_turns > len(conversation.turns).One minor defensive note:
if max_turns else conversation_turnsuses a truthy check rather thanis not None. Ifmax_turns=0is ever passed by a future caller, it would silently fall back toconversation_turnsinstead of0. Currently unreachable since the only caller passingmax_turnsguardsturns_to_send <= 0before calling this method, so this is safe today but fragile against future refactors.🛡️ Defensive fix
- max_turns=min(max_turns, conversation_turns) - if max_turns - else conversation_turns, + max_turns=min(max_turns, conversation_turns) + if max_turns is not None + else conversation_turns,🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/timing/strategies/user_centric_rate.py` around lines 181 - 196, In the user-centric rate strategy, the clamp in the user creation path is correct, but the `max_turns` fallback in the `User` constructor call is fragile because it uses a truthy check. Update the logic in the method that builds the `User` instance so it distinguishes `None` from valid numeric values, preserving an explicit `0` if it is ever passed while still clamping against `len(sampled.metadata.turns)`.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@src/aiperf/timing/strategies/user_centric_rate.py`:
- Around line 181-196: In the user-centric rate strategy, the clamp in the user
creation path is correct, but the `max_turns` fallback in the `User` constructor
call is fragile because it uses a truthy check. Update the logic in the method
that builds the `User` instance so it distinguishes `None` from valid numeric
values, preserving an explicit `0` if it is ever passed while still clamping
against `len(sampled.metadata.turns)`.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: e817a92f-57be-4e72-8c59-ac29addd8655
📒 Files selected for processing (2)
src/aiperf/timing/strategies/user_centric_rate.pytests/unit/timing/strategies/test_user_centric_rate.py
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
Typed buckets: probability gains allow_inf_nan=False (inf/inf normalized to NaN in the sampler); validate_probability_distribution also rejects a sum that overflows to inf. Legacy string format: SequenceLengthPair rejects NaN/inf probability and stddevs, and the total-weight check requires a finite positive sum. Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
1 participant
Summary
Makes every typed
SamplingDistributionin YAML configs actually sample at runtime, adds a percentile-targeting distribution, and introduces conversation-class workload modeling: sticky ISL/OSL buckets with per-bucket first-turn seed context, plus distribution-valued prefix context sizes.Distribution wiring (first 8 commits)
PromptGenerator.generatereceivesstddev=0so variance comes from the per-turn draw only).PercentileDistribution—{p50, p99[, mean]}: lognormal fit for two targets, deterministic mixture solver whenmeanis also pinned. Solved eagerly at config-validation time; infeasible targets fail inaiperf config validate, not mid-benchmark.prompts.first_turn_isl— starting-context distribution for growing multi-turn conversations: turn 1 samples it, turns 2+ sampleisl.sequence_distributionprobabilities no longer must sum to 100;probability: 0disables a bucket. Same relaxation in the legacy--seq-diststring parser.turns,turn_delay, imagewidth/height, audiolength, and rankings shapes.Conversation classes (last 4 commits)
sequence_distributionnow draws ONE bucket per conversation and keeps it for every turn, so ISL/OSL stay tied for the conversation's life (e.g. huge-context conversations keep tiny outputs). Single-turn workloads are byte-identical to the previous per-turn draw (RNG-stream replay test enforces this).first_turn_isl: each bucket sizes its own seed context (turn 1) separately from per-turn growth (isl). Top-levelfirst_turn_isl+sequence_distributionis now a config error pointing at the per-bucket field.prefix_prompts.shared_system_length/user_context_lengthaccept any distribution — shared system sampled once per run (prompt identical across sessions, preserving its KV-cache purpose), user context sampled once per session.Behavior change
Multi-turn conversations with
sequence_distributionpreviously drew a fresh bucket per turn; they now hold one bucket per conversation. Documented with a callout indocs/tutorials/sequence-distributions.md. Single-turn behavior is unchanged (test-enforced byte-compat).Testing
docs/tutorials/yaml-distributions.md(new) +sequence-distributions.mdupdated; JSON schema and generated artifacts regenerated.🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
p50/p99with optionalmean) includingmin/maxclamping.first_turn_isl, including per-bucket overrides) for multi-turn runs.Documentation
Bug Fixes
max_turnsis clamped to the sampled conversation length.