FIX: Decode video frames and audio off the event loop in video scorers

pyrit/score/video_scorer.py

@@ -1,12 +1,14 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+import asyncio
 import logging
 import random
 import tempfile
 import uuid
 from pathlib import Path
 
+from pyrit.common.async_compatibility import run_legacy_sync_async
 from pyrit.memory import CentralMemory
 from pyrit.models import MessagePiece, MessageScorable, Score, ScoringExpectation
 from pyrit.score.audio_transcript_scorer import AudioTranscriptHelper
@@ -167,7 +169,7 @@ async def _score_frames_async(
             raise FileNotFoundError(f"Video file not found: {video_path}")
 
         # Extract frames from video
-        frames = self._extract_frames(video_path)
+        frames = await asyncio.to_thread(self._extract_frames, video_path)
         if not frames:
             raise ValueError("No frames extracted from video for scoring.")
 
@@ -258,6 +260,30 @@ def _extract_frames(self, video_path: str) -> list[str]:
 
         return frame_paths
 
+    async def _extract_video_audio_async(self, video_path: str) -> str | None:
+        """
+        Extract audio off-loop and clean up its result before propagating cancellation.
+
+        Returns:
+            str | None: Extracted audio path, or None if extraction fails.
+
+        Raises:
+            asyncio.CancelledError: After extraction and temporary file cleanup finish.
+        """
+        audio_path: str | None = None
+
+        def extract_audio() -> str | None:
+            nonlocal audio_path
+            audio_path = AudioTranscriptHelper.extract_audio_from_video(video_path)
+            return audio_path
+
+        try:
+            return await run_legacy_sync_async(extract_audio)
+        except asyncio.CancelledError:
+            if audio_path and audio_path != video_path:
+                await run_legacy_sync_async(Path(audio_path).unlink, missing_ok=True)
+            raise
+
     async def _score_video_audio_async(
         self, *, message_piece: MessagePiece, expectation: ScoringExpectation | None, audio_scorer: Scorer | None = None
     ) -> list[Score]:
@@ -283,7 +309,7 @@ async def _score_video_audio_async(
 
         # Use BaseAudioTranscriptScorer's static method to extract audio
 
-        audio_path = AudioTranscriptHelper.extract_audio_from_video(video_path)
+        audio_path = await self._extract_video_audio_async(video_path)
         if not audio_path:
             logger.warning("Video does not have any audio! Skipping audio scoring.")
             return []

tests/unit/score/test_video_scorer.py

@@ -1,13 +1,16 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+import asyncio
 import os
+import threading
 import uuid
 from pathlib import Path
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import numpy as np
 import pytest
+from unit.async_utils import TEST_WATCHDOG_SECONDS, settle_tasks_async, wait_for_completion_async, wait_for_event_async
 from unit.mocks import get_mock_scorer_identifier
 
 from pyrit.models import AnswerMatches, ComponentIdentifier, MessagePiece, Score, ScoreStatus, ScoringExpectation
@@ -547,6 +550,36 @@ async def test_video_true_false_scorer_with_audio_scorer(video_converter_sample_
         assert "visual" in scores[0].score_rationale.lower() or "audio" in scores[0].score_rationale.lower()
 
 
+@pytest.mark.skipif(not is_opencv_installed(), reason="opencv is not installed")
+async def test_video_scorer_decodes_frames_and_audio_off_the_event_loop(video_converter_sample_video):
+    event_loop_thread = threading.get_ident()
+    decode_threads: dict[str, int] = {}
+    scorer = VideoTrueFalseScorer(
+        image_capable_scorer=MockTrueFalseScorer(return_value=True),
+        audio_scorer=MockAudioTrueFalseScorer(return_value=True),
+        num_sampled_frames=3,
+    )
+    extract_frames = scorer._video_helper._extract_frames
+
+    def record_frames(video_path: str) -> list[str]:
+        decode_threads["frames"] = threading.get_ident()
+        return extract_frames(video_path)
+
+    def record_audio(video_path: str) -> str | None:
+        decode_threads["audio"] = threading.get_ident()
+        return None
+
+    with (
+        patch.object(scorer._video_helper, "_extract_frames", side_effect=record_frames),
+        patch.object(AudioTranscriptHelper, "extract_audio_from_video", side_effect=record_audio),
+    ):
+        scores = await scorer._score_piece_with_expectation_async(video_converter_sample_video, expectation=None)
+
+    assert scores[0].score_value == "true"
+    assert decode_threads["frames"] != event_loop_thread
+    assert decode_threads["audio"] != event_loop_thread
+
+
 @pytest.mark.skipif(not is_opencv_installed(), reason="opencv is not installed")
 async def test_video_audio_scorer_cleans_up_extracted_audio(tmp_path, video_converter_sample_video):
     """_score_video_audio_async unlinks the temp audio file after successful scoring."""
@@ -569,6 +602,61 @@ async def test_video_audio_scorer_cleans_up_extracted_audio(tmp_path, video_conv
     assert not extracted_audio.exists()
 
 
+@pytest.mark.parametrize("extraction_result", ["temporary", "source", "none"])
+@pytest.mark.parametrize("cancel_twice", [False, True])
+async def test_video_audio_scorer_cancellation_during_extraction_async(
+    tmp_path: Path, video_converter_sample_video: MessagePiece, extraction_result: str, cancel_twice: bool
+) -> None:
+    loop = asyncio.get_running_loop()
+    started = asyncio.Event()
+    release = threading.Event()
+    source_path = Path(video_converter_sample_video.converted_value)
+    source_path.touch(exist_ok=True)
+    extracted_audio = tmp_path / "cancelled_audio.wav"
+    audio_scorer = MockAudioTrueFalseScorer()
+    scorer = VideoTrueFalseScorer(image_capable_scorer=MockTrueFalseScorer(), audio_scorer=audio_scorer)
+
+    def extract_audio(video_path: str) -> str | None:
+        loop.call_soon_threadsafe(started.set)
+        if not release.wait(TEST_WATCHDOG_SECONDS):
+            raise TimeoutError("The test did not release audio extraction.")
+        if extraction_result == "source":
+            return video_path
+        if extraction_result == "none":
+            return None
+        extracted_audio.write_bytes(b"audio")
+        return str(extracted_audio)
+
+    with (
+        patch.object(AudioTranscriptHelper, "extract_audio_from_video", side_effect=extract_audio),
+        patch.object(audio_scorer, "_score_batch_nested_async", new_callable=AsyncMock) as score_audio,
+    ):
+        task = asyncio.create_task(
+            scorer._video_helper._score_video_audio_async(
+                message_piece=video_converter_sample_video, expectation=None, audio_scorer=audio_scorer
+            )
+        )
+        try:
+            await wait_for_event_async(event=started, tasks=[task])
+            task.cancel()
+            await asyncio.sleep(0)
+            assert not task.done()
+            if cancel_twice:
+                task.cancel()
+                await asyncio.sleep(0)
+                assert not task.done()
+            release.set()
+            with pytest.raises(asyncio.CancelledError):
+                await wait_for_completion_async(future=task)
+        finally:
+            release.set()
+            await settle_tasks_async(task)
+
+    assert not extracted_audio.exists()
+    assert source_path.exists()
+    score_audio.assert_not_awaited()
+
+
 @pytest.mark.skipif(not is_opencv_installed(), reason="opencv is not installed")
 async def test_video_audio_scorer_keeps_source_when_extraction_returns_it(video_converter_sample_video):
     """_score_video_audio_async must not delete the video when extraction hands back the input path.