feat: add per-lane request budgets to agentic cache warmup

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` `<float>`
 
-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.
+<br/>_Constraints: > 0_
+
+#### `--warmup-requests-per-lane` `<int>`
+
+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.
 <br/>_Constraints: > 0_
 
 #### `--num-warmup-sessions` `<int>`

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.
 

src/aiperf/common/config/loadgen_config.py

@@ -503,14 +503,31 @@ 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",),
             group=Groups.LOAD_GENERATOR,
         ),
     ] = 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

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
 

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
 

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)

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

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,
     )