diff --git a/pyrit/score/video_scorer.py b/pyrit/score/video_scorer.py index 7cac6a1f06..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 @@ -169,7 +170,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 +211,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. @@ -222,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 @@ -232,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 99148e5fdf..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" @@ -654,6 +811,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.