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