From 08ec06877e94961b54a0ca0374c168a6bb7fa1bd Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Wed, 29 Jul 2026 11:41:15 -0500 Subject: [PATCH] feat(agentic): add deterministic warmup request budgets Port the per-lane warmup quota and robust warmup-to-profiling handoff onto the AgentX v1.0 lineage. Signed-off-by: Cam Quilici --- docs/cli-options.md | 7 +- docs/tutorials/agentx-mvp.md | 7 + src/aiperf/common/config/loadgen_config.py | 20 ++- src/aiperf/common/config/user_config.py | 12 +- src/aiperf/credit/callback_handler.py | 69 +++++--- src/aiperf/credit/issuer.py | 33 +++- src/aiperf/dataset/loader/weka_synth_buf.py | 17 +- src/aiperf/timing/config.py | 27 ++- src/aiperf/timing/phase/runner.py | 41 +++-- .../timing/strategies/agentic_replay.py | 166 +++++++++++++++--- tests/unit/credit/test_callback_handler.py | 34 ++++ tests/unit/credit/test_issuer.py | 43 +++++ .../loader/test_weka_synth_buf_turn_delta.py | 30 ++++ ...est_runner_agentic_replay_warmup_target.py | 16 ++ .../timing/strategies/test_agentic_replay.py | 71 ++++++++ .../test_phase_config_agentic_replay.py | 13 ++ tools/ergonomics_baseline.json | 5 + tools/ruff_baseline.json | 15 ++ 18 files changed, 543 insertions(+), 83 deletions(-) diff --git a/docs/cli-options.md b/docs/cli-options.md index d606e1bb48..ed73238595 100644 --- a/docs/cli-options.md +++ b/docs/cli-options.md @@ -934,7 +934,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_ #### `--num-warmup-sessions` `` diff --git a/docs/tutorials/agentx-mvp.md b/docs/tutorials/agentx-mvp.md index 160826a6fa..aca05163d8 100644 --- a/docs/tutorials/agentx-mvp.md +++ b/docs/tutorials/agentx-mvp.md @@ -266,6 +266,13 @@ limited to one output token. When the 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 issues +exactly that many warmup requests, including snapshot priming. For example, +16 lanes with 10 requests per lane produce 160 warmup requests. + +The duration and per-lane request options are mutually exclusive. + These requests remain part of `warmup`, so they are excluded from exported request metrics. The option is disabled by default. diff --git a/src/aiperf/common/config/loadgen_config.py b/src/aiperf/common/config/loadgen_config.py index 7784c47014..1406d5971b 100644 --- a/src/aiperf/common/config/loadgen_config.py +++ b/src/aiperf/common/config/loadgen_config.py @@ -503,7 +503,8 @@ def parse_concurrency_list( "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",), @@ -511,6 +512,22 @@ def parse_concurrency_list( ), ] = 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.LOAD_GENERATOR, + ), + ] = None + warmup_num_sessions: Annotated[ int | None, Field( @@ -917,6 +934,7 @@ def disable_warmup(self) -> None: self.warmup_request_count = None self.warmup_duration = None self.agentic_cache_warmup_duration = None + self.warmup_requests_per_lane = None self.warmup_num_sessions = None # Warmup load parameters diff --git a/src/aiperf/common/config/user_config.py b/src/aiperf/common/config/user_config.py index 810685d902..dcdd92e010 100644 --- a/src/aiperf/common/config/user_config.py +++ b/src/aiperf/common/config/user_config.py @@ -1215,12 +1215,18 @@ def validate_warmup_grace_period(self) -> Self: @model_validator(mode="after") def validate_agentic_cache_warmup(self) -> Self: """Restrict accelerated cache warmup to agentic replay.""" - if self.loadgen.agentic_cache_warmup_duration is None: + has_duration = self.loadgen.agentic_cache_warmup_duration is not None + has_request_budget = self.loadgen.warmup_requests_per_lane is not None + if not has_duration and not has_request_budget: return self + if has_duration and has_request_budget: + raise ValueError( + "--warmup-requests-per-lane and " + "--agentic-cache-warmup-duration are mutually exclusive." + ) if self.timing_mode != TimingMode.AGENTIC_REPLAY: raise ValueError( - "--agentic-cache-warmup-duration requires the agentic_replay " - "timing mode." + "agentic cache warmup requires the agentic_replay timing mode." ) return self diff --git a/src/aiperf/credit/callback_handler.py b/src/aiperf/credit/callback_handler.py index 695f0f89b0..4e0d376877 100644 --- a/src/aiperf/credit/callback_handler.py +++ b/src/aiperf/credit/callback_handler.py @@ -280,8 +280,10 @@ async def _handle_warmup_failure( return self._warmup_abort_triggered = True _logger.warning( - lambda: f"Terminal warmup failure for trace {credit.conversation_id}; " - f"aborting run early (broadcasting ProfileCancelCommand)." + lambda: ( + f"Terminal warmup failure for trace {credit.conversation_id}; " + f"aborting run early (broadcasting ProfileCancelCommand)." + ) ) try: await self._on_warmup_abort() @@ -316,16 +318,20 @@ async def on_credit_return( handler = self._phase_handlers.get(phase) if not handler: _logger.debug( - lambda: f"Credit return for unregistered phase {phase}, " - f"credit_id={credit.id}, worker={worker_id}" + lambda: ( + f"Credit return for unregistered phase {phase}, " + f"credit_id={credit.id}, worker={worker_id}" + ) ) return # Late arrivals after phase complete are logged but don't affect counts if handler.lifecycle.is_complete: _logger.warning( - lambda: f"Credit return after phase {phase} complete, " - f"credit_id={credit.id}, worker={worker_id}" + lambda: ( + f"Credit return after phase {phase} complete, " + f"credit_id={credit.id}, worker={worker_id}" + ) ) return @@ -399,9 +405,11 @@ async def on_credit_return( ) except Exception as exc: # noqa: BLE001 _logger.warning( - lambda exc=exc: f"BranchOrchestrator child-completion " - f"hook failed for x_correlation_id=" - f"{credit.x_correlation_id}: {exc}" + lambda exc=exc: ( + f"BranchOrchestrator child-completion " + f"hook failed for x_correlation_id=" + f"{credit.x_correlation_id}: {exc}" + ) ) observe_credit_return = getattr(handler.strategy, "observe_credit_return", None) @@ -450,23 +458,26 @@ 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. + 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 ) @@ -477,9 +488,11 @@ async def on_credit_return( ) except Exception as exc: # noqa: BLE001 _logger.warning( - lambda exc=exc: f"BranchOrchestrator on_child_stopped " - f"hook failed for x_correlation_id=" - f"{credit.x_correlation_id}: {exc}" + lambda exc=exc: ( + f"BranchOrchestrator on_child_stopped " + f"hook failed for x_correlation_id=" + f"{credit.x_correlation_id}: {exc}" + ) ) # WARMUP terminal-failure accumulation + live early-abort (agentic replay). @@ -587,7 +600,9 @@ def _release_slots_for_return( in_flight = handler.progress.in_flight_sessions if in_flight > 0: _logger.debug( - lambda: f"Releasing {in_flight} in-flight session slots for phase {phase}" + lambda: ( + f"Releasing {in_flight} in-flight session slots for phase {phase}" + ) ) for _ in range(in_flight): concurrency.release_session_slot(phase) @@ -611,8 +626,10 @@ async def on_first_token(self, first_token: FirstToken) -> None: if not handler: _logger.debug( - lambda: f"TTFT for unregistered phase {phase}, " - f"credit_id={first_token.credit_id}" + lambda: ( + f"TTFT for unregistered phase {phase}, " + f"credit_id={first_token.credit_id}" + ) ) return diff --git a/src/aiperf/credit/issuer.py b/src/aiperf/credit/issuer.py index 875295d27c..e8e877931b 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 @@ -113,6 +114,15 @@ def __init__( self._issuing_stopped = False self.replay_gate = ReplayIssueGate(replay_barrier) self._max_tokens_override: int | None = None + self._turn_admission: Callable[[TurnToSend], bool] | None = None + + 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.""" + return self._turn_admission is None or self._turn_admission(turn) def set_max_tokens_override(self, max_tokens: int | None) -> None: """Override generation length for every subsequently issued credit.""" @@ -268,7 +278,6 @@ async def _issue_credit_ready(self, turn: TurnToSend) -> bool: ) if not acquired: return False - self._open_session_tree(turn) # Prefill concurrency: one slot per request, released when TTFT arrives. # Limits concurrent prompt processing which is the GPU-intensive phase. @@ -281,6 +290,15 @@ async def _issue_credit_ready(self, turn: TurnToSend) -> bool: self._concurrency_manager.release_session_slot(self._phase) return False + if not self._is_turn_admitted(turn): + self._concurrency_manager.release_prefill_slot(self._phase) + if needs_session_slot: + self._concurrency_manager.release_session_slot(self._phase) + return False + + if needs_session_slot: + self._open_session_tree(turn) + # Slots acquired - proceed with credit issuance return await self._issue_credit_internal(turn) @@ -322,7 +340,6 @@ async def try_issue_credit(self, turn: TurnToSend) -> bool | None: ) if not acquired: return None # No slot - credit not issued - self._open_session_tree(turn) acquired = self._concurrency_manager.try_acquire_prefill_slot( self._phase, can_proceed_fn @@ -333,6 +350,15 @@ async def try_issue_credit(self, turn: TurnToSend) -> bool | None: self._concurrency_manager.release_session_slot(self._phase) return None # No slot - credit not issued + if not self._is_turn_admitted(turn): + self._concurrency_manager.release_prefill_slot(self._phase) + if needs_session_slot: + self._concurrency_manager.release_session_slot(self._phase) + return False + + if needs_session_slot: + self._open_session_tree(turn) + return await self._issue_credit_internal(turn) async def _issue_credit_internal(self, turn: TurnToSend) -> bool: @@ -442,6 +468,9 @@ async def _dispatch_child_turn_ready(self, turn: TurnToSend) -> bool: self._phase, can_proceed_fn ): return False + if not self._is_turn_admitted(turn): + self._concurrency_manager.release_prefill_slot(self._phase) + 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 f9afe0cc74..9ed97dd729 100644 --- a/src/aiperf/dataset/loader/weka_synth_buf.py +++ b/src/aiperf/dataset/loader/weka_synth_buf.py @@ -442,7 +442,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. @@ -452,6 +452,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``. @@ -460,9 +464,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/timing/config.py b/src/aiperf/timing/config.py index d9d33b7e3c..87c2f27818 100644 --- a/src/aiperf/timing/config.py +++ b/src/aiperf/timing/config.py @@ -221,6 +221,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.", + ) def _agentic_warmup_grace_period(loadgen: LoadGeneratorConfig) -> float | None: @@ -241,12 +248,25 @@ def _agentic_warmup_grace_period(loadgen: LoadGeneratorConfig) -> float | None: def _build_agentic_warmup_config(loadgen: LoadGeneratorConfig) -> CreditPhaseConfig: cache_warmup_duration = loadgen.agentic_cache_warmup_duration + requests_per_lane = loadgen.warmup_requests_per_lane + cache_warmup_request_cap = ( + loadgen.concurrency * requests_per_lane + if loadgen.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 = loadgen.concurrency return CreditPhaseConfig( phase=CreditPhase.WARMUP, timing_mode=TimingMode.AGENTIC_REPLAY, - total_expected_requests=( - None if cache_warmup_duration is not None else loadgen.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=loadgen.concurrency, @@ -257,6 +277,7 @@ def _build_agentic_warmup_config(loadgen: LoadGeneratorConfig) -> CreditPhaseCon seamless=False, grace_period_sec=_agentic_warmup_grace_period(loadgen), 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 76202785bb..77aa059a80 100644 --- a/src/aiperf/timing/phase/runner.py +++ b/src/aiperf/timing/phase/runner.py @@ -116,9 +116,12 @@ def __init__( self._conversation_source = conversation_source self._user_config = user_config self._session_tree_registry = session_tree_registry - cache_warmup_enabled = isinstance( - getattr(config, "agentic_cache_warmup_duration_sec", None), - int | float, + 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. @@ -148,16 +151,22 @@ def __init__( config.timing_mode == TimingMode.AGENTIC_REPLAY and config.phase == CreditPhase.WARMUP and isinstance(conversation_source, TrajectorySource) - and not cache_warmup_enabled ): - trajectory_count = conversation_source.warmup_credit_count - if ( - trajectory_count > 0 - and trajectory_count != config.total_expected_requests - ): + requests_per_lane = getattr(config, "warmup_requests_per_lane", None) + if cache_warmup_enabled and requests_per_lane is not None: + request_cap = requests_per_lane * len(conversation_source.trajectories) self._config = self._config.model_copy( - update={"total_expected_requests": trajectory_count} + update={"total_expected_requests": request_cap} ) + elif not cache_warmup_enabled: + trajectory_count = conversation_source.warmup_credit_count + if ( + trajectory_count > 0 + and trajectory_count != config.total_expected_requests + ): + self._config = self._config.model_copy( + update={"total_expected_requests": trajectory_count} + ) self._phase_publisher = phase_publisher self._credit_router = credit_router self._concurrency_manager = concurrency_manager @@ -848,11 +857,13 @@ def _release_tree_slots(self) -> None: return released = self._session_tree_registry.release_all(self._config.phase) self.info( - lambda: f"Session-tree slots for phase {self._config.phase}: " - f"peak_open={self._session_tree_registry.peak_open} " - f"(target concurrency {self._config.concurrency}); " - f"released {released} still-open at teardown; " - f"late_events={self._session_tree_registry.late_events}" + lambda: ( + f"Session-tree slots for phase {self._config.phase}: " + f"peak_open={self._session_tree_registry.peak_open} " + f"(target concurrency {self._config.concurrency}); " + f"released {released} still-open at teardown; " + f"late_events={self._session_tree_registry.late_events}" + ) ) def _release_stuck_slots(self) -> None: diff --git a/src/aiperf/timing/strategies/agentic_replay.py b/src/aiperf/timing/strategies/agentic_replay.py index 90bc0064c1..8cef7a4006 100644 --- a/src/aiperf/timing/strategies/agentic_replay.py +++ b/src/aiperf/timing/strategies/agentic_replay.py @@ -174,6 +174,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 @@ -261,21 +272,25 @@ 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 (PROFILING).""" + """True when per-tree session-slot accounting is engaged.""" 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: @@ -379,6 +394,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) @@ -403,6 +423,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: @@ -467,9 +530,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 @@ -573,6 +646,31 @@ 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." + ) + + if not prepared: + if self._cache_warmup_enabled: + await self._start_accelerated_warmup() + return + self.info("WARMUP execute: no requests precede t*; nothing to warm") + self.credit_issuer.mark_sending_complete() + return + # Pass 2: dispatch. if not spread: self.info( @@ -613,7 +711,7 @@ async def _finish_initial_warmup_dispatch( self, prepared: list[tuple[TurnToSend, float | None]] ) -> None: """Finish burst warmup or enter cache pressure when no priming exists.""" - if self._cache_warmup_duration is None: + if not self._cache_warmup_enabled: if not self.lifecycle.is_sending_complete: self.lifecycle.mark_sending_complete() return @@ -622,30 +720,37 @@ async def _finish_initial_warmup_dispatch( async def _start_accelerated_warmup_if_empty( self, prepared: list[tuple[TurnToSend, float | None]] ) -> None: - if not prepared and self._cache_warmup_duration is not None: + if not prepared and self._cache_warmup_enabled: await self._start_accelerated_warmup() 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 for trajectory in self.conversation_source.trajectories: self._seed_trajectory_replay_prefix(trajectory) self.credit_issuer.replay_gate.activate() self.credit_issuer.set_max_tokens_override(_WARMUP_MAX_TOKENS) + 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) @@ -902,6 +1007,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 ) @@ -1297,7 +1405,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) @@ -1337,16 +1445,16 @@ async def _dispatch_next_turn(self, credit: Credit) -> None: async def _issue_child_continuation_or_drain(self, turn: TurnToSend) -> None: """Single child-issuance chokepoint (dataflow-inspired IssuanceAuthority). - Dispatch a DAG child continuation via ``dispatch_child_turn`` (which - returns a clean True-iff-on-wire, avoiding ``issue_credit``'s overloaded - False) and, on ANY refusal (e.g. the ``--request-count`` wire cap), notify - ``BranchOrchestrator.on_child_stopped`` so the parent's join drains - deterministically rather than deadlocking on a child whose remaining - turns will never be issued. Centralizing here means no dispatch site can - "forget" to drain on refusal. + Dispatch a DAG child continuation via ``dispatch_child_turn``. Terminal + refusals notify the orchestrator so the parent's join drains. Warmup + quota refusals remain resumable across the profiling handoff. """ 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 _spawn_from_recycle_or_id( diff --git a/tests/unit/credit/test_callback_handler.py b/tests/unit/credit/test_callback_handler.py index cba2efb97e..a0fba2bf97 100644 --- a/tests/unit/credit/test_callback_handler.py +++ b/tests/unit/credit/test_callback_handler.py @@ -872,6 +872,40 @@ async def test_cache_warmup_handoff_allows_paused_dag_work( assert mock_progress.all_credits_returned_event.is_set() mock_orchestrator.has_pending_branch_work.assert_called_once_with() + async def test_cache_warmup_handoff_preserves_non_final_child( + self, + dag_handler, + mock_progress, + mock_lifecycle, + mock_stop_checker, + mock_strategy, + mock_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_orchestrator.has_pending_branch_work = MagicMock(return_value=True) + mock_orchestrator.intercept = AsyncMock(return_value=False) + mock_orchestrator.on_child_stopped = AsyncMock() + dag_handler.register_phase( + phase=CreditPhase.WARMUP, + progress=mock_progress, + lifecycle=mock_lifecycle, + stop_checker=mock_stop_checker, + strategy=mock_strategy, + ) + + credit = make_dag_credit( + phase=CreditPhase.WARMUP, + turn_index=1, + num_turns=7, + agent_depth=1, + ) + await dag_handler.on_credit_return("worker-1", make_credit_return(credit)) + + mock_strategy.handle_credit_return.assert_awaited_once_with(credit, error=None) + mock_orchestrator.on_child_stopped.assert_not_awaited() + async def test_child_leaf_reached_called_on_child_final_turn( self, registered_dag_handler, mock_orchestrator ): diff --git a/tests/unit/credit/test_issuer.py b/tests/unit/credit/test_issuer.py index d822e98651..d5868ac8a0 100644 --- a/tests/unit/credit/test_issuer.py +++ b/tests/unit/credit/test_issuer.py @@ -220,6 +220,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 after acquisition releases both root-owned 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 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_synth_buf_turn_delta.py b/tests/unit/dataset/loader/test_weka_synth_buf_turn_delta.py index 616cfde38a..fdcebfaf4f 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 @@ -142,6 +142,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/timing/phase/test_runner_agentic_replay_warmup_target.py b/tests/unit/timing/phase/test_runner_agentic_replay_warmup_target.py index 785bac0a72..2b3f009637 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 @@ -140,6 +140,22 @@ async def test_concurrency_below_pool_size_uses_concurrency(self) -> None: runner = _make_runner(config, src) assert runner._config.total_expected_requests == 4 + async def test_cache_warmup_target_uses_actual_lane_count(self) -> None: + """The request budget uses the actual wrap-filled trajectory lane count.""" + src = MagicMock(spec=TrajectorySource) + src.dataset_metadata = _make_dataset_metadata({"a": 2, "b": 2, "c": 2}) + 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_short_traces_skipped_below_concurrency_wrap_fills(self) -> None: """Pool of 6 with one 1-turn trace, concurrency=8: wrap-fill to 8 lanes. diff --git a/tests/unit/timing/strategies/test_agentic_replay.py b/tests/unit/timing/strategies/test_agentic_replay.py index c7011108c6..8af8c2fe81 100644 --- a/tests/unit/timing/strategies/test_agentic_replay.py +++ b/tests/unit/timing/strategies/test_agentic_replay.py @@ -99,6 +99,7 @@ def _make_strategy( user_config: 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, MagicMock, TrajectorySource]: src = _build_real_trajectory_source( @@ -108,6 +109,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 = () @@ -309,6 +311,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) diff --git a/tests/unit/timing/test_phase_config_agentic_replay.py b/tests/unit/timing/test_phase_config_agentic_replay.py index 5b8dc1e573..f2dfbc5fca 100644 --- a/tests/unit/timing/test_phase_config_agentic_replay.py +++ b/tests/unit/timing/test_phase_config_agentic_replay.py @@ -22,6 +22,7 @@ def _ar_user_config( cfg.loadgen.warmup_request_count = None cfg.loadgen.warmup_duration = None cfg.loadgen.agentic_cache_warmup_duration = None + cfg.loadgen.warmup_requests_per_lane = None cfg.loadgen.warmup_num_sessions = None cfg.loadgen.warmup_concurrency = None cfg.loadgen.warmup_prefill_concurrency = None @@ -106,6 +107,18 @@ 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: + cfg = _ar_user_config(concurrency=16) + cfg.loadgen.warmup_requests_per_lane = 10 + + warmup = _build_warmup_config(cfg) + + 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 + + def test_cache_warmup_grace_uses_short_duration_without_benchmark_grace() -> None: cfg = _ar_user_config(concurrency=10, benchmark_grace_period=None) cfg.loadgen.agentic_cache_warmup_duration = 2.0 diff --git a/tools/ergonomics_baseline.json b/tools/ergonomics_baseline.json index 0f35d5d5b1..3851644e0b 100644 --- a/tools/ergonomics_baseline.json +++ b/tools/ergonomics_baseline.json @@ -106,6 +106,11 @@ "src/aiperf/credit/callback_handler.py", "" ], + [ + "file-size", + "src/aiperf/credit/issuer.py", + "" + ], [ "file-size", "src/aiperf/credit/sticky_router.py", diff --git a/tools/ruff_baseline.json b/tools/ruff_baseline.json index 322f0dbfcc..ccbf82c50f 100644 --- a/tools/ruff_baseline.json +++ b/tools/ruff_baseline.json @@ -1572,6 +1572,11 @@ "src/aiperf/timing/phase/runner.py", "PhaseRunner.run" ], + [ + "PLR0912", + "src/aiperf/timing/strategies/agentic_replay.py", + "AgenticReplayStrategy._execute_warmup" + ], [ "PLR0912", "src/aiperf/timing/strategies/request_rate.py", @@ -1797,6 +1802,16 @@ "src/aiperf/timing/phase/runner.py", "PhaseRunner.run" ], + [ + "PLR0915", + "src/aiperf/timing/strategies/agentic_replay.py", + "AgenticReplayStrategy.__init__" + ], + [ + "PLR0915", + "src/aiperf/timing/strategies/agentic_replay.py", + "AgenticReplayStrategy._execute_warmup" + ], [ "PLR0915", "src/aiperf/transports/aiohttp_client.py",