From 45dc189d8d76a82f9d1287fc378c16fb195f0476 Mon Sep 17 00:00:00 2001 From: Gao Han Date: Mon, 8 Dec 2025 14:59:53 +0800 Subject: [PATCH] [Bugfix] Support passing vllm cli args to online serving in vLLM-Omni (#206) Signed-off-by: Gaohan123 Signed-off-by: Gaohan123 --- vllm_omni/entrypoints/async_omni.py | 51 ++++++++------- vllm_omni/entrypoints/cli/serve.py | 35 ++++++++++ vllm_omni/entrypoints/omni_stage.py | 2 +- vllm_omni/entrypoints/openai/api_server.py | 5 +- vllm_omni/entrypoints/utils.py | 74 ++++++++++++++++++++-- 5 files changed, 131 insertions(+), 36 deletions(-) diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 50a6ff748..63ccc2c4f 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -3,6 +3,7 @@ import multiprocessing as mp import os import socket import time +from argparse import Namespace from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Any, Optional, Union @@ -61,21 +62,23 @@ class AsyncOmni(EngineClient): Args: model: Model name or path to load - stage_configs_path: Optional path to YAML file containing stage - configurations. If None, configurations are loaded from the model. - log_stats: Whether to enable statistics logging - log_file: Optional path prefix for log files. If provided, logs will - be written to files with stage-specific suffixes. - init_sleep_seconds: Number of seconds to sleep between starting - each stage process during initialization - shm_threshold_bytes: Threshold in bytes for using shared memory - for IPC. Objects larger than this threshold will use shared memory. - batch_timeout: Timeout in seconds for batching requests within a stage - init_timeout: Timeout in seconds for waiting for all stages to initialize + cli_args: Namespace object containing command-line arguments. + Expected attributes include: + - stage_configs_path: Optional path to YAML file containing stage + configurations. If None, configurations are loaded from the model. + - log_stats: Whether to enable statistics logging + - log_file: Optional path prefix for log files. If provided, logs will + be written to files with stage-specific suffixes. + - init_sleep_seconds: Number of seconds to sleep between starting + each stage process during initialization + - shm_threshold_bytes: Threshold in bytes for using shared memory + for IPC. Objects larger than this threshold will use shared memory. + - batch_timeout: Timeout in seconds for batching requests within a stage + - init_timeout: Timeout in seconds for waiting for all stages to initialize **kwargs: Additional keyword arguments passed to stage engines Example: - >>> async_llm = AsyncOmni(model="Qwen/Qwen2.5-Omni-7B") + >>> async_llm = AsyncOmni(model="Qwen/Qwen2.5-Omni-7B", cli_args=args) >>> async for output in async_llm.generate( ... prompt="Hello", ... request_id="req-1", @@ -87,33 +90,29 @@ class AsyncOmni(EngineClient): def __init__( self, model: str, - stage_configs_path: Optional[str] = None, - log_stats: bool = False, - log_file: Optional[str] = None, - init_sleep_seconds: int = 30, - shm_threshold_bytes: int = 65536, - batch_timeout: int = 10, - init_timeout: int = 60000, + cli_args: Namespace, **kwargs: Any, ): - self.batch_timeout = batch_timeout - self._enable_stats: bool = bool(log_stats) + self.batch_timeout = cli_args.batch_timeout + self._enable_stats: bool = bool(cli_args.log_stats) - if stage_configs_path is None: - self.stage_configs = load_stage_configs_from_model(model) + base_engine_args = AsyncOmniEngineArgs.from_cli_args(cli_args).__dict__.copy() + + if cli_args.stage_configs_path is None: + self.stage_configs = load_stage_configs_from_model(model, base_engine_args) else: - self.stage_configs = load_stage_configs_from_yaml(stage_configs_path) + self.stage_configs = load_stage_configs_from_yaml(cli_args.stage_configs_path, base_engine_args) self.stage_list: list[OmniStage] = [] self.default_sampling_params_list: list[SamplingParams] = [] # Optional file handler for orchestrator - self._log_file = log_file + self._log_file = cli_args.log_file if self._log_file: remove_old_logs(self._log_file, len(self.stage_configs)) configure_orchestrator_logger(logger, self._log_file) self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file) - self._initialize_stages(model, init_sleep_seconds, shm_threshold_bytes, init_timeout) + self._initialize_stages(model, cli_args.init_sleep_seconds, cli_args.shm_threshold_bytes, cli_args.init_timeout) def _initialize_stages( self, diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 4606029ab..321352759 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -60,6 +60,41 @@ class OmniServeCommand(CLISubcommand): default=None, help="Path to the stage configs file. If not specified, the stage configs will be loaded from the model.", ) + serve_parser.add_argument( + "--init-sleep-seconds", + type=int, + default=30, + help="The number of seconds to sleep before initializing the stages.", + ) + serve_parser.add_argument( + "--init-timeout", + type=int, + default=60000, + help="The timeout for initializing the stages.", + ) + serve_parser.add_argument( + "--shm-threshold-bytes", + type=int, + default=65536, + help="The threshold for the shared memory size.", + ) + serve_parser.add_argument( + "--log-stats", + action="store_true", + help="Enable logging the stats.", + ) + serve_parser.add_argument( + "--log-file", + type=str, + default=None, + help="The path to the log file.", + ) + serve_parser.add_argument( + "--batch-timeout", + type=int, + default=10, + help="The timeout for the batch.", + ) return serve_parser diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 5150ea27e..facf26a7d 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -735,7 +735,7 @@ async def _stage_worker_async( stage_id, list(engine_args.keys()), ) - omni_engine_args = AsyncOmniEngineArgs(model=model, **engine_args) + omni_engine_args = AsyncOmniEngineArgs(**engine_args) usage_context = UsageContext.OPENAI_API_SERVER vllm_config = omni_engine_args.create_engine_config(usage_context=usage_context) stage_engine = AsyncOmniStageLLM.from_vllm_config( diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index f40e579bc..c9b8dd2a1 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -185,10 +185,7 @@ async def build_async_omni_from_stage_config( async_omni: Optional[EngineClient] = None try: - if getattr(args, "stage_configs_path", None): - async_omni = AsyncOmni(model=args.model, stage_configs_path=args.stage_configs_path) - else: - async_omni = AsyncOmni(model=args.model) + async_omni = AsyncOmni(model=args.model, cli_args=args) # # Don't keep the dummy data in memory # await async_llm.reset_mm_cache() diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 57a6b7554..fcd8246f8 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -1,5 +1,8 @@ import os +from collections import Counter +from dataclasses import asdict, is_dataclass from pathlib import Path +from typing import Any from omegaconf import OmegaConf from vllm.transformers_utils.config import get_config @@ -10,7 +13,56 @@ from vllm_omni.utils import detect_device_type PROJECT_ROOT = Path(__file__).parent.parent.parent -def load_stage_configs_from_model(model: str) -> list: +def _convert_dataclasses_to_dict(obj: Any) -> Any: + """Recursively convert non-serializable objects to OmegaConf-compatible types. + + This is needed because OmegaConf cannot handle: + - Dataclass objects with Literal type annotations (e.g., StructuredOutputsConfig) + - Counter objects (from collections or vllm.utils) + - Set objects + - Other non-primitive types + """ + # IMPORTANT: Check Counter BEFORE dict, since Counter is a subclass of dict + # Handle Counter objects (convert to dict) + # Check by class name first to catch both collections.Counter and vllm.utils.Counter + if hasattr(obj, "__class__") and obj.__class__.__name__ == "Counter": + try: + return dict(obj) + except (TypeError, ValueError): + # If Counter can't be converted to dict, return empty dict + return {} + # Also check isinstance for collections.Counter (must be before dict check) + if isinstance(obj, Counter): + return dict(obj) + # Handle set objects (convert to list) + if isinstance(obj, set): + return list(obj) + # Handle dataclass objects + # Note: asdict() recursively converts nested dataclasses but not Counter objects, + # so we need to recursively process the result + if is_dataclass(obj): + result = asdict(obj) + # Recursively process the result to convert any Counter objects + return _convert_dataclasses_to_dict(result) + # Handle dictionaries (recurse into values) + # Note: This must come AFTER Counter check since Counter is a dict subclass + if isinstance(obj, dict): + return {k: _convert_dataclasses_to_dict(v) for k, v in obj.items()} + # Handle lists and tuples (recurse into items) + if isinstance(obj, (list, tuple)): + return type(obj)(_convert_dataclasses_to_dict(item) for item in obj) + # Try to convert any dict-like object (has keys/values methods) to dict + if hasattr(obj, "keys") and hasattr(obj, "values") and not isinstance(obj, (str, bytes)): + try: + return {k: _convert_dataclasses_to_dict(v) for k, v in obj.items()} + except (TypeError, ValueError, AttributeError): + # If conversion fails, return as-is + return obj + # Primitive types and other objects that OmegaConf can handle + return obj + + +def load_stage_configs_from_model(model: str, base_engine_args: dict = {}) -> list: """Load stage configurations from model's default config file. Loads stage configurations based on the model type and device type. @@ -35,7 +87,9 @@ def load_stage_configs_from_model(model: str) -> list: device_config_file = f"vllm_omni/model_executor/stage_configs/{device_type}/{model_type}.yaml" device_config_path = PROJECT_ROOT / device_config_file if os.path.exists(device_config_path): - stage_configs = load_stage_configs_from_yaml(config_path=str(device_config_path)) + stage_configs = load_stage_configs_from_yaml( + config_path=str(device_config_path), base_engine_args=base_engine_args + ) return stage_configs # Fall back to default config @@ -43,11 +97,11 @@ def load_stage_configs_from_model(model: str) -> list: stage_config_path = PROJECT_ROOT / stage_config_file if not os.path.exists(stage_config_path): raise FileNotFoundError(f"Stage config file {stage_config_path} not found") - stage_configs = load_stage_configs_from_yaml(config_path=str(stage_config_path)) + stage_configs = load_stage_configs_from_yaml(config_path=str(stage_config_path), base_engine_args=base_engine_args) return stage_configs -def load_stage_configs_from_yaml(config_path: str) -> list: +def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict = {}) -> list: """Load stage configurations from a YAML file. Args: @@ -57,4 +111,14 @@ def load_stage_configs_from_yaml(config_path: str) -> list: List of stage configuration dictionaries from the file's stage_args """ config_data = OmegaConf.load(config_path) - return config_data.stage_args + stage_args = config_data.stage_args + # Convert any nested dataclass objects to dicts before creating OmegaConf + base_engine_args = _convert_dataclasses_to_dict(base_engine_args) + base_engine_args = OmegaConf.create(base_engine_args) + for stage_arg in stage_args: + base_engine_args_tmp = base_engine_args.copy() + # Update base_engine_args with stage-specific engine_args if they exist + if hasattr(stage_arg, "engine_args") and stage_arg.engine_args is not None: + base_engine_args_tmp = OmegaConf.merge(base_engine_args_tmp, stage_arg.engine_args) + stage_arg.engine_args = base_engine_args_tmp + return stage_args