From 91d7322840a749e7a10a4ffd334e1c2bbe779fed Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Thu, 30 Jul 2026 16:32:32 -0500 Subject: [PATCH 1/9] fix(agentx): preserve warmup state across handoff Signed-off-by: Cam Quilici --- src/aiperf/credit/dispatch.py | 44 +++++ src/aiperf/credit/issuer.py | 52 +++--- src/aiperf/timing/branch_orchestrator.py | 56 +++---- src/aiperf/timing/replay_dependencies.py | 27 ++- .../timing/strategies/agentic_replay.py | 53 ++++-- src/aiperf/timing/strategies/request_rate.py | 10 +- tests/integration/test_weka_flat_split_e2e.py | 77 +++++++++ tests/unit/credit/test_issuer.py | 28 +++- .../timing/strategies/test_agentic_replay.py | 155 ++++++++++++++++-- ...st_branch_orchestrator_warmup_intercept.py | 95 ++++++++++- .../timing/test_replay_barrier_coordinator.py | 16 ++ 11 files changed, 515 insertions(+), 98 deletions(-) create mode 100644 src/aiperf/credit/dispatch.py diff --git a/src/aiperf/credit/dispatch.py b/src/aiperf/credit/dispatch.py new file mode 100644 index 0000000000..3d734f8f4b --- /dev/null +++ b/src/aiperf/credit/dispatch.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Typed outcomes for credit admission and DAG-child dispatch.""" + +from __future__ import annotations + +from enum import StrEnum + + +class TurnAdmission(StrEnum): + """Final admission decision after concurrency slots are acquired.""" + + ADMIT = "admit" + DEFER = "defer" + REJECT = "reject" + + @classmethod + def normalize(cls, result: TurnAdmission | bool) -> TurnAdmission: + """Convert legacy Boolean admission callbacks to an explicit decision.""" + if isinstance(result, cls): + return result + return cls.ADMIT if result else cls.REJECT + + +class ChildDispatchResult(StrEnum): + """Lifecycle disposition of a DAG-child dispatch attempt.""" + + ISSUED = "issued" + DEFERRED = "deferred" + REJECTED = "rejected" + + @classmethod + def normalize( + cls, result: ChildDispatchResult | bool | None + ) -> ChildDispatchResult: + """Convert legacy issuer results without conflating deferral and refusal.""" + if isinstance(result, cls): + return result + return cls.ISSUED if result is True else cls.REJECTED + + @property + def preserves_tracking(self) -> bool: + """Whether the orchestrator must retain this child's dependency edges.""" + return self is not ChildDispatchResult.REJECTED diff --git a/src/aiperf/credit/issuer.py b/src/aiperf/credit/issuer.py index 4a899239c3..e3ba0a2094 100644 --- a/src/aiperf/credit/issuer.py +++ b/src/aiperf/credit/issuer.py @@ -23,6 +23,7 @@ from aiperf.common.aiperf_logger import AIPerfLogger from aiperf.common.enums import CreditPhase from aiperf.common.phase import phase_runtime_key +from aiperf.credit.dispatch import ChildDispatchResult, TurnAdmission from aiperf.credit.structs import Credit, TurnToSend from aiperf.timing.replay_dependencies import ReplayIssueGate from aiperf.timing.url_samplers import URLSelectionStrategyProtocol @@ -126,17 +127,21 @@ def __init__( ) self._issuing_stopped = False self._max_tokens_override: int | None = None - self._turn_admission: Callable[[TurnToSend], bool] | None = None + self._turn_admission: Callable[[TurnToSend], TurnAdmission | bool] | None = None self.replay_gate = ReplayIssueGate(replay_barrier) - def set_turn_admission(self, callback: Callable[[TurnToSend], bool]) -> None: + def set_turn_admission( + self, callback: Callable[[TurnToSend], TurnAdmission | bool] + ) -> None: """Install a synchronous final admission check for every turn.""" self._turn_admission = callback - def _is_turn_admitted(self, turn: TurnToSend) -> bool: + def _turn_admission_result(self, turn: TurnToSend) -> TurnAdmission: """Run the optional final admission check.""" callback = getattr(self, "_turn_admission", None) - return callback is None or callback(turn) + if callback is None: + return TurnAdmission.ADMIT + return TurnAdmission.normalize(callback(turn)) def set_max_tokens_override(self, max_tokens: int | None) -> None: """Override generation length for every subsequently issued credit.""" @@ -353,7 +358,7 @@ 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): + if self._turn_admission_result(turn) is not TurnAdmission.ADMIT: self._concurrency_manager.release_prefill_slot(self._phase_key) if needs_session_slot: self._concurrency_manager.release_session_slot(self._phase_key) @@ -417,7 +422,7 @@ 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): + if self._turn_admission_result(turn) is not TurnAdmission.ADMIT: self._concurrency_manager.release_prefill_slot(self._phase_key) if needs_session_slot: self._concurrency_manager.release_session_slot(self._phase_key) @@ -491,25 +496,22 @@ async def _issue_credit_internal(self, turn: TurnToSend) -> bool: return not is_final_credit - async def dispatch_first_turn(self, sampled_session: SampledSession) -> bool: + async def dispatch_first_turn( + self, sampled_session: SampledSession + ) -> ChildDispatchResult: """Dispatch the first turn of a mid-run DAG child session. Thin wrapper around ``dispatch_child_turn`` that builds the first ``TurnToSend`` from the sampled session. - - Returns True if the credit was sent on the wire (orchestrator - should expect a return), False otherwise (orchestrator should - roll back its tracking via ``BranchOrchestrator.on_child_stopped`` - / per-child rollback). """ return await self.dispatch_child_turn(sampled_session.build_first_turn()) - async def dispatch_child_turn(self, turn: TurnToSend) -> bool: + async def dispatch_child_turn(self, turn: TurnToSend) -> ChildDispatchResult: """Dispatch a DAG child turn (first or continuation). - Returns True if the credit was sent on the wire (caller should - expect a return), False otherwise (caller should roll back its - tracking via ``BranchOrchestrator.on_child_stopped``). + ``DEFERRED`` means the turn is retained by a replay barrier or phase + handoff, so callers must preserve its orchestrator bookkeeping. + ``REJECTED`` is terminal and permits callers to drain that bookkeeping. We avoid the overloaded ``issue_credit`` / ``try_issue_credit`` False (which conflates "gate refused, not issued" with "issued, @@ -528,28 +530,32 @@ async def dispatch_child_turn(self, turn: TurnToSend) -> bool: turn, lambda: self._dispatch_child_turn_ready(turn), child_refusal_cleanup=True, + retained_result=ChildDispatchResult.DEFERRED, ) - async def _dispatch_child_turn_ready(self, turn: TurnToSend) -> bool: + async def _dispatch_child_turn_ready(self, turn: TurnToSend) -> ChildDispatchResult: """Dispatch a child after its recorded predecessor frontier completes.""" if self._issuing_stopped: - return False + return ChildDispatchResult.REJECTED can_proceed_fn = self._stop_checker.can_send_child_turn if not can_proceed_fn(): - return False + return ChildDispatchResult.REJECTED # Children inherit the parent's session slot; wait for prefill # capacity so temporary saturation does not delete sibling branches. if not await self._concurrency_manager.acquire_prefill_slot( self._phase_key, can_proceed_fn ): - return False - if not self._is_turn_admitted(turn): + return ChildDispatchResult.REJECTED + admission = self._turn_admission_result(turn) + if admission is not TurnAdmission.ADMIT: self._concurrency_manager.release_prefill_slot(self._phase_key) - return False + if admission is TurnAdmission.DEFER: + return ChildDispatchResult.DEFERRED + return ChildDispatchResult.REJECTED if turn.counts_toward_phase_target: turn = _struct_replace(turn, counts_toward_phase_target=False) await self._issue_credit_internal(turn) - return True + return ChildDispatchResult.ISSUED async def dispatch_join_turn(self, pending: PendingBranchJoin) -> bool: """Dispatch a parent's gated turn after all its children complete. diff --git a/src/aiperf/timing/branch_orchestrator.py b/src/aiperf/timing/branch_orchestrator.py index ab9049d154..d6fbd63e12 100644 --- a/src/aiperf/timing/branch_orchestrator.py +++ b/src/aiperf/timing/branch_orchestrator.py @@ -104,6 +104,7 @@ ) from aiperf.common.environment import Environment from aiperf.common.models.branch_stats import BranchStats +from aiperf.credit.dispatch import ChildDispatchResult __all__ = [ "BranchOrchestrator", @@ -542,15 +543,14 @@ async def dispatch_pre_session_branches(self) -> None: ) self.stats.children_errored += 1 continue - issued = await self._issuer.dispatch_first_turn(child_session) - if issued: + result = ChildDispatchResult.normalize( + await self._issuer.dispatch_first_turn(child_session) + ) + if result.preserves_tracking: self.stats.children_spawned += 1 else: - # ``dispatch_first_turn`` -> ``dispatch_child_turn`` - # only returns False when a stop condition refuses - # the child before or during prefill-slot acquisition. - # Exceptions are caught above. Tally as truncated, - # not errored. + # Terminal refusal before the child reaches the wire. + # Exceptions are caught above; tally this as truncated. self.stats.children_truncated += 1 self._pre_dispatched_branches.add( (conv.conversation_id, branch.branch_id) @@ -974,8 +974,8 @@ async def _spawn_children_and_register_gates( # Dispatch children. A SPAWN child whose recorded first request # starts after the branch spawn dispatches via a delayed background # task at that offset; everything else dispatches immediately. - # try_issue_credit returning False/None rolls back per-child - # bookkeeping (shared between both paths). + # Terminal refusals roll back per-child bookkeeping. Deferred children + # retain it so their joins survive phase handoff. immediate_children: list = [] for child in all_children: offset_ms = dispatch_offset_by_corr.get(child.x_correlation_id, 0.0) @@ -991,7 +991,9 @@ async def _spawn_children_and_register_gates( return_exceptions=True, ) for child, result in zip(immediate_children, results, strict=True): - if result is not True: + if isinstance(result, BaseException) or ( + ChildDispatchResult.normalize(result) is ChildDispatchResult.REJECTED + ): self._rollback_failed_first_turn(child, result, parent_corr) # The parent's NEXT turn (about to be re-evaluated by # _maybe_suspend_parent on return) is the only gate that may need @@ -1004,9 +1006,7 @@ async def _spawn_children_and_register_gates( def _rollback_failed_first_turn(self, child, result, parent_corr: str) -> None: """Undo per-child bookkeeping for a turn-0 dispatch that didn't land. - Shared by the immediate gather path and the delayed-dispatch tasks so - both classify results identically (BaseException -> errored, False -> - truncated, None -> silent no-op). + Shared by the immediate gather path and the delayed-dispatch tasks. """ child_corr = child.x_correlation_id child_mode = self._child_modes.pop(child_corr, None) @@ -1040,13 +1040,12 @@ def _rollback_failed_first_turn(self, child, result, parent_corr: str) -> None: # over len(all_children)); a turn-0 dispatch that never landed must # decrement it too, or the tree's slot would never drain. self._tree_descendant_done(child_corr) - # Three-way classification of non-True dispatch results: + # Classification of terminal dispatch results: # * BaseException -> genuine error (mirror commit 05d02720b # which fixed the analogous bug in # ``dispatch_pre_session_branches``). - # * False -> ``dispatch_child_turn`` stop-condition refusal; - # not an error. - # * None -> issuer suppressed silently; observable no-op. + # * REJECTED (including legacy False/None) -> stop-condition + # refusal; not an error. if isinstance(result, BaseException): logger.error( "dispatch_first_turn failed for child %s", @@ -1054,10 +1053,8 @@ def _rollback_failed_first_turn(self, child, result, parent_corr: str) -> None: exc_info=result, ) self.stats.children_errored += 1 - elif result is False: + elif ChildDispatchResult.normalize(result) is ChildDispatchResult.REJECTED: self.stats.children_truncated += 1 - elif result is None: - pass else: logger.warning( "dispatch_first_turn returned unexpected value %r for child %s", @@ -1239,9 +1236,9 @@ async def _dispatch_first_turn_after_offset( ) -> None: """Delayed-dispatch task body: sleep, then dispatch + settle. - A post-sleep stop-condition refusal (issuer returns False) rolls back - exactly like an immediate refusal. Dispatch and settlement run under - the parent lock, matching the intercept path's locking. + A post-sleep terminal refusal rolls back exactly like an immediate + refusal. Dispatch and settlement run under the parent lock, matching + the intercept path's locking. """ await self._sleep_offset_ms(offset_ms) if self._cleaning_up: @@ -1251,7 +1248,9 @@ async def _dispatch_first_turn_after_offset( result = await self._dispatch_first_turn(child) except Exception as exc: result = exc - if result is not True: + if isinstance(result, BaseException) or ( + ChildDispatchResult.normalize(result) is ChildDispatchResult.REJECTED + ): self._rollback_failed_first_turn(child, result, parent_corr) await self._finalize_failed_dispatches(parent_corr) @@ -1427,15 +1426,14 @@ async def _release_blocked_join(self, pending: PendingBranchJoin) -> None: else: self.stats.joins_suppressed += 1 - async def _dispatch_first_turn(self, child_sampled_session) -> bool: + async def _dispatch_first_turn(self, child_sampled_session) -> ChildDispatchResult: """Dispatch a child's turn-0 via the credit issuer. - Returns True on successful dispatch, False when the issuer declined - because a stop condition fired. Callers use this to roll back - orchestrator bookkeeping when dispatch doesn't actually land a credit. + Legacy Boolean issuer results are normalized at this boundary so the + orchestrator handles every child through one explicit lifecycle type. """ result = await self._issuer.dispatch_first_turn(child_sampled_session) - return bool(result) + return ChildDispatchResult.normalize(result) async def on_child_leaf_reached(self, child_x_correlation_id: str) -> None: """Called when a child session reaches its final turn (or terminates early).""" diff --git a/src/aiperf/timing/replay_dependencies.py b/src/aiperf/timing/replay_dependencies.py index 27817e05bc..f965e0a359 100644 --- a/src/aiperf/timing/replay_dependencies.py +++ b/src/aiperf/timing/replay_dependencies.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING from aiperf.common.aiperf_logger import AIPerfLogger +from aiperf.credit.dispatch import ChildDispatchResult if TYPE_CHECKING: from aiperf.common.models import DatasetMetadata @@ -137,8 +138,8 @@ class _PendingDispatch: turn: TurnToSend """The turn queued for dispatch once its barrier clears.""" - issue: Callable[[], Awaitable[bool]] - """Coroutine factory that issues the credit; returns True on acceptance.""" + issue: Callable[[], Awaitable[bool | ChildDispatchResult]] + """Coroutine factory that resolves the credit's dispatch disposition.""" on_refused: Callable[[], Awaitable[None]] | None """Optional callback run when the dispatch is refused or cancelled.""" @@ -195,10 +196,11 @@ def pause_releases(self) -> None: async def submit( self, turn: TurnToSend, - issue: Callable[[], Awaitable[bool]], + issue: Callable[[], Awaitable[bool | ChildDispatchResult]], *, on_refused: Callable[[], Awaitable[None]] | None = None, - ) -> bool: + retained_result: bool | ChildDispatchResult = True, + ) -> bool | ChildDispatchResult: """Issue now when ready, otherwise retain one deferred dispatch.""" if not self._active: return await issue() @@ -216,7 +218,7 @@ async def submit( state.pending[key] = _PendingDispatch( turn=turn, issue=issue, on_refused=on_refused ) - return True + return retained_result def complete(self, credit: Credit) -> None: """Record any terminal request outcome and release newly ready work.""" @@ -344,7 +346,8 @@ async def _dispatch_pending(pending: _PendingDispatch) -> None: "Barrier-released replay dispatch failed for %r", pending.turn ) issued = False - if not issued and pending.on_refused is not None: + rejected = issued is False or issued is ChildDispatchResult.REJECTED + if rejected and pending.on_refused is not None: await pending.on_refused() @@ -374,10 +377,11 @@ def pause_releases(self) -> None: async def submit( self, turn: TurnToSend, - issue: Callable[[], Awaitable[bool]], + issue: Callable[[], Awaitable[bool | ChildDispatchResult]], *, child_refusal_cleanup: bool = False, - ) -> bool: + retained_result: bool | ChildDispatchResult = True, + ) -> bool | ChildDispatchResult: if self._coordinator is None: return await issue() on_refused = None @@ -386,7 +390,12 @@ async def submit( async def on_refused() -> None: await self._child_refused(turn.x_correlation_id) - return await self._coordinator.submit(turn, issue, on_refused=on_refused) + return await self._coordinator.submit( + turn, + issue, + on_refused=on_refused, + retained_result=retained_result, + ) def activate(self) -> None: if self._coordinator is not None: diff --git a/src/aiperf/timing/strategies/agentic_replay.py b/src/aiperf/timing/strategies/agentic_replay.py index 1ec9c27923..638e79ce2a 100644 --- a/src/aiperf/timing/strategies/agentic_replay.py +++ b/src/aiperf/timing/strategies/agentic_replay.py @@ -70,6 +70,7 @@ from aiperf.common.mixins import AIPerfLoggerMixin from aiperf.common.scenario.base import TrajectoryWarmupFailedError from aiperf.common.scenario.context_overflow import is_context_overflow_response +from aiperf.credit.dispatch import ChildDispatchResult, TurnAdmission from aiperf.credit.structs import TurnToSend from aiperf.timing.conversation_source import SampledSession from aiperf.timing.replay_dependencies import ReplayResumeBoundary @@ -190,7 +191,7 @@ def __init__( self._cache_warmup_requests_by_lane: Counter[int] = Counter() self._cache_warmup_request_budget_reached = False self._baseline_warmup_admitted = 0 - self._quota_handoff_starts: dict[str, TurnToSend] = {} + self._quota_handoff_turns: dict[tuple[str, str, int], TurnToSend] = {} self._baseline_warmup_returns: dict[str, Credit] = {} self._baseline_correlations: set[str] = set() self._baseline_warmup_turns: set[tuple[str, int]] = set() @@ -459,7 +460,7 @@ def _cache_warmup_lane(self, turn_or_credit: TurnToSend | Credit) -> int: ) return lane - def _admit_cache_warmup_turn(self, turn: TurnToSend) -> bool: + def _admit_cache_warmup_turn(self, turn: TurnToSend) -> TurnAdmission: """Admit mandatory primers, then enforce the additional per-lane quota. Snapshot reconstruction can require multiple primers on one lane when @@ -476,16 +477,18 @@ def _admit_cache_warmup_turn(self, turn: TurnToSend) -> bool: ) in self._baseline_warmup_turns if is_baseline: self._baseline_warmup_admitted += 1 - return True + return TurnAdmission.ADMIT 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 + state_key = ( + turn.conversation_id, + turn.x_correlation_id, + turn.turn_index, + ) + self._quota_handoff_turns[state_key] = turn + return TurnAdmission.DEFER 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] @@ -501,7 +504,7 @@ def _admit_cache_warmup_turn(self, turn: TurnToSend) -> bool: f"{len(self.conversation_source.trajectories)} lanes; " "draining requests" ) - return True + return TurnAdmission.ADMIT async def execute_phase(self) -> None: """Dispatch initial credits for the phase.""" @@ -760,6 +763,13 @@ async def _start_accelerated_warmup(self) -> None: if self._accelerated_warmup_started: return assert self._cache_warmup_enabled + # Baseline primers return before accelerated replay is active. A parent + # already blocked on a child join may never return again in WARMUP, so + # seed its last nonterminal credit into the profiling handoff now. + for credit in self._baseline_warmup_returns.values(): + if credit.is_final_turn: + continue + self._handoff_credits[credit.x_correlation_id] = credit self._accelerated_warmup_started = True self.credit_issuer.set_max_tokens_override(_WARMUP_MAX_TOKENS) for trajectory in self.conversation_source.trajectories: @@ -904,16 +914,18 @@ async def _handle_accelerated_warmup_return(self, credit: Credit) -> None: async def _issue_child_continuation_or_drain(self, turn: TurnToSend) -> None: """Dispatch a DAG child continuation, draining terminal refusals. - ``dispatch_child_turn`` returns True iff the turn reached the wire; on - 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 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. A deferred accelerated-warmup + turn is 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) + result = ChildDispatchResult.normalize( + await self.credit_issuer.dispatch_child_turn(turn) + ) if ( - not on_wire + result is ChildDispatchResult.REJECTED and self.branch_orchestrator is not None and not self.allows_pending_branch_handoff_after_sending_complete ): @@ -1074,7 +1086,8 @@ 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(): + for turn in self._quota_handoff_turns.values(): + root_correlation_id = turn.effective_root_correlation_id pending_by_root.setdefault(root_correlation_id, ()) pending_by_root[root_correlation_id] += (turn,) pending_turns_getter = getattr( @@ -1479,6 +1492,10 @@ async def _handle_warmup_return(self, credit: Credit) -> None: await self._handle_accelerated_warmup_return(credit) return self._baseline_warmup_returns[credit.x_correlation_id] = credit + if not credit.is_final_turn: + self._handoff_returned_at_ns[credit.x_correlation_id] = ( + time.perf_counter_ns() + ) if ( len(self._baseline_warmup_returns) >= self.conversation_source.warmup_credit_count diff --git a/src/aiperf/timing/strategies/request_rate.py b/src/aiperf/timing/strategies/request_rate.py index d3a3fda0b2..85f2cdce93 100644 --- a/src/aiperf/timing/strategies/request_rate.py +++ b/src/aiperf/timing/strategies/request_rate.py @@ -11,6 +11,7 @@ from aiperf.common.constants import MILLIS_PER_SECOND, NANOS_PER_SECOND from aiperf.common.mixins import AIPerfLoggerMixin from aiperf.common.utils import yield_to_event_loop +from aiperf.credit.dispatch import ChildDispatchResult from aiperf.credit.structs import Credit, TurnToSend from aiperf.plugin import plugins from aiperf.plugin.enums import PluginType @@ -301,11 +302,14 @@ async def _issue_child_continuation_or_release( credit NOT on wire" with "credit issued, was the final one". Calling ``on_child_stopped`` in the latter case prematurely drains a child whose return is still in flight, leading to a deadlock when that - return arrives at an empty join. ``dispatch_child_turn`` returns - True iff the credit was actually sent on the wire. + return arrives at an empty join. ``dispatch_child_turn`` gives the + caller an explicit issued/deferred/rejected disposition. """ + result = ChildDispatchResult.normalize( + await self._credit_issuer.dispatch_child_turn(turn) + ) if ( - not await self._credit_issuer.dispatch_child_turn(turn) + result is ChildDispatchResult.REJECTED and self._branch_orchestrator is not None ): try: diff --git a/tests/integration/test_weka_flat_split_e2e.py b/tests/integration/test_weka_flat_split_e2e.py index 0b866c0446..9250e08f8d 100644 --- a/tests/integration/test_weka_flat_split_e2e.py +++ b/tests/integration/test_weka_flat_split_e2e.py @@ -414,6 +414,45 @@ def _sibling_children_overlap(plays: list[_Play]) -> bool: return False +def _assert_complete_play_order_and_joins(plays: list[_Play]) -> None: + """Assert source turn order and every known synthetic SPAWN_JOIN frontier.""" + join_specs = { + "trace_alpha": {1: ("fa:001",), 2: ("fa:000",)}, + "trace_beta": {1: ("fa:000", "fa:001")}, + "trace_gamma": {1: ("fa:000",)}, + } + checked_trace_ids: set[str] = set() + for play in plays: + if play.trace_id not in FANOUT_EXPECTED or not _is_complete_play( + play, FANOUT_EXPECTED + ): + continue + for records in (play.root, *play.children.values()): + assert [m.turn_index for m in records] == list(range(len(records))) + for previous, current in zip(records, records[1:], strict=False): + assert current.request_start_ns >= previous.request_end_ns, ( + f"{play.trace_id}: turn {current.turn_index} started before " + f"turn {previous.turn_index} completed in " + f"{current.conversation_id}" + ) + children_by_suffix = { + _child_suffix(cid): records for cid, records in play.children.items() + } + for root_turn_index, child_suffixes in join_specs[play.trace_id].items(): + gated = play.root[root_turn_index] + for suffix in child_suffixes: + child_last = children_by_suffix[suffix][-1] + assert gated.request_start_ns >= child_last.request_end_ns, ( + f"{play.trace_id}: root turn {root_turn_index} crossed " + f"SPAWN_JOIN before child {suffix} completed" + ) + checked_trace_ids.add(play.trace_id) + assert checked_trace_ids == set(FANOUT_EXPECTED), ( + "expected a fully ordered, join-valid play for every synthetic trace; " + f"checked {sorted(checked_trace_ids)}" + ) + + async def test_fanout_dir_splits_and_replays_recorded_concurrency( tmp_path: Path, mock_server_factory: MockServerFactory ) -> None: @@ -478,6 +517,44 @@ async def test_fanout_dir_splits_and_replays_recorded_concurrency( assert branch_stats.children_errored == 0, branch_stats +@pytest.mark.parametrize("warmup_requests_per_lane", [1, 10, 100]) +async def test_cache_warmup_handoff_preserves_order_and_spawn_joins( + tmp_path: Path, + mock_server_factory: MockServerFactory, + warmup_requests_per_lane: int, +) -> None: + """Warmup quotas of different sizes hand profiling a complete DAG, never an independently resumed parent.""" + corpus = _write_fanout_corpus(tmp_path / "traces") + async with mock_server_factory(fast=True, workers=4) as server: + result = await _run_weka_profile( + input_dir=corpus, + artifact_dir=tmp_path / "artifacts", + url=server.url, + duration=3.0, + concurrency=3, + extra_args=[ + "--scenario", + "inferencex-agentx-mvp", + "--unsafe-override", + "--ignore-trace-delays", + "--warmup-requests-per-lane", + str(warmup_requests_per_lane), + ], + timeout=240.0, + ) + _assert_success( + result, f"cache warmup handoff ({warmup_requests_per_lane} per lane)" + ) + + plays = _collect_plays(result) + _assert_complete_play_order_and_joins(plays) + branch_stats = result.json.branch_stats + assert branch_stats is not None + assert branch_stats.children_errored == 0, branch_stats + assert branch_stats.parents_suspended >= 1, branch_stats + assert branch_stats.parents_resumed >= 1, branch_stats + + async def test_spawn_join_gates_main_turn_until_worker_chain_completes( tmp_path: Path, mock_server_factory: MockServerFactory ) -> None: diff --git a/tests/unit/credit/test_issuer.py b/tests/unit/credit/test_issuer.py index 3e5de4ae4e..f6ab354e42 100644 --- a/tests/unit/credit/test_issuer.py +++ b/tests/unit/credit/test_issuer.py @@ -12,6 +12,7 @@ import pytest from aiperf.common.enums import CreditPhase +from aiperf.credit.dispatch import ChildDispatchResult, TurnAdmission from aiperf.credit.issuer import CreditIssuer from aiperf.credit.structs import TurnToSend from aiperf.timing.session_tree import SessionTreeRegistry @@ -251,7 +252,32 @@ async def test_child_turn_admission_refusal_releases_prefill_slot( result = await credit_issuer.dispatch_child_turn(child) - assert result is False + assert result is ChildDispatchResult.REJECTED + 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() + + async def test_child_turn_phase_deferral_preserves_distinct_result( + self, credit_issuer, mock_concurrency, mock_progress, mock_router + ): + """A handoff deferral is not collapsed into a terminal child rejection.""" + credit_issuer.set_turn_admission(lambda _turn: TurnAdmission.DEFER) + 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 ChildDispatchResult.DEFERRED mock_concurrency.release_prefill_slot.assert_called_once_with( CreditPhase.PROFILING ) diff --git a/tests/unit/timing/strategies/test_agentic_replay.py b/tests/unit/timing/strategies/test_agentic_replay.py index 10ac070fed..c46431294c 100644 --- a/tests/unit/timing/strategies/test_agentic_replay.py +++ b/tests/unit/timing/strategies/test_agentic_replay.py @@ -21,6 +21,7 @@ TurnMetadata, ) from aiperf.common.scenario.base import TrajectoryWarmupFailedError +from aiperf.credit.dispatch import TurnAdmission from aiperf.credit.structs import Credit, TurnToSend from aiperf.dataset.dataset_samplers import SequentialSampler from aiperf.plugin.enums import DatasetSamplingStrategy @@ -286,6 +287,52 @@ 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_accelerated_warmup_handoff_keeps_nonterminal_baseline_return(): + trajectory = Trajectory(conversation_id="trace_0", start_turn_index=1) + strategy, _, _, _ = _make_strategy( + phase=CreditPhase.WARMUP, + trajectories=[trajectory], + cache_warmup_requests_per_lane=1, + ) + baseline = _make_credit( + conversation_id="trace_0", + x_correlation_id=trajectory.x_correlation_id, + turn_index=1, + num_turns=4, + phase=CreditPhase.WARMUP, + ) + strategy._baseline_warmup_returns[baseline.x_correlation_id] = baseline + strategy._dispatch_accelerated_trajectory = AsyncMock() + + await strategy._start_accelerated_warmup() + + assert strategy._handoff_credits == {baseline.x_correlation_id: baseline} + + +@pytest.mark.asyncio +async def test_accelerated_warmup_handoff_drops_terminal_baseline_return(): + trajectory = Trajectory(conversation_id="trace_0", start_turn_index=3) + strategy, _, _, _ = _make_strategy( + phase=CreditPhase.WARMUP, + trajectories=[trajectory], + cache_warmup_requests_per_lane=1, + ) + baseline = _make_credit( + conversation_id="trace_0", + x_correlation_id=trajectory.x_correlation_id, + turn_index=3, + num_turns=4, + phase=CreditPhase.WARMUP, + ) + strategy._baseline_warmup_returns[baseline.x_correlation_id] = baseline + strategy._dispatch_accelerated_trajectory = AsyncMock() + + await strategy._start_accelerated_warmup() + + assert strategy._handoff_credits == {} + + @pytest.mark.asyncio async def test_cache_warmup_request_budget_is_enforced_per_lane(): trajectories = [ @@ -319,15 +366,95 @@ async def test_cache_warmup_request_budget_is_enforced_per_lane(): strategy.conversation_source.warmup_credit_count ) - 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 + assert admission(lane_0) is TurnAdmission.ADMIT + assert admission(lane_0) is TurnAdmission.ADMIT + assert admission(lane_0) is TurnAdmission.DEFER + assert admission(lane_1) is TurnAdmission.ADMIT + assert admission(lane_1) is TurnAdmission.ADMIT + assert admission(lane_1) is TurnAdmission.DEFER issuer.replay_gate.pause_releases.assert_called_once_with() +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("agent_depth", "turn_index"), + [ + pytest.param(0, 0, id="root-start"), + pytest.param(0, 2, id="root-continuation"), + pytest.param(1, 0, id="child-start"), + pytest.param(1, 2, id="child-continuation"), + ], +) +async def test_cache_warmup_handoff_preserves_every_quota_refused_turn_once( + agent_depth: int, + turn_index: int, +): + 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=1, + ) + issuer.set_turn_admission = MagicMock() + + await strategy.setup_phase() + + admission = issuer.set_turn_admission.call_args.args[0] + quota_consumer = TurnToSend( + conversation_id="quota-consumer", + x_correlation_id="quota-consumer-root", + turn_index=1, + num_turns=4, + ) + is_child = agent_depth > 0 + refused_turn = TurnToSend( + conversation_id="trace_0::child" if is_child else "trace_0", + x_correlation_id="lane-0-child" if is_child else "lane-0-root", + turn_index=turn_index, + num_turns=4, + agent_depth=agent_depth, + parent_correlation_id="lane-0-root" if is_child else None, + root_correlation_id="lane-0-root" if is_child else None, + branch_mode=ConversationBranchMode.SPAWN, + ) + strategy._correlation_to_lane[quota_consumer.x_correlation_id] = 0 + strategy._correlation_to_lane[refused_turn.x_correlation_id] = 0 + strategy._root_to_lane[quota_consumer.effective_root_correlation_id] = 0 + strategy._root_to_lane[refused_turn.effective_root_correlation_id] = 0 + + if turn_index > 0: + strategy._handoff_credits[refused_turn.x_correlation_id] = _make_credit( + conversation_id=refused_turn.conversation_id, + x_correlation_id=refused_turn.x_correlation_id, + turn_index=turn_index - 1, + num_turns=refused_turn.num_turns, + phase=CreditPhase.WARMUP, + agent_depth=refused_turn.agent_depth, + parent_correlation_id=refused_turn.parent_correlation_id, + root_correlation_id=refused_turn.root_correlation_id, + branch_mode=refused_turn.branch_mode, + ) + + assert admission(quota_consumer) is TurnAdmission.ADMIT + assert admission(refused_turn) is TurnAdmission.DEFER + assert admission(refused_turn) is TurnAdmission.DEFER + issuer.replay_gate.pause_releases.assert_not_called() + + states = strategy._build_handoff_states(finalized_at_ns=0) + + assert len(states[0]) == 1 + state = states[0][0] + assert state.conversation_id == refused_turn.conversation_id + assert state.x_correlation_id == refused_turn.x_correlation_id + assert state.next_turn_index == refused_turn.turn_index + assert state.agent_depth == refused_turn.agent_depth + assert state.parent_correlation_id == refused_turn.parent_correlation_id + assert state.root_correlation_id == refused_turn.root_correlation_id + assert state.branch_mode == refused_turn.branch_mode + + @pytest.mark.asyncio async def test_cache_warmup_quota_is_additional_to_mandatory_lane_primers(): trajectories = [ @@ -404,10 +531,10 @@ async def test_cache_warmup_quota_is_additional_to_mandatory_lane_primers(): } ) - assert admission(lane_0_first) is True - assert admission(lane_0_second) is True + assert admission(lane_0_first) is TurnAdmission.ADMIT + assert admission(lane_0_second) is TurnAdmission.ADMIT issuer.replay_gate.pause_releases.assert_not_called() - assert admission(lane_1) is True + assert admission(lane_1) is TurnAdmission.ADMIT issuer.replay_gate.pause_releases.assert_not_called() assert ( admission( @@ -418,7 +545,7 @@ async def test_cache_warmup_quota_is_additional_to_mandatory_lane_primers(): num_turns=4, ) ) - is True + is TurnAdmission.ADMIT ) issuer.replay_gate.pause_releases.assert_not_called() assert ( @@ -430,7 +557,7 @@ async def test_cache_warmup_quota_is_additional_to_mandatory_lane_primers(): num_turns=4, ) ) - is True + is TurnAdmission.ADMIT ) issuer.replay_gate.pause_releases.assert_called_once_with() assert ( @@ -442,7 +569,7 @@ async def test_cache_warmup_quota_is_additional_to_mandatory_lane_primers(): num_turns=4, ) ) - is False + is TurnAdmission.DEFER ) assert ( admission( @@ -455,7 +582,7 @@ async def test_cache_warmup_quota_is_additional_to_mandatory_lane_primers(): agent_depth=1, ) ) - is False + is TurnAdmission.DEFER ) assert ( admission( @@ -466,7 +593,7 @@ async def test_cache_warmup_quota_is_additional_to_mandatory_lane_primers(): num_turns=4, ) ) - is False + is TurnAdmission.DEFER ) diff --git a/tests/unit/timing/test_branch_orchestrator_warmup_intercept.py b/tests/unit/timing/test_branch_orchestrator_warmup_intercept.py index acad648666..7b1b380014 100644 --- a/tests/unit/timing/test_branch_orchestrator_warmup_intercept.py +++ b/tests/unit/timing/test_branch_orchestrator_warmup_intercept.py @@ -8,7 +8,21 @@ import pytest -from aiperf.common.enums import ConversationBranchMode, CreditPhase +from aiperf.common.enums import ( + ConversationBranchMode, + CreditPhase, + PrerequisiteKind, +) +from aiperf.common.models import ( + ConversationBranchInfo, + ConversationMetadata, + DatasetMetadata, + TurnMetadata, + TurnPrerequisite, +) +from aiperf.credit.dispatch import ChildDispatchResult +from aiperf.credit.structs import Credit +from aiperf.plugin.enums import DatasetSamplingStrategy from aiperf.timing.branch_orchestrator import BranchOrchestrator @@ -95,3 +109,82 @@ async def test_profiling_credit_with_same_source_does_process(): assert cs.start_branch_child.call_count == 2 assert issuer.dispatch_first_turn.await_count == 2 assert orch.stats.children_spawned == 2 + + +@pytest.mark.asyncio +async def test_quota_deferred_child_starts_preserve_parent_join_for_handoff(): + """Warmup quota deferral retains the exact parent-to-child join graph.""" + branch_id = "root:spawn" + child_ids = ["child-a", "child-b", "child-c"] + branch = ConversationBranchInfo( + branch_id=branch_id, + child_conversation_ids=child_ids, + mode=ConversationBranchMode.SPAWN, + ) + parent = ConversationMetadata( + conversation_id="root", + turns=[ + TurnMetadata(branch_ids=[branch_id]), + TurnMetadata( + prerequisites=[ + TurnPrerequisite( + kind=PrerequisiteKind.SPAWN_JOIN, + branch_id=branch_id, + ) + ] + ), + ], + branches=[branch], + ) + children = [ + ConversationMetadata( + conversation_id=child_id, + turns=[TurnMetadata()], + is_root=False, + agent_depth=1, + parent_conversation_id="root", + ) + for child_id in child_ids + ] + source = MagicMock() + source.dataset_metadata = DatasetMetadata( + conversations=[parent, *children], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + ) + source.get_metadata.return_value = parent + + def start_child(*, child_conversation_id: str, **_kwargs): + child = MagicMock() + child.conversation_id = child_conversation_id + child.x_correlation_id = f"corr-{child_conversation_id}" + return child + + source.start_branch_child.side_effect = start_child + issuer = MagicMock() + issuer.dispatch_first_turn = AsyncMock(return_value=ChildDispatchResult.DEFERRED) + orchestrator = BranchOrchestrator( + conversation_source=source, + credit_issuer=issuer, + allow_accelerated_warmup=True, + ) + orchestrator.start_accelerated_warmup() + parent_credit = Credit( + id=0, + phase=CreditPhase.WARMUP, + conversation_id="root", + x_correlation_id="root-corr", + turn_index=0, + num_turns=2, + issued_at_ns=0, + branch_mode=ConversationBranchMode.SPAWN, + ) + + assert await orchestrator.intercept(parent_credit) is True + + blocked, memberships = orchestrator.snapshot_annotations() + assert blocked == {"root-corr": 1} + assert memberships == { + f"corr-{child_id}": [(branch_id, 1)] for child_id in child_ids + } + assert orchestrator.stats.children_spawned == 3 + assert orchestrator.stats.children_truncated == 0 diff --git a/tests/unit/timing/test_replay_barrier_coordinator.py b/tests/unit/timing/test_replay_barrier_coordinator.py index 5834653786..4422d4d8f1 100644 --- a/tests/unit/timing/test_replay_barrier_coordinator.py +++ b/tests/unit/timing/test_replay_barrier_coordinator.py @@ -12,6 +12,7 @@ ReplayTurnReference, TurnMetadata, ) +from aiperf.credit.dispatch import ChildDispatchResult from aiperf.credit.structs import Credit, TurnToSend from aiperf.plugin.enums import DatasetSamplingStrategy from aiperf.timing.replay_dependencies import ( @@ -174,6 +175,21 @@ async def test_pending_turns_exposes_deferred_dispatch_for_phase_handoff() -> No assert coordinator.pending_turns_by_root() == {} +@pytest.mark.asyncio +async def test_retained_child_dispatch_reports_deferred_not_rejected() -> None: + coordinator = ReplayBarrierCoordinator(_metadata()) + coordinator.activate() + + result = await coordinator.submit( + _turn("d"), + lambda: _record_issue([], "d"), + retained_result=ChildDispatchResult.DEFERRED, + ) + + assert result is ChildDispatchResult.DEFERRED + assert coordinator.pending_turns("root") == (_turn("d"),) + + @pytest.mark.asyncio async def test_paused_releases_retain_ready_pending_for_phase_handoff() -> None: coordinator = ReplayBarrierCoordinator(_metadata()) From 895d8fa2d0645cb0419e85c9c6063883b201cbb1 Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Fri, 31 Jul 2026 01:32:46 -0500 Subject: [PATCH 2/9] fix(agentx): preserve spawn-join replay timing Signed-off-by: Cam Quilici --- docs/benchmark-modes/dag.md | 2 +- src/aiperf/timing/branch_orchestrator.py | 237 +++++++++++++---- src/aiperf/timing/phase/runner.py | 11 +- .../timing/strategies/agentic_replay.py | 48 ++-- tests/integration/test_weka_flat_split_e2e.py | 102 +++++++- .../timing/strategies/test_agentic_replay.py | 6 +- .../test_branch_orchestrator_delayed.py | 245 ++++++++++++++++++ 7 files changed, 575 insertions(+), 76 deletions(-) diff --git a/docs/benchmark-modes/dag.md b/docs/benchmark-modes/dag.md index 270d5dc23d..591259b148 100644 --- a/docs/benchmark-modes/dag.md +++ b/docs/benchmark-modes/dag.md @@ -193,7 +193,7 @@ SPAWN targets may be referenced from multiple parents — the child conversation ### Join semantics -DAG-style conversations can declare that a turn dispatches only after children from a prior SPAWN branch complete. Gating is declared via a `TurnPrerequisite(kind=SPAWN_JOIN, branch_id=...)` on the *consuming* turn rather than on the spawning branch. The runtime builds a `(conversation_id, branch_id) -> gated_turn_index` index at phase init; when `BranchOrchestrator.intercept()` sees a spawning turn complete, it resolves the gate from the index and suspends the parent until every outstanding child drains. `CreditIssuer.dispatch_join_turn` then issues the parent's gated turn — reusing the parent's already-held session slot (the gated turn has `turn_index > 0`, so session-slot acquisition is naturally skipped). +DAG-style conversations can declare that a turn dispatches only after children from a prior SPAWN branch complete. Gating is declared via a `TurnPrerequisite(kind=SPAWN_JOIN, branch_id=...)` on the *consuming* turn rather than on the spawning branch. The runtime builds a `(conversation_id, branch_id) -> gated_turn_index` index at phase init; when `BranchOrchestrator.intercept()` sees a spawning turn complete, it resolves the gate and suspends the parent. The gated turn dispatches at the later of its normal recorded replay deadline and the completion of every outstanding child, so a fast child cannot erase the parent's authored think time and a slow child does not add that delay twice. `CreditIssuer.dispatch_join_turn` reuses the parent's already-held session slot (the gated turn has `turn_index > 0`, so session-slot acquisition is naturally skipped). For v1, the orchestrator honors these gate shapes: diff --git a/src/aiperf/timing/branch_orchestrator.py b/src/aiperf/timing/branch_orchestrator.py index d6fbd63e12..2a63993e3a 100644 --- a/src/aiperf/timing/branch_orchestrator.py +++ b/src/aiperf/timing/branch_orchestrator.py @@ -94,7 +94,7 @@ import time from collections import defaultdict from dataclasses import dataclass, field -from typing import Any +from typing import TYPE_CHECKING, Any from aiperf.common.enums import ( CacheBustTarget, @@ -106,6 +106,9 @@ from aiperf.common.models.branch_stats import BranchStats from aiperf.credit.dispatch import ChildDispatchResult +if TYPE_CHECKING: + from aiperf.common.loop_scheduler import LoopScheduler + __all__ = [ "BranchOrchestrator", "BranchStats", @@ -195,12 +198,25 @@ class PendingBranchJoin: # cache-bust for that one turn). parent_cache_bust_marker: str | None = None parent_cache_bust_target: CacheBustTarget = CacheBustTarget.NONE + replay_deadline_elapsed: bool = True + """Whether the gated turn's recorded replay deadline has elapsed. + + A blocked join releases at the later of this deadline and its last child + completion. Untimed joins keep the default ``True``. + """ + replay_deadline_armed: bool = False + """Whether a replay-deadline timer has already been registered.""" @property def is_satisfied(self) -> bool: """True when every prereq's expected completions have all arrived.""" return all(s.is_done for s in self.outstanding.values()) + @property + def can_release(self) -> bool: + """True when both child prerequisites and replay timing are ready.""" + return self.is_satisfied and self.replay_deadline_elapsed + @property def total_outstanding(self) -> int: """Total outstanding children across all prereqs (for diagnostics).""" @@ -241,6 +257,7 @@ def __init__( session_tree_registry=None, cache_bust_ledger=None, allow_accelerated_warmup: bool = False, + scheduler: LoopScheduler | None = None, ) -> None: self._cs = conversation_source self._issuer = credit_issuer @@ -260,6 +277,7 @@ def __init__( # the per-tree outstanding count. self._session_tree_registry = session_tree_registry self._allow_accelerated_warmup = allow_accelerated_warmup + self._scheduler = scheduler self._accelerated_warmup_started = False self._handoff_snapshot_taken = False # child x_correlation_id -> its tree's root_correlation_id, so the @@ -644,6 +662,7 @@ def seed_snapshot( states, *, cache_bust_markers: dict[str, str | None] | None = None, + join_release_delays_ms: dict[str, float] | None = None, ) -> None: """Seed join bookkeeping from an agentic replay wall-clock snapshot. @@ -656,50 +675,72 @@ def seed_snapshot( return states_by_corr = {state.x_correlation_id: state for state in states} + join_release_delays_ms = join_release_delays_ms or {} children_by_parent: dict[str, list] = defaultdict(list) for state in states: if state.agent_depth > 0 and state.parent_correlation_id is not None: children_by_parent[state.parent_correlation_id].append(state) for parent_corr, child_states in children_by_parent.items(): - parent_state = states_by_corr.get(parent_corr) - parent_meta = None - if parent_state is not None: - parent_meta = self._cs.get_metadata(parent_state.conversation_id) - - tracked_children = 0 - for child_state in child_states: - self._child_modes[child_state.x_correlation_id] = ( - child_state.branch_mode - ) - self._register_fork_routing(parent_corr, child_state.branch_mode) - entries = self._seeded_join_entries( + self._seed_snapshot_parent( + parent_corr=parent_corr, + child_states=child_states, + parent_state=states_by_corr.get(parent_corr), + cache_bust_markers=cache_bust_markers, + join_release_delay_ms=join_release_delays_ms.get(parent_corr, 0.0), + ) + + def _seed_snapshot_parent( + self, + *, + parent_corr: str, + child_states: list, + parent_state, + cache_bust_markers: dict[str, str | None] | None, + join_release_delay_ms: float, + ) -> None: + """Restore one snapshot parent's children, joins, and tree accounting.""" + parent_meta = ( + self._cs.get_metadata(parent_state.conversation_id) + if parent_state is not None + else None + ) + for child_state in child_states: + self._child_modes[child_state.x_correlation_id] = child_state.branch_mode + self._register_fork_routing(parent_corr, child_state.branch_mode) + self._child_to_join[child_state.x_correlation_id] = ( + self._seeded_join_entries( parent_corr=parent_corr, parent_state=parent_state, parent_meta=parent_meta, child_state=child_state, cache_bust_markers=cache_bust_markers, ) + ) + self._child_root[child_state.x_correlation_id] = ( + child_state.root_correlation_id or parent_corr + ) - self._child_to_join[child_state.x_correlation_id] = entries - self._child_root[child_state.x_correlation_id] = ( - child_state.root_correlation_id or parent_corr - ) - tracked_children += 1 - - if tracked_children: - # Per-parent drain accounting stays keyed on the direct parent - # (``has_pending_branch_work`` / intercept drain). Session-tree - # accounting keys on the depth-0 root — same as live spawn. - self._descendant_counts[parent_corr] = ( - self._descendant_counts.get(parent_corr, 0) + tracked_children - ) - roots: dict[str, int] = defaultdict(int) - for child_state in child_states: - roots[self._child_root[child_state.x_correlation_id]] += 1 - for tree_root, n in roots.items(): - self._register_tree_descendants(tree_root, n) - self.stats.children_spawned += tracked_children + if parent_state is not None and parent_state.waiting_on_children: + pending = self._active_joins.get(parent_corr) + if pending is not None: + self._arm_join_replay_deadline(pending, join_release_delay_ms) + + tracked_children = len(child_states) + if not tracked_children: + return + # Per-parent drain accounting stays keyed on the direct parent + # (``has_pending_branch_work`` / intercept drain). Session-tree + # accounting keys on the depth-0 root — same as live spawn. + self._descendant_counts[parent_corr] = ( + self._descendant_counts.get(parent_corr, 0) + tracked_children + ) + roots: dict[str, int] = defaultdict(int) + for child_state in child_states: + roots[self._child_root[child_state.x_correlation_id]] += 1 + for tree_root, count in roots.items(): + self._register_tree_descendants(tree_root, count) + self.stats.children_spawned += tracked_children def _ensure_seeded_join( self, @@ -1112,7 +1153,7 @@ async def _finalize_failed_dispatches( if ( active is not None and active.gated_turn_index == next_turn_index - and not active.is_satisfied + and not active.can_release ): parent_will_suspend = True @@ -1134,7 +1175,7 @@ async def _finalize_failed_dispatches( # gate with no child-leaf decrement to fire it -> the suspended parent # deadlocks until drain-timeout. Pop and dispatch it here on the same # machinery so the parent resumes. - if active is not None and active.is_satisfied: + if active is not None and active.can_release: self._active_joins.pop(parent_corr, None) drained_gates.append(active) # If no successful children AND no gated turns, release the @@ -1150,14 +1191,8 @@ async def _finalize_failed_dispatches( # (does not race / double-decrement registered children). if self._sticky_router is not None: self._sticky_router.evict_unclaimed_sticky(parent_corr) - if ( - not any_child_tracked_for_parent(self._child_to_join, parent_corr) - and not self._future_joins.get(parent_corr) - and parent_corr in self._descendant_counts - and self._descendant_counts[parent_corr] <= 0 - ): - self._release_slot(parent_corr) - del self._descendant_counts[parent_corr] + if not any_child_tracked_for_parent(self._child_to_join, parent_corr): + self._release_parent_slot_if_drained(parent_corr) # Dispatch each drained gate's gated turn immediately. The gate was # satisfied with zero outstanding children (every child rolled back), # so no child-leaf decrement will ever fire it; without this the @@ -1349,7 +1384,7 @@ def _maybe_suspend_parent(self, credit) -> bool: if ( active is not None and active.gated_turn_index == next_idx - and not active.is_satisfied + and not active.can_release ): return True @@ -1367,8 +1402,99 @@ def _maybe_suspend_parent(self, credit) -> bool: # would otherwise double-count in cleanup diagnostics. self._pop_future_join(parent_corr, next_idx) self.stats.parents_suspended += 1 + self._arm_join_replay_deadline( + future, + self._gated_turn_delay_ms(credit.conversation_id, next_idx), + ) return True + def _gated_turn_delay_ms(self, conversation_id: str, gated_idx: int) -> float: + """Return the gated turn's end-to-start delay from dataset metadata.""" + try: + metadata = self._cs.get_metadata(conversation_id) + except (AttributeError, KeyError): + return 0.0 + if gated_idx <= 0 or gated_idx >= len(metadata.turns): + return 0.0 + + gated = metadata.turns[gated_idx] + delay_ms = _as_timestamp_ms(getattr(gated, "delay_ms", None)) + if delay_ms is not None: + return max(0.0, delay_ms) + + previous = metadata.turns[gated_idx - 1] + previous_ts_ms = _as_timestamp_ms(getattr(previous, "timestamp_ms", None)) + gated_ts_ms = _as_timestamp_ms(getattr(gated, "timestamp_ms", None)) + if previous_ts_ms is None or gated_ts_ms is None: + return 0.0 + previous_api_ms = _as_timestamp_ms(getattr(previous, "api_time_ms", None)) + return max(0.0, gated_ts_ms - previous_ts_ms - (previous_api_ms or 0.0)) + + def _arm_join_replay_deadline( + self, pending: PendingBranchJoin, delay_ms: float + ) -> None: + """Arm the independent replay-time side of a blocked parent join.""" + if pending.replay_deadline_armed: + return + # Accelerated warmup deliberately removes recorded replay delays for + # every request while retaining dependency ordering. A join has the + # same two independent readiness conditions as profiling -- replay + # time and child completion -- but its replay-time condition is + # immediately ready in zero-idle warmup. Child completion is still + # mandatory, so this cannot release a parent ahead of its subagents. + if self._accelerated_warmup_started: + delay_ms = 0.0 + if delay_ms <= 0.0 or self._scheduler is None: + pending.replay_deadline_elapsed = True + return + pending.replay_deadline_armed = True + pending.replay_deadline_elapsed = False + self._scheduler.schedule_later( + delay_ms / 1000.0, + self._on_join_replay_deadline( + pending.parent_x_correlation_id, + pending.gated_turn_index, + ), + ) + + async def _on_join_replay_deadline( + self, parent_corr: str, gated_idx: int | None + ) -> None: + """Release a blocked parent if its children finished before its timer.""" + if self._cleaning_up or gated_idx is None: + return + pending = self._active_joins.get(parent_corr) + if pending is None or pending.gated_turn_index != gated_idx: + return + pending.replay_deadline_elapsed = True + if not pending.is_satisfied: + return + self._active_joins.pop(parent_corr, None) + self._release_parent_slot_if_drained(parent_corr) + await self._release_blocked_join(pending) + self._notify_drain() + + async def expire_replay_deadlines(self) -> None: + """Drain active joins after phase scheduling has stopped. + + ``PhaseRunner`` cancels the shared scheduler when the phase stops + sending. Any join timer still pending at that boundary can no longer + become a legal replay dispatch, so mark its timing condition complete + and let the issuer's normal stop condition suppress it once its child + condition is also complete. This preserves the two-condition join + state machine without leaving cancelled timers as phantom DAG work. + """ + releasable: list[PendingBranchJoin] = [] + for parent_corr, pending in list(self._active_joins.items()): + pending.replay_deadline_elapsed = True + if pending.can_release: + self._active_joins.pop(parent_corr, None) + self._release_parent_slot_if_drained(parent_corr) + releasable.append(pending) + for pending in releasable: + await self._release_blocked_join(pending) + self._notify_drain() + async def _satisfy_prerequisite( self, parent_corr: str, @@ -1408,8 +1534,10 @@ async def _satisfy_prerequisite( outstanding.completed.add(child_corr) if not pending.is_satisfied: return None - if pending.is_blocked: + if pending.is_blocked and pending.can_release: return self._active_joins.pop(parent_corr, None) + if pending.is_blocked: + return None # Satisfied before the parent arrived — pop the future entry and # let the parent breeze through when it reaches the turn. self._pop_future_join(parent_corr, gated_idx) @@ -1500,16 +1628,7 @@ async def _handle_child_done( # number of gates satisfied. if parent in self._descendant_counts: self._descendant_counts[parent] -= 1 - # If no active/future joins remain and count reached zero, - # release the slot (mirrors prior behavior for the - # no-join/no-child terminal path). - if ( - self._descendant_counts[parent] <= 0 - and parent not in self._active_joins - and parent not in self._future_joins - ): - self._release_slot(parent) - del self._descendant_counts[parent] + self._release_parent_slot_if_drained(parent) # Decrement the child's TREE outstanding count; the registry releases the # tree's session slot (and recycles its lane) once root + all descendants # have drained. Keyed on the tree root, not the direct parent, so a @@ -1517,6 +1636,16 @@ async def _handle_child_done( self._tree_descendant_done(child_corr) self._notify_drain() # cap-suppressed joins finalize w/o credit return + def _release_parent_slot_if_drained(self, parent_corr: str) -> None: + """Release branch state once no descendants or gated turns remain.""" + if ( + self._descendant_counts.get(parent_corr, 1) <= 0 + and parent_corr not in self._active_joins + and parent_corr not in self._future_joins + ): + self._release_slot(parent_corr) + del self._descendant_counts[parent_corr] + async def on_child_errored(self, child_x_correlation_id: str) -> None: """Called when a child session errors mid-branch. diff --git a/src/aiperf/timing/phase/runner.py b/src/aiperf/timing/phase/runner.py index 40e3c7d0f4..fbe58a325b 100644 --- a/src/aiperf/timing/phase/runner.py +++ b/src/aiperf/timing/phase/runner.py @@ -286,6 +286,7 @@ def _maybe_construct_branch_orchestrator( session_tree_registry=self._session_tree_registry, cache_bust_ledger=getattr(conversation_source, "cache_bust_ledger", None), allow_accelerated_warmup=self._cache_warmup_enabled, + scheduler=self._scheduler, ) def _wire_replay_gate(self) -> None: @@ -994,13 +995,21 @@ async def _wait_for_sending_complete( f"Error waiting for phase {self._config.phase} to send all credits: {e!r}" ) finally: + preserve_branch_handoff = self._preserve_replay_gate_until_finalize( + strategy + ) if not self._lifecycle.is_sending_complete: self._lifecycle.mark_sending_complete(timeout_triggered=timed_out) self._progress.freeze_sent_counts() self._scheduler.cancel_all_pending() + if ( + self._branch_orchestrator is not None + and not preserve_branch_handoff + ): + await self._branch_orchestrator.expire_replay_deadlines() self._progress.all_credits_sent_event.set() - if not self._preserve_replay_gate_until_finalize(strategy): + if not preserve_branch_handoff: await self._credit_issuer.replay_gate.cancel( notify_refused=self._config.phase == CreditPhase.PROFILING ) diff --git a/src/aiperf/timing/strategies/agentic_replay.py b/src/aiperf/timing/strategies/agentic_replay.py index 638e79ce2a..92b9c3be2c 100644 --- a/src/aiperf/timing/strategies/agentic_replay.py +++ b/src/aiperf/timing/strategies/agentic_replay.py @@ -1633,10 +1633,11 @@ async def _dispatch_snapshot_for_profiling( per-lane start offset differs. Gated parents (``waiting_on_children``) are not dispatched here; their - join is seeded with the orchestrator and their gated turn fires when - the blocking children complete during PROFILING. No stream completes - during WARMUP (warmup only ever sends a non-terminal turn), so there - is no warmup-continuation or terminal-root recycle step. + join is seeded with the orchestrator and their gated turn fires at the + later of its recorded replay deadline and the blocking-child completion + frontier. No stream completes during WARMUP (warmup only ever sends a + non-terminal turn), so there is no warmup-continuation or terminal-root + recycle step. """ snapshot = self._get_snapshot(trajectory) for state in snapshot.states: @@ -1647,13 +1648,39 @@ async def _dispatch_snapshot_for_profiling( lane, ) + dispatchable = [s for s in snapshot.states if not s.waiting_on_children] + # Compute one normalized replay timeline for dispatchable streams and + # gated parents alike. A parent's join deadline must use the same + # leading-idle shift / optional burst anchor as every other request in + # its trajectory; only its child gate is additional. + leading_shift_ms = self._leading_idle_shift_ms( + s.next_dispatch_offset_ms for s in dispatchable + ) + offset_by_corr = { + s.x_correlation_id: s.next_dispatch_offset_ms - leading_shift_ms + for s in snapshot.states + } + if self._burst_phase_starts and dispatchable: + t0_offset_ms = min( + offset_by_corr[state.x_correlation_id] for state in dispatchable + ) + else: + t0_offset_ms = 0.0 + if self.branch_orchestrator is not None: self.branch_orchestrator.seed_snapshot( snapshot.states, cache_bust_markers=self._session_marker, + join_release_delays_ms={ + state.x_correlation_id: max( + 0.0, + offset_by_corr[state.x_correlation_id] - t0_offset_ms, + ) + for state in snapshot.states + if state.waiting_on_children + }, ) - dispatchable = [s for s in snapshot.states if not s.waiting_on_children] # A lane needs its own session credit when it dispatches no # slot-acquiring depth-0 root credit at PROFILING start. Two cases: # - rootless: the root's turns are all before t*, so the snapshot has @@ -1725,17 +1752,6 @@ async def _dispatch_snapshot_for_profiling( # Spread (default): t0 = 0, each lane fires at its (shifted) offset from # t*. Burst (--burst-phase-starts): t0 = the lane's min offset, anchoring # the earliest post-t* request at profiling-0. - leading_shift_ms = self._leading_idle_shift_ms( - s.next_dispatch_offset_ms for s in dispatchable - ) - offset_by_corr = { - s.x_correlation_id: s.next_dispatch_offset_ms - leading_shift_ms - for s in dispatchable - } - if self._burst_phase_starts and offset_by_corr: - t0_offset_ms = min(offset_by_corr.values()) - else: - t0_offset_ms = 0.0 for state in dispatchable: session = self.conversation_source.session_for_state(state) turn = self._build_turn_for_session(session, state.next_turn_index) diff --git a/tests/integration/test_weka_flat_split_e2e.py b/tests/integration/test_weka_flat_split_e2e.py index 9250e08f8d..ca36c60e44 100644 --- a/tests/integration/test_weka_flat_split_e2e.py +++ b/tests/integration/test_weka_flat_split_e2e.py @@ -70,11 +70,12 @@ def _req( api_time: float, out: int = 8, model: str = TRACE_MODEL, + request_type: str = "n", ) -> dict: """One hash-aligned top-level normal request (in == len(hash_ids) * 64).""" return { "t": t, - "type": "n", + "type": request_type, "model": model, "in": len(hash_ids) * BLOCK_SIZE, "out": out, @@ -154,6 +155,38 @@ def _write_join_trace(target_dir: Path) -> Path: return target_dir +def _write_issue_1231_trace(target_dir: Path) -> Path: + """Scaled form of issue #1231's root-child-root join timing. + + Source seconds are divided by 1000 so the integration test stays fast: + the previous root ends at 31.042ms, the child ends at 47.658ms, and the + gated root is due at 229.088ms. That preserves both the parent's 198.046ms + end-to-start delay and the source-aligned 181.430ms post-child residual. + """ + _write_trace( + target_dir, + "trace_issue_1231", + [ + _req( + 0.0, + [1, 2, 3], + api_time=0.031042, + out=1, + request_type="s", + ), + _req(0.021181, [1, 2, 50, 51], api_time=0.026477, out=1), + _req( + 0.229088, + [1, 2, 3, 4], + api_time=0.005489, + out=1, + request_type="s", + ), + ], + ) + return target_dir + + def _write_background_trace(target_dir: Path) -> Path: """One trace whose worker chain ends after the last main turn, so the loader emits an ``is_background=True`` branch (no SPAWN_JOIN) the runtime must still send and drain cleanly.""" _write_trace( @@ -414,13 +447,20 @@ def _sibling_children_overlap(plays: list[_Play]) -> bool: return False -def _assert_complete_play_order_and_joins(plays: list[_Play]) -> None: - """Assert source turn order and every known synthetic SPAWN_JOIN frontier.""" +def _assert_complete_play_order_and_joins( + plays: list[_Play], *, assert_timing: bool = False +) -> None: + """Assert source order, SPAWN_JOIN frontiers, and optional replay delays.""" join_specs = { "trace_alpha": {1: ("fa:001",), 2: ("fa:000",)}, "trace_beta": {1: ("fa:000", "fa:001")}, "trace_gamma": {1: ("fa:000",)}, } + residual_delay_ms = { + "trace_alpha": {1: 980.0, 2: 530.0}, + "trace_beta": {1: 900.0}, + "trace_gamma": {1: 900.0}, + } checked_trace_ids: set[str] = set() for play in plays: if play.trace_id not in FANOUT_EXPECTED or not _is_complete_play( @@ -446,6 +486,19 @@ def _assert_complete_play_order_and_joins(plays: list[_Play]) -> None: f"{play.trace_id}: root turn {root_turn_index} crossed " f"SPAWN_JOIN before child {suffix} completed" ) + if assert_timing: + previous = play.root[root_turn_index - 1] + replay_gap_ms = ( + gated.request_start_ns - previous.request_end_ns + ) / 1_000_000.0 + assert replay_gap_ms >= ( + residual_delay_ms[play.trace_id][root_turn_index] - 25.0 + ), ( + f"{play.trace_id}: root turn {root_turn_index} replayed " + f"after {replay_gap_ms:.1f}ms, shorter than its recorded " + f"{residual_delay_ms[play.trace_id][root_turn_index]:.1f}ms " + "end-to-start delay" + ) checked_trace_ids.add(play.trace_id) assert checked_trace_ids == set(FANOUT_EXPECTED), ( "expected a fully ordered, join-valid play for every synthetic trace; " @@ -500,6 +553,7 @@ async def test_fanout_dir_splits_and_replays_recorded_concurrency( f"{sorted(_play_signature(p) for p in complete)}; all plays: " f"{sorted(_play_signature(p) for p in plays)}" ) + _assert_complete_play_order_and_joins(plays, assert_timing=True) # Recorded concurrency reproduced: two sibling worker sessions of one # play overlap in wall-clock (each request takes >= ttft=150ms and the @@ -615,6 +669,48 @@ async def test_spawn_join_gates_main_turn_until_worker_chain_completes( assert branch_stats.children_errored == 0, branch_stats +async def test_issue_1231_shape_preserves_parent_replay_deadline( + tmp_path: Path, mock_server_factory: MockServerFactory +) -> None: + """The exact #1231 timing shape waits on replay time and child completion.""" + corpus = _write_issue_1231_trace(tmp_path / "traces") + async with mock_server_factory(ttft=30.0, itl=0.0, workers=2) as server: + result = await _run_weka_profile( + input_dir=corpus, + artifact_dir=tmp_path / "artifacts", + url=server.url, + duration=1.0, + concurrency=1, + ) + _assert_success(result, "issue #1231 timing shape") + + complete = [ + play + for play in _collect_plays(result) + if play.trace_id == "trace_issue_1231" + and len(play.root) == 2 + and sorted(len(records) for records in play.children.values()) == [1] + ] + assert complete, "no complete replay of the #1231 root-child-root shape" + play = complete[0] + previous, gated = play.root + child_last = next(iter(play.children.values()))[-1] + + assert [record.source_outer_idx for record in play.root] == [0, 2] + assert child_last.source_outer_idx == 1 + assert gated.request_start_ns >= child_last.request_end_ns + parent_gap_ms = (gated.request_start_ns - previous.request_end_ns) / 1_000_000 + assert parent_gap_ms >= 173.0, ( + f"gated parent replayed after {parent_gap_ms:.1f}ms; expected the " + "198.046ms recorded parent delay within scheduler tolerance" + ) + post_child_gap_ms = (gated.request_start_ns - child_last.request_end_ns) / 1_000_000 + assert post_child_gap_ms >= 140.0, ( + f"gated parent resumed only {post_child_gap_ms:.1f}ms after the child; " + "the issue #1231 residual delay was dropped" + ) + + async def test_background_worker_chain_runs_after_root_and_drains_cleanly( tmp_path: Path, aiperf_mock_server: AIPerfMockServer ) -> None: diff --git a/tests/unit/timing/strategies/test_agentic_replay.py b/tests/unit/timing/strategies/test_agentic_replay.py index c46431294c..690eb19b56 100644 --- a/tests/unit/timing/strategies/test_agentic_replay.py +++ b/tests/unit/timing/strategies/test_agentic_replay.py @@ -1281,6 +1281,7 @@ async def test_profiling_snapshot_dispatches_inflight_child_and_seeds_join(): conversation_id="trace_0", x_correlation_id="parent", next_turn_index=2, + next_dispatch_offset_ms=181_430.0, agent_depth=0, waiting_on_children=True, join_target_turn_index=2, @@ -1289,7 +1290,7 @@ async def test_profiling_snapshot_dispatches_inflight_child_and_seeds_join(): conversation_id="trace_0::sa:0", x_correlation_id="child", next_turn_index=1, - next_dispatch_offset_ms=500.0, + next_dispatch_offset_ms=0.0, agent_depth=1, parent_correlation_id="parent", join_target_turn_index=2, @@ -1370,6 +1371,9 @@ async def capture(turn): assert seeded_states[1].x_correlation_id == "child" assert seeded_states[0].waiting_on_children is True assert seeded_states[1].parent_correlation_id == seeded_states[0].x_correlation_id + assert branch_orchestrator.seed_snapshot.call_args.kwargs[ + "join_release_delays_ms" + ] == {"parent": pytest.approx(181_430.0)} @pytest.mark.asyncio diff --git a/tests/unit/timing/test_branch_orchestrator_delayed.py b/tests/unit/timing/test_branch_orchestrator_delayed.py index aee84361a2..9fa631ba45 100644 --- a/tests/unit/timing/test_branch_orchestrator_delayed.py +++ b/tests/unit/timing/test_branch_orchestrator_delayed.py @@ -20,6 +20,7 @@ import pytest from aiperf.common.enums import ConversationBranchMode, CreditPhase, PrerequisiteKind +from aiperf.common.loop_scheduler import LoopScheduler from aiperf.common.models import ( ConversationBranchInfo, ConversationMetadata, @@ -29,6 +30,7 @@ ) from aiperf.plugin.enums import DatasetSamplingStrategy from aiperf.timing.branch_orchestrator import BranchOrchestrator +from aiperf.timing.trajectory_source import ConversationState from tests.unit.timing._shared_helpers import _mk_conv, _mk_source @@ -73,6 +75,249 @@ def _k5_metadata() -> list[ConversationMetadata]: return [root, c0, c1] +def _timed_k1_metadata( + *, + delay_ms: float | None = 198_046.0, +) -> list[ConversationMetadata]: + """Parent and child matching the two-readiness join shape from #1231. + + Times are normalized by subtracting the source root start (110.893s): + the previous parent ends at 31.042s, the child ends at 47.658s, and the + gated parent is due at 229.088s. Thus the parent's ordinary end-to-start + delay is 198.046s and its source-aligned post-child residual is 181.430s. + """ + branch = ConversationBranchInfo( + branch_id="root:0", + child_conversation_ids=["child"], + mode=ConversationBranchMode.SPAWN, + ) + root = _mk_conv( + "root", + [ + TurnMetadata( + timestamp_ms=0.0, + api_time_ms=31_042.0, + branch_ids=[branch.branch_id], + ), + TurnMetadata( + timestamp_ms=229_088.0, + delay_ms=delay_ms, + prerequisites=[ + TurnPrerequisite( + kind=PrerequisiteKind.SPAWN_JOIN, + branch_id=branch.branch_id, + ) + ], + ), + ], + [branch], + ) + child = _mk_conv( + "child", + [TurnMetadata(timestamp_ms=21_181.0, api_time_ms=26_477.0)], + [], + ) + return [root, child] + + +def _timed_join_orchestrator( + *, + delay_ms: float | None = 198_046.0, + scheduler: MagicMock | LoopScheduler | None = None, + allow_accelerated_warmup: bool = False, +) -> tuple[BranchOrchestrator, MagicMock, MagicMock]: + cs = _mk_source(_timed_k1_metadata(delay_ms=delay_ms)) + child_session = MagicMock( + x_correlation_id="corr-child", + metadata=cs.get_metadata("child"), + effective_root_correlation_id="corr-root", + ) + cs.start_branch_child.return_value = child_session + issuer = MagicMock() + issuer.dispatch_first_turn = AsyncMock(return_value=True) + issuer.dispatch_join_turn = AsyncMock(return_value=True) + scheduler = scheduler or MagicMock() + orchestrator = BranchOrchestrator( + conversation_source=cs, + credit_issuer=issuer, + allow_accelerated_warmup=allow_accelerated_warmup, + scheduler=scheduler, + ) + return orchestrator, issuer, scheduler + + +@pytest.mark.asyncio +async def test_join_fast_child_waits_for_recorded_parent_deadline() -> None: + """A fast child cannot erase the parent's recorded residual think time.""" + orch, issuer, scheduler = _timed_join_orchestrator() + + assert await orch.intercept(_mk_credit("root", "corr-root", 0)) is True + delay_s, deadline_coro = scheduler.schedule_later.call_args.args + assert delay_s == pytest.approx(198.046) + metadata = orch._cs.get_metadata("root") + child = orch._cs.get_metadata("child") + source_child_end_ms = child.turns[0].timestamp_ms + child.turns[0].api_time_ms + assert metadata.turns[1].timestamp_ms - source_child_end_ms == pytest.approx( + 181_430.0 + ) + + await orch.on_child_leaf_reached("corr-child") + issuer.dispatch_join_turn.assert_not_awaited() + assert orch._active_joins["corr-root"].is_satisfied + + await deadline_coro + issuer.dispatch_join_turn.assert_awaited_once() + assert orch.stats.parents_resumed == 1 + + +@pytest.mark.asyncio +async def test_join_slow_child_releases_at_child_completion_without_extra_delay() -> ( + None +): + """A child finishing after the deadline releases immediately, not delay later.""" + orch, issuer, scheduler = _timed_join_orchestrator() + + assert await orch.intercept(_mk_credit("root", "corr-root", 0)) is True + deadline_coro = scheduler.schedule_later.call_args.args[1] + await deadline_coro + issuer.dispatch_join_turn.assert_not_awaited() + + await orch.on_child_leaf_reached("corr-child") + issuer.dispatch_join_turn.assert_awaited_once() + assert orch.stats.parents_resumed == 1 + + +@pytest.mark.asyncio +async def test_accelerated_warmup_compresses_join_time_but_keeps_child_gate() -> None: + """Zero-idle warmup removes replay time, never the subagent dependency.""" + orch, issuer, scheduler = _timed_join_orchestrator(allow_accelerated_warmup=True) + orch.start_accelerated_warmup() + + assert await orch.intercept(_mk_credit("root", "corr-root", 0)) is True + scheduler.schedule_later.assert_not_called() + pending = orch._active_joins["corr-root"] + assert pending.replay_deadline_elapsed + assert not pending.is_satisfied + issuer.dispatch_join_turn.assert_not_awaited() + + await orch.on_child_leaf_reached("corr-child") + + issuer.dispatch_join_turn.assert_awaited_once() + assert orch.stats.parents_resumed == 1 + + +@pytest.mark.asyncio +async def test_join_derives_deadline_from_timestamps_when_delay_is_absent() -> None: + """Timestamp/api_time metadata retains the same end-to-start invariant.""" + orch, _, scheduler = _timed_join_orchestrator(delay_ms=None) + + assert await orch.intercept(_mk_credit("root", "corr-root", 0)) is True + delay_s, deadline_coro = scheduler.schedule_later.call_args.args + assert delay_s == pytest.approx(198.046) + deadline_coro.close() + + +@pytest.mark.asyncio +async def test_join_deadline_uses_shared_scheduler_idle_cap() -> None: + """The system-idle cap may advance time but never bypass the child gate.""" + scheduler = LoopScheduler() + orch, issuer, _ = _timed_join_orchestrator( + delay_ms=60_000.0, + scheduler=scheduler, + ) + + assert await orch.intercept(_mk_credit("root", "corr-root", 0)) is True + assert scheduler.pending_count == 1 + assert scheduler.cap_pending_delay(0.0) == pytest.approx(60.0, abs=0.02) + handle_id, (handle, deadline_coro) = next(iter(scheduler._handles.items())) + assert handle.when() - scheduler._loop.time() <= 0.02 + scheduler._handles.pop(handle_id) + handle.cancel() + await deadline_coro + issuer.dispatch_join_turn.assert_not_awaited() + + await orch.on_child_leaf_reached("corr-child") + issuer.dispatch_join_turn.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_phase_stop_drains_join_whose_child_already_finished() -> None: + """Cancelling a future deadline cannot strand a satisfied active join.""" + orch, issuer, _ = _timed_join_orchestrator() + issuer.dispatch_join_turn.return_value = False + + assert await orch.intercept(_mk_credit("root", "corr-root", 0)) is True + await orch.on_child_leaf_reached("corr-child") + assert orch.has_pending_branch_work() + + await orch.expire_replay_deadlines() + + issuer.dispatch_join_turn.assert_awaited_once() + assert orch.stats.joins_suppressed == 1 + assert not orch.has_pending_branch_work() + + +@pytest.mark.asyncio +async def test_phase_stop_then_child_completion_drains_join() -> None: + """A child returning after phase stop still closes the gated parent.""" + orch, issuer, _ = _timed_join_orchestrator() + issuer.dispatch_join_turn.return_value = False + + assert await orch.intercept(_mk_credit("root", "corr-root", 0)) is True + await orch.expire_replay_deadlines() + issuer.dispatch_join_turn.assert_not_awaited() + + await orch.on_child_leaf_reached("corr-child") + + issuer.dispatch_join_turn.assert_awaited_once() + assert orch.stats.joins_suppressed == 1 + assert not orch.has_pending_branch_work() + + +@pytest.mark.asyncio +async def test_snapshot_join_uses_absolute_replay_offset_and_child_gate() -> None: + """A parent already blocked at t* keeps its profiling-relative deadline.""" + cs = _mk_source(_timed_k1_metadata()) + issuer = MagicMock() + issuer.dispatch_join_turn = AsyncMock(return_value=True) + scheduler = MagicMock() + orch = BranchOrchestrator( + conversation_source=cs, + credit_issuer=issuer, + scheduler=scheduler, + ) + parent = ConversationState( + conversation_id="root", + x_correlation_id="corr-root", + next_turn_index=1, + waiting_on_children=True, + join_target_turn_index=1, + ) + child = ConversationState( + conversation_id="child", + x_correlation_id="corr-child", + next_turn_index=0, + agent_depth=1, + parent_correlation_id="corr-root", + root_correlation_id="corr-root", + join_target_turn_index=1, + branch_id="root:0", + branch_mode=ConversationBranchMode.SPAWN, + ) + + orch.seed_snapshot( + (parent, child), + join_release_delays_ms={"corr-root": 181_430.0}, + ) + delay_s, deadline_coro = scheduler.schedule_later.call_args.args + assert delay_s == pytest.approx(181.430) + + await orch.on_child_leaf_reached("corr-child") + issuer.dispatch_join_turn.assert_not_awaited() + await deadline_coro + issuer.dispatch_join_turn.assert_awaited_once() + + @pytest.mark.asyncio async def test_delayed_join_k5_parent_progresses(): """Spawn at T=0, gate at T=5. Parent returns from turns 0..3 without From f8481a4f909196d3de56677b5e5f04f63246180d Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Fri, 31 Jul 2026 02:57:50 -0500 Subject: [PATCH 3/9] fix(agentx): enforce global idle cap across replay waits Signed-off-by: Cam Quilici --- src/aiperf/common/loop_scheduler.py | 27 ++++++- src/aiperf/timing/branch_orchestrator.py | 39 ++++++---- .../timing/strategies/agentic_replay.py | 71 +++++++++++++++++-- tests/integration/test_weka_flat_split_e2e.py | 49 +++++++++++++ tests/unit/common/test_loop_scheduler.py | 21 ++++++ .../timing/strategies/test_agentic_replay.py | 50 ++++++++++++- .../test_branch_orchestrator_delayed.py | 9 ++- ...est_branch_orchestrator_dispatch_offset.py | 42 ++++++++++- 8 files changed, 277 insertions(+), 31 deletions(-) diff --git a/src/aiperf/common/loop_scheduler.py b/src/aiperf/common/loop_scheduler.py index 9d0c999455..860609abe0 100644 --- a/src/aiperf/common/loop_scheduler.py +++ b/src/aiperf/common/loop_scheduler.py @@ -319,6 +319,18 @@ def cap_pending_delay(self, max_delay_sec: float) -> float: return 0.0 now = self._loop.time() + # ``cap_pending_delay(0)`` deliberately moves due work into the + # event loop's ready queue with ``call_soon``. A second cap may run + # before that callback gets a turn. Such an asyncio.Handle has no + # ``when()`` and already represents the minimum possible delay, so + # there is nothing left to advance. + # Use the timer protocol rather than ``isinstance(TimerHandle)``: + # uvloop supplies its own compatible timer-handle implementation. + if any( + not callable(getattr(handle, "when", None)) + for handle, _ in self._handles.values() + ): + return 0.0 earliest = min(handle.when() for handle, _ in self._handles.values()) shift_sec = earliest - now - max_delay_sec if shift_sec <= 0: @@ -330,9 +342,18 @@ def cap_pending_delay(self, max_delay_sec: float) -> float: target = max(now, handle.when() - shift_sec) handle.cancel() handle_container = [None] - replacement = self._loop.call_at( - target, self._safe_callback, handle_container, coro - ) + if target <= now: + # A timer advanced exactly to ``loop.time()`` can be stranded + # in some event-loop implementations until another timed wake + # occurs. Queue it as ready work explicitly; this also states + # the max_delay_sec=0 contract directly. + replacement = self._loop.call_soon( + self._safe_callback, handle_container, coro + ) + else: + replacement = self._loop.call_at( + target, self._safe_callback, handle_container, coro + ) handle_container[0] = replacement self._track_handle_and_return_id(replacement, coro) return shift_sec diff --git a/src/aiperf/timing/branch_orchestrator.py b/src/aiperf/timing/branch_orchestrator.py index 2a63993e3a..c8680d2b78 100644 --- a/src/aiperf/timing/branch_orchestrator.py +++ b/src/aiperf/timing/branch_orchestrator.py @@ -18,12 +18,12 @@ ----------------------------- A SPAWN child whose recorded first request starts after the branch spawn (child turn-0 ``timestamp_ms`` past the branch ``start_timestamp_ms``) -dispatches via a delayed background task at that offset instead of firing -immediately, reproducing the recorded in-subagent timing (e.g. a weka -overflow stream whose first request landed minutes after the subagent -spawned). Join gates and descendant counts are registered before the sleep, -so gated parents wait for sleeping children; ``cleanup()`` cancels pending -sleepers. Datasets without timing (``--ignore-trace-delays``) carry None +dispatches through the shared replay scheduler at that offset instead of +firing immediately, reproducing the recorded in-subagent timing (e.g. a +weka overflow stream whose first request landed minutes after the subagent +spawned). Join gates and descendant counts are registered before scheduling, +so gated parents wait for delayed children; ``cleanup()`` cancels pending +dispatches. Datasets without timing (``--ignore-trace-delays``) carry None timestamps and keep the immediate-dispatch behavior. Sticky-routing locality (FORK mode) @@ -304,8 +304,10 @@ def __init__( self._fail_fast = Environment.DAG.FAIL_FAST self._cleaning_up: bool = False # SPAWN children whose recorded first request starts after the branch - # spawn dispatch via delayed background tasks (see - # _start_delayed_first_turn). cleanup() cancels pending sleepers. + # spawn dispatch through the shared replay scheduler (see + # _start_delayed_first_turn), so the system-idle cap can advance those + # timers uniformly with every other replay timer. The task set is a + # compatibility fallback for isolated callers that provide no scheduler. self._delayed_dispatch_tasks: set[asyncio.Task] = set() # Drain observer: sync callback fired after state mutations that may # drain has_pending_branch_work() to False. Wired by @@ -1013,8 +1015,8 @@ async def _spawn_children_and_register_gates( state.registered = True # Dispatch children. A SPAWN child whose recorded first request - # starts after the branch spawn dispatches via a delayed background - # task at that offset; everything else dispatches immediately. + # starts after the branch spawn dispatches through the shared replay + # scheduler at that offset; everything else dispatches immediately. # Terminal refusals roll back per-child bookkeeping. Deferred children # retain it so their joins survive phase handoff. immediate_children: list = [] @@ -1249,12 +1251,19 @@ def _start_delayed_first_turn( """Schedule a SPAWN child's turn-0 dispatch at its recorded offset. Join gates, descendant counts, and ``_child_to_join`` were registered - at spawn time, so a gated parent keeps waiting while the child sleeps - and ``has_pending_branch_work()`` stays True (the phase drain's - existing timeout backstop bounds the wait if the run ends mid-sleep). - ``cleanup()`` cancels pending sleepers and clears all bookkeeping - itself, so a cancelled task performs no rollback of its own. + at spawn time, so a gated parent keeps waiting while the timer is + pending and ``has_pending_branch_work()`` stays True. Production uses + the shared replay scheduler, making this authored offset subject to the + same uniform system-idle advancement as turn and join timers. Isolated + callers without a scheduler retain the legacy task/sleep fallback. """ + if self._scheduler is not None: + self._scheduler.schedule_later( + offset_ms / 1000.0, + self._dispatch_first_turn_after_offset(child, 0.0, parent_corr), + ) + self.stats.children_delayed += 1 + return task = asyncio.create_task( self._dispatch_first_turn_after_offset(child, offset_ms, parent_corr) ) diff --git a/src/aiperf/timing/strategies/agentic_replay.py b/src/aiperf/timing/strategies/agentic_replay.py index 92b9c3be2c..065e129080 100644 --- a/src/aiperf/timing/strategies/agentic_replay.py +++ b/src/aiperf/timing/strategies/agentic_replay.py @@ -83,6 +83,7 @@ ) _WARMUP_MAX_TOKENS = 1 +_IDLE_WATCHDOG_EPSILON_SECONDS = 1e-6 if TYPE_CHECKING: from aiperf.common.loop_scheduler import LoopScheduler @@ -253,6 +254,7 @@ def __init__( self._system_idle_jump_count = 0 self._system_idle_seconds_skipped = 0.0 self._system_idle_started_at: float | None = None + self._system_idle_watchdog: asyncio.TimerHandle | None = None if self._system_idle_gap_cap_seconds is not None: 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* @@ -525,24 +527,35 @@ def enforce_system_idle_cap(self, in_flight_requests: int | None = None) -> None return # Credit returns start a new idle interval. Scheduler-drain callbacks # recheck the same interval after a scheduled turn is barrier-retained. - follows_request_return = in_flight_requests is not None + in_flight_requests, follows_request_return = ( + self._resolve_in_flight_request_count(in_flight_requests) + ) if in_flight_requests is None: - if self._progress is None: - return - in_flight_requests = self._progress.in_flight + return if in_flight_requests > 0: self._system_idle_started_at = None + self._cancel_system_idle_watchdog() return now = time.monotonic() if follows_request_return or self._system_idle_started_at is None: self._system_idle_started_at = now - if self.scheduler.running_count > 0: - return idle_elapsed = max(0.0, now - self._system_idle_started_at) remaining_idle_budget = max( 0.0, self._system_idle_gap_cap_seconds - idle_elapsed ) + if remaining_idle_budget <= _IDLE_WATCHDOG_EPSILON_SECONDS: + remaining_idle_budget = 0.0 + if remaining_idle_budget > 0: + self._arm_system_idle_watchdog(remaining_idle_budget) + + # A just-fired replay callback may still be transitioning into a real + # request. Give it the remainder of the idle budget to do so, but do + # not let a control-plane coroutine extend true request-idle time past + # the configured cap. The watchdog re-enters here at the deadline and + # progress.in_flight remains the authoritative wire-activity signal. + if self.scheduler.running_count > 0 and remaining_idle_budget > 0: + return shifted = self.scheduler.cap_pending_delay(remaining_idle_budget) if shifted <= 0: @@ -556,6 +569,51 @@ def enforce_system_idle_cap(self, in_flight_requests: int | None = None) -> None ) ) + def _resolve_in_flight_request_count( + self, supplied_count: int | None + ) -> tuple[int | None, bool]: + """Resolve wire activity and whether the observation is a return.""" + if supplied_count is not None: + return supplied_count, True + if self._progress is None: + return None, False + return self._progress.in_flight, False + + def _arm_system_idle_watchdog(self, delay_seconds: float) -> None: + """Guarantee an idle-cap recheck even when no scheduler task drains. + + Request-return and scheduler-drain callbacks remain the fast path. The + independent control timer closes the gap where one of those callbacks + observes a transient running scheduler coroutine and no later state + transition re-enters the cap logic. It deliberately lives outside the + replay scheduler so ``cap_pending_delay`` cannot advance its own guard. + """ + if self._system_idle_watchdog is not None: + return + try: + loop = asyncio.get_running_loop() + except RuntimeError: + # Some direct unit/API callers exercise the synchronous fast path + # without an event loop. They still get the immediate schedule + # shift below; only the independent future recheck is unavailable. + return + self._system_idle_watchdog = loop.call_later( + max(0.0, delay_seconds), + self._run_system_idle_watchdog, + ) + + def _run_system_idle_watchdog(self) -> None: + """Clear the one-shot handle and re-evaluate actual request idleness.""" + self._system_idle_watchdog = None + self.enforce_system_idle_cap() + + def _cancel_system_idle_watchdog(self) -> None: + """Cancel the independent idle guard when wire activity resumes.""" + if self._system_idle_watchdog is None: + return + self._system_idle_watchdog.cancel() + self._system_idle_watchdog = None + def _capped_warmup_lead_ms(self, lead_ms: float) -> float: """Clamp a WARMUP lead (t* - the warmed turn) to the idle-gap cap. @@ -934,6 +992,7 @@ async def _issue_child_continuation_or_drain(self, turn: TurnToSend) -> None: async def finalize_phase(self) -> None: """Persist the drained accelerated-warmup DAG for profiling.""" if self._system_idle_gap_cap_seconds is not None: + self._cancel_system_idle_watchdog() self.scheduler.set_drain_observer(None) self.info( "Global system-idle cap summary: " diff --git a/tests/integration/test_weka_flat_split_e2e.py b/tests/integration/test_weka_flat_split_e2e.py index ca36c60e44..777e670bf5 100644 --- a/tests/integration/test_weka_flat_split_e2e.py +++ b/tests/integration/test_weka_flat_split_e2e.py @@ -711,6 +711,55 @@ async def test_issue_1231_shape_preserves_parent_replay_deadline( ) +async def test_system_idle_cap_advances_issue_1231_join_without_bypassing_child( + tmp_path: Path, mock_server_factory: MockServerFactory +) -> None: + """The global idle guard advances a gated replay deadline, never its join.""" + corpus = _write_issue_1231_trace(tmp_path / "traces") + async with mock_server_factory(ttft=30.0, itl=0.0, workers=2) as server: + result = await _run_weka_profile( + input_dir=corpus, + artifact_dir=tmp_path / "artifacts", + url=server.url, + duration=1.0, + concurrency=1, + extra_args=[ + "--scenario", + "inferencex-agentx-mvp", + "--unsafe-override", + "--system-idle-gap-cap-seconds", + "0.05", + ], + ) + _assert_success(result, "issue #1231 timing shape with global idle cap") + + complete = [ + play + for play in _collect_plays(result) + if play.trace_id == "trace_issue_1231" + and len(play.root) == 2 + and sorted(len(records) for records in play.children.values()) == [1] + ] + assert complete, "no complete replay of the capped #1231 join shape" + play = complete[0] + previous, gated = play.root + child_last = next(iter(play.children.values()))[-1] + + assert gated.request_start_ns >= child_last.request_end_ns, ( + "the system-idle cap bypassed the required child-completion gate" + ) + parent_gap_ms = (gated.request_start_ns - previous.request_end_ns) / 1_000_000 + post_child_gap_ms = (gated.request_start_ns - child_last.request_end_ns) / 1_000_000 + assert 35.0 <= post_child_gap_ms <= 100.0, ( + f"capped join resumed {post_child_gap_ms:.1f}ms after the child; " + "expected the 50ms whole-system idle guard within process jitter" + ) + assert parent_gap_ms < 173.0, ( + f"capped join still waited {parent_gap_ms:.1f}ms from its parent; " + "the replay deadline was not advanced" + ) + + async def test_background_worker_chain_runs_after_root_and_drains_cleanly( tmp_path: Path, aiperf_mock_server: AIPerfMockServer ) -> None: diff --git a/tests/unit/common/test_loop_scheduler.py b/tests/unit/common/test_loop_scheduler.py index 862ee46bfd..89408d17fa 100644 --- a/tests/unit/common/test_loop_scheduler.py +++ b/tests/unit/common/test_loop_scheduler.py @@ -160,6 +160,27 @@ async def noop(): assert scheduler.cap_pending_delay(2.0) == 0.0 scheduler.cancel_all_pending() + async def test_cap_pending_delay_zero_queues_timer_as_ready_work( + self, scheduler: LoopScheduler + ): + """A zero idle budget fires instead of stranding a timer at loop.time().""" + fired = asyncio.Event() + + async def mark_fired(): + fired.set() + + scheduler.schedule_later(60.0, mark_fired()) + assert scheduler.cap_pending_delay(0.0) == pytest.approx(60.0, abs=0.02) + # Re-entrant idle observations before the ready callback runs are a + # no-op, not an AttributeError from treating Handle as TimerHandle. + assert scheduler.cap_pending_delay(0.0) == 0.0 + + await asyncio.wait_for(fired.wait(), timeout=0.1) + await yield_to_scheduler() + + assert scheduler.pending_count == 0 + assert scheduler.running_count == 0 + async def test_cap_pending_delay_rejects_negative_cap( self, scheduler: LoopScheduler ): diff --git a/tests/unit/timing/strategies/test_agentic_replay.py b/tests/unit/timing/strategies/test_agentic_replay.py index 690eb19b56..7d87959a68 100644 --- a/tests/unit/timing/strategies/test_agentic_replay.py +++ b/tests/unit/timing/strategies/test_agentic_replay.py @@ -2361,6 +2361,50 @@ async def test_global_idle_cap_shifts_real_delayed_continuation_to_ten_seconds() scheduler.cancel_all_pending() +@pytest.mark.asyncio +async def test_global_idle_watchdog_bounds_persistent_control_plane_work(): + """A scheduler coroutine cannot mask an otherwise idle server past the cap.""" + trajectories = [Trajectory(conversation_id="trace_0", start_turn_index=0)] + scheduler = LoopScheduler() + progress = MagicMock() + progress.in_flight = 0 + run = _make_run(target=CacheBustTarget.FIRST_TURN_PREFIX) + run.cfg.get_profiling_phases()[0].system_idle_gap_cap_seconds = 0.05 + strategy, _, _, _ = _make_strategy( + phase=CreditPhase.PROFILING, + trajectories=trajectories, + scheduler=scheduler, + run=run, + progress=progress, + ) + await strategy.setup_phase() + + control_release = asyncio.Event() + replay_fired = asyncio.Event() + + async def persistent_control_work() -> None: + await control_release.wait() + + async def replay_dispatch() -> None: + replay_fired.set() + + scheduler.execute_async(persistent_control_work()) + scheduler.schedule_later(1.0, replay_dispatch()) + started_at = time.monotonic() + + strategy.enforce_system_idle_cap(in_flight_requests=0) + await asyncio.wait_for(replay_fired.wait(), timeout=0.2) + + assert time.monotonic() - started_at == pytest.approx(0.05, abs=0.03) + assert strategy._system_idle_jump_count == 1 + assert strategy._system_idle_seconds_skipped == pytest.approx(0.95, abs=0.04) + + control_release.set() + await asyncio.sleep(0) + await strategy.finalize_phase() + scheduler.cancel_all() + + @pytest.mark.asyncio async def test_global_idle_cap_rechecks_after_barrier_defers_near_timer( time_traveler_no_patch_sleep, @@ -2481,7 +2525,11 @@ async def submit_fa3_turn2() -> None: assert issued == [(fa3, 2)] assert issued_at[(fa3, 2)] - idle_started_at == pytest.approx(0.05, abs=0.005) - assert strategy._system_idle_jump_count == 1 + # Depending on event-loop timer granularity, the watchdog may make one + # additional sub-millisecond adjustment at the exact deadline. The + # invariant is the total schedule shift and dispatch time, not the number + # of internal timer rewrites. + assert strategy._system_idle_jump_count >= 1 assert strategy._system_idle_seconds_skipped == pytest.approx(2_171.478, abs=0.03) barrier.complete( _make_credit( diff --git a/tests/unit/timing/test_branch_orchestrator_delayed.py b/tests/unit/timing/test_branch_orchestrator_delayed.py index 9fa631ba45..da715d757b 100644 --- a/tests/unit/timing/test_branch_orchestrator_delayed.py +++ b/tests/unit/timing/test_branch_orchestrator_delayed.py @@ -15,6 +15,7 @@ from __future__ import annotations +import asyncio from unittest.mock import AsyncMock, MagicMock import pytest @@ -229,11 +230,9 @@ async def test_join_deadline_uses_shared_scheduler_idle_cap() -> None: assert await orch.intercept(_mk_credit("root", "corr-root", 0)) is True assert scheduler.pending_count == 1 assert scheduler.cap_pending_delay(0.0) == pytest.approx(60.0, abs=0.02) - handle_id, (handle, deadline_coro) = next(iter(scheduler._handles.items())) - assert handle.when() - scheduler._loop.time() <= 0.02 - scheduler._handles.pop(handle_id) - handle.cancel() - await deadline_coro + for _ in range(3): + await asyncio.sleep(0) + assert scheduler.pending_count == 0 issuer.dispatch_join_turn.assert_not_awaited() await orch.on_child_leaf_reached("corr-child") diff --git a/tests/unit/timing/test_branch_orchestrator_dispatch_offset.py b/tests/unit/timing/test_branch_orchestrator_dispatch_offset.py index 15b5d0f6f1..f3e11fc2fc 100644 --- a/tests/unit/timing/test_branch_orchestrator_dispatch_offset.py +++ b/tests/unit/timing/test_branch_orchestrator_dispatch_offset.py @@ -10,6 +10,7 @@ import pytest from aiperf.common.enums import ConversationBranchMode, PrerequisiteKind +from aiperf.common.loop_scheduler import LoopScheduler from aiperf.common.models import ( ConversationBranchInfo, ConversationMetadata, @@ -56,6 +57,7 @@ def _mk_harness( conversations: list[ConversationMetadata], *, dispatch_result: bool = True, + scheduler: LoopScheduler | None = None, ): """(orchestrator, conversation_source, issuer) with real metadata models.""" by_id = {c.conversation_id: c for c in conversations} @@ -78,10 +80,48 @@ def _fake_child(*, child_conversation_id, **kwargs): issuer.dispatch_first_turn = AsyncMock(return_value=dispatch_result) issuer.dispatch_join_turn = AsyncMock(return_value=True) - orch = BranchOrchestrator(conversation_source=cs, credit_issuer=issuer) + orch = BranchOrchestrator( + conversation_source=cs, + credit_issuer=issuer, + scheduler=scheduler, + ) return orch, cs, issuer +@pytest.mark.asyncio +async def test_offset_child_uses_shared_scheduler_for_idle_cap( + time_traveler_no_patch_sleep, +) -> None: + """Production offsets are replay timers, not untracked asyncio sleeps.""" + parent = _parent_conv([_spawn_branch("b0", ["kid"], start_timestamp_ms=1_000.0)]) + scheduler_errors: list[BaseException] = [] + scheduler = LoopScheduler( + exception_handler=lambda task: scheduler_errors.append(task.exception()) + ) + orch, _, issuer = _mk_harness( + [parent, _child_conv("kid", 5_000.0)], + scheduler=scheduler, + ) + + await orch.intercept(_mk_credit("parent", "P", 0)) + + assert scheduler.pending_count == 1 + assert orch._delayed_dispatch_tasks == set() + issuer.dispatch_first_turn.assert_not_awaited() + assert scheduler.cap_pending_delay(0.01) == pytest.approx(3.99, abs=0.02) + handle, _ = next(iter(scheduler._handles.values())) + assert handle.when() - scheduler._loop.time() < 0.02 + + await asyncio.sleep(0.10) + + assert scheduler_errors == [] + assert scheduler.pending_count == 0 + assert scheduler.running_count == 0 + issuer.dispatch_first_turn.assert_awaited_once() + assert scheduler.pending_count == 0 + scheduler.cancel_all() + + def _mk_credit(conv_id: str, corr_id: str, turn_index: int): return MagicMock( x_correlation_id=corr_id, From deb5421c7ff3cd39f78228f469e23e25233ac6c6 Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Thu, 30 Jul 2026 11:48:08 -0500 Subject: [PATCH 4/9] fix(agentx): enforce per-trajectory runtime idle caps Signed-off-by: Cam Quilici --- .../semianalysis-agentx-faq.md | 18 +- docs/cli-options.md | 4 +- docs/tutorials/agentx-mvp.md | 33 +- docs/tutorials/weka-trace.md | 2 +- src/aiperf/common/loop_scheduler.py | 106 +++- .../common/scenario/inferencex_agentx_mvp.py | 1 - src/aiperf/config/dataset/config.py | 26 +- src/aiperf/config/flags/_converter_dataset.py | 10 +- src/aiperf/config/flags/cli_config.py | 17 +- .../config/schema/aiperf-config.schema.json | 8 +- .../dataset/loader/weka_parallel_convert.py | 72 +-- src/aiperf/dataset/loader/weka_trace.py | 458 ++---------------- src/aiperf/timing/branch_orchestrator.py | 9 + src/aiperf/timing/phase/runner.py | 23 +- src/aiperf/timing/replay_dependencies.py | 105 +++- .../timing/strategies/agentic_replay.py | 124 ++--- .../test_agentic_replay_e2e.py | 2 +- .../test_agentx_trace_idle_gap_cap.py | 142 ++++++ tests/integration/test_weka_flat_split_e2e.py | 90 ++++ .../common/scenario/test_scenario_registry.py | 2 +- .../scenario/test_scenario_validator.py | 9 +- ...scenario_validator_advanced_adversarial.py | 13 +- tests/unit/common/test_loop_scheduler.py | 68 +++ .../test_weka_flat_split_parallel_adv.py | 28 -- .../loader/test_weka_flat_split_serial_adv.py | 91 +--- .../dataset/loader/test_weka_pathological.py | 87 ---- .../loader/test_weka_tool_shaped_messages.py | 1 - .../loader/test_weka_tool_turn_detection.py | 1 - tests/unit/dataset/loader/test_weka_trace.py | 198 -------- .../loader/test_weka_trace_parallel.py | 3 - .../test_agentic_hotpath_invariants.py | 91 +--- .../timing/strategies/test_agentic_replay.py | 113 +++-- tests/unit/timing/test_factories.py | 2 +- .../timing/test_replay_barrier_coordinator.py | 96 ++++ 34 files changed, 887 insertions(+), 1166 deletions(-) create mode 100644 tests/integration/test_agentx_trace_idle_gap_cap.py diff --git a/docs/benchmark-modes/semianalysis-agentx-faq.md b/docs/benchmark-modes/semianalysis-agentx-faq.md index eb6c87c23c..aae226b1fb 100644 --- a/docs/benchmark-modes/semianalysis-agentx-faq.md +++ b/docs/benchmark-modes/semianalysis-agentx-faq.md @@ -66,7 +66,9 @@ transparency. is rejected), and a fresh `--random-seed` is filled in — pin one yourself for reproducible run-to-run comparisons ([§7](#7-reading-the-results-metrics-validity-and-submission-requirements)). Timing mode is locked to agentic-replay. -- **Forbidden — do not pass:** `--ignore-trace-delays`, `--trace-idle-gap-cap-seconds`, +- **Optional timing control:** `--trace-idle-gap-cap-seconds` limits observed + whole-tree idle time independently for each root and all of its descendants. +- **Forbidden — do not pass:** `--ignore-trace-delays`, `--inter-turn-delay-cap-seconds`, `--synthesis-max-isl` (input truncation), and the rate/schedule flags `--request-rate` / `--arrival-pattern` / `--user-centric-rate` / `--fixed-schedule` / `--adaptive-scale`. @@ -282,10 +284,11 @@ the benchmark waits the recorded "think time"/gap before sending the next turn. request-start to request-start. This is always the case for weka trace replay; there is no flag to change it. Replay dispatches each turn only after the previous one completes, so start-to-start deltas would double-count the server's own response time and make every session -drift later turn by turn. Individual trace gaps are not capped. If the entire replay has no active -or ready request, AIPerf uniformly shifts every pending request timer so the next request arrives -within 10 seconds. This avoids benchmarking dead air without changing the recorded spacing inside -one trace while other sessions keep the system busy. +drift later turn by turn. Individual trace gaps are not capped by default; pass +`--trace-idle-gap-cap-seconds S` to advance a trajectory's pending timers when its root and +all descendants have had no request in flight for `S` seconds. This runtime guard does not +rewrite dataset timestamps or bypass spawn/join dependencies. If the entire replay has no active or ready request, AIPerf uniformly +shifts every pending request timer so the next request arrives within 10 seconds. ### Q: What does a single session look like on the wire? A sequence of chat-completions requests that grow turn over turn (the prompt prefix accumulates), @@ -793,7 +796,7 @@ the context-overflow bound is evaluated during the run. | SemiAnalysis corpus | A pinned `*_weka_*` `--public-dataset` alias, or `weka_hf` pinned to `semianalysisai/cc-traces-weka-062126` | A non-pinned dataset, local `weka_trace`, or another `--hf-weka-dataset` under the scenario; omitting a dataset entirely (CLI synthetic default) | Explicit wrong / unpinned loader: refuses to start (or `submission_valid: false` via `--unsafe-override`). Missing/synthetic: always refuses — `--unsafe-override` cannot bypass | | `ignore_eos` | Injected `ignore_eos=true` | Explicit `ignore_eos=false` | Refuses to start (or `submission_valid: false` via `--unsafe-override`) | | Streaming | `--streaming` auto-enabled | Explicit `--no-streaming` | Refuses to start (or `submission_valid: false` via `--unsafe-override`) | -| Honored trace delays | Recorded think-time gaps are preserved; only globally idle replay time is capped at 10s | `--ignore-trace-delays`, `--trace-idle-gap-cap-seconds`, or `--inter-turn-delay-cap-seconds` | Refuses to start (or `submission_valid: false` via `--unsafe-override`) | +| Honored trace delays | Recorded think-time gaps are preserved; `--trace-idle-gap-cap-seconds` may bound observed whole-tree runtime idle without rewriting them | `--ignore-trace-delays` or `--inter-turn-delay-cap-seconds` | Refuses to start (or `submission_valid: false` via `--unsafe-override`) | | No input truncation | Prompts built to the full recorded token counts | Truncating/capping input below the recorded length | Refuses to start (or `submission_valid: false` via `--unsafe-override`) | | Minimum duration | `--benchmark-duration` ≥ 900s (default 1800s) | A duration below 900s | Refuses to start (or `submission_valid: false` via `--unsafe-override`) | | First-turn-prefix cache-busting | Uniqueness marker placed as a first-turn prefix | Disabling cache-busting or relocating the marker | Refuses to start (or `submission_valid: false` via `--unsafe-override`) | @@ -1152,7 +1155,8 @@ The ones a serving engineer is most likely to use: | `--max-context-length` | Drops whole traces whose peak prompt+output exceeds your server's window, so you don't get guaranteed mid-run overflows. Blunter than a `_256k` corpus (removes entire traces, not just over-limit turns). If it would drop everything, the run errors instead of silently emptying the dataset. | | `--trajectory-start-min-ratio` / `--trajectory-start-max-ratio` | The window within each session where t\* is sampled — i.e. how deep into sessions the measured traffic sits. Defaults to the **full session (0.0–1.0)** under the `inferencex-agentx-mvp` scenario; 0.25–0.75 is the generic CLI default when not scenario-locked. | | `--system-idle-gap-cap-seconds` | Caps only globally idle replay time (scenario default 10s); all pending timers shift uniformly, preserving order and relative spacing. | -| `--trace-idle-gap-cap-seconds` / `--inter-turn-delay-cap-seconds` | Per-trace/per-turn timing compression; forbidden by the scenario so recorded think times and cache-TTL intervals remain faithful. | +| `--trace-idle-gap-cap-seconds` | Optional observed whole-tree runtime idle cap; unset by default and allowed by the scenario. | +| `--inter-turn-delay-cap-seconds` | Per-turn timing compression; forbidden by the scenario. | | `--cache-bust` | Where the per-session uniqueness marker goes; the scenario locks `first_turn_prefix`. | ### Validity thresholds (environment variables) diff --git a/docs/cli-options.md b/docs/cli-options.md index ba748a6fe8..5ef5ffe64e 100644 --- a/docs/cli-options.md +++ b/docs/cli-options.md @@ -1092,7 +1092,7 @@ AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches in #### `--trace-idle-gap-cap-seconds` `` -Hard ceiling (seconds) for idle gaps within each individual trace. For Weka trace replay, AIPerf looks at all parent and subagent request submission timestamps within one root trace, compresses long gaps between consecutive request submissions, and derives turn delays from the compressed per-trace timeline. Original request api_time values are not used to decide these idle gaps. When set for Weka, this takes precedence over `--inter-turn-delay-cap-seconds` so individual parent/subagent-line delays are not separately capped. Defaults to None (no per-trace idle-gap compression). +Hard ceiling (seconds) for observed runtime idle time within each AgentX trajectory tree. The idle clock covers initial future work and restarts when the final in-flight request across the root and all descendant streams completes. If no request from that tree reaches the wire before the cap, AIPerf uniformly advances only that tree's pending replay timers. Dataset timestamps are unchanged, and request order, spawn/join dependencies, and replay barriers remain authoritative. Defaults to None (no per-trace runtime idle cap).
_Constraints: ≥ 0.0_ #### `--system-idle-gap-cap-seconds` `` @@ -2627,7 +2627,7 @@ AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches in #### `--trace-idle-gap-cap-seconds` `` -Hard ceiling (seconds) for idle gaps within each individual trace. For Weka trace replay, AIPerf looks at all parent and subagent request submission timestamps within one root trace, compresses long gaps between consecutive request submissions, and derives turn delays from the compressed per-trace timeline. Original request api_time values are not used to decide these idle gaps. When set for Weka, this takes precedence over `--inter-turn-delay-cap-seconds` so individual parent/subagent-line delays are not separately capped. Defaults to None (no per-trace idle-gap compression). +Hard ceiling (seconds) for observed runtime idle time within each AgentX trajectory tree. The idle clock covers initial future work and restarts when the final in-flight request across the root and all descendant streams completes. If no request from that tree reaches the wire before the cap, AIPerf uniformly advances only that tree's pending replay timers. Dataset timestamps are unchanged, and request order, spawn/join dependencies, and replay barriers remain authoritative. Defaults to None (no per-trace runtime idle cap).
_Constraints: ≥ 0.0_ #### `--system-idle-gap-cap-seconds` `` diff --git a/docs/tutorials/agentx-mvp.md b/docs/tutorials/agentx-mvp.md index f8c2987fc4..14d22d2fc9 100644 --- a/docs/tutorials/agentx-mvp.md +++ b/docs/tutorials/agentx-mvp.md @@ -147,9 +147,9 @@ what you must set, what you may tune, and what you shouldn't touch: can omit all of them (AIPerf fills in exactly these values under `--scenario`), but writing them out makes the command document what runs. If you pass one of them with a *conflicting* value, AIPerf errors up front -rather than silently producing an invalid result. The scenario also forbids -`--trace-idle-gap-cap-seconds` and `--inter-turn-delay-cap-seconds`, preserving -the recorded timing within every trace. +rather than silently producing an invalid result. The scenario forbids +`--inter-turn-delay-cap-seconds`; `--trace-idle-gap-cap-seconds` remains an +optional CLI control and is unset by default. **Flags the scenario only defaults** — auto-filled when omitted, but an explicit value is honored *silently*, with no error and no change to the @@ -267,7 +267,8 @@ flag. | `--streaming` is on | Responses stream token-by-token (auto-enabled when unset; explicit `--no-streaming` errors) | The per-token latency metrics (TTFT, ITL) are core to this benchmark and need streaming responses. | | Replay delays are end-to-start (always on, no flag) | Each turn's replay delay is the recorded idle gap from the previous response's *end* to the next request's *start*, not the start-to-start delta. This is unconditional for weka trace replay — there is no toggle. | Replay dispatches each turn after the previous one completes, so start-to-start deltas would double-count the previous request's server time, making every session drift later turn by turn and overstating how many sessions overlap at once. | | `--ignore-trace-delays` is off | Trace-derived delays are preserved | The replay retains the captured agent pacing and KV-cache reuse intervals. | -| Per-trace and per-turn caps are forbidden | `--trace-idle-gap-cap-seconds` and `--inter-turn-delay-cap-seconds` must remain unset | Independently compressing each trace changes think-time and cache-TTL behavior even while other trajectories keep the server busy. | +| Per-trace idle-gap cap is optional | `--trace-idle-gap-cap-seconds` is unset by default and accepts an explicit CLI value | Setting it bounds observed runtime idle across the root and every descendant stream without rewriting dataset timestamps or bypassing spawn/join dependencies. | +| Per-turn cap is forbidden | `--inter-turn-delay-cap-seconds` must remain unset | Independently clamping parent and subagent delays can distort their relative timing. | | `--system-idle-gap-cap-seconds = 10` | When no request is active or ready, all pending replay timers shift uniformly so the next request arrives within 10 seconds | The benchmark avoids measuring long periods with no server work while preserving request order and relative spacing across every pending trajectory. | | `--cache-bust first_turn_prefix` | A unique per-conversation marker is injected at the start of the first user turn for every play (each dispatch of a trace, initial or recycled) | Without this, every time a trace is recycled the server's prefix cache would warm up further on identical content, and steady-state cache-hit rates would inflate the longer the run goes. The marker gives every recycled play a fresh prompt prefix. | | Loader is a pinned Weka with-subagents corpus | The dataset must be a with-subagents `--public-dataset` alias or `--hf-weka-dataset semianalysisai/cc-traces-weka-062126` (`weka_hf`). A local `weka_trace` directory is format-compatible but unpinned — the run refuses unless you pass `--unsafe-override` (which stamps `submission_valid: false`). The [Troubleshooting](#troubleshooting) entry for this lock lists the exact flag forms. | Submission validity requires a known public corpus identity; arbitrary local dirs are not hash-verifiable. | @@ -439,15 +440,21 @@ request metrics. After warmup, the profiling phase opens. Now you're measuring. Each trajectory keeps replaying its conversation from turn `k_i + 1` onward, honoring the original recorded inter-turn gaps as end-to-start delays counted from the -moment the previous turn completes. If every request completes while future -requests remain scheduled more than 10 seconds away, AIPerf shifts all pending -request timers earlier by the same amount. This bounds true system-idle time -without changing request order or the relative spacing among pending -trajectories. A scheduled turn may become retained by a cross-stream replay -barrier instead of reaching the wire; AIPerf re-evaluates the guard after that -scheduler task drains so a nearby blocked turn cannot hide a much later -dispatchable timer. The phase-end log reports how many global jumps occurred -and how many seconds they skipped. +moment the previous turn completes. With +`--trace-idle-gap-cap-seconds S`, a per-trajectory watchdog covers initial +future work and restarts when the last in-flight request across the root and +all descendant streams completes. +If that tree remains idle for `S` seconds, AIPerf uniformly advances only its +pending timers; dataset timestamps, relative timer spacing, request order, and +spawn/join gates remain unchanged. If every request completes while future +requests remain scheduled more than 10 seconds away, AIPerf +shifts all pending request timers earlier by the same amount. This bounds true +system-idle time without changing request order or the relative spacing among +pending trajectories. A scheduled turn may become retained by a cross-stream +replay barrier instead of reaching the wire; AIPerf re-evaluates the guard +after that scheduler task drains so a nearby blocked turn cannot hide a much +later dispatchable timer. The phase-end log reports how many global jumps +occurred and how many seconds they skipped. Concurrency here is **per session tree**: each lane holds one slot for a whole tree — the root conversation plus every subagent worker stream it diff --git a/docs/tutorials/weka-trace.md b/docs/tutorials/weka-trace.md index e189a64d28..1a5cecbb18 100644 --- a/docs/tutorials/weka-trace.md +++ b/docs/tutorials/weka-trace.md @@ -104,7 +104,7 @@ The HuggingFace path and the file-based `--input-file` path produce **byte-ident | `--public-dataset semianalysis_cc_traces_weka` (HuggingFace, legacy/default no-subagents alias) | Legacy/default pinned alias for the same current no-subagents corpus as `semianalysis_cc_traces_weka_no_subagents`. | | `--public-dataset semianalysis_cc_traces_weka_with_subagents` (HuggingFace, with subagents) | Rolling with-subagents alias tracking the current default corpus (062126) for zero-setup runs with full subagent SPAWN/JOIN topology. 393 traces. | -Existing Weka tunables work identically in both paths: `--synthesis-max-isl`, `--synthesis-max-osl`, `--inter-turn-delay-cap-seconds`, `--trace-idle-gap-cap-seconds`, `--ignore-trace-delays`, `--use-think-time-only`, `--cache-bust`, the per-trace model rewriting rules below — same flags, same behavior, same output bytes on the wire. For `--scenario inferencex-agentx-mvp`, the validator accepts a with-subagents `--public-dataset` alias or `weka_hf` constrained to `semianalysisai/cc-traces-weka-062126` for `submission_valid: true`; a local `weka_trace` directory requires `--unsafe-override` (stamps `submission_valid: false`) because corpus identity is not fingerprintable; it does not accept the no-subagents aliases. That scenario forbids per-turn and per-trace delay caps, preserves the original Weka timeline, and uses only the global `--system-idle-gap-cap-seconds` guard. +Existing Weka tunables work identically in both paths: `--synthesis-max-isl`, `--synthesis-max-osl`, `--inter-turn-delay-cap-seconds`, `--trace-idle-gap-cap-seconds`, `--ignore-trace-delays`, `--use-think-time-only`, `--cache-bust`, the per-trace model rewriting rules below — same flags, same behavior, same output bytes on the wire. `--trace-idle-gap-cap-seconds` is a runtime AgentX watchdog and does not rewrite either loader's dataset timestamps. For `--scenario inferencex-agentx-mvp`, the validator accepts a with-subagents `--public-dataset` alias or `weka_hf` constrained to `semianalysisai/cc-traces-weka-062126` for `submission_valid: true`; a local `weka_trace` directory requires `--unsafe-override` (stamps `submission_valid: false`) because corpus identity is not fingerprintable; it does not accept the no-subagents aliases. That scenario permits an explicit `--trace-idle-gap-cap-seconds`, forbids the per-turn `--inter-turn-delay-cap-seconds`, and retains the global `--system-idle-gap-cap-seconds` guard. For newly published compatible HuggingFace Weka trace corpora, pass `--hf-weka-dataset ` — it auto-selects the neutral `weka_hf` loader (passing `--public-dataset weka_hf` explicitly is equivalent): diff --git a/src/aiperf/common/loop_scheduler.py b/src/aiperf/common/loop_scheduler.py index 860609abe0..4a1588efe5 100644 --- a/src/aiperf/common/loop_scheduler.py +++ b/src/aiperf/common/loop_scheduler.py @@ -65,6 +65,10 @@ def __init__( self._handles: dict[ HandleId, tuple[asyncio.TimerHandle | asyncio.Handle, Coroutine] ] = {} + # Optional logical replay-tree ownership for pending timers. Kept in a + # parallel map so the long-standing ``_handles`` tuple contract remains + # stable for tests and diagnostics. + self._handle_groups: dict[HandleId, str] = {} self._exception_handler: Callable[[asyncio.Task], None] | None = ( exception_handler ) @@ -114,13 +118,19 @@ def _safe_callback( mutable list, then populate it after call_later() returns. """ if handle_container[0] is not None: - self._handles.pop(id(handle_container[0]), None) + handle_id = id(handle_container[0]) + self._handles.pop(handle_id, None) + self._handle_groups.pop(handle_id, None) task = self._loop.create_task(coro) self._tasks.add(task) task.add_done_callback(self._done_callback) def _track_handle_and_return_id( - self, handle: asyncio.TimerHandle | asyncio.Handle, coro: Coroutine + self, + handle: asyncio.TimerHandle | asyncio.Handle, + coro: Coroutine, + *, + group_id: str | None = None, ) -> HandleId: """Track a handle and coroutine and return the handle ID. @@ -131,6 +141,8 @@ def _track_handle_and_return_id( """ handle_id = id(handle) self._handles[handle_id] = (handle, coro) + if group_id is not None: + self._handle_groups[handle_id] = group_id return handle_id def execute_async(self, coro: Coroutine) -> asyncio.Task: @@ -146,7 +158,11 @@ def execute_async(self, coro: Coroutine) -> asyncio.Task: return task def schedule_later( - self, delay_sec: float, coro: Coroutine + self, + delay_sec: float, + coro: Coroutine, + *, + group_id: str | None = None, ) -> HandleId | asyncio.Task: """ Schedule coroutine after a relative delay (seconds). @@ -168,9 +184,15 @@ def schedule_later( delay_sec, self._safe_callback, handle_container, coro ) handle_container[0] = handle - return self._track_handle_and_return_id(handle, coro) + return self._track_handle_and_return_id(handle, coro, group_id=group_id) - def schedule_at(self, loop_time: float, coro: Coroutine) -> HandleId | asyncio.Task: + def schedule_at( + self, + loop_time: float, + coro: Coroutine, + *, + group_id: str | None = None, + ) -> HandleId | asyncio.Task: """ Schedule coroutine at an absolute loop.time() timestamp. @@ -188,10 +210,14 @@ def schedule_at(self, loop_time: float, coro: Coroutine) -> HandleId | asyncio.T loop_time, self._safe_callback, handle_container, coro ) handle_container[0] = handle - return self._track_handle_and_return_id(handle, coro) + return self._track_handle_and_return_id(handle, coro, group_id=group_id) def schedule_at_perf_sec( - self, perf_sec: float, coro: Coroutine + self, + perf_sec: float, + coro: Coroutine, + *, + group_id: str | None = None, ) -> HandleId | asyncio.Task: """ Schedule coroutine at an absolute perf_counter seconds timestamp. @@ -217,10 +243,14 @@ def schedule_at_perf_sec( if offset_sec <= 0: return self.execute_async(coro) - return self.schedule_at(cur_loop_time + offset_sec, coro) + return self.schedule_at(cur_loop_time + offset_sec, coro, group_id=group_id) def schedule_at_perf_ns( - self, perf_time_ns: int, coro: Coroutine + self, + perf_time_ns: int, + coro: Coroutine, + *, + group_id: str | None = None, ) -> HandleId | asyncio.Task: """ Schedule coroutine at an absolute perf_counter_ns timestamp. @@ -246,7 +276,7 @@ def schedule_at_perf_ns( if offset_sec <= 0: return self.execute_async(coro) - return self.schedule_at(cur_loop_time + offset_sec, coro) + return self.schedule_at(cur_loop_time + offset_sec, coro, group_id=group_id) def cancel_handle_id(self, handle_id: HandleId) -> bool: """Cancel a scheduled coroutine by its handle ID. @@ -260,6 +290,7 @@ def cancel_handle_id(self, handle_id: HandleId) -> bool: handle_data = self._handles.pop(handle_id, None) if handle_data is None: return False + self._handle_groups.pop(handle_id, None) handle, coro = handle_data handle.cancel() coro.close() @@ -274,6 +305,7 @@ def cancel_all_pending(self) -> None: """ items = list(self._handles.values()) self._handles.clear() + self._handle_groups.clear() # Cancel handles first, then close coroutines for handle, coro in items: @@ -336,9 +368,13 @@ def cap_pending_delay(self, max_delay_sec: float) -> float: if shift_sec <= 0: return 0.0 - items = list(self._handles.values()) + items = [ + (handle, coro, self._handle_groups.get(handle_id)) + for handle_id, (handle, coro) in self._handles.items() + ] self._handles.clear() - for handle, coro in items: + self._handle_groups.clear() + for handle, coro, group_id in items: target = max(now, handle.when() - shift_sec) handle.cancel() handle_container = [None] @@ -355,7 +391,51 @@ def cap_pending_delay(self, max_delay_sec: float) -> float: target, self._safe_callback, handle_container, coro ) handle_container[0] = replacement - self._track_handle_and_return_id(replacement, coro) + self._track_handle_and_return_id(replacement, coro, group_id=group_id) + return shift_sec + + def cap_pending_delay_for_group(self, group_id: str, max_delay_sec: float) -> float: + """Uniformly advance one logical group's pending replay timers. + + Other groups are untouched. Relative ordering and spacing among the + selected group's timers are preserved exactly. + """ + if max_delay_sec < 0: + raise ValueError("max_delay_sec must be non-negative") + selected = [ + (handle_id, handle, coro) + for handle_id, (handle, coro) in self._handles.items() + if self._handle_groups.get(handle_id) == group_id + ] + if not selected: + return 0.0 + + now = self._loop.time() + if any( + not callable(getattr(handle, "when", None)) for _, handle, _ in selected + ): + return 0.0 + earliest = min(handle.when() for _, handle, _ in selected) + shift_sec = earliest - now - max_delay_sec + if shift_sec <= 0: + return 0.0 + + for handle_id, handle, coro in selected: + self._handles.pop(handle_id, None) + self._handle_groups.pop(handle_id, None) + target = max(now, handle.when() - shift_sec) + handle.cancel() + handle_container = [None] + if target <= now: + replacement = self._loop.call_soon( + self._safe_callback, handle_container, coro + ) + else: + replacement = self._loop.call_at( + target, self._safe_callback, handle_container, coro + ) + handle_container[0] = replacement + self._track_handle_and_return_id(replacement, coro, group_id=group_id) return shift_sec @property diff --git a/src/aiperf/common/scenario/inferencex_agentx_mvp.py b/src/aiperf/common/scenario/inferencex_agentx_mvp.py index d5e3511e1e..601fb834cf 100644 --- a/src/aiperf/common/scenario/inferencex_agentx_mvp.py +++ b/src/aiperf/common/scenario/inferencex_agentx_mvp.py @@ -34,7 +34,6 @@ default_trajectory_start_min_ratio=0.0, default_trajectory_start_max_ratio=1.0, system_idle_gap_cap_seconds=10.0, - forbid_trace_idle_gap_cap=True, forbid_inter_turn_delay_cap=True, require_cache_bust=CacheBustTarget.FIRST_TURN_PREFIX, ) diff --git a/src/aiperf/config/dataset/config.py b/src/aiperf/config/dataset/config.py index fcbfd83132..7263fcb10a 100644 --- a/src/aiperf/config/dataset/config.py +++ b/src/aiperf/config/dataset/config.py @@ -510,15 +510,14 @@ class FileDataset(BaseConfig): Field( default=None, ge=0.0, - description="Hard ceiling (seconds) for idle gaps within each individual trace. " - "For Weka trace replay, AIPerf looks at all parent and subagent request " - "submission timestamps within one root trace, compresses long gaps between " - "consecutive request submissions, and derives turn delays from the " - "compressed per-trace timeline. Original request api_time values are not " - "used to decide these idle gaps. When set for Weka, this takes precedence over " - "`--inter-turn-delay-cap-seconds` so individual parent/subagent-line " - "delays are not separately capped. Defaults to None (no per-trace " - "idle-gap compression).", + description="Hard ceiling (seconds) for observed runtime idle time within " + "each AgentX trajectory tree. The idle clock covers initial future work and " + "restarts when the final in-flight request across the root and all descendant " + "streams completes. If no request from that tree reaches the wire before the " + "cap, AIPerf uniformly advances " + "only that tree's pending replay timers. Dataset timestamps are unchanged, " + "and request order, spawn/join dependencies, and replay barriers remain " + "authoritative. Defaults to None (no per-trace runtime idle cap).", ), ] @@ -794,11 +793,10 @@ class PublicDataset(BaseConfig): Field( default=None, ge=0.0, - description="Hard ceiling (seconds) for idle gaps within each " - "individual trace for HF-backed Weka replay (mirror of " - "``FileDataset.trace_idle_gap_cap_seconds``). When set, takes " - "precedence over ``inter_turn_delay_cap_seconds``. Defaults to " - "None (no per-trace idle-gap compression).", + description="Hard ceiling (seconds) for observed runtime idle time within " + "each AgentX trajectory tree for HF-backed Weka replay (mirror of " + "``FileDataset.trace_idle_gap_cap_seconds``). Dataset timestamps remain " + "unchanged. Defaults to None (no per-trace runtime idle cap).", ), ] diff --git a/src/aiperf/config/flags/_converter_dataset.py b/src/aiperf/config/flags/_converter_dataset.py index d907f63040..914e6c8b85 100644 --- a/src/aiperf/config/flags/_converter_dataset.py +++ b/src/aiperf/config/flags/_converter_dataset.py @@ -808,11 +808,11 @@ def _apply_inter_turn_delay_cap(d: dict[str, Any], cli: CLIConfig) -> None: def _apply_trace_delay_flags(d: dict[str, Any], cli: CLIConfig) -> None: """Route trace-replay delay knobs onto ``FileDataset``/``PublicDataset``. - ``--ignore-trace-delays``, ``--use-think-time-only``, and - ``--trace-idle-gap-cap-seconds`` live on - both FILE and PUBLIC dataset models and bake into ``Turn.delay`` / - ``Turn.timestamp`` at load time. Without this route the CLI flags are - silently dropped (YAML / scenario paths set the fields directly). + These flags live on both FILE and PUBLIC dataset models. Ignore/think-only + affect ``Turn.delay`` / ``Turn.timestamp`` at load time; the trace idle cap + is consumed by the AgentX runtime watchdog without rewriting either field. + Without this route the CLI flags are silently dropped (YAML / scenario + paths set the fields directly). Mutual exclusivity of ignore vs think-only is enforced by the dataset model validators after conversion. """ diff --git a/src/aiperf/config/flags/cli_config.py b/src/aiperf/config/flags/cli_config.py index d70b61aba3..0574f36fdb 100644 --- a/src/aiperf/config/flags/cli_config.py +++ b/src/aiperf/config/flags/cli_config.py @@ -2180,15 +2180,14 @@ def url(self) -> str: Field( default=None, ge=0.0, - description="Hard ceiling (seconds) for idle gaps within each individual trace. " - "For Weka trace replay, AIPerf looks at all parent and subagent request " - "submission timestamps within one root trace, compresses long gaps between " - "consecutive request submissions, and derives turn delays from the " - "compressed per-trace timeline. Original request api_time values are not " - "used to decide these idle gaps. When set for Weka, this takes precedence over " - "`--inter-turn-delay-cap-seconds` so individual parent/subagent-line " - "delays are not separately capped. Defaults to None (no per-trace " - "idle-gap compression).", + description="Hard ceiling (seconds) for observed runtime idle time within " + "each AgentX trajectory tree. The idle clock covers initial future work and " + "restarts when the final in-flight request across the root and all descendant " + "streams completes. If no request from that tree reaches the wire before the " + "cap, AIPerf uniformly advances " + "only that tree's pending replay timers. Dataset timestamps are unchanged, " + "and request order, spawn/join dependencies, and replay barriers remain " + "authoritative. Defaults to None (no per-trace runtime idle cap).", ), CLIParameter( name=("--trace-idle-gap-cap-seconds",), diff --git a/src/aiperf/config/schema/aiperf-config.schema.json b/src/aiperf/config/schema/aiperf-config.schema.json index 99c8b6371b..0cf8215eff 100644 --- a/src/aiperf/config/schema/aiperf-config.schema.json +++ b/src/aiperf/config/schema/aiperf-config.schema.json @@ -4532,7 +4532,7 @@ } ], "default": null, - "description": "Hard ceiling (seconds) for idle gaps within each individual trace. For Weka trace replay, AIPerf looks at all parent and subagent request submission timestamps within one root trace, compresses long gaps between consecutive request submissions, and derives turn delays from the compressed per-trace timeline. Original request api_time values are not used to decide these idle gaps. When set for Weka, this takes precedence over `--inter-turn-delay-cap-seconds` so individual parent/subagent-line delays are not separately capped. Defaults to None (no per-trace idle-gap compression).", + "description": "Hard ceiling (seconds) for observed runtime idle time within each AgentX trajectory tree. The idle clock covers initial future work and restarts when the final in-flight request across the root and all descendant streams completes. If no request from that tree reaches the wire before the cap, AIPerf uniformly advances only that tree's pending replay timers. Dataset timestamps are unchanged, and request order, spawn/join dependencies, and replay barriers remain authoritative. Defaults to None (no per-trace runtime idle cap).", "title": "Traceidlegapcapseconds" }, "ignoreTraceDelays": { @@ -4695,7 +4695,7 @@ } ], "default": null, - "description": "Hard ceiling (seconds) for idle gaps within each individual trace for HF-backed Weka replay (mirror of ``FileDataset.trace_idle_gap_cap_seconds``). When set, takes precedence over ``inter_turn_delay_cap_seconds``. Defaults to None (no per-trace idle-gap compression).", + "description": "Hard ceiling (seconds) for observed runtime idle time within each AgentX trajectory tree for HF-backed Weka replay (mirror of ``FileDataset.trace_idle_gap_cap_seconds``). Dataset timestamps remain unchanged. Defaults to None (no per-trace runtime idle cap).", "title": "Traceidlegapcapseconds" }, "ignoreTraceDelays": { @@ -13557,7 +13557,7 @@ } ], "default": null, - "description": "Hard ceiling (seconds) for idle gaps within each individual trace. For Weka trace replay, AIPerf looks at all parent and subagent request submission timestamps within one root trace, compresses long gaps between consecutive request submissions, and derives turn delays from the compressed per-trace timeline. Original request api_time values are not used to decide these idle gaps. When set for Weka, this takes precedence over `--inter-turn-delay-cap-seconds` so individual parent/subagent-line delays are not separately capped. Defaults to None (no per-trace idle-gap compression).", + "description": "Hard ceiling (seconds) for observed runtime idle time within each AgentX trajectory tree. The idle clock covers initial future work and restarts when the final in-flight request across the root and all descendant streams completes. If no request from that tree reaches the wire before the cap, AIPerf uniformly advances only that tree's pending replay timers. Dataset timestamps are unchanged, and request order, spawn/join dependencies, and replay barriers remain authoritative. Defaults to None (no per-trace runtime idle cap).", "title": "Traceidlegapcapseconds", "x-jinja2-supported": true }, @@ -18437,7 +18437,7 @@ } ], "default": null, - "description": "Hard ceiling (seconds) for idle gaps within each individual trace for HF-backed Weka replay (mirror of ``FileDataset.trace_idle_gap_cap_seconds``). When set, takes precedence over ``inter_turn_delay_cap_seconds``. Defaults to None (no per-trace idle-gap compression).", + "description": "Hard ceiling (seconds) for observed runtime idle time within each AgentX trajectory tree for HF-backed Weka replay (mirror of ``FileDataset.trace_idle_gap_cap_seconds``). Dataset timestamps remain unchanged. Defaults to None (no per-trace runtime idle cap).", "title": "Traceidlegapcapseconds", "x-jinja2-supported": true }, diff --git a/src/aiperf/dataset/loader/weka_parallel_convert.py b/src/aiperf/dataset/loader/weka_parallel_convert.py index cf6a5b0179..7c58c7429d 100644 --- a/src/aiperf/dataset/loader/weka_parallel_convert.py +++ b/src/aiperf/dataset/loader/weka_parallel_convert.py @@ -138,10 +138,6 @@ class _WekaNormalRequestPayload(TypedDict): source_kind: NotRequired[str | None] # Only present in parent normals (not in child requests): capped_output_length: NotRequired[int] - # Present when --trace-idle-gap-cap-seconds has rewritten the per-trace - # timeline before workers compute turns. - effective_t: NotRequired[float] - effective_delay_ms: NotRequired[float | None] # Theoretical prefix-cache values precomputed parent-side from the # per-trace shared seen-set (one hash namespace per trace file). theoretical_hit_blocks: int @@ -165,9 +161,7 @@ class _WekaSubagentMarkerPayload(TypedDict): system_tokens: int child_session_ids: list[str] sa_end_seconds: float - effective_sa_end_seconds: NotRequired[float] t: float - effective_t: NotRequired[float] class _WekaFlatChainMarkerPayload(TypedDict): @@ -177,9 +171,7 @@ class _WekaFlatChainMarkerPayload(TypedDict): chain_index: int first_outer_idx: int end_seconds: float - effective_end_seconds: NotRequired[float] t: float - effective_t: NotRequired[float] class _WekaParentPayload(TypedDict): @@ -433,20 +425,16 @@ def _process_task(task: _WekaTraceTask) -> _WekaProcessTaskResult: max_asst_blocks=parent_asst_block_caps[k], ) - if "effective_t" in req: - t_ms = req["effective_t"] * 1000.0 - delay_ms = req.get("effective_delay_ms") + t_ms = req["t"] * 1000.0 + if k == 0: + delay_ms = None + elif task.think_time_only and req.get("think_time") is not None: + delay_ms = req["think_time"] * 1000.0 else: - t_ms = req["t"] * 1000.0 - if k == 0: - delay_ms = None - elif task.think_time_only and req.get("think_time") is not None: - delay_ms = req["think_time"] * 1000.0 - else: - prev_req = normals[k - 1][1] - delay_ms = _end_to_start_delay_ms( - t_ms - prev_req["t"] * 1000.0, prev_req.get("api_time") - ) + prev_req = normals[k - 1][1] + delay_ms = _end_to_start_delay_ms( + t_ms - prev_req["t"] * 1000.0, prev_req.get("api_time") + ) if delay_ms is not None: delay_ms = delay_tracker.clamp(delay_ms) # Floor at 0, mirroring the serial parent loop in @@ -535,9 +523,7 @@ def _process_task(task: _WekaTraceTask) -> _WekaProcessTaskResult: defaultdict(list) ) group_order: list[tuple[int, int | None]] = [] - outer_to_t: dict[int, float] = { - oi: req.get("effective_t", req["t"]) for oi, req in normals - } + outer_to_t: dict[int, float] = {oi: req["t"] for oi, req in normals} dropped_agent_ids: set[str] = set() dropped_subagent_indices: set[int] = set() for subagent_index, (sa_outer_idx, sa_entry) in enumerate(parent["subagents"]): @@ -554,9 +540,7 @@ def _process_task(task: _WekaTraceTask) -> _WekaProcessTaskResult: for oi, pos in sorted(outer_to_turn_pos.items()): if oi <= sa_outer_idx: continue - sa_end_seconds = sa_entry.get( - "effective_sa_end_seconds", sa_entry["sa_end_seconds"] - ) + sa_end_seconds = sa_entry["sa_end_seconds"] if outer_to_t[oi] + _JOIN_EPSILON_SECONDS >= sa_end_seconds: join_turn = pos break @@ -580,8 +564,7 @@ def _process_task(task: _WekaTraceTask) -> _WekaProcessTaskResult: "is_background": is_background, "preceding_turn": preceding, "following_turn": join_turn, - "start_timestamp": min(e.get("effective_t", e["t"]) for e in entries) - * 1000.0, + "start_timestamp": min(e["t"] for e in entries) * 1000.0, } ) @@ -599,7 +582,7 @@ def _process_task(task: _WekaTraceTask) -> _WekaProcessTaskResult: (pos for oi, pos in outer_to_turn_pos.items() if oi < first_outer), default=0, ) - fp_end = marker.get("effective_end_seconds", marker["end_seconds"]) + fp_end = marker["end_seconds"] join_turn = None for oi, pos in sorted(outer_to_turn_pos.items()): if oi <= first_outer: @@ -621,12 +604,7 @@ def _process_task(task: _WekaTraceTask) -> _WekaProcessTaskResult: "is_background": join_turn is None, "preceding_turn": preceding, "following_turn": join_turn, - # Mapped spawn time under an idle-gap warp (effective_t), raw t - # otherwise. Mirrors the subagent branch above and the serial - # flat branch; raw t under a warp collapses the child dispatch - # offset to 0. - "start_timestamp": min(m.get("effective_t", m["t"]) for m in markers) - * 1000.0, + "start_timestamp": min(m["t"] for m in markers) * 1000.0, } ) @@ -680,20 +658,16 @@ def _process_task(task: _WekaTraceTask) -> _WekaProcessTaskResult: is_tool_result=is_tool_result, max_asst_blocks=child_asst_block_caps[k], ) - if "effective_t" in creq: - t_ms = creq["effective_t"] * 1000.0 - child_delay_ms = creq.get("effective_delay_ms") + t_ms = creq["t"] * 1000.0 + if k == 0: + child_delay_ms = None + elif task.think_time_only and creq.get("think_time") is not None: + child_delay_ms = creq["think_time"] * 1000.0 else: - t_ms = creq["t"] * 1000.0 - if k == 0: - child_delay_ms = None - elif task.think_time_only and creq.get("think_time") is not None: - child_delay_ms = creq["think_time"] * 1000.0 - else: - prev_creq = creqs[k - 1] - child_delay_ms = _end_to_start_delay_ms( - t_ms - prev_creq["t"] * 1000.0, prev_creq.get("api_time") - ) + prev_creq = creqs[k - 1] + child_delay_ms = _end_to_start_delay_ms( + t_ms - prev_creq["t"] * 1000.0, prev_creq.get("api_time") + ) if child_delay_ms is not None: child_delay_ms = delay_tracker.clamp(child_delay_ms) diff --git a/src/aiperf/dataset/loader/weka_trace.py b/src/aiperf/dataset/loader/weka_trace.py index c68e24cdc0..f60459b4c4 100644 --- a/src/aiperf/dataset/loader/weka_trace.py +++ b/src/aiperf/dataset/loader/weka_trace.py @@ -13,7 +13,7 @@ import math from collections import defaultdict -from dataclasses import dataclass, field +from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Any @@ -255,98 +255,6 @@ def _clamp_delay_ms(delay_ms: float, cap_seconds: float | None) -> float | None: return delay_ms -@dataclass(frozen=True) -class _RequestTiming: - timestamp_seconds: float - delay_ms: float | None - - -@dataclass(frozen=True) -class _IdleGap: - raw_start: float - raw_end: float - shift_before: float - cap_seconds: float - excess_seconds: float - - -@dataclass -class _TraceIdleTiming: - parent_by_outer_idx: dict[int, _RequestTiming] - # Keyed by (child_plan.session_id, request_idx_within_stream). Not id(req): - # the parallel reconstruction path pickles request objects to worker - # processes, where they materialize at fresh memory addresses and any - # id()-based dict misses with KeyError. (session_id, idx) is stable - # across the pickle round-trip. Detected flat chains land in the same - # map under their own session ids. - child_by_session_request: dict[tuple[str, int], _RequestTiming] - subagent_end_by_outer_idx: dict[int, float] - # Mapped subagent spawn time (warp.map(entry.t)), keyed by outer idx. The - # branch start_timestamp_ms must live on the same compressed timeline as - # the turns it links; the raw entry.t would land far past the parent's - # last turn whenever an idle gap is compressed. - subagent_start_by_outer_idx: dict[int, float] - flat_chain_end_by_session: dict[str, float] = field(default_factory=dict) - # Mapped flat-chain spawn time (warp.map of the chain's first request t), - # keyed by child session id. Same invariant as subagent_start_by_outer_idx: - # the flat-chain SPAWN branch start_timestamp_ms must live on the same - # compressed timeline as the chain's (already-warped) turns. The raw - # first-request t would land far past those turns whenever a leading idle - # gap is compressed, driving _child_dispatch_offset_ms negative -> clamped - # to 0, collapsing the recorded child dispatch offset. - flat_chain_start_by_session: dict[str, float] = field(default_factory=dict) - - -class _IdleGapTimeWarp: - """Compress request-start gaps in one trace and map raw seconds to adjusted seconds.""" - - def __init__(self, request_starts: list[float], cap_seconds: float): - self._gaps: list[_IdleGap] = [] - sorted_starts = sorted(request_starts) - if not sorted_starts: - return - - prev_start = sorted_starts[0] - cumulative_shift = 0.0 - for start in sorted_starts[1:]: - gap_seconds = start - prev_start - if gap_seconds > cap_seconds: - excess = gap_seconds - cap_seconds - self._gaps.append( - _IdleGap( - raw_start=prev_start, - raw_end=start, - shift_before=cumulative_shift, - cap_seconds=cap_seconds, - excess_seconds=excess, - ) - ) - cumulative_shift += excess - prev_start = start - - def map(self, t_seconds: float) -> float: - """Map a raw timestamp to the per-trace idle-gap-capped timeline. - - Each long request-start gap ``[a, b]`` is compressed by keeping the first - ``cap_seconds`` after request ``a`` intact and collapsing the remainder - to the cap boundary. Requests at or after ``b`` shift left by the - collapsed excess. Non-request events inside the collapsed tail, such as - subagent end markers, map to the same boundary so joins cannot wait past - the next shifted request. - """ - shift = 0.0 - for gap in self._gaps: - if t_seconds < gap.raw_start: - return t_seconds - gap.shift_before - if t_seconds < gap.raw_end: - local = t_seconds - gap.raw_start - if local <= gap.cap_seconds: - return t_seconds - gap.shift_before - return gap.raw_start - gap.shift_before + gap.cap_seconds - shift = gap.shift_before + gap.excess_seconds - return t_seconds - shift - - @dataclass class _ParentPlan: trace_id: str @@ -751,38 +659,6 @@ def _chain_init_tokens( return 0, 0 -def _populate_flat_chain_timing( - flat_plans_for_trace: list[_FlatChainPlan], - warp: _IdleGapTimeWarp, - child_by_session_request: dict[tuple[str, int], _RequestTiming], -) -> tuple[dict[str, float], dict[str, float]]: - """Warp flat-chain request timing onto the shared per-trace timeline. - - Mutates ``child_by_session_request`` in place (flat chains share the - subagent children's keyspace) and returns - ``(flat_chain_start_by_session, flat_chain_end_by_session)``: the warped - first-request time used for the SPAWN branch start and the warped chain-end - used for SPAWN_JOIN placement. Both must live on the compressed timeline so - the branch anchors share it with the chain's (already-warped) turns. - """ - flat_chain_start_by_session: dict[str, float] = {} - flat_chain_end_by_session: dict[str, float] = {} - for fp in flat_plans_for_trace: - prev_flat_t: float | None = None - prev_flat_api: float | None = None - for k, (_, req) in enumerate(fp.requests): - t = warp.map(req.t) - delay_ms = None if prev_flat_t is None else (t - prev_flat_t) * 1000.0 - delay_ms = _end_to_start_delay_ms(delay_ms, prev_flat_api) - child_by_session_request[(fp.session_id, k)] = _RequestTiming(t, delay_ms) - if k == 0: - flat_chain_start_by_session[fp.session_id] = t - prev_flat_t = t - prev_flat_api = req.api_time - flat_chain_end_by_session[fp.session_id] = warp.map(_flat_chain_end_seconds(fp)) - return flat_chain_start_by_session, flat_chain_end_by_session - - def _classify_turn_input( req: _NormalRequestT, prev_req: _NormalRequestT | None ) -> TurnInputKind | None: @@ -806,95 +682,6 @@ def _classify_turn_input( return None -def _build_trace_idle_timing( - *, - plan: _ParentPlan, - child_plans: list[_ChildPlan], - cap_seconds: float, - flat_plans: list[_FlatChainPlan] | None = None, -) -> _TraceIdleTiming: - """Build per-turn timing after capping request-start gaps in one root trace. - - The cap is per root trace, not global across the dataset. We collect every - request submission timestamp from the parent and all subagents, compress - any gap between consecutive starts above ``cap_seconds``, then derive parent - and child conversation delays from that adjusted timeline. - - Example with ``cap_seconds=60``: - - main request starts at t=0 - - subagent request starts at t=20 and originally takes 80s - - next main request starts at t=220 - The capped gap is based on request starts only: 20 -> 220 is 200s, so the - next main request shifts left by 140s to t=80. The original subagent - latency still matters for join placement, but it does not prevent this idle - gap from being compressed. - """ - flat_plans_for_trace = [ - fp for fp in (flat_plans or []) if fp.parent_trace_id == plan.trace_id - ] - request_starts: list[float] = [] - for _, req in plan.normals: - request_starts.append(req.t) - # Flat-chain requests were top-level rows before the split; including - # them keeps the warp's gap structure identical to the unsplit trace. - for fp in flat_plans_for_trace: - for _, req in fp.requests: - request_starts.append(req.t) - - child_plans_for_trace = _child_plans_for_active_subagents(plan, child_plans) - for cp in child_plans_for_trace: - for req in cp.requests: - request_starts.append(req.t) - - warp = _IdleGapTimeWarp(request_starts, cap_seconds) - parent_by_outer_idx: dict[int, _RequestTiming] = {} - prev_t: float | None = None - prev_api: float | None = None - for outer_idx, req in plan.normals: - t = warp.map(req.t) - delay_ms = None if prev_t is None else (t - prev_t) * 1000.0 - delay_ms = _end_to_start_delay_ms(delay_ms, prev_api) - parent_by_outer_idx[outer_idx] = _RequestTiming(t, delay_ms) - prev_t = t - prev_api = req.api_time - - child_by_session_request: dict[tuple[str, int], _RequestTiming] = {} - for cp in child_plans_for_trace: - prev_child_t: float | None = None - prev_child_api: float | None = None - for k, req in enumerate(cp.requests): - t = warp.map(req.t) - delay_ms = None if prev_child_t is None else (t - prev_child_t) * 1000.0 - delay_ms = _end_to_start_delay_ms(delay_ms, prev_child_api) - child_by_session_request[(cp.session_id, k)] = _RequestTiming(t, delay_ms) - prev_child_t = t - prev_child_api = req.api_time - - flat_chain_start_by_session, flat_chain_end_by_session = ( - _populate_flat_chain_timing( - flat_plans_for_trace, - warp, - child_by_session_request, - ) - ) - - subagent_end_by_outer_idx = { - outer_idx: warp.map(_sa_end_seconds(entry)) - for outer_idx, entry in plan.subagents - } - subagent_start_by_outer_idx = { - outer_idx: warp.map(entry.t) for outer_idx, entry in plan.subagents - } - return _TraceIdleTiming( - parent_by_outer_idx=parent_by_outer_idx, - child_by_session_request=child_by_session_request, - subagent_end_by_outer_idx=subagent_end_by_outer_idx, - subagent_start_by_outer_idx=subagent_start_by_outer_idx, - flat_chain_end_by_session=flat_chain_end_by_session, - flat_chain_start_by_session=flat_chain_start_by_session, - ) - - class WekaTraceLoader(HashIdsPromptSynthesisMixin, BaseFileLoader): """Dataset loader for Weka KV-cache-tester agentic coding trace files. @@ -1034,10 +821,6 @@ def __init__( self._delay_cap_tracker = DelayCapTracker( cap_seconds=self._inter_turn_delay_cap_seconds ) - # Per-trace idle-gap cap lives on FileDataset.trace_idle_gap_cap_seconds. - self._configured_trace_idle_gap_cap_seconds = getattr( - dataset, "trace_idle_gap_cap_seconds", None - ) self._tool_shaped_messages = Environment.DATASET.WEKA_TOOL_SHAPED_MESSAGES def _block_size_for_trace(self, trace: WekaTrace) -> int: @@ -1257,13 +1040,6 @@ def _cap_output(self, req: _NormalRequestT) -> int: capped = max_osl return capped if capped >= 1 else 1 - def _trace_idle_gap_cap_seconds(self) -> float | None: - """Optional per-trace idle-gap cap; robust to MagicMock test configs.""" - value = self._configured_trace_idle_gap_cap_seconds - if isinstance(value, int | float): - return float(value) - return None - def _build_reconstruction_plans( self, data: dict[str, list[WekaTrace]] ) -> _ReconstructionPlans: @@ -1540,25 +1316,6 @@ def _build_shared_metric_values( out[plan.trace_id] = compute_shared_prefix_cache_metrics(records) return out - def _build_trace_idle_timing_by_trace( - self, - parent_plans: list[_ParentPlan], - child_plans: list[_ChildPlan], - flat_plans: list[_FlatChainPlan] | None = None, - ) -> dict[str, _TraceIdleTiming]: - trace_idle_gap_cap_seconds = self._trace_idle_gap_cap_seconds() - if trace_idle_gap_cap_seconds is None: - return {} - return { - plan.trace_id: _build_trace_idle_timing( - plan=plan, - child_plans=child_plans, - cap_seconds=trace_idle_gap_cap_seconds, - flat_plans=flat_plans, - ) - for plan in parent_plans - } - def _build_model_map(self, trace: WekaTrace) -> dict[str, str]: """Map trace-side model names to ``endpoint.model_names``. @@ -1696,14 +1453,7 @@ def convert_to_conversations( ignore_delays = self._ignore_trace_delays think_time_only = self._use_think_time_only cap_seconds = self._inter_turn_delay_cap_seconds - trace_idle_gap_cap_seconds = self._trace_idle_gap_cap_seconds() - trace_idle_timing_by_trace = self._build_trace_idle_timing_by_trace( - parent_plans, child_plans, flat_plans - ) - turn_cap_seconds = ( - None if trace_idle_gap_cap_seconds is not None else cap_seconds - ) - self._delay_cap_tracker.cap_seconds = turn_cap_seconds + self._delay_cap_tracker.cap_seconds = cap_seconds _t0 = _time.monotonic() _t1 = _time.monotonic() @@ -1725,11 +1475,10 @@ def convert_to_conversations( data=data, ignore_delays=ignore_delays, think_time_only=think_time_only, - cap_seconds=turn_cap_seconds, + cap_seconds=cap_seconds, configured_workers=configured_workers, t_start=_t1, model_map_per_trace=model_map_per_trace, - trace_idle_timing_by_trace=trace_idle_timing_by_trace, metric_values_by_trace=metric_values_by_trace, flat_plans=flat_plans, ) @@ -1741,10 +1490,9 @@ def convert_to_conversations( dropped_per_trace=dropped_per_trace, ignore_delays=ignore_delays, think_time_only=think_time_only, - cap_seconds=turn_cap_seconds, + cap_seconds=cap_seconds, t_start=_t1, model_map_per_trace=model_map_per_trace, - trace_idle_timing_by_trace=trace_idle_timing_by_trace, metric_values_by_trace=metric_values_by_trace, flat_plans=flat_plans, ) @@ -1801,7 +1549,6 @@ def _reconstruct_serial( cap_seconds: float | None, t_start: float, model_map_per_trace: dict[str, dict[str, str]], - trace_idle_timing_by_trace: dict[str, _TraceIdleTiming], metric_values_by_trace: dict[str, dict[tuple[str, int], tuple[int, int]]], flat_plans: list[_FlatChainPlan] | None = None, ) -> list[Conversation]: @@ -1848,7 +1595,6 @@ def _reconstruct_serial( # endpoint accumulates across turns at request time, with # ``reset_context`` flagging non-monotonic LCP cuts. trace = data[plan.trace_id][0] - trace_idle_timing = trace_idle_timing_by_trace.get(plan.trace_id) conv = Conversation( session_id=plan.trace_id, context_mode=self._resolved_context_mode(), @@ -1904,21 +1650,16 @@ def _reconstruct_serial( ) # Turn.timestamp/delay are in milliseconds; weka traces record seconds. - if trace_idle_timing is not None: - timing = trace_idle_timing.parent_by_outer_idx[outer_idx] - t_ms = timing.timestamp_seconds * 1000.0 - delay_ms = timing.delay_ms + t_ms = req.t * 1000.0 + if k == 0: + delay_ms = None + elif think_time_only and req.think_time is not None: + delay_ms = req.think_time * 1000.0 else: - t_ms = req.t * 1000.0 - if k == 0: - delay_ms = None - elif think_time_only and req.think_time is not None: - delay_ms = req.think_time * 1000.0 - else: - prev_req = plan.normals[k - 1][1] - delay_ms = _end_to_start_delay_ms( - t_ms - prev_req.t * 1000.0, prev_req.api_time - ) + prev_req = plan.normals[k - 1][1] + delay_ms = _end_to_start_delay_ms( + t_ms - prev_req.t * 1000.0, prev_req.api_time + ) if delay_ms is not None: delay_ms = self._delay_cap_tracker.clamp(delay_ms) # Floor at 0: a negative inter-turn delay (corrupt @@ -2008,15 +1749,7 @@ def _reconstruct_serial( for cp in child_plans: if cp.parent_trace_id == plan.trace_id: child_sids_by_subagent[cp.subagent_index].append(cp.session_id) - if trace_idle_timing is not None: - outer_to_t: dict[int, float] = { - outer_idx: trace_idle_timing.parent_by_outer_idx[ - outer_idx - ].timestamp_seconds - for outer_idx, _ in plan.normals - } - else: - outer_to_t = {outer_idx: req.t for outer_idx, req in plan.normals} + outer_to_t = {outer_idx: req.t for outer_idx, req in plan.normals} for subagent_index, (sa_outer_idx, sa_entry) in enumerate(plan.subagents): preceding = max( @@ -2031,14 +1764,8 @@ def _reconstruct_serial( dropped_subagent_indices.add(subagent_index) continue - if trace_idle_timing is not None: - sa_end_t = trace_idle_timing.subagent_end_by_outer_idx[sa_outer_idx] - sa_start_t = trace_idle_timing.subagent_start_by_outer_idx[ - sa_outer_idx - ] - else: - sa_end_t = _sa_end_seconds(sa_entry) - sa_start_t = sa_entry.t + sa_end_t = _sa_end_seconds(sa_entry) + sa_start_t = sa_entry.t join_turn: int | None = None for oi, pos in sorted(outer_to_turn_pos.items()): if oi <= sa_outer_idx: @@ -2072,8 +1799,6 @@ def _reconstruct_serial( child_conversation_ids=child_sids, mode=ConversationBranchMode.SPAWN, is_background=is_background, - # Mapped spawn time (raw entry.t when no idle-gap warp), - # so the branch start shares the turns' compressed timeline. start_timestamp_ms=min(s for _, _, s, _ in entries) * 1000.0, ) ) @@ -2100,10 +1825,7 @@ def _reconstruct_serial( (pos for oi, pos in outer_to_turn_pos.items() if oi < first_outer), default=0, ) - if trace_idle_timing is not None: - fp_end = trace_idle_timing.flat_chain_end_by_session[fp.session_id] - else: - fp_end = _flat_chain_end_seconds(fp) + fp_end = _flat_chain_end_seconds(fp) join_turn = None for oi, pos in sorted(outer_to_turn_pos.items()): if oi <= first_outer: @@ -2119,18 +1841,7 @@ def _reconstruct_serial( for preceding, join_turn in flat_group_order: fps = flat_groups[(preceding, join_turn)] branch_id = f"{plan.trace_id}:flatspawn:{fps[0].chain_index}" - # Mapped flat-chain spawn time when the idle-gap warp is active - # (raw first-request t otherwise), so the branch start shares the - # chain turns' compressed timeline. Mirrors the subagent SPAWN - # branch above; using raw t under a warp would drive the recorded - # child dispatch offset negative -> clamped to 0. - if trace_idle_timing is not None: - flat_start_seconds = min( - trace_idle_timing.flat_chain_start_by_session[fp.session_id] - for fp in fps - ) - else: - flat_start_seconds = min(fp.requests[0][1].t for fp in fps) + flat_start_seconds = min(fp.requests[0][1].t for fp in fps) conv.branches.append( ConversationBranchInfo( branch_id=branch_id, @@ -2178,9 +1889,6 @@ def _reconstruct_serial( ignore_delays=ignore_delays, think_time_only=think_time_only, model_map=model_map_per_trace.get(cp.parent_trace_id, {}), - trace_idle_timing=trace_idle_timing_by_trace.get( - cp.parent_trace_id - ), metric_values=metric_values_by_trace[cp.parent_trace_id], ) ) @@ -2248,27 +1956,18 @@ def _reconstruct_serial( is_tool_result=is_tool_result, max_asst_blocks=child_asst_block_caps[k], ) - trace_idle_timing = trace_idle_timing_by_trace.get(cp.parent_trace_id) - if trace_idle_timing is not None: - timing = trace_idle_timing.child_by_session_request[ - (cp.session_id, k) - ] - t_ms = timing.timestamp_seconds * 1000.0 - child_delay_ms = timing.delay_ms + # Plan requests are already in root-trace coordinates, + # matching metric ordering and parent turn timestamps. + t_ms = creq.t * 1000.0 + if k == 0: + child_delay_ms = None + elif think_time_only and creq.think_time is not None: + child_delay_ms = creq.think_time * 1000.0 else: - # Plan requests are already in root-trace coordinates - # (normalized at expansion), matching the warp path, - # metric ordering, and parent turn timestamps. - t_ms = creq.t * 1000.0 - if k == 0: - child_delay_ms = None - elif think_time_only and creq.think_time is not None: - child_delay_ms = creq.think_time * 1000.0 - else: - prev_creq = cp.requests[k - 1] - child_delay_ms = _end_to_start_delay_ms( - t_ms - prev_creq.t * 1000.0, prev_creq.api_time - ) + prev_creq = cp.requests[k - 1] + child_delay_ms = _end_to_start_delay_ms( + t_ms - prev_creq.t * 1000.0, prev_creq.api_time + ) if child_delay_ms is not None: child_delay_ms = self._delay_cap_tracker.clamp(child_delay_ms) child_delta = child_recon.turn_delta() @@ -2308,7 +2007,6 @@ def _emit_flat_chain_conversation( ignore_delays: bool, think_time_only: bool, model_map: dict[str, str], - trace_idle_timing: _TraceIdleTiming | None, metric_values: dict[tuple[str, int], tuple[int, int]], ) -> Conversation: """Reconstruct one detected flat chain as a child Conversation. @@ -2369,7 +2067,6 @@ def _emit_flat_chain_conversation( fp=fp, k=k, req=req, - trace_idle_timing=trace_idle_timing, think_time_only=think_time_only, ) delta = recon.turn_delta() @@ -2412,30 +2109,23 @@ def _flat_turn_timing( fp: _FlatChainPlan, k: int, req: _NormalRequestT, - trace_idle_timing: _TraceIdleTiming | None, think_time_only: bool, ) -> tuple[float, float | None]: """(timestamp_ms, clamped delay_ms) for one flat-chain turn. - Same precedence as the subagent-child loop: warped per-trace timing - when the idle-gap cap is active, else raw per-chain deltas honoring - ``--use-think-time-only`` and the inter-turn delay cap. + Uses raw per-chain deltas while honoring ``--use-think-time-only`` + and the inter-turn delay cap. """ - if trace_idle_timing is not None: - timing = trace_idle_timing.child_by_session_request[(fp.session_id, k)] - t_ms = timing.timestamp_seconds * 1000.0 - delay_ms = timing.delay_ms + t_ms = req.t * 1000.0 + if k == 0: + delay_ms = None + elif think_time_only and req.think_time is not None: + delay_ms = req.think_time * 1000.0 else: - t_ms = req.t * 1000.0 - if k == 0: - delay_ms = None - elif think_time_only and req.think_time is not None: - delay_ms = req.think_time * 1000.0 - else: - prev_freq = fp.requests[k - 1][1] - delay_ms = _end_to_start_delay_ms( - t_ms - prev_freq.t * 1000.0, prev_freq.api_time - ) + prev_freq = fp.requests[k - 1][1] + delay_ms = _end_to_start_delay_ms( + t_ms - prev_freq.t * 1000.0, prev_freq.api_time + ) if delay_ms is not None: delay_ms = self._delay_cap_tracker.clamp(delay_ms) return t_ms, delay_ms @@ -2450,7 +2140,6 @@ def _build_parallel_reconstruction_tasks( think_time_only: bool, cap_seconds: float | None, model_map_per_trace: dict[str, dict[str, str]], - trace_idle_timing_by_trace: dict[str, _TraceIdleTiming], metric_values_by_trace: dict[str, dict[tuple[str, int], tuple[int, int]]], flat_plans: list[_FlatChainPlan] | None = None, ): @@ -2463,13 +2152,8 @@ def _build_parallel_reconstruction_tasks( for fp in flat_plans or []: flat_plans_by_trace[fp.parent_trace_id].append(fp) - # Drop the same child_plans the serial path drops at line ~1172. - # _build_trace_idle_timing only populates timing for active subagents - # (via _child_plans_for_active_subagents), so without this skip the - # lookup below KeyErrors on any subagent that appears before the - # first normal parent turn -- a real condition in the 256k-capped - # corpus where reshifted timelines can leave the first subagent - # without a preceding normal. + # Drop the same child plans as the serial path when a subagent has no + # preceding retained parent request to spawn it. dropped_per_trace: dict[str, set[int]] = { plan.trace_id: _dropped_subagent_indices(plan) for plan in parent_plans } @@ -2479,7 +2163,6 @@ def _build_parallel_reconstruction_tasks( for cp in child_plans: if cp.subagent_index in dropped_per_trace.get(cp.parent_trace_id, set()): continue - trace_idle_timing = trace_idle_timing_by_trace.get(cp.parent_trace_id) child_metric_values = metric_values_by_trace[cp.parent_trace_id] requests_dicts: list[_WekaNormalRequestPayload] = [] for k, creq in enumerate(cp.requests): @@ -2505,12 +2188,6 @@ def _build_parallel_reconstruction_tasks( creq, cp.requests[k - 1] if k else None ), } - if trace_idle_timing is not None: - timing = trace_idle_timing.child_by_session_request[ - (cp.session_id, k) - ] - req_payload["effective_t"] = timing.timestamp_seconds - req_payload["effective_delay_ms"] = timing.delay_ms requests_dicts.append(req_payload) children_by_trace[cp.parent_trace_id].append( { @@ -2532,21 +2209,15 @@ def _build_parallel_reconstruction_tasks( for plan in parent_plans: for fp in flat_plans_by_trace.get(plan.trace_id, []): children_by_trace[fp.parent_trace_id].append( - self._parallel_flat_child_payload( - fp, trace_idle_timing_by_trace, metric_values_by_trace - ) + self._parallel_flat_child_payload(fp, metric_values_by_trace) ) tasks: list[_WekaTraceTask] = [] for plan in parent_plans: trace = data[plan.trace_id][0] parent_payload: dict[str, Any] = { - "normals": self._parallel_parent_normals( - plan, trace_idle_timing_by_trace, metric_values_by_trace - ), - "subagents": self._parallel_subagents( - plan, sids_by_subagent, trace_idle_timing_by_trace - ), + "normals": self._parallel_parent_normals(plan, metric_values_by_trace), + "subagents": self._parallel_subagents(plan, sids_by_subagent), "tool_tokens": trace.tool_tokens, "system_tokens": trace.system_tokens, } @@ -2555,7 +2226,6 @@ def _build_parallel_reconstruction_tasks( plan=plan, trace=trace, flat_for_trace=flat_plans_by_trace.get(plan.trace_id, []), - trace_idle_timing=trace_idle_timing_by_trace.get(plan.trace_id), ) tasks.append( _WekaTraceTask( @@ -2579,7 +2249,6 @@ def _apply_flat_parent_payload_extras( plan: _ParentPlan, trace: WekaTrace, flat_for_trace: list[_FlatChainPlan], - trace_idle_timing: _TraceIdleTiming | None, ) -> None: """Add flat-chain branch markers to the parent payload.""" if not flat_for_trace: @@ -2593,24 +2262,12 @@ def _apply_flat_parent_payload_extras( "end_seconds": _flat_chain_end_seconds(fp), "t": fp.requests[0][1].t, } - if trace_idle_timing is not None: - marker["effective_end_seconds"] = ( - trace_idle_timing.flat_chain_end_by_session[fp.session_id] - ) - # Mapped spawn time -> _process_task reads it for the branch - # start_timestamp (m.get("effective_t", m["t"])); without it the - # parallel flat branch start silently reverts to raw seconds - # under a warp (mirrors the subagent marker's effective_t). - marker["effective_t"] = trace_idle_timing.flat_chain_start_by_session[ - fp.session_id - ] markers.append(marker) parent_payload["flat_markers"] = markers def _parallel_flat_child_payload( self, fp: _FlatChainPlan, - trace_idle_timing_by_trace: dict[str, _TraceIdleTiming], metric_values_by_trace: dict[str, dict[tuple[str, int], tuple[int, int]]], ) -> dict[str, Any]: """Build the worker child payload for one detected flat chain. @@ -2625,7 +2282,6 @@ def _parallel_flat_child_payload( _WekaNormalRequestPayload, ) - trace_idle_timing = trace_idle_timing_by_trace.get(fp.parent_trace_id) flat_metric_values = metric_values_by_trace[fp.parent_trace_id] requests_dicts: list[_WekaNormalRequestPayload] = [] for k, (outer_idx, req) in enumerate(fp.requests): @@ -2649,10 +2305,6 @@ def _parallel_flat_child_payload( "theoretical_hit_blocks": hit_blocks, "theoretical_total_blocks": total_blocks, } - if trace_idle_timing is not None: - timing = trace_idle_timing.child_by_session_request[(fp.session_id, k)] - req_payload["effective_t"] = timing.timestamp_seconds - req_payload["effective_delay_ms"] = timing.delay_ms requests_dicts.append(req_payload) return { "session_id": fp.session_id, @@ -2667,14 +2319,12 @@ def _parallel_flat_child_payload( def _parallel_parent_normals( self, plan: _ParentPlan, - trace_idle_timing_by_trace: dict[str, _TraceIdleTiming], metric_values_by_trace: dict[str, dict[tuple[str, int], tuple[int, int]]], ): from aiperf.dataset.loader.weka_parallel_convert import ( _WekaNormalRequestPayload, ) - trace_idle_timing = trace_idle_timing_by_trace.get(plan.trace_id) trace_metric_values = metric_values_by_trace[plan.trace_id] normals_dicts: list[tuple[int, _WekaNormalRequestPayload]] = [] for k, (outer_idx, req) in enumerate(plan.normals): @@ -2694,10 +2344,6 @@ def _parallel_parent_normals( "theoretical_hit_blocks": hit_blocks, "theoretical_total_blocks": total_blocks, } - if trace_idle_timing is not None: - timing = trace_idle_timing.parent_by_outer_idx[outer_idx] - req_payload["effective_t"] = timing.timestamp_seconds - req_payload["effective_delay_ms"] = timing.delay_ms normals_dicts.append((outer_idx, req_payload)) return normals_dicts @@ -2705,13 +2351,11 @@ def _parallel_subagents( self, plan: _ParentPlan, sids_by_subagent: dict[tuple[str, int], list[str]], - trace_idle_timing_by_trace: dict[str, _TraceIdleTiming], ): from aiperf.dataset.loader.weka_parallel_convert import ( _WekaSubagentMarkerPayload, ) - trace_idle_timing = trace_idle_timing_by_trace.get(plan.trace_id) subagents_dicts: list[tuple[int, _WekaSubagentMarkerPayload]] = [] for sa_index, (outer_idx, sa) in enumerate(plan.subagents): sa_payload: _WekaSubagentMarkerPayload = { @@ -2724,16 +2368,6 @@ def _parallel_subagents( "sa_end_seconds": _sa_end_seconds(sa), "t": sa.t, } - if trace_idle_timing is not None: - sa_payload["effective_sa_end_seconds"] = ( - trace_idle_timing.subagent_end_by_outer_idx[outer_idx] - ) - # Mapped spawn time -> _process_task reads it for the branch - # start_timestamp (e.get("effective_t", e["t"])); without it the - # parallel branch start silently reverts to raw seconds under a warp. - sa_payload["effective_t"] = ( - trace_idle_timing.subagent_start_by_outer_idx[outer_idx] - ) subagents_dicts.append((outer_idx, sa_payload)) return subagents_dicts @@ -2749,7 +2383,6 @@ def _reconstruct_parallel( configured_workers: int, t_start: float, model_map_per_trace: dict[str, dict[str, str]], - trace_idle_timing_by_trace: dict[str, _TraceIdleTiming], metric_values_by_trace: dict[str, dict[tuple[str, int], tuple[int, int]]], flat_plans: list[_FlatChainPlan] | None = None, ) -> list[Conversation]: @@ -2785,7 +2418,6 @@ def _reconstruct_parallel( think_time_only=think_time_only, cap_seconds=cap_seconds, model_map_per_trace=model_map_per_trace, - trace_idle_timing_by_trace=trace_idle_timing_by_trace, metric_values_by_trace=metric_values_by_trace, flat_plans=flat_plans, ) diff --git a/src/aiperf/timing/branch_orchestrator.py b/src/aiperf/timing/branch_orchestrator.py index c8680d2b78..cae0e92401 100644 --- a/src/aiperf/timing/branch_orchestrator.py +++ b/src/aiperf/timing/branch_orchestrator.py @@ -186,6 +186,7 @@ class PendingBranchJoin: parent_num_turns: int parent_agent_depth: int = 0 parent_parent_correlation_id: str | None = None + parent_root_correlation_id: str | None = None gated_turn_index: int | None = None outstanding: dict[str, PrereqState] = field(default_factory=dict) parent_branch_mode: ConversationBranchMode = ConversationBranchMode.FORK @@ -770,6 +771,9 @@ def _ensure_seeded_join( parent_num_turns=len(parent_meta.turns), parent_agent_depth=parent_state.agent_depth, parent_parent_correlation_id=parent_state.parent_correlation_id, + parent_root_correlation_id=( + parent_state.root_correlation_id or parent_state.x_correlation_id + ), gated_turn_index=gated_idx, parent_branch_mode=parent_state.branch_mode, parent_has_forks_on_gated_turn=has_forks, @@ -1261,6 +1265,7 @@ def _start_delayed_first_turn( self._scheduler.schedule_later( offset_ms / 1000.0, self._dispatch_first_turn_after_offset(child, 0.0, parent_corr), + group_id=child.effective_root_correlation_id, ) self.stats.children_delayed += 1 return @@ -1320,6 +1325,7 @@ def _ensure_future_join( parent_num_turns=len(parent_meta.turns), parent_agent_depth=credit.agent_depth, parent_parent_correlation_id=credit.parent_correlation_id, + parent_root_correlation_id=credit.effective_root_correlation_id, gated_turn_index=gated_idx, parent_branch_mode=getattr( credit, "branch_mode", ConversationBranchMode.FORK @@ -1464,6 +1470,9 @@ def _arm_join_replay_deadline( pending.parent_x_correlation_id, pending.gated_turn_index, ), + group_id=( + pending.parent_root_correlation_id or pending.parent_x_correlation_id + ), ) async def _on_join_replay_deadline( diff --git a/src/aiperf/timing/phase/runner.py b/src/aiperf/timing/phase/runner.py index fbe58a325b..52ccf72383 100644 --- a/src/aiperf/timing/phase/runner.py +++ b/src/aiperf/timing/phase/runner.py @@ -197,7 +197,11 @@ def __init__( counter=self._progress.counter, ) self._replay_barrier = ( - ReplayBarrierCoordinator(self._conversation_source.dataset_metadata) + ReplayBarrierCoordinator( + self._conversation_source.dataset_metadata, + scheduler=self._scheduler, + root_idle_gap_cap_seconds=self._root_idle_gap_cap_seconds(), + ) if ( self._config.timing_mode == TimingMode.AGENTIC_REPLAY and self._conversation_source.dataset_metadata is not None @@ -216,6 +220,23 @@ def __init__( self._baseline_start_ns: int | None = None self._baseline_end_ns: int | None = None + def _root_idle_gap_cap_seconds(self) -> float | None: + """Per-trace runtime idle cap for AgentX profiling. + + The cap measures actual whole-tree idle time after the final in-flight + request completes. Dataset timestamps remain untouched; only pending + runtime timers for that root and its descendants may be advanced. + """ + if ( + self._config.phase != CreditPhase.PROFILING + or self._config.timing_mode != TimingMode.AGENTIC_REPLAY + or self._run is None + ): + return None + dataset = self._run.cfg.get_default_dataset() + cap = getattr(dataset, "trace_idle_gap_cap_seconds", None) + return float(cap) if isinstance(cap, int | float) else None + def _build_credit_issuer( self, url_selection_strategy: URLSelectionStrategyProtocol | None ) -> CreditIssuer: diff --git a/src/aiperf/timing/replay_dependencies.py b/src/aiperf/timing/replay_dependencies.py index f965e0a359..5d147ece66 100644 --- a/src/aiperf/timing/replay_dependencies.py +++ b/src/aiperf/timing/replay_dependencies.py @@ -15,6 +15,7 @@ from aiperf.credit.dispatch import ChildDispatchResult if TYPE_CHECKING: + from aiperf.common.loop_scheduler import LoopScheduler from aiperf.common.models import DatasetMetadata from aiperf.credit.structs import Credit, TurnToSend @@ -152,12 +153,22 @@ class _RootBarrierState: """Keys of requests on this tree that have recorded completion.""" pending: dict[ReplayTurnKey, _PendingDispatch] """Dispatches keyed by request, waiting on their predecessors to complete.""" + in_flight: int = 0 + """Requests from this runtime tree currently on the wire.""" + idle_watchdog: asyncio.TimerHandle | None = None + """Per-tree idle-cap callback, armed only while no request is in flight.""" class ReplayBarrierCoordinator: """Release requests only after their recorded frontier has completed.""" - def __init__(self, dataset_metadata: DatasetMetadata) -> None: + def __init__( + self, + dataset_metadata: DatasetMetadata, + *, + scheduler: LoopScheduler | None = None, + root_idle_gap_cap_seconds: float | None = None, + ) -> None: self._predecessors: dict[ReplayTurnKey, tuple[ReplayTurnKey, ...]] = {} for conversation in dataset_metadata.conversations: for turn_index, turn in enumerate(conversation.turns): @@ -170,6 +181,39 @@ def __init__(self, dataset_metadata: DatasetMetadata) -> None: self._dispatch_tasks: set[asyncio.Task] = set() self._active = False self._releases_paused = False + self._scheduler = scheduler + self._root_idle_gap_cap_seconds = root_idle_gap_cap_seconds + self._root_idle_jumps = 0 + self._root_idle_seconds_skipped = 0.0 + + def _root_state(self, root_id: str) -> _RootBarrierState: + return self._roots.setdefault( + root_id, _RootBarrierState(completed=set(), pending={}) + ) + + def observe_issued(self, credit: Credit) -> None: + """Track one request reaching the wire and cancel its idle watchdog.""" + if not self._active: + return + state = self._root_state(credit.effective_root_correlation_id) + state.in_flight += 1 + if state.idle_watchdog is not None: + state.idle_watchdog.cancel() + state.idle_watchdog = None + + def observe_idle_root(self, root_id: str) -> None: + """Start monitoring a known-idle tree with pending runtime work. + + Profiling snapshots may begin with every stream scheduled in the + future, before any request from that runtime root has reached the wire. + Registering the root after those timers are installed lets the same + completion-driven watchdog cover that initial idle interval. + """ + if not self._active: + return + state = self._root_state(root_id) + if state.in_flight == 0: + self._arm_root_idle_watchdog(root_id, state) def activate(self) -> None: """Enable barriers after baseline cache priming completes.""" @@ -205,9 +249,7 @@ async def submit( if not self._active: return await issue() root_id = turn.effective_root_correlation_id - state = self._roots.setdefault( - root_id, _RootBarrierState(completed=set(), pending={}) - ) + state = self._root_state(root_id) key = ReplayTurnKey(turn.conversation_id, turn.turn_index) if self._ready(state, key) and not self._releases_paused: return await issue() @@ -225,9 +267,8 @@ def complete(self, credit: Credit) -> None: if not self._active: return root_id = credit.effective_root_correlation_id - state = self._roots.setdefault( - root_id, _RootBarrierState(completed=set(), pending={}) - ) + state = self._root_state(root_id) + state.in_flight = max(0, state.in_flight - 1) state.completed.add(ReplayTurnKey(credit.conversation_id, credit.turn_index)) if self._releases_paused: return @@ -237,10 +278,45 @@ def complete(self, credit: Credit) -> None: task = asyncio.create_task(self._dispatch_pending(pending)) self._dispatch_tasks.add(task) task.add_done_callback(self._dispatch_tasks.discard) + if state.in_flight == 0: + self._arm_root_idle_watchdog(root_id, state) + + def _arm_root_idle_watchdog(self, root_id: str, state: _RootBarrierState) -> None: + """Advance only this tree's timers after a fully idle capped gap.""" + cap = self._root_idle_gap_cap_seconds + if ( + cap is None + or cap < 0 + or self._scheduler is None + or state.idle_watchdog is not None + ): + return + loop = asyncio.get_running_loop() + state.idle_watchdog = loop.call_later(cap, self._enforce_root_idle_cap, root_id) + + def _enforce_root_idle_cap(self, root_id: str) -> None: + state = self._roots.get(root_id) + if state is None: + return + state.idle_watchdog = None + if state.in_flight != 0 or self._scheduler is None: + return + shifted = self._scheduler.cap_pending_delay_for_group(root_id, 0.0) + if shifted <= 0: + return + self._root_idle_jumps += 1 + self._root_idle_seconds_skipped += shifted + _logger.info( + "Per-trace idle cap advanced replay root %s by %.3fs", + root_id, + shifted, + ) def close_root(self, root_id: str) -> None: """Discard completed runtime state when a recycled tree drains.""" - self._roots.pop(root_id, None) + state = self._roots.pop(root_id, None) + if state is not None and state.idle_watchdog is not None: + state.idle_watchdog.cancel() def seed_completed_prefixes( self, @@ -248,9 +324,7 @@ def seed_completed_prefixes( boundaries: tuple[ReplayResumeBoundary, ...], ) -> None: """Seed exact pre-resume history before any turn can be submitted.""" - state = self._roots.setdefault( - root_id, _RootBarrierState(completed=set(), pending={}) - ) + state = self._root_state(root_id) if state.pending: raise RuntimeError( f"Cannot seed replay history after dispatch for root={root_id!r}" @@ -313,6 +387,9 @@ async def cancel_pending(self, *, notify_refused: bool) -> None: """Cancel retained dispatches during phase teardown.""" callbacks = [] for state in self._roots.values(): + if state.idle_watchdog is not None: + state.idle_watchdog.cancel() + state.idle_watchdog = None if notify_refused: callbacks.extend( pending.on_refused @@ -405,6 +482,10 @@ def complete(self, credit: Credit) -> None: if self._coordinator is not None: self._coordinator.complete(credit) + def observe_idle_root(self, root_correlation_id: str) -> None: + if self._coordinator is not None: + self._coordinator.observe_idle_root(root_correlation_id) + def close_root(self, root_correlation_id: str) -> None: if self._coordinator is not None: self._coordinator.close_root(root_correlation_id) @@ -439,5 +520,7 @@ async def cancel(self, *, notify_refused: bool) -> None: await self._coordinator.cancel_pending(notify_refused=notify_refused) async def observe_issued(self, credit: Credit) -> None: + if self._coordinator is not None: + self._coordinator.observe_issued(credit) if self._credit_issued is not None: await self._credit_issued(credit) diff --git a/src/aiperf/timing/strategies/agentic_replay.py b/src/aiperf/timing/strategies/agentic_replay.py index 065e129080..ecf12fcf7b 100644 --- a/src/aiperf/timing/strategies/agentic_replay.py +++ b/src/aiperf/timing/strategies/agentic_replay.py @@ -59,7 +59,6 @@ import time import uuid from collections import Counter, deque -from collections.abc import Iterable from typing import TYPE_CHECKING from msgspec.structs import replace as _struct_replace @@ -257,24 +256,6 @@ def __init__( self._system_idle_watchdog: asyncio.TimerHandle | None = None if self._system_idle_gap_cap_seconds is not None: 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: 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 - # every idle subagent onto t=cap. - # None when no cap is set (raw faithful timing -- the user's choice). - # v2: trace_idle_gap_cap_seconds lives on the (file) dataset config. - default_dataset = run.cfg.get_default_dataset() if run is not None else None - idle_cap_s = getattr(default_dataset, "trace_idle_gap_cap_seconds", None) - self._phase_offset_cap_ms: float | None = ( - idle_cap_s * MILLIS_PER_SECOND - if isinstance(idle_cap_s, int | float) - else None - ) # Wrap-fill + cache_bust=NONE produces byte-identical traffic across # shared-trace lanes. agentx-mvp auto-locks cache_bust=first_turn_prefix @@ -615,55 +596,20 @@ def _cancel_system_idle_watchdog(self) -> None: self._system_idle_watchdog = None def _capped_warmup_lead_ms(self, lead_ms: float) -> float: - """Clamp a WARMUP lead (t* - the warmed turn) to the idle-gap cap. + """Clamp a WARMUP lead to the global system-idle guard. Warmup is a cache-priming pass, not faithful replay: warming a stream that last fired hours before t* that far ahead would make the warmup - phase itself take hours. Clamping each lead bounds the priming window - to the cap; the relative timing among warmup requests carries no replay - meaning (each primes its own session, all converge at t*), so an - independent per-lead clamp is fine here. No-op when no cap is set. - - PROFILING dispatch offsets are NOT clamped this way -- see - :meth:`_leading_idle_shift_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 - a trajectory's leading idle (t* -> its earliest post-t* request) is - capped to the idle-gap cap. - - The idle-gap warp (applied at load time across ALL of a trace's request - timestamps) already bounds every inter-request gap to the cap, so the - dispatch offsets carry faithful relative spacing. The ONE gap the warp - cannot see is t* -> first request: t* is the trajectory sampling - instant, not a request. A single uniform shift caps that leading gap - while keeping every stream's recorded offset relative to the others - intact -- unlike a per-stream ``min(offset, cap)`` clamp, which collapses - every offset above the cap onto it, firing every idle subagent at the - same instant. - - Returns 0 when no cap is set or the leading idle is already within it. + phase itself take hours. The per-trajectory trace cap is deliberately + absent here: it is a profiling runtime watchdog, not a phase-start + schedule transform. """ - if self._phase_offset_cap_ms is None: - return 0.0 - present = [o for o in offsets if o is not None] - if not present: - return 0.0 - return max(0.0, min(present) - self._phase_offset_cap_ms) + if self._system_idle_gap_cap_seconds is None: + return lead_ms + return min( + lead_ms, + self._system_idle_gap_cap_seconds * MILLIS_PER_SECOND, + ) async def _execute_warmup(self) -> None: """Warm turn n-1 of every session active (mid-flight) at t*. @@ -1227,8 +1173,6 @@ def _handoff_residual_delay_ms( if returned_at_ns is not None: elapsed_ms = max(0.0, (finalized_at_ns - returned_at_ns) / 1_000_000.0) delay_ms = max(0.0, delay_ms - elapsed_ms) - if self._phase_offset_cap_ms is not None: - delay_ms = min(delay_ms, self._phase_offset_cap_ms) return delay_ms def _handoff_base_delay_ms(self, credit: Credit) -> float: @@ -1392,8 +1336,7 @@ def _profiling_spread_seconds(self) -> float: """Window over which each trajectory's FIRST request fires, in seconds. Per trajectory the first request is its earliest dispatchable stream - (``min`` of the lane offsets after the leading-idle shift; t0=0 spread, - or t0=lane-min burst). This returns max-minus-min of those per-trajectory + (t0=0 spread, or t0=lane-min burst). This returns max-minus-min of those per-trajectory first-request offsets -- the ramp-in window. Later streams in a trajectory (subagents that spawn further out at their own faithful offsets) are excluded: including them would report the full per- @@ -1409,12 +1352,7 @@ def _profiling_spread_seconds(self) -> float: ] if not dispatchable: continue - leading_shift_ms = self._leading_idle_shift_ms( - s.next_dispatch_offset_ms for s in dispatchable - ) - lane_offsets = [ - s.next_dispatch_offset_ms - leading_shift_ms for s in dispatchable - ] + lane_offsets = [s.next_dispatch_offset_ms for s in dispatchable] lane_t0 = min(lane_offsets) if self._burst_phase_starts else 0.0 first_offsets.append(min(lane_offsets) - lane_t0) if not first_offsets: @@ -1582,7 +1520,11 @@ async def _dispatch_next_turn(self, credit: Credit) -> None: else self.credit_issuer.issue_credit(turn) ) if next_meta.delay_ms is not None and next_meta.delay_ms > 0: - self.scheduler.schedule_later(next_meta.delay_ms / MILLIS_PER_SECOND, coro) + self.scheduler.schedule_later( + next_meta.delay_ms / MILLIS_PER_SECOND, + coro, + group_id=credit.effective_root_correlation_id, + ) else: await coro @@ -1708,16 +1650,12 @@ async def _dispatch_snapshot_for_profiling( ) dispatchable = [s for s in snapshot.states if not s.waiting_on_children] - # Compute one normalized replay timeline for dispatchable streams and - # gated parents alike. A parent's join deadline must use the same - # leading-idle shift / optional burst anchor as every other request in - # its trajectory; only its child gate is additional. - leading_shift_ms = self._leading_idle_shift_ms( - s.next_dispatch_offset_ms for s in dispatchable - ) + # Compute one replay timeline for dispatchable streams and gated + # parents alike. A parent's join deadline uses the same optional burst + # anchor as every other request in its trajectory; only its child gate + # is additional. offset_by_corr = { - s.x_correlation_id: s.next_dispatch_offset_ms - leading_shift_ms - for s in snapshot.states + s.x_correlation_id: s.next_dispatch_offset_ms for s in snapshot.states } if self._burst_phase_starts and dispatchable: t0_offset_ms = min( @@ -1801,16 +1739,9 @@ async def _dispatch_snapshot_for_profiling( ) if not self._has_tree_registry and not has_root_state: self._rootless_lane_outstanding[lane] = len(dispatchable) - # Cap the leading idle (t* -> the trajectory's earliest post-t* request) - # by shifting EVERY stream left by the same excess, so an idle-at-t* - # trajectory ramps in within the cap instead of going dead for the whole - # run -- while preserving the recorded spacing among streams. A per-stream - # min(offset, cap) clamp would instead collapse every idle subagent onto - # t=cap. The idle-gap warp already bounded inter-request gaps at load - # time; this fixes the one gap it cannot see (t* is not a request). - # Spread (default): t0 = 0, each lane fires at its (shifted) offset from - # t*. Burst (--burst-phase-starts): t0 = the lane's min offset, anchoring - # the earliest post-t* request at profiling-0. + # Spread (default): t0 = 0, each lane fires at its recorded offset from + # t*. Burst (--burst-phase-starts): t0 = the lane's minimum offset, + # anchoring the earliest post-t* request at profiling-0. for state in dispatchable: session = self.conversation_source.session_for_state(state) turn = self._build_turn_for_session(session, state.next_turn_index) @@ -1823,10 +1754,15 @@ async def _dispatch_snapshot_for_profiling( self.scheduler.schedule_later( delay_s, self.credit_issuer.issue_credit(turn), + group_id=turn.effective_root_correlation_id, ) else: await self.credit_issuer.issue_credit(turn) + root_correlation_id = self._lane_root_corr(snapshot) + if root_correlation_id is not None: + self.credit_issuer.replay_gate.observe_idle_root(root_correlation_id) + def _get_snapshot(self, trajectory: Trajectory) -> TrajectorySnapshot: """Return the persistent sampled snapshot for a trajectory lane. diff --git a/tests/component_integration/test_agentic_replay_e2e.py b/tests/component_integration/test_agentic_replay_e2e.py index f2c3ec9256..d4cf1e588f 100644 --- a/tests/component_integration/test_agentic_replay_e2e.py +++ b/tests/component_integration/test_agentic_replay_e2e.py @@ -216,7 +216,7 @@ def _make_running_scheduler() -> MagicMock: scheduler = MagicMock() scheduled: list[asyncio.Task] = [] - def _schedule_later(_delay, coro): + def _schedule_later(_delay, coro, **_kwargs): scheduled.append(asyncio.ensure_future(coro)) scheduler.schedule_later.side_effect = _schedule_later diff --git a/tests/integration/test_agentx_trace_idle_gap_cap.py b/tests/integration/test_agentx_trace_idle_gap_cap.py new file mode 100644 index 0000000000..218619ad6a --- /dev/null +++ b/tests/integration/test_agentx_trace_idle_gap_cap.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""End-to-end coverage for the AgentX per-trace idle-gap cap.""" + +from __future__ import annotations + +import json +from collections import defaultdict +from pathlib import Path + +import pytest + +from tests.harness.utils import AIPerfCLI, AIPerfMockServer, AIPerfResults + + +def _write_trace(path: Path, trace_id: str, starts: list[float]) -> None: + """Write one root-only Weka trace with the requested start schedule.""" + block_size = 16 + requests = [ + { + "t": start, + "type": "n", + "model": "mock-model", + "in": (turn_index + 1) * block_size + 8, + "out": 2, + "hash_ids": list(range(1, turn_index + 2)), + "input_types": ["text"], + "output_types": ["text"], + "stop": "end_turn", + "api_time": 0.05, + "think_time": max(0.0, start - starts[turn_index - 1] - 0.05) + if turn_index + else 0.0, + } + for turn_index, start in enumerate(starts) + ] + trace = { + "id": trace_id, + "models": ["mock-model"], + "block_size": block_size, + "hash_id_scope": "local", + "requests": requests, + } + path.write_text(json.dumps(trace), encoding="utf-8") + + +@pytest.fixture +def varied_weka_traces(tmp_path: Path) -> Path: + """Create several traces with different mixtures of short and long gaps.""" + trace_dir = tmp_path / "weka" + trace_dir.mkdir() + schedules = { + "alternating_2s": [0.0, 0.2, 0.4, 2.4, 2.6, 4.6, 4.8, 6.8], + "alternating_1_5s": [0.0, 0.15, 1.65, 1.85, 3.35, 3.55, 5.05, 5.25], + "alternating_2_5s": [0.0, 0.3, 0.6, 3.1, 3.4, 5.9, 6.2, 8.7], + "mixed_1s_3s": [0.0, 0.1, 0.2, 1.2, 1.3, 4.3, 4.4, 7.4], + } + for trace_id, starts in schedules.items(): + _write_trace(trace_dir / f"{trace_id}.json", trace_id, starts) + return trace_dir + + +def _tree_idle_gaps_seconds(result: AIPerfResults) -> list[float]: + """Return periods with no request active in each profiling trajectory tree.""" + assert result.jsonl is not None + by_root = defaultdict(list) + for record in result.jsonl: + root_id = record.metadata.root_correlation_id + assert root_id is not None + by_root[root_id].append(record.metadata) + + gaps: list[float] = [] + for records in by_root.values(): + records.sort(key=lambda record: record.request_start_ns) + latest_end_ns = records[0].request_end_ns + for record in records[1:]: + if record.request_start_ns > latest_end_ns: + gaps.append((record.request_start_ns - latest_end_ns) / 1e9) + latest_end_ns = max(latest_end_ns, record.request_end_ns) + return gaps + + +@pytest.mark.integration +@pytest.mark.asyncio +@pytest.mark.parametrize("cap_seconds", [None, 0.1, 0.25, 0.5]) +async def test_agentx_trace_idle_gap_cap_controls_replay_timing( + cli: AIPerfCLI, + aiperf_mock_server: AIPerfMockServer, + varied_weka_traces: Path, + cap_seconds: float | None, +) -> None: + """Replay varied traces through the mock server and verify the configured cap.""" + cap_arg = ( + "" if cap_seconds is None else f"--trace-idle-gap-cap-seconds {cap_seconds}" + ) + result = await cli.run( + f""" + aiperf profile + --scenario inferencex-agentx-mvp + --unsafe-override + --model mock-model + --url {aiperf_mock_server.url} + --endpoint-type chat + --streaming + --extra-inputs ignore_eos:true + --custom-dataset-type weka_trace + --input-file {varied_weka_traces} + --no-fixed-schedule + --concurrency 4 + --benchmark-duration 6 + --trajectory-start-min-ratio 0.25 + --trajectory-start-max-ratio 0.25 + --random-seed 42 + --workers-max 4 + --ui simple + {cap_arg} + """, + timeout=300.0, + ) + + assert result.exit_code == 0, ( + f"AIPerf failed; stderr=\n{result.stderr}\n\nstdout=\n{result.stdout}" + ) + assert result.jsonl is not None + source_traces = { + record.metadata.source_trace_id + for record in result.jsonl + if record.metadata.source_trace_id is not None + } + assert source_traces == { + "alternating_2s", + "alternating_1_5s", + "alternating_2_5s", + "mixed_1s_3s", + } + + gaps = _tree_idle_gaps_seconds(result) + assert len(gaps) >= 8 + if cap_seconds is None: + assert max(gaps) > 1.0 + else: + assert max(gaps) <= cap_seconds + 0.35 diff --git a/tests/integration/test_weka_flat_split_e2e.py b/tests/integration/test_weka_flat_split_e2e.py index 777e670bf5..4c5f8aa77a 100644 --- a/tests/integration/test_weka_flat_split_e2e.py +++ b/tests/integration/test_weka_flat_split_e2e.py @@ -187,6 +187,31 @@ def _write_issue_1231_trace(target_dir: Path) -> Path: return target_dir +def _write_independent_stream_drift_trace(target_dir: Path) -> Path: + """A scaled multi-stream trace whose runtime clock drift can exceed its cap. + + Every authored request start is at most 300ms after the preceding request + in the whole tree. One worker's first response is deliberately much slower + than its recorded ``api_time``, however, so its independent continuation + timer would otherwise leave the entire tree idle for over 300ms. + """ + _write_trace( + target_dir, + "trace_stream_drift", + [ + _req(0.00, [1, 2, 3], api_time=0.005, out=1), + _req(0.05, [1, 2, 10], api_time=0.005, out=1), + _req(0.10, [1, 2, 20], api_time=0.005, out=100), + _req(0.35, [1, 2, 10, 11], api_time=0.005, out=1), + _req(0.65, [1, 2, 10, 11, 12], api_time=0.005, out=1), + _req(0.95, [1, 2, 10, 11, 12, 13], api_time=0.005, out=1), + _req(1.20, [1, 2, 20, 21], api_time=0.005, out=1), + _req(1.50, [1, 2, 3, 4], api_time=0.005, out=1), + ], + ) + return target_dir + + def _write_background_trace(target_dir: Path) -> Path: """One trace whose worker chain ends after the last main turn, so the loader emits an ``is_background=True`` branch (no SPAWN_JOIN) the runtime must still send and drain cleanly.""" _write_trace( @@ -659,6 +684,71 @@ async def test_spawn_join_gates_main_turn_until_worker_chain_completes( "no complete gated play executed within the benchmark duration; " "cannot verify SPAWN_JOIN ordering" ) + branch_stats = result.json.branch_stats + assert branch_stats is not None + assert branch_stats.parents_suspended >= 1, ( + f"the parent never suspended on the gate (join was a no-op): {branch_stats}" + ) + assert branch_stats.parents_resumed >= 1, branch_stats + assert branch_stats.children_errored == 0, branch_stats + + +async def test_runtime_idle_cap_handles_independent_stream_clock_drift( + tmp_path: Path, mock_server_factory: MockServerFactory +) -> None: + """A slow child cannot recreate a tree-idle gap beyond the dataset cap.""" + corpus = _write_independent_stream_drift_trace(tmp_path / "traces") + cap_seconds = 0.3 + async with mock_server_factory(ttft=10.0, itl=5.0, workers=4) as server: + result = await _run_weka_profile( + input_dir=corpus, + artifact_dir=tmp_path / "artifacts", + url=server.url, + duration=4.0, + concurrency=1, + extra_args=[ + "--scenario", + "inferencex-agentx-mvp", + "--unsafe-override", + "--trace-idle-gap-cap-seconds", + str(cap_seconds), + "--warmup-requests-per-lane", + "1", + ], + timeout=240.0, + ) + _assert_success(result, "independent-stream runtime idle cap") + + complete_plays = [ + play + for play in _collect_plays(result) + if play.trace_id == "trace_stream_drift" + and [md.turn_index for md in play.root] == [0, 1] + and sorted(len(records) for records in play.children.values()) == [2, 4] + ] + assert complete_plays, "the profiling phase never completed the synthetic tree" + + for play in complete_plays: + records = sorted( + [*play.root, *(md for child in play.children.values() for md in child)], + key=lambda md: md.request_start_ns, + ) + latest_end_ns = records[0].request_end_ns + for current in records[1:]: + if current.request_start_ns > latest_end_ns: + gap_seconds = (current.request_start_ns - latest_end_ns) / 1e9 + assert gap_seconds <= cap_seconds + 0.15, ( + f"runtime tree was completely idle for {gap_seconds:.3f}s " + f"despite a {cap_seconds:.3f}s trace cap" + ) + latest_end_ns = max(latest_end_ns, current.request_end_ns) + + for session in (play.root, *play.children.values()): + assert [md.turn_index for md in session] == list(range(len(session))) + child_last_end = max( + child[-1].request_end_ns for child in play.children.values() + ) + assert play.root[-1].request_start_ns >= child_last_end branch_stats = result.json.branch_stats assert branch_stats is not None diff --git a/tests/unit/common/scenario/test_scenario_registry.py b/tests/unit/common/scenario/test_scenario_registry.py index 652eab23d0..56df54c1b5 100644 --- a/tests/unit/common/scenario/test_scenario_registry.py +++ b/tests/unit/common/scenario/test_scenario_registry.py @@ -36,7 +36,7 @@ def test_inferencex_agentx_mvp_registered(): assert spec.inter_turn_delay_cap_seconds is None assert spec.trace_idle_gap_cap_seconds is None assert spec.system_idle_gap_cap_seconds == 10.0 - assert spec.forbid_trace_idle_gap_cap is True + assert spec.forbid_trace_idle_gap_cap is False assert spec.forbid_inter_turn_delay_cap is True diff --git a/tests/unit/common/scenario/test_scenario_validator.py b/tests/unit/common/scenario/test_scenario_validator.py index 4a01d6c68a..ed343c5259 100644 --- a/tests/unit/common/scenario/test_scenario_validator.py +++ b/tests/unit/common/scenario/test_scenario_validator.py @@ -490,7 +490,7 @@ def test_system_idle_gap_cap_explicit_other_value_raises() -> None: assert "system-idle-gap-cap-seconds" in str(exc.value) -def test_trace_idle_gap_cap_explicit_other_value_raises() -> None: +def test_trace_idle_gap_cap_explicit_value_is_honored() -> None: run = _build_run( streaming=True, extra={"ignore_eos": True}, @@ -501,9 +501,10 @@ def test_trace_idle_gap_cap_explicit_other_value_raises() -> None: "trace_idle_gap_cap_seconds": 30.0, }, ) - with pytest.raises(ScenarioLockError) as exc: - apply_scenario(run) - assert "trace-idle-gap-cap-seconds" in str(exc.value) + outcome = apply_scenario(run) + assert outcome.violations == [] + assert outcome.submission_valid is True + assert run.cfg.get_default_dataset().trace_idle_gap_cap_seconds == 30.0 def test_inter_turn_delay_cap_shipped_scenario_forbids_value() -> None: diff --git a/tests/unit/common/scenario/test_scenario_validator_advanced_adversarial.py b/tests/unit/common/scenario/test_scenario_validator_advanced_adversarial.py index bd0c3b8c5c..d6ea075650 100644 --- a/tests/unit/common/scenario/test_scenario_validator_advanced_adversarial.py +++ b/tests/unit/common/scenario/test_scenario_validator_advanced_adversarial.py @@ -125,8 +125,8 @@ def test_ignore_eos_falsy_string_zero_violates() -> None: assert any(v.flag == "extra_inputs.ignore_eos" for v in exc_info.value.violations) -def test_trace_idle_gap_cap_explicit_old_default_is_forbidden() -> None: - """Even the old 10-second value now violates faithful trace timing.""" +def test_trace_idle_gap_cap_explicit_old_default_is_allowed() -> None: + """The prior 10-second per-trace cap remains a valid explicit choice.""" run = _build_run( streaming=True, extra={"ignore_eos": True}, @@ -137,11 +137,10 @@ def test_trace_idle_gap_cap_explicit_old_default_is_forbidden() -> None: "trace_idle_gap_cap_seconds": 10.0, }, ) - with pytest.raises(ScenarioLockError) as exc_info: - apply_scenario(run) - assert any( - v.flag == "--trace-idle-gap-cap-seconds" for v in exc_info.value.violations - ) + outcome = apply_scenario(run) + assert outcome.violations == [] + assert outcome.submission_valid is True + assert run.cfg.get_default_dataset().trace_idle_gap_cap_seconds == 10.0 def test_unsafe_override_with_no_violations_returns_submission_valid_true() -> None: diff --git a/tests/unit/common/test_loop_scheduler.py b/tests/unit/common/test_loop_scheduler.py index 89408d17fa..473dc0b30b 100644 --- a/tests/unit/common/test_loop_scheduler.py +++ b/tests/unit/common/test_loop_scheduler.py @@ -160,6 +160,74 @@ async def noop(): assert scheduler.cap_pending_delay(2.0) == 0.0 scheduler.cancel_all_pending() + async def test_cap_pending_delay_for_group_isolated_and_uniform( + self, scheduler: LoopScheduler + ): + """A trace-local cap cannot move another trace's replay schedule.""" + + async def noop(): + pass + + scheduler.schedule_later(5.0, noop(), group_id="root-a") + scheduler.schedule_later(8.0, noop(), group_id="root-a") + scheduler.schedule_later(6.0, noop(), group_id="root-b") + before_a = sorted( + handle.when() + for handle_id, (handle, _) in scheduler._handles.items() + if scheduler._handle_groups.get(handle_id) == "root-a" + ) + before_b = [ + handle.when() + for handle_id, (handle, _) in scheduler._handles.items() + if scheduler._handle_groups.get(handle_id) == "root-b" + ] + + shifted = scheduler.cap_pending_delay_for_group("root-a", 1.0) + + after_a = sorted( + handle.when() + for handle_id, (handle, _) in scheduler._handles.items() + if scheduler._handle_groups.get(handle_id) == "root-a" + ) + after_b = [ + handle.when() + for handle_id, (handle, _) in scheduler._handles.items() + if scheduler._handle_groups.get(handle_id) == "root-b" + ] + assert shifted == pytest.approx(4.0, abs=0.02) + assert after_a[1] - after_a[0] == pytest.approx(before_a[1] - before_a[0]) + assert after_b == before_b + scheduler.cancel_all_pending() + + async def test_global_cap_preserves_group_ownership(self, scheduler: LoopScheduler): + """A global idle jump cannot erase later per-trace cap ownership.""" + + async def noop(): + pass + + scheduler.schedule_later(5.0, noop(), group_id="root-a") + scheduler.schedule_later(8.0, noop(), group_id="root-b") + assert scheduler.cap_pending_delay(1.0) == pytest.approx(4.0, abs=0.02) + + groups = set(scheduler._handle_groups.values()) + assert groups == {"root-a", "root-b"} + root_a_before = next( + handle.when() + for handle_id, (handle, _) in scheduler._handles.items() + if scheduler._handle_groups.get(handle_id) == "root-a" + ) + + assert scheduler.cap_pending_delay_for_group("root-b", 0.0) == pytest.approx( + 4.0, abs=0.02 + ) + root_a_after = next( + handle.when() + for handle_id, (handle, _) in scheduler._handles.items() + if scheduler._handle_groups.get(handle_id) == "root-a" + ) + assert root_a_after == root_a_before + scheduler.cancel_all_pending() + async def test_cap_pending_delay_zero_queues_timer_as_ready_work( self, scheduler: LoopScheduler ): diff --git a/tests/unit/dataset/loader/test_weka_flat_split_parallel_adv.py b/tests/unit/dataset/loader/test_weka_flat_split_parallel_adv.py index 3ce92fad6a..23f2b4f1fc 100644 --- a/tests/unit/dataset/loader/test_weka_flat_split_parallel_adv.py +++ b/tests/unit/dataset/loader/test_weka_flat_split_parallel_adv.py @@ -267,34 +267,6 @@ def _by_sid(convs: list[Conversation]) -> dict[str, Conversation]: # Tests -def test_convert_fanout_idle_gap_warp_parallel_byte_identical(tmp_path, monkeypatch): - """Idle-gap warp redistributes the same start set after the split with byte-identical timestamps and delays across paths (spec 5.6).""" - reqs = _fanout_requests() - reqs[5]["t"] = 200.0 # main turn 3 after a 191s idle gap - reqs.append( - # third worker-chain request crossing two gaps; api_time absent (None) - _nreq(210.0, [1, 2, 50, 51, 52, 53], api_time=None) - ) - serial, _parallel = _run_both( - tmp_path, - monkeypatch, - [_trace("trace_warp", reqs)], - idle_gap_cap_seconds=5.0, - ) - - convs = _by_sid(serial) - root = convs["trace_warp"] - # Gaps: [2.5, 8.5] excess 1, [9, 200] excess 186, [200, 210] excess 5. - assert [t.timestamp for t in root.turns] == pytest.approx([0.0, 8000.0, 13000.0]) - # Delays are the end-to-start idle gaps (start-to-start minus prev api_time). - assert root.turns[1].delay == pytest.approx(7000.0) - assert root.turns[2].delay == pytest.approx(4000.0) - w0 = convs["trace_warp::fa:000"] - assert [t.timestamp for t in w0.turns] == pytest.approx([2000.0, 7500.0, 18000.0]) - assert w0.turns[1].delay == pytest.approx(0.0) # end-to-start floored at 0 - assert w0.turns[2].delay == pytest.approx(9500.0) - - def test_convert_nonmonotonic_parent_delay_floored_parallel_byte_identical( tmp_path, monkeypatch ): diff --git a/tests/unit/dataset/loader/test_weka_flat_split_serial_adv.py b/tests/unit/dataset/loader/test_weka_flat_split_serial_adv.py index a5c8168918..c6d3324205 100644 --- a/tests/unit/dataset/loader/test_weka_flat_split_serial_adv.py +++ b/tests/unit/dataset/loader/test_weka_flat_split_serial_adv.py @@ -297,87 +297,24 @@ def test_disjoint_batch_splits_into_independent_chains(caplog): ), "the nonce-poison guard was removed; no such warning should be logged" -# Idle-gap warp interaction (spec §5.6): flat-chain warped timestamps/delays -# and the warp gap structure must match the unsplit trace shifted equivalently. +# Runtime trace-idle caps must not rewrite the authored dataset timeline. -def test_idle_gap_warp_flat_chain_gap_structure_matches_unsplit_run(): - """With an idle-gap cap active, the main chain's warped timestamps match the legacy single-stream run over the same request starts (spec §5.6).""" - # Main founder at t=0, worker founder at t=20 (disjoint ns), main t2 at - # t=220. The 20->220 request-start gap (200s) caps to 60s -> 140s shift. +def test_trace_idle_gap_cap_does_not_rewrite_dataset_timestamps(): + """The trace cap is enforced only from observed runtime completions.""" requests = [ - _normal(0.0, [1, 2, 3], api_time=1.0), # main, t=0 - _normal(20.0, [900, 901], api_time=1.0, model=_HAIKU), # worker, t=20 - _normal(220.0, [1, 2, 3, 4], api_time=1.0), # main t2 -> warps to t=80 + _normal(0.0, [1, 2, 3], api_time=1.0), + _normal(20.0, [900, 901], api_time=1.0, model=_HAIKU), + _normal(220.0, [1, 2, 3, 4], api_time=1.0), ] - uc = _mk_user_config(trace_idle_gap_cap_seconds=60.0) - loader_split = _build_loader(uc) - convs_split = _convert(loader_split, _trace("flt_warp", requests)) - - uc2 = _mk_user_config(trace_idle_gap_cap_seconds=60.0) - loader_legacy = _build_loader(uc2) - import unittest.mock as _m - - with _m.patch.object(Environment.DATASET, "WEKA_SPLIT_FLATTENED_AGENTS", False): - convs_legacy = _convert(loader_legacy, _trace("flt_warp", requests)) - - root_split = convs_split["flt_warp"] - legacy = convs_legacy["flt_warp"] - # Main founder unwarped (gap is after it). - assert root_split.turns[0].timestamp == 0.0 - # Worker request start (t=20) is also <= cap boundary so it is unwarped; - # the legacy single stream sees it at t=20 too. The 200s gap 20->220 - # collapses to 60s, shifting t=220 -> t=80 in BOTH runs. - assert root_split.turns[1].timestamp == pytest.approx(80_000.0) - # Legacy keeps all 3 as one conversation; its turn at t=220 also warps to - # 80s. The main chain's warped end timestamp must agree across runs. - legacy_220 = next(t for t in legacy.turns if t.timestamp == pytest.approx(80_000.0)) - assert legacy_220.timestamp == pytest.approx(root_split.turns[1].timestamp) - # Worker conversation's single turn warps to t=20 (20_000 ms). - worker = convs_split["flt_warp::fa:000"] - assert worker.turns[0].timestamp == pytest.approx(20_000.0) - - -def test_idle_gap_warp_flat_branch_start_uses_mapped_time_not_raw(): - """A flat group whose workers begin after a compressed idle gap anchors its SPAWN branch on the warped first-request time, preserving the inter-worker dispatch stagger.""" - # Main founder at t=0; a 1000s idle gap; two disjoint-namespace workers at - # t=1000 and t=1002 (2s apart); main t1 at t=1003. Sorted request starts - # [0, 1000, 1002, 1003]: the 0->1000 gap (1000s) caps to 60s (940s excess), - # shifting everything at/after the gap left by 940s: - # worker A 1000 -> 60s, worker B 1002 -> 62s, main t1 1003 -> 63s. - requests = [ - _normal(0.0, [1, 2, 3], api_time=0.5), # main founder, outer 0 - _normal(1000.0, [900, 901], api_time=0.5, model=_HAIKU), # worker A - _normal(1002.0, [910, 911], api_time=0.5, model=_HAIKU), # worker B - _normal(1003.0, [1, 2, 3, 4], api_time=0.5), # main t1 - ] - uc = _mk_user_config(trace_idle_gap_cap_seconds=60.0) - loader = _build_loader(uc) - convs = _convert(loader, _trace("flt_warp_start", requests)) - root = convs["flt_warp_start"] - - flatspawn = [b for b in root.branches if "flatspawn" in b.branch_id] - assert len(flatspawn) == 1, "both workers share (preceding, join) -> one branch" - branch = flatspawn[0] - assert len(branch.child_conversation_ids) == 2 - - # Branch start is the WARPED min worker start (60s), NOT the raw 1000s. - assert branch.start_timestamp_ms == pytest.approx(60_000.0) - - # Worker turn-0 timestamps are warped (60s, 62s); the per-worker dispatch - # offset from the branch start stays non-negative AND preserves the recorded - # 2s stagger instead of collapsing both to 0. - child_ts = sorted( - convs[sid].turns[0].timestamp for sid in branch.child_conversation_ids - ) - assert child_ts == [pytest.approx(60_000.0), pytest.approx(62_000.0)] - offsets = sorted(ts - branch.start_timestamp_ms for ts in child_ts) - assert offsets[0] == pytest.approx(0.0) - assert offsets[1] == pytest.approx(2_000.0) - - -# Timing: per-chain delays, think_time_only and ignore_delays on flat-chain -# turns (spec §5.6). Delays never negative. + loader = _build_loader(_mk_user_config(trace_idle_gap_cap_seconds=60.0)) + convs = _convert(loader, _trace("runtime_cap_only", requests)) + + root = convs["runtime_cap_only"] + worker = convs["runtime_cap_only::fa:000"] + assert [turn.timestamp for turn in root.turns] == [0.0, 220_000.0] + assert root.turns[1].delay == 219_000.0 + assert worker.turns[0].timestamp == 20_000.0 def test_flat_chain_delays_are_within_chain_and_nonnegative(): diff --git a/tests/unit/dataset/loader/test_weka_pathological.py b/tests/unit/dataset/loader/test_weka_pathological.py index dd295aae42..cf62e78629 100644 --- a/tests/unit/dataset/loader/test_weka_pathological.py +++ b/tests/unit/dataset/loader/test_weka_pathological.py @@ -14,7 +14,6 @@ from aiperf.dataset.loader.weka_trace import ( WekaTraceLoader, _expand_subagent_to_child_plans, - _IdleGapTimeWarp, _sa_end_seconds, ) from aiperf.dataset.loader.weka_trace_models import ( @@ -158,74 +157,6 @@ def _inner_request(**overrides) -> WekaNormalRequest: # Regression: idle-gap-mapped subagent spawn time (fixed) -def test_idle_gap_branch_start_timestamp_uses_mapped_time_not_raw(tmp_path): - """SPAWN ``start_timestamp_ms`` must live on the same mapped timeline as every other turn, never exceeding the maximum mapped turn timestamp.""" - trace = _base_trace( - [ - _normal(0.0, [1]), - _normal(1000.0, [1, 2]), # 1000s start-gap -> compressed - _subagent(1005.0, "a", inner_hash_ids=[8]), - _normal(1006.0, [1, 2, 3]), - ], - trace_id="idle_branch", - ) - path = tmp_path / "t.json" - path.write_text(json.dumps(trace)) - uc = _mk_user_config(trace_idle_gap_cap_seconds=60.0) - loader = _make_loader(path, uc) - - convs = loader.convert_to_conversations(loader.load_dataset()) - root = next(c for c in convs if c.session_id == "idle_branch") - child = next(c for c in convs if c.session_id == "idle_branch::sa:a") - - max_mapped_ms = max(t.timestamp for t in root.turns) - branch = root.branches[0] - # The branch entered the timeline when the subagent spawned; on the mapped - # timeline that is the child's first-request timestamp, never ~940s later. - assert branch.start_timestamp_ms <= max_mapped_ms - assert branch.start_timestamp_ms == child.turns[0].timestamp - - -def test_parallel_subagent_payload_carries_mapped_spawn_time(tmp_path): - """The parallel marker payload must carry the mapped spawn time via ``effective_t``, not revert to raw seconds when an idle-gap warp shifts the timeline.""" - trace = _base_trace( - [ - _normal(0.0, [1]), - _normal(1000.0, [1, 2]), - _subagent(1005.0, "a", inner_hash_ids=[8]), - _normal(1006.0, [1, 2, 3]), - ], - trace_id="idle_parallel", - ) - path = tmp_path / "t.json" - path.write_text(json.dumps(trace)) - uc = _mk_user_config(trace_idle_gap_cap_seconds=60.0) - loader = _make_loader(path, uc) - - data = loader.load_dataset() - plans = loader._build_reconstruction_plans(data) - parent_plans, child_plans = plans.parent_plans, plans.child_plans - timing = loader._build_trace_idle_timing_by_trace(parent_plans, child_plans) - metric_values = loader._build_shared_metric_values( - parent_plans, child_plans, plans.flat_plans - ) - tasks = loader._build_parallel_reconstruction_tasks( - parent_plans=parent_plans, - child_plans=child_plans, - data=data, - ignore_delays=False, - think_time_only=False, - cap_seconds=None, - model_map_per_trace={"idle_parallel": {}}, - trace_idle_timing_by_trace=timing, - metric_values_by_trace=metric_values, - ) - _, marker = tasks[0].parent["subagents"][0] - # The mapped end time is plumbed through; the mapped spawn time must be too. - assert "effective_t" in marker - assert marker["effective_t"] != marker["t"] - - def test_sa_end_seconds_negative_duration_not_before_spawn(): """A subagent's recorded end time can never precede its own spawn, even with a corrupt negative ``duration_ms``.""" entry = _make_subagent_entry(t=10.0, duration_ms=-5000) @@ -264,24 +195,6 @@ def test_think_time_only_negative_think_time_not_negative_delay(tmp_path): # PASSING CHARACTERIZATIONS (surprising but intended / not invariant-breaking) -def test_idle_gap_exactly_equal_to_cap_is_not_compressed(): - """A request-start gap exactly equal to the cap is left untouched (``_IdleGapTimeWarp`` compresses only strict ``gap_seconds > cap_seconds``).""" - warp = _IdleGapTimeWarp([0.0, 60.0], cap_seconds=60.0) - assert warp.map(60.0) == 60.0 - # One microsecond over the cap does get compressed back to the boundary. - warp_over = _IdleGapTimeWarp([0.0, 60.001], cap_seconds=60.0) - assert warp_over.map(60.001) == pytest.approx(60.0) - - -def test_idle_gap_collapsed_tail_event_maps_to_cap_boundary(): - """A non-request event inside a collapsed gap tail is pinned to ``raw_start + cap`` so a join cannot wait past the next shifted request.""" - warp = _IdleGapTimeWarp([0.0, 20.0, 220.0], cap_seconds=60.0) - assert warp.map(80.0) == pytest.approx(80.0) # at the boundary - assert warp.map(150.0) == pytest.approx(80.0) # deep in the collapsed tail - assert warp.map(220.0) == pytest.approx(80.0) # the gap end - assert warp.map(300.0) == pytest.approx(160.0) # after: shifted left by excess - - def test_nested_chain_nan_api_time_treated_as_zero_duration(): """A NaN inner ``api_time`` is clamped to zero duration in chain detection, so a same-context continuation still extends the chain.""" entry = _make_subagent_entry( diff --git a/tests/unit/dataset/loader/test_weka_tool_shaped_messages.py b/tests/unit/dataset/loader/test_weka_tool_shaped_messages.py index 45d01ec546..0e51d51bba 100644 --- a/tests/unit/dataset/loader/test_weka_tool_shaped_messages.py +++ b/tests/unit/dataset/loader/test_weka_tool_shaped_messages.py @@ -323,7 +323,6 @@ def test_parallel_tool_shaping_matches_serial(tool_shaped_env, tmp_path): cap_seconds=None, t_start=0.0, model_map_per_trace=model_maps, - trace_idle_timing_by_trace={}, metric_values_by_trace=loader._build_shared_metric_values( parent_plans, child_plans, plans.flat_plans ), diff --git a/tests/unit/dataset/loader/test_weka_tool_turn_detection.py b/tests/unit/dataset/loader/test_weka_tool_turn_detection.py index 9528f82466..82bca5963b 100644 --- a/tests/unit/dataset/loader/test_weka_tool_turn_detection.py +++ b/tests/unit/dataset/loader/test_weka_tool_turn_detection.py @@ -259,7 +259,6 @@ def test_parallel_reconstruction_input_kind_matches_serial(tmp_path): cap_seconds=None, t_start=0.0, model_map_per_trace=model_maps, - trace_idle_timing_by_trace={}, metric_values_by_trace=loader._build_shared_metric_values( parent_plans, child_plans, plans.flat_plans ), diff --git a/tests/unit/dataset/loader/test_weka_trace.py b/tests/unit/dataset/loader/test_weka_trace.py index 3159d2cb47..380fe36769 100644 --- a/tests/unit/dataset/loader/test_weka_trace.py +++ b/tests/unit/dataset/loader/test_weka_trace.py @@ -365,83 +365,6 @@ def test_weka_zero_request_subagent_branch_targets_empty_child(tmp_path): assert child.agent_depth == 1 -def test_weka_trace_idle_gap_cap_ignores_dropped_orphan_subagent(tmp_path): - model = "claude-opus-4-5-20251101" - child_model = "claude-haiku-4-5-20251001" - - def normal(t: float, hash_ids: list[int], input_length: int) -> dict: - return { - "t": t, - "type": "n", - "model": model, - "in": input_length, - "out": 10, - "hash_ids": hash_ids, - "input_types": ["text"], - "output_types": ["text"], - "stop": "end_turn", - "api_time": 1.0, - "think_time": 0.0, - } - - trace = { - "id": "orphan_gap", - "models": [model, child_model], - "block_size": 64, - "hash_id_scope": "local", - "requests": [ - { - "t": 10.0, - "type": "subagent", - "agent_id": "orphan", - "subagent_type": "Explore", - "duration_ms": 1000, - "total_tokens": 10, - "tool_use_count": 1, - "status": "completed", - "requests": [ - { - "t": 10.0, - "type": "n", - "model": child_model, - "in": 64, - "out": 10, - "hash_ids": [8], - "input_types": ["text"], - "output_types": ["text"], - "stop": "end_turn", - "api_time": 1.0, - "think_time": 0.0, - } - ], - "models": [child_model], - "tool_tokens": 0, - "system_tokens": 0, - }, - normal(1000.0, [1], 64), - normal(2000.0, [1, 2], 128), - ], - } - path = tmp_path / "orphan_gap.json" - path.write_text(json.dumps(trace)) - uc = _mk_user_config(trace_idle_gap_cap_seconds=60.0) - loader = WekaTraceLoader(filename=str(path), run=uc) - _stub_prompt_generator_for_reconstructor(loader) - loader._tokenizer_name = "t" - loader._trust_remote_code = False - loader._tokenizer_revision = None - loader._block_size = 64 - - conversations = loader.convert_to_conversations(loader.load_dataset()) - - root = next(c for c in conversations if c.session_id == "orphan_gap") - assert [c.session_id for c in conversations] == ["orphan_gap"] - assert root.turns[0].timestamp == 1000.0 * 1000.0 - assert root.turns[1].timestamp == 1060.0 * 1000.0 - # end-to-start idle gap: 60.0s start-to-start minus 1.0s prev api_time. - assert root.turns[1].delay == 59.0 * 1000.0 - - def test_weka_parallel_child_conversation_metadata_is_non_root(monkeypatch): import aiperf.dataset.loader.weka_parallel_convert as parallel_convert @@ -489,7 +412,6 @@ def fake_run_parallel_weka_reconstruction(*args, **kwargs): configured_workers=1, t_start=0.0, model_map_per_trace={"trace_sa": {}}, - trace_idle_timing_by_trace={}, metric_values_by_trace=loader._build_shared_metric_values( parent_plans, child_plans ), @@ -916,126 +838,6 @@ def test_use_think_time_only_emits_recorded_think_time_as_delay(monkeypatch, tmp assert turns[2].delay == 9000.0 -def test_trace_idle_gap_cap_is_per_trace_and_uses_request_starts(tmp_path): - """Idle-gap capping uses parent+subagent request starts per root trace, so each trace compresses against its own gaps rather than globally.""" - - def normal( - *, - t: float, - in_tokens: int, - out_tokens: int, - hash_ids: list[int], - api_time: float, - think_time: float = 0.0, - model: str = "claude-opus-4-5-20251101", - ) -> dict: - return { - "t": t, - "type": "n", - "model": model, - "in": in_tokens, - "out": out_tokens, - "hash_ids": hash_ids, - "input_types": ["text"], - "output_types": ["text"], - "stop": "end_turn", - "api_time": api_time, - "think_time": think_time, - } - - trace_a = { - "id": "trace_idle_a", - "models": ["claude-opus-4-5-20251101", "claude-haiku-4-5-20251001"], - "block_size": 64, - "hash_id_scope": "local", - "requests": [ - normal(t=0.0, in_tokens=100, out_tokens=10, hash_ids=[1], api_time=10.0), - { - "t": 20.0, - "type": "subagent", - "agent_id": "agent_idle", - "subagent_type": "Explore", - "duration_ms": 80_000, - "total_tokens": 500, - "tool_use_count": 1, - "status": "completed", - "requests": [ - normal( - t=20.0, - in_tokens=80, - out_tokens=20, - hash_ids=[10], - api_time=80.0, - model="claude-haiku-4-5-20251001", - ) - ], - "models": ["claude-haiku-4-5-20251001"], - "tool_tokens": 0, - "system_tokens": 0, - }, - normal( - t=220.0, - in_tokens=200, - out_tokens=20, - hash_ids=[1, 2], - api_time=5.0, - think_time=999.0, - ), - ], - } - trace_b = { - "id": "trace_idle_b", - "models": ["claude-opus-4-5-20251101"], - "block_size": 64, - "hash_id_scope": "local", - "requests": [ - normal(t=150.0, in_tokens=100, out_tokens=10, hash_ids=[3], api_time=1.0), - normal( - t=220.0, in_tokens=150, out_tokens=10, hash_ids=[3, 4], api_time=1.0 - ), - ], - } - traces_dir = tmp_path / "traces" - traces_dir.mkdir() - (traces_dir / "a.json").write_text(json.dumps(trace_a)) - (traces_dir / "b.json").write_text(json.dumps(trace_b)) - - uc = _mk_user_config( - use_think_time_only=True, - inter_turn_delay_cap_seconds=60.0, - trace_idle_gap_cap_seconds=60.0, - ) - loader = WekaTraceLoader(filename=str(traces_dir), run=uc) - _stub_prompt_generator_for_reconstructor(loader) - loader._tokenizer_name = "t" - loader._trust_remote_code = False - loader._tokenizer_revision = None - loader._block_size = 64 - - convs = loader.convert_to_conversations(loader.load_dataset()) - conv_by_id = {conv.session_id: conv for conv in convs} - - trace_a_turns = conv_by_id["trace_idle_a"].turns - assert trace_a_turns[0].timestamp == 0.0 - assert trace_a_turns[1].timestamp == 80_000.0 - # The trace-wide idle-gap cap takes precedence over the old per-turn cap, - # so the start-to-start gap stays 80s (not clamped to 60s); the emitted - # delay is the end-to-start idle gap, 80s minus 10.0s prev api_time = 70s. - assert trace_a_turns[1].delay == 70_000.0 - assert conv_by_id["trace_idle_a::sa:agent_idle"].turns[0].timestamp == 20_000.0 - - # Trace B is compressed against its own request starts only: 150 -> 220 - # becomes 150 -> 210 after the 70s start gap is capped to 60s. - trace_b_turns = conv_by_id["trace_idle_b"].turns - assert trace_b_turns[0].timestamp == 150_000.0 - assert trace_b_turns[1].timestamp == 210_000.0 - - -# These hash_id_scope tests wire the REAL scope-sensitive HashIdRandomGenerator -# (not stub_hash_id_corpus_rng, which ignores set_trace_id) so a per-child scope -# regression is detectable: local scope means one hash_id namespace per trace file. - - def _wire_real_scope_rng(loader, *, block_size: int, seed: int = 1234) -> None: """Wire a MagicMock prompt_generator backed by the real scope-sensitive RNG, with a token-reflecting ``tokenizer.decode`` so raw_messages track decoded tokens.""" pg = MagicMock() diff --git a/tests/unit/dataset/loader/test_weka_trace_parallel.py b/tests/unit/dataset/loader/test_weka_trace_parallel.py index 9d0db702c9..ab99716345 100644 --- a/tests/unit/dataset/loader/test_weka_trace_parallel.py +++ b/tests/unit/dataset/loader/test_weka_trace_parallel.py @@ -284,7 +284,6 @@ def test_parallel_byte_equivalence_simple_fixture(tmp_path): model_map_per_trace={ tid: serial_loader._build_model_map(wekas[0]) for tid, wekas in data.items() }, - trace_idle_timing_by_trace={}, metric_values_by_trace=serial_loader._build_shared_metric_values( parent_plans, child_plans ), @@ -393,7 +392,6 @@ def test_parallel_byte_equivalence_with_subagent(tmp_path): model_map_per_trace={ tid: serial_loader._build_model_map(wekas[0]) for tid, wekas in data.items() }, - trace_idle_timing_by_trace={}, metric_values_by_trace=serial_loader._build_shared_metric_values( parent_plans, child_plans ), @@ -535,7 +533,6 @@ def test_directory_with_multiple_traces_parallel_path_byte_exact(tmp_path): model_map_per_trace={ tid: serial_loader._build_model_map(wekas[0]) for tid, wekas in data.items() }, - trace_idle_timing_by_trace={}, metric_values_by_trace=serial_loader._build_shared_metric_values( parent_plans, child_plans ), diff --git a/tests/unit/regression/test_agentic_hotpath_invariants.py b/tests/unit/regression/test_agentic_hotpath_invariants.py index 24817c11ae..da2ca68a4b 100644 --- a/tests/unit/regression/test_agentic_hotpath_invariants.py +++ b/tests/unit/regression/test_agentic_hotpath_invariants.py @@ -7,11 +7,9 @@ import asyncio import time from collections import defaultdict -from multiprocessing import shared_memory from pathlib import Path -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock -import numpy as np import pytest from aiperf.common.enums import CacheBustTarget, ConversationBranchMode, CreditPhase @@ -23,7 +21,6 @@ from aiperf.credit.messages import CreditReturn from aiperf.credit.sticky_router import StickyCreditRouter, _StickyEntry from aiperf.credit.structs import Credit -from aiperf.dataset.loader import weka_parallel_convert as wpc from aiperf.metrics.theoretical_prefix_cache import TheoreticalPrefixCacheAccumulator from aiperf.plugin import plugins from aiperf.plugin.enums import AccumulatorType, PluginType @@ -324,92 +321,6 @@ async def test_phase_keys_lane_credit_uses_runtime_index_not_enum() -> None: assert concurrency.release_session_slot.call_args.args[0] == phase_index -def test_parent_floor_process_task_nan_delay_stays_none() -> None: - """``_process_task`` must not ``max(None, 0.0)`` after NaN clamp → None.""" - corpus = np.arange(1024, dtype=np.int32) - shm = shared_memory.SharedMemory(create=True, size=corpus.nbytes, name=None) - try: - np.ndarray((len(corpus),), dtype=np.int32, buffer=shm.buf)[:] = corpus - tok = MagicMock() - tok.decode.side_effect = lambda toks: "x" * max(len(toks), 1) - - with patch( - "aiperf.dataset.loader.weka_parallel_convert.Tokenizer.from_pretrained", - return_value=tok, - ): - wpc._init_worker( - wpc._WekaWorkerInitArgs( - shm_name=shm.name, - corpus_len=len(corpus), - tokenizer_name="test-tok", - base_seed=0, - block_size=16, - bpe_stable_terminator_tokens=[], - ) - ) - - def _req( - *, - outer: int, - t: float, - delay: float | None, - hashes: list[int], - in_len: int, - ) -> tuple[int, dict]: - return ( - outer, - { - "hash_ids": hashes, - "input_length": in_len, - "output_length": 1, - "model": "m", - "t": t, - "think_time": None, - "capped_output_length": 1, - "theoretical_hit_blocks": 0, - "theoretical_total_blocks": len(hashes), - "effective_t": t, - "effective_delay_ms": delay, - }, - ) - - task = wpc._WekaTraceTask( - trace_id="nan-delay", - parent={ - "normals": [ - _req(outer=0, t=0.0, delay=None, hashes=[1], in_len=16), - _req( - outer=1, - t=1.0, - delay=float("nan"), - hashes=[1, 2], - in_len=32, - ), - ], - "subagents": [], - "tool_tokens": 0, - "system_tokens": 0, - }, - children=[], - cap_seconds=60.0, - ignore_delays=False, - think_time_only=False, - model_map={"m": "m"}, - block_size=16, - ) - result = wpc._process_task(task) - delay = result["parent_turns"][1]["delay"] - assert delay is None, ( - "NaN effective_delay_ms must stay None after " - f"clamp+floor, got {delay!r} (ungated max(None, 0.0) TypeErrors)" - ) - assert result["non_finite_count"] == 1 - finally: - wpc._worker_state = None - shm.close() - shm.unlink() - - def test_build_dataset_does_not_call_apply_file_block_size( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/tests/unit/timing/strategies/test_agentic_replay.py b/tests/unit/timing/strategies/test_agentic_replay.py index 7d87959a68..c5693ec8a2 100644 --- a/tests/unit/timing/strategies/test_agentic_replay.py +++ b/tests/unit/timing/strategies/test_agentic_replay.py @@ -715,7 +715,7 @@ def test_cache_warmup_handoff_preserves_residual_next_turn_delay() -> None: assert state.next_dispatch_offset_ms == pytest.approx(5_500.0) -def test_cache_warmup_handoff_uses_timestamp_fallback_and_idle_cap() -> None: +def test_cache_warmup_handoff_uses_timestamp_fallback_without_per_stream_cap() -> None: dataset = DatasetMetadata( conversations=[ ConversationMetadata( @@ -735,7 +735,6 @@ def test_cache_warmup_handoff_uses_timestamp_fallback_and_idle_cap() -> None: dataset=dataset, cache_warmup_duration=10.0, ) - strategy._phase_offset_cap_ms = 3_000.0 credit = _make_credit( conversation_id="trace_0", x_correlation_id="root", @@ -750,8 +749,71 @@ def test_cache_warmup_handoff_uses_timestamp_fallback_and_idle_cap() -> None: states_by_lane = strategy._build_handoff_states(finalized_at_ns=1_250_000_000) # Timestamp fallback gives 6000 - (1000 + 500) = 4500ms, less the 250ms - # handoff drain wait = 4250ms, then capped to the scenario idle cap. - assert states_by_lane[0][0].next_dispatch_offset_ms == pytest.approx(3_000.0) + # handoff drain wait = 4250ms. The raw stream offset stays intact. + assert states_by_lane[0][0].next_dispatch_offset_ms == pytest.approx(4_250.0) + + +def test_cache_warmup_handoff_preserves_raw_stream_offsets() -> None: + """The runtime watchdog does not rewrite handoff offsets.""" + dataset = DatasetMetadata( + conversations=[ + ConversationMetadata( + conversation_id="trace_0", + turns=[TurnMetadata(delay_ms=0.0), TurnMetadata(delay_ms=820_009.0)], + ), + ConversationMetadata( + conversation_id="trace_0::fa:000", + turns=[ + TurnMetadata(delay_ms=0.0), + TurnMetadata(delay_ms=9_356_213.0), + ], + is_root=False, + agent_depth=1, + parent_conversation_id="trace_0", + ), + ], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + ) + strategy, _, _, _ = _make_strategy( + phase=CreditPhase.WARMUP, + trajectories=[Trajectory(conversation_id="trace_0", start_turn_index=0)], + dataset=dataset, + cache_warmup_requests_per_lane=1, + ) + root = _make_credit( + conversation_id="trace_0", + x_correlation_id="root", + turn_index=0, + num_turns=2, + phase=CreditPhase.WARMUP, + ) + child = _make_credit( + conversation_id="trace_0::fa:000", + x_correlation_id="child", + turn_index=0, + num_turns=2, + phase=CreditPhase.WARMUP, + agent_depth=1, + parent_correlation_id="root", + root_correlation_id="root", + branch_mode=ConversationBranchMode.SPAWN, + ) + strategy._handoff_credits = {"root": root, "child": child} + strategy._handoff_returned_at_ns = { + "root": 1_000_000_000, + "child": 1_000_000_000, + } + strategy._root_to_lane["root"] = 0 + + states = strategy._build_handoff_states(finalized_at_ns=144_259_000_000)[0] + offsets = { + state.x_correlation_id: state.next_dispatch_offset_ms for state in states + } + + assert offsets == { + "root": pytest.approx(676_750.0), + "child": pytest.approx(9_212_954.0), + } def test_cache_warmup_handoff_elapsed_wait_can_exhaust_delay() -> None: @@ -987,8 +1049,8 @@ async def capture(turn): issuer.issue_credit.side_effect = capture scheduled: list[tuple[float, object]] = [] scheduler = MagicMock() - scheduler.schedule_later.side_effect = lambda delay, coro: scheduled.append( - (delay, coro) + scheduler.schedule_later.side_effect = ( + lambda delay, coro, **_kwargs: scheduled.append((delay, coro)) ) lifecycle = MagicMock() lifecycle.is_sending_complete = False @@ -1075,7 +1137,7 @@ async def capture(turn): issuer.issue_credit.side_effect = capture scheduled: list[float] = [] scheduler = MagicMock() - scheduler.schedule_later.side_effect = lambda d, c: scheduled.append(d) + scheduler.schedule_later.side_effect = lambda d, c, **_kwargs: scheduled.append(d) lifecycle = MagicMock() lifecycle.is_sending_complete = False @@ -1090,9 +1152,7 @@ async def capture(turn): credit_issuer=issuer, lifecycle=lifecycle, ) - # 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 + # Warmup priming honors the separate global system-idle guard. strategy._system_idle_gap_cap_seconds = 60.0 await strategy.setup_phase() @@ -1328,8 +1388,8 @@ async def capture(turn): issuer = AsyncMock() issuer.issue_credit.side_effect = capture scheduler = MagicMock() - scheduler.schedule_later.side_effect = lambda _delay, coro: asyncio.create_task( - coro + scheduler.schedule_later.side_effect = ( + lambda _delay, coro, **_kwargs: asyncio.create_task(coro) ) branch_orchestrator = MagicMock() @@ -1461,7 +1521,7 @@ async def capture(turn): issuer.issue_credit.side_effect = capture scheduled: list[tuple[float, object]] = [] - def fake_schedule_later(delay, coro): + def fake_schedule_later(delay, coro, **_kwargs): scheduled.append((delay, coro)) scheduler = MagicMock() @@ -1579,8 +1639,8 @@ async def capture(turn): issuer.issue_credit.side_effect = capture scheduled: list[tuple[float, object]] = [] scheduler = MagicMock() - scheduler.schedule_later.side_effect = lambda delay, coro: scheduled.append( - (delay, coro) + scheduler.schedule_later.side_effect = ( + lambda delay, coro, **_kwargs: scheduled.append((delay, coro)) ) cfg = MagicMock() @@ -1595,19 +1655,15 @@ async def capture(turn): lifecycle=MagicMock(), branch_orchestrator=MagicMock(), ) - strategy._phase_offset_cap_ms = 60_000.0 # agentx idle-gap cap of 60s - await strategy.setup_phase() await strategy.execute_phase() - # Leading idle (t* -> earliest child, 100s) capped to 60s via a 40s uniform - # shift; subagents keep their recorded 30s and 90s spacing -- NOT collapsed - # to a single 60s/60s/60s instant. + # The runtime-only trace cap does not alter phase-start offsets. assert issued == [] assert [delay for delay, _ in scheduled] == [ - pytest.approx(60.0), - pytest.approx(90.0), - pytest.approx(180.0), + pytest.approx(100.0), + pytest.approx(130.0), + pytest.approx(220.0), ] for _, coro in scheduled: await coro @@ -1646,14 +1702,11 @@ def _traj(cid: str, offsets: list[float]) -> Trajectory: strategy = AgenticReplayStrategy.__new__(AgenticReplayStrategy) strategy.conversation_source = src strategy._burst_phase_starts = False - strategy._phase_offset_cap_ms = 60_000.0 - assert strategy._profiling_spread_seconds() == pytest.approx(10.0) - # An idle-at-t* trajectory's first request is capped to the idle-gap cap, - # so the spread stays bounded by the cap rather than the raw leading idle. + # Runtime idle enforcement does not rewrite phase-start scheduling. src.trajectories.append(_traj("t2", [200_000.0, 9_000_000.0])) - assert strategy._profiling_spread_seconds() == pytest.approx(60.0) + assert strategy._profiling_spread_seconds() == pytest.approx(200.0) @pytest.mark.asyncio @@ -1702,8 +1755,8 @@ async def capture(turn): issuer.issue_credit.side_effect = capture scheduled: list[tuple[float, object]] = [] scheduler = MagicMock() - scheduler.schedule_later.side_effect = lambda delay, coro: scheduled.append( - (delay, coro) + scheduler.schedule_later.side_effect = ( + lambda delay, coro, **_kwargs: scheduled.append((delay, coro)) ) cfg = MagicMock() diff --git a/tests/unit/timing/test_factories.py b/tests/unit/timing/test_factories.py index 613e4c2007..9e411559a3 100644 --- a/tests/unit/timing/test_factories.py +++ b/tests/unit/timing/test_factories.py @@ -65,7 +65,7 @@ def cancel(self) -> None: class MockSched: tasks: list = field(default_factory=list) - def schedule_later(self, delay: float, coro) -> None: + def schedule_later(self, delay: float, coro, **_kwargs) -> None: self.tasks.append((delay, coro)) def cancel_all(self) -> None: diff --git a/tests/unit/timing/test_replay_barrier_coordinator.py b/tests/unit/timing/test_replay_barrier_coordinator.py index 4422d4d8f1..4028b12855 100644 --- a/tests/unit/timing/test_replay_barrier_coordinator.py +++ b/tests/unit/timing/test_replay_barrier_coordinator.py @@ -2,10 +2,12 @@ # SPDX-License-Identifier: Apache-2.0 import asyncio +from unittest.mock import MagicMock import pytest from aiperf.common.enums import CreditPhase +from aiperf.common.loop_scheduler import LoopScheduler from aiperf.common.models import ( ConversationMetadata, DatasetMetadata, @@ -134,6 +136,100 @@ async def test_runtime_roots_are_independent() -> None: assert issued == ["one:d"] +@pytest.mark.asyncio +@pytest.mark.parametrize("cap_seconds", [0.0, 0.01]) +async def test_idle_watchdog_advances_only_the_completed_runtime_root( + cap_seconds: float, +) -> None: + """A fully idle tree advances its own timers without touching another tree.""" + scheduler = MagicMock() + advanced = asyncio.Event() + + def advance(*_args) -> float: + advanced.set() + return 4.5 + + scheduler.cap_pending_delay_for_group.side_effect = advance + coordinator = ReplayBarrierCoordinator( + _metadata(), + scheduler=scheduler, + root_idle_gap_cap_seconds=cap_seconds, + ) + coordinator.activate() + + coordinator.observe_issued(_credit("a", "one")) + coordinator.observe_issued(_credit("a", "two")) + coordinator.complete(_credit("a", "one")) + await asyncio.wait_for(advanced.wait(), timeout=0.2) + + scheduler.cap_pending_delay_for_group.assert_called_once_with("one", 0.0) + + +@pytest.mark.asyncio +async def test_idle_watchdog_covers_initial_profiling_idle() -> None: + """A root with only future timers is idle before its first request.""" + scheduler = MagicMock() + advanced = asyncio.Event() + + def advance(*_args) -> float: + advanced.set() + return 4.5 + + scheduler.cap_pending_delay_for_group.side_effect = advance + coordinator = ReplayBarrierCoordinator( + _metadata(), + scheduler=scheduler, + root_idle_gap_cap_seconds=0.01, + ) + coordinator.activate() + + coordinator.observe_idle_root("one") + await asyncio.wait_for(advanced.wait(), timeout=0.2) + + scheduler.cap_pending_delay_for_group.assert_called_once_with("one", 0.0) + + +@pytest.mark.asyncio +async def test_initial_idle_watchdog_advances_real_group_timer() -> None: + scheduler = LoopScheduler() + coordinator = ReplayBarrierCoordinator( + _metadata(), + scheduler=scheduler, + root_idle_gap_cap_seconds=0.02, + ) + coordinator.activate() + fired = asyncio.Event() + + async def mark_fired() -> None: + fired.set() + + scheduler.schedule_later(0.2, mark_fired(), group_id="one") + started = asyncio.get_running_loop().time() + coordinator.observe_idle_root("one") + await asyncio.wait_for(fired.wait(), timeout=0.1) + + assert asyncio.get_running_loop().time() - started < 0.08 + + +@pytest.mark.asyncio +async def test_new_request_cancels_runtime_root_idle_watchdog() -> None: + """Overlapping work prevents an idle cap from advancing replay timers.""" + scheduler = MagicMock() + coordinator = ReplayBarrierCoordinator( + _metadata(), + scheduler=scheduler, + root_idle_gap_cap_seconds=0.05, + ) + coordinator.activate() + + coordinator.observe_issued(_credit("a")) + coordinator.complete(_credit("a")) + coordinator.observe_issued(_credit("b")) + await asyncio.sleep(0.06) + + scheduler.cap_pending_delay_for_group.assert_not_called() + + @pytest.mark.asyncio async def test_scalar_peak_would_slip_d_after_only_one_completion() -> None: coordinator = ReplayBarrierCoordinator(_metadata()) From b60d3a9a70c5b190a1c4c551c7d731a423a95f65 Mon Sep 17 00:00:00 2001 From: Anthony Casagrande Date: Thu, 30 Jul 2026 14:48:15 -0700 Subject: [PATCH 5/9] fix(warmup): stand down warmup prefix when cache-bust is active (#1228) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When cache-bust owns trajectory prefix isolation, preserve the same prefix across warmup and profiling so warmup primes the cache used by profiling. Upstream: 9b5f5f7282db3bf0e1e36a65771c76485b519ac2 中文：启用 cache-bust 时，由其负责轨迹前缀隔离，并在预热与性能测试阶段保持相同前缀，使预热能够填充性能测试实际复用的缓存。 Signed-off-by: Anthony Casagrande Signed-off-by: Cam Quilici --- src/aiperf/endpoints/openai_completions.py | 10 +++- src/aiperf/workers/worker.py | 25 +++++++- tests/unit/endpoints/conftest.py | 4 +- .../endpoints/test_completions_endpoint.py | 49 ++++++++++++++- tests/unit/workers/test_worker.py | 60 ++++++++++++++++++- 5 files changed, 142 insertions(+), 6 deletions(-) diff --git a/src/aiperf/endpoints/openai_completions.py b/src/aiperf/endpoints/openai_completions.py index e2f54dfe10..3f2977052c 100644 --- a/src/aiperf/endpoints/openai_completions.py +++ b/src/aiperf/endpoints/openai_completions.py @@ -4,7 +4,7 @@ from __future__ import annotations from aiperf.common.constants import WARMUP_SYSTEM_MESSAGE_PREFIX -from aiperf.common.enums import CreditPhase +from aiperf.common.enums import CacheBustTarget, CreditPhase from aiperf.common.models import ( BaseResponseData, InferenceServerResponse, @@ -39,7 +39,13 @@ def format_payload(self, request_info: RequestInfo) -> RequestOutputT: prompts = [ content for text in turn.texts for content in text.contents if content ] - if request_info.credit_phase == CreditPhase.WARMUP: + # Skipped when cache-bust owns prefix isolation: its markers are + # warmup-coherent by design (warmup primes the prefix profiling hits), + # and prefixing in front of the shared marker would break that prime. + # See ``Worker._system_message_for_phase``. + if request_info.credit_phase == CreditPhase.WARMUP and ( + request_info.cache_bust_target in (None, CacheBustTarget.NONE) + ): prompts = [ f"{WARMUP_SYSTEM_MESSAGE_PREFIX}\n{prompt}" for prompt in prompts ] diff --git a/src/aiperf/workers/worker.py b/src/aiperf/workers/worker.py index 241cddb6dd..2a8efbcb6a 100644 --- a/src/aiperf/workers/worker.py +++ b/src/aiperf/workers/worker.py @@ -1034,6 +1034,9 @@ async def _process_credit_with_session( self._system_message_for_phase( system_message=session.conversation.system_message, phase=credit.phase, + cache_bust_target=credit.cache_bust_target + if credit.cache_bust_marker is not None + else None, ), ) self._maybe_warn_cache_bust_silent_drop(session, credit) @@ -1337,10 +1340,30 @@ def _release_and_evict_for_terminal( @staticmethod def _system_message_for_phase( - *, system_message: str | None, phase: CreditPhase + *, + system_message: str | None, + phase: CreditPhase, + cache_bust_target: CacheBustTarget | None, ) -> str | None: + """Prefix warmup system messages so warmup cannot reuse profiling's + prefix cache — unless cache-bust already owns prefix isolation. + + Cache-bust is deliberately warmup-coherent: ``build_cache_bust_marker`` + keeps the phase out of its digest and the marker ledger survives the + WARMUP -> PROFILING boundary, so a trajectory's warmup turn ``k_i`` and + its first profiling turn ``k_i+1`` share one marker and warmup PRIMES + the cache profiling then hits. Injecting this prefix in front of that + shared marker diverges the two prefixes at token 0, so warmup primes an + entry profiling never touches — the round-trip is spent for nothing. + + Any non-NONE target already guarantees warmup cannot collide with + another trajectory's prefix (the marker digests per trajectory tree), so + standing down here costs no isolation. + """ if phase != CreditPhase.WARMUP: return system_message + if cache_bust_target not in (None, CacheBustTarget.NONE): + return system_message if not system_message: return WARMUP_SYSTEM_MESSAGE_PREFIX return f"{WARMUP_SYSTEM_MESSAGE_PREFIX}\n{system_message}" diff --git a/tests/unit/endpoints/conftest.py b/tests/unit/endpoints/conftest.py index f19389bc0f..e2b5401b8c 100644 --- a/tests/unit/endpoints/conftest.py +++ b/tests/unit/endpoints/conftest.py @@ -7,7 +7,7 @@ import pytest -from aiperf.common.enums import CreditPhase, ModelSelectionStrategy +from aiperf.common.enums import CacheBustTarget, CreditPhase, ModelSelectionStrategy from aiperf.common.models import Text, Turn from aiperf.common.models.model_endpoint_info import ( EndpointInfo, @@ -62,6 +62,7 @@ def create_request_info( conversation_id: str = "test-conversation", system_message: str | None = None, user_context_message: str | None = None, + cache_bust_target: CacheBustTarget | None = None, **turn_kwargs, ) -> RequestInfo: """Helper to create RequestInfo with all required fields. @@ -93,6 +94,7 @@ def create_request_info( conversation_id=conversation_id, system_message=system_message, user_context_message=user_context_message, + cache_bust_target=cache_bust_target, ) diff --git a/tests/unit/endpoints/test_completions_endpoint.py b/tests/unit/endpoints/test_completions_endpoint.py index 317a594878..ce54a5118e 100644 --- a/tests/unit/endpoints/test_completions_endpoint.py +++ b/tests/unit/endpoints/test_completions_endpoint.py @@ -2,9 +2,10 @@ # SPDX-License-Identifier: Apache-2.0 import pytest +from pytest import param from aiperf.common.constants import WARMUP_SYSTEM_MESSAGE_PREFIX -from aiperf.common.enums import CreditPhase +from aiperf.common.enums import CacheBustTarget, CreditPhase from aiperf.common.models import Text, Turn from aiperf.endpoints.openai_completions import CompletionsEndpoint from aiperf.plugin.enums import EndpointType @@ -95,6 +96,52 @@ def test_format_payload_warmup_single_prompt_prefixed_string( assert payload["prompt"] == f"{WARMUP_SYSTEM_MESSAGE_PREFIX}\nPrompt 1" + @pytest.mark.parametrize( + "target", + [ + param(CacheBustTarget.SYSTEM_PREFIX, id="system_prefix"), + param(CacheBustTarget.SYSTEM_SUFFIX, id="system_suffix"), + param(CacheBustTarget.FIRST_TURN_PREFIX, id="first_turn_prefix"), + param(CacheBustTarget.FIRST_TURN_SUFFIX, id="first_turn_suffix"), + ], + ) # fmt: skip + def test_format_payload_warmup_skips_prefix_when_cache_bust_active( + self, endpoint, model_endpoint, target + ): + """Cache-bust markers are warmup-coherent, so warmup primes the prefix + profiling hits. Prefixing in front of the shared marker would diverge + the two at token 0 and prime an entry profiling never touches. + """ + turn = Turn(texts=[Text(contents=["Prompt 1"])], model="completion-model") + request_info = create_request_info( + model_endpoint=model_endpoint, + turns=[turn], + credit_phase=CreditPhase.WARMUP, + cache_bust_target=target, + ) + + payload = endpoint.format_payload(request_info) + + assert payload["prompt"] == "Prompt 1" + + def test_format_payload_warmup_prefixes_when_cache_bust_target_none_enum( + self, endpoint, model_endpoint + ): + """``CacheBustTarget.NONE`` means cache-bust is off, so warmup still + needs its own prefix isolation. + """ + turn = Turn(texts=[Text(contents=["Prompt 1"])], model="completion-model") + request_info = create_request_info( + model_endpoint=model_endpoint, + turns=[turn], + credit_phase=CreditPhase.WARMUP, + cache_bust_target=CacheBustTarget.NONE, + ) + + payload = endpoint.format_payload(request_info) + + assert payload["prompt"] == f"{WARMUP_SYSTEM_MESSAGE_PREFIX}\nPrompt 1" + def test_format_payload_filters_empty_prompts(self, endpoint, model_endpoint): """Test that empty strings are filtered from prompts.""" turn = Turn( diff --git a/tests/unit/workers/test_worker.py b/tests/unit/workers/test_worker.py index 6bfb8bc7e5..19144e2784 100644 --- a/tests/unit/workers/test_worker.py +++ b/tests/unit/workers/test_worker.py @@ -7,7 +7,7 @@ import pytest from pytest import param -from aiperf.common.enums import CreditPhase +from aiperf.common.enums import CacheBustTarget, CreditPhase from aiperf.common.models import ( Conversation, ErrorDetails, @@ -467,6 +467,7 @@ def test_profiling_preserves_system_message(self): Worker._system_message_for_phase( system_message="existing system", phase=CreditPhase.PROFILING, + cache_bust_target=None, ) == "existing system" ) @@ -476,6 +477,7 @@ def test_warmup_sets_system_message_when_missing(self): Worker._system_message_for_phase( system_message=None, phase=CreditPhase.WARMUP, + cache_bust_target=None, ) == "warmup" ) @@ -485,10 +487,66 @@ def test_warmup_prefixes_existing_system_message(self): Worker._system_message_for_phase( system_message="existing system", phase=CreditPhase.WARMUP, + cache_bust_target=None, ) == "warmup\nexisting system" ) + def test_warmup_prefixes_when_cache_bust_target_none_enum(self): + """``CacheBustTarget.NONE`` is cache-bust-disabled, so the prefix applies.""" + assert ( + Worker._system_message_for_phase( + system_message="existing system", + phase=CreditPhase.WARMUP, + cache_bust_target=CacheBustTarget.NONE, + ) + == "warmup\nexisting system" + ) + + @pytest.mark.parametrize( + "target", + [ + param(CacheBustTarget.SYSTEM_PREFIX, id="system_prefix"), + param(CacheBustTarget.SYSTEM_SUFFIX, id="system_suffix"), + param(CacheBustTarget.FIRST_TURN_PREFIX, id="first_turn_prefix"), + param(CacheBustTarget.FIRST_TURN_SUFFIX, id="first_turn_suffix"), + ], + ) # fmt: skip + def test_warmup_skips_prefix_when_cache_bust_active(self, target): + """Cache-bust markers are warmup-coherent: warmup primes the prefix + profiling hits, so prefixing in front of the shared marker would break + the prime. Every non-NONE target already isolates per trajectory tree. + """ + assert ( + Worker._system_message_for_phase( + system_message="existing system", + phase=CreditPhase.WARMUP, + cache_bust_target=target, + ) + == "existing system" + ) + + @pytest.mark.parametrize( + "target", + [ + param(CacheBustTarget.SYSTEM_PREFIX, id="system_prefix"), + param(CacheBustTarget.FIRST_TURN_PREFIX, id="first_turn_prefix"), + ], + ) # fmt: skip + def test_warmup_leaves_system_message_none_when_cache_bust_active(self, target): + """No synthetic system message is invented under cache-bust — a + warmup-only ``system`` role would diverge the message array from + profiling's even before content is compared. + """ + assert ( + Worker._system_message_for_phase( + system_message=None, + phase=CreditPhase.WARMUP, + cache_bust_target=target, + ) + is None + ) + @pytest.mark.asyncio class TestCreateRequestInfo: From b719246fa671bca231e1c7f1dad7d5a910ae59be Mon Sep 17 00:00:00 2001 From: Anthony Casagrande Date: Thu, 30 Jul 2026 12:54:49 -0700 Subject: [PATCH 6/9] fix(config): preserve default profiling grace period (#1230) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Keep the default profiling grace period unless the user explicitly overrides it. Upstream: 5ad08166a1de3c2f84a86982c56f9f4862fabece 中文：除非用户显式覆盖，否则保留默认的性能测试宽限期。 Signed-off-by: Anthony Casagrande Signed-off-by: Cam Quilici --- .../config/flags/_converter_profiling.py | 5 +++ src/aiperf/config/flags/resolver.py | 8 +++++ .../cli_runner/test_request_count_override.py | 36 +++++++++++++++---- tests/unit/timing/test_timing_config.py | 21 +++++++++++ tools/ruff_baseline.json | 5 +++ 5 files changed, 69 insertions(+), 6 deletions(-) diff --git a/src/aiperf/config/flags/_converter_profiling.py b/src/aiperf/config/flags/_converter_profiling.py index 6243eaf8a8..ea8bc784f2 100644 --- a/src/aiperf/config/flags/_converter_profiling.py +++ b/src/aiperf/config/flags/_converter_profiling.py @@ -461,6 +461,11 @@ def build_profiling(cli: CLIConfig) -> dict[str, Any]: for output_key, attr_name in _PROF_FIELD_ROUTES: if attr_name in fields_set: prof[output_key] = getattr(cli, attr_name) + if ( + cli.benchmark_duration is not None + and "benchmark_grace_period" not in fields_set + ): + prof["grace_period"] = cli.benchmark_grace_period _apply_profiling_ramps(prof, cli) _apply_agentic_replay_fields(prof, cli) diff --git a/src/aiperf/config/flags/resolver.py b/src/aiperf/config/flags/resolver.py index e090abe1cc..aa80a46219 100644 --- a/src/aiperf/config/flags/resolver.py +++ b/src/aiperf/config/flags/resolver.py @@ -480,6 +480,14 @@ def _apply_phase_loadgen_overrides(merged: dict[str, Any], cli: CLIConfig) -> No continue target[key] = value + if ( + "benchmark_duration" in loadgen_set + and "benchmark_grace_period" not in loadgen_set + and cli.benchmark_duration is not None + and "grace_period" not in target + ): + target["grace_period"] = cli.benchmark_grace_period + _apply_agentic_replay_fields(target, cli) diff --git a/tests/unit/cli_runner/test_request_count_override.py b/tests/unit/cli_runner/test_request_count_override.py index 92cf37961b..1fa545b41e 100644 --- a/tests/unit/cli_runner/test_request_count_override.py +++ b/tests/unit/cli_runner/test_request_count_override.py @@ -42,6 +42,20 @@ def _warmup_requests(cfg) -> int | None: # noqa: ANN001 return None +def _profiling_duration(cfg) -> float | None: # noqa: ANN001 + for phase in cfg.benchmark.phases: + if phase.name == "profiling": + return getattr(phase, "duration", None) + return None + + +def _profiling_grace_period(cfg) -> float | None: # noqa: ANN001 + for phase in cfg.benchmark.phases: + if phase.name == "profiling": + return getattr(phase, "grace_period", None) + return None + + def test_request_count_overrides_template_phases_requests() -> None: """``--request-count 10`` with ``minimal.yaml`` overrides the YAML's 100.""" user = CLIConfig(**CLIConfig(request_count=10).model_dump(exclude_unset=True)) @@ -67,10 +81,20 @@ def test_no_loadgen_override_leaves_yaml_intact() -> None: assert _profiling_requests(cfg) == 100 -def test_concurrency_override_targets_profiling_phase() -> None: - """The same overlay rule applies to ``--concurrency``.""" - user = CLIConfig(**CLIConfig(concurrency=99).model_dump(exclude_unset=True)) +def test_benchmark_duration_gets_default_grace_period_with_config_file() -> None: + """``--benchmark-duration`` also gets the default grace period with YAML.""" + user = CLIConfig(**CLIConfig(benchmark_duration=5.0).model_dump(exclude_unset=True)) cfg = resolve_config(user, TEMPLATES_DIR / "minimal.yaml") - for phase in cfg.benchmark.phases: - if phase.name == "profiling": - assert getattr(phase, "concurrency", None) == 99 + assert _profiling_duration(cfg) == 5.0 + assert _profiling_grace_period(cfg) == 30.0 + + +def test_explicit_benchmark_grace_period_with_config_file_is_preserved() -> None: + """An explicit ``--benchmark-grace-period`` overrides the default.""" + user = CLIConfig( + **CLIConfig(benchmark_duration=5.0, benchmark_grace_period=15.0).model_dump( + exclude_unset=True + ) + ) + cfg = resolve_config(user, TEMPLATES_DIR / "minimal.yaml") + assert _profiling_grace_period(cfg) == 15.0 diff --git a/tests/unit/timing/test_timing_config.py b/tests/unit/timing/test_timing_config.py index c64b9c265f..5136edfd00 100644 --- a/tests/unit/timing/test_timing_config.py +++ b/tests/unit/timing/test_timing_config.py @@ -5,6 +5,7 @@ import pytest from pydantic import ValidationError +from pytest import param from aiperf.common.enums import CreditPhase from aiperf.config.flags.cli_config import CLIConfig @@ -342,6 +343,26 @@ def test_warmup_grace_period( warmup = next(pc for pc in cfg.phase_configs if pc.phase == CreditPhase.WARMUP) assert warmup.grace_period_sec == expected + @pytest.mark.parametrize( + "benchmark_grace_period,expected", + [ + param(None, 30.0, id="default"), + param(15.0, 15.0, id="explicit_positive"), + param(0.0, 0.0, id="explicit_zero"), + ], + ) # fmt: skip + def test_build_profiling_duration_grace_period_returns_expected_value( + self, benchmark_grace_period: float | None, expected: float + ) -> None: + kwargs: dict[str, Any] = {"benchmark_duration": 5.0} + if benchmark_grace_period is not None: + kwargs["benchmark_grace_period"] = benchmark_grace_period + cfg = _make_timing_config(**kwargs) + profiling = next( + pc for pc in cfg.phase_configs if pc.phase == CreditPhase.PROFILING + ) + assert profiling.grace_period_sec == expected + class TestPhaseRequestRate: """``_phase_request_rate`` must delegate to ``get_phase_rate`` so only diff --git a/tools/ruff_baseline.json b/tools/ruff_baseline.json index 03e2542207..3dfe8025fc 100644 --- a/tools/ruff_baseline.json +++ b/tools/ruff_baseline.json @@ -259,6 +259,11 @@ "src/aiperf/config/flags/_converter_dataset.py", "_build_prompts" ], + [ + "C901", + "src/aiperf/config/flags/resolver.py", + "_apply_phase_loadgen_overrides" + ], [ "C901", "src/aiperf/config/flags/variant_parser.py", From ed057829b78d25d79ce6f3b87763d48fe50363f5 Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Fri, 31 Jul 2026 14:50:31 -0500 Subject: [PATCH 7/9] fix(agentx): globally anchor profiling handoff MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Carry each stream's full recorded delay across the warmup barrier, then subtract one phase-wide minimum so the earliest profiling request starts immediately without changing relative timing, order, or join gates. 中文：在 warmup 屏障后保留每个流完整的记录延迟，再统一减去 profiling 阶段的全局最小值，使最早请求立即启动，同时保持相对时序、请求顺序和 join 门控不变。 Signed-off-by: Cam Quilici --- docs/cli-options.md | 4 +- docs/tutorials/agentx-mvp.md | 9 +- src/aiperf/config/flags/cli_config.py | 19 +-- src/aiperf/config/phases.py | 19 +-- .../config/schema/aiperf-config.schema.json | 48 +++--- .../timing/strategies/agentic_replay.py | 149 ++++++++---------- .../timing/strategies/test_agentic_replay.py | 137 +++++++++------- 7 files changed, 191 insertions(+), 194 deletions(-) diff --git a/docs/cli-options.md b/docs/cli-options.md index 5ef5ffe64e..3af9300901 100644 --- a/docs/cli-options.md +++ b/docs/cli-options.md @@ -1087,7 +1087,7 @@ AGENTIC_REPLAY only: upper bound (inclusive) on the random start position within #### `--burst-phase-starts` -AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay. +AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.
_Flag (no value required)_ #### `--trace-idle-gap-cap-seconds` `` @@ -2622,7 +2622,7 @@ AGENTIC_REPLAY only: upper bound (inclusive) on the random start position within #### `--burst-phase-starts` -AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay. +AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.
_Flag (no value required)_ #### `--trace-idle-gap-cap-seconds` `` diff --git a/docs/tutorials/agentx-mvp.md b/docs/tutorials/agentx-mvp.md index 14d22d2fc9..8cb42536fe 100644 --- a/docs/tutorials/agentx-mvp.md +++ b/docs/tutorials/agentx-mvp.md @@ -419,7 +419,9 @@ AIPerf continues the same live session trees for that duration with recorded idle delays removed and every request 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. +profiling from that exact state. Each stream's full next-turn delay is carried +across the phase boundary; time spent waiting for the global warmup drain does +not consume it. For repeatable warmup depth, use `--warmup-requests-per-lane REQUESTS` instead. After the mandatory snapshot @@ -440,7 +442,10 @@ request metrics. After warmup, the profiling phase opens. Now you're measuring. Each trajectory keeps replaying its conversation from turn `k_i + 1` onward, honoring the original recorded inter-turn gaps as end-to-start delays counted from the -moment the previous turn completes. With +moment the previous turn completes. At the phase boundary, AIPerf subtracts +one global minimum from the pending first-request offsets: the earliest +request starts immediately, while every other trajectory keeps the same +relative spacing. With `--trace-idle-gap-cap-seconds S`, a per-trajectory watchdog covers initial future work and restarts when the last in-flight request across the root and all descendant streams completes. diff --git a/src/aiperf/config/flags/cli_config.py b/src/aiperf/config/flags/cli_config.py index 0574f36fdb..75a8afc3b6 100644 --- a/src/aiperf/config/flags/cli_config.py +++ b/src/aiperf/config/flags/cli_config.py @@ -2157,17 +2157,14 @@ def url(self) -> str: Field( description="AGENTIC_REPLAY only: collapse the WARMUP-start and " "PROFILING-start dispatches into synchronized bursts instead of " - "spreading them by each request's recorded offset from t*. By " - "default (False) the phase starts are SPREAD: WARMUP requests are " - "aligned globally so every trajectory reaches its t* at the same " - "instant (the warmup end), and each lane's first PROFILING request " - "waits out its recorded gap after t* -- reproducing the recorded " - "arrival pattern at both phase boundaries. The rest of the replay " - "(inter-turn delays) is timing-faithful regardless of this flag; " - "it governs ONLY the burst-vs-spread of the two phase starts. Pass " - "--burst-phase-starts to fire each phase's first requests together " - "(faster concurrency ramp, synchronized start), e.g. for a " - "throughput-oriented run rather than a faithful arrival replay.", + "preserving recorded spacing. By default (False), WARMUP requests " + "are aligned globally so every trajectory reaches t* together. " + "PROFILING subtracts one phase-wide minimum from every first " + "request offset: the earliest request starts immediately and all " + "other trajectories retain their recorded spacing. Pass " + "--burst-phase-starts to subtract a separate minimum per lane, " + "making every lane start immediately. Subsequent inter-turn " + "delays are timing-faithful in either mode.", ), CLIParameter( name=("--burst-phase-starts",), diff --git a/src/aiperf/config/phases.py b/src/aiperf/config/phases.py index d4f7a0022a..c1a502c62c 100644 --- a/src/aiperf/config/phases.py +++ b/src/aiperf/config/phases.py @@ -321,17 +321,14 @@ class BasePhaseConfig(AdaptiveScalePhaseMixin, BaseConfig): default=False, description="AGENTIC_REPLAY only: collapse the WARMUP-start and " "PROFILING-start dispatches into synchronized bursts instead of " - "spreading them by each request's recorded offset from t*. By " - "default (False) the phase starts are SPREAD: WARMUP requests are " - "aligned globally so every trajectory reaches its t* at the same " - "instant (the warmup end), and each lane's first PROFILING request " - "waits out its recorded gap after t* -- reproducing the recorded " - "arrival pattern at both phase boundaries. The rest of the replay " - "(inter-turn delays) is timing-faithful regardless of this flag; " - "it governs ONLY the burst-vs-spread of the two phase starts. Pass " - "--burst-phase-starts to fire each phase's first requests together " - "(faster concurrency ramp, synchronized start), e.g. for a " - "throughput-oriented run rather than a faithful arrival replay.", + "preserving recorded spacing. By default (False), WARMUP requests " + "are aligned globally so every trajectory reaches t* together. " + "PROFILING subtracts one phase-wide minimum from every first " + "request offset: the earliest request starts immediately and all " + "other trajectories retain their recorded spacing. Pass " + "--burst-phase-starts to subtract a separate minimum per lane, " + "making every lane start immediately. Subsequent inter-turn " + "delays are timing-faithful in either mode.", ), ] diff --git a/src/aiperf/config/schema/aiperf-config.schema.json b/src/aiperf/config/schema/aiperf-config.schema.json index 0cf8215eff..122cd9ac39 100644 --- a/src/aiperf/config/schema/aiperf-config.schema.json +++ b/src/aiperf/config/schema/aiperf-config.schema.json @@ -1284,7 +1284,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, @@ -1732,7 +1732,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, @@ -2248,7 +2248,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, @@ -2778,7 +2778,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, @@ -3294,7 +3294,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, @@ -3787,7 +3787,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, @@ -5154,7 +5154,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, @@ -5602,7 +5602,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, @@ -6118,7 +6118,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, @@ -6648,7 +6648,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, @@ -7164,7 +7164,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, @@ -7657,7 +7657,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, @@ -8145,7 +8145,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, @@ -8593,7 +8593,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, @@ -9109,7 +9109,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, @@ -9639,7 +9639,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, @@ -10155,7 +10155,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, @@ -10648,7 +10648,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, @@ -11535,7 +11535,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, @@ -12296,7 +12296,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, @@ -14312,7 +14312,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, @@ -15134,7 +15134,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, @@ -17753,7 +17753,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, @@ -21382,7 +21382,7 @@ }, "burstPhaseStarts": { "default": false, - "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of spreading them by each request's recorded offset from t*. By default (False) the phase starts are SPREAD: WARMUP requests are aligned globally so every trajectory reaches its t* at the same instant (the warmup end), and each lane's first PROFILING request waits out its recorded gap after t* -- reproducing the recorded arrival pattern at both phase boundaries. The rest of the replay (inter-turn delays) is timing-faithful regardless of this flag; it governs ONLY the burst-vs-spread of the two phase starts. Pass --burst-phase-starts to fire each phase's first requests together (faster concurrency ramp, synchronized start), e.g. for a throughput-oriented run rather than a faithful arrival replay.", + "description": "AGENTIC_REPLAY only: collapse the WARMUP-start and PROFILING-start dispatches into synchronized bursts instead of preserving recorded spacing. By default (False), WARMUP requests are aligned globally so every trajectory reaches t* together. PROFILING subtracts one phase-wide minimum from every first request offset: the earliest request starts immediately and all other trajectories retain their recorded spacing. Pass --burst-phase-starts to subtract a separate minimum per lane, making every lane start immediately. Subsequent inter-turn delays are timing-faithful in either mode.", "title": "Burstphasestarts", "type": "boolean" }, diff --git a/src/aiperf/timing/strategies/agentic_replay.py b/src/aiperf/timing/strategies/agentic_replay.py index ecf12fcf7b..216f1813d4 100644 --- a/src/aiperf/timing/strategies/agentic_replay.py +++ b/src/aiperf/timing/strategies/agentic_replay.py @@ -38,12 +38,14 @@ degraded trajectory pool. PROFILING: each stream resumes at its first turn at/after t* -(``next_turn_index``). Default dispatch preserves the stream's recorded offset -from the replay boundary; ``--burst-phase-starts`` collapses each lane's first -eligible request to profiling-time 0. Accelerated cache warmup synthesizes a -new replay boundary at the warmup handoff and carries each live stream's -residual next-turn delay into profiling, so the handoff ramps instead of -firing every live stream at once. Subsequent turns honor trace inter-turn +(``next_turn_index``). Default dispatch subtracts one phase-wide minimum from +every stream's recorded offset: the earliest eligible request fires at +profiling-time 0 while all cross-trajectory spacing and ordering are preserved. +``--burst-phase-starts`` instead collapses each lane's first eligible request +to profiling-time 0. Accelerated cache warmup synthesizes a new replay boundary +at the warmup handoff and carries each live stream's full next-turn delay into +profiling; time spent waiting for the global warmup barrier does not consume +that delay. Subsequent turns honor trace inter-turn ``delay_ms`` from the original trace timeline. Gated parents fire their join turn when blocking children complete. When a root session reaches its final turn AND its whole tree (root + every descendant subagent) has drained, @@ -197,7 +199,6 @@ def __init__( self._baseline_warmup_turns: set[tuple[str, int]] = set() self._accelerated_warmup_started = False self._handoff_credits: dict[str, Credit] = {} - self._handoff_returned_at_ns: dict[str, int] = {} self._root_to_lane: dict[str, int] = {} # Cache-bust state. WARMUP and PROFILING construct distinct strategy @@ -228,13 +229,13 @@ def __init__( ) self._benchmark_id: str = run.benchmark_id if run is not None else "unknown" # ``--burst-phase-starts`` (BasePhaseConfig.burst_phase_starts on the - # profiling phase). Default False: the WARMUP and PROFILING phase starts - # are SPREAD by each request's recorded offset from t* (warmup globally - # t*-aligned, profiling preserving each lane's leading gap). When True, - # both phase starts collapse into synchronized bursts. Governs ONLY the - # two phase-start dispatch patterns; the rest of replay timing is - # faithful regardless. ``is True`` guards MagicMock/None test configs -> - # default (spread). + # profiling phase). Default False: WARMUP is globally t*-aligned and + # PROFILING subtracts one global minimum so its earliest request starts + # immediately while every other request keeps its relative offset. + # When True, each lane subtracts its own minimum and the phase starts as + # a synchronized burst. Governs ONLY the two phase-start dispatch + # patterns; the rest of replay timing is faithful regardless. ``is + # True`` guards MagicMock/None test configs -> default (spread). profiling_phase = None if run is not None: phases = run.cfg.get_profiling_phases() @@ -890,12 +891,8 @@ def observe_credit_return(self, credit: Credit) -> None: self._correlation_to_lane[credit.x_correlation_id] = lane if credit.is_final_turn: self._handoff_credits.pop(credit.x_correlation_id, None) - self._handoff_returned_at_ns.pop(credit.x_correlation_id, None) else: self._handoff_credits[credit.x_correlation_id] = credit - self._handoff_returned_at_ns[credit.x_correlation_id] = ( - time.perf_counter_ns() - ) async def _handle_accelerated_warmup_return(self, credit: Credit) -> None: """Issue the next compressed turn or recycle a completed tree.""" @@ -948,8 +945,7 @@ async def finalize_phase(self) -> None: ) if not self._accelerated_warmup_started: return - finalized_at_ns = time.perf_counter_ns() - states_by_lane = self._build_handoff_states(finalized_at_ns=finalized_at_ns) + states_by_lane = self._build_handoff_states() boundaries_by_lane = { lane: self._build_handoff_replay_boundaries(states) for lane, states in states_by_lane.items() @@ -962,12 +958,8 @@ async def finalize_phase(self) -> None: f"{sum(len(states) for states in states_by_lane.values())} live streams" ) - def _build_handoff_states( - self, *, finalized_at_ns: int | None = None - ) -> dict[int, list[ConversationState]]: + def _build_handoff_states(self) -> dict[int, list[ConversationState]]: """Convert drained credits and join annotations into lane states.""" - if finalized_at_ns is None: - finalized_at_ns = time.perf_counter_ns() blocked, child_annotations = self._handoff_annotations() states_by_lane: dict[int, list[ConversationState]] = { lane: [] for lane in range(len(self.conversation_source.trajectories)) @@ -976,7 +968,6 @@ def _build_handoff_states( states_by_lane, blocked=blocked, child_annotations=child_annotations, - finalized_at_ns=finalized_at_ns, ) self._add_pending_handoff_states( states_by_lane, seen_states, child_annotations=child_annotations @@ -1012,7 +1003,6 @@ def _add_returned_handoff_states( *, blocked: dict[str, int], child_annotations: dict[str, list[tuple[str | None, int | None]]], - finalized_at_ns: int, ) -> set[tuple[str, str, int]]: seen_states: set[tuple[str, str, int]] = set() for credit in self._handoff_credits.values(): @@ -1020,7 +1010,6 @@ def _add_returned_handoff_states( credit, blocked=blocked, child_annotations=child_annotations, - finalized_at_ns=finalized_at_ns, ) if handoff is None: continue @@ -1037,7 +1026,6 @@ def _returned_credit_handoff_state( *, blocked: dict[str, int], child_annotations: dict[str, list[tuple[str | None, int | None]]], - finalized_at_ns: int, ) -> tuple[int, ConversationState] | None: lane = self._root_to_lane.get(credit.effective_root_correlation_id) if lane is None or credit.turn_index + 1 >= credit.num_turns: @@ -1049,8 +1037,8 @@ def _returned_credit_handoff_state( conversation_id=credit.conversation_id, x_correlation_id=credit.x_correlation_id, next_turn_index=credit.turn_index + 1, - next_dispatch_offset_ms=self._handoff_residual_delay_ms( - credit, finalized_at_ns=finalized_at_ns + next_dispatch_offset_ms=self._handoff_delay_ms( + credit.conversation_id, credit.turn_index + 1 ), agent_depth=credit.agent_depth, parent_correlation_id=credit.parent_correlation_id, @@ -1130,7 +1118,9 @@ def _pending_turn_handoff_state( conversation_id=turn.conversation_id, x_correlation_id=turn.x_correlation_id, next_turn_index=turn.turn_index, - next_dispatch_offset_ms=0.0, + next_dispatch_offset_ms=self._handoff_delay_ms( + turn.conversation_id, turn.turn_index + ), agent_depth=turn.agent_depth, parent_correlation_id=turn.parent_correlation_id, root_correlation_id=turn.root_correlation_id, @@ -1157,47 +1147,25 @@ def _handoff_lane_for_turn( return lane return self._correlation_to_lane.get(turn.x_correlation_id) - def _handoff_residual_delay_ms( - self, credit: Credit, *, finalized_at_ns: int - ) -> float: - """Delay before a warmup-handoff stream should issue in PROFILING. - - Accelerated warmup intentionally removes idle gaps while it is building - KV pressure. At the profiling boundary, however, firing every drained - stream at once creates an artificial burst. Preserve the local replay - cadence by carrying the next turn's recorded delay forward, minus any - wall-clock time already spent waiting for the warmup drain/finalize. - """ - delay_ms = self._handoff_base_delay_ms(credit) - returned_at_ns = self._handoff_returned_at_ns.get(credit.x_correlation_id) - if returned_at_ns is not None: - elapsed_ms = max(0.0, (finalized_at_ns - returned_at_ns) / 1_000_000.0) - delay_ms = max(0.0, delay_ms - elapsed_ms) - return delay_ms - - def _handoff_base_delay_ms(self, credit: Credit) -> float: - """Recorded delay from a returned warmup credit to its next turn.""" + def _handoff_delay_ms(self, conversation_id: str, next_turn_index: int) -> float: + """Recorded end-to-start delay carried across the warmup barrier.""" try: - next_meta = self.conversation_source.get_next_turn_metadata(credit) - except (KeyError, ValueError): + metadata = self.conversation_source.get_metadata(conversation_id) + except KeyError: + return 0.0 + if next_turn_index < 0 or next_turn_index >= len(metadata.turns): return 0.0 + next_meta = metadata.turns[next_turn_index] delay_ms = _as_timestamp_ms(getattr(next_meta, "delay_ms", None)) if delay_ms is not None: return max(0.0, delay_ms) - - try: - metadata = self.conversation_source.get_metadata(credit.conversation_id) - except KeyError: - return 0.0 - if credit.turn_index < 0 or credit.turn_index + 1 >= len(metadata.turns): + if next_turn_index == 0: return 0.0 - previous_meta = metadata.turns[credit.turn_index] + previous_meta = metadata.turns[next_turn_index - 1] previous_ts_ms = _as_timestamp_ms(getattr(previous_meta, "timestamp_ms", None)) - next_ts_ms = _as_timestamp_ms( - getattr(metadata.turns[credit.turn_index + 1], "timestamp_ms", None) - ) + next_ts_ms = _as_timestamp_ms(getattr(next_meta, "timestamp_ms", None)) if previous_ts_ms is None or next_ts_ms is None: return 0.0 @@ -1301,6 +1269,7 @@ async def _execute_profiling(self) -> None: serializing over N credit round-trips. Subsequent turns and recycle-pool sessions are dispatched from handle_credit_return. """ + phase_t0_offset_ms = self._profiling_phase_t0_offset_ms() spread_s = self._profiling_spread_seconds() mode = "burst" if self._burst_phase_starts else "spread" self.info( @@ -1313,7 +1282,9 @@ async def _execute_profiling(self) -> None: # unreachable by the phase runner's cancellation. results = await asyncio.gather( *( - self._dispatch_one_profiling_trajectory(trajectory, lane) + self._dispatch_one_profiling_trajectory( + trajectory, lane, phase_t0_offset_ms + ) for lane, trajectory in enumerate(self.conversation_source.trajectories) ), return_exceptions=True, @@ -1359,12 +1330,29 @@ def _profiling_spread_seconds(self) -> float: return 0.0 return (max(first_offsets) - min(first_offsets)) / MILLIS_PER_SECOND + def _profiling_phase_t0_offset_ms(self) -> float: + """Global spread anchor that starts the earliest request at phase zero.""" + if self._burst_phase_starts: + return 0.0 + offsets: list[float] = [] + for trajectory in self.conversation_source.trajectories: + if trajectory.snapshot is None: + return 0.0 + offsets.extend( + state.next_dispatch_offset_ms + for state in trajectory.snapshot.states + if not state.waiting_on_children + ) + return min(offsets, default=0.0) + async def _dispatch_one_profiling_trajectory( - self, trajectory: Trajectory, lane: int + self, trajectory: Trajectory, lane: int, phase_t0_offset_ms: float ) -> None: """Dispatch one lane's initial PROFILING credit (run under gather).""" if trajectory.snapshot is not None: - await self._dispatch_snapshot_for_profiling(trajectory, lane) + await self._dispatch_snapshot_for_profiling( + trajectory, lane, phase_t0_offset_ms + ) return session = self.conversation_source.session_for(trajectory) @@ -1489,10 +1477,6 @@ async def _handle_warmup_return(self, credit: Credit) -> None: await self._handle_accelerated_warmup_return(credit) return self._baseline_warmup_returns[credit.x_correlation_id] = credit - if not credit.is_final_turn: - self._handoff_returned_at_ns[credit.x_correlation_id] = ( - time.perf_counter_ns() - ) if ( len(self._baseline_warmup_returns) >= self.conversation_source.warmup_credit_count @@ -1617,21 +1601,19 @@ async def _on_rootless_child_done(self, lane: int) -> None: await self._dispatch_recycled_on_lane(lane) async def _dispatch_snapshot_for_profiling( - self, trajectory: Trajectory, lane: int + self, trajectory: Trajectory, lane: int, phase_t0_offset_ms: float ) -> None: """Resume one trajectory's streams for PROFILING. Each stream profiles from turn ``next_turn_index`` (the first turn at or after t*; its predecessor, if any, was primed during WARMUP). - Dispatch anchoring depends on ``--burst-phase-starts``: by default - (spread) ``t0_offset = 0`` so each stream waits out its recorded offset - from t* -- the leading t*->first-request gap is preserved and lanes - ramp in. With ``--burst-phase-starts`` the trajectory's earliest - post-t* request is anchored at profiling-time 0 (subtracting T0, the - min offset) so the lane bursts at once. Relative timing among the - trajectory's streams and turns is identical either way -- only the - per-lane start offset differs. + Dispatch anchoring depends on ``--burst-phase-starts``. By default one + phase-wide minimum is subtracted from every stream: the earliest + request starts at profiling-time 0 and all cross-trajectory spacing is + preserved. With ``--burst-phase-starts`` each trajectory subtracts its + own minimum, so every lane starts at profiling-time 0. Relative timing + within each trajectory is identical either way. Gated parents (``waiting_on_children``) are not dispatched here; their join is seeded with the orchestrator and their gated turn fires at the @@ -1662,7 +1644,7 @@ async def _dispatch_snapshot_for_profiling( offset_by_corr[state.x_correlation_id] for state in dispatchable ) else: - t0_offset_ms = 0.0 + t0_offset_ms = phase_t0_offset_ms if self.branch_orchestrator is not None: self.branch_orchestrator.seed_snapshot( @@ -1739,9 +1721,8 @@ async def _dispatch_snapshot_for_profiling( ) if not self._has_tree_registry and not has_root_state: self._rootless_lane_outstanding[lane] = len(dispatchable) - # Spread (default): t0 = 0, each lane fires at its recorded offset from - # t*. Burst (--burst-phase-starts): t0 = the lane's minimum offset, - # anchoring the earliest post-t* request at profiling-0. + # Spread (default): every lane shares one phase-wide T0, preserving + # cross-trajectory offsets. Burst: each lane uses its own minimum. for state in dispatchable: session = self.conversation_source.session_for_state(state) turn = self._build_turn_for_session(session, state.next_turn_index) diff --git a/tests/unit/timing/strategies/test_agentic_replay.py b/tests/unit/timing/strategies/test_agentic_replay.py index c5693ec8a2..f468fba6cb 100644 --- a/tests/unit/timing/strategies/test_agentic_replay.py +++ b/tests/unit/timing/strategies/test_agentic_replay.py @@ -442,7 +442,7 @@ async def test_cache_warmup_handoff_preserves_every_quota_refused_turn_once( assert admission(refused_turn) is TurnAdmission.DEFER issuer.replay_gate.pause_releases.assert_not_called() - states = strategy._build_handoff_states(finalized_at_ns=0) + states = strategy._build_handoff_states() assert len(states[0]) == 1 state = states[0][0] @@ -677,7 +677,7 @@ async def test_cache_warmup_cutoff_stops_issuer_and_persists_next_turn(): assert snapshot.replay_resume_boundaries == (ReplayResumeBoundary("trace_0", 3),) -def test_cache_warmup_handoff_preserves_residual_next_turn_delay() -> None: +def test_cache_warmup_handoff_preserves_full_next_turn_delay() -> None: dataset = DatasetMetadata( conversations=[ ConversationMetadata( @@ -705,14 +705,13 @@ def test_cache_warmup_handoff_preserves_residual_next_turn_delay() -> None: phase=CreditPhase.WARMUP, ) strategy._handoff_credits[credit.x_correlation_id] = credit - strategy._handoff_returned_at_ns[credit.x_correlation_id] = 1_000_000_000 strategy._root_to_lane[credit.effective_root_correlation_id] = 0 - states_by_lane = strategy._build_handoff_states(finalized_at_ns=3_500_000_000) + states_by_lane = strategy._build_handoff_states() state = states_by_lane[0][0] assert state.next_turn_index == 2 - assert state.next_dispatch_offset_ms == pytest.approx(5_500.0) + assert state.next_dispatch_offset_ms == pytest.approx(8_000.0) def test_cache_warmup_handoff_uses_timestamp_fallback_without_per_stream_cap() -> None: @@ -743,14 +742,13 @@ def test_cache_warmup_handoff_uses_timestamp_fallback_without_per_stream_cap() - phase=CreditPhase.WARMUP, ) strategy._handoff_credits[credit.x_correlation_id] = credit - strategy._handoff_returned_at_ns[credit.x_correlation_id] = 1_000_000_000 strategy._root_to_lane[credit.effective_root_correlation_id] = 0 - states_by_lane = strategy._build_handoff_states(finalized_at_ns=1_250_000_000) + states_by_lane = strategy._build_handoff_states() - # Timestamp fallback gives 6000 - (1000 + 500) = 4500ms, less the 250ms - # handoff drain wait = 4250ms. The raw stream offset stays intact. - assert states_by_lane[0][0].next_dispatch_offset_ms == pytest.approx(4_250.0) + # Timestamp fallback gives 6000 - (1000 + 500) = 4500ms. Warmup barrier + # time does not consume the profiling delay. + assert states_by_lane[0][0].next_dispatch_offset_ms == pytest.approx(4_500.0) def test_cache_warmup_handoff_preserves_raw_stream_offsets() -> None: @@ -799,24 +797,20 @@ def test_cache_warmup_handoff_preserves_raw_stream_offsets() -> None: branch_mode=ConversationBranchMode.SPAWN, ) strategy._handoff_credits = {"root": root, "child": child} - strategy._handoff_returned_at_ns = { - "root": 1_000_000_000, - "child": 1_000_000_000, - } strategy._root_to_lane["root"] = 0 - states = strategy._build_handoff_states(finalized_at_ns=144_259_000_000)[0] + states = strategy._build_handoff_states()[0] offsets = { state.x_correlation_id: state.next_dispatch_offset_ms for state in states } assert offsets == { - "root": pytest.approx(676_750.0), - "child": pytest.approx(9_212_954.0), + "root": pytest.approx(820_009.0), + "child": pytest.approx(9_356_213.0), } -def test_cache_warmup_handoff_elapsed_wait_can_exhaust_delay() -> None: +def test_cache_warmup_handoff_barrier_wait_does_not_exhaust_delay() -> None: dataset = DatasetMetadata( conversations=[ ConversationMetadata( @@ -844,12 +838,11 @@ def test_cache_warmup_handoff_elapsed_wait_can_exhaust_delay() -> None: phase=CreditPhase.WARMUP, ) strategy._handoff_credits[credit.x_correlation_id] = credit - strategy._handoff_returned_at_ns[credit.x_correlation_id] = 1_000_000_000 strategy._root_to_lane[credit.effective_root_correlation_id] = 0 - states_by_lane = strategy._build_handoff_states(finalized_at_ns=4_000_000_000) + states_by_lane = strategy._build_handoff_states() - assert states_by_lane[0][0].next_dispatch_offset_ms == pytest.approx(0.0) + assert states_by_lane[0][0].next_dispatch_offset_ms == pytest.approx(2_000.0) def test_handoff_preserves_completed_streams_absent_from_live_state() -> None: @@ -878,16 +871,36 @@ def test_handoff_preserves_completed_streams_absent_from_live_state() -> None: def test_cache_warmup_handoff_preserves_pending_replay_barrier_turns() -> None: + dataset = DatasetMetadata( + conversations=[ + ConversationMetadata( + conversation_id="trace_0", + turns=[TurnMetadata(delay_ms=0.0)], + ), + ConversationMetadata( + conversation_id="trace_0::fa:001", + turns=[ + TurnMetadata(delay_ms=0.0), + TurnMetadata(delay_ms=30_000.0), + ], + is_root=False, + agent_depth=1, + parent_conversation_id="trace_0", + ), + ], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + ) strategy, issuer, _, _ = _make_strategy( phase=CreditPhase.WARMUP, trajectories=[Trajectory(conversation_id="trace_0", start_turn_index=0)], + dataset=dataset, cache_warmup_duration=10.0, ) pending_turn = TurnToSend( conversation_id="trace_0::fa:001", x_correlation_id="child", - turn_index=0, - num_turns=3, + turn_index=1, + num_turns=2, agent_depth=1, parent_correlation_id="root", root_correlation_id="root", @@ -895,14 +908,14 @@ def test_cache_warmup_handoff_preserves_pending_replay_barrier_turns() -> None: issuer.replay_gate.pending_turns_by_root.return_value = {"root": (pending_turn,)} strategy._root_to_lane["root"] = 0 - states_by_lane = strategy._build_handoff_states(finalized_at_ns=0) + states_by_lane = strategy._build_handoff_states() assert len(states_by_lane[0]) == 1 state = states_by_lane[0][0] assert state.conversation_id == "trace_0::fa:001" assert state.x_correlation_id == "child" - assert state.next_turn_index == 0 - assert state.next_dispatch_offset_ms == 0.0 + assert state.next_turn_index == 1 + assert state.next_dispatch_offset_ms == pytest.approx(30_000.0) assert state.agent_depth == 1 assert state.parent_correlation_id == "root" assert state.root_correlation_id == "root" @@ -1560,8 +1573,8 @@ def fake_schedule_later(delay, coro, **_kwargs): @pytest.mark.asyncio -async def test_profiling_idle_trajectory_caps_leading_idle_preserving_subagent_spacing(): - """A trajectory idle at t* caps only the leading idle and shifts every stream left uniformly, preserving subagent spacing: children at 100s/130s/220s with a 60s cap fire at 60s/90s/180s, not collapsed to 60s each.""" +async def test_profiling_global_anchor_preserves_subagent_spacing(): + """One phase anchor shifts 100s/130s/220s to 0s/30s/120s.""" ds = DatasetMetadata( conversations=[ ConversationMetadata( @@ -1658,12 +1671,10 @@ async def capture(turn): await strategy.setup_phase() await strategy.execute_phase() - # The runtime-only trace cap does not alter phase-start offsets. - assert issued == [] + assert issued == ["trace_0::sa:a:fa:000"] assert [delay for delay, _ in scheduled] == [ - pytest.approx(100.0), - pytest.approx(130.0), - pytest.approx(220.0), + pytest.approx(30.0), + pytest.approx(120.0), ] for _, coro in scheduled: await coro @@ -1710,40 +1721,47 @@ def _traj(cid: str, offsets: list[float]) -> Trajectory: @pytest.mark.asyncio -async def test_profiling_preserve_start_gap_delays_first_request_by_default(): - """By default (spread), a trajectory's first post-t* request waits out its recorded offset instead of firing at 0: a root resuming 8s after t* is scheduled 8s out, preserving the leading idle gap.""" +async def test_profiling_globally_anchors_earliest_request_preserving_spacing(): + """Offsets 30s/60s become 0s/30s under one phase-wide anchor.""" ds = DatasetMetadata( conversations=[ ConversationMetadata( conversation_id="trace_0", - turns=[ - TurnMetadata(timestamp_ms=0.0), - TurnMetadata(timestamp_ms=18_000.0), - ], + turns=[TurnMetadata(timestamp_ms=30_000.0)], + ), + ConversationMetadata( + conversation_id="trace_1", + turns=[TurnMetadata(timestamp_ms=60_000.0)], ), ], sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, ) - # t* = 10_000: turn 0 < t* (warmed), turn 1 (18s) is 8s out from t*. - root_state = ConversationState( - conversation_id="trace_0", - x_correlation_id="root", - next_turn_index=1, - next_dispatch_offset_ms=8_000.0, - ) - trajectory = Trajectory( - conversation_id="trace_0", - start_turn_index=1, - snapshot=TrajectorySnapshot( - t_star_ms=10_000.0, - states=(root_state,), - ), - ) + trajectories = [ + Trajectory( + conversation_id=conversation_id, + start_turn_index=0, + snapshot=TrajectorySnapshot( + t_star_ms=0.0, + states=( + ConversationState( + conversation_id=conversation_id, + x_correlation_id=correlation_id, + next_turn_index=0, + next_dispatch_offset_ms=offset_ms, + ), + ), + ), + ) + for conversation_id, correlation_id, offset_ms in ( + ("trace_0", "root-0", 30_000.0), + ("trace_1", "root-1", 60_000.0), + ) + ] src = TrajectorySource.__new__(TrajectorySource) src._dataset_metadata = ds src._dataset_sampler = MagicMock() src._metadata_lookup = {c.conversation_id: c for c in ds.conversations} - src.trajectories = [trajectory] + src.trajectories = trajectories issued: list[tuple[str, int]] = [] @@ -1761,7 +1779,7 @@ async def capture(turn): cfg = MagicMock() cfg.phase = CreditPhase.PROFILING - cfg.concurrency = 1 + cfg.concurrency = 2 strategy = AgenticReplayStrategy( config=cfg, conversation_source=src, @@ -1774,12 +1792,11 @@ async def capture(turn): await strategy.setup_phase() await strategy.execute_phase() - # Leading gap preserved: nothing fires inline; turn 1 scheduled 8s out. - assert issued == [] - assert [delay for delay, _ in scheduled] == [pytest.approx(8.0)] + assert issued == [("trace_0", 0)] + assert [delay for delay, _ in scheduled] == [pytest.approx(30.0)] for _, coro in scheduled: await coro - assert issued == [("trace_0", 1)] + assert issued == [("trace_0", 0), ("trace_1", 0)] @pytest.mark.asyncio From f9438058f5a9b8f6c7c02047efb95823a7d23be8 Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Sat, 1 Aug 2026 16:38:31 -0500 Subject: [PATCH 8/9] fix(agentx): preserve flattened timing across warmup handoff Keep one replay origin per active trajectory tree during accelerated warmup. Rebuild every surviving root, child, sidecar, and gated join on that shared dataset clock so non-burst profiling preserves cross-stream order and relative spacing. Retain the existing per-stream delay fallback for timestamp-less datasets and seed recycled warmup roots with a fresh origin. Signed-off-by: Cam Quilici --- docs/tutorials/agentx-mvp.md | 7 + .../timing/strategies/agentic_replay.py | 65 +++++++- .../timing/strategies/test_agentic_replay.py | 143 ++++++++++++++++++ 3 files changed, 211 insertions(+), 4 deletions(-) diff --git a/docs/tutorials/agentx-mvp.md b/docs/tutorials/agentx-mvp.md index 8cb42536fe..eb67d780cd 100644 --- a/docs/tutorials/agentx-mvp.md +++ b/docs/tutorials/agentx-mvp.md @@ -430,6 +430,13 @@ warmup wire requests. For example, `--concurrency 16` with `--warmup-requests-per-lane 10` adds 160 cache-pressure requests after the primers, with a strict 10-additional-request quota on every lane. +At handoff, timestamped root, subagent, and background streams are restored to +their next request's position on one shared per-trajectory dataset clock. The +earliest pending request starts profiling immediately; the remaining streams +retain their flattened cross-stream order and relative start spacing. For a +timestamp-less dataset, AIPerf falls back to each stream's recorded +end-to-start delay. + `--agentic-cache-warmup-duration` and `--warmup-requests-per-lane` are mutually exclusive: choose a time-bounded warmup or a deterministic request-bounded warmup. diff --git a/src/aiperf/timing/strategies/agentic_replay.py b/src/aiperf/timing/strategies/agentic_replay.py index 216f1813d4..cefe33156e 100644 --- a/src/aiperf/timing/strategies/agentic_replay.py +++ b/src/aiperf/timing/strategies/agentic_replay.py @@ -200,6 +200,10 @@ def __init__( self._accelerated_warmup_started = False self._handoff_credits: dict[str, Credit] = {} self._root_to_lane: dict[str, int] = {} + # Accelerated warmup removes idle delays. Keep each tree's sampled t* + # so handoff can restore every stream to one flattened dataset clock; + # otherwise pending child turn-0 requests all look due at offset zero. + self._replay_origin_ms_by_root: dict[str, float] = {} # Cache-bust state. WARMUP and PROFILING construct distinct strategy # instances (PhaseRunner builds a fresh AgenticReplayStrategy per @@ -376,6 +380,7 @@ def _on_tree_drained(self, root_corr: str, phase: CreditPhase | int) -> None: lane = self._correlation_to_lane.pop(root_corr, None) self._session_marker.pop(root_corr, None) self._root_to_lane.pop(root_corr, None) + self._replay_origin_ms_by_root.pop(root_corr, None) if lane is None: self.warning( lambda: ( @@ -671,6 +676,9 @@ async def _execute_warmup(self) -> None: continue t_star_ms = trajectory.snapshot.t_star_ms + root_correlation_id = self._lane_root_corr(trajectory.snapshot) + if root_correlation_id is not None: + self._replay_origin_ms_by_root[root_correlation_id] = t_star_ms for state in trajectory.snapshot.states: warm_index = state.warmup_turn_index if warm_index is None: @@ -830,6 +838,9 @@ async def _dispatch_accelerated_trajectory( return snapshot = trajectory.snapshot + root_correlation_id = self._lane_root_corr(snapshot) + if root_correlation_id is not None: + self._replay_origin_ms_by_root[root_correlation_id] = snapshot.t_star_ms for state in snapshot.states: self._correlation_to_lane[state.x_correlation_id] = lane self._root_to_lane[state.root_correlation_id or state.x_correlation_id] = ( @@ -1037,8 +1048,10 @@ def _returned_credit_handoff_state( conversation_id=credit.conversation_id, x_correlation_id=credit.x_correlation_id, next_turn_index=credit.turn_index + 1, - next_dispatch_offset_ms=self._handoff_delay_ms( - credit.conversation_id, credit.turn_index + 1 + next_dispatch_offset_ms=self._handoff_replay_offset_ms( + credit.effective_root_correlation_id, + credit.conversation_id, + credit.turn_index + 1, ), agent_depth=credit.agent_depth, parent_correlation_id=credit.parent_correlation_id, @@ -1118,8 +1131,10 @@ def _pending_turn_handoff_state( conversation_id=turn.conversation_id, x_correlation_id=turn.x_correlation_id, next_turn_index=turn.turn_index, - next_dispatch_offset_ms=self._handoff_delay_ms( - turn.conversation_id, turn.turn_index + next_dispatch_offset_ms=self._handoff_replay_offset_ms( + turn.effective_root_correlation_id, + turn.conversation_id, + turn.turn_index, ), agent_depth=turn.agent_depth, parent_correlation_id=turn.parent_correlation_id, @@ -1173,6 +1188,39 @@ def _handoff_delay_ms(self, conversation_id: str, next_turn_index: int) -> float previous_duration_ms = max(0.0, previous_api_ms or 0.0) return max(0.0, next_ts_ms - previous_ts_ms - previous_duration_ms) + def _handoff_replay_offset_ms( + self, + root_correlation_id: str, + conversation_id: str, + next_turn_index: int, + ) -> float: + """Place a surviving stream back on its tree's shared replay clock. + + Accelerated warmup compresses runtime waits, so elapsed wall time is + not a usable profiling deadline. Timestamped AgentX datasets already + provide a common clock for every stream in a tree. Subtracting the + tree's sampled replay origin preserves the original flattened order + and spacing across roots, children, sidecars, and gated joins. The + profiling phase later subtracts one phase-wide minimum so the earliest + request starts immediately without turning the rest into a burst. + + Timestamp-less datasets have no shared clock, so retain the legacy + per-stream end-to-start delay as the best available fallback. + """ + replay_origin_ms = self._replay_origin_ms_by_root.get(root_correlation_id) + if replay_origin_ms is not None: + try: + metadata = self.conversation_source.get_metadata(conversation_id) + except KeyError: + metadata = None + if metadata is not None and 0 <= next_turn_index < len(metadata.turns): + next_timestamp_ms = _as_timestamp_ms( + getattr(metadata.turns[next_turn_index], "timestamp_ms", None) + ) + if next_timestamp_ms is not None: + return max(0.0, next_timestamp_ms - replay_origin_ms) + return self._handoff_delay_ms(conversation_id, next_turn_index) + def _build_handoff_replay_boundaries( self, states: list[ConversationState] ) -> tuple[ReplayResumeBoundary, ...]: @@ -1546,6 +1594,7 @@ async def _spawn_from_recycle_or_id( # Prune so every early-return path leaves dicts clean. self._session_marker.pop(finished_correlation_id, None) self._root_to_lane.pop(finished_correlation_id, None) + self._replay_origin_ms_by_root.pop(finished_correlation_id, None) lane = self._release_lane_for(finished_correlation_id, finished_trace_id) await self._dispatch_recycled_on_lane(lane) @@ -1573,6 +1622,14 @@ async def _dispatch_recycled_on_lane(self, lane: int) -> None: self._correlation_to_lane[session.x_correlation_id] = lane self._root_to_lane[session.effective_root_correlation_id] = lane + if self.config.phase == CreditPhase.WARMUP: + first_timestamp_ms = _as_timestamp_ms( + getattr(session.metadata.turns[0], "timestamp_ms", None) + ) + if first_timestamp_ms is not None: + self._replay_origin_ms_by_root[ + session.effective_root_correlation_id + ] = first_timestamp_ms self._mint_marker_for_session( session.effective_root_correlation_id, next_trace_id, lane ) diff --git a/tests/unit/timing/strategies/test_agentic_replay.py b/tests/unit/timing/strategies/test_agentic_replay.py index f468fba6cb..4175e3b5ef 100644 --- a/tests/unit/timing/strategies/test_agentic_replay.py +++ b/tests/unit/timing/strategies/test_agentic_replay.py @@ -923,6 +923,149 @@ def test_cache_warmup_handoff_preserves_pending_replay_barrier_turns() -> None: issuer.replay_gate.pending_turns_by_root.assert_called_once_with() +def test_cache_warmup_handoff_restores_one_flattened_tree_timeline() -> None: + """Pending starts and continuations keep their original cross-stream order.""" + dataset = DatasetMetadata( + conversations=[ + ConversationMetadata( + conversation_id="trace_0", + turns=[ + TurnMetadata(timestamp_ms=0.0), + TurnMetadata(timestamp_ms=20_000.0), + TurnMetadata(timestamp_ms=220_000.0), + ], + ), + ConversationMetadata( + conversation_id="trace_0::fa:000", + turns=[ + TurnMetadata(timestamp_ms=110_000.0), + TurnMetadata(timestamp_ms=160_000.0), + ], + is_root=False, + agent_depth=1, + parent_conversation_id="trace_0", + ), + ConversationMetadata( + conversation_id="trace_0::fa:001", + turns=[TurnMetadata(timestamp_ms=140_000.0)], + is_root=False, + agent_depth=1, + parent_conversation_id="trace_0", + ), + ], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + ) + trajectory = Trajectory( + conversation_id="trace_0", + start_turn_index=1, + snapshot=TrajectorySnapshot( + t_star_ms=10_000.0, + states=( + ConversationState( + conversation_id="trace_0", + x_correlation_id="root", + root_correlation_id="root", + next_turn_index=1, + ), + ), + ), + ) + strategy, issuer, _, _ = _make_strategy( + phase=CreditPhase.WARMUP, + trajectories=[trajectory], + dataset=dataset, + cache_warmup_requests_per_lane=1, + ) + strategy._root_to_lane["root"] = 0 + strategy._replay_origin_ms_by_root["root"] = 10_000.0 + strategy._handoff_credits = { + "root": _make_credit( + conversation_id="trace_0", + x_correlation_id="root", + turn_index=1, + num_turns=3, + phase=CreditPhase.WARMUP, + root_correlation_id="root", + ), + "child-a": _make_credit( + conversation_id="trace_0::fa:000", + x_correlation_id="child-a", + turn_index=0, + num_turns=2, + phase=CreditPhase.WARMUP, + agent_depth=1, + parent_correlation_id="root", + root_correlation_id="root", + branch_mode=ConversationBranchMode.SPAWN, + ), + } + issuer.replay_gate.pending_turns_by_root.return_value = { + "root": ( + TurnToSend( + conversation_id="trace_0::fa:001", + x_correlation_id="child-b", + turn_index=0, + num_turns=1, + agent_depth=1, + parent_correlation_id="root", + root_correlation_id="root", + branch_mode=ConversationBranchMode.SPAWN, + ), + ) + } + + states = strategy._build_handoff_states()[0] + by_correlation = {state.x_correlation_id: state for state in states} + + # All offsets are measured from the same sampled t*=10s. In particular, + # child-b turn 0 is not reconstructed as a zero-delay request. + assert { + correlation_id: state.next_dispatch_offset_ms + for correlation_id, state in by_correlation.items() + } == { + "child-b": pytest.approx(130_000.0), + "child-a": pytest.approx(150_000.0), + "root": pytest.approx(210_000.0), + } + assert [ + state.x_correlation_id + for state in sorted(states, key=lambda item: item.next_dispatch_offset_ms) + ] == ["child-b", "child-a", "root"] + + +@pytest.mark.asyncio +async def test_recycled_warmup_root_gets_a_fresh_flattened_timeline_origin() -> None: + dataset = DatasetMetadata( + conversations=[ + ConversationMetadata( + conversation_id="trace_0", + turns=[ + TurnMetadata(timestamp_ms=50_000.0), + TurnMetadata(timestamp_ms=95_000.0), + ], + ) + ], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + ) + strategy, issuer, _, _ = _make_strategy( + phase=CreditPhase.WARMUP, + trajectories=[Trajectory(conversation_id="trace_0", start_turn_index=0)], + dataset=dataset, + cache_warmup_requests_per_lane=1, + ) + strategy.stop_checker.can_start_new_session.return_value = True + + await strategy._dispatch_recycled_on_lane(0) + + turn = issuer.issue_credit.await_args.args[0] + assert strategy._replay_origin_ms_by_root[ + turn.effective_root_correlation_id + ] == pytest.approx(50_000.0) + assert strategy._handoff_replay_offset_ms( + turn.effective_root_correlation_id, "trace_0", 1 + ) == pytest.approx(45_000.0) + + @pytest.mark.asyncio async def test_warmup_dispatch_uses_start_turn_index(): trajectories = [ From abf55f902cde0a3e8389c452a83ecbf5ba15dea9 Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Sat, 1 Aug 2026 18:38:37 -0500 Subject: [PATCH 9/9] fix(agentx): keep idle watchdogs active across barriers Register the global drain observer when its phase becomes active so warmup teardown cannot clear profiling ownership. When a per-root cap advances a dependency-blocked timer, retain the expired idle budget and immediately try the next timer without bypassing replay barriers. Signed-off-by: Cam Quilici --- src/aiperf/timing/replay_dependencies.py | 17 +++++++- .../timing/strategies/agentic_replay.py | 9 ++++- .../timing/strategies/test_agentic_replay.py | 34 ++++++++++++++++ .../timing/test_replay_barrier_coordinator.py | 40 +++++++++++++++++++ 4 files changed, 97 insertions(+), 3 deletions(-) diff --git a/src/aiperf/timing/replay_dependencies.py b/src/aiperf/timing/replay_dependencies.py index 5d147ece66..a1b15b749f 100644 --- a/src/aiperf/timing/replay_dependencies.py +++ b/src/aiperf/timing/replay_dependencies.py @@ -157,6 +157,8 @@ class _RootBarrierState: """Requests from this runtime tree currently on the wire.""" idle_watchdog: asyncio.TimerHandle | None = None """Per-tree idle-cap callback, armed only while no request is in flight.""" + idle_cap_expired: bool = False + """Whether this idle interval has already consumed its full cap budget.""" class ReplayBarrierCoordinator: @@ -197,6 +199,7 @@ def observe_issued(self, credit: Credit) -> None: return state = self._root_state(credit.effective_root_correlation_id) state.in_flight += 1 + state.idle_cap_expired = False if state.idle_watchdog is not None: state.idle_watchdog.cancel() state.idle_watchdog = None @@ -260,6 +263,13 @@ async def submit( state.pending[key] = _PendingDispatch( turn=turn, issue=issue, on_refused=on_refused ) + # A cap-expired timer can land here because its recorded predecessor + # has not completed yet. The tree is still fully idle, so immediately + # advance the next timer in this tree instead of waiting another full + # cap interval (or never re-arming at all). ``observe_issued`` cancels + # this retry as soon as any candidate actually reaches the wire. + if state.in_flight == 0 and state.idle_cap_expired: + self._arm_root_idle_watchdog(root_id, state) return retained_result def complete(self, credit: Credit) -> None: @@ -279,6 +289,7 @@ def complete(self, credit: Credit) -> None: self._dispatch_tasks.add(task) task.add_done_callback(self._dispatch_tasks.discard) if state.in_flight == 0: + state.idle_cap_expired = False self._arm_root_idle_watchdog(root_id, state) def _arm_root_idle_watchdog(self, root_id: str, state: _RootBarrierState) -> None: @@ -292,7 +303,10 @@ def _arm_root_idle_watchdog(self, root_id: str, state: _RootBarrierState) -> Non ): return loop = asyncio.get_running_loop() - state.idle_watchdog = loop.call_later(cap, self._enforce_root_idle_cap, root_id) + delay = 0.0 if state.idle_cap_expired else cap + state.idle_watchdog = loop.call_later( + delay, self._enforce_root_idle_cap, root_id + ) def _enforce_root_idle_cap(self, root_id: str) -> None: state = self._roots.get(root_id) @@ -301,6 +315,7 @@ def _enforce_root_idle_cap(self, root_id: str) -> None: state.idle_watchdog = None if state.in_flight != 0 or self._scheduler is None: return + state.idle_cap_expired = True shifted = self._scheduler.cap_pending_delay_for_group(root_id, 0.0) if shifted <= 0: return diff --git a/src/aiperf/timing/strategies/agentic_replay.py b/src/aiperf/timing/strategies/agentic_replay.py index cefe33156e..2bd8a8d082 100644 --- a/src/aiperf/timing/strategies/agentic_replay.py +++ b/src/aiperf/timing/strategies/agentic_replay.py @@ -259,8 +259,6 @@ def __init__( self._system_idle_seconds_skipped = 0.0 self._system_idle_started_at: float | None = None self._system_idle_watchdog: asyncio.TimerHandle | None = None - if self._system_idle_gap_cap_seconds is not None: - self.scheduler.set_drain_observer(self.enforce_system_idle_cap) # Wrap-fill + cache_bust=NONE produces byte-identical traffic across # shared-trace lanes. agentx-mvp auto-locks cache_bust=first_turn_prefix @@ -404,6 +402,13 @@ async def setup_phase(self) -> None: dataset's ``sampling_strategy`` and reuses every root about equally -- no strategy-side recycle queue to keep in sync. """ + # Strategies for later phases may be constructed before the current + # phase finalizes. Register the observer when this phase becomes + # active, not in ``__init__``; otherwise warmup teardown can clear the + # already-constructed profiling observer and a barrier-retained timer + # can leave the whole system idle past the global cap. + if self._system_idle_gap_cap_seconds is not None: + self.scheduler.set_drain_observer(self.enforce_system_idle_cap) if self._has_tree_registry: self._session_tree_registry.set_drain_callback(self._on_tree_drained) if ( diff --git a/tests/unit/timing/strategies/test_agentic_replay.py b/tests/unit/timing/strategies/test_agentic_replay.py index 4175e3b5ef..f92481e1b0 100644 --- a/tests/unit/timing/strategies/test_agentic_replay.py +++ b/tests/unit/timing/strategies/test_agentic_replay.py @@ -2618,6 +2618,40 @@ async def replay_dispatch() -> None: scheduler.cancel_all() +@pytest.mark.asyncio +async def test_active_phase_owns_global_idle_drain_observer() -> None: + """Warmup teardown cannot clear a preconstructed profiling observer.""" + scheduler = LoopScheduler() + run = _make_run(target=CacheBustTarget.FIRST_TURN_PREFIX) + run.cfg.get_profiling_phases()[0].system_idle_gap_cap_seconds = 0.05 + trajectories = [Trajectory(conversation_id="trace_0", start_turn_index=0)] + warmup, _, _, _ = _make_strategy( + phase=CreditPhase.WARMUP, + trajectories=trajectories, + scheduler=scheduler, + run=run, + ) + profiling, _, _, _ = _make_strategy( + phase=CreditPhase.PROFILING, + trajectories=trajectories, + scheduler=scheduler, + run=run, + ) + + # Construct both phases first, matching TimingManager lifecycle. Observer + # ownership follows setup/finalize of the active phase, not construction. + assert scheduler._drain_observer is None + await warmup.setup_phase() + assert getattr(scheduler._drain_observer, "__self__", None) is warmup + await warmup.finalize_phase() + assert scheduler._drain_observer is None + + await profiling.setup_phase() + assert getattr(scheduler._drain_observer, "__self__", None) is profiling + await profiling.finalize_phase() + assert scheduler._drain_observer is None + + @pytest.mark.asyncio async def test_global_idle_cap_rechecks_after_barrier_defers_near_timer( time_traveler_no_patch_sleep, diff --git a/tests/unit/timing/test_replay_barrier_coordinator.py b/tests/unit/timing/test_replay_barrier_coordinator.py index 4028b12855..c4bf850805 100644 --- a/tests/unit/timing/test_replay_barrier_coordinator.py +++ b/tests/unit/timing/test_replay_barrier_coordinator.py @@ -230,6 +230,46 @@ async def test_new_request_cancels_runtime_root_idle_watchdog() -> None: scheduler.cap_pending_delay_for_group.assert_not_called() +@pytest.mark.asyncio +async def test_idle_watchdog_retries_when_advanced_turn_is_barrier_blocked() -> None: + """A blocked candidate cannot consume the cap and strand later timers.""" + scheduler = LoopScheduler() + coordinator = ReplayBarrierCoordinator( + _metadata(), + scheduler=scheduler, + root_idle_gap_cap_seconds=0.02, + ) + coordinator.activate() + issued = asyncio.Event() + + async def submit_blocked() -> None: + await coordinator.submit(_turn("d"), lambda: _record_issue([], "d")) + + async def submit_ready() -> None: + async def mark_issued() -> bool: + issued.set() + return True + + await coordinator.submit( + _turn("a"), + mark_issued, + ) + + # The first cap advance lands on d, which is retained behind a/b/c. The + # watchdog must immediately retry and advance ready turn a as part of the + # same already-expired idle interval. + scheduler.schedule_later(1.0, submit_blocked(), group_id="root") + scheduler.schedule_later(2.0, submit_ready(), group_id="root") + started = asyncio.get_running_loop().time() + coordinator.observe_idle_root("root") + + await asyncio.wait_for(issued.wait(), timeout=0.15) + + assert asyncio.get_running_loop().time() - started < 0.10 + await coordinator.cancel_pending(notify_refused=False) + scheduler.cancel_all() + + @pytest.mark.asyncio async def test_scalar_peak_would_slip_d_after_only_one_completion() -> None: coordinator = ReplayBarrierCoordinator(_metadata())