diff --git a/docs/design/qwen3_omni_tts_performance_optimization.md b/docs/design/qwen3_omni_tts_performance_optimization.md index 2f18a1b1b..e4548a17c 100644 --- a/docs/design/qwen3_omni_tts_performance_optimization.md +++ b/docs/design/qwen3_omni_tts_performance_optimization.md @@ -411,6 +411,8 @@ Notes: - `runtime.max_batch_size` controls stage-level batching. - Thinker/Talker commonly use `enforce_eager: false` for CUDA Graph paths. - Code2Wav often remains eager (`enforce_eager: true`) depending on runtime behavior. +- Qwen3-Omni defaults `VLLM_USE_FLASHINFER_MOE_FP16=0`. The Triton has been more stable & faster + than the FlashInfer CUTLASS unquantized MoE backend on recent vLLM rebases. #### 2) Enable async chunk diff --git a/docs/mkdocs/hooks/generate_argparse.py b/docs/mkdocs/hooks/generate_argparse.py index e3cfb1b6a..0d18ad6c9 100644 --- a/docs/mkdocs/hooks/generate_argparse.py +++ b/docs/mkdocs/hooks/generate_argparse.py @@ -119,9 +119,10 @@ def extract_omni_serve_subparser_init(): exec_globals = { "_FlexibleArgumentParser": _FlexibleArgumentParser, "FlexibleArgumentParser": _FlexibleArgumentParser, + "TrackingArgumentParser": _FlexibleArgumentParser, # only needs add_argument, not tracking + "TrackingNamespace": None, # type annotation only "make_arg_parser": lambda parser: parser, # no-op for doc "_ensure_vllm_platform": lambda: None, # no-op for doc - "nullify_stage_engine_defaults": lambda parser: None, # no-op for doc "VLLM_SUBCMD_PARSER_EPILOG": "", "logger": logger, "DummySubparsers": DummySubparsers, diff --git a/docs/user_guide/examples/online_serving/qwen3_omni.md b/docs/user_guide/examples/online_serving/qwen3_omni.md index 413c80b7e..1973b198b 100644 --- a/docs/user_guide/examples/online_serving/qwen3_omni.md +++ b/docs/user_guide/examples/online_serving/qwen3_omni.md @@ -17,6 +17,9 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 The default deployment configuration situated at `vllm_omni/deploy/qwen3_omni_moe.yaml` is resolved and loaded automatically via the model registry, obviating the necessity for the `--deploy-config` flag in standard deployment topologies. +The bundled Qwen3-Omni setup defaults `VLLM_USE_FLASHINFER_MOE_FP16=0`. This keeps the Thinker & Talker on vLLM's +Triton unquantized MoE path and avoids the performance regression observed with the FlashInfer CUTLASS unquantized MoE +backend. Asynchronous chunk streaming is **enabled by default** within the bundled configuration. To explicitly utilize a custom deployment YAML, specify the configuration path: @@ -72,6 +75,12 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \ --stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}' ``` +To experiment with the FlashInfer FP16 MoE path, set `VLLM_USE_FLASHINFER_MOE_FP16=1` before launching the server: +```bash +VLLM_USE_FLASHINFER_MOE_FP16=1 \ +vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 +``` + For the stage-based CLI, you usually do **not** need `--stage-overrides` for that kind of change. Since each command launches one stage, just pass the knob directly on that stage command: diff --git a/examples/offline_inference/bagel/README.md b/examples/offline_inference/bagel/README.md index 9c2e31ff3..f2685d56f 100644 --- a/examples/offline_inference/bagel/README.md +++ b/examples/offline_inference/bagel/README.md @@ -195,6 +195,101 @@ stages: devices: "0,1" ``` +### VAE Patch Parallelism + +[VAE Patch Parallelism](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism/vae_patch_parallel.html) splits Bagel VAE **decode/encode** tiles across multiple GPUs on the **DiT stage**, reducing **per-GPU peak memory during VAE decode**. Use it when high-resolution `text2img` or `img2img` hits VAE OOM or large decode spikes. + +**Bagel-specific notes:** + +- Implemented in `BagelPipeline` via `DistributedAutoEncoder` (DiT stage only). +- **Single-stage** is the simplest path: one DiT process with TP + VAE patch parallel. +- **Two-stage**: enable on **stage 1 (DiT)** only; stage 0 (Thinker) keeps encoder-only `VAEEncoder` and does not use VAE patch parallel. +- You need a DiT `world_size` ≥ `vae_patch_parallel_size` (typically `tensor_parallel_size=2` on that stage). VAE PP reuses the DiT process group; it is not a standalone second-GPU VAE worker. + +**Single-stage via deploy YAML** (recommended for `end2end.py`): + +```yaml +pipeline: bagel_single_stage +async_chunk: false + +stages: + - stage_id: 0 + max_num_batched_tokens: 32768 + max_num_seqs: 1 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + devices: "0,1" + vae_use_tiling: true + parallel_config: + tensor_parallel_size: 2 + vae_patch_parallel_size: 2 + default_sampling_params: + seed: 52 +``` + +```bash +cd examples/offline_inference/bagel + +CUDA_VISIBLE_DEVICES=0,1 python end2end.py \ + --model /path/to/BAGEL-7B-MoT \ + --deploy-config /path/to/bagel_single_stage_vae_pp.yaml \ + --modality text2img \ + --prompts "A cute cat" \ + --steps 10 \ + --output ./out_vae_pp +``` + +**Single-stage via `Omni` kwargs** (same flags as online serving): + +```python +from vllm_omni.entrypoints.omni import Omni + +omni = Omni( + model="ByteDance-Seed/BAGEL-7B-MoT", + deploy_config="vllm_omni/deploy/bagel_single_stage.yaml", + tensor_parallel_size=2, + vae_patch_parallel_size=2, + vae_use_tiling=True, +) +# Then call omni.generate(...) as in end2end.py +``` + +**Two-stage (VAE PP on DiT only):** + +```yaml +stages: + - stage_id: 0 + devices: "0" + # AR Thinker — no vae_patch_parallel here + + - stage_id: 1 + devices: "0,1" + vae_use_tiling: true + parallel_config: + tensor_parallel_size: 2 + vae_patch_parallel_size: 2 +``` + +```bash +python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ + --deploy-config /path/to/bagel_vae_pp.yaml \ + --modality text2img \ + --prompts "A cute cat" +``` + +**Startup log checks:** + +```text +INFO ... vae_patch_parallel_size=2 requires vae_use_tiling; automatically enabling it. +``` + +| Setting | Role | +| :------ | :--- | +| `parallel_config.tensor_parallel_size` | DiT world size / TP (must be ≥ `vae_patch_parallel_size`) | +| `parallel_config.vae_patch_parallel_size` | Number of ranks for distributed VAE tiles (`1` = off) | +| `vae_use_tiling` | Enable spatial tiling (auto-enabled when `vae_patch_parallel_size > 1`) | + #### Hybrid Sharded Data Parallel (HSDP) For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage. diff --git a/examples/offline_inference/bagel/end2end.py b/examples/offline_inference/bagel/end2end.py index 3ce7ba8d5..088687253 100644 --- a/examples/offline_inference/bagel/end2end.py +++ b/examples/offline_inference/bagel/end2end.py @@ -121,9 +121,6 @@ def parse_args(): help="Temperature for text generation sampling (default: 0.3).", ) - from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults - - nullify_stage_engine_defaults(parser) args = parser.parse_args() return args @@ -154,6 +151,7 @@ def main(): from vllm_omni.entrypoints.omni import Omni omni_kwargs = vars(args).copy() + omni_kwargs.pop("model", None) deploy_config = args.deploy_config if args.think and deploy_config is None: deploy_config = "vllm_omni/deploy/bagel_think.yaml" @@ -164,9 +162,7 @@ def main(): if args.quantization: omni_kwargs["quantization_config"] = args.quantization - # Override CLI --model with the derived model_name. - omni_kwargs["model"] = model_name - omni = Omni(**omni_kwargs) + omni = Omni(model=model_name, **omni_kwargs) formatted_prompts = [] for p in prompts: diff --git a/examples/offline_inference/dynin_omni/end2end.py b/examples/offline_inference/dynin_omni/end2end.py index beba6be0c..66047934d 100644 --- a/examples/offline_inference/dynin_omni/end2end.py +++ b/examples/offline_inference/dynin_omni/end2end.py @@ -18,8 +18,6 @@ import numpy as np import torch from PIL import Image -from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults - TASK_CHOICES = ("t2t", "t2i", "t2s", "i2i", "i2t", "s2t", "v2t") TASK_DEFAULT_RUNTIME = { @@ -972,8 +970,6 @@ def parse_args(repo_root: Path) -> argparse.Namespace: parser.add_argument("--vq-model-audio-local-files-only", action=argparse.BooleanOptionalAction, default=None) parser.add_argument("--disable-hf-xet", action=argparse.BooleanOptionalAction, default=True) - - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/offline_inference/hunyuan_image3/end2end.py b/examples/offline_inference/hunyuan_image3/end2end.py index e8cd8a36f..3ba69c2e2 100644 --- a/examples/offline_inference/hunyuan_image3/end2end.py +++ b/examples/offline_inference/hunyuan_image3/end2end.py @@ -107,9 +107,6 @@ def parse_args(): ), ) - from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults - - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py index a742b535f..ca87b9e9a 100644 --- a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py +++ b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py @@ -19,7 +19,6 @@ from vllm import SamplingParams from vllm.multimodal.image import convert_image_mode from vllm_omni import Omni -from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults DEFAULT_SYSTEM = "You are a helpful assistant." DEFAULT_QUESTION = "Please summarize the content of this image." @@ -49,7 +48,6 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable diffusion pipeline profiler to display stage durations.", ) - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py index bd1282a21..a4c41fee1 100644 --- a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py +++ b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py @@ -29,7 +29,6 @@ from PIL import Image from vllm.sampling_params import SamplingParams from vllm_omni import Omni -from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") logger = logging.getLogger(__name__) @@ -118,7 +117,6 @@ def parse_args() -> argparse.Namespace: ) p.add_argument("--out", type=str, default="output.png", help="Path to save the generated image.") p.add_argument("--trust-remote-code", action="store_true", help="Trust remote code when loading the model.") - nullify_stage_engine_defaults(p) args = p.parse_args() if not args.prompt: args.prompt = ["A stylish woman with sunglasses riding a motorcycle in NYC."] diff --git a/examples/offline_inference/mimo_audio/end2end.py b/examples/offline_inference/mimo_audio/end2end.py index d3728dcc5..2fe46c972 100644 --- a/examples/offline_inference/mimo_audio/end2end.py +++ b/examples/offline_inference/mimo_audio/end2end.py @@ -21,11 +21,10 @@ from message_convert import ( to_prompt, ) from vllm import SamplingParams -from vllm.utils.argparse_utils import FlexibleArgumentParser -from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniTokensPrompt +from vllm_omni.utils.tracking_parser import TrackingArgumentParser SEED = 42 @@ -343,7 +342,7 @@ def main(args): def parse_args(): - parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models") + parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models") parser.add_argument( "--model-name", "-m", @@ -439,7 +438,6 @@ def parse_args(): "vllm_omni/deploy/mimo_audio.yaml based on the HF model_type.", ) - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/offline_inference/ming_flash_omni/end2end.py b/examples/offline_inference/ming_flash_omni/end2end.py index de59d1b66..9fe907ae5 100644 --- a/examples/offline_inference/ming_flash_omni/end2end.py +++ b/examples/offline_inference/ming_flash_omni/end2end.py @@ -17,14 +17,13 @@ from vllm.assets.image import ImageAsset from vllm.assets.video import VideoAsset, video_to_ndarrays from vllm.multimodal.image import convert_image_mode from vllm.multimodal.media.audio import load_audio -from vllm.utils.argparse_utils import FlexibleArgumentParser import vllm_omni -from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults from vllm_omni.entrypoints.omni import Omni # Imports the processor also registers itself from vllm_omni.transformers_utils.processors.ming import MingFlashOmniProcessor # noqa: F401 +from vllm_omni.utils.tracking_parser import TrackingArgumentParser SEED = 42 MODEL_NAME = "Jonathan1909/Ming-flash-omni-2.0" @@ -397,7 +396,7 @@ def main(args): def parse_args(): - parser = FlexibleArgumentParser(description="Ming-flash-omni 2.0 offline inference example") + parser = TrackingArgumentParser(description="Ming-flash-omni 2.0 offline inference example") parser.add_argument( "--query-type", "-q", @@ -496,7 +495,6 @@ def parse_args(): help="Output directory for results.", ) - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py index ad118b2b1..8290d27b5 100644 --- a/examples/offline_inference/qwen2_5_omni/end2end.py +++ b/examples/offline_inference/qwen2_5_omni/end2end.py @@ -19,10 +19,9 @@ from vllm.assets.video import VideoAsset, video_to_ndarrays from vllm.multimodal.image import convert_image_mode from vllm.multimodal.media.audio import load_audio from vllm.sampling_params import SamplingParams -from vllm.utils.argparse_utils import FlexibleArgumentParser -from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults from vllm_omni.entrypoints.omni import Omni +from vllm_omni.utils.tracking_parser import TrackingArgumentParser SEED = 42 @@ -432,7 +431,7 @@ def main(args): def parse_args(): - parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models") + parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models") parser.add_argument( "--model", type=str, @@ -554,7 +553,6 @@ def parse_args(): default=False, help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.", ) - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/offline_inference/qwen3_omni/end2end.py b/examples/offline_inference/qwen3_omni/end2end.py index b10f09b58..ae719e1c3 100644 --- a/examples/offline_inference/qwen3_omni/end2end.py +++ b/examples/offline_inference/qwen3_omni/end2end.py @@ -19,10 +19,9 @@ from vllm.assets.image import ImageAsset from vllm.assets.video import VideoAsset, video_to_ndarrays from vllm.multimodal.image import convert_image_mode from vllm.multimodal.media.audio import load_audio -from vllm.utils.argparse_utils import FlexibleArgumentParser -from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults from vllm_omni.entrypoints.omni import Omni +from vllm_omni.utils.tracking_parser import TrackingArgumentParser SEED = 42 @@ -422,7 +421,7 @@ def main(args): def parse_args(): - parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models") + parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models") parser.add_argument( "--model", type=str, @@ -561,7 +560,6 @@ def parse_args(): help="Model dtype (auto, half, float16, bfloat16, float, float32).", ) - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/offline_inference/qwen3_omni/end2end_async_chunk.py b/examples/offline_inference/qwen3_omni/end2end_async_chunk.py index a5bb6c33d..7d051ac4a 100644 --- a/examples/offline_inference/qwen3_omni/end2end_async_chunk.py +++ b/examples/offline_inference/qwen3_omni/end2end_async_chunk.py @@ -39,10 +39,9 @@ from vllm.assets.image import ImageAsset from vllm.assets.video import VideoAsset, video_to_ndarrays from vllm.multimodal.image import convert_image_mode from vllm.multimodal.media.audio import load_audio -from vllm.utils.argparse_utils import FlexibleArgumentParser -from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.utils.tracking_parser import TrackingArgumentParser logger = logging.getLogger(__name__) @@ -383,7 +382,9 @@ async def run_all(args): print(f"[Info] Creating AsyncOmni with deploy_config={args.deploy_config}") async_omni = None try: - async_omni = AsyncOmni(**vars(args)) + # ``from_cli_args`` forwards only explicitly-passed CLI args so + # argparse defaults do not silently override deploy YAML values. + async_omni = AsyncOmni.from_cli_args(args, model=args.model) # Use default sampling params from stage config (they are pre-configured # in the YAML for each stage). @@ -453,7 +454,7 @@ async def run_all(args): def parse_args(): - parser = FlexibleArgumentParser( + parser = TrackingArgumentParser( description=( "Offline inference with async_chunk enabled via AsyncOmni. " "Downstream stages start before upstream stages finish, " @@ -589,7 +590,6 @@ def parse_args(): default=16000, help="Sampling rate for audio loading.", ) - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/offline_inference/sensenova_u1/end2end.py b/examples/offline_inference/sensenova_u1/end2end.py index ebb0d7f24..0716c74d5 100644 --- a/examples/offline_inference/sensenova_u1/end2end.py +++ b/examples/offline_inference/sensenova_u1/end2end.py @@ -29,7 +29,6 @@ import os from PIL import Image -from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -168,7 +167,6 @@ def parse_args(): help="Enable cache-dit summary logging after diffusion forward passes.", ) - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index 443b21acc..b65ffa0da 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -307,9 +307,6 @@ def parse_args() -> argparse.Namespace: help="Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full).", ) current_omni_platform.pre_register_and_update(parser) - from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults - - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py b/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py index 8a3309ca7..e5c266f7a 100644 --- a/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py +++ b/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py @@ -10,7 +10,6 @@ import soundfile as sf from vllm import SamplingParams from vllm.multimodal.media.audio import load_audio -from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults from vllm_omni.entrypoints.omni import Omni from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token @@ -65,7 +64,6 @@ def run_e2e(): required=True, help="Path to tokenizer directory (e.g., /CosyVoice-BlankEN).", ) - nullify_stage_engine_defaults(parser) args = parser.parse_args() # Ensure tokenizer directory exists if not os.path.exists(args.tokenizer): diff --git a/examples/offline_inference/text_to_speech/fish_speech/end2end.py b/examples/offline_inference/text_to_speech/fish_speech/end2end.py index e63b844e7..4d6855246 100644 --- a/examples/offline_inference/text_to_speech/fish_speech/end2end.py +++ b/examples/offline_inference/text_to_speech/fish_speech/end2end.py @@ -26,16 +26,14 @@ import torch os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" -from vllm.utils.argparse_utils import FlexibleArgumentParser - from vllm_omni import AsyncOmni, Omni -from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults from vllm_omni.model_executor.models.fish_speech.dac_utils import DAC_HOP_LENGTH, DAC_SAMPLE_RATE from vllm_omni.model_executor.models.fish_speech.prompt_utils import ( build_fish_text_only_prompt_ids, estimate_fish_voice_clone_prompt_len_from_normalized, normalize_fish_voice_clone_texts, ) +from vllm_omni.utils.tracking_parser import TrackingArgumentParser logger = logging.getLogger(__name__) @@ -211,7 +209,7 @@ async def main_streaming(args): def parse_args(): - parser = FlexibleArgumentParser( + parser = TrackingArgumentParser( description="Fish Speech S2 Pro offline inference with vLLM Omni", ) parser.add_argument( @@ -267,7 +265,6 @@ def parse_args(): default=False, help="Stream audio chunks as they arrive via AsyncOmni.", ) - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py b/examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py index 62a346621..2d6d078bb 100644 --- a/examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py +++ b/examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py @@ -30,10 +30,9 @@ from pathlib import Path import numpy as np import soundfile as sf import torch -from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm_omni import Omni -from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults +from vllm_omni.utils.tracking_parser import TrackingArgumentParser SAMPLE_RATE = 24_000 DEFAULT_TEXTS = ( @@ -43,7 +42,7 @@ DEFAULT_TEXTS = ( def parse_args(): - parser = FlexibleArgumentParser(description="Offline higgs-audio v2 inference") + parser = TrackingArgumentParser(description="Offline higgs-audio v2 inference") parser.add_argument( "--model", type=str, @@ -88,7 +87,6 @@ def parse_args(): default=None, help="Transcript of the reference clip. Required when --ref-audio is set.", ) - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/offline_inference/text_to_speech/ming_flash_omni_tts/end2end.py b/examples/offline_inference/text_to_speech/ming_flash_omni_tts/end2end.py index 18fdc9d61..58a3fa5ef 100644 --- a/examples/offline_inference/text_to_speech/ming_flash_omni_tts/end2end.py +++ b/examples/offline_inference/text_to_speech/ming_flash_omni_tts/end2end.py @@ -5,11 +5,11 @@ from typing import Any import soundfile as sf import torch -from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni.utils.tracking_parser import TrackingArgumentParser os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" -from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniTokensPrompt from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import ( @@ -68,7 +68,7 @@ def save_audio(mm: dict[str, Any], output_path: str) -> None: def parse_args(): - parser = FlexibleArgumentParser(description="Ming-flash-omni standalone talker offline e2e example") + parser = TrackingArgumentParser(description="Ming-flash-omni standalone talker offline e2e example") parser.add_argument("--model", type=str, default=MODEL_NAME, help="Model name or local path.") parser.add_argument( "--deploy-config", @@ -89,7 +89,6 @@ def parse_args(): parser.add_argument("--init-timeout", type=int, default=600, help="Engine init timeout in seconds.") parser.add_argument("--stage-init-timeout", type=int, default=300, help="Single stage init timeout in seconds.") - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/offline_inference/text_to_speech/moss_tts_nano/end2end.py b/examples/offline_inference/text_to_speech/moss_tts_nano/end2end.py index 3a27ad04d..91291743e 100644 --- a/examples/offline_inference/text_to_speech/moss_tts_nano/end2end.py +++ b/examples/offline_inference/text_to_speech/moss_tts_nano/end2end.py @@ -38,7 +38,8 @@ from pathlib import Path import soundfile as sf import torch from vllm import SamplingParams -from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni.utils.tracking_parser import TrackingArgumentParser # Prevent multiprocessing from re-importing CUDA in the wrong context. os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") @@ -155,7 +156,7 @@ def main(args) -> None: def parse_args(): - parser = FlexibleArgumentParser(description="MOSS-TTS-Nano offline inference") + parser = TrackingArgumentParser(description="MOSS-TTS-Nano offline inference") parser.add_argument("--text", default="Hello, this is MOSS-TTS-Nano speaking.", help="Text to synthesize.") parser.add_argument( "--ref-audio", diff --git a/examples/offline_inference/text_to_speech/omnivoice/end2end.py b/examples/offline_inference/text_to_speech/omnivoice/end2end.py index f3c23e60b..619f8160b 100644 --- a/examples/offline_inference/text_to_speech/omnivoice/end2end.py +++ b/examples/offline_inference/text_to_speech/omnivoice/end2end.py @@ -21,7 +21,6 @@ import os import numpy as np import soundfile as sf -from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -86,7 +85,6 @@ def run_e2e(): default=None, help="Random seed for generation", ) - nullify_stage_engine_defaults(parser) args = parser.parse_args() if not os.path.exists(args.stage_config): diff --git a/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py b/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py index d9f96838a..8150338d5 100644 --- a/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py +++ b/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py @@ -15,10 +15,8 @@ import torch os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" -from vllm.utils.argparse_utils import FlexibleArgumentParser - from vllm_omni import AsyncOmni, Omni -from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults +from vllm_omni.utils.tracking_parser import TrackingArgumentParser logger = logging.getLogger(__name__) @@ -435,7 +433,7 @@ async def main_streaming(args): def parse_args(): - parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models") + parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models") parser.add_argument( "--query-type", "-q", @@ -548,7 +546,6 @@ def parse_args(): help="Number of prompts per batch (default: 1, sequential).", ) - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/offline_inference/text_to_speech/voxcpm2/end2end.py b/examples/offline_inference/text_to_speech/voxcpm2/end2end.py index 7d77d8173..feba91b7e 100644 --- a/examples/offline_inference/text_to_speech/voxcpm2/end2end.py +++ b/examples/offline_inference/text_to_speech/voxcpm2/end2end.py @@ -12,17 +12,16 @@ from pathlib import Path import soundfile as sf import torch -from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm_omni import Omni -from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults +from vllm_omni.utils.tracking_parser import TrackingArgumentParser REPO_ROOT = Path(__file__).resolve().parents[4] SAMPLE_RATE = 48_000 def parse_args(): - parser = FlexibleArgumentParser(description="Offline VoxCPM2 native AR inference") + parser = TrackingArgumentParser(description="Offline VoxCPM2 native AR inference") parser.add_argument( "--model", type=str, @@ -60,7 +59,6 @@ def parse_args(): default=None, help="Optional transcript of --ref-audio (enables continuation mode).", ) - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/offline_inference/text_to_speech/voxtral_tts/end2end.py b/examples/offline_inference/text_to_speech/voxtral_tts/end2end.py index dbcf42da5..c740197b5 100644 --- a/examples/offline_inference/text_to_speech/voxtral_tts/end2end.py +++ b/examples/offline_inference/text_to_speech/voxtral_tts/end2end.py @@ -27,11 +27,10 @@ except ImportError: "pip install -e ." ) from vllm import SamplingParams -from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm_omni import AsyncOmni -from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults from vllm_omni.entrypoints.omni import Omni +from vllm_omni.utils.tracking_parser import TrackingArgumentParser logger = logging.getLogger(__name__) @@ -230,7 +229,7 @@ def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir def parse_args() -> Namespace: - parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with Voxtral TTS") + parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with Voxtral TTS") parser.add_argument( "--model", type=str, @@ -313,7 +312,6 @@ def parse_args() -> Namespace: default=None, help="Quantization method (e.g. 'fp8'). Applied to the language model only.", ) - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/online_serving/bagel/README.md b/examples/online_serving/bagel/README.md index ea8abe43f..b7a54da35 100644 --- a/examples/online_serving/bagel/README.md +++ b/examples/online_serving/bagel/README.md @@ -68,6 +68,67 @@ vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --tensor-parallel-size Or set `tensor_parallel_size` per stage in a custom deploy YAML. +### VAE Patch Parallelism + +[VAE Patch Parallelism](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism/vae_patch_parallel.html) distributes Bagel VAE **decode/encode** across multiple GPUs by splitting latent tiles. It lowers **per-GPU peak memory during VAE decode**, which helps high-resolution `text2img` / `img2img` when VAE becomes a bottleneck. + +**Scope for Bagel:** + +| Topology | VAE patch parallel | +| :------- | :----------------- | +| **Single-stage** (DiT only) | Supported on stage 0 (`BagelPipeline` + `DistributedAutoEncoder`) | +| **Two-stage** | Supported on **stage 1 (DiT)** only; stage 0 (Thinker) uses encoder-only VAE and is unrelated | + +**Requirements:** + +- `vae_patch_parallel_size > 1` and a distributed VAE (`DistributedAutoEncoder` on the DiT pipeline). +- The DiT process group must have at least `vae_patch_parallel_size` ranks. In practice this means the diffusion stage `world_size` must be ≥ 2 (commonly `tensor_parallel_size=2` on that stage). +- `vae_use_tiling` must be enabled. If you set `vae_patch_parallel_size > 1` and omit tiling, the registry auto-enables `vae_use_tiling` at startup. + +VAE patch parallel **reuses the DiT process group** (`dit_group`); it does not create a separate VAE-only worker pool. It is not a substitute for single-GPU VAE tiling (`vae_pp=1`). + +**Online serving (single-stage, 2 GPUs):** + +```bash +CUDA_VISIBLE_DEVICES=0,1 vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \ + --deploy-config vllm_omni/deploy/bagel_single_stage.yaml \ + --tensor-parallel-size 2 \ + --vae-patch-parallel-size 2 \ + --vae-use-tiling +``` + +**Online serving (two-stage, VAE PP on DiT stage 1):** use a custom deploy YAML, for example: + +```yaml +stages: + - stage_id: 0 + devices: "0" + # Thinker (AR) — no VAE patch parallel here + + - stage_id: 1 + devices: "0,1" + vae_use_tiling: true + parallel_config: + tensor_parallel_size: 2 + vae_patch_parallel_size: 2 +``` + +```bash +vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \ + --deploy-config /path/to/bagel_vae_pp.yaml +``` + +**Verify it is active** (check server logs at startup): + +```text +INFO ... vae_patch_parallel_size=2 requires vae_use_tiling; automatically enabling it. +``` + +| CLI flag | Default | Description | +| :------- | :------ | :---------- | +| `--vae-patch-parallel-size` | `1` | Number of DiT ranks used for VAE tile parallelism. Set to `2` or higher to enable. Should be ≤ DiT process group size (typically match `--tensor-parallel-size` on the diffusion stage). | +| `--vae-use-tiling` | off | Enable VAE spatial tiling. Required for VAE patch parallel (auto-enabled when `vae_patch_parallel_size > 1`). | + #### Hybrid Sharded Data Parallel (HSDP) For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage. @@ -323,6 +384,8 @@ python openai_chat_client.py \ | `stages[].gpu_memory_utilization` | per-stage | Fraction of GPU memory to use | | `stages[].enforce_eager` | per-stage | Disable CUDA graphs | | `stages[].tensor_parallel_size` | per-stage | TP degree for this stage | +| `stages[].parallel_config.vae_patch_parallel_size` | per-stage (DiT) | VAE tile parallelism degree (DiT stage only) | +| `stages[].vae_use_tiling` | per-stage (DiT) | Enable VAE tiling (required for VAE patch parallel) | | `connectors` | top-level | Define available connector instances (SHM, Mooncake) | | `platforms` | top-level | Platform-specific overrides (e.g. `xpu`) | diff --git a/examples/online_serving/mimo_audio/openai_chat_completion_client_for_multimodal_generation.py b/examples/online_serving/mimo_audio/openai_chat_completion_client_for_multimodal_generation.py index 4053f810d..3030fa3e3 100644 --- a/examples/online_serving/mimo_audio/openai_chat_completion_client_for_multimodal_generation.py +++ b/examples/online_serving/mimo_audio/openai_chat_completion_client_for_multimodal_generation.py @@ -12,7 +12,8 @@ from typing import Any import requests from openai import OpenAI from vllm.assets.audio import AudioAsset -from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni.utils.tracking_parser import TrackingArgumentParser # Modify OpenAI's API key and API base to use vLLM's API server. openai_api_key = "EMPTY" @@ -493,7 +494,7 @@ def run_multimodal_generation(args) -> None: def parse_args(): - parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models") + parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models") parser.add_argument( "--query-type", "-q", diff --git a/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py b/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py index 2609945c3..4f8fce8c9 100644 --- a/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py +++ b/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py @@ -6,7 +6,8 @@ from typing import NamedTuple import requests from openai import OpenAI from vllm.assets.audio import AudioAsset -from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni.utils.tracking_parser import TrackingArgumentParser SEED = 42 @@ -478,7 +479,7 @@ def run_multimodal_generation(args, client: OpenAI) -> None: def parse_args(): - parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models") + parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models") parser.add_argument( "--query-type", "-q", diff --git a/tests/comfyui/test_comfyui_integration.py b/tests/comfyui/test_comfyui_integration.py index 80e86d824..f532c350f 100644 --- a/tests/comfyui/test_comfyui_integration.py +++ b/tests/comfyui/test_comfyui_integration.py @@ -30,12 +30,12 @@ from PIL import Image from pytest_mock import MockerFixture from vllm import SamplingParams from vllm.outputs import CompletionOutput, RequestOutput -from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm_omni.entrypoints.async_omni import AsyncOmni as RealAsyncOmni from vllm_omni.entrypoints.cli.serve import OmniServeCommand from vllm_omni.inputs.data import OmniSamplingParams from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.utils.tracking_parser import TrackingArgumentParser pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -444,7 +444,7 @@ def mock_async_omni( @pytest.fixture def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni): """Set up a API server in background process from command line with parametrized model name and mocked AsyncOmni.""" - parser = FlexibleArgumentParser() + parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="command") cmd = OmniServeCommand() cmd.subparser_init(subparsers) diff --git a/tests/e2e/online_serving/test_bagel_expansion.py b/tests/e2e/online_serving/test_bagel_expansion.py index 2feaa3f07..a8f52c311 100644 --- a/tests/e2e/online_serving/test_bagel_expansion.py +++ b/tests/e2e/online_serving/test_bagel_expansion.py @@ -9,6 +9,8 @@ Coverage: - Ulysses-SP - Ring-Attention - Layerwise Offloading +- Hybrid Sharded Data Parallel +- Tensor Parallelism + VAE Patch Parallelism assert_diffusion_response validates successful generation and the expected 512x512 resolution. @@ -38,12 +40,28 @@ BAGEL_HSDP_2_DEPLOY = modify_stage_config( BAGEL_CI_DEPLOY, updates={"stages": {0: {"devices": "0"}, 1: {"devices": "0,1"}}}, ) +BAGEL_TP_VAE_PP_2_DEPLOY = modify_stage_config( + BAGEL_CI_DEPLOY, + updates={ + "stages": { + 0: {"devices": "0"}, + 1: { + "devices": "0,1", + "parallel_config": { + "tensor_parallel_size": 2, + "vae_patch_parallel_size": 2, + }, + }, + }, + }, +) def _get_diffusion_feature_cases(model: str): """Return L4 diffusion feature cases for Bagel. TeaCache, Cache-DiT, CFG-Parallel, - Ulysses-SP, Ring-Attention, Layerwise Offloading. + Ulysses-SP, Ring-Attention, Layerwise Offloading, + Hybrid Sharded Data Parallel, Tensor Parallelism, VAE Patch Parallelism. """ return [ @@ -135,6 +153,15 @@ def _get_diffusion_feature_cases(model: str): id="parallel_hsdp_2", marks=HSDP_2_FEATURE_MARKS, ), + # Tensor Parallelism (TP) + VAE Patch Parallelism (size=2) + pytest.param( + OmniServerParams( + model=model, + stage_config_path=BAGEL_TP_VAE_PP_2_DEPLOY, + ), + id="tp_vae_patch_parallel_2", + marks=PARALLEL_2_FEATURE_MARKS, + ), ] @@ -157,6 +184,7 @@ def test_bagel( - Ring-Attention (degree=2) - Layerwise Offloading - Hybrid Sharded Data Parallel (size=2) + - Tensor Parallelism (TP) + VAE Patch Parallelism (size=2) Validation is delegated to assert_diffusion_response in tests/helpers/assertions.py, which checks output dimensions and basic correctness. diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 51be7d4ce..bc163873b 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -640,38 +640,6 @@ def test_initialize_llm_replica_passes_stage_init_timeout_to_complete_stage_hand assert captured_timeout == 302 -def test_async_omni_engine_reads_tokenizer_from_engine_args(monkeypatch): - import vllm_omni.engine.async_omni_engine as engine_mod - - class DummyThread: - def __init__(self, *args, **kwargs): - pass - - def start(self): - pass - - def is_alive(self): - return False - - fake_engine_args = types.SimpleNamespace( - _explicit_fields=frozenset({"tokenizer"}), - explicit_kwargs=lambda: {"tokenizer": "/tokenizer/from-engine-args"}, - ) - - monkeypatch.setattr(engine_mod.threading, "Thread", DummyThread) - monkeypatch.setattr( - AsyncOmniEngine, - "_resolve_stage_configs", - lambda self, model, kwargs: ("dummy-config", [types.SimpleNamespace(engine_args={})]), - ) - monkeypatch.setattr(AsyncOmniEngine, "_wait_for_orchestrator_init", lambda *_, **__: None) - - engine = AsyncOmniEngine("dummy-model", engine_args=fake_engine_args) - - assert engine.tokenizer == "/tokenizer/from-engine-args" - engine.shutdown() - - def test_build_engine_args_cli_tokenizer_overrides_inferred_base_tokenizer(tmp_path): from vllm_omni.engine.stage_init_utils import build_engine_args_dict diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index 04ddbe22d..2be2e9fb2 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -403,39 +403,6 @@ class TestSingleStageModeDetection: assert engine.single_stage_mode is True assert engine._single_stage_id_filter is None - def test_engine_args_create_only_forwards_explicit_fields(self, mocker: MockerFixture): - from vllm_omni.engine.arg_utils import OmniEngineArgs - - captured: dict[str, Any] = {} - - def fake_resolve(self, model: str, kwargs: dict[str, Any]): - captured.update(kwargs) - return "/fake/path", [_make_stage_cfg(0)] - - mocker.patch.object(AsyncOmniEngine, "_resolve_stage_configs", fake_resolve) - mocker.patch.object(AsyncOmniEngine, "_bootstrap_orchestrator") - mock_thread_cls = mocker.patch("threading.Thread") - mock_future_cls = mocker.patch("concurrent.futures.Future") - mock_future = mocker.Mock() - mock_future.result.return_value = mocker.Mock() - mock_future_cls.return_value = mock_future - mock_thread = mocker.Mock() - mock_thread.is_alive.return_value = False - mock_thread_cls.return_value = mock_thread - - ea = OmniEngineArgs.create(model="ignored", gpu_memory_utilization=0.5) - AsyncOmniEngine(model="fake-model", engine_args=ea) - - assert captured["gpu_memory_utilization"] == 0.5 - assert "model" not in captured - assert "max_num_seqs" not in captured - - def test_bare_engine_args_rejected(self, mocker: MockerFixture): - from vllm_omni.engine.arg_utils import OmniEngineArgs - - with pytest.raises(TypeError, match="OmniEngineArgs.create"): - self._make_engine_no_thread(mocker, engine_args=OmniEngineArgs(model="fake-model")) - def test_master_address_and_port_stored(self, mocker: MockerFixture): engine = self._make_engine_no_thread( mocker, diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index ab2229673..eba1878ec 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -4,12 +4,11 @@ from types import SimpleNamespace import pytest -from vllm.utils.argparse_utils import FlexibleArgumentParser -from vllm_omni.config.stage_config import deploy_override_field_names from vllm_omni.diffusion.data import AttentionConfig from vllm_omni.engine.async_omni_engine import AsyncOmniEngine -from vllm_omni.entrypoints.cli.serve import OmniServeCommand, _create_default_diffusion_stage_cfg +from vllm_omni.entrypoints.cli.serve import OmniServeCommand +from vllm_omni.utils.tracking_parser import TrackingArgumentParser pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -34,15 +33,6 @@ def test_default_stage_config_includes_cache_backend(): assert engine_args["model_stage"] == "diffusion" -def test_default_stage_config_ignores_none_deploy_overrides(): - """Ensure nullified deploy override defaults do not alter diffusion defaults.""" - baseline = AsyncOmniEngine._create_default_diffusion_stage_cfg({})[0] - nullified_overrides = {name: None for name in deploy_override_field_names()} - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(nullified_overrides)[0] - - assert stage_cfg == baseline - - def test_default_cache_config_used_when_missing(): """Ensure default cache_config is synthesized when only backend is given.""" stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( @@ -221,7 +211,7 @@ def test_default_stage_config_whitelist_none_fallback(): def test_serve_cli_accepts_ulysses_mode(): """Ensure diffusion serve CLI exposes ulysses_mode and wires it to parallel_config.""" - parser = FlexibleArgumentParser() + parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="command") OmniServeCommand().subparser_init(subparsers) @@ -237,7 +227,8 @@ def test_serve_cli_accepts_ulysses_mode(): ] ) - stage_cfg = _create_default_diffusion_stage_cfg(args)[0] + explicit_kwargs = args.get_explicit_kwargs_dict() + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] parallel_config = stage_cfg["engine_args"]["parallel_config"] assert args.ulysses_mode == "advanced_uaa" @@ -247,7 +238,7 @@ def test_serve_cli_accepts_ulysses_mode(): def test_serve_cli_accepts_diffusion_pipeline_profiler_flag(): """Ensure diffusion serve CLI exposes the profiler switch.""" - parser = FlexibleArgumentParser() + parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="command") OmniServeCommand().subparser_init(subparsers) @@ -260,7 +251,8 @@ def test_serve_cli_accepts_diffusion_pipeline_profiler_flag(): ] ) - stage_cfg = _create_default_diffusion_stage_cfg(args)[0] + explicit_kwargs = args.get_explicit_kwargs_dict() + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] assert args.enable_diffusion_pipeline_profiler is True assert stage_cfg["engine_args"]["enable_diffusion_pipeline_profiler"] is True @@ -268,7 +260,7 @@ def test_serve_cli_accepts_diffusion_pipeline_profiler_flag(): def test_serve_cli_accepts_diffusion_attention_backend(): """Ensure diffusion serve CLI exposes the shorthand backend flag.""" - parser = FlexibleArgumentParser() + parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="command") OmniServeCommand().subparser_init(subparsers) @@ -282,7 +274,8 @@ def test_serve_cli_accepts_diffusion_attention_backend(): ] ) - stage_cfg = _create_default_diffusion_stage_cfg(args)[0] + explicit_kwargs = args.get_explicit_kwargs_dict() + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] diffusion_attention_config = stage_cfg["engine_args"]["diffusion_attention_config"] assert args.diffusion_attention_backend == "FLASH_ATTN" @@ -293,7 +286,7 @@ def test_serve_cli_accepts_diffusion_attention_backend(): def test_serve_cli_accepts_additional_config(): """Ensure diffusion serve CLI exposes additional_config and forwards it to stage config.""" - parser = FlexibleArgumentParser() + parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="command") OmniServeCommand().subparser_init(subparsers) @@ -307,7 +300,8 @@ def test_serve_cli_accepts_additional_config(): ] ) - stage_cfg = _create_default_diffusion_stage_cfg(args)[0] + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(vars(args))[0] + engine_args = stage_cfg["engine_args"] assert args.additional_config == {"torchair_graph_config": {"enabled": True}} diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index bbef876d4..915efee5d 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -1,6 +1,5 @@ from __future__ import annotations -import argparse import queue from collections.abc import Callable from types import SimpleNamespace @@ -176,83 +175,6 @@ def _patch_engine(monkeypatch: pytest.MonkeyPatch, engine: FakeAsyncOmniEngine) ) -def test_direct_omni_with_nullified_parser_only_nulls_untyped_override_fields( - monkeypatch: pytest.MonkeyPatch, -): - from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults - from vllm_omni.entrypoints.omni import Omni - - captured: dict[str, Any] = {} - - def fake_engine(*args: Any, **kwargs: Any) -> FakeAsyncOmniEngine: - captured.update(kwargs) - return FakeAsyncOmniEngine() - - monkeypatch.setattr("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", fake_engine) - monkeypatch.setattr("vllm_omni.entrypoints.omni_base.omni_snapshot_download", lambda model: model) - - parser = argparse.ArgumentParser() - parser.add_argument("--gpu-memory-utilization", type=float, default=0.9) - parser.add_argument("--batch-timeout", type=int, default=10) - nullify_stage_engine_defaults(parser) - args = parser.parse_args([]) - args.model = "fake-model" - - Omni(**vars(args)) - - assert captured["gpu_memory_utilization"] is None - assert captured["batch_timeout"] == 10 - assert "_cli_explicit_keys" not in captured - - -def test_from_cli_args_warns_and_forwards_without_internal_keys( - monkeypatch: pytest.MonkeyPatch, -): - captured: dict[str, Any] = {} - - def fake_engine(*args: Any, **kwargs: Any) -> FakeAsyncOmniEngine: - captured.update(kwargs) - return FakeAsyncOmniEngine() - - monkeypatch.setattr("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", fake_engine) - monkeypatch.setattr("vllm_omni.entrypoints.omni_base.omni_snapshot_download", lambda model: model) - - args = argparse.Namespace(model="fake-model", gpu_memory_utilization=0.9, _cli_explicit_keys={"model"}) - with pytest.deprecated_call(match="from_cli_args"): - Omni.from_cli_args(args) - - assert captured["gpu_memory_utilization"] == 0.9 - assert "_cli_explicit_keys" not in captured - - -def test_deprecated_from_cli_args_preserves_legacy_parser_nulling( - monkeypatch: pytest.MonkeyPatch, -): - from vllm_omni.entrypoints.omni import Omni - - captured: dict[str, Any] = {} - - def fake_engine(*args: Any, **kwargs: Any) -> FakeAsyncOmniEngine: - captured.update(kwargs) - return FakeAsyncOmniEngine() - - monkeypatch.setattr("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", fake_engine) - monkeypatch.setattr("vllm_omni.entrypoints.omni_base.omni_snapshot_download", lambda model: model) - monkeypatch.setattr("sys.argv", ["prog"]) - - parser = argparse.ArgumentParser() - parser.add_argument("--gpu-memory-utilization", type=float, default=0.9) - parser.add_argument("--batch-timeout", type=int, default=10) - args = parser.parse_args([]) - args.model = "fake-model" - - with pytest.deprecated_call(match="from_cli_args"): - Omni.from_cli_args(args, parser=parser) - - assert captured["gpu_memory_utilization"] is None - assert captured["batch_timeout"] == 10 - - def _make_base(): from vllm_omni.entrypoints.omni_base import OmniBase diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 0915ddead..e05983d45 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -4,65 +4,56 @@ from __future__ import annotations import argparse from types import SimpleNamespace -from typing import Any import pytest from pytest_mock import MockerFixture from vllm_omni.entrypoints.cli.serve import OmniServeCommand, run_headless +from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -def test_serve_parser_accepts_no_async_chunk() -> None: - """``--no-async-chunk`` should parse after deploy-overriding parser - defaults are nullified.""" - try: - from vllm.utils.argparse_utils import FlexibleArgumentParser - except Exception as exc: - pytest.skip(f"Cannot build parser in this environment: {exc}") - - root = FlexibleArgumentParser() - subparsers = root.add_subparsers(dest="subcommand") +def test_serve_parser_accepts_no_async_chunk_and_marks_it_explicit() -> None: + """``--no-async-chunk`` should parse to ``async_chunk=False`` and mark the + shared deploy-level dest as explicitly provided by the user.""" + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") cmd = OmniServeCommand() cmd.subparser_init(subparsers) argv = ["serve", "fake-model", "--omni", "--no-async-chunk"] - args = root.parse_args(argv) - + args = parser.parse_args(argv) assert args.async_chunk is False - -# --------------------------------------------------------------------------- -# run_headless validation -# --------------------------------------------------------------------------- + explicit = args.get_explicit_kwargs_dict() + assert args.get_explicit_kwargs_dict() + assert not explicit["async_chunk"] -def _make_headless_args(**overrides: Any) -> argparse.Namespace: - """Build an argparse.Namespace shaped like the headless CLI passes in. - - Defaults pass every validation gate so individual tests can mutate just - the field they're exercising. - """ - defaults = dict( - model="fake-model", - stage_id=0, - replica_id=0, - omni_master_address="127.0.0.1", - omni_master_port=26000, - omni_replica_address=None, - omni_dp_size_local=1, - api_server_count=0, - worker_backend="multi_process", - stage_configs_path=None, - deploy_config=None, - log_stats=False, - disable_log_stats=False, - stage_init_timeout=600, - tokenizer=None, +def _make_headless_args(**kwargs) -> TrackingNamespace: + defaults = { + "model": "fake-model", + "stage_id": 0, + "replica_id": 0, + "omni_master_address": "127.0.0.1", + "omni_master_port": 26000, + "omni_replica_address": None, + "omni_dp_size_local": 1, + "worker_backend": "multi_process", + "stage_configs_path": None, + "deploy_config": None, + "log_stats": False, + "disable_log_stats": False, + "stage_init_timeout": 600, + "tokenizer": None, + } + ns_kwargs = {**defaults, **kwargs} + ns = argparse.Namespace(**ns_kwargs) + return TrackingNamespace( + unfiltered_ns=ns, + explicit_keys=frozenset(ns.__dict__.keys()), ) - defaults.update(overrides) - return argparse.Namespace(**defaults) def test_run_headless_requires_stage_id() -> None: @@ -83,12 +74,6 @@ def test_run_headless_requires_master_port() -> None: run_headless(args) -def test_run_headless_rejects_multi_api_server_count() -> None: - args = _make_headless_args(api_server_count=2) - with pytest.raises(ValueError, match="api_server_count can't be set"): - run_headless(args) - - def test_run_headless_rejects_non_multiprocess_worker_backend() -> None: args = _make_headless_args(worker_backend="ray") with pytest.raises(ValueError, match="worker_backend=multi_process"): diff --git a/tests/test_arg_utils.py b/tests/test_arg_utils.py index 1f454783f..40a584fce 100644 --- a/tests/test_arg_utils.py +++ b/tests/test_arg_utils.py @@ -5,303 +5,41 @@ hold for the orchestrator/engine/server CLI flag partition.""" from __future__ import annotations -import logging -from dataclasses import dataclass, fields +from dataclasses import fields import pytest from vllm_omni.engine.arg_utils import ( SHARED_FIELDS, - derive_server_dests_from_vllm_parser, + OmniEngineArgs, internal_blacklist_keys, - orchestrator_args_from_argparse, orchestrator_field_names, - split_kwargs, ) -# --------------------------------------------------------------------------- -# Fake engine class for unit testing — avoids pulling in the full vllm -# EngineArgs and its heavy __post_init__ at test time. -# --------------------------------------------------------------------------- - - -@dataclass -class _FakeEngineArgs: - """Stand-in for OmniEngineArgs with a representative subset of fields.""" - - model: str = "" - stage_id: int = 0 - max_num_seqs: int = 64 - gpu_memory_utilization: float = 0.9 - async_chunk: bool = False # also in OrchestratorArgs → shared - log_stats: bool = False # also in OrchestratorArgs → shared - stage_configs_path: str | None = None - - -# ============================================================================ -# Invariant 1 — OrchestratorArgs and engine must not ambiguously overlap. -# ============================================================================ - - -def test_no_ambiguous_overlap_with_fake_engine(): - """OrchestratorArgs ∩ engine fields must be ⊆ SHARED_FIELDS.""" - orch = orchestrator_field_names() - engine = {f.name for f in fields(_FakeEngineArgs)} - overlap = orch & engine - unexpected = overlap - SHARED_FIELDS - assert not unexpected, ( - f"Fields declared in both OrchestratorArgs and the engine class " - f"but not in SHARED_FIELDS: {sorted(unexpected)}. These cause " - f"double-routing — either remove the duplicate declaration or add " - f"to SHARED_FIELDS if sharing is intentional." - ) +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def test_no_ambiguous_overlap_with_real_engine(): - """Same check, but against the real OmniEngineArgs.""" - try: - from vllm_omni.engine.arg_utils import OmniEngineArgs - except Exception as exc: - pytest.skip(f"OmniEngineArgs not importable: {exc}") + """OrchestratorArgs ∩ OmniEngineArgs must be ⊆ SHARED_FIELDS ∪ orchestrator-captured. + + Fields that the orchestrator intentionally captures to prevent uniform + propagation (async_chunk, tokenizer) are excluded — they exist in both + classes by design but flow through DeployConfig, not kwargs forwarding. + """ + # Fields on both classes by design: orchestrator captures them to prevent + # uniform per-stage propagation; redistributed via DeployConfig. + orchestrator_captured = {"async_chunk", "tokenizer"} orch = orchestrator_field_names() engine = {f.name for f in fields(OmniEngineArgs)} overlap = orch & engine - unexpected = overlap - SHARED_FIELDS + unexpected = overlap - SHARED_FIELDS - orchestrator_captured assert not unexpected, ( - f"Real OmniEngineArgs has ambiguous overlap with OrchestratorArgs: " + f"OmniEngineArgs has ambiguous overlap with OrchestratorArgs: " f"{sorted(unexpected)}. Update SHARED_FIELDS or remove duplication." ) -# ============================================================================ -# Invariant 2 — split_kwargs partitions correctly. -# ============================================================================ - - -def test_split_orchestrator_only(): - """Pure orchestrator fields go to OrchestratorArgs, not engine_kwargs.""" - raw = {"stage_init_timeout": 500, "worker_backend": "ray"} - orch, engine = split_kwargs(raw, engine_cls=_FakeEngineArgs) - assert orch.stage_init_timeout == 500 - assert orch.worker_backend == "ray" - assert "stage_init_timeout" not in engine - assert "worker_backend" not in engine - - -def test_split_engine_only(): - """Pure engine fields go to engine_kwargs, not OrchestratorArgs.""" - raw = {"max_num_seqs": 128, "gpu_memory_utilization": 0.85} - orch, engine = split_kwargs(raw, engine_cls=_FakeEngineArgs) - assert engine["max_num_seqs"] == 128 - assert engine["gpu_memory_utilization"] == 0.85 - # These fields don't exist on OrchestratorArgs at all. - - -def test_split_shared_fields_go_to_both(): - """Fields in SHARED_FIELDS are copied to both buckets.""" - raw = {"model": "Qwen/Qwen2.5-Omni-7B", "log_stats": True} - orch, engine = split_kwargs(raw, engine_cls=_FakeEngineArgs) - assert orch.log_stats is True - assert engine["model"] == "Qwen/Qwen2.5-Omni-7B" - assert engine["log_stats"] is True - - -def test_split_drops_unclassified(): - """Unclassified fields (uvicorn/server) are dropped silently.""" - raw = { - "max_num_seqs": 64, # engine - "host": "0.0.0.0", # unclassified (server) - "port": 8091, # unclassified (server) - "ssl_keyfile": "key.pem", # unclassified (server) - } - orch, engine = split_kwargs(raw, engine_cls=_FakeEngineArgs) - assert engine == {"max_num_seqs": 64} - assert "host" not in engine - assert "port" not in engine - assert "ssl_keyfile" not in engine - - -def test_split_mixed_real_world(): - """End-to-end: raw CLI kwargs with all three classes present.""" - raw = { - # orchestrator - "stage_init_timeout": 400, - "deploy_config": "/tmp/deploy.yaml", - "worker_backend": "multi_process", - "async_chunk": True, - # engine - "max_num_seqs": 32, - "gpu_memory_utilization": 0.8, - # shared - "model": "Qwen/Qwen3-Omni", - "log_stats": False, - # server / unclassified - "host": "0.0.0.0", - "port": 8091, - "api_key": "secret", - # None values - "ray_address": None, - } - orch, engine = split_kwargs(raw, engine_cls=_FakeEngineArgs) - - # Orchestrator side - assert orch.stage_init_timeout == 400 - assert orch.deploy_config == "/tmp/deploy.yaml" - assert orch.worker_backend == "multi_process" - assert orch.async_chunk is True - assert orch.log_stats is False # shared, read from raw - assert orch.ray_address is None # default preserved - - # Engine side - assert engine["max_num_seqs"] == 32 - assert engine["gpu_memory_utilization"] == 0.8 - assert engine["model"] == "Qwen/Qwen3-Omni" - assert engine["log_stats"] is False - assert "host" not in engine - assert "port" not in engine - assert "api_key" not in engine - # orchestrator-only keys never reach engine - assert "stage_init_timeout" not in engine - assert "deploy_config" not in engine - # async_chunk is in SHARED_FIELDS so it flows to both - assert "async_chunk" in engine - - -# ============================================================================ -# Invariant 3 — user-typed unclassifiable flags warn (don't fail silently). -# ============================================================================ - - -def test_user_typed_unclassified_warns(monkeypatch): - """If the user types a flag we can't route, warn — don't silently drop.""" - import io - - buf = io.StringIO() - vllm_logger = logging.getLogger("vllm") - orig_handlers = list(vllm_logger.handlers) - vllm_logger.handlers.clear() - handler = logging.StreamHandler(buf) - handler.setLevel(logging.WARNING) - vllm_logger.addHandler(handler) - try: - raw = {"bogus_flag": "value", "max_num_seqs": 64} - split_kwargs(raw, engine_cls=_FakeEngineArgs, user_typed={"bogus_flag"}) - assert "bogus_flag" in buf.getvalue(), f"Expected 'bogus_flag' in log output, got: {buf.getvalue()!r}" - finally: - vllm_logger.handlers.clear() - for h in orig_handlers: - vllm_logger.addHandler(h) - - -def test_unclassified_without_user_typed_silent(): - """Without user_typed, unclassified keys drop silently (argparse defaults - for server flags shouldn't spam logs on every launch).""" - import io - - vllm_logger = logging.getLogger("vllm") - orig_handlers = list(vllm_logger.handlers) - vllm_logger.handlers.clear() - buf = io.StringIO() - handler = logging.StreamHandler(buf) - handler.setLevel(logging.WARNING) - vllm_logger.addHandler(handler) - try: - raw = {"host": "0.0.0.0", "port": 8091, "max_num_seqs": 64} - split_kwargs(raw, engine_cls=_FakeEngineArgs, user_typed=None) - assert "host" not in buf.getvalue() - finally: - vllm_logger.handlers.clear() - for h in orig_handlers: - vllm_logger.addHandler(h) - - -# ============================================================================ -# Invariant 4 — CLI flag classification completeness. -# Catches new flags added without updating OrchestratorArgs or OmniEngineArgs. -# ============================================================================ - - -def test_all_omni_cli_flags_classified(): - """Every vllm-omni-added CLI flag must be classifiable. - - Runs ``OmniServeCommand.subparser_init`` and checks that every new - argument (compared to vllm's base parser) is either: - - a field on OrchestratorArgs, OR - - a field on OmniEngineArgs, OR - - in SHARED_FIELDS - """ - try: - from vllm.utils.argparse_utils import FlexibleArgumentParser - - from vllm_omni.engine.arg_utils import OmniEngineArgs - from vllm_omni.entrypoints.cli.serve import OmniServeCommand - except Exception as exc: - pytest.skip(f"Cannot build parser in this environment: {exc}") - - # Build the serve parser - root = FlexibleArgumentParser() - subparsers = root.add_subparsers() - cmd = OmniServeCommand() - try: - parser = cmd.subparser_init(subparsers) - except Exception as exc: - pytest.skip(f"subparser_init failed (dev env issue): {exc}") - - all_dests = {a.dest for a in parser._actions if a.dest and a.dest not in {"help", "model_tag"}} - - orch = orchestrator_field_names() - engine = {f.name for f in fields(OmniEngineArgs)} - server_derived = derive_server_dests_from_vllm_parser() - - unclassified = all_dests - orch - engine - SHARED_FIELDS - server_derived - # Some argparse-internal dests (suppressed, private) may not match — - # filter those out. - unclassified = {d for d in unclassified if not d.startswith("_")} - - assert not unclassified, ( - f"These CLI flags are not classified as " - f"orchestrator/engine/shared/server: {sorted(unclassified)}. " - f"Add them to OrchestratorArgs (if consumed by orchestrator), " - f"OmniEngineArgs (if consumed by per-stage engine), or the known-server " - f"allowlist (if they're vllm frontend flags). " - f"If intentional (e.g. a new CLI-only flag that doesn't map to either " - f"dataclass), add it to a KNOWN_UNROUTED allowlist." - ) - - -# ============================================================================ -# argparse interop (Phase 3). -# ============================================================================ - - -def test_orchestrator_args_from_argparse(): - """Can build OrchestratorArgs from an argparse.Namespace.""" - import argparse - - ns = argparse.Namespace( - stage_init_timeout=500, - deploy_config="/tmp/x.yaml", - max_num_seqs=64, # engine field — ignored - host="0.0.0.0", # server field — ignored - ) - orch = orchestrator_args_from_argparse(ns) - assert orch.stage_init_timeout == 500 - assert orch.deploy_config == "/tmp/x.yaml" - assert orch.worker_backend == "multi_process" # default - - -def test_derive_server_dests_returns_frozenset(): - """Server-dest derivation returns a frozenset (possibly empty).""" - result = derive_server_dests_from_vllm_parser() - assert isinstance(result, frozenset) - - -# ============================================================================ -# internal_blacklist_keys — single source of truth for per-stage forwarding. -# ============================================================================ - - def test_internal_blacklist_keys_derived_from_orchestrator(): """Blacklist is exactly OrchestratorArgs fields minus SHARED_FIELDS. @@ -311,214 +49,3 @@ def test_internal_blacklist_keys_derived_from_orchestrator(): """ blacklist = internal_blacklist_keys() assert blacklist == orchestrator_field_names() - SHARED_FIELDS - # Spot-check expected entries - assert "stage_init_timeout" in blacklist - assert "deploy_config" in blacklist - assert "async_chunk" not in blacklist # now in SHARED_FIELDS - # Shared fields must NOT appear — they flow to both orchestrator and engine - assert "model" not in blacklist - assert "log_stats" not in blacklist - - -# ============================================================================ -# Boundary value analysis — edge cases around split_kwargs. -# ============================================================================ - - -def test_split_empty_kwargs(): - """Empty kwargs yields default OrchestratorArgs and empty engine dict.""" - orch, engine = split_kwargs({}, engine_cls=_FakeEngineArgs) - assert orch.stage_init_timeout == 300 # dataclass default - assert orch.worker_backend == "multi_process" # dataclass default - assert engine == {} - - -def test_split_all_none_values_preserved_on_orchestrator(): - """None values for orchestrator fields are kept (represents 'not set').""" - raw = {"ray_address": None, "deploy_config": None, "max_num_seqs": None} - orch, engine = split_kwargs(raw, engine_cls=_FakeEngineArgs) - assert orch.ray_address is None - assert orch.deploy_config is None - # Engine-side None still passes through; caller decides semantics downstream. - assert engine.get("max_num_seqs") is None - - -def test_split_user_typed_with_empty_kwargs_no_warn(caplog): - """user_typed non-empty but kwargs empty — no warnings emitted.""" - with caplog.at_level(logging.WARNING, logger="vllm_omni.engine.arg_utils"): - split_kwargs({}, engine_cls=_FakeEngineArgs, user_typed={"nothing"}) - assert not caplog.records - - -def test_ambiguous_field_strict_raises(): - """strict=True raises ValueError on overlap outside SHARED_FIELDS.""" - - # deploy_config is on OrchestratorArgs; declaring it on the engine class - # too (without adding to SHARED_FIELDS) creates an ambiguous route. - @dataclass - class _AmbiguousEngine: - deploy_config: str | None = None - - with pytest.raises(ValueError, match="both OrchestratorArgs and"): - split_kwargs({"deploy_config": "x"}, engine_cls=_AmbiguousEngine, strict=True) - - -def test_ambiguous_field_non_strict_routes_to_orchestrator(): - """strict=False logs ERROR but routes the ambiguous field to orchestrator.""" - - @dataclass - class _AmbiguousEngine: - deploy_config: str | None = None - - import io - - vllm_logger = logging.getLogger("vllm") - orig_handlers = list(vllm_logger.handlers) - vllm_logger.handlers.clear() - buf = io.StringIO() - handler = logging.StreamHandler(buf) - handler.setLevel(logging.ERROR) - vllm_logger.addHandler(handler) - try: - orch, engine = split_kwargs({"deploy_config": "x"}, engine_cls=_AmbiguousEngine, strict=False) - assert orch.deploy_config == "x" - assert "deploy_config" not in engine - assert "both OrchestratorArgs" in buf.getvalue() - finally: - vllm_logger.handlers.clear() - for h in orig_handlers: - vllm_logger.addHandler(h) - - -# Sentinel-default precedence invariants (#3035) - - -def _build_full_serve_parser(): - from vllm.utils.argparse_utils import FlexibleArgumentParser - - try: - from vllm.entrypoints.openai.cli_args import make_arg_parser - except ImportError: - pytest.skip("vllm parser not importable") - return make_arg_parser(FlexibleArgumentParser()) - - -def test_nullify_stage_engine_defaults_resets_inherited_defaults(): - import argparse - - from vllm_omni.config.stage_config import deploy_override_field_names - from vllm_omni.engine.arg_utils import ( - nullify_stage_engine_defaults, - ) - - parser = _build_full_serve_parser() - nullify_stage_engine_defaults(parser) - - override_dests = deploy_override_field_names() - offenders = [ - (a.dest, a.default) - for a in parser._actions - if a.dest not in ("help", "version") - and a.option_strings - and a.dest in override_dests - and a.default is not None - and a.default is not argparse.SUPPRESS - ] - assert not offenders, f"Stage flags with non-None defaults after nullify: {offenders}" - - -def test_non_override_flags_keep_real_defaults_after_nullify(): - import argparse - - from vllm_omni.config.stage_config import deploy_override_field_names - from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults - - parser = argparse.ArgumentParser() - parser.add_argument("--batch-timeout", type=int, default=10, help="Batch timeout.") - parser.add_argument("--max-num-seqs", type=int, default=64, help="Max num seqs.") - nullify_stage_engine_defaults(parser) - - assert "batch_timeout" not in deploy_override_field_names() - - batch_timeout = next(a for a in parser._actions if a.dest == "batch_timeout") - max_num_seqs = next(a for a in parser._actions if a.dest == "max_num_seqs") - assert batch_timeout.default == 10 - assert max_num_seqs.default is None - - -def test_help_text_preserves_default_after_nullify(): - # Real defaults must stay visible in --help even though parser stores None. - import argparse - - from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults - - parser = argparse.ArgumentParser() - parser.add_argument("--max-num-seqs", type=int, default=42, help="Example knob.") - nullify_stage_engine_defaults(parser) - - action = next(a for a in parser._actions if a.dest == "max_num_seqs") - assert action.default is None - assert "(default: 42)" in action.help - - -_OMNIENGINEARGS_USER_INPUT_FIELDS = frozenset( - { - "model_arch", - "engine_output_type", - "hf_config_name", - "custom_process_next_stage_input_func", - "subtalker_sampling_params", - "omni_kv_config", - "quantization_config", - "worker_type", - "task_type", - "worker_cls", - "omni_master_address", - "omni_master_port", - "stage_configs_path", - "output_modalities", - "custom_pipeline_args", - } -) - - -def test_omniengineargs_user_input_fields_default_to_none(): - try: - from vllm_omni.engine.arg_utils import OmniEngineArgs - except Exception as exc: - pytest.skip(f"OmniEngineArgs not importable: {exc}") - - offenders = [ - (f.name, f.default) - for f in fields(OmniEngineArgs) - if f.name in _OMNIENGINEARGS_USER_INPUT_FIELDS - and f.default is not dataclasses.MISSING - and f.default is not None - ] - assert not offenders, f"User-input fields with non-None defaults: {offenders}" - - -def test_omniengineargs_create_tracks_explicit_fields(): - try: - from vllm_omni.engine.arg_utils import OmniEngineArgs - except Exception as exc: - pytest.skip(f"OmniEngineArgs not importable: {exc}") - - ea = OmniEngineArgs.create(model="x", gpu_memory_utilization=0.5) - assert ea._explicit_fields == frozenset({"model", "gpu_memory_utilization"}) - assert ea.explicit_kwargs() == {"model": "x", "gpu_memory_utilization": 0.5} - - -def test_omniengineargs_bare_constructor_has_no_explicit_tracking(): - try: - from vllm_omni.engine.arg_utils import OmniEngineArgs - except Exception as exc: - pytest.skip(f"OmniEngineArgs not importable: {exc}") - - ea = OmniEngineArgs(model="x") - assert not hasattr(ea, "_explicit_fields") - assert "model" in ea.explicit_kwargs() - - -# dataclasses already imported via ``from dataclasses import dataclass, fields`` -import dataclasses # noqa: E402 -- needed for MISSING sentinel above diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index ed0a6fa1d..9119d78cd 100644 --- a/tests/test_config_factory.py +++ b/tests/test_config_factory.py @@ -29,7 +29,6 @@ from vllm_omni.config.stage_config import ( _deep_merge_stage, _resolve_scheduler, build_stage_runtime_overrides, - deploy_override_field_names, load_deploy_config, merge_pipeline_deploy, register_pipeline, @@ -39,6 +38,8 @@ from vllm_omni.engine.arg_utils import SHARED_FIELDS, EngineArgs, internal_black pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + class TestStageType: """Tests for StageType enum.""" @@ -156,23 +157,6 @@ class TestStageConfig: omega_config = config.to_omegaconf() assert omega_config.engine_args.max_num_seqs == 32 - def test_to_omegaconf_omits_none_deploy_overrides_for_engine_args(self): - """None deploy overrides must fall through to EngineArgs defaults.""" - - config = StageConfig( - stage_id=0, - model_stage="thinker", - runtime_overrides={name: None for name in deploy_override_field_names()}, - ) - - omega_config = config.to_omegaconf() - engine_args = dict(omega_config.engine_args) - - assert "devices" not in engine_args - assert "max_batch_size" not in engine_args - for name in deploy_override_field_names() - {"devices"}: - assert name not in engine_args - def test_to_omegaconf_diffusion_parallel_overrides_replace_nested_values(self): config = StageConfig( stage_id=1, @@ -1018,64 +1002,7 @@ class TestPipelineRegistry: class TestDeployConfigLoading: - def test_deploy_override_fields_include_deploy_schema_fields(self): - expected_fields = { - "active_stream_window", - "async_chunk", - # StageDeployConfig: stage placement and runtime fields. - "devices", - # StageDeployConfig: vLLM EngineArgs fields. - "async_scheduling", - "compilation_config", - "config_format", - "custom_voice_dir", - "data_parallel_size", - "devices", - "disable_hybrid_kv_cache_manager", - "enable_flashinfer_autotune", - "enforce_eager", - "gpu_memory_utilization", - "load_format", - "max_model_len", - "max_num_batched_tokens", - "max_num_seqs", - "mm_processor_cache_gb", - "profiler_config", - "skip_mm_profiling", - "subtalker_sampling_params", - "tensor_parallel_size", - "tokenizer_mode", - # StageDeployConfig: diffusion parallel_config deploy override fields. - "cfg_parallel_size", - "enable_expert_parallel", - "hsdp_replicate_size", - "hsdp_shard_size", - "ring_degree", - "sequence_parallel_size", - "ulysses_degree", - "ulysses_mode", - "use_hsdp", - "vae_patch_parallel_size", - # DeployConfig: pipeline-wide engine settings. - "data_parallel_size", - "distributed_executor_backend", - "dtype", - "enable_chunked_prefill", - "enable_prefix_caching", - "pipeline_parallel_size", - "quantization", - "trust_remote_code", - } - - actual_fields = deploy_override_field_names() - assert expected_fields == actual_fields, ( - f"added={actual_fields - expected_fields}, removed={expected_fields - actual_fields}" - ) - def test_custom_voice_dir_is_pipeline_wide_engine_arg(self, tmp_path): - import vllm_omni.model_executor.models.qwen3_tts.pipeline # noqa: F401 - from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy - deploy_path = tmp_path / "qwen3_tts_custom_voice.yaml" custom_voice_dir = tmp_path / "voices" deploy_path.write_text( @@ -1913,16 +1840,6 @@ class TestSentinelDefaultPrecedence: named = [p for p in sig.parameters.values() if p.kind != p.VAR_KEYWORD] assert "cli_explicit_keys" not in {p.name for p in named} - def test_cli_explicit_keys_kwarg_emits_deprecation(self): - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") - StageConfigFactory._create_from_registry( - "qwen3_omni_moe", - cli_overrides={}, - cli_explicit_keys={"max_num_seqs"}, - ) - assert any(issubclass(x.category, DeprecationWarning) for x in w) - def test_async_chunk_dispatches_processors(self): """A single ``qwen3_tts`` pipeline picks per-chunk vs end-to-end processors based on ``deploy.async_chunk``, without needing a diff --git a/tests/utils/test_tracking_parser.py b/tests/utils/test_tracking_parser.py new file mode 100644 index 000000000..537454661 --- /dev/null +++ b/tests/utils/test_tracking_parser.py @@ -0,0 +1,545 @@ +"""Tests for TrackingArgumentParser and related utilities.""" + +import argparse +import json +from unittest.mock import patch + +import pytest +import yaml +from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni.config.stage_config import ( + _PIPELINE_REGISTRY, + DeployConfig, + PipelineConfig, + StageConfigFactory, + StageDeployConfig, + StagePipelineConfig, +) +from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace + +### Fake pipeline/deploy config for integration tests + +_TEST_MODEL = "_test_tracking" + +_TEST_PIPELINE = PipelineConfig( + model_type=_TEST_MODEL, + stages=( + StagePipelineConfig(stage_id=0, model_stage="stage0"), + StagePipelineConfig(stage_id=1, model_stage="stage1"), + StagePipelineConfig(stage_id=2, model_stage="stage2"), + ), +) + +_TEST_DEPLOY = DeployConfig( + async_chunk=False, + stages=[ + StageDeployConfig(stage_id=0, gpu_memory_utilization=0.5, max_num_seqs=32), + StageDeployConfig(stage_id=1, gpu_memory_utilization=0.6, max_num_seqs=64), + StageDeployConfig(stage_id=2, gpu_memory_utilization=0.1, max_num_seqs=16), + ], +) + + +@pytest.fixture() +def mock_stages(monkeypatch): + """Register a fake pipeline and mock deploy YAML loading.""" + monkeypatch.setitem(_PIPELINE_REGISTRY._loaded, _TEST_MODEL, _TEST_PIPELINE) + monkeypatch.setattr( + "vllm_omni.config.stage_config.load_deploy_config", + lambda _path: _TEST_DEPLOY, + ) + return __file__ + + +### Tests for TrackingNamespace +def test_tracking_namespaces_cant_be_nested(): + """Ensure tracking namespaces explode if we try to nest them.""" + track_ns = TrackingNamespace( + unfiltered_ns=argparse.Namespace(foo="bar"), + explicit_keys=frozenset(), + ) + + with pytest.raises(ValueError): + TrackingNamespace( + unfiltered_ns=track_ns, + explicit_keys=frozenset(), + ) + + +def test_tracking_namespaces_init(): + """Check simple initialization for tracking namespaces.""" + unfiltered_ns = argparse.Namespace(foo="bar") + tracked_ns = TrackingNamespace( + unfiltered_ns=unfiltered_ns, + explicit_keys=frozenset({"foo"}), + ) + assert tracked_ns.foo == "bar" + assert tracked_ns.explicit_keys == frozenset({"foo"}) + + +def test_tracking_filtering(): + """Ensure tracking namespaces are filterable.""" + unfiltered_ns = argparse.Namespace(foo="bar", baz="foobar") + tracked_ns = TrackingNamespace( + unfiltered_ns=unfiltered_ns, + explicit_keys=frozenset({"foo"}), + ) + assert tracked_ns.foo == "bar" + assert tracked_ns.baz == "foobar" + assert tracked_ns.explicit_keys == frozenset({"foo"}) + # baz gets dropped because it's not marked in explicit_keys + assert tracked_ns.get_explicit_kwargs_dict() == {"foo": "bar"} + + +def test_setattr_writes_through_to_unfiltered_ns(): + """Ensure mutating an attribute on TrackingNamespace forwards to + get_explicit_kwargs_dict() and vars().""" + unfiltered_ns = argparse.Namespace(model="original") + tracked_ns = TrackingNamespace( + unfiltered_ns=unfiltered_ns, + explicit_keys=frozenset({"model"}), + ) + assert tracked_ns.model == "original" + assert tracked_ns.get_explicit_kwargs_dict() == {"model": "original"} + assert vars(tracked_ns) == {"model": "original"} + + # Ensure if we update the namespace, it's forwarded correctly + tracked_ns.model = "updated" + assert tracked_ns.model == "updated" + assert tracked_ns.get_explicit_kwargs_dict() == {"model": "updated"} + assert vars(tracked_ns) == {"model": "updated"} + + +### Tests for simple cases (no nested parsers or groups) +def test_vars_does_not_expose_internals(): + """Ensure vars on a TrackingArgumentParser is identical to running on the real one.""" + tracking = TrackingArgumentParser() + flexible = FlexibleArgumentParser() + namespaces = [] + + # For both parser, register two args, but only pass one + for parser in [tracking, flexible]: + parser.add_argument("--foo", type=int, default=42) + parser.add_argument("--bar", type=str, default="x") + namespaces.append(parser.parse_args(["--foo", "100"])) + + # Ensure that vars are the same result, since it isn't filtered + tracked_ns, flexible_ns = namespaces + tracked_kwargs = vars(tracked_ns) + real_kwargs = vars(flexible_ns) + assert tracked_kwargs == real_kwargs == {"foo": 100, "bar": "x"} + assert tracked_ns.explicit_keys == {"foo"} + + +def test_default_not_detected(): + """Ensure omitted defaults aren't in explicit keys and take defaults.""" + p = TrackingArgumentParser() + p.add_argument("--foo", type=int, default=42) + ns = p.parse_args([]) + assert ns.explicit_keys == set() + assert isinstance(ns, TrackingNamespace) + assert ns.foo == 42 + + +@pytest.mark.parametrize("val", ["42", "100"]) +def test_explicit_value_equal_to_default(val): + """Ensure explicit keys correctly handles passed values.""" + p = TrackingArgumentParser() + p.add_argument("--foo", type=int, default=42) + ns = p.parse_args(["--foo", val]) + assert ns.explicit_keys == {"foo"} + assert isinstance(ns, TrackingNamespace) + assert ns.foo == int(val) + + +def test_equals_syntax(): + """Ensure equals syntax is handled correctly wrt explicit keys.""" + p = TrackingArgumentParser() + p.add_argument("--foo", type=int, default=42) + ns = p.parse_args(["--foo=100"]) + assert ns.explicit_keys == {"foo"} + assert isinstance(ns, TrackingNamespace) + assert ns.foo == 100 + + +def test_explicit_none_via_store_const(): + """Ensure passing `None` to the namespace, it isn't filtered from explicit_keys.""" + parser = TrackingArgumentParser() + parser.add_argument("--foo", action="store_const", const=None, default="Something else") + ns = parser.parse_args(["--foo"]) + # User explicitly passed --foo, so it should be in the explicit keys even though it's None + assert ns.explicit_keys == {"foo"} + assert isinstance(ns, TrackingNamespace) + assert ns.foo is None + + +def test_multiple_args_mixed(): + """Ensure that explicit keys are correct when some are passed and others aren't.""" + p = TrackingArgumentParser() + p.add_argument("--foo", type=int, default=1) + p.add_argument("--bar", type=str, default="x") + p.add_argument("--baz", type=float, default=0.5) + ns = p.parse_args(["--foo", "10", "--baz", "0.9"]) + assert ns.explicit_keys == {"foo", "baz"} + assert isinstance(ns, TrackingNamespace) + assert ns.bar == "x" + + +def test_store_true_default(): + """Ensure that store true is handled correctly when omitted.""" + p = TrackingArgumentParser() + p.add_argument("--verbose", action="store_true") + ns = p.parse_args([]) + assert ns.explicit_keys == set() + assert isinstance(ns, TrackingNamespace) + assert ns.verbose is False + + +def test_store_true_explicit(): + """Ensure that store true is handled correctly when passed.""" + p = TrackingArgumentParser() + p.add_argument("--verbose", action="store_true") + ns = p.parse_args(["--verbose"]) + assert ns.explicit_keys == {"verbose"} + assert isinstance(ns, TrackingNamespace) + assert ns.verbose is True + + +def test_store_false_default(): + """Ensure that store false is handled correctly when omitted.""" + p = TrackingArgumentParser() + p.add_argument("--disable-x", action="store_false", dest="enable_x", default=True) + ns = p.parse_args([]) + assert ns.explicit_keys == set() + assert isinstance(ns, TrackingNamespace) + assert ns.enable_x is True + + +def test_store_false_explicit(): + """Ensure that store false is handled correctly when passed.""" + p = TrackingArgumentParser() + p.add_argument("--disable-x", action="store_false") + ns = p.parse_args(["--disable-x"]) + assert ns.explicit_keys == {"disable_x"} + assert isinstance(ns, TrackingNamespace) + assert ns.disable_x is False + + +def test_dest_is_reflected_in_explicit_keys(): + """Ensure that explicit keys use dest correctly.""" + p = TrackingArgumentParser() + p.add_argument("--foo", type=int, dest="bar") + ns = p.parse_args(["--foo", "100"]) + assert ns.explicit_keys == {"bar"} + assert isinstance(ns, TrackingNamespace) + assert ns.bar == 100 + + +def test_boolean_optional_action_default(): + """Check that boolean optional actions are handled correctly when omitted.""" + p = TrackingArgumentParser() + p.add_argument("--flag", action=argparse.BooleanOptionalAction) + ns = p.parse_args([]) + assert ns.explicit_keys == set() + assert isinstance(ns, TrackingNamespace) + assert ns.flag is None + + +def test_boolean_optional_action_positive(): + """Check that boolean optional actions are handled correctly.""" + p = TrackingArgumentParser() + p.add_argument("--flag", action=argparse.BooleanOptionalAction, default=None) + ns = p.parse_args(["--flag"]) + assert ns.explicit_keys == {"flag"} + assert isinstance(ns, TrackingNamespace) + assert ns.flag is True + + +def test_no_option_strings_are_handled(): + """Ensure --no- sets in the explicit keys correctly.""" + p = TrackingArgumentParser() + p.add_argument("--flag", action=argparse.BooleanOptionalAction, default=None) + ns = p.parse_args(["--no-flag"]) + assert ns.explicit_keys == {"flag"} + assert isinstance(ns, TrackingNamespace) + assert ns.flag is False + + +def test_json_type(): + """Ensure json type is handled correctly.""" + p = TrackingArgumentParser() + p.add_argument("--cfg", type=json.loads, default="{}") + ns = p.parse_args(["--cfg", '{"a": 1}']) + assert "cfg" in ns.explicit_keys + assert isinstance(ns, TrackingNamespace) + assert ns.cfg == {"a": 1} + + +def test_choices(): + """Ensure choices are handled correctly.""" + p = TrackingArgumentParser() + p.add_argument("--mode", choices=["fast", "slow"], default="fast") + ns = p.parse_args(["--mode", "slow"]) + assert "mode" in ns.explicit_keys + assert ns is not None + assert ns.mode == "slow" + + +def test_explicit_positional_arg(): + """Ensure positional args are handled correctly when provided.""" + p = TrackingArgumentParser() + p.add_argument("name", nargs="?", default=None) + ns = p.parse_args(["hello"]) + assert "name" in ns.explicit_keys + assert isinstance(ns, TrackingNamespace) + assert ns.name == "hello" + + +def test_omitted_positional_arg(): + """Ensure positional args are handled correctly when omitted.""" + p = TrackingArgumentParser() + p.add_argument("name", nargs="?", default=None) + ns = p.parse_args([]) + assert "name" not in ns.explicit_keys + assert isinstance(ns, TrackingNamespace) + assert ns.name is None + + +def test_explicit_nargs(): + """Ensure that variable num args are handled correctly when provided.""" + p = TrackingArgumentParser() + p.add_argument("--items", nargs="*", default=None) + ns = p.parse_args(["--items", "a", "b"]) + assert "items" in ns.explicit_keys + assert isinstance(ns, TrackingNamespace) + assert ns.items == ["a", "b"] + + +def test_omitted_nargs(): + """Ensure that variable num args are handled correctly when omitted.""" + p = TrackingArgumentParser() + p.add_argument("--items", nargs="*", default=None) + ns = p.parse_args([]) + assert "items" not in ns.explicit_keys + assert isinstance(ns, TrackingNamespace) + assert ns.items is None + + +def test_parse_known_args_tracking(): + """Ensure parse_known_args is also trackable""" + p = TrackingArgumentParser() + p.add_argument("--foo", type=int, default=42) + ns, remaining = p.parse_known_args(["--foo", "10", "--unknown", "val"]) + assert ns.explicit_keys == {"foo"} + assert isinstance(ns, TrackingNamespace) + assert ns.foo == 10 + assert remaining == ["--unknown", "val"] + + +### Tests for group handling +def test_group_arg_default(): + """Ensure that that groups with defaults are handled correctly.""" + p = TrackingArgumentParser() + g = p.add_argument_group("TestGroup") + g.add_argument("--bar", type=str, default="baz") + ns = p.parse_args([]) + assert ns.explicit_keys == set() + assert isinstance(ns, TrackingNamespace) + assert ns.bar == "baz" + + +def test_group_arg_explicit(): + """Ensure that that groups are handled correctly.""" + p = TrackingArgumentParser() + g = p.add_argument_group("TestGroup") + g.add_argument("--bar", type=str, default="baz") + ns = p.parse_args(["--bar", "qux"]) + assert ns.explicit_keys == {"bar"} + assert isinstance(ns, TrackingNamespace) + assert ns.bar == "qux" + + +def test_multiple_groups(): + """Test multiple groups behave correctly.""" + p = TrackingArgumentParser() + g1 = p.add_argument_group("Group1") + g2 = p.add_argument_group("Group2") + g1.add_argument("--a", type=int, default=1) + g2.add_argument("--b", type=int, default=2) + ns = p.parse_args(["--b", "20"]) + assert ns.explicit_keys == {"b"} + assert isinstance(ns, TrackingNamespace) + assert ns.a == 1 + assert ns.b == 20 + + +def test_omitted_mutually_exclusive_group(): + """Ensure that that mutually exclusive groups with defaults are handled correctly.""" + p = TrackingArgumentParser() + grp = p.add_mutually_exclusive_group() + grp.add_argument("--json", action="store_true") + grp.add_argument("--text", action="store_true") + ns = p.parse_args([]) + assert ns.explicit_keys == set() + assert isinstance(ns, TrackingNamespace) + + +def test_mutually_exclusive_group(): + """Ensure that that mutually exclusive groups are handled correctly.""" + p = TrackingArgumentParser() + grp = p.add_mutually_exclusive_group() + grp.add_argument("--json", action="store_true") + grp.add_argument("--text", action="store_true") + ns = p.parse_args(["--json"]) + assert ns.explicit_keys == {"json"} + assert isinstance(ns, TrackingNamespace) + assert ns.json is True + assert ns.text is False + + +### Tests for subparser handling +def test_subparser_explicit_detection(): + """Ensure that subparsers are handled correctly.""" + p = TrackingArgumentParser() + sub = p.add_subparsers() + child = sub.add_parser("foo") + child.add_argument("--bar", type=str) + ns = p.parse_args(["foo", "--bar", "baz"]) + assert isinstance(child, TrackingArgumentParser) + assert isinstance(ns, TrackingNamespace) + assert ns.explicit_keys == {"bar"} + assert ns.bar == "baz" + + +def test_subparser_group_args(): + """Ensure that subparsers with groups are handled correctly.""" + p = TrackingArgumentParser() + sub = p.add_subparsers() + child = sub.add_parser("foo") + g = child.add_argument_group("Config") + g.add_argument("--port", type=int, default=8000) + g.add_argument("--host", type=str, default="localhost") + ns = p.parse_args(["foo", "--port", "9000"]) + assert ns.explicit_keys == {"port"} + assert isinstance(ns, TrackingNamespace) + assert ns.host == "localhost" + + +### Tests for specific behaviors against FlexibleArgumentParser +def test_config_file_args_detected(tmp_path): + """Ensure config is handled correctly for vLLM's FlexibleArgumentParser.""" + p = TrackingArgumentParser() + p.add_argument("--foo", type=int) + p.add_argument("--bar", type=int) + cfg = tmp_path / "test.yaml" + cfg.write_text(yaml.dump({"foo": 100})) + ns = p.parse_args(["--config", str(cfg)]) + assert isinstance(ns, TrackingNamespace) + assert ns.explicit_keys == {"foo"} + assert ns.foo == 100 + + +def test_cli_overrides_config(tmp_path): + """Ensure tracking parser handles config vs cli overrides correctly.""" + p = TrackingArgumentParser() + p.add_argument("--foo", type=int) + cfg = tmp_path / "test.yaml" + cfg.write_text(yaml.dump({"foo": 100})) + ns = p.parse_args(["--config", str(cfg), "--foo", "200"]) + assert isinstance(ns, TrackingNamespace) + assert "foo" in ns.explicit_keys + assert ns.foo == 200 + + +### Integration tests for arg resolution through StageConfigFactory +def test_explicit_cli_arg_reaches_runtime_overrides(mock_stages): + """Explicitly passed CLI values reach runtime_overrides on all stages.""" + p = TrackingArgumentParser() + p.add_argument("--max-num-seqs", type=int, default=64) + ns = p.parse_args(["--max-num-seqs", "999"]) + + explicit_kwargs = ns.get_explicit_kwargs_dict() + stages = StageConfigFactory._create_from_registry( + _TEST_MODEL, + explicit_kwargs, + deploy_config_path=mock_stages, + ) + for stage in stages: + assert stage.runtime_overrides.get("max_num_seqs") == 999 + + +def test_omitted_default_not_in_runtime_overrides(mock_stages): + """Omitted defaults are overridden by deploy config values""" + p = TrackingArgumentParser() + p.add_argument("--max-num-seqs", type=int, default=64) + ns = p.parse_args([]) + + explicit_kwargs = ns.get_explicit_kwargs_dict() + stages = StageConfigFactory._create_from_registry( + _TEST_MODEL, + explicit_kwargs, + deploy_config_path=mock_stages, + ) + for stage in stages: + assert stage.runtime_overrides == {} + + +def test_config_file_args_reach_runtime_overrides(mock_stages): + """Args from --config YAML must be treated as explicitly passed and + flow through to runtime_overrides.""" + p = TrackingArgumentParser() + p.add_argument("--max-num-seqs", type=int, default=64) + p.add_argument("--gpu-memory-utilization", type=float, default=0.9) + with patch( + "vllm.utils.argparse_utils.FlexibleArgumentParser.load_config_file", return_value=["--max-num-seqs", "999"] + ): + ns = p.parse_args(["--config", "fake.yaml"]) + + explicit_kwargs = ns.get_explicit_kwargs_dict() + stages = StageConfigFactory._create_from_registry( + _TEST_MODEL, + explicit_kwargs, + deploy_config_path=mock_stages, + ) + for stage in stages: + assert stage.runtime_overrides.get("max_num_seqs") == 999 + assert "gpu_memory_utilization" not in stage.runtime_overrides + + +def test_per_stage_override_routes_correctly(mock_stages): + """Ensure stage__ only affects the targeted stage.""" + p = TrackingArgumentParser() + p.add_argument("--stage-0-gpu-memory-utilization", type=float) + ns = p.parse_args(["--stage-0-gpu-memory-utilization", "0.42"]) + + explicit_kwargs = ns.get_explicit_kwargs_dict() + stages = StageConfigFactory._create_from_registry( + _TEST_MODEL, + explicit_kwargs, + deploy_config_path=mock_stages, + ) + assert stages[0].runtime_overrides == {"gpu_memory_utilization": 0.42} + assert stages[1].runtime_overrides == {} + assert stages[2].runtime_overrides == {} + + +def test_explicit_args_omitted_from_yaml(mock_stages): + """Ensure only passed args end up in runtime overrides (regardless of whether + they are defined in the yaml config).""" + p = TrackingArgumentParser() + p.add_argument("--enforce-eager", action="store_true") + p.add_argument("--max-num-seqs", type=int, default=64) + p.add_argument("--gpu-memory-utilization", type=float, default=0.9) + # NOTE: enforce eager is not set in the mock config yaml. + ns = p.parse_args(["--enforce-eager"]) + + explicit_kwargs = ns.get_explicit_kwargs_dict() + stages = StageConfigFactory._create_from_registry( + _TEST_MODEL, + explicit_kwargs, + deploy_config_path=mock_stages, + ) + + for stage in stages: + assert stage.runtime_overrides == {"enforce_eager": True} diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 5254a5396..13dd0803d 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -30,15 +30,6 @@ def get_pipeline_path(model_dir: str, filename: str) -> Path: logger = init_logger(__name__) -def _warn_deprecated_kwargs(kwargs: dict[str, Any]) -> None: - if "cli_explicit_keys" in kwargs: - warnings.warn( - "cli_explicit_keys= is deprecated and ignored. Remove the kwarg.", - DeprecationWarning, - stacklevel=3, - ) - - _STAGE_OVERRIDE_PATTERN = re.compile(r"^stage_(\d+)_(.+)$") @@ -837,15 +828,6 @@ _PIPELINE_WIDE_ENGINE_FIELDS: tuple[str, ...] = ( ) -def deploy_override_field_names() -> frozenset[str]: - """Return deploy-schema fields whose CLI defaults must not override YAML.""" - return ( - frozenset(_STAGE_DEPLOY_FIELDS) - | frozenset(_PIPELINE_WIDE_ENGINE_FIELDS) - | frozenset({"async_chunk", "devices"}) - ) - - def _build_engine_args( ps: StagePipelineConfig, ds: StageDeployConfig | None, @@ -1169,14 +1151,11 @@ class StageConfigFactory: model: str, cli_overrides: dict[str, Any] | None = None, deploy_config_path: str | None = None, - **deprecated_kwargs: Any, ) -> list[StageConfig] | None: """Load pipeline + deploy config, merge with CLI overrides. Checks _PIPELINE_REGISTRY first (new path), falls back to legacy YAML. """ - _warn_deprecated_kwargs(deprecated_kwargs) - if cli_overrides is None: cli_overrides = {} @@ -1256,15 +1235,12 @@ class StageConfigFactory: model_type: str, cli_overrides: dict[str, Any], deploy_config_path: str | None = None, - **deprecated_kwargs: Any, ) -> list[StageConfig]: """Create StageConfigs from pipeline registry + deploy YAML. Precedence: caller-typed (non-None) value > deploy YAML > StageDeployConfig dataclass default. """ - _warn_deprecated_kwargs(deprecated_kwargs) - # Resolve deploy config path if deploy_config_path is None: deploy_path = _DEPLOY_DIR / f"{model_type}.yaml" diff --git a/vllm_omni/diffusion/models/bagel/autoencoder.py b/vllm_omni/diffusion/models/bagel/autoencoder.py index 0980f25cd..006e927b0 100644 --- a/vllm_omni/diffusion/models/bagel/autoencoder.py +++ b/vllm_omni/diffusion/models/bagel/autoencoder.py @@ -14,6 +14,16 @@ from dataclasses import dataclass import torch from einops import rearrange from torch import Tensor, nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( + DistributedOperator, + DistributedVaeMixin, + GridSpec, + TileTask, +) + +logger = init_logger(__name__) @dataclass @@ -322,3 +332,229 @@ class AutoEncoder(nn.Module): def forward(self, x: Tensor) -> Tensor: return self.decode(self.encode(x)) + + +class DistributedAutoEncoder(AutoEncoder, DistributedVaeMixin): + def __init__(self, params: AutoEncoderParams): + super().__init__(params) + self.init_distributed() + + self.use_tiling = False + self.tile_sample_min_height = 512 + self.tile_sample_min_width = 512 + self.tile_sample_stride_height = 448 + self.tile_sample_stride_width = 448 + self.spatial_compression_ratio = params.downsample + + def blend_v(self, above: Tensor, current: Tensor, blend_extent: int) -> Tensor: + blend_extent = min(above.shape[-2], current.shape[-2], blend_extent) + if blend_extent <= 0: + return current + + for y in range(blend_extent): + alpha = y / blend_extent + current[:, :, y, :] = above[:, :, -blend_extent + y, :] * (1 - alpha) + current[:, :, y, :] * alpha + return current + + def blend_h(self, left: Tensor, current: Tensor, blend_extent: int) -> Tensor: + blend_extent = min(left.shape[-1], current.shape[-1], blend_extent) + if blend_extent <= 0: + return current + + for x in range(blend_extent): + alpha = x / blend_extent + current[:, :, :, x] = left[:, :, :, -blend_extent + x] * (1 - alpha) + current[:, :, :, x] * alpha + return current + + def decode_tile_split(self, z: Tensor) -> tuple[list[TileTask], GridSpec]: + _, _, latent_height, latent_width = z.shape + + sample_height = latent_height * self.spatial_compression_ratio + sample_width = latent_width * self.spatial_compression_ratio + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + + blend_height = self.tile_sample_min_height - self.tile_sample_stride_height + blend_width = self.tile_sample_min_width - self.tile_sample_stride_width + + tiletask_list = [] + for i in range(0, latent_height, tile_latent_stride_height): + for j in range(0, latent_width, tile_latent_stride_width): + tile = z[ + :, + :, + i : i + tile_latent_min_height, + j : j + tile_latent_min_width, + ] + tiletask_list.append( + TileTask( + tile_id=len(tiletask_list), + grid_coord=(i // tile_latent_stride_height, j // tile_latent_stride_width), + tensor=tile, + workload=tile.shape[-2] * tile.shape[-1], + ) + ) + grid_spec = GridSpec( + split_dims=(2, 3), + grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1), + tile_spec={ + "sample_height": sample_height, + "sample_width": sample_width, + "blend_height": blend_height, + "blend_width": blend_width, + }, + output_dtype=z.dtype, + ) + return tiletask_list, grid_spec + + def decode_tile_exec(self, task: TileTask) -> Tensor: + tile = task.tensor + tile = tile / self.scale_factor + self.shift_factor + return self.decoder(tile) + + def decode_tile_merge(self, coord_tensor_map: dict[tuple[int, ...], Tensor], grid_spec: GridSpec) -> Tensor: + grid_h, grid_w = grid_spec.grid_shape + result_rows = [] + + for i in range(grid_h): + result_row = [] + for j in range(grid_w): + tile = coord_tensor_map[(i, j)] + if i > 0: + tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"]) + if j > 0: + tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"]) + + result_row.append( + tile[ + :, + :, + : self.tile_sample_stride_height, + : self.tile_sample_stride_width, + ] + ) + result_rows.append(torch.cat(result_row, dim=-1)) + + dec = torch.cat(result_rows, dim=-2)[ + :, + :, + : grid_spec.tile_spec["sample_height"], + : grid_spec.tile_spec["sample_width"], + ] + return dec + + def decode(self, z: Tensor) -> Tensor: + if not self.is_distributed_enabled(): + return super().decode(z) + + logger.debug("Bagel VAE decode running with distributed executor") + return self.distributed_executor.execute( + z, + DistributedOperator( + split=self.decode_tile_split, + exec=self.decode_tile_exec, + merge=self.decode_tile_merge, + ), + broadcast_result=True, + ) + + def encode_tile_split(self, x: Tensor) -> tuple[list[TileTask], GridSpec]: + _, _, height, width = x.shape + + latent_height = height // self.spatial_compression_ratio + latent_width = width // self.spatial_compression_ratio + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + + blend_height = tile_latent_min_height - tile_latent_stride_height + blend_width = tile_latent_min_width - tile_latent_stride_width + + tiletask_list = [] + for i in range(0, height, self.tile_sample_stride_height): + for j in range(0, width, self.tile_sample_stride_width): + tile = x[ + :, + :, + i : i + self.tile_sample_min_height, + j : j + self.tile_sample_min_width, + ] + tiletask_list.append( + TileTask( + tile_id=len(tiletask_list), + grid_coord=(i // self.tile_sample_stride_height, j // self.tile_sample_stride_width), + tensor=tile, + workload=tile.shape[-2] * tile.shape[-1], + ) + ) + grid_spec = GridSpec( + split_dims=(2, 3), + grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1), + tile_spec={ + "latent_height": latent_height, + "latent_width": latent_width, + "blend_height": blend_height, + "blend_width": blend_width, + "tile_latent_stride_height": tile_latent_stride_height, + "tile_latent_stride_width": tile_latent_stride_width, + }, + output_dtype=next(self.parameters()).dtype, + ) + return tiletask_list, grid_spec + + def encode_tile_exec(self, task: TileTask) -> Tensor: + tile = task.tensor + z = self.reg(self.encoder(tile)) + z = self.scale_factor * (z - self.shift_factor) + return z + + def encode_tile_merge(self, coord_tensor_map: dict[tuple[int, ...], Tensor], grid_spec: GridSpec) -> Tensor: + grid_h, grid_w = grid_spec.grid_shape + result_rows = [] + + for i in range(grid_h): + result_row = [] + for j in range(grid_w): + tile = coord_tensor_map[(i, j)] + if i > 0: + tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"]) + if j > 0: + tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"]) + + result_row.append( + tile[ + :, + :, + : grid_spec.tile_spec["tile_latent_stride_height"], + : grid_spec.tile_spec["tile_latent_stride_width"], + ] + ) + result_rows.append(torch.cat(result_row, dim=-1)) + + enc = torch.cat(result_rows, dim=-2)[ + :, + :, + : grid_spec.tile_spec["latent_height"], + : grid_spec.tile_spec["latent_width"], + ] + return enc + + def encode(self, x: Tensor) -> Tensor: + if not self.is_distributed_enabled(): + return super().encode(x) + logger.debug("Bagel VAE encode running with distributed executor") + return self.distributed_executor.execute( + x, + DistributedOperator( + split=self.encode_tile_split, + exec=self.encode_tile_exec, + merge=self.encode_tile_merge, + ), + broadcast_result=True, + ) diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 26451615e..02caed9e3 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -32,7 +32,7 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPi from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific -from .autoencoder import AutoEncoder, AutoEncoderParams +from .autoencoder import AutoEncoder, AutoEncoderParams, DistributedAutoEncoder from .bagel_transformer import Bagel, NaiveCache, Qwen2MoTConfig, Qwen2MoTForCausalLM logger = init_logger(__name__) @@ -253,7 +253,7 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf ) self.transformer = self.language_model.model ae_params: AutoEncoderParams = default_ae_params() - self.vae = AutoEncoder(ae_params) + self.vae = DistributedAutoEncoder(ae_params) self.bagel = Bagel( language_model=self.language_model, diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py index abdca8b1a..e61b59c56 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py @@ -43,7 +43,6 @@ from vllm_omni.diffusion.distributed.sp_plan import ( ) from vllm_omni.diffusion.forward_context import get_forward_context from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm -from vllm_omni.diffusion.layers.norm import RMSNorm from vllm_omni.diffusion.layers.rope import RotaryEmbedding logger = init_logger(__name__) @@ -538,8 +537,8 @@ class QwenImageCrossAttention(nn.Module): self.query_num_heads = self.to_qkv.num_heads self.kv_num_heads = self.to_qkv.num_kv_heads - self.norm_q = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity() - self.norm_k = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity() + self.norm_q = nn.RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = nn.RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity() self.inner_dim = out_dim if out_dim is not None else head_dim * self.total_num_heads @@ -576,8 +575,8 @@ class QwenImageCrossAttention(nn.Module): prefix=_join_prefix(prefix, "to_out"), ) - self.norm_added_q = RMSNorm(head_dim, eps=eps) - self.norm_added_k = RMSNorm(head_dim, eps=eps) + self.norm_added_q = nn.RMSNorm(head_dim, eps=eps) + self.norm_added_k = nn.RMSNorm(head_dim, eps=eps) self.attn = Attention( num_heads=self.query_num_heads, @@ -992,7 +991,7 @@ class QwenImageTransformer2DModel(CachedTransformer): quant_config=quant_config, ) - self.txt_norm = RMSNorm(joint_attention_dim, eps=1e-6) + self.txt_norm = nn.RMSNorm(joint_attention_dim, eps=1e-6) # Entry projections (image/text) are kept full precision — # small sensitive layers at the network boundary (see #2728). diff --git a/vllm_omni/diffusion/models/sd3/pipeline_sd3.py b/vllm_omni/diffusion/models/sd3/pipeline_sd3.py index 995f0739d..7eeffdde4 100644 --- a/vllm_omni/diffusion/models/sd3/pipeline_sd3.py +++ b/vllm_omni/diffusion/models/sd3/pipeline_sd3.py @@ -402,12 +402,12 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro truncation=True, add_special_tokens=True, return_tensors="pt", - ).to(self.device) + ) text_input_ids = text_inputs.input_ids untruncated_ids = self.tokenizer_3(prompt, padding="longest", return_tensors="pt").input_ids if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): - removed_text = self.tokenizer_3.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1]) + removed_text = self.tokenizer_3.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1]) logger.warning( "The following part of your input was truncated because `max_sequence_length` is set to " f" {max_sequence_length} tokens: {removed_text}" diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 7fcb7d0ef..3e5ce5536 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -1,5 +1,4 @@ import argparse -import dataclasses import json import os import tempfile @@ -191,30 +190,6 @@ class OmniEngineArgs(EngineArgs): load_omni_general_plugins() super().__post_init__() - @classmethod - def from_cli_args(cls, args: argparse.Namespace) -> "OmniEngineArgs": - attrs = [attr.name for attr in dataclasses.fields(cls)] - engine_args = cls(**{attr: getattr(args, attr) for attr in attrs if hasattr(args, attr)}) - engine_args._explicit_fields = frozenset( - attr for attr in attrs if hasattr(args, attr) and getattr(args, attr) is not None - ) - return engine_args - - @classmethod - def create(cls, **explicit: Any) -> "OmniEngineArgs": - """Tracks caller-set fields for ``Omni(..., engine_args=ea)``.""" - ea = cls(**explicit) - ea._explicit_fields = frozenset(explicit.keys()) - return ea - - def explicit_kwargs(self) -> dict[str, Any]: - explicit = getattr(self, "_explicit_fields", None) - if explicit is None: - return { - f.name: getattr(self, f.name) for f in dataclasses.fields(self) if getattr(self, f.name) is not None - } - return {k: getattr(self, k) for k in explicit} - def _ensure_omni_models_registered(self): if hasattr(self, "_omni_models_registered"): return True @@ -420,8 +395,7 @@ class OrchestratorArgs: forwarding). Fields that BOTH orchestrator and engine genuinely need (e.g. ``model``, - ``log_stats``) should be listed in ``SHARED_FIELDS`` below; ``split_kwargs`` - will copy them to both buckets. + ``log_stats``) should be listed in ``SHARED_FIELDS`` below. """ # === Lifecycle === @@ -530,163 +504,3 @@ def internal_blacklist_keys() -> frozenset[str]: dataclass — this function updates automatically. """ return orchestrator_field_names() - SHARED_FIELDS - - -def split_kwargs( - kwargs: dict[str, Any], - *, - engine_cls: type | None = None, - user_typed: set[str] | None = None, - strict: bool = False, -) -> tuple[OrchestratorArgs, dict[str, Any]]: - """Partition CLI kwargs into (orchestrator, engine) buckets. - - Args: - kwargs: Raw dict, typically ``vars(args)``. - engine_cls: Engine dataclass used to whitelist-filter the engine - bucket. Defaults to ``OmniEngineArgs``. Pass a custom class - for testing. - user_typed: Keys the user actually typed on the command line. Used - to warn when a user-typed flag is unclassifiable. - strict: If True, raise ``ValueError`` on ambiguous (double-classified - but not in ``SHARED_FIELDS``) fields. Default False to keep the - rollout non-breaking; flip to True in tests and CI. - - Returns: - ``(orchestrator_args, engine_kwargs)``. ``engine_kwargs`` has already - been whitelist-filtered against ``engine_cls`` — safe to pass directly - to ``engine_cls(**engine_kwargs)``. - """ - if engine_cls is None: - engine_cls = OmniEngineArgs - - orch_fields = orchestrator_field_names() - engine_fields = {f.name for f in fields(engine_cls)} - - orch_kwargs: dict[str, Any] = {} - engine_candidate: dict[str, Any] = {} - shared_values: dict[str, Any] = {} - unclassified: dict[str, Any] = {} - - for key, value in kwargs.items(): - in_orch = key in orch_fields - in_engine = key in engine_fields - is_shared = key in SHARED_FIELDS - - if is_shared: - shared_values[key] = value - elif in_orch and in_engine: - # Declared in both but not marked shared → ambiguous. - msg = ( - f"Field {key!r} is defined on both OrchestratorArgs and " - f"{engine_cls.__name__} but is not in SHARED_FIELDS. " - f"This causes double-routing. Either remove the duplicate or " - f"add {key!r} to SHARED_FIELDS if the sharing is intentional." - ) - if strict: - raise ValueError(msg) - logger.error(msg) - # Default: treat as orchestrator-only to preserve existing behavior. - orch_kwargs[key] = value - elif in_orch: - orch_kwargs[key] = value - elif in_engine: - engine_candidate[key] = value - else: - unclassified[key] = value - - # Warn on user-typed but unclassifiable flags so we don't silently drop - # something the user cared about (fixes the class of bug that spawned #873). - if unclassified and user_typed: - user_typed_unknown = sorted(k for k in unclassified if k in user_typed) - if user_typed_unknown: - logger.warning( - "CLI flags not consumed by vllm-omni and dropped before " - "per-stage engine construction: %s. If these are vllm " - "frontend/uvicorn flags (host, port, ssl_*, api_key, …) this " - "is expected; otherwise check your spelling.", - user_typed_unknown, - ) - - # Engine bucket: shared + engine-only. We do NOT pass through unclassified - # fields — that's exactly the server/uvicorn noise we want to shed. - engine_kwargs = {**shared_values, **engine_candidate} - - # Construct the orchestrator dataclass. Shared fields that OrchestratorArgs - # also declares get copied into its constructor. - orch_init: dict[str, Any] = dict(orch_kwargs) - for key, value in shared_values.items(): - if key in orch_fields: - orch_init[key] = value - orch_args = OrchestratorArgs(**orch_init) - - return orch_args, engine_kwargs - - -def derive_server_dests_from_vllm_parser() -> frozenset[str]: - """Derive the set of argparse dests that belong to vllm's frontend/server. - - Returns every dest registered by ``make_arg_parser`` that is NOT a field - of ``OmniEngineArgs`` and NOT a field of ``OrchestratorArgs``. Useful for - CI tests to assert all CLI flags are classifiable without maintaining - a hardcoded server list. - - Returns empty frozenset if vllm's parser cannot be built (e.g. in a - minimal test environment). - """ - try: - from vllm.entrypoints.openai.cli_args import make_arg_parser - from vllm.utils.argparse_utils import FlexibleArgumentParser - except ImportError: - logger.debug("Cannot import vllm parser — server-dest derivation skipped") - return frozenset() - - try: - parser = make_arg_parser(FlexibleArgumentParser()) - all_dests = {a.dest for a in parser._actions if a.dest and a.dest != "help"} - except Exception as exc: - logger.debug("Failed to build vllm parser: %s", exc) - return frozenset() - - engine_fields = {f.name for f in fields(OmniEngineArgs)} - orch_fields = orchestrator_field_names() - - return frozenset(all_dests - engine_fields - orch_fields - SHARED_FIELDS) - - -def orchestrator_args_from_argparse(args: Any) -> OrchestratorArgs: - """Build an ``OrchestratorArgs`` from an ``argparse.Namespace``. - - Only copies attributes that exist on the namespace — missing fields fall - back to the dataclass default. Useful when the full parser is already - built and ``vars(args)`` would include noise. - """ - kwargs: dict[str, Any] = {} - for f in fields(OrchestratorArgs): - if hasattr(args, f.name): - value = getattr(args, f.name) - if value is not None or f.default is None: - kwargs[f.name] = value - return OrchestratorArgs(**kwargs) - - -def nullify_stage_engine_defaults(parser: argparse.ArgumentParser) -> None: - """Reset stage-level engine flag defaults to ``None``; preserve real - default in help text. Only deploy-YAML override fields are touched. - Idempotent.""" - from vllm_omni.config.stage_config import deploy_override_field_names - - override_dests = deploy_override_field_names() - - for action in parser._actions: - if action.dest in ("help", "version") or not action.option_strings: - continue - if action.dest not in override_dests: - continue - if action.default is None or action.default is argparse.SUPPRESS: - continue - if action.help and "(default:" not in action.help and "%(default)" not in action.help: - action.help = f"{action.help} (default: {action.default})" - action.default = None - - parser._omni_nullified = True # type: ignore[attr-defined] diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index dbfbe15e9..d5abe7809 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -21,7 +21,7 @@ import weakref from collections.abc import Awaitable, Callable, Mapping, Sequence from contextlib import ExitStack from dataclasses import asdict, dataclass -from typing import TYPE_CHECKING, Any, Literal, cast +from typing import Any, Literal, cast import janus import torch @@ -113,9 +113,6 @@ from vllm_omni.inputs.data import OmniSamplingParams from vllm_omni.metrics.prometheus import OmniRequestCounter from vllm_omni.platforms import current_omni_platform -if TYPE_CHECKING: - from vllm_omni.engine.arg_utils import OmniEngineArgs - logger = init_logger(__name__) _STARTUP_POLL_INTERVAL_S = 1.0 @@ -269,16 +266,17 @@ class AsyncOmniEngine: def __init__( self, model: str, - engine_args: OmniEngineArgs | None = None, stage_init_timeout: int = 300, init_timeout: int = 600, diffusion_batch_size: int = 1, single_stage_mode: bool = False, transfer_emitter: Any = None, log_stats: bool = False, + tokenizer: str | None = None, **kwargs: Any, ) -> None: self.model = model + self.tokenizer = tokenizer self.diffusion_batch_size = diffusion_batch_size startup_timeout = int(init_timeout) # Forwarded into Orchestrator so its _forward_to_next_stage path can @@ -294,22 +292,6 @@ class AsyncOmniEngine: logger.info(f"[AsyncOmniEngine] Initializing with model {model}") - # Merge tracked engine_args fields into kwargs; explicit kwargs take priority. - if engine_args is not None: - if not hasattr(engine_args, "_explicit_fields"): - raise TypeError( - "engine_args=OmniEngineArgs(...) is ambiguous under " - "sentinel-default precedence. Use " - "OmniEngineArgs.create(**explicit) or pass explicit kwargs " - "directly." - ) - ea_dict = engine_args.explicit_kwargs() - # Remove model since it is passed as a positional arg already. - ea_dict.pop("model", None) - kwargs = {**ea_dict, **kwargs} - - self.tokenizer: str | None = kwargs.get("tokenizer") - # ------------------------------------------------------------------ # # Single-stage mode detection # # ------------------------------------------------------------------ # @@ -2072,7 +2054,6 @@ class AsyncOmniEngine: stage_configs_path = kwargs.get("stage_configs_path", None) deploy_config_path = kwargs.pop("deploy_config", None) stage_overrides_json = kwargs.pop("stage_overrides", None) - kwargs.pop("_cli_explicit_keys", None) explicit_stage_configs = kwargs.pop("stage_configs", None) if explicit_stage_configs is not None: logger.warning( diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index d870e0950..81ac09048 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -482,6 +482,15 @@ def prepare_engine_environment() -> None: pass +def _maybe_set_qwen3_omni_moe_env(engine_args_dict: dict[str, Any]) -> None: + if ( + engine_args_dict.get("model_arch") == "Qwen3OmniMoeForConditionalGeneration" + and "VLLM_USE_FLASHINFER_MOE_FP16" not in os.environ + ): + os.environ["VLLM_USE_FLASHINFER_MOE_FP16"] = "0" + logger.info("[stage_init] Set VLLM_USE_FLASHINFER_MOE_FP16=0 for Qwen3-Omni stage") + + def split_devices_for_replicas( devices_str: str | None, num_replicas: int, @@ -762,6 +771,10 @@ def build_engine_args_dict( default_sp = _to_dict(getattr(stage_config, "default_sampling_params", {})) engine_args_dict["has_sampling_extra_args"] = bool(default_sp.get("extra_args")) + # TODO: Remove this after the performance regression is fixed + # Set VLLM_USE_FLASHINFER_MOE_FP16=0 for Qwen3-Omni to avoid performance regression + _maybe_set_qwen3_omni_moe_env(engine_args_dict) + return engine_args_dict diff --git a/vllm_omni/entrypoints/cli/benchmark/main.py b/vllm_omni/entrypoints/cli/benchmark/main.py index 865064d1e..d9d2268e8 100644 --- a/vllm_omni/entrypoints/cli/benchmark/main.py +++ b/vllm_omni/entrypoints/cli/benchmark/main.py @@ -9,7 +9,7 @@ from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG from vllm_omni.entrypoints.cli.benchmark.base import OmniBenchmarkSubcommandBase if typing.TYPE_CHECKING: - from vllm.utils.argparse_utils import FlexibleArgumentParser + from vllm_omni.utils.tracking_parser import TrackingArgumentParser class OmniBenchmarkSubcommand(CLISubcommand): @@ -25,7 +25,7 @@ class OmniBenchmarkSubcommand(CLISubcommand): def validate(self, args: argparse.Namespace) -> None: pass - def subparser_init(self, subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser: + def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgumentParser: bench_parser = subparsers.add_parser( self.name, description=self.help, usage=f"vllm {self.name} [options]" ) diff --git a/vllm_omni/entrypoints/cli/main.py b/vllm_omni/entrypoints/cli/main.py index 81b9e8061..1a42b08f7 100644 --- a/vllm_omni/entrypoints/cli/main.py +++ b/vllm_omni/entrypoints/cli/main.py @@ -24,10 +24,10 @@ def main(): os.environ["VLLM_LOGGING_COLOR"] = "1" from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG, cli_env_setup - from vllm.utils.argparse_utils import FlexibleArgumentParser import vllm_omni.entrypoints.cli.benchmark.main import vllm_omni.entrypoints.cli.serve + from vllm_omni.utils.tracking_parser import TrackingArgumentParser CMD_MODULES = [ vllm_omni.entrypoints.cli.serve, @@ -40,7 +40,7 @@ def main(): _ensure_vllm_platform() - parser = FlexibleArgumentParser( + parser = TrackingArgumentParser( description="vLLM OMNI CLI", epilog=VLLM_SUBCMD_PARSER_EPILOG.format(subcmd="[subcommand]"), ) diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index d59f3dd48..c7f9bb5f3 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -19,11 +19,10 @@ from vllm.entrypoints.cli.types import CLISubcommand from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG from vllm.logger import init_logger -from vllm.utils.argparse_utils import FlexibleArgumentParser -from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults from vllm_omni.entrypoints.cli.logo import log_logo from vllm_omni.entrypoints.openai.api_server import omni_run_server +from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace logger = init_logger(__name__) @@ -84,10 +83,10 @@ class OmniServeCommand(CLISubcommand): name = "serve" # Parser stashed at subparser_init so ``cmd`` can resolve each user-typed # flag to its real ``dest`` via the parser's action table. - _parser: FlexibleArgumentParser | None = None + _parser: TrackingArgumentParser @staticmethod - def cmd(args: argparse.Namespace) -> None: + def cmd(args: TrackingNamespace) -> None: if not os.environ.get("VLLM_DISABLE_LOG_LOGO"): os.environ["VLLM_DISABLE_LOG_LOGO"] = "1" log_logo() @@ -176,7 +175,7 @@ class OmniServeCommand(CLISubcommand): return validate_parsed_serve_args(args) - def subparser_init(self, subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser: + def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgumentParser: serve_parser = subparsers.add_parser( self.name, description=DESCRIPTION, @@ -682,22 +681,10 @@ class OmniServeCommand(CLISubcommand): # sandboxed globals dict via ``DummySelf``) doesn't fail on a NameError. type(self)._parser = serve_parser - nullify_stage_engine_defaults(serve_parser) return serve_parser -def _create_default_diffusion_stage_cfg(args: argparse.Namespace) -> list[dict[str, Any]]: - """Create default diffusion stage configuration. - - Uses AsyncOmniEngine's implementation which doesn't have OmegaConf - compatibility issues. - """ - from vllm_omni.engine.async_omni_engine import AsyncOmniEngine - - return AsyncOmniEngine._create_default_diffusion_stage_cfg(vars(args)) - - -def run_headless(args: argparse.Namespace) -> None: +def run_headless(args: TrackingNamespace) -> None: """Run a single stage in headless mode. Honors ``--omni-dp-size-local``: launches that many replicas locally for @@ -738,43 +725,39 @@ def run_headless(args: argparse.Namespace) -> None: stage_id: int | None = args.stage_id omni_master_address: str | None = args.omni_master_address omni_master_port: int | None = args.omni_master_port + worker_backend: str | None = args.worker_backend + stage_configs_path: str | None = args.stage_configs_path omni_replica_address: str | None = getattr(args, "omni_replica_address", None) omni_dp_size_local: int = max(1, int(getattr(args, "omni_dp_size_local", 1) or 1)) + if not model: + raise ValueError("Failed to pass model from kwargs") if stage_id is None: raise ValueError("--stage-id is required in headless mode") + if omni_master_address is None or omni_master_port is None: + raise ValueError("--omni-master-address and --omni-master-port are required in headless mode") + if worker_backend != "multi_process": + raise ValueError("headless mode requires worker_backend=multi_process") + + # Filter down to a dict of things explicitly requested by the user + args_dict = args.get_explicit_kwargs_dict() # ``--replica-id`` is deprecated and ignored — replica ids are # auto-assigned by ``OmniMasterServer`` so headless processes carry # no knowledge of their per-replica id at launch time. Warn (don't # error) when the operator still supplies it so existing launchers # keep working with a single log line. - explicit_cli_keys: set[str] = getattr(args, "_cli_explicit_keys", set()) or set() - if "replica_id" in explicit_cli_keys: + if "replica_id" in args_dict: logger.warning( "[Headless] --replica-id is deprecated and ignored " "(supplied value: %s). Replica ids are auto-assigned by the " "master server.", args.replica_id, ) - if omni_master_address is None or omni_master_port is None: - raise ValueError("--omni-master-address and --omni-master-port are required in headless mode") - api_server_count = args.api_server_count or 0 - if api_server_count > 1: - raise ValueError("api_server_count can't be set in headless mode") - if args.worker_backend != "multi_process": - raise ValueError("headless mode requires worker_backend=multi_process") - args_dict = vars(args).copy() - args_dict.pop("_cli_explicit_keys", None) - # Forward ``--deploy-config`` so the headless reads the same YAML the - # head was launched with — otherwise ``load_and_resolve_stage_configs`` - # falls back to ``vllm_omni/deploy/.yaml`` and the headless's - # view of ``stage.runtime.devices`` diverges from the head's, breaking - # the per-replica device split. config_path, stage_configs = load_and_resolve_stage_configs( model, - args_dict.get("stage_configs_path"), + stage_configs_path, args_dict, deploy_config_path=args_dict.get("deploy_config"), ) diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index f538645b4..b10dd777d 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -1,13 +1,10 @@ from __future__ import annotations -import argparse import os -import sys import time -import warnings import weakref from collections.abc import Sequence -from typing import TYPE_CHECKING, Any, Literal +from typing import Any, Literal import huggingface_hub from vllm.logger import init_logger @@ -29,9 +26,7 @@ from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetric from vllm_omni.metrics.transfer import OmniTransferMetrics from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific from vllm_omni.outputs import OmniRequestOutput - -if TYPE_CHECKING: - from vllm_omni.engine.arg_utils import OmniEngineArgs +from vllm_omni.utils.tracking_parser import TrackingNamespace logger = init_logger(__name__) @@ -105,48 +100,44 @@ class OmniBase(PDDisaggregationMixin): @classmethod def from_cli_args( cls, - args: argparse.Namespace, - *, - parser: argparse.ArgumentParser | None = None, - **overrides: Any, + args: TrackingNamespace, + model: str | None = None, ) -> OmniBase: - """Deprecated argparse builder. - - Build from argparse. If ``parser`` is passed and not yet nullified, - un-typed engine fields are reset to ``None``. New callers should - nullify deploy-overriding parser defaults with - ``nullify_stage_engine_defaults(parser)`` and construct Omni/AsyncOmni - directly. + """Build from a TrackingNamespace parsed by TrackingArgumentParser. + Only args that are explicitly passed to parse_args are forwarded. """ - warnings.warn( - "`from_cli_args()` is deprecated. Nullify deploy-overriding parser defaults " - "with `nullify_stage_engine_defaults(parser)` and construct Omni/AsyncOmni " - "directly from `vars(args)`.", - DeprecationWarning, - stacklevel=2, - ) - kwargs: dict[str, Any] = {k: v for k, v in vars(args).items() if not k.startswith("_")} + if not isinstance(args, TrackingNamespace): + raise TypeError( + f"expected args to be of type TrackingNamespace, got {type(args)}. " + "Hint: did you parse your args with TrackingArgumentParser?" + ) - if parser is not None and not getattr(parser, "_omni_nullified", False): - from vllm_omni.config.stage_config import deploy_override_field_names - from vllm_omni.entrypoints.utils import detect_explicit_cli_keys + explicit_kwargs = args.get_explicit_kwargs_dict() + args_model = explicit_kwargs.pop("model", None) or args.model + if model is not None and args_model is not None and model != args_model: + raise ValueError( + f"explicit model kwarg and args.model were both provided, but do not match [{model} != {args_model}]" + ) - explicit = detect_explicit_cli_keys(sys.argv[1:], parser) or set() - override_dests = deploy_override_field_names() - for key in list(kwargs): - if key in override_dests and key not in explicit: - kwargs[key] = None + if model is None and args_model is None: + raise ValueError( + "model must be explicitly passed as a parsed arg in the TrackingNamespace or directly provided." + ) - kwargs.update(overrides) - return cls(**kwargs) + resolved_model = model or args_model + return cls(model=resolved_model, **explicit_kwargs) def __init__( self, model: str, **kwargs: Any, ) -> None: - engine_args: OmniEngineArgs | None = kwargs.pop("engine_args", None) - + if "engine_args" in kwargs: + logger.warning( + "engine_args were passed as a kwarg to an Omni instance; this is not supported. " + "You should instead, pass the keyword arguments used to initialize the engine args " + "directly to this object's initializer." + ) stage_init_timeout = kwargs.pop("stage_init_timeout", 300) init_timeout = kwargs.pop("init_timeout", 600) log_stats = kwargs.pop("log_stats", False) @@ -180,7 +171,6 @@ class OmniBase(PDDisaggregationMixin): st = time.time() self.engine = AsyncOmniEngine( model=model, - engine_args=engine_args, init_timeout=init_timeout, stage_init_timeout=stage_init_timeout, diffusion_batch_size=diffusion_batch_size, diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index ca7182edb..2c7388531 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -136,6 +136,7 @@ from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_reques from vllm_omni.entrypoints.openai.video_api_utils import decode_input_reference from vllm_omni.entrypoints.openpi.serving import ServingRealtimeRobotOpenPI from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt +from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace logger = init_logger(__name__) router = APIRouter() @@ -510,7 +511,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, @asynccontextmanager async def build_async_omni( - args: Namespace, + args: TrackingNamespace, *, disable_frontend_multiprocessing: bool | None = None, client_config: dict[str, Any] | None = None, @@ -550,7 +551,7 @@ async def build_async_omni( @asynccontextmanager async def build_async_omni_from_stage_config( - args: Namespace, + args: TrackingNamespace, *, disable_frontend_multiprocessing: bool = False, ) -> AsyncIterator[EngineClient]: @@ -604,9 +605,9 @@ async def build_async_omni_from_stage_config( logger.debug("Pre-loading transformers_modules failed: %s", e) try: - kwargs = vars(args).copy() - kwargs.pop("model", None) - async_omni = AsyncOmni(model=args.model, **kwargs) + kwargs = args.get_explicit_kwargs_dict() + model = kwargs.pop("model", None) or args.model + async_omni = AsyncOmni(model=model, **kwargs) # # Don't keep the dummy data in memory # await async_llm.reset_mm_cache() @@ -3040,25 +3041,23 @@ async def omni_wakeup(request: OmniWakeupRequest, raw_request: Request): if __name__ == "__main__": - import argparse import asyncio from vllm.entrypoints.openai.cli_args import make_arg_parser - from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults - - parser = argparse.ArgumentParser(description="vLLM-Omni OpenAI-Compatible REST API server") + parser = TrackingArgumentParser(description="vLLM-Omni OpenAI-Compatible REST API server") parser = make_arg_parser(parser) registered_flags = set() for action in parser._actions: registered_flags.update(action.option_strings) + + # FIXME - this is broken on `main`. `make_arg_parser` does not handle omni engine args properly. if "--omni" not in registered_flags: parser.add_argument("--omni", action="store_true", default=False, help="Enable vLLM-Omni mode.") if "--enable-sleep-mode" not in registered_flags: parser.add_argument( "--enable-sleep-mode", action="store_true", default=False, help="Enable GPU memory pool for sleep mode." ) - nullify_stage_engine_defaults(parser) args = parser.parse_args() if not hasattr(args, "model_tag"): setattr(args, "model_tag", args.model) diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 031527895..dc3498dc9 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -1,4 +1,3 @@ -import argparse import os import types from collections import Counter @@ -24,81 +23,11 @@ PROJECT_ROOT = Path(__file__).parent.parent.parent logger = init_logger(__name__) -def _warn_deprecated_explicit_keys(kwargs: dict[str, Any]) -> None: - if "cli_explicit_keys" in kwargs: - import warnings - - warnings.warn( - "cli_explicit_keys= is deprecated and ignored. Remove the kwarg.", - DeprecationWarning, - stacklevel=3, - ) - - _DIFFUSERS_CLASS_TO_CONFIG: dict[str, str] = { "GlmImagePipeline": "glm_image", } -def detect_explicit_cli_keys( - argv: list[str], - parser: argparse.ArgumentParser | None = None, -) -> set[str]: - """Walk ``argv`` and return the set of ``dest`` attribute names the user - explicitly provided (e.g. ``--max-num-seqs 64`` → ``max_num_seqs``). - - Used to distinguish user-typed CLI args from argparse default values so - deploy YAMLs are not silently overridden by parser defaults. Shared - across online (``vllm serve``) and offline (scripts, examples, tests, - CI) entry points — offline callers that parse CLI args via argparse - should invoke this on ``sys.argv[1:]`` and pass the result through to - ``AsyncOmni`` / ``Omni`` via the ``_cli_explicit_keys`` kwarg. - - When ``parser`` is provided, each token is looked up in the parser's - action table to find its real ``dest``. This correctly handles flags - with ``dest=`` overrides, alias flags (e.g. ``--usp`` / - ``--ulysses-degree`` both mapping to ``ulysses_degree``), and - ``--disable-foo`` / ``store_false`` patterns that map to a differently - named dest. Callers with access to an ``argparse.ArgumentParser`` should - always pass it. - - When ``parser`` is ``None``, a name-based heuristic is used as a - fallback (hyphens → underscores, plus a ``no_`` prefix strip for - ``argparse.BooleanOptionalAction``). This is correct for simple flags - but silently misidentifies ``--disable-X``-style flags and explicit - ``dest=`` overrides, so prefer the parser-aware form. - """ - if parser is not None: - dest_map: dict[str, str] = {} - for action in parser._actions: - for opt in action.option_strings: - dest_map[opt] = action.dest - explicit: set[str] = set() - for tok in argv: - if not tok.startswith("--"): - continue - flag = tok.split("=", 1)[0] - dest = dest_map.get(flag) - if dest is not None: - explicit.add(dest) - return explicit - - # Fallback: name-based heuristic (legacy path for callers without a parser). - explicit = set() - for tok in argv: - if not tok.startswith("--"): - continue - name = tok[2:].split("=", 1)[0] - if not name: - continue - attr = name.replace("-", "_") - explicit.add(attr) - # BooleanOptionalAction: --no-foo records as dest `foo`, not `no_foo`. - if attr.startswith("no_"): - explicit.add(attr[3:]) - return explicit - - def inject_omni_kv_config(stage: Any, omni_conn_cfg: dict[str, Any], omni_from: str, omni_to: str) -> None: """Inject connector configuration into stage engine arguments.""" # Prepare omni_kv_config dict @@ -383,7 +312,6 @@ def load_stage_configs_from_model( base_engine_args: dict | None = None, deploy_config_path: str | None = None, stage_overrides: dict[str, dict[str, Any]] | None = None, - **deprecated_kwargs: Any, ) -> list: """Load stage configurations from model's default config file. @@ -402,8 +330,6 @@ def load_stage_configs_from_model( Returns: List of stage configuration dictionaries """ - _warn_deprecated_explicit_keys(deprecated_kwargs) - if base_engine_args is None: base_engine_args = {} @@ -577,7 +503,6 @@ def load_and_resolve_stage_configs( default_stage_cfg_factory: Any = None, deploy_config_path: str | None = None, stage_overrides: dict[str, dict[str, Any]] | None = None, - **deprecated_kwargs: Any, ) -> tuple[str, list]: """Load stage configurations from model or YAML file with fallback to defaults. @@ -621,8 +546,6 @@ def load_and_resolve_stage_configs( stage_configs_path, ) - _warn_deprecated_explicit_keys(deprecated_kwargs) - if deploy_config_path is not None: config_path = deploy_config_path stage_configs = load_stage_configs_from_model( diff --git a/vllm_omni/utils/tracking_parser.py b/vllm_omni/utils/tracking_parser.py new file mode 100644 index 000000000..c6dc5fee7 --- /dev/null +++ b/vllm_omni/utils/tracking_parser.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import argparse +from typing import Any + +from vllm.utils.argparse_utils import FlexibleArgumentParser + +UNSET = object() + + +class TrackingNamespace(argparse.Namespace): + """Proxy that wraps an argparse namespace with explicit keys, which + can be filtered down to a dict containing only explicitly passed values. + """ + + def __init__(self, unfiltered_ns: argparse.Namespace, explicit_keys: frozenset[str]) -> None: + # We never have nested tracking namespaces, but explicitly guard + # against them to prevent bad behavior with nested __dict__ overrides. + if isinstance(unfiltered_ns, TrackingNamespace): + raise ValueError("Tracking namespaces cannot be nested") + + self.unfiltered_ns = unfiltered_ns + self.explicit_keys = explicit_keys + + def __setattr__(self, name: str, value: Any) -> None: + if name in ("unfiltered_ns", "explicit_keys"): + object.__setattr__(self, name, value) + else: + setattr(self.unfiltered_ns, name, value) + + def get_explicit_kwargs_dict(self): + """Return a dict containing only the explicitly passed key-value pairs.""" + return {k: v for k, v in vars(self.unfiltered_ns).items() if k in self.explicit_keys} + + def __getattr__(self, name: str) -> Any: + return getattr(self.unfiltered_ns, name) + + @property + def __dict__(self): + # NOTE: We do this so that vars() etc forward directly into the encapsulated namespace, + # which makes this class a drop-in replacement for the original namespace, while also + # ensuring that updates to the encapsulated namespace are correctly reflected. + return self.unfiltered_ns.__dict__ + + +class TrackingGroup: + """Proxy that wraps an argument group and its corresponding shadow group.""" + + def __init__( + self, + real_group: argparse._ArgumentGroup, + shadow_group: argparse._ArgumentGroup, + ): + self._real = real_group + self._shadow = shadow_group + + def add_argument(self, *args: Any, **kwargs: Any) -> argparse.Action: + """Add an argument to the real group and to the shadow group.""" + action = self._real.add_argument(*args, **kwargs) + default_kwargs = {**kwargs, "default": UNSET} + self._shadow.add_argument(*args, **default_kwargs) + return action + + def __getattr__(self, name: str) -> Any: + # Any attribute access is forwarded to the real argument group. + return getattr(self._real, name) + + +class TrackingSubparsers: + """Proxy that wraps a subparser and its corresponding shadow subparser.""" + + def __init__( + self, + real_sub: argparse._SubParsersAction, + shadow_sub: argparse._SubParsersAction, + ): + self._real = real_sub + self._shadow = shadow_sub + + def add_parser(self, name, *args, **kwargs): + """Add a parser to the encapsulated real parser and its shadow.""" + real_parser = self._real.add_parser(name, *args, **kwargs) + # real_parser is a TrackingArgumentParser with its own _shadow. + # Reuse that shadow as the parent shadow's child — so when + # real_parser.add_argument() mirrors to real_parser._shadow, + # the parent's shadow sees it too. + self._shadow._name_parser_map[name] = real_parser._shadow + return real_parser + + def __getattr__(self, name: str) -> Any: + # Any attribute access is forwarded to the real subparser. + return getattr(self._real, name) + + +class TrackingArgumentParser(FlexibleArgumentParser): + """Drop-in replacement for FlexibleArgumentParser, which tracks keys that + were explicitly passed as args on the parser namespace. + + Unfortunately, Argparse does not provide an easy way of doing this without + depending on a lot of internal attributes, so we implement it by instead + using a 'shadow' parser, which is essentially a clone of the parser, where + defaults are overridden to `None`. By comparing the parser against its + shadow, we can tell which values were passed in a non-destructive manner. + """ + + def __init__(self, *args, **kwargs): + # NOTE: We have to define the shadow parser before calling init, + # with add_help=False, since otherwise init will call add_argument + # and delegate to the override on this class and cause problems. + shadow_kwargs = {**kwargs, "add_help": False} + self._shadow = FlexibleArgumentParser(*args, **shadow_kwargs) + super().__init__(*args, **kwargs) + + def add_argument(self, *args: Any, **kwargs: Any) -> argparse.Action: + """Add an arg to the parser & the shadow, where the latter has UNSET for the default.""" + action = super().add_argument(*args, **kwargs) + shadow_kwargs = {**kwargs, "default": UNSET} + self._shadow.add_argument(*args, **shadow_kwargs) + return action + + def add_argument_group(self, *args, **kwargs) -> TrackingGroup: + real_group = super().add_argument_group(*args, **kwargs) + shadow_group = self._shadow.add_argument_group(*args, **kwargs) + return TrackingGroup(real_group, shadow_group) + + def add_mutually_exclusive_group(self, *args, **kwargs) -> TrackingGroup: + real_group = super().add_mutually_exclusive_group(*args, **kwargs) + shadow_group: argparse._MutuallyExclusiveGroup = self._shadow.add_mutually_exclusive_group(*args, **kwargs) + return TrackingGroup(real_group, shadow_group) + + def add_subparsers(self, *args, **kwargs) -> TrackingSubparsers: + real_sub = super().add_subparsers(*args, **kwargs) + shadow_sub = self._shadow.add_subparsers(*args, **kwargs) + return TrackingSubparsers(real_sub, shadow_sub) + + def build_tracking_namespace(self, real_ns: argparse.Namespace, shadow_ns: argparse.Namespace) -> TrackingNamespace: + """Build a tracking namespace for the real / shadow namespaces.""" + explicit_keys = frozenset(k for k, v in vars(shadow_ns).items() if v is not UNSET) + return TrackingNamespace(real_ns, explicit_keys) + + def parse_args( + self, + args: list[str] | None = None, + namespace: argparse.Namespace | None = None, + ) -> TrackingNamespace: + """Parse the args on the real/shadow parser.""" + # Only the real parser should use the namespace if one is, + # given since shadow parser will set its own defaults to None. + real_ns = super().parse_args(args, namespace) + shadow_ns = self._shadow.parse_args(args) + if real_ns is None or shadow_ns is None: + raise ValueError("Parse args created empty namespaces") + + # If this is called through parse_known_args on self, we will already + # get a TrackingNamespace back, which will already have set the explicit + # keys through build_tracking_namespace, so no need to do it again. + if isinstance(real_ns, TrackingNamespace): + return real_ns + + return self.build_tracking_namespace(real_ns, shadow_ns) + + def parse_known_args( + self, + args: list[str] | None = None, + namespace: argparse.Namespace | None = None, + ) -> tuple[TrackingNamespace, list[str]]: + """Parse the known args on the real/shadow parser.""" + real_ns, remaining = super().parse_known_args(args, namespace) + shadow_ns, _ = self._shadow.parse_known_args(args) + tracked_ns = self.build_tracking_namespace(real_ns, shadow_ns) + + return tracked_ns, remaining