[Bugfix][TTS] Isolate concurrent StepAudio2 TRT builds (#6957)

Signed-off-by: BANANASJIM <bananasjim1@gmail.com>

tests/model_executor/models/step_audio2/test_step_audio2_dit_trt.py

@@ -0,0 +1,257 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import os
+import stat
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from threading import Barrier, Event, Lock
+from typing import BinaryIO
+from uuid import UUID
+
+import pytest
+
+from vllm_omni.model_executor.models.step_audio2 import step_audio2_dit_trt
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _write_artifact(name: str, temporary: Path, stream: BinaryIO, payload: bytes) -> None:
+    if name.endswith(".onnx"):
+        temporary.write_bytes(payload)
+    else:
+        stream.write(payload)
+
+
+@pytest.mark.parametrize("name", ["dit_chunk.onnx", "dit_chunk.plan"])
+def test_atomic_publication_supports_same_process_concurrency(tmp_path: Path, name: str) -> None:
+    destination = tmp_path / name
+    destination.write_bytes(b"existing")
+    payloads = [b"first" * 1024, b"second" * 1024]
+    ready = Barrier(3)
+    release = Event()
+    temporary_paths: list[Path] = []
+    paths_lock = Lock()
+
+    def publish(payload: bytes) -> None:
+        def write(temporary: Path, stream: BinaryIO) -> None:
+            _write_artifact(name, temporary, stream, payload)
+            with paths_lock:
+                temporary_paths.append(temporary)
+            ready.wait(timeout=10)
+            assert release.wait(timeout=10)
+
+        step_audio2_dit_trt._publish_atomically(
+            destination,
+            write,
+            requires_path_write=name.endswith(".onnx"),
+        )
+
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        futures = [executor.submit(publish, payload) for payload in payloads]
+        ready.wait(timeout=10)
+        try:
+            assert destination.read_bytes() == b"existing"
+            assert len(set(temporary_paths)) == 2
+            assert all(path.parent == destination.parent for path in temporary_paths)
+        finally:
+            release.set()
+        for future in futures:
+            future.result(timeout=10)
+
+    assert destination.read_bytes() in payloads
+    assert not list(tmp_path.glob(f"{destination.name}.tmp.*"))
+
+
+def test_atomic_publication_does_not_remove_unowned_temporary_file(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    destination = tmp_path / "dit_chunk.onnx"
+    destination.write_bytes(b"existing")
+    collision_id = UUID(int=1)
+    monkeypatch.setattr(step_audio2_dit_trt.uuid, "uuid4", lambda: collision_id)
+    temporary = destination.with_name(f"{destination.name}.tmp.{os.getpid()}.{collision_id.hex}")
+    temporary.write_bytes(b"foreign")
+    writer_called = False
+
+    def write(path: Path, stream: BinaryIO) -> None:
+        nonlocal writer_called
+        del path, stream
+        writer_called = True
+
+    with pytest.raises(FileExistsError):
+        step_audio2_dit_trt._publish_atomically(destination, write)
+
+    assert not writer_called
+    assert destination.read_bytes() == b"existing"
+    assert temporary.read_bytes() == b"foreign"
+
+
+def test_atomic_publication_cleans_up_after_writer_failure(tmp_path: Path) -> None:
+    destination = tmp_path / "dit_chunk.onnx"
+    destination.write_bytes(b"existing")
+    write_error = RuntimeError("injected writer failure")
+    temporary: Path | None = None
+
+    def fail_write(path: Path, stream: BinaryIO) -> None:
+        nonlocal temporary
+        temporary = path
+        stream.write(b"partial")
+        raise write_error
+
+    with pytest.raises(RuntimeError) as error:
+        step_audio2_dit_trt._publish_atomically(destination, fail_write)
+
+    assert error.value is write_error
+    assert destination.read_bytes() == b"existing"
+    assert temporary is not None
+    assert not temporary.exists()
+
+
+def test_atomic_publication_cleans_up_after_fstat_failure(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    destination = tmp_path / "dit_chunk.onnx"
+    fstat_error = OSError("injected fstat failure")
+    owned_id = UUID(int=2)
+    monkeypatch.setattr(step_audio2_dit_trt.uuid, "uuid4", lambda: owned_id)
+    temporary = destination.with_name(f"{destination.name}.tmp.{os.getpid()}.{owned_id.hex}")
+    writer_called = False
+
+    def fail_fstat(fd: int) -> os.stat_result:
+        raise fstat_error
+
+    def write(path: Path, stream: BinaryIO) -> None:
+        nonlocal writer_called
+        del path, stream
+        writer_called = True
+
+    monkeypatch.setattr(os, "fstat", fail_fstat)
+
+    with pytest.raises(OSError) as error:
+        step_audio2_dit_trt._publish_atomically(destination, write, requires_path_write=True)
+
+    assert error.value is fstat_error
+    assert not writer_called
+    assert not destination.exists()
+    assert not temporary.exists()
+
+
+def test_atomic_publication_restores_permissions_when_cleanup_fails(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    destination = tmp_path / "dit_chunk.onnx"
+    writer_error = RuntimeError("injected writer failure")
+    temporary_id = UUID(int=3)
+    temporary = destination.with_name(f"{destination.name}.tmp.{os.getpid()}.{temporary_id.hex}")
+    original_unlink = Path.unlink
+    cleanup_attempted = False
+
+    def fail_write(path: Path, stream: BinaryIO) -> None:
+        del stream
+        path.write_bytes(b"partial")
+        raise writer_error
+
+    def fail_unlink(path: Path, *, missing_ok: bool = False) -> None:
+        nonlocal cleanup_attempted
+        del missing_ok
+        assert path == temporary
+        cleanup_attempted = True
+        raise OSError("injected cleanup failure")
+
+    monkeypatch.setattr(step_audio2_dit_trt.uuid, "uuid4", lambda: temporary_id)
+    monkeypatch.setattr(Path, "unlink", fail_unlink)
+    previous_umask = os.umask(0o222)
+    try:
+        with pytest.raises(RuntimeError) as error:
+            step_audio2_dit_trt._publish_atomically(destination, fail_write, requires_path_write=True)
+    finally:
+        os.umask(previous_umask)
+
+    assert error.value is writer_error
+    assert cleanup_attempted
+    assert stat.S_IMODE(temporary.stat().st_mode) == 0o444
+    original_unlink(temporary)
+
+
+@pytest.mark.parametrize(
+    ("name", "requires_path_write", "expected_fstat_calls", "expected_fchmod_calls"),
+    [("dit_chunk.onnx", True, 1, 2), ("dit_chunk.plan", False, 0, 0)],
+)
+def test_atomic_publication_preserves_permissions_with_restrictive_umask(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+    name: str,
+    requires_path_write: bool,
+    expected_fstat_calls: int,
+    expected_fchmod_calls: int,
+) -> None:
+    destination = tmp_path / name
+    original_fstat = os.fstat
+    original_fchmod = os.fchmod
+    fstat_calls = 0
+    fchmod_calls = 0
+
+    def track_fstat(fd: int) -> os.stat_result:
+        nonlocal fstat_calls
+        fstat_calls += 1
+        if not requires_path_write:
+            raise OSError("plan publication must not require fstat")
+        return original_fstat(fd)
+
+    def track_fchmod(fd: int, mode: int) -> None:
+        nonlocal fchmod_calls
+        fchmod_calls += 1
+        if not requires_path_write:
+            raise OSError("plan publication must not require fchmod")
+        original_fchmod(fd, mode)
+
+    monkeypatch.setattr(os, "fstat", track_fstat)
+    monkeypatch.setattr(os, "fchmod", track_fchmod)
+    previous_umask = os.umask(0o222)
+    try:
+        step_audio2_dit_trt._publish_atomically(
+            destination,
+            lambda temporary, stream: _write_artifact(name, temporary, stream, b"complete"),
+            requires_path_write=requires_path_write,
+        )
+    finally:
+        os.umask(previous_umask)
+
+    assert destination.read_bytes() == b"complete"
+    assert stat.S_IMODE(destination.stat().st_mode) == 0o444
+    assert fstat_calls == expected_fstat_calls
+    assert fchmod_calls == expected_fchmod_calls
+
+
+def test_atomic_publication_cleans_up_after_replace_failure(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    destination = tmp_path / "dit_chunk.plan"
+    destination.write_bytes(b"existing")
+    replace_error = OSError("injected replace failure")
+    temporary: Path | None = None
+
+    def write(path: Path, stream: BinaryIO) -> None:
+        nonlocal temporary
+        temporary = path
+        stream.write(b"complete")
+
+    def fail_replace(source: os.PathLike[str], target: os.PathLike[str]) -> None:
+        raise replace_error
+
+    monkeypatch.setattr(os, "replace", fail_replace)
+
+    with pytest.raises(OSError) as error:
+        step_audio2_dit_trt._publish_atomically(destination, write)
+
+    assert error.value is replace_error
+    assert destination.read_bytes() == b"existing"
+    assert temporary is not None
+    assert not temporary.exists()

vllm_omni/model_executor/models/step_audio2/step_audio2_dit_trt.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """TensorRT acceleration for the Step-Audio2 / CosyVoice2 flow DiT estimator.
 
 Ports stepfun-ai/Step-Audio2#70 into vLLM-Omni. The upstream PR patches
@@ -24,8 +24,13 @@ Differences from the upstream PR:
 
 from __future__ import annotations
 
+import contextlib
 import os
+import stat
+import uuid
+from collections.abc import Callable
 from pathlib import Path
+from typing import BinaryIO
 
 import torch
 import torch.nn as nn
@@ -36,6 +41,52 @@ logger = init_logger(__name__)
 _TRT_CACHE_ENV = "MINICPMO_TOKEN2WAV_TRT_CACHE"
 
 
+def _publish_atomically(
+    path: Path,
+    writer: Callable[[Path, BinaryIO], None],
+    *,
+    requires_path_write: bool = False,
+) -> None:
+    """Publish an artifact through an exclusively owned same-directory temp."""
+    temporary = path.with_name(f"{path.name}.tmp.{os.getpid()}.{uuid.uuid4().hex}")
+    stream = temporary.open("xb")
+    published = False
+    try:
+        original_mode: int | None = None
+        mode_adjusted = False
+        write_completed = False
+        try:
+            if requires_path_write:
+                original_mode = stat.S_IMODE(os.fstat(stream.fileno()).st_mode)
+                mode_adjusted = not original_mode & stat.S_IWUSR
+                if mode_adjusted:
+                    os.fchmod(stream.fileno(), original_mode | stat.S_IWUSR)
+            writer(temporary, stream)
+            stream.flush()
+            if mode_adjusted:
+                os.fchmod(stream.fileno(), original_mode)
+            write_completed = True
+        finally:
+            if not stream.closed:
+                if not write_completed and mode_adjusted and original_mode is not None:
+                    with contextlib.suppress(OSError):
+                        os.fchmod(stream.fileno(), original_mode)
+                if write_completed:
+                    stream.close()
+                else:
+                    with contextlib.suppress(OSError):
+                        stream.close()
+        os.replace(temporary, path)
+        published = True
+    finally:
+        if not published:
+            try:
+                temporary.unlink(missing_ok=True)
+            except OSError:
+                with contextlib.suppress(Exception):
+                    logger.warning("Failed to remove temporary StepAudio2 build artifact %s", temporary, exc_info=True)
+
+
 class _DiTChunkWrapper(nn.Module):
     """ONNX-exportable streaming ``forward_chunk`` with explicit stacked caches.
 
@@ -131,24 +182,22 @@ def export_dit_chunk_onnx(
     logger.info("Exporting DiT chunk ONNX to %s", onnx_path)
     onnx_path = Path(onnx_path)
     onnx_path.parent.mkdir(parents=True, exist_ok=True)
-    # Process-unique temp name + atomic rename, so concurrent first builds on
-    # a shared cache dir never see (or clobber) a partially written file.
-    tmp = onnx_path.with_suffix(f".tmp.{os.getpid()}")
-    try:
+
+    def write_onnx(temporary: Path, stream: BinaryIO) -> None:
+        del stream
         with torch.no_grad():
             torch.onnx.export(
                 wrapper,
                 sample,
-                str(tmp),
+                str(temporary),
                 input_names=input_names,
                 output_names=output_names,
                 dynamic_axes=dynamic_axes,
                 opset_version=17,
                 dynamo=False,
             )
-        os.replace(tmp, onnx_path)
-    finally:
-        tmp.unlink(missing_ok=True)
+
+    _publish_atomically(onnx_path, write_onnx, requires_path_write=True)
 
 
 def convert_onnx_to_trt(
@@ -212,16 +261,12 @@ def convert_onnx_to_trt(
         raise RuntimeError("TensorRT engine build failed")
     plan_path = Path(plan_path)
     plan_path.parent.mkdir(parents=True, exist_ok=True)
-    # Process-unique temp name: concurrent first builds (multiple ranks on a
-    # shared cache dir) must not clobber each other's partial writes; the
-    # final os.replace stays atomic and last-writer-wins with equal content.
-    tmp = plan_path.with_suffix(f".tmp.{os.getpid()}")
-    try:
-        with open(tmp, "wb") as f:
-            f.write(engine_bytes)
-        os.replace(tmp, plan_path)
-    finally:
-        tmp.unlink(missing_ok=True)
+
+    def write_plan(temporary: Path, stream: BinaryIO) -> None:
+        del temporary
+        stream.write(engine_bytes)
+
+    _publish_atomically(plan_path, write_plan)
     logger.info("DiT TRT engine written to %s", plan_path)