[Bug Fix] GLM-Image stage device isolation and t2i prompt preprocessing in Omni runtime (#2137)

tests/engine/test_async_omni_engine_stage_init.py

@@ -0,0 +1,129 @@
+import os
+import types
+
+import pytest
+
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_initialize_stages_restores_device_visibility_after_diffusion_init(monkeypatch):
+    """Regression test for stage device env leakage across stage init.
+
+    Diffusion init mutates process-level CUDA visibility. Ensure AsyncOmniEngine
+    restores the previous value after diffusion stage setup.
+    """
+    import vllm_omni.engine.async_omni_engine as engine_mod
+    from vllm_omni.platforms import current_omni_platform
+
+    engine = object.__new__(AsyncOmniEngine)
+    engine.model = "dummy-model"
+    engine.config_path = "dummy-config"
+    engine.num_stages = 1
+    engine.async_chunk = False
+    engine.diffusion_batch_size = 1
+    engine.stage_configs = [types.SimpleNamespace(stage_id=0, stage_type="diffusion")]
+
+    env_var = current_omni_platform.device_control_env_var
+    old_env = os.environ.get(env_var)
+    os.environ[env_var] = "0,1"
+
+    diffusion_client = types.SimpleNamespace(is_comprehension=False)
+
+    metadata = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="diffusion",
+        runtime_cfg={"devices": "1"},
+        prompt_expand_func=None,
+    )
+
+    monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None)
+    monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None)
+    monkeypatch.setattr(engine_mod, "extract_stage_metadata", lambda _cfg: metadata)
+    monkeypatch.setattr(engine_mod, "get_stage_connector_spec", lambda **_: {})
+    monkeypatch.setattr(engine_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None))
+
+    def _fake_setup_stage_devices(_stage_id, _runtime_cfg):
+        # Simulate diffusion setup mutating process-global visibility.
+        current_omni_platform.set_device_control_env_var("1")
+
+    monkeypatch.setattr(engine_mod, "setup_stage_devices", _fake_setup_stage_devices)
+    monkeypatch.setattr(engine_mod, "_inject_kv_stage_info", lambda *_: None)
+    monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", lambda *_, **__: diffusion_client)
+    monkeypatch.setattr(
+        engine_mod,
+        "finalize_initialized_stages",
+        lambda stage_clients, _input_processor: (
+            stage_clients,
+            [types.SimpleNamespace()],
+            [{"final_output_type": "image"}],
+        ),
+    )
+
+    try:
+        engine._initialize_stages(stage_init_timeout=1)
+        assert os.environ.get(env_var) == "0,1"
+    finally:
+        if old_env is None:
+            os.environ.pop(env_var, None)
+        else:
+            os.environ[env_var] = old_env
+
+
+def test_attach_llm_stage_uses_omni_input_preprocessor(monkeypatch):
+    """Regression test for GLM-Image t2i preprocessing path.
+
+    Stage-0 InputProcessor must use OmniInputPreprocessor so text prompts with
+    mm_processor_kwargs go through multimodal preprocessing.
+    """
+    import vllm_omni.engine.async_omni_engine as engine_mod
+
+    class DummyStageEngineCoreClient:
+        def __init__(self, **kwargs):
+            self.kwargs = kwargs
+
+        def shutdown(self):
+            return None
+
+    class DummyOutputProcessor:
+        def __init__(self, **kwargs):
+            self.kwargs = kwargs
+
+    class DummyInputProcessor:
+        def __init__(self, vllm_config):
+            self.vllm_config = vllm_config
+            self.renderer = object()
+            self.input_preprocessor = object()
+
+    class DummyOmniInputPreprocessor:
+        def __init__(self, vllm_config, renderer=None):
+            self.vllm_config = vllm_config
+            self.renderer = renderer
+
+    monkeypatch.setattr(engine_mod, "StageEngineCoreClient", DummyStageEngineCoreClient)
+    monkeypatch.setattr(engine_mod, "MultimodalOutputProcessor", DummyOutputProcessor)
+    monkeypatch.setattr(engine_mod, "InputProcessor", DummyInputProcessor)
+    monkeypatch.setattr(engine_mod, "OmniInputPreprocessor", DummyOmniInputPreprocessor)
+
+    started = types.SimpleNamespace(
+        stage_id=0,
+        metadata=types.SimpleNamespace(stage_id=0, engine_output_type="token_ids"),
+        vllm_config=types.SimpleNamespace(model_config=types.SimpleNamespace(skip_tokenizer_init=True)),
+        executor_class=object,
+        engine_manager=object(),
+        coordinator=object(),
+        addresses=types.SimpleNamespace(
+            inputs=["inproc://input"],
+            outputs=["inproc://output"],
+            frontend_stats_publish_address=None,
+        ),
+    )
+
+    engine = object.__new__(AsyncOmniEngine)
+
+    _stage_client, _out_proc, _vllm_cfg, input_processor = engine._attach_llm_stage(started)
+
+    assert input_processor is not None
+    assert isinstance(input_processor.input_preprocessor, DummyOmniInputPreprocessor)
+    assert input_processor.input_preprocessor.renderer is input_processor.renderer

vllm_omni/engine/async_omni_engine.py

@@ -64,6 +64,7 @@ from vllm_omni.engine.stage_init_utils import (
 from vllm_omni.entrypoints.utils import (
     load_and_resolve_stage_configs,
 )
+from vllm_omni.inputs.preprocess import OmniInputPreprocessor
 from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
@@ -412,6 +413,13 @@ class AsyncOmniEngine:
             input_processor = None
             if started.stage_id == 0:
                 input_processor = InputProcessor(vllm_config=started.vllm_config)
+                # Use omni preprocessor so text-only prompts with
+                # mm_processor_kwargs (e.g. GLM-Image t2i target_h/target_w)
+                # still go through multimodal processor path.
+                input_processor.input_preprocessor = OmniInputPreprocessor(
+                    vllm_config=started.vllm_config,
+                    renderer=input_processor.renderer,
+                )
         except Exception:
             try:
                 stage_client.shutdown()
@@ -469,32 +477,33 @@ class AsyncOmniEngine:
                     omni_kv_connector = resolve_omni_kv_config_for_stage(omni_transfer_config, stage_id)
 
                     if metadata.stage_type == "diffusion":
-                        previous_visible_devices = os.environ.get(device_control_env)
-                        try:
-                            setup_stage_devices(stage_id, metadata.runtime_cfg)
-                            omni_conn_cfg, omni_from, omni_to = omni_kv_connector
-                            if omni_conn_cfg:
-                                from vllm_omni.entrypoints.utils import inject_omni_kv_config
+                        with llm_stage_launch_lock:
+                            previous_visible_devices = os.environ.get(device_control_env)
+                            try:
+                                setup_stage_devices(stage_id, metadata.runtime_cfg)
+                                omni_conn_cfg, omni_from, omni_to = omni_kv_connector
+                                if omni_conn_cfg:
+                                    from vllm_omni.entrypoints.utils import inject_omni_kv_config
 
-                                inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to)
-                            _inject_kv_stage_info(stage_cfg, stage_id)
-                            stage_clients[stage_id] = initialize_diffusion_stage(
-                                self.model,
-                                stage_cfg,
-                                metadata,
-                                batch_size=self.diffusion_batch_size,
-                            )
-                            logger.info(
-                                "[AsyncOmniEngine] Stage %s initialized (diffusion, batch_size=%d)",
-                                stage_id,
-                                self.diffusion_batch_size,
-                            )
-                            continue
-                        finally:
-                            if previous_visible_devices is None:
-                                current_omni_platform.unset_device_control_env_var()
-                            else:
-                                current_omni_platform.set_device_control_env_var(previous_visible_devices)
+                                    inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to)
+                                _inject_kv_stage_info(stage_cfg, stage_id)
+                                stage_clients[stage_id] = initialize_diffusion_stage(
+                                    self.model,
+                                    stage_cfg,
+                                    metadata,
+                                    batch_size=self.diffusion_batch_size,
+                                )
+                                logger.info(
+                                    "[AsyncOmniEngine] Stage %s initialized (diffusion, batch_size=%d)",
+                                    stage_id,
+                                    self.diffusion_batch_size,
+                                )
+                            finally:
+                                if previous_visible_devices is None:
+                                    current_omni_platform.unset_device_control_env_var()
+                                else:
+                                    current_omni_platform.set_device_control_env_var(previous_visible_devices)
+                        continue
 
                     llm_stage_ids.append(stage_id)
                     llm_launch_futures[stage_id] = launch_executor.submit(