diff --git a/pyrit/executor/attack/multi_turn/tree_of_attacks.py b/pyrit/executor/attack/multi_turn/tree_of_attacks.py index 639f0e2427..38b44db43b 100644 --- a/pyrit/executor/attack/multi_turn/tree_of_attacks.py +++ b/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) diff --git a/tests/unit/executor/attack/multi_turn/test_tree_of_attacks_node_executor.py b/tests/unit/executor/attack/multi_turn/test_tree_of_attacks_node_executor.py new file mode 100644 index 0000000000..896d4f2555 --- /dev/null +++ b/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")