FIX: Reject concurrent SDK scenario runs and make attempt accounting atomic (#3748)

pyrit/exceptions/__init__.py

@@ -22,6 +22,7 @@
         PyritException,
         RateLimitException,
         ScenarioPartialFailureException,
+        ScenarioRunInProgressException,
         ScorerLLMResponseBlockedException,
         get_retry_max_num_attempts,
         handle_bad_request_exception,
@@ -77,6 +78,7 @@
     "remove_markdown_json": "pyrit.exceptions.exceptions_helpers",
     "RetryCollector": "pyrit.exceptions.retry_collector",
     "ScenarioPartialFailureException": "pyrit.exceptions.exception_classes",
+    "ScenarioRunInProgressException": "pyrit.exceptions.exception_classes",
     "ScorerLLMResponseBlockedException": "pyrit.exceptions.exception_classes",
     "set_execution_context": "pyrit.exceptions.exception_context",
     "set_retry_collector": "pyrit.exceptions.retry_collector",

pyrit/exceptions/exception_classes.py

@@ -304,6 +304,23 @@ def __init__(
             self.__cause__ = self.incomplete_objectives[0][1]
 
 
+class ScenarioRunInProgressException(PyritException, RuntimeError):  # noqa: N818
+    """Raised when this process is already executing the same scenario result."""
+
+    def __init__(self, *, status_code: int = 409, message: str | None = None) -> None:
+        """
+        Initialize a scenario run-in-progress exception.
+
+        Args:
+            status_code (int): Status code for the error. Defaults to 409.
+            message (str | None): Error message describing the in-progress run.
+        """
+        super().__init__(
+            status_code=status_code,
+            message=message or "Scenario execution is already in progress.",
+        )
+
+
 class InvalidJsonException(PyritException):
     """Exception class for blocked content errors."""
 

pyrit/memory/memory_interface.py

@@ -7161,6 +7161,38 @@ def _execute_update_scenario_metadata_fields(
             flag_modified(entry, "scenario_metadata")
             session.commit()
 
+    def _execute_increment_scenario_number_tries(
+        self,
+        *,
+        scenario_result_id: str,
+    ) -> int:
+        """
+        Atomically increment and return the number_tries counter for a scenario result.
+
+        Args:
+            scenario_result_id (str): The ID of the scenario result.
+
+        Returns:
+            int: The new number_tries value after incrementing.
+
+        Raises:
+            ValueError: If the scenario result is not found in memory.
+        """
+        with closing(self._get_session()) as session:
+            updated_rows = (
+                session.query(ScenarioResultEntry)
+                .filter(ScenarioResultEntry.id == scenario_result_id)
+                .update(
+                    {ScenarioResultEntry.number_tries: ScenarioResultEntry.number_tries + 1},
+                    synchronize_session=False,
+                )
+            )
+            if not updated_rows:
+                raise ValueError(f"Scenario result with ID {scenario_result_id} not found in memory")
+            new_tries: int = session.query(ScenarioResultEntry.number_tries).filter_by(id=scenario_result_id).scalar()
+            session.commit()
+            return new_tries
+
     def _execute_get_scenario_result_header(self, *, scenario_result_id: str) -> ScenarioResult | None:
         """Return one ScenarioResult header without hydrating linked attack results."""
         with closing(self._get_session()) as session:
@@ -11757,6 +11789,51 @@ async def update_scenario_metadata_fields_async(
             self._execute_update_scenario_metadata_fields, scenario_result_id=scenario_result_id, fields=fields
         )
 
+    def increment_scenario_number_tries(
+        self,
+        *,
+        scenario_result_id: str,
+    ) -> int:
+        """
+        Use ``increment_scenario_number_tries_async``.
+
+        This synchronous API is deprecated and can block the caller.
+
+        Returns:
+            int: The new number_tries value after incrementing.
+        """
+        print_deprecation_message(
+            old_item="MemoryInterface.increment_scenario_number_tries",
+            new_item="MemoryInterface.increment_scenario_number_tries_async",
+            removed_in="1.4.0",
+        )
+        return self._execute_increment_scenario_number_tries(
+            scenario_result_id=scenario_result_id,
+        )
+
+    @legacy_sync_override(lambda: MemoryInterface.increment_scenario_number_tries)
+    async def increment_scenario_number_tries_async(
+        self,
+        *,
+        scenario_result_id: str,
+    ) -> int:
+        """
+        Atomically increment and return the number_tries counter for a scenario result.
+
+        Args:
+            scenario_result_id (str): The ID of the scenario result to update.
+
+        Returns:
+            int: The new number_tries value after incrementing.
+
+        Raises:
+            ValueError: If the scenario result is not found.
+        """
+        return await self._run_database_operation_async(
+            self._execute_increment_scenario_number_tries,
+            scenario_result_id=scenario_result_id,
+        )
+
     def get_scenario_result_header(self, *, scenario_result_id: str) -> ScenarioResult | None:
         """
         Use ``get_scenario_result_header_async``.

pyrit/scenario/core/scenario.py

@@ -11,6 +11,7 @@
 import ast
 import asyncio
 import logging
+import threading
 import uuid
 from abc import ABC, abstractmethod
 from collections.abc import Mapping, Sequence
@@ -23,7 +24,7 @@
 
 from pyrit.common import default_values, get_global_default_values
 from pyrit.common.utils import to_sha256
-from pyrit.exceptions import ScenarioPartialFailureException
+from pyrit.exceptions import ScenarioPartialFailureException, ScenarioRunInProgressException
 from pyrit.exceptions.exceptions_helpers import summarize_exception
 from pyrit.executor.attack import AttackExecutor, AttackExecutorResult
 from pyrit.executor.attack.core.attack_executor import ActiveAttackExecution, _observe_interrupted_executions
@@ -168,6 +169,16 @@ class Scenario(ABC):
     #: Maximum number of resolved concrete techniques; None leaves the selection unrestricted.
     MAX_CONCRETE_TECHNIQUES: ClassVar[int | None] = None
 
+    #: Active scenario result IDs currently executing in this process.
+    _active_runs: ClassVar[set[str]] = set()
+    _active_runs_lock: ClassVar[threading.Lock] = threading.Lock()
+
+    @classmethod
+    def _clear_active_runs(cls) -> None:
+        """Clear active scenario runs. Intended for test isolation."""
+        with cls._active_runs_lock:
+            cls._active_runs.clear()
+
     def __init_subclass__(cls, **kwargs: Any) -> None:
         """
         Enforce the keyword-only constructor contract on subclasses.
@@ -1689,6 +1700,7 @@ async def run_async(self) -> ScenarioResult:
                 attack results from all atomic attacks.
 
         Raises:
+            ScenarioRunInProgressException: If this scenario result run is already in progress in this process.
             asyncio.CancelledError: If the scenario task is cancelled. Completed results remain persisted
                 and a later call can resume the unfinished objectives.
             ValueError: If the scenario has no atomic attacks configured. If your scenario
@@ -1716,57 +1728,74 @@ async def run_async(self) -> ScenarioResult:
         # Type narrowing: create local variable that type checker knows is non-None
         scenario_result_id: str = self._scenario_result_id
 
-        # Implement retry logic
-        last_exception = None
-        for retry_attempt in range(self._max_retries + 1):  # +1 for initial attempt
-            try:
-                return await self._execute_scenario_async()
-            except asyncio.CancelledError:
+        with Scenario._active_runs_lock:
+            if scenario_result_id in Scenario._active_runs:
+                raise ScenarioRunInProgressException(
+                    status_code=409,
+                    message=f"Scenario '{self._name}' execution is already in progress for run '{scenario_result_id}'.",
+                )
+            Scenario._active_runs.add(scenario_result_id)
+
+        try:
+            await self._prepare_run_async()
+
+            # Implement retry logic
+            last_exception = None
+            for retry_attempt in range(self._max_retries + 1):  # +1 for initial attempt
                 try:
-                    (
-                        await self._memory.update_scenario_run_state_async(
-                            scenario_result_id=scenario_result_id,
-                            scenario_run_state=ScenarioRunState.CANCELLED,
-                            error_message="Scenario run was cancelled",
-                            error_type="CancelledError",
+                    return await self._execute_scenario_async()
+                except asyncio.CancelledError:
+                    try:
+                        (
+                            await self._memory.update_scenario_run_state_async(
+                                scenario_result_id=scenario_result_id,
+                                scenario_run_state=ScenarioRunState.CANCELLED,
+                                error_message="Scenario run was cancelled",
+                                error_type="CancelledError",
+                            )
                         )
+                    except Exception:
+                        logger.exception(f"Failed to persist cancellation state for scenario '{self._name}'")
+                    raise
+                except Exception as e:
+                    last_exception = e
+
+                    # Get current scenario to check number of tries
+                    scenario_results = await self._memory.get_scenario_results_async(
+                        scenario_result_ids=[scenario_result_id]
                     )
-                except Exception:
-                    logger.exception(f"Failed to persist cancellation state for scenario '{self._name}'")
-                raise
-            except Exception as e:
-                last_exception = e
-
-                # Get current scenario to check number of tries
-                scenario_results = await self._memory.get_scenario_results_async(
-                    scenario_result_ids=[scenario_result_id]
-                )
-                current_tries = scenario_results[0].number_tries if scenario_results else retry_attempt + 1
+                    current_tries = scenario_results[0].number_tries if scenario_results else retry_attempt + 1
 
-                # Check if we have more retries available
-                remaining_retries = self._max_retries - retry_attempt
+                    # Check if we have more retries available
+                    remaining_retries = self._max_retries - retry_attempt
 
-                if remaining_retries > 0:
+                    if remaining_retries > 0:
+                        logger.error(
+                            f"Scenario '{self._name}' failed on attempt {current_tries} with error: {str(e)}. "
+                            f"Retrying... ({remaining_retries} retries remaining)",
+                            exc_info=True,
+                        )
+                        # Continue to next iteration for retry
+                        continue
+                    # No more retries, log final failure
                     logger.error(
-                        f"Scenario '{self._name}' failed on attempt {current_tries} with error: {str(e)}. "
-                        f"Retrying... ({remaining_retries} retries remaining)",
+                        f"Scenario '{self._name}' failed after {current_tries} attempts "
+                        f"(initial + {self._max_retries} retries) with error: {str(e)}. Giving up.",
                         exc_info=True,
                     )
-                    # Continue to next iteration for retry
-                    continue
-                # No more retries, log final failure
-                logger.error(
-                    f"Scenario '{self._name}' failed after {current_tries} attempts "
-                    f"(initial + {self._max_retries} retries) with error: {str(e)}. Giving up.",
-                    exc_info=True,
-                )
-                (await self._mark_scenario_failed_async(scenario_result_id=scenario_result_id, error=e))
-                raise
+                    (await self._mark_scenario_failed_async(scenario_result_id=scenario_result_id, error=e))
+                    raise
+
+            # This should never be reached, but just in case
+            if last_exception:
+                raise last_exception
+            raise RuntimeError(f"Scenario '{self._name}' completed unexpectedly without result")
+        finally:
+            with Scenario._active_runs_lock:
+                Scenario._active_runs.discard(scenario_result_id)
 
-        # This should never be reached, but just in case
-        if last_exception:
-            raise last_exception
-        raise RuntimeError(f"Scenario '{self._name}' completed unexpectedly without result")
+    async def _prepare_run_async(self) -> None:  # noqa: B027
+        """Run subclass pre-work while this process owns the run."""
 
     async def _execute_scenario_async(self) -> ScenarioResult:
         """
@@ -1792,15 +1821,9 @@ async def _execute_scenario_async(self) -> ScenarioResult:
             raise ValueError("self._scenario_result_id is not initialized")
         scenario_result_id: str = self._scenario_result_id
 
-        # Increment number_tries at the start of each run
-        scenario_results = await self._memory.get_scenario_results_async(scenario_result_ids=[scenario_result_id])
-        if scenario_results:
-            current_scenario = scenario_results[0]
-            current_scenario.number_tries += 1
-            await self._memory.update_scenario_result_async(scenario_result=current_scenario)
-            logger.info(f"Scenario '{self._name}' attempt #{current_scenario.number_tries}")
-        else:
-            raise ValueError(f"Scenario result with ID {scenario_result_id} not found")
+        # Increment number_tries atomically at the start of each run
+        current_tries = await self._memory.increment_scenario_number_tries_async(scenario_result_id=scenario_result_id)
+        logger.info(f"Scenario '{self._name}' attempt #{current_tries}")
 
         # Mark scenario as in progress
         (

pyrit/scenario/scenarios/benchmark/adversarial.py

@@ -23,7 +23,6 @@
     BoundedDatasetSize,
     ObjectiveTargetEvaluationIdentifier,
     ScenarioDatasetSizeEstimate,
-    ScenarioResult,
     ScenarioRunSizeComponent,
     ScenarioRunSizeEstimate,
     ScenarioRunSizeEstimateCondition,
@@ -685,26 +684,21 @@ def _resolve_adversarial_targets(self, *, target_names: list[str]) -> list[tuple
 
         return resolved
 
-    async def run_async(self) -> ScenarioResult:
+    async def _prepare_run_async(self) -> None:
         """
-        Persist compatible cached results into this run, then execute uncached objectives.
+        Persist compatible cached results into this run before executing uncached objectives.
 
         Cached results are copied with new result IDs and attributed to the new
         scenario result. This keeps database-backed status, API responses, and
         later artifact exports complete without moving results away from their
         original runs.
-
-        Returns:
-            ScenarioResult: The persisted scenario result containing cached and
-            newly executed objective results.
         """
         try:
-            (await self._persist_precomputed_cached_results_async())
+            await self._persist_precomputed_cached_results_async()
         except Exception as error:
             if self._scenario_result_id:
-                (await self._mark_scenario_failed_async(scenario_result_id=self._scenario_result_id, error=error))
+                await self._mark_scenario_failed_async(scenario_result_id=self._scenario_result_id, error=error)
             raise
-        return await super().run_async()
 
     async def _apply_reusable_cached_results_async(self, *, atomic_attacks: list[AtomicAttack]) -> None:
         """