MAINT: Split TAP node execution scheduling (#2382)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: romanlutz

pyrit/executor/attack/multi_turn/tree_of_attacks.py

@@ -69,6 +69,7 @@
 from pyrit.score.true_false.true_false_inverter_scorer import TrueFalseInverterScorer
 
 if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
     from pathlib import Path
 
     from pyrit.models.literals import PromptDataType
@@ -1261,6 +1262,72 @@ def __str__(self) -> str:
     __repr__ = __str__
 
 
+class _TreeOfAttacksNodeExecutor:
+    """Execute independent tree nodes with bounded concurrency."""
+
+    def __init__(
+        self,
+        *,
+        batch_size: int,
+        logger: logging.Logger | logging.LoggerAdapter[logging.Logger],
+    ) -> None:
+        """
+        Initialize the node executor.
+
+        Args:
+            batch_size (int): Maximum number of nodes to execute concurrently.
+            logger (logging.Logger | logging.LoggerAdapter[logging.Logger]): Logger
+                used for execution progress.
+        """
+        self._batch_size = batch_size
+        self._logger = logger
+
+    async def execute_nodes_async(
+        self,
+        *,
+        nodes: list[_TreeOfAttacksNode],
+        objective: str,
+    ) -> AsyncIterator[tuple[int, list[_TreeOfAttacksNode]]]:
+        """
+        Execute nodes in ordered batches and yield each completed batch.
+
+        Node instances own all branch-specific mutable state. This executor only
+        schedules their existing execution protocol, so failures and cancellation
+        retain ``asyncio.gather`` semantics.
+
+        Args:
+            nodes (list[_TreeOfAttacksNode]): Nodes to execute.
+            objective (str): Objective passed to every node.
+
+        Yields:
+            tuple[int, list[_TreeOfAttacksNode]]: The batch start offset and nodes
+                after every node in that batch has completed.
+        """
+        for batch_start in range(0, len(nodes), self._batch_size):
+            batch_nodes = nodes[batch_start : batch_start + self._batch_size]
+            self._log_batch_start(batch_start=batch_start, batch_nodes=batch_nodes, total_nodes=len(nodes))
+
+            await asyncio.gather(*(node.send_prompt_async(objective=objective) for node in batch_nodes))
+
+            yield batch_start, batch_nodes
+
+    def _log_batch_start(
+        self,
+        *,
+        batch_start: int,
+        batch_nodes: list[_TreeOfAttacksNode],
+        total_nodes: int,
+    ) -> None:
+        """Log the batch and node dispatch order."""
+        batch_end = batch_start + len(batch_nodes)
+        self._logger.debug(
+            f"Processing batch {batch_start // self._batch_size + 1} "
+            f"(nodes {batch_start + 1}-{batch_end} of {total_nodes})"
+        )
+        for node_index in range(batch_start + 1, batch_end + 1):
+            self._logger.debug(f"Preparing prompt for node {node_index}/{total_nodes}")
+
+
 class TreeOfAttacksWithPruningAttack(AttackStrategy[TAPAttackContext, TAPAttackResult]):
     """
     Implement the Tree of Attacks with Pruning (TAP) attack strategy.
@@ -1402,6 +1469,10 @@ def __init__(
         super().__init__(objective_target=objective_target, logger=logger, context_type=TAPAttackContext)
 
         self._memory = CentralMemory.get_memory_instance()
+        self._node_executor = _TreeOfAttacksNodeExecutor(
+            batch_size=self._configuration.batch_size,
+            logger=self._logger,
+        )
 
         # Initialize adversarial configuration
         self._adversarial_chat = attack_adversarial_config.target
@@ -1885,25 +1956,10 @@ async def _send_prompts_to_all_nodes_async(self, context: TAPAttackContext) -> N
             context.tree_visualization.create_node(f"{context.executed_turns}: ", vis_id, parent=node._vis_node_id)
             node._vis_node_id = vis_id
 
-        # Process nodes in batches
-        for batch_start in range(0, len(context.nodes), self._configuration.batch_size):
-            batch_end = min(batch_start + self._configuration.batch_size, len(context.nodes))
-            batch_nodes = context.nodes[batch_start:batch_end]
-
-            self._logger.debug(
-                f"Processing batch {batch_start // self._configuration.batch_size + 1} "
-                f"(nodes {batch_start + 1}-{batch_end} of {len(context.nodes)})"
-            )
-
-            # Create tasks for parallel execution
-            tasks = []
-            for node_index, node in enumerate(batch_nodes, start=batch_start + 1):
-                self._logger.debug(f"Preparing prompt for node {node_index}/{len(context.nodes)}")
-                task = node.send_prompt_async(objective=context.objective)
-                tasks.append(task)
-
-            await asyncio.gather(*tasks)
-
+        async for batch_start, batch_nodes in self._node_executor.execute_nodes_async(
+            nodes=context.nodes,
+            objective=context.objective,
+        ):
             # Update visualization with results after batch completes
             for node_index, node in enumerate(batch_nodes, start=batch_start + 1):
                 result_string = self._format_node_result(node)

tests/unit/executor/attack/multi_turn/test_tree_of_attacks_node_executor.py

@@ -0,0 +1,103 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import asyncio
+import logging
+from dataclasses import dataclass
+from unittest.mock import AsyncMock
+
+import pytest
+
+from pyrit.executor.attack.multi_turn.tree_of_attacks import _TreeOfAttacksNodeExecutor
+
+
+@dataclass
+class _NodeState:
+    node_id: str
+    executions: int = 0
+    objective: str | None = None
+
+
+def _make_node(*, node_id: str, call_order: list[str]) -> tuple[_NodeState, AsyncMock]:
+    state = _NodeState(node_id=node_id)
+
+    async def execute_async(objective: str) -> None:
+        call_order.append(node_id)
+        state.executions += 1
+        state.objective = objective
+
+    return state, AsyncMock(side_effect=execute_async)
+
+
+async def test_execute_nodes_async_preserves_batch_order_and_node_state() -> None:
+    call_order: list[str] = []
+    states_and_methods = [_make_node(node_id=f"node-{index}", call_order=call_order) for index in range(5)]
+    nodes = []
+    for state, method in states_and_methods:
+        node = AsyncMock()
+        node.state = state
+        node.send_prompt_async = method
+        nodes.append(node)
+
+    executor = _TreeOfAttacksNodeExecutor(batch_size=2, logger=logging.getLogger(__name__))
+
+    completed_batches = [
+        (batch_start, [node.state.node_id for node in batch])
+        async for batch_start, batch in executor.execute_nodes_async(nodes=nodes, objective="objective")
+    ]
+
+    assert completed_batches == [(0, ["node-0", "node-1"]), (2, ["node-2", "node-3"]), (4, ["node-4"])]
+    assert call_order == ["node-0", "node-1", "node-2", "node-3", "node-4"]
+    assert [state.executions for state, _ in states_and_methods] == [1, 1, 1, 1, 1]
+    assert [state.objective for state, _ in states_and_methods] == ["objective"] * 5
+
+
+async def test_execute_nodes_async_stops_before_next_batch_on_failure() -> None:
+    first = AsyncMock()
+    first.send_prompt_async = AsyncMock(return_value=None)
+    failing = AsyncMock()
+    failing.send_prompt_async = AsyncMock(side_effect=RuntimeError("node failed"))
+    not_started = AsyncMock()
+    not_started.send_prompt_async = AsyncMock(return_value=None)
+    executor = _TreeOfAttacksNodeExecutor(batch_size=2, logger=logging.getLogger(__name__))
+
+    with pytest.raises(RuntimeError, match="node failed"):
+        async for _ in executor.execute_nodes_async(
+            nodes=[first, failing, not_started],
+            objective="objective",
+        ):
+            pass
+
+    first.send_prompt_async.assert_awaited_once_with(objective="objective")
+    failing.send_prompt_async.assert_awaited_once_with(objective="objective")
+    not_started.send_prompt_async.assert_not_awaited()
+
+
+async def test_execute_nodes_async_propagates_cancellation_to_active_nodes() -> None:
+    started = asyncio.Event()
+    cancelled = asyncio.Event()
+
+    async def wait_for_cancellation_async(objective: str) -> None:
+        started.set()
+        try:
+            await asyncio.Event().wait()
+        finally:
+            cancelled.set()
+
+    node = AsyncMock()
+    node.send_prompt_async = AsyncMock(side_effect=wait_for_cancellation_async)
+    executor = _TreeOfAttacksNodeExecutor(batch_size=1, logger=logging.getLogger(__name__))
+
+    async def consume_async() -> None:
+        async for _ in executor.execute_nodes_async(nodes=[node], objective="objective"):
+            pass
+
+    task = asyncio.create_task(consume_async())
+    await started.wait()
+    task.cancel()
+
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    assert cancelled.is_set()
+    node.send_prompt_async.assert_awaited_once_with(objective="objective")