338 lines
13 KiB
Python
338 lines
13 KiB
Python
|
# SPDX-License-Identifier: Apache-2.0
|
||
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||
|
|
"""Startup / loading benchmark for multi-stage omni models (MammothModa2 first).
|
||
|
|
|
||
|
|
Measures, from a single fresh process:
|
||
|
|
* python import time (torch + vllm_omni)
|
||
|
|
* Omni() construction wall time (all stages: weight load, profiling, KV cache, warmup)
|
||
|
|
* first request latency (includes JIT / lazy init that only happens on first run)
|
||
|
|
* steady-state request latency (repeat N-1 more requests)
|
||
|
|
* optional host RSS / device memory samples (--sample-memory, once a second)
|
||
|
|
|
||
|
|
Per-stage breakdown (weight loading, model-load memory, engine init) is emitted by the
|
||
|
|
stage engine-core subprocesses into the log; use parse_startup_log.py on the captured
|
||
|
|
stdout/stderr to merge them with the JSON written here.
|
||
|
|
|
||
|
|
Usage (keep the log: the per-stage breakdown is parsed from it afterwards):
|
||
|
|
python benchmarks/mammoth_moda2/bench_startup.py --model /path/MammothModa2-Preview \
|
||
|
|
--deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
|
||
|
|
--height 1024 --width 1024 --seed 42 \
|
||
|
|
--extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \
|
||
|
|
--repeat 2 --label nfs-warm --output-json out.json 2>&1 | tee run.log
|
||
|
|
python benchmarks/mammoth_moda2/parse_startup_log.py run.log --markdown
|
||
|
|
"""
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import platform
|
||
|
|
import subprocess
|
||
|
|
import threading
|
||
|
|
import time
|
||
|
|
|
||
|
|
T_PROCESS_START = time.perf_counter()
|
||
|
|
|
||
|
|
import torch # noqa: E402
|
||
|
|
|
||
|
|
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args # noqa: E402
|
||
|
|
from vllm_omni.entrypoints.omni import Omni # noqa: E402
|
||
|
|
from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params # noqa: E402
|
||
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams # noqa: E402
|
||
|
|
from vllm_omni.model_extras import ( # noqa: E402
|
||
|
|
build_text_to_image_prompt as build_model_text_to_image_prompt,
|
||
|
|
)
|
||
|
|
from vllm_omni.model_extras import ( # noqa: E402
|
||
|
|
get_extra_body_params,
|
||
|
|
get_model_class_name,
|
||
|
|
should_init_extra_args_for_non_diffusion_stages,
|
||
|
|
)
|
||
|
|
from vllm_omni.platforms import current_omni_platform # noqa: E402
|
||
|
|
|
||
|
|
T_IMPORTS_DONE = time.perf_counter()
|
||
|
|
|
||
|
|
|
||
|
|
def parse_args() -> argparse.Namespace:
|
||
|
|
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||
|
|
p.add_argument("--model", required=True)
|
||
|
|
p.add_argument("--deploy-config", default=None)
|
||
|
|
p.add_argument("--prompt", default="A stylish woman riding a motorcycle in NYC, movie poster style")
|
||
|
|
p.add_argument("--negative-prompt", default=None)
|
||
|
|
p.add_argument("--height", type=int, default=1024)
|
||
|
|
p.add_argument("--width", type=int, default=1024)
|
||
|
|
p.add_argument("--seed", type=int, default=42)
|
||
|
|
p.add_argument("--extra-body", type=json.loads, default=None, help="JSON object forwarded as model extra_body")
|
||
|
|
p.add_argument(
|
||
|
|
"--repeat", type=int, default=2, help="number of generate() calls; 1st = first-request, rest = steady"
|
||
|
|
)
|
||
|
|
p.add_argument("--label", default="run", help="free-form label, e.g. nfs-cold / nfs-warm / local-nvme")
|
||
|
|
p.add_argument("--output-json", default=None)
|
||
|
|
p.add_argument("--save-image", default=None, help="optional path to save the first request's image")
|
||
|
|
p.add_argument("--sample-memory", action="store_true", help="sample memory once a second (adds overhead)")
|
||
|
|
p.add_argument(
|
||
|
|
"--parallel-stage-init",
|
||
|
|
action="store_true",
|
||
|
|
help="initialize stages that share a GPU concurrently (VllmOmniOrchestratorConfig.parallel_stage_init)",
|
||
|
|
)
|
||
|
|
args = p.parse_args()
|
||
|
|
if args.repeat < 1:
|
||
|
|
p.error("--repeat must be at least 1")
|
||
|
|
return args
|
||
|
|
|
||
|
|
|
||
|
|
def gpu_info() -> dict:
|
||
|
|
info = {"torch": torch.__version__, "cuda": torch.version.cuda, "python": platform.python_version()}
|
||
|
|
try:
|
||
|
|
import vllm
|
||
|
|
|
||
|
|
info["vllm"] = vllm.__version__
|
||
|
|
except Exception: # pragma: no cover
|
||
|
|
pass
|
||
|
|
try:
|
||
|
|
import vllm_omni
|
||
|
|
|
||
|
|
info["vllm_omni"] = getattr(vllm_omni, "__version__", "src")
|
||
|
|
except Exception: # pragma: no cover
|
||
|
|
pass
|
||
|
|
try:
|
||
|
|
out = subprocess.run(
|
||
|
|
["nvidia-smi", "--query-gpu=name,driver_version,memory.total", "--format=csv,noheader"],
|
||
|
|
capture_output=True,
|
||
|
|
text=True,
|
||
|
|
timeout=10,
|
||
|
|
)
|
||
|
|
info["gpu"] = out.stdout.strip()
|
||
|
|
except Exception:
|
||
|
|
info["gpu"] = "n/a"
|
||
|
|
return info
|
||
|
|
|
||
|
|
|
||
|
|
class PeakSampler:
|
||
|
|
"""Poll host RSS of this process tree and GPU memory in use once a second (benchmark-side only)."""
|
||
|
|
|
||
|
|
def __init__(self, interval_s: float = 1.0):
|
||
|
|
self.interval_s = interval_s
|
||
|
|
self.host_rss_peak_gib = 0.0
|
||
|
|
self.gpu_used_peak_gib = 0.0
|
||
|
|
self._stop = threading.Event()
|
||
|
|
self._thread = threading.Thread(target=self._run, daemon=True)
|
||
|
|
|
||
|
|
def _sample(self) -> None:
|
||
|
|
try:
|
||
|
|
import psutil
|
||
|
|
|
||
|
|
root = psutil.Process()
|
||
|
|
rss = root.memory_info().rss + sum(
|
||
|
|
c.memory_info().rss for c in root.children(recursive=True) if c.is_running()
|
||
|
|
)
|
||
|
|
self.host_rss_peak_gib = max(self.host_rss_peak_gib, rss / 2**30)
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
try:
|
||
|
|
out = subprocess.run(
|
||
|
|
["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"],
|
||
|
|
capture_output=True,
|
||
|
|
text=True,
|
||
|
|
timeout=5,
|
||
|
|
)
|
||
|
|
used = max(float(x) for x in out.stdout.split() if x.strip())
|
||
|
|
self.gpu_used_peak_gib = max(self.gpu_used_peak_gib, used / 1024)
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
|
||
|
|
def _run(self) -> None:
|
||
|
|
while not self._stop.is_set():
|
||
|
|
self._sample()
|
||
|
|
self._stop.wait(self.interval_s)
|
||
|
|
|
||
|
|
def start(self) -> "PeakSampler":
|
||
|
|
self._thread.start()
|
||
|
|
return self
|
||
|
|
|
||
|
|
def stop(self) -> dict:
|
||
|
|
self._stop.set()
|
||
|
|
self._thread.join(timeout=10)
|
||
|
|
self._sample()
|
||
|
|
return {
|
||
|
|
"host_rss_peak_gib": round(self.host_rss_peak_gib, 2),
|
||
|
|
"gpu_used_peak_gib": round(self.gpu_used_peak_gib, 2),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def fs_type(path: str) -> str:
|
||
|
|
try:
|
||
|
|
out = subprocess.run(["stat", "-f", "-c", "%T", path], capture_output=True, text=True, timeout=10)
|
||
|
|
return out.stdout.strip()
|
||
|
|
except Exception:
|
||
|
|
return "unknown"
|
||
|
|
|
||
|
|
|
||
|
|
def build_request(omni: Omni, args: argparse.Namespace, generator: torch.Generator):
|
||
|
|
model_class_name = get_model_class_name(omni)
|
||
|
|
declared = get_extra_body_params(model_class_name)
|
||
|
|
|
||
|
|
prompt_dict = {"prompt": args.prompt}
|
||
|
|
if args.negative_prompt:
|
||
|
|
prompt_dict["negative_prompt"] = args.negative_prompt
|
||
|
|
prompt_dict = build_model_text_to_image_prompt(
|
||
|
|
model_class_name=model_class_name, prompt=prompt_dict, height=args.height, width=args.width
|
||
|
|
)
|
||
|
|
|
||
|
|
diffusion_params = OmniDiffusionSamplingParams(
|
||
|
|
height=args.height,
|
||
|
|
width=args.width,
|
||
|
|
seed=args.seed,
|
||
|
|
generator=generator,
|
||
|
|
)
|
||
|
|
user_extra = dict(args.extra_body or {})
|
||
|
|
if declared:
|
||
|
|
apply_declared_extra_args(diffusion_params, declared, user_extra)
|
||
|
|
else:
|
||
|
|
diffusion_params.extra_args.update({k: v for k, v in user_extra.items() if v is not None})
|
||
|
|
|
||
|
|
init_non_diffusion = should_init_extra_args_for_non_diffusion_stages(model_class_name)
|
||
|
|
sampling_params_list = [clone_sampling_params(p) for p in (omni.default_sampling_params_list or [])]
|
||
|
|
if not sampling_params_list:
|
||
|
|
sampling_params_list = [diffusion_params]
|
||
|
|
|
||
|
|
replaced = False
|
||
|
|
for idx, params in enumerate(sampling_params_list):
|
||
|
|
if isinstance(params, OmniDiffusionSamplingParams):
|
||
|
|
sampling_params_list[idx] = diffusion_params
|
||
|
|
replaced = True
|
||
|
|
elif init_non_diffusion and hasattr(params, "extra_args"):
|
||
|
|
params.extra_args = {**(params.extra_args or {}), **(diffusion_params.extra_args or {})}
|
||
|
|
if hasattr(params, "seed"):
|
||
|
|
params.seed = args.seed
|
||
|
|
# MammothModa2 AR stage: one visual token per grid cell + one EOL per row + 1 look-ahead.
|
||
|
|
info = prompt_dict.get("additional_information", {})
|
||
|
|
if idx == 0 and info.get("omni_task") == ["t2i"]:
|
||
|
|
ar_w = int(info.get("ar_width", [0])[0])
|
||
|
|
ar_h = int(info.get("ar_height", [0])[0])
|
||
|
|
if ar_w > 0 and ar_h > 0:
|
||
|
|
params.max_tokens = ar_h * (ar_w + 1) + 1
|
||
|
|
if not replaced and len(sampling_params_list) == 1:
|
||
|
|
sampling_params_list = [diffusion_params]
|
||
|
|
return prompt_dict, sampling_params_list
|
||
|
|
|
||
|
|
|
||
|
|
def stage_metrics(outputs) -> dict:
|
||
|
|
"""Pull vllm-omni's per-stage timing out of the request output when present."""
|
||
|
|
res = {}
|
||
|
|
for out in outputs or []:
|
||
|
|
metrics = getattr(out, "metrics", None)
|
||
|
|
if isinstance(metrics, dict) and "stage_metrics" in metrics:
|
||
|
|
for sid, m in metrics["stage_metrics"].items():
|
||
|
|
if not isinstance(m, dict):
|
||
|
|
continue
|
||
|
|
res[str(sid)] = {
|
||
|
|
k: m.get(k)
|
||
|
|
for k in ("stage_gen_time_ms", "num_tokens_in", "num_tokens_out", "vllm_ttft_ms", "vllm_tpot_ms")
|
||
|
|
if k in m
|
||
|
|
}
|
||
|
|
return res
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
args = parse_args()
|
||
|
|
result = {
|
||
|
|
"label": args.label,
|
||
|
|
"model": args.model,
|
||
|
|
"prompt": args.prompt,
|
||
|
|
"negative_prompt": args.negative_prompt,
|
||
|
|
"sample_memory": args.sample_memory,
|
||
|
|
"model_fs": fs_type(args.model),
|
||
|
|
"deploy_config": args.deploy_config,
|
||
|
|
"height": args.height,
|
||
|
|
"width": args.width,
|
||
|
|
"seed": args.seed,
|
||
|
|
"extra_body": args.extra_body,
|
||
|
|
"parallel_stage_init": args.parallel_stage_init,
|
||
|
|
"omp_num_threads": os.environ.get("OMP_NUM_THREADS"),
|
||
|
|
"env": gpu_info(),
|
||
|
|
"phases_s": {"imports": round(T_IMPORTS_DONE - T_PROCESS_START, 3)},
|
||
|
|
"requests": [],
|
||
|
|
}
|
||
|
|
|
||
|
|
omni_kwargs = {"model": args.model, "mode": "text-to-image", "log_stats": True}
|
||
|
|
if args.deploy_config:
|
||
|
|
omni_kwargs["deploy_config"] = args.deploy_config
|
||
|
|
if args.parallel_stage_init:
|
||
|
|
omni_kwargs["parallel_stage_init"] = True
|
||
|
|
|
||
|
|
sampler = PeakSampler().start() if args.sample_memory else None
|
||
|
|
t0 = time.perf_counter()
|
||
|
|
omni = None
|
||
|
|
try:
|
||
|
|
omni = Omni(**omni_kwargs)
|
||
|
|
t1 = time.perf_counter()
|
||
|
|
result["peak_mem_startup"] = sampler.stop() if sampler else {}
|
||
|
|
sampler = None
|
||
|
|
result["phases_s"]["engine_init"] = round(t1 - t0, 3)
|
||
|
|
result["phases_s"]["process_to_engine_ready"] = round(t1 - T_PROCESS_START, 3)
|
||
|
|
print(
|
||
|
|
f"BENCH engine_init={t1 - t0:.2f}s process_to_engine_ready={t1 - T_PROCESS_START:.2f}s "
|
||
|
|
f"peak_mem_startup={result['peak_mem_startup']}",
|
||
|
|
flush=True,
|
||
|
|
)
|
||
|
|
sampler = PeakSampler().start() if args.sample_memory else None
|
||
|
|
|
||
|
|
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
|
||
|
|
outputs = None
|
||
|
|
first_outputs = None
|
||
|
|
first_image_elapsed = 0.0
|
||
|
|
for i in range(args.repeat):
|
||
|
|
prompt_dict, sampling_params_list = build_request(omni, args, generator)
|
||
|
|
tr0 = time.perf_counter()
|
||
|
|
outputs = omni.generate(prompt_dict, sampling_params_list=sampling_params_list)
|
||
|
|
tr1 = time.perf_counter()
|
||
|
|
if first_outputs is None:
|
||
|
|
first_outputs = outputs
|
||
|
|
first_image_elapsed = tr1 - T_PROCESS_START
|
||
|
|
rec = {
|
||
|
|
"index": i,
|
||
|
|
"kind": "first" if i == 0 else "steady",
|
||
|
|
"latency_s": round(tr1 - tr0, 3),
|
||
|
|
"stages": stage_metrics(outputs),
|
||
|
|
}
|
||
|
|
result["requests"].append(rec)
|
||
|
|
print(f"BENCH request[{i}] {rec['kind']} latency={tr1 - tr0:.2f}s stages={rec['stages']}", flush=True)
|
||
|
|
|
||
|
|
result["peak_mem_requests"] = sampler.stop() if sampler else {}
|
||
|
|
sampler = None
|
||
|
|
steady = [r["latency_s"] for r in result["requests"] if r["kind"] == "steady"]
|
||
|
|
result["summary_s"] = {
|
||
|
|
"imports": result["phases_s"]["imports"],
|
||
|
|
"engine_init": result["phases_s"]["engine_init"],
|
||
|
|
"first_request": result["requests"][0]["latency_s"],
|
||
|
|
"steady_request_avg": round(sum(steady) / len(steady), 3) if steady else None,
|
||
|
|
"time_to_first_image_from_process_start": round(first_image_elapsed, 3),
|
||
|
|
}
|
||
|
|
|
||
|
|
for key in ("host_rss_peak_gib", "gpu_used_peak_gib"):
|
||
|
|
result["summary_s"][key] = (
|
||
|
|
max(result["peak_mem_startup"][key], result["peak_mem_requests"][key]) if args.sample_memory else None
|
||
|
|
)
|
||
|
|
|
||
|
|
from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs
|
||
|
|
|
||
|
|
images = extract_images_from_outputs(first_outputs)
|
||
|
|
if not images:
|
||
|
|
raise RuntimeError("First request returned no image")
|
||
|
|
if args.save_image:
|
||
|
|
os.makedirs(os.path.dirname(os.path.abspath(args.save_image)), exist_ok=True)
|
||
|
|
images[0].save(args.save_image)
|
||
|
|
|
||
|
|
print("BENCH_JSON " + json.dumps(result), flush=True)
|
||
|
|
if args.output_json:
|
||
|
|
with open(args.output_json, "w") as f:
|
||
|
|
json.dump(result, f, indent=2)
|
||
|
|
print(f"BENCH wrote {args.output_json}", flush=True)
|
||
|
|
finally:
|
||
|
|
if sampler:
|
||
|
|
sampler.stop()
|
||
|
|
if omni is not None:
|
||
|
|
omni.close()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|