From 711daacabf1b4aecbb1492623a5d4df84b158656 Mon Sep 17 00:00:00 2001 From: Jim Ban <77719403+BANANASJIM@users.noreply.github.com> Date: Wed, 2 Sep 2026 23:44:42 -0700 Subject: [PATCH] [Bugfix][TTS] Make CosyVoice3 TRT plan publication concurrency-safe (#6955) Signed-off-by: BANANASJIM --- .../cosyvoice3/test_flow_estimator_trt.py | 125 ++++++++++++++++++ .../models/cosyvoice3/flow_estimator_trt.py | 28 +++- 2 files changed, 148 insertions(+), 5 deletions(-) create mode 100644 tests/model_executor/models/cosyvoice3/test_flow_estimator_trt.py diff --git a/tests/model_executor/models/cosyvoice3/test_flow_estimator_trt.py b/tests/model_executor/models/cosyvoice3/test_flow_estimator_trt.py new file mode 100644 index 000000000..07c618102 --- /dev/null +++ b/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) == [] diff --git a/vllm_omni/model_executor/models/cosyvoice3/flow_estimator_trt.py b/vllm_omni/model_executor/models/cosyvoice3/flow_estimator_trt.py index ea2f84e39..36292cd43 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/flow_estimator_trt.py +++ b/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)