sskj-h3/sglang-profile/scripts/h3_profile/h3_profile_client.py

241 lines
11 KiB
Python
Executable File

#!/usr/bin/env python3
"""Submit reproducible MiniMax-H3 profiling requests to SGLang's video API."""
from __future__ import annotations
import argparse
import json
import os
import time
from pathlib import Path
from typing import Any
import requests
DEFAULT_PROMPT = (
"A cinematic mountain lake at sunrise, natural motion, realistic lighting, "
"stable composition, synchronized ambient sound."
)
def condition(kind: str, uri: Path, *, role: str = "reference", frame_index: int | None = None) -> dict[str, Any]:
item: dict[str, Any] = {"type": kind, "uri": str(uri), "role": role}
if frame_index is not None:
item["frame_index"] = frame_index
return item
def scenario_spec(name: str, args: argparse.Namespace) -> tuple[str, list[dict[str, Any]]]:
image = args.reference_image
distinct_images = []
if args.reference_images_dir.is_dir():
distinct_images = sorted(
p for p in args.reference_images_dir.iterdir()
if p.suffix.lower() in {".jpg", ".jpeg", ".png", ".webp"}
)
images = (distinct_images + [image] * 9)[:9]
video_1 = args.reference_video_1s
video_5 = args.reference_video_5s
video_10 = args.reference_video_10s
specs: dict[str, tuple[str, list[dict[str, Any]]]] = {
"F0": ("t2va", []),
"F1": ("fl2va", [condition("image", image, role="keyframe", frame_index=0)]),
"F2": ("fl2va", [condition("image", image, role="keyframe", frame_index=-1)]),
"F3": (
"fl2va",
[
condition("image", image, role="keyframe", frame_index=0),
condition("image", image, role="keyframe", frame_index=-1),
],
),
# Use the first image from the same normalized pool as R5/R6/R9 so
# cardinality comparisons do not mix input resolution/aspect changes.
"R1": ("ref2va", [condition("image", images[0])]),
"R5": ("ref2va", [condition("image", p) for p in images[:5]]),
"R6": ("ref2va", [condition("image", p) for p in images[:6]]),
"R9": ("ref2va", [condition("image", p) for p in images[:9]]),
"RV1": ("ref2va", [condition("video", video_1)]),
"RV5": ("ref2va", [condition("video", video_5)]),
"RV10": ("ref2va", [condition("video", video_10)]),
"RIA": ("ref2va", [condition("image", image), condition("audio", video_5)]),
"RVA": ("ref2va", [condition("video", video_5), condition("audio", video_5)]),
"RVA_EMBEDDED": ("ref2va", [condition("video_audio", video_5)]),
"RMIX": (
"ref2va",
[condition("image", image) for _ in range(5)]
+ [condition("video", video_5), condition("audio", video_5)],
),
}
try:
return specs[name]
except KeyError as exc:
raise SystemExit(f"unknown scenario {name}; choices={','.join(specs)}") from exc
def submit_and_wait(session: requests.Session, args: argparse.Namespace, payload: dict[str, Any]) -> dict[str, Any]:
start_epoch = time.time()
start = time.monotonic()
result: dict[str, Any] = {"started_at_epoch": start_epoch, "success": False}
try:
response = session.post(
f"http://{args.host}:{args.port}/v1/videos",
json=payload,
timeout=args.submit_timeout,
)
if response.status_code != 200:
raise RuntimeError(f"submit HTTP {response.status_code}: {response.text[:2000]}")
status = response.json()
result["submit_response"] = status
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
poll_count = 0
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_count += 1
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[:2000]}")
status = poll.json()
result["poll_count"] = poll_count
result["final_status"] = status
if status.get("status") != "completed":
raise RuntimeError(f"job failed: {status.get('error') or status}")
result["success"] = True
for key in ("inference_time_s", "peak_memory_mb", "file_path"):
result[key] = status.get(key)
except Exception as exc:
result["error"] = f"{type(exc).__name__}: {exc}"
result["latency_s"] = time.monotonic() - start
result["finished_at_epoch"] = time.time()
return result
def main() -> int:
parser = argparse.ArgumentParser()
parser.add_argument("--scenario", required=True)
parser.add_argument("--deployment", required=True)
parser.add_argument("--host", default="127.0.0.1")
parser.add_argument("--port", type=int, required=True)
parser.add_argument("--replica-index", type=int, default=0)
parser.add_argument("--model", default="/data/hf_models/MiniMax-H3")
parser.add_argument("--reference-image", type=Path, required=True)
parser.add_argument("--reference-images-dir", type=Path, required=True)
parser.add_argument("--reference-video-1s", type=Path, required=True)
parser.add_argument("--reference-video-5s", type=Path, required=True)
parser.add_argument("--reference-video-10s", type=Path, required=True)
parser.add_argument("--prompt", default=DEFAULT_PROMPT)
parser.add_argument("--steps", type=int, default=20)
parser.add_argument("--duration", type=float, default=5.0)
parser.add_argument("--short-edge", type=int, default=768)
parser.add_argument("--aspect-ratio", default="16:9")
parser.add_argument("--fps", type=int, default=24)
parser.add_argument("--seed", type=int, default=1101)
parser.add_argument("--repeats", type=int, default=1)
parser.add_argument("--warmup", type=int, default=0)
parser.add_argument("--warmup-steps", type=int, default=5)
parser.add_argument("--profile", action="store_true")
parser.add_argument("--profile-all-stages", action="store_true")
parser.add_argument("--num-profiled-timesteps", type=int, default=5)
parser.add_argument("--perf-dir", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--submit-timeout", type=float, default=120.0)
parser.add_argument("--poll-timeout", type=float, default=30.0)
parser.add_argument("--poll-interval", type=float, default=1.0)
parser.add_argument("--request-timeout", type=float, default=7200.0)
args = parser.parse_args()
task, conditions = scenario_spec(args.scenario, args)
for path in [args.reference_image, args.reference_video_1s, args.reference_video_5s, args.reference_video_10s]:
if not path.is_file():
raise SystemExit(f"required local material missing: {path}")
args.output.parent.mkdir(parents=True, exist_ok=True)
args.perf_dir.mkdir(parents=True, exist_ok=True)
completed: set[str] = set()
if args.output.is_file():
for line in args.output.read_text(encoding="utf-8").splitlines():
try:
row = json.loads(line)
if row.get("success"):
completed.add(row["request_id"])
except (json.JSONDecodeError, KeyError):
pass
def payload_for(request_id: str, index: int, steps: int, profiled: bool) -> dict[str, Any]:
perf_path = args.perf_dir / f"{request_id}.json"
return {
"model": args.model,
"prompt": args.prompt,
"num_outputs_per_prompt": 1,
"num_inference_steps": steps,
"flow_shift": 12.0,
"audio_flow_shift": 3.0,
"seed": args.seed + index,
"task": task,
"conditions": conditions,
"target": {
"short_edge": args.short_edge,
"aspect_ratio": args.aspect_ratio,
"duration_seconds": args.duration,
},
"profile": profiled,
"profile_all_stages": bool(profiled and args.profile_all_stages),
"num_profiled_timesteps": args.num_profiled_timesteps,
"perf_dump_path": str(perf_path),
}
failures = 0
with requests.Session() as session, args.output.open("a", encoding="utf-8", buffering=1) as output:
for i in range(args.warmup):
request_id = f"warmup-{args.deployment}-{args.scenario}-r{args.replica_index}-{i}"
result = submit_and_wait(session, args, payload_for(request_id, i, args.warmup_steps, False))
print(f"{request_id} success={result['success']} latency={result['latency_s']:.3f}", flush=True)
if not result["success"]:
raise SystemExit(f"warmup failed: {result.get('error')}")
for i in range(args.repeats):
request_id = f"{args.deployment}-{args.scenario}-r{args.replica_index}-n{i}"
if request_id in completed:
print(f"skip completed {request_id}", flush=True)
continue
payload = payload_for(request_id, i, args.steps, args.profile)
result = submit_and_wait(session, args, payload)
row = {
"request_id": request_id,
"deployment": args.deployment,
"scenario": args.scenario,
"task": task,
"replica_index": args.replica_index,
"port": args.port,
"steps": args.steps,
"duration_s": args.duration,
"short_edge": args.short_edge,
"aspect_ratio": args.aspect_ratio,
"seed": args.seed + i,
"condition_types": [x["type"] for x in conditions],
"condition_count": len(conditions),
"profile": args.profile,
"profile_all_stages": args.profile_all_stages,
"perf_dump_path": payload["perf_dump_path"],
**result,
}
output.write(json.dumps(row, ensure_ascii=False) + "\n")
failures += int(not result["success"])
print(
f"{request_id} success={result['success']} latency={result['latency_s']:.3f} "
f"inference={result.get('inference_time_s')} error={result.get('error')}",
flush=True,
)
return int(failures > 0)
if __name__ == "__main__":
raise SystemExit(main())