Skip to content

FIX Drain nested scoring and batch tasks on failure - #2856

Open
biefan (biefan) wants to merge 4 commits into
microsoft:mainfrom
biefan:fix/concurrent-scoring-lifetimes
Open

biefan (biefan) wants to merge 4 commits into
microsoft:mainfrom
biefan:fix/concurrent-scoring-lifetimes

Conversation

@biefan

@biefan biefan (biefan) commented Sep 25, 2026 •

Copy link
Copy Markdown
Contributor

Description

A failed scoring call can return while nested scorers or other items in its batch still run, allowing old work to overlap retries or write scores after the caller has reported failure.

The response-level protection from #2695 does not cover every concurrency boundary. This change applies a shared cleanup helper to:

Entry point Work whose lifetime must end before failure returns
Scorer.score_with_scorers_async() Independently persisted scoring roots
TrueFalseCompositeScorer Nested child scorers
Multi-piece MessageScorer Piece-level scoring tasks
score_batch_async() / batch_task_async() Other items in the current batch

The helper shields the initial gather and explicitly owns cancellation. It cancels unfinished tasks only when they are not already processing a cancellation request, then shields and awaits the drain. This matters when one child exits immediately and another is still awaiting asynchronous cleanup: caller cancellation must not reach that slow child twice. Further caller cancellation is propagated after cleanup, and cleanup errors do not replace the primary failure. The original gather's exception is retrieved after draining so a detached shield does not leave an unobserved future error.

Successful result ordering, completed score persistence, batch sizes and rate-limit validation are preserved. Later batches do not start after a failure. Completed requests or scores are not rolled back.

Tests and Documentation

  • SQLite-backed public-API regressions cover four scoring entry points with both regular exceptions and child cancellation. They assert that started work finishes cleanup before return and that no late scoring completes or persists.
  • The persistence assertion uses get_scores(score_type="true_false"), which actually queries storage. A separate real-score check confirmed that an unfiltered lookup returns [] while the filtered lookup returns the saved row.
  • Event-gated batch tests hold a slow child's cleanup open while a fast child exits. They cover caller-initiated and independently initiated child cancellation, each with and without another caller cancellation during cleanup. The original caller-cancel cases reproduced two cancellation requests; all four cases now verify one request, completed cleanup, no pending children, and no later batch.
  • Existing tests cover preservation of the primary error when sibling cleanup fails and caller cancellation during error cleanup.
  • Focused scorer/batch/prompt-sending tests: 284 passed.
  • make unit-test (Python 3.11, default dependencies): 20,425 passed, 146 skipped, 1 failed. The sole failure is the pre-existing test_get_seed_dataset_summaries_follows_a_trailing_blank_insensitive_collation in tests/unit/memory/memory_interface/test_interface_seed_prompts.py, previously reproduced on unmodified upstream in the same environment.
  • All applicable pre-commit hooks passed with all optional dependencies installed, including repository-wide ty check pyrit.
uv run pytest tests/unit/prompt_target/test_batch_helper.py tests/unit/score/test_concurrent_scoring_lifetime.py tests/unit/score/test_scorer.py -q

The common helper and root/batch docstrings document the cleanup contract. No notebooks changed and no live model endpoints were used.

Comment thread pyrit/common/task_utils.py Outdated
Comment thread tests/unit/score/test_concurrent_scoring_lifetime.py Outdated
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

# Conflicts:
#	pyrit/score/message_scorer.py
#	pyrit/score/true_false/true_false_composite_scorer.py

This branch has not been deployed

No deployments
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.

2 participants