diff --git a/docs/tutorials/agentic-code-generator.md b/docs/tutorials/agentic-code-generator.md index 0a8a0628f7..d71ccebdc7 100644 --- a/docs/tutorials/agentic-code-generator.md +++ b/docs/tutorials/agentic-code-generator.md @@ -96,6 +96,39 @@ If the sampled session would hit `max_prompt_tokens` before reaching the target: - With `allow_truncation: true`, the generator returns the partial session and marks it as `forced_retire`. +## Inter-Turn Delays + +`inter_turn_delay.context_scaling` defaults to `true` for compatibility with +existing workloads. After sampling the agentic/human mixture and applying +`inter_turn_delay.max`, the generator multiplies the delay by: + +`max(0.2, 1 - 0.8 * previous_input / max_prompt_tokens)` + +`previous_input` is the previous turn's cumulative input length. This is a +workload heuristic: fuller contexts produce shorter gaps. It is not an +inference-latency measurement. Set `"context_scaling": false` inside the +`inter_turn_delay` object to disable it. The setting is saved in the manifest +and generated config schema, and appears as 0/1 in the quality config summary. + +Component `min`/`max` bounds apply during rejection sampling. After context +scaling, delays below the selected component's `min` are clamped to exactly +that minimum, without redrawing. This can concentrate probability at the +minimum; the final distribution differs from rejection sampling after scaling. +The global `inter_turn_delay.max` is the final upper cap and takes precedence +if it is lower than that component's minimum; a global maximum of zero permits +zero delay. Turn 0 always has zero delay. + +The random draws and cap-before-scaling order are unchanged. With scaling +enabled and no component minima, an existing config and seed produce the same +trace. Configs with minima may now produce longer gaps where scaling previously +violated the bound. Disabling scaling intentionally changes delays and timestamps. + +Component means/medians describe the underlying lognormal distributions, not +targets for the final mixture. Rejection sampling, caps, and context scaling +change the output distribution. `quality.json` reports observed delay statistics +with null aggregate mean/median targets; the component values shown in reports +should not be interpreted as those aggregate targets. + ## Generate a Dataset Create a dataset with the built-in default configuration: diff --git a/src/aiperf/dataset/agentic_code_gen/configs/spec.json b/src/aiperf/dataset/agentic_code_gen/configs/spec.json index 1d710d3223..2d209b2bb3 100644 --- a/src/aiperf/dataset/agentic_code_gen/configs/spec.json +++ b/src/aiperf/dataset/agentic_code_gen/configs/spec.json @@ -93,7 +93,7 @@ } ], "default": null, - "description": "Hard lower bound (rejection sampled)", + "description": "Hard lower bound (rejection sampled). For inter-turn delays, values below this bound after context scaling are clamped, subject to the global delay maximum.", "title": "Min" }, "mu": { @@ -147,6 +147,12 @@ "title": "Agentic Fraction", "type": "number" }, + "context_scaling": { + "default": true, + "description": "Shorten delays as context fills using max(0.2, 1 - 0.8 * previous_input / max_prompt_tokens); disable to keep sampled delays. Delays are clamped to component minima after scaling, with the global max taking precedence. Enabled by default for workload compatibility.", + "title": "Context Scaling", + "type": "boolean" + }, "human_delay": { "$ref": "#/$defs/LognormalParams", "description": "Slow delay distribution (human think time)" @@ -217,7 +223,7 @@ } ], "default": null, - "description": "Hard lower bound (rejection sampled)", + "description": "Hard lower bound (rejection sampled). For inter-turn delays, values below this bound after context scaling are clamped, subject to the global delay maximum.", "title": "Min" }, "mu": { diff --git a/src/aiperf/dataset/agentic_code_gen/distributions.py b/src/aiperf/dataset/agentic_code_gen/distributions.py index 5c902cfd15..2e5a165ca2 100644 --- a/src/aiperf/dataset/agentic_code_gen/distributions.py +++ b/src/aiperf/dataset/agentic_code_gen/distributions.py @@ -88,17 +88,35 @@ def sample_lognormal( def sample_mixture_delay( - config: MixtureDelayConfig, rng: Generator, size: int = 1 + config: MixtureDelayConfig, + rng: Generator, + size: int = 1, + *, + delay_scale: float = 1.0, ) -> np.ndarray: """Sample from the two-component mixture delay model. For each sample, a Bernoulli draw selects agentic (fast) vs human (slow), - then the corresponding lognormal is sampled. + then the corresponding lognormal is sampled. The global cap is applied + before scaling to preserve historical workloads. Component minima are + reapplied by clamping after scaling, without additional random draws. + This can concentrate probability at a component's minimum. The global cap + overrides a conflicting minimum. """ + if not 0 < delay_scale <= 1: + raise ValueError("delay_scale must be in (0, 1]") is_agentic = rng.random(size=size) < config.agentic_fraction agentic_samples = sample_lognormal(config.agentic_delay, rng, size=size) human_samples = sample_lognormal(config.human_delay, rng, size=size) samples = np.where(is_agentic, agentic_samples, human_samples) + if config.max is not None: + samples = np.minimum(samples, config.max) + samples *= delay_scale + if config.agentic_delay.min is not None or config.human_delay.min is not None: + minimums = np.where( + is_agentic, config.agentic_delay.min or 0.0, config.human_delay.min or 0.0 + ) + samples = np.maximum(samples, minimums) if config.max is not None: samples = np.minimum(samples, config.max) return samples diff --git a/src/aiperf/dataset/agentic_code_gen/models.py b/src/aiperf/dataset/agentic_code_gen/models.py index 8a39da34c5..43ff8eb04e 100644 --- a/src/aiperf/dataset/agentic_code_gen/models.py +++ b/src/aiperf/dataset/agentic_code_gen/models.py @@ -92,7 +92,13 @@ class LognormalParams(AIPerfBaseModel): mean: float = Field(gt=0.0, description="Real-space mean (derived)") median: float = Field(gt=0.0, description="Real-space median (derived)") min: float | None = Field( - default=None, gt=0.0, description="Hard lower bound (rejection sampled)" + default=None, + gt=0.0, + description=( + "Hard lower bound (rejection sampled). For inter-turn delays, " + "values below this bound after context scaling are clamped, " + "subject to the global delay maximum." + ), ) max: float | None = Field( default=None, gt=0.0, description="Hard upper bound (rejection sampled)" @@ -144,6 +150,15 @@ class MixtureDelayConfig(AIPerfBaseModel): A Bernoulli draw selects which component to sample from. """ + context_scaling: bool = Field( + default=True, + description=( + "Shorten delays as context fills using max(0.2, 1 - 0.8 * " + "previous_input / max_prompt_tokens); disable to keep sampled delays. " + "Delays are clamped to component minima after scaling, with the global max " + "taking precedence. Enabled by default for workload compatibility." + ), + ) agentic_fraction: float = Field( default=0.7, ge=0.0, diff --git a/src/aiperf/dataset/agentic_code_gen/reporting/metrics.py b/src/aiperf/dataset/agentic_code_gen/reporting/metrics.py index 6a73f225d5..64768839f2 100644 --- a/src/aiperf/dataset/agentic_code_gen/reporting/metrics.py +++ b/src/aiperf/dataset/agentic_code_gen/reporting/metrics.py @@ -330,6 +330,9 @@ def config_summary(config: SessionDistributionConfig) -> dict[str, float | int]: "inter_turn_delay_agentic_fraction": config.inter_turn_delay.agentic_fraction, "inter_turn_delay_agentic_mean_ms": config.inter_turn_delay.agentic_delay.mean, "inter_turn_delay_human_mean_ms": config.inter_turn_delay.human_delay.mean, + "inter_turn_delay_context_scaling": int( + config.inter_turn_delay.context_scaling + ), "turn_mode_enabled": int(config.turns is not None), } if config.reset is not None: diff --git a/src/aiperf/dataset/agentic_code_gen/session_synthesizer.py b/src/aiperf/dataset/agentic_code_gen/session_synthesizer.py index 9857c0876c..fb8592e90f 100644 --- a/src/aiperf/dataset/agentic_code_gen/session_synthesizer.py +++ b/src/aiperf/dataset/agentic_code_gen/session_synthesizer.py @@ -140,14 +140,14 @@ def _sample_output_length(self) -> int: ) def _sample_delay_ms(self, prev_input: int) -> float: - delay_ms = float( - sample_mixture_delay(self._config.inter_turn_delay, self._rng, size=1)[0] - ) + config = self._config.inter_turn_delay context_ratio = prev_input / self._config.max_prompt_tokens - delay_ms *= max(0.2, 1.0 - 0.8 * context_ratio) - if self._config.inter_turn_delay.max is not None: - delay_ms = min(delay_ms, self._config.inter_turn_delay.max) - return delay_ms + delay_scale = ( + max(0.2, 1.0 - 0.8 * context_ratio) if config.context_scaling else 1.0 + ) + return float( + sample_mixture_delay(config, self._rng, size=1, delay_scale=delay_scale)[0] + ) def _sample_new_tokens(self) -> int: new_tokens = int( diff --git a/tests/unit/dataset/agentic_code_gen/test_distributions.py b/tests/unit/dataset/agentic_code_gen/test_distributions.py index 787b8ee39a..511b1732f2 100644 --- a/tests/unit/dataset/agentic_code_gen/test_distributions.py +++ b/tests/unit/dataset/agentic_code_gen/test_distributions.py @@ -6,9 +6,11 @@ from __future__ import annotations import math +from statistics import NormalDist import numpy as np import pytest +from pytest import param from aiperf.dataset.agentic_code_gen.distributions import ( fit_from_samples, @@ -150,6 +152,79 @@ def test_clipping_works(self) -> None: class TestSampleMixtureDelay: + @pytest.mark.parametrize( + "delay_scale", + [param(0.2, id="full-context"), param(0.4, id="partial-context")], + ) # fmt: skip + def test_scaled_lognormal_has_clamped_probability_at_minimum( + self, delay_scale: float + ) -> None: + """The floor collects the truncated lognormal's sub-threshold mass.""" + config = MixtureDelayConfig( + agentic_fraction=0.0, + human_delay=LognormalParams( + mean=450_000, median=84_000, min=30_000, max=1_800_000 + ), + ) + log_distribution = NormalDist( + mu=math.log(84_000), sigma=math.sqrt(2 * math.log(450_000 / 84_000)) + ) + cdf_min = log_distribution.cdf(math.log(30_000)) + cdf_max = log_distribution.cdf(math.log(1_800_000)) + cdf_threshold = log_distribution.cdf(math.log(30_000 / delay_scale)) + expected_floor_fraction = (cdf_threshold - cdf_min) / (cdf_max - cdf_min) + + samples = sample_mixture_delay( + config, np.random.default_rng(42), size=20_000, delay_scale=delay_scale + ) + + assert samples.min() == 30_000 + assert samples.max() <= 1_800_000 * delay_scale + floor_fraction = np.count_nonzero(samples == 30_000) / len(samples) + assert floor_fraction == pytest.approx(expected_floor_fraction, abs=0.015) + + def test_scaled_mixture_applies_each_selected_component_minimum(self) -> None: + config = MixtureDelayConfig( + agentic_fraction=0.5, + agentic_delay=LognormalParams(mean=1000, median=1000, min=900), + human_delay=LognormalParams(mean=5000, median=5000, min=4500), + ) + + samples = sample_mixture_delay( + config, np.random.default_rng(42), size=100, delay_scale=0.2 + ) + + assert set(samples) == {900.0, 4500.0} + + def test_zero_global_cap_overrides_component_minima(self) -> None: + config = MixtureDelayConfig( + agentic_delay=LognormalParams(mean=1000, median=1000, min=900), + human_delay=LognormalParams(mean=5000, median=5000, min=4500), + max=0, + ) + + samples = sample_mixture_delay( + config, np.random.default_rng(42), size=100, delay_scale=0.2 + ) + + assert np.array_equal(samples, np.zeros(100)) + + @pytest.mark.parametrize( + "delay_scale", + [ + param(0, id="zero"), + param(-1, id="negative"), + param(1.1, id="above-one"), + param(float("nan"), id="nan"), + param(float("inf"), id="inf"), + ], + ) # fmt: skip + def test_invalid_delay_scale_is_rejected(self, delay_scale: float) -> None: + with pytest.raises(ValueError, match="delay_scale must be in"): + sample_mixture_delay( + MixtureDelayConfig(), np.random.default_rng(42), delay_scale=delay_scale + ) + def test_returns_correct_shape(self) -> None: config = MixtureDelayConfig( agentic_fraction=0.7, diff --git a/tests/unit/dataset/agentic_code_gen/test_roundtrip.py b/tests/unit/dataset/agentic_code_gen/test_roundtrip.py index aee3b561b4..7dad18b338 100644 --- a/tests/unit/dataset/agentic_code_gen/test_roundtrip.py +++ b/tests/unit/dataset/agentic_code_gen/test_roundtrip.py @@ -10,12 +10,15 @@ import orjson import pytest +from pytest import param from aiperf.dataset.agentic_code_gen.config import load_config from aiperf.dataset.agentic_code_gen.distributions import lognormal_from_mean_median from aiperf.dataset.agentic_code_gen.models import ( CacheLayerConfig, Layer15GroupConfig, + LognormalParams, + MixtureDelayConfig, ResetConfig, SessionDistributionConfig, ) @@ -25,6 +28,43 @@ class TestRoundtrip: + @pytest.mark.parametrize( + "context_scaling, expected_delay", + [param(True, 900, id="scaled"), param(False, 1000, id="unscaled")], + ) # fmt: skip + def test_delay_config_survives_manifest_and_trace_roundtrip( + self, + tmp_path: Path, + turns_config: SessionDistributionConfig, + context_scaling: bool, + expected_delay: float, + ) -> None: + turns_config.cache.layer1_tokens = 2000 + turns_config.inter_turn_delay = MixtureDelayConfig( + context_scaling=context_scaling, + agentic_fraction=1.0, + agentic_delay=LognormalParams(mean=1000, median=1000, min=900), + ) + sessions = SessionSynthesizer(turns_config, seed=42).synthesize_sessions(3) + + jsonl, manifest, quality = write_dataset( + sessions, tmp_path / "run", turns_config, seed=42 + ) + + restored_config = load_config(str(manifest)) + assert restored_config.inter_turn_delay.context_scaling is context_scaling + replay = SessionSynthesizer(restored_config, seed=42).synthesize_sessions(3) + assert replay == sessions + turns = [ + MooncakeTrace(**orjson.loads(line)) + for line in jsonl.read_bytes().splitlines() + ] + delays = [turn.delay for turn in turns if turn.delay is not None] + assert len(delays) == 9 + assert delays == pytest.approx([expected_delay] * 9) + summary = orjson.loads(quality.read_bytes())["config_summary"] + assert summary["inter_turn_delay_context_scaling"] == int(context_scaling) + def test_config_to_mooncake_roundtrip(self, tmp_path: Path) -> None: config = SessionDistributionConfig() synth = SessionSynthesizer(config, seed=42) diff --git a/tests/unit/dataset/agentic_code_gen/test_session_synthesizer.py b/tests/unit/dataset/agentic_code_gen/test_session_synthesizer.py index a66eea5e3c..6264758184 100644 --- a/tests/unit/dataset/agentic_code_gen/test_session_synthesizer.py +++ b/tests/unit/dataset/agentic_code_gen/test_session_synthesizer.py @@ -7,6 +7,7 @@ import numpy as np import pytest +from pytest import param from aiperf.dataset.agentic_code_gen.distributions import lognormal_from_mean_median from aiperf.dataset.agentic_code_gen.models import ( @@ -160,6 +161,68 @@ def test_sessions_have_at_least_one_turn( assert len(session.turns) >= 1 +class TestInterTurnDelay: + @pytest.mark.parametrize( + "agentic_fraction, expected", + [param(1.0, 900, id="agentic"), param(0.0, 4500, id="human")], + ) # fmt: skip + @pytest.mark.parametrize( + "context_ratio", + [param(0.75, id="partly-full"), param(1.0, id="full")], + ) # fmt: skip + def test_context_scaling_preserves_selected_component_minimum( + self, agentic_fraction: float, expected: float, context_ratio: float + ) -> None: + config = SessionDistributionConfig( + inter_turn_delay=MixtureDelayConfig( + agentic_fraction=agentic_fraction, + agentic_delay=LognormalParams(mean=1000, median=1000, min=900), + human_delay=LognormalParams(mean=5000, median=5000, min=4500), + ) + ) + synth = SessionSynthesizer(config, seed=42) + + delay = synth._sample_delay_ms(int(config.max_prompt_tokens * context_ratio)) + + assert delay == pytest.approx(expected) + + @pytest.mark.parametrize( + "context_ratio", + [ + param(0.0, id="empty"), + param(0.75, id="partly-full"), + param(1.0, id="full"), + ], + ) # fmt: skip + def test_disabled_context_scaling_preserves_sampled_delay( + self, context_ratio: float + ) -> None: + config = SessionDistributionConfig( + inter_turn_delay=MixtureDelayConfig( + context_scaling=False, + agentic_fraction=1.0, + agentic_delay=LognormalParams(mean=1000, median=1000), + ) + ) + synth = SessionSynthesizer(config, seed=42) + + delay = synth._sample_delay_ms(int(config.max_prompt_tokens * context_ratio)) + + assert delay == pytest.approx(1000) + + def test_global_cap_takes_precedence_over_component_minimum(self) -> None: + config = SessionDistributionConfig( + inter_turn_delay=MixtureDelayConfig( + agentic_fraction=0.0, + human_delay=LognormalParams(mean=5000, median=5000, min=4500), + max=600, + ) + ) + synth = SessionSynthesizer(config, seed=42) + + assert synth._sample_delay_ms(config.max_prompt_tokens) == 600 + + class TestExplicitTurnMode: def _overflowing_turn_mode_config( self, allow_truncation: bool