From 13dfb56a6860066d2a6adfc086e80f076f499043 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Mon, 10 Nov 2025 11:58:47 +0800 Subject: [PATCH] Refactor end-to-end inference script and enhance argument parsing. Added support for loading prompts from a .pt file, introduced new command-line arguments for initialization and output handling, and improved error handling for prompt loading. Removed deprecated files related to previous implementations. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 齐保元 --- .../qwen_2_5_omni/end2end.py | 117 +++++- .../qwen_2_5_omni/extract_prompts.py | 48 --- .../hf_qwen25_omni_batch_no_template.py | 279 ------------- .../qwen_2_5_omni/pipelined_end2end.py | 265 ------------- vllm_omni/__init__.py | 4 +- vllm_omni/entrypoints/log_utils.py | 22 ++ vllm_omni/entrypoints/omni_llm.py | 345 +++++++++++++++- vllm_omni/entrypoints/omni_stage.py | 73 +--- vllm_omni/entrypoints/pipeline_utils.py | 37 +- vllm_omni/entrypoints/pipelined_omni_llm.py | 370 ------------------ 10 files changed, 521 insertions(+), 1039 deletions(-) delete mode 100644 examples/offline_inference/qwen_2_5_omni/extract_prompts.py delete mode 100644 examples/offline_inference/qwen_2_5_omni/hf_qwen25_omni_batch_no_template.py delete mode 100644 examples/offline_inference/qwen_2_5_omni/pipelined_end2end.py delete mode 100644 vllm_omni/entrypoints/pipelined_omni_llm.py diff --git a/examples/offline_inference/qwen_2_5_omni/end2end.py b/examples/offline_inference/qwen_2_5_omni/end2end.py index fa5df3e0b..8c5da7c08 100644 --- a/examples/offline_inference/qwen_2_5_omni/end2end.py +++ b/examples/offline_inference/qwen_2_5_omni/end2end.py @@ -7,8 +7,8 @@ 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" @@ -49,7 +49,10 @@ def parse_args(): "--hf-revision", default=None, help="Optional HF revision (branch/tag/commit)." ) parser.add_argument( - "--prompts", required=True, nargs="+", help="Input text prompts." + "--prompts", nargs="+", default=None, help="Input text prompts." + ) + parser.add_argument( + "--pt-prompts", type=str, default=None, help="Path to .pt file containing List[str] prompts (overrides --prompts when prompt_type=text)." ) parser.add_argument( "--voice-type", default="default", help="Voice type, e.g., m02, f030, default." @@ -69,6 +72,12 @@ def parse_args(): "--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") @@ -103,13 +112,40 @@ def parse_args(): parser.add_argument("--use-torchvision", action="store_true") parser.add_argument("--tokenize", action="store_true") parser.add_argument( - "--output-wav", default="output.wav", help="Output wav file path." + "--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)", + ) args = parser.parse_args() return args @@ -117,7 +153,29 @@ def parse_args(): def main(): args = parse_args() model_name = args.model - omni_llm = OmniLLM(model=model_name) + # If pt-prompts provided and prompt_type is text, load and override args.prompts + try: + if getattr(args, "pt_prompts", None) and args.prompt_type == "text": + loaded = torch.load(args.pt_prompts) + if not isinstance(loaded, list) or (len(loaded) > 0 and not isinstance(loaded[0], str)): + raise ValueError("--pt-prompts must point to a .pt containing List[str]") + args.prompts = loaded + print(f"[Info] Loaded {len(args.prompts)} prompts from {args.pt_prompts}") + except Exception as e: + print(f"[Error] Failed to load pt-prompts: {e}") + raise + + if args.prompts is None: + raise ValueError("No prompts provided. Use --prompts ... or --pt-prompts (with --prompt_type text)") + omni_llm = PipelinedOmniLLM( + 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, + ) thinker_sampling_params = SamplingParams( temperature=0.0, # Deterministic - no randomness top_p=1.0, # Disable nucleus sampling @@ -128,13 +186,13 @@ def main(): repetition_penalty=1.1, ) talker_sampling_params = SamplingParams( - temperature=0.0, # Deterministic - no randomness - top_p=1.0, # Disable nucleus sampling - top_k=-1, # Disable top-k sampling + temperature=0.9, + top_p=0.8, + top_k=40, max_tokens=2048, seed=SEED, # Fixed seed for sampling detokenize=True, - repetition_penalty=1.1, + repetition_penalty=1.05, stop_token_ids=[8294], ) code2wav_sampling_params = SamplingParams( @@ -156,24 +214,51 @@ def main(): prompt = [make_omni_prompt(args, prompt) for prompt in args.prompts] omni_outputs = omni_llm.generate(prompt, sampling_params_list) - os.makedirs(args.output_wav, exist_ok=True) + # 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 = output.request_id + request_id = int(output.request_id) text_output = output.outputs[0].text - print(f"Request ID: {request_id}, Text Output: {text_output}") + # 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") + lines.append("transformers_text_output:\n\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 = output.request_id + request_id = int(output.request_id) audio_tensor = output.multimodal_output["audio"] - output_wav = os.path.join( - args.output_wav, f"output_{output.request_id}.wav" - ) + 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}") + # Save aligned text file with empty text (no text output in this stage) + 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\n") + lines.append("transformers_text_output:\n\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}, Saved audio to {output_wav} and text to {out_txt}") if __name__ == "__main__": diff --git a/examples/offline_inference/qwen_2_5_omni/extract_prompts.py b/examples/offline_inference/qwen_2_5_omni/extract_prompts.py deleted file mode 100644 index a1bc7a78f..000000000 --- a/examples/offline_inference/qwen_2_5_omni/extract_prompts.py +++ /dev/null @@ -1,48 +0,0 @@ -#!/usr/bin/env python3 -import argparse -from typing import Optional -import os -import torch - - -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="输入 .lst 文件路径") - parser.add_argument("--output", "-o", required=True, help="输出文件路径") - parser.add_argument("--topk", "-k", type=int, default=100, help="提取前 K 个 prompt（默认 100）") - parser.add_argument("--pt-output", "-p", type=str, default=None, help="可选：输出 .pt 文件路径（默认与 --output 同名 .pt）") - args = parser.parse_args() - - prompts = [] - with open(args.input, "r", 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") - - # 另存为 PyTorch 序列化的 .pt，内容为字符串列表 - pt_path = args.pt_output if args.pt_output else os.path.splitext(args.output)[0] + ".pt" - torch.save(prompts, pt_path) - - -if __name__ == "__main__": - main() - - diff --git a/examples/offline_inference/qwen_2_5_omni/hf_qwen25_omni_batch_no_template.py b/examples/offline_inference/qwen_2_5_omni/hf_qwen25_omni_batch_no_template.py deleted file mode 100644 index 510c60f82..000000000 --- a/examples/offline_inference/qwen_2_5_omni/hf_qwen25_omni_batch_no_template.py +++ /dev/null @@ -1,279 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- - -""" -Batch text generation with Qwen 2.5 Omni using Transformers without chat templates. - -- Processes 10 fixed prompts in batches of 3 -- Does NOT use tokenizer.apply_chat_template (raw prompts only) -- Prints only generated continuations for each prompt -""" - -import os -import argparse -from typing import Iterable, List -import time -import torch - -import random -import numpy as np -import torch -from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor -from tqdm import tqdm -import json -import soundfile as sf - -PROMPTS: List[str] = [ - "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words.", - "Generate a friendly greeting message suitable for a voice assistant. Answer in 15 words.", - "Summarize the benefits of on-device inference for edge deployments. Answer in 15 words.", - "Offer troubleshooting steps for degraded audio quality in a TTS pipeline. Answer in 15 words.", - "List best practices for optimizing latency in real-time voice synthesis. Answer in 15 words.", - "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words.", - "Generate a friendly greeting message suitable for a voice assistant. Answer in 15 words.", - "Summarize the benefits of on-device inference for edge deployments. Answer in 15 words.", - "Offer troubleshooting steps for degraded audio quality in a TTS pipeline. Answer in 15 words.", - "List best practices for optimizing latency in real-time voice synthesis. Answer in 15 words.", -] -def make_text_prompt(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."} - ], - }, - { - "role": "user", - "content": [ - {"type": "text", "text": prompt}, - ], - }] - -def chunked(items: List[str], size: int) -> Iterable[List[str]]: - for i in range(0, len(items), size): - yield items[i : i + size] - - -def main(): - parser = argparse.ArgumentParser( - description="Qwen 2.5 Omni batch generation without chat templates", - ) - parser.add_argument( - "--model-id", - type=str, - default=os.getenv("QWEN_OMNI_MODEL_ID", "Qwen/Qwen2.5-Omni"), - help="HF model repo id (e.g., Qwen/Qwen2.5-Omni)", - ) - parser.add_argument( - "--max-new-tokens", - type=int, - default=2048, - help="Maximum number of new tokens to generate", - ) - parser.add_argument( - "--temperature", type=float, default=0.0, help="Sampling temperature" - ) - parser.add_argument("--top-p", type=float, default=1.0, help="Top-p sampling") - parser.add_argument( - "--repetition-penalty", - type=float, - default=1.1, - help="Repetition penalty (matches pipelined_end2end defaults)", - ) - parser.add_argument( - "--stop-token-ids", - type=str, - default="8294", - help="Comma-separated stop token ids (added to eos_token_id)", - ) - parser.add_argument( - "--seed", - type=int, - default=42, - help="Random seed for reproducibility (sampling)", - ) - parser.add_argument( - "--pt-prompts", - type=str, - default=None, - help="Path to .pt file containing List[str] prompts (overrides built-in PROMPTS)", - ) - parser.add_argument( - "--tokenize", - action="store_true", - help="If set, utils.make_text_prompt will return token ids; keep unset for raw text", - ) - parser.add_argument( - "--output-dir", - type=str, - default="outputs", - help="Output directory to save aligned comparison files (same folder as pipelined_end2end)", - ) - args = parser.parse_args() - - if torch.cuda.is_available(): - torch_dtype = torch.bfloat16 - else: - torch_dtype = torch.float32 - - # Match pipelined_end2end.py deterministic settings - if args.seed is not None: - os.environ["PYTHONHASHSEED"] = str(args.seed) - os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" - try: - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed(args.seed) - torch.cuda.manual_seed_all(args.seed) - torch.backends.cudnn.deterministic = True - torch.backends.cudnn.benchmark = False - except Exception: - pass - - print(f"[Info] Loading model: {args.model_id}") - processor = Qwen2_5OmniProcessor.from_pretrained(args.model_id) - model = Qwen2_5OmniForConditionalGeneration.from_pretrained( - args.model_id, - trust_remote_code=True, - torch_dtype=torch_dtype, - device_map="auto", - attn_implementation="flash_attention_2", - ) - - # Load prompts: prefer .pt if provided - if args.pt_prompts is not None: - loaded = torch.load(args.pt_prompts) - if not isinstance(loaded, list) or (len(loaded) > 0 and not isinstance(loaded[0], str)): - raise ValueError("--pt-prompts must point to a .pt containing List[str]") - prompts: List[str] = loaded - print(f"[Info] Loaded {len(prompts)} prompts from {args.pt_prompts}") - else: - prompts = PROMPTS - - # Ensure utils.make_text_prompt can find expected fields - # utils expects args.model and args.tokenize - setattr(args, "model", args.model_id) - if not hasattr(args, "tokenize"): - setattr(args, "tokenize", False) - - print(f"[Info] Running single-sample generation (batch size fixed to 1) ...\n") - global_index = 0 - total_requests = len(prompts) - sum_request_time_ms = 0.0 - wall_t0 = time.time() - stats_path = os.path.join(args.output_dir, "hf_transformers.stats.jsonl") - - for rid, prompt in tqdm(list(enumerate(prompts)), total=len(prompts), desc="Generating"): - _t0 = time.time() - # Build single-sample inputs - tp = make_text_prompt(prompt) - tp = processor.apply_chat_template(tp, add_generation_prompt=True, tokenize=False) - inputs = processor(text=tp, return_tensors="pt", padding=True, use_audio_in_video=True) - inputs = inputs.to(model.device).to(model.dtype) - - with torch.no_grad(): - text_ids, audio = model.generate( - **inputs, - use_audio_in_video=True, - ) - - # Decode - try: - texts = processor.batch_decode(text_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False) # type: ignore[arg-type] - except Exception: - try: - texts = processor.tokenizer.batch_decode(text_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False) # type: ignore[attr-defined] - except Exception: - texts = [str(x) for x in text_ids] - - text_str = texts[0] if len(texts) > 0 else "" - os.makedirs(args.output_dir, exist_ok=True) - out_txt = os.path.join(args.output_dir, f"{rid:05d}.txt") - - content = None - if os.path.exists(out_txt): - try: - with open(out_txt, "r", encoding="utf-8") as f: - content = f.read() - except Exception: - content = None - if content and "transformers_text_output:" in content: - try: - parts = content.split("transformers_text_output:\n", 1) - new_content = parts[0] + "transformers_text_output:\n" + text_str.strip() + "\n" - with open(out_txt, "w", encoding="utf-8") as f: - f.write(new_content) - except Exception: - pass - else: - lines = [ - "Prompt:\n", - str(prompt) + "\n", - "vllm_text_output:\n\n", - "transformers_text_output:\n", - text_str.strip() + "\n", - ] - try: - with open(out_txt, "w", encoding="utf-8") as f: - f.writelines(lines) - except Exception: - pass - - # Save wav per request - try: - if audio is not None: - if isinstance(audio, (list, tuple)): - wav = audio[0] - else: - # Tensor or array: take first sample if batched, else as-is - if hasattr(audio, "ndim") and hasattr(audio, "shape") and len(audio.shape) > 1: - wav = audio[0] - else: - wav = audio - wav_path = os.path.join(args.output_dir, f"output_{rid}.wav") - sf.write(wav_path, torch.as_tensor(wav).reshape(-1).detach().cpu().numpy(), samplerate=24000) - except Exception: - pass - - # Stats per request - req_ms = (time.time() - _t0) * 1000.0 - sum_request_time_ms += req_ms - try: - rec = { - "type": "transformers_request_stats", - "request_id": rid, - "batch_size": 1, - "per_request_time_ms": req_ms, - "num_tokens_out": int(text_ids.shape[-1]) if hasattr(text_ids, "shape") else None, - } - with open(stats_path, "a", encoding="utf-8") as f: - f.write(json.dumps(rec, ensure_ascii=False) + "\n") - except Exception: - pass - - global_index += 1 - - print("[Done] Generated", global_index, "responses.") - wall_ms = (time.time() - wall_t0) * 1000.0 - avg_e2e_ms = (sum_request_time_ms / total_requests) if total_requests > 0 else 0.0 - print(f"[E2E] total_requests={total_requests}, wall_time_ms={wall_ms:.2f}, avg_e2e_per_request_ms={avg_e2e_ms:.2f}") - try: - summary = { - "type": "transformers_summary", - "total_requests": total_requests, - "wall_time_ms": wall_ms, - "sum_request_time_ms": sum_request_time_ms, - "avg_e2e_per_request_ms": avg_e2e_ms, - } - with open(stats_path, "a", encoding="utf-8") as f: - f.write(json.dumps(summary, ensure_ascii=False) + "\n") - except Exception: - pass - - -if __name__ == "__main__": - main() - - diff --git a/examples/offline_inference/qwen_2_5_omni/pipelined_end2end.py b/examples/offline_inference/qwen_2_5_omni/pipelined_end2end.py deleted file mode 100644 index 2e394e5f0..000000000 --- a/examples/offline_inference/qwen_2_5_omni/pipelined_end2end.py +++ /dev/null @@ -1,265 +0,0 @@ -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.pipelined_omni_llm import PipelinedOmniLLM - -_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( - "--pt-prompts", type=str, default=None, help="Path to .pt file containing List[str] prompts (overrides --prompts when prompt_type=text)." - ) - 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=True, - help="Enable writing detailed statistics (default: enabled)", - ) - args = parser.parse_args() - return args - - -def main(): - args = parse_args() - model_name = args.model - # If pt-prompts provided and prompt_type is text, load and override args.prompts - try: - if getattr(args, "pt_prompts", None) and args.prompt_type == "text": - loaded = torch.load(args.pt_prompts) - if not isinstance(loaded, list) or (len(loaded) > 0 and not isinstance(loaded[0], str)): - raise ValueError("--pt-prompts must point to a .pt containing List[str]") - args.prompts = loaded - print(f"[Info] Loaded {len(args.prompts)} prompts from {args.pt_prompts}") - except Exception as e: - print(f"[Error] Failed to load pt-prompts: {e}") - raise - - if args.prompts is None: - raise ValueError("No prompts provided. Use --prompts ... or --pt-prompts (with --prompt_type text)") - omni_llm = PipelinedOmniLLM( - model=model_name, - log_stats=args.enable_stats, - log_file="omni_llm_pipeline.log", - init_sleep_seconds=args.init_sleep_seconds, - batch_timeout=args.batch_timeout, - init_timeout=args.init_timeout, - shm_threshold_bytes=args.shm_threshold_bytes, - ) - 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, - ] - - prompt = [make_omni_prompt(args, prompt) for prompt in args.prompts] - omni_outputs = omni_llm.generate(prompt, sampling_params_list) - - # 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") - lines.append("transformers_text_output:\n\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 - ) - # Save aligned text file with empty text (no text output in this stage) - 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\n") - lines.append("transformers_text_output:\n\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}, Saved audio to {output_wav} and text to {out_txt}") - - -if __name__ == "__main__": - main() diff --git a/vllm_omni/__init__.py b/vllm_omni/__init__.py index afeb09c1b..6860b462c 100644 --- a/vllm_omni/__init__.py +++ b/vllm_omni/__init__.py @@ -16,11 +16,11 @@ from . import patch # noqa: F401 from .config import OmniModelConfig # Main entry points -from .entrypoints.omni_llm import OmniLLM +from .entrypoints.omni_llm import BaseOmniLLM __all__ = [ # Main components - "OmniLLM", + "BaseOmniLLM", # "AsyncOmniLLM", # TODO: add AsyncOmniLLM back when it is implemented # Configuration "OmniModelConfig", diff --git a/vllm_omni/entrypoints/log_utils.py b/vllm_omni/entrypoints/log_utils.py index 6a17c24f0..f0c8355df 100644 --- a/vllm_omni/entrypoints/log_utils.py +++ b/vllm_omni/entrypoints/log_utils.py @@ -194,6 +194,28 @@ def log_stage_batch_stats(stats_file: Optional[str], stage_id: int, batch_size: }) +def compute_and_log_stage_request_stats(stats_file: Optional[str], stage_id: int, request_id: Any, + batch_size: int, engine_outputs: List[Any], + stage_gen_time_ms: float, rx_transfer_bytes: int, + rx_decode_time_ms: float) -> None: + """Compute per-request metrics and log them in one call.""" + num_tokens = count_tokens_from_outputs(engine_outputs) + tokens_per_s = (num_tokens * 1000.0 / stage_gen_time_ms) if stage_gen_time_ms > 0 else 0.0 + rx_mbps = (float(rx_transfer_bytes) * 8.0) / (max(float(rx_decode_time_ms), 1e-6) * 1000.0) if rx_transfer_bytes > 0 else 0.0 + log_stage_request_stats( + stats_file, + stage_id, + request_id, + int(batch_size), + int(num_tokens), + float(stage_gen_time_ms), + float(tokens_per_s), + int(rx_transfer_bytes), + float(rx_decode_time_ms), + float(rx_mbps), + ) + + # ----------------- Aggregation helpers for orchestrator ----------------- def record_stage_metrics(per_request: Dict[int, Dict[str, Any]], diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index 57bd12958..458e000e5 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/vllm_omni/entrypoints/omni_llm.py @@ -1,8 +1,15 @@ from typing import Any, Optional, Sequence, Union +import logging +import multiprocessing as mp +import os +import queue +import sys +import time import cloudpickle from pydantic import ValidationError +# External library imports (vLLM) from vllm.config import CompilationConfig, StructuredOutputsConfig, is_init_field from vllm.engine.arg_utils import HfOverrides from vllm.entrypoints.llm import LLM @@ -13,17 +20,30 @@ from vllm.sampling_params import SamplingParams from vllm.usage.usage_lib import UsageContext from vllm.utils import Counter from vllm.v1.engine.llm_engine import LLMEngine + +# Internal imports (our code) from vllm_omni.engine.arg_utils import OmniEngineArgs from vllm_omni.engine.output_processor import MultimodalOutputProcessor from vllm_omni.engine.processor import OmniProcessor from vllm_omni.entrypoints.omni_stage import OmniStage from vllm_omni.entrypoints.utils import load_stage_configs_from_model +from vllm_omni.entrypoints.log_utils import ( + remove_old_logs, + configure_orchestrator_logger, + init_stats_paths, + OrchestratorMetrics, +) +from vllm_omni.entrypoints.pipeline_utils import ( + maybe_load_from_ipc as _load, + encode_for_ipc as _encode, + serialize_obj as _ser, +) from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) -class OmniLLM: +class BaseOmniLLM: def __init__( self, model: str, stage_configs=None, log_stats: bool = False, **kwargs ): @@ -90,6 +110,329 @@ class OmniLLM: return engine_outputs +class OmniLLM(BaseOmniLLM): + """Multi-process pipelined OmniLLM. + + - Per-stage process with copy-based IPC (Queues) + - window=-1 across stages (downstream starts when upstream finishes) + - max_inflight=1 per stage (serial within a stage), but pipeline across stages + """ + + def __init__(self, model: str, + stage_configs=None, + log_stats: bool = False, + log_file: Optional[str] = None, + init_sleep_seconds: int = 20, + shm_threshold_bytes: int = 65536, + batch_timeout: int = 10, + init_timeout: int = 300, + **kwargs): + self.batch_timeout = batch_timeout + self._enable_stats: bool = bool(log_stats) + # Do NOT call super().__init__ to avoid creating OmniStageLLM instances in parent. + if stage_configs is None: + self.stage_configs = load_stage_configs_from_model(model) + else: + self.stage_configs = stage_configs + + self.stage_list: list[OmniStage] = [] + for stage_config in self.stage_configs: + # Only construct lightweight OmniStage without setting engine in parent + stage = OmniStage(stage_config) + self.stage_list.append(stage) + logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list)) + + self._ctx = mp.get_context("spawn") + self._stage_in_queues: list[mp.Queue] = [] + self._stage_out_queues: list[mp.Queue] = [] + # Optional file handler for orchestrator + self._log_file = log_file + if self._log_file: + remove_old_logs(self._log_file, len(self.stage_list)) + configure_orchestrator_logger(logger, self._log_file) + + self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file) + + 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) + # 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: + 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) + 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) + 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, + ) + logger.debug("[Orchestrator] Stage-%s process started", stage_id) + + def close(self) -> None: + for q in self._stage_in_queues: + try: + q.put_nowait(None) + except Exception as e: + logger.warning("[Orchestrator] Failed to send shutdown signal to stage input queue: %s", e) + for stage in self.stage_list: + try: + stage.stop_stage_worker() + except Exception as e: + logger.warning("[Orchestrator] Failed to stop stage worker: %s", e) + + def __del__(self) -> None: # best-effort + try: + self.close() + except Exception as e: + logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True) + + def generate( + self, + prompts: Union[PromptType, Sequence[PromptType]], + sampling_params_list: Optional[ + Union[SamplingParams, Sequence[SamplingParams]] + ] = None, + ) -> list[OmniRequestOutput]: + logger.debug("[Orchestrator] generate() called") + if sampling_params_list is None: + raise ValueError("sampling_params_list is required for pipelined generation") + if len(sampling_params_list) != len(self.stage_list): + raise ValueError( + f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}" + ) + + # Normalize prompts to a list for per-request iteration + if not isinstance(prompts, (list, tuple)): + request_prompts: list[PromptType] = [prompts] + else: + request_prompts = list(prompts) + + final_outputs: list[OmniRequestOutput] = [] + + # 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)} + + # Track per-request start time for end-to-end timing + _req_start_ts: dict[int, float] = {} + _wall_start_ts: float = time.time() + _last_finish_ts: float = _wall_start_ts + + # Determine the final stage for E2E stats (highest stage_id with final_output=True; fallback to last stage) + final_stage_id_for_e2e = -1 + try: + for _sid, _st in enumerate(self.stage_list): + if getattr(_st, "final_output", False): + final_stage_id_for_e2e = max(final_stage_id_for_e2e, _sid) + if final_stage_id_for_e2e < 0: + final_stage_id_for_e2e = len(self.stage_list) - 1 + except Exception as e: + logger.debug("[Orchestrator] Failed to determine final stage for E2E; falling back to last: %s", e, exc_info=True) + final_stage_id_for_e2e = len(self.stage_list) - 1 + # Metrics/aggregation helper + metrics = OrchestratorMetrics(num_stages, self._enable_stats, self._stats_file, self._overall_stats_file, _wall_start_ts) + + # Seed stage-0 queue with all requests + logger.debug("[Orchestrator] Seeding %d requests into stage-0", len(request_prompts)) + # Mark first input time for stage-0 + metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time() + + for req_id, prompt in request_id_to_prompt.items(): + sp0: SamplingParams = sampling_params_list[0] # type: ignore[index] + task = { + "request_id": req_id, + "engine_inputs": prompt, + "sampling_params": sp0, + } + self.stage_list[0].submit(task) + _req_start_ts[req_id] = time.time() + logger.debug("[Orchestrator] Enqueued request %s to stage-0", req_id) + + # For each stage, forward results to next stage; collect finals at the end + # We pipeline by continually polling output queues in stage order + remaining_by_stage: list[int] = [len(request_prompts)] + [0] * (num_stages - 1) + completed_requests = 0 + total_requests = len(request_prompts) + + logger.debug("[Orchestrator] Entering scheduling loop: total_requests=%d, stages=%d", total_requests, num_stages) + while completed_requests < total_requests: + made_progress = False + for stage_id, stage in enumerate(self.stage_list): + result = stage.try_collect() + if result is None: + continue + + made_progress = True + req_id = result.get("request_id") + if "error" in result: + logger.error("Stage %s error on request %s: %s", stage_id, req_id, result["error"]) + continue + + if result.get("type") == "stage_ready": + #Only happens when stage is initialized slower than expected, so we wait for a short time and try again + time.sleep(0.05) + continue + + engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm") + # Mark last output time for this stage whenever we receive outputs + metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time()) + try: + _m = result.get("metrics") + if _m is not None: + metrics.on_stage_metrics(stage_id, req_id, _m) + except Exception as e: + logger.exception("[Orchestrator] Failed to process metrics for stage %s, req %s: %s", stage_id, req_id, e) + logger.debug("[Orchestrator] Stage-%s completed request %s; forwarding or finalizing", stage_id, req_id) + stage.set_engine_outputs(engine_outputs) + + if getattr(stage, "final_output", False): + final_outputs.append( + OmniRequestOutput( + stage_id=stage_id, + final_output_type=stage.final_output_type, # type: ignore[attr-defined] + request_output=engine_outputs, + ) + ) + logger.debug("[Orchestrator] Request %s finalized at stage-%s", req_id, stage_id) + + # End-to-end timing and time-per-token for final output (only once per request at the designated final stage) + try: + rid_int = int(req_id) if isinstance(req_id, (int, str)) and str(req_id).isdigit() else req_id + if stage_id == final_stage_id_for_e2e and rid_int not in metrics.e2e_done: + metrics.on_finalize_request(stage_id, req_id, engine_outputs, _req_start_ts.get(req_id, _wall_start_ts)) + except Exception as e: + logger.exception("[Orchestrator] Finalize request handling error for req %s at stage %s: %s", req_id, stage_id, e) + + next_stage_id = stage_id + 1 + if next_stage_id < num_stages: + next_stage: OmniStage = self.stage_list[next_stage_id] + next_inputs = next_stage.process_engine_inputs(self.stage_list, [request_id_to_prompt[req_id]]) + sp_next: SamplingParams = sampling_params_list[next_stage_id] # type: ignore[index] + try: + # Measure transfer size and time (encode + enqueue) + size_bytes = 0 + try: + size_bytes = len(_ser(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, next_stage_id) + remaining_by_stage[next_stage_id] += 1 + else: + completed_requests += 1 + logger.debug("[Orchestrator] Request %s fully completed (%d/%d)", req_id, completed_requests, total_requests) + + if not made_progress: + time.sleep(0.005) + logger.debug("[Orchestrator] All requests completed") + + # Summarize and print stats + try: + summary = metrics.build_and_log_summary(final_stage_id_for_e2e) + logger.info("[Summary] %s", summary) + except Exception as e: + logger.exception("[Orchestrator] Failed to build/log summary: %s", e) + + return final_outputs + + def _wait_for_stages_ready(self, timeout: int = 120) -> None: + deadline = time.time() + max(0, int(timeout)) + num_stages = len(self.stage_list) + while len(self._stages_ready) < num_stages and time.time() < deadline: + progressed = False + for stage_id, stage in enumerate(self.stage_list): + if stage_id in self._stages_ready: + continue + result = stage.try_collect() + if result is None: + continue + progressed = True + if result.get("type") == "stage_ready": + self._stages_ready.add(stage_id) + logger.debug("[Orchestrator] Stage-%s reported ready", stage_id) + else: + # No user data should arrive before seeding; ignore other messages + pass + if not progressed: + time.sleep(0.01) + if len(self._stages_ready) < num_stages: + not_ready = sorted(set(range(num_stages)) - set(self._stages_ready)) + logger.warning( + "[Orchestrator] Initialization timeout: only %s/%s stages are ready; not ready: %s", + len(self._stages_ready), num_stages, not_ready, + ) + # Provide actionable suggestions before shutdown + try: + suggestions = [ + "Verify GPU/device assignment in config (mrs.devices) is correct.", + "Check GPU/host memory availability; reduce model or batch size if needed.", + "Check model weights path and network reachability (if loading remotely).", + "Increase initialization wait time (init_sleep_seconds or call-site timeout).", + ] + if getattr(self, "_log_file", None): + suggestions.append( + f"Inspect per-stage log files for details: {self._log_file}.stage.log" + ) + logger.error( + "[Orchestrator] Stage initialization failed, shutting down. Suggestions:\n- %s", + "\n- ".join(suggestions), + ) + except Exception: + # Best-effort logging of suggestions + logger.error( + "[Orchestrator] Stage initialization failed and an error occurred while logging suggestions", + ) + + # Attempt graceful shutdown of all stages before exiting + try: + self.close() + except Exception: + pass + + # Terminate the current process with non-zero exit code + try: + sys.exit(1) + except SystemExit: + raise + except Exception: + os._exit(1) + + class OmniStageLLM(LLM): def __init__( self, diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 6f40064ea..06fb12041 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -14,7 +14,7 @@ from typing import Any, Dict, List, Optional, Union from vllm_omni.entrypoints.pipeline_utils import ( _to_dict, set_stage_gpu_devices, - maybe_load_from_ipc, + maybe_load_from_ipc_with_metrics, maybe_dump_to_shm, ) @@ -186,13 +186,14 @@ def _stage_worker( import logging as _logging from vllm_omni.entrypoints.omni_llm import OmniStageLLM # noqa: WPS433 from vllm_omni.entrypoints.log_utils import ( # noqa: WPS433 - log_stage_request_stats, log_stage_running_avg, log_stage_batch_stats, + compute_and_log_stage_request_stats, + count_tokens_from_outputs, ) import queue as _queue import time as _time - from vllm_omni.entrypoints.pipeline_utils import append_jsonl as _append_jsonl, serialize_obj as _ser + # no inline JSONL/serialization imports; logging handled by utilities stage_id = stage_payload["stage_id"] engine_args = stage_payload.get("engine_args", {}) @@ -277,17 +278,11 @@ def _stage_worker( _in_flight_ms_by_rid[rid] = 0.0 except Exception: _in_flight_ms_by_rid[rid] = 0.0 - _t0 = _time.time() - ein = maybe_load_from_ipc(t, obj_key="engine_inputs", shm_key="engine_inputs_shm") - _t1 = _time.time() - _rx_decode_ms_by_rid[rid] = (_t1 - _t0) * 1000.0 - try: - if "engine_inputs_shm" in t: - _rx_bytes_by_rid[rid] = int(t["engine_inputs_shm"]["size"]) # type: ignore[index] - else: - _rx_bytes_by_rid[rid] = len(_ser(ein)) - except Exception: - _rx_bytes_by_rid[rid] = 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)) batch_request_ids.append(rid) if isinstance(ein, list): batch_engine_inputs.extend(ein) @@ -326,41 +321,12 @@ def _stage_worker( idx += 1 # Per-request stats logging and aggregates - def _count_tokens(_ros: List[Any]) -> int: - total = 0 - for _ro in _ros: - try: - outs = getattr(_ro, "outputs", None) - if outs and len(outs) > 0: - tokens = getattr(outs[0], "token_ids", None) - if tokens is not None: - total += len(tokens) - except Exception: - pass - return total - for rid in batch_request_ids: _r_outputs = req_to_outputs.get(rid, []) - _num_tokens = _count_tokens(_r_outputs) + _num_tokens = count_tokens_from_outputs(_r_outputs) _agg_total_tokens += _num_tokens _agg_total_gen_time_ms += _gen_ms _tokens_per_s = (_num_tokens * 1000.0 / _gen_ms) if _gen_ms > 0 else 0.0 - if _stats_file: - try: - _append_jsonl(_stats_file, { - "type": "stage_request_stats", - "stage_id": stage_id, - "request_id": rid, - "batch_size": len(batch_tasks), - "num_tokens_out": int(_num_tokens), - "stage_gen_time_ms": _gen_ms, - "tokens_per_s": _tokens_per_s, - "rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)), - "rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)), - "rx_mbps": (float(_rx_bytes_by_rid.get(rid, 0)) * 8.0) / (max(float(_rx_decode_ms_by_rid.get(rid, 0.0)), 1e-6) * 1000.0), - }) - except Exception: - pass if _stats_file: _avg_tokens_per_s = (_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms) if _agg_total_gen_time_ms > 0 else 0.0 @@ -373,7 +339,7 @@ def _stage_worker( try: use_shm, payload = maybe_dump_to_shm(r_outputs, shm_threshold_bytes) _metrics = { - "num_tokens_out": int(_count_tokens(r_outputs)), + "num_tokens_out": int(count_tokens_from_outputs(r_outputs)), "stage_gen_time_ms": _gen_ms, "batch_id": int(_batch_seq), "rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)), @@ -381,22 +347,15 @@ def _stage_worker( "rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)), } if _stats_file: - _num_tokens = int(_metrics["num_tokens_out"]) # type: ignore[index] - _tokens_per_s = (_num_tokens * 1000.0 / _gen_ms) if _gen_ms > 0 else 0.0 - _rx_b = int(_metrics["rx_transfer_bytes"]) # type: ignore[index] - _rx_ms = float(_metrics["rx_decode_time_ms"]) # type: ignore[index] - _rx_mbps = (float(_rx_b) * 8.0) / (max(float(_rx_ms), 1e-6) * 1000.0) if _rx_b > 0 else 0.0 - log_stage_request_stats( + compute_and_log_stage_request_stats( _stats_file, stage_id, rid, len(batch_tasks), - _num_tokens, + r_outputs, float(_gen_ms), - float(_tokens_per_s), - _rx_b, - _rx_ms, - float(_rx_mbps), + int(_metrics["rx_transfer_bytes"]), # type: ignore[index] + float(_metrics["rx_decode_time_ms"]), # type: ignore[index] ) if use_shm: out_q.put({ @@ -418,7 +377,7 @@ def _stage_worker( "stage_id": stage_id, "engine_outputs": r_outputs, "metrics": { - "num_tokens_out": int(_count_tokens(r_outputs)), + "num_tokens_out": int(count_tokens_from_outputs(r_outputs)), "stage_gen_time_ms": _gen_ms, "rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)), "rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)), diff --git a/vllm_omni/entrypoints/pipeline_utils.py b/vllm_omni/entrypoints/pipeline_utils.py index 6a2384ec9..16f8c6331 100644 --- a/vllm_omni/entrypoints/pipeline_utils.py +++ b/vllm_omni/entrypoints/pipeline_utils.py @@ -170,12 +170,47 @@ def maybe_dump_to_shm(obj: Any, threshold: int) -> Tuple[bool, Any]: def maybe_load_from_ipc(container: Dict[str, Any], obj_key: str, shm_key: str) -> Any: - """Load object from container that may carry SHM or inline object.""" + """Load object from container that may carry SHM or inline object. + + Deprecated: prefer `maybe_load_from_ipc_with_metrics` to also obtain + decode-time and size metrics. + """ if shm_key in container: return pickle.loads(shm_read_bytes(container[shm_key])) return container[obj_key] +def maybe_load_from_ipc_with_metrics(container: Dict[str, Any], obj_key: str, shm_key: str) -> tuple[Any, Dict[str, float]]: + """Load object and return (object, metrics) with RX bytes and decode time. + + Metrics keys: + - rx_transfer_bytes: int + - rx_decode_time_ms: float + """ + import time as _time # local import to avoid overhead at module import + t0 = _time.time() + if shm_key in container: + meta = container[shm_key] # type: ignore[index] + payload = shm_read_bytes(meta) + obj = pickle.loads(payload) + try: + rx_bytes = int(meta.get("size", len(payload))) # type: ignore[call-arg] + except Exception: + rx_bytes = len(payload) + else: + obj = container[obj_key] + try: + rx_bytes = len(serialize_obj(obj)) + except Exception: + rx_bytes = 0 + t1 = _time.time() + rx_decode_ms = (t1 - t0) * 1000.0 + return obj, { + "rx_transfer_bytes": int(rx_bytes), + "rx_decode_time_ms": float(rx_decode_ms), + } + + def encode_for_ipc(obj: Any, threshold: int, obj_key: str, shm_key: str) -> Dict[str, Any]: """Return a dict payload for IPC: inline (obj_key) or SHM (shm_key). diff --git a/vllm_omni/entrypoints/pipelined_omni_llm.py b/vllm_omni/entrypoints/pipelined_omni_llm.py deleted file mode 100644 index da13b0dee..000000000 --- a/vllm_omni/entrypoints/pipelined_omni_llm.py +++ /dev/null @@ -1,370 +0,0 @@ -from __future__ import annotations - -from typing import Any, Dict, List, Optional, Sequence, Union - -import os -import sys -import queue -import time -import multiprocessing as mp -import logging - -from vllm.inputs import PromptType -from vllm.logger import init_logger -from vllm.sampling_params import SamplingParams - -from vllm_omni.entrypoints.omni_llm import OmniLLM, OmniStageLLM -from vllm_omni.entrypoints.omni_stage import OmniStage -from vllm_omni.entrypoints.utils import load_stage_configs_from_model -from vllm_omni.entrypoints.log_utils import ( - remove_old_logs, - configure_orchestrator_logger, - init_stats_paths, - record_stage_metrics, - aggregate_rx_and_maybe_total, - record_sender_transfer_agg, - count_tokens_from_outputs, - build_stage_summary, - build_transfer_summary, - log_transfer_tx, - log_transfer_rx, - log_transfer_total, - log_orchestrator_e2e, - log_orchestrator_summary, - log_overall_summary, - OrchestratorMetrics, -) -from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.entrypoints.pipeline_utils import ( - maybe_load_from_ipc as _load, - encode_for_ipc as _encode, - serialize_obj as _ser, - append_jsonl as _append_jsonl, -) - -logger = init_logger(__name__) - - -class PipelinedOmniLLM(OmniLLM): - """Multi-process pipelined OmniLLM. - - - Per-stage process with copy-based IPC (Queues) - - window=-1 across stages (downstream starts when upstream finishes) - - max_inflight=1 per stage (serial within a stage), but pipeline across stages - """ - - def __init__(self, model: str, - stage_configs=None, - log_stats: bool = True, - log_file: Optional[str] = None, - init_sleep_seconds: int = 20, - shm_threshold_bytes: int = 65536, - batch_timeout: int = 10, - init_timeout: int = 300, - **kwargs): - self.batch_timeout = batch_timeout - self._enable_stats: bool = bool(log_stats) - # Do NOT call super().__init__ to avoid creating OmniStageLLM instances in parent. - if stage_configs is None: - self.stage_configs = load_stage_configs_from_model(model) - else: - self.stage_configs = stage_configs - - self.stage_list: List[OmniStage] = [] - for stage_config in self.stage_configs: - # Only construct lightweight OmniStage without setting engine in parent - stage = OmniStage(stage_config) - self.stage_list.append(stage) - logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list)) - - self._ctx = mp.get_context("spawn") - self._stage_in_queues: List[mp.Queue] = [] - self._stage_out_queues: List[mp.Queue] = [] - # Optional file handler for orchestrator - self._log_file = log_file - if self._log_file: - remove_old_logs(self._log_file, len(self.stage_list)) - configure_orchestrator_logger(logger, self._log_file) - - self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file) - - 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) - # 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: - 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) - 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) - 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, - ) - logger.debug("[Orchestrator] Stage-%s process started", stage_id) - - def close(self) -> None: - for q in self._stage_in_queues: - try: - q.put_nowait(None) - except Exception as e: - logger.warning("[Orchestrator] Failed to send shutdown signal to stage input queue: %s", e) - for stage in self.stage_list: - try: - stage.stop_stage_worker() - except Exception as e: - logger.warning("[Orchestrator] Failed to stop stage worker: %s", e) - - def __del__(self) -> None: # best-effort - try: - self.close() - except Exception as e: - logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True) - - def generate( - self, - prompts: Union[PromptType, Sequence[PromptType]], - sampling_params_list: Optional[ - Union[SamplingParams, Sequence[SamplingParams]] - ] = None, - ) -> List[OmniRequestOutput]: - logger.debug("[Orchestrator] generate() called") - if sampling_params_list is None: - raise ValueError("sampling_params_list is required for pipelined generation") - if len(sampling_params_list) != len(self.stage_list): - raise ValueError( - f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}" - ) - - # Normalize prompts to a list for per-request iteration - if not isinstance(prompts, (list, tuple)): - request_prompts: List[PromptType] = [prompts] - else: - request_prompts = list(prompts) - - final_outputs: List[OmniRequestOutput] = [] - - # 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)} - - # Track per-request start time for end-to-end timing - _req_start_ts: Dict[int, float] = {} - _wall_start_ts: float = time.time() - _last_finish_ts: float = _wall_start_ts - - # Determine the final stage for E2E stats (highest stage_id with final_output=True; fallback to last stage) - final_stage_id_for_e2e = -1 - try: - for _sid, _st in enumerate(self.stage_list): - if getattr(_st, "final_output", False): - final_stage_id_for_e2e = max(final_stage_id_for_e2e, _sid) - if final_stage_id_for_e2e < 0: - final_stage_id_for_e2e = len(self.stage_list) - 1 - except Exception as e: - logger.debug("[Orchestrator] Failed to determine final stage for E2E; falling back to last: %s", e, exc_info=True) - final_stage_id_for_e2e = len(self.stage_list) - 1 - # Metrics/aggregation helper - metrics = OrchestratorMetrics(num_stages, self._enable_stats, self._stats_file, self._overall_stats_file, _wall_start_ts) - - # Seed stage-0 queue with all requests - logger.debug("[Orchestrator] Seeding %d requests into stage-0", len(request_prompts)) - # Mark first input time for stage-0 - metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time() - - for req_id, prompt in request_id_to_prompt.items(): - sp0: SamplingParams = sampling_params_list[0] # type: ignore[index] - task = { - "request_id": req_id, - "engine_inputs": prompt, - "sampling_params": sp0, - } - self.stage_list[0].submit(task) - _req_start_ts[req_id] = time.time() - logger.debug("[Orchestrator] Enqueued request %s to stage-0", req_id) - - # For each stage, forward results to next stage; collect finals at the end - # We pipeline by continually polling output queues in stage order - remaining_by_stage: List[int] = [len(request_prompts)] + [0] * (num_stages - 1) - completed_requests = 0 - total_requests = len(request_prompts) - - logger.debug("[Orchestrator] Entering scheduling loop: total_requests=%d, stages=%d", total_requests, num_stages) - while completed_requests < total_requests: - made_progress = False - for stage_id, stage in enumerate(self.stage_list): - result = stage.try_collect() - if result is None: - continue - - made_progress = True - req_id = result.get("request_id") - if "error" in result: - logger.error("Stage %s error on request %s: %s", stage_id, req_id, result["error"]) - continue - - if result.get("type") == "stage_ready": - #Only happens when stage is initialized slower than expected, so we wait for a short time and try again - time.sleep(0.05) - continue - - engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm") - # Mark last output time for this stage whenever we receive outputs - metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time()) - try: - _m = result.get("metrics") - if _m is not None: - metrics.on_stage_metrics(stage_id, req_id, _m) - except Exception as e: - logger.exception("[Orchestrator] Failed to process metrics for stage %s, req %s: %s", stage_id, req_id, e) - logger.debug("[Orchestrator] Stage-%s completed request %s; forwarding or finalizing", stage_id, req_id) - stage.set_engine_outputs(engine_outputs) - - if getattr(stage, "final_output", False): - final_outputs.append( - OmniRequestOutput( - stage_id=stage_id, - final_output_type=stage.final_output_type, # type: ignore[attr-defined] - request_output=engine_outputs, - ) - ) - logger.debug("[Orchestrator] Request %s finalized at stage-%s", req_id, stage_id) - - # End-to-end timing and time-per-token for final output (only once per request at the designated final stage) - try: - rid_int = int(req_id) if isinstance(req_id, (int, str)) and str(req_id).isdigit() else req_id - if stage_id == final_stage_id_for_e2e and rid_int not in metrics.e2e_done: - metrics.on_finalize_request(stage_id, req_id, engine_outputs, _req_start_ts.get(req_id, _wall_start_ts)) - except Exception as e: - logger.exception("[Orchestrator] Finalize request handling error for req %s at stage %s: %s", req_id, stage_id, e) - - next_stage_id = stage_id + 1 - if next_stage_id < num_stages: - next_stage: OmniStage = self.stage_list[next_stage_id] - next_inputs = next_stage.process_engine_inputs(self.stage_list, [request_id_to_prompt[req_id]]) - sp_next: SamplingParams = sampling_params_list[next_stage_id] # type: ignore[index] - try: - # Measure transfer size and time (encode + enqueue) - size_bytes = 0 - try: - size_bytes = len(_ser(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, next_stage_id) - remaining_by_stage[next_stage_id] += 1 - else: - completed_requests += 1 - logger.debug("[Orchestrator] Request %s fully completed (%d/%d)", req_id, completed_requests, total_requests) - - if not made_progress: - time.sleep(0.005) - logger.debug("[Orchestrator] All requests completed") - - # Summarize and print stats - try: - summary = metrics.build_and_log_summary(final_stage_id_for_e2e) - logger.info("[Summary] %s", summary) - except Exception as e: - logger.exception("[Orchestrator] Failed to build/log summary: %s", e) - - return final_outputs - - def _wait_for_stages_ready(self, timeout: int = 120) -> None: - deadline = time.time() + max(0, int(timeout)) - num_stages = len(self.stage_list) - while len(self._stages_ready) < num_stages and time.time() < deadline: - progressed = False - for stage_id, stage in enumerate(self.stage_list): - if stage_id in self._stages_ready: - continue - result = stage.try_collect() - if result is None: - continue - progressed = True - if result.get("type") == "stage_ready": - self._stages_ready.add(stage_id) - logger.debug("[Orchestrator] Stage-%s reported ready", stage_id) - else: - # No user data should arrive before seeding; ignore other messages - pass - if not progressed: - time.sleep(0.01) - if len(self._stages_ready) < num_stages: - not_ready = sorted(set(range(num_stages)) - set(self._stages_ready)) - logger.warning( - "[Orchestrator] Initialization timeout: only %s/%s stages are ready; not ready: %s", - len(self._stages_ready), num_stages, not_ready, - ) - # Provide actionable suggestions before shutdown - try: - suggestions = [ - "Verify GPU/device assignment in config (mrs.devices) is correct.", - "Check GPU/host memory availability; reduce model or batch size if needed.", - "Check model weights path and network reachability (if loading remotely).", - "Increase initialization wait time (init_sleep_seconds or call-site timeout).", - ] - if getattr(self, "_log_file", None): - suggestions.append( - f"Inspect per-stage log files for details: {self._log_file}.stage.log" - ) - logger.error( - "[Orchestrator] Stage initialization failed, shutting down. Suggestions:\n- %s", - "\n- ".join(suggestions), - ) - except Exception: - # Best-effort logging of suggestions - logger.error( - "[Orchestrator] Stage initialization failed and an error occurred while logging suggestions", - ) - - # Attempt graceful shutdown of all stages before exiting - try: - self.close() - except Exception: - pass - - # Terminate the current process with non-zero exit code - try: - sys.exit(1) - except SystemExit: - raise - except Exception: - os._exit(1) - -