diff --git a/tests/diffusion/test_stage_diffusion_proc.py b/tests/diffusion/test_stage_diffusion_proc.py index 697465097..90e753c9b 100644 --- a/tests/diffusion/test_stage_diffusion_proc.py +++ b/tests/diffusion/test_stage_diffusion_proc.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import asyncio import time @@ -7,6 +7,8 @@ from dataclasses import asdict, dataclass import pytest +import vllm_omni.diffusion.stage_diffusion_proc as stage_diffusion_proc +import vllm_omni.plugins as omni_plugins from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -26,6 +28,33 @@ BASE_INFER_STEPS = 10 DELAY_BASE = 0.01 +def test_run_diffusion_proc_sets_lifecycle_before_loading_plugins(monkeypatch: pytest.MonkeyPatch) -> None: + events: list[str] = [] + + class StopProcessError(Exception): + pass + + class TestStageDiffusionProc(StageDiffusionProc): + def __init__(self, model, od_config): + events.append("proc") + raise StopProcessError + + monkeypatch.setattr(omni_plugins, "load_omni_general_plugins", lambda: events.append("plugins")) + monkeypatch.setattr(stage_diffusion_proc, "set_death_signal", lambda _: events.append("death_signal")) + monkeypatch.setattr(stage_diffusion_proc.signal, "signal", lambda *_: events.append("signal_handler")) + + with pytest.raises(StopProcessError): + TestStageDiffusionProc.run_diffusion_proc( + model="test-model", + od_config=None, + handshake_address="test-address", + local_client=True, + headless=False, + ) + + assert events == ["death_signal", "signal_handler", "signal_handler", "plugins", "proc"] + + class MockDiffusionEngine: async def step_streaming(self, request): def simulate_step_delay(height, width, num_inference_steps) -> float: diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index 4419e51dc..dad453faa 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/vllm_omni/diffusion/stage_diffusion_proc.py @@ -598,6 +598,8 @@ class StageDiffusionProc: - ``omni_replica_id``: cluster-unique replica id within the stage (logging / metrics only). """ + from vllm_omni.plugins import load_omni_general_plugins + shutdown_requested = False set_death_signal(signal.SIGTERM) @@ -611,6 +613,11 @@ class StageDiffusionProc: signal.signal(signal.SIGTERM, signal_handler) signal.signal(signal.SIGINT, signal_handler) + # ``spawn`` starts this process with a fresh interpreter, so plugin + # side effects (for example, custom diffusion loader hooks) must be + # applied again before the engine is constructed. + load_omni_general_plugins() + proc = cls(model, od_config) coord_client: OmniCoordClientForStage | None = None handshake_ctx: zmq.Context | None = None