diff --git a/README.md b/README.md index 75c429edc..be1f140fa 100755 --- a/README.md +++ b/README.md @@ -31,6 +31,8 @@ TraceLens is a Python library for **automated performance analysis of training a **TraceLens Agent**: Receive a prioritized human-readable optimization report, derived through an agentic workflow, covering compute kernels, system bottlenecks, and kernel fusion opportunities with root-cause reasoning and concrete resolutions. +**Trace corpus indexing**: Build a searchable SQLite catalog of TraceLens reports so you can find traces by op, category, or kernel name without reopening every raw file. Scanner and importer sit behind a storage interface; SQLite is the first backend. + --- ## Quick Start @@ -79,6 +81,16 @@ TraceLens_compare_perf_reports_pytorch \ -o comparison.xlsx ``` +Index traces and existing TraceLens CSV reports (see [Index a corpus of traces](docs/how-to/trace-index.md)): + +```bash +TraceLens_trace_index --db trace_index.sqlite append \ + --trace-path /path/to/rank0_trace.json.gz \ + --report-dir path/to/perf_report_csvs +TraceLens_trace_index --db trace_index.sqlite build --traces-file traces.txt +TraceLens_trace_index --db trace_index.sqlite search Cijk +``` + For multi-rank runs, generate a collective-communication report across ranks (see [Generate a collective-communication report](docs/how-to/collective-report.md)): ```bash @@ -139,6 +151,7 @@ Each format's linked doc covers its full CLI reference. For PyTorch report compa | Multi-Rank Collective Report | [docs/how-to/collective-report.md](docs/how-to/collective-report.md) | | Performance Report Columns | [docs/reference/perf-report-columns.md](docs/reference/perf-report-columns.md) | | TraceLens Agent | [docs/how-to/agent.md](docs/how-to/agent.md) | +| TraceIndex | [docs/how-to/trace-index.md](docs/how-to/trace-index.md) | --- diff --git a/TraceLens/Agent/Analysis/category_analyses/analysis_utils.py b/TraceLens/Agent/Analysis/category_analyses/analysis_utils.py index 812cccc98..26f54f723 100644 --- a/TraceLens/Agent/Analysis/category_analyses/analysis_utils.py +++ b/TraceLens/Agent/Analysis/category_analyses/analysis_utils.py @@ -27,6 +27,7 @@ import numpy as np import pandas as pd +from TraceLens.PerfModel.kernel_library import classify_kernel_library from TraceLens.PerfModel.utils import torch_dtype_map TARGET_HIGH = 100.0 @@ -44,26 +45,6 @@ # Efficiency Boundary for grouping _EFF_BUCKET_BOUNDARIES = (30, 60) -_OP_NAME_LIBRARY_RULES = [ - ("aiter::", "AITER"), - ("rocm_aiter", "AITER"), - ("fbgemm", "FBGEMM"), - ("miopen", "MIOpen"), - ("triton", "Triton"), -] -_KERNEL_NAME_LIBRARY_RULES = [ - ("aiter", "AITER"), - ("ck_tile::", "CK"), - ("ck_tile6kentry", "CK"), - ("FmhaFwd", "CK"), - ("FmhaBwd", "CK"), - ("Cijk_", "Tensile"), - ("wvSplitK", "rocBLAS"), - ("splitKreduce", "rocBLAS"), - ("rocprim::", "rocPRIM"), - ("triton_", "Triton"), - ("void at::native::", "PyTorch Native"), -] _COMPARATIVE_SPEEDUP_COL = "speedup (trace2/trace1)" _COMPARATIVE_DELTA_COL = "delta_us (trace2 - trace1)" _COMPARATIVE_T2_TIME_COL = "lca_total_kernel_time_trace2_us" @@ -1093,15 +1074,5 @@ def run_category_analysis( print(f"Metrics written to: {output_path}") -# Category-specific helper functions -def classify_kernel_library(op_name: str, kernel_details: str = "") -> Optional[str]: - """Identify the GPU library backing an operation from its name or kernel strings.""" - op_lower = op_name.lower() - for marker, lib in _OP_NAME_LIBRARY_RULES: - if marker in op_lower: - return lib - kd = str(kernel_details) if kernel_details and not pd.isna(kernel_details) else "" - for marker, lib in _KERNEL_NAME_LIBRARY_RULES: - if marker in kd: - return lib - return None +# Category-specific helper functions — classify_kernel_library lives in +# TraceLens.PerfModel.kernel_library and is imported above. diff --git a/TraceLens/PerfModel/kernel_library.py b/TraceLens/PerfModel/kernel_library.py new file mode 100644 index 000000000..11ebd0f8b --- /dev/null +++ b/TraceLens/PerfModel/kernel_library.py @@ -0,0 +1,79 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Shared kernel-library classification for perf reports, agents, and TraceIndex.""" + +from __future__ import annotations + +from typing import Any, Optional + +from TraceLens.PerfModel import kernel_name_parser as knp + +_OP_NAME_LIBRARY_RULES = [ + ("aiter::", "AITER"), + ("rocm_aiter", "AITER"), + ("fbgemm", "FBGEMM"), + ("miopen", "MIOpen"), + ("triton", "Triton"), +] + +_KERNEL_NAME_LIBRARY_RULES = [ + ("aiter", "AITER"), + ("ck_tile::", "CK"), + ("ck_tile6kentry", "CK"), + ("FmhaFwd", "CK"), + ("FmhaBwd", "CK"), + ("Cijk_", "Tensile"), + ("wvSplitK", "rocBLAS"), + ("splitKreduce", "rocBLAS"), + ("rocprim::", "rocPRIM"), + ("triton_", "Triton"), + ("void at::native::", "PyTorch Native"), + ("nccl", "RCCL/NCCL"), + ("rccl", "RCCL/NCCL"), + ("composable", "CK"), +] + +# Ordered GEMM detectors; extend when kernel_name_parser gains more libraries (#805). +_GEMM_LIBRARY_CHECKS = ( + (knp.is_rocm_gemm, "Tensile"), + (knp.is_cuda_gemm, "nvjet"), +) + + +def _coerce_kernel_details(kernel_details: Any) -> str: + if kernel_details in (None, ""): + return "" + if isinstance(kernel_details, float) and kernel_details != kernel_details: + return "" + return str(kernel_details) + + +def _library_from_gemm_kernel_name(kernel_name: str) -> Optional[str]: + for detector, library in _GEMM_LIBRARY_CHECKS: + if detector(kernel_name): + return library + return None + + +def classify_kernel_library(op_name: str, kernel_details: Any = "") -> Optional[str]: + """Identify the GPU library backing an operation from its name or kernel strings.""" + op_lower = (op_name or "").lower() + for marker, library in _OP_NAME_LIBRARY_RULES: + if marker in op_lower: + return library + + kd = _coerce_kernel_details(kernel_details) + if kd: + gemm_library = _library_from_gemm_kernel_name(kd) + if gemm_library is not None: + return gemm_library + + kd_lower = kd.lower() + for marker, library in _KERNEL_NAME_LIBRARY_RULES: + if marker.lower() in kd_lower: + return library + return None diff --git a/TraceLens/TraceIndex/__init__.py b/TraceLens/TraceIndex/__init__.py new file mode 100644 index 000000000..1a1238632 --- /dev/null +++ b/TraceLens/TraceIndex/__init__.py @@ -0,0 +1,31 @@ +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Trace corpus indexing helpers.""" + +from .core import ( + append_trace, + build_traces, + execute_read_query, + generate_report_and_import, + search_index, +) +from .models import SearchHit, TraceRecord, TraceReport +from .sqlite_store import SQLiteTraceIndexStore +from .store import TraceIndexStore + +__all__ = [ + "TraceIndexStore", + "SQLiteTraceIndexStore", + "SearchHit", + "TraceRecord", + "TraceReport", + "append_trace", + "build_traces", + "execute_read_query", + "generate_report_and_import", + "search_index", +] diff --git a/TraceLens/TraceIndex/cli.py b/TraceLens/TraceIndex/cli.py new file mode 100644 index 000000000..dfd193836 --- /dev/null +++ b/TraceLens/TraceIndex/cli.py @@ -0,0 +1,225 @@ +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Command line entry point for TraceIndex.""" + +import argparse +import json +from pathlib import Path +from typing import List, Optional + +from TraceLens.TraceIndex.importer import ( + append_trace, + build_traces, + report_dir_for_trace, +) +from TraceLens.TraceIndex.sqlite_store import SQLiteTraceIndexStore +from TraceLens.TraceIndex.utils import collect_trace_paths + +DEFAULT_DB = Path("trace_index.sqlite") + + +def print_json(payload: object) -> None: + print(json.dumps(payload, indent=2, default=str)) + + +def create_store(args: argparse.Namespace): + if args.backend != "sqlite": + raise ValueError("only the sqlite backend is currently implemented") + return SQLiteTraceIndexStore(args.db) + + +def append_cmd(args: argparse.Namespace) -> int: + store = create_store(args) + try: + report_dir = args.report_dir + generated = report_dir is None + if generated: + report_dir = report_dir_for_trace(args.trace_path, args.report_root) + trace_id = append_trace( + store, + trace_path=args.trace_path, + report_dir=None if generated else report_dir, + root=args.root, + force=args.force, + enable_pseudo_ops=args.enable_pseudo_ops, + report_root=args.report_root, + ) + print_json( + { + "backend": args.backend, + "db": args.db, + "trace_id": trace_id, + "trace_path": args.trace_path, + "report_dir": report_dir, + "generated_report": generated, + } + ) + return 0 + finally: + store.close() + + +def build_cmd(args: argparse.Namespace) -> int: + trace_paths = collect_trace_paths(args.traces_file, args.trace_path) + if not trace_paths: + raise SystemExit( + "build requires --traces-file and/or one or more --trace-path values" + ) + store = create_store(args) + try: + result = build_traces( + store, + trace_paths, + report_root=args.report_root, + root=args.root, + force=args.force, + enable_pseudo_ops=args.enable_pseudo_ops, + ) + print_json( + { + "backend": args.backend, + "db": args.db, + "imported": result["imported"], + "failed": result["failed"], + } + ) + return 1 if result["failed"] else 0 + finally: + store.close() + + +def search_cmd(args: argparse.Namespace) -> int: + store = create_store(args) + try: + store.init_schema() + rows = [ + hit._asdict() + for hit in store.search(" ".join(args.terms), limit=args.limit) + ] + print_json({"backend": args.backend, "rows": rows}) + return 0 + finally: + store.close() + + +def sqlite_sql_cmd(args: argparse.Namespace) -> int: + store = create_store(args) + try: + store.init_schema() + rows = store.execute_read_query(args.sql, limit=args.limit) + print_json({"backend": args.backend, "rows": rows}) + return 0 + finally: + store.close() + + +def serve_cmd(args: argparse.Namespace) -> int: + from TraceLens.TraceIndex.server import serve + + serve( + db_path=args.db, + host=args.host, + port=args.port, + default_limit=args.default_limit, + max_limit=args.max_limit, + ) + return 0 + + +def add_generate_args(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "--report-root", + type=Path, + default=None, + help="Directory for generated CSV reports (default: trace_index_reports/)", + ) + parser.add_argument("--root", type=Path, default=None) + parser.add_argument( + "--force", + action="store_true", + help="Regenerate CSV reports even if they already exist", + ) + parser.add_argument("--enable-pseudo-ops", action="store_true") + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Build and query a TraceIndex catalog of traces." + ) + parser.add_argument( + "--backend", + choices=["sqlite"], + default="sqlite", + help="TraceIndex storage backend", + ) + parser.add_argument("--db", type=Path, default=DEFAULT_DB, help="SQLite DB path") + sub = parser.add_subparsers(dest="command") + + append = sub.add_parser( + "append", + help="Append one trace to the catalog, optionally from an existing CSV report", + ) + append.add_argument("--trace-path", type=Path, required=True) + append.add_argument( + "--report-dir", + type=Path, + default=None, + help="Existing TraceLens CSV report directory. If omitted, generate a training PyTorch report.", + ) + add_generate_args(append) + append.set_defaults(func=append_cmd) + + build = sub.add_parser( + "build", + help="Create or open the catalog and append a batch of traces", + ) + build.add_argument( + "--traces-file", + type=Path, + default=None, + help="Text file with one trace path per line (# comments allowed)", + ) + build.add_argument( + "--trace-path", + type=Path, + action="append", + default=[], + help="Trace path to include. Repeatable, can be combined with --traces-file", + ) + add_generate_args(build) + build.set_defaults(func=build_cmd) + + search = sub.add_parser("search", help="Full-text search indexed traces") + search.add_argument("terms", nargs="+") + search.add_argument("--limit", type=int, default=50) + search.set_defaults(func=search_cmd) + + sql = sub.add_parser("sqlite-sql", help="Run a read-only SQLite query") + sql.add_argument("sql") + sql.add_argument("--limit", type=int, default=500) + sql.set_defaults(func=sqlite_sql_cmd) + + serve = sub.add_parser("serve", help="Serve read-only HTTP SQL access") + serve.add_argument("--host", default="127.0.0.1") + serve.add_argument("--port", type=int, default=8765) + serve.add_argument("--default-limit", type=int, default=500) + serve.add_argument("--max-limit", type=int, default=5000) + serve.set_defaults(func=serve_cmd) + + return parser + + +def main(argv: Optional[List[str]] = None) -> int: + parser = build_parser() + args = parser.parse_args(argv) + if args.command is None: + parser.error("a command is required") + return args.func(args) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/TraceLens/TraceIndex/core.py b/TraceLens/TraceIndex/core.py new file mode 100644 index 000000000..e223da48d --- /dev/null +++ b/TraceLens/TraceIndex/core.py @@ -0,0 +1,108 @@ +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Compatibility facade for the default TraceIndex backend.""" + +from pathlib import Path +from typing import Any, Dict, List, Optional, Sequence + +from TraceLens.TraceIndex.importer import ( + append_trace as append_trace_with_store, + build_traces as build_traces_with_store, + generate_report_and_import as generate_report_and_import_with_store, +) +from TraceLens.TraceIndex.sqlite_store import SQLiteTraceIndexStore + + +def append_trace( + db_path: Path, + trace_path: Path, + report_dir: Optional[Path] = None, + root: Optional[Path] = None, + force: bool = False, + enable_pseudo_ops: bool = False, + report_root: Optional[Path] = None, +) -> int: + store = SQLiteTraceIndexStore(db_path) + try: + return append_trace_with_store( + store, + trace_path=trace_path, + report_dir=report_dir, + root=root, + force=force, + enable_pseudo_ops=enable_pseudo_ops, + report_root=report_root, + ) + finally: + store.close() + + +def build_traces( + db_path: Path, + trace_paths: List[Path], + report_root: Optional[Path] = None, + root: Optional[Path] = None, + force: bool = False, + enable_pseudo_ops: bool = False, +) -> Dict[str, Any]: + store = SQLiteTraceIndexStore(db_path) + try: + return build_traces_with_store( + store, + trace_paths, + report_root=report_root, + root=root, + force=force, + enable_pseudo_ops=enable_pseudo_ops, + ) + finally: + store.close() + + +def generate_report_and_import( + db_path: Path, + trace_path: Path, + report_dir: Optional[Path] = None, + root: Optional[Path] = None, + force: bool = False, + enable_pseudo_ops: bool = False, +) -> int: + store = SQLiteTraceIndexStore(db_path) + try: + return generate_report_and_import_with_store( + store, + trace_path=trace_path, + report_dir=report_dir, + root=root, + force=force, + enable_pseudo_ops=enable_pseudo_ops, + ) + finally: + store.close() + + +def search_index(db_path: Path, terms: str, limit: int = 50) -> List[Dict[str, Any]]: + store = SQLiteTraceIndexStore(db_path) + try: + store.init_schema() + return [hit._asdict() for hit in store.search(terms, limit=limit)] + finally: + store.close() + + +def execute_read_query( + db_path: Path, + sql: str, + params: Optional[Sequence[Any]] = None, + limit: int = 500, +) -> List[Dict[str, Any]]: + store = SQLiteTraceIndexStore(db_path) + try: + store.init_schema() + return store.execute_read_query(sql, params=params, limit=limit) + finally: + store.close() diff --git a/TraceLens/TraceIndex/importer.py b/TraceLens/TraceIndex/importer.py new file mode 100644 index 000000000..1ebb20730 --- /dev/null +++ b/TraceLens/TraceIndex/importer.py @@ -0,0 +1,169 @@ +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Backend-neutral TraceLens report import workflow.""" + +import re +from pathlib import Path +from typing import Any, Dict, List, Optional + +from TraceLens.TraceIndex.models import TraceRecord, TraceReport +from TraceLens.TraceIndex.scanner import trace_record_from_path +from TraceLens.TraceIndex.store import TraceIndexStore +from TraceLens.TraceIndex.utils import normalize_path, read_csv_rows + +DEFAULT_REPORT_ROOT = Path("trace_index_reports") + +REPORT_SHEETS = ( + "unified_perf_summary", + "ops_summary_by_category", + "gpu_timeline", +) + + +def load_report_dir(report_dir: Path) -> TraceReport: + report_dir = report_dir.resolve() + return TraceReport( + report_dir=normalize_path(report_dir), + sheets={ + sheet_name: read_csv_rows(report_dir / ("%s.csv" % sheet_name)) + for sheet_name in REPORT_SHEETS + }, + ) + + +def synthetic_trace_record_for_report(report_dir: Path) -> TraceRecord: + report_dir = report_dir.resolve() + return TraceRecord( + root=None, + path=normalize_path(report_dir), + rel_path=report_dir.name, + name=report_dir.name, + size_bytes=None, + md5=None, + format="tracelens_report_dir", + rank=None, + top_dir=None, + parent_rel=None, + should_enrich=True, + skip_reason=None, + ) + + +def import_report_dir( + store: TraceIndexStore, + report_dir: Path, + trace_path: Optional[Path] = None, + root: Optional[Path] = None, +) -> int: + store.init_schema() + trace = ( + trace_record_from_path(trace_path, root=root) + if trace_path is not None + else synthetic_trace_record_for_report(report_dir) + ) + trace_id = store.upsert_trace(trace) + store.import_report(trace_id, load_report_dir(report_dir)) + return trace_id + + +def report_dir_for_trace(trace_path: Path, report_root: Optional[Path] = None) -> Path: + safe_name = re.sub(r"[^A-Za-z0-9_.-]+", "_", trace_path.name) + return (report_root or DEFAULT_REPORT_ROOT) / safe_name + + +def append_trace( + store: TraceIndexStore, + trace_path: Path, + report_dir: Optional[Path] = None, + root: Optional[Path] = None, + force: bool = False, + enable_pseudo_ops: bool = False, + report_root: Optional[Path] = None, +) -> int: + """Append one trace to the catalog. + + If ``report_dir`` is set, import that existing CSV report. Otherwise generate + a training PyTorch CSV report and import it. + """ + if report_dir is not None: + return import_report_dir(store, report_dir, trace_path=trace_path, root=root) + return generate_report_and_import( + store, + trace_path=trace_path, + report_dir=report_dir_for_trace(trace_path, report_root), + root=root, + force=force, + enable_pseudo_ops=enable_pseudo_ops, + ) + + +def build_traces( + store: TraceIndexStore, + trace_paths: List[Path], + report_root: Optional[Path] = None, + root: Optional[Path] = None, + force: bool = False, + enable_pseudo_ops: bool = False, +) -> Dict[str, Any]: + """Generate reports and append a batch of traces. Continues after failures.""" + imported: List[Dict[str, Any]] = [] + failed: List[Dict[str, Any]] = [] + for trace_path in trace_paths: + report_dir = report_dir_for_trace(trace_path, report_root) + try: + trace_id = generate_report_and_import( + store, + trace_path=trace_path, + report_dir=report_dir, + root=root, + force=force, + enable_pseudo_ops=enable_pseudo_ops, + ) + imported.append( + { + "trace_id": trace_id, + "trace_path": normalize_path(trace_path), + "report_dir": normalize_path(report_dir), + } + ) + except Exception as exc: + failed.append( + { + "trace_path": normalize_path(trace_path), + "error": repr(exc), + } + ) + return {"imported": imported, "failed": failed} + + +def generate_report_and_import( + store: TraceIndexStore, + trace_path: Path, + report_dir: Optional[Path] = None, + root: Optional[Path] = None, + force: bool = False, + enable_pseudo_ops: bool = False, +) -> int: + if report_dir is None: + report_dir = report_dir_for_trace(trace_path) + report_dir = report_dir.resolve() + unified_csv = report_dir / "unified_perf_summary.csv" + if force or not unified_csv.exists(): + report_dir.mkdir(parents=True, exist_ok=True) + from TraceLens.Reporting.generate_perf_report_pytorch import ( # noqa: PLC0415 + generate_perf_report_pytorch, + ) + + generate_perf_report_pytorch( + profile_json_path=str(trace_path), + output_xlsx_path=None, + output_csvs_dir=str(report_dir), + kernel_summary=True, + include_first_occurrence_time=True, + enable_pseudo_ops=enable_pseudo_ops, + ) + return import_report_dir(store, report_dir, trace_path=trace_path, root=root) diff --git a/TraceLens/TraceIndex/models.py b/TraceLens/TraceIndex/models.py new file mode 100644 index 000000000..355d4756e --- /dev/null +++ b/TraceLens/TraceIndex/models.py @@ -0,0 +1,36 @@ +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Backend-neutral data objects for TraceIndex.""" + +from typing import Dict, List, NamedTuple, Optional + + +class TraceRecord(NamedTuple): + root: Optional[str] + path: str + rel_path: Optional[str] + name: str + size_bytes: Optional[int] + md5: Optional[str] + format: str + rank: Optional[int] + top_dir: Optional[str] + parent_rel: Optional[str] + should_enrich: bool + skip_reason: Optional[str] + + +class TraceReport(NamedTuple): + report_dir: str + sheets: Dict[str, List[Dict[str, str]]] + + +class SearchHit(NamedTuple): + trace_id: int + rel_path: Optional[str] + kind: str + hit: str diff --git a/TraceLens/TraceIndex/scanner.py b/TraceLens/TraceIndex/scanner.py new file mode 100644 index 000000000..042708e17 --- /dev/null +++ b/TraceLens/TraceIndex/scanner.py @@ -0,0 +1,147 @@ +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Trace metadata extraction for catalog ingest.""" + +import gzip +import hashlib +import re +from pathlib import Path +from typing import Optional + +from TraceLens.TraceIndex.models import TraceRecord +from TraceLens.TraceIndex.utils import normalize_path, rel_to + +RANK_RE = re.compile(r"(?:^|[^A-Za-z])rank[-_]?(\d+)(?:[^0-9]|$)", re.IGNORECASE) + +SKIP_PARTS_EXACT = { + ".git", + "__pycache__", + "node_modules", + "_perf_report_csvs", + "perf_report_csvs", + "gap_analysis", + "capture_traces", + "graph_capture", +} +SKIP_PARTS_CONTAINS = ( + "_perf_report_csvs", + "perf_report", + "gap_analysis", + "capture_traces", + "graph_capture", +) + + +def is_json_gz(path: Path) -> bool: + return path.name.lower().endswith(".json.gz") + + +def classify_skip(path: Path, root: Path) -> Optional[str]: + try: + rel_parts = [part.lower() for part in path.relative_to(root).parts[:-1]] + except ValueError: + rel_parts = [] + for part in rel_parts: + if part in SKIP_PARTS_EXACT: + return part + for token in SKIP_PARTS_CONTAINS: + if token in part: + return token + name = path.name.lower() + if name.endswith((".xlsx", ".csv", ".log", ".jsonl", ".md", ".txt")): + return "derived_or_log" + return None + + +def read_prefix(path: Path, max_bytes: int) -> str: + try: + if is_json_gz(path): + with gzip.open(path, "rb") as f: + data = f.read(max_bytes) + else: + with path.open("rb") as f: + data = f.read(max_bytes) + except (OSError, EOFError, gzip.BadGzipFile): + return "" + return data.decode("utf-8", errors="ignore") + + +def content_md5(path: Path, chunk_size: int = 8 * 1024 * 1024) -> Optional[str]: + hasher = hashlib.md5() + try: + with path.open("rb") as f: + for chunk in iter(lambda: f.read(chunk_size), b""): + hasher.update(chunk) + except OSError: + return None + return hasher.hexdigest() + + +def detect_format(path: Path, prefix: str) -> str: + name = path.name.lower() + suffix = path.suffix.lower() + if suffix == ".pftrace": + return "pftrace" + if suffix == ".rpd": + return "rocprof_rpd" + if name.endswith(".xplane.pb"): + return "xplane_pb" + if '"traceEvents"' in prefix: + return "kineto_chrome_json_gz" if is_json_gz(path) else "kineto_chrome_json" + if '"rocprofiler-sdk-tool"' in prefix: + return "rocprofv3_json_gz" if is_json_gz(path) else "rocprofv3_json" + if is_json_gz(path): + return "json_gz_unknown" + if suffix == ".json": + return "json_unknown" + return "trace_named_unknown" + + +def extract_rank(path: Path) -> Optional[int]: + match = RANK_RE.search(normalize_path(path)) + if not match: + return None + try: + return int(match.group(1)) + except ValueError: + return None + + +def trace_record_from_path( + trace_path: Path, + root: Optional[Path] = None, + peek_bytes: int = 2 * 1024 * 1024, + compute_md5: bool = False, +) -> TraceRecord: + trace_path = trace_path.resolve() + root = root.resolve() if root is not None else trace_path.parent.resolve() + prefix = read_prefix(trace_path, peek_bytes) + trace_format = detect_format(trace_path, prefix) + skip_reason = classify_skip(trace_path, root) + should_enrich = skip_reason is None and trace_format in { + "kineto_chrome_json", + "kineto_chrome_json_gz", + "rocprofv3_json", + "rocprofv3_json_gz", + "pftrace", + } + rel_path = rel_to(trace_path, root) + rel_parts = Path(rel_path).parts + return TraceRecord( + root=normalize_path(root), + path=normalize_path(trace_path), + rel_path=rel_path, + name=trace_path.name, + size_bytes=trace_path.stat().st_size if trace_path.exists() else None, + md5=content_md5(trace_path) if compute_md5 else None, + format=trace_format, + rank=extract_rank(trace_path), + top_dir=rel_parts[0] if rel_parts else "", + parent_rel=normalize_path(Path(*rel_parts[:-1])) if len(rel_parts) > 1 else "", + should_enrich=should_enrich, + skip_reason=skip_reason, + ) diff --git a/TraceLens/TraceIndex/server.py b/TraceLens/TraceIndex/server.py new file mode 100644 index 000000000..92f1ad83f --- /dev/null +++ b/TraceLens/TraceIndex/server.py @@ -0,0 +1,190 @@ +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Read-only HTTP query server for the SQLite TraceIndex backend.""" + +import json +import sqlite3 +import time +from http import HTTPStatus +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Any, List, Type +from urllib.parse import parse_qs, urlparse + +from TraceLens.TraceIndex.sqlite_store import is_read_only_sql + + +def json_bytes(payload: object) -> bytes: + return json.dumps(payload, indent=2, default=str).encode("utf-8") + + +def make_handler( + db_path: Path, default_limit: int, max_limit: int +) -> Type[BaseHTTPRequestHandler]: + class Handler(BaseHTTPRequestHandler): + server_version = "TraceIndexQuery/0.1" + + def log_message(self, fmt: str, *args: Any) -> None: + print("%s - %s" % (self.address_string(), fmt % args), flush=True) + + def send_json( + self, payload: object, status: HTTPStatus = HTTPStatus.OK + ) -> None: + body = json_bytes(payload) + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def read_json(self) -> dict: + length = int(self.headers.get("Content-Length", "0")) + if length <= 0: + return {} + return json.loads(self.rfile.read(length).decode("utf-8")) + + def connect(self) -> sqlite3.Connection: + uri = "file:%s?mode=ro" % db_path + conn = sqlite3.connect(uri, uri=True, timeout=30) + conn.row_factory = sqlite3.Row + conn.execute("PRAGMA query_only=ON") + conn.execute("PRAGMA temp_store=MEMORY") + return conn + + def do_GET(self) -> None: + parsed = urlparse(self.path) + if parsed.path == "/": + self.send_json( + { + "endpoints": { + "GET /health": "server and DB status", + "GET /tables": "table row counts", + "POST /query": { + "sql": "SELECT ...", + "params": [], + "limit": 500, + }, + } + } + ) + return + if parsed.path == "/health": + self.handle_health() + return + if parsed.path == "/tables": + self.handle_tables() + return + if parsed.path == "/query": + params = parse_qs(parsed.query) + sql = params.get("sql", [""])[0] + limit = params.get("limit", [default_limit])[0] + self.handle_query(sql, [], limit) + return + self.send_json({"error": "not found"}, HTTPStatus.NOT_FOUND) + + def do_POST(self) -> None: + parsed = urlparse(self.path) + if parsed.path != "/query": + self.send_json({"error": "not found"}, HTTPStatus.NOT_FOUND) + return + try: + payload = self.read_json() + except Exception as exc: + self.send_json({"error": repr(exc)}, HTTPStatus.BAD_REQUEST) + return + self.handle_query( + str(payload.get("sql", "")), + payload.get("params", []), + payload.get("limit", default_limit), + ) + + def handle_query(self, sql: str, params: List[Any], limit: Any) -> None: + try: + if not is_read_only_sql(sql): + self.send_json( + { + "error": "only single read-only SELECT/WITH/PRAGMA statements are allowed" + }, + HTTPStatus.BAD_REQUEST, + ) + return + limit = max(1, min(int(limit), max_limit)) + start = time.perf_counter() + conn = self.connect() + try: + # The server is a read-only SQL endpoint. sql is gated by + # is_read_only_sql and the connection is opened with mode=ro. + # codeql[py/sql-injection] + rows = conn.execute(sql, params).fetchmany(limit + 1) + elapsed_ms = (time.perf_counter() - start) * 1000 + returned = rows[:limit] + self.send_json( + { + "elapsed_ms": round(elapsed_ms, 3), + "limit": limit, + "truncated": len(rows) > limit, + "rows": [dict(row) for row in returned], + } + ) + finally: + conn.close() + except Exception as exc: + self.send_json({"error": repr(exc)}, HTTPStatus.BAD_REQUEST) + + def handle_health(self) -> None: + self.send_json( + { + "ok": db_path.exists(), + "db_path": str(db_path), + "db_size_mb": ( + round(db_path.stat().st_size / 1048576, 2) + if db_path.exists() + else None + ), + } + ) + + def handle_tables(self) -> None: + conn = self.connect() + try: + tables = [row["name"] for row in conn.execute(""" + SELECT name + FROM sqlite_master + WHERE type IN ('table', 'view') + AND name NOT LIKE 'sqlite_%' + AND name NOT LIKE '%_data' + AND name NOT LIKE '%_idx' + AND name NOT LIKE '%_docsize' + AND name NOT LIKE '%_config' + ORDER BY name + """)] + counts = {} + for table in tables: + try: + counts[table] = conn.execute( + 'SELECT COUNT(*) AS n FROM "%s"' % table + ).fetchone()["n"] + except sqlite3.DatabaseError as exc: + counts[table] = repr(exc) + self.send_json({"tables": counts}) + finally: + conn.close() + + return Handler + + +def serve( + db_path: Path, + host: str = "127.0.0.1", + port: int = 8765, + default_limit: int = 500, + max_limit: int = 5000, +) -> None: + handler = make_handler(db_path, default_limit, max_limit) + server = ThreadingHTTPServer((host, port), handler) + print("serving db=%s on http://%s:%s" % (db_path, host, port), flush=True) + server.serve_forever() diff --git a/TraceLens/TraceIndex/sqlite_store.py b/TraceLens/TraceIndex/sqlite_store.py new file mode 100644 index 000000000..e06ec0245 --- /dev/null +++ b/TraceLens/TraceIndex/sqlite_store.py @@ -0,0 +1,834 @@ +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""SQLite TraceIndex backend.""" + +import json +import sqlite3 +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional, Sequence + +from TraceLens.TraceIndex.models import SearchHit, TraceRecord, TraceReport +from TraceLens.TraceIndex.store import TraceIndexStore +from TraceLens.TraceIndex.utils import ( + as_bool_int, + as_duration_us, + as_float, + as_int, + as_optional_bool_int, + as_text, + first_value, + kernel_flags, + parse_repr, + search_text, + to_json, + utc_now, +) + + +def is_read_only_sql(sql: str) -> bool: + stripped = sql.strip().lower() + if not stripped: + return False + if ";" in stripped.rstrip(";"): + return False + return stripped.startswith(("select", "with", "pragma")) + + +class SQLiteTraceIndexStore(TraceIndexStore): + def __init__(self, db_path: Path): + self.db_path = db_path + self.conn = self._connect(db_path) + + def _connect(self, db_path: Path) -> sqlite3.Connection: + db_path.parent.mkdir(parents=True, exist_ok=True) + conn = sqlite3.connect(str(db_path), timeout=60) + conn.row_factory = sqlite3.Row + conn.execute("PRAGMA busy_timeout=60000") + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("PRAGMA synchronous=NORMAL") + conn.execute("PRAGMA foreign_keys=ON") + return conn + + def init_schema(self) -> None: + self.conn.executescript(""" + CREATE TABLE IF NOT EXISTS traces ( + id INTEGER PRIMARY KEY, + root TEXT, + path TEXT NOT NULL UNIQUE, + rel_path TEXT, + name TEXT, + size_bytes INTEGER, + md5 TEXT, + format TEXT, + rank INTEGER, + top_dir TEXT, + parent_rel TEXT, + should_enrich INTEGER NOT NULL DEFAULT 1, + skip_reason TEXT, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL + ); + + CREATE INDEX IF NOT EXISTS idx_trace_index_traces_md5 ON traces(md5); + CREATE INDEX IF NOT EXISTS idx_trace_index_traces_top_dir ON traces(top_dir); + CREATE INDEX IF NOT EXISTS idx_trace_index_traces_format ON traces(format); + CREATE INDEX IF NOT EXISTS idx_trace_index_traces_should_enrich ON traces(should_enrich); + + CREATE TABLE IF NOT EXISTS report_imports ( + id INTEGER PRIMARY KEY, + trace_id INTEGER NOT NULL REFERENCES traces(id) ON DELETE CASCADE, + report_dir TEXT NOT NULL, + imported_at TEXT NOT NULL, + sheets_json TEXT NOT NULL + ); + + CREATE TABLE IF NOT EXISTS unified_perf_rows ( + id INTEGER PRIMARY KEY, + trace_id INTEGER NOT NULL REFERENCES traces(id) ON DELETE CASCADE, + source_row INTEGER NOT NULL, + name TEXT, + op_category TEXT, + operation_count INTEGER, + kernel_time_sum_us REAL, + kernel_time_mean_us REAL, + kernel_time_median_us REAL, + kernel_time_std_us REAL, + kernel_time_min_us REAL, + kernel_time_max_us REAL, + op_duration_us REAL, + tflops_mean REAL, + tflops_median REAL, + tbs_mean REAL, + tbs_median REAL, + gflops REAL, + data_moved_mb REAL, + flops_per_byte REAL, + compute_spec TEXT, + has_perf_model INTEGER, + overlap_pct REAL, + perf_params_json TEXT, + kernel_details_json TEXT, + raw_row_json TEXT + ); + + CREATE INDEX IF NOT EXISTS idx_trace_index_unified_trace ON unified_perf_rows(trace_id); + CREATE INDEX IF NOT EXISTS idx_trace_index_unified_category ON unified_perf_rows(op_category); + CREATE INDEX IF NOT EXISTS idx_trace_index_unified_name ON unified_perf_rows(name); + + CREATE TABLE IF NOT EXISTS op_kernels ( + id INTEGER PRIMARY KEY, + trace_id INTEGER NOT NULL REFERENCES traces(id) ON DELETE CASCADE, + unified_row_id INTEGER REFERENCES unified_perf_rows(id) ON DELETE CASCADE, + kernel_name TEXT NOT NULL, + parent_op_name TEXT, + op_category TEXT, + stream INTEGER, + count INTEGER, + total_duration_us REAL, + mean_duration_us REAL, + median_duration_us REAL, + min_duration_us REAL, + max_duration_us REAL, + library TEXT, + is_tensile INTEGER NOT NULL DEFAULT 0, + is_transpose INTEGER NOT NULL DEFAULT 0, + is_layout_conversion INTEGER NOT NULL DEFAULT 0, + details_json TEXT + ); + + CREATE INDEX IF NOT EXISTS idx_trace_index_op_kernels_trace ON op_kernels(trace_id); + CREATE INDEX IF NOT EXISTS idx_trace_index_op_kernels_unified ON op_kernels(unified_row_id); + CREATE INDEX IF NOT EXISTS idx_trace_index_op_kernels_name ON op_kernels(kernel_name); + CREATE INDEX IF NOT EXISTS idx_trace_index_op_kernels_tensile ON op_kernels(is_tensile); + + CREATE TABLE IF NOT EXISTS gemm_perf ( + unified_row_id INTEGER PRIMARY KEY REFERENCES unified_perf_rows(id) ON DELETE CASCADE, + trace_id INTEGER NOT NULL REFERENCES traces(id) ON DELETE CASCADE, + m INTEGER, + n INTEGER, + k INTEGER, + batch INTEGER, + dtype TEXT, + transpose TEXT, + tflops_mean REAL, + tflops_median REAL + ); + + CREATE INDEX IF NOT EXISTS idx_trace_index_gemm_trace ON gemm_perf(trace_id); + CREATE INDEX IF NOT EXISTS idx_trace_index_gemm_tflops ON gemm_perf(tflops_mean); + + CREATE TABLE IF NOT EXISTS sdpa_perf ( + unified_row_id INTEGER PRIMARY KEY REFERENCES unified_perf_rows(id) ON DELETE CASCADE, + trace_id INTEGER NOT NULL REFERENCES traces(id) ON DELETE CASCADE, + batch INTEGER, + heads INTEGER, + seq_q INTEGER, + seq_kv INTEGER, + head_dim INTEGER, + dtype TEXT, + causal INTEGER, + tflops_mean REAL, + tflops_median REAL + ); + + CREATE INDEX IF NOT EXISTS idx_trace_index_sdpa_trace ON sdpa_perf(trace_id); + + CREATE TABLE IF NOT EXISTS conv_perf ( + unified_row_id INTEGER PRIMARY KEY REFERENCES unified_perf_rows(id) ON DELETE CASCADE, + trace_id INTEGER NOT NULL REFERENCES traces(id) ON DELETE CASCADE, + conv_nd TEXT, + input_shape_json TEXT, + filter_shape_json TEXT, + input_channels INTEGER, + output_channels INTEGER, + groups INTEGER, + kernel_h INTEGER, + kernel_w INTEGER, + is_depthwise INTEGER, + is_transposed_conv INTEGER + ); + + CREATE INDEX IF NOT EXISTS idx_trace_index_conv_trace ON conv_perf(trace_id); + CREATE INDEX IF NOT EXISTS idx_trace_index_conv_depthwise ON conv_perf(is_depthwise); + + CREATE TABLE IF NOT EXISTS op_category_rows ( + id INTEGER PRIMARY KEY, + trace_id INTEGER NOT NULL REFERENCES traces(id) ON DELETE CASCADE, + category TEXT NOT NULL, + operation_count INTEGER, + kernel_time_sum_us REAL, + percent REAL, + raw_row_json TEXT + ); + + CREATE INDEX IF NOT EXISTS idx_trace_index_category_trace ON op_category_rows(trace_id); + + CREATE TABLE IF NOT EXISTS gpu_timeline_rows ( + id INTEGER PRIMARY KEY, + trace_id INTEGER NOT NULL REFERENCES traces(id) ON DELETE CASCADE, + type TEXT NOT NULL, + time_ms REAL, + percent REAL, + raw_row_json TEXT + ); + + CREATE INDEX IF NOT EXISTS idx_trace_index_timeline_trace ON gpu_timeline_rows(trace_id); + CREATE INDEX IF NOT EXISTS idx_trace_index_timeline_type ON gpu_timeline_rows(type); + + CREATE TABLE IF NOT EXISTS trace_summary ( + trace_id INTEGER PRIMARY KEY REFERENCES traces(id) ON DELETE CASCADE, + total_duration_us REAL, + top_categories_json TEXT, + max_gemm_tflops REAL, + max_sdpa_tflops REAL, + imported_at TEXT NOT NULL + ); + + CREATE VIRTUAL TABLE IF NOT EXISTS trace_search_FTS5 USING fts5( + trace_id UNINDEXED, + kind UNINDEXED, + text, + tokenize='unicode61' + ); + """) + self.conn.commit() + + def upsert_trace(self, trace: TraceRecord) -> int: + now = utc_now() + self.conn.execute( + """ + INSERT INTO traces( + root, path, rel_path, name, size_bytes, md5, format, rank, top_dir, + parent_rel, should_enrich, skip_reason, created_at, updated_at + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(path) DO UPDATE SET + root = excluded.root, + rel_path = excluded.rel_path, + name = excluded.name, + size_bytes = excluded.size_bytes, + md5 = COALESCE(excluded.md5, traces.md5), + format = excluded.format, + rank = excluded.rank, + top_dir = excluded.top_dir, + parent_rel = excluded.parent_rel, + should_enrich = excluded.should_enrich, + skip_reason = excluded.skip_reason, + updated_at = excluded.updated_at + """, + ( + trace.root, + trace.path, + trace.rel_path, + trace.name, + trace.size_bytes, + trace.md5, + trace.format, + trace.rank, + trace.top_dir, + trace.parent_rel, + int(trace.should_enrich), + trace.skip_reason, + now, + now, + ), + ) + row = self.conn.execute( + "SELECT id FROM traces WHERE path = ?", (trace.path,) + ).fetchone() + self.conn.commit() + return int(row["id"]) + + def import_report(self, trace_id: int, report: TraceReport) -> None: + self._clear_trace_payload(trace_id) + unified_summary = self._import_unified_rows( + trace_id, report.sheets.get("unified_perf_summary", []) + ) + top_categories_json = self._import_category_rows( + trace_id, report.sheets.get("ops_summary_by_category", []) + ) + total_duration_us = self._import_gpu_timeline_rows( + trace_id, report.sheets.get("gpu_timeline", []) + ) + + self.conn.execute( + """ + INSERT INTO trace_summary( + trace_id, total_duration_us, top_categories_json, max_gemm_tflops, + max_sdpa_tflops, imported_at + ) + VALUES (?, ?, ?, ?, ?, ?) + """, + ( + trace_id, + total_duration_us, + top_categories_json, + unified_summary["max_gemm_tflops"], + unified_summary["max_sdpa_tflops"], + utc_now(), + ), + ) + self.conn.execute( + """ + INSERT INTO report_imports(trace_id, report_dir, imported_at, sheets_json) + VALUES (?, ?, ?, ?) + """, + ( + trace_id, + report.report_dir, + utc_now(), + json.dumps( + [name for name, rows in report.sheets.items() if rows], + sort_keys=True, + ), + ), + ) + self._insert_search(trace_id, "trace", [report.report_dir]) + self.conn.commit() + + def search(self, terms: str, limit: int = 50) -> List[SearchHit]: + rows = self.conn.execute( + """ + SELECT t.id AS trace_id, t.rel_path, s.kind, + snippet(trace_search_FTS5, 2, '[', ']', '...', 12) AS hit + FROM trace_search_FTS5 s + JOIN traces t ON t.id = s.trace_id + WHERE trace_search_FTS5 MATCH ? + LIMIT ? + """, + (terms, limit), + ).fetchall() + return [ + SearchHit( + trace_id=int(row["trace_id"]), + rel_path=row["rel_path"], + kind=row["kind"], + hit=row["hit"], + ) + for row in rows + ] + + def execute_read_query( + self, + sql: str, + params: Optional[Sequence[Any]] = None, + limit: int = 500, + ) -> List[Dict[str, Any]]: + if not is_read_only_sql(sql): + raise ValueError( + "only a single read-only SELECT/WITH/PRAGMA statement is allowed" + ) + self.conn.execute("PRAGMA query_only=ON") + # Callers pass a single SELECT/WITH/PRAGMA; writes are rejected above. + # codeql[py/sql-injection] + rows = self.conn.execute(sql, params or ()).fetchmany(limit) + return [dict(row) for row in rows] + + def close(self) -> None: + self.conn.close() + + def _clear_trace_payload(self, trace_id: int) -> None: + for table in ( + "report_imports", + "gemm_perf", + "sdpa_perf", + "conv_perf", + "op_kernels", + "op_category_rows", + "gpu_timeline_rows", + "trace_summary", + "unified_perf_rows", + ): + self.conn.execute("DELETE FROM %s WHERE trace_id = ?" % table, (trace_id,)) + self.conn.execute( + "DELETE FROM trace_search_FTS5 WHERE trace_id = ?", (trace_id,) + ) + + def _insert_search(self, trace_id: int, kind: str, parts: Iterable[Any]) -> None: + text = search_text(*parts) + if text: + self.conn.execute( + "INSERT INTO trace_search_FTS5(trace_id, kind, text) VALUES (?, ?, ?)", + (trace_id, kind, text), + ) + + def _import_unified_rows( + self, + trace_id: int, + rows: Sequence[Dict[str, str]], + ) -> Dict[str, Optional[float]]: + max_gemm_tflops = None + max_sdpa_tflops = None + for source_row, row in enumerate(rows): + name = as_text(first_value(row, ["name", "Name", "op_name"])) + op_category = as_text( + first_value( + row, ["op category", "op_category", "category", "Categories"] + ) + ) + tflops_mean = as_float( + first_value(row, ["TFLOPS/s_mean", "tflops_mean", "TFLOPS_mean"]) + ) + tflops_median = as_float( + first_value(row, ["TFLOPS/s_median", "tflops_median", "TFLOPS_median"]) + ) + tflops_for_summary = ( + tflops_mean if tflops_mean is not None else tflops_median + ) + if ( + op_category + and "gemm" in op_category.lower() + and tflops_for_summary is not None + ): + max_gemm_tflops = max( + max_gemm_tflops or tflops_for_summary, tflops_for_summary + ) + if ( + op_category + and "sdpa" in op_category.lower() + and tflops_for_summary is not None + ): + max_sdpa_tflops = max( + max_sdpa_tflops or tflops_for_summary, tflops_for_summary + ) + params = parse_repr(first_value(row, ["perf_params", "Perf Params"])) + kernel_details = parse_repr( + first_value(row, ["kernel_details_summary", "trunc_kernel_details"]) + ) + cursor = self.conn.execute( + """ + INSERT INTO unified_perf_rows( + trace_id, source_row, name, op_category, operation_count, + kernel_time_sum_us, kernel_time_mean_us, kernel_time_median_us, + kernel_time_std_us, kernel_time_min_us, kernel_time_max_us, + op_duration_us, tflops_mean, tflops_median, tbs_mean, tbs_median, + gflops, data_moved_mb, flops_per_byte, compute_spec, + has_perf_model, overlap_pct, perf_params_json, kernel_details_json, + raw_row_json + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + trace_id, + source_row, + name, + op_category, + as_int(first_value(row, ["operation_count", "Count", "count"])), + as_duration_us( + row, + [ + "Kernel Time (us)_sum", + "Kernel Time (µs)_sum", + "total_direct_kernel_time_sum", + "total_subtree_kernel_time_sum", + ], + [ + "total_direct_kernel_time_ms", + "total_subtree_kernel_time_ms", + ], + ), + as_float( + first_value( + row, + [ + "Kernel Time (us)_mean", + "Kernel Time (µs)_mean", + "total_direct_kernel_time_mean", + "total_subtree_kernel_time_mean", + ], + ) + ), + as_float( + first_value( + row, + [ + "Kernel Time (us)_median", + "Kernel Time (µs)_median", + "total_direct_kernel_time_median", + "total_subtree_kernel_time_median", + ], + ) + ), + as_float( + first_value( + row, ["Kernel Time (us)_std", "Kernel Time (µs)_std"] + ) + ), + as_float( + first_value( + row, ["Kernel Time (us)_min", "Kernel Time (µs)_min"] + ) + ), + as_float( + first_value( + row, ["Kernel Time (us)_max", "Kernel Time (µs)_max"] + ) + ), + as_float( + first_value( + row, + [ + "op_duration_us", + "CPU duration (us)", + "CPU duration (µs)", + ], + ) + ), + tflops_mean, + tflops_median, + as_float(first_value(row, ["TB/s_mean", "tbs_mean"])), + as_float(first_value(row, ["TB/s_median", "tbs_median"])), + as_float(first_value(row, ["GFLOPS", "gflops"])), + as_float(first_value(row, ["Data Moved (MB)", "data_moved_mb"])), + as_float(first_value(row, ["FLOPs/Byte", "flops_per_byte"])), + as_text(first_value(row, ["Compute Spec", "compute_spec"])), + as_bool_int(first_value(row, ["has_perf_model", "Has Perf Model"])), + as_float(first_value(row, ["overlap_pct", "Overlap (%)"])), + to_json(params), + to_json(kernel_details), + to_json(dict(row)), + ), + ) + unified_row_id = int(cursor.lastrowid) + self._import_kernels_from_details( + trace_id, unified_row_id, name, op_category, kernel_details + ) + self._maybe_insert_gemm( + trace_id, unified_row_id, params, tflops_mean, tflops_median + ) + self._maybe_insert_sdpa( + trace_id, unified_row_id, params, tflops_mean, tflops_median + ) + self._maybe_insert_conv(trace_id, unified_row_id, params) + self._insert_search( + trace_id, + "op", + [ + name, + op_category, + first_value( + row, ["kernel_details_summary", "trunc_kernel_details"] + ), + ], + ) + return {"max_gemm_tflops": max_gemm_tflops, "max_sdpa_tflops": max_sdpa_tflops} + + def _import_kernels_from_details( + self, + trace_id: int, + unified_row_id: int, + parent_op_name: Optional[str], + op_category: Optional[str], + kernel_details: Any, + ) -> None: + if not isinstance(kernel_details, list): + return + for detail in kernel_details: + if not isinstance(detail, dict): + continue + kernel_name = as_text(detail.get("name") or detail.get("Kernel name")) + if not kernel_name: + continue + library, is_tensile, is_transpose, is_layout = kernel_flags( + kernel_name, parent_op_name or "" + ) + self.conn.execute( + """ + INSERT INTO op_kernels( + trace_id, unified_row_id, kernel_name, parent_op_name, op_category, + stream, count, total_duration_us, mean_duration_us, + median_duration_us, min_duration_us, max_duration_us, library, + is_tensile, is_transpose, is_layout_conversion, details_json + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + trace_id, + unified_row_id, + kernel_name, + parent_op_name, + op_category, + as_int(detail.get("stream")), + as_int(detail.get("count")), + as_float(detail.get("total_duration_us")), + as_float(detail.get("mean_duration_us")), + as_float(detail.get("median_duration_us")), + as_float(detail.get("min_duration_us")), + as_float(detail.get("max_duration_us")), + library, + is_tensile, + is_transpose, + is_layout, + to_json(detail), + ), + ) + self._insert_search( + trace_id, + "kernel", + [kernel_name, library, parent_op_name, op_category], + ) + + def _maybe_insert_gemm( + self, + trace_id: int, + unified_row_id: int, + params: Any, + tflops_mean: Optional[float], + tflops_median: Optional[float], + ) -> None: + if not isinstance(params, dict): + return + if not {"M", "N", "K"}.intersection(params): + return + self.conn.execute( + """ + INSERT OR REPLACE INTO gemm_perf( + unified_row_id, trace_id, m, n, k, batch, dtype, transpose, + tflops_mean, tflops_median + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + unified_row_id, + trace_id, + as_int(params.get("M")), + as_int(params.get("N")), + as_int(params.get("K")), + as_int(params.get("B")), + ( + str(params.get("dtype_A_B")) + if params.get("dtype_A_B") is not None + else None + ), + ( + str(params.get("transpose")) + if params.get("transpose") is not None + else None + ), + tflops_mean, + tflops_median, + ), + ) + + def _maybe_insert_sdpa( + self, + trace_id: int, + unified_row_id: int, + params: Any, + tflops_mean: Optional[float], + tflops_median: Optional[float], + ) -> None: + if not isinstance(params, dict): + return + if not {"N_Q", "N_KV", "d_h_qk", "d_h_v"}.intersection(params): + return + self.conn.execute( + """ + INSERT OR REPLACE INTO sdpa_perf( + unified_row_id, trace_id, batch, heads, seq_q, seq_kv, + head_dim, dtype, causal, tflops_mean, tflops_median + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + unified_row_id, + trace_id, + as_int(params.get("B")), + as_int(params.get("H_Q")), + as_int(params.get("N_Q")), + as_int(params.get("N_KV")), + as_int(params.get("d_h_qk") or params.get("d_h_v")), + ( + str(params.get("dtype_A_B")) + if params.get("dtype_A_B") is not None + else None + ), + as_optional_bool_int(params.get("causal")), + tflops_mean, + tflops_median, + ), + ) + + def _maybe_insert_conv( + self, + trace_id: int, + unified_row_id: int, + params: Any, + ) -> None: + if not isinstance(params, dict): + return + if "convNd" not in params and "filter_shape" not in params: + return + input_shape = params.get("input_shape") + filter_shape = params.get("filter_shape") + groups = as_int(params.get("groups")) or 1 + input_channels = None + output_channels = None + kernel_h = None + kernel_w = None + if isinstance(input_shape, (list, tuple)) and len(input_shape) >= 2: + input_channels = as_int(input_shape[1]) + if isinstance(filter_shape, (list, tuple)) and len(filter_shape) >= 4: + output_channels = as_int(filter_shape[0]) + kernel_h = as_int(filter_shape[-2]) + kernel_w = as_int(filter_shape[-1]) + filter_c_per_group = None + if isinstance(filter_shape, (list, tuple)) and len(filter_shape) >= 2: + filter_c_per_group = as_int(filter_shape[1]) + is_depthwise = int( + bool(input_channels) + and groups == input_channels + and filter_c_per_group == 1 + and bool(output_channels) + and output_channels % input_channels == 0 + ) + self.conn.execute( + """ + INSERT OR REPLACE INTO conv_perf( + unified_row_id, trace_id, conv_nd, input_shape_json, filter_shape_json, + input_channels, output_channels, groups, kernel_h, kernel_w, + is_depthwise, is_transposed_conv + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + unified_row_id, + trace_id, + params.get("convNd"), + to_json(input_shape), + to_json(filter_shape), + input_channels, + output_channels, + groups, + kernel_h, + kernel_w, + is_depthwise, + as_optional_bool_int(params.get("transposed_conv")), + ), + ) + + def _import_category_rows( + self, + trace_id: int, + rows: Sequence[Dict[str, str]], + ) -> str: + top_categories = [] + for row in rows: + category = as_text( + first_value(row, ["op category", "category", "Categories", "name"]) + ) + if not category: + continue + kernel_time = as_duration_us( + row, + [ + "Kernel Time (us)_sum", + "Kernel Time (µs)_sum", + "total_direct_kernel_time_sum", + "total_subtree_kernel_time_sum", + ], + [ + "total_direct_kernel_time_ms", + "total_subtree_kernel_time_ms", + ], + ) + self.conn.execute( + """ + INSERT INTO op_category_rows( + trace_id, category, operation_count, kernel_time_sum_us, percent, raw_row_json + ) + VALUES (?, ?, ?, ?, ?, ?) + """, + ( + trace_id, + category, + as_int(first_value(row, ["operation_count", "Count", "count"])), + kernel_time, + as_float( + first_value( + row, + ["Percentage (%)", "percent", "Percent of total time (%)"], + ) + ), + json.dumps(row, sort_keys=True), + ), + ) + top_categories.append( + {"category": category, "kernel_time_sum_us": kernel_time or 0.0} + ) + self._insert_search(trace_id, "category", [category]) + top_categories.sort(key=lambda item: item["kernel_time_sum_us"], reverse=True) + return json.dumps(top_categories[:5], sort_keys=True) + + def _import_gpu_timeline_rows( + self, + trace_id: int, + rows: Sequence[Dict[str, str]], + ) -> Optional[float]: + total_duration_us = None + for row in rows: + metric_type = as_text(first_value(row, ["type", "metric"])) + if not metric_type: + continue + time_ms = as_float(first_value(row, ["time ms", "time_ms"])) + if metric_type == "total_time" and time_ms is not None: + total_duration_us = time_ms * 1000.0 + self.conn.execute( + """ + INSERT INTO gpu_timeline_rows(trace_id, type, time_ms, percent, raw_row_json) + VALUES (?, ?, ?, ?, ?) + """, + ( + trace_id, + metric_type, + time_ms, + as_float(first_value(row, ["percent", "Percentage (%)"])), + json.dumps(row, sort_keys=True), + ), + ) + self._insert_search(trace_id, "timeline", [metric_type]) + return total_duration_us diff --git a/TraceLens/TraceIndex/store.py b/TraceLens/TraceIndex/store.py new file mode 100644 index 000000000..70f5f13c3 --- /dev/null +++ b/TraceLens/TraceIndex/store.py @@ -0,0 +1,48 @@ +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Storage interface for TraceIndex backends.""" + +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Optional, Sequence + +from TraceLens.TraceIndex.models import SearchHit, TraceRecord, TraceReport + + +class TraceIndexStore(ABC): + """Persistence boundary for TraceIndex. + + New backends should implement this interface without changing importer + or CLI workflow code. + """ + + @abstractmethod + def init_schema(self) -> None: + pass + + @abstractmethod + def upsert_trace(self, trace: TraceRecord) -> int: + pass + + @abstractmethod + def import_report(self, trace_id: int, report: TraceReport) -> None: + pass + + @abstractmethod + def search(self, terms: str, limit: int = 50) -> List[SearchHit]: + pass + + @abstractmethod + def execute_read_query( + self, + sql: str, + params: Optional[Sequence[Any]] = None, + limit: int = 500, + ) -> List[Dict[str, Any]]: + pass + + def close(self) -> None: + pass diff --git a/TraceLens/TraceIndex/utils.py b/TraceLens/TraceIndex/utils.py new file mode 100644 index 000000000..15e5898fc --- /dev/null +++ b/TraceLens/TraceIndex/utils.py @@ -0,0 +1,224 @@ +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Shared TraceIndex helpers that are independent of a storage backend.""" + +import ast +import csv +import json +import math +import re +import sys +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, List, Optional, Sequence, Tuple + +from TraceLens.PerfModel import kernel_name_parser as knp +from TraceLens.PerfModel.kernel_library import classify_kernel_library + + +def utc_now() -> str: + return datetime.now(timezone.utc).isoformat(timespec="seconds") + + +def normalize_path(path: Path) -> str: + return str(path).replace("\\", "/") + + +def rel_to(path: Path, root: Path) -> str: + try: + return normalize_path(path.relative_to(root)) + except ValueError: + return normalize_path(path) + + +def _set_csv_field_size_limit() -> None: + limit = sys.maxsize + while True: + try: + csv.field_size_limit(limit) + return + except OverflowError: + limit = int(limit / 10) + + +_set_csv_field_size_limit() + + +def read_csv_rows(path: Path) -> List[Dict[str, str]]: + if not path.exists() or path.stat().st_size == 0: + return [] + with path.open("r", encoding="utf-8-sig", newline="") as f: + return list(csv.DictReader(f)) + + +def first_value(row: Dict[str, Any], names: Sequence[str], default: Any = None) -> Any: + lower_map = {key.lower(): key for key in row.keys()} + for name in names: + key = lower_map.get(name.lower()) + if key is None: + continue + value = row.get(key) + if value not in (None, "", "nan", "NaN"): + return value + return default + + +def as_text(value: Any) -> Optional[str]: + if value in (None, "", "nan", "NaN"): + return None + return str(value) + + +def as_float(value: Any) -> Optional[float]: + if value in (None, "", "nan", "NaN"): + return None + try: + return float(value) + except (TypeError, ValueError): + return None + + +def as_int(value: Any) -> Optional[int]: + number = as_float(value) + return int(number) if number is not None else None + + +def as_duration_us( + row: Dict[str, Any], + us_names: Sequence[str], + ms_names: Sequence[str] = (), +) -> Optional[float]: + us_value = as_float(first_value(row, us_names)) + if us_value is not None: + return us_value + ms_value = as_float(first_value(row, ms_names)) + if ms_value is not None: + return ms_value * 1000.0 + return None + + +def as_bool_int(value: Any) -> int: + optional = as_optional_bool_int(value) + return 0 if optional is None else optional + + +def as_optional_bool_int(value: Any) -> Optional[int]: + if isinstance(value, bool): + return int(value) + if value in (None, "", "nan", "NaN"): + return None + text = str(value).strip().lower() + if text in {"1", "true", "yes", "y"}: + return 1 + if text in {"0", "false", "no", "n"}: + return 0 + return None + + +def search_text(*parts: Any) -> str: + return " ".join(str(part) for part in parts if part not in (None, "", "nan", "NaN")) + + +NP_SCALAR_RE = re.compile(r"\b(?:np|numpy)\.(?:float|int)(?:16|32|64)?\(([^()]+)\)") + + +def clean_python_repr(text: str) -> str: + cleaned = NP_SCALAR_RE.sub(r"\1", text) + return cleaned.replace("nan", "None") + + +def parse_repr(text: Any) -> Any: + if not text: + return None + if not isinstance(text, str): + return text + try: + return ast.literal_eval(clean_python_repr(text)) + except (SyntaxError, ValueError, MemoryError): + return None + + +def json_safe(value: Any) -> Any: + if value is None: + return None + if isinstance(value, float) and math.isnan(value): + return None + if isinstance(value, (str, int, float, bool)): + return value + if isinstance(value, (list, tuple)): + return [json_safe(item) for item in value] + if isinstance(value, dict): + return {str(key): json_safe(item) for key, item in value.items()} + item = getattr(value, "item", None) + if callable(item): + try: + return json_safe(item()) + except Exception: + pass + return str(value) + + +def to_json(value: Any) -> Optional[str]: + if value is None: + return None + return json.dumps(json_safe(value), sort_keys=True) + + +def kernel_flags(name: str, op_name: str = "") -> Tuple[Optional[str], int, int, int]: + """Derive catalog flags for a GPU kernel using shared library classification.""" + low = name.lower() + library = classify_kernel_library(op_name, name) + + is_tensile = int( + knp.is_rocm_gemm(name) or library == "Tensile" or "tensile" in low + ) + is_transpose = int("transpose" in low or "permute" in low) + parsed = knp.gemm_name_parser(name) + if parsed: + transpose = parsed.get("transpose") + if transpose and any(transpose): + is_transpose = 1 + + is_layout = int( + is_transpose + or "contiguous" in low + or "copy" in low + or "cast" in low + or "convert" in low + ) + return library, is_tensile, is_transpose, is_layout + + +def read_traces_file(path: Path) -> List[Path]: + """Read one trace path per line. Blank lines and ``#`` comments are ignored.""" + traces: List[Path] = [] + for raw_line in path.read_text(encoding="utf-8").splitlines(): + line = raw_line.strip() + if not line or line.startswith("#"): + continue + traces.append(Path(line).expanduser()) + return traces + + +def collect_trace_paths( + traces_file: Optional[Path] = None, + trace_paths: Optional[Sequence[Path]] = None, +) -> List[Path]: + paths: List[Path] = [] + if traces_file is not None: + paths.extend(read_traces_file(traces_file)) + if trace_paths: + paths.extend(trace_paths) + unique: List[Path] = [] + seen = set() + for path in paths: + key = normalize_path(path) + if key in seen: + continue + seen.add(key) + unique.append(path) + return unique diff --git a/docs/how-to/generate-perf-report-pytorch-inference.md b/docs/how-to/generate-perf-report-pytorch-inference.md index a918b125a..e00d29c61 100644 --- a/docs/how-to/generate-perf-report-pytorch-inference.md +++ b/docs/how-to/generate-perf-report-pytorch-inference.md @@ -481,6 +481,7 @@ Run the tool with `--help` for the complete, version-specific argument list. - [Inference performance analysis in TraceLens](../conceptual/inference-analysis.md) - [Generate a PyTorch performance report](./generate-perf-report-pytorch.md) +- [Index a corpus of traces](./trace-index.md) - [Compare two traces](./compare-traces.md) - [Agentically orchestrate and generate optimization recommendations](./agent.md) - [API reference](../reference/api-reference.md) diff --git a/docs/how-to/generate-perf-report-pytorch.md b/docs/how-to/generate-perf-report-pytorch.md index c5888af32..674a930d7 100644 --- a/docs/how-to/generate-perf-report-pytorch.md +++ b/docs/how-to/generate-perf-report-pytorch.md @@ -210,6 +210,8 @@ The following table describes all optional arguments. [collective-communication report](./collective-report.md). - Isolate a single operation into a reproducer with [EventReplay](./event-replay.md). +- Catalog many reports for search with + [Index a corpus of traces](./trace-index.md). - Analyze [JAX](./generate-perf-report-jax.md) or [rocprof](./generate-perf-report-rocprof.md) traces. diff --git a/docs/how-to/generate-reports.md b/docs/how-to/generate-reports.md index 4bb918bb7..45d8034c8 100644 --- a/docs/how-to/generate-reports.md +++ b/docs/how-to/generate-reports.md @@ -11,7 +11,7 @@ See LICENSE for license information. :keywords: TraceLens, performance report, PyTorch profiler, JAX, rocprofv3, collective communication, ROCm, GPU trace analysis ``` -TraceLens generates structured performance reports from GPU trace files produced by PyTorch, JAX, and the AMD ROCm profiler. Choose the guide for your trace format: +TraceLens generates structured performance reports from GPU trace files produced by PyTorch, JAX, and the AMD ROCm profiler. Choose the topic for your trace format: - [Generate a PyTorch performance report](./generate-perf-report-pytorch.md) - [Generate a PyTorch inference performance report](./generate-perf-report-pytorch-inference.md) @@ -19,6 +19,10 @@ TraceLens generates structured performance reports from GPU trace files produced - [Generate a rocprof performance report](./generate-perf-report-rocprof.md) - [Generate a collective-communication report](./collective-report.md) +After you have CSV report directories, append them with +`TraceLens_trace_index append --trace-path … --report-dir …`. +See [Index a corpus of traces](./trace-index.md). + ## Related topics - [What is TraceLens?](../what-is-tracelens.md) diff --git a/docs/how-to/sdk-analysis.md b/docs/how-to/sdk-analysis.md index cc4d3c8eb..2ab757bb0 100644 --- a/docs/how-to/sdk-analysis.md +++ b/docs/how-to/sdk-analysis.md @@ -278,5 +278,6 @@ pretty-printer. - [The Trace2Tree data model](../conceptual/trace2tree.md) - [Model op performance without a trace](./perf-model-without-trace.md) - [Generate a PyTorch performance report](./generate-perf-report-pytorch.md) +- [Index a corpus of traces](./trace-index.md) - [Replay a single operation in TraceLens](./event-replay.md) - [API reference](../reference/api-reference.md) diff --git a/docs/how-to/trace-index.md b/docs/how-to/trace-index.md new file mode 100644 index 000000000..012c48579 --- /dev/null +++ b/docs/how-to/trace-index.md @@ -0,0 +1,229 @@ + + + +# Index a corpus of traces in TraceLens +```{meta} +:description: Learn how to catalog profiler traces and TraceLens CSV reports into a searchable index. SQLite is the first backend; the table schema is the shared query surface. +:keywords: TraceLens, TraceIndex, corpus search, catalog schema, SQLite, unified_perf_summary, kernel summary, full-text search, performance report +``` + +This topic shows how to build a queryable catalog of profiler traces and +TraceLens CSV reports so you can search a corpus without reopening every raw +file. TraceIndex stores summaries and paths back to the source traces; it +doesn't replace the traces themselves. + +The tables below are the catalog schema — the shared query surface. Notebooks, +SQL, and later storage backends should use these table and column names so +catalogs stay interchangeable across users. SQLite is the first backend, not a +prototype: it ships with Python, writes one file, and needs no extra service. +Other backends can implement the same schema later. + +## Before you begin + +- TraceLens installed (see [Install TraceLens](../install/install.md)). +- Profiler traces, and optionally existing TraceLens CSV report directories + (for example from + [Generate a PyTorch performance report](./generate-perf-report-pytorch.md)). + +## Append a trace + +Append one trace to the catalog. Pass `--report-dir` when you already have a +CSV report. This is the usual path for inference, rocprof, and pftrace +reports you generated separately: + +```bash +TraceLens_trace_index --db trace_index.sqlite append \ + --trace-path /path/to/traces/rank0_trace.json.gz \ + --report-dir ./rank0_perf_report_csvs +``` + +If you omit `--report-dir`, TraceIndex generates a training PyTorch CSV report +and then imports it: + +```bash +TraceLens_trace_index --db trace_index.sqlite append \ + --trace-path /path/to/traces/rank0_trace.json.gz +``` + +## Build a catalog from a list of traces + +`--db` creates the SQLite file if it doesn't exist. `build` walks a list of +trace paths, generates a training PyTorch report for each, and appends it. +Use a text file (one path per line; `#` starts a comment) and/or repeat +`--trace-path`: + +```bash +TraceLens_trace_index --db trace_index.sqlite build \ + --traces-file traces.txt \ + --report-root ./trace_index_reports +``` + +A failed trace is recorded and the rest of the list still runs. For inference, +rocprof, or pftrace, generate the CSV reports first, then `append` each trace +with `--report-dir`. + +## Search and query + +Full-text search (FTS) over indexed ops, kernels, categories, and timeline +labels: + +```bash +TraceLens_trace_index --backend sqlite --db trace_index.sqlite search attention +TraceLens_trace_index --backend sqlite --db trace_index.sqlite search Cijk +``` + +Run a single read-only SQL statement: + +```bash +TraceLens_trace_index --backend sqlite --db trace_index.sqlite sqlite-sql \ + "SELECT op_category, COUNT(*) AS rows FROM unified_perf_rows GROUP BY op_category" +``` + +## What the catalog stores + +The following tables are the catalog schema. There are ten relational tables +plus a full-text search (FTS) virtual table. SQLite holds this comfortably for +typical TraceLens corpora (hundreds of traces, hundreds of thousands of kernel +rows). The practical limit is one writer at a time, not row count. + +The following diagram shows how those tables relate. Every fact table points at +`traces`. Kernel rows and GEMM / SDPA / convolution satellites also point at +the `unified_perf_rows` row they came from. + +```mermaid +erDiagram + traces ||--o{ report_imports : "trace_id" + traces ||--o{ unified_perf_rows : "trace_id" + traces ||--o{ op_kernels : "trace_id" + traces ||--o{ op_category_rows : "trace_id" + traces ||--o{ gpu_timeline_rows : "trace_id" + traces ||--o| trace_summary : "trace_id" + traces ||--o{ gemm_perf : "trace_id" + traces ||--o{ sdpa_perf : "trace_id" + traces ||--o{ conv_perf : "trace_id" + traces ||--o{ trace_search_FTS5 : "trace_id" + unified_perf_rows ||--o{ op_kernels : "unified_row_id" + unified_perf_rows ||--o| gemm_perf : "unified_row_id" + unified_perf_rows ||--o| sdpa_perf : "unified_row_id" + unified_perf_rows ||--o| conv_perf : "unified_row_id" +``` + +| Table | Contents | +|---|---| +| `traces` | One row per indexed trace | +| `report_imports` | Import history for TraceLens CSV report directories | +| `unified_perf_rows` | Rows from `unified_perf_summary.csv`. `perf_params_json` and `kernel_details_json` are parsed JSON, not the Python `repr` from the CSV | +| `op_kernels` | One row per kernel exploded from `kernel_details_summary` on a unified row, with `unified_row_id` and Tensile / layout flags | +| `gemm_perf` | GEMM shapes (`M` / `N` / `K` / batch / dtype / transpose) parsed from `perf_params` | +| `sdpa_perf` | Attention shapes (`B`, `H_Q`, `N_Q`, `N_KV`, head dim, causal) parsed from `perf_params` | +| `conv_perf` | Convolution shapes, groups, and depthwise / transposed flags parsed from `perf_params` | +| `op_category_rows` | Rows from `ops_summary_by_category.csv` | +| `gpu_timeline_rows` | Rows from `gpu_timeline.csv` | +| `trace_summary` | Per-trace summary metrics derived during import | +| `trace_search_FTS5` | Full-text search over traces, ops, kernels, categories, and timeline labels | + +Query GEMM / SDPA / convolution shapes from the satellite tables, or with +`json_extract` on the parsed JSON columns. For example +`SELECT m, n, k FROM gemm_perf` or +`SELECT json_extract(perf_params_json, '$.M') FROM unified_perf_rows`. + +## Example queries + +Because shapes are first-class columns in the satellite tables, questions that +would otherwise mean reopening every trace become a single SQL filter. Run these +with `sqlite-sql` or the HTTP server, and `JOIN traces` to get the file to open. + +### Do any traces have depthwise convolution? + +```sql +SELECT t.name, c.input_channels, c.output_channels, c.groups, + c.kernel_h, c.kernel_w, COUNT(*) AS rows +FROM conv_perf c +JOIN traces t ON t.id = c.trace_id +WHERE c.is_depthwise = 1 AND c.groups > 1 +GROUP BY t.id, c.input_channels, c.output_channels, c.groups, c.kernel_h, c.kernel_w +ORDER BY rows DESC; +``` + +| trace | Cin | Cout | groups | Kh | Kw | rows | +|---|---:|---:|---:|---:|---:|---:| +| `diffusion_model_trace.json.gz` | 3072 | 3072 | 3072 | 5 | 5 | 2 | +| `diffusion_model_trace.json.gz` | 8192 | 8192 | 8192 | 3 | 3 | 2 | +| `diffusion_model_trace.json.gz` | 1536 | 1536 | 1536 | 5 | 5 | 1 | +| `diffusion_model_trace.json.gz` | 4096 | 4096 | 4096 | 3 | 3 | 1 | + +Channels equal groups (true depthwise) at 3×3 and 5×5 with widths 1536 / 3072 / +4096 / 8192, all in a single diffusion capture. If you're looking for a +depthwise-conv workload, that's the file to open — found without reopening any +trace. + +### What are the longest attention sequences in the catalog? + +```sql +SELECT t.name, p.seq_q, p.seq_kv, p.heads, p.head_dim, p.dtype +FROM sdpa_perf p +JOIN traces t ON t.id = p.trace_id +ORDER BY p.seq_q DESC +LIMIT 8; +``` + +| trace | seq_q | seq_kv | heads | d | dtype | +|---|---:|---:|---:|---:|---| +| `video_traces_rank_5_step_3.json` | 118872 | 118809 | 3 | 128 | BF16 | +| `video_traces_rank_2_step_3.json` | 118872 | 118809 | 3 | 128 | BF16 | +| `video_traces_rank_7_step_3.json` | 118872 | 118809 | 3 | 128 | BF16 | +| `video_traces_rank_6_step_3.json` | 118872 | 118809 | 3 | 128 | BF16 | + +The longest attention here is a video/DiT-style shape: sequence ≈ 119k, 3 heads, +head dim 128, BF16 — not LLM decode. Because `seq_q` / `seq_kv` / `heads` / +`head_dim` are columns, "find long-context attention" is a range query. + +## Serve read-only SQL + +For notebook or browser workflows, serve the SQLite catalog over HTTP: + +```bash +TraceLens_trace_index --backend sqlite --db trace_index.sqlite serve --host 127.0.0.1 --port 8765 +``` + +The server exposes: + +- `GET /health` +- `GET /tables` +- `POST /query` with `{"sql": "SELECT ...", "params": [], "limit": 500}` + +```{note} +Only a single `SELECT`, `WITH`, or `PRAGMA` statement is accepted. The server +is read-only but isn't authenticated, so bind it to loopback unless you put it +behind your own access control. +``` + +## Python API + +```python +from pathlib import Path + +from TraceLens.TraceIndex import append_trace, build_traces, search_index + +db = Path("trace_index.sqlite") +append_trace( + db, + Path("rank0_trace.json.gz"), + report_dir=Path("rank0_perf_report_csvs"), +) +build_traces(db, [Path("a.json.gz"), Path("b.json.gz")]) +rows = search_index(db, "Cijk", limit=20) +``` + +## Related topics + +- [Install TraceLens](../install/install.md) +- [Generate a PyTorch performance report](./generate-perf-report-pytorch.md) +- [Generate a PyTorch inference performance report](./generate-perf-report-pytorch-inference.md) +- [Analyze traces with the TraceLens SDK](./sdk-analysis.md) +- [Performance report columns](../reference/perf-report-columns.md) +- [API reference](../reference/api-reference.md) diff --git a/docs/index.rst b/docs/index.rst index 9a91a928a..72d8bb848 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -41,6 +41,7 @@ The TraceLens source code is hosted at `github.com/AMD-AGI/TraceLens ` * :doc:`Fuse multi-rank traces ` * :doc:`Analyze traces with the SDK ` + * :doc:`Index a corpus of traces ` * :doc:`Model op performance without a trace ` * :doc:`Analyze collective communication ` * :doc:`Estimate kernel times with Origami ` diff --git a/docs/reference/api-reference.md b/docs/reference/api-reference.md index 4ae3abcc8..39df69a67 100644 --- a/docs/reference/api-reference.md +++ b/docs/reference/api-reference.md @@ -177,6 +177,23 @@ Split an inference trace into per-iteration or per-phase sub-traces. | `--divide-phases` | off | Store steady-state steps into `prefilldecodemix/` and `decode_only/` sub-folders. | | `--CONC`, `--OSL`, `--R` | None | Expected concurrency and output-sequence-length window parameters. | +### TraceLens_trace_index + +Catalog profiler traces and TraceLens CSV reports into a searchable index. +SQLite is the first backend; the tables are the shared catalog schema. See +[Index a corpus of traces](../how-to/trace-index.md) for the full workflow. + +Global options: `--backend` (default `sqlite`) and `--db` (default +`trace_index.sqlite`). + +| Command | Description | +|---------|-------------| +| `append --trace-path PATH` | Append one trace. Pass `--report-dir` to load an existing CSV report; otherwise generate a training PyTorch report. | +| `build --traces-file FILE` | Create or open the catalog and append a batch of traces (one path per line). Repeatable `--trace-path` is also accepted. | +| `search TERMS` | Full-text search over indexed ops, kernels, categories, and timeline labels. | +| `sqlite-sql SQL` | Run one read-only SQL statement. | +| `serve` | HTTP SQL endpoint on `127.0.0.1:8765` by default. | + ## Python SDK The SDK modules live under the `TraceLens` package and can be imported to build @@ -193,6 +210,7 @@ example notebook under `examples/`. | `EventReplay` | Extract and replay isolated operations. | [Replay a single operation](../how-to/event-replay.md), [`event_replayer_example.ipynb`](https://github.com/AMD-AGI/TraceLens/blob/main/examples/event_replayer_example.ipynb) | | `TraceFusion` | Merge multi-rank traces for Perfetto visualization. | [Fuse multi-rank traces](../how-to/trace-fusion.md), [`trace_fusion_example.py`](https://github.com/AMD-AGI/TraceLens/blob/main/examples/trace_fusion_example.py) | | `Reporting` | The report generators behind the CLI tools; importable to return pandas data frames. | [Generate a PyTorch performance report](../how-to/generate-perf-report-pytorch.md) | +| `TraceIndex` | Catalog traces into a searchable index (SQLite first backend). | [Index a corpus of traces](../how-to/trace-index.md) | | `TraceUtils` | Trace utilities, including inference-trace splitting. | — | For report-column definitions across all sheets, see the diff --git a/docs/sphinx/_toc.yml.in b/docs/sphinx/_toc.yml.in index 39746c18f..9c8f6c60b 100644 --- a/docs/sphinx/_toc.yml.in +++ b/docs/sphinx/_toc.yml.in @@ -46,6 +46,8 @@ subtrees: title: Fuse multi-rank traces - file: how-to/sdk-analysis title: Analyze traces with the SDK + - file: how-to/trace-index + title: Index a corpus of traces - file: how-to/perf-model-without-trace title: Model op performance without a trace - file: how-to/nccl-analysis diff --git a/docs/what-is-tracelens.md b/docs/what-is-tracelens.md index ecd4850bf..9c3dce97f 100644 --- a/docs/what-is-tracelens.md +++ b/docs/what-is-tracelens.md @@ -50,6 +50,8 @@ TraceLens provides these capabilities: - **Event replay:** Isolate any operation for focused debugging. TraceLens generates minimal, self-contained replay scripts from trace metadata, making it straightforward to share IP-safe reproducers with kernel developers. +- **Corpus catalog:** Index many traces into a shared catalog and search ops, + kernels, and GPU-timeline metrics without reopening each raw file. - **Extensible SDK:** Start with ready-to-use scripts, then build custom workflows with a flexible Python API. - **Agentic analysis:** Turn a raw trace into a prioritized, human-readable @@ -71,6 +73,8 @@ TraceLens is suited to these scenarios: quantify the effect of a code, library, or hardware change. - **Reproducer generation:** Extract a single operator into a standalone replay script to share with kernel or framework developers. +- **Corpus search:** Index a set of traces and query kernels, op categories, or + performance ranges across the set. - **Autonomous bottleneck triage:** Hand a trace to the TraceLens Agent and get back a ranked action list of the highest-impact optimizations, ready for review or for feeding into automated performance-tuning platforms. diff --git a/setup.py b/setup.py index 3a924c4d7..280fbc295 100755 --- a/setup.py +++ b/setup.py @@ -93,6 +93,7 @@ def _wheel_version(): "TraceLens_generate_perf_report_pftrace_memory_copy = TraceLens.Reporting.generate_perf_report_pftrace_memory_copy:main", "TraceLens_generate_perf_report_genesis = TraceLens.Reporting.generate_perf_report_genesis:main", "TraceLens_split_inference_trace = TraceLens.TraceUtils.split_inference_trace_annotation:main", + "TraceLens_trace_index = TraceLens.TraceIndex.cli:main", ], }, ) diff --git a/tests/test_kernel_library.py b/tests/test_kernel_library.py new file mode 100644 index 000000000..555808c6b --- /dev/null +++ b/tests/test_kernel_library.py @@ -0,0 +1,55 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +import pytest + +from TraceLens.PerfModel.kernel_library import classify_kernel_library +from TraceLens.TraceIndex.utils import kernel_flags + + +@pytest.mark.parametrize( + "op_name,kernel_details,expected", + [ + ("aiter::gemm", "", "AITER"), + ("triton_kernel_op", "", "Triton"), + ("aten::mm", "[{'name': 'Cijk_foo'}]", "Tensile"), + ("aten::mm", "void at::native::x", "PyTorch Native"), + ("aten::mm", "plain", None), + ("aten::mm", "ncclAllReduce", "RCCL/NCCL"), + ( + "aten::mm", + "Cijk_Alik_Bljk_BBS_BH_Bias_HAS_SAV_UserArgs_MT64x16x64_MI16x16x1", + "Tensile", + ), + ("aten::mm", "nvjet_tst_144x128_64x6_2x1_v_bz_bias_TNN", "nvjet"), + ], +) +def test_classify_kernel_library(op_name, kernel_details, expected): + assert classify_kernel_library(op_name, kernel_details) == expected + + +@pytest.mark.parametrize( + "kernel_name,op_name,expected_library,expected_tensile,expected_transpose", + [ + ( + "Cijk_Alik_Bljk_BBS_BH_Bias_HAS_SAV_UserArgs_MT64x16x64", + "aten::mm", + "Tensile", + 1, + 1, + ), + ("triton_red_fused_add", "aten::add", "Triton", 0, 0), + ("ncclAllReduceRing", "", "RCCL/NCCL", 0, 0), + ("custom_transpose_kernel", "", None, 0, 1), + ], +) +def test_kernel_flags( + kernel_name, op_name, expected_library, expected_tensile, expected_transpose +): + library, is_tensile, is_transpose, _is_layout = kernel_flags(kernel_name, op_name) + assert library == expected_library + assert is_tensile == expected_tensile + assert is_transpose == expected_transpose diff --git a/tests/test_trace_index.py b/tests/test_trace_index.py new file mode 100644 index 000000000..33d27cea4 --- /dev/null +++ b/tests/test_trace_index.py @@ -0,0 +1,466 @@ +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +import csv +import json +import threading +import urllib.error +import urllib.parse +import urllib.request +from contextlib import contextmanager +from http.server import ThreadingHTTPServer +from pathlib import Path + +import pytest + +from TraceLens.TraceIndex.core import ( + append_trace, + execute_read_query, + search_index, +) +from TraceLens.TraceIndex.cli import main as trace_index_main +from TraceLens.TraceIndex.importer import ( + build_traces as build_traces_with_store, +) +from TraceLens.TraceIndex.sqlite_store import SQLiteTraceIndexStore +from TraceLens.TraceIndex.server import make_handler +from TraceLens.TraceIndex.utils import ( + collect_trace_paths, + parse_repr, + read_traces_file, + to_json, +) + +FIXTURES = Path(__file__).resolve().parent / "traces" +TRAINING_REPORT_DIR = ( + FIXTURES / "mi300" / "Qwen_Qwen1.5-0.5B-Chat__1016005_perf_report_csvs" +) +INFERENCE_REPORT_DIR = FIXTURES / "inference" / "sglang_decode" / "perf_csvs" + + +def write_csv(path, rows): + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8", newline="") as f: + writer = csv.DictWriter(f, fieldnames=list(rows[0].keys())) + writer.writeheader() + writer.writerows(rows) + + +def write_stub_trace(path): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps({"traceEvents": []}), encoding="utf-8") + return path + + +def write_mini_report(report_dir): + write_csv( + report_dir / "unified_perf_summary.csv", + [ + { + "name": "aten::mm", + "op category": "GEMM", + "operation_count": "1", + "Kernel Time (us)_sum": "10.0", + }, + { + "name": "aten::add", + "op category": "elementwise", + "operation_count": "1", + "Kernel Time (us)_sum": "1.0", + }, + ], + ) + return report_dir + + +def seed_mini_catalog(tmp_path): + db_path = tmp_path / "trace_index.sqlite" + trace_path = write_stub_trace(tmp_path / "rank0_trace.json") + append_trace(db_path, trace_path, report_dir=write_mini_report(tmp_path / "report")) + return db_path + + +def request_json(url, method="GET", payload=None): + data = None + headers = {} + if payload is not None: + data = json.dumps(payload).encode("utf-8") + headers["Content-Type"] = "application/json" + req = urllib.request.Request(url, data=data, method=method, headers=headers) + try: + with urllib.request.urlopen(req, timeout=5) as resp: + return resp.status, json.loads(resp.read().decode("utf-8")) + except urllib.error.HTTPError as exc: + body = exc.read().decode("utf-8") + return exc.code, json.loads(body) + + +@contextmanager +def query_server(db_path): + handler = make_handler(db_path, default_limit=500, max_limit=5000) + server = ThreadingHTTPServer(("127.0.0.1", 0), handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + host, port = server.server_address + try: + yield "http://%s:%s" % (host, port) + finally: + server.shutdown() + server.server_close() + thread.join(timeout=2) + + +def test_parse_repr_strips_numpy_scalars_and_to_json(): + """parse_repr turns a CSV Python repr (with np scalars) into plain data that + to_json can serialize.""" + parsed = parse_repr("[{'name': 'k', 'total_duration_us': np.float64(1.5)}]") + assert parsed[0]["name"] == "k" + assert parsed[0]["total_duration_us"] == 1.5 + encoded = to_json(parse_repr("{'M': 128, 'bias': False}")) + assert encoded is not None + assert json.loads(encoded) == {"M": 128, "bias": False} + + +def test_trace_index_append_from_report_and_search(tmp_path): + """Appending a report explodes kernels into op_kernels and fills the + gemm/sdpa/conv satellites, and the parsed perf_params are queryable via + json_extract and FTS.""" + db_path = tmp_path / "trace_index.sqlite" + trace_root = tmp_path / "traces" + trace_path = write_stub_trace(trace_root / "model_a" / "rank0_trace.json") + + report_dir = tmp_path / "reports" / "trace_1" + write_csv( + report_dir / "unified_perf_summary.csv", + [ + { + "name": "aten::mm", + "op category": "GEMM", + "operation_count": "2", + "Kernel Time (us)_sum": "123.5", + "TFLOPS/s_mean": "98.1", + "perf_params": ( + "{'M': 128, 'N': 64, 'K': 32, 'B': 1, 'bias': False, " + "'dtype_A_B': ('c10::BFloat16', 'c10::BFloat16'), " + "'transpose': (True, False)}" + ), + "kernel_details_summary": ( + "[{'name': 'Cijk_test_kernel', 'stream': 0, 'count': 2, " + "'total_duration_us': 123.5, 'mean_duration_us': 61.75, " + "'median_duration_us': 61.75, 'min_duration_us': 60.0, " + "'max_duration_us': 63.5}]" + ), + }, + { + "name": "aten::scaled_dot_product_attention", + "op category": "SDPA_fwd", + "operation_count": "1", + "Kernel Time (us)_sum": "10.0", + "TFLOPS/s_mean": "40.0", + "perf_params": ( + "{'B': 2, 'H_Q': 8, 'N_Q': 128, 'N_KV': 128, " + "'d_h_qk': 64, 'd_h_v': 64, 'causal': True, " + "'dtype_A_B': ('c10::BFloat16', 'c10::BFloat16')}" + ), + }, + { + "name": "aten::convolution", + "op category": "CONV", + "operation_count": "1", + "Kernel Time (us)_sum": "20.0", + "perf_params": ( + "{'convNd': 'conv2d', 'input_shape': (2, 32, 56, 56), " + "'filter_shape': (32, 1, 3, 3), 'groups': 32, " + "'transposed_conv': False}" + ), + }, + ], + ) + write_csv( + report_dir / "ops_summary_by_category.csv", + [ + { + "op category": "GEMM", + "operation_count": "2", + "Kernel Time (us)_sum": "123.5", + "Percentage (%)": "80.0", + } + ], + ) + write_csv( + report_dir / "gpu_timeline.csv", + [ + {"type": "total_time", "time ms": "1.0", "percent": "100.0"}, + {"type": "computation_time", "time ms": "0.8", "percent": "80.0"}, + ], + ) + + trace_id = append_trace(db_path, trace_path, report_dir=report_dir, root=trace_root) + assert trace_id == 1 + + rows = execute_read_query( + db_path, + "SELECT name, op_category, kernel_time_sum_us FROM unified_perf_rows " + "ORDER BY source_row", + ) + assert rows[0] == { + "name": "aten::mm", + "op_category": "GEMM", + "kernel_time_sum_us": 123.5, + } + + gemm = execute_read_query(db_path, "SELECT m, n, k, batch FROM gemm_perf") + assert gemm == [{"m": 128, "n": 64, "k": 32, "batch": 1}] + + params = execute_read_query( + db_path, + "SELECT json_extract(perf_params_json, '$.M') AS m " + "FROM unified_perf_rows WHERE name = 'aten::mm'", + ) + assert params[0]["m"] == 128 + + kernels = execute_read_query( + db_path, + "SELECT kernel_name, unified_row_id, library, stream, parent_op_name " + "FROM op_kernels", + ) + assert kernels[0]["kernel_name"] == "Cijk_test_kernel" + assert kernels[0]["unified_row_id"] is not None + assert kernels[0]["library"] == "Tensile" + assert kernels[0]["stream"] == 0 + assert kernels[0]["parent_op_name"] == "aten::mm" + + sdpa = execute_read_query( + db_path, "SELECT seq_q, seq_kv, head_dim, causal FROM sdpa_perf" + ) + assert sdpa == [{"seq_q": 128, "seq_kv": 128, "head_dim": 64, "causal": 1}] + + conv = execute_read_query( + db_path, "SELECT groups, is_depthwise, is_transposed_conv FROM conv_perf" + ) + assert conv == [{"groups": 32, "is_depthwise": 1, "is_transposed_conv": 0}] + + search_rows = search_index(db_path, "Cijk", limit=10) + assert search_rows + assert search_rows[0]["trace_id"] == trace_id + + +def test_trace_index_rejects_write_sql(tmp_path): + """The read-only query path refuses non-SELECT statements.""" + db_path = tmp_path / "trace_index.sqlite" + with pytest.raises(ValueError): + execute_read_query(db_path, "DELETE FROM traces") + + +@pytest.mark.skipif( + not TRAINING_REPORT_DIR.exists(), + reason="checked-in Qwen training report CSVs are missing", +) +def test_import_real_training_report_maps_kernel_stream_and_times(tmp_path): + """On the checked-in Qwen training report, op_kernels/gemm/sdpa are populated + from real perf_params and kernel_details with correct shapes and stream.""" + db_path = tmp_path / "trace_index.sqlite" + trace_path = write_stub_trace(tmp_path / "qwen_trace.json") + trace_id = append_trace(db_path, trace_path, report_dir=TRAINING_REPORT_DIR) + + unified = execute_read_query( + db_path, + "SELECT name, kernel_time_sum_us FROM unified_perf_rows WHERE name = 'aten::mm'", + ) + assert unified + assert unified[0]["kernel_time_sum_us"] > 0 + + kernels = execute_read_query( + db_path, + "SELECT k.kernel_name, k.stream, k.total_duration_us, k.unified_row_id, " + "k.library FROM op_kernels k " + "JOIN unified_perf_rows u ON u.id = k.unified_row_id " + "WHERE k.kernel_name LIKE 'Cijk%' LIMIT 1", + ) + assert kernels + assert kernels[0]["stream"] == 0 + assert kernels[0]["total_duration_us"] > 0 + assert kernels[0]["unified_row_id"] is not None + assert kernels[0]["library"] == "Tensile" + + gemm = execute_read_query( + db_path, + "SELECT g.m, g.n, g.k FROM gemm_perf g " + "JOIN unified_perf_rows u ON u.id = g.unified_row_id " + "WHERE u.name = 'aten::mm' AND g.n = 2816 AND g.k = 1024 LIMIT 1", + ) + assert gemm + assert gemm[0]["m"] == 8944 + assert gemm[0]["n"] == 2816 + assert gemm[0]["k"] == 1024 + + params = execute_read_query( + db_path, + "SELECT json_extract(perf_params_json, '$.M') AS m " + "FROM unified_perf_rows WHERE name = 'aten::mm' LIMIT 1", + ) + assert params[0]["m"] == 8944 + + sdpa = execute_read_query( + db_path, + "SELECT seq_q, seq_kv, head_dim FROM sdpa_perf LIMIT 1", + ) + assert sdpa + assert sdpa[0]["seq_q"] is not None + + categories = execute_read_query( + db_path, + "SELECT category, kernel_time_sum_us FROM op_category_rows " + "WHERE kernel_time_sum_us IS NOT NULL ORDER BY kernel_time_sum_us DESC", + ) + assert categories + assert categories[0]["kernel_time_sum_us"] > 0 + assert trace_id == 1 + + +@pytest.mark.skipif( + not INFERENCE_REPORT_DIR.exists(), + reason="checked-in inference report CSVs are missing", +) +def test_import_real_inference_report_converts_category_kernel_time_ms(tmp_path): + """On the checked-in inference report, category kernel time in ms is converted + to microseconds during import.""" + db_path = tmp_path / "trace_index.sqlite" + trace_path = write_stub_trace(tmp_path / "decode_trace.json") + append_trace(db_path, trace_path, report_dir=INFERENCE_REPORT_DIR) + rows = execute_read_query( + db_path, + "SELECT category, kernel_time_sum_us FROM op_category_rows " + "WHERE category = 'GEMM'", + ) + assert rows + assert rows[0]["kernel_time_sum_us"] > 1000 + + +def test_read_traces_file_skips_comments_and_blanks(tmp_path): + """A traces-file drops blank/comment lines and collect_trace_paths merges and + de-duplicates paths.""" + traces_file = tmp_path / "traces.txt" + traces_file.write_text( + "# header\n" "\n" " /data/a.json.gz \n" "# skip me\n" "C:/traces/b.json\n", + encoding="utf-8", + ) + paths = read_traces_file(traces_file) + assert paths == [Path("/data/a.json.gz"), Path("C:/traces/b.json")] + combined = collect_trace_paths( + traces_file, [Path("/data/a.json.gz"), Path("c.json")] + ) + assert combined == [ + Path("/data/a.json.gz"), + Path("C:/traces/b.json"), + Path("c.json"), + ] + + +def test_cli_append_from_existing_report(tmp_path, capsys): + """The CLI append command imports an existing report dir without regenerating + it and reports the new trace_id.""" + db_path = tmp_path / "trace_index.sqlite" + trace_path = write_stub_trace(tmp_path / "rank0_trace.json") + report_dir = write_mini_report(tmp_path / "report") + + rc = trace_index_main( + [ + "--db", + str(db_path), + "append", + "--trace-path", + str(trace_path), + "--report-dir", + str(report_dir), + ] + ) + assert rc == 0 + payload = json.loads(capsys.readouterr().out) + assert payload["trace_id"] == 1 + assert payload["generated_report"] is False + rows = execute_read_query( + db_path, "SELECT name FROM unified_perf_rows ORDER BY source_row" + ) + assert [row["name"] for row in rows] == ["aten::mm", "aten::add"] + + +def test_build_traces_continues_after_failure(tmp_path): + """A batch build records per-trace failures and keeps processing the rest of + the list.""" + db_path = tmp_path / "trace_index.sqlite" + store = SQLiteTraceIndexStore(db_path) + try: + result = build_traces_with_store( + store, + [tmp_path / "missing_a.json", tmp_path / "missing_b.json"], + report_root=tmp_path / "reports", + ) + finally: + store.close() + assert result["imported"] == [] + assert len(result["failed"]) == 2 + assert "missing_a.json" in result["failed"][0]["trace_path"] + + +def test_query_server_health_query_and_guards(tmp_path): + """The read-only HTTP server serves health/tables/SQL and rejects writes.""" + db_path = seed_mini_catalog(tmp_path) + with query_server(db_path) as base: + status, root = request_json(base + "/") + assert status == 200 + assert "POST /query" in root["endpoints"] + + status, health = request_json(base + "/health") + assert status == 200 + assert health["ok"] is True + + status, tables = request_json(base + "/tables") + assert status == 200 + assert tables["tables"]["unified_perf_rows"] == 2 + + status, queried = request_json( + base + "/query", + method="POST", + payload={ + "sql": "SELECT name FROM unified_perf_rows ORDER BY source_row", + "limit": 10, + }, + ) + assert status == 200 + assert queried["truncated"] is False + assert [row["name"] for row in queried["rows"]] == ["aten::mm", "aten::add"] + + encoded = urllib.parse.urlencode( + {"sql": "SELECT COUNT(*) AS n FROM unified_perf_rows", "limit": "1"} + ) + status, get_query = request_json(base + "/query?" + encoded) + assert status == 200 + assert get_query["rows"][0]["n"] == 2 + + status, truncated = request_json( + base + "/query", + method="POST", + payload={"sql": "SELECT name FROM unified_perf_rows", "limit": 1}, + ) + assert status == 200 + assert truncated["truncated"] is True + assert len(truncated["rows"]) == 1 + + status, payload = request_json( + base + "/query", + method="POST", + payload={"sql": "DELETE FROM traces"}, + ) + assert status == 400 + assert "read-only" in payload["error"] + + status, missing = request_json(base + "/nope") + assert status == 404 + assert missing["error"] == "not found" + + status, post_missing = request_json(base + "/tables", method="POST", payload={}) + assert status == 404