From 631251a1f8573fc1fcc325041bf1b3bf347226be Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Wed, 13 May 2026 15:31:48 +0800 Subject: [PATCH] [Bugfix, rl] Diffusion worker SIGKILL under Ray actor (exitcode -9) (#3533) Signed-off-by: knlnguyen1802 Co-authored-by: Samit <285365963@qq.com> --- .buildkite/test-ready.yml | 20 + pyproject.toml | 1 + .../custom_pipeline/worker_extension.py | 64 +- .../offline_inference/rlhf_test/__init__.py | 0 .../rlhf_test/test_verl_omni_e2e.py | 588 ++++++++++++++++++ vllm_omni/engine/async_omni_engine.py | 81 ++- 6 files changed, 727 insertions(+), 27 deletions(-) create mode 100644 tests/e2e/offline_inference/rlhf_test/__init__.py create mode 100644 tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py diff --git a/.buildkite/test-ready.yml b/.buildkite/test-ready.yml index f36b01e13..0fc174d8b 100644 --- a/.buildkite/test-ready.yml +++ b/.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: - | diff --git a/pyproject.toml b/pyproject.toml index 5d4273331..1cbf03dac 100644 --- a/pyproject.toml +++ b/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", diff --git a/tests/e2e/offline_inference/custom_pipeline/worker_extension.py b/tests/e2e/offline_inference/custom_pipeline/worker_extension.py index 962bc0127..6ed29d5a8 100644 --- a/tests/e2e/offline_inference/custom_pipeline/worker_extension.py +++ b/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.""" diff --git a/tests/e2e/offline_inference/rlhf_test/__init__.py b/tests/e2e/offline_inference/rlhf_test/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py b/tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py new file mode 100644 index 000000000..9bcaf0e40 --- /dev/null +++ b/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") diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index bcb33e417..151671051 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/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)