FIX: Move streaming PCM file work off the event loop in PromptNormalizer

pyrit/prompt_normalizer/prompt_normalizer.py

@@ -392,7 +392,7 @@ async def convert_audio_async(
         if not converter_configurations:
             return raw_pcm
 
-        input_path = _write_pcm_to_temp_wav(
+        input_path = await _write_pcm_to_temp_wav_async(
             raw_pcm=raw_pcm,
             sample_rate_hz=sample_rate_hz,
             num_channels=num_channels,
@@ -411,7 +411,9 @@ async def convert_audio_async(
                 converter_configurations=converter_configurations,
                 message=message,
             )
-            actual_rate, actual_channels, actual_width, converted_pcm = _read_pcm_from_wav(piece.converted_value)
+            actual_rate, actual_channels, actual_width, converted_pcm = await asyncio.to_thread(
+                _read_pcm_from_wav, piece.converted_value
+            )
             if (actual_rate, actual_channels, actual_width) != (
                 sample_rate_hz,
                 num_channels,
@@ -425,7 +427,8 @@ async def convert_audio_async(
                 )
             return converted_pcm
         finally:
-            Path(input_path).unlink(missing_ok=True)
+            # The worker thread finishes the removal even if this await is cancelled.
+            await asyncio.to_thread(Path(input_path).unlink, missing_ok=True)
 
     async def _calc_hash_async(self, request: Message) -> None:
         """Add a request to the memory."""
@@ -500,6 +503,53 @@ async def add_prepended_conversation_to_memory_async(
         return prepended_conversation
 
 
+async def _write_pcm_to_temp_wav_async(
+    *,
+    raw_pcm: bytes,
+    sample_rate_hz: int,
+    num_channels: int,
+    sample_width_bytes: int,
+) -> str:
+    """
+    Write the PCM to a new temp WAV file in a worker thread and return its path.
+
+    A worker thread cannot be stopped, so a cancelled caller would never learn the path of the
+    file the thread goes on to create. In that case the file is removed once the write finishes.
+
+    Args:
+        raw_pcm (bytes): Raw PCM audio samples (no WAV header).
+        sample_rate_hz (int): Sample rate of the PCM in Hz.
+        num_channels (int): Channel count.
+        sample_width_bytes (int): Bytes per sample.
+
+    Returns:
+        str: The path of the new temp WAV file.
+
+    Raises:
+        asyncio.CancelledError: If the caller is cancelled before the write finishes.
+    """
+    write = asyncio.ensure_future(
+        asyncio.to_thread(
+            _write_pcm_to_temp_wav,
+            raw_pcm=raw_pcm,
+            sample_rate_hz=sample_rate_hz,
+            num_channels=num_channels,
+            sample_width_bytes=sample_width_bytes,
+        )
+    )
+    try:
+        return await asyncio.shield(write)
+    except asyncio.CancelledError:
+        write.add_done_callback(_remove_abandoned_temp_wav)
+        raise
+
+
+def _remove_abandoned_temp_wav(write: asyncio.Future[str]) -> None:
+    """Remove the temp WAV a cancelled caller left behind, if the write produced one."""
+    if not write.cancelled() and write.exception() is None:
+        Path(write.result()).unlink(missing_ok=True)
+
+
 def _write_pcm_to_temp_wav(
     *,
     raw_pcm: bytes,

tests/unit/prompt_normalizer/test_prompt_normalizer.py

@@ -5,6 +5,7 @@
 import codecs
 import os
 import tempfile
+import threading
 import wave
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
@@ -45,6 +46,7 @@
     SeedPrompt,
 )
 from pyrit.prompt_normalizer import NormalizerRequest, PromptNormalizer
+from pyrit.prompt_normalizer import prompt_normalizer as prompt_normalizer_module
 from pyrit.prompt_normalizer.converter_configuration import (
     ConverterConfiguration,
 )
@@ -1173,3 +1175,119 @@ async def swap_to_stereo(*, converter_configurations, message):
                 )
     finally:
         Path(wrong_channels_path).unlink(missing_ok=True)
+
+
+async def test_convert_audio_async_runs_file_work_off_the_event_loop(
+    mock_memory_instance, sample_pcm, dummy_audio_converter_config
+):
+    event_loop_thread = threading.get_ident()
+    file_work_threads: dict[str, int] = {}
+
+    def record_thread(name, function):
+        def run(*args, **kwargs):
+            file_work_threads[name] = threading.get_ident()
+            return function(*args, **kwargs)
+
+        return run
+
+    normalizer = PromptNormalizer()
+    with (
+        patch.object(
+            prompt_normalizer_module,
+            "_write_pcm_to_temp_wav",
+            side_effect=record_thread("write", prompt_normalizer_module._write_pcm_to_temp_wav),
+        ),
+        patch.object(
+            prompt_normalizer_module,
+            "_read_pcm_from_wav",
+            side_effect=record_thread("read", prompt_normalizer_module._read_pcm_from_wav),
+        ),
+        patch.object(Path, "unlink", autospec=True, side_effect=record_thread("unlink", Path.unlink)),
+        patch.object(normalizer, "convert_values_async", new_callable=AsyncMock),
+    ):
+        result = await normalizer.convert_audio_async(
+            raw_pcm=sample_pcm,
+            converter_configurations=[dummy_audio_converter_config],
+            sample_rate_hz=_AUDIO_SAMPLE_RATE_HZ,
+            num_channels=_AUDIO_NUM_CHANNELS,
+            sample_width_bytes=_AUDIO_SAMPLE_WIDTH_BYTES,
+        )
+
+    assert result == sample_pcm
+    assert set(file_work_threads) == {"write", "read", "unlink"}
+    assert event_loop_thread not in file_work_threads.values()
+
+
+async def test_convert_audio_async_cleans_up_temp_file_when_cancelled_during_conversion(
+    mock_memory_instance, sample_pcm, dummy_audio_converter_config
+):
+    captured_paths: list[str] = []
+    converting = asyncio.Event()
+
+    async def capture_then_wait(*, converter_configurations, message):
+        captured_paths.append(message.message_pieces[0].converted_value)
+        converting.set()
+        await asyncio.Event().wait()
+
+    normalizer = PromptNormalizer()
+    with patch.object(normalizer, "convert_values_async", side_effect=capture_then_wait):
+        task = asyncio.create_task(
+            normalizer.convert_audio_async(
+                raw_pcm=sample_pcm,
+                converter_configurations=[dummy_audio_converter_config],
+                sample_rate_hz=_AUDIO_SAMPLE_RATE_HZ,
+                num_channels=_AUDIO_NUM_CHANNELS,
+                sample_width_bytes=_AUDIO_SAMPLE_WIDTH_BYTES,
+            )
+        )
+        await converting.wait()
+        task.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+
+    assert len(captured_paths) == 1
+    assert not Path(captured_paths[0]).exists()
+
+
+async def test_convert_audio_async_removes_temp_file_written_after_cancellation(
+    mock_memory_instance, sample_pcm, dummy_audio_converter_config
+):
+    write_started = threading.Event()
+    release_write = threading.Event()
+    written_paths: list[str] = []
+    write_pcm_to_temp_wav = prompt_normalizer_module._write_pcm_to_temp_wav
+
+    def write_after_release(**kwargs):
+        write_started.set()
+        release_write.wait(timeout=5)
+        written_paths.append(write_pcm_to_temp_wav(**kwargs))
+        return written_paths[-1]
+
+    normalizer = PromptNormalizer()
+    with (
+        patch.object(prompt_normalizer_module, "_write_pcm_to_temp_wav", side_effect=write_after_release),
+        patch.object(normalizer, "convert_values_async", new_callable=AsyncMock) as mock_convert,
+    ):
+        task = asyncio.create_task(
+            normalizer.convert_audio_async(
+                raw_pcm=sample_pcm,
+                converter_configurations=[dummy_audio_converter_config],
+                sample_rate_hz=_AUDIO_SAMPLE_RATE_HZ,
+                num_channels=_AUDIO_NUM_CHANNELS,
+                sample_width_bytes=_AUDIO_SAMPLE_WIDTH_BYTES,
+            )
+        )
+        assert await asyncio.to_thread(write_started.wait, 5)
+        task.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+
+        release_write.set()
+        for _ in range(500):
+            if written_paths and not Path(written_paths[0]).exists():
+                break
+            await asyncio.sleep(0.01)
+
+    assert len(written_paths) == 1
+    assert not Path(written_paths[0]).exists()
+    mock_convert.assert_not_called()