From 15df8ec97f8667d76638e780480f2cca7e9d76e5 Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Thu, 26 Mar 2026 18:50:54 +0800 Subject: [PATCH] [Bug Fix] GLM-Image stage device isolation and t2i prompt preprocessing in Omni runtime (#2137) --- .../test_async_omni_engine_stage_init.py | 129 ++++++++++++++++++ vllm_omni/engine/async_omni_engine.py | 59 ++++---- 2 files changed, 163 insertions(+), 25 deletions(-) create mode 100644 tests/engine/test_async_omni_engine_stage_init.py diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py new file mode 100644 index 000000000..28b44e9bd --- /dev/null +++ b/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 diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 4951d2152..a4d87c96e 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/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(