[Bugfix] Support passing vllm cli args to online serving in vLLM-Omni (#206)

Signed-off-by: Gaohan123 <gaohan19@huawei.com>
Signed-off-by: Gaohan123 <hgaoaf@connect.ust.hk>

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,

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
 
 

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(

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()

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