[Bugfix, rl] Diffusion worker SIGKILL under Ray actor (exitcode -9) (#3533)

Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
Co-authored-by: Samit <285365963@qq.com>

.buildkite/test-ready.yml

@@ -217,6 +217,26 @@ steps:
   - group: ":card_index_dividers: E2E Test"
     depends_on: upload-ready-pipeline
     steps:
+      - label: "RLHF · VeRL-Omni E2E Test"
+        commands:
+          - |
+            timeout 20m bash -c "
+              pytest -s -v tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py
+            "
+        agents:
+          queue: "gpu_1_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
       - label: "Omni · Qwen3-Omni Test"
         commands:
           - |

pyproject.toml

@@ -55,6 +55,7 @@ dev = [
     "opencc>=1.2.0",
     "mistune>=3.2.0", # for example tests
     "torchmetrics>=1.4.0", # for accuracy similarity metrics
+    "ray[default]", # for RL test
     "jiwer>=3.0.0",
     "zhon>=2.0.0",
     "zhconv>=1.4.2",

tests/e2e/offline_inference/custom_pipeline/worker_extension.py

@@ -1,20 +1,76 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""Test-only worker extension for custom diffusion pipeline E2E tests."""
+"""Test worker extension that mirrors
+``verl_omni.workers.rollout.vllm_rollout.utils.vLLMOmniColocateWorkerExtension``
+
+Fidelity goals:
+
+* ``__new__`` calls ``set_death_signal()`` (Linux ``PR_SET_PDEATHSIG``) so
+  spawned vLLM workers die with the parent.
+* ``__new__`` calls ``VLLMOmniHijackForTest.hijack()`` which monkey-patches
+  ``vllm_omni.diffusion.lora.manager.DiffusionLoRAManager._load_adapter``
+  to accept in-memory LoRA tensors (verbatim port of
+  ``verl_omni.utils.vllm_omni.utils.VLLMOmniHijack.hijack``).
+
+The ``verl.utils.vllm.VLLMHijack.hijack()`` call (which patches vLLM's
+``LRUCacheWorkerLoRAManager._load_adapter`` for AR/text mode) is
+intentionally skipped because (a) ``verl`` is not importable from this
+repo and (b) the diffusion test path never reaches that manager.
+"""
 
 from __future__ import annotations
 
+import ctypes
+import logging
+import os
+import platform
+import signal
+
+from msgspec import field
+
 from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension
+from vllm_omni.lora.request import LoRARequest as OmniLoRARequest
+
+logger = logging.getLogger(__name__)
+
+
+class OmniTensorLoRARequestForTest(OmniLoRARequest):
+    peft_config: dict = field(default=None)
+    lora_tensors: dict = field(default=None)
+
+
+def set_death_signal() -> None:
+    """Verbatim port of ``verl.workers.rollout.vllm_rollout.utils.set_death_signal``."""
+    if platform.system() != "Linux":
+        return
+    try:
+        libc = ctypes.CDLL("libc.so.6")
+        libc.prctl(1, signal.SIGKILL)
+        if os.getppid() == 1:
+            os.kill(os.getpid(), signal.SIGKILL)
+    except Exception:  # noqa: BLE001
+        # libc.so.6 may not be present (musl, alpine) — best-effort only.
+        pass
+
+
+# ---------------------------------------------------------------------
+#  The worker extension itself
+# ---------------------------------------------------------------------
 
 
 class vLLMOmniColocateWorkerExtensionForTest(CustomPipelineWorkerExtension):
-    """Minimal worker extension used by tests.
+    """Mirror of ``vLLMOmniColocateWorkerExtension`` (verl-omni).
 
-    This intentionally stays lightweight: we only inherit the base custom
-    pipeline extension and add one test function.
+    The production ``__new__`` runs ``set_death_signal`` + ``VLLMOmniHijack.hijack``
+    on every vLLM worker process. Replicating both ensures this test
+    reproduces the same monkey-patched environment.
     """
 
+    def __new__(cls, **kwargs):
+        set_death_signal()
+        return super().__new__(cls)
+
     @staticmethod
     def test_extension_name() -> str:
         """Return a stable identifier for assertions in unit tests."""

tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py

@@ -0,0 +1,588 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""E2E test that follows the EXACT same flow as
+``verl-omni/tests/workers/rollout/rollout_vllm/test_vllm_omni_generate.py``
+but inlines all the ``verl`` / ``verl_omni`` symbols the flow depends on
+(stripped down to only what the diffusion-mode test actually exercises).
+
+Flow (identical to the verl-omni test):
+
+1. Build two ``OmegaConf`` configs (``DiffusionRolloutConfig`` /
+   ``DiffusionModelConfig`` shape).
+2. ``ServerCls = ray.remote(vLLMOmniHttpServerLocal)`` then
+   ``ServerCls.options(...).remote(config=..., model_config=...,
+   rollout_mode=RolloutMode.STANDALONE, workers=[], replica_rank=0,
+   node_rank=0, gpus_per_node=1, nnodes=1, cuda_visible_devices="0")``.
+3. ``ray.get(server.launch_server.remote())``.
+4. ``server.generate.remote(prompt_ids=..., sampling_params=..., request_id=...)``
+   and assert the returned ``DiffusionOutput``.
+
+The inlined classes / helpers mirror their ``verl`` / ``verl_omni``
+counterparts; bits not reached by this test (AR mode, LoRA,
+quantization, run_uvicorn, headless, profiler, port wiring) are removed.
+
+Usage:
+    pytest tests/e2e/offline_inference/rlhf_test/\
+test_ray_async_omni_qwen_image_generate.py -v -s
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import os
+from dataclasses import asdict
+from enum import Enum
+from functools import lru_cache
+from typing import Any
+from uuid import uuid4
+
+import pytest
+import ray
+import torch
+import torchvision.transforms as T
+from huggingface_hub import snapshot_download
+from omegaconf import DictConfig, OmegaConf
+from pydantic import BaseModel, ConfigDict
+from transformers import AutoTokenizer
+from vllm.entrypoints.openai.api_server import build_app
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+import vllm_omni.entrypoints.cli.serve
+from tests.helpers.mark import hardware_test
+from vllm_omni.engine.arg_utils import OmniEngineArgs
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.entrypoints.openai.api_server import omni_init_app_state
+from vllm_omni.inputs.data import OmniCustomPrompt, OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+logger = logging.getLogger(__name__)
+
+CUSTOM_PIPELINE_CLASS = (
+    "tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob.QwenImagePipelineWithLogProbForTest"
+)
+WORKER_EXTENSION_CLASS = (
+    "tests.e2e.offline_inference.custom_pipeline.worker_extension.vLLMOmniColocateWorkerExtensionForTest"
+)
+
+# Use your specified HF repo name/path directly here
+MODEL = "tiny-random/Qwen-Image"
+
+TOKENIZER_MODEL = "Qwen/Qwen2-1.5B-Instruct"
+
+
+def _resolve_model_path(repo_id: str) -> str:
+    """Resolve an HF repo ID to a local snapshot path, downloading if needed."""
+    # Allow overriding with a pre-existing local path (skips download).
+    if os.path.isdir(repo_id):
+        return repo_id
+    return snapshot_download(repo_id=repo_id)
+
+
+_MIN_PROMPT_TOKENS = 35
+
+
+# ---------------------------------------------------------------------
+#       Inlined: verl.workers.rollout.replica.RolloutMode
+# ---------------------------------------------------------------------
+
+
+class RolloutMode(Enum):
+    HYBRID = "hybrid"
+    COLOCATED = "colocated"
+    STANDALONE = "standalone"
+
+
+# ---------------------------------------------------------------------
+#       Inlined: verl_omni.workers.rollout.replica.DiffusionOutput
+# ---------------------------------------------------------------------
+
+
+class DiffusionOutput(BaseModel):
+    """Pydantic mirror of ``verl_omni.workers.rollout.replica.DiffusionOutput``."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    diffusion_output: Any
+    log_probs: Any | None = None
+    stop_reason: str | None = None
+    num_preempted: int | None = None
+    extra_fields: dict[str, Any] = {}
+
+
+# ---------------------------------------------------------------------
+#       Inlined: verl.utils.tokenizer.normalize_token_ids
+# ---------------------------------------------------------------------
+
+
+def normalize_token_ids(tokenized_output) -> list[int]:
+    token_ids = tokenized_output
+    if isinstance(tokenized_output, dict):
+        if "input_ids" in tokenized_output:
+            token_ids = tokenized_output["input_ids"]
+    elif hasattr(tokenized_output, "input_ids"):
+        token_ids = tokenized_output.input_ids
+
+    if hasattr(token_ids, "tolist"):
+        token_ids = token_ids.tolist()
+    if isinstance(token_ids, tuple):
+        token_ids = list(token_ids)
+    if isinstance(token_ids, list) and len(token_ids) == 1 and isinstance(token_ids[0], (list, tuple)):
+        token_ids = list(token_ids[0])
+    if not isinstance(token_ids, list):
+        raise TypeError(f"token_ids must be list-like, got {type(token_ids).__name__}")
+
+    out: list[int] = []
+    for tid in token_ids:
+        if hasattr(tid, "item"):
+            tid = tid.item()
+        out.append(int(tid))
+    return out
+
+
+# ---------------------------------------------------------------------
+#  Inlined: verl.workers.rollout.vllm_rollout.utils.build_cli_args_from_config
+# ---------------------------------------------------------------------
+
+
+def build_cli_args_from_config(config: dict[str, Any]) -> list[str]:
+    cli_args: list[str] = []
+    for k, v in config.items():
+        if v is None:
+            continue
+        if isinstance(v, bool):
+            if v:
+                cli_args.append(f"--{k}")
+        elif isinstance(v, list):
+            if not v:
+                continue
+            cli_args.append(f"--{k}")
+            cli_args.extend([str(item) for item in v])
+        else:
+            cli_args.append(f"--{k}")
+            cli_args.append(json.dumps(v) if isinstance(v, dict) else str(v))
+    return cli_args
+
+
+# ---------------------------------------------------------------------
+#                  Tokenizer helper (test-side)
+# ---------------------------------------------------------------------
+
+
+@lru_cache(maxsize=1)
+def _get_tokenizer():
+    return AutoTokenizer.from_pretrained(TOKENIZER_MODEL, trust_remote_code=True)
+
+
+def _tokenize_prompt(text: str) -> list[int]:
+    tokenizer = _get_tokenizer()
+    messages = [{"role": "user", "content": text}]
+    token_ids = normalize_token_ids(tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False))
+    assert len(token_ids) > _MIN_PROMPT_TOKENS, (
+        f"Prompt too short ({len(token_ids)} tokens, need >{_MIN_PROMPT_TOKENS}). "
+        "The pipeline drops the first 34 chat-template prefix tokens; "
+        "use a longer prompt so content tokens remain after the drop."
+    )
+    return token_ids
+
+
+# ---------------------------------------------------------------------
+#  Inlined: vLLMOmniHttpServer (verl-omni) + relevant vLLMHttpServer (verl)
+#  bits, stripped to the diffusion-mode path the test exercises.
+# ---------------------------------------------------------------------
+
+
+def _cfg_get(cfg: Any, key: str, default: Any = None) -> Any:
+    """Uniform getter that works for both DictConfig and dataclass-like objects."""
+    if isinstance(cfg, DictConfig):
+        return cfg.get(key, default)
+    return getattr(cfg, key, default)
+
+
+class vLLMOmniHttpServerLocal:
+    """Self-contained simulation of ``vLLMOmniHttpServer`` for diffusion mode.
+
+    Constructor signature matches ``vLLMHttpServer.__init__`` exactly so the
+    Ray-actor call site in the test is identical to the verl-omni one.
+    """
+
+    # ------------------------------------------------------------------ init
+    def __init__(
+        self,
+        config,
+        model_config,
+        rollout_mode: RolloutMode,
+        workers: list,
+        replica_rank: int,
+        node_rank: int,
+        gpus_per_node: int,
+        nnodes: int,
+        cuda_visible_devices: str,
+    ):
+        os.environ.setdefault("CUDA_VISIBLE_DEVICES", cuda_visible_devices)
+        os.environ["VERL_REPLICA_RANK"] = str(replica_rank)
+
+        # Resolve configs. verl uses omega_conf_to_dataclass; here we keep
+        # them as DictConfig — every read site below uses ``_cfg_get``.
+        self.config = OmegaConf.create(config) if not isinstance(config, DictConfig) else config
+        self.model_config = OmegaConf.create(model_config) if not isinstance(model_config, DictConfig) else model_config
+
+        self.rollout_mode = rollout_mode
+        self.workers = workers
+        self.replica_rank = replica_rank
+        self.node_rank = node_rank
+        self.gpus_per_node = gpus_per_node
+        self.nnodes = nnodes
+        self.global_steps = 0
+        self.engine: AsyncOmni | None = None
+
+        # Diffusion-mode parity with ``vLLMOmniHttpServer._post_init``.
+        self._to_tensor = T.PILToTensor()
+
+        logger.info(
+            "%s, replica_rank: %d, node_rank: %d, CUDA_VISIBLE_DEVICES: %s",
+            self.__class__.__name__,
+            replica_rank,
+            node_rank,
+            cuda_visible_devices,
+        )
+
+    # ---------------------------------------------------------- launch_server
+    async def launch_server(self) -> None:
+        """Diffusion-only mirror of ``vLLMHttpServer.launch_server``."""
+
+        # vLLMOmniHttpServer._get_engine_kwargs_key() returns "vllm_omni".
+        engine_kwargs_root = _cfg_get(self.config, "engine_kwargs", {}) or {}
+        if isinstance(engine_kwargs_root, (dict, DictConfig)):
+            engine_kwargs = engine_kwargs_root.get("vllm_omni", {}) or {}
+        else:
+            engine_kwargs = {}
+        engine_kwargs = {k: v for k, v in dict(engine_kwargs).items() if v is not None}
+
+        # vLLMOmniHttpServer._preprocess_engine_kwargs (diffusion path).
+        engine_kwargs.pop("output_mode", None)
+
+        # ``compilation_config`` handling — verbatim from parent launch_server.
+        compilation_config = engine_kwargs.pop("compilation_config", None) or {}
+        if isinstance(compilation_config, str):
+            compilation_config = json.loads(compilation_config)
+        if isinstance(compilation_config, DictConfig):
+            compilation_config = OmegaConf.to_container(compilation_config, resolve=True)
+        compilation_config.setdefault("cudagraph_mode", "FULL_AND_PIECEWISE")
+        compilation_config = json.dumps(compilation_config)
+
+        # Args dict (subset of parent's; the diffusion override drops AR-only
+        # entries via ``_get_override_generation_config`` returning ``{}``).
+        args = {
+            "dtype": _cfg_get(self.config, "dtype", "bfloat16"),
+            "load_format": _cfg_get(self.config, "load_format", "auto"),
+            "skip_tokenizer_init": False,
+            "distributed_executor_backend": "mp",
+            "worker_extension_cls": WORKER_EXTENSION_CLASS,
+            "trust_remote_code": _cfg_get(self.model_config, "trust_remote_code", True),
+            "max_model_len": _cfg_get(self.config, "max_model_len", 1058),
+            "max_num_seqs": _cfg_get(self.config, "max_num_seqs", 256),
+            "enable_chunked_prefill": _cfg_get(self.config, "enable_chunked_prefill", False),
+            "max_num_batched_tokens": _cfg_get(self.config, "max_num_batched_tokens", 8192),
+            "enable_prefix_caching": _cfg_get(self.config, "enable_prefix_caching", False),
+            "enable_sleep_mode": _cfg_get(self.config, "enable_sleep_mode", False),
+            "logprobs_mode": _cfg_get(self.config, "logprobs_mode", "processed_logprobs"),
+            "enforce_eager": _cfg_get(self.config, "enforce_eager", True),
+            "gpu_memory_utilization": _cfg_get(self.config, "gpu_memory_utilization", 0.8),
+            "disable_log_stats": _cfg_get(self.config, "disable_log_stats", True),
+            "tensor_parallel_size": _cfg_get(self.config, "tensor_model_parallel_size", 1),
+            "seed": self.replica_rank + int(_cfg_get(self.config, "seed", 0) or 0),
+            "override_generation_config": json.dumps({}),
+            "scheduling_policy": _cfg_get(self.config, "scheduling_policy", "fcfs"),
+            "compilation_config": compilation_config,
+            **engine_kwargs,
+        }
+
+        model_path = str(_cfg_get(self.model_config, "local_path", None) or _cfg_get(self.model_config, "path"))
+        server_args_list = ["serve", model_path] + build_cli_args_from_config(args)
+
+        # Same parser pipeline as parent launch_server (with vllm-omni's
+        # cmd_init injected, matching ``_get_cli_modules`` + ``_get_cli_description``).
+        parser = FlexibleArgumentParser(description="vLLM-Omni CLI")
+        subparsers = parser.add_subparsers(required=False, dest="subparser")
+        cmds: dict[str, Any] = {}
+        for cmd in vllm_omni.entrypoints.cli.serve.cmd_init():
+            cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd)
+            cmds[cmd.name] = cmd
+        parsed = parser.parse_args(args=server_args_list)
+        parsed.model = parsed.model_tag
+        if parsed.subparser in cmds:
+            cmds[parsed.subparser].validate(parsed)
+
+        # node_rank == 0 → run_server (the test never sets node_rank != 0).
+        await self.run_server(parsed)
+
+    # ------------------------------------------------------------- run_server
+    async def run_server(self, args: argparse.Namespace) -> None:
+        """Diffusion branch of ``vLLMOmniHttpServer.run_server``."""
+
+        engine_args = OmniEngineArgs.from_cli_args(args)
+        engine_args = asdict(engine_args)
+
+        # Diffusion path: pin the in-repo test pipeline. The production code
+        # calls ``VllmOmniPipelineBase.get_pipeline_path(architecture)``; we
+        # bypass the registry and bind the test pipeline directly.
+        engine_args["enable_dummy_pipeline"] = True
+        engine_args["custom_pipeline_args"] = {"pipeline_class": CUSTOM_PIPELINE_CLASS}
+
+        engine_client = AsyncOmni(**engine_args)
+        app = build_app(args)
+        await omni_init_app_state(engine_client, app.state, args)
+
+        self.engine = engine_client
+        # ``run_uvicorn`` intentionally skipped — the test drives the engine
+        # via Ray RPC, not HTTP.
+
+    # ---------------------------------------------------------------- generate
+    async def generate(
+        self,
+        prompt_ids: list[int],
+        sampling_params: dict[str, Any],
+        request_id: str,
+        image_data: list[Any] | None = None,
+        video_data: list[Any] | None = None,
+        negative_prompt_ids: list[int] | None = None,
+        priority: int = 0,  # noqa: ARG002 (signature parity)
+    ) -> DiffusionOutput:
+        # Verbatim copy of ``vLLMOmniHttpServer._generate_diffusion`` minus
+        # the LoRA branch (the test does not enable ``lora_as_adapter``).
+        prompt_ids = normalize_token_ids(prompt_ids)
+
+        multi_modal_data: dict[str, Any] = {}
+        if image_data is not None:
+            multi_modal_data["image"] = image_data
+        if video_data is not None:
+            multi_modal_data["video"] = video_data
+
+        custom_prompt: OmniCustomPrompt = {"prompt_ids": prompt_ids}
+        if negative_prompt_ids is not None:
+            custom_prompt["negative_prompt_ids"] = negative_prompt_ids
+        if multi_modal_data:
+            custom_prompt["extra_args"] = {"multi_modal_data": multi_modal_data}
+
+        sampling_kwargs: dict[str, Any] = {}
+        extra_args: dict[str, Any] = {}
+        for k, v in sampling_params.items():
+            if hasattr(OmniDiffusionSamplingParams, k):
+                sampling_kwargs[k] = v
+            else:
+                extra_args[k] = v
+        sampling_kwargs["extra_args"] = extra_args
+        diffusion_sampling_params = OmniDiffusionSamplingParams(**sampling_kwargs)
+
+        generator = self.engine.generate(
+            prompt=custom_prompt,
+            request_id=request_id,
+            sampling_params_list=[diffusion_sampling_params],
+        )
+
+        final_res: OmniRequestOutput | None = None
+        async for output in generator:
+            final_res = output
+        assert final_res is not None
+
+        diffusion_output = self._to_tensor(final_res.images[0]).float() / 255.0
+
+        mm_output = final_res.custom_output or {}
+
+        if sampling_params.get("logprobs", False):
+            all_log_probs = mm_output.get("all_log_probs")
+            log_probs = all_log_probs[0] if all_log_probs is not None else None
+        else:
+            log_probs = None
+
+        all_latents = mm_output.get("all_latents")
+        all_timesteps = mm_output.get("all_timesteps")
+        prompt_embeds = mm_output.get("prompt_embeds")
+        prompt_embeds_mask = mm_output.get("prompt_embeds_mask")
+        negative_prompt_embeds = mm_output.get("negative_prompt_embeds")
+        negative_prompt_embeds_mask = mm_output.get("negative_prompt_embeds_mask")
+
+        extra_fields = {
+            "all_latents": all_latents[0] if all_latents is not None else None,
+            "all_timesteps": all_timesteps[0] if all_timesteps is not None else None,
+            "prompt_embeds": prompt_embeds[0] if prompt_embeds is not None else None,
+            "prompt_embeds_mask": prompt_embeds_mask[0] if prompt_embeds_mask is not None else None,
+            "negative_prompt_embeds": negative_prompt_embeds[0] if negative_prompt_embeds is not None else None,
+            "negative_prompt_embeds_mask": (
+                negative_prompt_embeds_mask[0] if negative_prompt_embeds_mask is not None else None
+            ),
+            "global_steps": self.global_steps,
+        }
+
+        req_output = final_res.request_output
+        if req_output is not None and hasattr(req_output, "finish_reason"):
+            finish_reason = req_output.finish_reason or "stop"
+        else:
+            finish_reason = "stop"
+
+        if finish_reason == "abort":
+            stop_reason = "aborted"
+        elif finish_reason in ("stop", "length"):
+            stop_reason = "completed"
+        else:
+            stop_reason = finish_reason
+
+        num_preempted = None
+        if req_output is not None and hasattr(req_output, "num_preempted"):
+            num_preempted = req_output.num_preempted
+
+        return DiffusionOutput(
+            diffusion_output=diffusion_output,
+            log_probs=log_probs,
+            stop_reason=stop_reason,
+            num_preempted=num_preempted,
+            extra_fields=extra_fields,
+        )
+
+
+# ---------------------------------------------------------------------
+#                          Fixtures & tests
+# ---------------------------------------------------------------------
+
+
+@pytest.fixture
+def init_server():
+    """Create and launch a vLLMOmniHttpServerLocal Ray actor with Qwen-Image."""
+    model_path = _resolve_model_path(MODEL)
+
+    ray.init(
+        runtime_env={
+            "env_vars": {
+                "TOKENIZERS_PARALLELISM": "true",
+                "NCCL_DEBUG": "WARN",
+                "VLLM_LOGGING_LEVEL": "INFO",
+            }
+        },
+        ignore_reinit_error=True,
+    )
+
+    rollout_cfg = OmegaConf.create(
+        {
+            "name": "vllm_omni",
+            "mode": "async",
+            "tensor_model_parallel_size": 1,
+            "data_parallel_size": 1,
+            "pipeline_model_parallel_size": 1,
+            "gpu_memory_utilization": 0.8,
+            "max_num_batched_tokens": 8192,
+            "max_num_seqs": 256,
+            "max_model_len": 1058,
+            "dtype": "bfloat16",
+            "load_format": "auto",
+            "enforce_eager": True,
+            "enable_chunked_prefill": False,
+            "enable_prefix_caching": False,
+            "enable_sleep_mode": False,
+            "free_cache_engine": True,
+            "disable_log_stats": True,
+            "logprobs_mode": "processed_logprobs",
+            "scheduling_policy": "fcfs",
+            "seed": 0,
+            "n": 4,
+            "pipeline": {
+                "height": 512,
+                "width": 512,
+                "num_inference_steps": 10,
+            },
+        }
+    )
+
+    model_cfg = OmegaConf.create(
+        {
+            "path": str(model_path),
+            "local_path": str(model_path),
+            "tokenizer_path": TOKENIZER_MODEL,
+            "trust_remote_code": True,
+            "load_tokenizer": True,
+        }
+    )
+
+    ServerCls = ray.remote(vLLMOmniHttpServerLocal)
+    server = ServerCls.options(
+        runtime_env={
+            "env_vars": {
+                "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1",
+                "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1",
+                "NCCL_CUMEM_ENABLE": "0",
+            }
+        },
+        max_concurrency=10,
+        num_gpus=1,
+    ).remote(
+        config=rollout_cfg,
+        model_config=model_cfg,
+        rollout_mode=RolloutMode.STANDALONE,
+        workers=[],
+        replica_rank=0,
+        node_rank=0,
+        gpus_per_node=1,
+        nnodes=1,
+        cuda_visible_devices="0",
+    )
+
+    ray.get(server.launch_server.remote())
+
+    yield server
+
+    ray.shutdown()
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_generate(init_server):
+    """Concurrent generate() calls covering basic output, logprobs, and multi-request correctness."""
+    server = init_server
+
+    prompts = [
+        "a beautiful sunset over the ocean with vibrant orange and purple clouds "
+        "reflecting on the calm water surface near a rocky coastline",
+        "a fluffy orange cat sitting on a wooden windowsill looking outside at "
+        "a garden full of colorful flowers on a bright sunny afternoon",
+        "a majestic mountain landscape covered with fresh white snow under a "
+        "clear blue sky with pine trees in the foreground and a frozen lake",
+        "a futuristic city at night with neon lights glowing on tall glass "
+        "skyscrapers and flying vehicles soaring between the buildings",
+    ]
+
+    refs = []
+    for i, prompt in enumerate(prompts):
+        rid = f"test_{i}_{uuid4().hex[:8]}"
+        ref = server.generate.remote(
+            prompt_ids=_tokenize_prompt(prompt),
+            sampling_params={
+                "num_inference_steps": 10,
+                "true_cfg_scale": 4.0,
+                "height": 512,
+                "width": 512,
+                "logprobs": i == 0,  # first request includes logprobs
+            },
+            request_id=rid,
+        )
+        refs.append(ref)
+
+    results = ray.get(refs, timeout=600)
+
+    for i, output in enumerate(results):
+        assert isinstance(output, DiffusionOutput), f"Request {i}: expected DiffusionOutput"
+        img = output.diffusion_output
+        assert isinstance(img, torch.Tensor) and img.ndim == 3, (
+            f"Request {i}: expected CHW torch.Tensor, got {type(img).__name__} ndim={getattr(img, 'ndim', None)}"
+        )
+        c, h, w = img.shape
+        assert c == 3, f"Request {i}: expected 3 channels (CHW), got {c}"
+        assert h > 0 and w > 0, f"Request {i}: image dimensions must be positive"
+        assert output.stop_reason in ("completed", "aborted", None), (
+            f"Request {i}: unexpected stop_reason {output.stop_reason!r}"
+        )
+        assert 0.0 <= float(img[0, 0, 0]) <= 1.0, f"Request {i}: pixel values must be in [0, 1]"
+
+    print(f"All {len(prompts)} concurrent requests returned valid DiffusionOutput")

vllm_omni/engine/async_omni_engine.py

@@ -877,43 +877,78 @@ class AsyncOmniEngine:
         stage_plans: Sequence[LogicalStageInitPlan],
         stage_init_timeout: int,
     ) -> dict[int, list[Any | None]]:
-        """Initialize all stage replicas in parallel."""
+        """Initialize all stage replicas.
+
+        Diffusion replicas are launched **inline on the orchestrator thread**
+        (the long-lived daemon thread created in ``__init__``). Their
+        ``mp.Process`` workers are therefore parented by a thread whose
+        lifetime equals the engine's lifetime. Submitting diffusion init to a
+        scoped ``ThreadPoolExecutor`` causes the clone-parent Python thread to
+        be destroyed at the end of init, which under Ray's actor subreaper
+        leads the spawned ``DiffusionWorker`` processes to be silently
+        ``SIGKILL``ed (exitcode -9). See git blame on this method.
+
+        LLM replicas keep using the parallel init executor.
+        """
 
         stage_launch_lock = threading.Lock()
         initialized_clients_by_stage: dict[int, list[Any | None]] = {
             plan.stage_idx: [None] * len(plan.replicas) for plan in stage_plans
         }
-        total_replicas = sum(len(plan.replicas) for plan in stage_plans)
-        future_to_replica: dict[concurrent.futures.Future[Any], tuple[int, int]] = {}
         primary_exc: Exception | None = None
 
-        with concurrent.futures.ThreadPoolExecutor(
-            max_workers=max(1, total_replicas),
-            thread_name_prefix="stage-init",
-        ) as init_executor:
-            for plan in stage_plans:
-                for replica in plan.replicas:
+        # Partition replicas: diffusion runs inline on the caller's thread;
+        # LLM replicas are submitted to a scoped ThreadPoolExecutor.
+        diffusion_replicas: list[tuple[int, ReplicaInitPlan]] = []
+        llm_replicas: list[tuple[int, ReplicaInitPlan]] = []
+        for plan in stage_plans:
+            for replica in plan.replicas:
+                if replica.metadata.stage_type == "diffusion":
+                    diffusion_replicas.append((plan.stage_idx, replica))
+                else:
+                    llm_replicas.append((plan.stage_idx, replica))
+
+        # --- 1) Diffusion replicas: inline on the orchestrator thread. ---
+        for stage_idx, replica in diffusion_replicas:
+            try:
+                initialized_clients_by_stage[stage_idx][replica.replica_id] = self._initialize_replica(
+                    replica,
+                    stage_init_timeout,
+                    stage_launch_lock,
+                )
+            except Exception as exc:
+                primary_exc = exc
+                break
+
+        # --- 2) LLM replicas: parallel init via a scoped ThreadPoolExecutor. ---
+        if primary_exc is None and llm_replicas:
+            future_to_replica: dict[concurrent.futures.Future[Any], tuple[int, int]] = {}
+            with concurrent.futures.ThreadPoolExecutor(
+                max_workers=max(1, len(llm_replicas)),
+                thread_name_prefix="stage-init",
+            ) as init_executor:
+                for stage_idx, replica in llm_replicas:
                     future = init_executor.submit(
                         self._initialize_replica,
                         replica,
                         stage_init_timeout,
                         stage_launch_lock,
                     )
-                    future_to_replica[future] = (plan.stage_idx, replica.replica_id)
+                    future_to_replica[future] = (stage_idx, replica.replica_id)
 
-            for future in concurrent.futures.as_completed(future_to_replica):
-                stage_idx, replica_id = future_to_replica[future]
-                try:
-                    initialized_clients_by_stage[stage_idx][replica_id] = future.result()
-                except concurrent.futures.CancelledError:
-                    continue
-                except Exception as exc:
-                    if primary_exc is None:
-                        primary_exc = exc
-                        for other_future in future_to_replica:
-                            if other_future is future:
-                                continue
-                            other_future.cancel()
+                for future in concurrent.futures.as_completed(future_to_replica):
+                    stage_idx, replica_id = future_to_replica[future]
+                    try:
+                        initialized_clients_by_stage[stage_idx][replica_id] = future.result()
+                    except concurrent.futures.CancelledError:
+                        continue
+                    except Exception as exc:
+                        if primary_exc is None:
+                            primary_exc = exc
+                            for other_future in future_to_replica:
+                                if other_future is future:
+                                    continue
+                                other_future.cancel()
 
         if primary_exc is not None:
             setattr(primary_exc, "_initialized_clients_by_stage", initialized_clients_by_stage)