From fb521400a2e61eb4313e4151b5efe998f073eb5c Mon Sep 17 00:00:00 2001 From: debpalash Date: Tue, 19 May 2026 08:00:33 +0530 Subject: [PATCH] =?UTF-8?q?P0(dub):=20atomic=20WAV=20writes=20=E2=80=94=20?= =?UTF-8?q?closes=20#48?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Direct torchaudio.save(path, ...) writes bytes as the encoder produces them. SIGKILL, OOM kill, or Tauri sidecar reap mid-write leaves the file at `path` truncated. Downstream tools (ffmpeg in the dub mux, NLEs the user imports the WAV into) happily read truncated RIFF — the header appears first, then the data chunk gets cut short — and surface as silently corrupt audio later in the pipeline. That's the shape of #48. New helper services/audio_io.py:atomic_save_wav() writes to a sibling temp file in the same directory then os.replace() into place. POSIX rename(2) is atomic; os.replace() ports the same guarantee to Windows. Either the target ends up with a complete WAV or it keeps its previous contents (or never exists) — no third state. Migrated three call sites in api/routers/dub_generate.py: - L289: RVC per-segment write - L328: deferred batch write of all segments - L390: final mixed-track export Left L508 alone — it writes to BytesIO (in-memory response body), no atomicity needed. Implementation note (recorded as a docstring in audio_io.py): the temp file must end in `.wav`, not `.tmp`. torchaudio.save infers the output format from the path suffix and ignores the `format=` kwarg with the soundfile backend. A `.tmp` suffix raises "Unsupported format: tmp". The leading dot + target-name prefix still marks the file as transient. Tests in backend/tests/test_atomic_wav.py: - success path: writes valid WAV, no temp leaks, overwrites cleanly - atomicity: target unchanged when save raises (pre-existing target) - atomicity: target absent when save raises (new target path) - no temp leaks on failure - temp file lives in target_dir (cross-fs renames are not atomic) Co-Authored-By: Claude Opus 4.7 (1M context) --- backend/api/routers/dub_generate.py | 7 +- backend/services/audio_io.py | 83 ++++++++++++++++ backend/tests/test_atomic_wav.py | 146 ++++++++++++++++++++++++++++ 3 files changed, 233 insertions(+), 3 deletions(-) create mode 100644 backend/services/audio_io.py create mode 100644 backend/tests/test_atomic_wav.py diff --git a/backend/api/routers/dub_generate.py b/backend/api/routers/dub_generate.py index 981ef08a..dc43e83a 100644 --- a/backend/api/routers/dub_generate.py +++ b/backend/api/routers/dub_generate.py @@ -14,6 +14,7 @@ from core.tasks import task_manager from schemas.requests import DubRequest from services.model_manager import get_model, _gpu_pool from services.audio_dsp import apply_mastering, normalize_audio +from services.audio_io import atomic_save_wav from services.rvc import apply_rvc, is_enabled as rvc_is_enabled from services.incremental import segment_fingerprint from services.watermark import embed_watermark @@ -286,7 +287,7 @@ async def dub_generate(job_id: str, req: DubRequest): # when RVC is active (uncommon path). if rvc_is_enabled(): seg_wav_path = os.path.join(DUB_DIR, job_id, f"seg_{i}.wav") - torchaudio.save(seg_wav_path, audio_tensor, _model.sampling_rate) + atomic_save_wav(seg_wav_path, audio_tensor, _model.sampling_rate) try: await loop.run_in_executor(_gpu_pool, apply_rvc, seg_wav_path) rvc_wav, rvc_sr = torchaudio.load(seg_wav_path) @@ -325,7 +326,7 @@ async def dub_generate(job_id: str, req: DubRequest): try: # Apply invisible watermark before writing to disk _wav = embed_watermark(_wav, _sr) - torchaudio.save(seg_wav_path, _wav, _sr) + atomic_save_wav(seg_wav_path, _wav, _sr) except Exception as e: logger.warning("deferred seg write failed for %s: %s", _sid, e) if _fp is not None: @@ -387,7 +388,7 @@ async def dub_generate(job_id: str, req: DubRequest): _t_save_0 = time.perf_counter() # Apply invisible watermark to the final assembled track full_audio = embed_watermark(full_audio, sr) - torchaudio.save(track_path, full_audio, sr) + atomic_save_wav(track_path, full_audio, sr) _t_save = time.perf_counter() - _t_save_0 _t_mix = _t_save_0 - _t_loop_end job["dubbed_tracks"][lang_code] = { diff --git a/backend/services/audio_io.py b/backend/services/audio_io.py new file mode 100644 index 00000000..810aa442 --- /dev/null +++ b/backend/services/audio_io.py @@ -0,0 +1,83 @@ +"""Atomic disk writes for audio files. + +A direct ``torchaudio.save(path, ...)`` writes bytes to ``path`` as the +encoder produces them. If the process is killed mid-write (SIGKILL, OOM +kill, power loss, Tauri sidecar reap), the file at ``path`` exists but is +truncated. Downstream tools — ffmpeg in the dub mux step, NLEs the user +imports the WAV into — often happily read a truncated RIFF (the header +appears first, then the data chunk gets cut short), producing silently +corrupt audio later in the pipeline. That is the root of issue #48. + +``atomic_save_wav`` writes to a sibling temp file in the same directory and +``os.replace()`` it into place once encoding completes. POSIX guarantees +``rename(2)`` is atomic on the same filesystem; ``os.replace()`` makes the +same guarantee portable to Windows, including the case where the target +path already exists. + +Either the new file fully exists at ``target_path`` after the call returns, +or the target keeps its previous contents (or never existed). There is no +intermediate window where a partial WAV is visible at ``target_path``. + +Closes #48. +""" +from __future__ import annotations + +import logging +import os +import tempfile +from typing import Any + +import torch +import torchaudio + +logger = logging.getLogger("omnivoice.audio_io") + + +def atomic_save_wav( + target_path: str, + audio: torch.Tensor, + sample_rate: int, + **kwargs: Any, +) -> None: + """Write a WAV to ``target_path`` atomically. + + Implementation: write to a sibling temp file in the same directory, then + ``os.replace()`` into place. Cross-filesystem renames are *not* atomic + on POSIX, so the temp file must live next to the target — that is why + we use ``dir=target_dir`` instead of the system temp dir. + + Args: + target_path: Final destination. Parent directory must already exist. + audio: ``(channels, samples)`` tensor — the same shape + ``torchaudio.save`` expects. + sample_rate: WAV sample rate in Hz. + **kwargs: Forwarded to ``torchaudio.save``. + + Raises: + Whatever ``torchaudio.save`` raises. The temp file is unlinked on + failure so we do not leak ``.tmp`` files in ``DUB_DIR``. + """ + target_dir = os.path.dirname(target_path) or "." + target_base = os.path.basename(target_path) + # The temp file must end in ``.wav`` even though it is conceptually a + # ``.tmp`` file. torchaudio.save infers the output format from the path + # suffix and *ignores* the ``format=`` kwarg with the soundfile backend + # — a ``.tmp`` suffix raises ``ValueError: Unsupported format: tmp``. + # The leading dot + ``target_base`` prefix still marks the file as + # transient and groups it next to its target in directory listings. + fd, tmp_path = tempfile.mkstemp( + prefix=f".{target_base}.", + suffix=".wav", + dir=target_dir, + ) + os.close(fd) # torchaudio reopens by path; we just needed a unique name. + try: + torchaudio.save(tmp_path, audio, sample_rate, **kwargs) + os.replace(tmp_path, target_path) + except BaseException: + # BaseException so we clean up on KeyboardInterrupt + SystemExit too. + try: + os.unlink(tmp_path) + except OSError: + pass + raise diff --git a/backend/tests/test_atomic_wav.py b/backend/tests/test_atomic_wav.py new file mode 100644 index 00000000..19d666b6 --- /dev/null +++ b/backend/tests/test_atomic_wav.py @@ -0,0 +1,146 @@ +"""Tests for ``services.audio_io.atomic_save_wav`` — closes #48. + +The invariant we are protecting: when ``atomic_save_wav`` returns, the +target path either contains a complete, valid WAV or is unchanged. There +is no third state where a partial WAV is visible at the target path and +downstream tools (ffmpeg in the dub mux, NLEs the user imports the WAV +into) read truncated audio without an error. +""" +import os +import sys +from pathlib import Path + +import pytest +import torch +import torchaudio + +sys.path.insert(0, os.path.dirname(os.path.dirname(__file__))) + +from services.audio_io import atomic_save_wav # noqa: E402 + + +class TestSuccessPath: + def test_writes_valid_wav(self, tmp_path: Path): + target = tmp_path / "out.wav" + audio = torch.randn(1, 24000) # 1s mono @ 24kHz + atomic_save_wav(str(target), audio, 24000) + + assert target.exists() + loaded, sr = torchaudio.load(str(target)) + assert sr == 24000 + assert loaded.shape == audio.shape + + def test_no_temp_leaks_on_success(self, tmp_path: Path): + target = tmp_path / "out.wav" + atomic_save_wav(str(target), torch.zeros(1, 100), 24000) + + leaked = [p for p in tmp_path.glob(".*") if p.name.startswith(".")] + assert leaked == [], f"leaked temp files after success: {leaked}" + + def test_overwrites_existing_target(self, tmp_path: Path): + target = tmp_path / "out.wav" + # Pre-populate with a different-length WAV + torchaudio.save(str(target), torch.zeros(1, 1000), 24000) + old_samples = torchaudio.load(str(target))[0].shape[-1] + + new_audio = torch.randn(1, 5000) + atomic_save_wav(str(target), new_audio, 24000) + + loaded, _ = torchaudio.load(str(target)) + assert loaded.shape[-1] == 5000 + assert loaded.shape[-1] != old_samples + + +class TestAtomicity: + """The core invariant of #48: no partial files at the target path.""" + + def test_target_unchanged_when_save_raises( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ): + target = tmp_path / "out.wav" + original_bytes = b"PREVIOUS-CONTENT-DO-NOT-CORRUPT" + target.write_bytes(original_bytes) + + def explode(*args, **kwargs): + raise RuntimeError("simulated kill mid-write") + + # Patch the symbol *inside* the audio_io module, not the global — + # rebinding torchaudio.save would leak into other tests. + monkeypatch.setattr( + "services.audio_io.torchaudio.save", explode + ) + + with pytest.raises(RuntimeError, match="simulated kill"): + atomic_save_wav(str(target), torch.zeros(1, 100), 24000) + + assert target.read_bytes() == original_bytes, ( + "atomic_save_wav must not modify the target path on failure" + ) + + def test_no_temp_leaks_on_failure( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ): + target = tmp_path / "out.wav" + + def explode(*args, **kwargs): + raise RuntimeError("boom") + + monkeypatch.setattr( + "services.audio_io.torchaudio.save", explode + ) + + with pytest.raises(RuntimeError): + atomic_save_wav(str(target), torch.zeros(1, 100), 24000) + + leaked = [p for p in tmp_path.glob(".*") if p.name.startswith(".")] + assert leaked == [], f"leaked temp files after failure: {leaked}" + + def test_target_absent_when_save_raises_on_new_path( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ): + target = tmp_path / "never-existed.wav" + assert not target.exists() + + def explode(*args, **kwargs): + raise RuntimeError("boom") + + monkeypatch.setattr( + "services.audio_io.torchaudio.save", explode + ) + + with pytest.raises(RuntimeError): + atomic_save_wav(str(target), torch.zeros(1, 100), 24000) + + assert not target.exists(), ( + "atomic_save_wav must not create the target path on failure" + ) + + def test_temp_file_lives_in_target_dir(self, tmp_path: Path): + """Cross-fs renames are not atomic on POSIX. The temp file *must* + live next to the target so os.replace() stays a single rename(). + We assert this by intercepting torchaudio.save to inspect the path + it was handed. + """ + target = tmp_path / "out.wav" + captured: list[str] = [] + + # Capture the path torchaudio.save is called with, then call the + # real implementation so the test still ends in a valid WAV. + from services import audio_io as _aio + real_save = _aio.torchaudio.save + + def spy(path, *args, **kwargs): + captured.append(path) + return real_save(path, *args, **kwargs) + + import unittest.mock + with unittest.mock.patch.object(_aio.torchaudio, "save", side_effect=spy): + atomic_save_wav(str(target), torch.zeros(1, 100), 24000) + + assert len(captured) == 1 + tmp_used = captured[0] + assert os.path.dirname(tmp_used) == str(tmp_path), ( + f"temp file {tmp_used} not in target dir {tmp_path} — " + "cross-fs rename would break atomicity" + ) + assert tmp_used != str(target)