Skip to content

MAINT: Reducing scenario db reads #3500

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jbolor21 merged 1 commit into from
Oct 23, 2026
+225 −71
Merged

MAINT: Reducing scenario db reads #3500

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 15 additions & 25 deletions pyrit/backend/services/scenario_run_service.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -450,7 +450,7 @@ async def _schedule_prepared_run_async(
self, *, request: RunScenarioRequest, prepared: _PreparedRun, resumed_from_cancelled: bool
) -> ScenarioRunSummary:
"""
Validate persisted preparation and admit it to the FIFO scheduler.
Validate the persisted header before admission; build details only for the returned response.
Returns:
ScenarioRunSummary: The persisted active, queued, or terminal state.
Expand All @@ -461,10 +461,10 @@ async def _schedule_prepared_run_async(
raise ValueError("Scenario did not produce a scenario_result_id during initialization.")
if request.scenario_result_id and scenario_result_id != request.scenario_result_id:
raise ValueError("Scenario initialization changed the saved result ID; resume was not started.")
persisted = await self._memory.get_scenario_results_async(scenario_result_ids=[scenario_result_id])
if not persisted:
persisted = await self._memory.get_scenario_result_header_async(scenario_result_id=scenario_result_id)
if persisted is None:
raise RuntimeError(f"Scenario run {scenario_result_id} was not persisted during initialization.")
if not resumed_from_cancelled and persisted[0].scenario_run_state == ScenarioRunState.CANCELLED:
if not resumed_from_cancelled and persisted.scenario_run_state == ScenarioRunState.CANCELLED:
response = await self._build_response_async(
scenario_result_id=scenario_result_id,
active_error=None,
Expand All @@ -476,23 +476,13 @@ async def _schedule_prepared_run_async(
f"Scenario run {scenario_result_id} was not found in the database after initialization."
)
return response
if (
await self._build_response_async(
scenario_result_id=scenario_result_id,
active_error=None,
queue_position=None,
active_scenario_result_id=self._active_scenario_result_id,
)
is None
):
raise RuntimeError(f"Scenario run {scenario_result_id} was not found in the database after initialization.")
scheduled = _ActiveTask(
scenario_result_id=scenario_result_id,
scenario=scenario,
adversarial_target=prepared.adversarial_target,
scenario_name=persisted[0].scenario_name,
scenario_name=persisted.scenario_name,
scenario_registry_name=request.scenario_name,
created_at=persisted[0].creation_time,
created_at=persisted.creation_time,
enqueued_at=datetime.now(UTC),
)
await self._enqueue_run_async(scheduled=scheduled)
Expand Down Expand Up @@ -839,11 +829,11 @@ async def cancel_run_async(self, *, scenario_result_id: str) -> ScenarioRunSumma
Raises:
ValueError: If the run is already in a terminal state or not active.
"""
results = await self._memory.get_scenario_results_async(scenario_result_ids=[scenario_result_id])
if not results:
header = await self._memory.get_scenario_result_header_async(scenario_result_id=scenario_result_id)
if header is None:
return None

db_status = results[0].scenario_run_state
db_status = header.scenario_run_state
if self._is_terminal_state(db_status):
raise ValueError(f"Cannot cancel run in '{db_status}' state.")

Expand All @@ -870,9 +860,9 @@ async def cancel_run_async(self, *, scenario_result_id: str) -> ScenarioRunSumma
active.cancellation_error_type = "CancelledError"
task = active.task
else:
latest = await self._memory.get_scenario_results_async(scenario_result_ids=[scenario_result_id])
if latest and self._is_terminal_state(latest[0].scenario_run_state):
raise ValueError(f"Cannot cancel run in '{latest[0].scenario_run_state}' state.")
latest = await self._memory.get_scenario_result_header_async(scenario_result_id=scenario_result_id)
if latest is not None and self._is_terminal_state(latest.scenario_run_state):
raise ValueError(f"Cannot cancel run in '{latest.scenario_run_state}' state.")
await self._memory.try_update_scenario_run_state_async(
scenario_result_id=scenario_result_id,
expected_states={
Expand Down Expand Up @@ -1093,10 +1083,10 @@ async def _handoff_scheduler_async(self, *, scenario_result_id: str) -> None:
return
while self._queued_runs:
next_run = self._queued_runs[0]
persisted = await self._memory.get_scenario_results_async(
scenario_result_ids=[next_run.scenario_result_id]
persisted = await self._memory.get_scenario_result_header_async(
scenario_result_id=next_run.scenario_result_id
)
if not persisted or persisted[0].scenario_run_state != ScenarioRunState.QUEUED:
if persisted is None or persisted.scenario_run_state != ScenarioRunState.QUEUED:
self._queued_runs.popleft()
self._queue_revision += 1
continue
Expand Down
Loading
Loading