diff --git a/docs/mrs_design.md b/docs/mrs_design.md new file mode 100644 index 000000000..50e3a8c15 --- /dev/null +++ b/docs/mrs_design.md @@ -0,0 +1,188 @@ +## Multi-Request Streaming (MRS) on a Single Machine + +### 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. +- 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: 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 + - 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: 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) +- 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 +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` (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. + - 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 (pseudo-code) +```python +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) +``` + +### 12. Observability +- 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. + +### 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: 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 +- 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 `runtime`), `in_q/out_q`, subprocess handle, stats + - API: + - `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) + - build `OmniStageLLM(model, **engine_args)` + - while-loop: + - 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 output queue after `maybe_dump_to_shm` + - log start/finish; report `error` per request on exceptions + +Data flow (sketch) +``` +OmniLLM + ├─ 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) +- `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 +- `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 + +Migration (execution order) +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/examples/offline_inference/qwen_2_5_omni/README.md b/examples/offline_inference/qwen_2_5_omni/README.md index 770c5e169..734060b6b 100644 --- a/examples/offline_inference/qwen_2_5_omni/README.md +++ b/examples/offline_inference/qwen_2_5_omni/README.md @@ -63,13 +63,30 @@ uv pip install -r requirements.txt ``` ## Run examples (Qwen2.5-omni) - +### Multiple Prompts +Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can: +```bash +tar -xf /seedtts_testset.tar +cp seedtts_testset/en/meta.lst examples/offline_inference/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 +``` 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 fa5df3e0b..bbd4e0f53 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,7 @@ 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( "--voice-type", default="default", help="Voice type, e.g., m02, f030, default." @@ -69,6 +69,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 +109,46 @@ 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)", + ) + 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 @@ -117,7 +156,28 @@ def parse_args(): def main(): args = parse_args() model_name = args.model - omni_llm = OmniLLM(model=model_name) + try: + # Preferred: load from txt file (one prompt per line) + if getattr(args, "txt_prompts", None) and args.prompt_type == "text": + with open(args.txt_prompts, "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 prompts: {e}") + raise + + if args.prompts is None: + raise ValueError("No prompts provided. Use --prompts ... or --txt-prompts (with --prompt_type text)") + omni_llm = OmniLLM( + model=model_name, + log_stats=args.enable_stats, + log_file=("omni_llm_pipeline.log" if args.enable_stats else None), + init_sleep_seconds=args.init_sleep_seconds, + batch_timeout=args.batch_timeout, + init_timeout=args.init_timeout, + shm_threshold_bytes=args.shm_threshold_bytes, + ) thinker_sampling_params = SamplingParams( temperature=0.0, # Deterministic - no randomness top_p=1.0, # Disable nucleus sampling @@ -128,13 +188,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,20 +216,33 @@ 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") + 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 ) 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..7380e701b --- /dev/null +++ b/examples/offline_inference/qwen_2_5_omni/extract_prompts.py @@ -0,0 +1,43 @@ +#!/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="Input .lst file path") + parser.add_argument("--output", "-o", required=True, help="Output file path") + parser.add_argument("--topk", "-k", type=int, default=100, help="Extract the top K prompts (default: 100)") + args = parser.parse_args() + + prompts = [] + with open(args.input, "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") + + +if __name__ == "__main__": + main() + + diff --git a/examples/offline_inference/qwen_2_5_omni/run_multiple_prompts.sh b/examples/offline_inference/qwen_2_5_omni/run_multiple_prompts.sh new file mode 100644 index 000000000..ba07616ce --- /dev/null +++ b/examples/offline_inference/qwen_2_5_omni/run_multiple_prompts.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 \ + --txt-prompts top100.txt diff --git a/examples/offline_inference/qwen_2_5_omni/run.sh b/examples/offline_inference/qwen_2_5_omni/run_single_prompt.sh similarity index 78% rename from examples/offline_inference/qwen_2_5_omni/run.sh rename to examples/offline_inference/qwen_2_5_omni/run_single_prompt.sh index 7f02a672c..8dc2c74e5 100644 --- a/examples/offline_inference/qwen_2_5_omni/run.sh +++ b/examples/offline_inference/qwen_2_5_omni/run_single_prompt.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 \ + --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words." 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/log_utils.py b/vllm_omni/entrypoints/log_utils.py new file mode 100644 index 000000000..11b98ed81 --- /dev/null +++ b/vllm_omni/entrypoints/log_utils.py @@ -0,0 +1,570 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Optional + +import logging +import os +import time + +from vllm_omni.entrypoints.stage_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), + }) + + +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]], + 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_llm.py b/vllm_omni/entrypoints/omni_llm.py index 57bd12958..c1b12a1b2 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/vllm_omni/entrypoints/omni_llm.py @@ -1,8 +1,16 @@ 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 +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,37 +21,128 @@ 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.stage_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: - def __init__( - self, model: str, stage_configs=None, log_stats: bool = False, **kwargs - ): + """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.initialize_stage_configs(model) + self.stage_configs = load_stage_configs_from_model(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) + # 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) - 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) + 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): + # 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) + time.sleep(self._init_sleep_seconds) + + 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, @@ -52,42 +151,233 @@ class OmniLLM: Union[SamplingParams, Sequence[SamplingParams]] ] = None, ) -> list[OmniRequestOutput]: - """Generate text outputs for the given prompts.""" - final_outputs: 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, " - f"got {len(sampling_params_list)}" + f"Expected {len(self.stage_list)} sampling params, 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, + + # 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 _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 + 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 (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).", + ] + 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): diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 06b1f1b4c..8f95cf62e 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.stage_utils import ( + _to_dict, + set_stage_gpu_devices, + maybe_load_from_ipc_with_metrics, + 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) + runtime_cfg = _to_dict(getattr(self.stage_config, "runtime", {})) + stage_payload: Dict[str, Any] = { + "stage_id": self.stage_id, + "engine_args": engine_args, + "runtime": runtime_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 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 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 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 + 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,224 @@ 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 + from vllm_omni.entrypoints.log_utils import ( # noqa: WPS433 + 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 + # no inline JSONL/serialization imports; logging handled by utilities + + stage_id = stage_payload["stage_id"] + engine_args = stage_payload.get("engine_args", {}) + runtime_cfg = stage_payload.get("runtime", {}) + 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 + # Monotonic batch id per stage process for orchestrator dedup on time aggregation + _batch_seq = 0 + + # Device mapping + try: + 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) + + # 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(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: + 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 + 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) + 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: + _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): + 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 + for rid in batch_request_ids: + _r_outputs = req_to_outputs.get(rid, []) + _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: + _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: + 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_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)), + "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: + compute_and_log_stage_request_stats( + _stats_file, + stage_id, + rid, + len(batch_tasks), + r_outputs, + float(_gen_ms), + int(_metrics["rx_transfer_bytes"]), # type: ignore[index] + float(_metrics["rx_decode_time_ms"]), # type: ignore[index] + ) + 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_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)), + "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/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py new file mode 100644 index 000000000..16f8c6331 --- /dev/null +++ b/vllm_omni/entrypoints/stage_utils.py @@ -0,0 +1,239 @@ +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 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)) + 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 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) + 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 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): + 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 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) + + +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 as e: + logger.debug("Ensure parent dir failed for %s: %s", path, e) + 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. + + 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). + + 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/model_executor/models/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni.py index f820ef3d7..b8e3a2d9a 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 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( - 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..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,36 +1,50 @@ # stage config for running qwen2.5-omni with architecture of OmniLLM. stage_args: - stage_id: 0 + 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 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 + runtime: + 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 + runtime: + 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 runtime config (concise): default windows and stage edges +runtime: + 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..d91defd81 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 @@ -99,9 +101,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 = {} # Overlay custom prompt_embeds per request for the prompt portion; # collect additional_information (tensor/list) for prefill portion only @@ -122,54 +126,31 @@ 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 ) - # 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] @@ -240,6 +221,7 @@ class GPUARModelRunner(OmniGPUModelRunner): positions, intermediate_tensors, model_kwargs, + per_req_additional_information, ) @torch.inference_mode() @@ -355,6 +337,7 @@ class GPUARModelRunner(OmniGPUModelRunner): positions, intermediate_tensors, model_kwargs, + per_req_additional_information, ) = self._preprocess( scheduler_output, num_scheduled_tokens_np, @@ -395,15 +378,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 + # 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: + 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 + # 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]) + 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 +429,42 @@ class GPUARModelRunner(OmniGPUModelRunner): hidden_states, multimodal_outputs = self.extract_multimodal_outputs( hidden_states ) - + # 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) + and ( + "additional_information_update" in multimodal_outputs + or "additional_information_update_by_req_id" in multimodal_outputs + ) + ): + # Option A: list[dict] in batch order + 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) + # Option B: dict[str, dict] keyed by req_id + 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 +624,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