[Entrypoint][Refactor] vLLM-Omni Entrypoint Refactoring (#1908)

Signed-off-by: yinpe <11810305@mail.sustech.edu.cn>
Signed-off-by: Peiqi Yin <yinpeiqi809@gmail.com>
Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
Signed-off-by: wuhang <whlbx@hotmail.com>
Signed-off-by: linyueqian <linyueqian@outlook.com>
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: lishunyang <lishunyang12@163.com>
Signed-off-by: Chenguang ZHENG <645327136@qq.com>
Signed-off-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com>
Co-authored-by: yinpe <11810305@mail.sustech.edu.cn>
Co-authored-by: Peiqi Yin <yinpeiqi809@gmail.com>
Co-authored-by: wuhang <whlbx@hotmail.com>
Co-authored-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com>
Co-authored-by: linyueqian <linyueqian@outlook.com>
Co-authored-by: princepride <wangzhipeng628@gmail.com>
Co-authored-by: lishunyang <lishunyang12@163.com>

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -3,87 +3,51 @@
 
 import pytest
 
-from vllm_omni.entrypoints import utils as utils_module
-from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
-MODEL = "riverclouds/qwen_image_random"
+
+def test_default_stage_config_includes_cache_backend():
+    """Ensure cache knobs survive the default diffusion-stage builder."""
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        {
+            "cache_backend": "cache_dit",
+            "cache_config": '{"Fn_compute_blocks": 2}',
+            "vae_use_slicing": True,
+            "ulysses_degree": 2,
+        }
+    )[0]
+
+    engine_args = stage_cfg["engine_args"]
+    assert stage_cfg["stage_type"] == "diffusion"
+    assert engine_args["cache_backend"] == "cache_dit"
+    assert engine_args["cache_config"]["Fn_compute_blocks"] == 2
+    assert engine_args["vae_use_slicing"] is True
+    assert engine_args["parallel_config"].ulysses_degree == 2
+    assert engine_args["model_stage"] == "diffusion"
 
 
-def _noop_inline_engine(self, model, stage_config, kwargs):
-    self._inline_diffusion = False
-    self._inline_engine = None
+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(
+        {
+            "cache_backend": "cache_dit",
+        }
+    )[0]
 
-
-def test_default_stage_config_includes_cache_backend(monkeypatch):
-    """Ensure cache_backend/cache_config are preserved in default diffusion stage."""
-    monkeypatch.setattr(utils_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
-    monkeypatch.setattr(utils_module, "resolve_model_config_path", lambda model: None)
-    monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
-    monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
-    monkeypatch.setattr(AsyncOmni, "_init_inline_diffusion_engine", _noop_inline_engine)
-
-    omni = AsyncOmni(
-        model=MODEL,
-        cache_backend="cache_dit",
-        cache_config='{"Fn_compute_blocks": 2}',
-        vae_use_slicing=True,
-        ulysses_degree=2,
-    )
-
-    stage_cfg = omni.stage_configs[0]
-    engine_args = stage_cfg.engine_args
-
-    assert engine_args.get("cache_backend") == "cache_dit"
-    cache_config = engine_args.get("cache_config")
-    assert cache_config["Fn_compute_blocks"] == 2
-    assert engine_args.get("vae_use_slicing") is True
-    parallel_config = engine_args.get("parallel_config")
-    if hasattr(parallel_config, "get"):
-        ulysses_degree = parallel_config.get("ulysses_degree")
-    else:
-        ulysses_degree = getattr(parallel_config, "ulysses_degree", None)
-    assert ulysses_degree == 2
-
-
-def test_default_cache_config_used_when_missing(monkeypatch):
-    """Ensure default cache_config is applied when cache_backend is set."""
-    monkeypatch.setattr(utils_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
-    monkeypatch.setattr(utils_module, "resolve_model_config_path", lambda model: None)
-    monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
-    monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
-    monkeypatch.setattr(AsyncOmni, "_init_inline_diffusion_engine", _noop_inline_engine)
-
-    omni = AsyncOmni(
-        model=MODEL,
-        cache_backend="cache_dit",
-    )
-
-    engine_args = omni.stage_configs[0].engine_args
-    cache_config = engine_args.get("cache_config")
+    cache_config = stage_cfg["engine_args"]["cache_config"]
     assert cache_config is not None
     assert cache_config["Fn_compute_blocks"] == 1
 
 
-def test_default_stage_devices_from_sequence_parallel(monkeypatch):
-    """Ensure devices list reflects sequence parallel size when no parallel_config is provided."""
-    monkeypatch.setattr(utils_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
-    monkeypatch.setattr(utils_module, "resolve_model_config_path", lambda model: None)
-    monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
-    monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
-    monkeypatch.setattr(AsyncOmni, "_init_inline_diffusion_engine", _noop_inline_engine)
+def test_default_stage_devices_from_sequence_parallel():
+    """Ensure runtime devices reflect computed diffusion world size."""
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        {
+            "ulysses_degree": 2,
+            "ring_degree": 2,
+        }
+    )[0]
 
-    omni = AsyncOmni(
-        model=MODEL,
-        ulysses_degree=2,
-        ring_degree=2,
-    )
-
-    stage_cfg = omni.stage_configs[0]
-    runtime = stage_cfg.runtime
-    if hasattr(runtime, "get"):
-        devices = runtime.get("devices")
-    else:
-        devices = getattr(runtime, "devices", None)
-    assert devices == "0,1,2,3"
+    assert stage_cfg["runtime"]["devices"] == "0,1,2,3"