llm-torch-profiler-analysis

1"""Bundle one or more triage text reports into a single markdown document."""2 3from __future__ import annotations4 5import argparse6from collections import defaultdict7from datetime import datetime, timezone8from pathlib import Path9from typing import Dict, List, Optional, Sequence, Tuple10 11FRAMEWORK_LABELS = {12    "sglang": "SGLang",13    "vllm": "vLLM",14    "trtllm": "TensorRT-LLM",15    "tokenspeed": "TokenSpeed",16}17 18FRAMEWORK_ORDER = {"sglang": 0, "vllm": 1, "trtllm": 2, "tokenspeed": 3}19 20 21def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:22    parser = argparse.ArgumentParser(23        description=(24            "Render multiple profiler triage text outputs into one markdown file. "25            "Input files are expected to be the existing analysis_*.txt outputs "26            "already emitted by analyze_llm_torch_profile.py."27        )28    )29    parser.add_argument(30        "--analysis-root",31        type=str,32        default=None,33        help=(34            "Root directory to scan recursively for analysis_*.txt files. "35            "Parent directory names are used as model section ids."36        ),37    )38    parser.add_argument(39        "--analysis-file",40        action="append",41        default=[],42        help=(43            "Explicit analysis file entry. Use either PATH or LABEL=PATH. "44            "When LABEL is omitted, the parent directory name is used."45        ),46    )47    parser.add_argument(48        "--title",49        type=str,50        default="Unified LLM Torch Profiler Triage Bundle",51        help="Top-level markdown title.",52    )53    parser.add_argument(54        "--output",55        type=str,56        default=None,57        help="Write the bundled markdown to this file. Prints to stdout when omitted.",58    )59    parser.add_argument(60        "--include-toc",61        action=argparse.BooleanOptionalAction,62        default=True,63        help="Include a simple table of contents.",64    )65    args = parser.parse_args(argv)66    if not args.analysis_root and not args.analysis_file:67        parser.error("Provide at least one of --analysis-root or --analysis-file.")68    return args69 70 71def framework_key_from_path(path: Path) -> str:72    lowered = path.name.lower()73    if "sglang" in lowered:74        return "sglang"75    if "vllm" in lowered:76        return "vllm"77    if "trtllm" in lowered or "tensorrt" in lowered:78        return "trtllm"79    if "tokenspeed" in lowered or "token-speed" in lowered:80        return "tokenspeed"81    return "other"82 83 84def framework_label(framework_key: str) -> str:85    return FRAMEWORK_LABELS.get(framework_key, framework_key)86 87 88def discover_analysis_files(root: Path) -> List[Tuple[str, Path]]:89    entries: List[Tuple[str, Path]] = []90    for path in sorted(root.rglob("analysis*.txt")):91        entries.append((path.parent.name, path))92    return entries93 94 95def parse_explicit_entry(raw: str) -> Tuple[str, Path]:96    if "=" in raw:97        label, path_text = raw.split("=", 1)98        path = Path(path_text).expanduser().resolve()99        return label.strip(), path100    path = Path(raw).expanduser().resolve()101    return path.parent.name, path102 103 104def slugify(text: str) -> str:105    chars = []106    last_dash = False107    for char in text.lower():108        if char.isalnum():109            chars.append(char)110            last_dash = False111        elif not last_dash:112            chars.append("-")113            last_dash = True114    return "".join(chars).strip("-")115 116 117def extract_model_name(report_text: str) -> Optional[str]:118    for line in report_text.splitlines():119        if line.startswith("Model: "):120            return line.split("Model: ", 1)[1].strip()121    return None122 123 124def choose_model_display_name(125    current: Optional[str],126    candidate: Optional[str],127    *,128    label: str,129) -> str:130    if candidate and candidate != label:131        if not current or current == label:132            return candidate133        if len(candidate) > len(current):134            return candidate135        return current136    if current:137        return current138    return label139 140 141def normalize_report_text(report_text: str) -> str:142    text = report_text.replace("\r\n", "\n").strip()143    if not text:144        return "_Empty analysis output._"145    heading_map = {146        "Triage View": "#### Triage View",147        "Kernel Table": "#### Kernel Table",148        "Overlap Opportunity Table": "#### Overlap Opportunity Table",149        "Fuse Opportunity Table": "#### Fuse Opportunity Table",150    }151    normalized_lines = []152    for line in text.splitlines():153        normalized_lines.append(heading_map.get(line, line))154    return "\n".join(normalized_lines)155 156 157def build_bundle_markdown(158    *,159    title: str,160    labeled_paths: Sequence[Tuple[str, Path]],161    include_toc: bool,162) -> str:163    grouped: Dict[str, List[Tuple[str, Path, str]]] = defaultdict(list)164    model_display: Dict[str, str] = {}165 166    for label, path in labeled_paths:167        raw_text = path.read_text(encoding="utf-8")168        report_text = normalize_report_text(raw_text)169        model_name = extract_model_name(report_text)170        grouped[label].append((framework_key_from_path(path), path, report_text))171        model_display[label] = choose_model_display_name(172            model_display.get(label),173            model_name,174            label=label,175        )176 177    ordered_labels = sorted(178        grouped,179        key=lambda item: (model_display[item].lower(), item.lower()),180    )181 182    lines: List[str] = [f"# {title}", ""]183    lines.append(184        f"_Generated on {datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S UTC')}_"185    )186    lines.append("")187 188    if include_toc:189        lines.append("## Contents")190        lines.append("")191        for label in ordered_labels:192            lines.append(193                f"- [{model_display[label]}](#{slugify(model_display[label])})"194            )195        lines.append("")196 197    for label in ordered_labels:198        display_name = model_display[label]199        lines.append(f"## {display_name}")200        lines.append("")201        lines.append(f"Model id: `{label}`")202        lines.append("")203 204        records = sorted(205            grouped[label],206            key=lambda item: (207                FRAMEWORK_ORDER.get(item[0], 99),208                item[1].name.lower(),209            ),210        )211 212        for framework_key, path, report_text in records:213            lines.append(f"### {framework_label(framework_key)}")214            lines.append("")215            lines.append(f"Source: `{path}`")216            lines.append("")217            lines.append(report_text)218            lines.append("")219 220    return "\n".join(lines).rstrip() + "\n"221 222 223def main(argv: Optional[Sequence[str]] = None) -> int:224    args = parse_args(argv)225 226    labeled_paths: List[Tuple[str, Path]] = []227    if args.analysis_root:228        labeled_paths.extend(229            discover_analysis_files(Path(args.analysis_root).expanduser().resolve())230        )231    for raw_entry in args.analysis_file:232        labeled_paths.append(parse_explicit_entry(raw_entry))233 234    existing = []235    missing = []236    for label, path in labeled_paths:237        if path.is_file():238            existing.append((label, path))239        else:240            missing.append(str(path))241    if missing:242        raise SystemExit("Missing analysis files:\n" + "\n".join(missing))243    if not existing:244        raise SystemExit("No analysis files found.")245 246    markdown = build_bundle_markdown(247        title=args.title,248        labeled_paths=existing,249        include_toc=args.include_toc,250    )251 252    if args.output:253        output_path = Path(args.output).expanduser().resolve()254        output_path.parent.mkdir(parents=True, exist_ok=True)255        output_path.write_text(markdown, encoding="utf-8")256    else:257        print(markdown, end="")258    return 0259 260 261if __name__ == "__main__":262    raise SystemExit(main())263