diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 00000000..2de54d0f --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,95 @@ + + +# AGENTS.md + +The authoring contract for human or AI, changing this repo: how to make a change land +cleanly. + +TraceLens is a Python library for automated performance analysis of ML training and inference +workloads from profiler traces. It parses PyTorch, JAX, and rocprofv3 traces into a hierarchical +event tree, models roofline and compute performance, analyzes collective communication, diffs +traces, and drives an agentic optimization report. + +## Documentation + +| Topic | Source of truth | +|---|---| +| Overview, capabilities, supported formats | [`docs/what-is-tracelens.md`](docs/what-is-tracelens.md) | +| Module architecture, per-module guides | [`README.md`](README.md) § Documentation, [`docs/`](docs/) | +| Dev setup, branch and commit style, updating references | [`CONTRIBUTING.md`](CONTRIBUTING.md) | + +## Common commands + +```bash +pip install -e .[dev] # editable install with dev extras + +python -m pytest tests/ # full suite +python -m pytest tests/test_perf_report_regression.py # one suite + +black . # required before every PR; CI pins black==26.3.1 +python tests/update_copyright.py # fix missing copyright headers +``` + +The primary CLIs install as `console_scripts` (see `entry_points` in [`setup.py`](setup.py)). +Demo traces for local runs are bundled in `tests/traces`. + +## Authoring rules of engagement + +### Change scope and shape + +- **One concern per change.** A PR fixes one issue or adds one capability; if you must bundle, say + why in the description. Don't ride unrelated refactors in on a fix. +- **Smallest diff.** Change exactly what was asked. If you find yourself editing a file the request + did not name, stop and confirm. +- **Small, focused PRs.** Open an issue first for a new analyser or backend integration. +- **Update the docs with the change.** When a change alters a CLI flag, API; update the affected `README.md` and `docs/` in the same PR. + +### Reuse and structure + +- **Build on the owning layer.** Before adding a standard loader, parser, or regex, find the + canonical owner and use it. In `TraceLens/util.py`: `DataLoader.load_data` for trace JSON / + `json.gz`. `Trace2Tree`/`GPUEventAnalyser` for the event tree and timeline, + `TraceUtils/annotation_utils` for annotations. +- **Build new analysis on the existing tree, not a re-load.** Parse the trace and construct the tree once, + then pass that tree to every downstream analysis; don't re-load the JSON or rebuild the tree per + consumer. Reopen or reparse only when an analysis genuinely needs a different view that the existing tree cannot provide. +- **Group helpers; don't sprawl.** A pile of one-line functions across many files is a class you + haven't named yet: put a helper in the module whose responsibility matches its purpose, not + next to its first caller, and extend an existing helper over adding a parallel one. +- **Derive over hardcode.** drop references to a + concern from agents that don't own it, and template sections the downstream consumer never reads. + +### Correctness and honesty + +- **Catch narrowly; don't route around.** No new broad `except Exception` or bare `except`; catch + the specific error, or let it raise. No new feature flag or env toggle to route around a design + problem. A single computed source beats duplicated constants. Thresholds and regexes are calibrated + in one place; tune them at the source, never fork a second copy. Keep interfaces minimal. +- **Trust the caller.** Validate at the system boundary, then trust internal callers; redundant + re-checks and layered fallbacks hide the failure they were added to survive. +- **Review feedback is a hypothesis.** A comment can be right about the symptom and wrong about the + fix; converge on the design that is correct, not the one that is merely defensible. + +### Hygiene + +- **Delete, don't comment out.** No `# removed …` tombstones. +- **Comment why, not what,** and only where the reason is non-obvious. Never narrate the change: no + "previously this did X," no step or plan numbering. Match the surrounding comment and docstring + style; don't write unusually long comments or docstrings. +- **Follow the file layout.** New Python file: copyright banner → module docstring → imports + (fully-qualified `from TraceLens…`, stdlib then third-party then local, no `sys.path.insert`, + imports at the top and no mid-file imports, the only exception being conditional imports) → a + `# Constants` block with every threshold and regex → public functions, + then private (`_name`). +- **Public repo, vendor-neutral.** Never add private, confidential, or customer data. Keep code and + docs vendor-neutral, unless the surrounding code is already specific. Quoting an actual kernel + name from a trace is fine. Don't commit generated output, traces, or large binaries; no destructive + git operations without an explicit target. +- **English, nothing generated in git.** The repo is English: code, identifiers, comments, commit + messages, and docs, regardless of the language the work was discussed in. +- **Leave nothing behind.** Working notes and analysis write-ups are byproducts of the work, not + deliverables; don't commit them, least of all at the repo root, unless they were asked for. diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 00000000..487cda58 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,9 @@ + + +# CLAUDE.md + +Read [AGENTS.md](AGENTS.md). diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 90cf0fc1..e0cbfc2c 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -21,7 +21,6 @@ Thanks for your interest in improving **TraceLens** — a toolkit that parses Py - Read the [README](./README.md) to understand scope and architecture. - Search existing **issues** and **discussions** to avoid duplicates. - For new features and enhancements (new analyser, backend integration, refactor), **open an issue** first to align on approach. -- Prefer small, modular, focused PRs. - **Have a ready-made utility?** If your utility is already developed, you can raise a PR to add it directly to `examples/custom_workflows/`. This lets the community start using it right away while we plan a tighter integration into the core library. --- @@ -29,10 +28,6 @@ Thanks for your interest in improving **TraceLens** — a toolkit that parses Py ## Table of Contents - [Dev Setup](#dev-setup) -- [Project Structure (high level)](#project-structure-high-level) -- [Code Formatting with Black](#code-formatting-with-black) - - [Installing Black](#installing-black) - - [Using Black](#using-black) - [Branch Naming Convention](#branch-naming-convention) - [Types (type)](#types-type) - [Scope (optional)](#scope-optional) @@ -57,54 +52,6 @@ pip install -U pip pip install -e .[dev] ``` -## Project Structure (high level) - -```text -TraceLens/ -├── TraceLens/ -│ ├── Reporting/ # CLI tools for quick start utils -│ ├── Trace2Tree/ # Trace2Tree parses trace into tree data structure -│ ├── PerfModel/ # Op meta data parsing and performance modelling code (roofline, FLOPs/Byte, etc.) -│ ├── TreePerf/ # TreePerf uses Trace Tree and PerfModel to generate perf breakdowns and perf metrics TFLOPS/s, etc. -| | # This directory also contains GPUEventAnalyzer -│ ├── NcclAnalyser/ # Analysis of collective communications -│ ├── TraceFusion/ # Merging of multi‑rank traces into a global view -│ ├── TraceDiff/ # TraceDiff uses the Trace Tree format and does morphological comparison across traces -│ └── EventReplay/ # Extracts meta data and replays almost arbitrary operations -├── docs/ # tool-specific guides -├── examples/ # example traces, notebooks, scripts, custom-workflows -├── tests/ # unit & integration tests -└── setup.py -``` - -## Code Formatting with Black - -This project uses [Black](https://black.readthedocs.io/en/stable/) to automatically format Python code for consistency and readability. - -### Installing Black - -You can install Black using pip: - -```sh -pip install black -``` - -### Using Black - -To format all Python files in the project, run: - -```sh -black . -``` - -You can also format a specific file: - -```sh -black path/to/your_file.py -``` - -Please ensure your code is formatted with Black before submitting a pull request. - ## Branch Naming Convention Please follow this branch naming convention for all feature and bug fix branches: diff --git a/README.md b/README.md index bff50e88..8c1e0424 100755 --- a/README.md +++ b/README.md @@ -47,7 +47,7 @@ pip install git+https://github.com/AMD-AGI/TraceLens.git TraceLens analyses profiler traces from PyTorch, JAX, and AMD rocprofv3; see [Supported Profile Formats](#supported-profile-formats) for the full list. The instructions below cover collecting a PyTorch trace: - **Generic Eager Traces**: Instrument your loop with `torch.profiler.profile(...)`, enabling CPU-side call-stack and shape capture (`with_stack=True`, `record_shapes=True`). Profile a representative steady-state window (a handful of steps, post-warmup) and log the trace with `prof.export_chrome_trace(...)`. A single rank's trace is enough for per-rank analysis. The [PyTorch profiling walkthrough](notebooks/torch-profiling.ipynb) walks through this end to end. -- **Inference Traces with Graph Capture**: Collection has framework-specific requirements. Follow guidelines in [Generate a PyTorch inference report](docs/how-to/generate-perf-report-pytorch-inference.md). The [Profiling skill](TraceLens/Agent/Profiling/README.md) automates vLLM/SGLang/ATOM benchmarking and PyTorch profiler trace collection via [Magpie](https://github.com/AMD-AGI/Magpie), producing analysis-ready traces. +- **Traces with Graph Capture**: Collection has framework-specific requirements. Follow guidelines in [Generate a PyTorch inference report](docs/how-to/generate-perf-report-pytorch-inference.md). The [Profiling skill](TraceLens/Agent/Profiling/README.md) automates vLLM/SGLang/ATOM benchmarking and PyTorch profiler trace collection via [Magpie](https://github.com/AMD-AGI/Magpie), producing analysis-ready traces. To try out TraceLens without collecting your own trace, use the [demo traces](tests/traces) bundled in the repository. @@ -125,6 +125,32 @@ Each format's linked doc covers its full CLI reference. For PyTorch report compa --- +## Project Structure + +```text +TraceLens/ +├── TraceLens/ +│ ├── Reporting/ # CLI tools for quick start utils +│ ├── Trace2Tree/ # Trace2Tree parses trace into tree data structure +│ ├── PerfModel/ # Op meta data parsing and performance modelling code (roofline, FLOPs/Byte, etc.) +│ ├── TreePerf/ # TreePerf uses Trace Tree and PerfModel to generate perf breakdowns and perf metrics TFLOPS/s, etc. +│ ├── NcclAnalyser/ # Analysis of collective communications +│ ├── TraceFusion/ # Merging of multi-rank traces into a global view +│ ├── TraceDiff/ # TraceDiff uses the Trace Tree format and does morphological comparison across traces +│ ├── EventReplay/ # Extracts meta data and replays almost arbitrary operations +│ ├── TraceIndex/ # Build a searchable SQLite catalog of TraceLens reports +│ ├── TraceUtils/ # Shared trace loaders and annotation helpers +│ └── Agent/ # Agentic optimization report (Analysis) and trace collection (Profiling) +├── docs/ # tool-specific guides +├── examples/ # example traces, notebooks, scripts, custom-workflows +├── notebooks/ # end-to-end profiling and analysis walkthroughs +├── scripts/ # helper scripts +├── tests/ # unit & integration tests +└── setup.py +``` + +--- + ## Documentation | Module | Doc | @@ -146,16 +172,6 @@ Each format's linked doc covers its full CLI reference. For PyTorch report compa --- -## Development - -```bash -git clone https://github.com/AMD-AGI/TraceLens.git && cd TraceLens -pip install -e .[dev] -python -m pytest tests/ -v -``` - ---- - ## Contributing Contributions are welcome across the entire project, including new analysis modules, performance models, documentation, examples, and bug fixes. diff --git a/TraceLens/Agent/Analysis/README.md b/TraceLens/Agent/Analysis/README.md index 57c7f000..36ad2dc4 100644 --- a/TraceLens/Agent/Analysis/README.md +++ b/TraceLens/Agent/Analysis/README.md @@ -58,9 +58,9 @@ pip install -e . ### 2. Collect a trace -The orchestrator runs against a single PyTorch profiler trace (`.json` or `.json.gz`). Collection is workload-specific: +The orchestrator runs against a PyTorch profiler trace. Collection is workload-specific: -- **Generic Eager Traces**: Instrument your loop with `torch.profiler.profile(...)`, enabling CPU-side call-stack and shape capture (`with_stack=True`, `record_shapes=True`). Profile a representative steady-state window (a handful of steps, post-warmup) and log the trace with `prof.export_chrome_trace(...)`. A single rank's trace is enough for per-rank analysis. The [PyTorch profiling walkthrough](../../../notebooks/torch-profiling.ipynb) walks through this end to end. +- **Generic Eager Traces**: Instrument your loop with `torch.profiler.profile(...)`, enabling CPU-side call-stack and shape capture (`with_stack=True`, `record_shapes=True`). Profile a representative steady-state window (a handful of steps, post-warmup) and log the trace with `prof.export_chrome_trace(...)`. A single rank's trace is enough for per-rank analysis. See [Collect a trace](../../../docs/how-to/generate-perf-report-pytorch.md#collect-a-trace). - **Inference Traces with Graph Capture**: Collection has framework-specific requirements. Follow guidelines in [Generate a PyTorch inference report](../../../docs/how-to/generate-perf-report-pytorch-inference.md). The [Profiling skill](../Profiling/README.md) automates vLLM/SGLang/ATOM benchmarking and PyTorch profiler trace collection via [Magpie](https://github.com/AMD-AGI/Magpie), producing analysis-ready traces. For graph-mode workloads you produce two artifacts: a graph-replay trace and a graph-capture folder. In inference mode with execution mode `graph replay + capture`, TraceLens merges call-stack and shape information from the capture folder into the replay tree before analysis. ### 3. Establish a hardware performance baseline diff --git a/docs/how-to/agent.md b/docs/how-to/agent.md index e3e953d1..6b67961b 100644 --- a/docs/how-to/agent.md +++ b/docs/how-to/agent.md @@ -57,25 +57,10 @@ pip install git+https://github.com/AMD-AGI/TraceLens.git ### Collect a trace -The orchestrator runs against a single `torch.profiler` trace (`.json` or -`.json.gz`). Collection is workload-specific: - -- **Generic Eager Traces**: Instrument your loop with - `torch.profiler.profile(...)`, enabling CPU-side call-stack and shape capture - (`with_stack=True`, `record_shapes=True`). Profile a representative steady-state - window of a handful of post-warmup steps, then log the trace with - `prof.export_chrome_trace(...)`. A single rank's trace is enough for per-rank - analysis. -- **Inference Traces with Graph Capture**: Collection has framework-specific - requirements. Follow - [Generate a PyTorch inference performance report](./generate-perf-report-pytorch-inference.md). - The Profiling Skill automates - vLLM, SGLang, and ATOM benchmarking and PyTorch profiler trace collection using - Magpie, producing analysis-ready traces. For - graph-mode workloads you produce two artifacts: a graph-replay trace and a - graph-capture folder. In inference mode with execution mode - `graph replay + capture`, TraceLens merges call-stack and shape information from - the capture folder into the replay tree before analysis. +The orchestrator runs against a PyTorch profiler trace. Collection is workload-specific: + +- **Generic Eager Traces**: Instrument your loop with `torch.profiler.profile(...)`, enabling CPU-side call-stack and shape capture (`with_stack=True`, `record_shapes=True`). Profile a representative steady-state window (a handful of steps, post-warmup) and log the trace with `prof.export_chrome_trace(...)`. A single rank's trace is enough for per-rank analysis. See [Collect a trace](../../../docs/how-to/generate-perf-report-pytorch.md#collect-a-trace). +- **Inference Traces with Graph Capture**: Collection has framework-specific requirements. Follow guidelines in [Generate a PyTorch inference report](../../../docs/how-to/generate-perf-report-pytorch-inference.md). The [Profiling skill](../Profiling/README.md) automates vLLM/SGLang/ATOM benchmarking and PyTorch profiler trace collection via [Magpie](https://github.com/AMD-AGI/Magpie), producing analysis-ready traces. For graph-mode workloads you produce two artifacts: a graph-replay trace and a graph-capture folder. In inference mode with execution mode `graph replay + capture`, TraceLens merges call-stack and shape information from the capture folder into the replay tree before analysis. ### Establish a hardware baseline @@ -206,6 +191,7 @@ node, or an SSH plus container-exec wrapper for a containerized node. - [Generate a PyTorch performance report](./generate-perf-report-pytorch.md) - [Generate a PyTorch inference performance report](./generate-perf-report-pytorch-inference.md) +- [Inference performance analysis](../conceptual/inference-analysis.md) - [Analyze traces with the TraceLens SDK](./sdk-analysis.md) - [Trace2Tree data model](../conceptual/trace2tree.md) - [GEMM analysis](../conceptual/gemm-analysis.md) \ No newline at end of file diff --git a/docs/how-to/generate-perf-report-pytorch.md b/docs/how-to/generate-perf-report-pytorch.md index 900d51af..36bd8c17 100644 --- a/docs/how-to/generate-perf-report-pytorch.md +++ b/docs/how-to/generate-perf-report-pytorch.md @@ -21,10 +21,48 @@ Before generating a report, confirm you have the following: - [TraceLens installed](../install/install.md). - A `torch.profiler` Chrome trace (`.json` or `.json.gz`). -```{note} -If you don't have a trace yet, see the +If you don't have a trace yet, capture one as shown below. The [PyTorch profiling walkthrough](https://github.com/AMD-AGI/TraceLens/blob/main/notebooks/torch-profiling.ipynb) -for instructions on capturing one with `torch.profiler`. +walks through it end to end. + +## Collect a trace + +The quality of the TraceLens Analysis report depends on +the quality of the trace. + +```python +import torch + +def export(prof): + prof.export_chrome_trace("trace.json") + +with torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + schedule=torch.profiler.schedule(wait=1, warmup=1, active=2), + on_trace_ready=export, + record_shapes=True, + with_stack=True, +) as prof: + for _ in range(num_steps): # num_steps >= 4 to cover wait + warmup + active + model(inputs) + torch.cuda.synchronize() + prof.step() +``` + +Two flags on `torch.profiler.profile` change what TraceLens can report. + +| Flag | What it captures | Why TraceLens needs it | +|---|---|---| +| `record_shapes=True` | Input argument shapes and dtypes for each operator | Roofline and compute modeling need shapes to compute FLOPs, bytes, and arithmetic intensity. Without them, per-operator efficiency metrics are unavailable. | +| `with_stack=True` | The Python call stack above each operator | Required for the call-stack views — GPU time grouped by `nn.Module` or by the line of Python that launched the work. Reports run with `--include_call_stack` need it. | + +```{note} +Inference frameworks that run in HIP graph mode (vLLM, SGLang, ATOM, xDiT) +need a different capture path. See +[Generate a PyTorch inference performance report](./generate-perf-report-pytorch-inference.md). ``` ## Generate the report