From 33e03b44c5a54b4e9f9b639d19c92ac82dc8b758 Mon Sep 17 00:00:00 2001 From: wzliu Date: Tue, 18 Nov 2025 15:51:06 +0800 Subject: [PATCH] added omni mooncake connectors to distributed directory, integrate omni connector to omni-llm, rebased with main (twice). ray + mooncake connector work on 2 nodes, but needes several updates\n1.start ray inside vllm-omni \n2.placement group optimize \n3.fix close --- .gitignore | 1 + .../offline_inference/qwen2_5_omni/end2end.py | 13 + examples/offline_inference/tmp_test/README.md | 34 ++ .../offline_inference/tmp_test/end2end.py | 264 +++++++++++++ .../tmp_test/extract_prompts.py | 45 +++ .../tmp_test/processing_omni.py | 367 ++++++++++++++++++ .../offline_inference/tmp_test/ray_head.sh | 1 + .../offline_inference/tmp_test/ray_worker.sh | 1 + .../tmp_test/run_multiple_prompts.sh | 8 + .../tmp_test/run_single_prompt.sh | 11 + examples/offline_inference/tmp_test/utils.py | 312 +++++++++++++++ vllm_omni/distributed/__init__.py | 25 ++ vllm_omni/distributed/connectors/__init__.py | 44 +++ vllm_omni/distributed/connectors/adapter.py | 145 +++++++ vllm_omni/distributed/connectors/base.py | 119 ++++++ vllm_omni/distributed/connectors/config.py | 31 ++ .../connector_tests/run_single_prompt.sh | 9 + .../test_connector_integration.py | 78 ++++ .../connector_tests/test_epdg_pipeline.py | 332 ++++++++++++++++ vllm_omni/distributed/connectors/factory.py | 86 ++++ .../connectors/mooncake_connector.py | 157 ++++++++ vllm_omni/distributed/connectors/utils.py | 314 +++++++++++++++ vllm_omni/entrypoints/omni_llm.py | 287 +++++++++++--- vllm_omni/entrypoints/omni_stage.py | 35 +- vllm_omni/entrypoints/utils.py | 55 ++- .../stage_configs/qwen2_5_omni.yaml | 40 +- 26 files changed, 2736 insertions(+), 78 deletions(-) create mode 100644 examples/offline_inference/tmp_test/README.md create mode 100644 examples/offline_inference/tmp_test/end2end.py create mode 100644 examples/offline_inference/tmp_test/extract_prompts.py create mode 100644 examples/offline_inference/tmp_test/processing_omni.py create mode 100644 examples/offline_inference/tmp_test/ray_head.sh create mode 100644 examples/offline_inference/tmp_test/ray_worker.sh create mode 100644 examples/offline_inference/tmp_test/run_multiple_prompts.sh create mode 100644 examples/offline_inference/tmp_test/run_single_prompt.sh create mode 100644 examples/offline_inference/tmp_test/utils.py create mode 100644 vllm_omni/distributed/connectors/__init__.py create mode 100644 vllm_omni/distributed/connectors/adapter.py create mode 100644 vllm_omni/distributed/connectors/base.py create mode 100644 vllm_omni/distributed/connectors/config.py create mode 100644 vllm_omni/distributed/connectors/connector_tests/run_single_prompt.sh create mode 100644 vllm_omni/distributed/connectors/connector_tests/test_connector_integration.py create mode 100644 vllm_omni/distributed/connectors/connector_tests/test_epdg_pipeline.py create mode 100644 vllm_omni/distributed/connectors/factory.py create mode 100644 vllm_omni/distributed/connectors/mooncake_connector.py create mode 100644 vllm_omni/distributed/connectors/utils.py diff --git a/.gitignore b/.gitignore index 1c15c4395..ac8bee278 100644 --- a/.gitignore +++ b/.gitignore @@ -232,3 +232,4 @@ configs/development.yaml # Docker .dockerignore Dockerfile.dev +discussion diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py index 463db0aa8..0d266748a 100644 --- a/examples/offline_inference/qwen2_5_omni/end2end.py +++ b/examples/offline_inference/qwen2_5_omni/end2end.py @@ -283,6 +283,19 @@ def parse_args(): default=None, help="Path to a .txt file with one prompt per line (preferred).", ) + parser.add_argument( + "--worker-backend", + type=str, + default="process", + choices=["process","ray"], + help="backend" + ) + parser.add_argument( + "--ray-address", + type=str, + default=None, + help="Path to a .txt file with one prompt per line (preferred).", + return parser.parse_args() diff --git a/examples/offline_inference/tmp_test/README.md b/examples/offline_inference/tmp_test/README.md new file mode 100644 index 000000000..112ef5335 --- /dev/null +++ b/examples/offline_inference/tmp_test/README.md @@ -0,0 +1,34 @@ +# Offline Example of vLLM-omni for Qwen2.5-omni + +## 🛠️ Installation + +Please refer to [README.md](../../../README.md) + +## Run examples (Qwen2.5-omni) +### Multiple Prompts +Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can: +```bash +tar -xf /seedtts_testset.tar +cp seedtts_testset/en/meta.lst examples/offline_inference/qwen2_5_omni/meta.lst +python3 examples/offline_inference/qwen2_5_omni/extract_prompts.py \ + --input examples/offline_inference/qwen2_5_omni/meta.lst \ + --output examples/offline_inference/qwen2_5_omni/top100.txt \ + --topk 100 +``` +Get into the example folder +```bash +cd examples/offline_inference/qwen2_5_omni +``` +Then run the command below. +```bash +bash run_multiple_prompts.sh +``` +### Single Prompts +Get into the example folder +```bash +cd examples/offline_inference/qwen2_5_omni +``` +Then run the command below. +```bash +bash run_single_prompt.sh +``` diff --git a/examples/offline_inference/tmp_test/end2end.py b/examples/offline_inference/tmp_test/end2end.py new file mode 100644 index 000000000..7e2e1043a --- /dev/null +++ b/examples/offline_inference/tmp_test/end2end.py @@ -0,0 +1,264 @@ +import argparse +import os +import os as _os_env_toggle +import random + +import numpy as np +import soundfile as sf +import torch +from utils import make_omni_prompt +from vllm.sampling_params import SamplingParams + +from vllm_omni.entrypoints.omni_llm import OmniLLM + +_os_env_toggle.environ["VLLM_USE_V1"] = "1" + +SEED = 42 +# Set all random seeds +random.seed(SEED) +np.random.seed(SEED) +torch.manual_seed(SEED) +torch.cuda.manual_seed(SEED) +torch.cuda.manual_seed_all(SEED) + +# Make PyTorch deterministic +torch.backends.cudnn.deterministic = True +torch.backends.cudnn.benchmark = False + +# Set environment variables for deterministic behavior +os.environ["PYTHONHASHSEED"] = str(SEED) +os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--model", + required=True, + help="Path to merged model directory (will be created if downloading).", + ) + parser.add_argument("--thinker-model", type=str, default=None) + parser.add_argument("--talker-model", type=str, default=None) + parser.add_argument("--code2wav-model", type=str, default=None) + parser.add_argument( + "--hf-hub-id", + default="Qwen/Qwen2.5-Omni-7B", + help="Hugging Face repo id to download if needed.", + ) + parser.add_argument("--hf-revision", default=None, help="Optional HF revision (branch/tag/commit).") + parser.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.") + parser.add_argument("--voice-type", default="default", help="Voice type, e.g., m02, f030, default.") + parser.add_argument( + "--code2wav-dir", + default=None, + help="Path to code2wav folder (contains spk_dict.pt).", + ) + parser.add_argument("--dit-ckpt", default=None, help="Path to DiT checkpoint file (e.g., dit.pt).") + parser.add_argument("--bigvgan-ckpt", default=None, help="Path to BigVGAN checkpoint file.") + parser.add_argument("--dtype", default="bfloat16", choices=["float16", "bfloat16", "float32"]) + parser.add_argument("--max-model-len", type=int, default=32768) + parser.add_argument( + "--init-sleep-seconds", + type=int, + default=20, + help="Sleep seconds after starting each stage process to allow initialization (default: 20)", + ) + + parser.add_argument("--thinker-only", action="store_true") + parser.add_argument("--text-only", action="store_true") + parser.add_argument("--do-wave", action="store_true") + parser.add_argument( + "--prompt_type", + choices=[ + "text", + "audio", + "audio-long", + "audio-long-chunks", + "audio-long-expand-chunks", + "image", + "video", + "video-frames", + "audio-in-video", + "audio-in-video-v2", + "audio-multi-round", + "badcase-vl", + "badcase-text", + "badcase-image-early-stop", + "badcase-two-audios", + "badcase-two-videos", + "badcase-multi-round", + "badcase-voice-type", + "badcase-voice-type-v2", + "badcase-audio-tower-1", + "badcase-audio-only", + ], + default="text", + ) + parser.add_argument("--use-torchvision", action="store_true") + parser.add_argument("--tokenize", action="store_true") + parser.add_argument( + "--output-wav", + default="output.wav", + help="[Deprecated] Output wav directory (use --output-dir).", + ) + parser.add_argument( + "--output-dir", + default="outputs", + help="Output directory to save text and wav files together.", + ) + parser.add_argument( + "--thinker-hidden-states-dir", + default="thinker_hidden_states", + help="Path to thinker hidden states directory.", + ) + parser.add_argument( + "--batch-timeout", + type=int, + default=5, + help="Timeout for batching in seconds (default: 5)", + ) + parser.add_argument( + "--init-timeout", + type=int, + default=300, + help="Timeout for initializing stages in seconds (default: 300)", + ) + parser.add_argument( + "--shm-threshold-bytes", + type=int, + default=65536, + help="Threshold for using shared memory in bytes (default: 65536)", + ) + parser.add_argument( + "--enable-stats", + action="store_true", + default=False, + help="Enable writing detailed statistics (default: disabled)", + ) + parser.add_argument( + "--txt-prompts", + type=str, + default=None, + help="Path to a .txt file with one prompt per line (preferred).", + ) + parser.add_argument( + "--worker-backend", + type=str, + default="process", + choices=["process","ray"], + help="backend" + ) + parser.add_argument( + "--ray-address", + type=str, + default=None, + help="Path to a .txt file with one prompt per line (preferred)." + ) + + args = parser.parse_args() + return args + + +def main(): + args = parse_args() + model_name = args.model + try: + # Preferred: load from txt file (one prompt per line) + if getattr(args, "txt_prompts", None) and args.prompt_type == "text": + with open(args.txt_prompts, encoding="utf-8") as f: + lines = [ln.strip() for ln in f.readlines()] + args.prompts = [ln for ln in lines if ln != ""] + print(f"[Info] Loaded {len(args.prompts)} prompts from {args.txt_prompts}") + except Exception as e: + print(f"[Error] Failed to load prompts: {e}") + raise + + if args.prompts is None: + raise ValueError("No prompts provided. Use --prompts ... or --txt-prompts (with --prompt_type text)") + omni_llm = OmniLLM( + model=model_name, + log_stats=args.enable_stats, + log_file=("omni_llm_pipeline.log" if args.enable_stats else None), + init_sleep_seconds=args.init_sleep_seconds, + batch_timeout=args.batch_timeout, + init_timeout=args.init_timeout, + shm_threshold_bytes=args.shm_threshold_bytes, + worker_backend=args.worker_backend, + ray_address=args.ray_address + ) + thinker_sampling_params = SamplingParams( + temperature=0.0, # Deterministic - no randomness + top_p=1.0, # Disable nucleus sampling + top_k=-1, # Disable top-k sampling + max_tokens=2048, + seed=SEED, # Fixed seed for sampling + detokenize=True, + repetition_penalty=1.1, + ) + talker_sampling_params = SamplingParams( + temperature=0.9, + top_p=0.8, + top_k=40, + max_tokens=2048, + seed=SEED, # Fixed seed for sampling + detokenize=True, + repetition_penalty=1.05, + stop_token_ids=[8294], + ) + code2wav_sampling_params = SamplingParams( + temperature=0.0, # Deterministic - no randomness + top_p=1.0, # Disable nucleus sampling + top_k=-1, # Disable top-k sampling + max_tokens=2048, + seed=SEED, # Fixed seed for sampling + detokenize=True, + repetition_penalty=1.1, + ) + + sampling_params_list = [ + thinker_sampling_params, + talker_sampling_params, + code2wav_sampling_params, + ] + import time + for i in range(1): + t1 = time.time() + prompt = [make_omni_prompt(args, prompt) for prompt in args.prompts] + print(f"prompt:{prompt}") + omni_outputs = omni_llm.generate(prompt, sampling_params_list) + t2 = time.time() + print(f"==========> time:{t2-t1}") + + # Determine output directory: prefer --output-dir; fallback to --output-wav + output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav + os.makedirs(output_dir, exist_ok=True) + for stage_outputs in omni_outputs: + if stage_outputs.final_output_type == "text": + for output in stage_outputs.request_output: + request_id = int(output.request_id) + text_output = output.outputs[0].text + # Save aligned text file per request + prompt_text = args.prompts[request_id] + out_txt = os.path.join(output_dir, f"{request_id:05d}.txt") + lines = [] + lines.append("Prompt:\n") + lines.append(str(prompt_text) + "\n") + lines.append("vllm_text_output:\n") + lines.append(str(text_output).strip() + "\n") + try: + with open(out_txt, "w", encoding="utf-8") as f: + f.writelines(lines) + except Exception as e: + print(f"[Warn] Failed writing text file {out_txt}: {e}") + print(f"Request ID: {request_id}, Text saved to {out_txt}") + elif stage_outputs.final_output_type == "audio": + for output in stage_outputs.request_output: + request_id = int(output.request_id) + audio_tensor = output.multimodal_output["audio"] + output_wav = os.path.join(output_dir, f"output_{output.request_id}.wav") + sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000) + print(f"Request ID: {request_id}, Saved audio to {output_wav}") + + +if __name__ == "__main__": + main() diff --git a/examples/offline_inference/tmp_test/extract_prompts.py b/examples/offline_inference/tmp_test/extract_prompts.py new file mode 100644 index 000000000..3d940b995 --- /dev/null +++ b/examples/offline_inference/tmp_test/extract_prompts.py @@ -0,0 +1,45 @@ +#!/usr/bin/env python3 +import argparse +from typing import Optional + + +def extract_prompt(line: str) -> Optional[str]: + # 提取第一个 '|' 与第二个 '|' 之间的内容 + i = line.find("|") + if i == -1: + return None + j = line.find("|", i + 1) + if j == -1: + return None + return line[i + 1 : j].strip() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--input", "-i", required=True, help="Input .lst file path") + parser.add_argument("--output", "-o", required=True, help="Output file path") + parser.add_argument( + "--topk", + "-k", + type=int, + default=100, + help="Extract the top K prompts (default: 100)", + ) + args = parser.parse_args() + + prompts = [] + with open(args.input, encoding="utf-8", errors="ignore") as f: + for line in f: + if len(prompts) >= args.topk: + break + p = extract_prompt(line.rstrip("\n")) + if p: + prompts.append(p) + + with open(args.output, "w", encoding="utf-8") as f: + for p in prompts: + f.write(p + "\n") + + +if __name__ == "__main__": + main() diff --git a/examples/offline_inference/tmp_test/processing_omni.py b/examples/offline_inference/tmp_test/processing_omni.py new file mode 100644 index 000000000..a22220dd3 --- /dev/null +++ b/examples/offline_inference/tmp_test/processing_omni.py @@ -0,0 +1,367 @@ +from __future__ import annotations + +import base64 +import logging +import math +import os +import time +import warnings +from functools import lru_cache +from io import BytesIO + +import requests +import torch +import torchvision +from packaging import version +from PIL import Image +from torchvision import io, transforms +from torchvision.transforms import InterpolationMode + +logger = logging.getLogger(__name__) + +IMAGE_FACTOR = 28 +MIN_PIXELS = 4 * 28 * 28 +MAX_PIXELS = 16384 * 28 * 28 +MAX_RATIO = 200 + +VIDEO_MIN_PIXELS = 128 * 28 * 28 +VIDEO_MAX_PIXELS = 768 * 28 * 28 +VIDEO_TOTAL_PIXELS = 24576 * 28 * 28 +FRAME_FACTOR = 2 +FPS = 2.0 +FPS_MIN_FRAMES = 4 +FPS_MAX_FRAMES = 768 + +temporal_patch_size = 2 +spatial_patch_size = 14 +spatial_merge_size = 2 + + +def round_by_factor(number: int, factor: int) -> int: + """Returns the closest integer to 'number' that is divisible by 'factor'.""" + return round(number / factor) * factor + + +def ceil_by_factor(number: int, factor: int) -> int: + """Returns the smallest integer greater than or equal to 'number' that is + divisible by 'factor'.""" + return math.ceil(number / factor) * factor + + +def floor_by_factor(number: int, factor: int) -> int: + """Returns the largest integer less than or equal to 'number' that is + divisible by 'factor'.""" + return math.floor(number / factor) * factor + + +def smart_resize( + height: int, + width: int, + factor: int = IMAGE_FACTOR, + min_pixels: int = MIN_PIXELS, + max_pixels: int = MAX_PIXELS, +) -> tuple[int, int]: + """ + Rescales the image so that the following conditions are met: + + 1. Both dimensions (height and width) are divisible by 'factor'. + + 2. The total number of pixels is within the range ['min_pixels', 'max_pixels']. + + 3. The aspect ratio of the image is maintained as closely as possible. + """ + if max(height, width) / min(height, width) > MAX_RATIO: + raise ValueError( + f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}" + ) + h_bar = max(factor, round_by_factor(height, factor)) + w_bar = max(factor, round_by_factor(width, factor)) + if h_bar * w_bar > max_pixels: + beta = math.sqrt((height * width) / max_pixels) + h_bar = floor_by_factor(height / beta, factor) + w_bar = floor_by_factor(width / beta, factor) + elif h_bar * w_bar < min_pixels: + beta = math.sqrt(min_pixels / (height * width)) + h_bar = ceil_by_factor(height * beta, factor) + w_bar = ceil_by_factor(width * beta, factor) + return h_bar, w_bar + + +def fetch_image(ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACTOR) -> Image.Image: + if "image" in ele: + image = ele["image"] + else: + image = ele["image_url"] + image_obj = None + if isinstance(image, Image.Image): + image_obj = image + elif image.startswith("http://") or image.startswith("https://"): + image_obj = Image.open(requests.get(image, stream=True).raw) + elif image.startswith("file://"): + image_obj = Image.open(image[7:]) + elif image.startswith("data:image"): + if "base64," in image: + _, base64_data = image.split("base64,", 1) + data = base64.b64decode(base64_data) + image_obj = Image.open(BytesIO(data)) + else: + image_obj = Image.open(image) + if image_obj is None: + raise ValueError(f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}") + image = image_obj.convert("RGB") + # resize + if "resized_height" in ele and "resized_width" in ele: + resized_height, resized_width = smart_resize( + ele["resized_height"], + ele["resized_width"], + factor=size_factor, + ) + else: + width, height = image.size + min_pixels = ele.get("min_pixels", MIN_PIXELS) + max_pixels = ele.get("max_pixels", MAX_PIXELS) + resized_height, resized_width = smart_resize( + height, + width, + factor=size_factor, + min_pixels=min_pixels, + max_pixels=max_pixels, + ) + image = image.resize((resized_width, resized_height)) + + return image + + +def smart_nframes( + ele: dict, + total_frames: int, + video_fps: int | float, +) -> int: + """calculate the number of frames for video used for model inputs. + + Args: + ele (dict): a dict contains the configuration of video. + support either `fps` or `nframes`: + - nframes: the number of frames to extract for model inputs. + - fps: the fps to extract frames for model inputs. + - min_frames: the minimum number of frames of the video, + only used when fps is provided. + - max_frames: the maximum number of frames of the video, + only used when fps is provided. + total_frames (int): the original total number of frames of the video. + video_fps (int | float): the original fps of the video. + + Raises: + ValueError: nframes should in interval [FRAME_FACTOR, total_frames]. + + Returns: + int: the number of frames for video used for model inputs. + """ + assert not ("fps" in ele and "nframes" in ele), "Only accept either `fps` or `nframes`" + if "nframes" in ele: + nframes = round_by_factor(ele["nframes"], FRAME_FACTOR) + else: + fps = ele.get("fps", FPS) + min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES), FRAME_FACTOR) + max_frames = floor_by_factor( + ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)), + FRAME_FACTOR, + ) + nframes = total_frames / video_fps * fps + nframes = min(max(nframes, min_frames), max_frames) + nframes = round_by_factor(nframes, FRAME_FACTOR) + if not (FRAME_FACTOR <= nframes and nframes <= total_frames): + raise ValueError(f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}.") + return nframes + + +def _read_video_torchvision( + ele: dict, +) -> torch.Tensor: + """read video using torchvision.io.read_video + + Args: + ele (dict): a dict contains the configuration of video. + support keys: + - video: the path of video. support "file://", "http://", + "https://" and local path. + - video_start: the start time of video. + - video_end: the end time of video. + Returns: + torch.Tensor: the video tensor with shape (T, C, H, W). + """ + video_path = ele["video"] + if version.parse(torchvision.__version__) < version.parse("0.19.0"): + if "http://" in video_path or "https://" in video_path: + warnings.warn("torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0.") + if "file://" in video_path: + video_path = video_path[7:] + st = time.time() + video, audio, info = io.read_video( + video_path, + start_pts=ele.get("video_start", 0.0), + end_pts=ele.get("video_end", None), + pts_unit="sec", + output_format="TCHW", + ) + total_frames, video_fps = video.size(0), info["video_fps"] + total_duration = round(total_frames / video_fps, 3) + logger.info( + f"torchvision: {video_path=}, {total_frames=}, {video_fps=}, duration={total_duration}s, time={time.time() - st:.3f}s" + ) + nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps) + idx = torch.linspace(0, total_frames - 1, nframes).round().long() + video = video[idx] + return video, total_duration, nframes + + +def is_decord_available() -> bool: + import importlib.util + + return importlib.util.find_spec("decord") is not None + + +def _read_video_decord( + ele: dict, +) -> torch.Tensor: + """read video using decord.VideoReader + + Args: + ele (dict): a dict contains the configuration of video. + support keys: + - video: the path of video. support "file://", "http://", + "https://" and local path. + - video_start: the start time of video. + - video_end: the end time of video. + Returns: + torch.Tensor: the video tensor with shape (T, C, H, W). + """ + import decord + + video_path = ele["video"] + st = time.time() + vr = decord.VideoReader(video_path) + # TODO: support start_pts and end_pts + if "video_start" in ele or "video_end" in ele: + raise NotImplementedError("not support start_pts and end_pts in decord for now.") + total_frames, video_fps = len(vr), vr.get_avg_fps() + total_duration = round(total_frames / video_fps, 3) + logger.info(f"decord: {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s") + nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps) + idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist() + video = vr.get_batch(idx).asnumpy() + video = torch.tensor(video).permute(0, 3, 1, 2) # Convert to TCHW format + return video, total_duration, nframes + + +VIDEO_READER_BACKENDS = { + "decord": _read_video_decord, + "torchvision": _read_video_torchvision, +} + +FORCE_QWENVL_VIDEO_READER = os.getenv("FORCE_QWENVL_VIDEO_READER", None) + + +@lru_cache(maxsize=1) +def get_video_reader_backend() -> str: + if FORCE_QWENVL_VIDEO_READER is not None: + video_reader_backend = FORCE_QWENVL_VIDEO_READER + elif is_decord_available(): + video_reader_backend = "decord" + else: + video_reader_backend = "torchvision" + # print(f"qwen-vl-utils using {video_reader_backend} to read video.", + # file=sys.stderr) + return video_reader_backend + + +def fetch_video(ele: dict, image_factor: int = IMAGE_FACTOR) -> torch.Tensor | list[Image.Image]: + if isinstance(ele["video"], str): + video_reader_backend = get_video_reader_backend() + video, total_dur, nframes = VIDEO_READER_BACKENDS[video_reader_backend](ele) + frame_timestamps = total_dur * torch.arange(1, nframes + 1) / nframes + grid_timestamps = frame_timestamps[::FRAME_FACTOR] + second_per_grid = grid_timestamps[1] - grid_timestamps[0] + nframes, _, height, width = video.shape + min_pixels = ele.get("min_pixels", VIDEO_MIN_PIXELS) + total_pixels = ele.get("total_pixels", VIDEO_TOTAL_PIXELS) + max_pixels = max( + min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR), + int(min_pixels * 1.05), + ) + max_pixels = ele.get("max_pixels", max_pixels) + # min_pixels = (factor ** 2) * 52 + # max_pixels = (factor ** 2) * min(768, (16384 / nframes * temporal_patch_size)) + if "resized_height" in ele and "resized_width" in ele: + resized_height, resized_width = smart_resize( + ele["resized_height"], + ele["resized_width"], + factor=image_factor, + ) + else: + resized_height, resized_width = smart_resize( + height, + width, + factor=image_factor, + min_pixels=min_pixels, + max_pixels=max_pixels, + ) + video = transforms.functional.resize( + video, + [resized_height, resized_width], + interpolation=InterpolationMode.BICUBIC, + antialias=True, + ).float() + return video, total_dur, nframes, second_per_grid + else: + assert isinstance(ele["video"], (list, tuple)) + process_info = ele.copy() + process_info.pop("type", None) + process_info.pop("video", None) + images = [ + fetch_image({"image": video_element, **process_info}, size_factor=image_factor) + for video_element in ele["video"] + ] + nframes = ceil_by_factor(len(images), FRAME_FACTOR) + if len(images) < nframes: + images.extend([images[-1]] * (nframes - len(images))) + return images, None, None, None + + +def extract_vision_info(conversations: list[dict] | list[list[dict]]) -> list[dict]: + vision_infos = [] + if isinstance(conversations[0], dict): + conversations = [conversations] + for conversation in conversations: + for message in conversation: + if isinstance(message["content"], list): + for ele in message["content"]: + if ( + "image" in ele + or "image_url" in ele + or "video" in ele + or ele["type"] in ("image", "image_url", "video") + ): + vision_infos.append(ele) + return vision_infos + + +def process_vision_info( + conversations: list[dict] | list[list[dict]], +) -> tuple[list[Image.Image] | None, list[torch.Tensor | list[Image.Image]] | None]: + vision_infos = extract_vision_info(conversations) + # Read images or videos + image_inputs = [] + video_inputs = [] + for vision_info in vision_infos: + if "image" in vision_info or "image_url" in vision_info: + image_inputs.append(fetch_image(vision_info)) + elif "video" in vision_info: + video_inputs.append(fetch_video(vision_info)) + else: + raise ValueError("image, image_url or video should in content.") + if len(image_inputs) == 0: + image_inputs = None + if len(video_inputs) == 0: + video_inputs = None + return image_inputs, video_inputs diff --git a/examples/offline_inference/tmp_test/ray_head.sh b/examples/offline_inference/tmp_test/ray_head.sh new file mode 100644 index 000000000..e6e456db2 --- /dev/null +++ b/examples/offline_inference/tmp_test/ray_head.sh @@ -0,0 +1 @@ +ray start --head --port=6399 --num-gpus=1 diff --git a/examples/offline_inference/tmp_test/ray_worker.sh b/examples/offline_inference/tmp_test/ray_worker.sh new file mode 100644 index 000000000..5aa042757 --- /dev/null +++ b/examples/offline_inference/tmp_test/ray_worker.sh @@ -0,0 +1 @@ +ray start --address=10.90.67.86:6399 diff --git a/examples/offline_inference/tmp_test/run_multiple_prompts.sh b/examples/offline_inference/tmp_test/run_multiple_prompts.sh new file mode 100644 index 000000000..78d3dd54f --- /dev/null +++ b/examples/offline_inference/tmp_test/run_multiple_prompts.sh @@ -0,0 +1,8 @@ +python end2end.py --model Qwen/Qwen2.5-Omni-7B \ + --voice-type "m02" \ + --dit-ckpt none \ + --bigvgan-ckpt none \ + --output-wav output_audio \ + --prompt_type text \ + --init-sleep-seconds 0 \ + --txt-prompts top100.txt diff --git a/examples/offline_inference/tmp_test/run_single_prompt.sh b/examples/offline_inference/tmp_test/run_single_prompt.sh new file mode 100644 index 000000000..367466284 --- /dev/null +++ b/examples/offline_inference/tmp_test/run_single_prompt.sh @@ -0,0 +1,11 @@ +export PYTHONPATH=/workspace/omni/vllm-omni/:$PYTHONPATH +CUDA_VISIBLE_DEVICES=4 python end2end.py --model /workspace/Qwen2.5-Omni-7B/ \ + --voice-type "m02" \ + --dit-ckpt none \ + --bigvgan-ckpt none \ + --output-wav output_audio \ + --prompt_type text \ + --init-sleep-seconds 0 \ + --worker-backend ray \ + --ray-address auto \ + --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words." diff --git a/examples/offline_inference/tmp_test/utils.py b/examples/offline_inference/tmp_test/utils.py new file mode 100644 index 000000000..11df5c55a --- /dev/null +++ b/examples/offline_inference/tmp_test/utils.py @@ -0,0 +1,312 @@ +import tempfile +from typing import Optional, Union +from urllib.request import urlopen + +import librosa +import requests +import resampy +import soundfile as sf +import torch +import torchvision.io +from processing_omni import fetch_image, fetch_video +from transformers import AutoConfig, AutoProcessor +from vllm.inputs import TextPrompt + +from vllm_omni.inputs.data import OmniTokensPrompt + +# Simple caches to avoid repeated heavy HF loads per prompt +_PROCESSOR_CACHE: dict[str, "AutoProcessor"] = {} +_CONFIG_CACHE: dict[str, "AutoConfig"] = {} + + +def get_system_prompt(): + return { + "role": "system", + "content": [ + { + "type": "text", + "text": ( + "You are Qwen, a virtual human developed by the Qwen Team, " + "Alibaba Group, capable of perceiving auditory and visual inputs, " + "as well as generating text and speech." + ), + } + ], + } + + +def resample_wav_to_16khz(input_filepath): + data, original_sample_rate = sf.read(input_filepath) + # Only use the first channel + if len(data.shape) > 1: + data = data[:, 0] + # resample to 16kHz + data_resampled = resampy.resample(data, sr_orig=original_sample_rate, sr_new=16000) + return data_resampled + + +def fetch_and_read_video(args, video_url: str, fps=2): + def read_video_with_torchvision(video_file_name: str): + video, audio, info = torchvision.io.read_video( + video_file_name, + start_pts=0.0, + end_pts=None, + pts_unit="sec", + output_format="TCHW", + ) + + total_frames, video_fps = video.size(0), info["video_fps"] + total_duration = round(total_frames / video_fps, 3) + nframes = int(total_frames / video_fps * fps) + + frame_timestamps = total_duration * torch.arange(1, nframes + 1) / nframes + grid_timestamps = frame_timestamps[::2] + second_per_grid = grid_timestamps[1] - grid_timestamps[0] + + idx = torch.linspace(0, video.size(0) - 1, nframes).round().long() + video = video[idx] + + if args.legacy_omni_video: + return [video, total_duration, nframes, second_per_grid.item()] + else: + return video + + def read_video_with_transformers(video_file_name: Union[str, list[str]]): + video, total_duration, nframes, second_per_grid = fetch_video({"video": video_file_name}) + if total_duration is None and nframes is None: + nframes = len(video) + total_duration = 0.5 * nframes + second_per_grid = 1.0 + if args.legacy_omni_video: + return [video, total_duration, nframes, second_per_grid] + else: + return video + + def read_video(video_file_name: str): + if args.use_torchvision: + return read_video_with_torchvision(video_file_name) + else: + return read_video_with_transformers(video_file_name) + + if isinstance(video_url, str) and video_url.startswith("http"): + with tempfile.NamedTemporaryFile(delete=True) as temp_video_file: + resp = requests.get(video_url) + assert resp.status_code == requests.codes.ok, ( + f"Failed to fetch video from {video_url}, status_code:{resp.status_code}, resp:{resp}" + ) + + temp_video_file.write(urlopen(video_url).read()) + temp_video_file_path = temp_video_file.name + video_file_name = temp_video_file_path + return read_video(video_file_name) + else: + video_file_name = video_url + return read_video(video_file_name) + + +def make_inputs_qwen2_omni( + args, + messages: list[dict[str, Union[str, list[dict[str, str]]]]], + use_audio_in_video: Optional[bool] = False, + tokenize: bool = False, +) -> Union[OmniTokensPrompt, TextPrompt]: + from transformers import AutoConfig, AutoProcessor + + # Cached processor/config to prevent per-prompt reloading and repeated warnings + if args.model not in _PROCESSOR_CACHE: + _PROCESSOR_CACHE[args.model] = AutoProcessor.from_pretrained(args.model) + processor = _PROCESSOR_CACHE[args.model] + + config = _CONFIG_CACHE.get(args.model) + if config is None: + try: + config = AutoConfig.from_pretrained(args.model) + except Exception: + config = None + _CONFIG_CACHE[args.model] = config # cache even if None to avoid retry storms + + # Decide legacy flag only once based on config (default True if unknown) + if getattr(args, "legacy_omni_video", None) is None: + if config is not None and hasattr(config, "architectures"): + args.legacy_omni_video = "Qwen2_5OmniModel" not in config.architectures + else: + args.legacy_omni_video = True + + audios, images, videos = [], [], [] + for message in messages: + if not isinstance(message["content"], list): + message["content"] = [ + { + "type": "text", + "text": message["content"], + } + ] + index, num_contents = 0, len(message["content"]) + while index < num_contents: + ele = message["content"][index] + if "type" not in ele: + if "text" in ele: + ele["type"] = "text" + elif "audio" in ele: + ele["type"] = "audio" + elif "audio_url" in ele: + ele["type"] = "audio_url" + elif "image" in ele: + ele["type"] = "image" + elif "image_url" in ele: + ele["type"] = "image_url" + elif "video" in ele: + ele["type"] = "video" + elif "video_url" in ele: + ele["type"] = "video_url" + else: + raise ValueError(f"Unknown ele: {ele}") + + if ele["type"] == "audio" or ele["type"] == "audio_url": + if "audio_url" in ele: + audio_key = "audio_url" + with tempfile.NamedTemporaryFile(delete=True) as temp_audio_file: + temp_audio_file.write(urlopen(ele[audio_key]).read()) + temp_audio_file_path = temp_audio_file.name + audios.append(resample_wav_to_16khz(temp_audio_file_path)) + ele["audio"] = temp_audio_file_path + elif "audio" in ele: + audio_key = "audio" + audios.append(resample_wav_to_16khz(ele[audio_key])) + else: + raise ValueError(f"Unknown ele {ele}") + elif use_audio_in_video and (ele["type"] == "video" or ele["type"] == "video_url"): + # use video as audio as well + if "video_url" in ele: + audio_key = "video_url" + with tempfile.NamedTemporaryFile(delete=True) as temp_video_file: + temp_video_file.write(urlopen(ele[audio_key]).read()) + temp_video_file_path = temp_video_file.name + ele[audio_key] = temp_video_file_path + audios.append(librosa.load(temp_video_file_path, sr=16000)[0]) + videos.append(fetch_and_read_video(args, temp_video_file_path)) + ele["video"] = temp_video_file_path + elif "video" in ele: + audio_key = "video" + audios.append(librosa.load(ele[audio_key], sr=16000)[0]) + videos.append(fetch_and_read_video(args, audio_key)) + else: + raise ValueError(f"Unknown ele {ele}") + # insert a audio after the video + message["content"].insert( + index + 1, + { + "type": "audio", + "audio": ele[audio_key], + }, + ) + # no need to load the added audio again + index += 1 + elif ele["type"] == "video" or ele["type"] == "video_url": + if "video_url" in ele: + video_key = "video_url" + with tempfile.NamedTemporaryFile(delete=True) as temp_video_file: + temp_video_file.write(urlopen(ele["video_url"]).read()) + temp_video_file_path = temp_video_file.name + videos.append(fetch_and_read_video(args, temp_video_file)) + ele["video"] = temp_video_file_path + else: + video_key = "video" + videos.append(fetch_and_read_video(args, ele[video_key])) + elif ele["type"] == "image" or ele["type"] == "image_url": + images.append(fetch_image(ele)) + + # move to the next content + index += 1 + + prompt = processor.apply_chat_template( + messages, + tokenize=tokenize, + add_generation_prompt=True, + add_vision_id=True, + ) + + audios = audios if len(audios) > 0 else None + images = images if len(images) > 0 else None + videos = videos if len(videos) > 0 else None + + multi_modal_data = {} + if audios: + multi_modal_data["audio"] = audios + if images: + multi_modal_data["image"] = images + if videos: + multi_modal_data["video"] = videos + + if isinstance(prompt, list) and isinstance(prompt[0], (list, str)): + prompt = prompt[0] + + if tokenize: + return OmniTokensPrompt( + prompt_token_ids=prompt, + multi_modal_data=multi_modal_data, + ) + else: + return TextPrompt( + prompt=prompt, + multi_modal_data=multi_modal_data, + ) + + +def make_text_prompt(args, prompt): + messages = [ + get_system_prompt(), + { + "role": "user", + "content": [ + {"type": "text", "text": prompt}, + ], + }, + ] + + prompt = make_inputs_qwen2_omni(args, messages, tokenize=args.tokenize) + return prompt + + +def make_audio_in_video_v2_prompt(args): + messages = [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": ( + "You are Qwen, a virtual human developed by the Qwen Team, " + "Alibaba Group, capable of perceiving auditory and visual " + "inputs, as well as generating text and speech." + ), + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "video_url", + "video_url": ("https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-Omni/draw_small.mp4"), + }, + ], + }, + ] + prompt = make_inputs_qwen2_omni( + args, + messages, + use_audio_in_video=True, + tokenize=args.tokenize, + ) + return prompt + + +def make_omni_prompt(args, prompt=None) -> Union[OmniTokensPrompt, list[OmniTokensPrompt]]: + if args.prompt_type == "text": + prompt = make_text_prompt(args, prompt) + elif args.prompt_type == "audio-in-video-v2": + prompt = make_audio_in_video_v2_prompt(args) + else: + raise ValueError(f"Unsupported prompt type: {args.prompt_type}") + return prompt diff --git a/vllm_omni/distributed/__init__.py b/vllm_omni/distributed/__init__.py index e69de29bb..50580be70 100644 --- a/vllm_omni/distributed/__init__.py +++ b/vllm_omni/distributed/__init__.py @@ -0,0 +1,25 @@ +# SPDX-License-Identifier: Apache-2.0 + +from .connectors import ( + ConnectorSpec, OmniTransferConfig, + OmniConnectorBase, InMemoryOmniConnector, + OmniConnectorFactory, MooncakeConnector, + load_omni_transfer_config, create_simple_config, create_mooncake_config +) + +__all__ = [ + # Config + "ConnectorSpec", + "OmniTransferConfig", + + # Connectors + "OmniConnectorBase", + "InMemoryOmniConnector", + "OmniConnectorFactory", + "MooncakeConnector", + + # Utilities + "load_omni_transfer_config", + "create_simple_config", + "create_mooncake_config", +] diff --git a/vllm_omni/distributed/connectors/__init__.py b/vllm_omni/distributed/connectors/__init__.py new file mode 100644 index 000000000..8ce2f1396 --- /dev/null +++ b/vllm_omni/distributed/connectors/__init__.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 + +from .config import ConnectorSpec, OmniTransferConfig +from .base import OmniConnectorBase, InMemoryOmniConnector +from .factory import OmniConnectorFactory +from .mooncake_connector import MooncakeConnector +from .utils import ( + load_omni_transfer_config, + initialize_connectors_from_config, + create_simple_config, + create_mooncake_config, + get_connectors_config_for_stage, + initialize_orchestrator_connectors, + get_stage_connector_config, + build_stage_connectors, +) + +__all__ = [ + # Config + "ConnectorSpec", + "OmniTransferConfig", + + # Base classes and implementations + "OmniConnectorBase", + "InMemoryOmniConnector", + + # Factory + "OmniConnectorFactory", + + # Specific implementations + "MooncakeConnector", + + # Utilities + "load_omni_transfer_config", + "initialize_connectors_from_config", + "create_simple_config", + "create_mooncake_config", + "get_connectors_config_for_stage", + + # Manager helpers + "initialize_orchestrator_connectors", + "get_stage_connector_config", + "build_stage_connectors", +] diff --git a/vllm_omni/distributed/connectors/adapter.py b/vllm_omni/distributed/connectors/adapter.py new file mode 100644 index 000000000..747de0cd5 --- /dev/null +++ b/vllm_omni/distributed/connectors/adapter.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# temporary for compatibility with vllm_omni.entrypoints.omni_stage.py +# and vllm_omni.entrypoints.omni_llm.py + +import time +import logging +from typing import Any, Optional, Dict, Callable + +logger = logging.getLogger(__name__) + + +def try_send_via_connector( + connector: Any, + stage_id: int, + next_stage_id: int, + req_id: str, + next_inputs: Any, + sampling_params: Any, + original_prompt: Any, + next_stage_queue_submit_fn: Callable[[Dict[str, Any]], None], + metrics: Any +) -> bool: + """ + Attempts to send data via OmniConnector. + Returns True if successful, False otherwise. + Encapsulates the logic of preparing payload, sending via connector, + sending notification, and recording metrics. + """ + try: + t0 = time.time() + + # Prepare data for connector + payload_data = { + "engine_inputs": next_inputs, + "sampling_params": sampling_params, + "metadata": { + "original_prompt": original_prompt, + "stage_transition": f"{stage_id}->{next_stage_id}", + "timestamp": time.time() + } + } + + # Send data via connector + success, serialized_size = connector.put(str(stage_id), str(next_stage_id), str(req_id), payload_data) + + if success: + # Send lightweight notification via queue + notify_payload = { + "request_id": req_id, + "sampling_params": sampling_params, + "from_connector": True, + "from_stage": str(stage_id), + "to_stage": str(next_stage_id), + "sent_ts": time.time(), + } + next_stage_queue_submit_fn(notify_payload) + + t1 = time.time() + tx_ms = (t1 - t0) * 1000.0 + + metrics.on_forward( + stage_id, + next_stage_id, + req_id, + serialized_size, # Use size from connector + float(tx_ms), + True, # Mark as using connector + ) + return True + else: + # If put returned False, we let the caller handle fallback + return False + + except Exception as e: + logger.warning( + "[Orchestrator] OmniConnector failed for req %s: %s; falling back to queue", + req_id, + e, + ) + return False + + +def try_recv_via_connector( + task: Dict[str, Any], + connectors: Dict[Any, Any], + stage_id: int, +) -> tuple[Any, Optional[Dict[str, Any]]]: + """ + Attempts to resolve input data from either connector or IPC. + Returns (engine_inputs, rx_metrics) or (None, None) if failed/skipped. + """ + rid = task["request_id"] + + if task.get("from_connector"): + from_stage = task.get("from_stage") + to_stage = str(stage_id) + + if not from_stage: + logger.error( + "[Stage-%s] 'from_connector' is true but 'from_stage' is missing for request %s", stage_id, rid + ) + return None, None + + # Get connector for this edge + connector_key = (from_stage, to_stage) + connector = connectors.get(connector_key) + + if connector: + try: + # Get data from connector with timeout + _t_start = time.time() + payload_data = connector.get(from_stage, to_stage, str(rid)) + _t_end = time.time() + + if payload_data and isinstance(payload_data, dict): + ein = payload_data.get("engine_inputs") + # Use connector's serialization size for consistency + serialized_size = len(connector.serialize_obj(payload_data)) + decode_ms = (_t_end - _t_start) * 1000.0 + + rx_metrics = { + "rx_decode_time_ms": decode_ms, + "rx_transfer_bytes": serialized_size + } + return ein, rx_metrics + else: + logger.error( + "[Stage-%s] Failed to get data from connector for request %s or payload is empty", stage_id, rid + ) + return None, None + except Exception as e: + logger.error( + "[Stage-%s] Error retrieving data from connector for request %s: %s", stage_id, rid, e + ) + return None, None + else: + logger.error( + "[Stage-%s] No connector found for edge %s -> %s for request %s", stage_id, from_stage, to_stage, rid + ) + return None, None + else: + # Data comes from queue as usual + #return maybe_load_from_ipc_with_metrics_fn(task, obj_key="engine_inputs", shm_key="engine_inputs_shm") + return None, None + diff --git a/vllm_omni/distributed/connectors/base.py b/vllm_omni/distributed/connectors/base.py new file mode 100644 index 000000000..3b75bbff5 --- /dev/null +++ b/vllm_omni/distributed/connectors/base.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 + +import logging +from abc import ABC, abstractmethod +from typing import Any, Optional + +logger = logging.getLogger(__name__) + + +class OmniConnectorBase(ABC): + """Base class for all OmniConnectors.""" + + @abstractmethod + def put(self, from_stage: str, to_stage: str, request_id: str, data: Any) -> tuple[bool, int]: + """Store Python object, internal serialization handled by connector. + + Args: + from_stage: Source stage identifier + to_stage: Destination stage identifier + request_id: Unique request identifier + data: Python object to store + + Returns: + tuple: (success: bool, serialized_size: int) + """ + pass + + @abstractmethod + def get(self, from_stage: str, to_stage: str, request_id: str) -> Optional[Any]: + """Retrieve Python object, internal deserialization handled by connector. + + Args: + from_stage: Source stage identifier + to_stage: Destination stage identifier + request_id: Unique request identifier + + Returns: + Python object if found, None otherwise + """ + pass + + @abstractmethod + def cleanup(self, request_id: str) -> None: + """Clean up resources for a request.""" + pass + + @abstractmethod + def health(self) -> dict[str, Any]: + """Return health status and metrics.""" + pass + + @staticmethod + def serialize_obj(obj: Any) -> bytes: + """Serialize a Python object to bytes using torch.save (consistent with stage_utils).""" + import io + import torch + buf = io.BytesIO() + torch.save(obj, buf) + return buf.getvalue() + + @staticmethod + def deserialize_obj(data: bytes) -> Any: + """Deserialize bytes to Python object using torch.load (consistent with stage_utils).""" + import io + import torch + return torch.load(io.BytesIO(data), map_location="cpu", weights_only=False) + + +class InMemoryOmniConnector(OmniConnectorBase): + """Simple in-memory connector for testing and co-located stages.""" + + def __init__(self): + self._storage: dict[str, Any] = {} + self._metrics = { + "puts": 0, + "gets": 0, + "errors": 0, + } + + def _make_key(self, from_stage: str, to_stage: str, request_id: str) -> str: + return f"{from_stage}:{to_stage}:{request_id}" + + def put(self, from_stage: str, to_stage: str, request_id: str, data: Any) -> tuple[bool, int]: + try: + key = self._make_key(from_stage, to_stage, request_id) + self._storage[key] = data + serialized_size = len(self.serialize_obj(data)) + self._metrics["puts"] += 1 + logger.debug(f"InMemoryConnector: stored {key}") + return True, serialized_size + except Exception as e: + self._metrics["errors"] += 1 + logger.error(f"InMemoryConnector put failed: {e}") + return False, 0 + + def get(self, from_stage: str, to_stage: str, request_id: str) -> Optional[Any]: + try: + key = self._make_key(from_stage, to_stage, request_id) + data = self._storage.get(key) + if data is not None: + self._metrics["gets"] += 1 + return data + return None + except Exception as e: + logger.error(f"InMemoryConnector get failed: {e}") + return None + + def cleanup(self, request_id: str) -> None: + keys_to_remove = [k for k in self._storage.keys() if k.endswith(f":{request_id}")] + for key in keys_to_remove: + del self._storage[key] + logger.debug(f"InMemoryConnector: cleaned up {len(keys_to_remove)} keys for {request_id}") + + def health(self) -> dict[str, Any]: + return { + "status": "healthy", + "storage_size": len(self._storage), + **self._metrics, + } diff --git a/vllm_omni/distributed/connectors/config.py b/vllm_omni/distributed/connectors/config.py new file mode 100644 index 000000000..1a6e9e23d --- /dev/null +++ b/vllm_omni/distributed/connectors/config.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 + +import logging +from dataclasses import dataclass, field +from typing import Any, Optional + +logger = logging.getLogger(__name__) + + +@dataclass +class ConnectorSpec: + """Specification for a connector instance.""" + name: str # e.g., "MooncakeConnector", "InMemoryConnector" + extra: dict[str, Any] = field(default_factory=dict) # backend-specific config + + +@dataclass +class OmniTransferConfig: + """Top-level configuration for OmniConnector system.""" + # Direct mapping: (from_stage, to_stage) -> connector + connectors: dict[tuple[str, str], ConnectorSpec] = field(default_factory=dict) + default_connector: Optional[ConnectorSpec] = None + + def get_connector_for_edge(self, from_stage: str, to_stage: str) -> Optional[ConnectorSpec]: + """Get connector spec for a specific edge.""" + edge_key = (from_stage, to_stage) + return self.connectors.get(edge_key, self.default_connector) + + def has_connector_for_edge(self, from_stage: str, to_stage: str) -> bool: + """Check if there's a connector configured for the edge.""" + return self.get_connector_for_edge(from_stage, to_stage) is not None diff --git a/vllm_omni/distributed/connectors/connector_tests/run_single_prompt.sh b/vllm_omni/distributed/connectors/connector_tests/run_single_prompt.sh new file mode 100644 index 000000000..5e0e7cd09 --- /dev/null +++ b/vllm_omni/distributed/connectors/connector_tests/run_single_prompt.sh @@ -0,0 +1,9 @@ +export PYTHONPATH=/workspace/omni/vllm-omni/:$PYTHONPATH +CUDA_VISIBLE_DEVICES=0,1,2,3 python end2end.py --model /workspace/Qwen2.5-Omni-7B/ \ + --voice-type "m02" \ + --dit-ckpt none \ + --bigvgan-ckpt none \ + --output-wav output_audio \ + --prompt_type text \ + --init-sleep-seconds 0 \ + --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words." diff --git a/vllm_omni/distributed/connectors/connector_tests/test_connector_integration.py b/vllm_omni/distributed/connectors/connector_tests/test_connector_integration.py new file mode 100644 index 000000000..ecf1bd9cb --- /dev/null +++ b/vllm_omni/distributed/connectors/connector_tests/test_connector_integration.py @@ -0,0 +1,78 @@ +#!/usr/bin/env python3 +""" +Test OmniConnector integration with OmniLLM. +""" + +import sys +import os +sys.path.insert(0, os.path.dirname(__file__)) + +from vllm_omni.distributed.connectors import create_mooncake_config + +def test_connector_config(): + """Test connector configuration creation.""" + print("Testing connector configuration...") + + # Create Mooncake config + config = create_mooncake_config( + host="127.0.0.1", + metadata_server="http://127.0.0.1:8080/metadata", + master="127.0.0.1:50051" + ) + + print(f"Created config with {len(config.connectors)} connectors") + for edge, connector_spec in config.connectors.items(): + print(f" {edge}: {connector_spec.name}") + + return config + +def test_omni_llm_initialization(): + """Test OmniLLM with connector config.""" + print("\nTesting OmniLLM initialization with connectors...") + + try: + from vllm_omni.entrypoints.omni_llm import OmniLLM + + # Create connector config + connector_config = { + "connectors": { + "0_to_1": { + "name": "InMemoryConnector", # Use InMemory for testing + "extra": {} + } + } + } + + # Test connector initialization logic + omni = OmniLLM.__new__(OmniLLM) # Create without calling __init__ + omni.omni_transfer_config = connector_config + omni.connectors = {} + omni._initialize_connectors() + + # Check if connector was created + if ("0", "1") in omni.connectors: + print(f"✅ Connector created: {type(omni.connectors[('0', '1')]).__name__}") + return True + else: + print("❌ Connector not created") + return False + + except Exception as e: + print(f"Error: {e}") + return False + +if __name__ == "__main__": + print("OmniConnector Integration Test") + print("=" * 40) + + # Test config creation + config = test_connector_config() + + # Test OmniLLM integration + success = test_omni_llm_initialization() + + if success: + print("\n✅ Integration test passed!") + print("Connector system is ready for use with OmniLLM") + else: + print("\n❌ Integration test failed!") diff --git a/vllm_omni/distributed/connectors/connector_tests/test_epdg_pipeline.py b/vllm_omni/distributed/connectors/connector_tests/test_epdg_pipeline.py new file mode 100644 index 000000000..7598ba1f7 --- /dev/null +++ b/vllm_omni/distributed/connectors/connector_tests/test_epdg_pipeline.py @@ -0,0 +1,332 @@ +#!/usr/bin/env python3 +""" +Test EPD (Encode-Prefill-Decode) 3-stage pipeline with connectors. +Tests data flow between stages using configured connectors. +""" + +import os +import sys +from typing import Any, Dict, Optional +from dataclasses import dataclass +import uuid + +try: + import yaml +except ImportError: + print("PyYAML not available, using basic YAML parsing") + yaml = None + +# Add the connectors package to path +current_dir = os.path.dirname(os.path.abspath(__file__)) +parent_dir = os.path.dirname(current_dir) +sys.path.insert(0, parent_dir) + +# Import from parent package +try: + from connectors.config import ConnectorSpec, OmniTransferConfig + from connectors.base import OmniConnectorBase, InMemoryOmniConnector + from connectors.factory import OmniConnectorFactory + from connectors.utils import load_omni_transfer_config +except ImportError: + # Fallback for direct execution + sys.path.insert(0, current_dir) + from config import ConnectorSpec, OmniTransferConfig + from base import OmniConnectorBase, InMemoryOmniConnector + from factory import OmniConnectorFactory + from utils import load_omni_transfer_config + + +@dataclass +class StageResult: + """Result from a stage processing.""" + stage_name: str + request_id: str + output_data: Any + metadata: Dict[str, Any] + + +class EncodeStage: + """Simple Encode stage for testing.""" + + def __init__(self, stage_id: int, output_connector: Optional[OmniConnectorBase] = None): + self.stage_id = stage_id + self.stage_name = "encode" + self.output_connector = output_connector + + def process(self, request_id: str, input_data: Any) -> StageResult: + """Process input data and optionally send to next stage.""" + print(f"[{self.stage_name}] Processing request {request_id}") + + # Simple encoding simulation + encoded_data = { + "original_text": input_data, + "encoded_tokens": [ord(c) for c in str(input_data)][:10], # Simple tokenization + "embeddings": [0.1 * i for i in range(10)], # Mock embeddings + "stage": self.stage_name + } + + result = StageResult( + stage_name=self.stage_name, + request_id=request_id, + output_data=encoded_data, + metadata={"tokens_count": len(encoded_data["encoded_tokens"])} + ) + + # Send to next stage if connector is available + if self.output_connector: + success = self.output_connector.put( + str(self.stage_id), "1", request_id, encoded_data + ) + print(f"[{self.stage_name}] Sent data to stage 1 via connector: {success}") + + return result + + +class PrefillStage: + """Simple Prefill stage for testing.""" + + def __init__(self, stage_id: int, + input_connector: Optional[OmniConnectorBase] = None, + output_connector: Optional[OmniConnectorBase] = None): + self.stage_id = stage_id + self.stage_name = "prefill" + self.input_connector = input_connector + self.output_connector = output_connector + + def process(self, request_id: str) -> Optional[StageResult]: + """Process data from previous stage.""" + print(f"[{self.stage_name}] Processing request {request_id}") + + # Try to get data from input connector + input_data = None + if self.input_connector: + input_data = self.input_connector.get("0", str(self.stage_id), request_id) + + if input_data is None: + print(f"[{self.stage_name}] No input data received for {request_id}") + return None + + print(f"[{self.stage_name}] Received data from encode") + + # Simple prefill simulation + prefills_data = { + **input_data, + "kv_cache": {"layer_0": [0.2 * i for i in range(20)]}, # Mock KV cache + "stage": self.stage_name + } + + result = StageResult( + stage_name=self.stage_name, + request_id=request_id, + output_data=prefills_data, + metadata={"kv_size": len(prefills_data["kv_cache"])} + ) + + # Send to next stage if connector is available + if self.output_connector: + success = self.output_connector.put( + str(self.stage_id), "2", request_id, prefills_data + ) + print(f"[{self.stage_name}] Sent data to stage 2 via connector: {success}") + + return result + + +class DecodeStage: + """Simple Decode stage for testing (final stage in 3-stage pipeline).""" + + def __init__(self, stage_id: int, + input_connector: Optional[OmniConnectorBase] = None): + self.stage_id = stage_id + self.stage_name = "decode" + self.input_connector = input_connector + + def process(self, request_id: str) -> Optional[StageResult]: + """Process data from previous stage (final output).""" + print(f"[{self.stage_name}] Processing request {request_id}") + + # Try to get data from input connector + input_data = None + if self.input_connector: + input_data = self.input_connector.get("1", str(self.stage_id), request_id) + + if input_data is None: + print(f"[{self.stage_name}] No input data received for {request_id}") + return None + + print(f"[{self.stage_name}] Received data from prefill") + + # Simple decode simulation (final output) + decoded_data = { + **input_data, + "hidden_states": [0.3 * i for i in range(15)], # Mock hidden states + "final_output": f"Decoded text from: {input_data.get('original_text', 'unknown')}", + "stage": self.stage_name + } + + result = StageResult( + stage_name=self.stage_name, + request_id=request_id, + output_data=decoded_data, + metadata={"hidden_size": len(decoded_data["hidden_states"]), "final": True} + ) + + return result + + +class GeneratorStage: + """Simple Generator stage for testing (removed for 3-stage pipeline).""" + pass # Not used in current 3-stage setup + + +def load_config_from_yaml(yaml_path: str) -> OmniTransferConfig: + """Load connector configuration from YAML file.""" + if yaml is None: + # Fallback: create a simple test config + print("Using fallback InMemoryConnector config for testing") + return create_test_config() + + with open(yaml_path, 'r', encoding='utf-8') as f: + config_dict = yaml.safe_load(f) + + try: + config = load_omni_transfer_config(config_dict=config_dict) + print(f"Successfully loaded config from YAML with {len(config.connectors)} connectors") + return config + except Exception as e: + print(f"Failed to load YAML config: {e}, using fallback config") + return create_test_config() + + +def create_test_config() -> OmniTransferConfig: + """Create a test configuration with InMemoryConnectors.""" + from .config import ConnectorSpec + + connectors = {} + # Create connectors for EPD (3-stage) pipeline: 0->1->2 + pipeline_stages = [("0", "1"), ("1", "2")] + + for from_stage, to_stage in pipeline_stages: + connectors[(from_stage, to_stage)] = ConnectorSpec( + name="InMemoryConnector", + extra={} + ) + + return OmniTransferConfig(connectors=connectors) + + +def create_stages_with_connectors(config: OmniTransferConfig) -> Dict[str, Any]: + """Create EPD (3-stage) pipeline with appropriate connectors.""" + + # Create connector instances based on config + connectors = {} + for edge_key, connector_spec in config.connectors.items(): + from_stage, to_stage = edge_key + connectors[edge_key] = OmniConnectorFactory.create_connector(connector_spec) + + # Create stages with connectors (3-stage pipeline: 0->1->2) + stages = {} + + # Encode stage (stage 0) + encode_output_connector = connectors.get(("0", "1")) # to prefill + stages["encode"] = EncodeStage(0, encode_output_connector) + + # Prefill stage (stage 1) + prefill_input_connector = connectors.get(("0", "1")) # from encode + prefill_output_connector = connectors.get(("1", "2")) # to decode + stages["prefill"] = PrefillStage(1, prefill_input_connector, prefill_output_connector) + + # Decode stage (stage 2) - final stage + decode_input_connector = connectors.get(("1", "2")) # from prefill + stages["decode"] = DecodeStage(2, decode_input_connector) + + return stages + + +def test_pipeline_with_config(yaml_path: str): + """Test the EPDG pipeline with connector configuration.""" + + print("🔧 Loading configuration from YAML...") + try: + config = load_config_from_yaml(yaml_path) + print(f"✓ Loaded config with {len(config.connectors)} connector configurations") + except Exception as e: + print(f"✗ Failed to load config: {e}") + return False + + print("\n🏗️ Creating EPDG stages with connectors...") + try: + stages = create_stages_with_connectors(config) + print("✓ Created all stages with connectors") + except Exception as e: + print(f"✗ Failed to create stages: {e}") + return False + + print("\n🚀 Testing data flow through pipeline...") + + # Test request + request_id = f"test_req_{uuid.uuid4()}" + input_text = "Hello, EPDG pipeline!" + + try: + # Stage 1: Encode + print(f"\n--- Stage 1: Encode ---") + encode_result = stages["encode"].process(request_id, input_text) + print(f"✓ Encode completed: {encode_result.metadata}") + + # Stage 2: Prefill + print(f"\n--- Stage 2: Prefill ---") + prefill_result = stages["prefill"].process(request_id) + if prefill_result: + print(f"✓ Prefill completed: {prefill_result.metadata}") + else: + print("✗ Prefill failed - no input data") + return False + + # Stage 3: Decode + print(f"\n--- Stage 3: Decode ---") + decode_result = stages["decode"].process(request_id) + if decode_result: + print(f"✓ Decode completed: {decode_result.metadata}") + else: + print("✗ Decode failed - no input data") + return False + + # Decode stage produces final output in 3-stage pipeline + + except Exception as e: + print(f"✗ Pipeline test failed: {e}") + import traceback + traceback.print_exc() + return False + + print("\n🎉 Pipeline test completed successfully!") + return True + + +def main(): + """Main test function.""" + # Path to the YAML config file + yaml_path = "/workspace/omni/vllm-omni/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml" + + if not os.path.exists(yaml_path): + print(f"✗ YAML config file not found: {yaml_path}") + return False + + print("Testing EPDG Pipeline with Connectors") + print("=" * 50) + print(f"Config file: {yaml_path}") + + success = test_pipeline_with_config(yaml_path) + + if success: + print("\n✅ All tests passed!") + else: + print("\n❌ Tests failed!") + + return success + + +if __name__ == "__main__": + success = main() + sys.exit(0 if success else 1) diff --git a/vllm_omni/distributed/connectors/factory.py b/vllm_omni/distributed/connectors/factory.py new file mode 100644 index 000000000..d4df1b87f --- /dev/null +++ b/vllm_omni/distributed/connectors/factory.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 + +import logging +from collections.abc import Callable +from typing import Any + +try: + from .base import OmniConnectorBase + from .config import ConnectorSpec +except ImportError: + # Fallback for direct execution + import sys + import os + sys.path.insert(0, os.path.dirname(os.path.dirname(__file__))) + from connectors.base import OmniConnectorBase + from connectors.config import ConnectorSpec + +logger = logging.getLogger(__name__) + + +class OmniConnectorFactory: + """Factory for creating OmniConnectors.""" + + _registry: dict[str, Callable[[dict[str, Any]], OmniConnectorBase]] = {} + + @classmethod + def register_connector( + cls, + name: str, + constructor: Callable[[dict[str, Any]], OmniConnectorBase] + ) -> None: + """Register a connector constructor.""" + if name in cls._registry: + raise ValueError(f"Connector '{name}' is already registered.") + cls._registry[name] = constructor + logger.debug(f"Registered connector: {name}") + + @classmethod + def create_connector(cls, spec: ConnectorSpec) -> OmniConnectorBase: + """Create a connector from specification.""" + if spec.name not in cls._registry: + raise ValueError(f"Unknown connector: {spec.name}. Available: {list(cls._registry.keys())}") + + constructor = cls._registry[spec.name] + try: + connector = constructor(spec.extra) + logger.info(f"Created connector: {spec.name}") + return connector + except Exception as e: + logger.error(f"Failed to create connector {spec.name}: {e}") + raise + + @classmethod + def list_registered_connectors(cls) -> list[str]: + """List all registered connector names.""" + return list(cls._registry.keys()) + + +# Register built-in connectors with lazy imports +def _create_mooncake_connector(config: dict[str, Any]) -> OmniConnectorBase: + try: + from .mooncake_connector import MooncakeConnector + except ImportError: + # Fallback import + import sys + import os + sys.path.insert(0, os.path.dirname(os.path.dirname(__file__))) + from connectors.mooncake_connector import MooncakeConnector + return MooncakeConnector(config) + + +def _create_in_memory_connector(config: dict[str, Any]) -> OmniConnectorBase: + try: + from .base import InMemoryOmniConnector + except ImportError: + # Fallback import + import sys + import os + sys.path.insert(0, os.path.dirname(os.path.dirname(__file__))) + from connectors.base import InMemoryOmniConnector + return InMemoryOmniConnector() + + +# Register connectors +OmniConnectorFactory.register_connector("MooncakeConnector", _create_mooncake_connector) +OmniConnectorFactory.register_connector("InMemoryConnector", _create_in_memory_connector) diff --git a/vllm_omni/distributed/connectors/mooncake_connector.py b/vllm_omni/distributed/connectors/mooncake_connector.py new file mode 100644 index 000000000..9370af750 --- /dev/null +++ b/vllm_omni/distributed/connectors/mooncake_connector.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 + +import io +import logging +import time +from typing import Any, Optional + +import torch + +from .base import OmniConnectorBase + +logger = logging.getLogger(__name__) + +try: + from mooncake.store import MooncakeDistributedStore, ReplicateConfig +except ImportError: + try: + from mooncake import MooncakeDistributedStore, ReplicateConfig + except ImportError: + logger.warning("Mooncake not available, MooncakeOmniConnector will not work") + MooncakeDistributedStore = None + ReplicateConfig = None + + +class MooncakeConnector(OmniConnectorBase): + """Mooncake-based distributed connector for OmniConnector.""" + + def __init__(self, config: dict[str, Any]): + if MooncakeDistributedStore is None: + raise ImportError("Mooncake not available") + + self.config = config + self.host = config.get("host", "127.0.0.1") + self.metadata = config.get("metadata_server", "http://127.0.0.1:8080/metadata") + self.master = config.get("master", "127.0.0.1:50051") + self.segment = config.get("segment", 512 * 1024 * 1024) # 512MB + self.localbuf = config.get("localbuf", 64 * 1024 * 1024) # 64MB + self.proto = config.get("proto", "tcp") + self.rdma = config.get("rdma", "") + + self.store: Optional[MooncakeDistributedStore] = None + self.pin: Optional[ReplicateConfig] = None + + self._metrics = { + "puts": 0, + "gets": 0, + "bytes_transferred": 0, + "errors": 0, + "timeouts": 0, + } + + self._init_store() + + def _make_key(self, rid: str, from_stage: str, to_stage: str) -> str: + """Generate store key for request between stages.""" + return f"{rid}/{from_stage}_to_{to_stage}" + + def _init_store(self): + """Initialize Mooncake store.""" + try: + self.store = MooncakeDistributedStore() + rc = self.store.setup( + self.host, self.metadata, self.segment, + self.localbuf, self.proto, self.rdma, self.master + ) + if rc != 0: + raise RuntimeError(f"Mooncake setup failed: {rc}") + + self.pin = ReplicateConfig() + self.pin.with_soft_pin = True + logger.info("MooncakeConnector initialized successfully") + except Exception as e: + logger.error(f"Failed to initialize Mooncake store: {e}") + raise + + # Use base class serialization methods for consistency + + def put(self, from_stage: str, to_stage: str, request_id: str, data: Any) -> tuple[bool, int]: + if not self.store: + logger.error("Store not initialized") + return False, 0 + + try: + serialized_data = self.serialize_obj(data) + key = self._make_key(request_id, from_stage, to_stage) + self.store.put(key, serialized_data, self.pin) + + self._metrics["puts"] += 1 + self._metrics["bytes_transferred"] += len(serialized_data) + + logger.debug(f"MooncakeConnector: stored {key}, {len(serialized_data)} bytes") + print(f"========= mooncake put {from_stage}-{to_stage}-{request_id}, with size {len(serialized_data)}") + return True, len(serialized_data) + + except Exception as e: + self._metrics["errors"] += 1 + logger.error(f"MooncakeConnector put failed: {e}") + return False, 0 + + def get(self, from_stage: str, to_stage: str, request_id: str) -> Optional[Any]: + if not self.store: + logger.error("Store not initialized") + return None + + retries = 20 + sleep_s = 0.05 + key = self._make_key(request_id, from_stage, to_stage) + + for attempt in range(retries): + try: + raw_data = self.store.get(key) + + if raw_data: + data = self.deserialize_obj(raw_data) + self._metrics["gets"] += 1 + print(f"========= mooncake get {from_stage}-{to_stage}-{request_id}, with size {len(raw_data)}") + return data + + except Exception as e: + logger.debug(f"MooncakeConnector get attempt {attempt} failed: {e}") + + if attempt < retries - 1: + time.sleep(sleep_s) + + self._metrics["timeouts"] += 1 + logger.warning(f"MooncakeConnector: timeout waiting for {key}") + return None + + def cleanup(self, request_id: str) -> None: + if not self.store: + return + + # Note: Mooncake doesn't have explicit delete, data will be garbage collected + # We could implement a cleanup mechanism by storing deletion markers + logger.debug(f"MooncakeConnector: cleanup requested for {request_id} (no-op)") + + def health(self) -> dict[str, Any]: + if not self.store: + return {"status": "unhealthy", "error": "Store not initialized"} + + return { + "status": "healthy", + "host": self.host, + "metadata_server": self.metadata, + "master": self.master, + **self._metrics, + } + + def close(self): + """Clean shutdown.""" + if self.store: + try: + self.store.close() + self.store = None + logger.info("MooncakeConnector closed") + except Exception as e: + logger.error(f"Error closing Mooncake store: {e}") diff --git a/vllm_omni/distributed/connectors/utils.py b/vllm_omni/distributed/connectors/utils.py new file mode 100644 index 000000000..56464ff3d --- /dev/null +++ b/vllm_omni/distributed/connectors/utils.py @@ -0,0 +1,314 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Utilities for OmniConnector configuration and validation.""" + +import json +import logging +from pathlib import Path +from typing import Any, Dict, Optional, Tuple, Union + +from vllm_omni.distributed.connectors.config import OmniTransferConfig, ConnectorSpec +from vllm_omni.distributed.connectors.factory import OmniConnectorFactory +from vllm_omni.distributed.connectors.base import OmniConnectorBase + +logger = logging.getLogger(__name__) + + +def initialize_connectors_from_config( + config_path: Optional[Union[str, Path]] = None +) -> tuple[Optional[OmniTransferConfig], dict[tuple[str, str], OmniConnectorBase]]: + """ + Initialize connectors from configuration file. + + Returns: + tuple: (OmniTransferConfig, dict of {(from, to): connector_instance}) + """ + try: + transfer_config = load_omni_transfer_config(config_path) + except Exception as e: + logger.warning(f"Failed to load OmniTransferConfig from {config_path}: {e}") + return None, {} + + if not transfer_config: + logger.info("No OmniTransferConfig provided, using default IPC") + return None, {} + + # 使用统一的连接器创建逻辑 + connectors = create_connectors_from_config(transfer_config.connectors) + return transfer_config, connectors + + +def create_connectors_from_config( + connectors_config: dict[tuple[str, str], ConnectorSpec] +) -> dict[tuple[str, str], OmniConnectorBase]: + """通用连接器创建逻辑，从配置字典创建连接器实例。""" + connectors = {} + try: + for edge_key, connector_spec in connectors_config.items(): + connector = OmniConnectorFactory.create_connector(connector_spec) + connectors[edge_key] = connector + logger.info(f"Created connector for {edge_key[0]} -> {edge_key[1]}: {type(connector).__name__}") + except Exception as e: + logger.error(f"Failed to initialize connectors: {e}") + + return connectors + + +def get_connectors_config_for_stage( + transfer_config: Optional[OmniTransferConfig], + stage_id: Union[str, int] +) -> dict[str, Any]: + """ + Extract connector configurations relevant for a specific stage worker. + + Returns a dict compatible with worker initialization: + { + "from_stage_X": { + "spec": { + "name": "ConnectorName", + "extra": {...} + } + }, + ... + } + """ + if not transfer_config: + return {} + + stage_connectors_config = {} + target_stage = str(stage_id) + + # Iterate through all configured edges + for (from_stage, to_stage), spec in transfer_config.connectors.items(): + # We only care about incoming edges for the worker process + # (Worker needs to create connectors to receive data) + if to_stage == target_stage: + stage_connectors_config[f"from_stage_{from_stage}"] = { + "spec": { + "name": spec.name, + "extra": spec.extra + } + } + + return stage_connectors_config + + +def load_omni_transfer_config( + config_path: Optional[Union[str, Path]] = None, + config_dict: Optional[dict[str, Any]] = None +) -> Optional[OmniTransferConfig]: + """Load OmniTransferConfig from file or dict.""" + if config_path is None and config_dict is None: + return None + + if config_path is not None: + config_path = Path(config_path) + if not config_path.exists(): + raise FileNotFoundError(f"Config file not found: {config_path}") + + with open(config_path, 'r', encoding='utf-8') as f: + if config_path.suffix.lower() == '.json': + config_dict = json.load(f) + elif config_path.suffix.lower() in ['.yaml', '.yml']: + try: + import yaml + config_dict = yaml.safe_load(f) + except ImportError: + raise ImportError("PyYAML required for YAML config files") + else: + raise ValueError(f"Unsupported config file format: {config_path.suffix}") + + if config_dict is None: + return None + + # Parse connectors + connectors = {} + runtime_config = config_dict.get('runtime', {}) + + # Parse global connectors (from runtime.connectors) + global_connectors = runtime_config.get('connectors', {}) + for conn_name, conn_config in global_connectors.items(): + connector = ConnectorSpec( + name=conn_config['name'], + extra=conn_config.get('extra', {}) + ) + # Store globally defined connectors for reference + # These will be resolved to specific edges later + + # Parse stage-level connectors + for stage_config in config_dict.get('stage_args', []): + stage_id = str(stage_config['stage_id']) + + # Input connectors + for input_key, conn_ref in stage_config.get('input_connectors', {}).items(): + if isinstance(conn_ref, str): + # Reference to global connector + if conn_ref in global_connectors: + conn_config = global_connectors[conn_ref] + connector = ConnectorSpec( + name=conn_config['name'], + extra=conn_config.get('extra', {}) + ) + else: + raise ValueError(f"Undefined connector reference: {conn_ref}") + else: + # Inline connector definition + connector = ConnectorSpec( + name=conn_ref['name'], + extra=conn_ref.get('extra', {}) + ) + + # Parse from_stage from key (e.g., "from_stage_0" -> "0") + from_stage = input_key.replace("from_stage_", "") + edge_key = (from_stage, stage_id) + connectors[edge_key] = connector + + # Output connectors + for output_key, conn_ref in stage_config.get('output_connectors', {}).items(): + if isinstance(conn_ref, str): + # Reference to global connector + if conn_ref in global_connectors: + conn_config = global_connectors[conn_ref] + connector = ConnectorSpec( + name=conn_config['name'], + extra=conn_config.get('extra', {}) + ) + else: + raise ValueError(f"Undefined connector reference: {conn_ref}") + else: + # Inline connector definition + connector = ConnectorSpec( + name=conn_ref['name'], + extra=conn_ref.get('extra', {}) + ) + + # Parse to_stage from key (e.g., "to_stage_1" -> "1") + to_stage = output_key.replace("to_stage_", "") + edge_key = (stage_id, to_stage) + connectors[edge_key] = connector + + config = OmniTransferConfig(connectors=connectors) + + logger.info(f"Loaded OmniTransferConfig with {len(connectors)} connector configurations") + return config + + +def create_simple_config( + connector_name: str = "InMemoryConnector", + connector_config: Optional[dict[str, Any]] = None +) -> OmniTransferConfig: + """Create a simple OmniTransferConfig for testing.""" + if connector_config is None: + connector_config = {} + + # Create connectors for typical EPDG pipeline + connectors = {} + pipeline_edges = [ + ("0", "1"), # thinker -> talker + ("1", "2"), # talker -> code2wav + ] + + for from_stage, to_stage in pipeline_edges: + connectors[(from_stage, to_stage)] = ConnectorSpec( + name=connector_name, + extra=connector_config + ) + + return OmniTransferConfig(connectors=connectors) + + +def create_mooncake_config( + host: str = "127.0.0.1", + metadata_server: str = "http://127.0.0.1:8080/metadata", + master: str = "127.0.0.1:50051", + segment: int = 512 * 1024 * 1024, + localbuf: int = 64 * 1024 * 1024, +) -> OmniTransferConfig: + """Create OmniTransferConfig for Mooncake-based setup.""" + connector_config = { + "host": host, + "metadata_server": metadata_server, + "master": master, + "segment": segment, + "localbuf": localbuf, + } + + return create_simple_config("MooncakeConnector", connector_config) + + +# High-level management functions (moved from manager.py) + +def initialize_orchestrator_connectors( + config_path: Optional[str], +) -> Tuple[Optional[OmniTransferConfig], Dict[tuple[str, str], OmniConnectorBase]]: + """Initialize connectors shared at orchestrator level.""" + try: + transfer_config, connectors = initialize_connectors_from_config(config_path) + return transfer_config, connectors + except Exception as exc: # pragma: no cover - defensive logging + logger.error("Unexpected error initializing connectors: %s", exc) + return None, {} + + +def get_stage_connector_config( + transfer_config: Optional[OmniTransferConfig], + stage_id: int, +) -> dict[str, Any]: + """Return the serialized connector config payload for a specific stage.""" + if transfer_config is None: + return {} + + try: + return get_connectors_config_for_stage(transfer_config, stage_id) + except Exception as exc: # pragma: no cover - defensive logging + logger.warning( + "Failed to build connector config for stage %s: %s. Using IPC fallback.", + stage_id, + exc, + ) + return {} + + +def build_stage_connectors( + stage_id: int, + connectors_config: dict[str, Any], +) -> Optional[dict[tuple[str, str], Any]]: + """Instantiate OmniConnectors for a stage based on config.""" + if not connectors_config: + return {} + + logger.info( + "[Stage-%s] Initializing OmniConnectors with config keys: %s", + stage_id, + list(connectors_config.keys()), + ) + + from .factory import OmniConnectorFactory + from .config import ConnectorSpec + + connectors: dict[tuple[str, str], Any] = {} + try: + # 将字典格式的配置转换为ConnectorSpec对象 + stage_connector_specs = {} + for input_key, config in connectors_config.items(): + if not input_key.startswith("from_stage_"): + continue + + from_stage = input_key.replace("from_stage_", "") + spec_dict = config.get("spec", {}) + if not spec_dict: + continue + + connector_spec = ConnectorSpec( + name=spec_dict.get("name", "InMemoryConnector"), + extra=spec_dict.get("extra", {}), + ) + stage_connector_specs[(str(from_stage), str(stage_id))] = connector_spec + + # 使用统一的连接器创建逻辑 + connectors = create_connectors_from_config(stage_connector_specs) + except Exception as exc: # pragma: no cover - defensive logging + logger.exception("[Stage-%s] Failed to initialize connectors: %s", stage_id, exc) + return None + + return connectors diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index 94d93ef76..51fb18ebb 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/vllm_omni/entrypoints/omni_llm.py @@ -1,6 +1,7 @@ import multiprocessing as mp import os import time +import uuid from collections.abc import Sequence from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Any, Optional, Union @@ -20,6 +21,11 @@ from vllm.utils import Counter from vllm.v1.engine.llm_engine import LLMEngine # Internal imports (our code) +from vllm_omni.distributed.connectors.adapter import try_send_via_connector +from vllm_omni.distributed.connectors import ( + initialize_orchestrator_connectors, + get_stage_connector_config, +) from vllm_omni.engine.arg_utils import OmniEngineArgs from vllm_omni.engine.output_processor import MultimodalOutputProcessor from vllm_omni.engine.processor import OmniProcessor @@ -31,9 +37,14 @@ from vllm_omni.entrypoints.log_utils import ( ) from vllm_omni.entrypoints.omni_stage import OmniStage from vllm_omni.entrypoints.stage_utils import encode_for_ipc as _encode +from vllm_omni.entrypoints.stage_utils import _to_dict from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load from vllm_omni.entrypoints.stage_utils import serialize_obj as _set -from vllm_omni.entrypoints.utils import load_stage_configs_from_model, load_stage_configs_from_yaml +from vllm_omni.entrypoints.utils import ( + load_stage_configs_from_model, + load_stage_configs_from_yaml, + resolve_model_config_path, +) from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) @@ -84,14 +95,23 @@ class OmniLLM: init_timeout: int = 300, **kwargs: Any, ): + self.worker_backend = kwargs.get("worker_backend", "process") + self.ray_address = kwargs.get("ray_address", None) self.batch_timeout = batch_timeout self._enable_stats: bool = bool(log_stats) + self._ray_pg = None + # Do NOT call super().__init__ to avoid creating OmniStageLLM instances in parent. if stage_configs_path is None: - self.stage_configs = load_stage_configs_from_model(model) + self.config_path = resolve_model_config_path(model) + self.stage_configs = load_stage_configs_from_yaml(self.config_path) else: + self.config_path = stage_configs_path self.stage_configs = load_stage_configs_from_yaml(stage_configs_path) + # Initialize connectors + self.omni_transfer_config,self.connectors = initialize_orchestrator_connectors(self.config_path) + # Optional file handler for orchestrator self._log_file = log_file @@ -124,36 +144,140 @@ class OmniLLM: self.stage_list = [st for _, st in results] logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list)) - self._ctx = mp.get_context("spawn") + if self.worker_backend == "ray": + try: + import ray + from ray.util.queue import Queue as RayQueue + except ImportError: + raise ImportError("ray is required for worker_backend='ray'") + self._queue_cls = lambda: RayQueue(maxsize=0) + else: + self._ctx = mp.get_context("spawn") + self._queue_cls = lambda: self._ctx.Queue(maxsize=0) + self._stage_in_queues: list[mp.Queue] = [] self._stage_out_queues: list[mp.Queue] = [] self._init_sleep_seconds = max(0, int(init_sleep_seconds)) self._shm_threshold_bytes = max(0, int(shm_threshold_bytes)) - self._start_stage_processes(model) + self._start_stages(model) # Wait for all stages to report readiness before seeding self._stages_ready: set[int] = set() self._wait_for_stages_ready(timeout=init_timeout) - def _start_stage_processes(self, model: str) -> None: + def _start_stages(self, model: str) -> None: + if self.worker_backend == "ray": + self._start_stage_actors_ray(model) + else: + self._start_stage_processes_mp(model) + + def _start_stage_processes_mp(self, model: str) -> None: for stage_id, stage in enumerate(self.stage_list): # Use unbounded queues to avoid deadlock when seeding many requests - in_q: mp.Queue = self._ctx.Queue(maxsize=0) - out_q: mp.Queue = self._ctx.Queue(maxsize=0) + in_q: mp.Queue = self._queue_cls() + out_q: mp.Queue = self._queue_cls() self._stage_in_queues.append(in_q) self._stage_out_queues.append(out_q) # Attach queues and start Stage-owned worker process stage.attach_queues(in_q, out_q) + + # Build connectors config for this stage + stage_connectors_config = get_stage_connector_config( + self.omni_transfer_config, + stage_id, + ) + stage.init_stage_worker( model, log_file=self._log_file, shm_threshold_bytes=self._shm_threshold_bytes, ctx=self._ctx, batch_timeout=self.batch_timeout, + connectors_config=stage_connectors_config, ) logger.debug("[Orchestrator] Stage-%s process started", stage_id) time.sleep(self._init_sleep_seconds) + def _start_stage_actors_ray(self, model: str) -> None: + import ray + import os + from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + + # Initialize Ray Cluster + if not ray.is_initialized(): + # Create a dummy parallel config just to reuse vLLM's init logic if needed + # or just call ray.init() + # Pass current PYTHONPATH to workers to ensure they can find vllm_omni + runtime_env = { + "env_vars": {"PYTHONPATH": os.environ.get("PYTHONPATH", "")} + } + ray.init(address=self.ray_address, ignore_reinit_error=True, runtime_env=runtime_env) + + # Create Placement Group + # We allocate 1 bundle per stage. Assume 1 GPU per stage for now. + # Must include CPU=1 because default Actor request includes 1 CPU. + bundles = [{"GPU": 1.0, "CPU": 1.0} for _ in self.stage_list] + self._ray_pg = ray.util.placement_group(bundles, strategy="SPREAD") + ray.get(self._ray_pg.ready()) + logger.info("[Orchestrator] Ray Placement Group created") + + # Define Actor Wrapper + @ray.remote(num_gpus=1) + class OmniStageRayWorker: + def run(self, func, *args, **kwargs): + return func(*args, **kwargs) + + from vllm_omni.entrypoints.omni_stage import _stage_worker + + for stage_id, stage in enumerate(self.stage_list): + in_q = self._queue_cls() + out_q = self._queue_cls() + self._stage_in_queues.append(in_q) + self._stage_out_queues.append(out_q) + stage.attach_queues(in_q, out_q) + + stage_connectors_config = get_stage_connector_config( + self.omni_transfer_config, + stage_id, + ) + + # Prepare payload + # Reuse logic from OmniStage.init_stage_worker but do not start process + engine_args = _to_dict(stage.engine_args) + runtime_cfg = _to_dict(getattr(stage.stage_config, "runtime", {})) + stage_payload = { + "stage_id": stage.stage_id, + "engine_args": engine_args, + "runtime": runtime_cfg, + "shm_threshold_bytes": sys.maxsize, # Disable SHM for Ray mode by setting huge threshold + "connectors_config": stage_connectors_config or {}, + } + + # Start Actor + # Inject PYTHONPATH again for the specific actor to be safe + worker_actor = OmniStageRayWorker.options( + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=self._ray_pg, + placement_group_bundle_index=stage_id + ), + runtime_env={"env_vars": {"PYTHONPATH": os.environ.get("PYTHONPATH", "")}} + ).remote() + + # Start worker loop asynchronously + worker_actor.run.remote( + _stage_worker, + model=model, + stage_payload=stage_payload, + in_q=in_q, + out_q=out_q, + log_file=self._log_file, + batch_timeout=self.batch_timeout, + ) + + stage._ray_actor = worker_actor + logger.debug("[Orchestrator] Stage-%s Ray Actor started", stage_id) + time.sleep(self._init_sleep_seconds) + def close(self) -> None: """Close all stage processes and clean up resources. @@ -171,10 +295,21 @@ class OmniLLM: ) for stage in self.stage_list: try: - stage.stop_stage_worker() + if self.worker_backend == "ray" and hasattr(stage, "_ray_actor"): + import ray + # Ray kill + ray.kill(stage._ray_actor) + else: + stage.stop_stage_worker() except Exception as e: logger.warning("[Orchestrator] Failed to stop stage worker: %s", e) + if self.worker_backend == "ray" and self._ray_pg: + try: + ray.util.remove_placement_group(self._ray_pg) + except Exception as e: + logger.warning("[Orchestrator] Failed to remove placement group: %s", e) + def __del__(self) -> None: # best-effort try: self.close() @@ -237,11 +372,14 @@ class OmniLLM: # Orchestrator keeps stage objects for input derivation num_stages = len(self.stage_list) - # Map from request_id to original prompt - request_id_to_prompt: dict[int, PromptType] = {i: p for i, p in enumerate(request_prompts)} + # Generate globally unique request IDs and map them to original prompts + request_ids: list[str] = [f"{i}_{uuid.uuid4()}" for i in range(len(request_prompts))] + request_id_to_prompt: dict[str, PromptType] = { + rid: p for rid, p in zip(request_ids, request_prompts) + } # Track per-request start time for end-to-end timing - _req_start_ts: dict[int, float] = {} + _req_start_ts: dict[str, float] = {} _wall_start_ts: float = time.time() # Determine the final stage for E2E stats (highest stage_id with final_output=True; fallback to last stage) @@ -387,51 +525,90 @@ class OmniLLM: ) continue sp_next: SamplingParams = sampling_params_list[next_stage_id] # type: ignore[index] - try: - # Measure transfer size and time (encode + enqueue) - size_bytes = 0 + + # Check if we have a connector for this edge + connector_key = (str(stage_id), str(next_stage_id)) + connector = self.connectors.get(connector_key) + print(f"connectors: {self.connectors}") + + sent_via_connector = False + if connector: + sent_via_connector = try_send_via_connector( + connector=connector, + stage_id=stage_id, + next_stage_id=next_stage_id, + req_id=req_id, + next_inputs=next_inputs, + sampling_params=sp_next, + original_prompt=request_id_to_prompt[req_id], + next_stage_queue_submit_fn=self.stage_list[next_stage_id].submit, + metrics=metrics + ) + + if not sent_via_connector: + # Use original queue mechanism try: - size_bytes = len(_set(next_inputs)) - except Exception: - size_bytes = 0 - t0 = time.time() - ipc_payload = _encode( - next_inputs, - getattr(self, "_shm_threshold_bytes", 65536), - obj_key="engine_inputs", - shm_key="engine_inputs_shm", - ) - ipc_payload.update( - { - "request_id": req_id, - "sampling_params": sp_next, - "sent_ts": time.time(), - } - ) - self.stage_list[next_stage_id].submit(ipc_payload) - t1 = time.time() - tx_ms = (t1 - t0) * 1000.0 - metrics.on_forward( - stage_id, - next_stage_id, - req_id, - int(size_bytes), - float(tx_ms), - bool("engine_inputs_shm" in ipc_payload), - ) - except Exception as e: - logger.warning( - "[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload", - req_id, - e, - ) - self.stage_list[next_stage_id].submit( - { - "request_id": req_id, - "engine_inputs": next_inputs, - "sampling_params": sp_next, - } - ) + if self.worker_backend == "ray": + # Ray mode: Skip SHM fallback, just send object via Queue (Ray handles serialization) + # This prevents SHM access errors across nodes + ipc_payload = { + "request_id": req_id, + "engine_inputs": next_inputs, + "sampling_params": sp_next, + "sent_ts": time.time(), + } + size_bytes = 0 # Approximation or calculate size if needed + try: + # Approximate size for stats + size_bytes = len(_set(next_inputs)) + except: + pass + t0 = time.time() + # Direct submit to Ray Queue + else: + size_bytes = 0 + try: + size_bytes = len(_set(next_inputs)) + except Exception: + size_bytes = 0 + t0 = time.time() + ipc_payload = _encode( + next_inputs, + getattr(self, "_shm_threshold_bytes", 65536), + obj_key="engine_inputs", + shm_key="engine_inputs_shm", + ) + ipc_payload.update( + { + "request_id": req_id, + "sampling_params": sp_next, + "sent_ts": time.time(), + } + ) + self.stage_list[next_stage_id].submit(ipc_payload) + t1 = time.time() + tx_ms = (t1 - t0) * 1000.0 + metrics.on_forward( + stage_id, + next_stage_id, + req_id, + int(size_bytes), + float(tx_ms), + bool("engine_inputs_shm" in ipc_payload), + ) + except Exception as e: + logger.warning( + "[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload", + req_id, + e, + ) + self.stage_list[next_stage_id].submit( + { + "request_id": req_id, + "engine_inputs": next_inputs, + "sampling_params": sp_next, + } + ) logger.debug( "[Orchestrator] Forwarded request %s to stage-%s", req_id, diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 5150ea27e..2a81a3500 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -163,6 +163,7 @@ class OmniStage: shm_threshold_bytes: int = 65536, ctx: Optional[mp.context.BaseContext] = None, batch_timeout: int = 10, + connectors_config: Optional[dict] = None, # New parameter ) -> None: """Initialize and start the stage worker process. @@ -192,6 +193,7 @@ class OmniStage: "engine_args": engine_args, "runtime": runtime_cfg, "shm_threshold_bytes": self._shm_threshold_bytes, + "connectors_config": connectors_config or {}, } if is_async: self._proc = ctx.Process( @@ -340,6 +342,8 @@ def _stage_worker( log_stage_batch_stats, log_stage_running_avg, ) + from vllm_omni.distributed.connectors.adapter import try_recv_via_connector + from vllm_omni.distributed.connectors.utils import build_stage_connectors from vllm_omni.entrypoints.omni_llm import OmniStageLLM # no inline JSONL/serialization imports; logging handled by utilities @@ -347,6 +351,7 @@ def _stage_worker( engine_args = stage_payload.get("engine_args", {}) runtime_cfg = stage_payload.get("runtime", {}) shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536)) + connectors_config = stage_payload.get("connectors_config", {}) # Per-stage logger: clear inherited handlers to avoid broken parent streams try: @@ -440,6 +445,18 @@ def _stage_worker( ) stage_engine = OmniStageLLM(model=model, **engine_args) _logging.getLogger(__name__).debug("[Stage-%s] Engine initialized", stage_id) + + # Initialize OmniConnectors if configured + connectors = {} + if connectors_config: + built_connectors = build_stage_connectors( + stage_id=stage_id, + connectors_config=connectors_config, + ) + if built_connectors is None: + return + connectors = built_connectors + # Signal readiness to orchestrator try: out_q.put({"type": "stage_ready", "stage_id": stage_id}) @@ -496,9 +513,21 @@ def _stage_worker( _in_flight_ms_by_rid[rid] = 0.0 except Exception: _in_flight_ms_by_rid[rid] = 0.0 - ein, _rx_metrics = maybe_load_from_ipc_with_metrics(t, obj_key="engine_inputs", shm_key="engine_inputs_shm") - _rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0)) - _rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0)) + + # Resolve input data (via Connector or IPC) + ein, _rx_metrics = try_recv_via_connector( + task=t, + connectors=connectors, + stage_id=stage_id + ) + + if ein is None or _rx_metrics is None: + ein, _rx_metrics = maybe_load_from_ipc_with_metrics(task, obj_key="engine_inputs", shm_key="engine_inputs_shm") + + if _rx_metrics: + _rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0)) + _rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0)) + batch_request_ids.append(rid) if isinstance(ein, list): batch_engine_inputs.extend(ein) diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 57a6b7554..3c47761b0 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -10,6 +10,41 @@ from vllm_omni.utils import detect_device_type PROJECT_ROOT = Path(__file__).parent.parent.parent +def resolve_model_config_path(model: str) -> str: + """Resolve the stage config file path from the model name. + + Resolves stage configuration path based on the model type and device type. + First tries to find a device-specific YAML file from stage_configs/{device_type}/ + directory. If not found, falls back to the default config file. + + Args: + model: Model name or path (used to determine model_type) + + Returns: + String path to the stage configuration file + + Raises: + FileNotFoundError: If no stage config file exists for the model type + """ + hf_config = get_config(model, trust_remote_code=True) + model_type = hf_config.model_type + device_type = detect_device_type() + + # Try device-specific config first + if device_type != "cuda": + device_config_file = f"vllm_omni/model_executor/stage_configs/{device_type}/{model_type}.yaml" + device_config_path = PROJECT_ROOT / device_config_file + if os.path.exists(device_config_path): + return str(device_config_path) + + # Fall back to default config + stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml" + stage_config_path = PROJECT_ROOT / stage_config_file + if not os.path.exists(stage_config_path): + raise FileNotFoundError(f"Stage config file {stage_config_path} not found") + return str(stage_config_path) + + def load_stage_configs_from_model(model: str) -> list: """Load stage configurations from model's default config file. @@ -26,24 +61,8 @@ def load_stage_configs_from_model(model: str) -> list: Raises: FileNotFoundError: If no stage config file exists for the model type """ - hf_config = get_config(model, trust_remote_code=True) - model_type = hf_config.model_type - device_type = detect_device_type() - - # Try device-specific config first - if device_type != "cuda": - device_config_file = f"vllm_omni/model_executor/stage_configs/{device_type}/{model_type}.yaml" - device_config_path = PROJECT_ROOT / device_config_file - if os.path.exists(device_config_path): - stage_configs = load_stage_configs_from_yaml(config_path=str(device_config_path)) - return stage_configs - - # Fall back to default config - stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml" - stage_config_path = PROJECT_ROOT / stage_config_file - if not os.path.exists(stage_config_path): - raise FileNotFoundError(f"Stage config file {stage_config_path} not found") - stage_configs = load_stage_configs_from_yaml(config_path=str(stage_config_path)) + stage_config_path = resolve_model_config_path(model) + stage_configs = load_stage_configs_from_yaml(config_path=stage_config_path) return stage_configs diff --git a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml index c378a5b48..7cc6bc433 100644 --- a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml @@ -26,6 +26,9 @@ stage_args: seed: 42 detokenize: True repetition_penalty: 1.1 + # Distributed connector configuration (optional) + output_connectors: + to_stage_1: connector_thinker2talker # thinker -> talker: use local mooncake - stage_id: 1 runtime: process: true @@ -52,17 +55,22 @@ stage_args: detokenize: True repetition_penalty: 1.05 stop_token_ids: [8294] + # Distributed connector configuration (optional) + input_connectors: + from_stage_0: connector_thinker2talker # talker <- thinker: use local mooncake + output_connectors: + to_stage_2: connector_talker2generator # talker -> code2wav: use remote mooncake (example) - stage_id: 2 runtime: process: true - devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU + devices: "2" # Example: use a different GPU than the previous stage; use "0" if single GPU max_batch_size: 1 engine_args: model_stage: code2wav model_arch: Qwen2_5OmniForConditionalGeneration worker_cls: vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler - gpu_memory_utilization: 0.15 + gpu_memory_utilization: 0.3 enforce_eager: true trust_remote_code: true enable_prefix_caching: false @@ -78,6 +86,9 @@ stage_args: seed: 42 detokenize: True repetition_penalty: 1.1 + # Distributed connector configuration (optional) + input_connectors: + from_stage_1: connector_talker2generator # code2wav <- talker: use remote mooncake # Top-level runtime config (concise): default windows and stage edges runtime: @@ -85,6 +96,31 @@ runtime: defaults: window_size: -1 # Simplified: trigger downstream only after full upstream completion max_inflight: 1 # Simplified: process serially within each stage + + # Distributed connectors configuration (optional) + connectors: + # Local mooncake for same-node communication + connector_thinker2talker: + name: MooncakeConnector + extra: + host: "127.0.0.1" + metadata_server: "http://10.90.67.86:8080/metadata" + master: "10.90.67.86:50051" + segment: 512000000 # 512MB + localbuf: 64000000 # 64MB + proto: "tcp" + + # Remote mooncake for cross-node communication (example) + connector_talker2generator: + name: MooncakeConnector + extra: + host: "127.0.0.1" + metadata_server: "http://10.90.67.86:8080/metadata" + master: "10.90.67.86:50051" + segment: 1073741824 # 1GB + localbuf: 134217728 # 128MB + proto: "tcp" + edges: - from: 0 # thinker → talker: trigger only after receiving full input (-1) to: 1