[Bugfix][TTS] Make CosyVoice3 TRT plan publication concurrency-safe (#6955)

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

tests/model_executor/models/cosyvoice3/test_flow_estimator_trt.py

@@ -0,0 +1,125 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+import pytest
+
+from vllm_omni.model_executor.models.cosyvoice3 import flow_estimator_trt
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _temporary_plans(plan_path: Path) -> list[Path]:
+    return list(plan_path.parent.glob(f"{plan_path.name}.tmp.*"))
+
+
+def test_write_plan_cleans_up_after_replace_failure(tmp_path, monkeypatch):
+    plan_path = tmp_path / "flow.plan"
+    plan_path.write_bytes(b"existing plan")
+    replace_error = OSError("replace failed")
+
+    def fail_replace(source, destination):
+        raise replace_error
+
+    monkeypatch.setattr(flow_estimator_trt.os, "replace", fail_replace)
+
+    with pytest.raises(OSError) as exc_info:
+        flow_estimator_trt._write_plan_atomically(b"new plan", str(plan_path))
+
+    assert exc_info.value is replace_error
+    assert plan_path.read_bytes() == b"existing plan"
+    assert _temporary_plans(plan_path) == []
+
+
+def test_write_plan_preserves_replace_error_when_cleanup_fails(tmp_path, monkeypatch):
+    plan_path = tmp_path / "flow.plan"
+    replace_error = OSError("replace failed")
+
+    def fail_replace(source, destination):
+        raise replace_error
+
+    def fail_unlink(path):
+        raise PermissionError("cleanup failed")
+
+    monkeypatch.setattr(flow_estimator_trt.os, "replace", fail_replace)
+    monkeypatch.setattr(flow_estimator_trt.os, "unlink", fail_unlink)
+
+    with pytest.raises(OSError) as exc_info:
+        flow_estimator_trt._write_plan_atomically(b"new plan", str(plan_path))
+
+    assert exc_info.value is replace_error
+    assert len(_temporary_plans(plan_path)) == 1
+
+
+def test_write_plan_cleans_up_after_write_failure(tmp_path, monkeypatch):
+    plan_path = tmp_path / "flow.plan"
+    write_error = OSError("write failed")
+    real_open = open
+
+    class FailingWriter:
+        def __init__(self, path, mode):
+            self.file = real_open(path, mode)
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc_value, traceback):
+            self.file.close()
+
+        def write(self, data):
+            self.file.write(data[:1])
+            raise write_error
+
+    monkeypatch.setattr(flow_estimator_trt, "open", FailingWriter, raising=False)
+
+    with pytest.raises(OSError) as exc_info:
+        flow_estimator_trt._write_plan_atomically(b"new plan", str(plan_path))
+
+    assert exc_info.value is write_error
+    assert not plan_path.exists()
+    assert _temporary_plans(plan_path) == []
+
+
+def test_write_plan_does_not_remove_a_colliding_temporary_file(tmp_path, monkeypatch):
+    plan_path = tmp_path / "flow.plan"
+    token = "0" * 32
+    temporary_path = Path(f"{plan_path}.tmp.{flow_estimator_trt.os.getpid()}.{token}")
+    temporary_path.write_bytes(b"another writer")
+    monkeypatch.setattr(flow_estimator_trt.uuid, "uuid4", lambda: flow_estimator_trt.uuid.UUID(hex=token))
+
+    with pytest.raises(FileExistsError):
+        flow_estimator_trt._write_plan_atomically(b"new plan", str(plan_path))
+
+    assert temporary_path.read_bytes() == b"another writer"
+    assert not plan_path.exists()
+
+
+def test_write_plan_supports_concurrent_publication(tmp_path, monkeypatch):
+    plan_path = tmp_path / "flow.plan"
+    payloads = (b"a" * 4096, b"b" * 4096)
+    barrier = threading.Barrier(len(payloads))
+    source_paths = []
+    source_paths_lock = threading.Lock()
+    real_replace = flow_estimator_trt.os.replace
+
+    def synchronized_replace(source, destination):
+        with source_paths_lock:
+            source_paths.append(Path(source))
+        barrier.wait(timeout=5)
+        real_replace(source, destination)
+
+    monkeypatch.setattr(flow_estimator_trt.os, "replace", synchronized_replace)
+
+    with ThreadPoolExecutor(max_workers=len(payloads)) as executor:
+        futures = [
+            executor.submit(flow_estimator_trt._write_plan_atomically, payload, str(plan_path)) for payload in payloads
+        ]
+        for future in futures:
+            future.result(timeout=10)
+
+    assert len(set(source_paths)) == len(payloads)
+    assert plan_path.read_bytes() in payloads
+    assert _temporary_plans(plan_path) == []

vllm_omni/model_executor/models/cosyvoice3/flow_estimator_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 engine for the CosyVoice3 flow-decoder (CFM) DiT estimator.
 
 The estimator is the per-step network the conditional flow-matching ODE solver
@@ -26,6 +26,7 @@ from __future__ import annotations
 
 import os
 import queue
+import uuid
 
 import torch
 from vllm.logger import init_logger
@@ -52,6 +53,26 @@ def _is_fp16_onnx(onnx_path: str) -> bool:
     return "fp16" in name or "autocast" in name
 
 
+def _write_plan_atomically(engine_bytes, plan_path: str) -> None:
+    tmp = f"{plan_path}.tmp.{os.getpid()}.{uuid.uuid4().hex}"
+    tmp_created = False
+    try:
+        with open(tmp, "xb") as f:
+            tmp_created = True
+            f.write(engine_bytes)
+        os.replace(tmp, plan_path)
+        tmp_created = False
+    except BaseException:
+        if tmp_created:
+            try:
+                os.unlink(tmp)
+            except FileNotFoundError:
+                pass
+            except OSError:
+                logger.warning("Failed to remove temporary TensorRT plan %s", tmp, exc_info=True)
+        raise
+
+
 def _convert_onnx_to_trt(onnx_path: str, plan_path: str, strongly_typed: bool) -> None:
     import tensorrt as trt
 
@@ -96,10 +117,7 @@ def _convert_onnx_to_trt(onnx_path: str, plan_path: str, strongly_typed: bool) -
     engine_bytes = builder.build_serialized_network(network, config)
     if engine_bytes is None:
         raise RuntimeError(f"TensorRT failed to build flow-estimator engine from {onnx_path}")
-    tmp = plan_path + ".tmp"
-    with open(tmp, "wb") as f:
-        f.write(engine_bytes)
-    os.replace(tmp, plan_path)
+    _write_plan_atomically(engine_bytes, plan_path)
     logger.info("Wrote flow-estimator TensorRT engine to %s", plan_path)