From 8f115164c236e140755b82f26edc7d8567de7757 Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Mon, 27 Jul 2026 09:04:08 -0700 Subject: [PATCH] [Bugfix][Connectors] Flush processor tail on terminal async-chunk sends (#5414) Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> --- .../test_chunk_transfer_adapter.py | 31 +++++++ .../minicpmo_4_5/test_code2wav_model_dir.py | 88 +++++++++++++++++++ .../chunk_transfer_adapter.py | 8 +- .../minicpmo_4_5/minicpmo_4_5_code2wav.py | 32 +++++-- 4 files changed, 152 insertions(+), 7 deletions(-) create mode 100644 tests/model_executor/models/minicpmo_4_5/test_code2wav_model_dir.py diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 1a045f306..3ab2528fb 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -252,6 +252,37 @@ def test_save_async_uses_confirmed_tokens_for_async_scheduler_watermark(build_ad assert len(adapter._pending_save_reqs) == 1 +def test_send_single_request_terminal_chunk_still_flushes_processor(build_adapter, monkeypatch): + """A terminal stop is not a segment boundary (#5383), but the producer-side + processor must still receive the flush signal on the terminal chunk. + Passing only ``is_segment_finished`` starved processors of their final + accumulated payload once terminal stops stopped setting it (#5413: the + downstream stage got ``meta.finished`` with the tail data missing). + """ + adapter, connector = build_adapter(stage_id=0) + request = _req("req-terminal", RequestStatus.FINISHED_STOPPED, external_req_id="ext-terminal") + + seen_flush_flags = [] + + def recording_processor(**kwargs): + seen_flush_flags.append(kwargs["is_finished"]) + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.tensor([1, 2, 3], dtype=torch.long)), + ) + + adapter.custom_process_next_stage_input_func = recording_processor + monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None) + + adapter._send_single_request( + {"multimodal_output": None, "request": request, "is_finished": True, "is_segment_finished": False} + ) + + assert seen_flush_flags == [True] + sent_payload = connector.put.call_args.kwargs["data"] + assert bool(sent_payload.meta.finished.item()) is True + assert bool(sent_payload.meta.is_segment_finished.item()) is False + + def test_send_single_request_struct_without_meta_does_not_crash(build_adapter, monkeypatch): """Producer may return a struct with ``meta=None`` (e.g. payload that carries only ``embed`` or ``codes``). The sender's ``meta is not None`` diff --git a/tests/model_executor/models/minicpmo_4_5/test_code2wav_model_dir.py b/tests/model_executor/models/minicpmo_4_5/test_code2wav_model_dir.py new file mode 100644 index 000000000..5aa78262a --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_code2wav_model_dir.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for MiniCPM-o 4.5 code2wav model-dir resolution (#5442). + +In hub/CI deployments ``model_config.model`` is a repo id rather than a local +directory, so asset lookups must resolve to the downloaded snapshot. The +resolution must stay lazy: constructing the model with a fake path (as the +CPU unit tests do) must not touch the hub. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import huggingface_hub +import pytest + +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import ( + MiniCPMO45Code2Wav, + _resolve_model_dir, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _no_hub(monkeypatch): + def _fail(*args, **kwargs): # pragma: no cover - must not be reached + raise AssertionError("snapshot_download must not be called here") + + monkeypatch.setattr(huggingface_hub, "snapshot_download", _fail) + + +def test_local_directory_is_returned_unchanged(tmp_path, monkeypatch): + _no_hub(monkeypatch) + assert _resolve_model_dir(str(tmp_path)) == str(tmp_path) + + +def test_repo_id_resolves_via_snapshot_download(tmp_path, monkeypatch): + calls = {} + + def _fake_snapshot_download(model_ref, revision=None, allow_patterns=None): + calls["model_ref"] = model_ref + calls["revision"] = revision + calls["allow_patterns"] = allow_patterns + return str(tmp_path / "snapshot") + + monkeypatch.setattr(huggingface_hub, "snapshot_download", _fake_snapshot_download) + resolved = _resolve_model_dir("openbmb/MiniCPM-o-4_5", revision="abc123") + assert resolved == str(tmp_path / "snapshot") + assert calls["model_ref"] == "openbmb/MiniCPM-o-4_5" + assert calls["revision"] == "abc123" + assert calls["allow_patterns"] == ["assets/*"] + + +def test_snapshot_download_failure_propagates(monkeypatch): + def _raise(*args, **kwargs): + raise FileNotFoundError("offline and not cached") + + monkeypatch.setattr(huggingface_hub, "snapshot_download", _raise) + with pytest.raises(FileNotFoundError): + _resolve_model_dir("openbmb/MiniCPM-o-4_5") + + +def test_init_with_fake_path_does_not_resolve(monkeypatch): + """Mirrors the CPU-test construction: fake model path, no ``revision``.""" + _no_hub(monkeypatch) + config = SimpleNamespace( + model_config=SimpleNamespace( + model="/fake/model", + stage_connector_config=None, + ) + ) + model = MiniCPMO45Code2Wav(vllm_config=config) + assert model.model_path == "/fake/model" + assert model._default_prompt_wav == "/fake/model/assets/HT_ref_audio.wav" + + +def test_default_prompt_wav_follows_resolved_model_path(monkeypatch): + _no_hub(monkeypatch) + config = SimpleNamespace( + model_config=SimpleNamespace( + model="openbmb/MiniCPM-o-4_5", + stage_connector_config=None, + ) + ) + model = MiniCPMO45Code2Wav(vllm_config=config) + model.model_path = "/resolved/snapshot" + assert model._default_prompt_wav == "/resolved/snapshot/assets/HT_ref_audio.wav" diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index 1a41372a8..83def2dc0 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -317,7 +317,13 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): multimodal_output=multimodal_output, request=request, # Existing processors use is_finished as a flush signal. - is_finished=is_segment_finished, + # Terminal stops no longer count as segment boundaries + # (is_segment_finished is False when the request finishes, + # see #5383), but the processor must still flush its + # accumulated tail on the terminal chunk — otherwise the + # downstream stage receives the finished marker without + # the final payload (#5413). + is_finished=is_segment_finished or is_finished, ) except Exception as e: diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py index a9ffa6e0a..3eac2e49d 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py @@ -30,6 +30,20 @@ from .batched_token2wav import ( logger = init_logger(__name__) +def _resolve_model_dir(model_ref: str, revision: str | None = None) -> str: + """Resolve ``model_ref`` to a local directory containing the repo assets. + + ``model_config.model`` is a filesystem path in local deployments but a + Hugging Face repo id in hub/CI deployments; the prompt-audio and + token2wav asset lookups need a real directory either way. + """ + if Path(model_ref).is_dir(): + return model_ref + from huggingface_hub import snapshot_download + + return snapshot_download(model_ref, revision=revision, allow_patterns=["assets/*"]) + + def _batch_error(reason: str, **details: Any) -> RuntimeError: payload = {"reason": reason, **details} return RuntimeError(f"MiniCPMO45Code2WavBatchError {json.dumps(payload, sort_keys=True)}") @@ -118,6 +132,7 @@ class MiniCPMO45Code2Wav(nn.Module): del prefix self.vllm_config = vllm_config self.model_path = str(vllm_config.model_config.model) + self._model_revision = getattr(vllm_config.model_config, "revision", None) self.backend: BatchedToken2Wav | None = None self._states: dict[str, _RequestState] = {} self._owned_prompt_wavs: dict[str, tuple[str, str]] = {} @@ -126,12 +141,13 @@ class MiniCPMO45Code2Wav(nn.Module): if self._min_batch_size < 1: raise ValueError("MiniCPM-o Code2Wav code2wav_min_batch_size must be >= 1") self._default_prompt_id = str(extra.get("prompt_cache_id", "HT_ref_audio")) - self._default_prompt_wav = str( - extra.get( - "prompt_wav", - Path(self.model_path) / "assets" / "HT_ref_audio.wav", - ) - ) + self._prompt_wav_override = extra.get("prompt_wav") + + @property + def _default_prompt_wav(self) -> str: + if self._prompt_wav_override is not None: + return str(self._prompt_wav_override) + return str(Path(self.model_path) / "assets" / "HT_ref_audio.wav") def _extra_config(self) -> dict[str, Any]: model_config = getattr(self.vllm_config, "model_config", None) @@ -600,6 +616,10 @@ class MiniCPMO45Code2Wav(nn.Module): from stepaudio2.token2wav import Token2wav extra = self._extra_config() + # Hub repo ids only need to become local directories once the vocoder + # assets are actually read; unit tests construct this model with fake + # paths and must not trigger a hub download (#5442). + self.model_path = _resolve_model_dir(self.model_path, self._model_revision) prompt_path = Path(self._default_prompt_wav) if not prompt_path.is_file(): raise FileNotFoundError(f"MiniCPM-o Code2Wav prompt audio not found: {prompt_path}")