llm-torch-profiler-analysis

1#!/usr/bin/env python32"""Run a small correctness and latency probe against an LLM server."""3 4from __future__ import annotations5 6import argparse7import json8import math9import statistics10import time11from pathlib import Path12from typing import Any, Dict, List, Optional13from urllib import request14 15from profile_common import extract_openai_chat_text16 17DEFAULT_PROMPTS = [18    "Introduce Shanghai in one short sentence.",19    "What is 2+2? Answer briefly.",20    "Write one short haiku about GPUs.",21]22 23 24def parse_args() -> argparse.Namespace:25    parser = argparse.ArgumentParser(26        description=(27            "Send a few short requests to an LLM server and record latency plus "28            "sample outputs."29        )30    )31    parser.add_argument(32        "--framework",33        required=True,34        choices=("sglang", "vllm", "trtllm", "tokenspeed"),35        help="Serving framework.",36    )37    parser.add_argument(38        "--url",39        required=True,40        help="Server base URL, for example http://127.0.0.1:30000.",41    )42    parser.add_argument(43        "--model",44        default=None,45        help="OpenAI model id. Auto-discovered for vLLM, TensorRT-LLM, and TokenSpeed when omitted.",46    )47    parser.add_argument(48        "--requests",49        type=int,50        default=6,51        help="How many probe requests to send.",52    )53    parser.add_argument(54        "--max-tokens",55        type=int,56        default=48,57        help="Generation length for each request.",58    )59    parser.add_argument(60        "--timeout",61        type=float,62        default=180.0,63        help="Per-request timeout in seconds.",64    )65    parser.add_argument(66        "--prompt",67        action="append",68        default=[],69        help="Optional prompt override. Repeat to add more prompts.",70    )71    parser.add_argument(72        "--output",73        default=None,74        help="Optional JSON output path.",75    )76    return parser.parse_args()77 78 79def post_json(url: str, payload: Dict[str, Any], timeout: float) -> Dict[str, Any]:80    req = request.Request(81        url=url,82        data=json.dumps(payload).encode("utf-8"),83        headers={"Content-Type": "application/json"},84        method="POST",85    )86    with request.urlopen(req, timeout=timeout) as resp:87        raw = resp.read()88    return json.loads(raw.decode("utf-8")) if raw else {}89 90 91def get_json(url: str, timeout: float) -> Dict[str, Any]:92    req = request.Request(url=url, method="GET")93    with request.urlopen(req, timeout=timeout) as resp:94        raw = resp.read()95    return json.loads(raw.decode("utf-8")) if raw else {}96 97 98def discover_openai_model(base_url: str, timeout: float) -> str:99    payload = get_json(base_url.rstrip("/") + "/v1/models", timeout=timeout)100    data = payload.get("data")101    if not isinstance(data, list) or not data:102        raise RuntimeError(f"No models returned by {base_url.rstrip('/')}/v1/models")103    first = data[0]104    if isinstance(first, dict) and first.get("id"):105        return str(first["id"])106    raise RuntimeError(f"Malformed /v1/models payload from {base_url.rstrip('/')}")107 108 109def p95(values: List[float]) -> Optional[float]:110    if not values:111        return None112    ordered = sorted(values)113    index = max(0, math.ceil(len(ordered) * 0.95) - 1)114    return ordered[index]115 116 117def sglang_request(base_url: str, prompt: str, max_tokens: int, timeout: float) -> str:118    payload = {119        "text": prompt,120        "sampling_params": {121            "temperature": 0.0,122            "max_new_tokens": max_tokens,123        },124        "stream": False,125    }126    body = post_json(base_url.rstrip("/") + "/generate", payload, timeout=timeout)127    return str(body.get("text", ""))128 129 130def openai_request(131    base_url: str,132    model: str,133    prompt: str,134    max_tokens: int,135    timeout: float,136) -> Dict[str, str]:137    payload = {138        "model": model,139        "messages": [{"role": "user", "content": prompt}],140        "temperature": 0.0,141        "max_tokens": max_tokens,142        "stream": False,143    }144    body = post_json(145        base_url.rstrip("/") + "/v1/chat/completions",146        payload,147        timeout=timeout,148    )149    text, source = extract_openai_chat_text(body)150    return {"text": text, "source": source}151 152 153def run_probe(args: argparse.Namespace) -> Dict[str, Any]:154    prompts = args.prompt or list(DEFAULT_PROMPTS)155    model = args.model156    if args.framework in {"vllm", "trtllm", "tokenspeed"} and not model:157        model = discover_openai_model(args.url, timeout=args.timeout)158 159    latencies: List[float] = []160    samples: List[Dict[str, Any]] = []161    errors: List[Dict[str, str]] = []162 163    for request_idx in range(args.requests):164        prompt = prompts[request_idx % len(prompts)]165        start = time.time()166        try:167            if args.framework == "sglang":168                text = sglang_request(169                    args.url,170                    prompt,171                    max_tokens=args.max_tokens,172                    timeout=args.timeout,173                )174                source = "generate.text"175            else:176                assert model is not None177                result = openai_request(178                    args.url,179                    model,180                    prompt,181                    max_tokens=args.max_tokens,182                    timeout=args.timeout,183                )184                text = result["text"]185                source = result["source"]186            elapsed = time.time() - start187            latencies.append(elapsed)188            samples.append(189                {190                    "prompt": prompt,191                    "latency_s": round(elapsed, 3),192                    "content": text[:240],193                    "source": source,194                    "non_empty": bool(text.strip()),195                }196            )197        except Exception as exc:  # pragma: no cover - runtime probe path198            errors.append({"prompt": prompt, "error": repr(exc)})199 200    return {201        "framework": args.framework,202        "url": args.url,203        "model": model,204        "requests": args.requests,205        "success": len(samples),206        "errors": len(errors),207        "all_non_empty": (208            all(sample["non_empty"] for sample in samples) if samples else False209        ),210        "avg_latency_s": round(statistics.mean(latencies), 3) if latencies else None,211        "p95_latency_s": round(p95(latencies), 3) if latencies else None,212        "samples": samples[:3],213        "error_samples": errors[:3],214    }215 216 217def main() -> int:218    args = parse_args()219    summary = run_probe(args)220    rendered = json.dumps(summary, ensure_ascii=False, indent=2)221    print(rendered)222    if args.output:223        output_path = Path(args.output).expanduser().resolve()224        output_path.parent.mkdir(parents=True, exist_ok=True)225        output_path.write_text(rendered + "\n", encoding="utf-8")226    return 0227 228 229if __name__ == "__main__":230    raise SystemExit(main())231