diff --git a/docs/cli-options.md b/docs/cli-options.md index 2b21c65073..60b4bb83da 100644 --- a/docs/cli-options.md +++ b/docs/cli-options.md @@ -1125,7 +1125,12 @@ The maximum duration in seconds for the warmup phase. If not set, it will use th #### `--agentic-cache-warmup-duration` `` -Additional agentic replay warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs, then drains and resumes profiling from the resulting trajectory state using each live stream's residual next-turn delay. +Additional agentic replay warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs, then drains and resumes profiling from the resulting trajectory state using each live stream's residual next-turn delay. Mutually exclusive with --warmup-requests-per-lane. +
_Constraints: > 0_ + +#### `--warmup-requests-per-lane` `` + +Deterministic agentic cache-pressure warmup request budget per concurrency lane. For example, 10 with concurrency 16 caps warmup at 160 wire requests, including initial snapshot priming. Mutually exclusive with --agentic-cache-warmup-duration.
_Constraints: > 0_ #### `--agentic-warmup-grace-period` `` @@ -2655,7 +2660,12 @@ The maximum duration in seconds for the warmup phase. If not set, it will use th #### `--agentic-cache-warmup-duration` `` -Additional agentic replay warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs, then drains and resumes profiling from the resulting trajectory state using each live stream's residual next-turn delay. +Additional agentic replay warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs, then drains and resumes profiling from the resulting trajectory state using each live stream's residual next-turn delay. Mutually exclusive with --warmup-requests-per-lane. +
_Constraints: > 0_ + +#### `--warmup-requests-per-lane` `` + +Deterministic agentic cache-pressure warmup request budget per concurrency lane. For example, 10 with concurrency 16 caps warmup at 160 wire requests, including initial snapshot priming. Mutually exclusive with --agentic-cache-warmup-duration.
_Constraints: > 0_ #### `--agentic-warmup-grace-period` `` diff --git a/docs/tutorials/agentx-mvp.md b/docs/tutorials/agentx-mvp.md index 36a743e2da..778452ce36 100644 --- a/docs/tutorials/agentx-mvp.md +++ b/docs/tutorials/agentx-mvp.md @@ -417,6 +417,17 @@ duration expires, it stops issuing new requests, drains requests already on the wire, snapshots each live root, subagent, and unresolved join, and starts profiling from that exact state. +For repeatable warmup depth, use +`--warmup-requests-per-lane REQUESTS` instead. Each concurrency lane is allowed +exactly that many warmup wire requests, including its initial snapshot-priming +requests. For example, `--concurrency 16` with +`--warmup-requests-per-lane 10` produces 160 warmup requests, with a strict +10-request quota on every lane. + +`--agentic-cache-warmup-duration` and `--warmup-requests-per-lane` are mutually +exclusive: choose a time-bounded warmup or a deterministic request-bounded +warmup. + These requests remain part of warmup, so they are excluded from exported request metrics. diff --git a/src/aiperf/config/config.py b/src/aiperf/config/config.py index d9c5bfd96b..331b847621 100644 --- a/src/aiperf/config/config.py +++ b/src/aiperf/config/config.py @@ -609,7 +609,8 @@ def validate_cache_bust_compatibility(self) -> Self: def validate_agentic_cache_warmup(self) -> Self: """Restrict accelerated cache warmup to the agentic_replay timing mode. - ``--agentic-cache-warmup-duration`` is consumed solely by + The mutually exclusive duration and deterministic request-budget modes + are consumed solely by ``aiperf.timing.config._build_agentic_warmup_config``, which only runs when the profiling phases resolve to AGENTIC_REPLAY. On any other run the value is silently dropped, so an unguarded flag is a no-op the user @@ -630,19 +631,34 @@ def validate_agentic_cache_warmup(self) -> Self: from aiperf.timing.config import _is_agentic_replay profiling_phases = self.get_profiling_phases() - if not any( + has_duration = any( getattr(phase, "agentic_cache_warmup_duration", None) is not None for phase in profiling_phases - ): + ) + has_request_budget = any( + getattr(phase, "warmup_requests_per_lane", None) is not None + for phase in profiling_phases + ) + if not has_duration and not has_request_budget: return self + if any( + getattr(phase, "warmup_requests_per_lane", None) is not None + and getattr(phase, "agentic_cache_warmup_duration", None) is not None + for phase in profiling_phases + ): + raise ValueError( + "--warmup-requests-per-lane and " + "--agentic-cache-warmup-duration are mutually exclusive." + ) + if self.scenario is not None: from aiperf.common.scenario.registry import get_scenario scenario_timing_mode = get_scenario(self.scenario).timing_mode if scenario_timing_mode != TimingMode.AGENTIC_REPLAY: raise ValueError( - "--agentic-cache-warmup-duration requires the agentic_replay " + "agentic cache warmup requires the agentic_replay " f"timing mode; scenario {self.scenario!r} locks " f"timing_mode={scenario_timing_mode}." ) @@ -650,7 +666,7 @@ def validate_agentic_cache_warmup(self) -> Self: if not _is_agentic_replay(profiling_phases): raise ValueError( - "--agentic-cache-warmup-duration requires the agentic_replay " + "agentic cache warmup requires the agentic_replay " "timing mode (set today by --scenario inferencex-agentx-mvp); " "the profiling phase(s) are not agentic_replay." ) diff --git a/src/aiperf/config/flags/_converter_profiling.py b/src/aiperf/config/flags/_converter_profiling.py index c65a161d14..6243eaf8a8 100644 --- a/src/aiperf/config/flags/_converter_profiling.py +++ b/src/aiperf/config/flags/_converter_profiling.py @@ -54,6 +54,7 @@ "burst_phase_starts", "system_idle_gap_cap_seconds", "agentic_cache_warmup_duration", + "warmup_requests_per_lane", "agentic_warmup_grace_period", ) diff --git a/src/aiperf/config/flags/cli_config.py b/src/aiperf/config/flags/cli_config.py index 847e77f508..ee4e0b78b8 100644 --- a/src/aiperf/config/flags/cli_config.py +++ b/src/aiperf/config/flags/cli_config.py @@ -2280,7 +2280,8 @@ def url(self) -> str: "After the normal snapshot warmup drains, AIPerf continues the live " "trajectories without recorded idle delays and with one-token outputs, " "then drains and resumes profiling from the resulting trajectory state " - "using each live stream's residual next-turn delay.", + "using each live stream's residual next-turn delay. Mutually exclusive " + "with --warmup-requests-per-lane.", ), CLIParameter( name=("--agentic-cache-warmup-duration",), @@ -2288,6 +2289,22 @@ def url(self) -> str: ), ] = None + warmup_requests_per_lane: Annotated[ + int | None, + Field( + gt=0, + description="Deterministic agentic cache-pressure warmup request " + "budget per concurrency lane. For example, 10 with concurrency 16 " + "caps warmup at 160 wire requests, including initial snapshot " + "priming. Mutually exclusive with " + "--agentic-cache-warmup-duration.", + ), + CLIParameter( + name=("--warmup-requests-per-lane",), + group=Groups.WARMUP, + ), + ] = None + agentic_warmup_grace_period: Annotated[ float | None, Field( diff --git a/src/aiperf/config/phases.py b/src/aiperf/config/phases.py index ed547e6eef..1633b59222 100644 --- a/src/aiperf/config/phases.py +++ b/src/aiperf/config/phases.py @@ -360,7 +360,21 @@ class BasePhaseConfig(AdaptiveScalePhaseMixin, BaseConfig): "continues the live trajectories without recorded idle delays and with " "one-token outputs for this long, then drains and resumes profiling " "from the resulting trajectory state. Read off the profiling phase by " - "``timing.config._build_agentic_warmup_config``. None disables it.", + "``timing.config._build_agentic_warmup_config``. Mutually exclusive " + "with warmup_requests_per_lane. None disables it.", + ), + ] + + warmup_requests_per_lane: Annotated[ + int | None, + Field( + default=None, + gt=0, + description="AGENTIC_REPLAY only: deterministic cache-pressure " + "warmup request budget per concurrency lane. The total warmup " + "wire-request cap is this value multiplied by the number of live " + "trajectory lanes, including the initial snapshot-priming " + "requests. Mutually exclusive with agentic_cache_warmup_duration.", ), ] diff --git a/src/aiperf/config/schema/aiperf-config.schema.json b/src/aiperf/config/schema/aiperf-config.schema.json index 816887b645..5a4ae0eed7 100644 --- a/src/aiperf/config/schema/aiperf-config.schema.json +++ b/src/aiperf/config/schema/aiperf-config.schema.json @@ -1313,9 +1313,23 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration" }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane" + }, "agenticWarmupGracePeriod": { "anyOf": [ { @@ -1747,9 +1761,23 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration" }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane" + }, "agenticWarmupGracePeriod": { "anyOf": [ { @@ -2249,9 +2277,23 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration" }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane" + }, "agenticWarmupGracePeriod": { "anyOf": [ { @@ -2765,9 +2807,23 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration" }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane" + }, "agenticWarmupGracePeriod": { "anyOf": [ { @@ -3267,9 +3323,23 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration" }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane" + }, "agenticWarmupGracePeriod": { "anyOf": [ { @@ -3746,9 +3816,23 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration" }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane" + }, "agenticWarmupGracePeriod": { "anyOf": [ { @@ -5099,9 +5183,23 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration" }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane" + }, "agenticWarmupGracePeriod": { "anyOf": [ { @@ -5533,9 +5631,23 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration" }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane" + }, "agenticWarmupGracePeriod": { "anyOf": [ { @@ -6035,9 +6147,23 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration" }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane" + }, "agenticWarmupGracePeriod": { "anyOf": [ { @@ -6551,9 +6677,23 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration" }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane" + }, "agenticWarmupGracePeriod": { "anyOf": [ { @@ -7053,9 +7193,23 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration" }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane" + }, "agenticWarmupGracePeriod": { "anyOf": [ { @@ -7532,9 +7686,23 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration" }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane" + }, "agenticWarmupGracePeriod": { "anyOf": [ { @@ -8006,9 +8174,23 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration" }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane" + }, "agenticWarmupGracePeriod": { "anyOf": [ { @@ -8440,9 +8622,23 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration" }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane" + }, "agenticWarmupGracePeriod": { "anyOf": [ { @@ -8942,9 +9138,23 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration" }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane" + }, "agenticWarmupGracePeriod": { "anyOf": [ { @@ -9458,9 +9668,23 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration" }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane" + }, "agenticWarmupGracePeriod": { "anyOf": [ { @@ -9960,9 +10184,23 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration" }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane" + }, "agenticWarmupGracePeriod": { "anyOf": [ { @@ -10439,9 +10677,23 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration" }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane" + }, "agenticWarmupGracePeriod": { "anyOf": [ { @@ -11333,10 +11585,35 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration", "x-jinja2-supported": true }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "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": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane", + "x-jinja2-supported": true + }, "agenticWarmupGracePeriod": { "anyOf": [ { @@ -12069,10 +12346,35 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration", "x-jinja2-supported": true }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "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": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane", + "x-jinja2-supported": true + }, "agenticWarmupGracePeriod": { "anyOf": [ { @@ -14060,10 +14362,35 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration", "x-jinja2-supported": true }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "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": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane", + "x-jinja2-supported": true + }, "agenticWarmupGracePeriod": { "anyOf": [ { @@ -14857,10 +15184,35 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration", "x-jinja2-supported": true }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "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": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane", + "x-jinja2-supported": true + }, "agenticWarmupGracePeriod": { "anyOf": [ { @@ -17451,10 +17803,35 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration", "x-jinja2-supported": true }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "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": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane", + "x-jinja2-supported": true + }, "agenticWarmupGracePeriod": { "anyOf": [ { @@ -21055,10 +21432,35 @@ } ], "default": null, - "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. None disables it.", + "description": "AGENTIC_REPLAY only: additional cache-pressure warmup duration in seconds. After the normal snapshot warmup drains, AIPerf continues the live trajectories without recorded idle delays and with one-token outputs for this long, then drains and resumes profiling from the resulting trajectory state. Read off the profiling phase by ``timing.config._build_agentic_warmup_config``. Mutually exclusive with warmup_requests_per_lane. None disables it.", "title": "Agenticcachewarmupduration", "x-jinja2-supported": true }, + "warmupRequestsPerLane": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "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": "AGENTIC_REPLAY only: deterministic cache-pressure warmup request budget per concurrency lane. The total warmup wire-request cap is this value multiplied by the number of live trajectory lanes, including the initial snapshot-priming requests. Mutually exclusive with agentic_cache_warmup_duration.", + "title": "Warmuprequestsperlane", + "x-jinja2-supported": true + }, "agenticWarmupGracePeriod": { "anyOf": [ { diff --git a/src/aiperf/credit/callback_handler.py b/src/aiperf/credit/callback_handler.py index 972f7a5608..8b7053d31f 100644 --- a/src/aiperf/credit/callback_handler.py +++ b/src/aiperf/credit/callback_handler.py @@ -492,23 +492,27 @@ async def on_credit_return( # behind ``can_send_child_turn`` instead — the phase-level # sending-complete flag is driven by root sampling exhaustion, not # by DAG work, but the global ``--request-count`` cap still - # applies. When the cap blocks a non-final child continuation, we - # notify the orchestrator (``on_child_stopped``) so the parent's - # join still drains instead of deadlocking on a child whose - # remaining turns will never be issued. Final-turn child returns - # are always passed through (the strategy is a no-op for them, but - # observer hooks still need to fire). + # applies. In terminal phases, a blocked non-final child notifies the + # orchestrator (``on_child_stopped``) so the parent's join can drain. + # Strategies requesting stopped returns (such as quota warmup) instead + # preserve resumable children across a phase handoff. Final-turn child + # returns are always passed through (the strategy is a no-op for them, + # but observer hooks still need to fire). + wants_stopped_returns = ( + getattr(handler.strategy, "wants_returns_after_sending_complete", False) + is True + ) is_child = credit.agent_depth > 0 if not is_child: - wants_stopped_returns = ( - getattr(handler.strategy, "wants_returns_after_sending_complete", False) - is True - ) if handler.stop_checker.can_send_any_turn() or wants_stopped_returns: await handler.strategy.handle_credit_return( credit, error=credit_return.error ) - elif credit.is_final_turn or handler.stop_checker.can_send_child_turn(): + elif ( + credit.is_final_turn + or handler.stop_checker.can_send_child_turn() + or wants_stopped_returns + ): await handler.strategy.handle_credit_return( credit, error=credit_return.error ) diff --git a/src/aiperf/credit/issuer.py b/src/aiperf/credit/issuer.py index 2023f25d9c..4a899239c3 100644 --- a/src/aiperf/credit/issuer.py +++ b/src/aiperf/credit/issuer.py @@ -15,6 +15,7 @@ from __future__ import annotations import time +from collections.abc import Callable from typing import TYPE_CHECKING from msgspec.structs import replace as _struct_replace @@ -125,8 +126,18 @@ def __init__( ) self._issuing_stopped = False self._max_tokens_override: int | None = None + self._turn_admission: Callable[[TurnToSend], bool] | None = None self.replay_gate = ReplayIssueGate(replay_barrier) + def set_turn_admission(self, callback: Callable[[TurnToSend], bool]) -> None: + """Install a synchronous final admission check for every turn.""" + self._turn_admission = callback + + def _is_turn_admitted(self, turn: TurnToSend) -> bool: + """Run the optional final admission check.""" + callback = getattr(self, "_turn_admission", None) + return callback is None or callback(turn) + def set_max_tokens_override(self, max_tokens: int | None) -> None: """Override generation length for every subsequently issued credit.""" self._max_tokens_override = max_tokens @@ -342,6 +353,12 @@ async def _issue_credit_ready(self, turn: TurnToSend) -> bool: self._concurrency_manager.release_session_slot(self._phase_key) return False + if not self._is_turn_admitted(turn): + self._concurrency_manager.release_prefill_slot(self._phase_key) + if needs_session_slot: + self._concurrency_manager.release_session_slot(self._phase_key) + return False + # Both slots held: register the tree before issuing so drain/teardown # own the slot release. Must not run before prefill succeeds. if needs_session_slot: @@ -400,6 +417,12 @@ async def try_issue_credit(self, turn: TurnToSend) -> bool | None: self._concurrency_manager.release_session_slot(self._phase_key) return None # No slot - credit not issued + if not self._is_turn_admitted(turn): + self._concurrency_manager.release_prefill_slot(self._phase_key) + if needs_session_slot: + self._concurrency_manager.release_session_slot(self._phase_key) + return False + if needs_session_slot: self._open_session_tree(turn) @@ -520,6 +543,9 @@ async def _dispatch_child_turn_ready(self, turn: TurnToSend) -> bool: self._phase_key, can_proceed_fn ): return False + if not self._is_turn_admitted(turn): + self._concurrency_manager.release_prefill_slot(self._phase_key) + return False if turn.counts_toward_phase_target: turn = _struct_replace(turn, counts_toward_phase_target=False) await self._issue_credit_internal(turn) diff --git a/src/aiperf/dataset/loader/weka_synth_buf.py b/src/aiperf/dataset/loader/weka_synth_buf.py index 1b7b49ba65..329a47d752 100644 --- a/src/aiperf/dataset/loader/weka_synth_buf.py +++ b/src/aiperf/dataset/loader/weka_synth_buf.py @@ -431,7 +431,7 @@ def _assert_trailing_user(self) -> None: def turn_delta(self) -> TurnDelta: """Compute the raw_messages to emit for the just-completed turn. - Three cases: + Four cases: 1. First call after ``init_turn_0`` (``_emitted_segment_count == 0``): emit ALL current segments, ``reset_context=False``. This is turn 0's baseline state. @@ -441,6 +441,10 @@ def turn_delta(self) -> TurnDelta: 3. Disturbance touched a previously-emitted segment (index ``< _emitted_segment_count``): emit ALL current segments, ``reset_context=True``. + 4. An equal-context retry appended no segments: re-emit ALL current + segments with ``reset_context=True``. An empty delta would + otherwise render as an invalid empty user message instead of the + recorded repeated request. Updates ``_emitted_segment_count`` to ``len(self._segments)`` on return. Clears ``_last_disturbance_at`` to ``None``. @@ -449,9 +453,16 @@ def turn_delta(self) -> TurnDelta: self._last_disturbance_at is not None and self._last_disturbance_at < self._emitted_segment_count ) - if self._emitted_segment_count == 0 or disturbed_emitted: + unchanged_retry = ( + self._emitted_segment_count > 0 + and self._emitted_segment_count == len(self._segments) + and not disturbed_emitted + ) + if self._emitted_segment_count == 0 or disturbed_emitted or unchanged_retry: source = self._segments - reset = self._emitted_segment_count != 0 and disturbed_emitted + reset = self._emitted_segment_count != 0 and ( + disturbed_emitted or unchanged_retry + ) else: source = self._segments[self._emitted_segment_count :] reset = False diff --git a/src/aiperf/orchestrator/strategies.py b/src/aiperf/orchestrator/strategies.py index ada0694182..1dae177069 100644 --- a/src/aiperf/orchestrator/strategies.py +++ b/src/aiperf/orchestrator/strategies.py @@ -321,6 +321,7 @@ def _disable_warmup(self, config: BenchmarkConfig) -> BenchmarkConfig: config.phases = [p for p in config.phases if not p.exclude_from_results] for phase in config.get_profiling_phases(): phase.agentic_cache_warmup_duration = None + phase.warmup_requests_per_lane = None return config diff --git a/src/aiperf/timing/config.py b/src/aiperf/timing/config.py index aaa85ed328..c0a982d7a9 100644 --- a/src/aiperf/timing/config.py +++ b/src/aiperf/timing/config.py @@ -351,6 +351,13 @@ class CreditPhaseConfig(AIPerfBaseModel): description="Duration of the accelerated cache-pressure substage for " "agentic replay warmup.", ) + warmup_requests_per_lane: int | None = Field( + default=None, + gt=0, + description="Deterministic cache-pressure warmup wire-request budget " + "per live agentic replay lane. Mutually exclusive with " + "agentic_cache_warmup_duration_sec.", + ) artifact_dir: Path | None = Field( default=None, @@ -614,15 +621,25 @@ def _build_agentic_warmup_config(phase: PhaseConfig) -> CreditPhaseConfig | None concurrency = getattr(phase, "concurrency", None) grace_period = _agentic_warmup_grace_period(phase) cache_warmup_duration = getattr(phase, "agentic_cache_warmup_duration", None) + requests_per_lane = getattr(phase, "warmup_requests_per_lane", None) + cache_warmup_request_cap = ( + concurrency * requests_per_lane + if concurrency is not None and requests_per_lane is not None + else None + ) + if cache_warmup_request_cap is not None: + total_expected_requests = cache_warmup_request_cap + elif cache_warmup_duration is not None: + total_expected_requests = None + else: + total_expected_requests = concurrency return CreditPhaseConfig( phase=CreditPhase.WARMUP, timing_mode=TimingMode.AGENTIC_REPLAY, - # An accelerated cache-pressure warmup is strategy-terminated (the - # strategy emits ``mark_sending_complete`` when the duration elapses), - # so leave the request cap open instead of sizing it to concurrency. - total_expected_requests=( - None if cache_warmup_duration is not None else concurrency - ), + # Duration mode is strategy-terminated by its timer. Count mode uses + # the generic request-count stop condition as a global backstop while + # the agentic strategy independently enforces each lane's quota. + total_expected_requests=total_expected_requests, expected_duration_sec=None, expected_num_sessions=None, concurrency=concurrency, @@ -633,6 +650,7 @@ def _build_agentic_warmup_config(phase: PhaseConfig) -> CreditPhaseConfig | None seamless=False, grace_period_sec=grace_period if grace_period is not None else float("inf"), agentic_cache_warmup_duration_sec=cache_warmup_duration, + warmup_requests_per_lane=requests_per_lane, ) diff --git a/src/aiperf/timing/phase/runner.py b/src/aiperf/timing/phase/runner.py index b27de01b71..5d012657d0 100644 --- a/src/aiperf/timing/phase/runner.py +++ b/src/aiperf/timing/phase/runner.py @@ -131,9 +131,12 @@ def __init__( self._branch_orchestrator = branch_orchestrator self._run = run self._session_tree_registry = session_tree_registry - self._cache_warmup_enabled = isinstance( - getattr(config, "agentic_cache_warmup_duration_sec", None), - int | float, + self._cache_warmup_enabled = ( + isinstance( + getattr(config, "agentic_cache_warmup_duration_sec", None), + int | float, + ) + or getattr(config, "warmup_requests_per_lane", None) is not None ) # For FIXED_SCHEDULE mode, use actual dataset size instead of config values. @@ -150,22 +153,24 @@ def __init__( elif ( config.timing_mode == TimingMode.AGENTIC_REPLAY and config.phase == CreditPhase.WARMUP - and not self._cache_warmup_enabled ): - # AGENTIC_REPLAY warmup dispatches one priming credit per warmable - # stream (root + each mid-flight subagent at t*), which exceeds the - # `concurrency` placeholder when lanes hold multiple streams. Without - # this re-anchor the concurrency-sized barrier fires early and cancels - # the closest-to-t* priming credits -- under-priming the server cache - # and masking warmup failures for the cancelled streams. Re-anchor the - # barrier to the actual dispatch count (``warmup_credit_count`` - # promises exactly this). Single-stream lanes already equal - # concurrency, so this is a no-op for them. - warmup_count = getattr(conversation_source, "warmup_credit_count", None) - if warmup_count: + requests_per_lane = getattr(config, "warmup_requests_per_lane", None) + if self._cache_warmup_enabled and requests_per_lane is not None: + lane_count = len(getattr(conversation_source, "trajectories", ())) + request_cap = requests_per_lane * lane_count self._config = config.model_copy( - update={"total_expected_requests": warmup_count} + update={"total_expected_requests": request_cap} ) + elif not self._cache_warmup_enabled: + # AGENTIC_REPLAY warmup dispatches one priming credit per + # warmable stream (root + each mid-flight subagent at t*), which + # exceeds the `concurrency` placeholder when lanes hold multiple + # streams. Re-anchor the barrier to the actual dispatch count. + warmup_count = getattr(conversation_source, "warmup_credit_count", None) + if warmup_count: + self._config = config.model_copy( + update={"total_expected_requests": warmup_count} + ) self._phase_publisher = phase_publisher self._credit_router = credit_router self._concurrency_manager = concurrency_manager diff --git a/src/aiperf/timing/strategies/agentic_replay.py b/src/aiperf/timing/strategies/agentic_replay.py index ccb11cd8cd..bacc7b9927 100644 --- a/src/aiperf/timing/strategies/agentic_replay.py +++ b/src/aiperf/timing/strategies/agentic_replay.py @@ -179,6 +179,17 @@ def __init__( if isinstance(cache_warmup_duration, int | float) else None ) + cache_warmup_requests_per_lane = getattr( + config, "warmup_requests_per_lane", None + ) + self._cache_warmup_requests_per_lane: int | None = ( + int(cache_warmup_requests_per_lane) + if isinstance(cache_warmup_requests_per_lane, int) + else None + ) + self._cache_warmup_requests_by_lane: Counter[int] = Counter() + self._cache_warmup_request_budget_reached = False + self._quota_handoff_starts: dict[str, TurnToSend] = {} self._baseline_warmup_returns: dict[str, Credit] = {} self._baseline_correlations: set[str] = set() self._accelerated_warmup_started = False @@ -243,8 +254,9 @@ def __init__( self.scheduler.set_drain_observer(self.enforce_system_idle_cap) # Idle-gap cap (ms) for the t* boundary the load-time warp can't see (t* # is the sampling instant, not a request). Consumed two ways: - # - WARMUP: clamps each warmup lead so priming doesn't start hours - # early (``_capped_warmup_lead_ms``); priming spacing is meaningless. + # - WARMUP: this cap and the global system-idle cap both clamp each + # warmup lead so priming doesn't start hours early + # (``_capped_warmup_lead_ms``); priming spacing is meaningless. # - PROFILING: a single uniform shift caps the leading idle (t* -> # earliest stream) while preserving recorded inter-stream spacing # (``_leading_idle_shift_ms``); a per-stream clamp would collapse @@ -275,6 +287,14 @@ def __init__( sum(self._lanes_per_trace.values()), ) + @property + def _cache_warmup_enabled(self) -> bool: + """Whether either accelerated cache-pressure warmup mode is active.""" + return ( + self._cache_warmup_duration is not None + or self._cache_warmup_requests_per_lane is not None + ) + @property def _has_tree_registry(self) -> bool: """True when per-tree session-slot accounting is engaged. @@ -283,17 +303,13 @@ def _has_tree_registry(self) -> bool: engages it because it opens trees and spawns descendants during WARMUP. """ return self._session_tree_registry is not None and ( - self.config.phase == CreditPhase.PROFILING - or self._cache_warmup_duration is not None + self.config.phase == CreditPhase.PROFILING or self._cache_warmup_enabled ) @property def wants_returns_after_sending_complete(self) -> bool: """Pressure warmup returns must be observed to build the handoff state.""" - return ( - self.config.phase == CreditPhase.WARMUP - and self._cache_warmup_duration is not None - ) + return self.config.phase == CreditPhase.WARMUP and self._cache_warmup_enabled @property def allows_pending_branch_handoff_after_sending_complete(self) -> bool: @@ -398,6 +414,11 @@ async def setup_phase(self) -> None: """ if self._has_tree_registry: self._session_tree_registry.set_drain_callback(self._on_tree_drained) + if ( + self.config.phase == CreditPhase.WARMUP + and self._cache_warmup_requests_per_lane is not None + ): + self.credit_issuer.set_turn_admission(self._admit_cache_warmup_turn) if self.config.phase == CreditPhase.PROFILING: for trajectory in self.conversation_source.trajectories: self._seed_trajectory_replay_prefix(trajectory) @@ -422,6 +443,49 @@ async def setup_phase(self) -> None: f"fresh root once their background subagents drain" ) + def _cache_warmup_lane(self, turn_or_credit: TurnToSend | Credit) -> int: + """Resolve a warmup turn or credit to its stable trajectory lane.""" + lane = self._root_to_lane.get(turn_or_credit.effective_root_correlation_id) + if lane is None: + lane = self._correlation_to_lane.get(turn_or_credit.x_correlation_id) + if lane is None: + raise RuntimeError( + "Agentic cache warmup could not resolve a request to a " + f"trajectory lane: correlation_id={turn_or_credit.x_correlation_id!r}, " + "root_correlation_id=" + f"{turn_or_credit.effective_root_correlation_id!r}" + ) + return lane + + def _admit_cache_warmup_turn(self, turn: TurnToSend) -> bool: + """Atomically reserve one request from a lane's deterministic quota.""" + assert self._cache_warmup_requests_per_lane is not None + lane = self._cache_warmup_lane(turn) + if ( + self._cache_warmup_requests_by_lane[lane] + >= self._cache_warmup_requests_per_lane + ): + if turn.agent_depth == 0 and ( + turn.turn_index == 0 or turn.is_session_start + ): + self._quota_handoff_starts[turn.effective_root_correlation_id] = turn + return False + self._cache_warmup_requests_by_lane[lane] += 1 + if not self._cache_warmup_request_budget_reached and all( + self._cache_warmup_requests_by_lane[lane_index] + >= self._cache_warmup_requests_per_lane + for lane_index in range(len(self.conversation_source.trajectories)) + ): + self._cache_warmup_request_budget_reached = True + self.credit_issuer.replay_gate.pause_releases() + self.info( + "WARMUP cache pressure request budget reached: " + f"{self._cache_warmup_requests_per_lane} requests on each of " + f"{len(self.conversation_source.trajectories)} lanes; " + "draining requests" + ) + return True + async def execute_phase(self) -> None: """Dispatch initial credits for the phase.""" if self.config.phase == CreditPhase.WARMUP: @@ -485,9 +549,19 @@ def _capped_warmup_lead_ms(self, lead_ms: float) -> float: PROFILING dispatch offsets are NOT clamped this way -- see :meth:`_leading_idle_shift_ms`. """ - if self._phase_offset_cap_ms is not None: - return min(lead_ms, self._phase_offset_cap_ms) - return lead_ms + caps_ms = [ + cap + for cap in ( + self._phase_offset_cap_ms, + ( + self._system_idle_gap_cap_seconds * MILLIS_PER_SECOND + if self._system_idle_gap_cap_seconds is not None + else None + ), + ) + if cap is not None + ] + return min(lead_ms, *caps_ms) if caps_ms else lead_ms def _leading_idle_shift_ms(self, offsets: Iterable[float]) -> float: """Excess to subtract UNIFORMLY from every PROFILING dispatch offset so @@ -591,13 +665,30 @@ async def _execute_warmup(self) -> None: self._baseline_correlations.add(turn.x_correlation_id) self._root_to_lane[turn.effective_root_correlation_id] = lane + if self._cache_warmup_requests_per_lane is not None: + baseline_counts = Counter( + self._cache_warmup_lane(turn) for turn, _ in prepared + ) + oversized = { + lane: count + for lane, count in baseline_counts.items() + if count > self._cache_warmup_requests_per_lane + } + if oversized: + raise ValueError( + "Agentic cache warmup requests-per-lane budget is smaller " + "than the required snapshot-priming dispatch count for " + f"lane(s) {oversized}; increase " + "--warmup-requests-per-lane." + ) + # Nothing to warm: every lane's first request is at/after t* (no turn # precedes t*), so no warmup credit will dispatch. The count path that # normally drives completion is triggered by credit dispatch/return, so # with zero credits it would never fire and the warmup phase would hang # waiting on a barrier sized to concurrency. Finalize immediately. if not prepared: - if self._cache_warmup_duration is not None: + if self._cache_warmup_enabled: # No baseline priming to wait for; jump straight to the # accelerated cache-pressure substage. await self._start_accelerated_warmup() @@ -644,37 +735,41 @@ async def _execute_warmup(self) -> None: async def _finish_initial_warmup_dispatch(self) -> None: """Mark sending complete for burst warmup with no cache-pressure stage. - When a cache-pressure duration is set the accelerated substage is driven - by baseline credit returns (``_handle_warmup_return``), so there is - nothing to finalize here in that case. + When either cache-pressure mode is set, the accelerated substage is + driven by baseline credit returns (``_handle_warmup_return``), so there + is nothing to finalize here. """ - if ( - self._cache_warmup_duration is None - and not self.lifecycle.is_sending_complete - ): + if not self._cache_warmup_enabled and not self.lifecycle.is_sending_complete: self.lifecycle.mark_sending_complete() async def _start_accelerated_warmup(self) -> None: """Continue the sampled trajectories under compressed warmup traffic.""" if self._accelerated_warmup_started: return - assert self._cache_warmup_duration is not None + assert self._cache_warmup_enabled self._accelerated_warmup_started = True self.credit_issuer.set_max_tokens_override(_WARMUP_MAX_TOKENS) for trajectory in self.conversation_source.trajectories: self._seed_trajectory_replay_prefix(trajectory) self.credit_issuer.replay_gate.activate() + if self._cache_warmup_duration is not None: + limit = f"for {self._cache_warmup_duration:.1f}s" + else: + limit = ( + f"until each lane reaches " + f"{self._cache_warmup_requests_per_lane} requests" + ) self.info( - "WARMUP cache pressure: replaying live trajectories for " - f"{self._cache_warmup_duration:.1f}s with zero idle delay and " - f"max_tokens={_WARMUP_MAX_TOKENS}" + "WARMUP cache pressure: replaying live trajectories " + f"{limit} with zero idle delay and max_tokens={_WARMUP_MAX_TOKENS}" ) if self.branch_orchestrator is not None: self.branch_orchestrator.start_accelerated_warmup() - self.scheduler.schedule_later( - self._cache_warmup_duration, - self._finish_accelerated_warmup(), - ) + if self._cache_warmup_duration is not None: + self.scheduler.schedule_later( + self._cache_warmup_duration, + self._finish_accelerated_warmup(), + ) results = await asyncio.gather( *( self._dispatch_accelerated_trajectory(trajectory, lane) @@ -794,15 +889,21 @@ async def _handle_accelerated_warmup_return(self, credit: Credit) -> None: await self.credit_issuer.issue_credit(turn) async def _issue_child_continuation_or_drain(self, turn: TurnToSend) -> None: - """Dispatch a DAG child continuation, draining the join on refusal. + """Dispatch a DAG child continuation, draining terminal refusals. ``dispatch_child_turn`` returns True iff the turn reached the wire; on - any refusal (e.g. the ``--request-count`` cap) notify the orchestrator - so the parent's join drains deterministically instead of deadlocking on - a child whose remaining turns will never be issued. + a terminal refusal notify the orchestrator so the parent's join drains + deterministically instead of deadlocking on a child whose remaining + turns will never be issued. Accelerated warmup refusals are different: + the remaining child and its active join are persisted for profiling, so + marking the child stopped here would release the parent prematurely. """ on_wire = await self.credit_issuer.dispatch_child_turn(turn) - if not on_wire and self.branch_orchestrator is not None: + if ( + not on_wire + and self.branch_orchestrator is not None + and not self.allows_pending_branch_handoff_after_sending_complete + ): await self.branch_orchestrator.on_child_stopped(turn.x_correlation_id) async def finalize_phase(self) -> None: @@ -960,6 +1061,9 @@ def _pending_handoff_turns_by_root(self) -> dict[str, tuple[TurnToSend, ...]]: pending_by_root = dict(pending_by_root_getter()) else: pending_by_root = {} + for root_correlation_id, turn in self._quota_handoff_starts.items(): + pending_by_root.setdefault(root_correlation_id, ()) + pending_by_root[root_correlation_id] += (turn,) pending_turns_getter = getattr( self.credit_issuer.replay_gate, "pending_turns", None ) @@ -1356,7 +1460,7 @@ async def handle_credit_return( async def _handle_warmup_return(self, credit: Credit) -> None: """Advance baseline warmup into the optional cache-pressure stage.""" - if self._cache_warmup_duration is None: + if not self._cache_warmup_enabled: return if self._accelerated_warmup_started: await self._handle_accelerated_warmup_return(credit) @@ -1373,11 +1477,12 @@ async def _dispatch_next_turn(self, credit: Credit) -> None: DAG child continuations (``agent_depth > 0``) go through the single child-issuance chokepoint (``_issue_child_continuation_or_drain``) so a - ``--request-count`` cap refusal is routed to ``on_child_stopped`` (drain - the parent join) instead of being silently swallowed by the discarded - ``issue_credit`` return -- including on the delayed (``delay_ms``) path, - where the refusal would otherwise fire long after the callback handler - decided the child could proceed. Root continuations keep ``issue_credit``. + terminal refusal is routed to ``on_child_stopped`` (drain the parent + join) instead of being silently swallowed by the discarded + ``issue_credit`` return. A warmup cutoff is preserved for profiling + handoff instead. This applies equally to delayed continuations, whose + refusal may happen after the callback handler decided the child could + proceed. Root continuations keep ``issue_credit``. """ next_meta = self.conversation_source.get_next_turn_metadata(credit) turn = TurnToSend.from_previous_credit(credit, next_meta) diff --git a/tests/component_integration/test_agentic_replay_cli_e2e.py b/tests/component_integration/test_agentic_replay_cli_e2e.py index d92019be7c..d610aaa95c 100644 --- a/tests/component_integration/test_agentic_replay_cli_e2e.py +++ b/tests/component_integration/test_agentic_replay_cli_e2e.py @@ -189,13 +189,21 @@ def test_agentic_replay_cli_scenario_unsafe_override_runs_to_completion( @pytest.mark.component_integration +@pytest.mark.parametrize( + "warmup_option", + [ + "--agentic-cache-warmup-duration 2", + "--warmup-requests-per-lane 2", + ], +) def test_agentic_replay_cli_cache_warmup_runs_to_completion( cli: AIPerfCLI, weka_small_dir: Path, + warmup_option: str, ) -> None: - """E2E smoke that ``--agentic-cache-warmup-duration`` runs the warmup substage, drain, and profiling handoff without deadlocking the replay barriers.""" + """Both cache-pressure modes complete their warmup, drain, and handoff.""" cmd = _build_command(weka_small_dir, scenario=True, unsafe_override=True) - cmd += " --agentic-cache-warmup-duration 2" + cmd += f" {warmup_option}" result = cli.run_sync(cmd, timeout=defaults.timeout) assert result.exit_code == 0, ( diff --git a/tests/unit/cli_runner/test_agentic_replay_phase_override.py b/tests/unit/cli_runner/test_agentic_replay_phase_override.py index 5b3c550854..cfecc46f20 100644 --- a/tests/unit/cli_runner/test_agentic_replay_phase_override.py +++ b/tests/unit/cli_runner/test_agentic_replay_phase_override.py @@ -68,6 +68,18 @@ def test_agentic_cache_warmup_duration_overrides_yaml_profiling_phase( assert _profiling_phase(cfg).agentic_cache_warmup_duration == 30.0 +def test_agentic_cache_warmup_request_budget_overrides_yaml_profiling_phase( + tmp_path: pathlib.Path, +) -> None: + cfg = resolve_config( + _cli(warmup_requests_per_lane=10), + _agentic_yaml(tmp_path), + ) + phase = _profiling_phase(cfg) + assert phase.agentic_cache_warmup_duration is None + assert phase.warmup_requests_per_lane == 10 + + def test_agentic_replay_sibling_flags_override_yaml_profiling_phase() -> None: """The four sibling agentic-replay phase flags overlay the profiling phase.""" cfg = resolve_config( diff --git a/tests/unit/config/test_validators.py b/tests/unit/config/test_validators.py index 70140fe8c1..40ec8efade 100644 --- a/tests/unit/config/test_validators.py +++ b/tests/unit/config/test_validators.py @@ -263,6 +263,27 @@ def test_agentic_cache_warmup_with_explicit_agentic_timing_mode_accepted() -> No assert cfg.benchmark.phases[0].agentic_cache_warmup_duration == 30.0 +def test_agentic_cache_warmup_request_budget_without_duration_accepted() -> None: + cfg = _make( + phases=_agentic_phase( + warmup_requests_per_lane=10, + timing_mode="agentic_replay", + ) + ) + assert cfg.benchmark.phases[0].warmup_requests_per_lane == 10 + + +def test_agentic_cache_warmup_modes_are_mutually_exclusive() -> None: + with pytest.raises(ValueError, match="mutually exclusive"): + _make( + phases=_agentic_phase( + agentic_cache_warmup_duration=30.0, + warmup_requests_per_lane=10, + timing_mode="agentic_replay", + ) + ) + + def test_no_agentic_cache_warmup_duration_accepted() -> None: cfg = _make(phases=_agentic_phase()) assert cfg.benchmark.phases[0].agentic_cache_warmup_duration is None diff --git a/tests/unit/credit/test_callback_handler.py b/tests/unit/credit/test_callback_handler.py index 7aa13f1c08..eabd9b99ed 100644 --- a/tests/unit/credit/test_callback_handler.py +++ b/tests/unit/credit/test_callback_handler.py @@ -790,6 +790,42 @@ async def test_cache_warmup_handoff_allows_paused_dag_work( mock_branch_orchestrator.has_pending_branch_work.assert_called_once_with() +@pytest.mark.asyncio +async def test_cache_warmup_handoff_preserves_non_final_child( + callback_handler, + mock_progress, + mock_lifecycle, + mock_stop_checker, + mock_strategy, + mock_branch_orchestrator, +): + """A quota-stopped warmup child remains live for profiling handoff.""" + mock_stop_checker.can_send_child_turn = MagicMock(return_value=False) + mock_strategy.wants_returns_after_sending_complete = True + mock_branch_orchestrator.has_pending_branch_work = MagicMock(return_value=True) + mock_branch_orchestrator.intercept = AsyncMock(return_value=False) + mock_branch_orchestrator.on_child_stopped = AsyncMock() + callback_handler.set_branch_orchestrator(mock_branch_orchestrator) + callback_handler.register_phase( + phase=CreditPhase.WARMUP, + progress=mock_progress, + lifecycle=mock_lifecycle, + stop_checker=mock_stop_checker, + strategy=mock_strategy, + ) + + credit = make_credit( + phase=CreditPhase.WARMUP, + turn_index=1, + num_turns=7, + agent_depth=1, + ) + await callback_handler.on_credit_return("worker-1", make_credit_return(credit)) + + mock_strategy.handle_credit_return.assert_awaited_once_with(credit, error=None) + mock_branch_orchestrator.on_child_stopped.assert_not_awaited() + + # ============================================================================= # Test: Credit Return - Unregistered/Complete Phase # ============================================================================= diff --git a/tests/unit/credit/test_issuer.py b/tests/unit/credit/test_issuer.py index d316d13e94..3e5de4ae4e 100644 --- a/tests/unit/credit/test_issuer.py +++ b/tests/unit/credit/test_issuer.py @@ -216,6 +216,49 @@ async def test_issue_credit_returns_false_when_final_credit( assert result is False + async def test_turn_admission_refusal_releases_acquired_slots( + self, credit_issuer, mock_concurrency, mock_progress, mock_router + ): + """A quota refusal happens after acquisition and must release both slots.""" + credit_issuer.set_turn_admission(lambda _turn: False) + + result = await credit_issuer.issue_credit(make_turn()) + + assert result is False + mock_concurrency.release_prefill_slot.assert_called_once_with( + CreditPhase.PROFILING + ) + mock_concurrency.release_session_slot.assert_called_once_with( + CreditPhase.PROFILING + ) + mock_progress.increment_sent.assert_not_called() + mock_router.send_credit.assert_not_called() + + async def test_child_turn_admission_refusal_releases_prefill_slot( + self, credit_issuer, mock_concurrency, mock_progress, mock_router + ): + """DAG children share the same quota gate without owning a session slot.""" + credit_issuer.set_turn_admission(lambda _turn: False) + child = TurnToSend( + conversation_id="child", + x_correlation_id="child-corr", + turn_index=0, + num_turns=1, + agent_depth=1, + parent_correlation_id="parent-corr", + root_correlation_id="root-corr", + ) + + result = await credit_issuer.dispatch_child_turn(child) + + assert result is False + mock_concurrency.release_prefill_slot.assert_called_once_with( + CreditPhase.PROFILING + ) + mock_concurrency.release_session_slot.assert_not_called() + mock_progress.increment_sent.assert_not_called() + mock_router.send_credit.assert_not_called() + # ============================================================================= # Test: Slot Acquisition Failures diff --git a/tests/unit/dataset/loader/test_weka_pathological.py b/tests/unit/dataset/loader/test_weka_pathological.py index 93829601f2..dd295aae42 100644 --- a/tests/unit/dataset/loader/test_weka_pathological.py +++ b/tests/unit/dataset/loader/test_weka_pathological.py @@ -501,6 +501,29 @@ def test_duplicate_hash_ids_in_request_inflate_theoretical_hit_to_full(tmp_path) ) +def test_equal_context_retry_reemits_nonempty_full_prompt(tmp_path): + """An identical retry must not become an invalid empty user message.""" + trace = _base_trace( + [ + _normal(0.0, [1, 2], in_tokens=128), + _normal(1.0, [1, 2], in_tokens=128), + ], + trace_id="equal_retry", + ) + path = tmp_path / "t.json" + path.write_text(json.dumps(trace)) + loader = _make_loader(path, _mk_user_config()) + + convs = loader.convert_to_conversations(loader.load_dataset()) + first, retry = convs[0].turns + + assert first.raw_messages + assert retry.raw_messages == first.raw_messages + assert retry.reset_context is True + assert retry.raw_messages[-1]["role"] == "user" + assert retry.raw_messages[-1]["content"] + + def test_empty_requests_trace_reconstructs_empty_conversation(tmp_path): """A trace with zero requests yields a single empty Conversation, no crash.""" trace = _base_trace([], trace_id="empty_trace") diff --git a/tests/unit/dataset/loader/test_weka_synth_buf_turn_delta.py b/tests/unit/dataset/loader/test_weka_synth_buf_turn_delta.py index ec1a353399..0f52818109 100644 --- a/tests/unit/dataset/loader/test_weka_synth_buf_turn_delta.py +++ b/tests/unit/dataset/loader/test_weka_synth_buf_turn_delta.py @@ -119,6 +119,36 @@ def test_turn_delta_case_1_strict_append_emits_only_new_segments(): assert r._last_disturbance_at is None +def test_turn_delta_equal_context_retry_reemits_full_context(): + """An unchanged retry resets to the full prompt instead of emitting ``[]``.""" + r = _make_recon() + hash_ids = [1, 2] + in_tokens = 2 * BLOCK_SIZE + r.init_turn_0( + hash_ids=hash_ids, + in_tokens=in_tokens, + tool_tokens=0, + system_tokens=0, + seed="t:0", + ) + first = r.turn_delta() + + r.advance_turn( + prev_hash_ids=hash_ids, + prev_in_tokens=in_tokens, + prev_out_tokens=BLOCK_SIZE, + curr_hash_ids=hash_ids, + curr_in_tokens=in_tokens, + seed="t:1", + ) + retry = r.turn_delta() + + assert retry.reset_context is True + assert retry.delta_messages == first.delta_messages + assert retry.delta_messages + assert retry.delta_messages[-1]["role"] == "user" + + def test_turn_delta_case_1_strict_append_three_turns_chain(): """Three sequential strict-append advances: each delta is incremental.""" r = _make_recon() diff --git a/tests/unit/orchestrator/test_strategies.py b/tests/unit/orchestrator/test_strategies.py index 72dbf38b7f..c26cbaf130 100644 --- a/tests/unit/orchestrator/test_strategies.py +++ b/tests/unit/orchestrator/test_strategies.py @@ -339,6 +339,39 @@ def test_disable_warmup_clears_agentic_cache_warmup_duration(self): # Original config untouched (deep copy). assert config.get_profiling_phases()[0].agentic_cache_warmup_duration == 30.0 + def test_disable_warmup_clears_per_lane_request_budget(self): + from aiperf.plugin.enums import TimingMode + + strategy = FixedTrialsStrategy(num_trials=3, disable_warmup_after_first=True) + config = _make_config( + phases=[ + { + "name": "profiling", + "type": "concurrency", + "requests": 100, + "concurrency": 1, + "timing_mode": TimingMode.AGENTIC_REPLAY, + "warmup_requests_per_lane": 10, + }, + ], + ) + + first_config = strategy.get_next_config(config, []) + assert first_config.get_profiling_phases()[0].warmup_requests_per_lane == 10 + + results = [ + RunResult( + label="run_0001", + success=True, + summary_metrics={"ttft": JsonMetricResult(unit="ms", avg=100.0)}, + artifacts_path=Path("/tmp/run_0001"), + ) + ] + second_config = strategy.get_next_config(config, results) + + assert second_config.get_profiling_phases()[0].warmup_requests_per_lane is None + assert config.get_profiling_phases()[0].warmup_requests_per_lane == 10 + def test_get_run_path(self): """Test get_run_path returns correct path structure.""" strategy = FixedTrialsStrategy(num_trials=3) diff --git a/tests/unit/timing/phase/test_runner_agentic_replay_warmup_target.py b/tests/unit/timing/phase/test_runner_agentic_replay_warmup_target.py index cf11975dae..24fb2ed124 100644 --- a/tests/unit/timing/phase/test_runner_agentic_replay_warmup_target.py +++ b/tests/unit/timing/phase/test_runner_agentic_replay_warmup_target.py @@ -158,6 +158,22 @@ async def test_warmup_target_reanchored_to_warmup_credit_count(self) -> None: runner = _make_runner(_warmup_config(concurrency=2), src) assert runner._config.total_expected_requests == 6 + async def test_cache_warmup_target_uses_actual_lane_count(self) -> None: + """The placeholder is re-anchored to the actual wrap-filled trajectory lanes.""" + src = MagicMock() + src.dataset_metadata = None + src.trajectories = [MagicMock(), MagicMock(), MagicMock()] + config = _warmup_config(concurrency=4).model_copy( + update={ + "warmup_requests_per_lane": 10, + "total_expected_requests": 40, + } + ) + + runner = _make_runner(config, src) + + assert runner._config.total_expected_requests == 30 + async def test_profiling_not_reanchored_to_warmup_count(self) -> None: """PROFILING must NOT be re-anchored to warmup_credit_count.""" src = MagicMock() diff --git a/tests/unit/timing/strategies/test_agentic_replay.py b/tests/unit/timing/strategies/test_agentic_replay.py index 9ff5530017..eecb656255 100644 --- a/tests/unit/timing/strategies/test_agentic_replay.py +++ b/tests/unit/timing/strategies/test_agentic_replay.py @@ -76,6 +76,7 @@ def _make_strategy( run: object | None = None, dataset: DatasetMetadata | None = None, cache_warmup_duration: float | None = None, + cache_warmup_requests_per_lane: int | None = None, progress: MagicMock | None = None, ) -> tuple[ AgenticReplayStrategy, AsyncMock, LoopScheduler | MagicMock, TrajectorySource @@ -87,6 +88,7 @@ def _make_strategy( cfg.phase = phase cfg.concurrency = len(trajectories) cfg.agentic_cache_warmup_duration_sec = cache_warmup_duration + cfg.warmup_requests_per_lane = cache_warmup_requests_per_lane issuer = issuer if issuer is not None else AsyncMock() issuer.replay_gate = MagicMock() issuer.replay_gate.completed_prefixes.return_value = () @@ -284,6 +286,75 @@ async def test_cache_warmup_starts_after_baseline_and_removes_idle_delay(): assert scheduler.schedule_later.call_args.args[0] == 600.0 +@pytest.mark.asyncio +async def test_cache_warmup_request_budget_is_enforced_per_lane(): + trajectories = [ + Trajectory(conversation_id=f"trace_{i}", start_turn_index=0) for i in range(2) + ] + strategy, issuer, _, _ = _make_strategy( + phase=CreditPhase.WARMUP, + trajectories=trajectories, + cache_warmup_requests_per_lane=2, + ) + issuer.set_turn_admission = MagicMock() + + await strategy.setup_phase() + + admission = issuer.set_turn_admission.call_args.args[0] + lane_0 = TurnToSend( + conversation_id="trace_0", + x_correlation_id=trajectories[0].x_correlation_id, + turn_index=0, + num_turns=4, + ) + lane_1 = TurnToSend( + conversation_id="trace_1", + x_correlation_id=trajectories[1].x_correlation_id, + turn_index=0, + num_turns=4, + ) + strategy._correlation_to_lane[lane_0.x_correlation_id] = 0 + strategy._correlation_to_lane[lane_1.x_correlation_id] = 1 + + assert admission(lane_0) is True + assert admission(lane_0) is True + assert admission(lane_0) is False + assert admission(lane_1) is True + assert admission(lane_1) is True + assert admission(lane_1) is False + issuer.replay_gate.pause_releases.assert_called_once_with() + + +@pytest.mark.asyncio +async def test_count_cache_warmup_starts_without_duration_timer(): + trajectory = Trajectory(conversation_id="trace_0", start_turn_index=1) + strategy, issuer, scheduler, _ = _make_strategy( + phase=CreditPhase.WARMUP, + trajectories=[trajectory], + cache_warmup_requests_per_lane=3, + ) + issuer.set_turn_admission = MagicMock() + + await strategy.setup_phase() + await strategy.execute_phase() + baseline = issuer.issue_credit.await_args_list[0].args[0] + await strategy.handle_credit_return( + _make_credit( + conversation_id="trace_0", + x_correlation_id=baseline.x_correlation_id, + turn_index=1, + num_turns=4, + phase=CreditPhase.WARMUP, + ) + ) + + pressure = issuer.issue_credit.await_args_list[1].args[0] + assert pressure.turn_index == 2 + assert pressure.max_tokens_override == 1 + issuer.set_max_tokens_override.assert_called_once_with(1) + scheduler.schedule_later.assert_not_called() + + @pytest.mark.asyncio async def test_cache_warmup_cutoff_stops_issuer_and_persists_next_turn(): trajectory = Trajectory(conversation_id="trace_0", start_turn_index=1) @@ -747,8 +818,10 @@ async def capture(turn): credit_issuer=issuer, lifecycle=lifecycle, ) - # Idle-gap cap of 60s (what the agentx scenario sets). - strategy._phase_offset_cap_ms = 60_000.0 + # The AgentX scenario sets the global system-idle cap, not the per-trace + # timestamp-warp cap. Warmup priming must honor that real configuration. + strategy._phase_offset_cap_ms = None + strategy._system_idle_gap_cap_seconds = 60.0 await strategy.setup_phase() await strategy.execute_phase() diff --git a/tests/unit/timing/strategies/test_agentic_replay_child_continuation.py b/tests/unit/timing/strategies/test_agentic_replay_child_continuation.py index e4c7bb45e2..039fb24905 100644 --- a/tests/unit/timing/strategies/test_agentic_replay_child_continuation.py +++ b/tests/unit/timing/strategies/test_agentic_replay_child_continuation.py @@ -19,9 +19,14 @@ def _make_strategy( branch_orchestrator: MagicMock | None = None, dispatch_result: bool = True, delay_ms: float | None = None, + phase: CreditPhase = CreditPhase.PROFILING, + accelerated_warmup: bool = False, ) -> tuple[AgenticReplayStrategy, MagicMock, MagicMock]: """Build a strategy with only the attributes ``_dispatch_next_turn`` reads.""" strategy = AgenticReplayStrategy.__new__(AgenticReplayStrategy) + strategy.config = MagicMock(phase=phase) + strategy._cache_warmup_duration = None + strategy._cache_warmup_requests_per_lane = 10 if accelerated_warmup else None conversation_source = MagicMock() conversation_source.get_next_turn_metadata.return_value = TurnMetadata( @@ -104,6 +109,25 @@ async def test_child_at_cap_routes_to_on_child_stopped() -> None: orch.on_child_stopped.assert_awaited_once_with("child-xcid") +@pytest.mark.asyncio +async def test_child_at_warmup_quota_is_preserved_for_profiling_handoff() -> None: + """A resumable warmup refusal must not prematurely satisfy the parent join.""" + orch = MagicMock() + orch.on_child_stopped = AsyncMock() + strategy, issuer, _ = _make_strategy( + branch_orchestrator=orch, + dispatch_result=False, + phase=CreditPhase.WARMUP, + accelerated_warmup=True, + ) + + await strategy._dispatch_next_turn(_child_credit()) + + issuer.dispatch_child_turn.assert_awaited_once() + issuer.issue_credit.assert_not_called() + orch.on_child_stopped.assert_not_called() + + @pytest.mark.asyncio async def test_child_at_cap_without_orchestrator_swallows_silently() -> None: """No orchestrator wired: refusal must not raise.""" diff --git a/tests/unit/timing/test_phase_config_agentic_replay.py b/tests/unit/timing/test_phase_config_agentic_replay.py index 6a5b6ef0e0..7baf752b74 100644 --- a/tests/unit/timing/test_phase_config_agentic_replay.py +++ b/tests/unit/timing/test_phase_config_agentic_replay.py @@ -148,6 +148,28 @@ def test_cache_warmup_uses_strategy_controlled_stop() -> None: assert warmup.grace_period_sec == 300.0 +def test_cache_warmup_request_budget_scales_with_concurrency() -> None: + """The deterministic budget becomes an exact global backstop while the strategy enforces the same quota independently on each lane.""" + phase = _PHASE_ADAPTER.validate_python( + { + "name": "profiling", + "type": "concurrency", + "concurrency": 16, + "duration": 900, + "timing_mode": TimingMode.AGENTIC_REPLAY, + "warmup_requests_per_lane": 10, + } + ) + + warmup = _build_agentic_warmup_config(phase) + + assert warmup is not None + assert warmup.total_expected_requests == 160 + assert warmup.agentic_cache_warmup_duration_sec is None + assert warmup.warmup_requests_per_lane == 10 + assert warmup.grace_period_sec == float("inf") + + def test_cache_warmup_grace_uses_short_duration_without_benchmark_grace() -> None: phase = _PHASE_ADAPTER.validate_python( {