From 7ee18816f7d1101a13f0518c68ce267159ee7a62 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Thu, 27 Aug 2026 22:15:01 -0700 Subject: [PATCH] fix(agentic): settle DAG state after wire drain MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Preserve terminal DAG semantics while allowing the benchmark grace period to drain accepted wire work. Reconcile scheduler, branch, credit, and phase state only after the transport is quiet, and cover the terminal races with focused unit and integration tests. 中文：在 benchmark grace period 内先排空已接收的网络请求，再统一收敛调度器、分支、credit 和 phase 状态，既保留 DAG 的终态语义，也避免传输层已完成时残留调度子任务导致错误判定；补充针对终态竞态的单元与集成测试。 Signed-off-by: Yangmin Li --- docs/benchmark-modes/dag.md | 2 +- src/aiperf/common/loop_scheduler.py | 41 ++++-- src/aiperf/credit/callback_handler.py | 8 ++ src/aiperf/timing/branch_orchestrator.py | 126 +++++++++++++++++- src/aiperf/timing/phase/runner.py | 21 +++ .../timing/strategies/agentic_replay.py | 83 +++++++++--- tests/integration/test_weka_flat_split_e2e.py | 53 +++++++- tests/unit/common/test_loop_scheduler.py | 13 ++ tests/unit/credit/test_callback_handler.py | 39 ++++++ tests/unit/timing/phase/test_runner.py | 56 ++++++++ .../timing/strategies/test_agentic_replay.py | 7 +- .../test_agentic_replay_child_continuation.py | 48 +++++++ tests/unit/timing/test_branch_orchestrator.py | 38 ++++++ ...est_branch_orchestrator_dispatch_offset.py | 38 ++++++ 14 files changed, 537 insertions(+), 36 deletions(-) diff --git a/docs/benchmark-modes/dag.md b/docs/benchmark-modes/dag.md index 591259b148..8bcb61918e 100644 --- a/docs/benchmark-modes/dag.md +++ b/docs/benchmark-modes/dag.md @@ -275,7 +275,7 @@ If you are using `--concurrency` as a hard cap to protect a fragile server, size Children are dispatched reactively by `BranchOrchestrator` at credit-return time, not by the phase's `TimingStrategy` loop, and do not consume entries from the `DatasetSampler`. Their stop-condition behavior splits by intent: -- **`--request-count` (`RequestCountStopCondition`): HONORED for children.** It is a literal wire-request cap and applies to every credit on the wire. When the cap fires mid-tree, an in-flight child's remaining turns will be elided — `BranchStats.children_truncated` records the child, and `BranchStats.joins_suppressed` counts any parent join that was released without firing because the gated child was capped. Cancellation and duration timeouts honor the same rule. +- **`--request-count` (`RequestCountStopCondition`): HONORED for children.** It is a literal wire-request cap and applies to every credit on the wire. When the cap fires mid-tree, an in-flight child's remaining turns will be elided — `BranchStats.children_truncated` records the child, and `BranchStats.joins_suppressed` counts any parent join that was released without firing because the gated child was capped. Cancellation and duration timeouts honor the same rule. At a duration boundary, any child turn still waiting on a recorded replay timer—either the first turn of a SPAWN child or a later snapshot/continuation turn—is cancelled before the shared replay scheduler is swept. Its pre-registered join and descendant bookkeeping is settled so already-issued wire requests can drain without a phantom DAG wait. Timer IDs remain stable when an idle-gap cap advances the replay clock, allowing this cleanup to target the same logical timer after rescheduling. After sending stops and every frozen wire request returns, any residual join-only state is terminally truncated because it can no longer produce a legal request; this lets the phase finish without consuming the grace-period timeout. - **`--num-conversations` (`SessionCountStopCondition`): BYPASSED for children.** It targets sampler-plan completion ("run N full conversations") — children belong to a conversation tree and should run as part of their parent's session, not be truncated mid-tree. The wire-cap intent is served by `--request-count` instead. ### `--num-conversations` autodefault for `dag_jsonl` diff --git a/src/aiperf/common/loop_scheduler.py b/src/aiperf/common/loop_scheduler.py index 4a1588efe5..0acf24d0b7 100644 --- a/src/aiperf/common/loop_scheduler.py +++ b/src/aiperf/common/loop_scheduler.py @@ -108,6 +108,7 @@ def _safe_callback( self, handle_container: list[asyncio.TimerHandle | asyncio.Handle], coro: Coroutine, + tracked_handle_id: HandleId | None = None, ) -> None: """ Timer callback: transition coroutine from pending to running. @@ -118,7 +119,11 @@ def _safe_callback( mutable list, then populate it after call_later() returns. """ if handle_container[0] is not None: - handle_id = id(handle_container[0]) + handle_id = ( + tracked_handle_id + if tracked_handle_id is not None + else id(handle_container[0]) + ) self._handles.pop(handle_id, None) self._handle_groups.pop(handle_id, None) task = self._loop.create_task(coro) @@ -131,15 +136,19 @@ def _track_handle_and_return_id( coro: Coroutine, *, group_id: str | None = None, + tracked_handle_id: HandleId | None = None, ) -> HandleId: """Track a handle and coroutine and return the handle ID. - We use the handle ID as the key to avoid recursion when handles contain circular refs (handle_container). + New timers use the handle ID as their logical key to avoid recursion + when handles contain circular refs (handle_container). Replay idle-cap + rescheduling preserves that logical key even though it replaces the + physical timer. Also, we pass the handle id back to the caller instead of the original handle to ensure they do not keep a reference to the handle, and do not attempt to cancel it without stopping the coroutine. Returning the handle ID forces the caller to use the cancel_handle_id() method to cancel the coroutine. """ - handle_id = id(handle) + handle_id = tracked_handle_id if tracked_handle_id is not None else id(handle) self._handles[handle_id] = (handle, coro) if group_id is not None: self._handle_groups[handle_id] = group_id @@ -369,12 +378,12 @@ def cap_pending_delay(self, max_delay_sec: float) -> float: return 0.0 items = [ - (handle, coro, self._handle_groups.get(handle_id)) + (handle_id, handle, coro, self._handle_groups.get(handle_id)) for handle_id, (handle, coro) in self._handles.items() ] self._handles.clear() self._handle_groups.clear() - for handle, coro, group_id in items: + for handle_id, handle, coro, group_id in items: target = max(now, handle.when() - shift_sec) handle.cancel() handle_container = [None] @@ -384,14 +393,19 @@ def cap_pending_delay(self, max_delay_sec: float) -> float: # 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 + self._safe_callback, handle_container, coro, handle_id ) else: replacement = self._loop.call_at( - target, self._safe_callback, handle_container, coro + target, self._safe_callback, handle_container, coro, handle_id ) handle_container[0] = replacement - self._track_handle_and_return_id(replacement, coro, group_id=group_id) + self._track_handle_and_return_id( + replacement, + coro, + group_id=group_id, + tracked_handle_id=handle_id, + ) return shift_sec def cap_pending_delay_for_group(self, group_id: str, max_delay_sec: float) -> float: @@ -428,14 +442,19 @@ def cap_pending_delay_for_group(self, group_id: str, max_delay_sec: float) -> fl handle_container = [None] if target <= now: replacement = self._loop.call_soon( - self._safe_callback, handle_container, coro + self._safe_callback, handle_container, coro, handle_id ) else: replacement = self._loop.call_at( - target, self._safe_callback, handle_container, coro + target, self._safe_callback, handle_container, coro, handle_id ) handle_container[0] = replacement - self._track_handle_and_return_id(replacement, coro, group_id=group_id) + self._track_handle_and_return_id( + replacement, + coro, + group_id=group_id, + tracked_handle_id=handle_id, + ) return shift_sec @property diff --git a/src/aiperf/credit/callback_handler.py b/src/aiperf/credit/callback_handler.py index 8b7053d31f..2c7d017950 100644 --- a/src/aiperf/credit/callback_handler.py +++ b/src/aiperf/credit/callback_handler.py @@ -574,6 +574,14 @@ def _signal_all_credits_returned_if_ready( if allows_pending_branch_handoff else handler.progress.check_all_returned_or_cancelled() ) + if ( + self._branch_orchestrator is not None + and handler.lifecycle.is_sending_complete + and not allows_pending_branch_handoff + and all_wire_requests_returned + and self._branch_orchestrator.has_pending_branch_work() + ): + self._branch_orchestrator.truncate_pending_after_wire_drain() if ( self._branch_orchestrator is not None and not handler.progress.all_credits_returned_event.is_set() diff --git a/src/aiperf/timing/branch_orchestrator.py b/src/aiperf/timing/branch_orchestrator.py index cae0e92401..21f126f41b 100644 --- a/src/aiperf/timing/branch_orchestrator.py +++ b/src/aiperf/timing/branch_orchestrator.py @@ -23,8 +23,11 @@ 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. +dispatches. At a phase deadline, pending child dispatches are cancelled through +the orchestrator before the shared scheduler is swept so their pre-registered +join and descendant bookkeeping is rolled back. Datasets without timing +(``--ignore-trace-delays``) carry None timestamps and keep the +immediate-dispatch behavior. Sticky-routing locality (FORK mode) ----------------------------------- @@ -310,6 +313,11 @@ def __init__( # 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() + # child correlation ID -> (shared-scheduler handle ID, child, parent). + # A pending timer owns already-registered join/descendant bookkeeping; + # closing its coroutine through LoopScheduler alone would bypass the + # rollback body and leave has_pending_branch_work() true forever. + self._scheduled_delayed_dispatches: dict[str, tuple[int, object, str]] = {} # Drain observer: sync callback fired after state mutations that may # drain has_pending_branch_work() to False. Wired by # CreditCallbackHandler.set_branch_orchestrator to re-evaluate the @@ -1262,11 +1270,19 @@ def _start_delayed_first_turn( callers without a scheduler retain the legacy task/sleep fallback. """ if self._scheduler is not None: - self._scheduler.schedule_later( + handle_id = 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, ) + # offset_ms is strictly positive in this path, so schedule_later + # returns a pending handle ID rather than a running Task. + assert isinstance(handle_id, int) + self._scheduled_delayed_dispatches[child.x_correlation_id] = ( + handle_id, + child, + parent_corr, + ) self.stats.children_delayed += 1 return task = asyncio.create_task( @@ -1289,6 +1305,7 @@ async def _dispatch_first_turn_after_offset( refusal. Dispatch and settlement run under the parent lock, matching the intercept path's locking. """ + self._scheduled_delayed_dispatches.pop(child.x_correlation_id, None) await self._sleep_offset_ms(offset_ms) if self._cleaning_up: return @@ -1303,6 +1320,37 @@ async def _dispatch_first_turn_after_offset( self._rollback_failed_first_turn(child, result, parent_corr) await self._finalize_failed_dispatches(parent_corr) + async def cancel_pending_delayed_dispatches(self) -> None: + """Cancel scheduler-pending child turns and roll back DAG bookkeeping. + + PhaseRunner calls this before LoopScheduler.cancel_all_pending() at a + sending boundary. A coroutine that has already started is deliberately + left alone: it either reaches the wire and drains normally or observes + the phase stop condition and performs its existing refusal rollback. + """ + if self._scheduler is None: + return + + cancelled_by_parent: defaultdict[str, list[object]] = defaultdict(list) + # Do not await until every still-pending handle has been cancelled. This + # keeps a timer from firing between its cancellation and bookkeeping + # capture on the single event-loop thread. + for child_corr, (handle_id, child, parent_corr) in list( + self._scheduled_delayed_dispatches.items() + ): + if not self._scheduler.cancel_handle_id(handle_id): + continue + self._scheduled_delayed_dispatches.pop(child_corr, None) + cancelled_by_parent[parent_corr].append(child) + + for parent_corr, children in cancelled_by_parent.items(): + async with self._parent_locks[parent_corr]: + for child in children: + self._rollback_failed_first_turn( + child, ChildDispatchResult.REJECTED, parent_corr + ) + await self._finalize_failed_dispatches(parent_corr) + def _ensure_future_join( self, credit, @@ -1761,6 +1809,74 @@ def has_pending_branch_work(self) -> bool: return any(count > 0 for count in self._descendant_counts.values()) return False + def truncate_pending_after_wire_drain(self) -> None: + """Settle scheduling-only DAG state after a terminal phase cutoff. + + Once sending is complete and every frozen wire request has returned, + no pending join or descendant can legally produce another request. + Such state represents a continuation whose timer/barrier was cancelled + at the phase boundary, not live server work. Drain it synchronously so + ``all_credits_returned_event`` can fire without waiting for the grace + timeout. The caller owns both preconditions; this method deliberately + has no force/timeout policy of its own. + """ + if self._cleaning_up or not self.has_pending_branch_work(): + return + + tracked_children = list(self._child_to_join.items()) + pending_child_ids = set(self._child_root) | { + child_corr for child_corr, _ in tracked_children + } + parent_ids: set[str] = set() + + for child_corr, entries in tracked_children: + self._child_to_join.pop(child_corr, None) + if entries: + parent = entries[0].parent_correlation_id + parent_ids.add(parent) + child_mode = self._child_modes.pop(child_corr, None) + if ( + child_mode == ConversationBranchMode.FORK + and self._sticky_router is not None + ): + self._sticky_router.release_child_routing(parent) + if self._descendant_counts.get(parent, 0) > 0: + self._descendant_counts[parent] -= 1 + self._tree_descendant_done(child_corr) + + # ``_child_root`` is normally a subset of ``_child_to_join``. Settle + # any orphaned registry entries too; keeping them would leak a tree + # slot even though no wire callback can arrive for the child. + for child_corr in pending_child_ids: + if child_corr in self._child_root: + self._tree_descendant_done(child_corr) + self._child_modes.pop(child_corr, None) + + active_joins = len(self._active_joins) + future_joins = sum(len(gates) for gates in self._future_joins.values()) + parent_ids.update(self._active_joins) + parent_ids.update(self._future_joins) + self._active_joins.clear() + self._future_joins.clear() + + for parent in parent_ids: + self._release_parent_slot_if_drained(parent) + # A non-zero residue here has no remaining child owner. It is stale + # logical accounting at the terminal boundary, so clear it explicitly. + self._descendant_counts.clear() + self._parent_locks.clear() + + self.stats.children_truncated += len(tracked_children) + self.stats.joins_suppressed += active_joins + future_joins + logger.info( + "Terminal wire drain truncated pending DAG state: children=%d " + "active_joins=%d future_joins=%d", + len(tracked_children), + active_joins, + future_joins, + ) + self._notify_drain() + def snapshot_branch_stats(self) -> BranchStats: """Return a deep copy of the current branch stats. @@ -1778,6 +1894,10 @@ def cleanup(self) -> None: for task in self._delayed_dispatch_tasks: task.cancel() self._delayed_dispatch_tasks.clear() + if self._scheduler is not None: + for handle_id, _, _ in self._scheduled_delayed_dispatches.values(): + self._scheduler.cancel_handle_id(handle_id) + self._scheduled_delayed_dispatches.clear() s = self.stats logger.info( "BranchOrchestrator stats: spawned=%d completed=%d errored=%d " diff --git a/src/aiperf/timing/phase/runner.py b/src/aiperf/timing/phase/runner.py index 52ccf72383..b15347e4c4 100644 --- a/src/aiperf/timing/phase/runner.py +++ b/src/aiperf/timing/phase/runner.py @@ -1022,6 +1022,19 @@ async def _wait_for_sending_complete( if not self._lifecycle.is_sending_complete: self._lifecycle.mark_sending_complete(timeout_triggered=timed_out) self._progress.freeze_sent_counts() + if ( + self._branch_orchestrator is not None + and not preserve_branch_handoff + ): + await self._branch_orchestrator.cancel_pending_delayed_dispatches() + cancel_pending_child_turns = getattr( + strategy, "cancel_pending_child_turns", None + ) + if ( + cancel_pending_child_turns is not None + and not preserve_branch_handoff + ): + await cancel_pending_child_turns() self._scheduler.cancel_all_pending() if ( self._branch_orchestrator is not None @@ -1075,6 +1088,14 @@ async def _wait_for_returning_complete( if allows_pending_branch_handoff else self._progress.check_all_returned_or_cancelled() ) + if ( + all_wire_requests_returned + and not allows_pending_branch_handoff + and self._lifecycle.is_sending_complete + and self._branch_orchestrator is not None + and self._branch_orchestrator.has_pending_branch_work() + ): + self._branch_orchestrator.truncate_pending_after_wire_drain() if all_wire_requests_returned and ( allows_pending_branch_handoff or self._branch_orchestrator is None diff --git a/src/aiperf/timing/strategies/agentic_replay.py b/src/aiperf/timing/strategies/agentic_replay.py index 2bd8a8d082..c7882d0dd6 100644 --- a/src/aiperf/timing/strategies/agentic_replay.py +++ b/src/aiperf/timing/strategies/agentic_replay.py @@ -200,6 +200,11 @@ def __init__( self._accelerated_warmup_started = False self._handoff_credits: dict[str, Credit] = {} self._root_to_lane: dict[str, int] = {} + # Profiling child turns can be parked in LoopScheduler for recorded + # inter-turn delays. The orchestrator has already registered those + # children, so a duration-boundary scheduler sweep must notify it when + # a never-started turn is discarded or the parent join never drains. + self._scheduled_child_turns: 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. @@ -948,6 +953,48 @@ async def _issue_child_continuation_or_drain(self, turn: TurnToSend) -> None: ): await self.branch_orchestrator.on_child_stopped(turn.x_correlation_id) + async def _run_scheduled_child_turn(self, turn: TurnToSend) -> None: + """Run one delayed child turn after dropping its pending marker.""" + self._scheduled_child_turns.pop(turn.x_correlation_id, None) + await self._issue_child_continuation_or_drain(turn) + + def _schedule_child_turn(self, turn: TurnToSend, delay_s: float) -> None: + """Schedule a child turn whose cancellation still owns DAG cleanup.""" + handle_id = self.scheduler.schedule_later( + delay_s, + self._run_scheduled_child_turn(turn), + group_id=turn.effective_root_correlation_id, + ) + # Positive-delay LoopScheduler calls return an integer handle. Some + # isolated strategy tests supply an eager task instead; that work has + # already started and therefore is not pending-boundary cleanup state. + if isinstance(handle_id, int): + self._scheduled_child_turns[turn.x_correlation_id] = handle_id + + async def cancel_pending_child_turns(self) -> None: + """Cancel never-started profiling child turns and drain their DAG state. + + All scheduler cancellations happen before the first await so a timer + cannot fire between cancellation and bookkeeping capture on the event + loop thread. A handle that already fired is left alone: its wrapper + will issue the turn or route the terminal refusal through the normal + child-drain chokepoint. + """ + if self.allows_pending_branch_handoff_after_sending_complete: + return + + cancelled: list[str] = [] + for child_corr, handle_id in list(self._scheduled_child_turns.items()): + if not self.scheduler.cancel_handle_id(handle_id): + continue + self._scheduled_child_turns.pop(child_corr, None) + cancelled.append(child_corr) + + if self.branch_orchestrator is None: + return + for child_corr in cancelled: + await self.branch_orchestrator.on_child_stopped(child_corr) + async def finalize_phase(self) -> None: """Persist the drained accelerated-warmup DAG for profiling.""" if self._system_idle_gap_cap_seconds is not None: @@ -1551,19 +1598,20 @@ async def _dispatch_next_turn(self, credit: Credit) -> None: next_meta = self.conversation_source.get_next_turn_metadata(credit) turn = TurnToSend.from_previous_credit(credit, next_meta) - coro = ( - self._issue_child_continuation_or_drain(turn) - if turn.agent_depth > 0 - 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, - group_id=credit.effective_root_correlation_id, - ) + delay_s = next_meta.delay_ms / MILLIS_PER_SECOND + if turn.agent_depth > 0: + self._schedule_child_turn(turn, delay_s) + else: + self.scheduler.schedule_later( + delay_s, + self.credit_issuer.issue_credit(turn), + group_id=credit.effective_root_correlation_id, + ) + elif turn.agent_depth > 0: + await self._issue_child_continuation_or_drain(turn) else: - await coro + await self.credit_issuer.issue_credit(turn) async def _spawn_from_recycle_or_id( self, @@ -1794,11 +1842,14 @@ async def _dispatch_snapshot_for_profiling( offset_by_corr[state.x_correlation_id] - t0_offset_ms ) / MILLIS_PER_SECOND if delay_s > 0: - self.scheduler.schedule_later( - delay_s, - self.credit_issuer.issue_credit(turn), - group_id=turn.effective_root_correlation_id, - ) + if turn.agent_depth > 0: + self._schedule_child_turn(turn, delay_s) + else: + 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) diff --git a/tests/integration/test_weka_flat_split_e2e.py b/tests/integration/test_weka_flat_split_e2e.py index 4c5f8aa77a..4ce92a5ba5 100644 --- a/tests/integration/test_weka_flat_split_e2e.py +++ b/tests/integration/test_weka_flat_split_e2e.py @@ -212,6 +212,21 @@ def _write_independent_stream_drift_trace(target_dir: Path) -> Path: return target_dir +def _write_duration_cutoff_child_trace(target_dir: Path) -> Path: + """A child whose next recorded turn lies far beyond the phase duration.""" + _write_trace( + target_dir, + "trace_duration_cutoff", + [ + _req(0.0, [1, 2, 3], api_time=0.01, out=1), + _req(0.05, [1, 2, 10], api_time=0.01, out=1), + _req(60.05, [1, 2, 10, 11], api_time=0.01, out=1), + _req(60.10, [1, 2, 3, 4], api_time=0.01, 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( @@ -266,10 +281,8 @@ async def _run_weka_profile( "--no-fixed-schedule", "--benchmark-duration", str(duration), - # Bounded grace: normal drains finish in ~1-2s; the bound also caps - # the cost of a worst-case DAG drain stall (a delay-scheduled child - # turn cancelled by the deadline's cancel_all_pending leaves - # has_pending_branch_work stuck until the grace timeout fires). + # Bounded grace: normal drains finish in ~1-2s; the bound also prevents + # a future DAG teardown regression from wedging the integration suite. "--benchmark-grace-period", "20", "--random-seed", @@ -634,6 +647,38 @@ async def test_cache_warmup_handoff_preserves_order_and_spawn_joins( assert branch_stats.parents_resumed >= 1, branch_stats +async def test_duration_cutoff_cancels_delayed_child_without_grace_stall( + tmp_path: Path, mock_server_factory: MockServerFactory +) -> None: + """A discarded child continuation drains its join at the send boundary.""" + corpus = _write_duration_cutoff_child_trace(tmp_path / "traces") + async with mock_server_factory(fast=True, 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", + "--warmup-requests-per-lane", + "1", + "--trajectory-start-min-ratio", + "0", + "--trajectory-start-max-ratio", + "0", + ], + # The configured grace is 20s. Exiting under this bound proves the + # phase did not wait for the grace backstop after wire drain. + timeout=18.0, + ) + + _assert_success(result, "duration-cutoff delayed child teardown") + assert "grace_period_timeout=True" not in _combined_log_text(result) + + 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/common/test_loop_scheduler.py b/tests/unit/common/test_loop_scheduler.py index 473dc0b30b..158a49c128 100644 --- a/tests/unit/common/test_loop_scheduler.py +++ b/tests/unit/common/test_loop_scheduler.py @@ -160,6 +160,19 @@ async def noop(): assert scheduler.cap_pending_delay(2.0) == 0.0 scheduler.cancel_all_pending() + async def test_handle_id_survives_timer_reschedule(self, scheduler: LoopScheduler): + """Owners can still cancel a timer after an idle-cap clock jump.""" + + async def should_not_run(): + raise AssertionError("cancelled timer ran") + + handle_id = scheduler.schedule_later(60.0, should_not_run()) + assert isinstance(handle_id, int) + scheduler.cap_pending_delay(10.0) + + assert scheduler.cancel_handle_id(handle_id) is True + assert scheduler.pending_count == 0 + async def test_cap_pending_delay_for_group_isolated_and_uniform( self, scheduler: LoopScheduler ): diff --git a/tests/unit/credit/test_callback_handler.py b/tests/unit/credit/test_callback_handler.py index eabd9b99ed..ab6832179d 100644 --- a/tests/unit/credit/test_callback_handler.py +++ b/tests/unit/credit/test_callback_handler.py @@ -98,6 +98,45 @@ def registered_handler( return callback_handler +def test_terminal_wire_drain_settles_dag_before_signalling_completion( + callback_handler, + mock_progress, + mock_lifecycle, + mock_stop_checker, + mock_strategy, + mock_branch_orchestrator, +): + """The last wire return truncates only scheduling residue, then completes.""" + pending = True + + def _has_pending() -> bool: + return pending + + def _truncate() -> None: + nonlocal pending + pending = False + + mock_progress.in_flight = 0 + mock_progress.check_all_returned_or_cancelled.return_value = True + mock_lifecycle.is_sending_complete = True + mock_branch_orchestrator.has_pending_branch_work.side_effect = _has_pending + mock_branch_orchestrator.truncate_pending_after_wire_drain.side_effect = _truncate + callback_handler.register_phase( + phase=CreditPhase.PROFILING, + progress=mock_progress, + lifecycle=mock_lifecycle, + stop_checker=mock_stop_checker, + strategy=mock_strategy, + ) + callback_handler.set_branch_orchestrator(mock_branch_orchestrator) + context = callback_handler._phase_handlers[CreditPhase.PROFILING] + + callback_handler._signal_all_credits_returned_if_ready(context) + + mock_branch_orchestrator.truncate_pending_after_wire_drain.assert_called_once_with() + assert mock_progress.all_credits_returned_event.is_set() + + def make_credit( credit_id: int = 1, conversation_id: str = "conv1", diff --git a/tests/unit/timing/phase/test_runner.py b/tests/unit/timing/phase/test_runner.py index 3ec7cd22ed..584abb17e3 100644 --- a/tests/unit/timing/phase/test_runner.py +++ b/tests/unit/timing/phase/test_runner.py @@ -204,6 +204,62 @@ async def runner( class TestPhaseRunnerLifecycle: + async def test_sending_boundary_rolls_back_delayed_children_before_scheduler_sweep( + self, runner: PhaseRunner + ) -> None: + events: list[str] = [] + orchestrator = MagicMock() + orchestrator.cancel_pending_delayed_dispatches = AsyncMock( + side_effect=lambda: events.append("cancel_children") + ) + orchestrator.expire_replay_deadlines = AsyncMock( + side_effect=lambda: events.append("expire_joins") + ) + strategy = MockStrategy() + strategy.cancel_pending_child_turns = AsyncMock( + side_effect=lambda: events.append("cancel_continuations") + ) + runner._branch_orchestrator = orchestrator + runner._scheduler.cancel_all_pending = MagicMock( + side_effect=lambda: events.append("cancel_scheduler") + ) + runner._lifecycle.start() + runner._progress.all_credits_sent_event.set() + + await runner._wait_for_sending_complete(strategy) + + assert events == [ + "cancel_children", + "cancel_continuations", + "cancel_scheduler", + "expire_joins", + ] + + async def test_pre_return_wait_truncates_dag_after_wire_already_drained( + self, runner: PhaseRunner + ) -> None: + pending = True + + def _has_pending() -> bool: + return pending + + def _truncate() -> None: + nonlocal pending + pending = False + + orchestrator = MagicMock() + orchestrator.has_pending_branch_work.side_effect = _has_pending + orchestrator.truncate_pending_after_wire_drain.side_effect = _truncate + runner._branch_orchestrator = orchestrator + runner._progress.check_all_returned_or_cancelled = MagicMock(return_value=True) + runner._lifecycle.start() + runner._lifecycle.mark_sending_complete(timeout_triggered=True) + + await runner._wait_for_returning_complete(MockStrategy()) + + orchestrator.truncate_pending_after_wire_drain.assert_called_once_with() + assert runner._progress.all_credits_returned_event.is_set() + async def test_baseline_boundary_capture_is_fire_and_forget( self, runner: PhaseRunner, pub: MagicMock ) -> None: diff --git a/tests/unit/timing/strategies/test_agentic_replay.py b/tests/unit/timing/strategies/test_agentic_replay.py index f92481e1b0..f8ea437679 100644 --- a/tests/unit/timing/strategies/test_agentic_replay.py +++ b/tests/unit/timing/strategies/test_agentic_replay.py @@ -1675,6 +1675,7 @@ async def capture(turn): issuer = AsyncMock() issuer.issue_credit.side_effect = capture + issuer.dispatch_child_turn.side_effect = capture scheduled: list[tuple[float, object]] = [] def fake_schedule_later(delay, coro, **_kwargs): @@ -1683,6 +1684,7 @@ def fake_schedule_later(delay, coro, **_kwargs): scheduler = MagicMock() scheduler.schedule_later.side_effect = fake_schedule_later branch_orchestrator = MagicMock() + branch_orchestrator.on_child_stopped = AsyncMock() cfg = MagicMock() cfg.phase = CreditPhase.PROFILING @@ -1793,6 +1795,7 @@ async def capture(turn): issuer = AsyncMock() issuer.issue_credit.side_effect = capture + issuer.dispatch_child_turn.side_effect = capture scheduled: list[tuple[float, object]] = [] scheduler = MagicMock() scheduler.schedule_later.side_effect = ( @@ -1802,6 +1805,8 @@ async def capture(turn): cfg = MagicMock() cfg.phase = CreditPhase.PROFILING cfg.concurrency = 1 + branch_orchestrator = MagicMock() + branch_orchestrator.on_child_stopped = AsyncMock() strategy = AgenticReplayStrategy( config=cfg, conversation_source=src, @@ -1809,7 +1814,7 @@ async def capture(turn): stop_checker=MagicMock(), credit_issuer=issuer, lifecycle=MagicMock(), - branch_orchestrator=MagicMock(), + branch_orchestrator=branch_orchestrator, ) await strategy.setup_phase() await strategy.execute_phase() diff --git a/tests/unit/timing/strategies/test_agentic_replay_child_continuation.py b/tests/unit/timing/strategies/test_agentic_replay_child_continuation.py index 039fb24905..bcd6e8c69c 100644 --- a/tests/unit/timing/strategies/test_agentic_replay_child_continuation.py +++ b/tests/unit/timing/strategies/test_agentic_replay_child_continuation.py @@ -38,11 +38,13 @@ def _make_strategy( credit_issuer.issue_credit = AsyncMock(return_value=True) scheduler = MagicMock() + scheduler.schedule_later.return_value = 123 strategy.conversation_source = conversation_source strategy.credit_issuer = credit_issuer strategy.scheduler = scheduler strategy.branch_orchestrator = branch_orchestrator + strategy._scheduled_child_turns = {} return strategy, credit_issuer, scheduler @@ -167,6 +169,52 @@ async def test_child_delayed_schedules_chokepoint_coro_not_issue_credit() -> Non orch.on_child_stopped.assert_awaited_once_with("child-xcid") +@pytest.mark.asyncio +async def test_duration_boundary_cancels_delayed_child_and_drains_orchestrator() -> ( + None +): + """A never-started continuation must not leave its parent join pending.""" + orch = MagicMock() + orch.on_child_stopped = AsyncMock() + strategy, issuer, scheduler = _make_strategy( + branch_orchestrator=orch, delay_ms=60_000.0 + ) + + await strategy._dispatch_next_turn(_child_credit()) + _, coro = scheduler.schedule_later.call_args.args + + def cancel_handle(handle_id: int) -> bool: + assert handle_id == 123 + coro.close() + return True + + scheduler.cancel_handle_id.side_effect = cancel_handle + await strategy.cancel_pending_child_turns() + + scheduler.cancel_handle_id.assert_called_once_with(123) + issuer.dispatch_child_turn.assert_not_awaited() + orch.on_child_stopped.assert_awaited_once_with("child-xcid") + assert strategy._scheduled_child_turns == {} + + +@pytest.mark.asyncio +async def test_fired_delayed_child_is_not_drained_by_pending_cleanup() -> None: + """Once a timer fires, its live issue/refusal path owns DAG settlement.""" + orch = MagicMock() + orch.on_child_stopped = AsyncMock() + strategy, _, scheduler = _make_strategy(branch_orchestrator=orch, delay_ms=250.0) + + await strategy._dispatch_next_turn(_child_credit()) + _, coro = scheduler.schedule_later.call_args.args + scheduler.cancel_handle_id.return_value = False + + await strategy.cancel_pending_child_turns() + + orch.on_child_stopped.assert_not_awaited() + await coro + assert strategy._scheduled_child_turns == {} + + # Root continuation keeps issue_credit diff --git a/tests/unit/timing/test_branch_orchestrator.py b/tests/unit/timing/test_branch_orchestrator.py index 6aa2283147..445d6e1f20 100644 --- a/tests/unit/timing/test_branch_orchestrator.py +++ b/tests/unit/timing/test_branch_orchestrator.py @@ -232,6 +232,44 @@ async def test_no_join_case_releases_slot_when_descendants_drain(): assert released == ["parent"] +def test_terminal_wire_drain_truncates_scheduling_only_dag_state(): + """No logical DAG residue may hold a duration phase after wire drain.""" + sticky_router = MagicMock() + registry = MagicMock() + orch = BranchOrchestrator( + conversation_source=MagicMock(), + credit_issuer=MagicMock(), + sticky_router=sticky_router, + session_tree_registry=registry, + ) + pending = _mk_pending_for_parent( + "parent", + gated_turn_index=2, + prereq_key="SPAWN_JOIN:b0", + outstanding={"child"}, + ) + pending.is_blocked = True + orch._active_joins["parent"] = pending + orch._child_to_join["child"] = [ + ChildJoinEntry( + parent_correlation_id="parent", + gated_turn_index=2, + prereq_key="SPAWN_JOIN:b0", + ) + ] + orch._child_modes["child"] = ConversationBranchMode.FORK + orch._child_root["child"] = "root" + orch._descendant_counts["parent"] = 1 + + orch.truncate_pending_after_wire_drain() + + assert not orch.has_pending_branch_work() + assert orch.stats.children_truncated == 1 + assert orch.stats.joins_suppressed == 1 + sticky_router.release_child_routing.assert_called_once_with("parent") + registry.on_descendant_done.assert_called_once_with("root") + + @pytest.mark.asyncio async def test_leaf_for_unknown_child_is_noop(): orch = BranchOrchestrator( diff --git a/tests/unit/timing/test_branch_orchestrator_dispatch_offset.py b/tests/unit/timing/test_branch_orchestrator_dispatch_offset.py index f3e11fc2fc..b554a7d529 100644 --- a/tests/unit/timing/test_branch_orchestrator_dispatch_offset.py +++ b/tests/unit/timing/test_branch_orchestrator_dispatch_offset.py @@ -122,6 +122,44 @@ async def test_offset_child_uses_shared_scheduler_for_idle_cap( scheduler.cancel_all() +@pytest.mark.asyncio +async def test_phase_boundary_cancels_pending_scheduler_child_and_drains_dag() -> None: + """A duration cutoff must roll back a delayed child before scheduler sweep. + + Regression for the full-replay failure where LoopScheduler closed the + never-started coroutine but BranchOrchestrator retained its child, join, + and descendant entries until the phase grace timeout. + """ + parent = _parent_conv( + [_spawn_branch("b0", ["kid"], start_timestamp_ms=0.0, is_background=False)], + gated_turn=1, + ) + scheduler = LoopScheduler() + orch, _, issuer = _mk_harness( + [parent, _child_conv("kid", 60_000.0)], + scheduler=scheduler, + ) + + assert await orch.intercept(_mk_credit("parent", "P", 0)) is True + # One timer is the child dispatch and one is the gated parent's replay + # deadline. The branch-specific cancellation must remove only the former. + assert scheduler.pending_count == 2 + assert orch.has_pending_branch_work() + assert "P" in orch._active_joins + + await orch.cancel_pending_delayed_dispatches() + assert scheduler.pending_count == 1 + scheduler.cancel_all_pending() + await orch.expire_replay_deadlines() + + issuer.dispatch_first_turn.assert_not_awaited() + issuer.dispatch_join_turn.assert_awaited_once() + assert scheduler.pending_count == 0 + assert orch.stats.children_truncated == 1 + assert orch._scheduled_delayed_dispatches == {} + assert not orch.has_pending_branch_work() + + def _mk_credit(conv_id: str, corr_id: str, turn_index: int): return MagicMock( x_correlation_id=corr_id,