[Bugfix] Forward all vllm-omni serve command parameters to model (#985)

Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com>
Signed-off-by: Junhong Liu <ljh_lbj@163.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docs/cli/serve.md

@@ -0,0 +1,9 @@
+# vllm-omni serve
+
+## JSON CLI Arguments
+
+--8<-- "docs/cli/json_tip.inc.md"
+
+## Arguments
+
+--8<-- "docs/generated/argparse_omni/omni_serve.inc.md"

docs/mkdocs/hooks/generate_argparse.py

@@ -0,0 +1,257 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import ast
+import logging
+import re
+import sys
+from argparse import SUPPRESS, Action, ArgumentParser, HelpFormatter, _ArgumentGroup
+from collections.abc import Iterable
+from importlib.machinery import ModuleSpec
+from pathlib import Path
+from typing import Literal
+from unittest.mock import MagicMock
+
+from pydantic_core import core_schema
+
+
+class _FlexibleArgumentParser(ArgumentParser):
+    """Fallback parser for docs when vllm is unavailable.
+
+    Accepts the 'deprecated' kwarg used by vllm CLI and emits warnings
+    if a deprecated argument is actually provided.
+    """
+
+    _deprecated: set[Action] = set()
+    _deprecated_warned: set[str] = set()
+
+    if sys.version_info < (3, 13):
+
+        def parse_known_args(self, args=None, namespace=None):
+            namespace, args = super().parse_known_args(args, namespace)
+            for action in _FlexibleArgumentParser._deprecated:
+                if (
+                    hasattr(namespace, dest := action.dest)
+                    and getattr(namespace, dest) != action.default
+                    and dest not in _FlexibleArgumentParser._deprecated_warned
+                ):
+                    _FlexibleArgumentParser._deprecated_warned.add(dest)
+                    logger.warning("argument '%s' is deprecated", dest)
+            return namespace, args
+
+    def add_argument(self, *args, **kwargs):
+        deprecated = kwargs.pop("deprecated", False)
+        action = super().add_argument(*args, **kwargs)
+        if deprecated:
+            _FlexibleArgumentParser._deprecated.add(action)
+        return action
+
+    class _FlexibleArgumentGroup(_ArgumentGroup):
+        def add_argument(self, *args, **kwargs):
+            deprecated = kwargs.pop("deprecated", False)
+            action = super().add_argument(*args, **kwargs)
+            if deprecated:
+                _FlexibleArgumentParser._deprecated.add(action)
+            return action
+
+    def add_argument_group(self, *args, **kwargs):
+        group = self._FlexibleArgumentGroup(self, *args, **kwargs)
+        self._action_groups.append(group)
+        return group
+
+
+logger = logging.getLogger("mkdocs")
+
+# Define root and doc output directories
+ROOT_DIR = Path(__file__).parent.parent.parent.parent
+ARGPARSE_DOC_DIR = ROOT_DIR / "docs/generated/argparse_omni"
+
+# Ensure the repo root is in sys.path for dynamic imports
+sys.path.insert(0, str(ROOT_DIR))
+
+
+class PydanticMagicMock(MagicMock):
+    """`MagicMock` that's able to generate pydantic-core schemas and avoids _mock_name=None errors."""
+
+    def __init__(self, *args, **kwargs):
+        name = kwargs.pop("name", None)
+        super().__init__(*args, **kwargs)
+        self.__spec__ = ModuleSpec(name, None)
+        self._mock_name = name or "mock"
+
+    def __get_pydantic_core_schema__(self, source_type, handler):
+        return core_schema.any_schema()
+
+
+# --- Static extraction for CLI argument docs ---
+
+
+def extract_omni_serve_subparser_init():
+    """
+    Statically parse vllm_omni/entrypoints/cli/serve.py to extract the subparser_init method
+    and return a callable that adds arguments to a parser. This avoids import and mock issues.
+    """
+    serve_path = ROOT_DIR / "vllm_omni" / "entrypoints" / "cli" / "serve.py"
+    with open(serve_path, encoding="utf-8") as f:
+        source = f.read()
+    tree = ast.parse(source, filename=str(serve_path))
+    # Find class OmniServeCommand
+    for node in tree.body:
+        if isinstance(node, ast.ClassDef) and node.name == "OmniServeCommand":
+            for item in node.body:
+                if isinstance(item, ast.FunctionDef) and item.name == "subparser_init":
+                    # We'll exec this function body in a dummy context
+                    func_code = ast.Module(body=[item], type_ignores=[])
+                    code = compile(func_code, filename=str(serve_path), mode="exec")
+                    # Prepare dummy context
+                    local_vars = {}
+
+                    # Provide a dummy subparsers with add_parser returning our parser
+                    class DummySubparsers:
+                        def add_parser(self, name, **kwargs):
+                            return _FlexibleArgumentParser(prog=name)
+
+                    dummy_subparsers = DummySubparsers()
+                    # Provide globals for exec
+                    # Extract DESCRIPTION from the source file (assume it's a triple-quoted string at the top)
+                    m = re.search(r'DESCRIPTION\s*=\s*([ru]?""".*?""")', source, re.DOTALL)
+                    DESCRIPTION = m.group(1)[3:-3] if m else ""
+                    exec_globals = {
+                        "_FlexibleArgumentParser": _FlexibleArgumentParser,
+                        "FlexibleArgumentParser": _FlexibleArgumentParser,
+                        "make_arg_parser": lambda parser: parser,  # no-op for doc
+                        "VLLM_SUBCMD_PARSER_EPILOG": "",
+                        "logger": logger,
+                        "DummySubparsers": DummySubparsers,
+                        "argparse": __import__("argparse"),
+                        "DESCRIPTION": DESCRIPTION,
+                    }
+                    exec(code, exec_globals, local_vars)
+                    # Get the function
+                    subparser_init = local_vars["subparser_init"]
+
+                    # Return a callable that mimics OmniServeCommand().subparser_init
+                    def parser_factory():
+                        class DummySelf:
+                            name = "serve"
+
+                        return subparser_init(DummySelf(), dummy_subparsers)
+
+                    return parser_factory
+    raise RuntimeError("Could not statically extract OmniServeCommand.subparser_init")
+
+
+OmniServeCommand = type("OmniServeCommand", (), {"subparser_init": staticmethod(extract_omni_serve_subparser_init())})
+
+
+class MarkdownFormatter(HelpFormatter):
+    """Custom formatter that generates markdown for argument groups."""
+
+    def __init__(self, prog: str, starting_heading_level: int = 3):
+        super().__init__(prog, max_help_position=sys.maxsize, width=sys.maxsize)
+
+        self._section_heading_prefix = "#" * starting_heading_level
+        self._argument_heading_prefix = "#" * (starting_heading_level + 1)
+        self._markdown_output = []
+
+    def start_section(self, heading: str):
+        if heading not in {"positional arguments", "options"}:
+            heading_md = f"\n{self._section_heading_prefix} {heading}\n\n"
+            self._markdown_output.append(heading_md)
+
+    def end_section(self):
+        pass
+
+    def add_text(self, text: str):
+        if text:
+            self._markdown_output.append(f"{text.strip()}\n\n")
+
+    def add_usage(self, usage, actions, groups, prefix=None):
+        pass
+
+    def add_arguments(self, actions: Iterable[Action]):
+        for action in actions:
+            if len(action.option_strings) == 0 or "--help" in action.option_strings:
+                continue
+
+            option_strings = f"`{'`, `'.join(action.option_strings)}`"
+            heading_md = f"{self._argument_heading_prefix} {option_strings}\n\n"
+            self._markdown_output.append(heading_md)
+
+            if choices := action.choices:
+                choices = f"`{'`, `'.join(str(c) for c in choices)}`"
+                self._markdown_output.append(f"Possible choices: {choices}\n\n")
+            elif (metavar := action.metavar) and isinstance(metavar, (list, tuple)):
+                metavar = f"`{'`, `'.join(str(m) for m in metavar)}`"
+                self._markdown_output.append(f"Possible choices: {metavar}\n\n")
+
+            if action.help:
+                self._markdown_output.append(f"{action.help}\n\n")
+
+            if (default := action.default) != SUPPRESS:
+                # Make empty string defaults visible
+                if default == "":
+                    default = '""'
+                self._markdown_output.append(f"Default: `{default}`\n\n")
+
+    def format_help(self):
+        """Return the formatted help as markdown."""
+        return "".join(self._markdown_output)
+
+
+# Function to create parser using subparser_init style CLI class
+
+
+def create_parser_subparser_init(subcmd_class):
+    """
+    Create an argparse parser using subparser_init style CLI class, with MarkdownFormatter.
+    Compatible with both AST-extracted static parser_factory and real class instance method.
+    """
+
+    class DummySubparsers:
+        def add_parser(self, name, **kwargs):
+            return _FlexibleArgumentParser(prog=name)
+
+    dummy_subparsers = DummySubparsers()
+    # If subparser_init is a zero-arg static function (AST-extracted), call it directly
+    subparser_init = getattr(subcmd_class, "subparser_init", None)
+    if subparser_init is not None:
+        import inspect
+
+        sig = inspect.signature(subparser_init)
+        if len(sig.parameters) == 0:
+            parser = subparser_init()
+        else:
+            parser = subcmd_class().subparser_init(dummy_subparsers)
+    else:
+        parser = subcmd_class().subparser_init(dummy_subparsers)
+    parser.formatter_class = MarkdownFormatter
+    return parser
+
+
+def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
+    """
+    Entry point for doc generation. Builds doc directory and outputs markdown for each CLI command.
+    """
+    logger.info("Generating vllm-omni argparse documentation")
+    logger.debug("Root directory: %s", ROOT_DIR.resolve())
+    logger.debug("Output directory: %s", ARGPARSE_DOC_DIR.resolve())
+    if not ARGPARSE_DOC_DIR.exists():
+        ARGPARSE_DOC_DIR.mkdir(parents=True)
+
+    # Register all CLI parsers; you can easily add more commands here
+    parsers = {
+        "omni_serve": create_parser_subparser_init(OmniServeCommand),
+        # "another_cmd": create_parser_subparser_init(AnotherCommandClass),
+    }
+
+    for stem, parser in parsers.items():
+        doc_path = ARGPARSE_DOC_DIR / f"{stem}.inc.md"
+        with open(doc_path, "w", encoding="utf-8") as f:
+            f.write(super(type(parser), parser).format_help())
+        logger.info("Argparse generated: %s", doc_path.relative_to(ROOT_DIR))
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    on_startup("build", False)

mkdocs.yml

@@ -59,6 +59,7 @@ hooks:
   - docs/mkdocs/hooks/generate_api_readme.py
   - docs/mkdocs/hooks/url_schemes.py
   - docs/mkdocs/hooks/generate_examples.py
+  - docs/mkdocs/hooks/generate_argparse.py
 
 # Exclude include files from navigation warnings
 exclude_docs: |

tests/entrypoints/test_utils.py

@@ -0,0 +1,271 @@
+"""Unit tests for vllm_omni.entrypoints.utils module."""
+
+from collections import Counter
+from dataclasses import dataclass
+from unittest.mock import patch
+
+import pytest
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.engine.arg_utils import OmniEngineArgs
+from vllm_omni.entrypoints.utils import (
+    _convert_dataclasses_to_dict,
+    _filter_dict_like_object,
+    filter_dataclass_kwargs,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestFilterDictLikeObject:
+    """Test suite for _filter_dict_like_object function."""
+
+    def test_simple_dict(self):
+        """Test filtering a simple dictionary with no callables."""
+        input_dict = {"key1": "value1", "key2": 42, "key3": [1, 2, 3]}
+        result = _filter_dict_like_object(input_dict)
+
+        assert result == input_dict
+        assert isinstance(result, dict)
+
+    def test_dict_with_nested_values(self):
+        """Test filtering dict with nested dict and list values."""
+        input_dict = {
+            "level1": {
+                "level2": {"key": "value"},
+                "list": [1, 2, 3],
+            },
+            "simple": "string",
+        }
+
+        result = _filter_dict_like_object(input_dict)
+
+        # Nested dicts and lists should be recursively processed
+        assert result["simple"] == "string"
+        assert isinstance(result["level1"], dict)
+
+    def test_dict_with_dataclass_values(self):
+        """Test filtering dict containing dataclass values."""
+
+        @dataclass
+        class TestDataclass:
+            field1: str
+            field2: int
+
+        obj = TestDataclass(field1="test", field2=42)
+        input_dict = {"data": obj, "normal": "value"}
+
+        result = _filter_dict_like_object(input_dict)
+
+        # Dataclass should be converted to dict by recursive _convert_dataclasses_to_dict
+        assert "data" in result
+        assert "normal" in result
+        assert result["normal"] == "value"
+
+    def test_dict_with_counter_values(self):
+        """Test filtering dict containing Counter objects."""
+        counter_obj = Counter({"a": 1, "b": 2})
+        input_dict = {"counter": counter_obj, "normal": "value"}
+
+        result = _filter_dict_like_object(input_dict)
+
+        # Counter should be converted to regular dict
+        assert "counter" in result
+        assert "normal" in result
+        assert result["normal"] == "value"
+
+    def test_empty_dict(self):
+        """Test filtering an empty dictionary."""
+        result = _filter_dict_like_object({})
+        assert result == {}
+        assert isinstance(result, dict)
+
+    def test_dict_with_set_values(self):
+        """Test filtering dict with set values."""
+        input_dict = {"set_key": {1, 2, 3}, "normal": "value"}
+
+        result = _filter_dict_like_object(input_dict)
+
+        assert "set_key" in result
+        assert "normal" in result
+        # Set should be converted to list by _convert_dataclasses_to_dict
+        assert result["normal"] == "value"
+
+    def test_dict_with_none_values(self):
+        """Test filtering dict with None values."""
+        input_dict = {"key1": None, "key2": "value", "key3": 0}
+
+        result = _filter_dict_like_object(input_dict)
+
+        assert result == input_dict
+
+    def test_dict_with_mixed_types(self):
+        """Test filtering dict with mixed value types."""
+        input_dict = {
+            "string": "hello",
+            "int": 42,
+            "float": 3.14,
+            "bool": True,
+            "none": None,
+            "list": [1, 2, 3],
+            "tuple": (1, 2, 3),
+            "set": {1, 2, 3},
+            "dict": {"nested": "value"},
+        }
+
+        result = _filter_dict_like_object(input_dict)
+
+        assert "string" in result
+        assert "int" in result
+        assert "float" in result
+        assert "bool" in result
+        assert "none" in result
+        assert "list" in result
+        assert "tuple" in result
+        assert "set" in result
+        assert "dict" in result
+
+    def test_dict_preserves_key_types(self):
+        """Test that dict key types are preserved."""
+        input_dict = {
+            "string_key": "value1",
+            42: "value2",
+            (1, 2): "value3",  # tuple as key
+        }
+
+        result = _filter_dict_like_object(input_dict)
+
+        # Keys should remain the same
+        assert "string_key" in result
+        assert 42 in result
+        assert (1, 2) in result
+
+    def test_dict_with_recursive_structure(self):
+        """Test filtering dict with recursive/complex nested structure."""
+        input_dict = {
+            "level1": {
+                "level2": {
+                    "level3": {"key": "value"},
+                    "callable": lambda x: x,
+                }
+            },
+            "normal": "value",
+        }
+
+        with patch("vllm_omni.entrypoints.utils.logger"):
+            result = _filter_dict_like_object(input_dict)
+
+            # Normal key should exist
+            assert "normal" in result
+            # Level1 should exist
+            assert "level1" in result
+
+    def test_integration_with_convert_dataclasses(self):
+        """Test that _filter_dict_like_object integrates properly with _convert_dataclasses_to_dict."""
+
+        @dataclass
+        class Config:
+            name: str
+            count: int
+
+        input_dict = {
+            "config": Config(name="test", count=5),
+            "func": lambda x: x,
+            "normal": "value",
+        }
+
+        with patch("vllm_omni.entrypoints.utils.logger"):
+            result = _filter_dict_like_object(input_dict)
+
+            # Callable should be filtered
+            assert "func" not in result
+            # Config should be converted to dict
+            assert "config" in result
+            assert "normal" in result
+
+
+class TestConvertDataclassesToDict:
+    """Test suite for _convert_dataclasses_to_dict function."""
+
+    def test_uses_filter_dict_like_object(self):
+        """Test that _convert_dataclasses_to_dict uses _filter_dict_like_object for dicts."""
+        input_dict = {
+            "normal": "value",
+            "callable": lambda x: x,
+        }
+
+        with patch("vllm_omni.entrypoints.utils.logger"):
+            result = _convert_dataclasses_to_dict(input_dict)
+
+            # Callable should be filtered out by _filter_dict_like_object
+            assert "normal" in result
+            assert "callable" not in result
+
+
+class TestFilterDataclassKwargs:
+    """Test basic functionality of filter_dataclass_kwargs."""
+
+    def test_simple_filtering(self):
+        """Test basic dataclass kwargs filtering."""
+
+        @dataclass
+        class SimpleConfig:
+            name: str
+            count: int
+
+        kwargs = {"name": "test", "count": 42, "invalid": "should_be_removed"}
+        result = filter_dataclass_kwargs(SimpleConfig, kwargs)
+
+        assert "name" in result
+        assert "count" in result
+        assert "invalid" not in result
+
+    def test_invalid_dataclass_raises_error(self):
+        """Test that non-dataclass raises ValueError."""
+        with pytest.raises(ValueError, match="is not a dataclass"):
+            filter_dataclass_kwargs(dict, {})
+
+    def test_invalid_kwargs_type_raises_error(self):
+        """Test that non-dict kwargs raises ValueError."""
+
+        @dataclass
+        class SimpleConfig:
+            name: str
+
+        with pytest.raises(ValueError, match="kwargs must be a dictionary"):
+            filter_dataclass_kwargs(SimpleConfig, "invalid")
+
+    def test_filters_omni_engine_args_unknown_fields(self):
+        """Test that OmniEngineArgs kwargs are filtered to valid fields only."""
+        kwargs = {
+            "model": "dummy",
+            "stage_id": 1,
+            "engine_output_type": "image",
+            "unknown_field": "drop_me",
+        }
+
+        result = filter_dataclass_kwargs(OmniEngineArgs, kwargs)
+
+        assert "model" in result
+        assert "stage_id" in result
+        assert "engine_output_type" in result
+        assert "unknown_field" not in result
+
+    def test_filters_omni_diffusion_config_union_dataclass(self):
+        """Test that OmniDiffusionConfig filters nested dataclass in Union fields."""
+        kwargs = {
+            "model": "dummy",
+            "cache_config": {
+                "rel_l1_thresh": 0.3,
+                "extra_param": "should_drop",
+            },
+            "unknown_top": "drop_me",
+        }
+
+        result = filter_dataclass_kwargs(OmniDiffusionConfig, kwargs)
+
+        assert "model" in result
+        assert "cache_config" in result
+        assert "unknown_top" not in result
+        assert result["cache_config"]["rel_l1_thresh"] == 0.3
+        assert "extra_param" not in result["cache_config"]

vllm_omni/entrypoints/omni.py

@@ -38,7 +38,6 @@ from vllm_omni.entrypoints.omni_stage import OmniStage
 from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK, OmniStageTaskType
 from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
 from vllm_omni.entrypoints.utils import (
-    build_base_engine_args,
     get_final_stage_id_for_e2e,
     inject_omni_kv_config,
     load_and_resolve_stage_configs,
@@ -255,7 +254,6 @@ class OmniBase:
         # TODO(wuhang):
         # Remove kwargs as parameters in the future.
         # Use dataclass directly for engine args.
-        base_engine_args = build_base_engine_args(kwargs)
 
         stage_configs_path = kwargs.get("stage_configs_path", None)
 
@@ -263,7 +261,7 @@ class OmniBase:
         config_path, stage_configs = load_and_resolve_stage_configs(
             model,
             stage_configs_path,
-            base_engine_args,
+            kwargs,
             default_stage_cfg_factory=lambda: self._create_default_diffusion_stage_cfg(kwargs),
         )
 

vllm_omni/entrypoints/omni_llm.py

@@ -22,6 +22,7 @@ from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.engine.input_processor import OmniInputProcessor
 from vllm_omni.engine.output_processor import MultimodalOutputProcessor
 from vllm_omni.entrypoints.utils import (
+    filter_dataclass_kwargs,
     load_stage_configs_from_model,
     load_stage_configs_from_yaml,
     resolve_model_config_path,
@@ -151,7 +152,7 @@ class OmniLLM(LLM):
             compilation_config=compilation_config_instance,
             structured_outputs_config=structured_outputs_instance,
             omni_kv_config=omni_kv_config,
-            **kwargs,
+            **filter_dataclass_kwargs(OmniEngineArgs, kwargs),
         )
 
         # Create the Engine (autoselects V0 vs V1)

vllm_omni/entrypoints/omni_stage.py

@@ -36,7 +36,7 @@ from vllm_omni.distributed.omni_connectors import build_stage_connectors
 from vllm_omni.distributed.omni_connectors.adapter import try_recv_via_connector
 from vllm_omni.distributed.omni_connectors.connectors.base import OmniConnectorBase
 from vllm_omni.distributed.ray_utils.utils import kill_ray_actor, start_ray_actor
-from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs
+from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs, OmniEngineArgs
 from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion
 from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM
 from vllm_omni.entrypoints.omni_diffusion import OmniDiffusion
@@ -49,7 +49,7 @@ from vllm_omni.entrypoints.stage_utils import (
     maybe_dump_to_shm,
     set_stage_devices,
 )
-from vllm_omni.entrypoints.utils import detect_pid_host
+from vllm_omni.entrypoints.utils import detect_pid_host, filter_dataclass_kwargs
 from vllm_omni.entrypoints.zmq_utils import (
     ZmqQueue,
     create_zmq_queue,
@@ -212,7 +212,8 @@ def _resolve_worker_cls(engine_args: dict[str, Any]) -> None:
     worker_type = engine_args.get("worker_type", None)
     if not worker_type:
         return
-    if engine_args.get("worker_cls"):
+    worker_cls = engine_args.get("worker_cls")
+    if worker_cls is not None and worker_cls != "auto":
         return
     from vllm_omni.platforms import current_omni_platform
 
@@ -268,6 +269,12 @@ class OmniStage:
         self.is_comprehension = getattr(stage_config, "is_comprehension", False)
         # Support for different stage types: "llm" (default) or "diffusion"
         self.stage_type: Literal["llm", "diffusion"] = getattr(stage_config, "stage_type", "llm")
+        if (
+            "stage_id" in stage_config.engine_args
+            and stage_config.engine_args.stage_id != self.stage_id
+            and self.stage_id is not None
+        ):
+            stage_config.engine_args.stage_id = self.stage_id
         if hasattr(stage_config, "custom_process_input_func"):
             # Import the module specified in the config (already a full module path)
             module_path, func_name = stage_config.custom_process_input_func.rsplit(".", 1)
@@ -746,6 +753,7 @@ def _stage_worker(
             engine_args["stage_connector_spec"] = stage_connector_spec
             engine_args["stage_id"] = stage_id
         if stage_type == "diffusion":
+            engine_args = filter_dataclass_kwargs(OmniDiffusionConfig, engine_args)
             engine_args.pop("model_stage", None)
             engine_args.pop("model", None)
             stage_engine = OmniDiffusion(
@@ -755,6 +763,8 @@ def _stage_worker(
                 **engine_args,
             )
         else:
+            engine_args = filter_dataclass_kwargs(OmniEngineArgs, engine_args)
+            engine_args.pop("model", None)
             # Default to LLM engine
             stage_engine = OmniLLM(model=model, **engine_args)
 
@@ -1162,6 +1172,7 @@ async def _stage_worker_async(
             engine_args["stage_id"] = stage_id
         if stage_type == "diffusion":
             # For diffusion, we need to extract diffusion-specific config
+            engine_args = filter_dataclass_kwargs(OmniDiffusionConfig, engine_args)
             od_config = _build_od_config(engine_args, model)
 
             # Inject omni config for worker to access stage info
@@ -1178,6 +1189,8 @@ async def _stage_worker_async(
             )
             vllm_config = None  # Diffusion doesn't use vllm_config
         else:
+            engine_args = filter_dataclass_kwargs(AsyncOmniEngineArgs, engine_args)
+            engine_args.pop("model", None)
             omni_engine_args = AsyncOmniEngineArgs(model=model, **engine_args)
             usage_context = UsageContext.OPENAI_API_SERVER
             vllm_config = omni_engine_args.create_engine_config(usage_context=usage_context)

vllm_omni/entrypoints/utils.py

@@ -1,8 +1,9 @@
 import os
+import types
 from collections import Counter
-from dataclasses import asdict, is_dataclass
+from dataclasses import asdict, fields, is_dataclass
 from pathlib import Path
-from typing import Any
+from typing import Any, get_args, get_origin
 
 from omegaconf import OmegaConf
 from vllm.logger import init_logger
@@ -68,6 +69,52 @@ def _try_get_class_name_from_diffusers_config(model: str) -> str | None:
     return None
 
 
+def _filter_dict_like_object(obj: dict | Any) -> dict:
+    """Filter dict-like object by removing callables and recursively converting values.
+
+    Converts dict-like objects to regular dicts while filtering out callable values
+    that are incompatible with OmegaConf. Recursively processes values through
+    _convert_dataclasses_to_dict for nested object conversion.
+
+    Args:
+        obj: Dict or dict-like object to filter
+
+    Returns:
+        Regular dict with callables filtered out and values recursively converted
+
+    Raises:
+        TypeError: If obj doesn't support .items() method
+        ValueError: If dict conversion fails unexpectedly
+    """
+
+    def _is_callable_value(value: Any) -> bool:
+        if callable(value):
+            return True
+        return isinstance(
+            value,
+            (
+                types.FunctionType,
+                types.MethodType,
+                types.BuiltinFunctionType,
+                types.BuiltinMethodType,
+            ),
+        )
+
+    result = {}
+    filtered_keys = []
+    for k, v in obj.items():
+        if _is_callable_value(v):
+            filtered_keys.append(str(k))
+        else:
+            result[k] = _convert_dataclasses_to_dict(v)
+    if filtered_keys:
+        logger.warning(
+            f"Filtered out {len(filtered_keys)} callable object(s) from base_engine_args "
+            f"that are not compatible with OmegaConf: {filtered_keys}. "
+        )
+    return result
+
+
 def _convert_dataclasses_to_dict(obj: Any) -> Any:
     """Recursively convert non-serializable objects to OmegaConf-compatible types.
 
@@ -75,6 +122,7 @@ def _convert_dataclasses_to_dict(obj: Any) -> Any:
     - Dataclass objects with Literal type annotations (e.g., StructuredOutputsConfig)
     - Counter objects (from collections or vllm.utils)
     - Set objects
+    - Callable objects (functions, methods, etc.)
     - Other non-primitive types
     """
     # IMPORTANT: Check Counter BEFORE dict, since Counter is a subclass of dict
@@ -99,17 +147,21 @@ def _convert_dataclasses_to_dict(obj: Any) -> Any:
         result = asdict(obj)
         # Recursively process the result to convert any Counter objects
         return _convert_dataclasses_to_dict(result)
-    # Handle dictionaries (recurse into values)
+    # Handle dictionaries (recurse into values) and filter out callables(cause error in OmegaConf.create)
     # 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()}
+        return _filter_dict_like_object(obj)
+    # Handle callable objects (functions, methods, etc.) - skip them
+    # Note: This comes after dict/list checks to avoid misclassifying dict-like objects
+    if callable(obj):
+        return None
     # Handle lists and tuples (recurse into items)
     if isinstance(obj, (list, tuple)):
-        return type(obj)(_convert_dataclasses_to_dict(item) for item in obj)
+        return type(obj)(_convert_dataclasses_to_dict(item) for item in obj if not callable(item))
     # 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()}
+            return _filter_dict_like_object(obj)
         except (TypeError, ValueError, AttributeError):
             # If conversion fails, return as-is
             return obj
@@ -239,7 +291,7 @@ def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict | None
 def load_and_resolve_stage_configs(
     model: str,
     stage_configs_path: str | None,
-    base_engine_args: dict | None,
+    kwargs: dict | None,
     default_stage_cfg_factory: Any = None,
 ) -> tuple[str, list]:
     """Load stage configurations from model or YAML file with fallback to defaults.
@@ -247,7 +299,7 @@ def load_and_resolve_stage_configs(
     Args:
         model: Model name or path
         stage_configs_path: Optional path to YAML file containing stage configurations
-        base_engine_args: Base engine arguments to merge with stage configs
+        kwargs: Engine arguments to merge with stage configs
         default_stage_cfg_factory: Optional callable that takes no args and returns
             default stage config list when no configs are found
 
@@ -256,7 +308,7 @@ def load_and_resolve_stage_configs(
     """
     if stage_configs_path is None:
         config_path = resolve_model_config_path(model)
-        stage_configs = load_stage_configs_from_model(model, base_engine_args=base_engine_args)
+        stage_configs = load_stage_configs_from_model(model, base_engine_args=kwargs)
         if not stage_configs:
             if default_stage_cfg_factory is not None:
                 default_stage_cfg = default_stage_cfg_factory()
@@ -265,7 +317,7 @@ def load_and_resolve_stage_configs(
                 stage_configs = []
     else:
         config_path = stage_configs_path
-        stage_configs = load_stage_configs_from_yaml(stage_configs_path, base_engine_args=base_engine_args)
+        stage_configs = load_stage_configs_from_yaml(stage_configs_path, base_engine_args=kwargs)
 
     return config_path, stage_configs
 
@@ -316,6 +368,69 @@ def get_final_stage_id_for_e2e(
     return final_stage_id_for_e2e
 
 
+def filter_dataclass_kwargs(cls: Any, kwargs: dict) -> dict:
+    """Filter kwargs to only include fields defined in the dataclass.
+
+    Args:
+        cls: Dataclass type
+        kwargs: Keyword arguments to filter
+
+    Returns:
+        Filtered keyword arguments containing only valid dataclass fields
+    """
+    if not is_dataclass(cls):
+        raise ValueError(f"{cls} is not a dataclass")
+    if not isinstance(kwargs, dict):
+        raise ValueError("kwargs must be a dictionary")
+
+    def _filter_value(value: Any, annotation: Any) -> Any:
+        """Recursively filter nested dict/list values based on dataclass annotations."""
+        if annotation is None:
+            return value
+
+        origin = get_origin(annotation)
+        if origin is None:
+            if isinstance(annotation, type) and is_dataclass(annotation) and isinstance(value, dict):
+                return filter_dataclass_kwargs(annotation, value)
+            return value
+
+        if origin in (list, tuple, set):
+            args = get_args(annotation)
+            inner = args[0] if args else None
+            if isinstance(value, (list, tuple, set)):
+                return type(value)(_filter_value(v, inner) for v in value)
+            return value
+
+        if origin is dict:
+            args = get_args(annotation)
+            val_type = args[1] if len(args) > 1 else None
+            if isinstance(value, dict):
+                return {k: _filter_value(v, val_type) for k, v in value.items()}
+            return value
+
+        if origin is types.UnionType or origin is getattr(types, "UnionType", None):
+            for arg in get_args(annotation):
+                if isinstance(arg, type) and is_dataclass(arg) and isinstance(value, dict):
+                    return filter_dataclass_kwargs(arg, value)
+                # Try container-style filtering for union members
+                filtered = _filter_value(value, arg)
+                if filtered is not value:
+                    return filtered
+            return value
+
+        return value
+
+    valid_fields = {f.name: f for f in fields(cls) if f.init}
+    filtered_kwargs = {}
+    for k, v in kwargs.items():
+        if k not in valid_fields:
+            continue
+        field = valid_fields[k]
+        filtered_kwargs[k] = _filter_value(v, field.type)
+
+    return filtered_kwargs
+
+
 # TODO(wuhang): Remove after PR #1115.
 def build_base_engine_args(source: Any) -> dict[str, Any] | None:
     """Build base engine args with tokenizer and parallel configuration.