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,