Skip to content

FIX: Keep the backend responsive while starting scenario runs - #2522

Merged
varunj-msft merged 5 commits into
microsoft:mainfrom
varunj-msft:varunj-msft/v1.1.0-Release-Scenario-Start-Responsiveness
Sep 3, 2026
Merged

varunj-msft merged 5 commits into
microsoft:mainfrom
varunj-msft:varunj-msft/v1.1.0-Release-Scenario-Start-Responsiveness

Conversation

@varunj-msft

Copy link
Copy Markdown
Contributor

Description

POST /runs initializes everything eagerly, on purpose, so that configuration errors reach the caller instead of disappearing into a background task. The problem is where that work ran.

It ran directly on the event loop, and it is slow and almost entirely synchronous — loading the default datasets alone takes minutes. For that whole window the server answered nothing at all. Health probes timed out, and the CLI reported the server as unavailable even though it was alive and simply busy. That is the failure mode behind the current End to End Tests failures, where the client gives up before the server has any chance to reply.

The main fix moves the eager initialization onto a worker thread with asyncio.to_thread, following the pattern initializer_service already uses for the same reason. The semaphore, the active-task registry and the create_task hand-off all deliberately stay on the server loop: asyncio.run cancels whatever is still pending when it closes its loop, so a background task created inside the worker thread would be destroyed the instant initialization finished. That shape silently cancels every run, and it is specifically avoided here.

Two concurrency-permit leaks are fixed along the way. Both are on paths the previous except Exception: release; raise could not reach:

  • A client disconnecting during initialization raises CancelledError, which is a BaseException and so was never caught.
  • The missing scenario_result_id check sat outside the try block entirely.

Either one leaked a permit, and three such failures exhausted the concurrency limit and wedged the server for the rest of the session. For the E2E suite that matters, because one session-scoped backend serves every scenario in the run. The permit is now released from a finally block until ownership transfers to the background task, tracked with an explicit release_on_exit flag so it is released exactly once and never twice.

The response is also built before the task hand-off, so a lookup failure can no longer leave a run executing that the caller never received an id for. The active_tasks entry is unwound on that path too.

Finally, the start_scenario_run route docstring said "Returns immediately", which was not true before this change and is still not true after it. It now describes what actually happens.

Part of the v1.1.0 release wave with #2510, #2511 and #2512.

Tests and Documentation

Six new tests in tests/unit/backend/test_scenario_run_service.py:

  • test_start_run_keeps_event_loop_responsive counts heartbeats on the loop during a slow start. A blocked loop yields zero.
  • test_start_run_background_task_survives_handoff asserts the run actually executes. This is the test that catches the "silently cancels every run" shape.
  • test_start_run_releases_semaphore_when_cancelled_during_init covers CancelledError being a BaseException.
  • test_start_run_releases_semaphore_when_result_id_missing covers the check that used to sit outside the try.
  • test_start_run_cleans_up_when_response_lookup_fails asserts no stranded permit and no stranded active_tasks entry.
  • test_start_run_releases_semaphore_exactly_once_on_success guards against the obvious over-correction of double-releasing.

The first two matter as a pair rather than individually: a responsiveness fix that cancels every run would pass the responsiveness test on its own, so the hand-off test is what makes the first one meaningful.

test_start_run_exceeds_concurrent_limit needed a fix. It was passing for the wrong reason: it relied on the event loop never yielding during start, so the mocked runs completed and handed their permits straight back before the limit could ever be reached. Now that start yields, the test holds its background runs open, which is what a real run does.

Ran pytest tests/unit/backend/test_scenario_run_service.py: 70 passed.

Documentation: the start_scenario_run route docstring is corrected in this PR. JupyText was not run and is not applicable: no notebooks or code samples are affected, and the public API is unchanged.

@hannahwestra25 hannahwestra25 self-assigned this Sep 1, 2026
Comment thread pyrit/backend/services/scenario_run_service.py Outdated
Comment thread pyrit/backend/services/scenario_run_service.py Outdated
@varunj-msft
varunj-msft force-pushed the varunj-msft/v1.1.0-Release-Scenario-Start-Responsiveness branch from 1b69719 to 6eae54c Compare September 1, 2026 18:04
@richlundeen

Copy link
Copy Markdown
Contributor

(GHCP Generated): FYI only - this PR should not wait on the Scenario stack. Stacked PR #2376 introduces FIFO scheduling in ScenarioRunService and queues fully initialized runs, so it overlaps this service. The event-loop responsiveness fix here is still needed. If #2522 merges first, we will rebase the stack and preserve the worker-thread preparation while adapting the semaphore-specific cleanup that FIFO scheduling supersedes.

@varunj-msft
varunj-msft force-pushed the varunj-msft/v1.1.0-Release-Scenario-Start-Responsiveness branch from 6eae54c to 5e00ea8 Compare September 1, 2026 22:18
Comment thread pyrit/backend/services/scenario_run_service.py
Comment thread pyrit/backend/services/scenario_run_service.py
Comment thread pyrit/backend/services/scenario_run_service.py Outdated
@varunj-msft
varunj-msft force-pushed the varunj-msft/v1.1.0-Release-Scenario-Start-Responsiveness branch from 5e00ea8 to 3aebe70 Compare September 1, 2026 23:57
Comment thread pyrit/backend/services/scenario_run_service.py
@varunj-msft
varunj-msft force-pushed the varunj-msft/v1.1.0-Release-Scenario-Start-Responsiveness branch from 3aebe70 to fdca5fd Compare September 2, 2026 17:17
Comment thread pyrit/backend/services/scenario_run_service.py
Comment thread pyrit/backend/services/scenario_run_service.py Outdated
Scenario initialization loads the default datasets, which takes minutes of mostly
synchronous work. Running it on the event loop wedged the backend, so every later
scenario failed its health probe. Initialization now runs on a worker thread.

- Offload initialization to a single-worker executor so the loop stays free, and
  serialize preparations because the in-memory backend shares one DBAPI connection.
- Hold the concurrency permit until an abandoned preparation thread actually stops,
  including when it finishes as the cancellation lands, and terminalize the run it
  already stored instead of leaving it in CREATED.
- Drain initialization's own teardown tasks before closing the throwaway loop. If
  anything outlives the drain, mark the run failed and refuse the start rather than
  returning a scenario that holds dead async resources.
- Do not start a run that was cancelled while it was still initializing.
- Serialize in-memory SQLite sessions so a preparation thread and a status poll
  cannot interleave on the shared connection and lose writes.
@varunj-msft
varunj-msft force-pushed the varunj-msft/v1.1.0-Release-Scenario-Start-Responsiveness branch from fdca5fd to b38012c Compare September 2, 2026 19:52
Comment thread pyrit/backend/services/scenario_run_service.py
Comment thread pyrit/backend/services/scenario_run_service.py Outdated
Comment thread pyrit/backend/services/scenario_run_service.py Outdated
test_file_backed_database_is_not_serialized left the SQLite engine holding
an open handle on locking.db when the TemporaryDirectory context exited.
POSIX allows unlinking an open file, so this passed on Linux and macOS, but
Windows raised PermissionError (WinError 32) and failed the whole matrix.

The isolated_memory_factory fixture does dispose the engine, but that runs at
fixture teardown, after the with block has already tried to remove the
directory. Dispose inside the block instead; Engine.dispose() is idempotent,
so the fixture's later call is still safe.
Comment thread pyrit/backend/services/scenario_run_service.py Outdated
Three review findings on the scenario start path:

The drain only inspected the tasks that existed when it started. A task that
spawned another one before finishing left the child out of the wait set, so
the drain reported success and asyncio.run then cancelled the child while the
scenario was handed back as healthy. Rebuild the set after every wait and
share one deadline across the whole drain, so a chain of tasks cannot extend
the budget either.

Preparation failures wrote the run state unconditionally, so a cancellation
that landed while preparation was draining was overwritten by FAILED, and an
abandoned preparation could stamp CANCELLED over a state a real failure had
already recorded. Add try_update_scenario_run_state, which compares and writes
in one UPDATE, and use it at both sites. A read followed by a write cannot
close this: preparation runs on a worker thread while cancellation runs on the
loop thread, and the sqlite connection lock only covers in-memory databases.

Resuming a cancelled run is deliberate, but the run keeps its stored state
while it initializes, so the new cancelled-during-initialization check treated
it as a run the caller had given up on and refused to start it. Read the state
before preparation to tell the two apart, using the header so a run with many
attack results does not pay for a full hydration on the event loop.
The reST role hook rejects Sphinx cross-reference roles because PyRIT renders
docstrings with MyST, so :meth:`update_scenario_run_state` would have shown up
as raw literal text in the built docs. Use double backticks like the rest of
the file.

ty rejected the update mapping because Query.update takes
Dict[_DMLColumnArgument, Any] and dict is invariant in its key type, so
dict[str, Any] is not assignable even though string column names are what the
call actually passes. Widen the annotation rather than switching to ORM
attribute keys; the runtime behaviour is unchanged.

Also assert the expected states on the two cancellation paths. They checked
that the run was marked CANCELLED but not that the write was guarded, so a
wrong or missing guard there would not have failed a test.
Overriding the status on a resumed run only changed that one field, so the
response still carried the previous run's error, error type and completion
time alongside a CREATED status. Drop the override: the row is reported as it
is stored, which is what main does today, and the scenario moves it on when it
starts. Resuming a cancelled run still starts, which is the part that was
broken.

cancel_run_async read the state, waited up to five seconds for the task, then
wrote CANCELLED unconditionally. A run that failed or completed during that
wait had its outcome replaced, so a drain failure came back to the caller as a
user cancellation. Guard that write the same way as the others and let the
re-read report whichever state won.
@varunj-msft
varunj-msft added this pull request to the merge queue Sep 3, 2026
Merged via the queue into microsoft:main with commit ee50d98 Sep 3, 2026
54 checks passed
@varunj-msft
varunj-msft deleted the varunj-msft/v1.1.0-Release-Scenario-Start-Responsiveness branch September 3, 2026 02:46
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants