fix(agentic): settle DAG state after wire drain / 网络排空后收敛 DAG 状态

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`

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

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

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 "

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