llm-torch-profiler-analysis

1"""Generate a TensorRT-LLM py_executor override for stable torch-profiler capture."""2 3from __future__ import annotations4 5import argparse6from dataclasses import dataclass7from pathlib import Path8 9START_MARKER = "torch_profiler = torch.profiler.profile("10 11 12@dataclass13class ProfileCallSpan:14    start: int15    end: int16    block: str17 18 19def parse_args() -> argparse.Namespace:20    parser = argparse.ArgumentParser(21        description=(22            "Create a py_executor.py override that enables with_stack=True for "23            "TensorRT-LLM torch-profiler traces."24        )25    )26    parser.add_argument("--source", required=True, help="Original py_executor.py path.")27    parser.add_argument("--output", required=True, help="Override file path to write.")28    return parser.parse_args()29 30 31def find_profile_call_span(text: str) -> ProfileCallSpan:32    start = text.find(START_MARKER)33    if start == -1:34        raise SystemExit("Could not find torch profiler setup in source file.")35 36    open_paren = text.find("(", start)37    if open_paren == -1:38        raise SystemExit("Malformed torch profiler setup in source file.")39 40    depth = 041    for index in range(open_paren, len(text)):42        char = text[index]43        if char == "(":44            depth += 145        elif char == ")":46            depth -= 147            if depth == 0:48                return ProfileCallSpan(49                    start=start,50                    end=index + 1,51                    block=text[start : index + 1],52                )53    raise SystemExit("Could not find the end of the torch profiler call.")54 55 56def inject_with_stack(block: str) -> str:57    if "with_stack=" in block:58        return block59 60    lines = block.splitlines()61    if not lines:62        raise SystemExit("Unexpected torch profiler block format.")63 64    last_line = lines[-1]65    if not last_line.strip():66        raise SystemExit("Unexpected torch profiler block terminator.")67 68    if last_line.strip() == ")":69        if len(lines) < 2:70            raise SystemExit("Could not find the last torch profiler argument line.")71        last_arg_index = len(lines) - 272        last_arg_line = lines[last_arg_index]73        indent = last_arg_line[: len(last_arg_line) - len(last_arg_line.lstrip())]74        if not last_arg_line.rstrip().endswith(","):75            lines[last_arg_index] = last_arg_line.rstrip() + ","76        lines.insert(len(lines) - 1, f"{indent}with_stack=True")77        return "\n".join(lines)78 79    if not last_line.rstrip().endswith(")"):80        raise SystemExit("Unexpected torch profiler block terminator.")81 82    indent = last_line[: len(last_line) - len(last_line.lstrip())]83    last_arg_text = last_line.rstrip()[:-1].rstrip()84    if not last_arg_text.endswith(","):85        last_arg_text += ","86    lines[-1] = last_arg_text87    lines.append(f"{indent}with_stack=True)")88    return "\n".join(lines)89 90 91def inject_rank0_trace_guard(text: str) -> str:92    needle = (93        "        enable_torch_trace = bool(torch_trace_path and profile_start_stop)\n"94    )95    replacement = (96        "        # Multi-rank PyTorch backend workers race on the same chrome-trace "97        "path.\n"98        "        # Keep the full torch-profiler trace on rank 0 and let the other "99        "ranks\n"100        "        # continue with CUDA-profiler gating only.\n"101        "        enable_torch_trace = bool(\n"102        "            torch_trace_path and profile_start_stop and self.dist.rank == 0\n"103        "        )\n"104    )105    if replacement in text:106        return text107    if needle not in text:108        raise SystemExit("Could not find enable_torch_trace assignment in source file.")109    return text.replace(needle, replacement, 1)110 111 112def main() -> int:113    args = parse_args()114    source = Path(args.source).expanduser().resolve()115    output = Path(args.output).expanduser().resolve()116    text = source.read_text(encoding="utf-8")117    span = find_profile_call_span(text)118    patched_block = inject_with_stack(span.block)119    patched = (120        text121        if patched_block == span.block122        else (text[: span.start] + patched_block + text[span.end :])123    )124    patched = inject_rank0_trace_guard(patched)125    output.parent.mkdir(parents=True, exist_ok=True)126    output.write_text(patched, encoding="utf-8")127    print(output)128    return 0129 130 131if __name__ == "__main__":132    raise SystemExit(main())133