From f9ecc4347df70444cb74ecdabb14959cabc9aa56 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Alex=20Castillo=20Gonz=C3=A1lez?= Date: Sat, 31 Oct 2026 13:35:31 -0400 Subject: [PATCH 1/2] fix: clean up extracted video frames when scoring is cancelled MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Alex Castillo González --- pyrit/score/video_scorer.py | 29 ++++++++++++++++- tests/unit/score/test_video_scorer.py | 47 +++++++++++++++++++++++++++ 2 files changed, 75 insertions(+), 1 deletion(-) diff --git a/pyrit/score/video_scorer.py b/pyrit/score/video_scorer.py index 7cac6a1f06..c1b6c86d35 100644 --- a/pyrit/score/video_scorer.py +++ b/pyrit/score/video_scorer.py @@ -169,7 +169,7 @@ async def _score_frames_async( raise FileNotFoundError(f"Video file not found: {video_path}") # Extract frames from video - frames = await asyncio.to_thread(self._extract_frames, video_path) + frames = await self._extract_video_frames_async(video_path) if not frames: raise ValueError("No frames extracted from video for scoring.") @@ -210,6 +210,33 @@ async def _score_frames_async( return frame_scores, len(frames) + async def _extract_video_frames_async(self, video_path: str) -> list[str]: + """ + Extract frames off-loop and clean up their files before propagating cancellation. + + Returns: + list[str]: Temporary file paths for the extracted frames. + + Raises: + asyncio.CancelledError: After extraction and temporary file cleanup finish. + """ + frame_paths: list[str] = [] + + def extract_frames() -> list[str]: + nonlocal frame_paths + frame_paths = self._extract_frames(video_path) + return frame_paths + + def remove_frames() -> None: + for frame_path in frame_paths: + Path(frame_path).unlink(missing_ok=True) + + try: + return await run_legacy_sync_async(extract_frames) + except asyncio.CancelledError: + await run_legacy_sync_async(remove_frames) + raise + def _extract_frames(self, video_path: str) -> list[str]: """ Extract a specified number of image frames from a video file. diff --git a/tests/unit/score/test_video_scorer.py b/tests/unit/score/test_video_scorer.py index 99148e5fdf..f440cf74bf 100644 --- a/tests/unit/score/test_video_scorer.py +++ b/tests/unit/score/test_video_scorer.py @@ -654,6 +654,53 @@ def extract_audio(video_path: str) -> str | None: score_audio.assert_not_awaited() +@pytest.mark.parametrize("cancel_twice", [False, True]) +async def test_video_frame_scorer_cancellation_during_extraction_cleans_up_frames_async( + tmp_path: Path, video_converter_sample_video: MessagePiece, cancel_twice: bool +) -> None: + loop = asyncio.get_running_loop() + started = asyncio.Event() + release = threading.Event() + Path(video_converter_sample_video.converted_value).touch(exist_ok=True) + frame_paths = [tmp_path / "cancelled_frame_0.png", tmp_path / "cancelled_frame_1.png"] + image_scorer = MockTrueFalseScorer() + scorer = VideoTrueFalseScorer(image_capable_scorer=image_scorer) + + def extract_frames(video_path: str) -> list[str]: + loop.call_soon_threadsafe(started.set) + if not release.wait(TEST_WATCHDOG_SECONDS): + raise TimeoutError("The test did not release frame extraction.") + for frame_path in frame_paths: + frame_path.write_bytes(b"frame") + return [str(frame_path) for frame_path in frame_paths] + + with ( + patch.object(scorer._video_helper, "_extract_frames", side_effect=extract_frames), + patch.object(image_scorer, "_score_batch_nested_async", new_callable=AsyncMock) as score_frames, + ): + task = asyncio.create_task( + scorer._video_helper._score_frames_async(message_piece=video_converter_sample_video, expectation=None) + ) + 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 any(frame_path.exists() for frame_path in frame_paths) + score_frames.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. From 2a914edf9a61d0990583ee2b5564b549e2875e49 Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Sun, 1 Nov 2026 23:13:43 -0800 Subject: [PATCH 2/2] FIX: clean up partially extracted video frames Keep ownership of every frame file until extraction succeeds. Remove files on extraction errors, avoid creating files for failed encodes, and reject failed frame writes. Cover partial failures and repeated cancellation with isolated real-video tests. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- pyrit/score/video_scorer.py | 43 +++---- tests/unit/score/test_video_scorer.py | 157 ++++++++++++++++++++++++++ 2 files changed, 181 insertions(+), 19 deletions(-) diff --git a/pyrit/score/video_scorer.py b/pyrit/score/video_scorer.py index c1b6c86d35..809604353a 100644 --- a/pyrit/score/video_scorer.py +++ b/pyrit/score/video_scorer.py @@ -6,6 +6,7 @@ import random import tempfile import uuid +from contextlib import ExitStack from pathlib import Path from pyrit.common.async_compatibility import run_legacy_sync_async @@ -249,6 +250,7 @@ def _extract_frames(self, video_path: str) -> list[str]: Raises: ModuleNotFoundError: If OpenCV is not installed. + OSError: If an extracted frame cannot be written. Files created by this extraction are removed on failure. """ try: import cv2 @@ -259,31 +261,34 @@ def _extract_frames(self, video_path: str) -> list[str]: frame_paths = [] video_capture = cv2.VideoCapture(video_path) - try: - total_frames = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT)) - if total_frames > 0: - # Choose up to num_sampled_frames random unique frame indices - frame_indices = sorted(random.sample(range(total_frames), min(self.num_sampled_frames, total_frames))) - for frame_index in frame_indices: - # Set the video position to the selected frame - video_capture.set(cv2.CAP_PROP_POS_FRAMES, frame_index) - ret, frame = video_capture.read() - if not ret: - logger.warning(f"Failed to read frame at index {frame_index}") - continue - - # Create a temporary file for the frame - with tempfile.NamedTemporaryFile(suffix=f"_frame_{frame_index}.png", delete=False) as temp_file: - # Encode and write frame to temporary file + with ExitStack() as cleanup: + try: + total_frames = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT)) + if total_frames > 0: + frame_indices = sorted( + random.sample(range(total_frames), min(self.num_sampled_frames, total_frames)) + ) + for frame_index in frame_indices: + video_capture.set(cv2.CAP_PROP_POS_FRAMES, frame_index) + ret, frame = video_capture.read() + if not ret: + logger.warning(f"Failed to read frame at index {frame_index}") + continue + ret, _ = cv2.imencode(".png", frame) if not ret: logger.warning(f"Failed to encode frame at index {frame_index}") continue - cv2.imwrite(temp_file.name, frame) + with tempfile.NamedTemporaryFile(suffix=f"_frame_{frame_index}.png", delete=False) as temp_file: + cleanup.callback(Path(temp_file.name).unlink, missing_ok=True) + if not cv2.imwrite(temp_file.name, frame): + raise OSError(f"Failed to write frame at index {frame_index}") frame_paths.append(temp_file.name) - finally: - video_capture.release() + finally: + video_capture.release() + # Successful extraction transfers ownership to the caller. + cleanup.pop_all() return frame_paths diff --git a/tests/unit/score/test_video_scorer.py b/tests/unit/score/test_video_scorer.py index f440cf74bf..24c72408c9 100644 --- a/tests/unit/score/test_video_scorer.py +++ b/tests/unit/score/test_video_scorer.py @@ -3,8 +3,10 @@ import asyncio import os +import tempfile import threading import uuid +from collections.abc import Iterator from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch @@ -20,6 +22,7 @@ from pyrit.score.scorer_prompt_validator import ScorerPromptValidator from pyrit.score.true_false.true_false_scorer import MessageTrueFalseScorer from pyrit.score.true_false.video_true_false_scorer import VideoTrueFalseScorer +from pyrit.score.video_scorer import VideoHelper def is_opencv_installed(): @@ -140,6 +143,160 @@ def _make_score(*, score_type: str, score_value: str | None, message_piece_id: u ) +@pytest.fixture +def isolated_frame_extraction( + tmp_path: Path, video_converter_sample_video: MessagePiece +) -> Iterator[tuple[VideoHelper, Path]]: + source = Path(video_converter_sample_video.converted_value) + original_video = source.read_bytes() + unrelated = tmp_path / "unrelated.png" + unrelated.write_bytes(b"caller-owned image") + helper = VideoHelper(image_capable_scorer=MockTrueFalseScorer(), num_sampled_frames=3) + with ( + patch.object(tempfile, "tempdir", str(tmp_path)), + patch("pyrit.score.video_scorer.random.sample", return_value=[0, 1, 2]), + ): + yield helper, source + assert source.read_bytes() == original_video + assert unrelated.read_bytes() == b"caller-owned image" + + +@pytest.mark.skipif(not is_opencv_installed(), reason="opencv is not installed") +@pytest.mark.parametrize("failure", ["encode", "write", "write_false", "create"]) +def test_extract_frames_partial_failure_removes_owned_files( + isolated_frame_extraction: tuple[VideoHelper, Path], failure: str +) -> None: + import cv2 + + helper, source = isolated_frame_extraction + capture = cv2.VideoCapture(str(source)) + encode = cv2.imencode + write = cv2.imwrite + create = tempfile.NamedTemporaryFile + + def encode_frame(extension: str, frame: np.ndarray) -> tuple[bool, np.ndarray]: + if failure == "encode" and encode_mock.call_count == 2: + raise OSError("Injected encode failure") + return encode(extension, frame) + + def write_frame(path: str, frame: np.ndarray) -> bool: + if failure in ("write", "write_false") and write_mock.call_count == 2: + if failure == "write_false": + return False + raise OSError("Injected write failure") + return bool(write(path, frame)) + + with ( + patch.object(cv2, "VideoCapture", return_value=capture), + patch.object(cv2, "imencode", side_effect=encode_frame) as encode_mock, + patch.object(cv2, "imwrite", side_effect=write_frame) as write_mock, + patch.object(tempfile, "NamedTemporaryFile", side_effect=create) as create_mock, + ): + if failure == "create": + create_mock.side_effect = [ + create(dir=source.parent, suffix="_frame_0.png", delete=False), + OSError("create"), + ] + with pytest.raises(OSError): + helper._extract_frames(str(source)) + + assert not capture.isOpened() + assert not list(source.parent.glob("*_frame_*.png")) + + +@pytest.mark.skipif(not is_opencv_installed(), reason="opencv is not installed") +@pytest.mark.parametrize("failed_frames", [{1}, {0, 1, 2}]) +def test_extract_frames_failed_encodes_do_not_create_files( + isolated_frame_extraction: tuple[VideoHelper, Path], failed_frames: set[int] +) -> None: + import cv2 + + helper, source = isolated_frame_extraction + encode = cv2.imencode + + def encode_frame(extension: str, frame: np.ndarray) -> tuple[bool, np.ndarray]: + if encode_mock.call_count - 1 in failed_frames: + return False, np.empty(0, dtype=np.uint8) + return encode(extension, frame) + + with patch.object(cv2, "imencode", side_effect=encode_frame) as encode_mock: + frames = helper._extract_frames(str(source)) + + assert set(frames) == {str(path) for path in source.parent.glob("*_frame_*.png")} + assert len(frames) == 3 - len(failed_frames) + assert all(cv2.imread(path) is not None for path in frames) + for path in frames: + Path(path).unlink() + + +@pytest.mark.skipif(not is_opencv_installed(), reason="opencv is not installed") +@pytest.mark.parametrize("cancel_twice", [False, True]) +@pytest.mark.parametrize("failure", ["write", "encode_false"]) +async def test_partial_frame_extraction_cancellation_removes_owned_files_async( + isolated_frame_extraction: tuple[VideoHelper, Path], failure: str, cancel_twice: bool +) -> None: + import cv2 + + helper, source = isolated_frame_extraction + loop = asyncio.get_running_loop() + started = asyncio.Event() + release = threading.Event() + capture = cv2.VideoCapture(str(source)) + encode = cv2.imencode + write = cv2.imwrite + extract = helper._extract_frames + + def hold_worker() -> None: + loop.call_soon_threadsafe(started.set) + if not release.wait(TEST_WATCHDOG_SECONDS): + raise TimeoutError("The test did not release partial frame extraction.") + + def write_frame(path: str, frame: np.ndarray) -> bool: + if failure == "write" and write_mock.call_count == 2: + hold_worker() + raise OSError("Injected partial write failure") + return bool(write(path, frame)) + + def encode_frame(extension: str, frame: np.ndarray) -> tuple[bool, np.ndarray]: + if failure == "encode_false" and encode_mock.call_count == 2: + return False, np.empty(0, dtype=np.uint8) + return encode(extension, frame) + + def extract_frames(video_path: str) -> list[str]: + paths = extract(video_path) + hold_worker() + return paths + + with ( + patch.object(cv2, "VideoCapture", return_value=capture), + patch.object(cv2, "imwrite", side_effect=write_frame) as write_mock, + patch.object(cv2, "imencode", side_effect=encode_frame) as encode_mock, + patch.object(helper, "_extract_frames", side_effect=extract_frames), + ): + task = asyncio.create_task(helper._extract_video_frames_async(str(source))) + try: + await wait_for_event_async(event=started, tasks=[task]) + assert list(source.parent.glob("*_frame_*.png")) + 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) as caught: + await wait_for_completion_async(future=task) + if failure == "write": + assert isinstance(caught.value.__cause__, OSError) + finally: + release.set() + await settle_tasks_async(task) + + assert not capture.isOpened() + assert not list(source.parent.glob("*_frame_*.png")) + + @pytest.mark.parametrize("float_scale", [False, True]) async def test_video_forwards_explicit_conditions_and_templates_async(float_scale: bool, tmp_path: Path) -> None: video_path = tmp_path / "typed-video.mp4"