From dd3decbe5285707925dc9cb6e9e0f6c4e75cd666 Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Mon, 24 Aug 2026 14:25:56 -0400 Subject: [PATCH 01/18] Add TraceIndex corpus search feature. Co-authored-by: Cursor --- README.md | 11 + TraceLens/TraceIndex/__init__.py | 23 + TraceLens/TraceIndex/cli.py | 152 ++++++ TraceLens/TraceIndex/core.py | 791 +++++++++++++++++++++++++++++++ TraceLens/TraceIndex/server.py | 157 ++++++ docs/TraceIndex.md | 104 ++++ setup.py | 1 + tests/test_trace_index.py | 101 ++++ 8 files changed, 1340 insertions(+) create mode 100644 TraceLens/TraceIndex/__init__.py create mode 100644 TraceLens/TraceIndex/cli.py create mode 100644 TraceLens/TraceIndex/core.py create mode 100644 TraceLens/TraceIndex/server.py create mode 100644 docs/TraceIndex.md create mode 100644 tests/test_trace_index.py diff --git a/README.md b/README.md index 75c429edc..cfaad14b5 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. + --- ## Quick Start @@ -79,6 +81,14 @@ TraceLens_compare_perf_reports_pytorch \ -o comparison.xlsx ``` +Index a directory of traces and existing TraceLens CSV reports (see [TraceIndex](docs/TraceIndex.md)): + +```bash +TraceLens_trace_index --backend sqlite --db trace_index.sqlite scan --root /path/to/traces +TraceLens_trace_index --backend sqlite --db trace_index.sqlite import-report --report-dir path/to/perf_report_csvs +TraceLens_trace_index --backend sqlite --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 +149,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/TraceIndex.md](docs/TraceIndex.md) | --- diff --git a/TraceLens/TraceIndex/__init__.py b/TraceLens/TraceIndex/__init__.py new file mode 100644 index 000000000..945ae8253 --- /dev/null +++ b/TraceLens/TraceIndex/__init__.py @@ -0,0 +1,23 @@ +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Trace corpus indexing helpers.""" + +from .core import ( + execute_read_query, + generate_report_and_import, + import_report_dir, + scan_traces, + search_index, +) + +__all__ = [ + "execute_read_query", + "generate_report_and_import", + "import_report_dir", + "scan_traces", + "search_index", +] diff --git a/TraceLens/TraceIndex/cli.py b/TraceLens/TraceIndex/cli.py new file mode 100644 index 000000000..172f9942d --- /dev/null +++ b/TraceLens/TraceIndex/cli.py @@ -0,0 +1,152 @@ +############################################################################### +# 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.core import ( + execute_read_query, + generate_report_and_import, + import_report_dir, + scan_traces, + search_index, +) + + +DEFAULT_DB = Path("trace_index.sqlite") + + +def print_json(payload: object) -> None: + print(json.dumps(payload, indent=2, default=str)) + + +def scan_cmd(args: argparse.Namespace) -> int: + count = scan_traces( + db_path=args.db, + root=args.root, + peek_mb=args.peek_mb, + compute_md5=args.compute_md5, + ) + print_json({"db": args.db, "root": args.root, "candidate_traces": count}) + return 0 + + +def import_report_cmd(args: argparse.Namespace) -> int: + trace_id = import_report_dir( + db_path=args.db, + report_dir=args.report_dir, + trace_path=args.trace_path, + root=args.root, + ) + print_json({"db": args.db, "trace_id": trace_id, "report_dir": args.report_dir}) + return 0 + + +def build_cmd(args: argparse.Namespace) -> int: + trace_id = generate_report_and_import( + db_path=args.db, + trace_path=args.trace_path, + report_dir=args.report_dir, + root=args.root, + force=args.force, + enable_pseudo_ops=args.enable_pseudo_ops, + ) + print_json({"db": args.db, "trace_id": trace_id, "trace_path": args.trace_path}) + return 0 + + +def search_cmd(args: argparse.Namespace) -> int: + rows = search_index(args.db, " ".join(args.terms), limit=args.limit) + print_json({"rows": rows}) + return 0 + + +def sql_cmd(args: argparse.Namespace) -> int: + rows = execute_read_query(args.db, args.sql, limit=args.limit) + print_json({"rows": rows}) + return 0 + + +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 build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Build and query a SQLite index of TraceLens reports." + ) + parser.add_argument("--db", type=Path, default=DEFAULT_DB, help="SQLite DB path") + sub = parser.add_subparsers(dest="command") + + scan = sub.add_parser("scan", help="Catalog trace-like files under a root") + scan.add_argument("--root", type=Path, required=True) + scan.add_argument("--peek-mb", type=int, default=2) + scan.add_argument("--compute-md5", action="store_true") + scan.set_defaults(func=scan_cmd) + + import_report = sub.add_parser( + "import-report", + help="Import an existing TraceLens CSV report directory", + ) + import_report.add_argument("--report-dir", type=Path, required=True) + import_report.add_argument("--trace-path", type=Path, default=None) + import_report.add_argument("--root", type=Path, default=None) + import_report.set_defaults(func=import_report_cmd) + + build = sub.add_parser( + "build", + help="Generate a TraceLens CSV report for one trace, then import it", + ) + build.add_argument("--trace-path", type=Path, required=True) + build.add_argument("--report-dir", type=Path, default=None) + build.add_argument("--root", type=Path, default=None) + build.add_argument("--force", action="store_true") + build.add_argument("--enable-pseudo-ops", action="store_true") + 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("sql", help="Run a read-only SQL query") + sql.add_argument("sql") + sql.add_argument("--limit", type=int, default=500) + sql.set_defaults(func=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..2fb1f8648 --- /dev/null +++ b/TraceLens/TraceIndex/core.py @@ -0,0 +1,791 @@ +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""SQLite-backed index for searchable TraceLens trace corpora.""" + +import csv +import gzip +import hashlib +import json +import os +import re +import sqlite3 +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional, Sequence + + +TRACE_NAME_RE = re.compile(r"trace|profile|pytorch_profile|rocprof", re.IGNORECASE) +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 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 iter_files(root: Path) -> Iterable[Path]: + stack = [root] + while stack: + current = stack.pop() + try: + with os.scandir(current) as entries: + dirs = [] + for entry in entries: + try: + if entry.is_dir(follow_symlinks=False): + dirs.append(Path(entry.path)) + elif entry.is_file(follow_symlinks=False): + yield Path(entry.path) + except OSError: + continue + stack.extend(reversed(dirs)) + except OSError: + continue + + +def is_json_gz(path: Path) -> bool: + return path.name.lower().endswith(".json.gz") + + +def is_candidate(path: Path) -> bool: + name = path.name.lower() + suffix = path.suffix.lower() + if is_json_gz(path): + return True + if suffix in {".json", ".pftrace", ".rpd"}: + return True + if name.endswith(".xplane.pb"): + return True + if ".pt.trace" in name or ".trace." in name: + return True + return bool(TRACE_NAME_RE.search(path.name)) + + +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 connect(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_db(conn: sqlite3.Connection) -> None: + 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 kernel_summary ( + id INTEGER PRIMARY KEY, + trace_id INTEGER NOT NULL REFERENCES traces(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, + is_tensile INTEGER NOT NULL DEFAULT 0, + is_transpose INTEGER NOT NULL DEFAULT 0, + is_layout_conversion INTEGER NOT NULL DEFAULT 0, + raw_row_json TEXT + ); + + CREATE INDEX IF NOT EXISTS idx_trace_index_kernel_trace ON kernel_summary(trace_id); + CREATE INDEX IF NOT EXISTS idx_trace_index_kernel_name ON kernel_summary(kernel_name); + CREATE INDEX IF NOT EXISTS idx_trace_index_kernel_tensile ON kernel_summary(is_tensile); + + 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' + ); + """ + ) + + +def add_trace( + conn: sqlite3.Connection, + trace_path: Path, + root: Optional[Path] = None, + peek_bytes: int = 2 * 1024 * 1024, + compute_md5: bool = False, +) -> int: + 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) + skip_reason = classify_skip(trace_path, root) + should_enrich = int(skip_reason is None and detect_format(trace_path, prefix) 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 + now = utc_now() + 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 + """, + ( + normalize_path(root), + normalize_path(trace_path), + rel_path, + trace_path.name, + trace_path.stat().st_size if trace_path.exists() else None, + content_md5(trace_path) if compute_md5 else None, + detect_format(trace_path, prefix), + extract_rank(trace_path), + rel_parts[0] if rel_parts else "", + normalize_path(Path(*rel_parts[:-1])) if len(rel_parts) > 1 else "", + should_enrich, + skip_reason, + now, + now, + ), + ) + row = conn.execute("SELECT id FROM traces WHERE path = ?", (normalize_path(trace_path),)).fetchone() + return int(row["id"]) + + +def scan_traces( + db_path: Path, + root: Path, + peek_mb: int = 2, + compute_md5: bool = False, +) -> int: + conn = connect(db_path) + init_db(conn) + root = root.resolve() + count = 0 + for path in iter_files(root): + if not is_candidate(path): + continue + add_trace( + conn, + path, + root=root, + peek_bytes=peek_mb * 1024 * 1024, + compute_md5=compute_md5, + ) + count += 1 + conn.commit() + conn.close() + return count + + +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_bool_int(value: Any) -> int: + if isinstance(value, bool): + return int(value) + if value in (None, "", "nan", "NaN"): + return 0 + text = str(value).strip().lower() + return int(text in {"1", "true", "yes", "y"}) + + +def search_text(*parts: Any) -> str: + return " ".join(str(part) for part in parts if part not in (None, "", "nan", "NaN")) + + +def clear_trace_payload(conn: sqlite3.Connection, trace_id: int) -> None: + for table in ( + "report_imports", + "unified_perf_rows", + "kernel_summary", + "op_category_rows", + "gpu_timeline_rows", + "trace_summary", + ): + conn.execute(f"DELETE FROM {table} WHERE trace_id = ?", (trace_id,)) + conn.execute("DELETE FROM trace_search_FTS5 WHERE trace_id = ?", (trace_id,)) + + +def insert_search(conn: sqlite3.Connection, trace_id: int, kind: str, parts: Iterable[Any]) -> None: + text = search_text(*parts) + if text: + conn.execute( + "INSERT INTO trace_search_FTS5(trace_id, kind, text) VALUES (?, ?, ?)", + (trace_id, kind, text), + ) + + +def import_unified_rows( + conn: sqlite3.Connection, + 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) + 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_float(first_value(row, ["Kernel Time (us)_sum", "Kernel Time (µs)_sum", "total_direct_kernel_time_sum", "total_subtree_kernel_time_sum"])), + 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 (%)"])), + as_text(first_value(row, ["perf_params", "Perf Params"])), + as_text(first_value(row, ["kernel_details_summary", "trunc_kernel_details"])), + json.dumps(row, sort_keys=True), + ), + ) + insert_search(conn, 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_kernel_summary_rows( + conn: sqlite3.Connection, + trace_id: int, + rows: Sequence[Dict[str, str]], +) -> None: + for row in rows: + kernel_name = as_text(first_value(row, ["Kernel name", "kernel_name", "name"])) + if not kernel_name: + continue + parent_op_name = as_text(first_value(row, ["Parent cpu_op", "parent_op_name", "Launcher"])) + op_category = as_text(first_value(row, ["Parent op category", "op_category", "category"])) + conn.execute( + """ + INSERT INTO kernel_summary( + trace_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, is_tensile, is_transpose, + is_layout_conversion, raw_row_json + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + trace_id, + kernel_name, + parent_op_name, + op_category, + as_int(first_value(row, ["stream", "Stream"])), + as_int(first_value(row, ["Kernel duration (us)_count", "Kernel duration (µs)_count", "count"])), + as_float(first_value(row, ["Kernel duration (us)_sum", "Kernel duration (µs)_sum", "total_us"])), + as_float(first_value(row, ["Kernel duration (us)_mean", "Kernel duration (µs)_mean", "mean_us"])), + as_float(first_value(row, ["Kernel duration (us)_median", "Kernel duration (µs)_median", "median_us"])), + as_float(first_value(row, ["Kernel duration (us)_min", "Kernel duration (µs)_min", "min_us"])), + as_float(first_value(row, ["Kernel duration (us)_max", "Kernel duration (µs)_max", "max_us"])), + int("cijk" in kernel_name.lower() or "tensile" in kernel_name.lower()), + int("transpose" in kernel_name.lower()), + int("layout" in kernel_name.lower() or "permute" in kernel_name.lower()), + json.dumps(row, sort_keys=True), + ), + ) + insert_search(conn, trace_id, "kernel", [kernel_name, parent_op_name, op_category]) + + +def import_category_rows( + conn: sqlite3.Connection, + 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_float(first_value(row, ["Kernel Time (us)_sum", "Kernel Time (µs)_sum", "total_direct_kernel_time_sum", "total_subtree_kernel_time_sum"])) + 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}) + insert_search(conn, 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( + conn: sqlite3.Connection, + 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 + 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), + ), + ) + insert_search(conn, trace_id, "timeline", [metric_type]) + return total_duration_us + + +def ensure_trace_row( + conn: sqlite3.Connection, + trace_path: Optional[Path], + report_dir: Path, + root: Optional[Path], +) -> int: + if trace_path is not None: + return add_trace(conn, trace_path, root=root) + synthetic_path = normalize_path(report_dir.resolve()) + now = utc_now() + conn.execute( + """ + INSERT INTO traces(path, rel_path, name, format, should_enrich, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(path) DO UPDATE SET updated_at = excluded.updated_at + """, + (synthetic_path, report_dir.name, report_dir.name, "tracelens_report_dir", 1, now, now), + ) + row = conn.execute("SELECT id FROM traces WHERE path = ?", (synthetic_path,)).fetchone() + return int(row["id"]) + + +def import_report_dir( + db_path: Path, + report_dir: Path, + trace_path: Optional[Path] = None, + root: Optional[Path] = None, +) -> int: + report_dir = report_dir.resolve() + conn = connect(db_path) + init_db(conn) + trace_id = ensure_trace_row(conn, trace_path, report_dir, root) + clear_trace_payload(conn, trace_id) + + sheet_rows = { + "unified_perf_summary": read_csv_rows(report_dir / "unified_perf_summary.csv"), + "kernel_summary": read_csv_rows(report_dir / "kernel_summary.csv"), + "ops_summary_by_category": read_csv_rows(report_dir / "ops_summary_by_category.csv"), + "gpu_timeline": read_csv_rows(report_dir / "gpu_timeline.csv"), + } + unified_summary = import_unified_rows(conn, trace_id, sheet_rows["unified_perf_summary"]) + import_kernel_summary_rows(conn, trace_id, sheet_rows["kernel_summary"]) + top_categories_json = import_category_rows(conn, trace_id, sheet_rows["ops_summary_by_category"]) + total_duration_us = import_gpu_timeline_rows(conn, trace_id, sheet_rows["gpu_timeline"]) + + 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(), + ), + ) + conn.execute( + """ + INSERT INTO report_imports(trace_id, report_dir, imported_at, sheets_json) + VALUES (?, ?, ?, ?) + """, + ( + trace_id, + normalize_path(report_dir), + utc_now(), + json.dumps([name for name, rows in sheet_rows.items() if rows], sort_keys=True), + ), + ) + insert_search(conn, trace_id, "trace", [trace_path or report_dir, report_dir.name]) + conn.commit() + conn.close() + return trace_id + + +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: + if report_dir is None: + safe_name = re.sub(r"[^A-Za-z0-9_.-]+", "_", trace_path.name) + report_dir = db_path.resolve().parent / "trace_index_reports" / safe_name + 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(db_path, report_dir, trace_path=trace_path, root=root) + + +def search_index(db_path: Path, terms: str, limit: int = 50) -> List[Dict[str, Any]]: + conn = connect(db_path) + init_db(conn) + rows = 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() + conn.close() + return [dict(row) for row in rows] + + +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")) + + +def execute_read_query( + db_path: Path, + 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") + conn = connect(db_path) + init_db(conn) + conn.execute("PRAGMA query_only=ON") + rows = conn.execute(sql, params or ()).fetchmany(limit) + conn.close() + return [dict(row) for row in rows] diff --git a/TraceLens/TraceIndex/server.py b/TraceLens/TraceIndex/server.py new file mode 100644 index 000000000..d9f1f343f --- /dev/null +++ b/TraceLens/TraceIndex/server.py @@ -0,0 +1,157 @@ +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Read-only HTTP query server for TraceIndex SQLite databases.""" + +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.core 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 = int(params.get("limit", [str(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() + self.handle_query( + str(payload.get("sql", "")), + payload.get("params", []), + int(payload.get("limit", default_limit)), + ) + 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: + with self.connect() as conn: + 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}) + + def handle_query(self, sql: str, params: List[Any], limit: int) -> None: + 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(limit, max_limit)) + start = time.perf_counter() + with self.connect() as conn: + 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], + } + ) + + 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/docs/TraceIndex.md b/docs/TraceIndex.md new file mode 100644 index 000000000..a1e63e9df --- /dev/null +++ b/docs/TraceIndex.md @@ -0,0 +1,104 @@ + + +# TraceIndex + +TraceIndex builds a SQLite catalog of trace files and imports key TraceLens CSV +report tables so a corpus can be searched without reopening every raw trace. +It is useful when you have many profiler captures and want to answer questions +like: + +- Which traces contain GEMM, SDPA, convolution, collective, or short-kernel heavy workloads? +- Which traces contain a specific backend kernel name? +- Which trace should I open next in TraceLens or Perfetto? + +TraceIndex does not replace raw traces. It stores searchable summaries and paths +back to the source traces. + +## Quick Start + +Catalog trace-like files under a directory: + +```bash +TraceLens_trace_index --db trace_index.sqlite scan --root /path/to/traces +``` + +Generate a TraceLens report for one PyTorch trace and import it into the index: + +```bash +TraceLens_trace_index --db trace_index.sqlite build \ + --trace-path /path/to/traces/rank0_trace.json.gz \ + --report-dir ./trace_index_reports/rank0 +``` + +If you already have a TraceLens CSV report directory, import it directly: + +```bash +TraceLens_trace_index --db trace_index.sqlite import-report \ + --trace-path /path/to/traces/rank0_trace.json.gz \ + --report-dir ./rank0_perf_report_csvs +``` + +Search the full-text index: + +```bash +TraceLens_trace_index --db trace_index.sqlite search attention +TraceLens_trace_index --db trace_index.sqlite search Cijk +``` + +Run a read-only SQL query: + +```bash +TraceLens_trace_index --db trace_index.sqlite sql \ + "SELECT op_category, COUNT(*) AS rows FROM unified_perf_rows GROUP BY op_category" +``` + +## Imported Tables + +TraceIndex imports the stable report tables that are most useful for corpus +search: + +| Table | Contents | +|---|---| +| `traces` | One row per trace-like file or imported report directory | +| `report_imports` | Import history for TraceLens CSV report directories | +| `unified_perf_rows` | Rows from `unified_perf_summary.csv` | +| `kernel_summary` | Rows from `kernel_summary.csv`, including basic kernel flags | +| `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 Server + +For notebook or browser workflows, serve read-only SQL access locally: + +```bash +TraceLens_trace_index --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}` + +Only single `SELECT`, `WITH`, or `PRAGMA` statements are accepted. The server is +read-only but does not implement authentication, 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 import_report_dir, scan_traces, search_index + +db = Path("trace_index.sqlite") +scan_traces(db, Path("/path/to/traces")) +import_report_dir(db, Path("rank0_perf_report_csvs"), trace_path=Path("rank0_trace.json.gz")) +rows = search_index(db, "Cijk", limit=20) +``` 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_trace_index.py b/tests/test_trace_index.py new file mode 100644 index 000000000..10f8bb176 --- /dev/null +++ b/tests/test_trace_index.py @@ -0,0 +1,101 @@ +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +import csv +import json + +import pytest + +from TraceLens.TraceIndex.core import ( + execute_read_query, + import_report_dir, + scan_traces, + search_index, +) + + +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 test_trace_index_scan_import_and_search(tmp_path): + db_path = tmp_path / "trace_index.sqlite" + trace_root = tmp_path / "traces" + trace_path = trace_root / "model_a" / "rank0_trace.json" + trace_path.parent.mkdir(parents=True) + trace_path.write_text(json.dumps({"traceEvents": []}), encoding="utf-8") + + assert scan_traces(db_path, trace_root) == 1 + + 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", + "kernel_details_summary": "Cijk_test_kernel", + } + ], + ) + write_csv( + report_dir / "kernel_summary.csv", + [ + { + "Kernel name": "Cijk_test_kernel", + "Parent cpu_op": "aten::mm", + "Parent op category": "GEMM", + "Kernel duration (us)_count": "2", + "Kernel duration (us)_sum": "123.5", + } + ], + ) + 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 = import_report_dir(db_path, report_dir, trace_path=trace_path, 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", + ) + assert rows == [ + {"name": "aten::mm", "op_category": "GEMM", "kernel_time_sum_us": 123.5} + ] + + 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): + db_path = tmp_path / "trace_index.sqlite" + with pytest.raises(ValueError): + execute_read_query(db_path, "DELETE FROM traces") From 93e1582147dbea48917413dc54a2493e8044a670 Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Mon, 24 Aug 2026 14:27:01 -0400 Subject: [PATCH 02/18] Refactor TraceIndex around storage backends. Co-authored-by: Cursor --- README.md | 2 +- TraceLens/TraceIndex/__init__.py | 8 + TraceLens/TraceIndex/cli.py | 115 ++-- TraceLens/TraceIndex/core.py | 783 ++------------------------- TraceLens/TraceIndex/importer.py | 100 ++++ TraceLens/TraceIndex/models.py | 36 ++ TraceLens/TraceIndex/scanner.py | 209 +++++++ TraceLens/TraceIndex/server.py | 4 +- TraceLens/TraceIndex/sqlite_store.py | 485 +++++++++++++++++ TraceLens/TraceIndex/store.py | 48 ++ TraceLens/TraceIndex/utils.py | 79 +++ docs/TraceIndex.md | 54 +- tests/test_trace_index.py | 39 ++ 13 files changed, 1159 insertions(+), 803 deletions(-) create mode 100644 TraceLens/TraceIndex/importer.py create mode 100644 TraceLens/TraceIndex/models.py create mode 100644 TraceLens/TraceIndex/scanner.py create mode 100644 TraceLens/TraceIndex/sqlite_store.py create mode 100644 TraceLens/TraceIndex/store.py create mode 100644 TraceLens/TraceIndex/utils.py diff --git a/README.md b/README.md index cfaad14b5..39855919f 100755 --- a/README.md +++ b/README.md @@ -31,7 +31,7 @@ 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. +**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. --- diff --git a/TraceLens/TraceIndex/__init__.py b/TraceLens/TraceIndex/__init__.py index 945ae8253..adbc72b21 100644 --- a/TraceLens/TraceIndex/__init__.py +++ b/TraceLens/TraceIndex/__init__.py @@ -13,8 +13,16 @@ scan_traces, search_index, ) +from .models import SearchHit, TraceRecord, TraceReport +from .sqlite_store import SQLiteTraceIndexStore +from .store import TraceIndexStore __all__ = [ + "TraceIndexStore", + "SQLiteTraceIndexStore", + "SearchHit", + "TraceRecord", + "TraceReport", "execute_read_query", "generate_report_and_import", "import_report_dir", diff --git a/TraceLens/TraceIndex/cli.py b/TraceLens/TraceIndex/cli.py index 172f9942d..41302ddcd 100644 --- a/TraceLens/TraceIndex/cli.py +++ b/TraceLens/TraceIndex/cli.py @@ -11,13 +11,9 @@ from pathlib import Path from typing import List, Optional -from TraceLens.TraceIndex.core import ( - execute_read_query, - generate_report_and_import, - import_report_dir, - scan_traces, - search_index, -) +from TraceLens.TraceIndex.importer import generate_report_and_import, import_report_dir +from TraceLens.TraceIndex.scanner import scan_traces +from TraceLens.TraceIndex.sqlite_store import SQLiteTraceIndexStore DEFAULT_DB = Path("trace_index.sqlite") @@ -27,51 +23,79 @@ 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 scan_cmd(args: argparse.Namespace) -> int: - count = scan_traces( - db_path=args.db, - root=args.root, - peek_mb=args.peek_mb, - compute_md5=args.compute_md5, - ) - print_json({"db": args.db, "root": args.root, "candidate_traces": count}) - return 0 + store = create_store(args) + try: + count = scan_traces( + store, + root=args.root, + peek_mb=args.peek_mb, + compute_md5=args.compute_md5, + ) + print_json({"backend": args.backend, "db": args.db, "root": args.root, "candidate_traces": count}) + return 0 + finally: + store.close() def import_report_cmd(args: argparse.Namespace) -> int: - trace_id = import_report_dir( - db_path=args.db, - report_dir=args.report_dir, - trace_path=args.trace_path, - root=args.root, - ) - print_json({"db": args.db, "trace_id": trace_id, "report_dir": args.report_dir}) - return 0 + store = create_store(args) + try: + trace_id = import_report_dir( + store, + report_dir=args.report_dir, + trace_path=args.trace_path, + root=args.root, + ) + print_json({"backend": args.backend, "db": args.db, "trace_id": trace_id, "report_dir": args.report_dir}) + return 0 + finally: + store.close() def build_cmd(args: argparse.Namespace) -> int: - trace_id = generate_report_and_import( - db_path=args.db, - trace_path=args.trace_path, - report_dir=args.report_dir, - root=args.root, - force=args.force, - enable_pseudo_ops=args.enable_pseudo_ops, - ) - print_json({"db": args.db, "trace_id": trace_id, "trace_path": args.trace_path}) - return 0 + store = create_store(args) + try: + trace_id = generate_report_and_import( + store, + trace_path=args.trace_path, + report_dir=args.report_dir, + root=args.root, + force=args.force, + enable_pseudo_ops=args.enable_pseudo_ops, + ) + print_json({"backend": args.backend, "db": args.db, "trace_id": trace_id, "trace_path": args.trace_path}) + return 0 + finally: + store.close() def search_cmd(args: argparse.Namespace) -> int: - rows = search_index(args.db, " ".join(args.terms), limit=args.limit) - print_json({"rows": rows}) - return 0 - - -def sql_cmd(args: argparse.Namespace) -> int: - rows = execute_read_query(args.db, args.sql, limit=args.limit) - print_json({"rows": rows}) - return 0 + 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: @@ -89,8 +113,9 @@ def serve_cmd(args: argparse.Namespace) -> int: def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( - description="Build and query a SQLite index of TraceLens reports." + description="Build and query a TraceIndex catalog of TraceLens reports." ) + 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") @@ -125,10 +150,10 @@ def build_parser() -> argparse.ArgumentParser: search.add_argument("--limit", type=int, default=50) search.set_defaults(func=search_cmd) - sql = sub.add_parser("sql", help="Run a read-only SQL query") + 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=sql_cmd) + 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") diff --git a/TraceLens/TraceIndex/core.py b/TraceLens/TraceIndex/core.py index 2fb1f8648..ffdb763bc 100644 --- a/TraceLens/TraceIndex/core.py +++ b/TraceLens/TraceIndex/core.py @@ -4,371 +4,17 @@ # See LICENSE for license information. ############################################################################### -"""SQLite-backed index for searchable TraceLens trace corpora.""" +"""Compatibility facade for the default TraceIndex backend.""" -import csv -import gzip -import hashlib -import json -import os -import re -import sqlite3 -from datetime import datetime, timezone from pathlib import Path -from typing import Any, Dict, Iterable, List, Optional, Sequence +from typing import Any, Dict, List, Optional, Sequence - -TRACE_NAME_RE = re.compile(r"trace|profile|pytorch_profile|rocprof", re.IGNORECASE) -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", +from TraceLens.TraceIndex.importer import ( + generate_report_and_import as generate_report_and_import_with_store, ) - - -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 iter_files(root: Path) -> Iterable[Path]: - stack = [root] - while stack: - current = stack.pop() - try: - with os.scandir(current) as entries: - dirs = [] - for entry in entries: - try: - if entry.is_dir(follow_symlinks=False): - dirs.append(Path(entry.path)) - elif entry.is_file(follow_symlinks=False): - yield Path(entry.path) - except OSError: - continue - stack.extend(reversed(dirs)) - except OSError: - continue - - -def is_json_gz(path: Path) -> bool: - return path.name.lower().endswith(".json.gz") - - -def is_candidate(path: Path) -> bool: - name = path.name.lower() - suffix = path.suffix.lower() - if is_json_gz(path): - return True - if suffix in {".json", ".pftrace", ".rpd"}: - return True - if name.endswith(".xplane.pb"): - return True - if ".pt.trace" in name or ".trace." in name: - return True - return bool(TRACE_NAME_RE.search(path.name)) - - -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 connect(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_db(conn: sqlite3.Connection) -> None: - 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 kernel_summary ( - id INTEGER PRIMARY KEY, - trace_id INTEGER NOT NULL REFERENCES traces(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, - is_tensile INTEGER NOT NULL DEFAULT 0, - is_transpose INTEGER NOT NULL DEFAULT 0, - is_layout_conversion INTEGER NOT NULL DEFAULT 0, - raw_row_json TEXT - ); - - CREATE INDEX IF NOT EXISTS idx_trace_index_kernel_trace ON kernel_summary(trace_id); - CREATE INDEX IF NOT EXISTS idx_trace_index_kernel_name ON kernel_summary(kernel_name); - CREATE INDEX IF NOT EXISTS idx_trace_index_kernel_tensile ON kernel_summary(is_tensile); - - 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' - ); - """ - ) - - -def add_trace( - conn: sqlite3.Connection, - trace_path: Path, - root: Optional[Path] = None, - peek_bytes: int = 2 * 1024 * 1024, - compute_md5: bool = False, -) -> int: - 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) - skip_reason = classify_skip(trace_path, root) - should_enrich = int(skip_reason is None and detect_format(trace_path, prefix) 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 - now = utc_now() - 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 - """, - ( - normalize_path(root), - normalize_path(trace_path), - rel_path, - trace_path.name, - trace_path.stat().st_size if trace_path.exists() else None, - content_md5(trace_path) if compute_md5 else None, - detect_format(trace_path, prefix), - extract_rank(trace_path), - rel_parts[0] if rel_parts else "", - normalize_path(Path(*rel_parts[:-1])) if len(rel_parts) > 1 else "", - should_enrich, - skip_reason, - now, - now, - ), - ) - row = conn.execute("SELECT id FROM traces WHERE path = ?", (normalize_path(trace_path),)).fetchone() - return int(row["id"]) +from TraceLens.TraceIndex.importer import import_report_dir as import_report_dir_with_store +from TraceLens.TraceIndex.scanner import scan_traces as scan_traces_with_store +from TraceLens.TraceIndex.sqlite_store import SQLiteTraceIndexStore, is_read_only_sql def scan_traces( @@ -377,287 +23,16 @@ def scan_traces( peek_mb: int = 2, compute_md5: bool = False, ) -> int: - conn = connect(db_path) - init_db(conn) - root = root.resolve() - count = 0 - for path in iter_files(root): - if not is_candidate(path): - continue - add_trace( - conn, - path, + store = SQLiteTraceIndexStore(db_path) + try: + return scan_traces_with_store( + store, root=root, - peek_bytes=peek_mb * 1024 * 1024, + peek_mb=peek_mb, compute_md5=compute_md5, ) - count += 1 - conn.commit() - conn.close() - return count - - -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_bool_int(value: Any) -> int: - if isinstance(value, bool): - return int(value) - if value in (None, "", "nan", "NaN"): - return 0 - text = str(value).strip().lower() - return int(text in {"1", "true", "yes", "y"}) - - -def search_text(*parts: Any) -> str: - return " ".join(str(part) for part in parts if part not in (None, "", "nan", "NaN")) - - -def clear_trace_payload(conn: sqlite3.Connection, trace_id: int) -> None: - for table in ( - "report_imports", - "unified_perf_rows", - "kernel_summary", - "op_category_rows", - "gpu_timeline_rows", - "trace_summary", - ): - conn.execute(f"DELETE FROM {table} WHERE trace_id = ?", (trace_id,)) - conn.execute("DELETE FROM trace_search_FTS5 WHERE trace_id = ?", (trace_id,)) - - -def insert_search(conn: sqlite3.Connection, trace_id: int, kind: str, parts: Iterable[Any]) -> None: - text = search_text(*parts) - if text: - conn.execute( - "INSERT INTO trace_search_FTS5(trace_id, kind, text) VALUES (?, ?, ?)", - (trace_id, kind, text), - ) - - -def import_unified_rows( - conn: sqlite3.Connection, - 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) - 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_float(first_value(row, ["Kernel Time (us)_sum", "Kernel Time (µs)_sum", "total_direct_kernel_time_sum", "total_subtree_kernel_time_sum"])), - 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 (%)"])), - as_text(first_value(row, ["perf_params", "Perf Params"])), - as_text(first_value(row, ["kernel_details_summary", "trunc_kernel_details"])), - json.dumps(row, sort_keys=True), - ), - ) - insert_search(conn, 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_kernel_summary_rows( - conn: sqlite3.Connection, - trace_id: int, - rows: Sequence[Dict[str, str]], -) -> None: - for row in rows: - kernel_name = as_text(first_value(row, ["Kernel name", "kernel_name", "name"])) - if not kernel_name: - continue - parent_op_name = as_text(first_value(row, ["Parent cpu_op", "parent_op_name", "Launcher"])) - op_category = as_text(first_value(row, ["Parent op category", "op_category", "category"])) - conn.execute( - """ - INSERT INTO kernel_summary( - trace_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, is_tensile, is_transpose, - is_layout_conversion, raw_row_json - ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """, - ( - trace_id, - kernel_name, - parent_op_name, - op_category, - as_int(first_value(row, ["stream", "Stream"])), - as_int(first_value(row, ["Kernel duration (us)_count", "Kernel duration (µs)_count", "count"])), - as_float(first_value(row, ["Kernel duration (us)_sum", "Kernel duration (µs)_sum", "total_us"])), - as_float(first_value(row, ["Kernel duration (us)_mean", "Kernel duration (µs)_mean", "mean_us"])), - as_float(first_value(row, ["Kernel duration (us)_median", "Kernel duration (µs)_median", "median_us"])), - as_float(first_value(row, ["Kernel duration (us)_min", "Kernel duration (µs)_min", "min_us"])), - as_float(first_value(row, ["Kernel duration (us)_max", "Kernel duration (µs)_max", "max_us"])), - int("cijk" in kernel_name.lower() or "tensile" in kernel_name.lower()), - int("transpose" in kernel_name.lower()), - int("layout" in kernel_name.lower() or "permute" in kernel_name.lower()), - json.dumps(row, sort_keys=True), - ), - ) - insert_search(conn, trace_id, "kernel", [kernel_name, parent_op_name, op_category]) - - -def import_category_rows( - conn: sqlite3.Connection, - 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_float(first_value(row, ["Kernel Time (us)_sum", "Kernel Time (µs)_sum", "total_direct_kernel_time_sum", "total_subtree_kernel_time_sum"])) - 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}) - insert_search(conn, 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( - conn: sqlite3.Connection, - 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 - 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), - ), - ) - insert_search(conn, trace_id, "timeline", [metric_type]) - return total_duration_us - - -def ensure_trace_row( - conn: sqlite3.Connection, - trace_path: Optional[Path], - report_dir: Path, - root: Optional[Path], -) -> int: - if trace_path is not None: - return add_trace(conn, trace_path, root=root) - synthetic_path = normalize_path(report_dir.resolve()) - now = utc_now() - conn.execute( - """ - INSERT INTO traces(path, rel_path, name, format, should_enrich, created_at, updated_at) - VALUES (?, ?, ?, ?, ?, ?, ?) - ON CONFLICT(path) DO UPDATE SET updated_at = excluded.updated_at - """, - (synthetic_path, report_dir.name, report_dir.name, "tracelens_report_dir", 1, now, now), - ) - row = conn.execute("SELECT id FROM traces WHERE path = ?", (synthetic_path,)).fetchone() - return int(row["id"]) + finally: + store.close() def import_report_dir( @@ -666,56 +41,16 @@ def import_report_dir( trace_path: Optional[Path] = None, root: Optional[Path] = None, ) -> int: - report_dir = report_dir.resolve() - conn = connect(db_path) - init_db(conn) - trace_id = ensure_trace_row(conn, trace_path, report_dir, root) - clear_trace_payload(conn, trace_id) - - sheet_rows = { - "unified_perf_summary": read_csv_rows(report_dir / "unified_perf_summary.csv"), - "kernel_summary": read_csv_rows(report_dir / "kernel_summary.csv"), - "ops_summary_by_category": read_csv_rows(report_dir / "ops_summary_by_category.csv"), - "gpu_timeline": read_csv_rows(report_dir / "gpu_timeline.csv"), - } - unified_summary = import_unified_rows(conn, trace_id, sheet_rows["unified_perf_summary"]) - import_kernel_summary_rows(conn, trace_id, sheet_rows["kernel_summary"]) - top_categories_json = import_category_rows(conn, trace_id, sheet_rows["ops_summary_by_category"]) - total_duration_us = import_gpu_timeline_rows(conn, trace_id, sheet_rows["gpu_timeline"]) - - conn.execute( - """ - INSERT INTO trace_summary( - trace_id, total_duration_us, top_categories_json, max_gemm_tflops, - max_sdpa_tflops, imported_at + store = SQLiteTraceIndexStore(db_path) + try: + return import_report_dir_with_store( + store, + report_dir=report_dir, + trace_path=trace_path, + root=root, ) - VALUES (?, ?, ?, ?, ?, ?) - """, - ( - trace_id, - total_duration_us, - top_categories_json, - unified_summary["max_gemm_tflops"], - unified_summary["max_sdpa_tflops"], - utc_now(), - ), - ) - conn.execute( - """ - INSERT INTO report_imports(trace_id, report_dir, imported_at, sheets_json) - VALUES (?, ?, ?, ?) - """, - ( - trace_id, - normalize_path(report_dir), - utc_now(), - json.dumps([name for name, rows in sheet_rows.items() if rows], sort_keys=True), - ), - ) - insert_search(conn, trace_id, "trace", [trace_path or report_dir, report_dir.name]) - conn.commit() - conn.close() - return trace_id + finally: + store.close() def generate_report_and_import( @@ -726,53 +61,27 @@ def generate_report_and_import( force: bool = False, enable_pseudo_ops: bool = False, ) -> int: - if report_dir is None: - safe_name = re.sub(r"[^A-Za-z0-9_.-]+", "_", trace_path.name) - report_dir = db_path.resolve().parent / "trace_index_reports" / safe_name - 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, + 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, ) - return import_report_dir(db_path, report_dir, trace_path=trace_path, root=root) + finally: + store.close() def search_index(db_path: Path, terms: str, limit: int = 50) -> List[Dict[str, Any]]: - conn = connect(db_path) - init_db(conn) - rows = 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() - conn.close() - return [dict(row) for row in rows] - - -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")) + 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( @@ -781,11 +90,9 @@ def execute_read_query( 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") - conn = connect(db_path) - init_db(conn) - conn.execute("PRAGMA query_only=ON") - rows = conn.execute(sql, params or ()).fetchmany(limit) - conn.close() - return [dict(row) for row in rows] + 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..f007ded4c --- /dev/null +++ b/TraceLens/TraceIndex/importer.py @@ -0,0 +1,100 @@ +############################################################################### +# 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 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 + + +REPORT_SHEETS = ( + "unified_perf_summary", + "kernel_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 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: + safe_name = re.sub(r"[^A-Za-z0-9_.-]+", "_", trace_path.name) + report_dir = Path("trace_index_reports") / safe_name + 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..eebf7b846 --- /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 Any, 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..1465ed157 --- /dev/null +++ b/TraceLens/TraceIndex/scanner.py @@ -0,0 +1,209 @@ +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Trace discovery and metadata extraction.""" + +import gzip +import hashlib +import os +import re +from pathlib import Path +from typing import Iterable, Optional + +from TraceLens.TraceIndex.models import TraceRecord +from TraceLens.TraceIndex.store import TraceIndexStore +from TraceLens.TraceIndex.utils import normalize_path, rel_to + + +TRACE_NAME_RE = re.compile(r"trace|profile|pytorch_profile|rocprof", re.IGNORECASE) +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 iter_files(root: Path) -> Iterable[Path]: + stack = [root] + while stack: + current = stack.pop() + try: + with os.scandir(current) as entries: + dirs = [] + for entry in entries: + try: + if entry.is_dir(follow_symlinks=False): + dirs.append(Path(entry.path)) + elif entry.is_file(follow_symlinks=False): + yield Path(entry.path) + except OSError: + continue + stack.extend(reversed(dirs)) + except OSError: + continue + + +def is_json_gz(path: Path) -> bool: + return path.name.lower().endswith(".json.gz") + + +def is_candidate(path: Path) -> bool: + name = path.name.lower() + suffix = path.suffix.lower() + if is_json_gz(path): + return True + if suffix in {".json", ".pftrace", ".rpd"}: + return True + if name.endswith(".xplane.pb"): + return True + if ".pt.trace" in name or ".trace." in name: + return True + return bool(TRACE_NAME_RE.search(path.name)) + + +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, + ) + + +def scan_traces( + store: TraceIndexStore, + root: Path, + peek_mb: int = 2, + compute_md5: bool = False, +) -> int: + store.init_schema() + root = root.resolve() + count = 0 + for path in iter_files(root): + if not is_candidate(path): + continue + store.upsert_trace( + trace_record_from_path( + path, + root=root, + peek_bytes=peek_mb * 1024 * 1024, + compute_md5=compute_md5, + ) + ) + count += 1 + return count diff --git a/TraceLens/TraceIndex/server.py b/TraceLens/TraceIndex/server.py index d9f1f343f..70327fead 100644 --- a/TraceLens/TraceIndex/server.py +++ b/TraceLens/TraceIndex/server.py @@ -4,7 +4,7 @@ # See LICENSE for license information. ############################################################################### -"""Read-only HTTP query server for TraceIndex SQLite databases.""" +"""Read-only HTTP query server for the SQLite TraceIndex backend.""" import json import sqlite3 @@ -15,7 +15,7 @@ from typing import Any, List, Type from urllib.parse import parse_qs, urlparse -from TraceLens.TraceIndex.core import is_read_only_sql +from TraceLens.TraceIndex.sqlite_store import is_read_only_sql def json_bytes(payload: object) -> bytes: diff --git a/TraceLens/TraceIndex/sqlite_store.py b/TraceLens/TraceIndex/sqlite_store.py new file mode 100644 index 000000000..a6cc2c5e7 --- /dev/null +++ b/TraceLens/TraceIndex/sqlite_store.py @@ -0,0 +1,485 @@ +############################################################################### +# 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_float, + as_int, + as_text, + first_value, + search_text, + 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 kernel_summary ( + id INTEGER PRIMARY KEY, + trace_id INTEGER NOT NULL REFERENCES traces(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, + is_tensile INTEGER NOT NULL DEFAULT 0, + is_transpose INTEGER NOT NULL DEFAULT 0, + is_layout_conversion INTEGER NOT NULL DEFAULT 0, + raw_row_json TEXT + ); + + CREATE INDEX IF NOT EXISTS idx_trace_index_kernel_trace ON kernel_summary(trace_id); + CREATE INDEX IF NOT EXISTS idx_trace_index_kernel_name ON kernel_summary(kernel_name); + CREATE INDEX IF NOT EXISTS idx_trace_index_kernel_tensile ON kernel_summary(is_tensile); + + 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", [])) + self._import_kernel_summary_rows(trace_id, report.sheets.get("kernel_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") + 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", + "unified_perf_rows", + "kernel_summary", + "op_category_rows", + "gpu_timeline_rows", + "trace_summary", + ): + 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) + 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_float(first_value(row, ["Kernel Time (us)_sum", "Kernel Time (µs)_sum", "total_direct_kernel_time_sum", "total_subtree_kernel_time_sum"])), + 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 (%)"])), + as_text(first_value(row, ["perf_params", "Perf Params"])), + as_text(first_value(row, ["kernel_details_summary", "trunc_kernel_details"])), + json.dumps(row, sort_keys=True), + ), + ) + 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_kernel_summary_rows( + self, + trace_id: int, + rows: Sequence[Dict[str, str]], + ) -> None: + for row in rows: + kernel_name = as_text(first_value(row, ["Kernel name", "kernel_name", "name"])) + if not kernel_name: + continue + parent_op_name = as_text(first_value(row, ["Parent cpu_op", "parent_op_name", "Launcher"])) + op_category = as_text(first_value(row, ["Parent op category", "op_category", "category"])) + self.conn.execute( + """ + INSERT INTO kernel_summary( + trace_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, is_tensile, is_transpose, + is_layout_conversion, raw_row_json + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + trace_id, + kernel_name, + parent_op_name, + op_category, + as_int(first_value(row, ["stream", "Stream"])), + as_int(first_value(row, ["Kernel duration (us)_count", "Kernel duration (µs)_count", "count"])), + as_float(first_value(row, ["Kernel duration (us)_sum", "Kernel duration (µs)_sum", "total_us"])), + as_float(first_value(row, ["Kernel duration (us)_mean", "Kernel duration (µs)_mean", "mean_us"])), + as_float(first_value(row, ["Kernel duration (us)_median", "Kernel duration (µs)_median", "median_us"])), + as_float(first_value(row, ["Kernel duration (us)_min", "Kernel duration (µs)_min", "min_us"])), + as_float(first_value(row, ["Kernel duration (us)_max", "Kernel duration (µs)_max", "max_us"])), + int("cijk" in kernel_name.lower() or "tensile" in kernel_name.lower()), + int("transpose" in kernel_name.lower()), + int("layout" in kernel_name.lower() or "permute" in kernel_name.lower()), + json.dumps(row, sort_keys=True), + ), + ) + self._insert_search(trace_id, "kernel", [kernel_name, parent_op_name, op_category]) + + 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_float(first_value(row, ["Kernel Time (us)_sum", "Kernel Time (µs)_sum", "total_direct_kernel_time_sum", "total_subtree_kernel_time_sum"])) + 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..326b3cdf6 --- /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 scanner, + 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..e9fe76a41 --- /dev/null +++ b/TraceLens/TraceIndex/utils.py @@ -0,0 +1,79 @@ +############################################################################### +# 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 csv +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, List, Optional, Sequence + + +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 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_bool_int(value: Any) -> int: + if isinstance(value, bool): + return int(value) + if value in (None, "", "nan", "NaN"): + return 0 + text = str(value).strip().lower() + return int(text in {"1", "true", "yes", "y"}) + + +def search_text(*parts: Any) -> str: + return " ".join(str(part) for part in parts if part not in (None, "", "nan", "NaN")) diff --git a/docs/TraceIndex.md b/docs/TraceIndex.md index a1e63e9df..9b134acaa 100644 --- a/docs/TraceIndex.md +++ b/docs/TraceIndex.md @@ -6,30 +6,33 @@ See LICENSE for license information. # TraceIndex -TraceIndex builds a SQLite catalog of trace files and imports key TraceLens CSV -report tables so a corpus can be searched without reopening every raw trace. -It is useful when you have many profiler captures and want to answer questions -like: +TraceIndex builds a queryable catalog of trace files and imports key TraceLens +CSV report tables so a corpus can be searched without reopening every raw +trace. It is useful when you have many profiler captures and want to answer +questions like: - Which traces contain GEMM, SDPA, convolution, collective, or short-kernel heavy workloads? - Which traces contain a specific backend kernel name? - Which trace should I open next in TraceLens or Perfetto? TraceIndex does not replace raw traces. It stores searchable summaries and paths -back to the source traces. +back to the source traces. The first backend is SQLite because it needs no +service to run and works well for local or single-team workflows. The scanner, +importer, and storage interface are separated so other backends can be added by +implementing the same store contract. ## Quick Start Catalog trace-like files under a directory: ```bash -TraceLens_trace_index --db trace_index.sqlite scan --root /path/to/traces +TraceLens_trace_index --backend sqlite --db trace_index.sqlite scan --root /path/to/traces ``` Generate a TraceLens report for one PyTorch trace and import it into the index: ```bash -TraceLens_trace_index --db trace_index.sqlite build \ +TraceLens_trace_index --backend sqlite --db trace_index.sqlite build \ --trace-path /path/to/traces/rank0_trace.json.gz \ --report-dir ./trace_index_reports/rank0 ``` @@ -37,7 +40,7 @@ TraceLens_trace_index --db trace_index.sqlite build \ If you already have a TraceLens CSV report directory, import it directly: ```bash -TraceLens_trace_index --db trace_index.sqlite import-report \ +TraceLens_trace_index --backend sqlite --db trace_index.sqlite import-report \ --trace-path /path/to/traces/rank0_trace.json.gz \ --report-dir ./rank0_perf_report_csvs ``` @@ -45,14 +48,14 @@ TraceLens_trace_index --db trace_index.sqlite import-report \ Search the full-text index: ```bash -TraceLens_trace_index --db trace_index.sqlite search attention -TraceLens_trace_index --db trace_index.sqlite search Cijk +TraceLens_trace_index --backend sqlite --db trace_index.sqlite search attention +TraceLens_trace_index --backend sqlite --db trace_index.sqlite search Cijk ``` -Run a read-only SQL query: +Run a read-only SQLite query: ```bash -TraceLens_trace_index --db trace_index.sqlite sql \ +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" ``` @@ -72,12 +75,29 @@ search: | `trace_summary` | Per-trace summary metrics derived during import | | `trace_search_FTS5` | Full-text search over traces, ops, kernels, categories, and timeline labels | -## Query Server +## Backend Model + +TraceIndex separates workflow code from persistence: + +| Module | Responsibility | +|---|---| +| `models.py` | Backend-neutral trace, report, and search-result objects | +| `store.py` | `TraceIndexStore` interface for persistence/search backends | +| `scanner.py` | File discovery and trace metadata extraction | +| `importer.py` | TraceLens CSV report loading and import workflow | +| `sqlite_store.py` | SQLite schema, inserts, full-text search, and read-only SQL | +| `cli.py` | User-facing commands that call scanner/importer/store APIs | + +To add another backend, implement `TraceIndexStore` and wire it into the CLI +backend selector. Backend-neutral commands such as `scan`, `build`, +`import-report`, and `search` should not need to change. + +## SQLite Query Server For notebook or browser workflows, serve read-only SQL access locally: ```bash -TraceLens_trace_index --db trace_index.sqlite serve --host 127.0.0.1 --port 8765 +TraceLens_trace_index --backend sqlite --db trace_index.sqlite serve --host 127.0.0.1 --port 8765 ``` The server exposes: @@ -86,9 +106,9 @@ The server exposes: - `GET /tables` - `POST /query` with `{"sql": "SELECT ...", "params": [], "limit": 500}` -Only single `SELECT`, `WITH`, or `PRAGMA` statements are accepted. The server is -read-only but does not implement authentication, so bind it to loopback unless -you put it behind your own access control. +Only single SQLite `SELECT`, `WITH`, or `PRAGMA` statements are accepted. The +server is read-only but does not implement authentication, so bind it to +loopback unless you put it behind your own access control. ## Python API diff --git a/tests/test_trace_index.py b/tests/test_trace_index.py index 10f8bb176..a950ea342 100644 --- a/tests/test_trace_index.py +++ b/tests/test_trace_index.py @@ -15,6 +15,9 @@ scan_traces, search_index, ) +from TraceLens.TraceIndex.importer import import_report_dir as import_report_dir_with_store +from TraceLens.TraceIndex.scanner import scan_traces as scan_traces_with_store +from TraceLens.TraceIndex.sqlite_store import SQLiteTraceIndexStore def write_csv(path, rows): @@ -99,3 +102,39 @@ def test_trace_index_rejects_write_sql(tmp_path): db_path = tmp_path / "trace_index.sqlite" with pytest.raises(ValueError): execute_read_query(db_path, "DELETE FROM traces") + + +def test_trace_index_store_boundary_supports_scan_import_and_search(tmp_path): + db_path = tmp_path / "trace_index.sqlite" + trace_root = tmp_path / "traces" + trace_path = trace_root / "rank0_trace.json" + trace_path.parent.mkdir(parents=True) + trace_path.write_text(json.dumps({"traceEvents": []}), encoding="utf-8") + + report_dir = tmp_path / "reports" / "trace_1" + write_csv( + report_dir / "unified_perf_summary.csv", + [ + { + "name": "aten::scaled_dot_product_attention", + "op category": "SDPA_fwd", + "operation_count": "1", + "Kernel Time (us)_sum": "10.0", + } + ], + ) + + store = SQLiteTraceIndexStore(db_path) + try: + assert scan_traces_with_store(store, trace_root) == 1 + trace_id = import_report_dir_with_store( + store, + report_dir, + trace_path=trace_path, + root=trace_root, + ) + assert trace_id == 1 + hits = store.search("scaled", limit=10) + assert hits[0].kind == "op" + finally: + store.close() From 3c0b8b5ff60b47f47aff18119a3094612e52f411 Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Mon, 24 Aug 2026 14:30:27 -0400 Subject: [PATCH 03/18] Fix TraceIndex CSV aliases for real TraceLens reports. Map Kernel stream and category kernel-time_ms columns, cover import with checked-in report fixtures, and Black-format the new files. --- TraceLens/TraceIndex/cli.py | 40 +++- TraceLens/TraceIndex/core.py | 6 +- TraceLens/TraceIndex/importer.py | 1 - TraceLens/TraceIndex/scanner.py | 1 - TraceLens/TraceIndex/server.py | 48 +++-- TraceLens/TraceIndex/sqlite_store.py | 286 ++++++++++++++++++++++----- TraceLens/TraceIndex/utils.py | 14 ++ tests/test_trace_index.py | 65 +++++- 8 files changed, 389 insertions(+), 72 deletions(-) diff --git a/TraceLens/TraceIndex/cli.py b/TraceLens/TraceIndex/cli.py index 41302ddcd..e501e3b21 100644 --- a/TraceLens/TraceIndex/cli.py +++ b/TraceLens/TraceIndex/cli.py @@ -15,7 +15,6 @@ from TraceLens.TraceIndex.scanner import scan_traces from TraceLens.TraceIndex.sqlite_store import SQLiteTraceIndexStore - DEFAULT_DB = Path("trace_index.sqlite") @@ -38,7 +37,14 @@ def scan_cmd(args: argparse.Namespace) -> int: peek_mb=args.peek_mb, compute_md5=args.compute_md5, ) - print_json({"backend": args.backend, "db": args.db, "root": args.root, "candidate_traces": count}) + print_json( + { + "backend": args.backend, + "db": args.db, + "root": args.root, + "candidate_traces": count, + } + ) return 0 finally: store.close() @@ -53,7 +59,14 @@ def import_report_cmd(args: argparse.Namespace) -> int: trace_path=args.trace_path, root=args.root, ) - print_json({"backend": args.backend, "db": args.db, "trace_id": trace_id, "report_dir": args.report_dir}) + print_json( + { + "backend": args.backend, + "db": args.db, + "trace_id": trace_id, + "report_dir": args.report_dir, + } + ) return 0 finally: store.close() @@ -70,7 +83,14 @@ def build_cmd(args: argparse.Namespace) -> int: force=args.force, enable_pseudo_ops=args.enable_pseudo_ops, ) - print_json({"backend": args.backend, "db": args.db, "trace_id": trace_id, "trace_path": args.trace_path}) + print_json( + { + "backend": args.backend, + "db": args.db, + "trace_id": trace_id, + "trace_path": args.trace_path, + } + ) return 0 finally: store.close() @@ -80,7 +100,10 @@ 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)] + rows = [ + hit._asdict() + for hit in store.search(" ".join(args.terms), limit=args.limit) + ] print_json({"backend": args.backend, "rows": rows}) return 0 finally: @@ -115,7 +138,12 @@ def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( description="Build and query a TraceIndex catalog of TraceLens reports." ) - parser.add_argument("--backend", choices=["sqlite"], default="sqlite", help="TraceIndex storage backend") + 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") diff --git a/TraceLens/TraceIndex/core.py b/TraceLens/TraceIndex/core.py index ffdb763bc..401fcdadd 100644 --- a/TraceLens/TraceIndex/core.py +++ b/TraceLens/TraceIndex/core.py @@ -12,9 +12,11 @@ from TraceLens.TraceIndex.importer import ( generate_report_and_import as generate_report_and_import_with_store, ) -from TraceLens.TraceIndex.importer import import_report_dir as import_report_dir_with_store +from TraceLens.TraceIndex.importer import ( + import_report_dir as import_report_dir_with_store, +) from TraceLens.TraceIndex.scanner import scan_traces as scan_traces_with_store -from TraceLens.TraceIndex.sqlite_store import SQLiteTraceIndexStore, is_read_only_sql +from TraceLens.TraceIndex.sqlite_store import SQLiteTraceIndexStore def scan_traces( diff --git a/TraceLens/TraceIndex/importer.py b/TraceLens/TraceIndex/importer.py index f007ded4c..315025abd 100644 --- a/TraceLens/TraceIndex/importer.py +++ b/TraceLens/TraceIndex/importer.py @@ -15,7 +15,6 @@ from TraceLens.TraceIndex.store import TraceIndexStore from TraceLens.TraceIndex.utils import normalize_path, read_csv_rows - REPORT_SHEETS = ( "unified_perf_summary", "kernel_summary", diff --git a/TraceLens/TraceIndex/scanner.py b/TraceLens/TraceIndex/scanner.py index 1465ed157..b01eded46 100644 --- a/TraceLens/TraceIndex/scanner.py +++ b/TraceLens/TraceIndex/scanner.py @@ -17,7 +17,6 @@ from TraceLens.TraceIndex.store import TraceIndexStore from TraceLens.TraceIndex.utils import normalize_path, rel_to - TRACE_NAME_RE = re.compile(r"trace|profile|pytorch_profile|rocprof", re.IGNORECASE) RANK_RE = re.compile(r"(?:^|[^A-Za-z])rank[-_]?(\d+)(?:[^0-9]|$)", re.IGNORECASE) diff --git a/TraceLens/TraceIndex/server.py b/TraceLens/TraceIndex/server.py index 70327fead..ea797904e 100644 --- a/TraceLens/TraceIndex/server.py +++ b/TraceLens/TraceIndex/server.py @@ -22,14 +22,18 @@ 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]: +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: + 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") @@ -59,7 +63,11 @@ def do_GET(self) -> None: "endpoints": { "GET /health": "server and DB status", "GET /tables": "table row counts", - "POST /query": {"sql": "SELECT ...", "params": [], "limit": 500}, + "POST /query": { + "sql": "SELECT ...", + "params": [], + "limit": 500, + }, } } ) @@ -98,16 +106,17 @@ def handle_health(self) -> None: { "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, + "db_size_mb": ( + round(db_path.stat().st_size / 1048576, 2) + if db_path.exists() + else None + ), } ) def handle_tables(self) -> None: with self.connect() as conn: - tables = [ - row["name"] - for row in conn.execute( - """ + tables = [row["name"] for row in conn.execute(""" SELECT name FROM sqlite_master WHERE type IN ('table', 'view') @@ -117,20 +126,25 @@ def handle_tables(self) -> None: 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"] + 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}) def handle_query(self, sql: str, params: List[Any], limit: int) -> None: if not is_read_only_sql(sql): - self.send_json({"error": "only single read-only SELECT/WITH/PRAGMA statements are allowed"}, HTTPStatus.BAD_REQUEST) + self.send_json( + { + "error": "only single read-only SELECT/WITH/PRAGMA statements are allowed" + }, + HTTPStatus.BAD_REQUEST, + ) return limit = max(1, min(limit, max_limit)) start = time.perf_counter() @@ -150,7 +164,13 @@ def handle_query(self, sql: str, params: List[Any], limit: int) -> None: 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: +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) diff --git a/TraceLens/TraceIndex/sqlite_store.py b/TraceLens/TraceIndex/sqlite_store.py index a6cc2c5e7..6b066219d 100644 --- a/TraceLens/TraceIndex/sqlite_store.py +++ b/TraceLens/TraceIndex/sqlite_store.py @@ -15,6 +15,7 @@ from TraceLens.TraceIndex.store import TraceIndexStore from TraceLens.TraceIndex.utils import ( as_bool_int, + as_duration_us, as_float, as_int, as_text, @@ -49,8 +50,7 @@ def _connect(self, db_path: Path) -> sqlite3.Connection: return conn def init_schema(self) -> None: - self.conn.executescript( - """ + self.conn.executescript(""" CREATE TABLE IF NOT EXISTS traces ( id INTEGER PRIMARY KEY, root TEXT, @@ -177,8 +177,7 @@ def init_schema(self) -> None: text, tokenize='unicode61' ); - """ - ) + """) self.conn.commit() def upsert_trace(self, trace: TraceRecord) -> int: @@ -221,16 +220,26 @@ def upsert_trace(self, trace: TraceRecord) -> int: now, ), ) - row = self.conn.execute("SELECT id FROM traces WHERE path = ?", (trace.path,)).fetchone() + 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", [])) - self._import_kernel_summary_rows(trace_id, report.sheets.get("kernel_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", [])) + unified_summary = self._import_unified_rows( + trace_id, report.sheets.get("unified_perf_summary", []) + ) + self._import_kernel_summary_rows( + trace_id, report.sheets.get("kernel_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( """ @@ -258,7 +267,10 @@ def import_report(self, trace_id: int, report: TraceReport) -> None: trace_id, report.report_dir, utc_now(), - json.dumps([name for name, rows in report.sheets.items() if rows], sort_keys=True), + json.dumps( + [name for name, rows in report.sheets.items() if rows], + sort_keys=True, + ), ), ) self._insert_search(trace_id, "trace", [report.report_dir]) @@ -293,7 +305,9 @@ def execute_read_query( 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") + raise ValueError( + "only a single read-only SELECT/WITH/PRAGMA statement is allowed" + ) self.conn.execute("PRAGMA query_only=ON") rows = self.conn.execute(sql, params or ()).fetchmany(limit) return [dict(row) for row in rows] @@ -311,7 +325,9 @@ def _clear_trace_payload(self, trace_id: int) -> None: "trace_summary", ): 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,)) + 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) @@ -330,14 +346,36 @@ def _import_unified_rows( 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) + 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 + ) self.conn.execute( """ INSERT INTO unified_perf_rows( @@ -357,13 +395,66 @@ def _import_unified_rows( name, op_category, as_int(first_value(row, ["operation_count", "Count", "count"])), - as_float(first_value(row, ["Kernel Time (us)_sum", "Kernel Time (µs)_sum", "total_direct_kernel_time_sum", "total_subtree_kernel_time_sum"])), - 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)"])), + 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"])), @@ -375,11 +466,25 @@ def _import_unified_rows( as_bool_int(first_value(row, ["has_perf_model", "Has Perf Model"])), as_float(first_value(row, ["overlap_pct", "Overlap (%)"])), as_text(first_value(row, ["perf_params", "Perf Params"])), - as_text(first_value(row, ["kernel_details_summary", "trunc_kernel_details"])), + as_text( + first_value( + row, ["kernel_details_summary", "trunc_kernel_details"] + ) + ), json.dumps(row, sort_keys=True), ), ) - self._insert_search(trace_id, "op", [name, op_category, first_value(row, ["kernel_details_summary", "trunc_kernel_details"])]) + 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_kernel_summary_rows( @@ -388,11 +493,17 @@ def _import_kernel_summary_rows( rows: Sequence[Dict[str, str]], ) -> None: for row in rows: - kernel_name = as_text(first_value(row, ["Kernel name", "kernel_name", "name"])) + kernel_name = as_text( + first_value(row, ["Kernel name", "kernel_name", "name"]) + ) if not kernel_name: continue - parent_op_name = as_text(first_value(row, ["Parent cpu_op", "parent_op_name", "Launcher"])) - op_category = as_text(first_value(row, ["Parent op category", "op_category", "category"])) + parent_op_name = as_text( + first_value(row, ["Parent cpu_op", "parent_op_name", "Launcher"]) + ) + op_category = as_text( + first_value(row, ["Parent op category", "op_category", "category"]) + ) self.conn.execute( """ INSERT INTO kernel_summary( @@ -408,20 +519,82 @@ def _import_kernel_summary_rows( kernel_name, parent_op_name, op_category, - as_int(first_value(row, ["stream", "Stream"])), - as_int(first_value(row, ["Kernel duration (us)_count", "Kernel duration (µs)_count", "count"])), - as_float(first_value(row, ["Kernel duration (us)_sum", "Kernel duration (µs)_sum", "total_us"])), - as_float(first_value(row, ["Kernel duration (us)_mean", "Kernel duration (µs)_mean", "mean_us"])), - as_float(first_value(row, ["Kernel duration (us)_median", "Kernel duration (µs)_median", "median_us"])), - as_float(first_value(row, ["Kernel duration (us)_min", "Kernel duration (µs)_min", "min_us"])), - as_float(first_value(row, ["Kernel duration (us)_max", "Kernel duration (µs)_max", "max_us"])), - int("cijk" in kernel_name.lower() or "tensile" in kernel_name.lower()), + as_int(first_value(row, ["Kernel stream", "stream", "Stream"])), + as_int( + first_value( + row, + [ + "Kernel duration (us)_count", + "Kernel duration (µs)_count", + "count", + ], + ) + ), + as_float( + first_value( + row, + [ + "Kernel duration (us)_sum", + "Kernel duration (µs)_sum", + "total_us", + ], + ) + ), + as_float( + first_value( + row, + [ + "Kernel duration (us)_mean", + "Kernel duration (µs)_mean", + "mean_us", + ], + ) + ), + as_float( + first_value( + row, + [ + "Kernel duration (us)_median", + "Kernel duration (µs)_median", + "median_us", + ], + ) + ), + as_float( + first_value( + row, + [ + "Kernel duration (us)_min", + "Kernel duration (µs)_min", + "min_us", + ], + ) + ), + as_float( + first_value( + row, + [ + "Kernel duration (us)_max", + "Kernel duration (µs)_max", + "max_us", + ], + ) + ), + int( + "cijk" in kernel_name.lower() + or "tensile" in kernel_name.lower() + ), int("transpose" in kernel_name.lower()), - int("layout" in kernel_name.lower() or "permute" in kernel_name.lower()), + int( + "layout" in kernel_name.lower() + or "permute" in kernel_name.lower() + ), json.dumps(row, sort_keys=True), ), ) - self._insert_search(trace_id, "kernel", [kernel_name, parent_op_name, op_category]) + self._insert_search( + trace_id, "kernel", [kernel_name, parent_op_name, op_category] + ) def _import_category_rows( self, @@ -430,10 +603,24 @@ def _import_category_rows( ) -> str: top_categories = [] for row in rows: - category = as_text(first_value(row, ["op category", "category", "Categories", "name"])) + category = as_text( + first_value(row, ["op category", "category", "Categories", "name"]) + ) if not category: continue - kernel_time = as_float(first_value(row, ["Kernel Time (us)_sum", "Kernel Time (µs)_sum", "total_direct_kernel_time_sum", "total_subtree_kernel_time_sum"])) + 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( @@ -446,11 +633,18 @@ def _import_category_rows( category, as_int(first_value(row, ["operation_count", "Count", "count"])), kernel_time, - as_float(first_value(row, ["Percentage (%)", "percent", "Percent of total 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}) + 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) diff --git a/TraceLens/TraceIndex/utils.py b/TraceLens/TraceIndex/utils.py index e9fe76a41..83498e04b 100644 --- a/TraceLens/TraceIndex/utils.py +++ b/TraceLens/TraceIndex/utils.py @@ -66,6 +66,20 @@ def as_int(value: Any) -> Optional[int]: 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: if isinstance(value, bool): return int(value) diff --git a/tests/test_trace_index.py b/tests/test_trace_index.py index a950ea342..c292852e7 100644 --- a/tests/test_trace_index.py +++ b/tests/test_trace_index.py @@ -6,6 +6,7 @@ import csv import json +from pathlib import Path import pytest @@ -15,10 +16,18 @@ scan_traces, search_index, ) -from TraceLens.TraceIndex.importer import import_report_dir as import_report_dir_with_store +from TraceLens.TraceIndex.importer import ( + import_report_dir as import_report_dir_with_store, +) from TraceLens.TraceIndex.scanner import scan_traces as scan_traces_with_store from TraceLens.TraceIndex.sqlite_store import SQLiteTraceIndexStore +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) @@ -82,7 +91,9 @@ def test_trace_index_scan_import_and_search(tmp_path): ], ) - trace_id = import_report_dir(db_path, report_dir, trace_path=trace_path, root=trace_root) + trace_id = import_report_dir( + db_path, report_dir, trace_path=trace_path, root=trace_root + ) assert trace_id == 1 rows = execute_read_query( @@ -138,3 +149,53 @@ def test_trace_index_store_boundary_supports_scan_import_and_search(tmp_path): assert hits[0].kind == "op" finally: store.close() + + +@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): + db_path = tmp_path / "trace_index.sqlite" + trace_id = import_report_dir(db_path, 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 kernel_name, stream, total_duration_us FROM kernel_summary " + "WHERE kernel_name LIKE 'Cijk%' LIMIT 1", + ) + assert kernels + assert kernels[0]["stream"] == 0 + assert kernels[0]["total_duration_us"] > 0 + + 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): + db_path = tmp_path / "trace_index.sqlite" + import_report_dir(db_path, 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 From 5fd3ce6209b3bc3a5e77fa7d45050e879fdf8a7a Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Tue, 25 Aug 2026 14:54:26 -0400 Subject: [PATCH 04/18] Move TraceIndex docs into the ROCm how-to tree. Co-authored-by: Cursor --- README.md | 4 +- docs/TraceIndex.md | 124 ---------------- .../generate-perf-report-pytorch-inference.md | 1 + docs/how-to/generate-perf-report-pytorch.md | 2 + docs/how-to/generate-reports.md | 5 +- docs/how-to/sdk-analysis.md | 1 + docs/how-to/trace-index.md | 138 ++++++++++++++++++ docs/index.rst | 1 + docs/reference/api-reference.md | 19 +++ docs/sphinx/_toc.yml.in | 2 + 10 files changed, 170 insertions(+), 127 deletions(-) delete mode 100644 docs/TraceIndex.md create mode 100644 docs/how-to/trace-index.md diff --git a/README.md b/README.md index 39855919f..e0c278914 100755 --- a/README.md +++ b/README.md @@ -81,7 +81,7 @@ TraceLens_compare_perf_reports_pytorch \ -o comparison.xlsx ``` -Index a directory of traces and existing TraceLens CSV reports (see [TraceIndex](docs/TraceIndex.md)): +Index a directory of traces and existing TraceLens CSV reports (see [Index a corpus of traces](docs/how-to/trace-index.md)): ```bash TraceLens_trace_index --backend sqlite --db trace_index.sqlite scan --root /path/to/traces @@ -149,7 +149,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/TraceIndex.md](docs/TraceIndex.md) | +| TraceIndex | [docs/how-to/trace-index.md](docs/how-to/trace-index.md) | --- diff --git a/docs/TraceIndex.md b/docs/TraceIndex.md deleted file mode 100644 index 9b134acaa..000000000 --- a/docs/TraceIndex.md +++ /dev/null @@ -1,124 +0,0 @@ - - -# TraceIndex - -TraceIndex builds a queryable catalog of trace files and imports key TraceLens -CSV report tables so a corpus can be searched without reopening every raw -trace. It is useful when you have many profiler captures and want to answer -questions like: - -- Which traces contain GEMM, SDPA, convolution, collective, or short-kernel heavy workloads? -- Which traces contain a specific backend kernel name? -- Which trace should I open next in TraceLens or Perfetto? - -TraceIndex does not replace raw traces. It stores searchable summaries and paths -back to the source traces. The first backend is SQLite because it needs no -service to run and works well for local or single-team workflows. The scanner, -importer, and storage interface are separated so other backends can be added by -implementing the same store contract. - -## Quick Start - -Catalog trace-like files under a directory: - -```bash -TraceLens_trace_index --backend sqlite --db trace_index.sqlite scan --root /path/to/traces -``` - -Generate a TraceLens report for one PyTorch trace and import it into the index: - -```bash -TraceLens_trace_index --backend sqlite --db trace_index.sqlite build \ - --trace-path /path/to/traces/rank0_trace.json.gz \ - --report-dir ./trace_index_reports/rank0 -``` - -If you already have a TraceLens CSV report directory, import it directly: - -```bash -TraceLens_trace_index --backend sqlite --db trace_index.sqlite import-report \ - --trace-path /path/to/traces/rank0_trace.json.gz \ - --report-dir ./rank0_perf_report_csvs -``` - -Search the full-text index: - -```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 read-only SQLite query: - -```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" -``` - -## Imported Tables - -TraceIndex imports the stable report tables that are most useful for corpus -search: - -| Table | Contents | -|---|---| -| `traces` | One row per trace-like file or imported report directory | -| `report_imports` | Import history for TraceLens CSV report directories | -| `unified_perf_rows` | Rows from `unified_perf_summary.csv` | -| `kernel_summary` | Rows from `kernel_summary.csv`, including basic kernel flags | -| `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 | - -## Backend Model - -TraceIndex separates workflow code from persistence: - -| Module | Responsibility | -|---|---| -| `models.py` | Backend-neutral trace, report, and search-result objects | -| `store.py` | `TraceIndexStore` interface for persistence/search backends | -| `scanner.py` | File discovery and trace metadata extraction | -| `importer.py` | TraceLens CSV report loading and import workflow | -| `sqlite_store.py` | SQLite schema, inserts, full-text search, and read-only SQL | -| `cli.py` | User-facing commands that call scanner/importer/store APIs | - -To add another backend, implement `TraceIndexStore` and wire it into the CLI -backend selector. Backend-neutral commands such as `scan`, `build`, -`import-report`, and `search` should not need to change. - -## SQLite Query Server - -For notebook or browser workflows, serve read-only SQL access locally: - -```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}` - -Only single SQLite `SELECT`, `WITH`, or `PRAGMA` statements are accepted. The -server is read-only but does not implement authentication, 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 import_report_dir, scan_traces, search_index - -db = Path("trace_index.sqlite") -scan_traces(db, Path("/path/to/traces")) -import_report_dir(db, Path("rank0_perf_report_csvs"), trace_path=Path("rank0_trace.json.gz")) -rows = search_index(db, "Cijk", limit=20) -``` 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..158ff4171 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,9 @@ 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, you can catalog them for corpus search. +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..4bafcc21b --- /dev/null +++ b/docs/how-to/trace-index.md @@ -0,0 +1,138 @@ + + + +# Index a corpus of traces in TraceLens +```{meta} +:description: Learn how to catalog profiler traces and TraceLens CSV reports into a searchable SQLite index without reopening every raw trace. +:keywords: TraceLens, TraceIndex, corpus search, 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 catalog is a SQLite file. It needs no extra service and works for local or +single-team workflows. + +## Before you begin + +- TraceLens installed (see [Install TraceLens](../install/install.md)). +- A directory of profiler traces, or an existing TraceLens CSV report directory + (for example from + [Generate a PyTorch performance report](./generate-perf-report-pytorch.md)). + +## Catalog traces + +Scan a directory for trace-like files and record them in the catalog: + +```bash +TraceLens_trace_index --backend sqlite --db trace_index.sqlite scan --root /path/to/traces +``` + +This pass peeks at file headers. It doesn't parse whole traces or run +TraceLens analysis. + +## Import a report + +If you already have a TraceLens CSV report directory, import it. This is the +usual ingest path, including for inference reports you generated separately: + +```bash +TraceLens_trace_index --backend sqlite --db trace_index.sqlite import-report \ + --trace-path /path/to/traces/rank0_trace.json.gz \ + --report-dir ./rank0_perf_report_csvs +``` + +`--trace-path` is optional. If you omit it, the report directory is cataloged +as its own row. + +To generate a training PyTorch CSV report and import it in one step: + +```bash +TraceLens_trace_index --backend sqlite --db trace_index.sqlite build \ + --trace-path /path/to/traces/rank0_trace.json.gz \ + --report-dir ./trace_index_reports/rank0 +``` + +`build` calls the training PyTorch report generator. For inference, rocprof, or +pftrace reports, generate the CSV directory with the matching report command, +then use `import-report`. + +## 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 table lists the SQLite tables filled on import. + +| Table | Contents | +|---|---| +| `traces` | One row per trace-like file or imported report directory | +| `report_imports` | Import history for TraceLens CSV report directories | +| `unified_perf_rows` | Rows from `unified_perf_summary.csv` | +| `kernel_summary` | Rows from `kernel_summary.csv`, including Tensile and layout flags | +| `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 | + +## 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 import_report_dir, scan_traces, search_index + +db = Path("trace_index.sqlite") +scan_traces(db, Path("/path/to/traces")) +import_report_dir(db, Path("rank0_perf_report_csvs"), trace_path=Path("rank0_trace.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..f1c932d09 100644 --- a/docs/reference/api-reference.md +++ b/docs/reference/api-reference.md @@ -177,6 +177,24 @@ 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 SQLite +index. 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 | +|---------|-------------| +| `scan --root PATH` | Record trace-like files under a directory. Doesn't parse whole traces. | +| `import-report --report-dir DIR` | Import an existing CSV report directory. `--trace-path` is optional. | +| `build --trace-path PATH` | Generate a training PyTorch CSV report, then import it. | +| `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 +211,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 and CSV reports into SQLite for search. | [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 From 158098d341ae2463940d63300aee9e0bf8c58c2d Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Tue, 25 Aug 2026 15:52:56 -0400 Subject: [PATCH 05/18] Replace TraceIndex scan/import-report with append and batch build. Co-authored-by: Cursor --- README.md | 10 ++- TraceLens/TraceIndex/__init__.py | 4 + TraceLens/TraceIndex/cli.py | 128 ++++++++++++++++++------------- TraceLens/TraceIndex/core.py | 50 +++++++++++- TraceLens/TraceIndex/importer.py | 77 ++++++++++++++++++- TraceLens/TraceIndex/models.py | 2 +- TraceLens/TraceIndex/utils.py | 31 ++++++++ docs/how-to/generate-reports.md | 3 +- docs/how-to/trace-index.md | 74 ++++++++++-------- docs/reference/api-reference.md | 13 ++-- docs/what-is-tracelens.md | 4 + tests/test_trace_index.py | 87 ++++++++++++++++++--- 12 files changed, 369 insertions(+), 114 deletions(-) diff --git a/README.md b/README.md index e0c278914..be1f140fa 100755 --- a/README.md +++ b/README.md @@ -81,12 +81,14 @@ TraceLens_compare_perf_reports_pytorch \ -o comparison.xlsx ``` -Index a directory of traces and existing TraceLens CSV reports (see [Index a corpus of traces](docs/how-to/trace-index.md)): +Index traces and existing TraceLens CSV reports (see [Index a corpus of traces](docs/how-to/trace-index.md)): ```bash -TraceLens_trace_index --backend sqlite --db trace_index.sqlite scan --root /path/to/traces -TraceLens_trace_index --backend sqlite --db trace_index.sqlite import-report --report-dir path/to/perf_report_csvs -TraceLens_trace_index --backend sqlite --db trace_index.sqlite search Cijk +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)): diff --git a/TraceLens/TraceIndex/__init__.py b/TraceLens/TraceIndex/__init__.py index adbc72b21..5bf2aec5f 100644 --- a/TraceLens/TraceIndex/__init__.py +++ b/TraceLens/TraceIndex/__init__.py @@ -7,6 +7,8 @@ """Trace corpus indexing helpers.""" from .core import ( + append_trace, + build_traces, execute_read_query, generate_report_and_import, import_report_dir, @@ -23,6 +25,8 @@ "SearchHit", "TraceRecord", "TraceReport", + "append_trace", + "build_traces", "execute_read_query", "generate_report_and_import", "import_report_dir", diff --git a/TraceLens/TraceIndex/cli.py b/TraceLens/TraceIndex/cli.py index e501e3b21..dfd193836 100644 --- a/TraceLens/TraceIndex/cli.py +++ b/TraceLens/TraceIndex/cli.py @@ -11,9 +11,13 @@ from pathlib import Path from typing import List, Optional -from TraceLens.TraceIndex.importer import generate_report_and_import, import_report_dir -from TraceLens.TraceIndex.scanner import scan_traces +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") @@ -28,43 +32,30 @@ def create_store(args: argparse.Namespace): return SQLiteTraceIndexStore(args.db) -def scan_cmd(args: argparse.Namespace) -> int: +def append_cmd(args: argparse.Namespace) -> int: store = create_store(args) try: - count = scan_traces( + 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, - root=args.root, - peek_mb=args.peek_mb, - compute_md5=args.compute_md5, - ) - print_json( - { - "backend": args.backend, - "db": args.db, - "root": args.root, - "candidate_traces": count, - } - ) - return 0 - finally: - store.close() - - -def import_report_cmd(args: argparse.Namespace) -> int: - store = create_store(args) - try: - trace_id = import_report_dir( - store, - report_dir=args.report_dir, 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, - "report_dir": args.report_dir, + "trace_path": args.trace_path, + "report_dir": report_dir, + "generated_report": generated, } ) return 0 @@ -73,12 +64,17 @@ def import_report_cmd(args: argparse.Namespace) -> int: 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: - trace_id = generate_report_and_import( + result = build_traces( store, - trace_path=args.trace_path, - report_dir=args.report_dir, + trace_paths, + report_root=args.report_root, root=args.root, force=args.force, enable_pseudo_ops=args.enable_pseudo_ops, @@ -87,11 +83,11 @@ def build_cmd(args: argparse.Namespace) -> int: { "backend": args.backend, "db": args.db, - "trace_id": trace_id, - "trace_path": args.trace_path, + "imported": result["imported"], + "failed": result["failed"], } ) - return 0 + return 1 if result["failed"] else 0 finally: store.close() @@ -134,9 +130,25 @@ def serve_cmd(args: argparse.Namespace) -> int: 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 TraceLens reports." + description="Build and query a TraceIndex catalog of traces." ) parser.add_argument( "--backend", @@ -147,30 +159,38 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--db", type=Path, default=DEFAULT_DB, help="SQLite DB path") sub = parser.add_subparsers(dest="command") - scan = sub.add_parser("scan", help="Catalog trace-like files under a root") - scan.add_argument("--root", type=Path, required=True) - scan.add_argument("--peek-mb", type=int, default=2) - scan.add_argument("--compute-md5", action="store_true") - scan.set_defaults(func=scan_cmd) - - import_report = sub.add_parser( - "import-report", - help="Import an existing TraceLens CSV report directory", + 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.", ) - import_report.add_argument("--report-dir", type=Path, required=True) - import_report.add_argument("--trace-path", type=Path, default=None) - import_report.add_argument("--root", type=Path, default=None) - import_report.set_defaults(func=import_report_cmd) + add_generate_args(append) + append.set_defaults(func=append_cmd) build = sub.add_parser( "build", - help="Generate a TraceLens CSV report for one trace, then import it", + 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", ) - build.add_argument("--trace-path", type=Path, required=True) - build.add_argument("--report-dir", type=Path, default=None) - build.add_argument("--root", type=Path, default=None) - build.add_argument("--force", action="store_true") - build.add_argument("--enable-pseudo-ops", action="store_true") + add_generate_args(build) build.set_defaults(func=build_cmd) search = sub.add_parser("search", help="Full-text search indexed traces") diff --git a/TraceLens/TraceIndex/core.py b/TraceLens/TraceIndex/core.py index 401fcdadd..db592ea6d 100644 --- a/TraceLens/TraceIndex/core.py +++ b/TraceLens/TraceIndex/core.py @@ -10,9 +10,9 @@ 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.importer import ( import_report_dir as import_report_dir_with_store, ) from TraceLens.TraceIndex.scanner import scan_traces as scan_traces_with_store @@ -55,6 +55,52 @@ def import_report_dir( store.close() +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, diff --git a/TraceLens/TraceIndex/importer.py b/TraceLens/TraceIndex/importer.py index 315025abd..d20230e40 100644 --- a/TraceLens/TraceIndex/importer.py +++ b/TraceLens/TraceIndex/importer.py @@ -8,13 +8,15 @@ import re from pathlib import Path -from typing import Optional +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", "kernel_summary", @@ -69,6 +71,76 @@ def import_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, @@ -78,8 +150,7 @@ def generate_report_and_import( enable_pseudo_ops: bool = False, ) -> int: if report_dir is None: - safe_name = re.sub(r"[^A-Za-z0-9_.-]+", "_", trace_path.name) - report_dir = Path("trace_index_reports") / safe_name + 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(): diff --git a/TraceLens/TraceIndex/models.py b/TraceLens/TraceIndex/models.py index eebf7b846..355d4756e 100644 --- a/TraceLens/TraceIndex/models.py +++ b/TraceLens/TraceIndex/models.py @@ -6,7 +6,7 @@ """Backend-neutral data objects for TraceIndex.""" -from typing import Any, Dict, List, NamedTuple, Optional +from typing import Dict, List, NamedTuple, Optional class TraceRecord(NamedTuple): diff --git a/TraceLens/TraceIndex/utils.py b/TraceLens/TraceIndex/utils.py index 83498e04b..171b02ccc 100644 --- a/TraceLens/TraceIndex/utils.py +++ b/TraceLens/TraceIndex/utils.py @@ -91,3 +91,34 @@ def as_bool_int(value: Any) -> int: def search_text(*parts: Any) -> str: return " ".join(str(part) for part in parts if part not in (None, "", "nan", "NaN")) + + +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-reports.md b/docs/how-to/generate-reports.md index 158ff4171..45d8034c8 100644 --- a/docs/how-to/generate-reports.md +++ b/docs/how-to/generate-reports.md @@ -19,7 +19,8 @@ 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, you can catalog them for corpus search. +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 diff --git a/docs/how-to/trace-index.md b/docs/how-to/trace-index.md index 4bafcc21b..5d01513fa 100644 --- a/docs/how-to/trace-index.md +++ b/docs/how-to/trace-index.md @@ -7,8 +7,8 @@ See LICENSE for license information. # Index a corpus of traces in TraceLens ```{meta} -:description: Learn how to catalog profiler traces and TraceLens CSV reports into a searchable SQLite index without reopening every raw trace. -:keywords: TraceLens, TraceIndex, corpus search, SQLite, unified_perf_summary, kernel summary, full-text search, performance report +: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 @@ -16,52 +16,55 @@ 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 catalog is a SQLite file. It needs no extra service and works for local or -single-team workflows. +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)). -- A directory of profiler traces, or an existing TraceLens CSV report directory +- Profiler traces, and optionally existing TraceLens CSV report directories (for example from [Generate a PyTorch performance report](./generate-perf-report-pytorch.md)). -## Catalog traces +## Append a trace -Scan a directory for trace-like files and record them in the catalog: +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 --backend sqlite --db trace_index.sqlite scan --root /path/to/traces +TraceLens_trace_index --db trace_index.sqlite append \ + --trace-path /path/to/traces/rank0_trace.json.gz \ + --report-dir ./rank0_perf_report_csvs ``` -This pass peeks at file headers. It doesn't parse whole traces or run -TraceLens analysis. - -## Import a report - -If you already have a TraceLens CSV report directory, import it. This is the -usual ingest path, including for inference reports you generated separately: +If you omit `--report-dir`, TraceIndex generates a training PyTorch CSV report +and then imports it: ```bash -TraceLens_trace_index --backend sqlite --db trace_index.sqlite import-report \ - --trace-path /path/to/traces/rank0_trace.json.gz \ - --report-dir ./rank0_perf_report_csvs +TraceLens_trace_index --db trace_index.sqlite append \ + --trace-path /path/to/traces/rank0_trace.json.gz ``` -`--trace-path` is optional. If you omit it, the report directory is cataloged -as its own row. +## Build a catalog from a list of traces -To generate a training PyTorch CSV report and import it in one step: +`--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 --backend sqlite --db trace_index.sqlite build \ - --trace-path /path/to/traces/rank0_trace.json.gz \ - --report-dir ./trace_index_reports/rank0 +TraceLens_trace_index --db trace_index.sqlite build \ + --traces-file traces.txt \ + --report-root ./trace_index_reports ``` -`build` calls the training PyTorch report generator. For inference, rocprof, or -pftrace reports, generate the CSV directory with the matching report command, -then use `import-report`. +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 @@ -82,11 +85,14 @@ TraceLens_trace_index --backend sqlite --db trace_index.sqlite sqlite-sql \ ## What the catalog stores -The following table lists the SQLite tables filled on import. +The following tables are the catalog schema. There are seven 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. | Table | Contents | |---|---| -| `traces` | One row per trace-like file or imported report directory | +| `traces` | One row per indexed trace | | `report_imports` | Import history for TraceLens CSV report directories | | `unified_perf_rows` | Rows from `unified_perf_summary.csv` | | `kernel_summary` | Rows from `kernel_summary.csv`, including Tensile and layout flags | @@ -120,11 +126,15 @@ behind your own access control. ```python from pathlib import Path -from TraceLens.TraceIndex import import_report_dir, scan_traces, search_index +from TraceLens.TraceIndex import append_trace, build_traces, search_index db = Path("trace_index.sqlite") -scan_traces(db, Path("/path/to/traces")) -import_report_dir(db, Path("rank0_perf_report_csvs"), trace_path=Path("rank0_trace.json.gz")) +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) ``` diff --git a/docs/reference/api-reference.md b/docs/reference/api-reference.md index f1c932d09..39df69a67 100644 --- a/docs/reference/api-reference.md +++ b/docs/reference/api-reference.md @@ -179,18 +179,17 @@ Split an inference trace into per-iteration or per-phase sub-traces. ### TraceLens_trace_index -Catalog profiler traces and TraceLens CSV reports into a searchable SQLite -index. See [Index a corpus of traces](../how-to/trace-index.md) for the full -workflow. +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 | |---------|-------------| -| `scan --root PATH` | Record trace-like files under a directory. Doesn't parse whole traces. | -| `import-report --report-dir DIR` | Import an existing CSV report directory. `--trace-path` is optional. | -| `build --trace-path PATH` | Generate a training PyTorch CSV report, then import it. | +| `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. | @@ -211,7 +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 and CSV reports into SQLite for search. | [Index a corpus of traces](../how-to/trace-index.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/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/tests/test_trace_index.py b/tests/test_trace_index.py index c292852e7..73569aef8 100644 --- a/tests/test_trace_index.py +++ b/tests/test_trace_index.py @@ -11,16 +11,18 @@ import pytest from TraceLens.TraceIndex.core import ( + append_trace, execute_read_query, import_report_dir, - scan_traces, search_index, ) +from TraceLens.TraceIndex.cli import main as trace_index_main from TraceLens.TraceIndex.importer import ( + build_traces as build_traces_with_store, import_report_dir as import_report_dir_with_store, ) -from TraceLens.TraceIndex.scanner import scan_traces as scan_traces_with_store from TraceLens.TraceIndex.sqlite_store import SQLiteTraceIndexStore +from TraceLens.TraceIndex.utils import collect_trace_paths, read_traces_file FIXTURES = Path(__file__).resolve().parent / "traces" TRAINING_REPORT_DIR = ( @@ -37,15 +39,13 @@ def write_csv(path, rows): writer.writerows(rows) -def test_trace_index_scan_import_and_search(tmp_path): +def test_trace_index_append_from_report_and_search(tmp_path): db_path = tmp_path / "trace_index.sqlite" trace_root = tmp_path / "traces" trace_path = trace_root / "model_a" / "rank0_trace.json" trace_path.parent.mkdir(parents=True) trace_path.write_text(json.dumps({"traceEvents": []}), encoding="utf-8") - assert scan_traces(db_path, trace_root) == 1 - report_dir = tmp_path / "reports" / "trace_1" write_csv( report_dir / "unified_perf_summary.csv", @@ -91,9 +91,7 @@ def test_trace_index_scan_import_and_search(tmp_path): ], ) - trace_id = import_report_dir( - db_path, report_dir, trace_path=trace_path, root=trace_root - ) + trace_id = append_trace(db_path, trace_path, report_dir=report_dir, root=trace_root) assert trace_id == 1 rows = execute_read_query( @@ -115,7 +113,7 @@ def test_trace_index_rejects_write_sql(tmp_path): execute_read_query(db_path, "DELETE FROM traces") -def test_trace_index_store_boundary_supports_scan_import_and_search(tmp_path): +def test_trace_index_store_boundary_supports_append_and_search(tmp_path): db_path = tmp_path / "trace_index.sqlite" trace_root = tmp_path / "traces" trace_path = trace_root / "rank0_trace.json" @@ -137,7 +135,6 @@ def test_trace_index_store_boundary_supports_scan_import_and_search(tmp_path): store = SQLiteTraceIndexStore(db_path) try: - assert scan_traces_with_store(store, trace_root) == 1 trace_id = import_report_dir_with_store( store, report_dir, @@ -199,3 +196,73 @@ def test_import_real_inference_report_converts_category_kernel_time_ms(tmp_path) ) assert rows assert rows[0]["kernel_time_sum_us"] > 1000 + + +def test_read_traces_file_skips_comments_and_blanks(tmp_path): + 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): + db_path = tmp_path / "trace_index.sqlite" + trace_path = tmp_path / "rank0_trace.json" + trace_path.write_text(json.dumps({"traceEvents": []}), encoding="utf-8") + report_dir = tmp_path / "report" + write_csv( + report_dir / "unified_perf_summary.csv", + [ + { + "name": "aten::mm", + "op category": "GEMM", + "operation_count": "1", + "Kernel Time (us)_sum": "10.0", + } + ], + ) + + 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") + assert rows[0]["name"] == "aten::mm" + + +def test_build_traces_continues_after_failure(tmp_path): + 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"] From 7b50c4f6827cfd28fda078a531fa37ea4a9377f4 Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Tue, 25 Aug 2026 16:14:54 -0400 Subject: [PATCH 06/18] Docs: add TraceIndex catalog schema diagram. Co-authored-by: Cursor --- docs/how-to/trace-index.md | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/docs/how-to/trace-index.md b/docs/how-to/trace-index.md index 5d01513fa..9db96c51d 100644 --- a/docs/how-to/trace-index.md +++ b/docs/how-to/trace-index.md @@ -90,6 +90,20 @@ 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`. + +```mermaid +erDiagram + traces ||--o{ report_imports : "trace_id" + traces ||--o{ unified_perf_rows : "trace_id" + traces ||--o{ kernel_summary : "trace_id" + traces ||--o{ op_category_rows : "trace_id" + traces ||--o{ gpu_timeline_rows : "trace_id" + traces ||--o| trace_summary : "trace_id" + traces ||--o{ trace_search_FTS5 : "trace_id" +``` + | Table | Contents | |---|---| | `traces` | One row per indexed trace | From 2316779dc8759c740b42a5f626e023b88c3c6477 Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Tue, 25 Aug 2026 16:46:13 -0400 Subject: [PATCH 07/18] TraceIndex: explode kernels per op and add GEMM/SDPA/conv shape tables Parse perf_params and kernel_details_summary from Python repr into JSON so they are queryable. Kernel rows now come from each unified op's kernel_details_summary with a unified_row_id FK, replacing the flat kernel_summary.csv import. Add gemm_perf, sdpa_perf, and conv_perf satellite tables keyed on unified_row_id for shape-based SQL. Update the how-to schema diagram and tests accordingly. Co-authored-by: Cursor --- TraceLens/TraceIndex/importer.py | 1 - TraceLens/TraceIndex/sqlite_store.py | 355 +++++++++++++++++++-------- TraceLens/TraceIndex/utils.py | 101 +++++++- docs/how-to/trace-index.md | 24 +- tests/test_trace_index.py | 138 +++++++++-- 5 files changed, 489 insertions(+), 130 deletions(-) diff --git a/TraceLens/TraceIndex/importer.py b/TraceLens/TraceIndex/importer.py index d20230e40..1ebb20730 100644 --- a/TraceLens/TraceIndex/importer.py +++ b/TraceLens/TraceIndex/importer.py @@ -19,7 +19,6 @@ REPORT_SHEETS = ( "unified_perf_summary", - "kernel_summary", "ops_summary_by_category", "gpu_timeline", ) diff --git a/TraceLens/TraceIndex/sqlite_store.py b/TraceLens/TraceIndex/sqlite_store.py index 6b066219d..afe570916 100644 --- a/TraceLens/TraceIndex/sqlite_store.py +++ b/TraceLens/TraceIndex/sqlite_store.py @@ -18,9 +18,13 @@ 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, ) @@ -118,6 +122,7 @@ def init_schema(self) -> None: CREATE TABLE IF NOT EXISTS kernel_summary ( 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, @@ -128,16 +133,68 @@ def init_schema(self) -> None: 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, - raw_row_json TEXT + details_json TEXT ); CREATE INDEX IF NOT EXISTS idx_trace_index_kernel_trace ON kernel_summary(trace_id); + CREATE INDEX IF NOT EXISTS idx_trace_index_kernel_unified ON kernel_summary(unified_row_id); CREATE INDEX IF NOT EXISTS idx_trace_index_kernel_name ON kernel_summary(kernel_name); CREATE INDEX IF NOT EXISTS idx_trace_index_kernel_tensile ON kernel_summary(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, @@ -231,9 +288,6 @@ def import_report(self, trace_id: int, report: TraceReport) -> None: unified_summary = self._import_unified_rows( trace_id, report.sheets.get("unified_perf_summary", []) ) - self._import_kernel_summary_rows( - trace_id, report.sheets.get("kernel_summary", []) - ) top_categories_json = self._import_category_rows( trace_id, report.sheets.get("ops_summary_by_category", []) ) @@ -318,11 +372,14 @@ def close(self) -> None: def _clear_trace_payload(self, trace_id: int) -> None: for table in ( "report_imports", - "unified_perf_rows", + "gemm_perf", + "sdpa_perf", + "conv_perf", "kernel_summary", "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( @@ -376,7 +433,11 @@ def _import_unified_rows( max_sdpa_tflops = max( max_sdpa_tflops or tflops_for_summary, tflops_for_summary ) - self.conn.execute( + 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, @@ -465,15 +526,22 @@ def _import_unified_rows( 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 (%)"])), - as_text(first_value(row, ["perf_params", "Perf Params"])), - as_text( - first_value( - row, ["kernel_details_summary", "trunc_kernel_details"] - ) - ), - json.dumps(row, sort_keys=True), + 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", @@ -487,115 +555,198 @@ def _import_unified_rows( ) return {"max_gemm_tflops": max_gemm_tflops, "max_sdpa_tflops": max_sdpa_tflops} - def _import_kernel_summary_rows( + def _import_kernels_from_details( self, trace_id: int, - rows: Sequence[Dict[str, str]], + unified_row_id: int, + parent_op_name: Optional[str], + op_category: Optional[str], + kernel_details: Any, ) -> None: - for row in rows: - kernel_name = as_text( - first_value(row, ["Kernel name", "kernel_name", "name"]) - ) + 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 - parent_op_name = as_text( - first_value(row, ["Parent cpu_op", "parent_op_name", "Launcher"]) - ) - op_category = as_text( - first_value(row, ["Parent op category", "op_category", "category"]) - ) + library, is_tensile, is_transpose, is_layout = kernel_flags(kernel_name) self.conn.execute( """ INSERT INTO kernel_summary( - trace_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, is_tensile, is_transpose, - is_layout_conversion, raw_row_json + 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 (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( trace_id, + unified_row_id, kernel_name, parent_op_name, op_category, - as_int(first_value(row, ["Kernel stream", "stream", "Stream"])), - as_int( - first_value( - row, - [ - "Kernel duration (us)_count", - "Kernel duration (µs)_count", - "count", - ], - ) - ), - as_float( - first_value( - row, - [ - "Kernel duration (us)_sum", - "Kernel duration (µs)_sum", - "total_us", - ], - ) - ), - as_float( - first_value( - row, - [ - "Kernel duration (us)_mean", - "Kernel duration (µs)_mean", - "mean_us", - ], - ) - ), - as_float( - first_value( - row, - [ - "Kernel duration (us)_median", - "Kernel duration (µs)_median", - "median_us", - ], - ) - ), - as_float( - first_value( - row, - [ - "Kernel duration (us)_min", - "Kernel duration (µs)_min", - "min_us", - ], - ) - ), - as_float( - first_value( - row, - [ - "Kernel duration (us)_max", - "Kernel duration (µs)_max", - "max_us", - ], - ) - ), - int( - "cijk" in kernel_name.lower() - or "tensile" in kernel_name.lower() - ), - int("transpose" in kernel_name.lower()), - int( - "layout" in kernel_name.lower() - or "permute" in kernel_name.lower() - ), - json.dumps(row, sort_keys=True), + 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, parent_op_name, op_category] + 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, diff --git a/TraceLens/TraceIndex/utils.py b/TraceLens/TraceIndex/utils.py index 171b02ccc..3394fc44b 100644 --- a/TraceLens/TraceIndex/utils.py +++ b/TraceLens/TraceIndex/utils.py @@ -6,10 +6,15 @@ """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 +from typing import Any, Dict, List, Optional, Sequence, Tuple def utc_now() -> str: @@ -27,6 +32,19 @@ def rel_to(path: Path, root: Path) -> str: 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 [] @@ -81,18 +99,95 @@ def as_duration_us( 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 0 + return None text = str(value).strip().lower() - return int(text in {"1", "true", "yes", "y"}) + 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) -> Tuple[Optional[str], int, int, int]: + low = name.lower() + is_tensile = int("cijk" in low or "tensile" in low) + is_transpose = int("transpose" in low or "permute" in low) + is_layout = int( + is_transpose + or "contiguous" in low + or "copy" in low + or "cast" in low + or "convert" in low + ) + library = None + if is_tensile: + library = "Tensile" + elif "triton" in low: + library = "Triton" + elif "ck" in low or "composable" in low: + library = "CK" + elif "nccl" in low or "rccl" in low: + library = "RCCL/NCCL" + 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] = [] diff --git a/docs/how-to/trace-index.md b/docs/how-to/trace-index.md index 9db96c51d..504f28222 100644 --- a/docs/how-to/trace-index.md +++ b/docs/how-to/trace-index.md @@ -85,13 +85,14 @@ TraceLens_trace_index --backend sqlite --db trace_index.sqlite sqlite-sql \ ## What the catalog stores -The following tables are the catalog schema. There are seven relational tables +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`. +`traces`. Kernel rows and GEMM / SDPA / convolution satellites also point at +the `unified_perf_rows` row they came from. ```mermaid erDiagram @@ -101,20 +102,35 @@ erDiagram 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{ kernel_summary : "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` | -| `kernel_summary` | Rows from `kernel_summary.csv`, including Tensile and layout flags | +| `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 | +| `kernel_summary` | 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`. + ## Serve read-only SQL For notebook or browser workflows, serve the SQLite catalog over HTTP: diff --git a/tests/test_trace_index.py b/tests/test_trace_index.py index 73569aef8..7ee917100 100644 --- a/tests/test_trace_index.py +++ b/tests/test_trace_index.py @@ -22,7 +22,12 @@ import_report_dir as import_report_dir_with_store, ) from TraceLens.TraceIndex.sqlite_store import SQLiteTraceIndexStore -from TraceLens.TraceIndex.utils import collect_trace_paths, read_traces_file +from TraceLens.TraceIndex.utils import ( + collect_trace_paths, + parse_repr, + read_traces_file, + to_json, +) FIXTURES = Path(__file__).resolve().parent / "traces" TRAINING_REPORT_DIR = ( @@ -39,6 +44,15 @@ def write_csv(path, rows): writer.writerows(rows) +def test_parse_repr_strips_numpy_scalars_and_to_json(): + 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): db_path = tmp_path / "trace_index.sqlite" trace_root = tmp_path / "traces" @@ -56,20 +70,41 @@ def test_trace_index_append_from_report_and_search(tmp_path): "operation_count": "2", "Kernel Time (us)_sum": "123.5", "TFLOPS/s_mean": "98.1", - "kernel_details_summary": "Cijk_test_kernel", - } - ], - ) - write_csv( - report_dir / "kernel_summary.csv", - [ + "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}]" + ), + }, { - "Kernel name": "Cijk_test_kernel", - "Parent cpu_op": "aten::mm", - "Parent op category": "GEMM", - "Kernel duration (us)_count": "2", - "Kernel duration (us)_sum": "123.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( @@ -96,11 +131,45 @@ def test_trace_index_append_from_report_and_search(tmp_path): rows = execute_read_query( db_path, - "SELECT name, op_category, kernel_time_sum_us FROM unified_perf_rows", + "SELECT name, op_category, kernel_time_sum_us FROM unified_perf_rows " + "ORDER BY source_row", ) - assert rows == [ - {"name": "aten::mm", "op_category": "GEMM", "kernel_time_sum_us": 123.5} - ] + 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 kernel_summary", + ) + 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 @@ -165,12 +234,41 @@ def test_import_real_training_report_maps_kernel_stream_and_times(tmp_path): kernels = execute_read_query( db_path, - "SELECT kernel_name, stream, total_duration_us FROM kernel_summary " - "WHERE kernel_name LIKE 'Cijk%' LIMIT 1", + "SELECT k.kernel_name, k.stream, k.total_duration_us, k.unified_row_id, " + "k.library FROM kernel_summary 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, From 5ed8d9d22fcf66234376679362ad1ae4da96de89 Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Tue, 25 Aug 2026 16:55:30 -0400 Subject: [PATCH 08/18] TraceIndex: rename kernel_summary table to op_kernels The table holds one row per kernel belonging to a specific op (via unified_row_id), not the aggregate kernel_summary.csv sheet. Rename it to op_kernels to avoid colliding with that report sheet name and to reflect that each row is a per-op kernel. Update indexes, ingest, docs diagram, and tests. Co-authored-by: Cursor --- TraceLens/TraceIndex/sqlite_store.py | 14 +++++++------- docs/how-to/trace-index.md | 6 +++--- tests/test_trace_index.py | 4 ++-- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/TraceLens/TraceIndex/sqlite_store.py b/TraceLens/TraceIndex/sqlite_store.py index afe570916..763417bdd 100644 --- a/TraceLens/TraceIndex/sqlite_store.py +++ b/TraceLens/TraceIndex/sqlite_store.py @@ -119,7 +119,7 @@ def init_schema(self) -> None: 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 kernel_summary ( + 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, @@ -140,10 +140,10 @@ def init_schema(self) -> None: details_json TEXT ); - CREATE INDEX IF NOT EXISTS idx_trace_index_kernel_trace ON kernel_summary(trace_id); - CREATE INDEX IF NOT EXISTS idx_trace_index_kernel_unified ON kernel_summary(unified_row_id); - CREATE INDEX IF NOT EXISTS idx_trace_index_kernel_name ON kernel_summary(kernel_name); - CREATE INDEX IF NOT EXISTS idx_trace_index_kernel_tensile ON kernel_summary(is_tensile); + 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, @@ -375,7 +375,7 @@ def _clear_trace_payload(self, trace_id: int) -> None: "gemm_perf", "sdpa_perf", "conv_perf", - "kernel_summary", + "op_kernels", "op_category_rows", "gpu_timeline_rows", "trace_summary", @@ -574,7 +574,7 @@ def _import_kernels_from_details( library, is_tensile, is_transpose, is_layout = kernel_flags(kernel_name) self.conn.execute( """ - INSERT INTO kernel_summary( + 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, diff --git a/docs/how-to/trace-index.md b/docs/how-to/trace-index.md index 504f28222..1553d1a29 100644 --- a/docs/how-to/trace-index.md +++ b/docs/how-to/trace-index.md @@ -98,7 +98,7 @@ the `unified_perf_rows` row they came from. erDiagram traces ||--o{ report_imports : "trace_id" traces ||--o{ unified_perf_rows : "trace_id" - traces ||--o{ kernel_summary : "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" @@ -106,7 +106,7 @@ erDiagram traces ||--o{ sdpa_perf : "trace_id" traces ||--o{ conv_perf : "trace_id" traces ||--o{ trace_search_FTS5 : "trace_id" - unified_perf_rows ||--o{ kernel_summary : "unified_row_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" @@ -117,7 +117,7 @@ erDiagram | `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 | -| `kernel_summary` | One row per kernel exploded from `kernel_details_summary` on a unified row, with `unified_row_id` and Tensile / layout flags | +| `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` | diff --git a/tests/test_trace_index.py b/tests/test_trace_index.py index 7ee917100..d57e77f0b 100644 --- a/tests/test_trace_index.py +++ b/tests/test_trace_index.py @@ -153,7 +153,7 @@ def test_trace_index_append_from_report_and_search(tmp_path): kernels = execute_read_query( db_path, "SELECT kernel_name, unified_row_id, library, stream, parent_op_name " - "FROM kernel_summary", + "FROM op_kernels", ) assert kernels[0]["kernel_name"] == "Cijk_test_kernel" assert kernels[0]["unified_row_id"] is not None @@ -235,7 +235,7 @@ def test_import_real_training_report_maps_kernel_stream_and_times(tmp_path): kernels = execute_read_query( db_path, "SELECT k.kernel_name, k.stream, k.total_duration_us, k.unified_row_id, " - "k.library FROM kernel_summary k " + "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", ) From 7f8837069554dc1525f5332ee7809584829c1b10 Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Tue, 25 Aug 2026 17:03:29 -0400 Subject: [PATCH 09/18] TraceIndex: add one-line docstrings to trace index tests Document what each test in test_trace_index.py verifies so the suite reads as a table of contents for the catalog behavior. Co-authored-by: Cursor --- tests/test_trace_index.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/test_trace_index.py b/tests/test_trace_index.py index d57e77f0b..56b5165b1 100644 --- a/tests/test_trace_index.py +++ b/tests/test_trace_index.py @@ -45,6 +45,8 @@ def write_csv(path, rows): 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 @@ -54,6 +56,9 @@ def test_parse_repr_strips_numpy_scalars_and_to_json(): 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 = trace_root / "model_a" / "rank0_trace.json" @@ -177,12 +182,15 @@ def test_trace_index_append_from_report_and_search(tmp_path): 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") def test_trace_index_store_boundary_supports_append_and_search(tmp_path): + """Driving SQLiteTraceIndexStore directly (the storage boundary) imports a + report and returns an op-kind FTS hit.""" db_path = tmp_path / "trace_index.sqlite" trace_root = tmp_path / "traces" trace_path = trace_root / "rank0_trace.json" @@ -222,6 +230,8 @@ def test_trace_index_store_boundary_supports_append_and_search(tmp_path): 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_id = import_report_dir(db_path, TRAINING_REPORT_DIR) @@ -285,6 +295,8 @@ def test_import_real_training_report_maps_kernel_stream_and_times(tmp_path): 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" import_report_dir(db_path, INFERENCE_REPORT_DIR) rows = execute_read_query( @@ -297,6 +309,8 @@ def test_import_real_inference_report_converts_category_kernel_time_ms(tmp_path) 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", @@ -315,6 +329,8 @@ def test_read_traces_file_skips_comments_and_blanks(tmp_path): 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 = tmp_path / "rank0_trace.json" trace_path.write_text(json.dumps({"traceEvents": []}), encoding="utf-8") @@ -351,6 +367,8 @@ def test_cli_append_from_existing_report(tmp_path, capsys): 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: From 136f6e33cf6ce0c5dec6320c2a31f80ac1fadcc3 Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Tue, 25 Aug 2026 17:26:29 -0400 Subject: [PATCH 10/18] Docs: add TraceIndex example queries for conv and SDPA shapes Add an example-queries section to the trace-index how-to showing depthwise convolution and longest-context attention lookups against the shape satellite tables, so shape questions are SQL filters rather than trace reopens. Co-authored-by: Cursor --- docs/how-to/trace-index.md | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/docs/how-to/trace-index.md b/docs/how-to/trace-index.md index 1553d1a29..fe6ed8744 100644 --- a/docs/how-to/trace-index.md +++ b/docs/how-to/trace-index.md @@ -131,6 +131,32 @@ Query GEMM / SDPA / convolution shapes from the satellite tables, or with `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. + +Find depthwise convolutions (channels equal groups): + +```sql +SELECT t.name, c.input_channels, c.output_channels, c.groups, c.kernel_h, c.kernel_w +FROM conv_perf c +JOIN traces t ON t.id = c.trace_id +WHERE c.is_depthwise = 1 AND c.groups > 1 +ORDER BY c.groups DESC; +``` + +Find the longest-context attention shapes: + +```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; +``` + ## Serve read-only SQL For notebook or browser workflows, serve the SQLite catalog over HTTP: From 71817f02403d2c8838c192741e753e582a8fd3b6 Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Tue, 25 Aug 2026 17:28:43 -0400 Subject: [PATCH 11/18] Docs: add example query outputs and takeaways for conv and SDPA Include sample result tables and short takeaways for the depthwise-convolution and longest-attention example queries so the how-to shows the payoff of the shape satellite tables, not just the SQL. Co-authored-by: Cursor --- docs/how-to/trace-index.md | 33 +++++++++++++++++++++++++++++---- 1 file changed, 29 insertions(+), 4 deletions(-) diff --git a/docs/how-to/trace-index.md b/docs/how-to/trace-index.md index fe6ed8744..012c48579 100644 --- a/docs/how-to/trace-index.md +++ b/docs/how-to/trace-index.md @@ -137,17 +137,31 @@ 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. -Find depthwise convolutions (channels equal groups): +### Do any traces have depthwise convolution? ```sql -SELECT t.name, c.input_channels, c.output_channels, c.groups, c.kernel_h, c.kernel_w +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 -ORDER BY c.groups DESC; +GROUP BY t.id, c.input_channels, c.output_channels, c.groups, c.kernel_h, c.kernel_w +ORDER BY rows DESC; ``` -Find the longest-context attention shapes: +| 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 @@ -157,6 +171,17 @@ 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: From fa3d69b6b28fa6f9d5b5f8db5ebb466904ee6e04 Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Tue, 25 Aug 2026 17:42:18 -0400 Subject: [PATCH 12/18] TraceIndex: drop scan and public import_report_dir Catalog ingest is append/build only. Remove the unused directory scanner and stop exporting import_report_dir; existing CSV reports go through append --report-dir. Tests for real report fixtures now use that public path. Co-authored-by: Cursor --- TraceLens/TraceIndex/__init__.py | 4 -- TraceLens/TraceIndex/core.py | 38 ------------------- TraceLens/TraceIndex/scanner.py | 65 +------------------------------- TraceLens/TraceIndex/store.py | 4 +- tests/test_trace_index.py | 9 +++-- 5 files changed, 10 insertions(+), 110 deletions(-) diff --git a/TraceLens/TraceIndex/__init__.py b/TraceLens/TraceIndex/__init__.py index 5bf2aec5f..1a1238632 100644 --- a/TraceLens/TraceIndex/__init__.py +++ b/TraceLens/TraceIndex/__init__.py @@ -11,8 +11,6 @@ build_traces, execute_read_query, generate_report_and_import, - import_report_dir, - scan_traces, search_index, ) from .models import SearchHit, TraceRecord, TraceReport @@ -29,7 +27,5 @@ "build_traces", "execute_read_query", "generate_report_and_import", - "import_report_dir", - "scan_traces", "search_index", ] diff --git a/TraceLens/TraceIndex/core.py b/TraceLens/TraceIndex/core.py index db592ea6d..e223da48d 100644 --- a/TraceLens/TraceIndex/core.py +++ b/TraceLens/TraceIndex/core.py @@ -13,48 +13,10 @@ 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, - import_report_dir as import_report_dir_with_store, ) -from TraceLens.TraceIndex.scanner import scan_traces as scan_traces_with_store from TraceLens.TraceIndex.sqlite_store import SQLiteTraceIndexStore -def scan_traces( - db_path: Path, - root: Path, - peek_mb: int = 2, - compute_md5: bool = False, -) -> int: - store = SQLiteTraceIndexStore(db_path) - try: - return scan_traces_with_store( - store, - root=root, - peek_mb=peek_mb, - compute_md5=compute_md5, - ) - finally: - store.close() - - -def import_report_dir( - db_path: Path, - report_dir: Path, - trace_path: Optional[Path] = None, - root: Optional[Path] = None, -) -> int: - store = SQLiteTraceIndexStore(db_path) - try: - return import_report_dir_with_store( - store, - report_dir=report_dir, - trace_path=trace_path, - root=root, - ) - finally: - store.close() - - def append_trace( db_path: Path, trace_path: Path, diff --git a/TraceLens/TraceIndex/scanner.py b/TraceLens/TraceIndex/scanner.py index b01eded46..042708e17 100644 --- a/TraceLens/TraceIndex/scanner.py +++ b/TraceLens/TraceIndex/scanner.py @@ -4,20 +4,17 @@ # See LICENSE for license information. ############################################################################### -"""Trace discovery and metadata extraction.""" +"""Trace metadata extraction for catalog ingest.""" import gzip import hashlib -import os import re from pathlib import Path -from typing import Iterable, Optional +from typing import Optional from TraceLens.TraceIndex.models import TraceRecord -from TraceLens.TraceIndex.store import TraceIndexStore from TraceLens.TraceIndex.utils import normalize_path, rel_to -TRACE_NAME_RE = re.compile(r"trace|profile|pytorch_profile|rocprof", re.IGNORECASE) RANK_RE = re.compile(r"(?:^|[^A-Za-z])rank[-_]?(\d+)(?:[^0-9]|$)", re.IGNORECASE) SKIP_PARTS_EXACT = { @@ -39,44 +36,10 @@ ) -def iter_files(root: Path) -> Iterable[Path]: - stack = [root] - while stack: - current = stack.pop() - try: - with os.scandir(current) as entries: - dirs = [] - for entry in entries: - try: - if entry.is_dir(follow_symlinks=False): - dirs.append(Path(entry.path)) - elif entry.is_file(follow_symlinks=False): - yield Path(entry.path) - except OSError: - continue - stack.extend(reversed(dirs)) - except OSError: - continue - - def is_json_gz(path: Path) -> bool: return path.name.lower().endswith(".json.gz") -def is_candidate(path: Path) -> bool: - name = path.name.lower() - suffix = path.suffix.lower() - if is_json_gz(path): - return True - if suffix in {".json", ".pftrace", ".rpd"}: - return True - if name.endswith(".xplane.pb"): - return True - if ".pt.trace" in name or ".trace." in name: - return True - return bool(TRACE_NAME_RE.search(path.name)) - - def classify_skip(path: Path, root: Path) -> Optional[str]: try: rel_parts = [part.lower() for part in path.relative_to(root).parts[:-1]] @@ -182,27 +145,3 @@ def trace_record_from_path( should_enrich=should_enrich, skip_reason=skip_reason, ) - - -def scan_traces( - store: TraceIndexStore, - root: Path, - peek_mb: int = 2, - compute_md5: bool = False, -) -> int: - store.init_schema() - root = root.resolve() - count = 0 - for path in iter_files(root): - if not is_candidate(path): - continue - store.upsert_trace( - trace_record_from_path( - path, - root=root, - peek_bytes=peek_mb * 1024 * 1024, - compute_md5=compute_md5, - ) - ) - count += 1 - return count diff --git a/TraceLens/TraceIndex/store.py b/TraceLens/TraceIndex/store.py index 326b3cdf6..70f5f13c3 100644 --- a/TraceLens/TraceIndex/store.py +++ b/TraceLens/TraceIndex/store.py @@ -15,8 +15,8 @@ class TraceIndexStore(ABC): """Persistence boundary for TraceIndex. - New backends should implement this interface without changing scanner, - importer, or CLI workflow code. + New backends should implement this interface without changing importer + or CLI workflow code. """ @abstractmethod diff --git a/tests/test_trace_index.py b/tests/test_trace_index.py index 56b5165b1..bee262a0a 100644 --- a/tests/test_trace_index.py +++ b/tests/test_trace_index.py @@ -13,7 +13,6 @@ from TraceLens.TraceIndex.core import ( append_trace, execute_read_query, - import_report_dir, search_index, ) from TraceLens.TraceIndex.cli import main as trace_index_main @@ -233,7 +232,9 @@ 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_id = import_report_dir(db_path, TRAINING_REPORT_DIR) + trace_path = tmp_path / "qwen_trace.json" + trace_path.write_text(json.dumps({"traceEvents": []}), encoding="utf-8") + trace_id = append_trace(db_path, trace_path, report_dir=TRAINING_REPORT_DIR) unified = execute_read_query( db_path, @@ -298,7 +299,9 @@ 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" - import_report_dir(db_path, INFERENCE_REPORT_DIR) + trace_path = tmp_path / "decode_trace.json" + trace_path.write_text(json.dumps({"traceEvents": []}), encoding="utf-8") + 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 " From c127a9176de94314475542c75c10048933dfc670 Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Tue, 25 Aug 2026 17:45:45 -0400 Subject: [PATCH 13/18] TraceIndex: suppress CodeQL on the gated read-only SQL path The HTTP query server and sqlite-sql intentionally run caller-provided SELECT statements. They are already restricted to a single read-only statement and opened with SQLite query_only/mode=ro. Add CodeQL suppressions so that expected use is not reported as SQL injection. Co-authored-by: Cursor --- TraceLens/TraceIndex/server.py | 3 +++ TraceLens/TraceIndex/sqlite_store.py | 2 ++ 2 files changed, 5 insertions(+) diff --git a/TraceLens/TraceIndex/server.py b/TraceLens/TraceIndex/server.py index ea797904e..d5f268b92 100644 --- a/TraceLens/TraceIndex/server.py +++ b/TraceLens/TraceIndex/server.py @@ -149,6 +149,9 @@ def handle_query(self, sql: str, params: List[Any], limit: int) -> None: limit = max(1, min(limit, max_limit)) start = time.perf_counter() with self.connect() as conn: + # 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] diff --git a/TraceLens/TraceIndex/sqlite_store.py b/TraceLens/TraceIndex/sqlite_store.py index 763417bdd..38746aad5 100644 --- a/TraceLens/TraceIndex/sqlite_store.py +++ b/TraceLens/TraceIndex/sqlite_store.py @@ -363,6 +363,8 @@ def execute_read_query( "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] From 043e26b8897ca7f392ec06fc458ba9e5d76388cb Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Wed, 26 Aug 2026 15:13:25 -0400 Subject: [PATCH 14/18] TraceIndex: add HTTP query-server tests and close DB connections Cover health, tables, GET/POST /query, write rejection, 404s, and result truncation. Close SQLite connections in the handler; a connection context manager does not close the DB and leaked handles in tests. Co-authored-by: Cursor --- TraceLens/TraceIndex/server.py | 10 ++- tests/test_trace_index.py | 131 +++++++++++++++++++++++++++++++++ 2 files changed, 139 insertions(+), 2 deletions(-) diff --git a/TraceLens/TraceIndex/server.py b/TraceLens/TraceIndex/server.py index d5f268b92..929c1bc4f 100644 --- a/TraceLens/TraceIndex/server.py +++ b/TraceLens/TraceIndex/server.py @@ -115,7 +115,8 @@ def handle_health(self) -> None: ) def handle_tables(self) -> None: - with self.connect() as conn: + conn = self.connect() + try: tables = [row["name"] for row in conn.execute(""" SELECT name FROM sqlite_master @@ -136,6 +137,8 @@ def handle_tables(self) -> None: except sqlite3.DatabaseError as exc: counts[table] = repr(exc) self.send_json({"tables": counts}) + finally: + conn.close() def handle_query(self, sql: str, params: List[Any], limit: int) -> None: if not is_read_only_sql(sql): @@ -148,7 +151,8 @@ def handle_query(self, sql: str, params: List[Any], limit: int) -> None: return limit = max(1, min(limit, max_limit)) start = time.perf_counter() - with self.connect() as conn: + 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] @@ -163,6 +167,8 @@ def handle_query(self, sql: str, params: List[Any], limit: int) -> None: "rows": [dict(row) for row in returned], } ) + finally: + conn.close() return Handler diff --git a/tests/test_trace_index.py b/tests/test_trace_index.py index bee262a0a..786b6820f 100644 --- a/tests/test_trace_index.py +++ b/tests/test_trace_index.py @@ -6,6 +6,11 @@ import csv import json +import threading +import urllib.error +import urllib.parse +import urllib.request +from http.server import ThreadingHTTPServer from pathlib import Path import pytest @@ -21,6 +26,7 @@ import_report_dir as import_report_dir_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, @@ -43,6 +49,56 @@ def write_csv(path, rows): writer.writerows(rows) +def seed_mini_catalog(tmp_path): + db_path = tmp_path / "trace_index.sqlite" + trace_path = tmp_path / "rank0_trace.json" + trace_path.write_text(json.dumps({"traceEvents": []}), encoding="utf-8") + report_dir = tmp_path / "report" + 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", + }, + ], + ) + append_trace(db_path, trace_path, report_dir=report_dir) + 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) + + +def start_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 + return server, thread, "http://%s:%s" % (host, port) + + 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.""" @@ -385,3 +441,78 @@ def test_build_traces_continues_after_failure(tmp_path): assert result["imported"] == [] assert len(result["failed"]) == 2 assert "missing_a.json" in result["failed"][0]["trace_path"] + + +def test_query_server_health_tables_and_read_sql(tmp_path): + """The read-only HTTP server reports health/tables and runs a SELECT.""" + db_path = seed_mini_catalog(tmp_path) + server, thread, base = start_query_server(db_path) + try: + 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 + finally: + server.shutdown() + server.server_close() + thread.join(timeout=2) + + +def test_query_server_rejects_write_sql_and_unknown_paths(tmp_path): + """Write SQL and unknown routes are rejected; SELECT limit truncation is reported.""" + db_path = seed_mini_catalog(tmp_path) + server, thread, base = start_query_server(db_path) + try: + 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 + + 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 + finally: + server.shutdown() + server.server_close() + thread.join(timeout=2) From 6ff26598aaa0b3b592cd695c31bce4c4d3325f50 Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Sun, 30 Aug 2026 18:21:59 -0400 Subject: [PATCH 15/18] TraceIndex: refactor test helpers and consolidate query-server tests Extract write_stub_trace/write_mini_report fixtures, use a query_server context manager for HTTP server lifecycle, and merge the two server test functions. Drop redundant store-boundary test covered by append/import paths. Co-authored-by: Cursor --- tests/test_trace_index.py | 138 ++++++++++++-------------------------- 1 file changed, 43 insertions(+), 95 deletions(-) diff --git a/tests/test_trace_index.py b/tests/test_trace_index.py index 786b6820f..33d27cea4 100644 --- a/tests/test_trace_index.py +++ b/tests/test_trace_index.py @@ -10,6 +10,7 @@ import urllib.error import urllib.parse import urllib.request +from contextlib import contextmanager from http.server import ThreadingHTTPServer from pathlib import Path @@ -23,7 +24,6 @@ from TraceLens.TraceIndex.cli import main as trace_index_main from TraceLens.TraceIndex.importer import ( build_traces as build_traces_with_store, - import_report_dir as import_report_dir_with_store, ) from TraceLens.TraceIndex.sqlite_store import SQLiteTraceIndexStore from TraceLens.TraceIndex.server import make_handler @@ -49,11 +49,13 @@ def write_csv(path, rows): writer.writerows(rows) -def seed_mini_catalog(tmp_path): - db_path = tmp_path / "trace_index.sqlite" - trace_path = tmp_path / "rank0_trace.json" - trace_path.write_text(json.dumps({"traceEvents": []}), encoding="utf-8") - report_dir = tmp_path / "report" +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", [ @@ -71,7 +73,13 @@ def seed_mini_catalog(tmp_path): }, ], ) - append_trace(db_path, trace_path, report_dir=report_dir) + 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 @@ -90,13 +98,19 @@ def request_json(url, method="GET", payload=None): return exc.code, json.loads(body) -def start_query_server(db_path): +@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 - return server, thread, "http://%s:%s" % (host, port) + 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(): @@ -116,9 +130,7 @@ def test_trace_index_append_from_report_and_search(tmp_path): json_extract and FTS.""" db_path = tmp_path / "trace_index.sqlite" trace_root = tmp_path / "traces" - trace_path = trace_root / "model_a" / "rank0_trace.json" - trace_path.parent.mkdir(parents=True) - trace_path.write_text(json.dumps({"traceEvents": []}), encoding="utf-8") + trace_path = write_stub_trace(trace_root / "model_a" / "rank0_trace.json") report_dir = tmp_path / "reports" / "trace_1" write_csv( @@ -243,43 +255,6 @@ def test_trace_index_rejects_write_sql(tmp_path): execute_read_query(db_path, "DELETE FROM traces") -def test_trace_index_store_boundary_supports_append_and_search(tmp_path): - """Driving SQLiteTraceIndexStore directly (the storage boundary) imports a - report and returns an op-kind FTS hit.""" - db_path = tmp_path / "trace_index.sqlite" - trace_root = tmp_path / "traces" - trace_path = trace_root / "rank0_trace.json" - trace_path.parent.mkdir(parents=True) - trace_path.write_text(json.dumps({"traceEvents": []}), encoding="utf-8") - - report_dir = tmp_path / "reports" / "trace_1" - write_csv( - report_dir / "unified_perf_summary.csv", - [ - { - "name": "aten::scaled_dot_product_attention", - "op category": "SDPA_fwd", - "operation_count": "1", - "Kernel Time (us)_sum": "10.0", - } - ], - ) - - store = SQLiteTraceIndexStore(db_path) - try: - trace_id = import_report_dir_with_store( - store, - report_dir, - trace_path=trace_path, - root=trace_root, - ) - assert trace_id == 1 - hits = store.search("scaled", limit=10) - assert hits[0].kind == "op" - finally: - store.close() - - @pytest.mark.skipif( not TRAINING_REPORT_DIR.exists(), reason="checked-in Qwen training report CSVs are missing", @@ -288,8 +263,7 @@ 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 = tmp_path / "qwen_trace.json" - trace_path.write_text(json.dumps({"traceEvents": []}), encoding="utf-8") + 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( @@ -355,8 +329,7 @@ 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 = tmp_path / "decode_trace.json" - trace_path.write_text(json.dumps({"traceEvents": []}), encoding="utf-8") + 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, @@ -391,20 +364,8 @@ 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 = tmp_path / "rank0_trace.json" - trace_path.write_text(json.dumps({"traceEvents": []}), encoding="utf-8") - report_dir = tmp_path / "report" - write_csv( - report_dir / "unified_perf_summary.csv", - [ - { - "name": "aten::mm", - "op category": "GEMM", - "operation_count": "1", - "Kernel Time (us)_sum": "10.0", - } - ], - ) + trace_path = write_stub_trace(tmp_path / "rank0_trace.json") + report_dir = write_mini_report(tmp_path / "report") rc = trace_index_main( [ @@ -421,8 +382,10 @@ def test_cli_append_from_existing_report(tmp_path, capsys): 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") - assert rows[0]["name"] == "aten::mm" + 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): @@ -443,11 +406,10 @@ def test_build_traces_continues_after_failure(tmp_path): assert "missing_a.json" in result["failed"][0]["trace_path"] -def test_query_server_health_tables_and_read_sql(tmp_path): - """The read-only HTTP server reports health/tables and runs a SELECT.""" +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) - server, thread, base = start_query_server(db_path) - try: + with query_server(db_path) as base: status, root = request_json(base + "/") assert status == 200 assert "POST /query" in root["endpoints"] @@ -478,17 +440,16 @@ def test_query_server_health_tables_and_read_sql(tmp_path): status, get_query = request_json(base + "/query?" + encoded) assert status == 200 assert get_query["rows"][0]["n"] == 2 - finally: - server.shutdown() - server.server_close() - thread.join(timeout=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 -def test_query_server_rejects_write_sql_and_unknown_paths(tmp_path): - """Write SQL and unknown routes are rejected; SELECT limit truncation is reported.""" - db_path = seed_mini_catalog(tmp_path) - server, thread, base = start_query_server(db_path) - try: status, payload = request_json( base + "/query", method="POST", @@ -503,16 +464,3 @@ def test_query_server_rejects_write_sql_and_unknown_paths(tmp_path): status, post_missing = request_json(base + "/tables", method="POST", payload={}) assert status == 404 - - 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 - finally: - server.shutdown() - server.server_close() - thread.join(timeout=2) From 359a62580e7546ef890fbd0f441a485686e88cbc Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Sun, 30 Aug 2026 18:27:06 -0400 Subject: [PATCH 16/18] Docs: explain report-root and add TraceIndex catalog schema reference Address Kyle review: document --report-root for generated reports, add trace-index-catalog-schema.md for SQL column lookup (separate from perf-report-columns CSV reference), and cross-link from the how-to. Co-authored-by: Cursor --- docs/how-to/trace-index.md | 14 ++ docs/index.rst | 1 + docs/reference/trace-index-catalog-schema.md | 227 +++++++++++++++++++ docs/sphinx/_toc.yml.in | 2 + 4 files changed, 244 insertions(+) create mode 100644 docs/reference/trace-index-catalog-schema.md diff --git a/docs/how-to/trace-index.md b/docs/how-to/trace-index.md index 012c48579..019b00efb 100644 --- a/docs/how-to/trace-index.md +++ b/docs/how-to/trace-index.md @@ -62,6 +62,14 @@ TraceLens_trace_index --db trace_index.sqlite build \ --report-root ./trace_index_reports ``` +`--report-root` is only used when TraceIndex **generates** CSV reports (the +default for `build`, or for `append` without `--report-dir`). TraceIndex writes +one subdirectory per trace under that root, named from the trace filename (for +example `rank0_trace.json.gz` becomes +`./trace_index_reports/rank0_trace.json.gz/`). Re-running `build` or `append` +reuses an existing report directory unless you pass `--force`. If you already +have CSV reports elsewhere, skip generation and use `append --report-dir` instead. + 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`. @@ -90,6 +98,11 @@ 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. +For column-level detail (every SQL column, units, and which CSV sheet it came +from), see [TraceIndex catalog schema](../reference/trace-index-catalog-schema.md). +For columns inside the TraceLens CSV reports themselves, see +[Performance report columns](../reference/perf-report-columns.md). + 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. @@ -226,4 +239,5 @@ rows = search_index(db, "Cijk", limit=20) - [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) +- [TraceIndex catalog schema](../reference/trace-index-catalog-schema.md) - [API reference](../reference/api-reference.md) diff --git a/docs/index.rst b/docs/index.rst index 72d8bb848..00b389e37 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -66,6 +66,7 @@ The TraceLens source code is hosted at `github.com/AMD-AGI/TraceLens ` * :doc:`Performance report columns ` + * :doc:`TraceIndex catalog schema ` * :doc:`Compatibility matrix ` For information on contributing to TraceLens, see the diff --git a/docs/reference/trace-index-catalog-schema.md b/docs/reference/trace-index-catalog-schema.md new file mode 100644 index 000000000..6d4b5afaa --- /dev/null +++ b/docs/reference/trace-index-catalog-schema.md @@ -0,0 +1,227 @@ + + +# TraceIndex catalog schema +```{meta} +:description: Column reference for TraceIndex SQLite catalog tables — the SQL query surface for searching traces, ops, kernels, and parsed shapes. +:keywords: TraceIndex, TraceDB, SQLite, catalog schema, unified_perf_rows, gemm_perf, sdpa_perf, conv_perf, op_kernels, full-text search +``` + +This topic lists every column in the TraceIndex SQLite catalog. Use it when +writing `sqlite-sql` or HTTP `/query` statements against a catalog built with +[Index a corpus of traces](../how-to/trace-index.md). + +TraceIndex imports three CSV sheets from each TraceLens report directory: +`unified_perf_summary.csv`, `ops_summary_by_category.csv`, and `gpu_timeline.csv`. +It also derives satellite shape tables and kernel rows during import. For the +meaning of fields inside the CSV reports, see +[Performance report columns](./perf-report-columns.md). + +## Unit conventions + +- **Time in catalog tables**: microseconds (`*_us`) unless the column name ends + in `_ms` (for example `gpu_timeline_rows.time_ms`). +- **JSON columns**: parsed Python structures stored as JSON text (`*_json`). + Query with `json_extract(column, '$.key')` when no typed satellite column exists. + +## traces + +One row per indexed trace file (or synthetic row for a report-only import). + +| Column | Type | Description | +|--------|------|-------------| +| `id` | INTEGER | Primary key; referenced as `trace_id` elsewhere | +| `root` | TEXT | Optional root directory passed to `append` / `build` | +| `path` | TEXT | Absolute or normalized path to the trace (unique) | +| `rel_path` | TEXT | Path relative to `root`, when set | +| `name` | TEXT | Basename of the trace file | +| `size_bytes` | INTEGER | File size when known | +| `md5` | TEXT | Content hash when computed | +| `format` | TEXT | Trace format hint (for example `chrome`, `pftrace`) | +| `rank` | INTEGER | Rank parsed from filename when present | +| `top_dir` | TEXT | Top-level directory under `root` | +| `parent_rel` | TEXT | Parent directory relative to `root` | +| `should_enrich` | INTEGER | 1 if shape enrichment is expected | +| `skip_reason` | TEXT | Why enrichment was skipped, if any | +| `created_at` | TEXT | ISO timestamp | +| `updated_at` | TEXT | ISO timestamp | + +## report_imports + +Import audit trail for each CSV report directory loaded into a trace. + +| Column | Type | Description | +|--------|------|-------------| +| `id` | INTEGER | Primary key | +| `trace_id` | INTEGER | FK → `traces.id` | +| `report_dir` | TEXT | Path to the TraceLens CSV report directory | +| `imported_at` | TEXT | ISO timestamp | +| `sheets_json` | TEXT | JSON list of sheet names imported | + +## unified_perf_rows + +Rows from `unified_perf_summary.csv`, normalized for SQL. Kernel time columns are +always stored in **microseconds** even when the CSV used milliseconds. + +| Column | Type | CSV source (typical) | Description | +|--------|------|----------------------|-------------| +| `id` | INTEGER | — | Primary key | +| `trace_id` | INTEGER | — | FK → `traces.id` | +| `source_row` | INTEGER | — | 1-based row index in the CSV | +| `name` | TEXT | `name` | Operation name | +| `op_category` | TEXT | `op category` | Category label (GEMM, SDPA_fwd, …) | +| `operation_count` | INTEGER | `operation_count` | Invocation count | +| `kernel_time_sum_us` | REAL | `Kernel Time (us)_sum` | Total kernel time (µs) | +| `kernel_time_mean_us` | REAL | `Kernel Time (us)_mean` | Mean kernel time (µs) | +| `kernel_time_median_us` | REAL | `Kernel Time (us)_median` | Median kernel time (µs) | +| `kernel_time_std_us` | REAL | `Kernel Time (us)_std` | Std dev (µs) | +| `kernel_time_min_us` | REAL | `Kernel Time (us)_min` | Min (µs) | +| `kernel_time_max_us` | REAL | `Kernel Time (us)_max` | Max (µs) | +| `op_duration_us` | REAL | `op duration (us)` | Op wall time (µs) | +| `tflops_mean` | REAL | `TFLOPS_mean` | Mean achieved TFLOPS | +| `tflops_median` | REAL | `TFLOPS_median` | Median TFLOPS | +| `tbs_mean` | REAL | `TB/s_mean` | Mean memory throughput | +| `tbs_median` | REAL | `TB/s_median` | Median memory throughput | +| `gflops` | REAL | `GFLOPS` | Reported GFLOPS | +| `data_moved_mb` | REAL | `data moved (MB)` | Data moved | +| `flops_per_byte` | REAL | `FLOPs/Byte` | Arithmetic intensity | +| `compute_spec` | TEXT | `compute spec` | Roofline / spec label | +| `has_perf_model` | INTEGER | `has perf model` | 1 if perf model attached | +| `overlap_pct` | REAL | `overlap %` | Overlap with other streams | +| `perf_params_json` | TEXT | `perf_params` | Parsed shape / params JSON | +| `kernel_details_json` | TEXT | `kernel_details_summary` | Parsed kernel list JSON | +| `raw_row_json` | TEXT | — | Full original CSV row as JSON | + +## op_kernels + +One row per kernel exploded from `kernel_details_summary` on a unified row. + +| Column | Type | Description | +|--------|------|-------------| +| `id` | INTEGER | Primary key | +| `trace_id` | INTEGER | FK → `traces.id` | +| `unified_row_id` | INTEGER | FK → `unified_perf_rows.id` | +| `kernel_name` | TEXT | HIP kernel symbol | +| `parent_op_name` | TEXT | Parent op from unified row | +| `op_category` | TEXT | Category copied from unified row | +| `stream` | INTEGER | ROCm stream id | +| `count` | INTEGER | Launch count | +| `total_duration_us` | REAL | Total time (µs) | +| `mean_duration_us` | REAL | Mean time (µs) | +| `median_duration_us` | REAL | Median time (µs) | +| `min_duration_us` | REAL | Min time (µs) | +| `max_duration_us` | REAL | Max time (µs) | +| `library` | TEXT | Heuristic label: Tensile, Triton, CK, RCCL/NCCL, … | +| `is_tensile` | INTEGER | 1 if name matches Tensile / Cijk pattern | +| `is_transpose` | INTEGER | 1 if transpose / permute in name | +| `is_layout_conversion` | INTEGER | 1 if layout / cast / copy pattern | +| `details_json` | TEXT | Per-kernel JSON from import | + +## gemm_perf + +Typed GEMM shapes parsed from `perf_params` on GEMM-category unified rows. + +| Column | Type | Description | +|--------|------|-------------| +| `unified_row_id` | INTEGER | PK, FK → `unified_perf_rows.id` | +| `trace_id` | INTEGER | FK → `traces.id` | +| `m`, `n`, `k` | INTEGER | GEMM dimensions | +| `batch` | INTEGER | Batch when present | +| `dtype` | TEXT | Data type string | +| `transpose` | TEXT | Transpose flags | +| `tflops_mean`, `tflops_median` | REAL | Copied from unified row | + +## sdpa_perf + +Typed attention shapes from SDPA-category unified rows. + +| Column | Type | Description | +|--------|------|-------------| +| `unified_row_id` | INTEGER | PK, FK → `unified_perf_rows.id` | +| `trace_id` | INTEGER | FK → `traces.id` | +| `batch` | INTEGER | Batch size | +| `heads` | INTEGER | Number of heads | +| `seq_q` | INTEGER | Query sequence length | +| `seq_kv` | INTEGER | Key/value sequence length | +| `head_dim` | INTEGER | Head dimension | +| `dtype` | TEXT | Data type | +| `causal` | INTEGER | 1 if causal mask | +| `tflops_mean`, `tflops_median` | REAL | Copied from unified row | + +## conv_perf + +Typed convolution shapes from CONV-category unified rows. + +| Column | Type | Description | +|--------|------|-------------| +| `unified_row_id` | INTEGER | PK, FK → `unified_perf_rows.id` | +| `trace_id` | INTEGER | FK → `traces.id` | +| `conv_nd` | TEXT | `"2d"` or `"3d"` | +| `input_shape_json` | TEXT | Input tensor shape JSON | +| `filter_shape_json` | TEXT | Filter shape JSON | +| `input_channels` | INTEGER | Input channels | +| `output_channels` | INTEGER | Output channels | +| `groups` | INTEGER | Group count | +| `kernel_h`, `kernel_w` | INTEGER | Spatial kernel size | +| `is_depthwise` | INTEGER | 1 if depthwise conv | +| `is_transposed_conv` | INTEGER | 1 if transposed conv | + +## op_category_rows + +Rows from `ops_summary_by_category.csv`. + +| Column | Type | CSV source | Description | +|--------|------|------------|-------------| +| `id` | INTEGER | — | Primary key | +| `trace_id` | INTEGER | — | FK → `traces.id` | +| `category` | TEXT | `category` | Op category name | +| `operation_count` | INTEGER | `operation_count` | Count | +| `kernel_time_sum_us` | REAL | kernel time column | Sum in µs | +| `percent` | REAL | `percent` | Fraction of total time | +| `raw_row_json` | TEXT | — | Full CSV row | + +## gpu_timeline_rows + +Rows from `gpu_timeline.csv`. + +| Column | Type | CSV source | Description | +|--------|------|------------|-------------| +| `id` | INTEGER | — | Primary key | +| `trace_id` | INTEGER | — | FK → `traces.id` | +| `type` | TEXT | `type` | Timeline bucket (computation_time, idle_time, …) | +| `time_ms` | REAL | `time ms` | Duration in **milliseconds** | +| `percent` | REAL | `percent` | Percent of wall time | +| `raw_row_json` | TEXT | — | Full CSV row | + +## trace_summary + +Per-trace rollups computed at import time. + +| Column | Type | Description | +|--------|------|-------------| +| `trace_id` | INTEGER | PK, FK → `traces.id` | +| `total_duration_us` | REAL | Total trace duration (µs) | +| `top_categories_json` | TEXT | JSON list of dominant categories | +| `max_gemm_tflops` | REAL | Peak GEMM TFLOPS in this trace | +| `max_sdpa_tflops` | REAL | Peak SDPA TFLOPS in this trace | +| `imported_at` | TEXT | ISO timestamp | + +## trace_search (FTS5) + +Full-text search virtual table. Query via the `search` CLI subcommand or +`MATCH` syntax; columns are not meant for ad-hoc SELECT listing. + +| Column | Type | Description | +|--------|------|-------------| +| `trace_id` | INTEGER | FK → `traces.id` (unindexed in FTS) | +| `kind` | TEXT | Hit kind: trace, op, kernel, category, timeline | +| `text` | TEXT | Searchable text | + +## Related topics + +- [Index a corpus of traces](../how-to/trace-index.md) +- [Performance report columns](./perf-report-columns.md) — CSV / Excel report sheets +- [API reference](./api-reference.md) diff --git a/docs/sphinx/_toc.yml.in b/docs/sphinx/_toc.yml.in index 9c8f6c60b..8e359e375 100644 --- a/docs/sphinx/_toc.yml.in +++ b/docs/sphinx/_toc.yml.in @@ -87,6 +87,8 @@ subtrees: title: API reference - file: reference/perf-report-columns title: Performance report columns + - file: reference/trace-index-catalog-schema + title: TraceIndex catalog schema - url: https://rocm.docs.amd.com/projects/hyperloom/en/latest/index.html title: Hyperloom From dde6efb42e74552f8e49f472c2d2123a162b3914 Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Sun, 30 Aug 2026 18:35:10 -0400 Subject: [PATCH 17/18] TraceIndex: unify query-server error handling in handle_query Move try/except into handle_query so GET and POST /query return the same 400 JSON errors. Cross-link perf-report-columns (semantics) with catalog schema (SQL mapping). Co-authored-by: Cursor --- TraceLens/TraceIndex/server.py | 76 ++++++++++---------- docs/reference/perf-report-columns.md | 15 ++++ docs/reference/trace-index-catalog-schema.md | 41 ++++++----- 3 files changed, 80 insertions(+), 52 deletions(-) diff --git a/TraceLens/TraceIndex/server.py b/TraceLens/TraceIndex/server.py index 929c1bc4f..92f1ad83f 100644 --- a/TraceLens/TraceIndex/server.py +++ b/TraceLens/TraceIndex/server.py @@ -81,7 +81,7 @@ def do_GET(self) -> None: if parsed.path == "/query": params = parse_qs(parsed.query) sql = params.get("sql", [""])[0] - limit = int(params.get("limit", [str(default_limit)])[0]) + limit = params.get("limit", [default_limit])[0] self.handle_query(sql, [], limit) return self.send_json({"error": "not found"}, HTTPStatus.NOT_FOUND) @@ -93,11 +93,45 @@ def do_POST(self) -> None: return try: payload = self.read_json() - self.handle_query( - str(payload.get("sql", "")), - payload.get("params", []), - int(payload.get("limit", default_limit)), - ) + 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) @@ -140,36 +174,6 @@ def handle_tables(self) -> None: finally: conn.close() - def handle_query(self, sql: str, params: List[Any], limit: int) -> None: - 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(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() - return Handler diff --git a/docs/reference/perf-report-columns.md b/docs/reference/perf-report-columns.md index b29c0a946..7616f0aed 100644 --- a/docs/reference/perf-report-columns.md +++ b/docs/reference/perf-report-columns.md @@ -12,6 +12,21 @@ See LICENSE for license information. This topic explains the columns in each sheet of the TraceLens performance report. The report is an Excel file with multiple sheets, and each sheet analyzes a different aspect of GPU performance. Use it as a lookup while you read a report generated with one of the how-to topics listed under [Related topics](#related-topics). +## TraceIndex catalog + +[TraceIndex](../how-to/trace-index.md) imports the same CSV reports into SQLite so you can +search a corpus with SQL. **Column semantics are defined in this topic** — the +catalog is another representation of the same perf metrics, not a different schema. + +TraceIndex currently imports `unified_perf_summary.csv`, `ops_summary_by_category.csv`, +and `gpu_timeline.csv`. It normalizes header spelling (for example `op category` → +`op_category`) and stores kernel times in microseconds in SQL even when the CSV used +milliseconds. It also derives typed shape tables (`gemm_perf`, `sdpa_perf`, `conv_perf`) +and exploded kernel rows (`op_kernels`). + +For the SQL table/column names and CSV → SQL mapping, see +[TraceIndex catalog schema](./trace-index-catalog-schema.md). + ## Overview The performance report Excel file contains multiple sheets analyzing different aspects of GPU performance. The core sheets are: diff --git a/docs/reference/trace-index-catalog-schema.md b/docs/reference/trace-index-catalog-schema.md index 6d4b5afaa..eca45da82 100644 --- a/docs/reference/trace-index-catalog-schema.md +++ b/docs/reference/trace-index-catalog-schema.md @@ -6,26 +6,35 @@ See LICENSE for license information. # TraceIndex catalog schema ```{meta} -:description: Column reference for TraceIndex SQLite catalog tables — the SQL query surface for searching traces, ops, kernels, and parsed shapes. -:keywords: TraceIndex, TraceDB, SQLite, catalog schema, unified_perf_rows, gemm_perf, sdpa_perf, conv_perf, op_kernels, full-text search +:description: SQL table and column names for TraceIndex catalogs, with CSV source mapping. Column semantics live in the performance report column reference. +:keywords: TraceIndex, TraceDB, SQLite, catalog schema, unified_perf_rows, gemm_perf, sdpa_perf, conv_perf, op_kernels ``` -This topic lists every column in the TraceIndex SQLite catalog. Use it when -writing `sqlite-sql` or HTTP `/query` statements against a catalog built with +This topic is the **SQL naming and mapping index** for TraceIndex catalogs built with [Index a corpus of traces](../how-to/trace-index.md). -TraceIndex imports three CSV sheets from each TraceLens report directory: -`unified_perf_summary.csv`, `ops_summary_by_category.csv`, and `gpu_timeline.csv`. -It also derives satellite shape tables and kernel rows during import. For the -meaning of fields inside the CSV reports, see -[Performance report columns](./perf-report-columns.md). +**Column meanings** (units, roofline fields, shape params, kernel details) are documented +in [Performance report columns](./perf-report-columns.md). That topic covers the CSV/Excel +sheets; this topic lists where those fields land in SQLite and which columns are +catalog-only. -## Unit conventions +## Imported CSV sheets → SQL tables -- **Time in catalog tables**: microseconds (`*_us`) unless the column name ends - in `_ms` (for example `gpu_timeline_rows.time_ms`). -- **JSON columns**: parsed Python structures stored as JSON text (`*_json`). - Query with `json_extract(column, '$.key')` when no typed satellite column exists. +| CSV file (report dir) | SQL table(s) | Semantics | +|---|---|---| +| `unified_perf_summary.csv` | `unified_perf_rows` | [unified / ops sheets](./perf-report-columns.md) | +| `ops_summary_by_category.csv` | `op_category_rows` | [ops_summary_by_category](./perf-report-columns.md) | +| `gpu_timeline.csv` | `gpu_timeline_rows` | [gpu_timeline](./perf-report-columns.md) | + +Derived during import (not separate CSV files): `op_kernels` (from `kernel_details_summary`), +`gemm_perf` / `sdpa_perf` / `conv_perf` (typed shapes from `perf_params`), +`trace_summary`, `trace_search` (FTS). + +## Unit conventions in SQL + +- Kernel/op times in relational tables: **microseconds** (`*_us`) unless the column ends in `_ms`. +- `gpu_timeline_rows.time_ms` stays in **milliseconds** (matches the CSV). +- `*_json` columns hold parsed structures; use `json_extract` when no typed satellite column exists. ## traces @@ -66,8 +75,8 @@ Import audit trail for each CSV report directory loaded into a trace. Rows from `unified_perf_summary.csv`, normalized for SQL. Kernel time columns are always stored in **microseconds** even when the CSV used milliseconds. -| Column | Type | CSV source (typical) | Description | -|--------|------|----------------------|-------------| +| Column | Type | CSV source (typical) | SQL name | +|--------|------|----------------------|----------| | `id` | INTEGER | — | Primary key | | `trace_id` | INTEGER | — | FK → `traces.id` | | `source_row` | INTEGER | — | 1-based row index in the CSV | From 83b6e9e69ef75ae8f185d1a67e77f04fff1cb1f8 Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Sun, 30 Aug 2026 18:57:49 -0400 Subject: [PATCH 18/18] Docs: unify metrics reference for CSV reports and TraceIndex SQL Rename perf-report-columns to tracelens-metrics.md with shared intro, an In performance reports section, and an In TraceIndex catalogs section. Remove the separate catalog schema page; update TOC and cross-links. Co-authored-by: Cursor --- README.md | 4 +- docs/conceptual/gemm-analysis.md | 2 +- docs/conceptual/shape-metadata.md | 2 +- docs/conceptual/torch-profiling-analysis.md | 2 +- .../triton-perf-model-walkthrough.md | 2 +- docs/how-to/generate-perf-report-pytorch.md | 2 +- docs/how-to/perf-model-without-trace.md | 2 +- docs/how-to/sdk-analysis.md | 2 +- docs/how-to/trace-index.md | 8 +- docs/index.rst | 3 +- docs/reference/api-reference.md | 2 +- docs/reference/trace-index-catalog-schema.md | 236 ---------------- ...report-columns.md => tracelens-metrics.md} | 261 ++++++++++++++++-- docs/sphinx/_toc.yml.in | 6 +- 14 files changed, 261 insertions(+), 273 deletions(-) delete mode 100644 docs/reference/trace-index-catalog-schema.md rename docs/reference/{perf-report-columns.md => tracelens-metrics.md} (81%) diff --git a/README.md b/README.md index be1f140fa..de75b42bd 100755 --- a/README.md +++ b/README.md @@ -69,7 +69,7 @@ Generate a performance analysis report from an eager execution PyTorch trace wit TraceLens_generate_perf_report_pytorch --profile_json_path path/to/your/trace.json ``` -This produces an Excel workbook with GPU timeline breakdown, ops summary, roofline metrics and more. For additional details, see [Generate a PyTorch performance report](docs/how-to/generate-perf-report-pytorch.md) and [Performance report column reference](docs/reference/perf-report-columns.md). For other input formats, see [Supported Profile Formats](#supported-profile-formats). +This produces an Excel workbook with GPU timeline breakdown, ops summary, roofline metrics and more. For additional details, see [Generate a PyTorch performance report](docs/how-to/generate-perf-report-pytorch.md) and [TraceLens metrics reference](docs/reference/tracelens-metrics.md). For other input formats, see [Supported Profile Formats](#supported-profile-formats). Compare two reports to quantify the impact of a change (see [Compare performance reports](docs/how-to/compare-perf-reports.md)): @@ -149,7 +149,7 @@ Each format's linked doc covers its full CLI reference. For PyTorch report compa | pftrace Reports | [docs/how-to/generate-perf-report-rocprof.md](docs/how-to/generate-perf-report-rocprof.md) | | Compare PyTorch Reports | [docs/how-to/compare-perf-reports.md](docs/how-to/compare-perf-reports.md) | | 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 metrics | [docs/reference/tracelens-metrics.md](docs/reference/tracelens-metrics.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/docs/conceptual/gemm-analysis.md b/docs/conceptual/gemm-analysis.md index 821b9d903..56c3328da 100644 --- a/docs/conceptual/gemm-analysis.md +++ b/docs/conceptual/gemm-analysis.md @@ -468,5 +468,5 @@ For hands-on usage, work through the [`examples/gemm_dim_eff.ipynb`](https://git - [Trace2Tree](../conceptual/trace2tree.md) - [Perf model walkthrough](../conceptual/triton-perf-model-walkthrough.md) - [Torch profiling analysis](../conceptual/torch-profiling-analysis.md) -- [Performance report columns](../reference/perf-report-columns.md) +- [TraceLens metrics](../reference/tracelens-metrics.md) - [Generate a performance report from PyTorch](../how-to/generate-perf-report-pytorch.md) diff --git a/docs/conceptual/shape-metadata.md b/docs/conceptual/shape-metadata.md index feb925279..74213d316 100644 --- a/docs/conceptual/shape-metadata.md +++ b/docs/conceptual/shape-metadata.md @@ -147,5 +147,5 @@ def get_backward_shapes(trace_events, backward_event): - [Trace2Tree](../conceptual/trace2tree.md) - [Understanding PyTorch traces](../conceptual/torch-profiling-analysis.md) - [Generate a performance report from PyTorch traces](../how-to/generate-perf-report-pytorch.md) -- [Performance report columns](../reference/perf-report-columns.md) +- [TraceLens metrics](../reference/tracelens-metrics.md) - [API reference](../reference/api-reference.md) diff --git a/docs/conceptual/torch-profiling-analysis.md b/docs/conceptual/torch-profiling-analysis.md index c16dce5b2..1b940a365 100644 --- a/docs/conceptual/torch-profiling-analysis.md +++ b/docs/conceptual/torch-profiling-analysis.md @@ -147,5 +147,5 @@ Together with the flows discussed earlier, this lets you connect a forward convo - [Trace2Tree](../conceptual/trace2tree.md) - [GEMM analysis](../conceptual/gemm-analysis.md) - [Generate a performance report from PyTorch traces](../how-to/generate-perf-report-pytorch.md) -- [Performance report columns](../reference/perf-report-columns.md) +- [TraceLens metrics](../reference/tracelens-metrics.md) - [What is TraceLens?](../what-is-tracelens.md) diff --git a/docs/conceptual/triton-perf-model-walkthrough.md b/docs/conceptual/triton-perf-model-walkthrough.md index 9537f5b03..f0a24aab2 100644 --- a/docs/conceptual/triton-perf-model-walkthrough.md +++ b/docs/conceptual/triton-perf-model-walkthrough.md @@ -556,5 +556,5 @@ seq_len=4096, dtype=bf16. Traced with PyTorch 2.11+rocm7.2 on AMD Instinct™ MI - [GEMM analysis](../conceptual/gemm-analysis.md) - [Trace2Tree](../conceptual/trace2tree.md) - [Generate a performance report from a PyTorch trace](../how-to/generate-perf-report-pytorch.md) -- [Performance report columns](../reference/perf-report-columns.md) +- [TraceLens metrics](../reference/tracelens-metrics.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 674a930d7..384c8ae33 100644 --- a/docs/how-to/generate-perf-report-pytorch.md +++ b/docs/how-to/generate-perf-report-pytorch.md @@ -77,7 +77,7 @@ For the GPU timeline, a low computation percentage with significant idle time indicates poor compute/communication overlap; use `--micro_idle_thresh_us` to split very short idle gaps into their own category. -See [Performance report column reference](../reference/perf-report-columns.md) +See [TraceLens metrics reference](../reference/tracelens-metrics.md) for what each column means. ## Roofline classification diff --git a/docs/how-to/perf-model-without-trace.md b/docs/how-to/perf-model-without-trace.md index 6ff451f52..a7084ab57 100644 --- a/docs/how-to/perf-model-without-trace.md +++ b/docs/how-to/perf-model-without-trace.md @@ -101,5 +101,5 @@ decode, bf16, AMD Instinct™ MI300X) — all without a single trace. - [Analyze traces with the TraceLens SDK](./sdk-analysis.md) - [GEMM analysis in TraceLens](../conceptual/gemm-analysis.md) -- [Performance report columns](../reference/perf-report-columns.md) +- [TraceLens metrics](../reference/tracelens-metrics.md) - [API reference](../reference/api-reference.md) diff --git a/docs/how-to/sdk-analysis.md b/docs/how-to/sdk-analysis.md index 2ab757bb0..b3176eff1 100644 --- a/docs/how-to/sdk-analysis.md +++ b/docs/how-to/sdk-analysis.md @@ -203,7 +203,7 @@ or memory-bound. For how these metrics are defined and modeled, see [GEMM analysis](../conceptual/gemm-analysis.md) and the -[performance report columns](../reference/perf-report-columns.md) reference. To +[TraceLens metrics](../reference/tracelens-metrics.md) reference. To compute the same metrics from shapes alone, without a trace, see [Model op performance without a trace](./perf-model-without-trace.md). The [`tree_perf_example.ipynb`](https://github.com/AMD-AGI/TraceLens/blob/main/examples/tree_perf_example.ipynb) diff --git a/docs/how-to/trace-index.md b/docs/how-to/trace-index.md index 019b00efb..9e830be6e 100644 --- a/docs/how-to/trace-index.md +++ b/docs/how-to/trace-index.md @@ -99,9 +99,8 @@ typical TraceLens corpora (hundreds of traces, hundreds of thousands of kernel rows). The practical limit is one writer at a time, not row count. For column-level detail (every SQL column, units, and which CSV sheet it came -from), see [TraceIndex catalog schema](../reference/trace-index-catalog-schema.md). -For columns inside the TraceLens CSV reports themselves, see -[Performance report columns](../reference/perf-report-columns.md). +from), see [TraceLens metrics](../reference/tracelens-metrics.md) — report sheets +in the CSV section, catalog tables in the SQLite section. The following diagram shows how those tables relate. Every fact table points at `traces`. Kernel rows and GEMM / SDPA / convolution satellites also point at @@ -238,6 +237,5 @@ rows = search_index(db, "Cijk", limit=20) - [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) -- [TraceIndex catalog schema](../reference/trace-index-catalog-schema.md) +- [TraceLens metrics](../reference/tracelens-metrics.md) - [API reference](../reference/api-reference.md) diff --git a/docs/index.rst b/docs/index.rst index 00b389e37..40cd3a015 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -65,8 +65,7 @@ The TraceLens source code is hosted at `github.com/AMD-AGI/TraceLens ` - * :doc:`Performance report columns ` - * :doc:`TraceIndex catalog schema ` + * :doc:`TraceLens metrics ` * :doc:`Compatibility matrix ` For information on contributing to TraceLens, see the diff --git a/docs/reference/api-reference.md b/docs/reference/api-reference.md index 39df69a67..35a187e29 100644 --- a/docs/reference/api-reference.md +++ b/docs/reference/api-reference.md @@ -214,7 +214,7 @@ example notebook under `examples/`. | `TraceUtils` | Trace utilities, including inference-trace splitting. | — | For report-column definitions across all sheets, see the -[Performance report column reference](./perf-report-columns.md). +[TraceLens metrics reference](./tracelens-metrics.md). ```{note} For a class- and function-level SDK reference generated directly from diff --git a/docs/reference/trace-index-catalog-schema.md b/docs/reference/trace-index-catalog-schema.md deleted file mode 100644 index eca45da82..000000000 --- a/docs/reference/trace-index-catalog-schema.md +++ /dev/null @@ -1,236 +0,0 @@ - - -# TraceIndex catalog schema -```{meta} -:description: SQL table and column names for TraceIndex catalogs, with CSV source mapping. Column semantics live in the performance report column reference. -:keywords: TraceIndex, TraceDB, SQLite, catalog schema, unified_perf_rows, gemm_perf, sdpa_perf, conv_perf, op_kernels -``` - -This topic is the **SQL naming and mapping index** for TraceIndex catalogs built with -[Index a corpus of traces](../how-to/trace-index.md). - -**Column meanings** (units, roofline fields, shape params, kernel details) are documented -in [Performance report columns](./perf-report-columns.md). That topic covers the CSV/Excel -sheets; this topic lists where those fields land in SQLite and which columns are -catalog-only. - -## Imported CSV sheets → SQL tables - -| CSV file (report dir) | SQL table(s) | Semantics | -|---|---|---| -| `unified_perf_summary.csv` | `unified_perf_rows` | [unified / ops sheets](./perf-report-columns.md) | -| `ops_summary_by_category.csv` | `op_category_rows` | [ops_summary_by_category](./perf-report-columns.md) | -| `gpu_timeline.csv` | `gpu_timeline_rows` | [gpu_timeline](./perf-report-columns.md) | - -Derived during import (not separate CSV files): `op_kernels` (from `kernel_details_summary`), -`gemm_perf` / `sdpa_perf` / `conv_perf` (typed shapes from `perf_params`), -`trace_summary`, `trace_search` (FTS). - -## Unit conventions in SQL - -- Kernel/op times in relational tables: **microseconds** (`*_us`) unless the column ends in `_ms`. -- `gpu_timeline_rows.time_ms` stays in **milliseconds** (matches the CSV). -- `*_json` columns hold parsed structures; use `json_extract` when no typed satellite column exists. - -## traces - -One row per indexed trace file (or synthetic row for a report-only import). - -| Column | Type | Description | -|--------|------|-------------| -| `id` | INTEGER | Primary key; referenced as `trace_id` elsewhere | -| `root` | TEXT | Optional root directory passed to `append` / `build` | -| `path` | TEXT | Absolute or normalized path to the trace (unique) | -| `rel_path` | TEXT | Path relative to `root`, when set | -| `name` | TEXT | Basename of the trace file | -| `size_bytes` | INTEGER | File size when known | -| `md5` | TEXT | Content hash when computed | -| `format` | TEXT | Trace format hint (for example `chrome`, `pftrace`) | -| `rank` | INTEGER | Rank parsed from filename when present | -| `top_dir` | TEXT | Top-level directory under `root` | -| `parent_rel` | TEXT | Parent directory relative to `root` | -| `should_enrich` | INTEGER | 1 if shape enrichment is expected | -| `skip_reason` | TEXT | Why enrichment was skipped, if any | -| `created_at` | TEXT | ISO timestamp | -| `updated_at` | TEXT | ISO timestamp | - -## report_imports - -Import audit trail for each CSV report directory loaded into a trace. - -| Column | Type | Description | -|--------|------|-------------| -| `id` | INTEGER | Primary key | -| `trace_id` | INTEGER | FK → `traces.id` | -| `report_dir` | TEXT | Path to the TraceLens CSV report directory | -| `imported_at` | TEXT | ISO timestamp | -| `sheets_json` | TEXT | JSON list of sheet names imported | - -## unified_perf_rows - -Rows from `unified_perf_summary.csv`, normalized for SQL. Kernel time columns are -always stored in **microseconds** even when the CSV used milliseconds. - -| Column | Type | CSV source (typical) | SQL name | -|--------|------|----------------------|----------| -| `id` | INTEGER | — | Primary key | -| `trace_id` | INTEGER | — | FK → `traces.id` | -| `source_row` | INTEGER | — | 1-based row index in the CSV | -| `name` | TEXT | `name` | Operation name | -| `op_category` | TEXT | `op category` | Category label (GEMM, SDPA_fwd, …) | -| `operation_count` | INTEGER | `operation_count` | Invocation count | -| `kernel_time_sum_us` | REAL | `Kernel Time (us)_sum` | Total kernel time (µs) | -| `kernel_time_mean_us` | REAL | `Kernel Time (us)_mean` | Mean kernel time (µs) | -| `kernel_time_median_us` | REAL | `Kernel Time (us)_median` | Median kernel time (µs) | -| `kernel_time_std_us` | REAL | `Kernel Time (us)_std` | Std dev (µs) | -| `kernel_time_min_us` | REAL | `Kernel Time (us)_min` | Min (µs) | -| `kernel_time_max_us` | REAL | `Kernel Time (us)_max` | Max (µs) | -| `op_duration_us` | REAL | `op duration (us)` | Op wall time (µs) | -| `tflops_mean` | REAL | `TFLOPS_mean` | Mean achieved TFLOPS | -| `tflops_median` | REAL | `TFLOPS_median` | Median TFLOPS | -| `tbs_mean` | REAL | `TB/s_mean` | Mean memory throughput | -| `tbs_median` | REAL | `TB/s_median` | Median memory throughput | -| `gflops` | REAL | `GFLOPS` | Reported GFLOPS | -| `data_moved_mb` | REAL | `data moved (MB)` | Data moved | -| `flops_per_byte` | REAL | `FLOPs/Byte` | Arithmetic intensity | -| `compute_spec` | TEXT | `compute spec` | Roofline / spec label | -| `has_perf_model` | INTEGER | `has perf model` | 1 if perf model attached | -| `overlap_pct` | REAL | `overlap %` | Overlap with other streams | -| `perf_params_json` | TEXT | `perf_params` | Parsed shape / params JSON | -| `kernel_details_json` | TEXT | `kernel_details_summary` | Parsed kernel list JSON | -| `raw_row_json` | TEXT | — | Full original CSV row as JSON | - -## op_kernels - -One row per kernel exploded from `kernel_details_summary` on a unified row. - -| Column | Type | Description | -|--------|------|-------------| -| `id` | INTEGER | Primary key | -| `trace_id` | INTEGER | FK → `traces.id` | -| `unified_row_id` | INTEGER | FK → `unified_perf_rows.id` | -| `kernel_name` | TEXT | HIP kernel symbol | -| `parent_op_name` | TEXT | Parent op from unified row | -| `op_category` | TEXT | Category copied from unified row | -| `stream` | INTEGER | ROCm stream id | -| `count` | INTEGER | Launch count | -| `total_duration_us` | REAL | Total time (µs) | -| `mean_duration_us` | REAL | Mean time (µs) | -| `median_duration_us` | REAL | Median time (µs) | -| `min_duration_us` | REAL | Min time (µs) | -| `max_duration_us` | REAL | Max time (µs) | -| `library` | TEXT | Heuristic label: Tensile, Triton, CK, RCCL/NCCL, … | -| `is_tensile` | INTEGER | 1 if name matches Tensile / Cijk pattern | -| `is_transpose` | INTEGER | 1 if transpose / permute in name | -| `is_layout_conversion` | INTEGER | 1 if layout / cast / copy pattern | -| `details_json` | TEXT | Per-kernel JSON from import | - -## gemm_perf - -Typed GEMM shapes parsed from `perf_params` on GEMM-category unified rows. - -| Column | Type | Description | -|--------|------|-------------| -| `unified_row_id` | INTEGER | PK, FK → `unified_perf_rows.id` | -| `trace_id` | INTEGER | FK → `traces.id` | -| `m`, `n`, `k` | INTEGER | GEMM dimensions | -| `batch` | INTEGER | Batch when present | -| `dtype` | TEXT | Data type string | -| `transpose` | TEXT | Transpose flags | -| `tflops_mean`, `tflops_median` | REAL | Copied from unified row | - -## sdpa_perf - -Typed attention shapes from SDPA-category unified rows. - -| Column | Type | Description | -|--------|------|-------------| -| `unified_row_id` | INTEGER | PK, FK → `unified_perf_rows.id` | -| `trace_id` | INTEGER | FK → `traces.id` | -| `batch` | INTEGER | Batch size | -| `heads` | INTEGER | Number of heads | -| `seq_q` | INTEGER | Query sequence length | -| `seq_kv` | INTEGER | Key/value sequence length | -| `head_dim` | INTEGER | Head dimension | -| `dtype` | TEXT | Data type | -| `causal` | INTEGER | 1 if causal mask | -| `tflops_mean`, `tflops_median` | REAL | Copied from unified row | - -## conv_perf - -Typed convolution shapes from CONV-category unified rows. - -| Column | Type | Description | -|--------|------|-------------| -| `unified_row_id` | INTEGER | PK, FK → `unified_perf_rows.id` | -| `trace_id` | INTEGER | FK → `traces.id` | -| `conv_nd` | TEXT | `"2d"` or `"3d"` | -| `input_shape_json` | TEXT | Input tensor shape JSON | -| `filter_shape_json` | TEXT | Filter shape JSON | -| `input_channels` | INTEGER | Input channels | -| `output_channels` | INTEGER | Output channels | -| `groups` | INTEGER | Group count | -| `kernel_h`, `kernel_w` | INTEGER | Spatial kernel size | -| `is_depthwise` | INTEGER | 1 if depthwise conv | -| `is_transposed_conv` | INTEGER | 1 if transposed conv | - -## op_category_rows - -Rows from `ops_summary_by_category.csv`. - -| Column | Type | CSV source | Description | -|--------|------|------------|-------------| -| `id` | INTEGER | — | Primary key | -| `trace_id` | INTEGER | — | FK → `traces.id` | -| `category` | TEXT | `category` | Op category name | -| `operation_count` | INTEGER | `operation_count` | Count | -| `kernel_time_sum_us` | REAL | kernel time column | Sum in µs | -| `percent` | REAL | `percent` | Fraction of total time | -| `raw_row_json` | TEXT | — | Full CSV row | - -## gpu_timeline_rows - -Rows from `gpu_timeline.csv`. - -| Column | Type | CSV source | Description | -|--------|------|------------|-------------| -| `id` | INTEGER | — | Primary key | -| `trace_id` | INTEGER | — | FK → `traces.id` | -| `type` | TEXT | `type` | Timeline bucket (computation_time, idle_time, …) | -| `time_ms` | REAL | `time ms` | Duration in **milliseconds** | -| `percent` | REAL | `percent` | Percent of wall time | -| `raw_row_json` | TEXT | — | Full CSV row | - -## trace_summary - -Per-trace rollups computed at import time. - -| Column | Type | Description | -|--------|------|-------------| -| `trace_id` | INTEGER | PK, FK → `traces.id` | -| `total_duration_us` | REAL | Total trace duration (µs) | -| `top_categories_json` | TEXT | JSON list of dominant categories | -| `max_gemm_tflops` | REAL | Peak GEMM TFLOPS in this trace | -| `max_sdpa_tflops` | REAL | Peak SDPA TFLOPS in this trace | -| `imported_at` | TEXT | ISO timestamp | - -## trace_search (FTS5) - -Full-text search virtual table. Query via the `search` CLI subcommand or -`MATCH` syntax; columns are not meant for ad-hoc SELECT listing. - -| Column | Type | Description | -|--------|------|-------------| -| `trace_id` | INTEGER | FK → `traces.id` (unindexed in FTS) | -| `kind` | TEXT | Hit kind: trace, op, kernel, category, timeline | -| `text` | TEXT | Searchable text | - -## Related topics - -- [Index a corpus of traces](../how-to/trace-index.md) -- [Performance report columns](./perf-report-columns.md) — CSV / Excel report sheets -- [API reference](./api-reference.md) diff --git a/docs/reference/perf-report-columns.md b/docs/reference/tracelens-metrics.md similarity index 81% rename from docs/reference/perf-report-columns.md rename to docs/reference/tracelens-metrics.md index 7616f0aed..07ddbac4b 100644 --- a/docs/reference/perf-report-columns.md +++ b/docs/reference/tracelens-metrics.md @@ -4,28 +4,38 @@ Copyright (c) 2025 - 2026 Advanced Micro Devices, Inc. All rights reserved. See LICENSE for license information. --> -# Performance report column reference +# TraceLens metrics reference ```{meta} -:description: Reference for every sheet and column in a TraceLens performance report, including the gpu_timeline breakdown, operation-analysis sheets, roofline metrics, and collective communication analysis. -:keywords: TraceLens, performance report, column reference, gpu_timeline, ops, ops_summary, ops_unique_args, roofline, GEMM, SDPA, convolution, NCCL, coll_analysis, GFLOPS, TFLOPS, arithmetic intensity, ROCm, AMD Instinct, MI300X +:description: Reference for TraceLens performance metrics — what they mean, and how they appear in CSV/Excel performance reports and TraceIndex SQLite catalogs. +:keywords: TraceLens, metrics, performance report, TraceIndex, catalog schema, gpu_timeline, unified_perf_summary, GEMM, SDPA, roofline, SQLite ``` -This topic explains the columns in each sheet of the TraceLens performance report. The report is an Excel file with multiple sheets, and each sheet analyzes a different aspect of GPU performance. Use it as a lookup while you read a report generated with one of the how-to topics listed under [Related topics](#related-topics). +TraceLens turns profiler traces into **metrics**: aggregated timings, roofline +efficiency, operation shapes, kernel details, timeline breakdowns, and collective +analysis. The same metrics appear in two representations: -## TraceIndex catalog +| Representation | Where | Typical use | +|---|---|---| +| **Performance report** | CSV directory or Excel workbook | Read one run, share with teammates | +| **TraceIndex catalog** | SQLite database | Search and SQL-query a corpus of runs | -[TraceIndex](../how-to/trace-index.md) imports the same CSV reports into SQLite so you can -search a corpus with SQL. **Column semantics are defined in this topic** — the -catalog is another representation of the same perf metrics, not a different schema. +This topic defines what the metrics **mean** (report sections below) and how they +are **named in SQL** (catalog section below). For building and querying catalogs, +see [Index a corpus of traces](../how-to/trace-index.md). -TraceIndex currently imports `unified_perf_summary.csv`, `ops_summary_by_category.csv`, -and `gpu_timeline.csv`. It normalizes header spelling (for example `op category` → -`op_category`) and stores kernel times in microseconds in SQL even when the CSV used -milliseconds. It also derives typed shape tables (`gemm_perf`, `sdpa_perf`, `conv_perf`) -and exploded kernel rows (`op_kernels`). +### Unit conventions + +- **Time**: microseconds (µs) in most op/kernel columns unless a column says `ms`. +- **Compute**: GFLOPS (billions of FLOPs), TFLOPS/s (trillions of FLOPs per second). +- **Memory**: MB (megabytes), GB/s, TB/s. + +TraceIndex stores kernel/op times as `*_us` in SQL even when the source CSV used +milliseconds for category rows. `gpu_timeline_rows.time_ms` stays in milliseconds. -For the SQL table/column names and CSV → SQL mapping, see -[TraceIndex catalog schema](./trace-index-catalog-schema.md). +## In performance reports (CSV and Excel) + +Performance reports are an Excel workbook (or equivalent CSV export). Each sheet +analyzes a different aspect of GPU performance. ## Overview @@ -928,10 +938,229 @@ A typical top-down analysis flow: 6. Check collectives (`coll_analysis`): analyze communication costs in distributed training. 7. Replay operations: reproduce and benchmark specific cases in isolation. +## In TraceIndex catalogs (SQLite) + +TraceIndex imports CSV reports into SQLite for corpus search. The metric +**semantics** are defined in the report sections above; this section lists SQL +table and column names and CSV → SQL mapping. + +## Imported CSV sheets → SQL tables + +| CSV file (report dir) | SQL table(s) | Semantics | +|---|---|---| +| `unified_perf_summary.csv` | `unified_perf_rows` | [Additional sheets](#additional-sheets) | +| `ops_summary_by_category.csv` | `op_category_rows` | [Operation analysis](#operation-analysis) | +| `gpu_timeline.csv` | `gpu_timeline_rows` | [The gpu_timeline sheet](#the-gpu_timeline-sheet) | + +Derived during import (not separate CSV files): `op_kernels` (from `kernel_details_summary`), +`gemm_perf` / `sdpa_perf` / `conv_perf` (typed shapes from `perf_params`), +`trace_summary`, `trace_search` (FTS). + +## Unit conventions in SQL + +- Kernel/op times in relational tables: **microseconds** (`*_us`) unless the column ends in `_ms`. +- `gpu_timeline_rows.time_ms` stays in **milliseconds** (matches the CSV). +- `*_json` columns hold parsed structures; use `json_extract` when no typed satellite column exists. + +## traces + +One row per indexed trace file (or synthetic row for a report-only import). + +| Column | Type | Description | +|--------|------|-------------| +| `id` | INTEGER | Primary key; referenced as `trace_id` elsewhere | +| `root` | TEXT | Optional root directory passed to `append` / `build` | +| `path` | TEXT | Absolute or normalized path to the trace (unique) | +| `rel_path` | TEXT | Path relative to `root`, when set | +| `name` | TEXT | Basename of the trace file | +| `size_bytes` | INTEGER | File size when known | +| `md5` | TEXT | Content hash when computed | +| `format` | TEXT | Trace format hint (for example `chrome`, `pftrace`) | +| `rank` | INTEGER | Rank parsed from filename when present | +| `top_dir` | TEXT | Top-level directory under `root` | +| `parent_rel` | TEXT | Parent directory relative to `root` | +| `should_enrich` | INTEGER | 1 if shape enrichment is expected | +| `skip_reason` | TEXT | Why enrichment was skipped, if any | +| `created_at` | TEXT | ISO timestamp | +| `updated_at` | TEXT | ISO timestamp | + +## report_imports + +Import audit trail for each CSV report directory loaded into a trace. + +| Column | Type | Description | +|--------|------|-------------| +| `id` | INTEGER | Primary key | +| `trace_id` | INTEGER | FK → `traces.id` | +| `report_dir` | TEXT | Path to the TraceLens CSV report directory | +| `imported_at` | TEXT | ISO timestamp | +| `sheets_json` | TEXT | JSON list of sheet names imported | + +## unified_perf_rows + +Rows from `unified_perf_summary.csv`, normalized for SQL. Kernel time columns are +always stored in **microseconds** even when the CSV used milliseconds. + +| Column | Type | CSV source (typical) | SQL name | +|--------|------|----------------------|----------| +| `id` | INTEGER | — | Primary key | +| `trace_id` | INTEGER | — | FK → `traces.id` | +| `source_row` | INTEGER | — | 1-based row index in the CSV | +| `name` | TEXT | `name` | Operation name | +| `op_category` | TEXT | `op category` | Category label (GEMM, SDPA_fwd, …) | +| `operation_count` | INTEGER | `operation_count` | Invocation count | +| `kernel_time_sum_us` | REAL | `Kernel Time (us)_sum` | Total kernel time (µs) | +| `kernel_time_mean_us` | REAL | `Kernel Time (us)_mean` | Mean kernel time (µs) | +| `kernel_time_median_us` | REAL | `Kernel Time (us)_median` | Median kernel time (µs) | +| `kernel_time_std_us` | REAL | `Kernel Time (us)_std` | Std dev (µs) | +| `kernel_time_min_us` | REAL | `Kernel Time (us)_min` | Min (µs) | +| `kernel_time_max_us` | REAL | `Kernel Time (us)_max` | Max (µs) | +| `op_duration_us` | REAL | `op duration (us)` | Op wall time (µs) | +| `tflops_mean` | REAL | `TFLOPS_mean` | Mean achieved TFLOPS | +| `tflops_median` | REAL | `TFLOPS_median` | Median TFLOPS | +| `tbs_mean` | REAL | `TB/s_mean` | Mean memory throughput | +| `tbs_median` | REAL | `TB/s_median` | Median memory throughput | +| `gflops` | REAL | `GFLOPS` | Reported GFLOPS | +| `data_moved_mb` | REAL | `data moved (MB)` | Data moved | +| `flops_per_byte` | REAL | `FLOPs/Byte` | Arithmetic intensity | +| `compute_spec` | TEXT | `compute spec` | Roofline / spec label | +| `has_perf_model` | INTEGER | `has perf model` | 1 if perf model attached | +| `overlap_pct` | REAL | `overlap %` | Overlap with other streams | +| `perf_params_json` | TEXT | `perf_params` | Parsed shape / params JSON | +| `kernel_details_json` | TEXT | `kernel_details_summary` | Parsed kernel list JSON | +| `raw_row_json` | TEXT | — | Full original CSV row as JSON | + +## op_kernels + +One row per kernel exploded from `kernel_details_summary` on a unified row. + +| Column | Type | Description | +|--------|------|-------------| +| `id` | INTEGER | Primary key | +| `trace_id` | INTEGER | FK → `traces.id` | +| `unified_row_id` | INTEGER | FK → `unified_perf_rows.id` | +| `kernel_name` | TEXT | HIP kernel symbol | +| `parent_op_name` | TEXT | Parent op from unified row | +| `op_category` | TEXT | Category copied from unified row | +| `stream` | INTEGER | ROCm stream id | +| `count` | INTEGER | Launch count | +| `total_duration_us` | REAL | Total time (µs) | +| `mean_duration_us` | REAL | Mean time (µs) | +| `median_duration_us` | REAL | Median time (µs) | +| `min_duration_us` | REAL | Min time (µs) | +| `max_duration_us` | REAL | Max time (µs) | +| `library` | TEXT | Heuristic label: Tensile, Triton, CK, RCCL/NCCL, … | +| `is_tensile` | INTEGER | 1 if name matches Tensile / Cijk pattern | +| `is_transpose` | INTEGER | 1 if transpose / permute in name | +| `is_layout_conversion` | INTEGER | 1 if layout / cast / copy pattern | +| `details_json` | TEXT | Per-kernel JSON from import | + +## gemm_perf + +Typed GEMM shapes parsed from `perf_params` on GEMM-category unified rows. + +| Column | Type | Description | +|--------|------|-------------| +| `unified_row_id` | INTEGER | PK, FK → `unified_perf_rows.id` | +| `trace_id` | INTEGER | FK → `traces.id` | +| `m`, `n`, `k` | INTEGER | GEMM dimensions | +| `batch` | INTEGER | Batch when present | +| `dtype` | TEXT | Data type string | +| `transpose` | TEXT | Transpose flags | +| `tflops_mean`, `tflops_median` | REAL | Copied from unified row | + +## sdpa_perf + +Typed attention shapes from SDPA-category unified rows. + +| Column | Type | Description | +|--------|------|-------------| +| `unified_row_id` | INTEGER | PK, FK → `unified_perf_rows.id` | +| `trace_id` | INTEGER | FK → `traces.id` | +| `batch` | INTEGER | Batch size | +| `heads` | INTEGER | Number of heads | +| `seq_q` | INTEGER | Query sequence length | +| `seq_kv` | INTEGER | Key/value sequence length | +| `head_dim` | INTEGER | Head dimension | +| `dtype` | TEXT | Data type | +| `causal` | INTEGER | 1 if causal mask | +| `tflops_mean`, `tflops_median` | REAL | Copied from unified row | + +## conv_perf + +Typed convolution shapes from CONV-category unified rows. + +| Column | Type | Description | +|--------|------|-------------| +| `unified_row_id` | INTEGER | PK, FK → `unified_perf_rows.id` | +| `trace_id` | INTEGER | FK → `traces.id` | +| `conv_nd` | TEXT | `"2d"` or `"3d"` | +| `input_shape_json` | TEXT | Input tensor shape JSON | +| `filter_shape_json` | TEXT | Filter shape JSON | +| `input_channels` | INTEGER | Input channels | +| `output_channels` | INTEGER | Output channels | +| `groups` | INTEGER | Group count | +| `kernel_h`, `kernel_w` | INTEGER | Spatial kernel size | +| `is_depthwise` | INTEGER | 1 if depthwise conv | +| `is_transposed_conv` | INTEGER | 1 if transposed conv | + +## op_category_rows + +Rows from `ops_summary_by_category.csv`. + +| Column | Type | CSV source | Description | +|--------|------|------------|-------------| +| `id` | INTEGER | — | Primary key | +| `trace_id` | INTEGER | — | FK → `traces.id` | +| `category` | TEXT | `category` | Op category name | +| `operation_count` | INTEGER | `operation_count` | Count | +| `kernel_time_sum_us` | REAL | kernel time column | Sum in µs | +| `percent` | REAL | `percent` | Fraction of total time | +| `raw_row_json` | TEXT | — | Full CSV row | + +## gpu_timeline_rows + +Rows from `gpu_timeline.csv`. + +| Column | Type | CSV source | Description | +|--------|------|------------|-------------| +| `id` | INTEGER | — | Primary key | +| `trace_id` | INTEGER | — | FK → `traces.id` | +| `type` | TEXT | `type` | Timeline bucket (computation_time, idle_time, …) | +| `time_ms` | REAL | `time ms` | Duration in **milliseconds** | +| `percent` | REAL | `percent` | Percent of wall time | +| `raw_row_json` | TEXT | — | Full CSV row | + +## trace_summary + +Per-trace rollups computed at import time. + +| Column | Type | Description | +|--------|------|-------------| +| `trace_id` | INTEGER | PK, FK → `traces.id` | +| `total_duration_us` | REAL | Total trace duration (µs) | +| `top_categories_json` | TEXT | JSON list of dominant categories | +| `max_gemm_tflops` | REAL | Peak GEMM TFLOPS in this trace | +| `max_sdpa_tflops` | REAL | Peak SDPA TFLOPS in this trace | +| `imported_at` | TEXT | ISO timestamp | + +## trace_search (FTS5) + +Full-text search virtual table. Query via the `search` CLI subcommand or +`MATCH` syntax; columns are not meant for ad-hoc SELECT listing. + +| Column | Type | Description | +|--------|------|-------------| +| `trace_id` | INTEGER | FK → `traces.id` (unindexed in FTS) | +| `kind` | TEXT | Hit kind: trace, op, kernel, category, timeline | +| `text` | TEXT | Searchable text | + ## Related topics +- [Index a corpus of traces](../how-to/trace-index.md) - [Generate a PyTorch performance report](../how-to/generate-perf-report-pytorch.md) - [Generate a JAX performance report](../how-to/generate-perf-report-jax.md) - [Generate a rocprofv3 performance report](../how-to/generate-perf-report-rocprof.md) - [Replay a traced operation](../how-to/event-replay.md) -- [API reference](../reference/api-reference.md) +- [API reference](./api-reference.md) + diff --git a/docs/sphinx/_toc.yml.in b/docs/sphinx/_toc.yml.in index 8e359e375..876b6ce97 100644 --- a/docs/sphinx/_toc.yml.in +++ b/docs/sphinx/_toc.yml.in @@ -85,10 +85,8 @@ subtrees: entries: - file: reference/api-reference title: API reference - - file: reference/perf-report-columns - title: Performance report columns - - file: reference/trace-index-catalog-schema - title: TraceIndex catalog schema + - file: reference/tracelens-metrics + title: TraceLens metrics - url: https://rocm.docs.amd.com/projects/hyperloom/en/latest/index.html title: Hyperloom