303 lines
12 KiB
Python
Executable File
303 lines
12 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Run or summarize a stratified MiniMax-H3 FL2VA/Ref2VA serving workload."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
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 build_plan(args: argparse.Namespace) -> list[dict[str, Any]]:
|
|
resolutions = [int(item) for item in args.resolutions.split(",") if item.strip()]
|
|
prompts = load_prompts(args.prompt_file, args.requests_per_resolution)
|
|
plan: list[dict[str, Any]] = []
|
|
# Interleave resolutions so any slow drift affects every bucket similarly.
|
|
for prompt_index, prompt in enumerate(prompts):
|
|
for short_edge in resolutions:
|
|
plan.append(
|
|
{
|
|
"request_id": f"{args.task}-r{short_edge}-p{prompt_index:02d}",
|
|
"task": args.task,
|
|
"short_edge": short_edge,
|
|
"prompt_index": prompt_index,
|
|
"prompt": prompt,
|
|
"seed": args.seed + prompt_index,
|
|
}
|
|
)
|
|
return plan
|
|
|
|
|
|
def make_payload(args: argparse.Namespace, item: dict[str, Any], steps: int) -> dict[str, Any]:
|
|
condition: dict[str, Any] = {
|
|
"type": "image",
|
|
"uri": str(args.reference_image),
|
|
"role": "keyframe" if args.task == "fl2va" else "reference",
|
|
}
|
|
if args.task == "fl2va":
|
|
condition["frame_index"] = 0
|
|
return {
|
|
"model": args.model,
|
|
"prompt": item["prompt"],
|
|
"num_outputs_per_prompt": 1,
|
|
"num_inference_steps": steps,
|
|
"flow_shift": args.flow_shift,
|
|
"audio_flow_shift": args.audio_flow_shift,
|
|
"seed": item["seed"],
|
|
"task": args.task,
|
|
"conditions": [condition],
|
|
"target": {
|
|
"short_edge": item["short_edge"],
|
|
"aspect_ratio": args.aspect_ratio,
|
|
"duration_seconds": args.duration_seconds,
|
|
},
|
|
}
|
|
|
|
|
|
def run_one(
|
|
session: requests.Session,
|
|
args: argparse.Namespace,
|
|
item: dict[str, Any],
|
|
steps: int,
|
|
) -> dict[str, Any]:
|
|
started_epoch = time.time()
|
|
started = time.monotonic()
|
|
result: dict[str, Any] = {
|
|
**item,
|
|
"replica_index": args.replica_index,
|
|
"port": args.port,
|
|
"num_inference_steps": steps,
|
|
"duration_seconds": args.duration_seconds,
|
|
"aspect_ratio": args.aspect_ratio,
|
|
"started_at_epoch": started_epoch,
|
|
"success": False,
|
|
"error": None,
|
|
}
|
|
try:
|
|
response = session.post(
|
|
f"http://{args.host}:{args.port}/v1/videos",
|
|
json=make_payload(args, item, steps),
|
|
timeout=args.submit_timeout,
|
|
)
|
|
if response.status_code != 200:
|
|
raise RuntimeError(f"submit HTTP {response.status_code}: {response.text[:1000]}")
|
|
status = response.json()
|
|
video_id = status.get("id")
|
|
if not video_id:
|
|
raise RuntimeError(f"submit response has no id: {status}")
|
|
result["video_id"] = video_id
|
|
deadline = time.monotonic() + args.request_timeout
|
|
while status.get("status") not in {"completed", "failed"}:
|
|
if time.monotonic() >= deadline:
|
|
raise TimeoutError(f"video job {video_id} exceeded {args.request_timeout}s")
|
|
time.sleep(args.poll_interval)
|
|
poll = session.get(
|
|
f"http://{args.host}:{args.port}/v1/videos/{video_id}",
|
|
timeout=args.poll_timeout,
|
|
)
|
|
if poll.status_code != 200:
|
|
raise RuntimeError(f"poll HTTP {poll.status_code}: {poll.text[:1000]}")
|
|
status = poll.json()
|
|
if status.get("status") != "completed":
|
|
raise RuntimeError(f"job failed: {status.get('error') or status}")
|
|
result["success"] = True
|
|
result["inference_time_s"] = status.get("inference_time_s")
|
|
result["peak_memory_mb"] = status.get("peak_memory_mb")
|
|
result["file_path"] = status.get("file_path")
|
|
except Exception as exc: # Keep the rest of the matrix running and record the cell failure.
|
|
result["error"] = f"{type(exc).__name__}: {exc}"
|
|
result["latency_s"] = time.monotonic() - started
|
|
result["finished_at_epoch"] = time.time()
|
|
return result
|
|
|
|
|
|
def run_command(args: argparse.Namespace) -> int:
|
|
if not args.reference_image.is_file():
|
|
raise SystemExit(f"reference image not found: {args.reference_image}")
|
|
full_plan = build_plan(args)
|
|
# Stratify by prompt index so every replica receives the same number of
|
|
# samples from every resolution. This avoids assigning an entire slow
|
|
# resolution bucket (for example 1080p) to only one replica.
|
|
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)
|
|
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}",
|
|
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)
|
|
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)
|
|
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 = {
|
|
"task": args.task,
|
|
"tp": args.tp,
|
|
"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.tp),
|
|
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("--model", default="/data/hf_models/MiniMax-H3")
|
|
run.add_argument("--task", choices=["fl2va", "ref2va"], required=True)
|
|
run.add_argument("--reference-image", type=Path, required=True)
|
|
run.add_argument("--prompt-file", type=Path, default=Path.home() / ".cache/sglang/vbench_subject_consistency.txt")
|
|
run.add_argument("--resolutions", default="480,720,768,1080")
|
|
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("--num-inference-steps", type=int, default=20)
|
|
run.add_argument("--warmup-requests", type=int, default=1)
|
|
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("--flow-shift", type=float, default=12.0)
|
|
run.add_argument("--audio-flow-shift", type=float, default=3.0)
|
|
run.add_argument("--seed", type=int, default=1101)
|
|
run.add_argument("--submit-timeout", type=float, default=120.0)
|
|
run.add_argument("--poll-timeout", type=float, default=30.0)
|
|
run.add_argument("--poll-interval", type=float, default=1.0)
|
|
run.add_argument("--request-timeout", type=float, default=3600.0)
|
|
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("--tp", 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())
|