FIX: Keep the backend responsive while starting scenario runs

pyrit/backend/routes/scenarios.py

@@ -146,7 +146,9 @@ async def start_scenario_run(request: RunScenarioRequest) -> ScenarioRunSummary:
     """
     Start a new scenario run as a background task.
 
-    Returns immediately with a scenario_result_id that can be polled for status.
+    Initialization runs eagerly so configuration errors surface here, then the run
+    itself continues in the background. Returns a scenario_result_id that can be
+    polled for status.
 
     Args:
         request: Scenario run configuration.

pyrit/backend/services/scenario_run_service.py

@@ -11,10 +11,12 @@
 import asyncio
 import base64
 import contextlib
+import functools
 import json
 import logging
 import uuid
 from collections.abc import Sequence
+from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from typing import Any
@@ -162,6 +164,10 @@ class ScenarioRunService:
     Keeps an in-memory dict only for active asyncio tasks (cancellation support).
     """
 
+    #: Seconds to let initialization's own background tasks (for example HTTP client teardown
+    #: scheduled from ``__del__``) finish before the initialization loop is torn down.
+    _INITIALIZATION_DRAIN_TIMEOUT = 5.0
+
     def __init__(self, *, max_concurrent_runs: int = _DEFAULT_MAX_CONCURRENT_RUNS) -> None:
         """Initialize the scenario run service."""
         self._max_concurrent_runs = max_concurrent_runs
@@ -170,6 +176,13 @@ def __init__(self, *, max_concurrent_runs: int = _DEFAULT_MAX_CONCURRENT_RUNS) -
         self._run_semaphore = asyncio.Semaphore(max_concurrent_runs)
         self._configuration_resolver = ScenarioConfigurationResolver()
 
+        # Initialization writes to CentralMemory, and the in-memory SQLite backend shares one
+        # DBAPI connection across every thread (StaticPool, sqlite_memory.py). Two preparations
+        # running at once would use that connection concurrently and lose or corrupt writes, so
+        # they are serialized onto a single worker. The event loop is still free while they run,
+        # which is the point of the offload.
+        self._prepare_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="pyrit-scenario-prep")
+
     async def start_run_async(self, *, request: RunScenarioRequest) -> ScenarioRunSummary:
         """
         Start a new scenario run as a background task.
@@ -197,46 +210,198 @@ async def start_run_async(self, *, request: RunScenarioRequest) -> ScenarioRunSu
 
         await self._run_semaphore.acquire()
 
-        # Perform all initialization eagerly — errors propagate to caller
+        # This frame owns the permit until the background task is created; every exit path
+        # before that hand-off has to release it, including cancellation, which is a
+        # BaseException and so is not caught by ``except Exception``.
+        release_on_exit = True
+        registered_run_id: str | None = None
         try:
-            scenario_class = self._configuration_resolver.resolve_scenario_class(scenario_name=request.scenario_name)
-            await self._run_initializers_async(request=request)
-            objective_target = self._configuration_resolver.resolve_target(target_name=request.target_name)
-            init_kwargs = self._configuration_resolver.resolve_configuration(
-                scenario_name=request.scenario_name,
-                scenario_class=scenario_class,
-                objective_target=objective_target,
-                techniques=request.techniques,
-                dataset_names=request.dataset_names,
-                max_dataset_size=request.max_dataset_size,
-                dataset_filters=request.dataset_filters,
-                include_baseline=request.include_baseline,
-                max_concurrency=request.max_concurrency,
-                max_retries=request.max_retries,
-                memory_labels=request.labels,
+            # Initialization loads the default datasets, which takes minutes, and is mostly
+            # synchronous work. Run it on a worker thread so the event loop stays free to
+            # answer health checks and status polls while a run is starting.
+            prepare_task = asyncio.get_running_loop().run_in_executor(
+                self._prepare_executor, functools.partial(self._prepare_run_blocking, request=request)
             )
-            scenario = await self._initialize_scenario_async(request=request, init_kwargs=init_kwargs)
-        except Exception:
-            self._run_semaphore.release()
-            raise
+            try:
+                scenario = await asyncio.shield(prepare_task)
+            except BaseException as exc:
+                # A worker thread cannot be killed, so it keeps initializing after this frame
+                # unwinds. Keep holding the permit until it actually finishes, otherwise the
+                # next caller is admitted while this run is still loading datasets and
+                # ``max_concurrent_runs`` stops bounding the work that is really running.
+                if not prepare_task.done():
+                    prepare_task.add_done_callback(self._release_abandoned_prepare)
+                    release_on_exit = False
+                elif isinstance(exc, asyncio.CancelledError):
+                    # The thread can finish just as the cancellation lands. A done future never
+                    # calls back, so cleaning up here is the only chance to release the permit
+                    # and terminalize the run that initialization already stored.
+                    release_on_exit = False
+                    try:
+                        self._release_abandoned_prepare(prepare_task)
+                    except Exception as cleanup_error:
+                        # The permit is released first, so it is already back even if the rest
+                        # failed. Never let cleanup replace the cancellation being propagated.
+                        logger.warning(f"Could not clean up after a cancelled scenario preparation: {cleanup_error}")
+                raise
+
+            # scenario_result_id is set during initialize_async
+            scenario_result_id = scenario._scenario_result_id
+            if scenario_result_id is None:
+                raise ValueError("Scenario did not produce a scenario_result_id during initialization.")
+
+            # Track active task
+            active = _ActiveTask(scenario_result_id=scenario_result_id, scenario=scenario)
+            self._active_tasks[scenario_result_id] = active
+            registered_run_id = scenario_result_id
+
+            # Build the response before spawning the task so that a failure here cannot leave
+            # a run executing that the caller never received an id for.
+            response = self.get_run(scenario_result_id=scenario_result_id)
+            if response is None:
+                raise RuntimeError(
+                    f"Scenario run {scenario_result_id} was not found in the database after initialization."
+                )
 
-        # scenario_result_id is set during initialize_async
-        scenario_result_id = scenario._scenario_result_id
-        if scenario_result_id is None:
-            raise ValueError("Scenario did not produce a scenario_result_id during initialization.")
+            # Spawn background task (only runs scenario.run_async). It releases the permit in
+            # its own finally, so ownership transfers here and this frame must not release it.
+            task = asyncio.create_task(self._execute_run_async(scenario_result_id=scenario_result_id))
+            active.task = task
+            release_on_exit = False
+            registered_run_id = None
+        finally:
+            if registered_run_id is not None:
+                self._active_tasks.pop(registered_run_id, None)
+            if release_on_exit:
+                self._run_semaphore.release()
 
-        # Track active task
-        active = _ActiveTask(scenario_result_id=scenario_result_id, scenario=scenario)
-        self._active_tasks[scenario_result_id] = active
+        return response
 
-        # Spawn background task (only runs scenario.run_async)
-        task = asyncio.create_task(self._execute_run_async(scenario_result_id=scenario_result_id))
-        active.task = task
+    def _release_abandoned_prepare(self, prepare_task: "asyncio.Future[Scenario]") -> None:
+        """
+        Clean up after an abandoned preparation thread has finished.
 
-        response = self.get_run(scenario_result_id=scenario_result_id)
-        if response is None:
-            raise RuntimeError(f"Scenario run {scenario_result_id} was not found in the database after initialization.")
-        return response
+        ``start_run_async`` hands ownership of the permit to this callback when it is
+        cancelled while the worker thread is still initializing, so the permit is only
+        released after the thread has genuinely stopped using the slot. A preparation that
+        succeeds anyway leaves behind a scenario result nobody will run, which is marked
+        cancelled here rather than left waiting in ``CREATED``.
+
+        Args:
+            prepare_task: The future wrapping the abandoned ``_prepare_run_blocking`` call.
+        """
+        self._run_semaphore.release()
+
+        if prepare_task.cancelled():
+            return
+        error = prepare_task.exception()
+        if error is not None:
+            logger.warning(f"Abandoned scenario preparation failed after the request was cancelled: {error}")
+            return
+
+        # Initialization already stored a CREATED scenario result, and nothing is going to run
+        # it now, so terminalize it rather than leaving a run that never starts.
+        scenario_result_id = prepare_task.result()._scenario_result_id
+        if scenario_result_id:
+            try:
+                self._memory.update_scenario_run_state(
+                    scenario_result_id=scenario_result_id,
+                    scenario_run_state=ScenarioRunState.CANCELLED,
+                    error_message="The start request was cancelled while the scenario was being initialized.",
+                )
+            except Exception as update_error:
+                logger.warning(
+                    f"Could not mark abandoned scenario run {scenario_result_id} as cancelled: {update_error}"
+                )
+        logger.warning("Abandoned scenario preparation completed after the request was cancelled.")
+
+    def _prepare_run_blocking(self, *, request: RunScenarioRequest) -> Scenario:
+        """
+        Run the eager initialization for a scenario run on the calling thread.
+
+        Exists so ``start_run_async`` can offload initialization onto a worker thread.
+        The scenario is executed later on the caller's event loop, so initialization must not
+        leave anything bound to the throwaway loop used here. Clients that schedule their own
+        teardown are given a moment to finish; anything still running after that would be
+        cancelled when the loop closes, so the start fails rather than handing back a scenario
+        that holds dead async resources.
+
+        Args:
+            request: The run request with scenario name, target, and options.
+
+        Returns:
+            Scenario: The initialized scenario.
+
+        Raises:
+            RuntimeError: If tasks are still running on the initialization loop after the drain.
+        """
+
+        async def prepare_async() -> Scenario:
+            scenario = await self._prepare_run_async(request=request)
+            await self._drain_initialization_tasks_async()
+            return scenario
+
+        return asyncio.run(prepare_async())
+
+    async def _drain_initialization_tasks_async(self) -> None:
+        """
+        Let initialization's background tasks finish before the initialization loop closes.
+
+        Initialization builds throwaway async clients, and some of them schedule their own
+        teardown from ``__del__``, so a task can appear purely because a garbage collection
+        landed late. Waiting for those is the difference between a scenario that starts and
+        one that fails at random.
+
+        Raises:
+            RuntimeError: If any task is still running after the drain timeout.
+        """
+        pending = [task for task in asyncio.all_tasks() if task is not asyncio.current_task()]
+        if not pending:
+            return
+
+        done, still_running = await asyncio.wait(pending, timeout=self._INITIALIZATION_DRAIN_TIMEOUT)
+        for task in done:
+            # Retrieve outcomes so a failed teardown task does not log "never retrieved" noise.
+            if not task.cancelled() and task.exception() is not None:
+                logger.debug(f"A scenario initialization task failed during teardown: {task.exception()}")
+
+        if still_running:
+            raise RuntimeError(
+                "Scenario initialization left background tasks on the initialization loop, which is "
+                "about to close. They would be cancelled and the scenario would hold dead async "
+                f"resources: {', '.join(sorted(task.get_name() for task in still_running))}"
+            )
+
+    async def _prepare_run_async(self, *, request: RunScenarioRequest) -> Scenario:
+        """
+        Resolve and initialize the scenario for a run request.
+
+        Args:
+            request: The run request with scenario name, target, and options.
+
+        Returns:
+            Scenario: The initialized scenario.
+
+        Raises:
+            ValueError: If scenario, target, initializer, or technique cannot be found.
+        """
+        scenario_class = self._configuration_resolver.resolve_scenario_class(scenario_name=request.scenario_name)
+        await self._run_initializers_async(request=request)
+        objective_target = self._configuration_resolver.resolve_target(target_name=request.target_name)
+        init_kwargs = self._configuration_resolver.resolve_configuration(
+            scenario_name=request.scenario_name,
+            scenario_class=scenario_class,
+            objective_target=objective_target,
+            techniques=request.techniques,
+            dataset_names=request.dataset_names,
+            max_dataset_size=request.max_dataset_size,
+            dataset_filters=request.dataset_filters,
+            include_baseline=request.include_baseline,
+            max_concurrency=request.max_concurrency,
+            max_retries=request.max_retries,
+            memory_labels=request.labels,
+        )
+        return await self._initialize_scenario_async(request=request, init_kwargs=init_kwargs)
 
     def get_run(self, *, scenario_result_id: str) -> ScenarioRunSummary | None:
         """

pyrit/memory/sqlite_memory.py

@@ -2,6 +2,8 @@
 # Licensed under the MIT license.
 
 import logging
+import threading
+import weakref
 from collections.abc import Sequence
 from contextlib import closing
 from datetime import datetime
@@ -72,6 +74,11 @@ def __init__(
             self.db_path = Path(db_path or Path(DB_DATA_PATH, self.DEFAULT_DB_FILE_NAME)).resolve()
         self.results_path = str(DB_DATA_PATH)
 
+        # An in-memory database shares a single DBAPI connection across every thread (see
+        # ``_create_engine``), so concurrent sessions would interleave on it. Serialize session
+        # lifetimes for that backend only; file-backed databases get a connection per checkout.
+        self._connection_lock: threading.RLock | None = threading.RLock() if self.db_path == ":memory:" else None
+
         self.engine = self._create_engine(has_echo=verbose)
         self.SessionFactory = sessionmaker(bind=self.engine)
         if not skip_schema_migration:
@@ -285,10 +292,43 @@ def get_session(self) -> Session:
         """
         Provide a SQLAlchemy session for transactional operations.
 
+        For an in-memory database every session borrows the same DBAPI connection, so the
+        session is handed out under a lock that is only released when it is closed. That keeps
+        a whole transaction, not just a single statement, isolated from the other threads.
+
         Returns:
             Session: A SQLAlchemy session bound to the engine.
         """
-        return self.SessionFactory()
+        session = self.SessionFactory()
+        connection_lock = self._connection_lock
+        if connection_lock is None:
+            return session
+
+        connection_lock.acquire()
+        close_session = session.close
+        released = False
+
+        def release_once() -> None:
+            # Also runs if the session is discarded without being closed, so one caller that
+            # forgets cannot leave the lock held and stall every other thread forever.
+            nonlocal released
+            if released:
+                return
+            released = True
+            try:
+                connection_lock.release()
+            except RuntimeError:
+                logger.warning("An in-memory session was discarded by a thread that did not open it.")
+
+        def close_and_release() -> None:
+            try:
+                close_session()
+            finally:
+                release_once()
+
+        session.close = close_and_release  # type: ignore[ty:invalid-assignment]
+        weakref.finalize(session, release_once)
+        return session
 
     def print_schema(self) -> None:
         """