vllm-omni/tests/config/test_deploy_profiles.py

"""Tests for named deploy profiles (``vllm_omni/deploy/<profile>/<model_type>.yaml``)."""

from __future__ import annotations

from pathlib import Path
from unittest import mock

import pytest

from vllm_omni.config.stage_config import _DEPLOY_DIR, load_deploy_config, resolve_deploy_yaml
from vllm_omni.entrypoints import utils as entry_utils
from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]

EDGE_QWEN3_TTS = _DEPLOY_DIR / "edge" / "qwen3_tts.yaml"


def test_edge_profile_file_exists():
    assert EDGE_QWEN3_TTS.exists()
    assert "edge" in entry_utils.available_deploy_profiles()
    assert "edge" in entry_utils.available_deploy_profiles("qwen3_tts")


def test_edge_profile_resolves_two_stages_on_one_device():
    cfg = load_deploy_config(EDGE_QWEN3_TTS)
    assert len(cfg.stages) == 2
    assert all(stage.devices == "0" for stage in cfg.stages)
    assert cfg.stages[0].max_num_seqs == 4
    assert cfg.stages[1].max_num_seqs == 4
    assert cfg.async_chunk is True
    # Inherited from the base config: stage 0 keeps its sampling defaults.
    assert cfg.stages[0].default_sampling_params is not None


def test_edge_profile_ramp_parses_and_ends_at_steady_chunk():
    raw = resolve_deploy_yaml(EDGE_QWEN3_TTS)
    extra = raw["connectors"]["connector_of_shared_memory"]["extra"]
    steady = int(extra["codec_chunk_frames"])
    ramp = parse_chunk_ramp(extra, steady=steady)
    assert ramp == [2, 4, 8, 16, 25]
    assert ramp[-1] == steady
    assert extra["decode_batch_max_size"] == 1
    assert extra["codec_left_context_frames"] == 72


def test_deploy_path_for_profile_lookup():
    assert entry_utils.deploy_path_for_profile("qwen3_tts", "edge") == EDGE_QWEN3_TTS
    assert entry_utils.deploy_path_for_profile("qwen3_tts", "no_such_profile") is None
    assert entry_utils.deploy_path_for_profile("no_such_model", "edge") is None
    assert entry_utils.deploy_path_for_profile("qwen3_tts", "") is None


def test_resolve_deploy_profile_path_uses_default_stem():
    with mock.patch.object(entry_utils, "resolve_model_config_path", return_value=str(_DEPLOY_DIR / "qwen3_tts.yaml")):
        assert entry_utils.resolve_deploy_profile_path("Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", "edge") == str(
            EDGE_QWEN3_TTS
        )
        with pytest.raises(FileNotFoundError, match="no_such_profile"):
            entry_utils.resolve_deploy_profile_path("Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", "no_such_profile")


def test_apply_deploy_profile_translates_and_rejects_conflicts():
    with mock.patch.object(entry_utils, "resolve_model_config_path", return_value=str(_DEPLOY_DIR / "qwen3_tts.yaml")):
        args = {"deploy_profile": "edge", "other": 1}
        out = entry_utils.apply_deploy_profile("Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", args)
        assert out["deploy_config"] == str(EDGE_QWEN3_TTS)
        assert "deploy_profile" not in out and out["other"] == 1
        with pytest.raises(ValueError, match="mutually exclusive"):
            entry_utils.apply_deploy_profile("m", {"deploy_profile": "edge", "deploy_config": "x.yaml"})
    # No profile: untouched.
    assert entry_utils.apply_deploy_profile("m", {"deploy_config": "x.yaml"}) == {"deploy_config": "x.yaml"}


def test_modify_stage_config_keeps_relative_base_config_resolvable():
    from tests.helpers.stage_config import modify_stage_config

    tmp_yaml = modify_stage_config(str(EDGE_QWEN3_TTS), updates={"stages": {0: {"max_num_seqs": 1}}})
    raw = resolve_deploy_yaml(Path(tmp_yaml))
    assert len(raw["stages"]) == 2
    cfg = load_deploy_config(Path(tmp_yaml))
    assert cfg.stages[0].max_num_seqs == 1
    assert cfg.stages[1].max_num_seqs == 4


@pytest.mark.core_model
@pytest.mark.cpu
def test_edge_profile_carries_parallel_stage_init():
    from vllm_omni.entrypoints.utils import apply_deploy_profile, profile_orchestrator_defaults

    args = apply_deploy_profile("Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", {"deploy_profile": "edge"})
    assert args["parallel_stage_init"] is True
    assert profile_orchestrator_defaults(args["deploy_config"]) == {"parallel_stage_init": True}
    # explicit values win over the profile default
    args = apply_deploy_profile(
        "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", {"deploy_profile": "edge", "parallel_stage_init": True}
    )
    assert args["parallel_stage_init"] is True
    # unknown block keys are ignored
    assert "init_timeout" not in args