From dd3decbe5285707925dc9cb6e9e0f6c4e75cd666 Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Mon, 24 Aug 2026 14:25:56 -0400 Subject: [PATCH 01/17] 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/17] 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/17] 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/17] 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/17] 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/17] 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/17] 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/17] 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/17] 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/17] 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/17] 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/17] 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/17] 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/17] 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/17] 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 7b6930ca6f70b1e6fedd1e9f2721e463d04a7e90 Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Sun, 30 Aug 2026 19:03:23 -0400 Subject: [PATCH 16/17] TraceIndex: unify query-server error handling in handle_query Co-authored-by: Cursor --- TraceLens/TraceIndex/server.py | 76 ++++++++++++++++++---------------- 1 file changed, 40 insertions(+), 36 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 From 1cf1294553266c49629a2f101b4bbb0eed2f7fd7 Mon Sep 17 00:00:00 2001 From: Adeem Jassani Date: Sun, 30 Aug 2026 19:11:22 -0400 Subject: [PATCH 17/17] PerfModel: unify kernel library classification for TraceIndex and agents Co-authored-by: Cursor --- .../category_analyses/analysis_utils.py | 35 +------- TraceLens/PerfModel/kernel_library.py | 79 +++++++++++++++++++ TraceLens/TraceIndex/sqlite_store.py | 4 +- TraceLens/TraceIndex/utils.py | 27 ++++--- tests/test_kernel_library.py | 55 +++++++++++++ 5 files changed, 156 insertions(+), 44 deletions(-) create mode 100644 TraceLens/PerfModel/kernel_library.py create mode 100644 tests/test_kernel_library.py diff --git a/TraceLens/Agent/Analysis/category_analyses/analysis_utils.py b/TraceLens/Agent/Analysis/category_analyses/analysis_utils.py index 812cccc98..26f54f723 100644 --- a/TraceLens/Agent/Analysis/category_analyses/analysis_utils.py +++ b/TraceLens/Agent/Analysis/category_analyses/analysis_utils.py @@ -27,6 +27,7 @@ import numpy as np import pandas as pd +from TraceLens.PerfModel.kernel_library import classify_kernel_library from TraceLens.PerfModel.utils import torch_dtype_map TARGET_HIGH = 100.0 @@ -44,26 +45,6 @@ # Efficiency Boundary for grouping _EFF_BUCKET_BOUNDARIES = (30, 60) -_OP_NAME_LIBRARY_RULES = [ - ("aiter::", "AITER"), - ("rocm_aiter", "AITER"), - ("fbgemm", "FBGEMM"), - ("miopen", "MIOpen"), - ("triton", "Triton"), -] -_KERNEL_NAME_LIBRARY_RULES = [ - ("aiter", "AITER"), - ("ck_tile::", "CK"), - ("ck_tile6kentry", "CK"), - ("FmhaFwd", "CK"), - ("FmhaBwd", "CK"), - ("Cijk_", "Tensile"), - ("wvSplitK", "rocBLAS"), - ("splitKreduce", "rocBLAS"), - ("rocprim::", "rocPRIM"), - ("triton_", "Triton"), - ("void at::native::", "PyTorch Native"), -] _COMPARATIVE_SPEEDUP_COL = "speedup (trace2/trace1)" _COMPARATIVE_DELTA_COL = "delta_us (trace2 - trace1)" _COMPARATIVE_T2_TIME_COL = "lca_total_kernel_time_trace2_us" @@ -1093,15 +1074,5 @@ def run_category_analysis( print(f"Metrics written to: {output_path}") -# Category-specific helper functions -def classify_kernel_library(op_name: str, kernel_details: str = "") -> Optional[str]: - """Identify the GPU library backing an operation from its name or kernel strings.""" - op_lower = op_name.lower() - for marker, lib in _OP_NAME_LIBRARY_RULES: - if marker in op_lower: - return lib - kd = str(kernel_details) if kernel_details and not pd.isna(kernel_details) else "" - for marker, lib in _KERNEL_NAME_LIBRARY_RULES: - if marker in kd: - return lib - return None +# Category-specific helper functions — classify_kernel_library lives in +# TraceLens.PerfModel.kernel_library and is imported above. diff --git a/TraceLens/PerfModel/kernel_library.py b/TraceLens/PerfModel/kernel_library.py new file mode 100644 index 000000000..11ebd0f8b --- /dev/null +++ b/TraceLens/PerfModel/kernel_library.py @@ -0,0 +1,79 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Shared kernel-library classification for perf reports, agents, and TraceIndex.""" + +from __future__ import annotations + +from typing import Any, Optional + +from TraceLens.PerfModel import kernel_name_parser as knp + +_OP_NAME_LIBRARY_RULES = [ + ("aiter::", "AITER"), + ("rocm_aiter", "AITER"), + ("fbgemm", "FBGEMM"), + ("miopen", "MIOpen"), + ("triton", "Triton"), +] + +_KERNEL_NAME_LIBRARY_RULES = [ + ("aiter", "AITER"), + ("ck_tile::", "CK"), + ("ck_tile6kentry", "CK"), + ("FmhaFwd", "CK"), + ("FmhaBwd", "CK"), + ("Cijk_", "Tensile"), + ("wvSplitK", "rocBLAS"), + ("splitKreduce", "rocBLAS"), + ("rocprim::", "rocPRIM"), + ("triton_", "Triton"), + ("void at::native::", "PyTorch Native"), + ("nccl", "RCCL/NCCL"), + ("rccl", "RCCL/NCCL"), + ("composable", "CK"), +] + +# Ordered GEMM detectors; extend when kernel_name_parser gains more libraries (#805). +_GEMM_LIBRARY_CHECKS = ( + (knp.is_rocm_gemm, "Tensile"), + (knp.is_cuda_gemm, "nvjet"), +) + + +def _coerce_kernel_details(kernel_details: Any) -> str: + if kernel_details in (None, ""): + return "" + if isinstance(kernel_details, float) and kernel_details != kernel_details: + return "" + return str(kernel_details) + + +def _library_from_gemm_kernel_name(kernel_name: str) -> Optional[str]: + for detector, library in _GEMM_LIBRARY_CHECKS: + if detector(kernel_name): + return library + return None + + +def classify_kernel_library(op_name: str, kernel_details: Any = "") -> Optional[str]: + """Identify the GPU library backing an operation from its name or kernel strings.""" + op_lower = (op_name or "").lower() + for marker, library in _OP_NAME_LIBRARY_RULES: + if marker in op_lower: + return library + + kd = _coerce_kernel_details(kernel_details) + if kd: + gemm_library = _library_from_gemm_kernel_name(kd) + if gemm_library is not None: + return gemm_library + + kd_lower = kd.lower() + for marker, library in _KERNEL_NAME_LIBRARY_RULES: + if marker.lower() in kd_lower: + return library + return None diff --git a/TraceLens/TraceIndex/sqlite_store.py b/TraceLens/TraceIndex/sqlite_store.py index 38746aad5..e06ec0245 100644 --- a/TraceLens/TraceIndex/sqlite_store.py +++ b/TraceLens/TraceIndex/sqlite_store.py @@ -573,7 +573,9 @@ def _import_kernels_from_details( kernel_name = as_text(detail.get("name") or detail.get("Kernel name")) if not kernel_name: continue - library, is_tensile, is_transpose, is_layout = kernel_flags(kernel_name) + library, is_tensile, is_transpose, is_layout = kernel_flags( + kernel_name, parent_op_name or "" + ) self.conn.execute( """ INSERT INTO op_kernels( diff --git a/TraceLens/TraceIndex/utils.py b/TraceLens/TraceIndex/utils.py index 3394fc44b..15e5898fc 100644 --- a/TraceLens/TraceIndex/utils.py +++ b/TraceLens/TraceIndex/utils.py @@ -16,6 +16,9 @@ from pathlib import Path from typing import Any, Dict, List, Optional, Sequence, Tuple +from TraceLens.PerfModel import kernel_name_parser as knp +from TraceLens.PerfModel.kernel_library import classify_kernel_library + def utc_now() -> str: return datetime.now(timezone.utc).isoformat(timespec="seconds") @@ -165,10 +168,21 @@ def to_json(value: Any) -> Optional[str]: return json.dumps(json_safe(value), sort_keys=True) -def kernel_flags(name: str) -> Tuple[Optional[str], int, int, int]: +def kernel_flags(name: str, op_name: str = "") -> Tuple[Optional[str], int, int, int]: + """Derive catalog flags for a GPU kernel using shared library classification.""" low = name.lower() - is_tensile = int("cijk" in low or "tensile" in low) + library = classify_kernel_library(op_name, name) + + is_tensile = int( + knp.is_rocm_gemm(name) or library == "Tensile" or "tensile" in low + ) is_transpose = int("transpose" in low or "permute" in low) + parsed = knp.gemm_name_parser(name) + if parsed: + transpose = parsed.get("transpose") + if transpose and any(transpose): + is_transpose = 1 + is_layout = int( is_transpose or "contiguous" in low @@ -176,15 +190,6 @@ def kernel_flags(name: str) -> Tuple[Optional[str], int, int, int]: 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 diff --git a/tests/test_kernel_library.py b/tests/test_kernel_library.py new file mode 100644 index 000000000..555808c6b --- /dev/null +++ b/tests/test_kernel_library.py @@ -0,0 +1,55 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +import pytest + +from TraceLens.PerfModel.kernel_library import classify_kernel_library +from TraceLens.TraceIndex.utils import kernel_flags + + +@pytest.mark.parametrize( + "op_name,kernel_details,expected", + [ + ("aiter::gemm", "", "AITER"), + ("triton_kernel_op", "", "Triton"), + ("aten::mm", "[{'name': 'Cijk_foo'}]", "Tensile"), + ("aten::mm", "void at::native::x", "PyTorch Native"), + ("aten::mm", "plain", None), + ("aten::mm", "ncclAllReduce", "RCCL/NCCL"), + ( + "aten::mm", + "Cijk_Alik_Bljk_BBS_BH_Bias_HAS_SAV_UserArgs_MT64x16x64_MI16x16x1", + "Tensile", + ), + ("aten::mm", "nvjet_tst_144x128_64x6_2x1_v_bz_bias_TNN", "nvjet"), + ], +) +def test_classify_kernel_library(op_name, kernel_details, expected): + assert classify_kernel_library(op_name, kernel_details) == expected + + +@pytest.mark.parametrize( + "kernel_name,op_name,expected_library,expected_tensile,expected_transpose", + [ + ( + "Cijk_Alik_Bljk_BBS_BH_Bias_HAS_SAV_UserArgs_MT64x16x64", + "aten::mm", + "Tensile", + 1, + 1, + ), + ("triton_red_fused_add", "aten::add", "Triton", 0, 0), + ("ncclAllReduceRing", "", "RCCL/NCCL", 0, 0), + ("custom_transpose_kernel", "", None, 0, 1), + ], +) +def test_kernel_flags( + kernel_name, op_name, expected_library, expected_tensile, expected_transpose +): + library, is_tensile, is_transpose, _is_layout = kernel_flags(kernel_name, op_name) + assert library == expected_library + assert is_tensile == expected_tensile + assert is_transpose == expected_transpose