355 lines
14 KiB
Python
Executable File
355 lines
14 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Balanced synchronous serving benchmark for MiniMax-H3 on vLLM-Omni."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
import mimetypes
|
|
import statistics
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import requests
|
|
|
|
|
|
DEFAULT_PROMPT = "A cinematic landscape with natural motion and realistic lighting."
|
|
|
|
|
|
def percentile(values: list[float], q: float) -> float:
|
|
if not values:
|
|
return 0.0
|
|
values = sorted(values)
|
|
pos = (len(values) - 1) * q
|
|
lo, hi = math.floor(pos), math.ceil(pos)
|
|
if lo == hi:
|
|
return values[lo]
|
|
return values[lo] * (hi - pos) + values[hi] * (pos - lo)
|
|
|
|
|
|
def load_prompts(path: Path, count: int) -> list[str]:
|
|
prompts: list[str] = []
|
|
if path.is_file():
|
|
prompts = [line.strip() for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
|
|
if not prompts:
|
|
prompts = [DEFAULT_PROMPT]
|
|
repeats = (count + len(prompts) - 1) // len(prompts)
|
|
return (prompts * repeats)[:count]
|
|
|
|
|
|
def parse_resolution_map(raw: str) -> dict[int, tuple[int, int]]:
|
|
mapping: dict[int, tuple[int, int]] = {}
|
|
for entry in raw.split(","):
|
|
short_edge, shape = entry.split(":", 1)
|
|
width, height = shape.lower().split("x", 1)
|
|
mapping[int(short_edge)] = (int(width), int(height))
|
|
return mapping
|
|
|
|
|
|
def build_plan(args: argparse.Namespace) -> list[dict[str, Any]]:
|
|
resolution_map = parse_resolution_map(args.resolution_map)
|
|
prompts = load_prompts(args.prompt_file, args.requests_per_resolution)
|
|
plan: list[dict[str, Any]] = []
|
|
for prompt_index, prompt in enumerate(prompts):
|
|
for short_edge, (width, height) in resolution_map.items():
|
|
plan.append(
|
|
{
|
|
"request_id": f"{args.task}-r{short_edge}-p{prompt_index:02d}",
|
|
"task": args.task,
|
|
"short_edge": short_edge,
|
|
"width": width,
|
|
"height": height,
|
|
"prompt_index": prompt_index,
|
|
"prompt": prompt,
|
|
"seed": args.seed + prompt_index,
|
|
}
|
|
)
|
|
return plan
|
|
|
|
|
|
def make_form(args: argparse.Namespace, item: dict[str, Any], steps: int) -> dict[str, str]:
|
|
extra_params = {
|
|
"task": args.task,
|
|
"duration": args.duration_seconds,
|
|
"audio_flow_shift": args.audio_flow_shift,
|
|
}
|
|
return {
|
|
"prompt": item["prompt"],
|
|
"width": str(item["width"]),
|
|
"height": str(item["height"]),
|
|
"aspect_ratio": args.aspect_ratio,
|
|
"fps": str(args.fps),
|
|
"num_inference_steps": str(steps),
|
|
"flow_shift": str(args.flow_shift),
|
|
"seed": str(item["seed"]),
|
|
"quality": args.quality,
|
|
"extra_params": json.dumps(extra_params, separators=(",", ":")),
|
|
}
|
|
|
|
|
|
def run_one(
|
|
session: requests.Session,
|
|
args: argparse.Namespace,
|
|
item: dict[str, Any],
|
|
steps: int,
|
|
save_video: bool,
|
|
) -> dict[str, Any]:
|
|
started_epoch = time.time()
|
|
started = time.monotonic()
|
|
output_path = args.video_dir / f"{item['request_id']}.mp4"
|
|
result: dict[str, Any] = {
|
|
**item,
|
|
"backend": "vllm-omni",
|
|
"replica_index": args.replica_index,
|
|
"port": args.port,
|
|
"gpus_per_instance": args.gpus_per_instance,
|
|
"dit_tensor_parallel_size": args.dit_tp,
|
|
"ulysses_degree": args.usp,
|
|
"num_inference_steps": steps,
|
|
"duration_seconds": args.duration_seconds,
|
|
"aspect_ratio": args.aspect_ratio,
|
|
"started_at_epoch": started_epoch,
|
|
"success": False,
|
|
"error": None,
|
|
}
|
|
try:
|
|
# T2VA (FastH3) carries no reference media; pass {} to keep multipart encoding.
|
|
if args.task == "t2va":
|
|
files: dict[str, Any] = {}
|
|
else:
|
|
mime_type = mimetypes.guess_type(args.reference_image.name)[0] or "application/octet-stream"
|
|
with args.reference_image.open("rb") as reference_handle:
|
|
files = {
|
|
"input_reference": (
|
|
args.reference_image.name,
|
|
reference_handle,
|
|
mime_type,
|
|
)
|
|
}
|
|
with session.post(
|
|
f"http://{args.host}:{args.port}/v1/videos/sync",
|
|
data=make_form(args, item, steps),
|
|
files=files,
|
|
timeout=(args.connect_timeout, args.request_timeout),
|
|
stream=True,
|
|
) as response:
|
|
if response.status_code != 200:
|
|
body = response.content[:2000].decode("utf-8", errors="replace")
|
|
raise RuntimeError(f"HTTP {response.status_code}: {body}")
|
|
result["content_type"] = response.headers.get("content-type")
|
|
result["stage_durations"] = response.headers.get("x-stage-durations")
|
|
result["peak_memory_mb"] = response.headers.get("x-peak-memory-mb")
|
|
byte_count = 0
|
|
first_bytes = b""
|
|
output_handle = output_path.open("wb") if save_video else None
|
|
try:
|
|
for chunk in response.iter_content(chunk_size=1024 * 1024):
|
|
if not chunk:
|
|
continue
|
|
if len(first_bytes) < 32:
|
|
first_bytes += chunk[: 32 - len(first_bytes)]
|
|
byte_count += len(chunk)
|
|
if output_handle is not None:
|
|
output_handle.write(chunk)
|
|
finally:
|
|
if output_handle is not None:
|
|
output_handle.close()
|
|
if byte_count < 12 or b"ftyp" not in first_bytes:
|
|
raise RuntimeError(
|
|
f"response is not a valid-looking MP4: bytes={byte_count}, "
|
|
f"content_type={result['content_type']}"
|
|
)
|
|
result["response_bytes"] = byte_count
|
|
if save_video:
|
|
result["file_path"] = str(output_path)
|
|
result["success"] = True
|
|
except Exception as exc:
|
|
result["error"] = f"{type(exc).__name__}: {exc}"
|
|
if save_video and output_path.exists():
|
|
output_path.unlink()
|
|
result["latency_s"] = time.monotonic() - started
|
|
result["finished_at_epoch"] = time.time()
|
|
return result
|
|
|
|
|
|
def run_command(args: argparse.Namespace) -> int:
|
|
if args.task != "t2va" and (
|
|
not args.reference_image or not args.reference_image.is_file()
|
|
):
|
|
raise SystemExit(f"reference image not found: {args.reference_image}")
|
|
full_plan = build_plan(args)
|
|
# Stratified allocation: every replica gets the same number of prompts
|
|
# from every resolution bucket.
|
|
shard = [
|
|
item
|
|
for item in full_plan
|
|
if item["prompt_index"] % args.num_replicas == args.replica_index
|
|
]
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.video_dir.mkdir(parents=True, exist_ok=True)
|
|
completed_ids: set[str] = set()
|
|
if args.output.is_file():
|
|
for line in args.output.read_text(encoding="utf-8").splitlines():
|
|
try:
|
|
completed_ids.add(json.loads(line)["request_id"])
|
|
except (json.JSONDecodeError, KeyError):
|
|
continue
|
|
shard = [item for item in shard if item["request_id"] not in completed_ids]
|
|
print(
|
|
f"task={args.task} replica={args.replica_index}/{args.num_replicas} "
|
|
f"requests={len(shard)} port={args.port} gpus={args.gpus_per_instance} "
|
|
f"dit_tp={args.dit_tp} usp={args.usp}",
|
|
flush=True,
|
|
)
|
|
failures = 0
|
|
with requests.Session() as session, args.output.open("a", encoding="utf-8", buffering=1) as output:
|
|
for warmup_index in range(args.warmup_requests):
|
|
warmup_item = (shard or full_plan)[warmup_index % len(shard or full_plan)].copy()
|
|
warmup_item["request_id"] = f"warmup-{warmup_index}-{warmup_item['request_id']}"
|
|
warmup = run_one(
|
|
session,
|
|
args,
|
|
warmup_item,
|
|
args.warmup_inference_steps,
|
|
save_video=False,
|
|
)
|
|
print(
|
|
f"warmup {warmup_index + 1}/{args.warmup_requests}: "
|
|
f"success={warmup['success']} latency={warmup['latency_s']:.2f}s "
|
|
f"error={warmup['error']}",
|
|
flush=True,
|
|
)
|
|
if not warmup["success"]:
|
|
raise SystemExit("warmup failed")
|
|
for index, item in enumerate(shard, start=1):
|
|
result = run_one(session, args, item, args.num_inference_steps, save_video=True)
|
|
output.write(json.dumps(result, ensure_ascii=False) + "\n")
|
|
output.flush()
|
|
failures += int(not result["success"])
|
|
print(
|
|
f"request {index}/{len(shard)} id={item['request_id']} "
|
|
f"success={result['success']} latency={result['latency_s']:.2f}s "
|
|
f"error={result['error']}",
|
|
flush=True,
|
|
)
|
|
return int(failures > 0)
|
|
|
|
|
|
def summarize_command(args: argparse.Namespace) -> int:
|
|
rows: list[dict[str, Any]] = []
|
|
for path in sorted(args.input_dir.glob("client_*/results.jsonl")):
|
|
for line in path.read_text(encoding="utf-8").splitlines():
|
|
if line.strip():
|
|
rows.append(json.loads(line))
|
|
successful = [row for row in rows if row.get("success")]
|
|
latencies = [float(row["latency_s"]) for row in successful]
|
|
started = [float(row["started_at_epoch"]) for row in rows]
|
|
finished = [float(row["finished_at_epoch"]) for row in rows]
|
|
wall_s = max(finished) - min(started) if started and finished else 0.0
|
|
buckets: dict[str, dict[str, Any]] = {}
|
|
for short_edge in sorted({int(row["short_edge"]) for row in rows}):
|
|
bucket_rows = [row for row in rows if int(row["short_edge"]) == short_edge]
|
|
bucket_success = [row for row in bucket_rows if row.get("success")]
|
|
bucket_latencies = [float(row["latency_s"]) for row in bucket_success]
|
|
buckets[str(short_edge)] = {
|
|
"requests": len(bucket_rows),
|
|
"completed": len(bucket_success),
|
|
"failed": len(bucket_rows) - len(bucket_success),
|
|
"latency_mean_s": statistics.fmean(bucket_latencies) if bucket_latencies else 0.0,
|
|
"latency_p95_s": percentile(bucket_latencies, 0.95),
|
|
}
|
|
summary = {
|
|
"backend": "vllm-omni",
|
|
"task": args.task,
|
|
"gpus_per_instance": args.gpus_per_instance,
|
|
"dit_tensor_parallel_size": args.dit_tp,
|
|
"ulysses_degree": args.usp,
|
|
"replicas": args.replicas,
|
|
"expected_requests": args.expected_requests,
|
|
"requests_recorded": len(rows),
|
|
"completed": len(successful),
|
|
"failed": len(rows) - len(successful),
|
|
"machine_wall_s": wall_s,
|
|
"machine_qps": len(successful) / wall_s if wall_s else 0.0,
|
|
"latency_mean_s": statistics.fmean(latencies) if latencies else 0.0,
|
|
"latency_p50_s": percentile(latencies, 0.50),
|
|
"latency_p95_s": percentile(latencies, 0.95),
|
|
"by_short_edge": buckets,
|
|
}
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
print(
|
|
"\t".join(
|
|
[
|
|
str(args.gpus_per_instance),
|
|
str(args.dit_tp),
|
|
str(args.usp),
|
|
str(args.replicas),
|
|
args.task,
|
|
str(args.expected_requests),
|
|
str(len(rows)),
|
|
str(len(successful)),
|
|
str(len(rows) - len(successful)),
|
|
f"{summary['machine_qps']:.8f}",
|
|
f"{summary['latency_mean_s']:.6f}",
|
|
f"{summary['latency_p95_s']:.6f}",
|
|
f"{wall_s:.3f}",
|
|
]
|
|
)
|
|
)
|
|
return int(len(rows) != args.expected_requests or len(successful) != len(rows))
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
subparsers = parser.add_subparsers(dest="command", required=True)
|
|
run = subparsers.add_parser("run")
|
|
run.add_argument("--host", default="127.0.0.1")
|
|
run.add_argument("--port", type=int, required=True)
|
|
run.add_argument("--task", choices=["t2va", "fl2va", "ref2va"], required=True)
|
|
run.add_argument("--reference-image", type=Path, default=None)
|
|
run.add_argument("--prompt-file", type=Path, required=True)
|
|
run.add_argument("--resolution-map", default="480:864x480,720:1280x736,768:1344x768,1080:1920x1088")
|
|
run.add_argument("--requests-per-resolution", type=int, default=8)
|
|
run.add_argument("--replica-index", type=int, required=True)
|
|
run.add_argument("--num-replicas", type=int, required=True)
|
|
run.add_argument("--gpus-per-instance", type=int, required=True)
|
|
run.add_argument("--dit-tp", type=int, required=True)
|
|
run.add_argument("--usp", type=int, required=True)
|
|
run.add_argument("--num-inference-steps", type=int, default=4)
|
|
run.add_argument("--warmup-requests", type=int, default=2)
|
|
run.add_argument("--warmup-inference-steps", type=int, default=5)
|
|
run.add_argument("--duration-seconds", type=float, default=5.0)
|
|
run.add_argument("--aspect-ratio", default="16:9")
|
|
run.add_argument("--fps", type=int, default=24)
|
|
run.add_argument("--flow-shift", type=float, default=12.0)
|
|
run.add_argument("--audio-flow-shift", type=float, default=3.0)
|
|
run.add_argument("--quality", default="lossless")
|
|
run.add_argument("--seed", type=int, default=1101)
|
|
run.add_argument("--connect-timeout", type=float, default=30.0)
|
|
run.add_argument("--request-timeout", type=float, default=14400.0)
|
|
run.add_argument("--video-dir", type=Path, required=True)
|
|
run.add_argument("--output", type=Path, required=True)
|
|
|
|
summarize = subparsers.add_parser("summarize")
|
|
summarize.add_argument("--input-dir", type=Path, required=True)
|
|
summarize.add_argument("--output", type=Path, required=True)
|
|
summarize.add_argument("--task", required=True)
|
|
summarize.add_argument("--gpus-per-instance", type=int, required=True)
|
|
summarize.add_argument("--dit-tp", type=int, required=True)
|
|
summarize.add_argument("--usp", type=int, required=True)
|
|
summarize.add_argument("--replicas", type=int, required=True)
|
|
summarize.add_argument("--expected-requests", type=int, required=True)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
return run_command(args) if args.command == "run" else summarize_command(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|