From 6a7a27801bcaa1634b09f28b6998c54e0f6e691d Mon Sep 17 00:00:00 2001 From: tzhouam Date: Fri, 7 Nov 2025 15:48:47 +0800 Subject: [PATCH 01/19] Port changes from feat/multi-request-stream (vs main) --- docs/mrs_design.md | 238 +++++++ .../qwen_2_5_omni/extract_prompts.py | 48 ++ .../hf_qwen25_omni_batch_no_template.py | 279 ++++++++ .../qwen_2_5_omni/pipelined_end2end.py | 264 ++++++++ vllm_omni/entrypoints/omni_stage.py | 334 +++++++++- vllm_omni/entrypoints/pipeline_utils.py | 202 ++++++ vllm_omni/entrypoints/pipelined_omni_llm.py | 621 ++++++++++++++++++ .../model_executor/models/qwen2_5_omni.py | 182 +++-- .../stage_configs/qwen2_5_omni.yaml | 34 +- .../stage_input_processors/qwen2_5_omni.py | 10 +- vllm_omni/worker/gpu_ar_model_runner.py | 154 +++-- 11 files changed, 2250 insertions(+), 116 deletions(-) create mode 100644 docs/mrs_design.md create mode 100644 examples/offline_inference/qwen_2_5_omni/extract_prompts.py create mode 100644 examples/offline_inference/qwen_2_5_omni/hf_qwen25_omni_batch_no_template.py create mode 100644 examples/offline_inference/qwen_2_5_omni/pipelined_end2end.py create mode 100644 vllm_omni/entrypoints/pipeline_utils.py create mode 100644 vllm_omni/entrypoints/pipelined_omni_llm.py diff --git a/docs/mrs_design.md b/docs/mrs_design.md new file mode 100644 index 000000000..2c463fa64 --- /dev/null +++ b/docs/mrs_design.md @@ -0,0 +1,238 @@ +## Multi-Request Streaming (MRS) on a Single Machine — Design Doc + +- Author: Team +- Date: 2025-10-27 +- Version: v0.3 (simplified: copy-based IPC / serial per stage / full-trigger) + +### 1. Background & Scope +- All processing runs on a single physical machine with multi-process, per-stage workers. No proxy or network transport involved. +- Current alignment with vllm-omni: `OmniLLM` supports multiple stages (`OmniStage`). GPU runners already expose streamable steps (prefill/decoding/diffusion), but the entry layer still collects lists and lacks intra-stage streaming and window scheduling. +- Goal: implement multi-stage, multi-request streaming (MRS) locally. Each stage outputs segments; downstream stages stitch and trigger compute based on configured windows. Shared memory and zero-copy strategies reduce data movement overhead. + +### 2. Key Constraints +- Multi-process per stage: each stage is an independent process with a while loop; device visibility can be configured (`CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`). +- Simple IPC (copy-based): use `multiprocessing.Queue`/Pipe for inter-process communication with CPU copies/serialization; do not rely on CUDA IPC/SHM zero-copy in this version. +- Full-trigger across stages: unify `window_size = -1` for all edges — downstream starts only after upstream completes for that request. +- Serial within a stage: at most one in-flight request per stage (max_inflight=1). No batching/concurrency inside a stage in this version. +- Cross-stage pipeline: different stages can process different requests concurrently (e.g., stage A handles request 1 while stage B handles request 0). + +### 3. Architecture Overview +- Processes & IPC queues + - Each "sub-stage" is an OS process (worker). The loop: take from input_queue → compute → put to output_queue. + - Inter-stage connection via IPC: control plane (Unix Domain Socket / Pipe) + data plane (shared memory ring buffer in future; copy-based in this version). + - MPMC queue semantics; control plane passes light-weight pointers/metadata; data plane carries large payloads. +- Device visibility + - Each stage sets `CUDA_VISIBLE_DEVICES` or calls `torch.cuda.set_device` to bind to GPU sets. + - A stage may use multiple GPUs internally (TP/PP/DP) but presents as a single stage unit. +- Simplified IPC: copy-based queues/pipes for data transfer; zero-copy is future work. +- Pipeline progression: when a stage finishes a request, it enqueues outputs to the downstream stage; if downstream is idle, it starts immediately. +- Scheduling & windows + - Each edge (stage→stage) defines a window policy; downstream triggers compute by consuming stitched window views. + - `-1` mode triggers only after upstream emits `is_last=True` (full aggregation). + - Backpressure propagates upstream via the control plane: once high-watermark is reached, upstream blocks or degrades behavior. + +### 4. Stage Breakdown +- AR (autoregressive) + - S0 Intake/Plan: parse request, assign request_id, output normalized tasks. + - S1 Prefill: tokenize/encode/KV warmup; may micro-batch across requests. + - S2 Decode Loop: step-by-step tokens; flush segments by `max_segment_tokens` and `min_flush_interval_ms`. + - S3 Postprocess: detokenize, logprobs/struct wrapping (CPU if necessary). + - S4 Consumer: local consumer callback/buffering. +- DiT/Diffusion + - D0 Intake/Plan: scheduler/seed init. + - D1 Sampling Loop: preview frames every k steps (downsample allowed). + - D2 Decode/Encode: VAE decode and JPEG/WEBP encode (prefer GPU→pinned CPU). + - D3 Postprocess/Consumer. + +### 5. Unified Events & Segment Structure (in-process) +- Base event: `{ request_id, sequence, event_type, ts, payload, device_handles }` +- ARSegmentEvent (segment): + - Fields: `tokens` (List[int] or GPU tensor view), optional `logprobs`, `is_last`, optional `hidden_states/aux` refs. + - Semantics: strictly ordered by `sequence` for the same `request_id`; `is_last` marks the final segment. +- ImageProgressEvent: `{ step, image_tensor|jpeg_ref, is_last }` (diffusion preview/completion). +- Completed/Error: terminal or error events affect only the current request. + +### 6. AR Segmentation & Windowing (simplified) +- Production (upstream stage) + - Accumulate tokens during decoding; flush a segment if any of the following holds: + - Reaches `max_segment_tokens` length + - Exceeds `min_flush_interval_ms` since last flush + - Request ended (`is_last=True`) + - Flush immediately via `output_queue.put(segment)`; do not wait for entire request completion. +- Consumption (downstream stage) + - Use unified `window_size = -1`: trigger only after receiving `is_last=True`; no window slicing. +- Backpressure & watermark + - Must-deliver events (token.delta, Completed) are not dropped; block upstream briefly on high watermark. + - For long sequences under `-1`, allow early segments to spill from GPU to pinned CPU to control VRAM pressure. + +#### Example config (YAML, simplified) +```yaml +mrs: + enabled: true + defaults: + window_size: -1 # Full-trigger across stages + max_inflight: 1 # Serial per stage +``` + +### 7. IPC Implementation (simplified: copy-based) +- Use `multiprocessing.Queue`/Pipe for inter-process communication (control + data). +- Data is serialized/copied via CPU; no CUDA IPC/SHM zero-copy in this version. +- Backpressure: blocking reads/writes based on queue capacity, blocking upstream on high watermark. + +### 8. Scheduling & Cancellation (simplified) +- Serial FIFO: each stage `max_inflight=1`; no cross-request micro-batching or concurrency. +- Cancel/timeout: ongoing work should exit early; remaining queued requests stay pending. +- Pipeline: when a stage finishes a request, it enqueues to the next stage; that stage immediately pulls the next request from its input queue, enabling cross-stage concurrency. + +#### Short sequence example (req0/req1, stage A→B) +1) t0: stage A handles req0 +2) t1: req0 completes on A → enters B; A immediately starts req1 +3) t2: B handles req0 while A handles req1 (parallel across stages) + +### 9. Integration Points (by file) +- `vllm_omni/entrypoints/omni_llm.py` + - Introduce a ProcessOrchestrator (replace threaded version): + - Spawn stage processes per config (set `CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`), create control/data channels, build window state machines. + - Add `enable_mrs` in `generate`; when enabled, do not collect `stage.engine.generate(...)` into lists; stream segments to downstream via IPC instead. + - Register `window_size` per edge; `-1` triggers downstream once `is_last` is observed. + - Manage process lifecycle: start/heartbeat/graceful shutdown/restart on crash; propagate cancel/timeout. +- `vllm_omni/outputs.py` + - Add in-process events: `ARSegmentEvent`, `ImageProgressEvent`, `CompletedEvent`, `ErrorEvent`. + - Keep `OmniRequestOutput` as the external representation of final stage outputs. +- `vllm_omni/worker/gpu_ar_model_runner.py` + - Near `valid_sampled_token_ids` parsing (after updating `req_state.output_token_ids`), aggregate new tokens into segments and emit; send final `is_last=True` and Completed. + - Replace current `.to("cpu").contiguous()` sync path: by default keep GPU refs and let downstream copy asynchronously from a pinned pool when needed. +- `vllm_omni/worker/gpu_diffusion_model_runner.py` + +### 9.1 Process Device Visibility +- Device binding: set `CUDA_VISIBLE_DEVICES` before process start, or call `torch.cuda.set_device` early in init; configs can be single/multi-GPU. +- Cross-device transfer (simplified): via CPU copies; zero-copy paths are out of scope for now. + +### 10. Worker Template & Window State (pseudo-code) +```python +def stage_worker(input_q, output_q, ctx): + while ctx.alive: + item = input_q.get() # blocking / with timeout + if item.is_cancelled(): + continue + with cuda_stream(ctx.stream): + out_segments = run_stage_compute(item) # produce segments + for seg in out_segments: + output_q.put(seg) + +class WindowState: + def __init__(self, window_size): + self.window_size = window_size + self.last_emitted = 0 + self.buffer = [] # references/paged views + def on_segment(self, seg): + self.buffer.append(seg) + if self.window_size == -1: + return self.emit_full_if_last(seg) + return self.emit_windows() + def emit_windows(self): + emits = [] + while self._pending_len() - self.last_emitted >= self.window_size: + emits.append(self._make_view(self.last_emitted, self.window_size)) + self.last_emitted += self.window_size + return emits + def emit_full_if_last(self, seg): + return [self._make_full_view()] if seg.is_last else [] +``` + +### 11. Configuration +- `mrs.enabled`, `mrs.ar.default_window_size`, `mrs.ar.min_flush_interval_ms`, `mrs.ar.max_segment_tokens`, `mrs.ar.scheduler_policy`. +- Devices/processes: + - `mrs.stages`: [{ name, process: true, devices: "0" | "1,2", extra_env: {} }] + - `mrs.edges`: [{ from, to, window_size }] +- Pinned/SHM pool sizing: estimate by concurrency and max sequence; tune watermarks and GC policies; configure CUDA IPC limits and cleanup thresholds if enabled. + +### 12. Observability +- Metrics: + - Stage-level: enqueue/dequeue rates, queue depth, blocking time on watermarks, window triggers/counts. + - Request-level: TTFT, TPST, completion latency, cancel rate. + - Resource-level: GPU utilization, VRAM usage/fragmentation, pinned usage. +- Logs/tracing: carry `request_id`, stage name, and `sequence`; sample decisions for downsampling/backpressure. + +### 13. Test Plan +- Unit: queue correctness, window state machine, cancel/timeout, memory pool rent/return. +- Performance: AR QPS/TPST vs. window/segment sizes; diffusion preview frequency vs. throughput. +- Stability: long-run memory leaks; backpressure trigger/recovery paths. +- Compatibility: behavior parity when MRS is disabled. + +### 14. Milestones +- M1: Process Orchestrator with simple IPC; serial per stage; full-trigger across stages (window=-1). +- M2: Hardening and observability; cancel/timeout and error paths. +- M3: Optional zero-copy and concurrency/micro-batching optimizations. + +### 15. Risks & Mitigations +— End-to-end latency: full-trigger sacrifices interactivity; start simple/correct, add windowed streaming later. +— CPU copy overhead: copy-based IPC for maintainability; add SHM/zero-copy later as an optimization path. +— Serial throughput: `max_inflight=1` limits per-stage throughput; enable batching/concurrency later as needed. + +### 16. Refactor: Sink LLM init and process into Stage (diagram) + +Objective: without changing external call patterns, encapsulate per-stage LLM init, process creation, and worker logic inside `Stage`, so `PipelinedOmniLLM` focuses on orchestration (seeding, polling, forwarding, collecting) while keeping multi-process, shared memory, and device mapping capabilities. + +Key changes (vs. current): +- Before: `PipelinedOmniLLM` created processes directly, passing `_stage_worker`; device setup, LLM init, batching, and SHM were implemented there. +- After: enhance `Stage` to start its own subprocess via `start_process()`, with `worker_main()` as the entry; device setup, LLM init, batching, and SHM live in `Stage`. + +Classes & responsibilities +- PipelinedOmniLLM (orchestrator) + - Build `Stage` list (preserve `OmniStage` init timing and `process_engine_inputs` flow) + - Connect adjacent `Stage` input/output queues; seed requests to stage 0 + - Poll stage outputs, decode results, call `process_engine_inputs`, then encode and forward to next stage + - Termination & cleanup: distribute shutdown signals; join/terminate subprocesses +- Stage (stage unit) + - Members: `stage_config` (with `engine_args` and `mrs`), `in_q/out_q`, subprocess handle, stats + - API: + - `start_process(in_q, out_q)`: spawn subprocess + - `stop_process()`: graceful exit + - `submit(ipc_payload)`: submit to `in_q` (may use shared memory) + - `try_collect() -> Optional[payload]`: non-blocking get from `out_q` + - `process_engine_inputs(stages, prompts)`: reuse/delegate existing logic + - Subprocess entry `worker_main()`: + - set_stage_gpu_devices (`CUDA_VISIBLE_DEVICES` & logical index mapping) + - build `OmniStageLLM(model, **engine_args)` + - while-loop: + - take first task + up to `mrs.max_batch_size-1` non-blocking; support `max_batch_size>=1` + - `maybe_load_from_ipc` per `engine_inputs` in batch + - `stage_engine.generate(batched_inputs, sampling_params)` (window `-1`) + - write per-request aggregated outputs to `out_q` after `maybe_dump_to_shm` + - log start/finish; report `error` per request on exceptions + +Data flow (sketch) +``` +PipelinedOmniLLM + ├─ Stage[0].start_process(in0, out0) + ├─ Stage[1].start_process(in1, out1) + ├─ ... + └─ Orchestrator loop: + - stage0.submit(encode(engine_inputs)) + - for each stage: + res = stage.try_collect() + outputs = decode(res) + stage.set_engine_outputs(outputs) + if !final: next_inputs = next_stage.process_engine_inputs(stages, original_prompt) + next_stage.submit(encode(next_inputs)) +``` + +Config (preserve & extend) +- `mrs.devices`: + - Comma-separated visibility list (e.g., "2,5,7"): set visible set; default logical index 0 (first GPU) + - Number/string: treat as logical index (0/1-based per implementation) mapping to physical GPUs +- `mrs.max_batch_size`: max requests taken per intake (default 1) +- `shm_threshold_bytes`: objects larger than this go via shared memory (default 64KB; can be injected into stages) + +Kept & strengthened +- Multi-process: each stage in its own process; parallel across stages; serial/batching inside a stage controlled by `max_batch_size` +- Shared memory: large objects via SHM with strict `memoryview -> close -> unlink` +- Device mapping: unified `set_stage_gpu_devices`, support comma lists and logical indices +- Windows: keep `window_size=-1` semantics (downstream triggers only after upstream finishes the request) + +Migration (execution order) +1) Implement enhanced `Stage` (e.g., `entrypoints/stage.py`) with `start_process/stop_process/submit/try_collect/worker_main`. +2) Make `PipelinedOmniLLM` hold `Stage` instances, replacing direct queues and `_stage_worker` with `Stage` APIs. +3) Reuse `pipeline_utils`: device/SHM/encode-decode utilities stay in utils. +4) Regression: single/multi-request, various `max_batch_size`, SHM threshold, device mapping, multi-stage parallelism, error paths, and shutdown flows. diff --git a/examples/offline_inference/qwen_2_5_omni/extract_prompts.py b/examples/offline_inference/qwen_2_5_omni/extract_prompts.py new file mode 100644 index 000000000..a1bc7a78f --- /dev/null +++ b/examples/offline_inference/qwen_2_5_omni/extract_prompts.py @@ -0,0 +1,48 @@ +#!/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 new file mode 100644 index 000000000..510c60f82 --- /dev/null +++ b/examples/offline_inference/qwen_2_5_omni/hf_qwen25_omni_batch_no_template.py @@ -0,0 +1,279 @@ +#!/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 new file mode 100644 index 000000000..9b034521e --- /dev/null +++ b/examples/offline_inference/qwen_2_5_omni/pipelined_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.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: + 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: + pass + 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: + pass + 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/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 06b1f1b4c..a76c1126f 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -1,9 +1,22 @@ """ Stage manager for orchestrating multiple engines in vLLM-omni. + +Enhanced to encapsulate per-stage process lifecycle and worker logic +(device setup, LLM init, batching, shared-memory IPC), while preserving +the original input processing utilities for cross-stage data wiring. """ import importlib -from typing import List, Union +import logging +import multiprocessing as mp +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_dump_to_shm, +) from vllm.inputs import TextPrompt from vllm.v1.engine import EngineCoreOutput @@ -46,6 +59,14 @@ class OmniStage: else: self.final_output_type = None + # Runtime orchestration state (added) + self._in_q: Optional[mp.Queue] = None + self._out_q: Optional[mp.Queue] = None + self._proc: Optional[mp.Process] = None + self._log_file: Optional[str] = None + self._shm_threshold_bytes: int = 65536 + self._logger = logging.getLogger(__name__) + def set_engine(self, engine: LLMEngine) -> None: """Initialize the engine for the stage.""" self.engine = engine @@ -58,6 +79,67 @@ class OmniStage: """Set the engine output for the stage.""" self.engine_outputs = engine_outputs + # ----------------- New Orchestration APIs ----------------- + def attach_queues(self, in_q: mp.Queue, out_q: mp.Queue) -> None: + self._in_q = in_q + self._out_q = out_q + + def init_stage_worker( + self, + model: str, + *, + log_file: Optional[str] = None, + shm_threshold_bytes: int = 65536, + ctx: Optional[mp.context.BaseContext] = None, + batch_timeout: int = 10, + ) -> None: + assert self._in_q is not None and self._out_q is not None, "Queues must be attached before start_process" + self._log_file = log_file + self._shm_threshold_bytes = shm_threshold_bytes + ctx = ctx or mp.get_context("spawn") + # Prepare lightweight dict config for worker + engine_args = _to_dict(self.engine_args) + mrs_cfg = _to_dict(getattr(self.stage_config, "mrs", {})) + stage_payload: Dict[str, Any] = { + "stage_id": self.stage_id, + "engine_args": engine_args, + "mrs": mrs_cfg, + "shm_threshold_bytes": self._shm_threshold_bytes, + } + self._proc = ctx.Process( + target=_stage_worker, + args=(model, stage_payload, self._in_q, self._out_q, self._log_file, batch_timeout), + ) + self._proc.start() + + def stop_stage_worker(self) -> None: + if self._in_q is not None: + try: + self._in_q.put_nowait(None) + except Exception: + pass + if self._proc is not None: + try: + self._proc.join(timeout=5) + except Exception: + pass + if self._proc.is_alive(): + try: + self._proc.terminate() + except Exception: + pass + + def submit(self, payload: Dict[str, Any]) -> None: + assert self._in_q is not None + self._in_q.put(payload) + + def try_collect(self) -> Optional[Dict[str, Any]]: + assert self._out_q is not None + try: + return self._out_q.get_nowait() + except Exception: + return None + def process_engine_inputs( self, stage_list, prompt: Union[OmniTokensPrompt, TextPrompt] = None ) -> List[Union[OmniTokensPrompt, TextPrompt]]: @@ -90,3 +172,253 @@ class OmniStage: return self.custom_process_input_func( stage_list, engine_input_source, prompt ) + + +def _stage_worker( + model: str, + stage_payload: Dict[str, Any], + in_q: mp.Queue, + out_q: mp.Queue, + log_file: Optional[str] = None, + batch_timeout: int = 10, +) -> None: + """Stage worker entry: device setup, LLM init, batching, SHM IPC.""" + import logging as _logging + from vllm_omni.entrypoints.omni_llm import OmniStageLLM # noqa: WPS433 + import queue as _queue + import time as _time + from vllm_omni.entrypoints.pipeline_utils import append_jsonl as _append_jsonl, serialize_obj as _ser + + stage_id = stage_payload["stage_id"] + engine_args = stage_payload.get("engine_args", {}) + mrs_cfg = stage_payload.get("mrs", {}) + shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536)) + + # Per-stage file logger (optional) + try: + if log_file: + stage_log = _logging.getLogger(__name__) + stage_log.setLevel(_logging.DEBUG) + fh = _logging.FileHandler(f"{log_file}.stage{stage_id}.log") + fh.setLevel(_logging.DEBUG) + fh.setFormatter(_logging.Formatter("%(asctime)s [PID:%(process)d] [Stage-%(stage)s] %(levelname)s: %(message)s")) + class _StageFilter(_logging.Filter): + def filter(self, record: _logging.LogRecord) -> bool: + setattr(record, "stage", stage_id) + return True + fh.addFilter(_StageFilter()) + stage_log.addHandler(fh) + except Exception: + pass + + # Stage stats JSONL file + _stats_file = f"{log_file}.stage{stage_id}.stats.jsonl" if log_file else None + + # Aggregates for running average + _agg_total_tokens = 0 + _agg_total_gen_time_ms = 0.0 + + # Device mapping + try: + set_stage_gpu_devices(stage_id, mrs_cfg.get("devices")) + except Exception as e: + _logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e) + + # Init LLM + _logging.getLogger(__name__).debug("[Stage-%s] Initializing engine with args keys=%s", stage_id, list(engine_args.keys())) + stage_engine = OmniStageLLM(model=model, **engine_args) + _logging.getLogger(__name__).debug("[Stage-%s] Engine initialized", stage_id) + # Signal readiness to orchestrator + try: + out_q.put({"type": "stage_ready", "stage_id": stage_id}) + except Exception: + pass + + # Batch processing loop + while True: + task = in_q.get() + _recv_dequeue_ts = _time.time() + if task is None: + _logging.getLogger(__name__).debug("[Stage-%s] Received shutdown signal", stage_id) + break + + max_batch_size = int(mrs_cfg.get("max_batch_size", 1) or 1) + batch_tasks: List[Dict[str, Any]] = [task] + if max_batch_size > 1: + while len(batch_tasks) < max_batch_size: + if not in_q.empty(): + extra = in_q.get(timeout=batch_timeout) + if extra is None: + in_q.put(None) + break + batch_tasks.append(extra) + else: + break + + batch_request_ids: List[Any] = [] + batch_engine_inputs: List[Any] = [] + _rx_bytes_by_rid: Dict[Any, int] = {} + _rx_decode_ms_by_rid: Dict[Any, float] = {} + _in_flight_ms_by_rid: Dict[Any, float] = {} + for t in batch_tasks: + rid = t["request_id"] + try: + sent_ts = float(t.get("sent_ts", None)) if isinstance(t, dict) else None + if sent_ts is not None: + _in_flight_ms_by_rid[rid] = (_recv_dequeue_ts - sent_ts) * 1000.0 + else: + _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 + batch_request_ids.append(rid) + if isinstance(ein, list): + batch_engine_inputs.extend(ein) + elif isinstance(ein, dict): + batch_engine_inputs.append(ein) + else: + _logging.getLogger(__name__).exception("[Stage-%s] Invalid engine input type: %s", stage_id, type(ein)) + sampling_params = batch_tasks[0]["sampling_params"] + _logging.getLogger(__name__).debug("[Stage-%s] Received batch size=%d, request_ids=%s", stage_id, len(batch_tasks), batch_request_ids) + print("--------------------------------", flush=True) + print(f"[Stage-{stage_id}] Received batch size={len(batch_tasks)}, request_ids={batch_request_ids}", flush=True) + print("--------------------------------", flush=True) + try: + gen_outputs: List[Any] = [] + _gen_t0 = _time.time() + for ro in stage_engine.generate(batch_engine_inputs, sampling_params, use_tqdm=False): + gen_outputs.append(ro) + _gen_t1 = _time.time() + _gen_ms = (_gen_t1 - _gen_t0) * 1000.0 + + # Group outputs per request id with fallback + req_to_outputs: Dict[Any, List[Any]] = {rid: [] for rid in batch_request_ids} + unmapped: List[Any] = [] + for ro in gen_outputs: + rid = getattr(ro, "request_id", None) + if rid in req_to_outputs: + req_to_outputs[rid].append(ro) + else: + unmapped.append(ro) + if unmapped: + idx = 0 + for ro in unmapped: + target_rid = batch_request_ids[idx % len(batch_request_ids)] + req_to_outputs[target_rid].append(ro) + 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) + _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: + try: + _avg_tokens_per_s = (_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms) if _agg_total_gen_time_ms > 0 else 0.0 + _append_jsonl(_stats_file, { + "type": "stage_running_avg", + "stage_id": stage_id, + "total_tokens": int(_agg_total_tokens), + "total_gen_time_ms": _agg_total_gen_time_ms, + "avg_tokens_per_s": _avg_tokens_per_s, + }) + _append_jsonl(_stats_file, { + "type": "stage_batch_stats", + "stage_id": stage_id, + "batch_size": len(batch_tasks), + "batch_gen_time_ms": _gen_ms, + "request_ids": list(batch_request_ids), + }) + except Exception: + pass + + # Emit per-request results + for rid in batch_request_ids: + r_outputs = req_to_outputs.get(rid, []) + try: + use_shm, payload = maybe_dump_to_shm(r_outputs, shm_threshold_bytes) + _metrics = { + "num_tokens_out": int(_count_tokens(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)), + "rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)), + } + if use_shm: + out_q.put({ + "request_id": rid, + "stage_id": stage_id, + "engine_outputs_shm": payload, + "metrics": _metrics, + }) + else: + out_q.put({ + "request_id": rid, + "stage_id": stage_id, + "engine_outputs": payload, + "metrics": _metrics, + }) + except Exception: + out_q.put({ + "request_id": rid, + "stage_id": stage_id, + "engine_outputs": r_outputs, + "metrics": { + "num_tokens_out": int(_count_tokens(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)), + "rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)), + }, + }) + _logging.getLogger(__name__).debug("[Stage-%s] Enqueued result for request %s to downstream", stage_id, rid) + except Exception as e: + _logging.getLogger(__name__).exception("[Stage-%s] Failed on batch %s: %s", stage_id, batch_request_ids, e) + for rid in batch_request_ids: + out_q.put({ + "request_id": rid, + "stage_id": stage_id, + "error": str(e), + }) diff --git a/vllm_omni/entrypoints/pipeline_utils.py b/vllm_omni/entrypoints/pipeline_utils.py new file mode 100644 index 000000000..cbf036332 --- /dev/null +++ b/vllm_omni/entrypoints/pipeline_utils.py @@ -0,0 +1,202 @@ +from __future__ import annotations + +from typing import Any, Dict, Optional, Tuple, Union + +import logging +import json +import os +import pickle +from multiprocessing import shared_memory as _shm + +import cloudpickle + +logger = logging.getLogger(__name__) + + +def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) -> None: + """Configure per-stage CUDA visibility and current device. + + Behavior + - Comma-separated string (e.g. "2,5,7"): set CUDA_VISIBLE_DEVICES exactly + to this list; logical index 0 is used as current device. + - Integer or digit-string: treat as logical index (0-based) into the current + CUDA_VISIBLE_DEVICES mapping; map to the physical device, and then set + CUDA_VISIBLE_DEVICES to this single device. + - None/"cpu": keep default visibility. + - Otherwise: set CUDA_VISIBLE_DEVICES to the provided single device string. + """ + try: + selected_physical: Optional[int] = None + logical_idx: Optional[int] = None + + if isinstance(devices, str) and "," in devices: + os.environ["CUDA_VISIBLE_DEVICES"] = devices + toks = [t.strip() for t in devices.split(",") if t.strip() != ""] + if toks: + try: + selected_physical = int(toks[0]) + logger.debug( + "[Stage-%s] Set CUDA_VISIBLE_DEVICES to %s; logical 0 -> physical %s", + stage_id, devices, selected_physical, + ) + except Exception: + selected_physical = None + elif isinstance(devices, (int, str)) and (isinstance(devices, int) or str(devices).isdigit()): + logical_idx = max(0, int(devices)) + vis = os.environ.get("CUDA_VISIBLE_DEVICES") + if vis: + try: + mapping = [int(x) for x in vis.split(",") if x.strip() != ""] + if 0 <= logical_idx < len(mapping): + selected_physical = mapping[logical_idx] + except Exception: + selected_physical = None + if selected_physical is None: + selected_physical = int(logical_idx) + os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical) + logger.debug( + "[Stage-%s] Logical index %d -> physical %s; set CUDA_VISIBLE_DEVICES to single device", + stage_id, logical_idx + 1, selected_physical, + ) + elif devices in (None, "cpu"): + logger.debug("[Stage-%s] Using default device visibility (devices=%s)", stage_id, devices) + else: + selected_physical = int(str(devices)) + os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical) + logger.debug("[Stage-%s] Set CUDA_VISIBLE_DEVICES to single device %s (fallback)", stage_id, selected_physical) + + try: + import torch # noqa: WPS433 + if torch.cuda.is_available(): + try: + torch.cuda.set_device(0) + except Exception: + pass + num = torch.cuda.device_count() + info = [] + for i in range(num): + total = torch.cuda.get_device_properties(i).total_memory + free, _ = torch.cuda.mem_get_info(i) + info.append({ + "idx": i, + "name": torch.cuda.get_device_name(i), + "total": int(total), + "free": int(free), + }) + logger.debug("[Stage-%s] CUDA devices visible=%s info=%s", stage_id, num, info) + except Exception: + pass + except Exception as e: + logger.warning("Failed to interpret devices for stage %s: %s", stage_id, e) + + +def serialize_obj(obj: Any) -> bytes: + """Serialize a Python object to bytes using cloudpickle.""" + return cloudpickle.dumps(obj) + + +def shm_write_bytes(payload: bytes) -> Dict[str, Any]: + """Write bytes into SharedMemory and return meta dict {name,size}. + + Caller should close the segment; the receiver should unlink. + """ + shm = _shm.SharedMemory(create=True, size=len(payload)) + mv = memoryview(shm.buf) + mv[: len(payload)] = payload + del mv + meta = {"name": shm.name, "size": len(payload)} + try: + shm.close() + except Exception: + pass + return meta + + +def shm_read_bytes(meta: Dict[str, Any]) -> bytes: + """Read bytes from SharedMemory by meta {name,size} and cleanup.""" + shm = _shm.SharedMemory(name=meta["name"]) # type: ignore[index] + mv = memoryview(shm.buf) + data = bytes(mv[: meta["size"]]) + del mv + try: + shm.close() + except Exception: + pass + try: + shm.unlink() + except Exception: + pass + return data + + +def _ensure_parent_dir(path: str) -> None: + """Ensure the parent directory for a file path exists (best-effort).""" + try: + parent = os.path.dirname(path) + if parent: + os.makedirs(parent, exist_ok=True) + except Exception: + pass + + +def append_jsonl(path: str, record: Dict[str, Any]) -> None: + """Append a JSON record as one line to a JSONL file (best-effort). + + This is safe to call from multiple processes when each process writes + to a distinct file. For concurrent writes to the same file, OS append + semantics typically suffice, but no additional locking is provided. + """ + try: + _ensure_parent_dir(path) + line = json.dumps(record, ensure_ascii=False) + fd = os.open(path, os.O_APPEND | os.O_CREAT | os.O_WRONLY, 0o644) + with os.fdopen(fd, "a", encoding="utf-8") as f: + f.write(line + "\n") + except Exception: + logger.exception("Failed to append JSONL to %s", path) + + +def maybe_dump_to_shm(obj: Any, threshold: int) -> Tuple[bool, Any]: + """Dump object to SHM if serialized size exceeds threshold. + + Returns (True, meta) when dumped; otherwise (False, original_obj). + """ + payload = serialize_obj(obj) + if len(payload) > threshold: + return True, shm_write_bytes(payload) + return False, obj + + +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.""" + if shm_key in container: + return pickle.loads(shm_read_bytes(container[shm_key])) + return container[obj_key] + + +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). + + When serialized size exceeds threshold, returns {shm_key: {name,size}}; + otherwise returns {obj_key: obj}. + """ + payload: Dict[str, Any] = {} + use_shm, data = maybe_dump_to_shm(obj, threshold) + if use_shm: + payload[shm_key] = data + else: + payload[obj_key] = data + return payload + + +# Convert OmegaConf/objects to plain dicts +def _to_dict(x: Any) -> Dict[str, Any]: + try: + if isinstance(x, dict): + return dict(x) + return OmegaConf.to_container(x, resolve=True) # type: ignore[arg-type] + except Exception: + try: + return dict(x) + except Exception: + return {} diff --git a/vllm_omni/entrypoints/pipelined_omni_llm.py b/vllm_omni/entrypoints/pipelined_omni_llm.py new file mode 100644 index 000000000..b874c2ee1 --- /dev/null +++ b/vllm_omni/entrypoints/pipelined_omni_llm.py @@ -0,0 +1,621 @@ +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.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: + try: + # Avoid duplicate handlers + has_file_handler = any(isinstance(h, logging.FileHandler) for h in logger.handlers) + if not has_file_handler: + fh = logging.FileHandler(self._log_file) + fh.setLevel(logging.DEBUG) + fh.setFormatter(logging.Formatter("%(asctime)s [PID:%(process)d] %(levelname)s: %(message)s")) + logger.addHandler(fh) + logger.setLevel(logging.DEBUG) + except Exception: + pass + + # Orchestrator stats JSONL file + self._stats_file: Optional[str] = None + if self._enable_stats and self._log_file: + try: + self._stats_file = f"{self._log_file}.orchestrator.stats.jsonl" + except Exception: + self._stats_file = None + # Overall stats JSONL file (per-request + summary) + self._overall_stats_file: Optional[str] = None + if self._enable_stats and self._log_file: + try: + self._overall_stats_file = f"{self._log_file}.overall.stats.jsonl" + except Exception: + self._overall_stats_file = None + + 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: + pass + for stage in self.stage_list: + try: + stage.stop_stage_worker() + except Exception: + pass + + def __del__(self) -> None: # best-effort + try: + self.close() + except Exception: + pass + + 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: + final_stage_id_for_e2e = len(self.stage_list) - 1 + # In-memory aggregators for this generate() call + stage_total_time_ms: List[float] = [0.0 for _ in range(num_stages)] + stage_total_tokens: List[int] = [0 for _ in range(num_stages)] + stage_req_counts: List[int] = [0 for _ in range(num_stages)] + transfer_agg: Dict[tuple[int, int], Dict[str, float]] = {} + # Per-edge per-request sender timing to combine with receiver timing later + transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]] = {} + e2e_total_ms: float = 0.0 + e2e_total_tokens: int = 0 + e2e_count: int = 0 + e2e_done: set[int] = set() + # Per-request overall aggregation + per_request: Dict[int, Dict[str, Any]] = {} + sum_per_request_transfer_ms: float = 0.0 + + # Seed stage-0 queue with all requests + logger.debug("[Orchestrator] Seeding %d requests into stage-0", len(request_prompts)) + + 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") + # Aggregate per-request metrics from stage worker if present + try: + _m = result.get("metrics") + if _m is not None: + stage_req_counts[stage_id] += 1 + stage_total_time_ms[stage_id] += float(_m.get("stage_gen_time_ms", 0.0)) + stage_total_tokens[stage_id] += int(_m.get("num_tokens_out", 0)) + # record per-request stage metrics + try: + rid_int = int(req_id) + pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}) + pr_stages = pr["stages"] # type: ignore[index] + pr_stages[stage_id] = { + "stage_gen_time_ms": float(_m.get("stage_gen_time_ms", 0.0)), + "num_tokens_out": int(_m.get("num_tokens_out", 0)), + } + except Exception: + pass + # Also aggregate receiver-side transfer decode time for (stage_id-1 -> stage_id) + try: + if stage_id > 0: + key = (stage_id - 1, stage_id) + agg = transfer_agg.get(key) + if agg is None: + agg = {"sum_bytes": 0.0, "sum_ms": 0.0, "count": 0.0, + "sum_rx_bytes": 0.0, "sum_rx_ms": 0.0, "rx_count": 0.0, + "sum_total_ms": 0.0, "total_count": 0.0} + transfer_agg[key] = agg + rx_b = float(_m.get("rx_transfer_bytes", 0.0)) + rx_ms = float(_m.get("rx_decode_time_ms", 0.0)) + in_flight_ms = float(_m.get("rx_in_flight_time_ms", 0.0)) + agg["sum_rx_bytes"] += rx_b + agg["sum_rx_ms"] += rx_ms + agg["rx_count"] += 1.0 + # If we have sender-side timing stored, compute combined (encode+enqueue + decode) + try: + edge_req = (stage_id - 1, stage_id, int(req_id)) + s = transfer_edge_req.get(edge_req) + if s is not None: + total_ms = float(s.get("tx_ms", 0.0)) + in_flight_ms + rx_ms + agg["sum_total_ms"] += total_ms + agg["total_count"] += 1.0 + # accumulate per-request transfer totals + try: + rid_int = int(req_id) + pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}) + pr["transfers_ms"] = float(pr.get("transfers_ms", 0.0)) + total_ms # type: ignore[index] + pr["transfers_bytes"] = int(pr.get("transfers_bytes", 0)) + int(rx_b) # type: ignore[index] + except Exception: + pass + if getattr(self, "_stats_file", None): + try: + size_b = float(s.get("size_bytes", rx_b)) + _append_jsonl(self._stats_file, { # type: ignore[arg-type] + "type": "transfer_total_stats", + "from_stage": stage_id - 1, + "to_stage": stage_id, + "request_id": req_id, + "size_bytes": int(size_b), + "tx_time_ms": float(s.get("tx_ms", 0.0)), + "in_flight_time_ms": in_flight_ms, + "rx_decode_time_ms": rx_ms, + "total_time_ms": total_ms, + "total_time_per_kb_ms": total_ms / max(size_b / 1024.0, 1e-6) if size_b > 0 else 0.0, + }) + except Exception: + pass + except Exception: + pass + # Emit per-request RX stats to JSONL + if getattr(self, "_stats_file", None): + try: + _append_jsonl(self._stats_file, { # type: ignore[arg-type] + "type": "transfer_rx_stats", + "from_stage": stage_id - 1, + "to_stage": stage_id, + "request_id": req_id, + "rx_bytes": int(rx_b), + "rx_decode_time_ms": rx_ms, + "in_flight_time_ms": in_flight_ms, + "rx_time_per_kb_ms": rx_ms / max(rx_b / 1024.0, 1e-6) if rx_b > 0 else 0.0, + }) + except Exception: + pass + except Exception: + pass + except Exception: + pass + 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: + if stage_id == final_stage_id_for_e2e and req_id not in e2e_done: + _t0 = _req_start_ts.get(req_id) + if _t0 is not None: + _t1 = time.time() + _last_finish_ts = max(_last_finish_ts, _t1) + _e2e_ms = (_t1 - _t0) * 1000.0 + # Count tokens from final stage outputs + def _count_tokens(_ros: List[Any]) -> int: # type: ignore[name-defined] + 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 + _num_tokens = _count_tokens(engine_outputs) + _time_per_token_ms = (_e2e_ms / _num_tokens) if _num_tokens > 0 else 0.0 + # Update E2E aggregators + e2e_total_ms += _e2e_ms + e2e_total_tokens += int(_num_tokens) + e2e_count += 1 + e2e_done.add(req_id) + # Write per-request overall record + try: + rid_int = int(req_id) + pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}) + per_req_record = { + "type": "overall_request", + "request_id": rid_int, + "e2e_time_ms": _e2e_ms, + "num_tokens_out": int(_num_tokens), + "transfers_total_time_ms": float(pr.get("transfers_ms", 0.0)), + "transfers_total_bytes": int(pr.get("transfers_bytes", 0)), + "stages": pr.get("stages", {}), + } + sum_per_request_transfer_ms += float(pr.get("transfers_ms", 0.0)) + if getattr(self, "_overall_stats_file", None): + try: + _append_jsonl(self._overall_stats_file, per_req_record) # type: ignore[arg-type] + except Exception: + pass + except Exception: + pass + if getattr(self, "_stats_file", None): + try: + _append_jsonl(self._stats_file, { # type: ignore[arg-type] + "type": "orchestrator_request_e2e", + "request_id": req_id, + "final_stage_id": stage_id, + "e2e_time_ms": _e2e_ms, + "num_tokens_out": int(_num_tokens), + "e2e_time_per_token_ms": _time_per_token_ms, + }) + except Exception: + pass + except Exception: + pass + + 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 + if self._enable_stats and getattr(self, "_stats_file", None): + try: + _append_jsonl(self._stats_file, { + "type": "transfer_stats", + "from_stage": stage_id, + "to_stage": next_stage_id, + "request_id": req_id, + "size_bytes": int(size_bytes), + "tx_time_ms": tx_ms, + "tx_mbps": (float(size_bytes) * 8.0) / (max(tx_ms, 1e-6) * 1000.0), + "used_shm": bool("engine_inputs_shm" in ipc_payload), + }) + except Exception: + pass + # Update in-memory transfer aggregator + try: + key = (stage_id, next_stage_id) + agg = transfer_agg.get(key) + if agg is None: + agg = {"sum_bytes": 0.0, "sum_ms": 0.0, "count": 0.0, + "sum_rx_bytes": 0.0, "sum_rx_ms": 0.0, "rx_count": 0.0, + "sum_total_ms": 0.0, "total_count": 0.0} + transfer_agg[key] = agg + agg["sum_bytes"] += float(size_bytes) + agg["sum_ms"] += float(tx_ms) + agg["count"] += 1.0 + # Store sender-side timing for per-request combination + transfer_edge_req[(stage_id, next_stage_id, int(req_id))] = { + "tx_ms": float(tx_ms), + "size_bytes": float(size_bytes), + } + except Exception: + pass + except Exception: + 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: + stage_summary: List[Dict[str, Any]] = [] + for sid in range(num_stages): + reqs = stage_req_counts[sid] + tokens = stage_total_tokens[sid] + total_ms = float(stage_total_time_ms[sid]) + avg_req = (total_ms / reqs) if reqs > 0 else 0.0 + avg_tok = (tokens * 1000.0 / total_ms) if total_ms > 0 else 0.0 + stage_summary.append({ + "stage_id": sid, + "requests": int(reqs), + "tokens": int(tokens), + "total_time_ms": total_ms, + "avg_time_per_request_ms": avg_req, + "avg_tokens_per_s": avg_tok, + }) + + transfer_summary: List[Dict[str, Any]] = [] + for (src, dst), agg in transfer_agg.items(): + sum_bytes = float(agg.get("sum_bytes", 0.0)) + sum_ms = float(agg.get("sum_ms", 0.0)) + samples = int(agg.get("count", 0.0)) + tx_mbps = (sum_bytes * 8.0) / (max(sum_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0 + sum_rx_bytes = float(agg.get("sum_rx_bytes", 0.0)) + sum_rx_ms = float(agg.get("sum_rx_ms", 0.0)) + samples_rx = int(agg.get("rx_count", 0.0)) + rx_mbps = (sum_rx_bytes * 8.0) / (max(sum_rx_ms, 1e-6) * 1000.0) if sum_rx_bytes > 0 else 0.0 + sum_total_ms = float(agg.get("sum_total_ms", 0.0)) + samples_total = int(agg.get("total_count", 0.0)) + total_mbps = (sum_bytes * 8.0) / (max(sum_total_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0 + transfer_summary.append({ + "from_stage": src, + "to_stage": dst, + "samples": samples, + "total_bytes": int(sum_bytes), + "total_time_ms": sum_ms, + "tx_mbps": tx_mbps, + "rx_samples": samples_rx, + "rx_total_bytes": int(sum_rx_bytes), + "rx_total_time_ms": sum_rx_ms, + "rx_mbps": rx_mbps, + "total_samples": samples_total, + "total_transfer_time_ms": sum_total_ms, + "total_mbps": total_mbps, + }) + + e2e_avg_req = (e2e_total_ms / e2e_count) if e2e_count > 0 else 0.0 + e2e_avg_tok = (e2e_total_tokens * 1000.0 / e2e_total_ms) if e2e_total_ms > 0 else 0.0 + wall_time_ms = max(0.0, (_last_finish_ts - _wall_start_ts) * 1000.0) + summary: Dict[str, Any] = { + "e2e_requests": int(e2e_count), + # 按你的需求：e2e_total_time_ms 直接使用墙钟时间（非各请求之和） + "e2e_total_time_ms": float(wall_time_ms), + # 额外保留各请求 E2E 之和，供需要时参考 + "e2e_sum_time_ms": float(e2e_total_ms), + "e2e_total_tokens": int(e2e_total_tokens), + "e2e_avg_time_per_request_ms": e2e_avg_req, + "e2e_avg_tokens_per_s": e2e_avg_tok, + "wall_time_ms": wall_time_ms, + "final_stage_id": final_stage_id_for_e2e, + "stages": stage_summary, + "transfers": transfer_summary, + } + logger.info("[Summary] %s", summary) + if self._enable_stats and getattr(self, "_stats_file", None): + try: + _append_jsonl(self._stats_file, {"type": "orchestrator_summary", **summary}) # type: ignore[arg-type] + except Exception: + pass + if self._enable_stats and getattr(self, "_overall_stats_file", None): + try: + _append_jsonl(self._overall_stats_file, {"type": "overall_summary", **summary}) # type: ignore[arg-type] + except Exception: + pass + except Exception: + pass + + 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) + + diff --git a/vllm_omni/model_executor/models/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni.py index 6f9ec9669..e78f77dad 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni.py @@ -269,12 +269,17 @@ class Qwen2_5OmniForConditionalGeneration( # 2) Talker (if codec not provided) if self.model_stage == "talker": + # Mixed-mode support: in一个step内同时支持 Prefill*n 与 Decode*n。 + # 规则： + # - 使用特殊token将 Prefill 片段包裹为 [BOS][PAD...][EOS] + # - Decode 片段则为“单一非特殊token” + # - 若提供了 additional_information（可为按请求分割的 list 或拼接后的 tensor + 形状列表）， + # 则对 Prefill 片段按请求重建 thinker→talker 的输入嵌入； + # - 对 Decode 片段，若提供了 per-request 的 decode 辅助嵌入（可选），则进行叠加；否则保留原嵌入。 if input_ids is None and additional_information is None: input_ids = torch.zeros( - inputs_embeds.shape[0], - dtype=torch.long, - device=inputs_embeds.device, + inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device ) additional_information = {} self.thinker_reply_part = torch.zeros_like(inputs_embeds) @@ -282,63 +287,109 @@ class Qwen2_5OmniForConditionalGeneration( else: is_profile = False + # Ensure we have base embeddings when only ids are provided + if inputs_embeds is None and input_ids is not None: + inputs_embeds = self.talker.get_input_embeddings(input_ids) + + # ------- Request-scoped additional information (no cross-request concat) ------- + request_ids: Optional[list[str]] = kwargs.get("request_ids") # ordered + request_token_spans: Optional[list[tuple[int,int]]] = kwargs.get("request_token_spans") + addi_by_req: Optional[dict] = kwargs.get("additional_information_by_req_id") + runtime_addi = kwargs.get("runtime_additional_information") + + # Normalize runtime_addi into a mapping by request_id for convenience + runtime_addi_by_req: dict[str, dict] = {} + if isinstance(request_ids, list) and isinstance(runtime_addi, list) and len(runtime_addi) == len(request_ids): + for i, rid in enumerate(request_ids): + if isinstance(rid, str) and isinstance(runtime_addi[i], dict): + runtime_addi_by_req[rid] = runtime_addi[i] + elif isinstance(request_ids, list) and isinstance(runtime_addi, dict): + for rid in request_ids: + if isinstance(rid, str) and isinstance(runtime_addi.get(rid), dict): + runtime_addi_by_req[rid] = runtime_addi[rid] + + # Containers to return per-request updates (e.g., thinker_reply_part_per_request) + update_by_req_id: dict[str, dict] = {} + + # ------- Prefill: span_len > 1 ------- if ( - input_ids is not None - and additional_information is not None - and not is_profile + not is_profile + and isinstance(request_ids, list) + and isinstance(request_token_spans, list) + and isinstance(addi_by_req, dict) ): - # read from additional_information dict - thinker_result = None - if additional_information is not None and isinstance( - additional_information, dict - ): - thinker_result = additional_information.get("thinker_result") - prompt_embeds = additional_information.get("prompt_embeds") - prompt_token_ids = additional_information.get("prompt_token_ids") - thinker_output_token_ids = additional_information.get( - "thinker_output_token_ids" - ) - else: - thinker_result = torch.zeros_like(inputs_embeds) - prompt_embeds = torch.zeros_like(inputs_embeds) - prompt_token_ids = torch.zeros( - inputs_embeds.shape[0], - dtype=torch.int64, - device=inputs_embeds.device, - ) - thinker_output_token_ids = torch.zeros( - inputs_embeds.shape[0], - dtype=torch.int64, - device=inputs_embeds.device, - ) + for idx_req, rid in enumerate(request_ids): + s, e = request_token_spans[idx_req] + span_len = int(e) - int(s) + if span_len <= 1: + continue + info = addi_by_req.get(rid, {}) if isinstance(rid, str) else {} + if not isinstance(info, dict): + info = {} + pe = info.get("prompt_embeds") # Tensor [P,H] + tr = info.get("thinker_result") # Tensor [K,H] + ptoks = info.get("prompt_token_ids") # list[int] + otoks = info.get("thinker_output_token_ids") # list[int] - if thinker_result is None: - thinker_result = torch.zeros_like(inputs_embeds) - self.thinker_reply_part = thinker_result.squeeze(0) - if self.thinker_reply_part.shape[1] > 1: - self.thinker_reply_part = self.thinker_reply_part[1:, :] - input_ids, inputs_embeds = self._thinker_to_talker_prefill( - voice_type=voice_type, - output_prompt_embeds=thinker_result, - output_token_ids=thinker_output_token_ids, - thinker_prompt_embeds=prompt_embeds, - prompt_token_ids=prompt_token_ids, - ) - elif not is_profile: - input_ids, inputs_embeds = self._thinker_to_talker_decode_one_step( - output_prompt_embeds=( - self.thinker_reply_part[:1] - if self.thinker_reply_part.shape[0] >= 1 - else torch.zeros(1, self.thinker_reply_part.shape[1]) - .to(self._module_device(self.model)) - .to(torch.bfloat16) - + (-1.25 * 2 ** (-123)) - ), - output_token_ids=input_ids, - ) + if not isinstance(pe, torch.Tensor): + pe = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model)) + if not isinstance(tr, torch.Tensor): + tr = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model)) + if not isinstance(ptoks, (list, torch.Tensor)): + ptoks = [] + if not isinstance(otoks, (list, torch.Tensor)): + otoks = [] - if self.thinker_reply_part.shape[0] >= 1: - self.thinker_reply_part = self.thinker_reply_part[1:, :] + req_input_ids, req_embeds = self._thinker_to_talker_prefill( + voice_type=voice_type, + output_prompt_embeds=tr.to(inputs_embeds.dtype).to(self._module_device(self.model)), + output_token_ids=otoks, + thinker_prompt_embeds=pe.to(inputs_embeds.dtype).to(self._module_device(self.model)), + prompt_token_ids=ptoks, + ) + seg_len = min(span_len, req_embeds.shape[0]) + inputs_embeds[s : s + seg_len] = req_embeds[:seg_len] + if isinstance(req_input_ids, torch.Tensor) and req_input_ids.numel() == seg_len: + input_ids[s : s + seg_len] = req_input_ids + + # Prepare per-request reply queue for subsequent decode: drop first row + if tr.ndim == 2 and tr.shape[0] > 0: + update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = tr[1:].detach().to("cpu").contiguous() + + # ------- Decode: span_len == 1 ------- + if not is_profile and isinstance(request_ids, list) and isinstance(request_token_spans, list): + for idx_req, rid in enumerate(request_ids): + s, e = request_token_spans[idx_req] + if (int(e) - int(s)) != 1: + continue + # choose step vector in priority order + step_vec = None + # A) runtime queue + q = None + if isinstance(rid, str): + q = runtime_addi_by_req.get(rid, {}).get("thinker_reply_part_per_request") + if isinstance(q, torch.Tensor) and q.numel() > 0: + step_vec = q[0:1] + new_q = q[1:].detach().to("cpu").contiguous() + update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = new_q + else: + # B) per-request provided decode vector (optional) + info = addi_by_req.get(rid, {}) if isinstance(addi_by_req, dict) else {} + dv = info.get("decode_output_prompt_embeds") if isinstance(info, dict) else None + if isinstance(dv, torch.Tensor) and dv.numel() > 0: + step_vec = dv[0:1] if dv.ndim == 2 else dv.view(1, -1) + elif hasattr(self, "thinker_reply_part") and isinstance(self.thinker_reply_part, torch.Tensor) and self.thinker_reply_part.numel() > 0: + # C) fallback shared pool + step_vec = self.thinker_reply_part[0:1] + self.thinker_reply_part = self.thinker_reply_part[1:] + + if isinstance(step_vec, torch.Tensor) and step_vec.numel() > 0: + one_id = input_ids[s : s + 1] + _, one_embed = self._thinker_to_talker_decode_one_step( + output_prompt_embeds=step_vec.to(inputs_embeds.dtype).to(self._module_device(self.model)), + output_token_ids=one_id, + ) + inputs_embeds[s] = one_embed[0] with torch.inference_mode(): talker_hidden = self.talker( @@ -346,8 +397,19 @@ class Qwen2_5OmniForConditionalGeneration( positions=positions[0], inputs_embeds=inputs_embeds, ) + multimodal_outputs: dict = None + # Return updates if any + if update_by_req_id: + multimodal_outputs = {"additional_information_update_by_req_id": update_by_req_id} - return OmniOutput(text_hidden_states=talker_hidden, multimodal_outputs=None) + if sampling_metadata is not None: + # the padding token id is set to text model's pad token id, which do not match with the talker model's word embedding size + sampling_metadata.prompt_token_ids[sampling_metadata.prompt_token_ids == 152064] = 8448 + + return OmniOutput( + text_hidden_states=talker_hidden, + multimodal_outputs=multimodal_outputs, + ) if self.model_stage == "code2wav": code = ( @@ -359,8 +421,12 @@ class Qwen2_5OmniForConditionalGeneration( device=inputs_embeds.device, ) ) + + code = code[:-1] if code[-1] == TALKER_CODEC_EOS_TOKEN_ID else code + code = code[1:] if code[0] == TALKER_CODEC_BOS_TOKEN_ID else code + audio_tensor = self.generate_audio( - code[:-1] if code[-1] == TALKER_CODEC_EOS_TOKEN_ID else code, voice_type + code, voice_type ) return OmniOutput( text_hidden_states=None, multimodal_outputs={"audio": audio_tensor} 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 337194f9d..2f69b3080 100644 --- a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml @@ -1,36 +1,50 @@ # stage config for running qwen2.5-omni with architecture of OmniLLM. stage_args: - stage_id: 0 + mrs: + process: true # Run this stage in a separate process + devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device) + max_batch_size: 1 engine_args: model_stage: thinker model_arch: Qwen2_5OmniForConditionalGeneration worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler - gpu_memory_utilization: 0.32 + gpu_memory_utilization: 0.8 enforce_eager: true # need to discuss trust_remote_code: true engine_output_type: latent # change the param name,such as pooling_output + enable_prefix_caching: false final_output: true final_output_type: text - stage_id: 1 + mrs: + process: true + devices: "1" + max_batch_size: 3 engine_args: model_stage: talker model_arch: Qwen2_5OmniForConditionalGeneration worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler - gpu_memory_utilization: 0.32 + gpu_memory_utilization: 0.8 enforce_eager: true trust_remote_code: true + enable_prefix_caching: false engine_output_type: latent engine_input_source: [0] custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker - stage_id: 2 + mrs: + process: true + devices: "0" # 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.3 + gpu_memory_utilization: 0.15 enforce_eager: true trust_remote_code: true enable_prefix_caching: false @@ -38,3 +52,17 @@ stage_args: engine_input_source: [1] final_output: true final_output_type: audio + +# Top-level MRS config (concise): default windows and stage edges +mrs: + enabled: true + defaults: + window_size: -1 # Simplified: trigger downstream only after full upstream completion + max_inflight: 1 # Simplified: process serially within each stage + edges: + - from: 0 # thinker → talker: trigger only after receiving full input (-1) + to: 1 + window_size: -1 + - from: 1 # talker → code2wav: trigger only after receiving full input (-1) + to: 2 + window_size: -1 diff --git a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py index 4dc10d4fd..ca5e02ba6 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py @@ -4,8 +4,9 @@ import torch from vllm.inputs import TextPrompt from vllm_omni.inputs.data import OmniTokensPrompt - - +TALKER_CODEC_PAD_TOKEN_ID = 8292 +TALKER_CODEC_START_TOKEN_ID = 8293 +TALKER_CODEC_END_TOKEN_ID = 8294 def thinker2talker( stage_list, engine_input_source, prompt: Union[OmniTokensPrompt, TextPrompt] = None ): @@ -33,8 +34,7 @@ def thinker2talker( ) talker_inputs.append( OmniTokensPrompt( - prompt_token_ids=[0] - * (len(prompt_token_ids) + 2), # add 2 for codec pad and start token + prompt_token_ids=[TALKER_CODEC_START_TOKEN_ID] + [TALKER_CODEC_PAD_TOKEN_ID] * (len(prompt_token_ids)) + [TALKER_CODEC_END_TOKEN_ID], additional_information={ "thinker_result": thinker_hidden_states[prompt_token_ids_len:].to( torch.float32 @@ -44,6 +44,8 @@ def thinker2talker( ), "prompt_token_ids": prompt_token_ids, "thinker_output_token_ids": thinker_output_ids, + "thinker_result_shape": list(thinker_hidden_states[prompt_token_ids_len:].shape), + "prompt_embeds_shape": list(thinker_hidden_states[:prompt_token_ids_len].shape), }, multi_modal_data=( multi_modal_data[thinker_output.request_id] diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 20a98e55a..2b1ce3509 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -99,9 +99,11 @@ class GPUARModelRunner(OmniGPUModelRunner): # TODO(woosuk): Avoid the copy. Optimize. self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled) - # Reset per-step additional information collector + # Reset per-step additional information collector (deprecated concat path) if hasattr(self, "_forward_additional_information"): self._forward_additional_information = None + # New: per-request additional information for this step + per_req_additional_information: dict[str, dict] = {} # Overlay custom prompt_embeds per request for the prompt portion; # collect additional_information (tensor/list) for prefill portion only @@ -126,50 +128,27 @@ class GPUARModelRunner(OmniGPUModelRunner): self.inputs_embeds[start_offset : start_offset + overlay_len].copy_( src ) - # For additional_information: handle arbitrary keys + # Build per-request additional information (no cross-request concat) if addi_cpu is not None and isinstance(addi_cpu, dict): - # Lazy init collector dict - if ( - not hasattr(self, "_forward_additional_information") - or self._forward_additional_information is None - ): - self._forward_additional_information = {} - # Process tensors (slice by scheduled prompt range) and - # lists (append per-request) + req_info: dict[str, object] = {} for k, v in addi_cpu.items(): if isinstance(v, torch.Tensor): - # Slice along token dimension for prefill part - try: - seg = v[ - num_computed_tokens : num_computed_tokens - + overlay_len - ].to( - dtype=self.dtype, - device=self.device, - non_blocking=True, - ) - except Exception: - # Fallback: move whole tensor if slicing fails - seg = v.to( - dtype=self.dtype, - device=self.device, - non_blocking=True, - ) - prev_val = self._forward_additional_information.get(k) - self._forward_additional_information[k] = ( - torch.cat([prev_val, seg], dim=0) - if isinstance(prev_val, torch.Tensor) - else seg.clone() - ) - elif isinstance(v, list): - prev_val = self._forward_additional_information.get(k) - if prev_val is None: - self._forward_additional_information[k] = [v] - elif isinstance(prev_val, list): - self._forward_additional_information[k].append(v) + # For prefill tokens, pass only the scheduled slice; for decode or no scheduled tokens, pass whole tensor + if overlay_len > 0: + try: + seg = v[ + num_computed_tokens : num_computed_tokens + overlay_len + ].detach().to("cpu").contiguous() + except Exception: + seg = v.detach().to("cpu").contiguous() + req_info[k] = seg else: - # Mixed types: wrap existing into list - self._forward_additional_information[k] = [prev_val, v] + req_info[k] = v.detach().to("cpu").contiguous() + elif isinstance(v, list): + req_info[k] = v + else: + req_info[k] = v + per_req_additional_information[req_id] = req_info input_ids = self.input_ids.gpu[:num_input_tokens] inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] @@ -395,15 +374,33 @@ class GPUARModelRunner(OmniGPUModelRunner): ): model_kwargs_extra = {} - # Only pass additional_information for the prefill part - if ( - hasattr(self, "_forward_additional_information") - and self._forward_additional_information is not None - and isinstance(self._forward_additional_information, dict) - ): - model_kwargs_extra["additional_information"] = ( - self._forward_additional_information + # Pass per-request additional information map for this step (no concat) + if per_req_additional_information: + model_kwargs_extra["additional_information_by_req_id"] = per_req_additional_information + # 始终传递每请求的运行期 additional_information（持久化在 request state） + try: + per_req_runtime_info = [] + for req_id in self.input_batch.req_ids: + req_state = self.requests.get(req_id) + info = ( + getattr(req_state, "additional_information_cpu", None) + if req_state is not None + else None + ) + per_req_runtime_info.append(info if isinstance(info, dict) else {}) + model_kwargs_extra["runtime_additional_information"] = ( + per_req_runtime_info ) + model_kwargs_extra["request_ids"] = self.input_batch.req_ids + # 传递每个请求在本步扁平序列中的 token 区间，便于模型按请求映射 decode/prefill + req_token_spans = [] + for req_index in range(len(self.input_batch.req_ids)): + start_offset = int(self.query_start_loc_cpu[req_index]) + sched_tokens = int(num_scheduled_tokens_np[req_index]) + req_token_spans.append((start_offset, start_offset + sched_tokens)) + model_kwargs_extra["request_token_spans"] = req_token_spans + except Exception: + pass model_output = self.model( input_ids=input_ids, positions=positions, @@ -428,8 +425,42 @@ class GPUARModelRunner(OmniGPUModelRunner): hidden_states, multimodal_outputs = self.extract_multimodal_outputs( hidden_states ) - + # 模型侧可能返回 per-request 的 additional_information 更新（与模型无关的通用通道）。 + # 约定：multimodal_outputs["additional_information_update"] 为按 batch 顺序的 list[dict]， + # runner 将其合并到对应 request 的 additional_information_cpu，供后续 decode 使用。 + try: + if ( + isinstance(multimodal_outputs, dict) + and ( + "additional_information_update" in multimodal_outputs + or "additional_information_update_by_req_id" in multimodal_outputs + ) + ): + # 方案 A：按 batch 顺序的 list[dict] + updates_list = multimodal_outputs.get( + "additional_information_update" + ) + if isinstance(updates_list, list): + for idx, upd in enumerate(updates_list): + if not isinstance(upd, dict) or idx >= len(self.input_batch.req_ids): + continue + req_id = self.input_batch.req_ids[idx] + self._merge_additional_information_update(req_id, upd) + # 方案 B：按 req_id 的 dict[str, dict] + updates_map = multimodal_outputs.get( + "additional_information_update_by_req_id" + ) + if isinstance(updates_map, dict): + for req_id, upd in updates_map.items(): + if not isinstance(upd, dict): + continue + if req_id not in self.requests: + continue + self._merge_additional_information_update(req_id, upd) + except Exception as e: + logger.error(f"Error merging for requests:{self.input_batch.req_ids} additional information update: {e}, with the multimodal_outputs as {multimodal_outputs}") if not self.broadcast_pp_output: + # Common case. if not get_pp_group().is_last_rank: # Return the intermediate tensors. @@ -589,3 +620,26 @@ class GPUARModelRunner(OmniGPUModelRunner): invalid_req_indices=invalid_req_indices, async_output_copy_stream=self.async_output_copy_stream, ) + + def _merge_additional_information_update(self, req_id: str, upd: dict) -> None: + req_state = self.requests.get(req_id) + if req_state is None: + return + existing = getattr(req_state, "additional_information_cpu", {}) + if not isinstance(existing, dict): + existing = {} + merged = dict(existing) + for k, v in upd.items(): + if isinstance(v, torch.Tensor): + merged[k] = v.detach().to("cpu").contiguous() + elif isinstance(v, list): + new_list = [] + for item in v: + if isinstance(item, torch.Tensor): + new_list.append(item.detach().to("cpu").contiguous()) + else: + new_list.append(item) + merged[k] = new_list + else: + merged[k] = v + setattr(req_state, "additional_information_cpu", merged) \ No newline at end of file From 3c2ad1bf1202ca619ed8b26a9adc9a3e300670f2 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Fri, 7 Nov 2025 16:25:48 +0800 Subject: [PATCH 02/19] debug ar model runner --- vllm_omni/worker/gpu_ar_model_runner.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 2b1ce3509..4c115bfc8 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -65,6 +65,7 @@ class GPUARModelRunner(OmniGPUModelRunner): torch.Tensor, Optional[IntermediateTensors], dict[str, Any], + Optional[dict[str, dict]], ]: num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens @@ -79,6 +80,7 @@ class GPUARModelRunner(OmniGPUModelRunner): # _prepare_inputs may reorder the batch, so we must gather multi # modal outputs after that to ensure the correct order + per_req_additional_information: Optional[dict[str, dict]] = None if ( self.supports_mm_inputs and get_pp_group().is_first_rank @@ -103,7 +105,7 @@ class GPUARModelRunner(OmniGPUModelRunner): if hasattr(self, "_forward_additional_information"): self._forward_additional_information = None # New: per-request additional information for this step - per_req_additional_information: dict[str, dict] = {} + per_req_additional_information = {} # Overlay custom prompt_embeds per request for the prompt portion; # collect additional_information (tensor/list) for prefill portion only @@ -219,6 +221,7 @@ class GPUARModelRunner(OmniGPUModelRunner): positions, intermediate_tensors, model_kwargs, + per_req_additional_information, ) @torch.inference_mode() @@ -334,6 +337,7 @@ class GPUARModelRunner(OmniGPUModelRunner): positions, intermediate_tensors, model_kwargs, + per_req_additional_information, ) = self._preprocess( scheduler_output, num_scheduled_tokens_np, From c5631ca753551e382bf95a4620ab318fa428d2fe Mon Sep 17 00:00:00 2001 From: tzhouam Date: Fri, 7 Nov 2025 18:05:00 +0800 Subject: [PATCH 03/19] Enhance error handling and logging in pipeline and stage processing. Added detailed logging for prompt loading failures, file writing issues, and stage shutdown processes. Introduced a new logging utility for orchestrator metrics and streamlined stats handling in the PipelinedOmniLLM class. --- .../qwen_2_5_omni/pipelined_end2end.py | 9 +- vllm_omni/entrypoints/log_utils.py | 548 ++++++++++++++++++ vllm_omni/entrypoints/omni_stage.py | 60 +- vllm_omni/entrypoints/pipeline_utils.py | 18 +- vllm_omni/entrypoints/pipelined_omni_llm.py | 349 ++--------- vllm_omni/worker/gpu_ar_model_runner.py | 18 +- 6 files changed, 657 insertions(+), 345 deletions(-) create mode 100644 vllm_omni/entrypoints/log_utils.py diff --git a/examples/offline_inference/qwen_2_5_omni/pipelined_end2end.py b/examples/offline_inference/qwen_2_5_omni/pipelined_end2end.py index 9b034521e..2e394e5f0 100644 --- a/examples/offline_inference/qwen_2_5_omni/pipelined_end2end.py +++ b/examples/offline_inference/qwen_2_5_omni/pipelined_end2end.py @@ -162,6 +162,7 @@ def main(): 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: @@ -233,8 +234,8 @@ def main(): try: with open(out_txt, "w", encoding="utf-8") as f: f.writelines(lines) - except Exception: - pass + 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: @@ -255,8 +256,8 @@ def main(): try: with open(out_txt, "w", encoding="utf-8") as f: f.writelines(lines) - except Exception: - pass + 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}") diff --git a/vllm_omni/entrypoints/log_utils.py b/vllm_omni/entrypoints/log_utils.py new file mode 100644 index 000000000..6a17c24f0 --- /dev/null +++ b/vllm_omni/entrypoints/log_utils.py @@ -0,0 +1,548 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Optional + +import logging +import os +import time + +from vllm_omni.entrypoints.pipeline_utils import append_jsonl as _append_jsonl + + +def remove_old_logs(log_file: Optional[str], num_stages: int) -> None: + try: + if not log_file: + return + try: + if os.path.exists(log_file): + os.remove(log_file) + except Exception: + pass + # Per-stage logs and stats + for sid in range(num_stages): + try: + p = f"{log_file}.stage{sid}.log" + if os.path.exists(p): + os.remove(p) + except Exception: + pass + try: + p = f"{log_file}.stage{sid}.stats.jsonl" + if os.path.exists(p): + os.remove(p) + except Exception: + pass + # Orchestrator stats files + try: + p = f"{log_file}.orchestrator.stats.jsonl" + if os.path.exists(p): + os.remove(p) + except Exception: + pass + try: + p = f"{log_file}.overall.stats.jsonl" + if os.path.exists(p): + os.remove(p) + except Exception: + pass + except Exception: + pass + + +def configure_orchestrator_logger(logger: logging.Logger, log_file: Optional[str]) -> None: + try: + if not log_file: + return + has_file_handler = any(isinstance(h, logging.FileHandler) for h in logger.handlers) + if not has_file_handler: + fh = logging.FileHandler(log_file) + fh.setLevel(logging.DEBUG) + fh.setFormatter(logging.Formatter("%(asctime)s [PID:%(process)d] %(levelname)s: %(message)s")) + logger.addHandler(fh) + logger.setLevel(logging.DEBUG) + except Exception: + pass + + +def init_stats_paths(enable_stats: bool, log_file: Optional[str]) -> tuple[Optional[str], Optional[str]]: + stats_file: Optional[str] = None + overall_file: Optional[str] = None + try: + if enable_stats and log_file: + stats_file = f"{log_file}.orchestrator.stats.jsonl" + overall_file = f"{log_file}.overall.stats.jsonl" + except Exception: + stats_file = None + overall_file = None + return stats_file, overall_file + + +def _safe_append_jsonl(path: Optional[str], record: Dict[str, Any]) -> None: + if not path: + return + try: + _append_jsonl(path, record) # type: ignore[arg-type] + except Exception: + pass + + +def log_transfer_tx(stats_file: Optional[str], from_stage: int, to_stage: int, request_id: Any, + size_bytes: int, tx_time_ms: float, used_shm: bool) -> None: + _safe_append_jsonl(stats_file, { + "type": "transfer_stats", + "from_stage": from_stage, + "to_stage": to_stage, + "request_id": request_id, + "size_bytes": int(size_bytes), + "tx_time_ms": float(tx_time_ms), + "tx_mbps": (float(size_bytes) * 8.0) / (max(tx_time_ms, 1e-6) * 1000.0), + "used_shm": bool(used_shm), + }) + + +def log_transfer_rx(stats_file: Optional[str], from_stage: int, to_stage: int, request_id: Any, + rx_bytes: int, rx_decode_time_ms: float, in_flight_time_ms: float) -> None: + _safe_append_jsonl(stats_file, { + "type": "transfer_rx_stats", + "from_stage": from_stage, + "to_stage": to_stage, + "request_id": request_id, + "rx_bytes": int(rx_bytes), + "rx_decode_time_ms": float(rx_decode_time_ms), + "in_flight_time_ms": float(in_flight_time_ms), + "rx_time_per_kb_ms": (float(rx_decode_time_ms) / max(float(rx_bytes) / 1024.0, 1e-6)) if rx_bytes > 0 else 0.0, + }) + + +def log_transfer_total(stats_file: Optional[str], from_stage: int, to_stage: int, request_id: Any, + size_bytes: int, tx_time_ms: float, in_flight_time_ms: float, + rx_decode_time_ms: float, total_time_ms: float) -> None: + _safe_append_jsonl(stats_file, { + "type": "transfer_total_stats", + "from_stage": from_stage, + "to_stage": to_stage, + "request_id": request_id, + "size_bytes": int(size_bytes), + "tx_time_ms": float(tx_time_ms), + "in_flight_time_ms": float(in_flight_time_ms), + "rx_decode_time_ms": float(rx_decode_time_ms), + "total_time_ms": float(total_time_ms), + "total_time_per_kb_ms": float(total_time_ms) / max(float(size_bytes) / 1024.0, 1e-6) if size_bytes > 0 else 0.0, + }) + + +def log_orchestrator_e2e(stats_file: Optional[str], request_id: Any, final_stage_id: int, + e2e_time_ms: float, num_tokens_out: int, e2e_time_per_token_ms: float) -> None: + _safe_append_jsonl(stats_file, { + "type": "orchestrator_request_e2e", + "request_id": request_id, + "final_stage_id": final_stage_id, + "e2e_time_ms": float(e2e_time_ms), + "num_tokens_out": int(num_tokens_out), + "e2e_time_per_token_ms": float(e2e_time_per_token_ms), + }) + + +def log_orchestrator_summary(stats_file: Optional[str], summary: Dict[str, Any]) -> None: + _safe_append_jsonl(stats_file, {"type": "orchestrator_summary", **summary}) + + +def log_overall_summary(overall_stats_file: Optional[str], summary: Dict[str, Any]) -> None: + _safe_append_jsonl(overall_stats_file, {"type": "overall_summary", **summary}) + + +def log_overall_record(overall_stats_file: Optional[str], record: Dict[str, Any]) -> None: + _safe_append_jsonl(overall_stats_file, record) + + +def log_stage_request_stats(stats_file: Optional[str], stage_id: int, request_id: Any, batch_size: int, + num_tokens_out: int, stage_gen_time_ms: float, tokens_per_s: float, + rx_transfer_bytes: int, rx_decode_time_ms: float, rx_mbps: float) -> None: + _safe_append_jsonl(stats_file, { + "type": "stage_request_stats", + "stage_id": stage_id, + "request_id": request_id, + "batch_size": int(batch_size), + "num_tokens_out": int(num_tokens_out), + "stage_gen_time_ms": float(stage_gen_time_ms), + "tokens_per_s": float(tokens_per_s), + "rx_transfer_bytes": int(rx_transfer_bytes), + "rx_decode_time_ms": float(rx_decode_time_ms), + "rx_mbps": float(rx_mbps), + }) + + +def log_stage_running_avg(stats_file: Optional[str], stage_id: int, total_tokens: int, + total_gen_time_ms: float, avg_tokens_per_s: float) -> None: + _safe_append_jsonl(stats_file, { + "type": "stage_running_avg", + "stage_id": stage_id, + "total_tokens": int(total_tokens), + "total_gen_time_ms": float(total_gen_time_ms), + "avg_tokens_per_s": float(avg_tokens_per_s), + }) + + +def log_stage_batch_stats(stats_file: Optional[str], stage_id: int, batch_size: int, + batch_gen_time_ms: float, request_ids: List[Any]) -> None: + _safe_append_jsonl(stats_file, { + "type": "stage_batch_stats", + "stage_id": stage_id, + "batch_size": int(batch_size), + "batch_gen_time_ms": float(batch_gen_time_ms), + "request_ids": list(request_ids), + }) + + +# ----------------- Aggregation helpers for orchestrator ----------------- + +def record_stage_metrics(per_request: Dict[int, Dict[str, Any]], + stage_req_counts: List[int], + stage_total_time_ms: List[float], + stage_total_tokens: List[int], + stage_id: int, + req_id: Any, + metrics: Dict[str, Any]) -> None: + try: + stage_req_counts[stage_id] += 1 + stage_total_tokens[stage_id] += int(metrics.get("num_tokens_out", 0)) + rid_int = int(req_id) + pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}) + pr_stages = pr["stages"] # type: ignore[index] + pr_stages[stage_id] = { + "stage_gen_time_ms": float(metrics.get("stage_gen_time_ms", 0.0)), + "num_tokens_out": int(metrics.get("num_tokens_out", 0)), + } + except Exception: + pass + + +def aggregate_rx_and_maybe_total(transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]], + transfer_agg: Dict[tuple[int, int], Dict[str, float]], + per_request: Dict[int, Dict[str, Any]], + stage_id: int, + req_id: Any, + rx_bytes: float, + rx_ms: float, + in_flight_ms: float) -> Optional[tuple[int, float, float]]: + try: + # Update RX aggregates for (stage_id-1 -> stage_id) + if stage_id > 0: + key = (stage_id - 1, stage_id) + agg = transfer_agg.get(key) + if agg is None: + agg = {"sum_bytes": 0.0, "sum_ms": 0.0, "count": 0.0, + "sum_rx_bytes": 0.0, "sum_rx_ms": 0.0, "rx_count": 0.0, + "sum_total_ms": 0.0, "total_count": 0.0} + transfer_agg[key] = agg + agg["sum_rx_bytes"] += float(rx_bytes) + agg["sum_rx_ms"] += float(rx_ms) + agg["rx_count"] += 1.0 + + # Try combine with sender-side timing if present + try: + rid_int = int(req_id) + except Exception: + rid_int = req_id # best effort + s = transfer_edge_req.get((stage_id - 1, stage_id, rid_int)) + if s is None: + return None + tx_ms = float(s.get("tx_ms", 0.0)) + size_b = float(s.get("size_bytes", rx_bytes)) + total_ms = tx_ms + float(in_flight_ms) + float(rx_ms) + agg["sum_total_ms"] += total_ms + agg["total_count"] += 1.0 + # accumulate per-request transfer totals + try: + pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}) + pr["transfers_ms"] = float(pr.get("transfers_ms", 0.0)) + total_ms # type: ignore[index] + pr["transfers_bytes"] = int(pr.get("transfers_bytes", 0)) + int(rx_bytes) # type: ignore[index] + except Exception: + pass + return int(size_b), float(tx_ms), float(total_ms) + return None + except Exception: + return None + + +def record_sender_transfer_agg(transfer_agg: Dict[tuple[int, int], Dict[str, float]], + transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]], + from_stage: int, + to_stage: int, + req_id: Any, + size_bytes: int, + tx_ms: float) -> None: + try: + key = (from_stage, to_stage) + agg = transfer_agg.get(key) + if agg is None: + agg = {"sum_bytes": 0.0, "sum_ms": 0.0, "count": 0.0, + "sum_rx_bytes": 0.0, "sum_rx_ms": 0.0, "rx_count": 0.0, + "sum_total_ms": 0.0, "total_count": 0.0} + transfer_agg[key] = agg + agg["sum_bytes"] += float(size_bytes) + agg["sum_ms"] += float(tx_ms) + agg["count"] += 1.0 + # Store sender-side timing for per-request combination + try: + rid_int = int(req_id) + except Exception: + rid_int = req_id + transfer_edge_req[(from_stage, to_stage, rid_int)] = { + "tx_ms": float(tx_ms), + "size_bytes": float(size_bytes), + } + except Exception: + pass + + +def count_tokens_from_outputs(engine_outputs: List[Any]) -> int: + total = 0 + for _ro in engine_outputs: + 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 + + +def build_stage_summary(stage_req_counts: List[int], + stage_total_tokens: List[int], + stage_total_time_ms: List[float]) -> List[Dict[str, Any]]: + summary: List[Dict[str, Any]] = [] + for sid in range(len(stage_req_counts)): + reqs = stage_req_counts[sid] + tokens = stage_total_tokens[sid] + total_ms = float(stage_total_time_ms[sid]) + avg_req = (total_ms / reqs) if reqs > 0 else 0.0 + avg_tok = (tokens * 1000.0 / total_ms) if total_ms > 0 else 0.0 + summary.append({ + "stage_id": sid, + "requests": int(reqs), + "tokens": int(tokens), + "total_time_ms": total_ms, + "avg_time_per_request_ms": avg_req, + "avg_tokens_per_s": avg_tok, + }) + return summary + + +def build_transfer_summary(transfer_agg: Dict[tuple[int, int], Dict[str, float]]) -> List[Dict[str, Any]]: + summary: List[Dict[str, Any]] = [] + for (src, dst), agg in transfer_agg.items(): + sum_bytes = float(agg.get("sum_bytes", 0.0)) + sum_ms = float(agg.get("sum_ms", 0.0)) + samples = int(agg.get("count", 0.0)) + tx_mbps = (sum_bytes * 8.0) / (max(sum_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0 + sum_rx_bytes = float(agg.get("sum_rx_bytes", 0.0)) + sum_rx_ms = float(agg.get("sum_rx_ms", 0.0)) + samples_rx = int(agg.get("rx_count", 0.0)) + rx_mbps = (sum_rx_bytes * 8.0) / (max(sum_rx_ms, 1e-6) * 1000.0) if sum_rx_bytes > 0 else 0.0 + sum_total_ms = float(agg.get("sum_total_ms", 0.0)) + samples_total = int(agg.get("total_count", 0.0)) + total_mbps = (sum_bytes * 8.0) / (max(sum_total_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0 + summary.append({ + "from_stage": src, + "to_stage": dst, + "samples": samples, + "total_bytes": int(sum_bytes), + "total_time_ms": sum_ms, + "tx_mbps": tx_mbps, + "rx_samples": samples_rx, + "rx_total_bytes": int(sum_rx_bytes), + "rx_total_time_ms": sum_rx_ms, + "rx_mbps": rx_mbps, + "total_samples": samples_total, + "total_transfer_time_ms": sum_total_ms, + "total_mbps": total_mbps, + }) + return summary + + +class OrchestratorMetrics: + def __init__(self, num_stages: int, enable_stats: bool, + stats_file: Optional[str], overall_stats_file: Optional[str], + wall_start_ts: float) -> None: + self.num_stages = int(num_stages) + self.enable_stats = bool(enable_stats) + self.stats_file = stats_file + self.overall_stats_file = overall_stats_file + self.stage_total_time_ms: List[float] = [0.0 for _ in range(self.num_stages)] + self.stage_total_tokens: List[int] = [0 for _ in range(self.num_stages)] + self.stage_req_counts: List[int] = [0 for _ in range(self.num_stages)] + self.transfer_agg: Dict[tuple[int, int], Dict[str, float]] = {} + self.transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]] = {} + self.e2e_total_ms: float = 0.0 + self.e2e_total_tokens: int = 0 + self.e2e_count: int = 0 + self.e2e_done: set[int] = set() + self.per_request: Dict[int, Dict[str, Any]] = {} + self.sum_per_request_transfer_ms: float = 0.0 + self.wall_start_ts: float = float(wall_start_ts) + self.last_finish_ts: float = float(wall_start_ts) + self.stage_seen_batches: Dict[int, set] = {sid: set() for sid in range(self.num_stages)} + self.stage_first_ts: List[Optional[float]] = [None for _ in range(self.num_stages)] + self.stage_last_ts: List[Optional[float]] = [None for _ in range(self.num_stages)] + + def on_stage_metrics(self, stage_id: int, req_id: Any, metrics: Dict[str, Any]) -> None: + record_stage_metrics( + self.per_request, + self.stage_req_counts, + self.stage_total_time_ms, + self.stage_total_tokens, + stage_id, + req_id, + metrics, + ) + try: + batch_id_raw = metrics.get("batch_id", None) + if batch_id_raw is not None: + batch_id = int(batch_id_raw) + if batch_id not in self.stage_seen_batches[stage_id]: + self.stage_total_time_ms[stage_id] += float(metrics.get("stage_gen_time_ms", 0.0)) + self.stage_seen_batches[stage_id].add(batch_id) + except Exception: + pass + rx_b = float(metrics.get("rx_transfer_bytes", 0.0)) + rx_ms = float(metrics.get("rx_decode_time_ms", 0.0)) + in_flight_ms = float(metrics.get("rx_in_flight_time_ms", 0.0)) + combined = aggregate_rx_and_maybe_total( + self.transfer_edge_req, + self.transfer_agg, + self.per_request, + stage_id, + req_id, + rx_b, + rx_ms, + in_flight_ms, + ) + if self.enable_stats and self.stats_file and stage_id > 0: + log_transfer_rx( + self.stats_file, + stage_id - 1, + stage_id, + req_id, + int(rx_b), + float(rx_ms), + float(in_flight_ms), + ) + if combined is not None: + size_b_c, tx_ms_c, total_ms_c = combined + log_transfer_total( + self.stats_file, + stage_id - 1, + stage_id, + req_id, + int(size_b_c), + float(tx_ms_c), + float(in_flight_ms), + float(rx_ms), + float(total_ms_c), + ) + + def on_forward(self, from_stage: int, to_stage: int, req_id: Any, + size_bytes: int, tx_ms: float, used_shm: bool) -> None: + # Mark first input time for the destination stage if not set + if self.stage_first_ts[to_stage] is None: + self.stage_first_ts[to_stage] = time.time() + if self.enable_stats and self.stats_file: + log_transfer_tx( + self.stats_file, + from_stage, + to_stage, + req_id, + int(size_bytes), + float(tx_ms), + bool(used_shm), + ) + record_sender_transfer_agg( + self.transfer_agg, + self.transfer_edge_req, + from_stage, + to_stage, + req_id, + int(size_bytes), + float(tx_ms), + ) + + def on_finalize_request(self, stage_id: int, req_id: Any, engine_outputs: List[Any], + req_start_ts: float) -> None: + try: + rid_int = int(req_id) + except Exception: + rid_int = req_id + _t0 = float(req_start_ts) + _t1 = time.time() + # Update last output time for this stage + prev_last = self.stage_last_ts[stage_id] + self.stage_last_ts[stage_id] = _t1 if prev_last is None else max(prev_last, _t1) + self.last_finish_ts = max(self.last_finish_ts, _t1) + e2e_ms = (_t1 - _t0) * 1000.0 + num_tokens = count_tokens_from_outputs(engine_outputs) + self.e2e_total_ms += e2e_ms + self.e2e_total_tokens += int(num_tokens) + self.e2e_count += 1 + self.e2e_done.add(rid_int) + pr = self.per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}) + per_req_record = { + "type": "overall_request", + "request_id": rid_int, + "e2e_time_ms": e2e_ms, + "num_tokens_out": int(num_tokens), + "transfers_total_time_ms": float(pr.get("transfers_ms", 0.0)), + "transfers_total_bytes": int(pr.get("transfers_bytes", 0)), + "stages": pr.get("stages", {}), + } + self.sum_per_request_transfer_ms += float(pr.get("transfers_ms", 0.0)) + if self.enable_stats and self.overall_stats_file: + log_overall_record(self.overall_stats_file, per_req_record) + if self.enable_stats and self.stats_file: + e2e_tpt = (e2e_ms / num_tokens) if num_tokens > 0 else 0.0 + log_orchestrator_e2e(self.stats_file, req_id, stage_id, e2e_ms, int(num_tokens), e2e_tpt) + + def build_and_log_summary(self, final_stage_id_for_e2e: int) -> Dict[str, Any]: + # Compute stage summary using wall time between first input and last output per stage + stage_summary: List[Dict[str, Any]] = [] + for sid in range(self.num_stages): + first_ts = self.stage_first_ts[sid] + last_ts = self.stage_last_ts[sid] + total_ms = (max(0.0, (last_ts - first_ts)) * 1000.0) if (first_ts is not None and last_ts is not None) else 0.0 + reqs = self.stage_req_counts[sid] + tokens = self.stage_total_tokens[sid] + avg_req = (total_ms / reqs) if reqs > 0 else 0.0 + avg_tok = (tokens * 1000.0 / total_ms) if total_ms > 0 else 0.0 + stage_summary.append({ + "stage_id": sid, + "requests": int(reqs), + "tokens": int(tokens), + "total_time_ms": float(total_ms), + "avg_time_per_request_ms": float(avg_req), + "avg_tokens_per_s": float(avg_tok), + }) + transfer_summary = build_transfer_summary(self.transfer_agg) + e2e_avg_req = (self.e2e_total_ms / self.e2e_count) if self.e2e_count > 0 else 0.0 + e2e_avg_tok = (self.e2e_total_tokens * 1000.0 / self.e2e_total_ms) if self.e2e_total_ms > 0 else 0.0 + wall_time_ms = max(0.0, (self.last_finish_ts - self.wall_start_ts) * 1000.0) + summary: Dict[str, Any] = { + "e2e_requests": int(self.e2e_count), + "e2e_total_time_ms": float(wall_time_ms), + "e2e_sum_time_ms": float(self.e2e_total_ms), + "e2e_total_tokens": int(self.e2e_total_tokens), + "e2e_avg_time_per_request_ms": e2e_avg_req, + "e2e_avg_tokens_per_s": e2e_avg_tok, + "wall_time_ms": wall_time_ms, + "final_stage_id": final_stage_id_for_e2e, + "stages": stage_summary, + "transfers": transfer_summary, + } + if self.enable_stats and self.stats_file: + log_orchestrator_summary(self.stats_file, summary) + if self.enable_stats and self.overall_stats_file: + log_overall_summary(self.overall_stats_file, summary) + return summary + + diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index a76c1126f..6f40064ea 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -116,18 +116,18 @@ class OmniStage: if self._in_q is not None: try: self._in_q.put_nowait(None) - except Exception: - pass + except Exception as e: + self._logger.warning("[Stage-%s] Failed to send shutdown to in_q: %s", self.stage_id, e) if self._proc is not None: try: self._proc.join(timeout=5) - except Exception: - pass + except Exception as e: + self._logger.debug("[Stage-%s] join() failed: %s", self.stage_id, e, exc_info=True) if self._proc.is_alive(): try: self._proc.terminate() - except Exception: - pass + except Exception as e: + self._logger.warning("[Stage-%s] terminate() failed: %s", self.stage_id, e) def submit(self, payload: Dict[str, Any]) -> None: assert self._in_q is not None @@ -185,6 +185,11 @@ def _stage_worker( """Stage worker entry: device setup, LLM init, batching, SHM IPC.""" 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, + ) import queue as _queue import time as _time from vllm_omni.entrypoints.pipeline_utils import append_jsonl as _append_jsonl, serialize_obj as _ser @@ -217,6 +222,8 @@ def _stage_worker( # Aggregates for running average _agg_total_tokens = 0 _agg_total_gen_time_ms = 0.0 + # Monotonic batch id per stage process for orchestrator dedup on time aggregation + _batch_seq = 0 # Device mapping try: @@ -294,6 +301,7 @@ def _stage_worker( print(f"[Stage-{stage_id}] Received batch size={len(batch_tasks)}, request_ids={batch_request_ids}", flush=True) print("--------------------------------", flush=True) try: + _batch_seq += 1 gen_outputs: List[Any] = [] _gen_t0 = _time.time() for ro in stage_engine.generate(batch_engine_inputs, sampling_params, use_tqdm=False): @@ -355,24 +363,9 @@ def _stage_worker( pass if _stats_file: - try: - _avg_tokens_per_s = (_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms) if _agg_total_gen_time_ms > 0 else 0.0 - _append_jsonl(_stats_file, { - "type": "stage_running_avg", - "stage_id": stage_id, - "total_tokens": int(_agg_total_tokens), - "total_gen_time_ms": _agg_total_gen_time_ms, - "avg_tokens_per_s": _avg_tokens_per_s, - }) - _append_jsonl(_stats_file, { - "type": "stage_batch_stats", - "stage_id": stage_id, - "batch_size": len(batch_tasks), - "batch_gen_time_ms": _gen_ms, - "request_ids": list(batch_request_ids), - }) - except Exception: - pass + _avg_tokens_per_s = (_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms) if _agg_total_gen_time_ms > 0 else 0.0 + log_stage_running_avg(_stats_file, stage_id, int(_agg_total_tokens), float(_agg_total_gen_time_ms), float(_avg_tokens_per_s)) + log_stage_batch_stats(_stats_file, stage_id, len(batch_tasks), float(_gen_ms), list(batch_request_ids)) # Emit per-request results for rid in batch_request_ids: @@ -382,10 +375,29 @@ def _stage_worker( _metrics = { "num_tokens_out": int(_count_tokens(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)), "rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)), "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( + _stats_file, + stage_id, + rid, + len(batch_tasks), + _num_tokens, + float(_gen_ms), + float(_tokens_per_s), + _rx_b, + _rx_ms, + float(_rx_mbps), + ) if use_shm: out_q.put({ "request_id": rid, diff --git a/vllm_omni/entrypoints/pipeline_utils.py b/vllm_omni/entrypoints/pipeline_utils.py index cbf036332..6a2384ec9 100644 --- a/vllm_omni/entrypoints/pipeline_utils.py +++ b/vllm_omni/entrypoints/pipeline_utils.py @@ -39,7 +39,8 @@ def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) -> "[Stage-%s] Set CUDA_VISIBLE_DEVICES to %s; logical 0 -> physical %s", stage_id, devices, selected_physical, ) - except Exception: + except Exception as e: + logger.debug("[Stage-%s] Failed to parse first CUDA device: %s", stage_id, e) selected_physical = None elif isinstance(devices, (int, str)) and (isinstance(devices, int) or str(devices).isdigit()): logical_idx = max(0, int(devices)) @@ -49,7 +50,8 @@ def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) -> mapping = [int(x) for x in vis.split(",") if x.strip() != ""] if 0 <= logical_idx < len(mapping): selected_physical = mapping[logical_idx] - except Exception: + except Exception as e: + logger.debug("[Stage-%s] Failed to map logical index via CUDA_VISIBLE_DEVICES: %s", stage_id, e) selected_physical = None if selected_physical is None: selected_physical = int(logical_idx) @@ -70,8 +72,8 @@ def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) -> if torch.cuda.is_available(): try: torch.cuda.set_device(0) - except Exception: - pass + except Exception as e: + logger.debug("[Stage-%s] torch.cuda.set_device(0) failed: %s", stage_id, e, exc_info=True) num = torch.cuda.device_count() info = [] for i in range(num): @@ -84,8 +86,8 @@ def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) -> "free": int(free), }) logger.debug("[Stage-%s] CUDA devices visible=%s info=%s", stage_id, num, info) - except Exception: - pass + except Exception as e: + logger.debug("[Stage-%s] Failed to query CUDA devices: %s", stage_id, e, exc_info=True) except Exception as e: logger.warning("Failed to interpret devices for stage %s: %s", stage_id, e) @@ -107,8 +109,8 @@ def shm_write_bytes(payload: bytes) -> Dict[str, Any]: meta = {"name": shm.name, "size": len(payload)} try: shm.close() - except Exception: - pass + except Exception as e: + logger.debug("Ensure parent dir failed for %s: %s", path, e) return meta diff --git a/vllm_omni/entrypoints/pipelined_omni_llm.py b/vllm_omni/entrypoints/pipelined_omni_llm.py index b874c2ee1..da13b0dee 100644 --- a/vllm_omni/entrypoints/pipelined_omni_llm.py +++ b/vllm_omni/entrypoints/pipelined_omni_llm.py @@ -16,6 +16,24 @@ 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, @@ -65,32 +83,10 @@ class PipelinedOmniLLM(OmniLLM): # Optional file handler for orchestrator self._log_file = log_file if self._log_file: - try: - # Avoid duplicate handlers - has_file_handler = any(isinstance(h, logging.FileHandler) for h in logger.handlers) - if not has_file_handler: - fh = logging.FileHandler(self._log_file) - fh.setLevel(logging.DEBUG) - fh.setFormatter(logging.Formatter("%(asctime)s [PID:%(process)d] %(levelname)s: %(message)s")) - logger.addHandler(fh) - logger.setLevel(logging.DEBUG) - except Exception: - pass + remove_old_logs(self._log_file, len(self.stage_list)) + configure_orchestrator_logger(logger, self._log_file) - # Orchestrator stats JSONL file - self._stats_file: Optional[str] = None - if self._enable_stats and self._log_file: - try: - self._stats_file = f"{self._log_file}.orchestrator.stats.jsonl" - except Exception: - self._stats_file = None - # Overall stats JSONL file (per-request + summary) - self._overall_stats_file: Optional[str] = None - if self._enable_stats and self._log_file: - try: - self._overall_stats_file = f"{self._log_file}.overall.stats.jsonl" - except Exception: - self._overall_stats_file = None + 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)) @@ -122,19 +118,19 @@ class PipelinedOmniLLM(OmniLLM): for q in self._stage_in_queues: try: q.put_nowait(None) - except Exception: - pass + 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: - pass + 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: - pass + except Exception as e: + logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True) def generate( self, @@ -178,25 +174,16 @@ class PipelinedOmniLLM(OmniLLM): 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: + 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 - # In-memory aggregators for this generate() call - stage_total_time_ms: List[float] = [0.0 for _ in range(num_stages)] - stage_total_tokens: List[int] = [0 for _ in range(num_stages)] - stage_req_counts: List[int] = [0 for _ in range(num_stages)] - transfer_agg: Dict[tuple[int, int], Dict[str, float]] = {} - # Per-edge per-request sender timing to combine with receiver timing later - transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]] = {} - e2e_total_ms: float = 0.0 - e2e_total_tokens: int = 0 - e2e_count: int = 0 - e2e_done: set[int] = set() - # Per-request overall aggregation - per_request: Dict[int, Dict[str, Any]] = {} - sum_per_request_transfer_ms: float = 0.0 + # 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] @@ -235,94 +222,14 @@ class PipelinedOmniLLM(OmniLLM): continue engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm") - # Aggregate per-request metrics from stage worker if present + # 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: - stage_req_counts[stage_id] += 1 - stage_total_time_ms[stage_id] += float(_m.get("stage_gen_time_ms", 0.0)) - stage_total_tokens[stage_id] += int(_m.get("num_tokens_out", 0)) - # record per-request stage metrics - try: - rid_int = int(req_id) - pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}) - pr_stages = pr["stages"] # type: ignore[index] - pr_stages[stage_id] = { - "stage_gen_time_ms": float(_m.get("stage_gen_time_ms", 0.0)), - "num_tokens_out": int(_m.get("num_tokens_out", 0)), - } - except Exception: - pass - # Also aggregate receiver-side transfer decode time for (stage_id-1 -> stage_id) - try: - if stage_id > 0: - key = (stage_id - 1, stage_id) - agg = transfer_agg.get(key) - if agg is None: - agg = {"sum_bytes": 0.0, "sum_ms": 0.0, "count": 0.0, - "sum_rx_bytes": 0.0, "sum_rx_ms": 0.0, "rx_count": 0.0, - "sum_total_ms": 0.0, "total_count": 0.0} - transfer_agg[key] = agg - rx_b = float(_m.get("rx_transfer_bytes", 0.0)) - rx_ms = float(_m.get("rx_decode_time_ms", 0.0)) - in_flight_ms = float(_m.get("rx_in_flight_time_ms", 0.0)) - agg["sum_rx_bytes"] += rx_b - agg["sum_rx_ms"] += rx_ms - agg["rx_count"] += 1.0 - # If we have sender-side timing stored, compute combined (encode+enqueue + decode) - try: - edge_req = (stage_id - 1, stage_id, int(req_id)) - s = transfer_edge_req.get(edge_req) - if s is not None: - total_ms = float(s.get("tx_ms", 0.0)) + in_flight_ms + rx_ms - agg["sum_total_ms"] += total_ms - agg["total_count"] += 1.0 - # accumulate per-request transfer totals - try: - rid_int = int(req_id) - pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}) - pr["transfers_ms"] = float(pr.get("transfers_ms", 0.0)) + total_ms # type: ignore[index] - pr["transfers_bytes"] = int(pr.get("transfers_bytes", 0)) + int(rx_b) # type: ignore[index] - except Exception: - pass - if getattr(self, "_stats_file", None): - try: - size_b = float(s.get("size_bytes", rx_b)) - _append_jsonl(self._stats_file, { # type: ignore[arg-type] - "type": "transfer_total_stats", - "from_stage": stage_id - 1, - "to_stage": stage_id, - "request_id": req_id, - "size_bytes": int(size_b), - "tx_time_ms": float(s.get("tx_ms", 0.0)), - "in_flight_time_ms": in_flight_ms, - "rx_decode_time_ms": rx_ms, - "total_time_ms": total_ms, - "total_time_per_kb_ms": total_ms / max(size_b / 1024.0, 1e-6) if size_b > 0 else 0.0, - }) - except Exception: - pass - except Exception: - pass - # Emit per-request RX stats to JSONL - if getattr(self, "_stats_file", None): - try: - _append_jsonl(self._stats_file, { # type: ignore[arg-type] - "type": "transfer_rx_stats", - "from_stage": stage_id - 1, - "to_stage": stage_id, - "request_id": req_id, - "rx_bytes": int(rx_b), - "rx_decode_time_ms": rx_ms, - "in_flight_time_ms": in_flight_ms, - "rx_time_per_kb_ms": rx_ms / max(rx_b / 1024.0, 1e-6) if rx_b > 0 else 0.0, - }) - except Exception: - pass - except Exception: - pass - except Exception: - pass + 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) @@ -338,67 +245,11 @@ class PipelinedOmniLLM(OmniLLM): # End-to-end timing and time-per-token for final output (only once per request at the designated final stage) try: - if stage_id == final_stage_id_for_e2e and req_id not in e2e_done: - _t0 = _req_start_ts.get(req_id) - if _t0 is not None: - _t1 = time.time() - _last_finish_ts = max(_last_finish_ts, _t1) - _e2e_ms = (_t1 - _t0) * 1000.0 - # Count tokens from final stage outputs - def _count_tokens(_ros: List[Any]) -> int: # type: ignore[name-defined] - 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 - _num_tokens = _count_tokens(engine_outputs) - _time_per_token_ms = (_e2e_ms / _num_tokens) if _num_tokens > 0 else 0.0 - # Update E2E aggregators - e2e_total_ms += _e2e_ms - e2e_total_tokens += int(_num_tokens) - e2e_count += 1 - e2e_done.add(req_id) - # Write per-request overall record - try: - rid_int = int(req_id) - pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}) - per_req_record = { - "type": "overall_request", - "request_id": rid_int, - "e2e_time_ms": _e2e_ms, - "num_tokens_out": int(_num_tokens), - "transfers_total_time_ms": float(pr.get("transfers_ms", 0.0)), - "transfers_total_bytes": int(pr.get("transfers_bytes", 0)), - "stages": pr.get("stages", {}), - } - sum_per_request_transfer_ms += float(pr.get("transfers_ms", 0.0)) - if getattr(self, "_overall_stats_file", None): - try: - _append_jsonl(self._overall_stats_file, per_req_record) # type: ignore[arg-type] - except Exception: - pass - except Exception: - pass - if getattr(self, "_stats_file", None): - try: - _append_jsonl(self._stats_file, { # type: ignore[arg-type] - "type": "orchestrator_request_e2e", - "request_id": req_id, - "final_stage_id": stage_id, - "e2e_time_ms": _e2e_ms, - "num_tokens_out": int(_num_tokens), - "e2e_time_per_token_ms": _time_per_token_ms, - }) - except Exception: - pass - except Exception: - pass + 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: @@ -427,40 +278,9 @@ class PipelinedOmniLLM(OmniLLM): self.stage_list[next_stage_id].submit(ipc_payload) t1 = time.time() tx_ms = (t1 - t0) * 1000.0 - if self._enable_stats and getattr(self, "_stats_file", None): - try: - _append_jsonl(self._stats_file, { - "type": "transfer_stats", - "from_stage": stage_id, - "to_stage": next_stage_id, - "request_id": req_id, - "size_bytes": int(size_bytes), - "tx_time_ms": tx_ms, - "tx_mbps": (float(size_bytes) * 8.0) / (max(tx_ms, 1e-6) * 1000.0), - "used_shm": bool("engine_inputs_shm" in ipc_payload), - }) - except Exception: - pass - # Update in-memory transfer aggregator - try: - key = (stage_id, next_stage_id) - agg = transfer_agg.get(key) - if agg is None: - agg = {"sum_bytes": 0.0, "sum_ms": 0.0, "count": 0.0, - "sum_rx_bytes": 0.0, "sum_rx_ms": 0.0, "rx_count": 0.0, - "sum_total_ms": 0.0, "total_count": 0.0} - transfer_agg[key] = agg - agg["sum_bytes"] += float(size_bytes) - agg["sum_ms"] += float(tx_ms) - agg["count"] += 1.0 - # Store sender-side timing for per-request combination - transfer_edge_req[(stage_id, next_stage_id, int(req_id))] = { - "tx_ms": float(tx_ms), - "size_bytes": float(size_bytes), - } - except Exception: - pass - except Exception: + 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, @@ -478,81 +298,10 @@ class PipelinedOmniLLM(OmniLLM): # Summarize and print stats try: - stage_summary: List[Dict[str, Any]] = [] - for sid in range(num_stages): - reqs = stage_req_counts[sid] - tokens = stage_total_tokens[sid] - total_ms = float(stage_total_time_ms[sid]) - avg_req = (total_ms / reqs) if reqs > 0 else 0.0 - avg_tok = (tokens * 1000.0 / total_ms) if total_ms > 0 else 0.0 - stage_summary.append({ - "stage_id": sid, - "requests": int(reqs), - "tokens": int(tokens), - "total_time_ms": total_ms, - "avg_time_per_request_ms": avg_req, - "avg_tokens_per_s": avg_tok, - }) - - transfer_summary: List[Dict[str, Any]] = [] - for (src, dst), agg in transfer_agg.items(): - sum_bytes = float(agg.get("sum_bytes", 0.0)) - sum_ms = float(agg.get("sum_ms", 0.0)) - samples = int(agg.get("count", 0.0)) - tx_mbps = (sum_bytes * 8.0) / (max(sum_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0 - sum_rx_bytes = float(agg.get("sum_rx_bytes", 0.0)) - sum_rx_ms = float(agg.get("sum_rx_ms", 0.0)) - samples_rx = int(agg.get("rx_count", 0.0)) - rx_mbps = (sum_rx_bytes * 8.0) / (max(sum_rx_ms, 1e-6) * 1000.0) if sum_rx_bytes > 0 else 0.0 - sum_total_ms = float(agg.get("sum_total_ms", 0.0)) - samples_total = int(agg.get("total_count", 0.0)) - total_mbps = (sum_bytes * 8.0) / (max(sum_total_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0 - transfer_summary.append({ - "from_stage": src, - "to_stage": dst, - "samples": samples, - "total_bytes": int(sum_bytes), - "total_time_ms": sum_ms, - "tx_mbps": tx_mbps, - "rx_samples": samples_rx, - "rx_total_bytes": int(sum_rx_bytes), - "rx_total_time_ms": sum_rx_ms, - "rx_mbps": rx_mbps, - "total_samples": samples_total, - "total_transfer_time_ms": sum_total_ms, - "total_mbps": total_mbps, - }) - - e2e_avg_req = (e2e_total_ms / e2e_count) if e2e_count > 0 else 0.0 - e2e_avg_tok = (e2e_total_tokens * 1000.0 / e2e_total_ms) if e2e_total_ms > 0 else 0.0 - wall_time_ms = max(0.0, (_last_finish_ts - _wall_start_ts) * 1000.0) - summary: Dict[str, Any] = { - "e2e_requests": int(e2e_count), - # 按你的需求：e2e_total_time_ms 直接使用墙钟时间（非各请求之和） - "e2e_total_time_ms": float(wall_time_ms), - # 额外保留各请求 E2E 之和，供需要时参考 - "e2e_sum_time_ms": float(e2e_total_ms), - "e2e_total_tokens": int(e2e_total_tokens), - "e2e_avg_time_per_request_ms": e2e_avg_req, - "e2e_avg_tokens_per_s": e2e_avg_tok, - "wall_time_ms": wall_time_ms, - "final_stage_id": final_stage_id_for_e2e, - "stages": stage_summary, - "transfers": transfer_summary, - } + summary = metrics.build_and_log_summary(final_stage_id_for_e2e) logger.info("[Summary] %s", summary) - if self._enable_stats and getattr(self, "_stats_file", None): - try: - _append_jsonl(self._stats_file, {"type": "orchestrator_summary", **summary}) # type: ignore[arg-type] - except Exception: - pass - if self._enable_stats and getattr(self, "_overall_stats_file", None): - try: - _append_jsonl(self._overall_stats_file, {"type": "overall_summary", **summary}) # type: ignore[arg-type] - except Exception: - pass - except Exception: - pass + except Exception as e: + logger.exception("[Orchestrator] Failed to build/log summary: %s", e) return final_outputs diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 4c115bfc8..d91defd81 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -126,7 +126,7 @@ class GPUARModelRunner(OmniGPUModelRunner): src = pe_cpu[ num_computed_tokens : num_computed_tokens + overlay_len ].to(dtype=self.dtype, device=self.device, non_blocking=True) - start_offset = int(self.query_start_loc_cpu[req_index]) + start_offset = int(self.query_start_loc.cpu[req_index]) self.inputs_embeds[start_offset : start_offset + overlay_len].copy_( src ) @@ -381,7 +381,7 @@ class GPUARModelRunner(OmniGPUModelRunner): # Pass per-request additional information map for this step (no concat) if per_req_additional_information: model_kwargs_extra["additional_information_by_req_id"] = per_req_additional_information - # 始终传递每请求的运行期 additional_information（持久化在 request state） + # Always pass per-request runtime additional_information (persisted in request state) try: per_req_runtime_info = [] for req_id in self.input_batch.req_ids: @@ -396,10 +396,10 @@ class GPUARModelRunner(OmniGPUModelRunner): per_req_runtime_info ) model_kwargs_extra["request_ids"] = self.input_batch.req_ids - # 传递每个请求在本步扁平序列中的 token 区间，便于模型按请求映射 decode/prefill + # Pass each request's token span within the flattened sequence for this step, enabling the model to map decode/prefill by request req_token_spans = [] for req_index in range(len(self.input_batch.req_ids)): - start_offset = int(self.query_start_loc_cpu[req_index]) + start_offset = int(self.query_start_loc.cpu[req_index]) sched_tokens = int(num_scheduled_tokens_np[req_index]) req_token_spans.append((start_offset, start_offset + sched_tokens)) model_kwargs_extra["request_token_spans"] = req_token_spans @@ -429,9 +429,9 @@ class GPUARModelRunner(OmniGPUModelRunner): hidden_states, multimodal_outputs = self.extract_multimodal_outputs( hidden_states ) - # 模型侧可能返回 per-request 的 additional_information 更新（与模型无关的通用通道）。 - # 约定：multimodal_outputs["additional_information_update"] 为按 batch 顺序的 list[dict]， - # runner 将其合并到对应 request 的 additional_information_cpu，供后续 decode 使用。 + # The model side may return per-request additional_information updates (model-agnostic channel). + # Convention: multimodal_outputs["additional_information_update"] is a list[dict] in batch order; + # the runner merges it into the corresponding request's additional_information_cpu for subsequent decode. try: if ( isinstance(multimodal_outputs, dict) @@ -440,7 +440,7 @@ class GPUARModelRunner(OmniGPUModelRunner): or "additional_information_update_by_req_id" in multimodal_outputs ) ): - # 方案 A：按 batch 顺序的 list[dict] + # Option A: list[dict] in batch order updates_list = multimodal_outputs.get( "additional_information_update" ) @@ -450,7 +450,7 @@ class GPUARModelRunner(OmniGPUModelRunner): continue req_id = self.input_batch.req_ids[idx] self._merge_additional_information_update(req_id, upd) - # 方案 B：按 req_id 的 dict[str, dict] + # Option B: dict[str, dict] keyed by req_id updates_map = multimodal_outputs.get( "additional_information_update_by_req_id" ) From 8e8783d78a53f2ffbf71834b91685084c8e51c27 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Mon, 10 Nov 2025 11:58:47 +0800 Subject: [PATCH 04/19] 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. --- .../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) - - From 08aef1f964dba8060dceb85c2cac541701ec4563 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Mon, 10 Nov 2025 14:41:44 +0800 Subject: [PATCH 05/19] remove text output for audio output Signed-off-by: tzhouam --- .../offline_inference/qwen_2_5_omni/end2end.py | 18 ++---------------- 1 file changed, 2 insertions(+), 16 deletions(-) diff --git a/examples/offline_inference/qwen_2_5_omni/end2end.py b/examples/offline_inference/qwen_2_5_omni/end2end.py index 8c5da7c08..8614d4653 100644 --- a/examples/offline_inference/qwen_2_5_omni/end2end.py +++ b/examples/offline_inference/qwen_2_5_omni/end2end.py @@ -167,7 +167,7 @@ def main(): if args.prompts is None: raise ValueError("No prompts provided. Use --prompts ... or --pt-prompts (with --prompt_type text)") - omni_llm = PipelinedOmniLLM( + omni_llm = OmniLLM( model=model_name, log_stats=args.enable_stats, log_file=("omni_llm_pipeline.log" if args.enable_stats else None), @@ -230,7 +230,6 @@ def main(): 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) @@ -245,20 +244,7 @@ def main(): 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}") + print(f"Request ID: {request_id}, Saved audio to {output_wav}") if __name__ == "__main__": From 36ada0478aebfdb9b0599af6e3b279911544b052 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Mon, 10 Nov 2025 15:07:01 +0800 Subject: [PATCH 06/19] add back the sleep time for the stage init to avoid single card conflict Signed-off-by: tzhouam --- vllm_omni/entrypoints/omni_llm.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index 458e000e5..63dab9fde 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/vllm_omni/entrypoints/omni_llm.py @@ -178,6 +178,7 @@ class OmniLLM(BaseOmniLLM): batch_timeout=self.batch_timeout, ) logger.debug("[Orchestrator] Stage-%s process started", stage_id) + time.sleep(self._init_sleep_seconds) def close(self) -> None: for q in self._stage_in_queues: From 0ef1fefd4522cb4767f77b8ccde0a9fc1121872e Mon Sep 17 00:00:00 2001 From: tzhouam Date: Mon, 10 Nov 2025 15:44:44 +0800 Subject: [PATCH 07/19] add prompt file Signed-off-by: tzhouam --- .../offline_inference/qwen_2_5_omni/top100.pt | Bin 0 -> 9319 bytes 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 examples/offline_inference/qwen_2_5_omni/top100.pt diff --git a/examples/offline_inference/qwen_2_5_omni/top100.pt b/examples/offline_inference/qwen_2_5_omni/top100.pt new file mode 100644 index 0000000000000000000000000000000000000000..017c60ac806dc45952b6bf652fd60ad362ece700 GIT binary patch literal 9319 zcmd6tU6UNudB+zpg4i*TIKd`1<{&#Pah7JqN?P$M5FoP_fZeeU$nqZ6E6rGvdV5UV#n)_BK z>i)DwtDW_kshl*L`4$c5mVf=&SkrnxTMx!&z7JiI1+=o2nNMA1w5V$rMKv*#nNCln-*A3-ZB$!%XoC$Rhd~)roK6Nxhj3c#=a1NTF}*wg~)WyZCD;el{FYx8kVUCCx)Clk{+U z37u|ui&Z9oB`;kF7}2)AAFys}+a7GNvqQhC18E+Enkmjon!mWWjYl7WTm{+20Bj7f0KYlGaHvSFP%R8ba#dz$^$bKO%~wuC=wGI;H?(9}~RWk6ZzhF`%Kb~j;W+b;f`G>@jQBYWSK?yv~rh73}f z*#l=5zBhS48*P|T---0}i!C|%Nb%)!CVhc< zaA%n-AS`Twd82v69ikhE7Nhfax1i8?8|-g94rqy+S<-y%6lDHQ`U-l!Lr}I(6bS7P zTpk>pUv$gdgs=3?j6J%TF=eV z+%90AF7y=?7#aAcOluVzwv3XsO)-NYISbtFGg~g|0Q>#jsqgU3^fG6-Qna{y>=YDJ zSF{{v9LN%QHi9RObb28VZ^v1QDpz*YQ%gy{l5I)VRq8s$c_Y+fkQZk9d>JhS9I^(lM+Y4 zWug}G#&Cp1!~R1{pDqsc}|T_2U<}A$;x6jx^)aY{Dx*Z@kEAgx zmG*R+a?%e&48g1|>duTx@QJS!J6K*AJAVAz1lzi%FIqTTW`H_T)2sez;axK-NOWh@ zJsC2xYwIW?dK?K+>+E=MPQvzzmDUj?T%P>+!| zHYlE|h(3?mWHCa1)cU!!%2dcmx1h>0)~N(xYjI$x5u7yNjl8ojlBAY^A;wXO&d(uT z)y}pYqhyF=QP`6Fsw@t+qT7}7zEU97_2pXQ&G1OlWBEGPs4Kh^6gV%Yd&(XTN8N?yk?#8z>r?QkT*y6mRZ0NWehzA720>47{m0lH?W>xgrUwg3 z3f;WGMY%0XZ`tXF)|VPR=H?SVZ_8mI<#H~VT*ieWpGlWbbQ;5J!bMIzis z#WWO5$^Qo>o{wRw99t`mS3`AJ79;`XmO$ED;)rsh&?TK%o$?awNwXPaYM(lN9V)uh zgvtU#3q2&PLp$m5G;B{=!yee;CJZ0pA^;+gC1l{FIk$R|FQlL1@;zKGzv3y<-jb7M z^prG~bzeeqv2@iLBj$5Jo~g1Cy~et}4}D~B0X{%gl&vH);JJ0$8=Ox+-~=Dpve$+N z3!=k@DaMNA5lK)@w53S0qoOPnsypjLt%bESd2t}}%Jj9dkyXhXNps94ncx8i=5JQ%@zwMOSi0VZ#O$}0K6N=A*imt$3S?mp z3d&_1Mi= zyBqe46hHT+aNH@S+5Aw60s0@Pgq1J?O@mFE%cmy!^_b<~q?Pu8IOoKqsUg|ZjK6I5 z$W#Hy#6cew8{-HiV>4nHm^5$1sEU=iHI)+a+?JZS=|0S=Lu;Zi^b!mm30vAq7%_DH zc$_p>(hYp^uF3;s_AOTC!w}shf5w$FL}&qFj*qDV-b}wI>Y}+DFGN@vnIEfhG3ncw zRuy@%8a=dSK`r`Ly32xVuBylglo}j{6OL>?lvP@<^vI~`X&TzYrAQc3 zH&(RlP*6`^c{@&thA~P|m8^jlDNSTBnU1(;nD558W9hQ*j1&=vVKfYp9ZJohDh+&+ z=IUt({r>G}4Zyb;7bIfNUinz<`ijq!2>b|StdBlmP$d?%AN41(q{A!CseB?3LhV(B zF&ZiW{yXX0Sa_dX@X`pxMzjdVZ9l7>5Kl|bAmIoHIoFxwL`tDa>F7_Ym<*FC~(Qt(`9#R zP9l(0J7Yk90J_bs=UQ58hD4Mtt*s+<$^3B`bbY1L-iyt_*-w103TxRmmokbhaZvgB z<+IYdV}KPQK(oJOFU-YwZQ*E|-yg;b1MSi$j@*PJ4)M-nLX66W7^`8A*#^X_W}QKc zo0BkU`QQX>yt@)9AEf8-@_TfQM1BQvlWK#b|LZ5}38-dMKw~xI8{)Pg(bsBpQk&JIA zkoj`t9K*qC+RXf!7M~>f(MqL#oW6mJ-kds~AW;di-$18xHzMmqvK}k=pk+om;^7=4 z$AJw^9aO9pxGY_foKI4RK5q=8e_D~6%reB~+yek?)PmF_9Py;f1K#CCTkS^X)y$DM zC<4F-&(JWaiUTInYC-!-iB1VMX+He|_xfRMYxm-4V`m*_Pbh?8A&WgH!vR}2iM@$V zGfXw*{jtAAaIgUOqg^95e-zc;>*0P3;zLGZFaS6&YEOclk?sJN#`x~mP)G9khiYkA z?C5$4#2J>@38}ZpKq5Pk7J@S0So0~Yg^7Or1@85e^a`$jKSFS+Zh(?9f|P)C_$`hL zi(Fv~UCem~0S0Mtc#t7$Ov%*hD?eM6pr5AippoKK4GTab2B<;FmZZ;&HbZ2jM<21o zB1kdboWY(ht7M|dAgzE!svTm)#(wvDe1J5>P2^?h8G;Sb;W^;2tWdhdSjqy5L%&aEvU?I#2?wES5NR$pDc^4!TM z{Kud7pZs=CyZ-Bc{o60^!-IN%{;p4~UAn!idPkDXXVj;YN>Oz6%K1NU zH!q#Pcv|JxO^Dob09ldw2 z-6w@_|Id@Vf%~X_IQBa6|4>dYz8im!A6xi);^0ps{iKY$A@syC86TKCGX4m+C&k>2 zg2#@D`Q05cx6yZU_1%7a3f_csb~UboH;e_8jCK0Wx|;rkH(560WiuuAt2_x?9;1?{B( literal 0 HcmV?d00001 From 519c140d2018ec7d3177fb813285da8fda226d71 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Mon, 10 Nov 2025 15:47:28 +0800 Subject: [PATCH 08/19] Update run script to include initialization sleep time and support for loading prompts from a .pt file. Signed-off-by: tzhouam --- examples/offline_inference/qwen_2_5_omni/run.sh | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/examples/offline_inference/qwen_2_5_omni/run.sh b/examples/offline_inference/qwen_2_5_omni/run.sh index 7f02a672c..bee28e2c4 100644 --- a/examples/offline_inference/qwen_2_5_omni/run.sh +++ b/examples/offline_inference/qwen_2_5_omni/run.sh @@ -1,8 +1,9 @@ export PYTHONPATH=/path/to/vllm-omni:$PYTHONPATH python end2end.py --model Qwen/Qwen2.5-Omni-7B \ - --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words." \ --voice-type "m02" \ --dit-ckpt none \ --bigvgan-ckpt none \ --output-wav output_audio \ - --prompt_type text + --prompt_type text \ + --init-sleep-seconds 0 \ + --pt-prompts top100.pt From c18fc40f104c665b020e76163aa6cc52e2610581 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Mon, 10 Nov 2025 15:47:28 +0800 Subject: [PATCH 09/19] Update run script to include initialization sleep time and support for loading prompts from a .pt file. Signed-off-by: tzhouam --- .../offline_inference/qwen_2_5_omni/run.sh | 5 ++-- .../offline_inference/qwen_2_5_omni/utils.py | 27 ++++++++++++++----- 2 files changed, 23 insertions(+), 9 deletions(-) diff --git a/examples/offline_inference/qwen_2_5_omni/run.sh b/examples/offline_inference/qwen_2_5_omni/run.sh index 7f02a672c..bee28e2c4 100644 --- a/examples/offline_inference/qwen_2_5_omni/run.sh +++ b/examples/offline_inference/qwen_2_5_omni/run.sh @@ -1,8 +1,9 @@ export PYTHONPATH=/path/to/vllm-omni:$PYTHONPATH python end2end.py --model Qwen/Qwen2.5-Omni-7B \ - --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words." \ --voice-type "m02" \ --dit-ckpt none \ --bigvgan-ckpt none \ --output-wav output_audio \ - --prompt_type text + --prompt_type text \ + --init-sleep-seconds 0 \ + --pt-prompts top100.pt diff --git a/examples/offline_inference/qwen_2_5_omni/utils.py b/examples/offline_inference/qwen_2_5_omni/utils.py index 730d2a077..42894c4a7 100644 --- a/examples/offline_inference/qwen_2_5_omni/utils.py +++ b/examples/offline_inference/qwen_2_5_omni/utils.py @@ -13,6 +13,10 @@ from processing_omni import fetch_image, fetch_video 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(): @@ -113,16 +117,25 @@ def make_inputs_qwen2_omni( from transformers import AutoConfig, AutoProcessor - processor = AutoProcessor.from_pretrained(args.model) + # 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] - try: - config = AutoConfig.from_pretrained(args.model) - if "Qwen2_5OmniModel" in config.architectures: - args.legacy_omni_video = False + 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 = not ("Qwen2_5OmniModel" in config.architectures) else: args.legacy_omni_video = True - except Exception: - args.legacy_omni_video = True audios, images, videos = [], [], [] for message in messages: From 5f3f53a39d2b3e98850664ed32e0bf980b64f20e Mon Sep 17 00:00:00 2001 From: tzhouam Date: Mon, 10 Nov 2025 16:36:55 +0800 Subject: [PATCH 10/19] Revert "Update run script to include initialization sleep time and support for loading prompts from a .pt file." This reverts commit c18fc40f104c665b020e76163aa6cc52e2610581. --- .../offline_inference/qwen_2_5_omni/run.sh | 5 ++-- .../offline_inference/qwen_2_5_omni/utils.py | 27 +++++-------------- 2 files changed, 9 insertions(+), 23 deletions(-) diff --git a/examples/offline_inference/qwen_2_5_omni/run.sh b/examples/offline_inference/qwen_2_5_omni/run.sh index bee28e2c4..7f02a672c 100644 --- a/examples/offline_inference/qwen_2_5_omni/run.sh +++ b/examples/offline_inference/qwen_2_5_omni/run.sh @@ -1,9 +1,8 @@ export PYTHONPATH=/path/to/vllm-omni:$PYTHONPATH python end2end.py --model Qwen/Qwen2.5-Omni-7B \ + --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words." \ --voice-type "m02" \ --dit-ckpt none \ --bigvgan-ckpt none \ --output-wav output_audio \ - --prompt_type text \ - --init-sleep-seconds 0 \ - --pt-prompts top100.pt + --prompt_type text diff --git a/examples/offline_inference/qwen_2_5_omni/utils.py b/examples/offline_inference/qwen_2_5_omni/utils.py index 42894c4a7..730d2a077 100644 --- a/examples/offline_inference/qwen_2_5_omni/utils.py +++ b/examples/offline_inference/qwen_2_5_omni/utils.py @@ -13,10 +13,6 @@ from processing_omni import fetch_image, fetch_video 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(): @@ -117,25 +113,16 @@ def make_inputs_qwen2_omni( 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] + processor = AutoProcessor.from_pretrained(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 = not ("Qwen2_5OmniModel" in config.architectures) + try: + config = AutoConfig.from_pretrained(args.model) + if "Qwen2_5OmniModel" in config.architectures: + args.legacy_omni_video = False else: args.legacy_omni_video = True + except Exception: + args.legacy_omni_video = True audios, images, videos = [], [], [] for message in messages: From da385277a210d3c0b4b588404f4bdf27a53279c4 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Mon, 10 Nov 2025 16:37:37 +0800 Subject: [PATCH 11/19] Revert "Revert "Update run script to include initialization sleep time and support for loading prompts from a .pt file."" This reverts commit 5f3f53a39d2b3e98850664ed32e0bf980b64f20e. --- .../offline_inference/qwen_2_5_omni/run.sh | 5 ++-- .../offline_inference/qwen_2_5_omni/utils.py | 27 ++++++++++++++----- 2 files changed, 23 insertions(+), 9 deletions(-) diff --git a/examples/offline_inference/qwen_2_5_omni/run.sh b/examples/offline_inference/qwen_2_5_omni/run.sh index 7f02a672c..bee28e2c4 100644 --- a/examples/offline_inference/qwen_2_5_omni/run.sh +++ b/examples/offline_inference/qwen_2_5_omni/run.sh @@ -1,8 +1,9 @@ export PYTHONPATH=/path/to/vllm-omni:$PYTHONPATH python end2end.py --model Qwen/Qwen2.5-Omni-7B \ - --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words." \ --voice-type "m02" \ --dit-ckpt none \ --bigvgan-ckpt none \ --output-wav output_audio \ - --prompt_type text + --prompt_type text \ + --init-sleep-seconds 0 \ + --pt-prompts top100.pt diff --git a/examples/offline_inference/qwen_2_5_omni/utils.py b/examples/offline_inference/qwen_2_5_omni/utils.py index 730d2a077..42894c4a7 100644 --- a/examples/offline_inference/qwen_2_5_omni/utils.py +++ b/examples/offline_inference/qwen_2_5_omni/utils.py @@ -13,6 +13,10 @@ from processing_omni import fetch_image, fetch_video 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(): @@ -113,16 +117,25 @@ def make_inputs_qwen2_omni( from transformers import AutoConfig, AutoProcessor - processor = AutoProcessor.from_pretrained(args.model) + # 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] - try: - config = AutoConfig.from_pretrained(args.model) - if "Qwen2_5OmniModel" in config.architectures: - args.legacy_omni_video = False + 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 = not ("Qwen2_5OmniModel" in config.architectures) else: args.legacy_omni_video = True - except Exception: - args.legacy_omni_video = True audios, images, videos = [], [], [] for message in messages: From c83d223e914c1c71623d70caff9e8df285d14232 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Mon, 10 Nov 2025 16:42:58 +0800 Subject: [PATCH 12/19] Squash 7 commits: 36ada0478aebfdb9b0599af6e3b279911544b052 0ef1fefd4522cb4767f77b8ccde0a9fc1121872e 519c140d2018ec7d3177fb813285da8fda226d71 c18fc40f104c665b020e76163aa6cc52e2610581 5f3f53a39d2b3e98850664ed32e0bf980b64f20e da385277a210d3c0b4b588404f4bdf27a53279c4 edefcc244bc84010248bfd5180265a58af31a5ce Signed-off-by: tzhouam --- .../offline_inference/qwen_2_5_omni/run.sh | 5 ++-- .../offline_inference/qwen_2_5_omni/top100.pt | Bin 0 -> 9319 bytes .../offline_inference/qwen_2_5_omni/utils.py | 27 +++++++++++++----- vllm_omni/entrypoints/omni_llm.py | 1 + 4 files changed, 24 insertions(+), 9 deletions(-) create mode 100644 examples/offline_inference/qwen_2_5_omni/top100.pt diff --git a/examples/offline_inference/qwen_2_5_omni/run.sh b/examples/offline_inference/qwen_2_5_omni/run.sh index 7f02a672c..bee28e2c4 100644 --- a/examples/offline_inference/qwen_2_5_omni/run.sh +++ b/examples/offline_inference/qwen_2_5_omni/run.sh @@ -1,8 +1,9 @@ export PYTHONPATH=/path/to/vllm-omni:$PYTHONPATH python end2end.py --model Qwen/Qwen2.5-Omni-7B \ - --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words." \ --voice-type "m02" \ --dit-ckpt none \ --bigvgan-ckpt none \ --output-wav output_audio \ - --prompt_type text + --prompt_type text \ + --init-sleep-seconds 0 \ + --pt-prompts top100.pt diff --git a/examples/offline_inference/qwen_2_5_omni/top100.pt b/examples/offline_inference/qwen_2_5_omni/top100.pt new file mode 100644 index 0000000000000000000000000000000000000000..017c60ac806dc45952b6bf652fd60ad362ece700 GIT binary patch literal 9319 zcmd6tU6UNudB+zpg4i*TIKd`1<{&#Pah7JqN?P$M5FoP_fZeeU$nqZ6E6rGvdV5UV#n)_BK z>i)DwtDW_kshl*L`4$c5mVf=&SkrnxTMx!&z7JiI1+=o2nNMA1w5V$rMKv*#nNCln-*A3-ZB$!%XoC$Rhd~)roK6Nxhj3c#=a1NTF}*wg~)WyZCD;el{FYx8kVUCCx)Clk{+U z37u|ui&Z9oB`;kF7}2)AAFys}+a7GNvqQhC18E+Enkmjon!mWWjYl7WTm{+20Bj7f0KYlGaHvSFP%R8ba#dz$^$bKO%~wuC=wGI;H?(9}~RWk6ZzhF`%Kb~j;W+b;f`G>@jQBYWSK?yv~rh73}f z*#l=5zBhS48*P|T---0}i!C|%Nb%)!CVhc< zaA%n-AS`Twd82v69ikhE7Nhfax1i8?8|-g94rqy+S<-y%6lDHQ`U-l!Lr}I(6bS7P zTpk>pUv$gdgs=3?j6J%TF=eV z+%90AF7y=?7#aAcOluVzwv3XsO)-NYISbtFGg~g|0Q>#jsqgU3^fG6-Qna{y>=YDJ zSF{{v9LN%QHi9RObb28VZ^v1QDpz*YQ%gy{l5I)VRq8s$c_Y+fkQZk9d>JhS9I^(lM+Y4 zWug}G#&Cp1!~R1{pDqsc}|T_2U<}A$;x6jx^)aY{Dx*Z@kEAgx zmG*R+a?%e&48g1|>duTx@QJS!J6K*AJAVAz1lzi%FIqTTW`H_T)2sez;axK-NOWh@ zJsC2xYwIW?dK?K+>+E=MPQvzzmDUj?T%P>+!| zHYlE|h(3?mWHCa1)cU!!%2dcmx1h>0)~N(xYjI$x5u7yNjl8ojlBAY^A;wXO&d(uT z)y}pYqhyF=QP`6Fsw@t+qT7}7zEU97_2pXQ&G1OlWBEGPs4Kh^6gV%Yd&(XTN8N?yk?#8z>r?QkT*y6mRZ0NWehzA720>47{m0lH?W>xgrUwg3 z3f;WGMY%0XZ`tXF)|VPR=H?SVZ_8mI<#H~VT*ieWpGlWbbQ;5J!bMIzis z#WWO5$^Qo>o{wRw99t`mS3`AJ79;`XmO$ED;)rsh&?TK%o$?awNwXPaYM(lN9V)uh zgvtU#3q2&PLp$m5G;B{=!yee;CJZ0pA^;+gC1l{FIk$R|FQlL1@;zKGzv3y<-jb7M z^prG~bzeeqv2@iLBj$5Jo~g1Cy~et}4}D~B0X{%gl&vH);JJ0$8=Ox+-~=Dpve$+N z3!=k@DaMNA5lK)@w53S0qoOPnsypjLt%bESd2t}}%Jj9dkyXhXNps94ncx8i=5JQ%@zwMOSi0VZ#O$}0K6N=A*imt$3S?mp z3d&_1Mi= zyBqe46hHT+aNH@S+5Aw60s0@Pgq1J?O@mFE%cmy!^_b<~q?Pu8IOoKqsUg|ZjK6I5 z$W#Hy#6cew8{-HiV>4nHm^5$1sEU=iHI)+a+?JZS=|0S=Lu;Zi^b!mm30vAq7%_DH zc$_p>(hYp^uF3;s_AOTC!w}shf5w$FL}&qFj*qDV-b}wI>Y}+DFGN@vnIEfhG3ncw zRuy@%8a=dSK`r`Ly32xVuBylglo}j{6OL>?lvP@<^vI~`X&TzYrAQc3 zH&(RlP*6`^c{@&thA~P|m8^jlDNSTBnU1(;nD558W9hQ*j1&=vVKfYp9ZJohDh+&+ z=IUt({r>G}4Zyb;7bIfNUinz<`ijq!2>b|StdBlmP$d?%AN41(q{A!CseB?3LhV(B zF&ZiW{yXX0Sa_dX@X`pxMzjdVZ9l7>5Kl|bAmIoHIoFxwL`tDa>F7_Ym<*FC~(Qt(`9#R zP9l(0J7Yk90J_bs=UQ58hD4Mtt*s+<$^3B`bbY1L-iyt_*-w103TxRmmokbhaZvgB z<+IYdV}KPQK(oJOFU-YwZQ*E|-yg;b1MSi$j@*PJ4)M-nLX66W7^`8A*#^X_W}QKc zo0BkU`QQX>yt@)9AEf8-@_TfQM1BQvlWK#b|LZ5}38-dMKw~xI8{)Pg(bsBpQk&JIA zkoj`t9K*qC+RXf!7M~>f(MqL#oW6mJ-kds~AW;di-$18xHzMmqvK}k=pk+om;^7=4 z$AJw^9aO9pxGY_foKI4RK5q=8e_D~6%reB~+yek?)PmF_9Py;f1K#CCTkS^X)y$DM zC<4F-&(JWaiUTInYC-!-iB1VMX+He|_xfRMYxm-4V`m*_Pbh?8A&WgH!vR}2iM@$V zGfXw*{jtAAaIgUOqg^95e-zc;>*0P3;zLGZFaS6&YEOclk?sJN#`x~mP)G9khiYkA z?C5$4#2J>@38}ZpKq5Pk7J@S0So0~Yg^7Or1@85e^a`$jKSFS+Zh(?9f|P)C_$`hL zi(Fv~UCem~0S0Mtc#t7$Ov%*hD?eM6pr5AippoKK4GTab2B<;FmZZ;&HbZ2jM<21o zB1kdboWY(ht7M|dAgzE!svTm)#(wvDe1J5>P2^?h8G;Sb;W^;2tWdhdSjqy5L%&aEvU?I#2?wES5NR$pDc^4!TM z{Kud7pZs=CyZ-Bc{o60^!-IN%{;p4~UAn!idPkDXXVj;YN>Oz6%K1NU zH!q#Pcv|JxO^Dob09ldw2 z-6w@_|Id@Vf%~X_IQBa6|4>dYz8im!A6xi);^0ps{iKY$A@syC86TKCGX4m+C&k>2 zg2#@D`Q05cx6yZU_1%7a3f_csb~UboH;e_8jCK0Wx|;rkH(560WiuuAt2_x?9;1?{B( literal 0 HcmV?d00001 diff --git a/examples/offline_inference/qwen_2_5_omni/utils.py b/examples/offline_inference/qwen_2_5_omni/utils.py index 730d2a077..42894c4a7 100644 --- a/examples/offline_inference/qwen_2_5_omni/utils.py +++ b/examples/offline_inference/qwen_2_5_omni/utils.py @@ -13,6 +13,10 @@ from processing_omni import fetch_image, fetch_video 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(): @@ -113,16 +117,25 @@ def make_inputs_qwen2_omni( from transformers import AutoConfig, AutoProcessor - processor = AutoProcessor.from_pretrained(args.model) + # 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] - try: - config = AutoConfig.from_pretrained(args.model) - if "Qwen2_5OmniModel" in config.architectures: - args.legacy_omni_video = False + 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 = not ("Qwen2_5OmniModel" in config.architectures) else: args.legacy_omni_video = True - except Exception: - args.legacy_omni_video = True audios, images, videos = [], [], [] for message in messages: diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index 458e000e5..63dab9fde 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/vllm_omni/entrypoints/omni_llm.py @@ -178,6 +178,7 @@ class OmniLLM(BaseOmniLLM): batch_timeout=self.batch_timeout, ) logger.debug("[Orchestrator] Stage-%s process started", stage_id) + time.sleep(self._init_sleep_seconds) def close(self) -> None: for q in self._stage_in_queues: From edae58b137074ac7002e3253974a6e92df330f6d Mon Sep 17 00:00:00 2001 From: tzhouam Date: Tue, 11 Nov 2025 10:03:26 +0800 Subject: [PATCH 13/19] modify code based on review suggestions Signed-off-by: tzhouam --- .../offline_inference/qwen_2_5_omni/run.sh | 2 +- vllm_omni/__init__.py | 4 +- vllm_omni/entrypoints/omni_llm.py | 102 +++++------------- vllm_omni/entrypoints/omni_stage.py | 10 +- .../model_executor/models/qwen2_5_omni.py | 14 +-- .../stage_configs/qwen2_5_omni.yaml | 10 +- 6 files changed, 44 insertions(+), 98 deletions(-) diff --git a/examples/offline_inference/qwen_2_5_omni/run.sh b/examples/offline_inference/qwen_2_5_omni/run.sh index bee28e2c4..e27d21143 100644 --- a/examples/offline_inference/qwen_2_5_omni/run.sh +++ b/examples/offline_inference/qwen_2_5_omni/run.sh @@ -6,4 +6,4 @@ python end2end.py --model Qwen/Qwen2.5-Omni-7B \ --output-wav output_audio \ --prompt_type text \ --init-sleep-seconds 0 \ - --pt-prompts top100.pt + --pt-prompts .pt diff --git a/vllm_omni/__init__.py b/vllm_omni/__init__.py index 6860b462c..afeb09c1b 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 BaseOmniLLM +from .entrypoints.omni_llm import OmniLLM __all__ = [ # Main components - "BaseOmniLLM", + "OmniLLM", # "AsyncOmniLLM", # TODO: add AsyncOmniLLM back when it is implemented # Configuration "OmniModelConfig", diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index 63dab9fde..c52e2e969 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/vllm_omni/entrypoints/omni_llm.py @@ -1,6 +1,7 @@ from typing import Any, Optional, Sequence, Union import logging import multiprocessing as mp +from concurrent.futures import ThreadPoolExecutor, as_completed import os import queue import sys @@ -43,74 +44,7 @@ from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) -class BaseOmniLLM: - def __init__( - self, model: str, stage_configs=None, log_stats: bool = False, **kwargs - ): - if stage_configs is None: - self.initialize_stage_configs(model) - else: - self.stage_configs = stage_configs - - self.stage_list = [] - self.initialize_stages(model) - - def initialize_stage_configs(self, model: str): - self.stage_configs = load_stage_configs_from_model(model) - - def initialize_stages(self, model: str): - for stage_config in self.stage_configs: - stage = OmniStage(stage_config) - stage_llm = OmniStageLLM(model=model, **stage_config.engine_args) - stage.set_engine(stage_llm) - self.stage_list.append(stage) - - def generate( - self, - prompts: Union[PromptType, Sequence[PromptType]], - sampling_params_list: Optional[ - Union[SamplingParams, Sequence[SamplingParams]] - ] = None, - ) -> list[OmniRequestOutput]: - """Generate text outputs for the given prompts.""" - final_outputs: list[OmniRequestOutput] = [] - if len(sampling_params_list) != len(self.stage_list): - raise ValueError( - f"Expected {len(self.stage_list)} sampling params, " - f"got {len(sampling_params_list)}" - ) - for stage_id, stage in enumerate(self.stage_list): - if stage_id > 0: - engine_inputs = stage.process_engine_inputs(self.stage_list, prompts) - else: - engine_inputs = prompts - engine_outputs = self._run_generation( - stage, sampling_params_list[stage_id], engine_inputs - ) - stage.set_engine_outputs(engine_outputs) - if hasattr(stage, "final_output") and stage.final_output: - final_outputs.append( - OmniRequestOutput( - stage_id=stage_id, - final_output_type=stage.final_output_type, - request_output=engine_outputs, - ) - ) - return final_outputs - - def _run_generation( - self, - stage: OmniStage, - sampling_params: SamplingParams, - prompts: Union[PromptType, Sequence[PromptType]], - ): - engine_outputs = [] - for ro in stage.engine.generate(prompts, sampling_params): - engine_outputs.append(ro) - return engine_outputs - - -class OmniLLM(BaseOmniLLM): +class OmniLLM: """Multi-process pipelined OmniLLM. - Per-stage process with copy-based IPC (Queues) @@ -135,16 +69,7 @@ class OmniLLM(BaseOmniLLM): 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: @@ -152,13 +77,34 @@ class OmniLLM(BaseOmniLLM): configure_orchestrator_logger(logger, self._log_file) self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file) + self._initialize_stages(model, init_sleep_seconds, shm_threshold_bytes, init_timeout) + def _initialize_stages(self, model: str, init_sleep_seconds: int, shm_threshold_bytes: int, init_timeout: int) -> None: + self.stage_list: list[OmniStage] = [] + # Build OmniStage instances in parallel, preserve original order + def _build_stage(idx_cfg: tuple[int, Any]) -> tuple[int, OmniStage]: + idx, cfg = idx_cfg + return idx, OmniStage(cfg) + + with ThreadPoolExecutor(max_workers=min(len(self.stage_configs), max(1, os.cpu_count() or 1))) as executor: + futures = [executor.submit(_build_stage, (idx, cfg)) for idx, cfg in enumerate(self.stage_configs)] + results: list[tuple[int, OmniStage]] = [] + for fut in as_completed(futures): + results.append(fut.result()) + results.sort(key=lambda x: x[0]) + self.stage_list = [st for _, st in results] + 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] = [] 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): @@ -400,7 +346,7 @@ class OmniLLM(BaseOmniLLM): # Provide actionable suggestions before shutdown try: suggestions = [ - "Verify GPU/device assignment in config (mrs.devices) is correct.", + "Verify GPU/device assignment in config (runtime.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).", diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 06fb12041..6ec93878e 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -99,11 +99,11 @@ class OmniStage: ctx = ctx or mp.get_context("spawn") # Prepare lightweight dict config for worker engine_args = _to_dict(self.engine_args) - mrs_cfg = _to_dict(getattr(self.stage_config, "mrs", {})) + runtime_cfg = _to_dict(getattr(self.stage_config, "runtime", {})) stage_payload: Dict[str, Any] = { "stage_id": self.stage_id, "engine_args": engine_args, - "mrs": mrs_cfg, + "runtime": runtime_cfg, "shm_threshold_bytes": self._shm_threshold_bytes, } self._proc = ctx.Process( @@ -197,7 +197,7 @@ def _stage_worker( stage_id = stage_payload["stage_id"] engine_args = stage_payload.get("engine_args", {}) - mrs_cfg = stage_payload.get("mrs", {}) + runtime_cfg = stage_payload.get("runtime", {}) shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536)) # Per-stage file logger (optional) @@ -228,7 +228,7 @@ def _stage_worker( # Device mapping try: - set_stage_gpu_devices(stage_id, mrs_cfg.get("devices")) + set_stage_gpu_devices(stage_id, runtime_cfg.get("devices")) except Exception as e: _logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e) @@ -250,7 +250,7 @@ def _stage_worker( _logging.getLogger(__name__).debug("[Stage-%s] Received shutdown signal", stage_id) break - max_batch_size = int(mrs_cfg.get("max_batch_size", 1) or 1) + max_batch_size = int(runtime_cfg.get("max_batch_size", 1) or 1) batch_tasks: List[Dict[str, Any]] = [task] if max_batch_size > 1: while len(batch_tasks) < max_batch_size: diff --git a/vllm_omni/model_executor/models/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni.py index e78f77dad..87df50504 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni.py @@ -269,13 +269,13 @@ class Qwen2_5OmniForConditionalGeneration( # 2) Talker (if codec not provided) if self.model_stage == "talker": - # Mixed-mode support: in一个step内同时支持 Prefill*n 与 Decode*n。 - # 规则： - # - 使用特殊token将 Prefill 片段包裹为 [BOS][PAD...][EOS] - # - Decode 片段则为“单一非特殊token” - # - 若提供了 additional_information（可为按请求分割的 list 或拼接后的 tensor + 形状列表）， - # 则对 Prefill 片段按请求重建 thinker→talker 的输入嵌入； - # - 对 Decode 片段，若提供了 per-request 的 decode 辅助嵌入（可选），则进行叠加；否则保留原嵌入。 + # Mixed-mode support: In a single step, both Prefill*n and Decode*n are supported. + # Rules: + # - Prefill segments are wrapped with special tokens: [BOS][PAD...][EOS] + # - Decode segments consist of a single non-special token. + # - If additional_information is provided (can be a list split by request or a concatenated tensor plus a list of shapes), + # then for each request, reconstruct the thinker→talker input embeddings for the Prefill segments; + # - For Decode segments, if per-request auxiliary decode embeddings are provided (optional), add them; otherwise, keep the original embedding. if input_ids is None and additional_information is None: input_ids = torch.zeros( 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 2f69b3080..620da3a89 100644 --- a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml @@ -1,7 +1,7 @@ # stage config for running qwen2.5-omni with architecture of OmniLLM. stage_args: - stage_id: 0 - mrs: + runtime: process: true # Run this stage in a separate process devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device) max_batch_size: 1 @@ -18,7 +18,7 @@ stage_args: final_output: true final_output_type: text - stage_id: 1 - mrs: + runtime: process: true devices: "1" max_batch_size: 3 @@ -35,7 +35,7 @@ stage_args: engine_input_source: [0] custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker - stage_id: 2 - mrs: + runtime: process: true devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU max_batch_size: 1 @@ -53,8 +53,8 @@ stage_args: final_output: true final_output_type: audio -# Top-level MRS config (concise): default windows and stage edges -mrs: +# Top-level runtime config (concise): default windows and stage edges +runtime: enabled: true defaults: window_size: -1 # Simplified: trigger downstream only after full upstream completion From b943abb22fdcfc2967d6bad49eea49121cd2ce19 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Tue, 11 Nov 2025 10:49:53 +0800 Subject: [PATCH 14/19] add data download instructions Signed-off-by: tzhouam --- .../offline_inference/qwen_2_5_omni/README.md | 11 +++- .../qwen_2_5_omni/extract_prompts.py | 48 ++++++++++++++++++ .../offline_inference/qwen_2_5_omni/top100.pt | Bin 9319 -> 0 bytes 3 files changed, 58 insertions(+), 1 deletion(-) create mode 100644 examples/offline_inference/qwen_2_5_omni/extract_prompts.py delete mode 100644 examples/offline_inference/qwen_2_5_omni/top100.pt diff --git a/examples/offline_inference/qwen_2_5_omni/README.md b/examples/offline_inference/qwen_2_5_omni/README.md index 770c5e169..08c5c3400 100644 --- a/examples/offline_inference/qwen_2_5_omni/README.md +++ b/examples/offline_inference/qwen_2_5_omni/README.md @@ -63,7 +63,16 @@ uv pip install -r requirements.txt ``` ## Run examples (Qwen2.5-omni) - +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/qwen_2_5_omni/meta.lst +python3 examples/offline_inference/qwen_2_5_omni/extract_prompts.py \ + --input examples/offline_inference/qwen_2_5_omni/meta.lst \ + --output examples/offline_inference/qwen_2_5_omni/top100.txt \ + --topk 100 \ + --pt-output examples/offline_inference/qwen_2_5_omni/top100.pt +``` Get into the example folder ```bash cd examples/offline_inference/qwen_2_5_omni diff --git a/examples/offline_inference/qwen_2_5_omni/extract_prompts.py b/examples/offline_inference/qwen_2_5_omni/extract_prompts.py new file mode 100644 index 000000000..a1bc7a78f --- /dev/null +++ b/examples/offline_inference/qwen_2_5_omni/extract_prompts.py @@ -0,0 +1,48 @@ +#!/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/top100.pt b/examples/offline_inference/qwen_2_5_omni/top100.pt deleted file mode 100644 index 017c60ac806dc45952b6bf652fd60ad362ece700..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 9319 zcmd6tU6UNudB+zpg4i*TIKd`1<{&#Pah7JqN?P$M5FoP_fZeeU$nqZ6E6rGvdV5UV#n)_BK z>i)DwtDW_kshl*L`4$c5mVf=&SkrnxTMx!&z7JiI1+=o2nNMA1w5V$rMKv*#nNCln-*A3-ZB$!%XoC$Rhd~)roK6Nxhj3c#=a1NTF}*wg~)WyZCD;el{FYx8kVUCCx)Clk{+U z37u|ui&Z9oB`;kF7}2)AAFys}+a7GNvqQhC18E+Enkmjon!mWWjYl7WTm{+20Bj7f0KYlGaHvSFP%R8ba#dz$^$bKO%~wuC=wGI;H?(9}~RWk6ZzhF`%Kb~j;W+b;f`G>@jQBYWSK?yv~rh73}f z*#l=5zBhS48*P|T---0}i!C|%Nb%)!CVhc< zaA%n-AS`Twd82v69ikhE7Nhfax1i8?8|-g94rqy+S<-y%6lDHQ`U-l!Lr}I(6bS7P zTpk>pUv$gdgs=3?j6J%TF=eV z+%90AF7y=?7#aAcOluVzwv3XsO)-NYISbtFGg~g|0Q>#jsqgU3^fG6-Qna{y>=YDJ zSF{{v9LN%QHi9RObb28VZ^v1QDpz*YQ%gy{l5I)VRq8s$c_Y+fkQZk9d>JhS9I^(lM+Y4 zWug}G#&Cp1!~R1{pDqsc}|T_2U<}A$;x6jx^)aY{Dx*Z@kEAgx zmG*R+a?%e&48g1|>duTx@QJS!J6K*AJAVAz1lzi%FIqTTW`H_T)2sez;axK-NOWh@ zJsC2xYwIW?dK?K+>+E=MPQvzzmDUj?T%P>+!| zHYlE|h(3?mWHCa1)cU!!%2dcmx1h>0)~N(xYjI$x5u7yNjl8ojlBAY^A;wXO&d(uT z)y}pYqhyF=QP`6Fsw@t+qT7}7zEU97_2pXQ&G1OlWBEGPs4Kh^6gV%Yd&(XTN8N?yk?#8z>r?QkT*y6mRZ0NWehzA720>47{m0lH?W>xgrUwg3 z3f;WGMY%0XZ`tXF)|VPR=H?SVZ_8mI<#H~VT*ieWpGlWbbQ;5J!bMIzis z#WWO5$^Qo>o{wRw99t`mS3`AJ79;`XmO$ED;)rsh&?TK%o$?awNwXPaYM(lN9V)uh zgvtU#3q2&PLp$m5G;B{=!yee;CJZ0pA^;+gC1l{FIk$R|FQlL1@;zKGzv3y<-jb7M z^prG~bzeeqv2@iLBj$5Jo~g1Cy~et}4}D~B0X{%gl&vH);JJ0$8=Ox+-~=Dpve$+N z3!=k@DaMNA5lK)@w53S0qoOPnsypjLt%bESd2t}}%Jj9dkyXhXNps94ncx8i=5JQ%@zwMOSi0VZ#O$}0K6N=A*imt$3S?mp z3d&_1Mi= zyBqe46hHT+aNH@S+5Aw60s0@Pgq1J?O@mFE%cmy!^_b<~q?Pu8IOoKqsUg|ZjK6I5 z$W#Hy#6cew8{-HiV>4nHm^5$1sEU=iHI)+a+?JZS=|0S=Lu;Zi^b!mm30vAq7%_DH zc$_p>(hYp^uF3;s_AOTC!w}shf5w$FL}&qFj*qDV-b}wI>Y}+DFGN@vnIEfhG3ncw zRuy@%8a=dSK`r`Ly32xVuBylglo}j{6OL>?lvP@<^vI~`X&TzYrAQc3 zH&(RlP*6`^c{@&thA~P|m8^jlDNSTBnU1(;nD558W9hQ*j1&=vVKfYp9ZJohDh+&+ z=IUt({r>G}4Zyb;7bIfNUinz<`ijq!2>b|StdBlmP$d?%AN41(q{A!CseB?3LhV(B zF&ZiW{yXX0Sa_dX@X`pxMzjdVZ9l7>5Kl|bAmIoHIoFxwL`tDa>F7_Ym<*FC~(Qt(`9#R zP9l(0J7Yk90J_bs=UQ58hD4Mtt*s+<$^3B`bbY1L-iyt_*-w103TxRmmokbhaZvgB z<+IYdV}KPQK(oJOFU-YwZQ*E|-yg;b1MSi$j@*PJ4)M-nLX66W7^`8A*#^X_W}QKc zo0BkU`QQX>yt@)9AEf8-@_TfQM1BQvlWK#b|LZ5}38-dMKw~xI8{)Pg(bsBpQk&JIA zkoj`t9K*qC+RXf!7M~>f(MqL#oW6mJ-kds~AW;di-$18xHzMmqvK}k=pk+om;^7=4 z$AJw^9aO9pxGY_foKI4RK5q=8e_D~6%reB~+yek?)PmF_9Py;f1K#CCTkS^X)y$DM zC<4F-&(JWaiUTInYC-!-iB1VMX+He|_xfRMYxm-4V`m*_Pbh?8A&WgH!vR}2iM@$V zGfXw*{jtAAaIgUOqg^95e-zc;>*0P3;zLGZFaS6&YEOclk?sJN#`x~mP)G9khiYkA z?C5$4#2J>@38}ZpKq5Pk7J@S0So0~Yg^7Or1@85e^a`$jKSFS+Zh(?9f|P)C_$`hL zi(Fv~UCem~0S0Mtc#t7$Ov%*hD?eM6pr5AippoKK4GTab2B<;FmZZ;&HbZ2jM<21o zB1kdboWY(ht7M|dAgzE!svTm)#(wvDe1J5>P2^?h8G;Sb;W^;2tWdhdSjqy5L%&aEvU?I#2?wES5NR$pDc^4!TM z{Kud7pZs=CyZ-Bc{o60^!-IN%{;p4~UAn!idPkDXXVj;YN>Oz6%K1NU zH!q#Pcv|JxO^Dob09ldw2 z-6w@_|Id@Vf%~X_IQBa6|4>dYz8im!A6xi);^0ps{iKY$A@syC86TKCGX4m+C&k>2 zg2#@D`Q05cx6yZU_1%7a3f_csb~UboH;e_8jCK0Wx|;rkH(560WiuuAt2_x?9;1?{B( From 12cc2cd245bbd1ac36bb6ecc74cec385bd3ecf67 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Tue, 11 Nov 2025 10:55:57 +0800 Subject: [PATCH 15/19] rename util Signed-off-by: tzhouam --- vllm_omni/entrypoints/{pipeline_utils.py => stage_utils.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename vllm_omni/entrypoints/{pipeline_utils.py => stage_utils.py} (100%) diff --git a/vllm_omni/entrypoints/pipeline_utils.py b/vllm_omni/entrypoints/stage_utils.py similarity index 100% rename from vllm_omni/entrypoints/pipeline_utils.py rename to vllm_omni/entrypoints/stage_utils.py From bc69d8e931c403ccc5a0da41e134b015b2f0ee92 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Tue, 11 Nov 2025 10:57:15 +0800 Subject: [PATCH 16/19] rename util further rename imports Signed-off-by: tzhouam --- vllm_omni/entrypoints/omni_llm.py | 2 +- vllm_omni/entrypoints/omni_stage.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index c52e2e969..c1b12a1b2 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/vllm_omni/entrypoints/omni_llm.py @@ -34,7 +34,7 @@ from vllm_omni.entrypoints.log_utils import ( init_stats_paths, OrchestratorMetrics, ) -from vllm_omni.entrypoints.pipeline_utils import ( +from vllm_omni.entrypoints.stage_utils import ( maybe_load_from_ipc as _load, encode_for_ipc as _encode, serialize_obj as _ser, diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 6ec93878e..8f95cf62e 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -11,7 +11,7 @@ import logging import multiprocessing as mp from typing import Any, Dict, List, Optional, Union -from vllm_omni.entrypoints.pipeline_utils import ( +from vllm_omni.entrypoints.stage_utils import ( _to_dict, set_stage_gpu_devices, maybe_load_from_ipc_with_metrics, From c12bd006bc0d97a6d5693d7fcea23f323aaefe8b Mon Sep 17 00:00:00 2001 From: tzhouam Date: Tue, 11 Nov 2025 11:56:08 +0800 Subject: [PATCH 17/19] Make Example Easier To Use Signed-off-by: tzhouam --- .../offline_inference/qwen_2_5_omni/README.md | 20 +++++++++----- .../qwen_2_5_omni/end2end.py | 26 ++++++++++--------- .../qwen_2_5_omni/extract_prompts.py | 11 +++----- .../{run.sh => run_multiple_prompts.sh} | 2 +- .../qwen_2_5_omni/run_single_prompt.sh | 9 +++++++ 5 files changed, 41 insertions(+), 27 deletions(-) rename examples/offline_inference/qwen_2_5_omni/{run.sh => run_multiple_prompts.sh} (86%) create mode 100644 examples/offline_inference/qwen_2_5_omni/run_single_prompt.sh diff --git a/examples/offline_inference/qwen_2_5_omni/README.md b/examples/offline_inference/qwen_2_5_omni/README.md index 08c5c3400..734060b6b 100644 --- a/examples/offline_inference/qwen_2_5_omni/README.md +++ b/examples/offline_inference/qwen_2_5_omni/README.md @@ -63,22 +63,30 @@ uv pip install -r requirements.txt ``` ## Run examples (Qwen2.5-omni) -Download dataset from seed_tts[https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit]. To get the prompt, you can: +### 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/qwen_2_5_omni/meta.lst python3 examples/offline_inference/qwen_2_5_omni/extract_prompts.py \ --input examples/offline_inference/qwen_2_5_omni/meta.lst \ --output examples/offline_inference/qwen_2_5_omni/top100.txt \ - --topk 100 \ - --pt-output examples/offline_inference/qwen_2_5_omni/top100.pt + --topk 100 ``` Get into the example folder ```bash cd examples/offline_inference/qwen_2_5_omni ``` -Modify PYTHONPATH in run.sh as your path of vllm_omni. Then run. +Modify PYTHONPATH in run_multiple_prompts.sh as your path of vllm_omni. Then run. ```bash -bash run.sh +bash run_multiple_prompts.sh +``` +### Single Prompts +Get into the example folder +```bash +cd examples/offline_inference/qwen_2_5_omni +``` +Modify PYTHONPATH in run_single_prompt.sh as your path of vllm_omni. Then run. +```bash +bash run_single_prompt.sh ``` -The output audio is saved in ./output_audio diff --git a/examples/offline_inference/qwen_2_5_omni/end2end.py b/examples/offline_inference/qwen_2_5_omni/end2end.py index 8614d4653..bbd4e0f53 100644 --- a/examples/offline_inference/qwen_2_5_omni/end2end.py +++ b/examples/offline_inference/qwen_2_5_omni/end2end.py @@ -51,9 +51,6 @@ def parse_args(): 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." ) @@ -146,6 +143,12 @@ def parse_args(): 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).", + ) args = parser.parse_args() return args @@ -153,20 +156,19 @@ def parse_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}") + # 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, "r", 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 pt-prompts: {e}") + print(f"[Error] Failed to load prompts: {e}") raise if args.prompts is None: - raise ValueError("No prompts provided. Use --prompts ... or --pt-prompts (with --prompt_type text)") + 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, diff --git a/examples/offline_inference/qwen_2_5_omni/extract_prompts.py b/examples/offline_inference/qwen_2_5_omni/extract_prompts.py index a1bc7a78f..7380e701b 100644 --- a/examples/offline_inference/qwen_2_5_omni/extract_prompts.py +++ b/examples/offline_inference/qwen_2_5_omni/extract_prompts.py @@ -18,10 +18,9 @@ def extract_prompt(line: str) -> Optional[str]: 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）") + 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 = [] @@ -37,10 +36,6 @@ def main() -> None: 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/run.sh b/examples/offline_inference/qwen_2_5_omni/run_multiple_prompts.sh similarity index 86% rename from examples/offline_inference/qwen_2_5_omni/run.sh rename to examples/offline_inference/qwen_2_5_omni/run_multiple_prompts.sh index e27d21143..ba07616ce 100644 --- a/examples/offline_inference/qwen_2_5_omni/run.sh +++ b/examples/offline_inference/qwen_2_5_omni/run_multiple_prompts.sh @@ -6,4 +6,4 @@ python end2end.py --model Qwen/Qwen2.5-Omni-7B \ --output-wav output_audio \ --prompt_type text \ --init-sleep-seconds 0 \ - --pt-prompts .pt + --txt-prompts top100.txt diff --git a/examples/offline_inference/qwen_2_5_omni/run_single_prompt.sh b/examples/offline_inference/qwen_2_5_omni/run_single_prompt.sh new file mode 100644 index 000000000..8dc2c74e5 --- /dev/null +++ b/examples/offline_inference/qwen_2_5_omni/run_single_prompt.sh @@ -0,0 +1,9 @@ +export PYTHONPATH=/path/to/vllm-omni:$PYTHONPATH +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 \ + --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words." From 26fcae5def1fe4a6750b886d569c31c21597f62a Mon Sep 17 00:00:00 2001 From: tzhouam Date: Tue, 11 Nov 2025 14:49:45 +0800 Subject: [PATCH 18/19] Debug and Update doc Signed-off-by: tzhouam --- docs/mrs_design.md | 61 +++++++++++++++--------------- vllm_omni/entrypoints/log_utils.py | 2 +- 2 files changed, 31 insertions(+), 32 deletions(-) diff --git a/docs/mrs_design.md b/docs/mrs_design.md index 2c463fa64..28b59ac25 100644 --- a/docs/mrs_design.md +++ b/docs/mrs_design.md @@ -1,4 +1,4 @@ -## Multi-Request Streaming (MRS) on a Single Machine — Design Doc +## Multi-Request Streaming (MRS) on a Single Machine - Author: Team - Date: 2025-10-27 @@ -67,7 +67,7 @@ #### Example config (YAML, simplified) ```yaml -mrs: +runtime: enabled: true defaults: window_size: -1 # Full-trigger across stages @@ -90,12 +90,11 @@ mrs: 3) t2: B handles req0 while A handles req1 (parallel across stages) ### 9. Integration Points (by file) -- `vllm_omni/entrypoints/omni_llm.py` - - Introduce a ProcessOrchestrator (replace threaded version): - - Spawn stage processes per config (set `CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`), create control/data channels, build window state machines. - - Add `enable_mrs` in `generate`; when enabled, do not collect `stage.engine.generate(...)` into lists; stream segments to downstream via IPC instead. - - Register `window_size` per edge; `-1` triggers downstream once `is_last` is observed. - - Manage process lifecycle: start/heartbeat/graceful shutdown/restart on crash; propagate cancel/timeout. +- `vllm_omni/entrypoints/omni_llm.py` (Orchestrator) + - Class `OmniLLM` orchestrates multi-process stages; constructs `OmniStage` instances in parallel and spawns per-stage workers. + - Spawns stage processes per config (set `CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`), creates control/data channels, builds simple full-trigger flow. + - Stats/logging are disabled by default; per-stage and orchestrator stats are only written when explicitly enabled. + - Manages process lifecycle: start/wait for readiness, graceful shutdown; forwards results between stages using copy-based IPC and optional SHM. - `vllm_omni/outputs.py` - Add in-process events: `ARSegmentEvent`, `ImageProgressEvent`, `CompletedEvent`, `ErrorEvent`. - Keep `OmniRequestOutput` as the external representation of final stage outputs. @@ -141,10 +140,10 @@ class WindowState: ``` ### 11. Configuration -- `mrs.enabled`, `mrs.ar.default_window_size`, `mrs.ar.min_flush_interval_ms`, `mrs.ar.max_segment_tokens`, `mrs.ar.scheduler_policy`. +- `runtime.enabled`, `runtime.ar.default_window_size`, `runtime.ar.min_flush_interval_ms`, `runtime.ar.max_segment_tokens`, `runtime.ar.scheduler_policy`. - Devices/processes: - - `mrs.stages`: [{ name, process: true, devices: "0" | "1,2", extra_env: {} }] - - `mrs.edges`: [{ from, to, window_size }] + - Per-stage: `runtime` block on each stage with fields like `process`, `devices`, `max_batch_size`. + - Top-level: `runtime.edges`: [{ from, to, window_size }] - Pinned/SHM pool sizing: estimate by concurrency and max sequence; tune watermarks and GC policies; configure CUDA IPC limits and cleanup thresholds if enabled. ### 12. Observability @@ -175,36 +174,36 @@ class WindowState: Objective: without changing external call patterns, encapsulate per-stage LLM init, process creation, and worker logic inside `Stage`, so `PipelinedOmniLLM` focuses on orchestration (seeding, polling, forwarding, collecting) while keeping multi-process, shared memory, and device mapping capabilities. Key changes (vs. current): -- Before: `PipelinedOmniLLM` created processes directly, passing `_stage_worker`; device setup, LLM init, batching, and SHM were implemented there. -- After: enhance `Stage` to start its own subprocess via `start_process()`, with `worker_main()` as the entry; device setup, LLM init, batching, and SHM live in `Stage`. +- Before: Orchestrator created processes directly, passing `_stage_worker`; device setup, LLM init, batching, and SHM were implemented there. +- After: enhanced `OmniStage` owns its subprocess (`init_stage_worker`), with `_stage_worker` as the entry; device setup, LLM init, batching, and SHM live in `OmniStage`. Classes & responsibilities -- PipelinedOmniLLM (orchestrator) - - Build `Stage` list (preserve `OmniStage` init timing and `process_engine_inputs` flow) +- OmniLLM (orchestrator) + - Build `OmniStage` list in parallel (preserve `process_engine_inputs` wiring) - Connect adjacent `Stage` input/output queues; seed requests to stage 0 - Poll stage outputs, decode results, call `process_engine_inputs`, then encode and forward to next stage - Termination & cleanup: distribute shutdown signals; join/terminate subprocesses - Stage (stage unit) - - Members: `stage_config` (with `engine_args` and `mrs`), `in_q/out_q`, subprocess handle, stats + - Members: `stage_config` (with `engine_args` and `runtime`), `in_q/out_q`, subprocess handle, stats - API: - - `start_process(in_q, out_q)`: spawn subprocess - - `stop_process()`: graceful exit - - `submit(ipc_payload)`: submit to `in_q` (may use shared memory) - - `try_collect() -> Optional[payload]`: non-blocking get from `out_q` + - `init_stage_worker(...)`: spawn subprocess + - `stop_stage_worker()`: graceful exit + - `submit(ipc_payload)`: submit to input queue (may use shared memory) + - `try_collect() -> Optional[payload]`: non-blocking get from output queue - `process_engine_inputs(stages, prompts)`: reuse/delegate existing logic - Subprocess entry `worker_main()`: - - set_stage_gpu_devices (`CUDA_VISIBLE_DEVICES` & logical index mapping) + - `set_stage_gpu_devices` (`CUDA_VISIBLE_DEVICES` & logical index mapping) - build `OmniStageLLM(model, **engine_args)` - while-loop: - - take first task + up to `mrs.max_batch_size-1` non-blocking; support `max_batch_size>=1` - - `maybe_load_from_ipc` per `engine_inputs` in batch + - take first task + up to `runtime.max_batch_size-1` non-blocking; support `max_batch_size>=1` + - `maybe_load_from_ipc_with_metrics` per `engine_inputs` in batch - `stage_engine.generate(batched_inputs, sampling_params)` (window `-1`) - - write per-request aggregated outputs to `out_q` after `maybe_dump_to_shm` + - write per-request aggregated outputs to output queue after `maybe_dump_to_shm` - log start/finish; report `error` per request on exceptions Data flow (sketch) ``` -PipelinedOmniLLM +OmniLLM ├─ Stage[0].start_process(in0, out0) ├─ Stage[1].start_process(in1, out1) ├─ ... @@ -219,11 +218,11 @@ PipelinedOmniLLM ``` Config (preserve & extend) -- `mrs.devices`: +- `runtime.devices`: - Comma-separated visibility list (e.g., "2,5,7"): set visible set; default logical index 0 (first GPU) - Number/string: treat as logical index (0/1-based per implementation) mapping to physical GPUs -- `mrs.max_batch_size`: max requests taken per intake (default 1) -- `shm_threshold_bytes`: objects larger than this go via shared memory (default 64KB; can be injected into stages) +- `runtime.max_batch_size`: max requests taken per intake (default 1) +- `shm_threshold_bytes`: objects larger than this go via shared memory (default 64KB; injected into stages) Kept & strengthened - Multi-process: each stage in its own process; parallel across stages; serial/batching inside a stage controlled by `max_batch_size` @@ -232,7 +231,7 @@ Kept & strengthened - Windows: keep `window_size=-1` semantics (downstream triggers only after upstream finishes the request) Migration (execution order) -1) Implement enhanced `Stage` (e.g., `entrypoints/stage.py`) with `start_process/stop_process/submit/try_collect/worker_main`. -2) Make `PipelinedOmniLLM` hold `Stage` instances, replacing direct queues and `_stage_worker` with `Stage` APIs. -3) Reuse `pipeline_utils`: device/SHM/encode-decode utilities stay in utils. +1) Enhanced `OmniStage` with `init_stage_worker/stop_stage_worker/submit/try_collect/_stage_worker`. +2) `OmniLLM` holds `OmniStage` instances, replacing the older pipelined class; uses Stage APIs. +3) Reuse `entrypoints/stage_utils`: device/SHM/encode-decode utilities live here. 4) Regression: single/multi-request, various `max_batch_size`, SHM threshold, device mapping, multi-stage parallelism, error paths, and shutdown flows. diff --git a/vllm_omni/entrypoints/log_utils.py b/vllm_omni/entrypoints/log_utils.py index f0c8355df..11b98ed81 100644 --- a/vllm_omni/entrypoints/log_utils.py +++ b/vllm_omni/entrypoints/log_utils.py @@ -6,7 +6,7 @@ import logging import os import time -from vllm_omni.entrypoints.pipeline_utils import append_jsonl as _append_jsonl +from vllm_omni.entrypoints.stage_utils import append_jsonl as _append_jsonl def remove_old_logs(log_file: Optional[str], num_stages: int) -> None: From 8f6280f4f65bcc24f0f429750cbcb0db5946e226 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Wed, 12 Nov 2025 15:51:56 +0800 Subject: [PATCH 19/19] adjust the doc Signed-off-by: tzhouam --- docs/mrs_design.md | 189 +++++++++++++++++---------------------------- 1 file changed, 70 insertions(+), 119 deletions(-) diff --git a/docs/mrs_design.md b/docs/mrs_design.md index 28b59ac25..50e3a8c15 100644 --- a/docs/mrs_design.md +++ b/docs/mrs_design.md @@ -1,9 +1,5 @@ ## Multi-Request Streaming (MRS) on a Single Machine -- Author: Team -- Date: 2025-10-27 -- Version: v0.3 (simplified: copy-based IPC / serial per stage / full-trigger) - ### 1. Background & Scope - All processing runs on a single physical machine with multi-process, per-stage workers. No proxy or network transport involved. - Current alignment with vllm-omni: `OmniLLM` supports multiple stages (`OmniStage`). GPU runners already expose streamable steps (prefill/decoding/diffusion), but the entry layer still collects lists and lacks intra-stage streaming and window scheduling. @@ -12,77 +8,30 @@ ### 2. Key Constraints - Multi-process per stage: each stage is an independent process with a while loop; device visibility can be configured (`CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`). - Simple IPC (copy-based): use `multiprocessing.Queue`/Pipe for inter-process communication with CPU copies/serialization; do not rely on CUDA IPC/SHM zero-copy in this version. -- Full-trigger across stages: unify `window_size = -1` for all edges — downstream starts only after upstream completes for that request. -- Serial within a stage: at most one in-flight request per stage (max_inflight=1). No batching/concurrency inside a stage in this version. - Cross-stage pipeline: different stages can process different requests concurrently (e.g., stage A handles request 1 while stage B handles request 0). ### 3. Architecture Overview - Processes & IPC queues - Each "sub-stage" is an OS process (worker). The loop: take from input_queue → compute → put to output_queue. - - Inter-stage connection via IPC: control plane (Unix Domain Socket / Pipe) + data plane (shared memory ring buffer in future; copy-based in this version). - - MPMC queue semantics; control plane passes light-weight pointers/metadata; data plane carries large payloads. + - Inter-stage connection via IPC: copy-based `multiprocessing.Queue` passing dict payloads; use shared memory for large objects. + - Each link is SPSC (single-producer/single-consumer): the upstream is the orchestrator and the downstream is a single stage process; queues are unbounded (maxsize=0) on the orchestrator side. - Device visibility - Each stage sets `CUDA_VISIBLE_DEVICES` or calls `torch.cuda.set_device` to bind to GPU sets. - A stage may use multiple GPUs internally (TP/PP/DP) but presents as a single stage unit. - Simplified IPC: copy-based queues/pipes for data transfer; zero-copy is future work. - Pipeline progression: when a stage finishes a request, it enqueues outputs to the downstream stage; if downstream is idle, it starts immediately. -- Scheduling & windows - - Each edge (stage→stage) defines a window policy; downstream triggers compute by consuming stitched window views. - - `-1` mode triggers only after upstream emits `is_last=True` (full aggregation). - - Backpressure propagates upstream via the control plane: once high-watermark is reached, upstream blocks or degrades behavior. - -### 4. Stage Breakdown -- AR (autoregressive) - - S0 Intake/Plan: parse request, assign request_id, output normalized tasks. - - S1 Prefill: tokenize/encode/KV warmup; may micro-batch across requests. - - S2 Decode Loop: step-by-step tokens; flush segments by `max_segment_tokens` and `min_flush_interval_ms`. - - S3 Postprocess: detokenize, logprobs/struct wrapping (CPU if necessary). - - S4 Consumer: local consumer callback/buffering. -- DiT/Diffusion - - D0 Intake/Plan: scheduler/seed init. - - D1 Sampling Loop: preview frames every k steps (downsample allowed). - - D2 Decode/Encode: VAE decode and JPEG/WEBP encode (prefer GPU→pinned CPU). - - D3 Postprocess/Consumer. - -### 5. Unified Events & Segment Structure (in-process) -- Base event: `{ request_id, sequence, event_type, ts, payload, device_handles }` -- ARSegmentEvent (segment): - - Fields: `tokens` (List[int] or GPU tensor view), optional `logprobs`, `is_last`, optional `hidden_states/aux` refs. - - Semantics: strictly ordered by `sequence` for the same `request_id`; `is_last` marks the final segment. -- ImageProgressEvent: `{ step, image_tensor|jpeg_ref, is_last }` (diffusion preview/completion). -- Completed/Error: terminal or error events affect only the current request. - -### 6. AR Segmentation & Windowing (simplified) -- Production (upstream stage) - - Accumulate tokens during decoding; flush a segment if any of the following holds: - - Reaches `max_segment_tokens` length - - Exceeds `min_flush_interval_ms` since last flush - - Request ended (`is_last=True`) - - Flush immediately via `output_queue.put(segment)`; do not wait for entire request completion. -- Consumption (downstream stage) - - Use unified `window_size = -1`: trigger only after receiving `is_last=True`; no window slicing. -- Backpressure & watermark - - Must-deliver events (token.delta, Completed) are not dropped; block upstream briefly on high watermark. - - For long sequences under `-1`, allow early segments to spill from GPU to pinned CPU to control VRAM pressure. - -#### Example config (YAML, simplified) -```yaml -runtime: - enabled: true - defaults: - window_size: -1 # Full-trigger across stages - max_inflight: 1 # Serial per stage -``` +- Scheduling + - A downstream stage triggers only after the upstream completes the request. + - Windowed segmentation/stitched triggering is not implemented; intra-stage streaming is not provided. ### 7. IPC Implementation (simplified: copy-based) - Use `multiprocessing.Queue`/Pipe for inter-process communication (control + data). - Data is serialized/copied via CPU; no CUDA IPC/SHM zero-copy in this version. -- Backpressure: blocking reads/writes based on queue capacity, blocking upstream on high watermark. +- Backpressure: queues are unbounded; pressure manifests as compute-rate differences. Optional SHM reduces large-object transfer cost; RX/decoding overhead is recorded for observability. ### 8. Scheduling & Cancellation (simplified) -- Serial FIFO: each stage `max_inflight=1`; no cross-request micro-batching or concurrency. -- Cancel/timeout: ongoing work should exit early; remaining queued requests stay pending. - Pipeline: when a stage finishes a request, it enqueues to the next stage; that stage immediately pulls the next request from its input queue, enabling cross-stage concurrency. +- Cancellation/timeout: explicit cancellation/timeouts are not provided; graceful shutdown uses a `None` sentinel sent to each stage input queue. #### Short sequence example (req0/req1, stage A→B) 1) t0: stage A handles req0 @@ -95,79 +44,77 @@ runtime: - Spawns stage processes per config (set `CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`), creates control/data channels, builds simple full-trigger flow. - Stats/logging are disabled by default; per-stage and orchestrator stats are only written when explicitly enabled. - Manages process lifecycle: start/wait for readiness, graceful shutdown; forwards results between stages using copy-based IPC and optional SHM. -- `vllm_omni/outputs.py` - - Add in-process events: `ARSegmentEvent`, `ImageProgressEvent`, `CompletedEvent`, `ErrorEvent`. - - Keep `OmniRequestOutput` as the external representation of final stage outputs. -- `vllm_omni/worker/gpu_ar_model_runner.py` - - Near `valid_sampled_token_ids` parsing (after updating `req_state.output_token_ids`), aggregate new tokens into segments and emit; send final `is_last=True` and Completed. - - Replace current `.to("cpu").contiguous()` sync path: by default keep GPU refs and let downstream copy asynchronously from a pinned pool when needed. -- `vllm_omni/worker/gpu_diffusion_model_runner.py` + - Stage readiness: each stage emits `{"type": "stage_ready"}` after initialization; the orchestrator waits for all stages or times out and logs diagnostic suggestions. ### 9.1 Process Device Visibility - Device binding: set `CUDA_VISIBLE_DEVICES` before process start, or call `torch.cuda.set_device` early in init; configs can be single/multi-GPU. - Cross-device transfer (simplified): via CPU copies; zero-copy paths are out of scope for now. -### 10. Worker Template & Window State (pseudo-code) +### 10. Worker Template (pseudo-code) ```python -def stage_worker(input_q, output_q, ctx): - while ctx.alive: - item = input_q.get() # blocking / with timeout - if item.is_cancelled(): - continue - with cuda_stream(ctx.stream): - out_segments = run_stage_compute(item) # produce segments - for seg in out_segments: - output_q.put(seg) - -class WindowState: - def __init__(self, window_size): - self.window_size = window_size - self.last_emitted = 0 - self.buffer = [] # references/paged views - def on_segment(self, seg): - self.buffer.append(seg) - if self.window_size == -1: - return self.emit_full_if_last(seg) - return self.emit_windows() - def emit_windows(self): - emits = [] - while self._pending_len() - self.last_emitted >= self.window_size: - emits.append(self._make_view(self.last_emitted, self.window_size)) - self.last_emitted += self.window_size - return emits - def emit_full_if_last(self, seg): - return [self._make_full_view()] if seg.is_last else [] +def stage_worker(input_q, output_q, runtime, shm_threshold): + # Device binding + set_stage_gpu_devices(runtime.devices) + engine = OmniStageLLM(model, **engine_args) + batch_seq = 0 + while True: + first = input_q.get() + if first is None: + break + # Batch intake (up to runtime.max_batch_size) + batch = [first] + while len(batch) < int(runtime.max_batch_size or 1) and not input_q.empty(): + nxt = input_q.get() + if nxt is None: + input_q.put(None) + break + batch.append(nxt) + # Decode IPC payload and sampling params + request_ids = [] + engine_inputs = [] + rx_bytes = {} + rx_decode_ms = {} + for t in batch: + ein, rx = maybe_load_from_ipc_with_metrics(t, "engine_inputs", "engine_inputs_shm") + request_ids.append(t["request_id"]) + engine_inputs.extend(ein if isinstance(ein, list) else [ein]) + rx_bytes[t["request_id"]] = rx["rx_transfer_bytes"] + rx_decode_ms[t["request_id"]] = rx["rx_decode_time_ms"] + # Generate and dispatch (grouped by request_id) + batch_seq += 1 + outputs = list(engine.generate(engine_inputs, batch[0]["sampling_params"], use_tqdm=False)) + grouped = group_by_request_id(outputs, request_ids) + for rid in request_ids: + r_out = grouped.get(rid, []) + use_shm, payload = maybe_dump_to_shm(r_out, shm_threshold) + msg = { + "request_id": rid, + "metrics": { + "batch_id": batch_seq, + "rx_transfer_bytes": int(rx_bytes.get(rid, 0)), + "rx_decode_time_ms": float(rx_decode_ms.get(rid, 0.0)), + # omitted: generation latency and token stats + }, + } + if use_shm: + msg["engine_outputs_shm"] = payload + else: + msg["engine_outputs"] = payload + output_q.put(msg) ``` -### 11. Configuration -- `runtime.enabled`, `runtime.ar.default_window_size`, `runtime.ar.min_flush_interval_ms`, `runtime.ar.max_segment_tokens`, `runtime.ar.scheduler_policy`. -- Devices/processes: - - Per-stage: `runtime` block on each stage with fields like `process`, `devices`, `max_batch_size`. - - Top-level: `runtime.edges`: [{ from, to, window_size }] -- Pinned/SHM pool sizing: estimate by concurrency and max sequence; tune watermarks and GC policies; configure CUDA IPC limits and cleanup thresholds if enabled. - ### 12. Observability -- Metrics: - - Stage-level: enqueue/dequeue rates, queue depth, blocking time on watermarks, window triggers/counts. - - Request-level: TTFT, TPST, completion latency, cancel rate. - - Resource-level: GPU utilization, VRAM usage/fragmentation, pinned usage. -- Logs/tracing: carry `request_id`, stage name, and `sequence`; sample decisions for downsampling/backpressure. - -### 13. Test Plan -- Unit: queue correctness, window state machine, cancel/timeout, memory pool rent/return. -- Performance: AR QPS/TPST vs. window/segment sizes; diffusion preview frequency vs. throughput. -- Stability: long-run memory leaks; backpressure trigger/recovery paths. -- Compatibility: behavior parity when MRS is disabled. - -### 14. Milestones -- M1: Process Orchestrator with simple IPC; serial per stage; full-trigger across stages (window=-1). -- M2: Hardening and observability; cancel/timeout and error paths. -- M3: Optional zero-copy and concurrency/micro-batching optimizations. +- Metrics (as implemented): + - per-request (emitted by stages): `num_tokens_out`, `stage_gen_time_ms`, `rx_transfer_bytes`, `rx_decode_time_ms`, `rx_in_flight_time_ms`, `batch_id` + - orchestrator aggregates: E2E latency, tokens/s (written only when stats are enabled) + - optional per-stage JSONL: `{log_file}.stage{stage_id}.stats.jsonl` +- Logs/tracing: + - optional per-stage log files: `{log_file}.stage{stage_id}.log` + - the orchestrator can log readiness, forward size/time, and summary information ### 15. Risks & Mitigations — End-to-end latency: full-trigger sacrifices interactivity; start simple/correct, add windowed streaming later. — CPU copy overhead: copy-based IPC for maintainability; add SHM/zero-copy later as an optimization path. -— Serial throughput: `max_inflight=1` limits per-stage throughput; enable batching/concurrency later as needed. ### 16. Refactor: Sink LLM init and process into Stage (diagram) @@ -223,12 +170,16 @@ Config (preserve & extend) - Number/string: treat as logical index (0/1-based per implementation) mapping to physical GPUs - `runtime.max_batch_size`: max requests taken per intake (default 1) - `shm_threshold_bytes`: objects larger than this go via shared memory (default 64KB; injected into stages) +- Stage link and inputs: + - `engine_input_source`: list of upstream stage_id(s) from which the downstream stage takes inputs (typically a single source) + - `custom_process_input_func`: function to construct downstream engine_inputs from upstream outputs +- Output flags: + - `final_output`, `final_output_type`: marks that a stage produces the final external output and its type Kept & strengthened - Multi-process: each stage in its own process; parallel across stages; serial/batching inside a stage controlled by `max_batch_size` - Shared memory: large objects via SHM with strict `memoryview -> close -> unlink` - Device mapping: unified `set_stage_gpu_devices`, support comma lists and logical indices -- Windows: keep `window_size=-1` semantics (downstream triggers only after upstream finishes the request) Migration (execution order) 1) Enhanced `OmniStage` with `init_stage_worker/stop_stage_worker/submit/try_collect/_stage_worker`.