# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Ray-serializable verl-omni HTTP server test double for RLHF E2E tests. The actor class must live in an importable module (not a pytest test file) so Ray workers can unpickle ``vLLMOmniHttpServerLocal`` via ``tests.e2e.features.helpers.verl_omni_server``. """ from __future__ import annotations import argparse import json import logging import os from dataclasses import asdict from enum import Enum from typing import Any import torchvision.transforms as T from omegaconf import DictConfig, OmegaConf from pydantic import BaseModel, ConfigDict from vllm.entrypoints.openai.api_server import build_app from vllm.utils.argparse_utils import FlexibleArgumentParser import vllm_omni.entrypoints.cli.serve 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.features.helpers.custom_pipeline.QwenImagePipelineWithLogProbForTest" WORKER_EXTENSION_CLASS = "tests.e2e.features.helpers.custom_pipeline.vLLMOmniColocateWorkerExtensionForTest" class RolloutMode(Enum): HYBRID = "hybrid" COLOCATED = "colocated" STANDALONE = "standalone" 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] = {} 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 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 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.""" 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) 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 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, ) async def launch_server(self) -> None: 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} engine_kwargs.pop("output_mode", None) 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 = { "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", 1), "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) 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) await self.run_server(parsed) async def run_server(self, args: argparse.Namespace) -> None: engine_args = OmniEngineArgs.from_cli_args(args) engine_args = asdict(engine_args) 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 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: 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 metadata = ( final_res.multimodal_output.get("metadata", {}) if isinstance(final_res.multimodal_output, dict) else {} ) prompt_embeddings = metadata.get("prompt_embeddings", {}) if isinstance(metadata, dict) else {} if sampling_params.get("logprobs", False): all_log_probs = final_res.trajectory_log_probs log_probs = all_log_probs[0] if all_log_probs is not None else None else: log_probs = None all_latents = final_res.trajectory_latents all_timesteps = final_res.trajectory_timesteps prompt_embeds = prompt_embeddings.get("prompt_embeds") prompt_embeds_mask = prompt_embeddings.get("prompt_embeds_mask") negative_prompt_embeds = prompt_embeddings.get("negative_prompt_embeds") negative_prompt_embeds_mask = prompt_embeddings.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 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, )