From e50e823b4970dc5148773749c3ee9dbe5fba8a29 Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Thu, 25 Jun 2026 00:41:03 -0500 Subject: [PATCH 1/2] feat(agentic): replay api-time overlap barriers Signed-off-by: Cam Quilici --- docs/architecture.md | 2 +- .../benchmark-modes/timing-modes-reference.md | 14 + src/aiperf/common/models/__init__.py | 2 + src/aiperf/common/models/dataset_models.py | 39 +++ src/aiperf/credit/issuer.py | 28 +- src/aiperf/dataset/loader/weka_trace.py | 63 +++- src/aiperf/dataset/mmap_cache.py | 7 +- src/aiperf/timing/branch_orchestrator.py | 64 +++- src/aiperf/timing/phase/runner.py | 18 + src/aiperf/timing/replay_dependencies.py | 312 ++++++++++++++++++ .../timing/strategies/agentic_replay.py | 6 + .../test_agentic_replay_overlap_barrier.py | 93 ++++++ .../abc_join_d.json | 14 + .../loader/test_weka_overlap_groups.py | 76 +++++ tests/unit/dataset/test_mmap_cache.py | 13 + .../test_branch_orchestrator_delayed.py | 54 ++- .../timing/test_replay_barrier_coordinator.py | 141 ++++++++ tests/unit/timing/test_replay_dependencies.py | 108 ++++++ 18 files changed, 1046 insertions(+), 8 deletions(-) create mode 100644 src/aiperf/timing/replay_dependencies.py create mode 100644 tests/component_integration/test_agentic_replay_overlap_barrier.py create mode 100644 tests/fixtures/weka_traces_overlap_groups/abc_join_d.json create mode 100644 tests/unit/dataset/loader/test_weka_overlap_groups.py create mode 100644 tests/unit/timing/test_replay_barrier_coordinator.py create mode 100644 tests/unit/timing/test_replay_dependencies.py diff --git a/docs/architecture.md b/docs/architecture.md index 999ccad2d..8416febd1 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -190,7 +190,7 @@ The `CreditIssuer` is timing-mode agnostic; the strategy in `src/aiperf/timing/s - **Fixed schedule** (`fixed_schedule.py`): replay trace timestamps from dataset metadata. - **Request-rate** (`request_rate.py`): issue at a target rate with constant / Poisson / gamma / concurrency-burst arrival patterns. - **User-centric rate** (`user_centric_rate.py`): each session is an independent user; turn gaps come from the trace. -- **Agentic replay** (`agentic_replay.py`): scenario-driven DAG replay where children are dispatched on parent completion via the `BranchOrchestrator`. +- **Agentic replay** (`agentic_replay.py`): scenario-driven DAG replay. Weka loaders persist interval-order predecessor frontiers derived from `[t, t + api_time]`; `ReplayBarrierCoordinator` enforces their fan-out/join barriers, while `BranchOrchestrator` starts branches from the parent send when their recorded intervals overlap and otherwise uses the normal completion path. #### Relationship to `--request-count`, `--num-conversations`, Concurrency diff --git a/docs/benchmark-modes/timing-modes-reference.md b/docs/benchmark-modes/timing-modes-reference.md index 4ff3eea09..9c7e17b0d 100644 --- a/docs/benchmark-modes/timing-modes-reference.md +++ b/docs/benchmark-modes/timing-modes-reference.md @@ -246,6 +246,20 @@ aiperf profile \ **How it works:** The strategy picks `--concurrency` distinct conversations as *trajectories*, samples a per-trajectory starting turn `k_i` somewhere between 25% and 75% of each conversation (the default `--trajectory-start-min-ratio` / `--trajectory-start-max-ratio` window, clamped to leave at least one profile turn after warmup), and warms each trajectory by dispatching that one turn before profiling starts. During profiling, each trajectory resumes from `k_i + 1` and replays the remaining turns honoring the trace's recorded request-start schedule after applying the per-trace idle-gap rule. The default `--trace-idle-gap-cap-seconds` is `None` (no compression); the `inferencex-agentx-mvp` scenario locks it to `60` so coffee-break request-start gaps don't distort steady-state while preserving local subagent overlap. +Weka replay also preserves the capture's fan-out/join shape. The loader compares +request intervals `[t, t + api_time]` within each logical agent or subagent +scope and records an explicit cross-stream completion frontier on every turn. +Requests whose intervals overlap have no ordering edge and may execute in +parallel. A later request waits until every request on its recorded predecessor +frontier reaches a terminal outcome. Exact interval-boundary touches are +sequential. Long transitive overlaps use this dependency frontier rather than +an overlap connected-component, so a long request can overlap multiple +sequential requests on another stream without launching all of them at once. +Branches that began while their spawning request was in flight are scheduled +from that request's send time; independent subagent scopes are not globally +joined. The same barriers remain active during accelerated cache-pressure +warmup, where idle delays are otherwise removed. + Concurrency is **per session tree**: each `--concurrency` lane holds one slot for a whole tree — the root conversation plus every subagent it spawns (children, subchildren, background `::fa:`/`::aux:` sidecars). A lane's slot is released, and the lane recycled into a fresh root, only once the entire tree drains (root terminal **and** all descendants returned) — so a background subagent that outlives its root does not free the lane early. Recycle then draws the next root from the dataset sampler (honoring the dataset's `sampling_strategy` — sequential / shuffle / random), starting it from turn 0. This keeps exactly `--concurrency` trees live at all times. The shared tree id (`root_correlation_id`) is persisted per record in `profile_export.jsonl`, so `aiperf analyze swim-lane` groups each tree under one lane and renders exactly `--concurrency` slots. **When to use:** A scenario-locked timing mode for multi-turn agentic-coding traces (currently WEKA), especially long runs where you want steady-state metrics rather than first-turn-only metrics. Pairs naturally with `--cache-bust first_turn_prefix` (auto-injected by the `inferencex-agentx-mvp` scenario) so recycled plays don't progressively warm the server's KV-cache prefix on identical content. diff --git a/src/aiperf/common/models/__init__.py b/src/aiperf/common/models/__init__.py index 48e0dcf64..b477c7b31 100644 --- a/src/aiperf/common/models/__init__.py +++ b/src/aiperf/common/models/__init__.py @@ -20,6 +20,7 @@ InputsFile, Media, MemoryMapClientMetadata, + ReplayTurnReference, SessionPayloads, Text, Turn, @@ -200,6 +201,7 @@ "Media", "MediaCounts", "MemoryMapClientMetadata", + "ReplayTurnReference", "MetricFamily", "MetricRecordInfo", "MetricRecordMetadata", diff --git a/src/aiperf/common/models/dataset_models.py b/src/aiperf/common/models/dataset_models.py index 49f4e6fd3..281b44d44 100644 --- a/src/aiperf/common/models/dataset_models.py +++ b/src/aiperf/common/models/dataset_models.py @@ -111,6 +111,13 @@ class Video(Media): media_type: ClassVar[MediaTypeT] = MediaType.VIDEO +class ReplayTurnReference(AIPerfBaseModel): + """Dataset-stable reference to one request in a replay dependency graph.""" + + conversation_id: str = Field(description="Referenced conversation ID.") + turn_index: int = Field(ge=0, description="Referenced turn index.") + + class TurnMetadata(AIPerfBaseModel): """Metadata of a turn.""" @@ -134,6 +141,13 @@ class TurnMetadata(AIPerfBaseModel): "without per-request timing." ), ) + replay_predecessors: list["ReplayTurnReference"] = Field( + default_factory=list, + description=( + "Cross-stream requests that reached a recorded terminal outcome before " + "this request began and must complete before agentic replay may issue it." + ), + ) branch_ids: list[str] = Field( default_factory=list, description="Branch IDs triggered after this turn completes (DAG projection).", @@ -207,6 +221,14 @@ class Turn(AIPerfBaseModel): "duration (not warped). None for loaders without per-request timing." ), ) + replay_predecessors: list["ReplayTurnReference"] = Field( + default_factory=list, + exclude=True, + description=( + "Explicit cross-stream completion frontier inferred from recorded " + "request intervals by trace-aware loaders." + ), + ) max_tokens: int | None = Field( default=None, description="Maximum number of tokens to generate for this turn." ) @@ -309,6 +331,7 @@ def metadata(self) -> TurnMetadata: timestamp_ms=self.timestamp, delay_ms=self.delay, api_time_ms=self.api_time_ms, + replay_predecessors=self.replay_predecessors, branch_ids=self.branch_ids, prerequisites=self.prerequisites, raw_messages_count=None @@ -372,6 +395,13 @@ class ConversationMetadata(AIPerfBaseModel): default=None, description="For DAG children: the parent conversation ID.", ) + replay_scope_id: str | None = Field( + default=None, + description=( + "Logical agent/subagent scope whose request intervals participate in " + "one replay dependency graph. Independent scopes are never joined." + ), + ) accuracy_ground_truth: str | None = Field( default=None, description="Ground-truth answer for this conversation (accuracy mode only). " @@ -499,6 +529,13 @@ def _reject_unimplemented_context_mode( default=None, description="For DAG children: the parent conversation ID.", ) + replay_scope_id: str | None = Field( + default=None, + exclude=True, + description=( + "Logical agent/subagent scope used to infer cross-stream replay barriers." + ), + ) accuracy_ground_truth: str | None = Field( default=None, description="Ground-truth answer for this conversation (accuracy mode only). " @@ -534,6 +571,7 @@ def metadata(self) -> ConversationMetadata: agent_depth=self.agent_depth, subagent_type=self.subagent_type, parent_conversation_id=self.parent_conversation_id, + replay_scope_id=self.replay_scope_id, accuracy_ground_truth=self.accuracy_ground_truth, accuracy_task=self.accuracy_task, ) @@ -554,6 +592,7 @@ def to_metadata(self) -> "ConversationMetadata": agent_depth=self.agent_depth, subagent_type=self.subagent_type, parent_conversation_id=self.parent_conversation_id, + replay_scope_id=self.replay_scope_id, accuracy_ground_truth=self.accuracy_ground_truth, accuracy_task=self.accuracy_task, ) diff --git a/src/aiperf/credit/issuer.py b/src/aiperf/credit/issuer.py index 5b1fc9ec1..107028556 100644 --- a/src/aiperf/credit/issuer.py +++ b/src/aiperf/credit/issuer.py @@ -22,6 +22,7 @@ from aiperf.common.aiperf_logger import AIPerfLogger from aiperf.common.enums import CreditPhase from aiperf.credit.structs import Credit, TurnToSend +from aiperf.timing.replay_dependencies import ReplayIssueGate from aiperf.timing.url_samplers import URLSelectionStrategyProtocol if TYPE_CHECKING: @@ -32,6 +33,7 @@ from aiperf.timing.phase.lifecycle import PhaseLifecycle from aiperf.timing.phase.progress_tracker import PhaseProgressTracker from aiperf.timing.phase.stop_conditions import StopConditionChecker + from aiperf.timing.replay_dependencies import ReplayBarrierCoordinator from aiperf.timing.request_cancellation import RequestCancellationSimulator from aiperf.timing.session_tree import SessionTreeRegistry @@ -71,6 +73,7 @@ def __init__( url_selection_strategy: URLSelectionStrategyProtocol | None = None, session_tree_registry: SessionTreeRegistry | None = None, session_tree_registry_enabled: bool | None = None, + replay_barrier: ReplayBarrierCoordinator | None = None, ) -> None: """Initialize credit issuer. @@ -108,6 +111,7 @@ def __init__( else None ) self._issuing_stopped = False + self.replay_gate = ReplayIssueGate(replay_barrier) def stop_issuing(self) -> None: """Refuse every subsequent root and child credit.""" @@ -217,6 +221,11 @@ async def issue_credit(self, turn: TurnToSend) -> bool: 5. Create and send Credit 6. If final credit: freeze counts + set event """ + gate = getattr(self, "replay_gate", ReplayIssueGate(None)) + return await gate.submit(turn, lambda: self._issue_credit_ready(turn)) + + async def _issue_credit_ready(self, turn: TurnToSend) -> bool: + """Issue a turn whose recorded predecessor frontier is complete.""" if self._issuing_stopped: return False @@ -369,6 +378,9 @@ async def _issue_credit_internal(self, turn: TurnToSend) -> bool: ) await self._credit_router.send_credit(credit=credit) + replay_gate = getattr(self, "replay_gate", None) + if replay_gate is not None: + await replay_gate.observe_issued(credit) if is_final_credit: self._progress.freeze_sent_counts() self._progress.all_credits_sent_event.set() @@ -407,6 +419,15 @@ async def dispatch_child_turn(self, turn: TurnToSend) -> bool: attempt would send the first sibling and permanently truncate every other sibling spawned in the same gather. """ + gate = getattr(self, "replay_gate", ReplayIssueGate(None)) + return await gate.submit( + turn, + lambda: self._dispatch_child_turn_ready(turn), + child_refusal_cleanup=True, + ) + + async def _dispatch_child_turn_ready(self, turn: TurnToSend) -> bool: + """Dispatch a child after its recorded predecessor frontier completes.""" if self._issuing_stopped: return False can_proceed_fn = self._stop_checker.can_send_child_turn @@ -467,5 +488,8 @@ async def dispatch_join_turn(self, pending: PendingBranchJoin) -> bool: cache_bust_marker=pending.parent_cache_bust_marker, cache_bust_target=pending.parent_cache_bust_target, ) - result = await self.try_issue_credit(turn) - return result is True + replay_gate = getattr(self, "replay_gate", None) + if replay_gate is None or not replay_gate.enabled: + result = await self.try_issue_credit(turn) + return result is True + return await self.issue_credit(turn) diff --git a/src/aiperf/dataset/loader/weka_trace.py b/src/aiperf/dataset/loader/weka_trace.py index c62dc7716..aaa5d4dbf 100644 --- a/src/aiperf/dataset/loader/weka_trace.py +++ b/src/aiperf/dataset/loader/weka_trace.py @@ -25,7 +25,7 @@ from aiperf.common.enums import ConversationContextMode, TurnInputKind from aiperf.common.environment import Environment from aiperf.common.exceptions import DatasetLoaderError -from aiperf.common.models import Conversation +from aiperf.common.models import Conversation, ReplayTurnReference from aiperf.dataset.generator.prompt import PromptGenerator from aiperf.dataset.loader._delay_cap import DelayCapTracker from aiperf.dataset.loader.base_loader import BaseFileLoader @@ -44,6 +44,56 @@ _JOIN_EPSILON_SECONDS = 1e-6 +def _replay_scope_for_session(session_id: str, parent_trace_id: str) -> str: + """Keep each captured agent/subagent's interval graph independent.""" + marker = "::sa:" + if marker not in session_id: + return parent_trace_id + trace_id, suffix = session_id.split(marker, 1) + agent_id = suffix + for worker_marker in (":aux:red:", ":aux:", ":fa:", ":wg:"): + if worker_marker in agent_id: + agent_id = agent_id.rsplit(worker_marker, 1)[0] + return f"{trace_id}{marker}{agent_id}" + + +def _install_replay_dependencies(conversations: list[Conversation]) -> None: + """Persist cross-stream interval frontiers on reconstructed Weka turns.""" + from aiperf.timing.replay_dependencies import ( + RecordedTurnInterval, + ReplayTurnKey, + infer_cross_stream_predecessors, + ) + + by_scope: dict[str, list[RecordedTurnInterval]] = defaultdict(list) + turns_by_key = {} + for conversation in conversations: + scope_id = conversation.replay_scope_id + if scope_id is None: + continue + for turn_index, turn in enumerate(conversation.turns): + key = ReplayTurnKey(conversation.session_id, turn_index) + turns_by_key[key] = turn + by_scope[scope_id].append( + RecordedTurnInterval( + key=key, + stream_id=conversation.session_id, + start_ms=turn.timestamp, + api_time_ms=turn.api_time_ms, + ) + ) + + for intervals in by_scope.values(): + for key, predecessors in infer_cross_stream_predecessors(intervals).items(): + turns_by_key[key].replay_predecessors = [ + ReplayTurnReference( + conversation_id=predecessor.conversation_id, + turn_index=predecessor.turn_index, + ) + for predecessor in predecessors + ] + + def _subagent_request_absolute_t( entry: WekaSubagentEntry, req: WekaNormalRequest ) -> float: @@ -1595,6 +1645,8 @@ def convert_to_conversations( # that share the same PromptGenerator. self.prompt_generator._cache.clear() + _install_replay_dependencies(conversations) + from aiperf.common.models import DatasetMetadata from aiperf.common.validators.orchestrator_v1 import ( validate_for_orchestrator_v1, @@ -1689,6 +1741,7 @@ def _reconstruct_serial( conv = Conversation( session_id=plan.trace_id, context_mode=self._resolved_context_mode(), + replay_scope_id=plan.trace_id, ) recon = ConversationReconstructor( block_size=plan.block_size, @@ -2036,6 +2089,9 @@ def _reconstruct_serial( is_root=False, agent_depth=1, parent_conversation_id=cp.parent_trace_id, + replay_scope_id=_replay_scope_for_session( + cp.session_id, cp.parent_trace_id + ), ) child_metric_values = metric_values_by_trace[cp.parent_trace_id] for k, creq in enumerate(cp.requests): @@ -2138,6 +2194,7 @@ def _emit_flat_chain_conversation( is_root=False, agent_depth=1, parent_conversation_id=fp.parent_trace_id, + replay_scope_id=fp.parent_trace_id, ) from aiperf.common.models import Turn @@ -2625,6 +2682,7 @@ def _reconstruct_parallel( parent_conv = Conversation( session_id=trace_id, context_mode=self._resolved_context_mode(), + replay_scope_id=trace_id, ) for t_dict in result["parent_turns"]: parent_conv.turns.append( @@ -2678,6 +2736,9 @@ def _reconstruct_parallel( parent_conversation_id=child.get( "parent_conversation_id", result["trace_id"] ), + replay_scope_id=_replay_scope_for_session( + child["session_id"], result["trace_id"] + ), ) for t_dict in child["turns"]: child_conv.turns.append( diff --git a/src/aiperf/dataset/mmap_cache.py b/src/aiperf/dataset/mmap_cache.py index c54a35745..b0ed4b41f 100644 --- a/src/aiperf/dataset/mmap_cache.py +++ b/src/aiperf/dataset/mmap_cache.py @@ -75,6 +75,11 @@ # Version 5 fixed the Conversation.metadata() projection of per-turn # theoretical prefix-cache block counts for realtime infinite-cache hit rate. MANIFEST_VERSION = ( + # v21: Weka traces carry explicit api_time interval-frontier metadata in + # DatasetMetadata (replay_scope_id + per-turn replay_predecessors). Cached + # manifests produced before v21 deserialize with empty defaults, silently + # disabling fan-out/join barriers even though dataset.dat remains usable. + # Rebuild so the manifest sidecar contains the inferred dependency graph. # v20: DAG datasets (any FORK/SPAWN branch) are no longer preformatted into # the PAYLOAD_BYTES mmap fast path -- they are delta-compressed and # accumulate context across the tree (FORK children seed from the parent's @@ -137,7 +142,7 @@ # v10: merge of the flattened-agent-splitting lineage and the # tool-shaping lineage (boundary-cut overhang strip; shaping decided at # first emission so reset re-emits reproduce the first-sent shape). - 20 + 21 ) MANIFEST_FILENAME = "manifest.json" INPUTS_JSON_FILENAME = "inputs.json" diff --git a/src/aiperf/timing/branch_orchestrator.py b/src/aiperf/timing/branch_orchestrator.py index d42cfb99a..24e964e53 100644 --- a/src/aiperf/timing/branch_orchestrator.py +++ b/src/aiperf/timing/branch_orchestrator.py @@ -279,6 +279,7 @@ def __init__( # children are not dispatched a second time when the parent's # turn 0 credit returns. self._pre_dispatched_branches: set[tuple[str, str]] = set() + self._overlap_dispatched_branches: set[tuple[str, str]] = set() self._fail_fast = Environment.DAG.FAIL_FAST self._cleaning_up: bool = False # SPAWN children whose recorded first request starts after the branch @@ -328,6 +329,14 @@ def start_accelerated_warmup(self) -> None: if self._allow_accelerated_warmup: self._accelerated_warmup_started = True + def close_replay_root(self, root_correlation_id: str) -> None: + """Discard per-instance overlap markers after a tree drains.""" + self._overlap_dispatched_branches = { + item + for item in self._overlap_dispatched_branches + if item[0] != root_correlation_id + } + def snapshot_annotations( self, ) -> tuple[dict[str, int], dict[str, tuple[str | None, int | None]]]: @@ -394,6 +403,44 @@ def get_branch_ids(self, credit) -> list[str]: return [] return list(meta.turns[credit.turn_index].branch_ids) + async def on_credit_issued(self, credit) -> None: + """Start branches that overlapped their spawning request in the capture.""" + if self._cleaning_up or credit.agent_depth > 0: + return + if credit.phase == CreditPhase.WARMUP and not self._accelerated_warmup_started: + return + parent_meta = self._cs.get_metadata(credit.conversation_id) + if getattr(parent_meta, "replay_scope_id", None) is None: + return + if credit.turn_index >= len(parent_meta.turns): + return + turn_meta = parent_meta.turns[credit.turn_index] + parent_start_ms = _as_timestamp_ms(turn_meta.timestamp_ms) + parent_api_ms = _as_timestamp_ms(turn_meta.api_time_ms) + if parent_start_ms is None or parent_api_ms is None or parent_api_ms <= 0: + return + parent_end_ms = parent_start_ms + parent_api_ms + branches_by_id = {branch.branch_id: branch for branch in parent_meta.branches} + overlapping = [ + branch_id + for branch_id in turn_meta.branch_ids + if (branch := branches_by_id.get(branch_id)) is not None + and (branch_start := self._branch_start_timestamp_ms(branch)) is not None + and branch_start < parent_end_ms + ] + if not overlapping: + return + parent_corr = credit.x_correlation_id + async with self._parent_locks[parent_corr]: + await self._spawn_children_and_register_gates( + credit, + overlapping, + dispatch_origin_ms=parent_start_ms, + ) + self._overlap_dispatched_branches.update( + (parent_corr, branch_id) for branch_id in overlapping + ) + def _marker_for_root(self, root_correlation_id: str | None) -> str | None: """Resolve the tree-root cache-bust marker for a spawned descendant. @@ -711,7 +758,11 @@ async def intercept(self, credit) -> bool: return self._maybe_suspend_parent(credit) async def _spawn_children_and_register_gates( - self, credit, branch_ids: list[str] + self, + credit, + branch_ids: list[str], + *, + dispatch_origin_ms: float | None = None, ) -> None: """Resolve branches, start children, and register future joins. @@ -753,6 +804,11 @@ async def _spawn_children_and_register_gates( # parent's turn-0 return to avoid double-dispatch. if (credit.conversation_id, b_id) in self._pre_dispatched_branches: continue + if ( + dispatch_origin_ms is None + and (parent_corr, b_id) in self._overlap_dispatched_branches + ): + continue branch_gates = gate_for_branch.get(branch.branch_id, []) # Background branches never gate the parent even if the dataset # authored a spawning turn for them (the validator would have @@ -796,7 +852,10 @@ async def _spawn_children_and_register_gates( per_child_branch_mode[child_corr] = branch.mode per_child_gates[child_corr] = list(branch_gates) dispatch_offset_by_corr[child_corr] = self._child_dispatch_offset_ms( - branch_start_ms, child + dispatch_origin_ms + if dispatch_origin_ms is not None + else branch_start_ms, + child, ) all_children.append(child) @@ -1519,6 +1578,7 @@ def cleanup(self) -> None: self._descendant_counts.clear() self._parent_locks.clear() self._pre_dispatched_branches.clear() + self._overlap_dispatched_branches.clear() def any_child_tracked_for_parent( diff --git a/src/aiperf/timing/phase/runner.py b/src/aiperf/timing/phase/runner.py index d91a1ec1e..1cb2a518c 100644 --- a/src/aiperf/timing/phase/runner.py +++ b/src/aiperf/timing/phase/runner.py @@ -25,6 +25,7 @@ from aiperf.timing.phase.progress_tracker import PhaseProgressTracker from aiperf.timing.phase.stop_conditions import StopConditionChecker from aiperf.timing.ramping import RampConfig, Ramper, RampType +from aiperf.timing.replay_dependencies import ReplayBarrierCoordinator from aiperf.timing.strategies.core import RateSettableProtocol from aiperf.timing.trajectory_source import TrajectorySource from aiperf.timing.url_samplers import URLSelectionStrategyProtocol @@ -172,6 +173,11 @@ def __init__( lifecycle=self._lifecycle, counter=self._progress.counter, ) + self._replay_barrier = ( + ReplayBarrierCoordinator(self._conversation_source.dataset_metadata) + if self._config.timing_mode == TimingMode.AGENTIC_REPLAY + else None + ) self._credit_issuer = CreditIssuer( phase=self._config.phase, stop_checker=self._stop_checker, @@ -185,6 +191,7 @@ def __init__( session_tree_registry_enabled=( self._config.phase == CreditPhase.PROFILING or cache_warmup_enabled ), + replay_barrier=self._replay_barrier, ) self._branch_orchestrator = BranchOrchestrator( conversation_source=self._conversation_source, @@ -206,6 +213,13 @@ def __init__( ), allow_accelerated_warmup=(cache_warmup_enabled), ) + self._credit_issuer.replay_gate.set_child_refused( + self._branch_orchestrator.on_child_stopped + ) + if self._replay_barrier is not None: + self._credit_issuer.replay_gate.set_credit_issued( + self._branch_orchestrator.on_credit_issued + ) self._callback_handler.set_branch_orchestrator(self._branch_orchestrator) # Execution state @@ -633,6 +647,10 @@ async def _wait_for_sending_complete(self) -> None: self._scheduler.cancel_all_pending() self._progress.all_credits_sent_event.set() + await self._credit_issuer.replay_gate.cancel( + notify_refused=self._config.phase == CreditPhase.PROFILING + ) + stats = self._progress.create_stats(self._lifecycle) self.notice(self._format_phase_sending_complete(stats)) await self._phase_publisher.publish_progress(stats) diff --git a/src/aiperf/timing/replay_dependencies.py b/src/aiperf/timing/replay_dependencies.py new file mode 100644 index 000000000..929946099 --- /dev/null +++ b/src/aiperf/timing/replay_dependencies.py @@ -0,0 +1,312 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Recorded interval-order dependencies for agentic replay.""" + +from __future__ import annotations + +import asyncio +import logging +import math +from collections import Counter +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from aiperf.common.models import DatasetMetadata + from aiperf.credit.structs import Credit, TurnToSend + +_logger = logging.getLogger(__name__) + + +@dataclass(frozen=True, slots=True, order=True) +class ReplayTurnKey: + """Stable dataset identity for one replayed request.""" + + conversation_id: str + turn_index: int + + +@dataclass(frozen=True, slots=True) +class RecordedTurnInterval: + """One request interval on a logical replay stream.""" + + key: ReplayTurnKey + stream_id: str + start_ms: float | None + api_time_ms: float | None + + @property + def normalized_interval(self) -> tuple[float, float] | None: + """Return ``[start, end]`` using the Weka duration fallback policy.""" + if self.start_ms is None or not math.isfinite(self.start_ms): + return None + duration_ms = self.api_time_ms + if duration_ms is None or not math.isfinite(duration_ms) or duration_ms < 0: + duration_ms = 0.0 + return self.start_ms, self.start_ms + duration_ms + + +def infer_cross_stream_predecessors( + intervals: list[RecordedTurnInterval], +) -> dict[ReplayTurnKey, tuple[ReplayTurnKey, ...]]: + """Infer the recorded completion frontier each request must join. + + Per-stream ordering remains owned by normal conversation replay. For every + other stream, a request depends on that stream's latest request known to + have completed by its recorded start. Overlapping intervals create no edge. + This represents transitive overlap precisely: a long request may overlap + several sequential requests on another stream without forcing those later + requests into one simultaneously launched connected component. + + Exact boundary touches are ordered. Equal starts are unordered, including + zero-width intervals. Missing or non-finite starts add no cross-stream edge; + missing, negative, or non-finite durations are deterministic zero-width + intervals, matching the Weka loader's request-end fallback. + """ + by_stream: dict[str, list[tuple[RecordedTurnInterval, float, float]]] = {} + for interval in intervals: + normalized = interval.normalized_interval + if normalized is None: + continue + start_ms, end_ms = normalized + by_stream.setdefault(interval.stream_id, []).append( + (interval, start_ms, end_ms) + ) + + dependencies: dict[ReplayTurnKey, tuple[ReplayTurnKey, ...]] = {} + for target in intervals: + target_interval = target.normalized_interval + if target_interval is None: + dependencies[target.key] = () + continue + target_start_ms, _ = target_interval + frontier: list[tuple[RecordedTurnInterval, float, float]] = [] + for stream_id, candidates in by_stream.items(): + if stream_id == target.stream_id: + continue + completed = [ + candidate + for candidate in candidates + if candidate[1] < target_start_ms and candidate[2] <= target_start_ms + ] + if not completed: + continue + latest = max( + completed, + key=lambda candidate: ( + candidate[2], + candidate[1], + candidate[0].key, + ), + ) + frontier.append(latest) + predecessors = [ + candidate[0].key + for candidate in frontier + if not any( + candidate[1] < later[1] and candidate[2] <= later[1] + for later in frontier + if later is not candidate + ) + ] + dependencies[target.key] = tuple(sorted(predecessors)) + return dependencies + + +@dataclass(slots=True) +class _PendingDispatch: + issue: Callable[[], Awaitable[bool]] + on_refused: Callable[[], Awaitable[None]] | None + + +@dataclass(slots=True) +class _RootBarrierState: + completed: set[ReplayTurnKey] + pending: dict[ReplayTurnKey, _PendingDispatch] + initialized: bool = False + + +class ReplayBarrierCoordinator: + """Release requests only after their recorded frontier has completed.""" + + def __init__(self, dataset_metadata: DatasetMetadata) -> None: + self._predecessors: dict[ReplayTurnKey, tuple[ReplayTurnKey, ...]] = {} + self._timestamps: dict[ReplayTurnKey, float | None] = {} + for conversation in dataset_metadata.conversations: + for turn_index, turn in enumerate(conversation.turns): + key = ReplayTurnKey(conversation.conversation_id, turn_index) + self._predecessors[key] = tuple( + ReplayTurnKey(ref.conversation_id, ref.turn_index) + for ref in turn.replay_predecessors + ) + timestamp = turn.timestamp_ms + self._timestamps[key] = ( + float(timestamp) + if isinstance(timestamp, int | float) and math.isfinite(timestamp) + else None + ) + self._roots: dict[str, _RootBarrierState] = {} + self._dispatch_tasks: set[asyncio.Task] = set() + self._active = False + + def activate(self) -> None: + """Enable barriers after baseline cache priming completes.""" + if self._active: + return + self._active = True + widths = Counter( + len(predecessors) + for predecessors in self._predecessors.values() + if predecessors + ) + _logger.info( + "Replay interval barriers active: %d requests, %d gated turns, " + "join-widths=%s", + len(self._predecessors), + sum(widths.values()), + dict(sorted(widths.items())), + ) + + async def submit( + self, + turn: TurnToSend, + issue: Callable[[], Awaitable[bool]], + *, + on_refused: Callable[[], Awaitable[None]] | None = None, + ) -> bool: + """Issue now when ready, otherwise retain one deferred dispatch.""" + if not self._active: + return await issue() + root_id = turn.effective_root_correlation_id + state = self._roots.setdefault( + root_id, _RootBarrierState(completed=set(), pending={}) + ) + key = ReplayTurnKey(turn.conversation_id, turn.turn_index) + if not state.initialized: + if turn.turn_index > 0: + self._seed_resumed_prefix(state, key) + state.initialized = True + if self._ready(state, key): + return await issue() + if key in state.pending: + raise RuntimeError( + f"Duplicate deferred replay dispatch for root={root_id!r}, turn={key!r}" + ) + state.pending[key] = _PendingDispatch(issue=issue, on_refused=on_refused) + return True + + def complete(self, credit: Credit) -> None: + """Record any terminal request outcome and release newly ready work.""" + if not self._active: + return + root_id = credit.effective_root_correlation_id + state = self._roots.setdefault( + root_id, _RootBarrierState(completed=set(), pending={}, initialized=True) + ) + state.completed.add(ReplayTurnKey(credit.conversation_id, credit.turn_index)) + ready = [key for key in state.pending if self._ready(state, key)] + for key in sorted(ready): + pending = state.pending.pop(key) + task = asyncio.create_task(self._dispatch_pending(pending)) + self._dispatch_tasks.add(task) + task.add_done_callback(self._dispatch_tasks.discard) + + def close_root(self, root_id: str) -> None: + """Discard completed runtime state when a recycled tree drains.""" + self._roots.pop(root_id, None) + + async def cancel_pending(self, *, notify_refused: bool) -> None: + """Cancel retained dispatches during phase teardown.""" + callbacks = [] + for state in self._roots.values(): + if notify_refused: + callbacks.extend( + pending.on_refused + for pending in state.pending.values() + if pending.on_refused is not None + ) + state.pending.clear() + for task in self._dispatch_tasks: + task.cancel() + self._dispatch_tasks.clear() + for callback in callbacks: + await callback() + + def _seed_resumed_prefix( + self, state: _RootBarrierState, first_key: ReplayTurnKey + ) -> None: + first_timestamp = self._timestamps.get(first_key) + if first_timestamp is None: + return + for key, timestamp in self._timestamps.items(): + if timestamp is not None and timestamp < first_timestamp: + state.completed.add(key) + + def _ready(self, state: _RootBarrierState, key: ReplayTurnKey) -> bool: + return all( + predecessor in state.completed + for predecessor in self._predecessors.get(key, ()) + ) + + @staticmethod + async def _dispatch_pending(pending: _PendingDispatch) -> None: + issued = await pending.issue() + if not issued and pending.on_refused is not None: + await pending.on_refused() + + +class ReplayIssueGate: + """Small CreditIssuer adapter around a replay barrier coordinator.""" + + def __init__(self, coordinator: ReplayBarrierCoordinator | None) -> None: + self._coordinator = coordinator + self._child_refused: Callable[[str], Awaitable[None]] | None = None + self._credit_issued: Callable[[Credit], Awaitable[None]] | None = None + + @property + def enabled(self) -> bool: + return self._coordinator is not None + + def set_child_refused(self, callback: Callable[[str], Awaitable[None]]) -> None: + self._child_refused = callback + + def set_credit_issued(self, callback: Callable[[Credit], Awaitable[None]]) -> None: + self._credit_issued = callback + + async def submit( + self, + turn: TurnToSend, + issue: Callable[[], Awaitable[bool]], + *, + child_refusal_cleanup: bool = False, + ) -> bool: + if self._coordinator is None: + return await issue() + on_refused = None + if child_refusal_cleanup and self._child_refused is not None: + + async def on_refused() -> None: + await self._child_refused(turn.x_correlation_id) + + return await self._coordinator.submit(turn, issue, on_refused=on_refused) + + def activate(self) -> None: + if self._coordinator is not None: + self._coordinator.activate() + + def complete(self, credit: Credit) -> None: + if self._coordinator is not None: + self._coordinator.complete(credit) + + def close_root(self, root_correlation_id: str) -> None: + if self._coordinator is not None: + self._coordinator.close_root(root_correlation_id) + + async def cancel(self, *, notify_refused: bool) -> None: + if self._coordinator is not None: + await self._coordinator.cancel_pending(notify_refused=notify_refused) + + async def observe_issued(self, credit: Credit) -> None: + 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 d6190e230..6c80e0ac2 100644 --- a/src/aiperf/timing/strategies/agentic_replay.py +++ b/src/aiperf/timing/strategies/agentic_replay.py @@ -290,6 +290,9 @@ def _on_tree_drained(self, root_corr: str, phase: CreditPhase) -> None: already released by the registry, so the recycled root's acquire keeps occupancy at exactly the configured concurrency. """ + self.credit_issuer.replay_gate.close_root(root_corr) + if self.branch_orchestrator is not None: + self.branch_orchestrator.close_replay_root(root_corr) lane = self._correlation_to_lane.pop(root_corr, None) self._session_marker.pop(root_corr, None) if lane is None: @@ -318,6 +321,7 @@ async def setup_phase(self) -> None: if self._has_tree_registry: self._session_tree_registry.set_drain_callback(self._on_tree_drained) if self.config.phase == CreditPhase.PROFILING: + self.credit_issuer.replay_gate.activate() if not self.conversation_source.trajectories: raise RuntimeError( "AgenticReplayStrategy PROFILING setup: trajectories empty. " @@ -522,6 +526,7 @@ async def _start_accelerated_warmup(self) -> None: return assert self._cache_warmup_duration is not None self._accelerated_warmup_started = True + self.credit_issuer.replay_gate.activate() self.info( "WARMUP cache pressure: replaying live trajectories for " f"{self._cache_warmup_duration:.1f}s with zero idle delay and " @@ -613,6 +618,7 @@ async def _dispatch_accelerated_trajectory( def observe_credit_return(self, credit: Credit) -> None: """Track the next live turn for the warmup-to-profile handoff.""" + self.credit_issuer.replay_gate.complete(credit) if not self._accelerated_warmup_started: return root_correlation_id = credit.effective_root_correlation_id diff --git a/tests/component_integration/test_agentic_replay_overlap_barrier.py b/tests/component_integration/test_agentic_replay_overlap_barrier.py new file mode 100644 index 000000000..2eaa4363d --- /dev/null +++ b/tests/component_integration/test_agentic_replay_overlap_barrier.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio + +import pytest + +from aiperf.common.enums import CreditPhase +from aiperf.common.models import DatasetMetadata +from aiperf.credit.structs import Credit, TurnToSend +from aiperf.plugin.enums import DatasetSamplingStrategy +from aiperf.timing.replay_dependencies import ReplayBarrierCoordinator +from tests.unit.dataset.loader.test_weka_overlap_groups import _load + +pytestmark = pytest.mark.component_integration + + +@pytest.mark.asyncio +async def test_loaded_abc_group_joins_before_d() -> None: + conversations = _load() + metadata = DatasetMetadata( + conversations=[conversation.to_metadata() for conversation in conversations], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + ) + coordinator = ReplayBarrierCoordinator(metadata) + coordinator.activate() + root = next( + conversation for conversation in metadata.conversations if conversation.is_root + ) + children = [ + conversation + for conversation in metadata.conversations + if not conversation.is_root + ] + issued: list[tuple[str, int]] = [] + + async def submit(conversation_id: str, turn_index: int) -> None: + turn = TurnToSend( + conversation_id=conversation_id, + x_correlation_id=f"corr:{conversation_id}", + turn_index=turn_index, + num_turns=len( + next( + c + for c in metadata.conversations + if c.conversation_id == conversation_id + ).turns + ), + root_correlation_id="root", + ) + + async def issue() -> bool: + issued.append((conversation_id, turn_index)) + return True + + await coordinator.submit(turn, issue) + + await submit(root.conversation_id, 0) + for child in children: + await submit(child.conversation_id, 0) + await submit(root.conversation_id, 1) + assert issued == [ + (root.conversation_id, 0), + *[(child.conversation_id, 0) for child in children], + ] + + coordinator.complete(_credit(root.conversation_id, 0, len(root.turns), "corr:root")) + await asyncio.sleep(0) + assert (root.conversation_id, 1) not in issued + coordinator.complete(_credit(children[0].conversation_id, 0, 1, "corr:child-0")) + await asyncio.sleep(0) + assert (root.conversation_id, 1) not in issued + coordinator.complete(_credit(children[1].conversation_id, 0, 1, "corr:child-1")) + await asyncio.sleep(0) + assert issued[-1] == (root.conversation_id, 1) + + +def _credit( + conversation_id: str, + turn_index: int, + num_turns: int, + correlation_id: str, +) -> Credit: + return Credit( + id=turn_index, + phase=CreditPhase.PROFILING, + conversation_id=conversation_id, + x_correlation_id=correlation_id, + turn_index=turn_index, + num_turns=num_turns, + issued_at_ns=0, + root_correlation_id="root", + ) diff --git a/tests/fixtures/weka_traces_overlap_groups/abc_join_d.json b/tests/fixtures/weka_traces_overlap_groups/abc_join_d.json new file mode 100644 index 000000000..37829850d --- /dev/null +++ b/tests/fixtures/weka_traces_overlap_groups/abc_join_d.json @@ -0,0 +1,14 @@ +{ + "id": "abc_join_d", + "models": ["m"], + "block_size": 64, + "hash_id_scope": "local", + "tool_tokens": 0, + "system_tokens": 0, + "requests": [ + {"t": 0.0, "type": "n", "model": "m", "in": 64, "out": 1, "hash_ids": [1], "api_time": 5.0}, + {"t": 0.1, "type": "n", "model": "m", "in": 64, "out": 1, "hash_ids": [2], "api_time": 5.0}, + {"t": 0.2, "type": "n", "model": "m", "in": 64, "out": 1, "hash_ids": [3], "api_time": 5.0}, + {"t": 6.0, "type": "n", "model": "m", "in": 128, "out": 1, "hash_ids": [1, 4], "api_time": 1.0} + ] +} diff --git a/tests/unit/dataset/loader/test_weka_overlap_groups.py b/tests/unit/dataset/loader/test_weka_overlap_groups.py new file mode 100644 index 000000000..428eaf301 --- /dev/null +++ b/tests/unit/dataset/loader/test_weka_overlap_groups.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path + +from aiperf.common.models import DatasetMetadata +from aiperf.dataset.loader.weka_trace import WekaTraceLoader +from aiperf.plugin.enums import DatasetSamplingStrategy +from tests.unit.dataset.loader.test_weka_trace import ( + _mk_user_config, + _stub_prompt_generator_for_reconstructor, +) + +FIXTURE = ( + Path(__file__).parents[3] + / "fixtures" + / "weka_traces_overlap_groups" + / "abc_join_d.json" +) + + +def _load(): + loader = WekaTraceLoader(filename=str(FIXTURE), user_config=_mk_user_config()) + _stub_prompt_generator_for_reconstructor(loader) + return loader.convert_to_conversations(loader.load_dataset()) + + +def test_loader_encodes_abc_parallel_frontier_before_d() -> None: + conversations = _load() + root = next(c for c in conversations if c.is_root) + children = [c for c in conversations if not c.is_root] + + assert len(root.turns) == 2 + assert len(children) == 2 + assert all(len(child.turns) == 1 for child in children) + assert {conversation.replay_scope_id for conversation in conversations} == { + "abc_join_d" + } + assert all(child.turns[0].replay_predecessors == [] for child in children) + assert { + (ref.conversation_id, ref.turn_index) + for ref in root.turns[1].replay_predecessors + } == {(child.session_id, 0) for child in children} + + +def test_loader_keeps_join_prerequisite_on_d() -> None: + root = next(c for c in _load() if c.is_root) + + assert root.turns[0].branch_ids + assert { + prerequisite.branch_id for prerequisite in root.turns[1].prerequisites + } == set(root.turns[0].branch_ids) + + +def test_overlap_frontier_survives_cached_metadata_round_trip() -> None: + conversations = _load() + metadata = DatasetMetadata( + conversations=[conversation.metadata() for conversation in conversations], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + ) + + restored = DatasetMetadata.model_validate_json(metadata.model_dump_json()) + root = next( + conversation for conversation in restored.conversations if conversation.is_root + ) + children = [ + conversation + for conversation in restored.conversations + if not conversation.is_root + ] + + assert root.replay_scope_id == "abc_join_d" + assert { + (reference.conversation_id, reference.turn_index) + for reference in root.turns[1].replay_predecessors + } == {(child.conversation_id, 0) for child in children} diff --git a/tests/unit/dataset/test_mmap_cache.py b/tests/unit/dataset/test_mmap_cache.py index dcf27a4ca..c795ae311 100644 --- a/tests/unit/dataset/test_mmap_cache.py +++ b/tests/unit/dataset/test_mmap_cache.py @@ -345,6 +345,19 @@ def test_lookup_version_mismatch_returns_none(self, tmp_path: Path) -> None: manifest_path.write_bytes(orjson.dumps(raw)) assert mmap_cache.lookup("oldver", compressed=False) is None + def test_lookup_rejects_pre_overlap_frontier_manifest(self, tmp_path: Path) -> None: + cache_root = mmap_cache.cache_dir() + _populate_entry(cache_root, cache_key="pre-overlap-frontier") + manifest_path = ( + cache_root / "pre-overlap-frontier" / mmap_cache.MANIFEST_FILENAME + ) + raw = orjson.loads(manifest_path.read_bytes()) + raw["version"] = 20 + manifest_path.write_bytes(orjson.dumps(raw)) + + assert mmap_cache.MANIFEST_VERSION == 21 + assert mmap_cache.lookup("pre-overlap-frontier", compressed=False) is None + def test_lookup_compressed_mismatch_returns_none(self, tmp_path: Path) -> None: cache_root = mmap_cache.cache_dir() _populate_entry(cache_root, cache_key="uncomp", compressed=False) diff --git a/tests/unit/timing/test_branch_orchestrator_delayed.py b/tests/unit/timing/test_branch_orchestrator_delayed.py index e50d805ea..b247aea10 100644 --- a/tests/unit/timing/test_branch_orchestrator_delayed.py +++ b/tests/unit/timing/test_branch_orchestrator_delayed.py @@ -19,7 +19,7 @@ import pytest -from aiperf.common.enums import ConversationBranchMode, PrerequisiteKind +from aiperf.common.enums import ConversationBranchMode, CreditPhase, PrerequisiteKind from aiperf.common.models import ( ConversationBranchInfo, ConversationMetadata, @@ -228,6 +228,58 @@ def _start( assert orch.stats.parents_resumed == 1 +@pytest.mark.asyncio +async def test_overlap_branch_dispatches_from_parent_start_not_return() -> None: + branch = ConversationBranchInfo( + branch_id="root:overlap", + child_conversation_ids=["child"], + mode=ConversationBranchMode.SPAWN, + start_timestamp_ms=0.0, + ) + root = _mk_conv( + "root", + [ + TurnMetadata( + timestamp_ms=0.0, + api_time_ms=5000.0, + branch_ids=[branch.branch_id], + ), + TurnMetadata( + timestamp_ms=6000.0, + prerequisites=[ + TurnPrerequisite( + kind=PrerequisiteKind.SPAWN_JOIN, + branch_id=branch.branch_id, + ) + ], + ), + ], + [branch], + ) + root.replay_scope_id = "root" + child = _mk_conv("child", [TurnMetadata(timestamp_ms=0.0)], []) + cs = _mk_source([root, child]) + child_session = MagicMock( + x_correlation_id="corr-child", + 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) + orch = BranchOrchestrator(conversation_source=cs, credit_issuer=issuer) + credit = _mk_credit("root", "corr-root", 0) + credit.phase = CreditPhase.PROFILING + credit.effective_root_correlation_id = "corr-root" + + await orch.on_credit_issued(credit) + + issuer.dispatch_first_turn.assert_awaited_once_with(child_session) + assert await orch.intercept(credit) is True + issuer.dispatch_first_turn.assert_awaited_once() + + @pytest.mark.asyncio async def test_delayed_join_stop_condition_fires_during_gap_suppresses_join(): """If the issuer reports ``dispatch_join_turn`` returned False (stop diff --git a/tests/unit/timing/test_replay_barrier_coordinator.py b/tests/unit/timing/test_replay_barrier_coordinator.py new file mode 100644 index 000000000..562189c05 --- /dev/null +++ b/tests/unit/timing/test_replay_barrier_coordinator.py @@ -0,0 +1,141 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio + +import pytest + +from aiperf.common.enums import CreditPhase +from aiperf.common.models import ( + ConversationMetadata, + DatasetMetadata, + ReplayTurnReference, + TurnMetadata, +) +from aiperf.credit.structs import Credit, TurnToSend +from aiperf.plugin.enums import DatasetSamplingStrategy +from aiperf.timing.replay_dependencies import ReplayBarrierCoordinator + + +def _metadata() -> DatasetMetadata: + abc = [ReplayTurnReference(conversation_id=name, turn_index=0) for name in "abc"] + return DatasetMetadata( + sampling_strategy=DatasetSamplingStrategy.RANDOM, + conversations=[ + ConversationMetadata( + conversation_id=name, turns=[TurnMetadata(timestamp_ms=0)] + ) + for name in "abc" + ] + + [ + ConversationMetadata( + conversation_id="d", + turns=[TurnMetadata(timestamp_ms=10, replay_predecessors=abc)], + ) + ], + ) + + +def _turn(name: str, root: str = "root") -> TurnToSend: + return TurnToSend( + conversation_id=name, + x_correlation_id=f"{root}:{name}", + turn_index=0, + num_turns=1, + root_correlation_id=root, + ) + + +def _credit(name: str, root: str = "root") -> Credit: + return Credit( + id=0, + phase=CreditPhase.PROFILING, + conversation_id=name, + x_correlation_id=f"{root}:{name}", + turn_index=0, + num_turns=1, + issued_at_ns=0, + root_correlation_id=root, + ) + + +@pytest.mark.asyncio +async def test_abc_issue_together_and_d_waits_for_every_member() -> None: + coordinator = ReplayBarrierCoordinator(_metadata()) + coordinator.activate() + issued: list[str] = [] + + async def submit(name: str) -> None: + await coordinator.submit( + _turn(name), + lambda name=name: _record_issue(issued, name), + ) + + await submit("a") + await submit("b") + await submit("c") + await submit("d") + assert issued == ["a", "b", "c"] + + coordinator.complete(_credit("a")) + await asyncio.sleep(0) + assert issued == ["a", "b", "c"] + coordinator.complete(_credit("b")) + await asyncio.sleep(0) + assert issued == ["a", "b", "c"] + coordinator.complete(_credit("c")) + await asyncio.sleep(0) + assert issued == ["a", "b", "c", "d"] + + +@pytest.mark.asyncio +async def test_any_terminal_outcome_releases_barrier() -> None: + coordinator = ReplayBarrierCoordinator(_metadata()) + coordinator.activate() + issued: list[str] = [] + for name in "abcd": + await coordinator.submit( + _turn(name), lambda name=name: _record_issue(issued, name) + ) + + for name in "abc": + coordinator.complete(_credit(name)) + await asyncio.sleep(0) + + assert issued[-1] == "d" + + +@pytest.mark.asyncio +async def test_runtime_roots_are_independent() -> None: + coordinator = ReplayBarrierCoordinator(_metadata()) + coordinator.activate() + issued: list[str] = [] + await coordinator.submit(_turn("d", "one"), lambda: _record_issue(issued, "one:d")) + await coordinator.submit(_turn("d", "two"), lambda: _record_issue(issued, "two:d")) + + for name in "abc": + coordinator.complete(_credit(name, "one")) + await asyncio.sleep(0) + + assert issued == ["one:d"] + + +@pytest.mark.asyncio +async def test_scalar_peak_would_slip_d_after_only_one_completion() -> None: + coordinator = ReplayBarrierCoordinator(_metadata()) + coordinator.activate() + issued: list[str] = [] + for name in "abcd": + await coordinator.submit( + _turn(name), lambda name=name: _record_issue(issued, name) + ) + + coordinator.complete(_credit("a")) + await asyncio.sleep(0) + + assert "d" not in issued + + +async def _record_issue(issued: list[str], name: str) -> bool: + issued.append(name) + return True diff --git a/tests/unit/timing/test_replay_dependencies.py b/tests/unit/timing/test_replay_dependencies.py new file mode 100644 index 000000000..26cbb3940 --- /dev/null +++ b/tests/unit/timing/test_replay_dependencies.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import math + +from aiperf.timing.replay_dependencies import ( + RecordedTurnInterval, + ReplayTurnKey, + infer_cross_stream_predecessors, +) + + +def _interval( + name: str, + stream: str, + start_ms: float | None, + api_time_ms: float | None, +) -> RecordedTurnInterval: + return RecordedTurnInterval( + key=ReplayTurnKey(name, 0), + stream_id=stream, + start_ms=start_ms, + api_time_ms=api_time_ms, + ) + + +def test_sequential_intervals_form_singleton_frontiers() -> None: + a = _interval("a", "a", 0.0, 10.0) + b = _interval("b", "b", 10.0, 10.0) + c = _interval("c", "c", 20.0, 10.0) + + result = infer_cross_stream_predecessors([a, b, c]) + + assert result[a.key] == () + assert result[b.key] == (a.key,) + assert result[c.key] == (b.key,) + + +def test_parallel_abc_all_join_before_d() -> None: + a = _interval("a", "a", 0.0, 10.0) + b = _interval("b", "b", 1.0, 8.0) + c = _interval("c", "c", 2.0, 7.0) + d = _interval("d", "d", 10.0, 1.0) + + result = infer_cross_stream_predecessors([a, b, c, d]) + + assert result[a.key] == () + assert result[b.key] == () + assert result[c.key] == () + assert result[d.key] == (a.key, b.key, c.key) + + +def test_boundary_touch_is_sequential_but_equal_starts_are_parallel() -> None: + a = _interval("a", "a", 0.0, 5.0) + b = _interval("b", "b", 5.0, 0.0) + c = _interval("c", "c", 5.0, 0.0) + + result = infer_cross_stream_predecessors([a, b, c]) + + assert result[b.key] == (a.key,) + assert result[c.key] == (a.key,) + assert b.key not in result[c.key] + assert c.key not in result[b.key] + + +def test_missing_and_invalid_durations_are_zero_width() -> None: + missing = _interval("missing", "missing", 1.0, None) + negative = _interval("negative", "negative", 2.0, -5.0) + non_finite = _interval("non-finite", "non-finite", 3.0, math.inf) + target = _interval("target", "target", 4.0, 1.0) + + result = infer_cross_stream_predecessors([missing, negative, non_finite, target]) + + assert result[negative.key] == (missing.key,) + assert result[non_finite.key] == (negative.key,) + assert result[target.key] == (non_finite.key,) + + +def test_missing_and_non_finite_starts_add_no_cross_stream_edges() -> None: + missing = _interval("missing", "missing", None, 1.0) + non_finite = _interval("non-finite", "non-finite", math.nan, 1.0) + target = _interval("target", "target", 10.0, 1.0) + + result = infer_cross_stream_predecessors([missing, non_finite, target]) + + assert result[target.key] == () + + +def test_transitive_overlap_uses_precise_frontier_not_connected_component() -> None: + long_a = _interval("long-a", "a", 0.0, 10.0) + short_b = _interval("short-b", "b", 1.0, 1.0) + later_c = _interval("later-c", "c", 3.0, 1.0) + + result = infer_cross_stream_predecessors([long_a, short_b, later_c]) + + assert result[short_b.key] == () + assert result[later_c.key] == (short_b.key,) + assert long_a.key not in result[later_c.key] + + +def test_same_stream_order_is_left_to_conversation_replay() -> None: + a = RecordedTurnInterval(ReplayTurnKey("stream", 0), "stream", 0.0, 5.0) + b = RecordedTurnInterval(ReplayTurnKey("stream", 1), "stream", 1.0, 5.0) + + result = infer_cross_stream_predecessors([a, b]) + + assert result[a.key] == () + assert result[b.key] == () From bde92ededb883cc63a1ddfeaf45a5127797d6bbb Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Thu, 25 Jun 2026 16:38:41 -0500 Subject: [PATCH 2/2] fix(agentic): seed exact replay resume prefixes Signed-off-by: Cam Quilici --- .../benchmark-modes/timing-modes-reference.md | 5 +- src/aiperf/timing/replay_dependencies.py | 97 +++++++++++++---- .../timing/strategies/agentic_replay.py | 95 ++++++++++++++++- src/aiperf/timing/trajectory_source.py | 27 ++++- .../timing/strategies/test_agentic_replay.py | 29 +++++ .../timing/test_replay_barrier_coordinator.py | 100 ++++++++++++++++-- tests/unit/timing/test_trajectory_source.py | 10 ++ .../test_trajectory_source_pathological.py | 5 + 8 files changed, 331 insertions(+), 37 deletions(-) diff --git a/docs/benchmark-modes/timing-modes-reference.md b/docs/benchmark-modes/timing-modes-reference.md index 9c7e17b0d..b0968e0f7 100644 --- a/docs/benchmark-modes/timing-modes-reference.md +++ b/docs/benchmark-modes/timing-modes-reference.md @@ -258,7 +258,10 @@ sequential requests on another stream without launching all of them at once. Branches that began while their spawning request was in flight are scheduled from that request's send time; independent subagent scopes are not globally joined. The same barriers remain active during accelerated cache-pressure -warmup, where idle delays are otherwise removed. +warmup, where idle delays are otherwise removed. When replay resumes at a +sampled `t*` or crosses from accelerated warmup into profiling, the snapshot +seeds each stream's exact completed prefix before any request can dispatch; +request arrival order therefore cannot weaken or deadlock a join barrier. Concurrency is **per session tree**: each `--concurrency` lane holds one slot for a whole tree — the root conversation plus every subagent it spawns (children, subchildren, background `::fa:`/`::aux:` sidecars). A lane's slot is released, and the lane recycled into a fresh root, only once the entire tree drains (root terminal **and** all descendants returned) — so a background subagent that outlives its root does not free the lane early. Recycle then draws the next root from the dataset sampler (honoring the dataset's `sampling_strategy` — sequential / shuffle / random), starting it from turn 0. This keeps exactly `--concurrency` trees live at all times. The shared tree id (`root_correlation_id`) is persisted per record in `profile_export.jsonl`, so `aiperf analyze swim-lane` groups each tree under one lane and renders exactly `--concurrency` slots. diff --git a/src/aiperf/timing/replay_dependencies.py b/src/aiperf/timing/replay_dependencies.py index 929946099..c2600ee67 100644 --- a/src/aiperf/timing/replay_dependencies.py +++ b/src/aiperf/timing/replay_dependencies.py @@ -27,6 +27,14 @@ class ReplayTurnKey: turn_index: int +@dataclass(frozen=True, slots=True, order=True) +class ReplayResumeBoundary: + """Completed prefix of one replay stream at a phase boundary.""" + + conversation_id: str + next_turn_index: int + + @dataclass(frozen=True, slots=True) class RecordedTurnInterval: """One request interval on a logical replay stream.""" @@ -124,7 +132,6 @@ class _PendingDispatch: class _RootBarrierState: completed: set[ReplayTurnKey] pending: dict[ReplayTurnKey, _PendingDispatch] - initialized: bool = False class ReplayBarrierCoordinator: @@ -132,7 +139,6 @@ class ReplayBarrierCoordinator: def __init__(self, dataset_metadata: DatasetMetadata) -> None: self._predecessors: dict[ReplayTurnKey, tuple[ReplayTurnKey, ...]] = {} - self._timestamps: dict[ReplayTurnKey, float | None] = {} for conversation in dataset_metadata.conversations: for turn_index, turn in enumerate(conversation.turns): key = ReplayTurnKey(conversation.conversation_id, turn_index) @@ -140,12 +146,6 @@ def __init__(self, dataset_metadata: DatasetMetadata) -> None: ReplayTurnKey(ref.conversation_id, ref.turn_index) for ref in turn.replay_predecessors ) - timestamp = turn.timestamp_ms - self._timestamps[key] = ( - float(timestamp) - if isinstance(timestamp, int | float) and math.isfinite(timestamp) - else None - ) self._roots: dict[str, _RootBarrierState] = {} self._dispatch_tasks: set[asyncio.Task] = set() self._active = False @@ -183,10 +183,6 @@ async def submit( root_id, _RootBarrierState(completed=set(), pending={}) ) key = ReplayTurnKey(turn.conversation_id, turn.turn_index) - if not state.initialized: - if turn.turn_index > 0: - self._seed_resumed_prefix(state, key) - state.initialized = True if self._ready(state, key): return await issue() if key in state.pending: @@ -202,7 +198,7 @@ def complete(self, credit: Credit) -> None: return root_id = credit.effective_root_correlation_id state = self._roots.setdefault( - root_id, _RootBarrierState(completed=set(), pending={}, initialized=True) + root_id, _RootBarrierState(completed=set(), pending={}) ) state.completed.add(ReplayTurnKey(credit.conversation_id, credit.turn_index)) ready = [key for key in state.pending if self._ready(state, key)] @@ -216,6 +212,56 @@ def close_root(self, root_id: str) -> None: """Discard completed runtime state when a recycled tree drains.""" self._roots.pop(root_id, None) + def seed_completed_prefixes( + self, + root_id: str, + 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={}) + ) + if state.pending: + raise RuntimeError( + f"Cannot seed replay history after dispatch for root={root_id!r}" + ) + for boundary in boundaries: + if boundary.next_turn_index < 0: + raise ValueError( + "Replay resume boundary must have a non-negative turn index" + ) + state.completed.update( + ReplayTurnKey(boundary.conversation_id, turn_index) + for turn_index in range(boundary.next_turn_index) + ) + + def completed_prefixes(self, root_id: str) -> tuple[ReplayResumeBoundary, ...]: + """Return the contiguous completed prefix of every replay stream.""" + state = self._roots.get(root_id) + if state is None: + return () + next_turn_by_conversation: dict[str, int] = {} + for key in state.completed: + next_turn_by_conversation[key.conversation_id] = max( + next_turn_by_conversation.get(key.conversation_id, 0), + key.turn_index + 1, + ) + for conversation_id, next_turn_index in next_turn_by_conversation.items(): + if any( + ReplayTurnKey(conversation_id, turn_index) not in state.completed + for turn_index in range(next_turn_index) + ): + raise RuntimeError( + "Replay completion history is not a contiguous stream prefix: " + f"root={root_id!r}, conversation={conversation_id!r}" + ) + return tuple( + ReplayResumeBoundary(conversation_id, next_turn_index) + for conversation_id, next_turn_index in sorted( + next_turn_by_conversation.items() + ) + ) + async def cancel_pending(self, *, notify_refused: bool) -> None: """Cancel retained dispatches during phase teardown.""" callbacks = [] @@ -233,16 +279,6 @@ async def cancel_pending(self, *, notify_refused: bool) -> None: for callback in callbacks: await callback() - def _seed_resumed_prefix( - self, state: _RootBarrierState, first_key: ReplayTurnKey - ) -> None: - first_timestamp = self._timestamps.get(first_key) - if first_timestamp is None: - return - for key, timestamp in self._timestamps.items(): - if timestamp is not None and timestamp < first_timestamp: - state.completed.add(key) - def _ready(self, state: _RootBarrierState, key: ReplayTurnKey) -> bool: return all( predecessor in state.completed @@ -303,6 +339,21 @@ def close_root(self, root_correlation_id: str) -> None: if self._coordinator is not None: self._coordinator.close_root(root_correlation_id) + def seed_completed_prefixes( + self, + root_correlation_id: str, + boundaries: tuple[ReplayResumeBoundary, ...], + ) -> None: + if self._coordinator is not None: + self._coordinator.seed_completed_prefixes(root_correlation_id, boundaries) + + def completed_prefixes( + self, root_correlation_id: str + ) -> tuple[ReplayResumeBoundary, ...]: + if self._coordinator is None: + return () + return self._coordinator.completed_prefixes(root_correlation_id) + async def cancel(self, *, notify_refused: bool) -> None: if self._coordinator is not None: await self._coordinator.cancel_pending(notify_refused=notify_refused) diff --git a/src/aiperf/timing/strategies/agentic_replay.py b/src/aiperf/timing/strategies/agentic_replay.py index 6c80e0ac2..81ecef3d9 100644 --- a/src/aiperf/timing/strategies/agentic_replay.py +++ b/src/aiperf/timing/strategies/agentic_replay.py @@ -64,6 +64,7 @@ from aiperf.common.scenario.context_overflow import is_context_overflow_response from aiperf.credit.structs import TurnToSend from aiperf.timing.conversation_source import SampledSession +from aiperf.timing.replay_dependencies import ReplayResumeBoundary from aiperf.timing.trajectory_source import ( ConversationState, Trajectory, @@ -274,11 +275,54 @@ def _lane_root_corr(self, snapshot: TrajectorySnapshot) -> str | None: root is present, or the shared parent of the background subagents when rootless). Falls back to a state's x_correlation_id for snapshots built without the field (older fixtures).""" - return next( - (s.root_correlation_id or s.x_correlation_id for s in snapshot.states), - None, + return ( + next( + ( + state.root_correlation_id + for state in snapshot.states + if state.root_correlation_id is not None + ), + None, + ) + or next( + ( + state.x_correlation_id + for state in snapshot.states + if state.agent_depth == 0 + ), + None, + ) + or next( + ( + state.parent_correlation_id + for state in snapshot.states + if state.parent_correlation_id is not None + ), + None, + ) + or next( + (state.x_correlation_id for state in snapshot.states), + None, + ) ) + def _seed_trajectory_replay_prefix(self, trajectory: Trajectory) -> None: + """Seed the exact completed history before a resumed phase dispatches.""" + if trajectory.snapshot is None: + root_correlation_id = trajectory.x_correlation_id + boundaries = ( + ReplayResumeBoundary( + trajectory.conversation_id, trajectory.start_turn_index + 1 + ), + ) + else: + root_correlation_id = self._lane_root_corr(trajectory.snapshot) + boundaries = trajectory.snapshot.replay_resume_boundaries + if root_correlation_id is not None: + self.credit_issuer.replay_gate.seed_completed_prefixes( + root_correlation_id, boundaries + ) + def _on_tree_drained(self, root_corr: str, phase: CreditPhase) -> None: """Registry drain callback: a session tree fully drained and freed its slot, so recycle its lane into a fresh root. @@ -321,6 +365,8 @@ async def setup_phase(self) -> None: if self._has_tree_registry: self._session_tree_registry.set_drain_callback(self._on_tree_drained) if self.config.phase == CreditPhase.PROFILING: + for trajectory in self.conversation_source.trajectories: + self._seed_trajectory_replay_prefix(trajectory) self.credit_issuer.replay_gate.activate() if not self.conversation_source.trajectories: raise RuntimeError( @@ -526,6 +572,8 @@ async def _start_accelerated_warmup(self) -> None: return assert self._cache_warmup_duration is not None self._accelerated_warmup_started = True + for trajectory in self.conversation_source.trajectories: + self._seed_trajectory_replay_prefix(trajectory) self.credit_issuer.replay_gate.activate() self.info( "WARMUP cache pressure: replaying live trajectories for " @@ -656,8 +704,12 @@ async def finalize_phase(self) -> None: if not self._accelerated_warmup_started: return 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() + } self.conversation_source.trajectories = self._build_handoff_trajectories( - states_by_lane + states_by_lane, boundaries_by_lane ) self.info( "WARMUP cache pressure handoff: persisted " @@ -699,13 +751,44 @@ def _build_handoff_states(self) -> dict[int, list[ConversationState]]: ) return states_by_lane + def _build_handoff_replay_boundaries( + self, states: list[ConversationState] + ) -> tuple[ReplayResumeBoundary, ...]: + """Merge live stream positions with terminal history from warmup.""" + next_turn_by_conversation = { + state.conversation_id: state.next_turn_index + for state in states + if state.next_turn_index > 0 + } + if states: + root_correlation_id = self._lane_root_corr( + TrajectorySnapshot(t_star_ms=0.0, states=tuple(states)) + ) + if root_correlation_id is not None: + for boundary in self.credit_issuer.replay_gate.completed_prefixes( + root_correlation_id + ): + next_turn_by_conversation[boundary.conversation_id] = max( + next_turn_by_conversation.get(boundary.conversation_id, 0), + boundary.next_turn_index, + ) + return tuple( + ReplayResumeBoundary(conversation_id, next_turn_index) + for conversation_id, next_turn_index in sorted( + next_turn_by_conversation.items() + ) + ) + def _build_handoff_trajectories( - self, states_by_lane: dict[int, list[ConversationState]] + self, + states_by_lane: dict[int, list[ConversationState]], + boundaries_by_lane: dict[int, tuple[ReplayResumeBoundary, ...]], ) -> list[Trajectory]: """Build the shared trajectory list consumed by profiling.""" rebuilt: list[Trajectory] = [] for lane, previous in enumerate(self.conversation_source.trajectories): states = states_by_lane[lane] + boundaries = boundaries_by_lane[lane] if not states: trace_id = self.conversation_source.next_recycle_conversation_id() if trace_id is None: @@ -719,6 +802,7 @@ def _build_handoff_trajectories( next_turn_index=0, ) ] + boundaries = () root_state = next( (state for state in states if state.agent_depth == 0), None ) @@ -744,6 +828,7 @@ def _build_handoff_trajectories( ), ) ), + replay_resume_boundaries=boundaries, ), x_correlation_id=( root_state.x_correlation_id diff --git a/src/aiperf/timing/trajectory_source.py b/src/aiperf/timing/trajectory_source.py index 0457cd39a..6c1cb787e 100644 --- a/src/aiperf/timing/trajectory_source.py +++ b/src/aiperf/timing/trajectory_source.py @@ -33,6 +33,7 @@ from aiperf.common.scenario.base import EmptyTracePoolError from aiperf.dataset.protocols import DatasetSamplingStrategyProtocol from aiperf.timing.conversation_source import ConversationSource, SampledSession +from aiperf.timing.replay_dependencies import ReplayResumeBoundary _logger = logging.getLogger(__name__) @@ -83,6 +84,7 @@ class TrajectorySnapshot: t_star_ms: float states: tuple[ConversationState, ...] + replay_resume_boundaries: tuple[ReplayResumeBoundary, ...] = () @dataclass(slots=True, frozen=True) @@ -569,6 +571,25 @@ def _trace_time_bounds(self, root_id: str) -> tuple[float, float] | None: return None return min(timestamps), max(timestamps) + def _replay_resume_boundaries( + self, root_id: str, t_star_ms: float + ) -> tuple[ReplayResumeBoundary, ...]: + """Return exact completed stream prefixes at the sampled instant.""" + boundaries: list[ReplayResumeBoundary] = [] + for conversation_id in sorted(self._collect_trace_conversation_ids(root_id)): + metadata = self._metadata_lookup.get(conversation_id) + if metadata is None: + continue + next_turn_index = _next_turn_index_at_or_after(metadata, t_star_ms) + completed_turns = ( + len(metadata.turns) if next_turn_index is None else next_turn_index + ) + if completed_turns > 0: + boundaries.append( + ReplayResumeBoundary(conversation_id, completed_turns) + ) + return tuple(boundaries) + def _warn_if_live_delta_snapshot_needs_prior_responses( self, root_id: str, snapshot: TrajectorySnapshot ) -> None: @@ -690,7 +711,11 @@ def _snapshot_for( return None if not any(not state.waiting_on_children for state in states): return None - return TrajectorySnapshot(t_star_ms=t_star_ms, states=tuple(states)) + return TrajectorySnapshot( + t_star_ms=t_star_ms, + states=tuple(states), + replay_resume_boundaries=self._replay_resume_boundaries(root_id, t_star_ms), + ) def _branch_runtimes(self, parent_meta) -> list[_BranchRuntime]: join_by_branch: dict[str, int] = {} diff --git a/tests/unit/timing/strategies/test_agentic_replay.py b/tests/unit/timing/strategies/test_agentic_replay.py index 597909577..af0b38aca 100644 --- a/tests/unit/timing/strategies/test_agentic_replay.py +++ b/tests/unit/timing/strategies/test_agentic_replay.py @@ -26,6 +26,7 @@ from aiperf.credit.structs import Credit from aiperf.dataset.dataset_samplers import SequentialSampler from aiperf.plugin.enums import DatasetSamplingStrategy +from aiperf.timing.replay_dependencies import ReplayResumeBoundary from aiperf.timing.strategies.agentic_replay import AgenticReplayStrategy from aiperf.timing.trajectory_source import ( ConversationState, @@ -101,6 +102,8 @@ def _make_strategy( cfg.concurrency = len(trajectories) cfg.agentic_cache_warmup_duration_sec = cache_warmup_duration issuer = issuer if issuer is not None else AsyncMock() + issuer.replay_gate = MagicMock() + issuer.replay_gate.completed_prefixes.return_value = () scheduler = scheduler if scheduler is not None else MagicMock() strategy = AgenticReplayStrategy( config=cfg, @@ -292,6 +295,32 @@ async def test_cache_warmup_cutoff_stops_issuer_and_persists_next_turn(): assert len(snapshot.states) == 1 assert snapshot.states[0].next_turn_index == 3 assert snapshot.states[0].x_correlation_id == baseline.x_correlation_id + assert snapshot.replay_resume_boundaries == (ReplayResumeBoundary("trace_0", 3),) + + +def test_handoff_preserves_completed_streams_absent_from_live_state() -> None: + strategy, issuer, _, _ = _make_strategy( + phase=CreditPhase.WARMUP, + trajectories=[Trajectory(conversation_id="trace_0", start_turn_index=0)], + cache_warmup_duration=10.0, + ) + issuer.replay_gate.completed_prefixes.return_value = ( + ReplayResumeBoundary("trace_0::aux:0", 1), + ) + live_state = ConversationState( + conversation_id="trace_0", + x_correlation_id="root", + root_correlation_id="root", + next_turn_index=3, + ) + + boundaries = strategy._build_handoff_replay_boundaries([live_state]) + + assert boundaries == ( + ReplayResumeBoundary("trace_0", 3), + ReplayResumeBoundary("trace_0::aux:0", 1), + ) + issuer.replay_gate.completed_prefixes.assert_called_once_with("root") @pytest.mark.asyncio diff --git a/tests/unit/timing/test_replay_barrier_coordinator.py b/tests/unit/timing/test_replay_barrier_coordinator.py index 562189c05..c46fe3622 100644 --- a/tests/unit/timing/test_replay_barrier_coordinator.py +++ b/tests/unit/timing/test_replay_barrier_coordinator.py @@ -14,7 +14,10 @@ ) from aiperf.credit.structs import Credit, TurnToSend from aiperf.plugin.enums import DatasetSamplingStrategy -from aiperf.timing.replay_dependencies import ReplayBarrierCoordinator +from aiperf.timing.replay_dependencies import ( + ReplayBarrierCoordinator, + ReplayResumeBoundary, +) def _metadata() -> DatasetMetadata: @@ -36,24 +39,34 @@ def _metadata() -> DatasetMetadata: ) -def _turn(name: str, root: str = "root") -> TurnToSend: +def _turn( + name: str, + root: str = "root", + turn_index: int = 0, + num_turns: int = 1, +) -> TurnToSend: return TurnToSend( conversation_id=name, x_correlation_id=f"{root}:{name}", - turn_index=0, - num_turns=1, + turn_index=turn_index, + num_turns=num_turns, root_correlation_id=root, ) -def _credit(name: str, root: str = "root") -> Credit: +def _credit( + name: str, + root: str = "root", + turn_index: int = 0, + num_turns: int = 1, +) -> Credit: return Credit( id=0, phase=CreditPhase.PROFILING, conversation_id=name, x_correlation_id=f"{root}:{name}", - turn_index=0, - num_turns=1, + turn_index=turn_index, + num_turns=num_turns, issued_at_ns=0, root_correlation_id=root, ) @@ -136,6 +149,79 @@ async def test_scalar_peak_would_slip_d_after_only_one_completion() -> None: assert "d" not in issued +@pytest.mark.asyncio +@pytest.mark.parametrize( + "submission_order", + [ + (("aux", 0, 1), ("flat", 1, 2)), + (("flat", 1, 2), ("aux", 0, 1)), + ], +) +async def test_resumed_prefix_is_exact_and_submission_order_independent( + submission_order: tuple[tuple[str, int, int], ...], +) -> None: + metadata = DatasetMetadata( + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + conversations=[ + ConversationMetadata( + conversation_id="flat", + turns=[ + TurnMetadata(timestamp_ms=0), + TurnMetadata( + timestamp_ms=20, + replay_predecessors=[ + ReplayTurnReference(conversation_id="aux", turn_index=0) + ], + ), + ], + ), + ConversationMetadata( + conversation_id="aux", + turns=[ + TurnMetadata( + timestamp_ms=10, + replay_predecessors=[ + ReplayTurnReference(conversation_id="flat", turn_index=0) + ], + ) + ], + is_root=False, + ), + ], + ) + coordinator = ReplayBarrierCoordinator(metadata) + coordinator.seed_completed_prefixes("root", (ReplayResumeBoundary("flat", 1),)) + coordinator.activate() + issued: list[tuple[str, int]] = [] + + async def record_issue(item: tuple[str, int]) -> bool: + issued.append(item) + return True + + for conversation_id, turn_index, num_turns in submission_order: + await coordinator.submit( + _turn( + conversation_id, + turn_index=turn_index, + num_turns=num_turns, + ), + lambda conversation_id=conversation_id, turn_index=turn_index: record_issue( + (conversation_id, turn_index) + ), + ) + + assert issued == [("aux", 0)] + + coordinator.complete(_credit("aux")) + await asyncio.sleep(0) + + assert issued == [("aux", 0), ("flat", 1)] + assert coordinator.completed_prefixes("root") == ( + ReplayResumeBoundary("aux", 1), + ReplayResumeBoundary("flat", 1), + ) + + async def _record_issue(issued: list[str], name: str) -> bool: issued.append(name) return True diff --git a/tests/unit/timing/test_trajectory_source.py b/tests/unit/timing/test_trajectory_source.py index abc08ba6e..d58a0b90b 100644 --- a/tests/unit/timing/test_trajectory_source.py +++ b/tests/unit/timing/test_trajectory_source.py @@ -189,6 +189,10 @@ def test_timestamped_snapshot_includes_inflight_subagent_and_gated_parent(): by_cid = {state.conversation_id: state for state in trajectory.snapshot.states} parent = by_cid["trace"] child = by_cid["trace::sa:agent_0"] + resume_boundaries = { + boundary.conversation_id: boundary.next_turn_index + for boundary in trajectory.snapshot.replay_resume_boundaries + } assert parent.waiting_on_children is True assert parent.next_turn_index == 2 @@ -197,6 +201,7 @@ def test_timestamped_snapshot_includes_inflight_subagent_and_gated_parent(): assert child.branch_id == branch_id assert child.branch_mode == ConversationBranchMode.SPAWN assert child.next_dispatch_offset_ms == pytest.approx(500.0) + assert resume_boundaries == {"trace": 2, "trace::sa:agent_0": 1} # Both active-at-t* sessions are warmed: the mid-flight child (turn 0) and # the gated parent (turn 1, priming its join turn). Gated parents are no # longer excluded from warmup. @@ -284,11 +289,16 @@ def test_timestamped_snapshot_after_spawning_turn_schedules_future_child_start() by_cid = {state.conversation_id: state for state in trajectory.snapshot.states} parent = by_cid["trace"] child = by_cid["trace::sa:agent_0"] + resume_boundaries = { + boundary.conversation_id: boundary.next_turn_index + for boundary in trajectory.snapshot.replay_resume_boundaries + } assert parent.waiting_on_children is True assert parent.next_turn_index == 2 assert child.next_turn_index == 0 assert child.next_dispatch_offset_ms == pytest.approx(500.0) + assert resume_boundaries == {"trace": 2} def test_next_recycle_conversation_id_uses_sampler_round_robin(): diff --git a/tests/unit/timing/test_trajectory_source_pathological.py b/tests/unit/timing/test_trajectory_source_pathological.py index 6c82beb84..b4ec641a3 100644 --- a/tests/unit/timing/test_trajectory_source_pathological.py +++ b/tests/unit/timing/test_trajectory_source_pathological.py @@ -296,6 +296,11 @@ def test_background_only_snapshot_keeps_child_but_start_index_defaults_to_zero() assert traj.snapshot is not None cids = {s.conversation_id for s in traj.snapshot.states} assert cids == {"trace::bg"} # only the background child survives + resume_boundaries = { + boundary.conversation_id: boundary.next_turn_index + for boundary in traj.snapshot.replay_resume_boundaries + } + assert resume_boundaries == {"trace": 2, "trace::bg": 1} assert traj.start_turn_index == 0 # sentinel default, no root state assert src.warmup_credit_count == 1 # the one ready child state