199 lines
8.2 KiB
Python
199 lines
8.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Separate paired untraced/traced requests through Omni's existing profiler RPC."""
|
|
|
|
import argparse
|
|
import asyncio
|
|
import hashlib
|
|
import os
|
|
import platform
|
|
import statistics
|
|
import sys
|
|
import time
|
|
import traceback
|
|
from pathlib import Path
|
|
|
|
from profile_local_text import save
|
|
|
|
|
|
async def run(args):
|
|
from vllm_omni import AsyncOmni
|
|
from vllm_omni.edge.local.engine import apply_runtime_env
|
|
from vllm_omni.edge.local.manifest import runtime_versions
|
|
|
|
args.out = args.out.resolve()
|
|
args.out.mkdir(parents=True, exist_ok=False)
|
|
trace_dir = (args.out / "traces").resolve()
|
|
report = {
|
|
"status": "running",
|
|
"start_unix": time.time(),
|
|
"argv": sys.argv,
|
|
"python": sys.executable,
|
|
"platform": platform.platform(),
|
|
"runtime": runtime_versions().to_dict(),
|
|
"environment": apply_runtime_env(),
|
|
"scope": "diagnostic paired profiling overhead; not steady-state latency percentiles",
|
|
"text_output_tokens": 32,
|
|
"requests": [],
|
|
}
|
|
config = {
|
|
"profiler": "torch",
|
|
"torch_profiler_dir": str(trace_dir),
|
|
"torch_profiler_with_stack": False,
|
|
"torch_profiler_record_shapes": False,
|
|
"torch_profiler_with_memory": False,
|
|
# Archive compression happens after writers exit, avoiding async gzip/hash races.
|
|
"torch_profiler_use_gzip": False,
|
|
"torch_profiler_dump_cuda_time_total": False,
|
|
}
|
|
report["profiler_config"] = config
|
|
save(args.out / "report.json", report)
|
|
omni = None
|
|
profiling = False
|
|
try:
|
|
params = None
|
|
if args.kind == "spark":
|
|
from vllm import SamplingParams
|
|
|
|
from vllm_omni.edge.local.plan import plan_text_session
|
|
|
|
plan = plan_text_session(
|
|
args.model, max_model_len=4096, max_num_seqs=4, max_num_batched_tokens=512, enforce_eager=True
|
|
)
|
|
report["plan"] = plan.to_dict()
|
|
if not plan.admitted:
|
|
report["status"] = "rejected"
|
|
return
|
|
kwargs = plan.engine_kwargs
|
|
prompt = "The garden has trees, flowers, and a small pond. Describe the garden in detail."
|
|
params = SamplingParams(temperature=0, max_tokens=32, ignore_eos=True)
|
|
else:
|
|
from transformers import AutoTokenizer
|
|
|
|
from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig
|
|
from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import Qwen3TTSPromptEmbedsBuilder
|
|
|
|
cfg = Qwen3TTSConfig.from_pretrained(args.model)
|
|
tok = AutoTokenizer.from_pretrained(args.model)
|
|
info = {
|
|
"task_type": ["CustomVoice"],
|
|
"text": ["Hello, this is a speech test."],
|
|
"language": ["English"],
|
|
"speaker": ["Vivian"],
|
|
}
|
|
length = Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
|
|
additional_information=info,
|
|
task_type="CustomVoice",
|
|
tokenize_prompt=lambda x: tok(x, padding=False)["input_ids"],
|
|
codec_language_id=cfg.talker_config.codec_language_id,
|
|
spk_is_dialect=cfg.talker_config.spk_is_dialect,
|
|
)
|
|
prompt = {"prompt_token_ids": [0] * length, "additional_information": info}
|
|
kwargs = {"deploy_profile": "edge", "stage_init_timeout": 180}
|
|
start = time.perf_counter()
|
|
omni = AsyncOmni(model=args.model, **kwargs, profiler_config=config)
|
|
report["startup_s"] = time.perf_counter() - start
|
|
|
|
async def request(phase, number):
|
|
start = time.perf_counter()
|
|
outputs = 0
|
|
finished = False
|
|
audio_samples = 0
|
|
async for output in omni.generate(prompt, sampling_params=params, request_id=f"{phase}-{number}"):
|
|
outputs += 1
|
|
finished = finished or output.finished
|
|
if args.kind == "tts" and output.outputs:
|
|
mm = output.outputs[0].multimodal_output
|
|
audio = mm.get("audio") if mm else None
|
|
parts = audio if isinstance(audio, list) else [audio]
|
|
audio_samples += sum(int(part.numel()) for part in parts if part is not None)
|
|
row = {
|
|
"phase": phase,
|
|
"index": number,
|
|
"wall_s": time.perf_counter() - start,
|
|
"events": outputs,
|
|
"finished": finished,
|
|
"audio_samples": audio_samples,
|
|
}
|
|
report["requests"].append(row)
|
|
save(args.out / "report.json", report)
|
|
if not finished or not outputs:
|
|
raise RuntimeError(f"Incomplete diagnostic request: {row}")
|
|
|
|
await request("warmup", 0)
|
|
for i in range(3):
|
|
await request("untraced_before", i)
|
|
rpc_started = time.perf_counter()
|
|
profiling = True
|
|
if args.kind == "tts":
|
|
# Both workers are rank 0. A shared prefix overwrites the first stage's trace.
|
|
report["stage_trace_prefixes"] = {str(stage): f"e2e-stage{stage}" for stage in (0, 1)}
|
|
report["start_rpc"] = [
|
|
await omni.start_profile(profile_prefix=prefix, stages=[int(stage)])
|
|
for stage, prefix in report["stage_trace_prefixes"].items()
|
|
]
|
|
else:
|
|
report["start_rpc"] = await omni.start_profile(profile_prefix="e2e-diagnostic")
|
|
report["start_rpc_s"] = time.perf_counter() - rpc_started
|
|
profiling = True
|
|
for i in range(3):
|
|
await request("traced", i)
|
|
rpc_started = time.perf_counter()
|
|
report["stop_rpc"] = await omni.stop_profile()
|
|
report["stop_rpc_s"] = time.perf_counter() - rpc_started
|
|
profiling = False
|
|
for i in range(3):
|
|
await request("untraced_after", i)
|
|
medians = {
|
|
phase: statistics.median(r["wall_s"] for r in report["requests"] if r["phase"] == phase)
|
|
for phase in ("untraced_before", "traced", "untraced_after")
|
|
}
|
|
report["median_wall_s"] = medians
|
|
report["traced_over_untraced_before"] = medians["traced"] / medians["untraced_before"]
|
|
report["traced_over_untraced_after"] = medians["traced"] / medians["untraced_after"]
|
|
if args.kind == "tts":
|
|
report["overhead_caveat"] = (
|
|
"Same prompt but speech generation can vary; compare audio_samples alongside wall-time ratios."
|
|
)
|
|
report["traces"] = [
|
|
{
|
|
"path": str(f.relative_to(args.out)),
|
|
"bytes": f.stat().st_size,
|
|
"sha256": hashlib.sha256(f.read_bytes()).hexdigest(),
|
|
}
|
|
for f in trace_dir.rglob("*")
|
|
if f.is_file()
|
|
]
|
|
if not report["traces"]:
|
|
raise RuntimeError("Profiler RPC returned but no trace artifact was written")
|
|
for prefix in report.get("stage_trace_prefixes", {}).values():
|
|
if not any(prefix in trace["path"] and ".json" in trace["path"] for trace in report["traces"]):
|
|
raise RuntimeError(f"Missing trace artifact for {prefix}")
|
|
report["status"] = "completed"
|
|
except Exception as error:
|
|
report.update(status="failed", error=repr(error), traceback=traceback.format_exc())
|
|
traceback.print_exc()
|
|
finally:
|
|
if profiling and omni is not None:
|
|
try:
|
|
await omni.stop_profile()
|
|
except Exception as error:
|
|
report["stop_error"] = repr(error)
|
|
if omni is not None:
|
|
omni.shutdown()
|
|
report["end_unix"] = time.time()
|
|
save(args.out / "report.json", report)
|
|
return report["status"] == "completed"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
p = argparse.ArgumentParser(description=__doc__)
|
|
p.add_argument("--kind", choices=["spark", "tts"], required=True)
|
|
p.add_argument("--model", required=True)
|
|
p.add_argument("--out", type=Path, required=True)
|
|
args = p.parse_args()
|
|
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
|
from vllm_omni.windows.aio import install_selector_policy
|
|
|
|
install_selector_policy()
|
|
raise SystemExit(0 if asyncio.run(run(args)) else 1)
|