[Vardhman Gupta (Kaap10)]

Description

Fixes #3748

When SDK callers execute concurrent tasks against the same Scenario instance or multiple Scenario instances sharing the same persisted scenario_result_id (via scenario_result_id resume), concurrent run_async() executions bypass the planned concurrency limit, duplicating target calls and attack results while racing on number_tries.

Changes:

• Scenario (pyrit/scenario/core/scenario.py):
  • Added class-level active run guard _active_runs: ClassVar[set[str]] protected by _active_runs_lock: ClassVar[threading.Lock].
  • In run_async(), acquires ownership of _scenario_result_id before starting the retry loop or scheduling target work. If already active in another task or thread, raises RuntimeError immediately with a clear diagnostic message.
  • Releases guard in the outermost finally: block of run_async(), guaranteeing cleanup on normal completion, unhandled exceptions, and cancellation.
  • Prior to releasing the guard upon cancellation, target worker tasks in _execute_atomic_attacks_parallel_async() are drained to ensure background calls have completed.
  • Replaced the vulnerable read-modify-write pattern on number_tries with atomic accounting via increment_scenario_number_tries_async().
• MemoryInterface (pyrit/memory/memory_interface.py):
  • Added increment_scenario_number_tries_async(scenario_result_id: str) to MemoryInterface and database implementations using @legacy_sync_override.
  • Executes atomic SQL update: UPDATE ScenarioResult SET number_tries = number_tries + 1 WHERE id = :id.
  • Added deprecated synchronous wrapper increment_scenario_number_tries(scenario_result_id: str) per PyRIT memory interface conventions.

cc Roman Lutz (@romanlutz)

Tests and Documentation

• Exact Issue Concurrent SDK run_async calls duplicate one persisted run and exceed its concurrency limit #3748 Reproduction Test: Added test_concurrent_run_async_rejects_second_execution_before_target_work in tests/unit/scenario/core/test_scenario.py. Verifies that two concurrent tasks with 6 objectives and max_concurrency=2 result in 1 task executing cleanly, the second task rejected with RuntimeError before target work, peak concurrency strictly <= 2, exactly 6 target calls, exactly 6 attack results, and number_tries == 1.
• Cross-Instance Same-Run Rejection: Added test_concurrent_run_async_on_different_instances_same_run_id_rejects_second to ensure multiple Scenario objects sharing the same run ID cannot run concurrently.
• Cancellation & Legitimate Resume: Added test_cancellation_cleans_up_guard_and_allows_legitimate_resume testing with real SQLite to verify cancellation safely releases the guard and enables subsequent legitimate resumption to complete the remaining objectives and reach number_tries = 2.
• Concurrency Isolation: Added test_concurrent_run_async_distinct_scenarios_run_in_parallel ensuring distinct scenarios run concurrently without mutual blocking.
• Atomic Accounting Tests: Added test_increment_scenario_number_tries_async, test_increment_scenario_number_tries_sync, test_increment_scenario_number_tries_missing_id_raises, and test_increment_scenario_number_tries_concurrent in tests/unit/memory/memory_interface/test_interface_scenario_results.py.
• Test Suite Verification:
  • uv run -m pytest tests/unit/scenario/core/test_scenario.py tests/unit/memory/memory_interface/test_interface_scenario_results.py: 167 passed in 25.97s.
  • Diff coverage on modified code: 98% (exceeds 90% requirement).
• Code Quality:
  • uv run -m ruff format: clean.
  • uv run -m ruff check: 0 errors.
• JupyText / Documentation: Documentation was not modified as internal execution concurrency guards and atomic attempt accounting do not change the public user-facing API surface.

[Vardhman Gupta (Kaap10)]

Thanks for the great suggestion Ismaeel Habeeb (@ismaeelakin)! I've updated test_concurrent_run_async_distinct_scenarios_run_in_parallel in commit 17066d466 with an event-based rendezvous barrier (entered1 and entered2 asyncio.Events) and a shared watchdog with task cancellation cleanup. This guarantees that both scenarios must enter execution concurrently before either can proceed, strictly verifying that unrelated scenario runs are never serialized.

[hannahwestra25]

Not atomic: this SELECTs, then writes the Python value back (SET number_tries=<n>), so concurrent callers lose updates. On file-backed SQLite, 10 concurrent calls all returned 1 and stored 1. A single UPDATE ... + 1 (the repo's Query.update pattern) fixes it; verified: returns 1–10, stores 10.

Suggested change

	with closing(self._get_session()) as session:
	entry = session.query(ScenarioResultEntry).filter_by(id=scenario_result_id).first()
	if not entry:
	raise ValueError(f"Scenario result with ID {scenario_result_id} not found in memory")
	entry.number_tries += 1
	new_tries = entry.number_tries
	session.commit()
	return new_tries
	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

[Vardhman Gupta (Kaap10)]

Updated to an atomic Query.update statement so concurrent increments are safely handled at the database level.

[hannahwestra25]

Passes for the wrong reason: sqlite_instance is :memory:, whose _transaction_lock serializes whole sessions, so these increments never interleave. On a file-backed DB this exact test fails today (stored 1, expected 10). Suggest:

async def test_increment_scenario_number_tries_concurrent(tmp_path: Path):
    memory = SQLiteMemory(db_path=tmp_path / "tries.db", silent=True)
    memory.disable_embedding()
    try:
        scenario = make_scenario_result(...)  # unchanged
        await memory.add_scenario_results_to_memory_async(scenario_results=[scenario])
        scenario_id = str(scenario.id)
        results = await asyncio.gather(
            *[memory.increment_scenario_number_tries_async(scenario_result_id=scenario_id) for _ in range(10)]
        )
        assert sorted(results) == list(range(1, 11))
    finally:
        await memory.dispose_engine_async()

[Vardhman Gupta (Kaap10)]

Switched the concurrent test to a file-backed SQLite database via tmp_path to verify true multi-session atomicity.

[hannahwestra25]

The guard sits in an overridable method. AdversarialBenchmark.run_async persists cached-result copies (fresh UUIDs) before super().run_async(), so a rejected concurrent caller has already duplicated rows, and on insert failure marks the shared run FAILED. Move pre-work into a hook inside the guard:

# Scenario.run_async, first statement inside the guarded try:
await self._prepare_run_async()

async def _prepare_run_async(self) -> None:
    """Run subclass pre-work while this process owns the run."""

# AdversarialBenchmark: replaces its run_async override
async def _prepare_run_async(self) -> None:
    try:
        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)
        raise

[Vardhman Gupta (Kaap10)]

Moved the subclass pre-work into an overridable _prepare_run_async hook called inside the active run guard.

[hannahwestra25]

Bare RuntimeError collides with the docstring (line 1697: "fails for any other reason"), so callers may mark the shared, still-running run FAILED. Mirror ScenarioPartialFailureException: a dedicated type that keeps except RuntimeError working; add it to Raises: and assert the type in tests.

# pyrit/exceptions/exception_classes.py (export from pyrit.exceptions)
class ScenarioRunInProgressException(PyritException, RuntimeError):  # noqa: N818
    """Raised when this process is already executing the same scenario result."""


# scenario.py
raise ScenarioRunInProgressException(
    status_code=409,
    message=f"Scenario '{self._name}' execution is already in progress for run '{scenario_result_id}'.",
)

[Vardhman Gupta (Kaap10)]

Added ScenarioRunInProgressException inheriting from PyritException and RuntimeError with status code 409, and updated the docstrings and tests.

[hannahwestra25]

Per test.instructions.md:47-48, hang-only deadlines and rendezvous waits should use TEST_WATCHDOG_SECONDS; 5s/10s can flake on busy -n 4 workers. Same at L2881. L2789's await two_completed.wait() is unbounded, so an early exit hangs CI.

from unit.async_utils import TEST_WATCHDOG_SECONDS

await asyncio.wait_for(entered_other.wait(), timeout=TEST_WATCHDOG_SECONDS)  # L2854
res1, res2 = await asyncio.wait_for(asyncio.gather(task1, task2), timeout=TEST_WATCHDOG_SECONDS)  # L2881
await asyncio.wait_for(two_completed.wait(), timeout=TEST_WATCHDOG_SECONDS)  # L2789

[Vardhman Gupta (Kaap10)]

Wrapped all rendezvous and completion waits with TEST_WATCHDOG_SECONDS per test instructions.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.