From b22dcc59f7c48b70474abea9c8cca2d6b20b486c Mon Sep 17 00:00:00 2001 From: tzhouam Date: Wed, 9 Sep 2026 10:27:33 +0800 Subject: [PATCH] [Edge][WP1] Add edge deploy profile and deploy_profile selector * vllm_omni/deploy/edge/qwen3_tts.yaml: single-device Qwen3-TTS profile with a ramped codec chunk schedule (2->4->8->16->25 frames) that removes the ~80 ms playback stall of the default 1->25 schedule, decode batch 1, max_num_seqs 4. * deploy_profile="edge" (Omni/AsyncOmni kwarg, --deploy-profile CLI, headless serve) resolves vllm_omni/deploy//.yaml; mutually exclusive with deploy_config; explicit profiles never fall back silently. * tests/helpers/stage_config.py: modify_stage_config anchors a relative base_config before writing the temp YAML so overlay-derived configs resolve. * async_omni_engine.py also carries the init-timeline hooks used by WP2. Co-Authored-By: Claude Fable 5.1 --- tests/config/test_deploy_profiles.py | 83 +++++++++++++++++++++++++++ tests/helpers/stage_config.py | 7 +++ vllm_omni/deploy/edge/qwen3_tts.yaml | 47 +++++++++++++++ vllm_omni/engine/arg_utils.py | 3 + vllm_omni/engine/async_omni_engine.py | 46 +++++++++++++++ vllm_omni/entrypoints/cli/serve.py | 11 ++++ vllm_omni/entrypoints/utils.py | 81 ++++++++++++++++++++++++++ 7 files changed, 278 insertions(+) create mode 100644 tests/config/test_deploy_profiles.py create mode 100644 vllm_omni/deploy/edge/qwen3_tts.yaml diff --git a/tests/config/test_deploy_profiles.py b/tests/config/test_deploy_profiles.py new file mode 100644 index 000000000..c11782280 --- /dev/null +++ b/tests/config/test_deploy_profiles.py @@ -0,0 +1,83 @@ +"""Tests for named deploy profiles (``vllm_omni/deploy//.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 diff --git a/tests/helpers/stage_config.py b/tests/helpers/stage_config.py index fed61937f..539b612d3 100644 --- a/tests/helpers/stage_config.py +++ b/tests/helpers/stage_config.py @@ -241,6 +241,13 @@ def modify_stage_config( # get_chunk_config and get_batch_token_config). int(time.time()) would collide # and the later write would overwrite the earlier YAML on disk. # Keep generated configs outside the repo and delete them when pytest exits. + # The generated YAML lives in a temp dir, so a relative ``base_config`` (e.g. + # ``../qwen3_tts.yaml`` in ``deploy/edge/``) must be anchored to the source + # file's directory before writing (same rule as ``_materialize_ci_overlay``). + base_config = config.get("base_config") + if isinstance(base_config, str) and base_config and not os.path.isabs(base_config): + config["base_config"] = str((path.parent / base_config).resolve()) + output_fd, output_path = tempfile.mkstemp(prefix=f"{path.stem}_", suffix=".yaml") atexit.register(Path(output_path).unlink, missing_ok=True) diff --git a/vllm_omni/deploy/edge/qwen3_tts.yaml b/vllm_omni/deploy/edge/qwen3_tts.yaml new file mode 100644 index 000000000..fbff8b975 --- /dev/null +++ b/vllm_omni/deploy/edge/qwen3_tts.yaml @@ -0,0 +1,47 @@ +# Qwen3-TTS edge profile: single device, small batch, ramped codec chunks. +# +# Select with ``Omni(model=..., deploy_profile="edge")`` or ``vllm serve --omni +# --deploy-profile edge``. Inherits everything from ../qwen3_tts.yaml; only the +# keys below differ. Top-level blocks other than ``stages``/``platforms`` are +# replaced wholesale by the overlay merge, so ``connectors`` is restated. +# +# Why these values (see benchmarks/tts/edge_results and analysis/experiments.md): +# * The default schedule emits 1 frame (80 ms) first and the next 25-frame +# chunk ~160 ms later, which stalls a player that starts at first audio by +# ~80 ms. ``codec_chunk_ramp`` grows chunks 2 -> 4 -> 8 -> 16 -> 25 frames so +# the audio buffered ahead always exceeds the wait for the next chunk. +# * ``decode_batch_max_size: 1`` and ``max_num_seqs: 4`` size the decoder and +# scheduler for 1-4 concurrent utterances instead of 64. +# * Both stages share logical device "0" (one GPU). Memory fractions are +# what the two 0.6B/1.7B stages actually need with these limits, not the +# server defaults; raise stage 1 if the KV-profile log reports OOM. +# * ``enforce_eager`` / ``parallel_stage_init`` are intentionally not set +# here yet; benchmarks/tts/init_profile.py results decide them (WP2). +base_config: ../qwen3_tts.yaml +async_chunk: true + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + codec_streaming: true + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + codec_chunk_frames: 25 + codec_left_context_frames: 72 + # Ramp overrides initial_codec_chunk_frames; last entry == codec_chunk_frames. + codec_chunk_ramp: [2, 4, 8, 16, 25] + decode_batch_max_size: 1 + +stages: + - stage_id: 0 + max_num_seqs: 4 + gpu_memory_utilization: 0.25 + # Small per-step batches keep first-chunk latency low at edge concurrency. + max_num_batched_tokens: 512 + devices: "0" + - stage_id: 1 + max_num_seqs: 4 + gpu_memory_utilization: 0.15 + devices: "0" diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 812f4c3e7..b3d894eb4 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -518,6 +518,9 @@ class OrchestratorArgs: # === Config Files === deploy_config: str | None = None + # Named deploy profile (e.g. ``edge``) resolved to ``deploy//.yaml`` + # by the orchestrator; mutually exclusive with ``deploy_config``. + deploy_profile: str | None = None stage_overrides: str | None = None # raw JSON string; parsed downstream # Optional composable-parallel strategy.yaml; orchestrator reads it, overlays # derived sizing onto merged stages, then drops it before per-stage engine args. diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 4cec16382..97c000c4b 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -14,6 +14,7 @@ import asyncio import concurrent.futures import copy import json +import os import queue import shutil import threading @@ -65,6 +66,7 @@ from vllm_omni.engine.duplex.runtime import ( load_duplex_runtime_extension, validate_duplex_runtime_extension, ) +from vllm_omni.engine.init_timeline import ENV_OUT, InitTimeline, worker_phase from vllm_omni.engine.messages import ( AbortRequestMessage, AbortResultMessage, @@ -198,6 +200,12 @@ class AsyncOmniEngine: self._omni_master_port, ) + # A named deploy profile (``deploy_profile="edge"``) resolves to a concrete + # ``deploy_config`` path before anything reads it. + from vllm_omni.entrypoints.utils import apply_deploy_profile + + apply_deploy_profile(model, kwargs) + # Stage resolution pops deploy_config, so get pipeline-wide settings beforehand. deploy_config_path = kwargs.get("deploy_config") # ``trust_remote_code`` is tri-state (bool | None): ``None`` means "not @@ -260,6 +268,19 @@ class AsyncOmniEngine: self._running_counter = OmniRequestCounter() self._engines_waiting_counter = OmniRequestCounter() + # Init timeline: on when VLLM_OMNI_INIT_TIMELINE is set (JSON dump) or + # log_stats is requested (INFO table only). Must exist before stages + # spawn so children inherit VLLM_OMNI_INIT_TIMELINE_DIR. + timeline_out = os.environ.get(ENV_OUT) + InitTimeline.start( + enabled=bool(timeline_out) or bool(log_stats), + run_dir=( + os.path.join(os.path.dirname(os.path.abspath(timeline_out)), "init_timeline_workers") + if timeline_out + else None + ), + ) + logger.info(f"[AsyncOmniEngine] Launching Orchestrator thread with {self.num_stages} stages") # Launch orchestrator background thread @@ -276,6 +297,7 @@ class AsyncOmniEngine: ) self.orchestrator_thread.start() self._wait_for_orchestrator_init(startup_future, startup_timeout) + self._report_init_timeline(partial=False) self._correlated_rpc_client = CorrelatedRpcClient( self.request_queue.sync_q, self.rpc_output_queue.sync_q, @@ -318,6 +340,10 @@ class AsyncOmniEngine: def _initialize_stages(self, stage_init_timeout: int) -> None: """Initialize stage clients/processors via StageRuntime and assign to self.""" + with worker_phase("stages_init"): + self._initialize_stages_impl(stage_init_timeout) + + def _initialize_stages_impl(self, stage_init_timeout: int) -> None: self._runtime = create_stage_runtime( stage_configs=self.stage_configs, model=self.model, @@ -465,6 +491,25 @@ class AsyncOmniEngine: asyncio.set_event_loop(None) loop.close() + def _report_init_timeline(self, partial: bool) -> None: + """Merge worker phase records and log/dump the init timeline (no-op when disabled).""" + tl = InitTimeline.current() + if tl is None or not tl.enabled: + return + try: + tl.finish() + tl.ingest_workers() + status = "PARTIAL (startup timed out)" if partial else "complete" + logger.info("[AsyncOmniEngine] init timeline %s\n%s", status, tl.format_table()) + if not partial: + logger.info("[AsyncOmniEngine] init timeline suggested timeouts: %s", tl.suggest_timeouts()) + out = os.environ.get(ENV_OUT) + if out: + tl.dump(out) + logger.info("[AsyncOmniEngine] init timeline written to %s", out) + except Exception as exc: # telemetry must never break start-up + logger.warning("[AsyncOmniEngine] init timeline reporting failed: %s", exc) + def _wait_for_orchestrator_init(self, startup_future: concurrent.futures.Future, startup_timeout: int) -> None: """ Wait for orchestrator startup future to return ready. Raises exception on any failures to the init process. @@ -480,6 +525,7 @@ class AsyncOmniEngine: "--stage-init-timeout values.", startup_timeout, ) + self._report_init_timeline(partial=True) self._try_shutdown("[AsyncOmniEngine] Failed to cleanup after orchestrator startup timeout") raise TimeoutError(f"Orchestrator did not become ready within {startup_timeout}s") try: diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index ea56acb1a..e07d0eb9b 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -275,6 +275,13 @@ class OmniServeCommand(CLISubcommand): default=None, help="Path to a deploy config YAML (new format with stages/engine_args).", ) + omni_config_group.add_argument( + "--deploy-profile", + type=str, + default=None, + help="Named deploy profile resolved to vllm_omni/deploy//.yaml " + "(e.g. 'edge'). Mutually exclusive with --deploy-config.", + ) omni_config_group.add_argument( "--strategy-config", type=str, @@ -1002,6 +1009,10 @@ def run_headless(args: TrackingNamespace) -> None: # standard launches resolve to the same per-stage device layout. stage_overrides = parse_stage_overrides(args_dict.get("stage_overrides")) + from vllm_omni.entrypoints.utils import apply_deploy_profile + + apply_deploy_profile(model, args_dict) + config_path, stage_configs, _ = load_and_resolve_stage_configs( model, args_dict, diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 4b91f752a..2bfd755e1 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -261,6 +261,87 @@ def _registry_default_deploy_path(model: str) -> str | None: return None +def deploy_path_for_profile(model_type: str, profile: str) -> Path | None: + """Return ``deploy//.yaml`` if it exists (pure path lookup). + + Looks first under the platform's default stage-config directory, then under + the packaged ``vllm_omni/deploy`` directory. ``model_type`` is the deploy + YAML stem (the same key ``resolve_model_config_path`` resolves to). + """ + if not profile or not model_type: + return None + file_name = f"{model_type}.yaml" + candidates = [] + try: + candidates.append(PROJECT_ROOT / current_omni_platform.get_default_stage_config_path() / profile / file_name) + except NotImplementedError: + pass + candidates.append(_DEPLOY_DIR / profile / file_name) + for candidate in candidates: + if candidate.exists(): + return candidate + return None + + +def available_deploy_profiles(model_type: str | None = None) -> list[str]: + """List profile directories under ``vllm_omni/deploy`` (optionally only those defining ``model_type``).""" + profiles = [] + if not _DEPLOY_DIR.exists(): + return profiles + for entry in sorted(_DEPLOY_DIR.iterdir()): + if not entry.is_dir() or entry.name.startswith(("_", ".")): + continue + if model_type is None or (entry / f"{model_type}.yaml").exists(): + profiles.append(entry.name) + return profiles + + +def resolve_deploy_profile_path(model: str, profile: str) -> str: + """Resolve a named deploy profile (e.g. ``edge``) for ``model`` to a YAML path. + + The model's default deploy YAML stem (from ``resolve_model_config_path``) + is the key under ``deploy//``. An explicit profile never silently + falls back to the default deploy config. + + Raises: + FileNotFoundError: if the profile does not define a YAML for the model. + """ + default_path = resolve_model_config_path(model) + model_type = Path(default_path).stem if default_path else None + if model_type is None: + raise FileNotFoundError(f"Cannot resolve a deploy profile for {model!r}: no default deploy config found.") + path = deploy_path_for_profile(model_type, profile) + if path is None: + raise FileNotFoundError( + f"Deploy profile {profile!r} has no config for model type {model_type!r}. " + f"Profiles defining it: {available_deploy_profiles(model_type)}; " + f"all profiles: {available_deploy_profiles()}." + ) + logger.info("[deploy_profile] %s -> %s", profile, path) + return str(path) + + +def apply_deploy_profile(model: str, args_dict: dict[str, Any]) -> dict[str, Any]: + """Translate ``deploy_profile`` into ``deploy_config`` in-place (used by CLI and engine).""" + profile = args_dict.pop("deploy_profile", None) + if not profile: + return args_dict + if args_dict.get("deploy_config"): + raise ValueError("`deploy_profile` and `deploy_config` are mutually exclusive; pass only one.") + if profile == "auto": + from vllm_omni.edge.auto_profile import materialize_auto_deploy + + default_path = resolve_model_config_path(model) + if default_path is None: + raise FileNotFoundError(f"Cannot resolve a deploy config for {model!r} (needed by deploy_profile='auto').") + edge_path = deploy_path_for_profile(Path(default_path).stem, "edge") + base = str(edge_path) if edge_path is not None else default_path + args_dict["deploy_config"], _ = materialize_auto_deploy(model, base) + return args_dict + args_dict["deploy_config"] = resolve_deploy_profile_path(model, profile) + return args_dict + + def resolve_model_config_path(model: str) -> str | None: """Resolve the stage/deploy config file path from the model name.