From 229495ee97cc4e522709a2cf41710f3f1bf46e90 Mon Sep 17 00:00:00 2001 From: Cemberk Date: Tue, 11 Aug 2026 09:32:15 -0500 Subject: [PATCH] feat(profiling): integrate TraceLens trace analysis Turn the trace artifacts a run leaves behind into TraceLens operator, kernel, roofline, and collective reports, and add on-demand torch.profiler capture so PyTorch workloads can produce those traces without editing the model script. Analysis runs in either of two places. On the host, `madengine report tracelens` keeps TraceLens' pinned protobuf and xprof out of workload images entirely. In-container, the `tracelens` tool installs into an isolated virtualenv for the same reason. SLURM and Kubernetes collection now gather torch_profiler_output/ and tracelens_output/ per node. The `scripts/` and `*.json` ignore rules are anchored to the repo root: unanchored they also matched packaged source, which silently excluded the five runtime scripts these tool definitions depend on. Co-authored-by: Cursor --- .gitignore | 7 +- docs/cli-reference.md | 80 +++ docs/profiling.md | 163 +++++ examples/profiling-configs/README.md | 29 +- .../torch_profiler_tracelens.json | 19 + .../tracelens_rocprofv3.json | 15 + pyproject.toml | 8 + src/madengine/cli/commands/report.py | 281 ++++++++ src/madengine/deployment/k8s_results.py | 10 +- src/madengine/deployment/k8s_scripts.py | 10 +- .../templates/kubernetes/job.yaml.j2 | 26 +- .../deployment/templates/slurm/job.sh.j2 | 4 +- src/madengine/reporting/tracelens_report.py | 233 +++++++ .../common/post_scripts/dynolog_stop.sh | 84 +++ .../scripts/common/post_scripts/trace.sh | 19 + .../scripts/common/post_scripts/tracelens.sh | 53 ++ .../common/pre_scripts/dynolog_start.sh | 63 ++ .../scripts/common/pre_scripts/trace.sh | 118 ++++ src/madengine/scripts/common/tools.json | 151 ++++ .../scripts/common/tools/dynolog_trigger.sh | 75 ++ .../scripts/common/tools/tracelens_analyze.py | 650 ++++++++++++++++++ tests/e2e/test_tracelens_workflows.py | 380 ++++++++++ .../test_tracelens_tools_config.py | 202 ++++++ tests/unit/test_tracelens_analyze.py | 335 +++++++++ tests/unit/test_tracelens_report.py | 282 ++++++++ 25 files changed, 3280 insertions(+), 17 deletions(-) create mode 100644 examples/profiling-configs/torch_profiler_tracelens.json create mode 100644 examples/profiling-configs/tracelens_rocprofv3.json create mode 100644 src/madengine/reporting/tracelens_report.py create mode 100644 src/madengine/scripts/common/post_scripts/dynolog_stop.sh create mode 100644 src/madengine/scripts/common/post_scripts/tracelens.sh create mode 100644 src/madengine/scripts/common/pre_scripts/dynolog_start.sh create mode 100644 src/madengine/scripts/common/tools/dynolog_trigger.sh create mode 100644 src/madengine/scripts/common/tools/tracelens_analyze.py create mode 100644 tests/e2e/test_tracelens_workflows.py create mode 100644 tests/integration/test_tracelens_tools_config.py create mode 100644 tests/unit/test_tracelens_analyze.py create mode 100644 tests/unit/test_tracelens_report.py diff --git a/.gitignore b/.gitignore index c824efdf..1edfc91d 100644 --- a/.gitignore +++ b/.gitignore @@ -128,9 +128,12 @@ venv/ .venv/ # model relatives +# Anchored to the repo root: these are a MAD project's own model files, and the +# directories madengine populates in the working directory during a run. Left +# unanchored they also swallow packaged source such as src/madengine/scripts/. docker/ -scripts/ -*.json +/scripts/ +/*.json .*_env/ .vscode/ diff --git a/docs/cli-reference.md b/docs/cli-reference.md index 842e0fdf..cf9e1791 100644 --- a/docs/cli-reference.md +++ b/docs/cli-reference.md @@ -442,6 +442,86 @@ madengine report to-email --directory ./results --verbose --- +##### `report tracelens` - Analyze GPU Traces + +Generate [TraceLens](https://github.com/AMD-AGI/TraceLens) performance reports from the trace artifacts a run left behind (`torch_profiler_output/`, `rocprof_output/`, `slurm_results/`, `k8s_results/`). + +Requires TraceLens: `pip install 'madengine[tracelens]'`. `--discover-only` works without it. + +**Usage:** + +```bash +madengine report tracelens [OPTIONS] +``` + +**Options:** + +| Option | Type | Default | Description | +|--------|------|---------|-------------| +| `--root` | TEXT | `"."` | Directory searched recursively for traces | +| `--output-dir` | TEXT | `tracelens_output` | Directory for generated reports | +| `--mode` | TEXT | `auto` | `auto`, `pytorch`, `rocprof`, `pftrace`, or `collective` | +| `--python` | TEXT | | Interpreter that has TraceLens installed | +| `--gpu-arch` | TEXT | | GPU arch (e.g. `MI300X`) enabling roofline bound classification | +| `--world-size` | INTEGER | trace count | Rank count for the collective report | +| `--max-traces` | INTEGER | `0` | Cap traces analyzed per kind (`0` means no cap) | +| `--discover-only` | FLAG | `False` | List discovered traces without running TraceLens | +| `--verbose` | FLAG | `False` | Enable verbose logging | + +**Examples:** + +```bash +# Analyze every trace found under the current directory +madengine report tracelens + +# See what would be analyzed first +madengine report tracelens --discover-only + +# Enable roofline bound classification +madengine report tracelens --gpu-arch MI300X + +# Multi-rank collective analysis of a distributed run +madengine report tracelens --root slurm_results --mode collective --world-size 8 +``` + +**Output:** Per-trace reports in `--output-dir`, plus `tracelens_summary.csv` and `tracelens_summary.json`. + +See the [Profiling Guide](profiling.md#tracelens---tracelens-trace-analysis) for which trace formats map to which report. + +--- + +##### `report tracelens-compare` - Diff TraceLens Reports + +Compare two or more TraceLens reports into a single diff workbook. The first report is the baseline; every metric gains `_diff` and `_pct` columns relative to it. + +**Usage:** + +```bash +madengine report tracelens-compare REPORT... [OPTIONS] +``` + +**Options:** + +| Option | Short | Type | Default | Description | +|--------|-------|------|---------|-------------| +| `--output` | `-o` | TEXT | `tracelens_comparison.xlsx` | Output comparison workbook | +| `--names` | | TEXT | | Display tag per report (repeat the flag) | +| `--python` | | TEXT | | Interpreter that has TraceLens installed | +| `--verbose` | `-v` | FLAG | `False` | Enable verbose logging | + +**Examples:** + +```bash +# Compare a baseline against a candidate run +madengine report tracelens-compare baseline.xlsx candidate.xlsx + +# Label each report and choose the output path +madengine report tracelens-compare a.xlsx b.xlsx \ + --names before --names after -o diff.xlsx +``` + +--- + ### `database` - Upload to MongoDB Upload CSV performance data to MongoDB database. diff --git a/docs/profiling.md b/docs/profiling.md index 5e421ec8..39352804 100644 --- a/docs/profiling.md +++ b/docs/profiling.md @@ -342,6 +342,167 @@ madengine run --tags your_model \ }' ``` +### torch_profiler_dynolog - On-Demand PyTorch Traces + +Capture `torch.profiler` (Kineto) traces from a running PyTorch workload without editing the model script. Instead of wrapping the command, this tool starts the [dynolog](https://github.com/facebookincubator/dynolog) daemon inside the container; PyTorch registers with it because the tool sets `KINETO_USE_DAEMON=1`, and a background script then requests a trace with `dyno gputrace`. + +```json +{ + "tools": [ + {"name": "torch_profiler_dynolog"} + ] +} +``` + +**Output:** `torch_profiler_output/libkineto_trace_.json` (one file per rank) + +**Requirements:** + +- The workload must be PyTorch >= 1.13. Nothing is captured from non-PyTorch models. +- Iteration-based capture counts `optimizer.step()` calls. Workloads without an optimizer (pure inference) should set `TORCH_PROFILE_ITERATIONS` to `0` to fall back to duration-based capture. +- The pre-script downloads the dynolog `.deb` from GitHub, so the container needs outbound HTTPS on the first run (x86_64 Debian/Ubuntu base image). Set `DYNOLOG_DEB_URL` to use a mirror, or bake `dynolog` and `dyno` into the image to skip the download entirely. + +**Environment Variables:** + +| Variable | Default | Purpose | +|----------|---------|---------| +| `TORCH_PROFILE_WARMUP_S` | `60` | Delay before the first trace request, so the workload reaches steady state | +| `TORCH_PROFILE_ITERATIONS` | `5` | Iterations to capture. `0` switches to `TORCH_PROFILE_DURATION_MS` | +| `TORCH_PROFILE_DURATION_MS` | `500` | Capture window when iteration counting is not used | +| `TORCH_PROFILE_RETRY_INTERVAL_S` | `15` | Wait between attempts while no PyTorch process has registered | +| `TORCH_PROFILE_MAX_ATTEMPTS` | `40` | Attempts before giving up | +| `TORCH_PROFILE_PROCESS_LIMIT` | `64` | Ranks to trace. Upstream defaults to 3, which silently drops most ranks | +| `TORCH_PROFILE_RECORD_SHAPES` | `1` | Record input shapes (needed for TraceLens per-operator analysis) | +| `TORCH_PROFILE_WITH_STACKS` | `1` | Record CPU call stacks | +| `TORCH_PROFILE_WITH_MODULES` | `1` | Record the `nn.Module` hierarchy | +| `TORCH_PROFILE_WITH_FLOPS` | `0` | Estimate FLOPs per operator | +| `TORCH_PROFILE_PROFILE_MEMORY` | `0` | Record allocator events | + +For a short-lived workload, shorten the warmup so the request lands while the model is still running: + +```json +{ + "tools": [ + { + "name": "torch_profiler_dynolog", + "env_vars": { + "TORCH_PROFILE_WARMUP_S": "20", + "TORCH_PROFILE_MAX_ATTEMPTS": "10" + } + } + ] +} +``` + +**No trace produced?** The teardown script reports how many traces it found. `no PyTorch process registered` means `dyno gputrace` never matched the workload: confirm the model is PyTorch, and raise `TORCH_PROFILE_WARMUP_S` and `TORCH_PROFILE_MAX_ATTEMPTS` for slow-starting jobs. + +### tracelens - TraceLens Trace Analysis + +[TraceLens](https://github.com/AMD-AGI/TraceLens) turns raw traces into operator, kernel, roofline, and collective reports. It analyzes traces the other profiling tools produce, so always pair it with a profiler; on its own it has nothing to read. + +Each trace format is routed to the matching TraceLens report: + +| Trace produced by | Format | TraceLens report | +|-------------------|--------|------------------| +| `torch_profiler_dynolog` | Kineto JSON | Per-operator, kernel summary, roofline, `nn.Module` breakdown | +| `rocprofv3_lightweight` | rocprofv3 JSON | Kernel summary and details | +| `rocprofv3_perfetto` | `.pftrace` | HIP activity, HIP API, and memory-copy reports | +| Multiple per-rank Kineto traces | Kineto JSON | Multi-rank collective report | + +**Unreadable formats:** TraceLens cannot read rocprofv3's default SQLite (`*_results.db`) or RPD (`.rpd`) databases. Those are listed in the summary as `SKIPPED` with a pointer at a preset that works — use `rocprofv3_lightweight` for JSON or `rocprofv3_perfetto` for `.pftrace`. For `rpd`, point TraceLens at the `trace.json` its post-script writes alongside the database. + +#### Analyzing on the Host (Recommended) + +Running analysis on the host keeps TraceLens' pinned `protobuf` and `xprof` out of your workload image: + +```bash +pip install 'madengine[tracelens]' + +# Analyze every trace a run left behind +madengine report tracelens + +# See what would be analyzed, without running TraceLens +madengine report tracelens --discover-only + +# Enable roofline bound classification +madengine report tracelens --gpu-arch MI300X + +# Analyze a distributed run's collected artifacts +madengine report tracelens --root slurm_results --mode collective --world-size 8 +``` + +**Output:** `tracelens_output/` with per-trace reports plus `tracelens_summary.csv` and `tracelens_summary.json` + +Quantify the effect of a change by diffing two runs' reports. The first report is the baseline; every metric gains `_diff` and `_pct` columns relative to it: + +```bash +madengine report tracelens-compare baseline.xlsx candidate.xlsx \ + --names before --names after -o diff.xlsx +``` + +#### Analyzing in the Container + +Stack the `tracelens` tool **after** a profiler to get reports as part of the run itself: + +```json +{ + "tools": [ + {"name": "rocprofv3_lightweight"}, + {"name": "tracelens"} + ] +} +``` + +**Output:** `tracelens_output/` copied to the working directory alongside the profiler's own output + +The pre-script installs TraceLens into an isolated virtualenv at `/opt/madengine-tracelens-venv` (no system site-packages), so its dependency pins cannot disturb the workload's Python environment. This needs outbound HTTPS to `github.com` on the first run. Analysis failures are reported as warnings and never fail a passing model run. + +Use a mode-specific variant to restrict analysis to one trace kind: + +| Tool | Analyzes | +|------|----------| +| `tracelens` | Every supported trace found (default) | +| `tracelens_pytorch` | Kineto traces only | +| `tracelens_rocprof` | rocprofv3 JSON only | +| `tracelens_pftrace` | Perfetto traces only | +| `tracelens_collective` | Multi-rank collective report only | + +**Full PyTorch pipeline** — capture and analyze in one run: + +```bash +madengine run --tags your_model \ + --additional-context '{ + "gpu_vendor": "AMD", + "guest_os": "UBUNTU", + "tools": [ + {"name": "torch_profiler_dynolog"}, + {"name": "tracelens"} + ] + }' +``` + +**Environment Variables:** + +| Variable | Default | Purpose | +|----------|---------|---------| +| `TRACELENS_VENV` | `/opt/madengine-tracelens-venv` | Virtualenv holding TraceLens | +| `TRACELENS_OUTPUT_DIR` | `tracelens_output` | Directory for generated reports | +| `TRACELENS_MODE` | `auto` | `auto`, `pytorch`, `rocprof`, `pftrace`, or `collective` | +| `TRACELENS_GPU_ARCH` | unset | GPU arch (e.g. `MI300X`) enabling roofline bound classification | +| `TRACELENS_WORLD_SIZE` | trace count | Rank count for the collective report | +| `TRACELENS_MAX_TRACES` | unset | Cap traces analyzed per kind | +| `TRACELENS_GIT_REF` | pinned commit | TraceLens revision to install | +| `TRACELENS_PIP_SPEC` | git URL at `TRACELENS_GIT_REF` | Full pip spec, for private mirrors | + +#### Distributed Runs + +SLURM and Kubernetes runs collect `torch_profiler_output/` and `tracelens_output/` alongside the other profiling directories, per node and per rank. Analyze the collected tree on the host by pointing `--root` at it: + +```bash +madengine report tracelens --root slurm_results +madengine report tracelens --root k8s_results --mode collective --world-size 16 +``` + ### rocblas_trace - rocBLAS Library Tracing Trace rocBLAS API calls and configurations: @@ -681,6 +842,8 @@ madengine run --tags model \ |------|---------------|---------| | `rocprof` | `rocprof_output/*` | GPU kernel traces, HIP API calls | | `rpd` | Various RPD files | ROCm profiler data | +| `torch_profiler_dynolog` | `torch_profiler_output/*.json` | Kineto traces, one per rank | +| `tracelens` | `tracelens_output/*` | TraceLens reports plus `tracelens_summary.csv` | | `rocblas_trace` | `library_trace.csv`, logs | rocBLAS API calls | | `miopen_trace` | `library_trace.csv`, logs | MIOpen API calls | | `tensile_trace` | `library_trace.csv`, logs | Tensile operations | diff --git a/examples/profiling-configs/README.md b/examples/profiling-configs/README.md index 51130ef4..a88c5271 100644 --- a/examples/profiling-configs/README.md +++ b/examples/profiling-configs/README.md @@ -110,7 +110,29 @@ madengine run --tags your_model \ --additional-context-file examples/profiling-configs/rocm_trace_lite_default.json ``` -### 7. Multi-Node Distributed (`rocprofv3_multi_node.json`) +### 7. TraceLens Analysis (`torch_profiler_tracelens.json`, `tracelens_rocprofv3.json`) + +**Use Case**: Turn raw traces into operator, kernel, roofline, and collective reports with [TraceLens](https://github.com/AMD-AGI/TraceLens). TraceLens only analyzes traces, so it is always stacked after a profiler. + +- **`torch_profiler_tracelens.json`** — captures `torch.profiler` (Kineto) traces on demand via dynolog, then generates the PyTorch operator and roofline reports. The workload must be PyTorch, and iteration-based capture needs an `optimizer.step()` loop. +- **`tracelens_rocprofv3.json`** — profiles with `rocprofv3_lightweight` (JSON) and generates the kernel summary reports. + +**Requirements / notes:** + +- The trace pre-script installs TraceLens into an isolated virtualenv at `/opt/madengine-tracelens-venv`, so its pinned `protobuf` and `xprof` cannot disturb the workload's Python environment. The container needs **HTTPS access to GitHub** on the first run. Override the revision with `TRACELENS_GIT_REF`, or the whole spec with `TRACELENS_PIP_SPEC`. +- TraceLens cannot read rocprofv3's default SQLite (`*_results.db`) or RPD databases. Use `rocprofv3_lightweight` (JSON) or `rocprofv3_perfetto` (`.pftrace`); unreadable artifacts are listed in `tracelens_output/tracelens_summary.csv` as `SKIPPED` with guidance. +- To keep TraceLens out of the workload image entirely, skip the in-container tool and analyze on the host instead: `pip install 'madengine[tracelens]'` then `madengine report tracelens`. See the [Profiling Guide](../../docs/profiling.md#tracelens---tracelens-trace-analysis). + +**Usage**: +```bash +madengine run --tags your_model \ + --additional-context-file examples/profiling-configs/torch_profiler_tracelens.json + +madengine run --tags your_model \ + --additional-context-file examples/profiling-configs/tracelens_rocprofv3.json +``` + +### 8. Multi-Node Distributed (`rocprofv3_multi_node.json`) **Use Case**: Large-scale distributed training on SLURM clusters @@ -217,6 +239,8 @@ The wrapper script auto-detects which profiler is available and formats the comm **Other:** `rocm_trace_lite` (RTL **lite** mode) and `rocm_trace_lite_default` (RTL **default** mode) — kernel dispatch SQLite trace via [rocm-trace-lite](https://sunway513.github.io/rocm-trace-lite/index.html), installed from **GitHub Release wheels** by the trace pre-script (not PyPI; see [Profiling Guide](../../docs/profiling.md)). Not a rocprofv3 preset; do not combine with `rocprof` / `rocprofv3_*` on the same run. +**PyTorch tracing and analysis:** `torch_profiler_dynolog` captures Kineto traces on demand, and `tracelens` (or the mode-specific `tracelens_pytorch`, `tracelens_rocprof`, `tracelens_pftrace`, `tracelens_collective`) generates TraceLens reports from whatever traces a co-selected profiler produced. Neither is a rocprofv3 preset; `tracelens` is analysis only and must be stacked after a profiler. + ## Counter Definition Files Counter files are located at `src/madengine/scripts/common/tools/counters/`: @@ -247,6 +271,9 @@ gpu_info_vram_profiler_output.csv # VRAM usage over time library_trace.csv # Library API calls (if library tracing enabled) rocm_trace_lite_output/trace.db # rocm-trace-lite (also trace.json.gz / trace_summary.txt as emitted by RTL) + +torch_profiler_output/*.json # Kineto traces, one per rank (torch_profiler_dynolog) +tracelens_output/ # TraceLens reports plus tracelens_summary.csv ``` ## Visualization diff --git a/examples/profiling-configs/torch_profiler_tracelens.json b/examples/profiling-configs/torch_profiler_tracelens.json new file mode 100644 index 00000000..89da2acf --- /dev/null +++ b/examples/profiling-configs/torch_profiler_tracelens.json @@ -0,0 +1,19 @@ +{ + "gpu_vendor": "AMD", + "guest_os": "UBUNTU", + "tools": [ + { + "name": "torch_profiler_dynolog", + "env_vars": { + "TORCH_PROFILE_ITERATIONS": "5", + "TORCH_PROFILE_WARMUP_S": "60" + } + }, + { + "name": "tracelens", + "env_vars": { + "TRACELENS_MODE": "pytorch" + } + } + ] +} diff --git a/examples/profiling-configs/tracelens_rocprofv3.json b/examples/profiling-configs/tracelens_rocprofv3.json new file mode 100644 index 00000000..f5b85df5 --- /dev/null +++ b/examples/profiling-configs/tracelens_rocprofv3.json @@ -0,0 +1,15 @@ +{ + "gpu_vendor": "AMD", + "guest_os": "UBUNTU", + "tools": [ + { + "name": "rocprofv3_lightweight" + }, + { + "name": "tracelens", + "env_vars": { + "TRACELENS_MODE": "rocprof" + } + } + ] +} diff --git a/pyproject.toml b/pyproject.toml index dd9c7566..cc845c66 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,14 @@ classifiers = [ "Operating System :: OS Independent", ] +[project.optional-dependencies] +# Host-side trace analysis for `madengine report tracelens`. Kept optional +# because TraceLens pins protobuf>=6.31 and xprof, which can conflict with other +# packages in a shared environment. +tracelens = [ + "TraceLens @ git+https://github.com/AMD-AGI/TraceLens.git@6f9bcdbf6cc9911eb650de57b345917ea4d31a17", +] + [project.scripts] madengine = "madengine.cli.app:cli_main" diff --git a/src/madengine/cli/commands/report.py b/src/madengine/cli/commands/report.py index 2bd348c0..8bb6ce43 100644 --- a/src/madengine/cli/commands/report.py +++ b/src/madengine/cli/commands/report.py @@ -10,9 +10,12 @@ import os from pathlib import Path +from typing import List, Optional import typer +from rich.markup import escape from rich.panel import Panel +from rich.table import Table try: from typing import Annotated # Python 3.9+ @@ -21,6 +24,12 @@ from madengine.reporting.csv_to_html import ConvertCsvToHtml from madengine.reporting.csv_to_email import ConvertCsvToEmail +from madengine.reporting.tracelens_report import ( + TraceLensNotInstalledError, + compare_tracelens_reports, + discover_traces, + generate_tracelens_reports, +) from ..constants import ExitCode from ..utils import console, setup_logging, create_args_namespace @@ -182,6 +191,278 @@ def to_email( raise typer.Exit(ExitCode.FAILURE) +def _print_tracelens_results(summary: dict) -> None: + """Render the analyzer's per-trace results as a table.""" + results = summary.get("results") or [] + if not results: + return + + table = Table(title="TraceLens reports", show_lines=False) + table.add_column("Status") + table.add_column("Trace", overflow="fold") + table.add_column("Kind") + table.add_column("Report") + table.add_column("Detail", overflow="fold") + + # Trace paths and TraceLens error text routinely contain square brackets + # (e.g. "[rocprofv3]"), which rich would otherwise parse as markup. + styles = {"SUCCESS": "green", "FAILURE": "red", "SKIPPED": "yellow"} + for result in results: + status = str(result.get("status", "")) + table.add_row( + f"[{styles.get(status, 'white')}]{status}[/]", + escape(str(result.get("trace_file", ""))), + escape(str(result.get("kind", ""))), + escape( + str(result.get("tracelens_tool", "")).replace( + "TraceLens_generate_perf_report_", "" + ) + ), + escape(str(result.get("detail", ""))), + ) + console.print(table) + + +@report_app.command("tracelens") +def tracelens( + root: Annotated[ + str, + typer.Option( + "--root", + "-r", + help="Directory to search recursively for trace artifacts", + ), + ] = ".", + output_dir: Annotated[ + str, + typer.Option("--output-dir", "-o", help="Directory for generated reports"), + ] = "tracelens_output", + mode: Annotated[ + str, + typer.Option( + "--mode", + help="Restrict analysis to one trace kind: auto, pytorch, rocprof, pftrace, collective", + ), + ] = "auto", + python: Annotated[ + Optional[str], + typer.Option("--python", help="Interpreter that has TraceLens installed"), + ] = None, + gpu_arch: Annotated[ + Optional[str], + typer.Option( + "--gpu-arch", + help="TraceLens GPU arch platform (e.g. MI300X) for roofline bound classification", + ), + ] = None, + world_size: Annotated[ + int, + typer.Option( + "--world-size", + help="Rank count for the collective report (default: number of traces found)", + ), + ] = 0, + max_traces: Annotated[ + int, + typer.Option("--max-traces", help="Cap traces analyzed per kind (0 = no cap)"), + ] = 0, + discover_only: Annotated[ + bool, + typer.Option( + "--discover-only", + help="List discovered traces without running TraceLens", + ), + ] = False, + verbose: Annotated[ + bool, typer.Option("--verbose", "-v", help="Enable verbose logging") + ] = False, +) -> None: + """ + 🔬 Generate TraceLens performance reports from collected GPU traces. + + Discovers trace artifacts a run left behind (rocprof_output/, + torch_profiler_output/, slurm_results/, k8s_results/) and generates the + matching TraceLens report for each: operator and roofline analysis for + torch.profiler traces, kernel summaries for rocprofv3 JSON, and + activity/API/memory-copy reports for pftrace. + + Requires TraceLens: pip install 'madengine[tracelens]' + + Examples: + madengine report tracelens + madengine report tracelens --discover-only + madengine report tracelens --gpu-arch MI300X + madengine report tracelens --root slurm_results --mode collective --world-size 8 + """ + setup_logging(verbose) + + valid_modes = ("auto", "pytorch", "rocprof", "pftrace", "collective") + if mode not in valid_modes: + console.print( + f"❌ [bold red]Error: invalid --mode '{mode}'. " + f"Choose one of: {', '.join(valid_modes)}[/bold red]" + ) + raise typer.Exit(ExitCode.INVALID_ARGS) + + if not os.path.isdir(root): + console.print(f"❌ [bold red]Error: directory not found: {root}[/bold red]") + raise typer.Exit(ExitCode.FAILURE) + + console.print( + Panel( + f"🔬 [bold cyan]TraceLens Analysis[/bold cyan]\n" + f"Search root: [yellow]{root}[/yellow]\n" + f"Output directory: [yellow]{output_dir}[/yellow]\n" + f"Mode: [yellow]{mode}[/yellow]", + title="TraceLens Report", + border_style="blue", + ) + ) + + try: + if discover_only: + summary = discover_traces(root=root, output_dir=output_dir) + discovered = summary.get("discovered") or {} + if not discovered and not summary.get("unsupported"): + console.print( + f"⚠️ [yellow]No trace artifacts found under {root}[/yellow]" + ) + for kind, count in discovered.items(): + console.print(f" [cyan]{kind}[/cyan]: {count} trace(s)") + for item in summary.get("unsupported") or []: + console.print( + f" [yellow]unsupported[/yellow]: {escape(item['path'])} — " + f"{escape(item['reason'])}" + ) + return + + summary = generate_tracelens_reports( + root=root, + output_dir=output_dir, + mode=mode, + python=python, + gpu_arch=gpu_arch, + world_size=world_size, + max_traces=max_traces, + ) + except TraceLensNotInstalledError as e: + console.print(f"❌ [bold red]{escape(str(e))}[/bold red]") + raise typer.Exit(ExitCode.FAILURE) + except Exception as e: + console.print(f"💥 [bold red]TraceLens analysis failed: {escape(str(e))}[/bold red]") + if verbose: + console.print_exception() + raise typer.Exit(ExitCode.FAILURE) + + _print_tracelens_results(summary) + + succeeded = int(summary.get("succeeded", 0)) + failed = int(summary.get("failed", 0)) + skipped = int(summary.get("skipped", 0)) + + if not succeeded and not failed: + console.print( + f"⚠️ [yellow]No supported trace artifacts found under {root}. " + "Stack a profiling tool (torch_profiler_dynolog, rocprofv3_lightweight, " + "rocprofv3_perfetto) on the run first.[/yellow]" + ) + return + + console.print( + f"📄 [bold]Reports written to:[/bold] [yellow]{output_dir}[/yellow] " + f"(summary: {os.path.join(output_dir, 'tracelens_summary.csv')})" + ) + if failed: + console.print( + f"⚠️ [yellow]{succeeded} report(s) generated, {failed} failed, " + f"{skipped} skipped[/yellow]" + ) + raise typer.Exit(ExitCode.FAILURE) + + console.print( + f"✅ [bold green]{succeeded} report(s) generated" + + (f", {skipped} skipped" if skipped else "") + + "[/bold green]" + ) + + +@report_app.command("tracelens-compare") +def tracelens_compare( + reports: Annotated[ + List[str], + typer.Argument( + help="Two or more TraceLens reports (.xlsx files or per-sheet CSV directories)" + ), + ], + output: Annotated[ + str, + typer.Option("--output", "-o", help="Output comparison workbook"), + ] = "tracelens_comparison.xlsx", + names: Annotated[ + Optional[List[str]], + typer.Option("--names", help="Display tag per report (repeat the flag)"), + ] = None, + python: Annotated[ + Optional[str], + typer.Option("--python", help="Interpreter that has TraceLens installed"), + ] = None, + verbose: Annotated[ + bool, typer.Option("--verbose", "-v", help="Enable verbose logging") + ] = False, +) -> None: + """ + ⚖️ Compare two or more TraceLens reports into a single diff workbook. + + The first report is the baseline; every metric gains ``_diff`` and ``_pct`` + columns relative to it. Use this to quantify the effect of a change across + two madengine runs. + + Examples: + madengine report tracelens-compare baseline.xlsx candidate.xlsx + madengine report tracelens-compare a.xlsx b.xlsx --names before --names after -o diff.xlsx + """ + setup_logging(verbose) + + missing = [r for r in reports if not os.path.exists(r)] + if missing: + console.print( + f"❌ [bold red]Error: report(s) not found: {', '.join(missing)}[/bold red]" + ) + raise typer.Exit(ExitCode.FAILURE) + + console.print( + Panel( + f"⚖️ [bold cyan]Comparing TraceLens Reports[/bold cyan]\n" + f"Reports: [yellow]{', '.join(reports)}[/yellow]\n" + f"Output: [yellow]{output}[/yellow]", + title="TraceLens Comparison", + border_style="blue", + ) + ) + + try: + summary = compare_tracelens_reports( + reports=reports, output=output, names=names or (), python=python + ) + except (TraceLensNotInstalledError, ValueError) as e: + console.print(f"❌ [bold red]{escape(str(e))}[/bold red]") + raise typer.Exit(ExitCode.FAILURE) + except Exception as e: + console.print(f"💥 [bold red]Comparison failed: {escape(str(e))}[/bold red]") + if verbose: + console.print_exception() + raise typer.Exit(ExitCode.FAILURE) + + if summary.get("status") == "SUCCESS": + console.print(f"✅ [bold green]Comparison written to: {output}[/bold green]") + else: + console.print( + f"💥 [bold red]Comparison failed: " + f"{escape(str(summary.get('detail', '')))}[/bold red]" + ) + raise typer.Exit(ExitCode.FAILURE) + + # Export the report app def report() -> typer.Typer: """Return the report sub-app.""" diff --git a/src/madengine/deployment/k8s_results.py b/src/madengine/deployment/k8s_results.py index 6da189b5..3638d1ba 100644 --- a/src/madengine/deployment/k8s_results.py +++ b/src/madengine/deployment/k8s_results.py @@ -522,6 +522,7 @@ def _collect_pod_artifacts(self, pod_name: str, dest_dir: Path) -> List[Dict]: {"pattern": "results*", "type": "profiling"}, {"pattern": "*.db", "type": "profiling"}, {"pattern": "trace.*", "type": "tracing"}, + {"pattern": "*.pftrace", "type": "tracing"}, {"pattern": "prof.csv", "type": "profiling"}, # Raw profiler output before post-script renames it {"pattern": "gpu_info_*.csv", "type": "profiling"}, {"pattern": "library_trace.csv", "type": "tracing"}, @@ -611,7 +612,14 @@ def _collect_pod_artifacts(self, pod_name: str, dest_dir: Path) -> List[Dict]: pass # File not found or not accessible - this is expected # Try to collect known output directories using kubectl cp directly (during sleep period) - output_directories = ["rocprof_output", "rpd_output", "trace_output"] + output_directories = [ + "rocprof_output", + "rpd_output", + "trace_output", + "rocm_trace_lite_output", + "torch_profiler_output", + "tracelens_output", + ] for dir_name in output_directories: try: local_dir = dest_dir / dir_name diff --git a/src/madengine/deployment/k8s_scripts.py b/src/madengine/deployment/k8s_scripts.py index 277661e0..dfe44abc 100644 --- a/src/madengine/deployment/k8s_scripts.py +++ b/src/madengine/deployment/k8s_scripts.py @@ -222,14 +222,20 @@ def _load_tool_wrapper_scripts(self, script_contents: Dict[str, str], script_contents[script_path] = f.read() self.console.print(f"[dim]Loaded tool post-script: {script_path}[/dim]") - for script_config in tool_def.get("pre_scripts", []): + # Pre/post scripts may invoke helper scripts under scripts/common/tools/ + # (e.g. dynolog_start.sh -> dynolog_trigger.sh, tracelens.sh -> + # tracelens_analyze.py). Those helpers are not named in tools.json, so + # bundle whatever the scripts reference. + for script_config in tool_def.get("pre_scripts", []) + tool_def.get( + "post_scripts", [] + ): script_path = script_config.get("path", "") if script_path: abs_script_path = madengine_root / script_path if abs_script_path.exists(): with open(abs_script_path, "r") as f: script_content = f.read() - tool_refs = re.findall(r'(?:\.\./)?scripts/common/tools/[\w_]+\.py', script_content) + tool_refs = re.findall(r'(?:\.\./)?scripts/common/tools/[\w_]+\.(?:py|sh)', script_content) for tool_ref in tool_refs: tool_script_path = tool_ref.strip('"\'').replace("../", "") abs_tool_path = madengine_root / tool_script_path diff --git a/src/madengine/deployment/templates/kubernetes/job.yaml.j2 b/src/madengine/deployment/templates/kubernetes/job.yaml.j2 index 320d049f..082157d9 100644 --- a/src/madengine/deployment/templates/kubernetes/job.yaml.j2 +++ b/src/madengine/deployment/templates/kubernetes/job.yaml.j2 @@ -386,11 +386,14 @@ spec: echo "✓ Copied rocprofv3 directory: $dir" fi done - # rocm-trace-lite (SQLite and optional Perfetto JSON beside trace.db) - if [ -d "rocm_trace_lite_output" ]; then - cp -r rocm_trace_lite_output /results/${HOSTNAME}/ 2>/dev/null || true - echo "✓ Copied rocm_trace_lite_output" - fi + # Per-tool output directories: rocm-trace-lite SQLite, RPD, Kineto + # traces from torch_profiler_dynolog, and TraceLens reports. + for tool_dir in rocm_trace_lite_output rpd_output torch_profiler_output tracelens_output; do + if [ -d "$tool_dir" ]; then + cp -r "$tool_dir" /results/${HOSTNAME}/ 2>/dev/null || true + echo "✓ Copied $tool_dir" + fi + done # Copy tool-specific outputs if ls gpu_info_*.csv 1> /dev/null 2>&1; then @@ -575,11 +578,14 @@ spec: echo "✓ Copied rocprofv3 directory: $dir" fi done - # rocm-trace-lite (SQLite under rocm_trace_lite_output/) - if [ -d "rocm_trace_lite_output" ]; then - cp -r rocm_trace_lite_output /results/${HOSTNAME}/ 2>/dev/null || true - echo "✓ Copied rocm_trace_lite_output" - fi + # Per-tool output directories: rocm-trace-lite SQLite, RPD, Kineto + # traces from torch_profiler_dynolog, and TraceLens reports. + for tool_dir in rocm_trace_lite_output rpd_output torch_profiler_output tracelens_output; do + if [ -d "$tool_dir" ]; then + cp -r "$tool_dir" /results/${HOSTNAME}/ 2>/dev/null || true + echo "✓ Copied $tool_dir" + fi + done # Copy GPU profiler outputs if ls gpu_info_*.csv 1> /dev/null 2>&1; then diff --git a/src/madengine/deployment/templates/slurm/job.sh.j2 b/src/madengine/deployment/templates/slurm/job.sh.j2 index 3b236d9f..9dd188a2 100644 --- a/src/madengine/deployment/templates/slurm/job.sh.j2 +++ b/src/madengine/deployment/templates/slurm/job.sh.j2 @@ -694,7 +694,9 @@ if [ $TASK_EXIT -eq 0 ]; then if [ -f "$WORKSPACE/run_directory/{{ multiple_results }}" ]; then cp "$WORKSPACE/run_directory/{{ multiple_results }}" "$NODE_COLLECTION_DIR/" 2>/dev/null || true; fi if [ -f "$WORKSPACE/{{ multiple_results }}" ]; then cp "$WORKSPACE/{{ multiple_results }}" "$NODE_COLLECTION_DIR/" 2>/dev/null || true; fi {% endif %} - if [ -d "$WORKSPACE/rocprof_output" ]; then cp -r "$WORKSPACE/rocprof_output" "$NODE_COLLECTION_DIR/" 2>/dev/null || true; fi + for d in rocprof_output rpd_output rocm_trace_lite_output torch_profiler_output tracelens_output; do + if [ -d "$WORKSPACE/$d" ]; then cp -r "$WORKSPACE/$d" "$NODE_COLLECTION_DIR/" 2>/dev/null || true; fi + done echo " ✓ Node ${SLURM_PROCID} artifacts copied" echo "========================================================================" else diff --git a/src/madengine/reporting/tracelens_report.py b/src/madengine/reporting/tracelens_report.py new file mode 100644 index 00000000..6aabc22d --- /dev/null +++ b/src/madengine/reporting/tracelens_report.py @@ -0,0 +1,233 @@ +"""Host-side TraceLens report generation for collected madengine trace artifacts. + +This module drives ``scripts/common/tools/tracelens_analyze.py``, the same +analyzer the in-container ``tracelens`` tool runs, against artifacts that a run +already copied back to the host working directory (``rocprof_output/``, +``torch_profiler_output/``, ``slurm_results/``, ``k8s_results/``, ...). + +Running TraceLens on the host keeps its pinned ``protobuf`` and ``xprof`` +dependencies out of the workload container. Install support with:: + + pip install 'madengine[tracelens]' + +Copyright (c) Advanced Micro Devices, Inc. All rights reserved. +""" + +import json +import logging +import os +import subprocess +import sys +import tempfile +from pathlib import Path +from typing import Dict, List, Optional, Sequence + +from madengine.utils.path_utils import get_madengine_root + +logger = logging.getLogger(__name__) + +ANALYZER_RELATIVE_PATH = Path("scripts") / "common" / "tools" / "tracelens_analyze.py" + +INSTALL_HINT = ( + "TraceLens is not importable by the selected interpreter. Install it with " + "\"pip install 'madengine[tracelens]'\", or point --python at an " + "environment that has it (or set TRACELENS_VENV)." +) + + +class TraceLensNotInstalledError(RuntimeError): + """Raised when the selected interpreter cannot import TraceLens.""" + + +def find_analyzer_script() -> Path: + """Return the path to the packaged trace analyzer script. + + Raises: + FileNotFoundError: If the script is missing from the installation. + """ + script = get_madengine_root() / ANALYZER_RELATIVE_PATH + if not script.is_file(): + raise FileNotFoundError( + f"Trace analyzer not found at {script}. The madengine installation " + "appears to be missing its bundled scripts." + ) + return script + + +def resolve_python(python: Optional[str] = None) -> str: + """Return the interpreter used to run TraceLens itself. + + Prefers an explicit ``python``, then ``$TRACELENS_VENV``, then the + interpreter running madengine. + """ + if python: + return python + venv = os.environ.get("TRACELENS_VENV", "").strip() + if venv: + for candidate in ( + Path(venv) / "bin" / "python3", + Path(venv) / "Scripts" / "python.exe", + ): + if candidate.is_file(): + return str(candidate) + return sys.executable + + +def check_tracelens_available(python: str) -> bool: + """Return True if ``python`` can import TraceLens.""" + try: + completed = subprocess.run( + [python, "-c", "import TraceLens"], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + except OSError: + return False + return completed.returncode == 0 + + +def _run_analyzer(args: Sequence[str], summary_path: Path) -> Dict[str, object]: + """Run the analyzer script and return the summary it wrote.""" + script = find_analyzer_script() + command = [sys.executable, str(script), *args] + logger.info("Running trace analyzer: %s", " ".join(command)) + completed = subprocess.run(command) + + summary: Dict[str, object] = {} + if summary_path.is_file(): + with open(summary_path, encoding="utf-8") as handle: + summary = json.load(handle) + summary["exit_code"] = completed.returncode + return summary + + +def generate_tracelens_reports( + root: str = ".", + output_dir: str = "tracelens_output", + mode: str = "auto", + python: Optional[str] = None, + gpu_arch: Optional[str] = None, + world_size: int = 0, + max_traces: int = 0, + extra_args: Sequence[str] = (), +) -> Dict[str, object]: + """Generate TraceLens reports for every supported trace found under ``root``. + + Args: + root: Directory to search recursively for trace artifacts. + output_dir: Directory that receives the generated reports. + mode: ``auto``, or one of ``pytorch``, ``rocprof``, ``pftrace``, + ``collective`` to restrict the run to a single trace kind. + python: Interpreter that has TraceLens installed. + gpu_arch: Bundled TraceLens GPU arch name (e.g. ``MI300X``) that enables + roofline bound classification on PyTorch reports. + world_size: Rank count for the multi-rank collective report. Inferred + from the number of discovered PyTorch traces when 0. + max_traces: Cap on traces analyzed per kind. 0 means no cap. + extra_args: Extra flags forwarded verbatim to every TraceLens command. + + Returns: + The analyzer summary, including ``results``, ``discovered``, counters, + and ``exit_code``. + + Raises: + TraceLensNotInstalledError: If the interpreter cannot import TraceLens. + """ + interpreter = resolve_python(python) + if not check_tracelens_available(interpreter): + raise TraceLensNotInstalledError(INSTALL_HINT) + + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + summary_path = output_path / "tracelens_summary.json" + + args: List[str] = [ + "--root", + root, + "--output-dir", + str(output_path), + "--mode", + mode, + "--python", + interpreter, + "--json-summary", + str(summary_path), + ] + if gpu_arch: + args += ["--gpu-arch", gpu_arch] + if world_size: + args += ["--world-size", str(world_size)] + if max_traces: + args += ["--max-traces", str(max_traces)] + args += list(extra_args) + + return _run_analyzer(args, summary_path) + + +def discover_traces(root: str = ".", output_dir: str = "tracelens_output") -> Dict[str, object]: + """List the trace artifacts under ``root`` without running TraceLens. + + Unlike :func:`generate_tracelens_reports`, this does not require TraceLens to + be installed. + """ + with tempfile.TemporaryDirectory() as tmp: + summary_path = Path(tmp) / "discovery.json" + return _run_analyzer( + [ + "--root", + root, + "--output-dir", + output_dir, + "--discover-only", + "--json-summary", + str(summary_path), + ], + summary_path, + ) + + +def compare_tracelens_reports( + reports: Sequence[str], + output: str = "tracelens_comparison.xlsx", + names: Sequence[str] = (), + python: Optional[str] = None, +) -> Dict[str, object]: + """Diff two or more TraceLens reports into a single comparison workbook. + + Args: + reports: TraceLens ``.xlsx`` reports or per-sheet CSV directories. The + first is treated as the baseline. + output: Output workbook path. + names: Display tags, one per report. + python: Interpreter that has TraceLens installed. + + Raises: + TraceLensNotInstalledError: If the interpreter cannot import TraceLens. + ValueError: If fewer than two reports were given. + """ + if len(reports) < 2: + raise ValueError("Comparing reports needs at least two inputs.") + if names and len(names) != len(reports): + raise ValueError( + f"Got {len(names)} names for {len(reports)} reports; counts must match." + ) + + interpreter = resolve_python(python) + if not check_tracelens_available(interpreter): + raise TraceLensNotInstalledError(INSTALL_HINT) + + with tempfile.TemporaryDirectory() as tmp: + summary_path = Path(tmp) / "comparison.json" + args: List[str] = [ + "--python", + interpreter, + "--compare", + *reports, + "--compare-output", + output, + "--json-summary", + str(summary_path), + ] + if names: + args += ["--compare-names", *names] + return _run_analyzer(args, summary_path) diff --git a/src/madengine/scripts/common/post_scripts/dynolog_stop.sh b/src/madengine/scripts/common/post_scripts/dynolog_stop.sh new file mode 100644 index 00000000..7fbc9472 --- /dev/null +++ b/src/madengine/scripts/common/post_scripts/dynolog_stop.sh @@ -0,0 +1,84 @@ +#!/usr/bin/env bash +# +# Copyright (c) Advanced Micro Devices, Inc. +# All rights reserved. +# +# Stop the dynolog daemon and its trace trigger, then report what was captured. +# Artifact collection into /myworkspace is handled by post_scripts/trace.sh. + +set -x + +echo "Stopping dynolog daemon..." + +DYNOLOG_PID_FILE="/tmp/madengine_dynolog.pid" +TRIGGER_PID_FILE="/tmp/madengine_dynolog_trigger.pid" +DYNOLOG_START_FILE="/tmp/madengine_dynolog.started" +RESULT_FILE="/tmp/madengine_dynolog_trigger.result" + +OUTPUT_DIR=${TORCH_PROFILE_OUTPUT_DIR:-torch_profiler_output} + +if [ ! -f "$DYNOLOG_START_FILE" ]; then + echo "⚠️ Warning: dynolog was not started - skipping" + exit 0 +fi + +stop_pid() { + local name=$1 + local pid_file=$2 + if [ ! -f "$pid_file" ]; then + echo "⚠️ Warning: $name PID file not found" + return 0 + fi + local pid + pid=$(cat "$pid_file") + if kill -0 "$pid" 2>/dev/null; then + kill -TERM "$pid" 2>/dev/null || true + local waited=0 + while kill -0 "$pid" 2>/dev/null && [ $waited -lt 20 ]; do + sleep 0.5 + waited=$((waited + 1)) + done + if kill -0 "$pid" 2>/dev/null; then + echo "⚠️ $name did not stop gracefully, force killing..." + kill -9 "$pid" 2>/dev/null || true + fi + echo "✓ $name stopped (PID: $pid)" + else + echo "⚠️ $name (PID: $pid) was no longer running" + fi + rm -f "$pid_file" +} + +# Stop the trigger first so it cannot issue a request against a dying daemon. +stop_pid "dynolog trace trigger" "$TRIGGER_PID_FILE" +stop_pid "dynolog daemon" "$DYNOLOG_PID_FILE" +rm -f "$DYNOLOG_START_FILE" + +# Kineto appends the process id to the requested filename, so a multi-rank run +# produces one file per rank. +trace_count=0 +if [ -d "$OUTPUT_DIR" ]; then + trace_count=$(find "$OUTPUT_DIR" -maxdepth 1 -type f \( -name '*.json' -o -name '*.json.gz' \) 2>/dev/null | wc -l) +fi + +if [ "$trace_count" -gt 0 ]; then + echo "✓ Captured ${trace_count} torch.profiler trace(s) in ${OUTPUT_DIR}" + ls -la "$OUTPUT_DIR" || true +else + echo "⚠️ No torch.profiler traces were captured in ${OUTPUT_DIR}" + if [ -f "$RESULT_FILE" ] && [ "$(cat "$RESULT_FILE")" = "no_process" ]; then + echo "⚠️ The trigger never matched a PyTorch process. Most likely causes:" + echo "⚠️ - the workload is not PyTorch, or predates PyTorch 1.13" + echo "⚠️ - the model run finished before TORCH_PROFILE_WARMUP_S elapsed" + fi +fi + +for log in /tmp/madengine_dynolog.log /tmp/madengine_dynolog_trigger.log; do + if [ -f "$log" ]; then + echo "=== $(basename "$log") ===" + tail -40 "$log" || true + echo "==========================" + fi +done + +echo "✓ dynolog cleanup complete" diff --git a/src/madengine/scripts/common/post_scripts/trace.sh b/src/madengine/scripts/common/post_scripts/trace.sh index 1e489861..dc3f56f4 100644 --- a/src/madengine/scripts/common/post_scripts/trace.sh +++ b/src/madengine/scripts/common/post_scripts/trace.sh @@ -154,6 +154,25 @@ rocm_trace_lite) cp -vLR --preserve=all "$OUTPUT" "$SAVESPACE" || echo "Note: rocm_trace_lite output directory may be empty" ;; +torch_profiler) + # OUTPUT is torch_profiler_output, the directory Kineto wrote traces into + # (see TORCH_PROFILE_OUTPUT_DIR in tools.json). + echo "torch.profiler post-script: Collecting Kineto traces under ${OUTPUT}..." + if ! compgen -G "${OUTPUT}/*.json*" > /dev/null; then + echo "WARNING: no traces found in ${OUTPUT} (see the dynolog_stop.sh output above)." + fi + cp -vLR --preserve=all "$OUTPUT" "$SAVESPACE" || echo "Note: torch_profiler output directory may be empty" + ;; + +tracelens) + # OUTPUT is tracelens_output, written by post_scripts/tracelens.sh. + echo "TraceLens post-script: Collecting reports under ${OUTPUT}..." + if [ ! -f "${OUTPUT}/tracelens_summary.csv" ]; then + echo "WARNING: ${OUTPUT}/tracelens_summary.csv not found (analysis may have failed)." + fi + cp -vLR --preserve=all "$OUTPUT" "$SAVESPACE" || echo "Note: tracelens output directory may be empty" + ;; + esac chmod -R a+rw "${SAVESPACE}/${OUTPUT}" diff --git a/src/madengine/scripts/common/post_scripts/tracelens.sh b/src/madengine/scripts/common/post_scripts/tracelens.sh new file mode 100644 index 00000000..58acfcee --- /dev/null +++ b/src/madengine/scripts/common/post_scripts/tracelens.sh @@ -0,0 +1,53 @@ +#!/usr/bin/env bash +# +# Copyright (c) Advanced Micro Devices, Inc. +# All rights reserved. +# +# Generate TraceLens reports for whatever trace artifacts the co-selected +# profiling tool produced. Stack this after a profiler, for example: +# --additional-context '{"tools":[{"name":"rocprofv3_perfetto"},{"name":"tracelens"}]}' + +set -x + +TRACELENS_VENV=${TRACELENS_VENV:-/opt/madengine-tracelens-venv} +OUTPUT_DIR=${TRACELENS_OUTPUT_DIR:-tracelens_output} +MODE=${TRACELENS_MODE:-auto} + +if [ ! -x "${TRACELENS_VENV}/bin/python3" ]; then + echo "Error: TraceLens venv missing at ${TRACELENS_VENV}. The tracelens pre-script must run first." >&2 + exit 1 +fi + +if [ -f "scripts/common/tools/tracelens_analyze.py" ]; then + ANALYZER="scripts/common/tools/tracelens_analyze.py" +elif [ -f "../scripts/common/tools/tracelens_analyze.py" ]; then + ANALYZER="../scripts/common/tools/tracelens_analyze.py" +else + echo "Error: Cannot find tracelens_analyze.py" >&2 + exit 1 +fi + +ARGS=( + --root . + --output-dir "$OUTPUT_DIR" + --mode "$MODE" + --python "${TRACELENS_VENV}/bin/python3" + --json-summary "${OUTPUT_DIR}/tracelens_summary.json" +) +if [ -n "${TRACELENS_GPU_ARCH:-}" ]; then + ARGS+=(--gpu-arch "$TRACELENS_GPU_ARCH") +fi +if [ -n "${TRACELENS_WORLD_SIZE:-}" ]; then + ARGS+=(--world-size "$TRACELENS_WORLD_SIZE") +fi +if [ -n "${TRACELENS_MAX_TRACES:-}" ]; then + ARGS+=(--max-traces "$TRACELENS_MAX_TRACES") +fi + +mkdir -p "$OUTPUT_DIR" + +# Analysis is reporting, not the workload: a TraceLens failure must not turn a +# passing model run into a failure. +if ! python3 "$ANALYZER" "${ARGS[@]}"; then + echo "WARNING: TraceLens analysis reported failures; see ${OUTPUT_DIR}/tracelens_summary.csv" +fi diff --git a/src/madengine/scripts/common/pre_scripts/dynolog_start.sh b/src/madengine/scripts/common/pre_scripts/dynolog_start.sh new file mode 100644 index 00000000..f5fd5a37 --- /dev/null +++ b/src/madengine/scripts/common/pre_scripts/dynolog_start.sh @@ -0,0 +1,63 @@ +#!/usr/bin/env bash +# +# Copyright (c) Advanced Micro Devices, Inc. +# All rights reserved. +# +# Start the dynolog daemon and arm a background trigger that requests a +# torch.profiler trace once the workload's PyTorch processes have registered. + +set -x + +echo "Starting dynolog daemon for on-demand torch.profiler tracing..." + +PORT=${DYNOLOG_PORT:-1778} +OUTPUT_DIR=${TORCH_PROFILE_OUTPUT_DIR:-torch_profiler_output} + +DYNOLOG_PID_FILE="/tmp/madengine_dynolog.pid" +TRIGGER_PID_FILE="/tmp/madengine_dynolog_trigger.pid" +DYNOLOG_START_FILE="/tmp/madengine_dynolog.started" + +if ! command -v dynolog >/dev/null 2>&1 || ! command -v dyno >/dev/null 2>&1; then + echo "Error: dynolog/dyno not on PATH. The dynolog pre-script must run first." + exit 1 +fi + +# Traces are written by the workload process itself (via Kineto), so the output +# directory has to exist before the trigger fires. +mkdir -p "$OUTPUT_DIR" + +# --enable_ipc_monitor is what allows the daemon to talk to PyTorch/Kineto. +nohup dynolog --enable_ipc_monitor --port "$PORT" \ + > /tmp/madengine_dynolog.log 2>&1 & +DYNOLOG_PID=$! +echo "$DYNOLOG_PID" > "$DYNOLOG_PID_FILE" + +# Give the daemon time to bind its port before the workload tries to register. +sleep 3 + +if ! kill -0 "$DYNOLOG_PID" 2>/dev/null; then + echo "Error: dynolog daemon exited immediately. Log follows:" + cat /tmp/madengine_dynolog.log || true + rm -f "$DYNOLOG_PID_FILE" + exit 1 +fi +echo "✓ dynolog daemon started (PID: $DYNOLOG_PID, port: $PORT)" + +# The trigger has to run alongside the workload, because `dyno gputrace` can only +# match PyTorch processes that have already started and registered. +if [ -f "scripts/common/tools/dynolog_trigger.sh" ]; then + TRIGGER_SCRIPT="scripts/common/tools/dynolog_trigger.sh" +elif [ -f "../scripts/common/tools/dynolog_trigger.sh" ]; then + TRIGGER_SCRIPT="../scripts/common/tools/dynolog_trigger.sh" +else + echo "Error: Cannot find dynolog_trigger.sh" + exit 1 +fi + +nohup bash "$TRIGGER_SCRIPT" > /tmp/madengine_dynolog_trigger.log 2>&1 & +TRIGGER_PID=$! +echo "$TRIGGER_PID" > "$TRIGGER_PID_FILE" +echo "✓ dynolog trace trigger armed (PID: $TRIGGER_PID)" + +touch "$DYNOLOG_START_FILE" +echo "✓ dynolog initialization complete" diff --git a/src/madengine/scripts/common/pre_scripts/trace.sh b/src/madengine/scripts/common/pre_scripts/trace.sh index f6b8e624..42ad9336 100644 --- a/src/madengine/scripts/common/pre_scripts/trace.sh +++ b/src/madengine/scripts/common/pre_scripts/trace.sh @@ -145,4 +145,122 @@ except (json.JSONDecodeError, KeyError, TypeError, ValueError): fi ;; +dynolog) + # dynolog is the profiling daemon that lets us drive torch.profiler on an + # unmodified workload: PyTorch/Kineto registers with it when KINETO_USE_DAEMON=1, + # and `dyno gputrace` then configures the profiler over IPC. + # https://github.com/facebookincubator/dynolog/blob/main/docs/pytorch_profiler.md + if command -v dynolog >/dev/null 2>&1 && command -v dyno >/dev/null 2>&1; then + echo "dynolog: dynolog and dyno already on PATH, skipping install." + exit 0 + fi + + # Only x86_64 debian packages are published upstream. + _arch=$(uname -m) + if [ "$_arch" != "x86_64" ]; then + echo "Error: dynolog pre-script only supports x86_64 (found $_arch)." >&2 + echo "Build dynolog from source and put dynolog/dyno on PATH, or use a" >&2 + echo "model-side torch.profiler instead." >&2 + exit 1 + fi + if ! command -v dpkg >/dev/null 2>&1; then + echo "Error: dynolog pre-script needs dpkg (Debian/Ubuntu base image)." >&2 + exit 1 + fi + + _DYNOLOG_PINNED_DEB='https://github.com/facebookincubator/dynolog/releases/download/v0.5.0/dynolog_0.5.0-0-amd64.deb' + _dynolog_deb="${DYNOLOG_DEB_URL:-$_DYNOLOG_PINNED_DEB}" + _dynolog_tmp="/tmp/dynolog.deb" + + if command -v curl >/dev/null 2>&1; then + curl -fsSL -o "$_dynolog_tmp" "$_dynolog_deb" + elif command -v wget >/dev/null 2>&1; then + wget -q -O "$_dynolog_tmp" "$_dynolog_deb" + else + echo "Error: dynolog pre-script needs curl or wget to download the package." >&2 + exit 1 + fi + + # The package ships a systemd unit; enabling it fails in a container, which is + # harmless because we run the daemon directly. Tolerate a non-zero dpkg exit + # and verify by checking for the binaries instead. + if [ "$(id -u)" -eq 0 ]; then + dpkg -i "$_dynolog_tmp" || apt-get install -f -y -qq || true + elif command -v sudo >/dev/null 2>&1; then + sudo dpkg -i "$_dynolog_tmp" || sudo apt-get install -f -y -qq || true + else + echo "Error: dynolog pre-script needs root or sudo to install the package." >&2 + exit 1 + fi + rm -f "$_dynolog_tmp" + + if ! command -v dynolog >/dev/null 2>&1 || ! command -v dyno >/dev/null 2>&1; then + echo "Error: dynolog package installed but dynolog/dyno are not on PATH." >&2 + exit 1 + fi + echo "dynolog: installed $(dynolog --help 2>&1 | head -1 || echo 'ok')" + + # Kineto's daemon registration landed in torch 1.13; warn rather than fail so + # the tool stays usable for diagnosing the environment. + if ! python3 -c 'import torch' 2>/dev/null; then + echo "Warning: torch is not importable here; on-demand tracing needs a PyTorch workload." >&2 + fi + ;; + +tracelens) + # TraceLens pins protobuf>=6.31 and xprof, which routinely conflicts with a + # workload's own torch/tensorboard stack. Install it into a fully isolated + # venv (no --system-site-packages) so the model environment is untouched. + _tl_venv="${TRACELENS_VENV:-/opt/madengine-tracelens-venv}" + _TRACELENS_PINNED_REF='6f9bcdbf6cc9911eb650de57b345917ea4d31a17' + _tl_ref="${TRACELENS_GIT_REF:-$_TRACELENS_PINNED_REF}" + _tl_spec="${TRACELENS_PIP_SPEC:-git+https://github.com/AMD-AGI/TraceLens.git@${_tl_ref}}" + + if [ -x "${_tl_venv}/bin/python3" ] && "${_tl_venv}/bin/python3" -c 'import TraceLens' 2>/dev/null; then + echo "TraceLens: already installed in ${_tl_venv}, skipping." + exit 0 + fi + + if ! python3 -m venv "$_tl_venv" 2>/dev/null; then + echo "python3 -m venv failed; attempting to install the venv module..." >&2 + if [ "$(id -u)" -eq 0 ] && command -v apt-get >/dev/null 2>&1; then + apt-get update -qq && DEBIAN_FRONTEND=noninteractive apt-get install -y -qq python3-venv + elif command -v sudo >/dev/null 2>&1 && command -v apt-get >/dev/null 2>&1; then + sudo apt-get update -qq && sudo DEBIAN_FRONTEND=noninteractive apt-get install -y -qq python3-venv + fi + if ! python3 -m venv "$_tl_venv"; then + echo "Error: could not create a virtualenv at ${_tl_venv}." >&2 + echo "Install python3-venv, or set TRACELENS_VENV to an existing venv." >&2 + exit 1 + fi + fi + + "${_tl_venv}/bin/python3" -m pip install --upgrade -q pip + # TRACELENS_PIP_SPEC may embed credentials for a private mirror; keep it out of + # the `set -x` trace and out of stderr. + _tl_restore_x=0 + case $- in *x*) _tl_restore_x=1 ;; esac + set +x + if ! "${_tl_venv}/bin/python3" -m pip install -q "$_tl_spec"; then + echo "Error: pip could not install TraceLens (spec omitted from logs)." >&2 + echo "Check network access, or override TRACELENS_PIP_SPEC / TRACELENS_GIT_REF." >&2 + [ "$_tl_restore_x" -eq 1 ] && set -x + exit 1 + fi + [ "$_tl_restore_x" -eq 1 ] && set -x + unset _tl_restore_x + "${_tl_venv}/bin/python3" -c 'import TraceLens; print("TraceLens import OK")' + + # .pftrace input needs traceconv. TraceLens downloads it on demand, which fails + # in an air-gapped container, so pre-stage it here when we still have network. + if ! command -v traceconv >/dev/null 2>&1 && command -v curl >/dev/null 2>&1; then + if curl -fsSL -o /usr/local/bin/traceconv https://get.perfetto.dev/traceconv 2>/dev/null; then + chmod +x /usr/local/bin/traceconv + echo "TraceLens: pre-staged traceconv for .pftrace input." + else + echo "TraceLens: could not pre-stage traceconv (only needed for .pftrace input)." + fi + fi + ;; + esac diff --git a/src/madengine/scripts/common/tools.json b/src/madengine/scripts/common/tools.json index 82869087..9199eaa5 100644 --- a/src/madengine/scripts/common/tools.json +++ b/src/madengine/scripts/common/tools.json @@ -213,6 +213,157 @@ } ] }, + "torch_profiler_dynolog": { + "pre_scripts": [ + { + "path": "scripts/common/pre_scripts/trace.sh", + "args": "dynolog" + }, + { + "path": "scripts/common/pre_scripts/dynolog_start.sh" + } + ], + "cmd": "", + "env_vars": { + "KINETO_USE_DAEMON": "1", + "KINETO_DAEMON_INIT_DELAY_S": "3", + "DYNOLOG_PORT": "1778", + "TORCH_PROFILE_OUTPUT_DIR": "torch_profiler_output", + "TORCH_PROFILE_LOG_FILE": "libkineto_trace.json", + "TORCH_PROFILE_ITERATIONS": "5", + "TORCH_PROFILE_WARMUP_S": "60", + "TORCH_PROFILE_RETRY_INTERVAL_S": "15", + "TORCH_PROFILE_MAX_ATTEMPTS": "40", + "TORCH_PROFILE_PROCESS_LIMIT": "64", + "TORCH_PROFILE_RECORD_SHAPES": "1", + "TORCH_PROFILE_WITH_STACKS": "1", + "TORCH_PROFILE_WITH_MODULES": "1" + }, + "post_scripts": [ + { + "path": "scripts/common/post_scripts/dynolog_stop.sh" + }, + { + "path": "scripts/common/post_scripts/trace.sh", + "args": "torch_profiler" + } + ] + }, + "tracelens": { + "pre_scripts": [ + { + "path": "scripts/common/pre_scripts/trace.sh", + "args": "tracelens" + } + ], + "cmd": "", + "env_vars": { + "TRACELENS_VENV": "/opt/madengine-tracelens-venv", + "TRACELENS_OUTPUT_DIR": "tracelens_output", + "TRACELENS_MODE": "auto" + }, + "post_scripts": [ + { + "path": "scripts/common/post_scripts/tracelens.sh" + }, + { + "path": "scripts/common/post_scripts/trace.sh", + "args": "tracelens" + } + ] + }, + "tracelens_pytorch": { + "pre_scripts": [ + { + "path": "scripts/common/pre_scripts/trace.sh", + "args": "tracelens" + } + ], + "cmd": "", + "env_vars": { + "TRACELENS_VENV": "/opt/madengine-tracelens-venv", + "TRACELENS_OUTPUT_DIR": "tracelens_output", + "TRACELENS_MODE": "pytorch" + }, + "post_scripts": [ + { + "path": "scripts/common/post_scripts/tracelens.sh" + }, + { + "path": "scripts/common/post_scripts/trace.sh", + "args": "tracelens" + } + ] + }, + "tracelens_rocprof": { + "pre_scripts": [ + { + "path": "scripts/common/pre_scripts/trace.sh", + "args": "tracelens" + } + ], + "cmd": "", + "env_vars": { + "TRACELENS_VENV": "/opt/madengine-tracelens-venv", + "TRACELENS_OUTPUT_DIR": "tracelens_output", + "TRACELENS_MODE": "rocprof" + }, + "post_scripts": [ + { + "path": "scripts/common/post_scripts/tracelens.sh" + }, + { + "path": "scripts/common/post_scripts/trace.sh", + "args": "tracelens" + } + ] + }, + "tracelens_pftrace": { + "pre_scripts": [ + { + "path": "scripts/common/pre_scripts/trace.sh", + "args": "tracelens" + } + ], + "cmd": "", + "env_vars": { + "TRACELENS_VENV": "/opt/madengine-tracelens-venv", + "TRACELENS_OUTPUT_DIR": "tracelens_output", + "TRACELENS_MODE": "pftrace" + }, + "post_scripts": [ + { + "path": "scripts/common/post_scripts/tracelens.sh" + }, + { + "path": "scripts/common/post_scripts/trace.sh", + "args": "tracelens" + } + ] + }, + "tracelens_collective": { + "pre_scripts": [ + { + "path": "scripts/common/pre_scripts/trace.sh", + "args": "tracelens" + } + ], + "cmd": "", + "env_vars": { + "TRACELENS_VENV": "/opt/madengine-tracelens-venv", + "TRACELENS_OUTPUT_DIR": "tracelens_output", + "TRACELENS_MODE": "collective" + }, + "post_scripts": [ + { + "path": "scripts/common/post_scripts/tracelens.sh" + }, + { + "path": "scripts/common/post_scripts/trace.sh", + "args": "tracelens" + } + ] + }, "rocblas_trace": { "env_vars": { "ROCBLAS_TRACE": "1" diff --git a/src/madengine/scripts/common/tools/dynolog_trigger.sh b/src/madengine/scripts/common/tools/dynolog_trigger.sh new file mode 100644 index 00000000..9999e36c --- /dev/null +++ b/src/madengine/scripts/common/tools/dynolog_trigger.sh @@ -0,0 +1,75 @@ +#!/usr/bin/env bash +# +# Copyright (c) Advanced Micro Devices, Inc. +# All rights reserved. +# +# Request a torch.profiler trace from the running workload via dynolog. +# +# Runs in the background for the lifetime of the model run. `dyno gputrace` can +# only configure PyTorch processes that have already registered with the daemon, +# and there is no way to know when the workload reaches steady state, so this +# polls with --fail-on-no-process until a request is accepted. + +set -u + +PORT=${DYNOLOG_PORT:-1778} +OUTPUT_DIR=${TORCH_PROFILE_OUTPUT_DIR:-torch_profiler_output} +LOG_NAME=${TORCH_PROFILE_LOG_FILE:-libkineto_trace.json} +ITERATIONS=${TORCH_PROFILE_ITERATIONS:-5} +DURATION_MS=${TORCH_PROFILE_DURATION_MS:-500} +WARMUP_S=${TORCH_PROFILE_WARMUP_S:-60} +RETRY_INTERVAL_S=${TORCH_PROFILE_RETRY_INTERVAL_S:-15} +MAX_ATTEMPTS=${TORCH_PROFILE_MAX_ATTEMPTS:-40} +# Upstream defaults to 3, which silently drops most ranks of a multi-GPU job. +PROCESS_LIMIT=${TORCH_PROFILE_PROCESS_LIMIT:-64} +JOB_ID=${TORCH_PROFILE_JOB_ID:-${SLURM_JOB_ID:-0}} + +RESULT_FILE="/tmp/madengine_dynolog_trigger.result" +rm -f "$RESULT_FILE" + +# TraceLens needs input shapes and CPU call stacks for per-op and roofline +# analysis, and modules for the nn.Module breakdown. +OPTS=() +[ "${TORCH_PROFILE_RECORD_SHAPES:-1}" = "1" ] && OPTS+=(--record-shapes) +[ "${TORCH_PROFILE_WITH_STACKS:-1}" = "1" ] && OPTS+=(--with-stacks) +[ "${TORCH_PROFILE_WITH_MODULES:-1}" = "1" ] && OPTS+=(--with-modules) +[ "${TORCH_PROFILE_WITH_FLOPS:-0}" = "1" ] && OPTS+=(--with-flops) +[ "${TORCH_PROFILE_PROFILE_MEMORY:-0}" = "1" ] && OPTS+=(--profile-memory) + +# Iteration-based capture needs an optimizer step hook; PyTorch falls back to a +# duration-based trace on its own when it cannot count iterations. +if [ "$ITERATIONS" -gt 0 ] 2>/dev/null; then + OPTS+=(--iterations "$ITERATIONS") +else + OPTS+=(--duration-ms "$DURATION_MS") +fi + +mkdir -p "$OUTPUT_DIR" +LOG_FILE="$(cd "$OUTPUT_DIR" && pwd)/${LOG_NAME}" + +echo "[dynolog-trigger] waiting ${WARMUP_S}s for the workload to reach steady state" +sleep "$WARMUP_S" + +attempt=0 +while [ "$attempt" -lt "$MAX_ATTEMPTS" ]; do + attempt=$((attempt + 1)) + echo "[dynolog-trigger] attempt ${attempt}/${MAX_ATTEMPTS}: requesting trace -> ${LOG_FILE}" + if dyno --port "$PORT" gputrace \ + --job-id "$JOB_ID" \ + --log-file "$LOG_FILE" \ + --process-limit "$PROCESS_LIMIT" \ + --fail-on-no-process \ + "${OPTS[@]}"; then + echo "[dynolog-trigger] trace request accepted on attempt ${attempt}" + echo "accepted" > "$RESULT_FILE" + exit 0 + fi + echo "[dynolog-trigger] no PyTorch process matched yet; retrying in ${RETRY_INTERVAL_S}s" + sleep "$RETRY_INTERVAL_S" +done + +echo "[dynolog-trigger] gave up after ${MAX_ATTEMPTS} attempts: no PyTorch process registered." +echo "[dynolog-trigger] Confirm the workload runs PyTorch >= 1.13 with KINETO_USE_DAEMON=1," +echo "[dynolog-trigger] and raise TORCH_PROFILE_WARMUP_S / TORCH_PROFILE_MAX_ATTEMPTS for slow starts." +echo "no_process" > "$RESULT_FILE" +exit 1 diff --git a/src/madengine/scripts/common/tools/tracelens_analyze.py b/src/madengine/scripts/common/tools/tracelens_analyze.py new file mode 100644 index 00000000..bfe5606e --- /dev/null +++ b/src/madengine/scripts/common/tools/tracelens_analyze.py @@ -0,0 +1,650 @@ +#!/usr/bin/env python3 +"""Discover GPU trace artifacts and generate TraceLens reports for them. + +Copyright (c) Advanced Micro Devices, Inc. All rights reserved. + +This script is the single implementation shared by both TraceLens execution +paths in madengine: + +* in-container, as the ``tracelens`` tool post-script (``post_scripts/tracelens.sh``) +* on the host, via ``madengine report tracelens`` + +It therefore uses only the Python standard library and never imports madengine. +TraceLens itself is invoked out-of-process through ``--python``, which lets the +caller point at an isolated virtualenv so that TraceLens' pinned ``protobuf`` +and ``xprof`` cannot disturb the workload's own Python environment. +""" + +import argparse +import csv +import glob +import gzip +import json +import os +import re +import shutil +import subprocess +import sys +from typing import Dict, List, Optional, Sequence, Tuple + +# Trace kinds, in discovery precedence order. The first pattern set that claims a +# file wins, so PyTorch traces are matched before the broader JSON patterns. +KIND_PYTORCH = "pytorch" +KIND_ROCPROF_JSON = "rocprof_json" +KIND_PFTRACE = "pftrace" +KIND_UNSUPPORTED = "unsupported" + +# Glob patterns are matched against paths relative to the discovery root. +_PYTORCH_PATTERNS = ( + "**/*.pt.trace.json", + "**/*.pt.trace.json.gz", + "**/libkineto_trace*.json", + "**/libkineto_trace*.json.gz", + "**/torch_profiler_output/*.json", + "**/torch_profiler_output/*.json.gz", +) +_ROCPROF_JSON_PATTERNS = ("**/*_results.json",) +_PFTRACE_PATTERNS = ("**/*.pftrace",) +# Formats madengine can produce that TraceLens cannot read. Reported with +# actionable guidance rather than silently ignored. +_UNSUPPORTED_PATTERNS = { + "**/*_results.db": ( + "rocprofv3 SQLite output is not readable by TraceLens. Re-run with a " + "preset that sets an explicit --output-format, e.g. " + "rocprofv3_lightweight (JSON) or rocprofv3_perfetto (pftrace)." + ), + "**/*.rpd": ( + "RPD databases are not readable by TraceLens. The rpd post-script also " + "writes a converted trace.json alongside it; point TraceLens at that." + ), + "**/*.pb": ( + "JAX XPlane protobuf traces need TraceLens_generate_perf_report_jax, " + "which madengine does not drive yet." + ), +} + +# Ambiguous names written by more than one madengine tool (rpd writes a Chrome +# trace here, rocm-trace-lite writes its own). Sniffed rather than assumed. +_AMBIGUOUS_NAMES = ("trace.json", "trace.json.gz") + +# Directories never worth walking: our own output, virtualenvs, VCS metadata. +_SKIP_DIRS = frozenset( + { + ".git", + "__pycache__", + "node_modules", + "site-packages", + "venv", + ".venv", + } +) + +# TraceLens console script -> module providing main(). The console script is +# preferred when present; the module is the fallback so the integration keeps +# working if entry points were not installed onto PATH. +_ENTRY_POINTS = { + "TraceLens_generate_perf_report_pytorch": "TraceLens.Reporting.generate_perf_report_pytorch", + "TraceLens_generate_perf_report_rocprof": "TraceLens.Reporting.generate_perf_report_rocprof", + "TraceLens_generate_perf_report_pftrace_hip_activity": "TraceLens.Reporting.generate_perf_report_pftrace_hip_activity", + "TraceLens_generate_perf_report_pftrace_hip_api": "TraceLens.Reporting.generate_perf_report_pftrace_hip_api", + "TraceLens_generate_perf_report_pftrace_memory_copy": "TraceLens.Reporting.generate_perf_report_pftrace_memory_copy", + "TraceLens_generate_multi_rank_collective_report_pytorch": "TraceLens.Reporting.generate_multi_rank_collective_report_pytorch", + "TraceLens_compare_perf_reports_pytorch": "TraceLens.Reporting.compare_perf_reports_pytorch", +} + +SUMMARY_CSV_FIELDS = ( + "trace_file", + "kind", + "tracelens_tool", + "status", + "output", + "detail", +) + + +def _read_head(path: str, size: int = 4096) -> str: + """Return the first ``size`` bytes of a plain or gzipped file as text.""" + opener = gzip.open if path.endswith(".gz") else open + try: + with opener(path, "rb") as handle: # type: ignore[operator] + return handle.read(size).decode("utf-8", errors="replace") + except OSError: + return "" + + +def _is_chrome_trace(path: str) -> bool: + """Return True if the file looks like a Chrome Trace Event JSON document.""" + return '"traceEvents"' in _read_head(path) + + +def _iter_files(root: str) -> List[str]: + """Return every file under ``root``, skipping uninteresting directories.""" + found: List[str] = [] + for dirpath, dirnames, filenames in os.walk(root): + dirnames[:] = [d for d in dirnames if d not in _SKIP_DIRS] + for name in filenames: + found.append(os.path.join(dirpath, name)) + return found + + +def _match(root: str, patterns: Sequence[str]) -> List[str]: + matches: List[str] = [] + for pattern in patterns: + matches.extend(glob.glob(os.path.join(root, pattern), recursive=True)) + return matches + + +def discover_traces( + root: str, exclude_dirs: Sequence[str] = () +) -> Tuple[Dict[str, List[str]], List[Tuple[str, str]]]: + """Classify trace artifacts under ``root`` by the TraceLens reader they need. + + Args: + root: Directory to search recursively. + exclude_dirs: Absolute or relative directories to omit from results, + typically the report output directory. + + Returns: + A ``(traces, unsupported)`` pair. ``traces`` maps a trace kind to sorted + file paths. ``unsupported`` is a list of ``(path, reason)`` for artifacts + that were found but cannot be analyzed. + """ + excluded = [os.path.abspath(d) for d in exclude_dirs] + + def is_excluded(path: str) -> bool: + absolute = os.path.abspath(path) + return any( + absolute == prefix or absolute.startswith(prefix + os.sep) + for prefix in excluded + ) + + claimed = set() + traces: Dict[str, List[str]] = {} + for kind, patterns in ( + (KIND_PYTORCH, _PYTORCH_PATTERNS), + (KIND_ROCPROF_JSON, _ROCPROF_JSON_PATTERNS), + (KIND_PFTRACE, _PFTRACE_PATTERNS), + ): + for path in _match(root, patterns): + if not os.path.isfile(path) or is_excluded(path) or path in claimed: + continue + claimed.add(path) + traces.setdefault(kind, []).append(path) + + # Sniff ambiguously named files that no pattern claimed. + for path in _iter_files(root): + if path in claimed or is_excluded(path): + continue + if os.path.basename(path) in _AMBIGUOUS_NAMES and _is_chrome_trace(path): + claimed.add(path) + traces.setdefault(KIND_PYTORCH, []).append(path) + + unsupported: List[Tuple[str, str]] = [] + for pattern, reason in _UNSUPPORTED_PATTERNS.items(): + for path in _match(root, [pattern]): + if os.path.isfile(path) and not is_excluded(path) and path not in claimed: + unsupported.append((path, reason)) + + for kind in traces: + traces[kind] = sorted(set(traces[kind])) + return traces, sorted(set(unsupported)) + + +def _resolve_python(python: Optional[str]) -> str: + """Return the interpreter used to run TraceLens.""" + if python: + return python + venv = os.environ.get("TRACELENS_VENV", "").strip() + if venv: + candidate = os.path.join(venv, "bin", "python3") + if os.path.isfile(candidate): + return candidate + return sys.executable or "python3" + + +def _build_command(python: str, script_name: str, args: Sequence[str]) -> List[str]: + """Return argv invoking a TraceLens entry point with ``args``. + + Prefers the installed console script (clearer logs, honours the package's + own entry-point wiring) and falls back to importing the module's ``main``. + """ + bindir = os.path.dirname(os.path.abspath(python)) + candidate = os.path.join(bindir, script_name) + if os.path.isfile(candidate) and os.access(candidate, os.X_OK): + return [candidate, *args] + on_path = shutil.which(script_name) + if on_path: + return [on_path, *args] + module = _ENTRY_POINTS[script_name] + return [python, "-c", f"from {module} import main; main()", *args] + + +def _run(command: Sequence[str], cwd: Optional[str] = None) -> Tuple[int, str]: + """Run ``command``, streaming nothing, returning ``(returncode, output)``.""" + printable = " ".join(command) + print(f" $ {printable}", flush=True) + try: + completed = subprocess.run( + list(command), + cwd=cwd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + universal_newlines=True, + ) + except OSError as exc: + return 1, str(exc) + output = completed.stdout or "" + if output: + for line in output.splitlines(): + print(f" {line}", flush=True) + return completed.returncode, output + + +def _failure_detail(returncode: int, output: str) -> str: + """Return a one-line explanation for a failed TraceLens invocation.""" + lines = [line for line in output.strip().splitlines() if line.strip()] + return lines[-1] if lines else f"exit code {returncode}" + + +def _report_stem(path: str, root: str) -> str: + """Return a filesystem-safe, collision-resistant name for a trace's reports.""" + relative = os.path.relpath(path, root) + for suffix in (".json.gz", ".pt.trace.json", ".pftrace", ".json"): + if relative.endswith(suffix): + relative = relative[: -len(suffix)] + break + return re.sub(r"[^A-Za-z0-9._-]+", "_", relative).strip("_") or "trace" + + +def _pytorch_args( + trace: str, out_base: str, gpu_arch: Optional[str], extra: Sequence[str] +) -> List[str]: + args = [ + "--profile_json_path", + trace, + "--output_xlsx_path", + f"{out_base}.xlsx", + "--output_csvs_dir", + f"{out_base}_csv", + "--enable_kernel_summary", + "--short_kernel_study", + ] + if gpu_arch: + args += ["--gpu_arch_platform", gpu_arch] + return args + list(extra) + + +def _rocprof_args(trace: str, out_base: str, extra: Sequence[str]) -> List[str]: + return [ + "--profile_json_path", + trace, + "--output_xlsx_path", + f"{out_base}.xlsx", + "--output_csvs_dir", + f"{out_base}_csv", + "--kernel_details", + "--short_kernel_study", + *extra, + ] + + +def _pftrace_jobs( + trace: str, out_base: str, extra: Sequence[str] +) -> List[Tuple[str, List[str]]]: + """Return the three complementary pftrace reports for one trace.""" + return [ + ( + "TraceLens_generate_perf_report_pftrace_hip_activity", + [ + "--trace_path", + trace, + "--output_csvs_dir", + f"{out_base}_activity_csv", + "--output_md_path", + f"{out_base}_activity.md", + *extra, + ], + ), + ( + "TraceLens_generate_perf_report_pftrace_hip_api", + [ + "--trace_path", + trace, + "--output_xlsx_path", + f"{out_base}_hip_api.xlsx", + "--output_csvs_dir", + f"{out_base}_hip_api_csv", + *extra, + ], + ), + ( + "TraceLens_generate_perf_report_pftrace_memory_copy", + [ + "--trace_path", + trace, + "--output_xlsx_path", + f"{out_base}_memory_copy.xlsx", + "--output_csvs_dir", + f"{out_base}_memory_copy_csv", + *extra, + ], + ), + ] + + +def _rank_regex() -> str: + """Return the rank-extraction regex covering madengine trace filenames. + + Matches dynolog output (``libkineto_trace_.json``, where madengine + renames per rank), torch.profiler defaults (``..._rank0_...``), and the + ``rank[N]`` form used by ``tensorboard_trace_handler``. + """ + return r"rank[\[\-_/]?(?P\d+)" + + +def _collective_args( + root: str, out_base: str, world_size: int, extra: Sequence[str] +) -> List[str]: + return [ + "--trace_glob", + os.path.join(root, "**", "*.json*"), + "--rank_regex", + _rank_regex(), + "--world_size", + str(world_size), + "--output_xlsx_path", + f"{out_base}.xlsx", + "--output_csvs_dir", + f"{out_base}_csv", + "--use_multiprocessing", + *extra, + ] + + +def analyze( + root: str, + output_dir: str, + mode: str = "auto", + python: Optional[str] = None, + gpu_arch: Optional[str] = None, + world_size: int = 0, + max_traces: int = 0, + extra_args: Sequence[str] = (), +) -> Dict[str, object]: + """Generate TraceLens reports for every supported trace under ``root``. + + Args: + root: Directory to search for trace artifacts. + output_dir: Directory that receives the generated reports. + mode: ``auto`` to analyze every discovered kind, or one of ``pytorch``, + ``rocprof``, ``pftrace``, ``collective`` to restrict the run. + python: Interpreter that has TraceLens installed. Defaults to + ``$TRACELENS_VENV/bin/python3`` when set, else the current one. + gpu_arch: Bundled TraceLens GPU arch name (e.g. ``MI300X``) enabling + roofline bound classification on PyTorch reports. + world_size: Rank count for the multi-rank collective report. When 0, it + is inferred from the number of discovered PyTorch traces. + max_traces: Cap on traces analyzed per kind. 0 means no cap. + extra_args: Extra flags forwarded verbatim to every TraceLens command. + + Returns: + A summary dict with ``results``, ``unsupported``, and counters. + """ + root = os.path.abspath(root) + output_dir = os.path.abspath(output_dir) + os.makedirs(output_dir, exist_ok=True) + + interpreter = _resolve_python(python) + traces, unsupported = discover_traces(root, exclude_dirs=[output_dir]) + + if max_traces > 0: + traces = {kind: paths[:max_traces] for kind, paths in traces.items()} + + wanted = { + "pytorch": {KIND_PYTORCH}, + "rocprof": {KIND_ROCPROF_JSON}, + "pftrace": {KIND_PFTRACE}, + "collective": {KIND_PYTORCH}, + "auto": {KIND_PYTORCH, KIND_ROCPROF_JSON, KIND_PFTRACE}, + }[mode] + + jobs: List[Tuple[str, str, str, List[str]]] = [] + for kind, paths in sorted(traces.items()): + if kind not in wanted: + continue + for trace in paths: + stem = _report_stem(trace, root) + out_base = os.path.join(output_dir, stem) + if kind == KIND_PYTORCH and mode != "collective": + jobs.append( + ( + trace, + kind, + "TraceLens_generate_perf_report_pytorch", + _pytorch_args(trace, out_base, gpu_arch, extra_args), + ) + ) + elif kind == KIND_ROCPROF_JSON: + jobs.append( + ( + trace, + kind, + "TraceLens_generate_perf_report_rocprof", + _rocprof_args(trace, out_base, extra_args), + ) + ) + elif kind == KIND_PFTRACE: + for tool, args in _pftrace_jobs(trace, out_base, extra_args): + jobs.append((trace, kind, tool, args)) + + # A multi-rank collective report needs at least two per-rank PyTorch traces. + pytorch_traces = traces.get(KIND_PYTORCH, []) + ranks = world_size or len(pytorch_traces) + if mode in ("auto", "collective") and len(pytorch_traces) > 1 and ranks > 1: + jobs.append( + ( + f"{len(pytorch_traces)} per-rank traces", + KIND_PYTORCH, + "TraceLens_generate_multi_rank_collective_report_pytorch", + _collective_args( + root, + os.path.join(output_dir, "multi_rank_collective"), + ranks, + extra_args, + ), + ) + ) + + results: List[Dict[str, str]] = [] + for trace, kind, tool, args in jobs: + print(f"[tracelens] {tool}: {trace}", flush=True) + code, output = _run(_build_command(interpreter, tool, args)) + results.append( + { + "trace_file": os.path.relpath(trace, root) + if os.path.exists(trace) + else trace, + "kind": kind, + "tracelens_tool": tool, + "status": "SUCCESS" if code == 0 else "FAILURE", + "output": os.path.relpath(output_dir, root), + "detail": "" if code == 0 else _failure_detail(code, output), + } + ) + + for path, reason in unsupported: + print(f"[tracelens] skipping {path}: {reason}", flush=True) + results.append( + { + "trace_file": os.path.relpath(path, root), + "kind": KIND_UNSUPPORTED, + "tracelens_tool": "", + "status": "SKIPPED", + "output": "", + "detail": reason, + } + ) + + summary_csv = os.path.join(output_dir, "tracelens_summary.csv") + with open(summary_csv, "w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=list(SUMMARY_CSV_FIELDS)) + writer.writeheader() + writer.writerows(results) + + return { + "root": root, + "output_dir": output_dir, + "mode": mode, + "python": interpreter, + "summary_csv": summary_csv, + "discovered": {kind: len(paths) for kind, paths in sorted(traces.items())}, + "succeeded": sum(1 for r in results if r["status"] == "SUCCESS"), + "failed": sum(1 for r in results if r["status"] == "FAILURE"), + "skipped": sum(1 for r in results if r["status"] == "SKIPPED"), + "results": results, + } + + +def compare( + reports: Sequence[str], + output: str, + names: Sequence[str] = (), + python: Optional[str] = None, +) -> Dict[str, object]: + """Diff two or more TraceLens reports into a single comparison workbook.""" + interpreter = _resolve_python(python) + args: List[str] = [*reports, "-o", output] + if names: + args += ["--names", *names] + code, out = _run( + _build_command(interpreter, "TraceLens_compare_perf_reports_pytorch", args) + ) + return { + "reports": list(reports), + "output": output, + "status": "SUCCESS" if code == 0 else "FAILURE", + "detail": "" if code == 0 else out.strip(), + } + + +def _parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Generate TraceLens reports for madengine trace artifacts." + ) + parser.add_argument( + "--root", default=".", help="Directory to search for traces (default: .)" + ) + parser.add_argument( + "--output-dir", + default="tracelens_output", + help="Directory for generated reports (default: tracelens_output)", + ) + parser.add_argument( + "--mode", + default="auto", + choices=["auto", "pytorch", "rocprof", "pftrace", "collective"], + help="Restrict analysis to one trace kind (default: auto)", + ) + parser.add_argument( + "--python", default=None, help="Interpreter that has TraceLens installed" + ) + parser.add_argument( + "--gpu-arch", + default=None, + help="TraceLens GPU arch platform for roofline bound classification", + ) + parser.add_argument( + "--world-size", + type=int, + default=0, + help="Rank count for the collective report (default: number of traces)", + ) + parser.add_argument( + "--max-traces", + type=int, + default=0, + help="Cap traces analyzed per kind (default: no cap)", + ) + parser.add_argument( + "--json-summary", default=None, help="Write the run summary as JSON here" + ) + parser.add_argument( + "--discover-only", + action="store_true", + help="List discovered traces without running TraceLens", + ) + parser.add_argument( + "--compare", + nargs="+", + default=None, + metavar="REPORT", + help="Compare existing TraceLens reports instead of analyzing traces", + ) + parser.add_argument( + "--compare-output", + default="tracelens_comparison.xlsx", + help="Output workbook for --compare (default: tracelens_comparison.xlsx)", + ) + parser.add_argument( + "--compare-names", nargs="+", default=(), help="Display tags for --compare" + ) + parser.add_argument( + "extra_args", + nargs="*", + help="Extra flags forwarded verbatim to every TraceLens command", + ) + return parser.parse_args(argv) + + +def main(argv: Optional[Sequence[str]] = None) -> int: + args = _parse_args(argv) + + if args.compare: + summary: Dict[str, object] = compare( + args.compare, args.compare_output, args.compare_names, args.python + ) + failed = summary["status"] != "SUCCESS" + elif args.discover_only: + traces, unsupported = discover_traces( + args.root, exclude_dirs=[args.output_dir] + ) + for kind, paths in sorted(traces.items()): + for path in paths: + print(f"{kind}\t{path}") + for path, reason in unsupported: + print(f"{KIND_UNSUPPORTED}\t{path}\t{reason}") + summary = { + "discovered": {kind: len(paths) for kind, paths in sorted(traces.items())}, + "unsupported": [{"path": p, "reason": r} for p, r in unsupported], + } + failed = False + else: + summary = analyze( + root=args.root, + output_dir=args.output_dir, + mode=args.mode, + python=args.python, + gpu_arch=args.gpu_arch, + world_size=args.world_size, + max_traces=args.max_traces, + extra_args=args.extra_args, + ) + if not summary["results"]: + print( + "[tracelens] No supported trace artifacts found under " + f"{os.path.abspath(args.root)}. Stack a profiling tool such as " + "torch_profiler_dynolog, rocprofv3_lightweight, or " + "rocprofv3_perfetto with the tracelens tool.", + flush=True, + ) + failed = bool(summary["failed"]) + + if args.json_summary: + with open(args.json_summary, "w", encoding="utf-8") as handle: + json.dump(summary, handle, indent=2) + + return 1 if failed else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/e2e/test_tracelens_workflows.py b/tests/e2e/test_tracelens_workflows.py new file mode 100644 index 00000000..f147603f --- /dev/null +++ b/tests/e2e/test_tracelens_workflows.py @@ -0,0 +1,380 @@ +"""End-to-end tests for the TraceLens integration. + +Two independent execution paths are covered: + +* the host-side reporting path (``madengine report tracelens``), which needs + neither Docker nor a GPU and therefore runs everywhere +* the in-container tool path (``tracelens`` / ``torch_profiler_dynolog`` + stacked onto a run), which needs Docker and a GPU + +Copyright (c) Advanced Micro Devices, Inc. All rights reserved. +""" + +# built-in modules +import csv +import gzip +import json +import os +import subprocess +import sys + +# third-party modules +import pytest + +# project modules +from tests.fixtures.utils import ( + BASE_DIR, + DEFAULT_CLEAN_FILES, + build_run_command, + clean_test_temp_files, + global_data, + is_nvidia, + requires_gpu, +) +from madengine.reporting.tracelens_report import ( + check_tracelens_available, + resolve_python, +) + + +def tracelens_installed() -> bool: + """Return True when the interpreter running the tests can import TraceLens.""" + try: + return check_tracelens_available(resolve_python()) + except Exception: + return False + + +def run_report_cli(*args: str) -> subprocess.CompletedProcess: + """Run ``madengine report ...`` in a wide, colourless console. + + Rich wraps and colourises output based on the terminal, which would make + substring assertions brittle. A wide COLUMNS plus NO_COLOR keeps the text + intact, and callers still normalise whitespace via :func:`flatten`. + PYTHONIOENCODING is needed because the CLI prints emoji, which a piped + stdout on a non-UTF-8 Windows console cannot encode. + """ + env = dict( + os.environ, + COLUMNS="300", + NO_COLOR="1", + TERM="dumb", + PYTHONIOENCODING="utf-8", + ) + return subprocess.run( + [sys.executable, "-m", "madengine.cli.app", "report", *args], + cwd=BASE_DIR, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + encoding="utf-8", + errors="replace", + timeout=300, + ) + + +def flatten(text: str) -> str: + """Collapse all whitespace so assertions survive console line wrapping.""" + return " ".join(text.split()) + + +def run_context(tools: list) -> dict: + """Return an additional-context dict selecting ``tools`` on an AMD host.""" + return {"gpu_vendor": "AMD", "guest_os": "UBUNTU", "tools": tools} + + +def write_pytorch_trace(path) -> None: + """Write a minimal Chrome Trace Event document, as Kineto would.""" + path.parent.mkdir(parents=True, exist_ok=True) + payload = { + "schemaVersion": 1, + "traceEvents": [ + { + "ph": "X", + "cat": "kernel", + "name": "void gemm_kernel(float*)", + "pid": 1, + "tid": 7, + "ts": 100, + "dur": 42, + "args": {"stream": 7, "grid": [8, 1, 1], "block": [256, 1, 1]}, + } + ], + } + if str(path).endswith(".gz"): + with gzip.open(path, "wt", encoding="utf-8") as handle: + json.dump(payload, handle) + else: + path.write_text(json.dumps(payload), encoding="utf-8") + + +def summary_rows(summary_csv: str) -> list: + """Read the analyzer summary CSV into a list of dict rows.""" + with open(summary_csv, "r", newline="", encoding="utf-8") as handle: + return list(csv.DictReader(handle)) + + +class TestTraceLensHostSideReporting: + """`madengine report tracelens` against traces already on the filesystem.""" + + def test_discover_only_classifies_every_trace_kind(self, tmp_path): + """Discovery reports each trace kind and explains unreadable formats.""" + write_pytorch_trace(tmp_path / "torch_profiler_output" / "libkineto_trace.json") + rocprof_dir = tmp_path / "rocprof_output" + rocprof_dir.mkdir() + (rocprof_dir / "run_results.json").write_text( + json.dumps({"rocprofiler-sdk-tool": []}), encoding="utf-8" + ) + (rocprof_dir / "run.pftrace").write_bytes(b"\x0a\x00perfetto-ish") + # rocprofv3's default SQLite output: found, but unusable by TraceLens. + (rocprof_dir / "run_results.db").write_bytes(b"SQLite format 3\x00") + + result = run_report_cli( + "tracelens", + "--root", + str(tmp_path), + "--output-dir", + str(tmp_path / "tracelens_output"), + "--discover-only", + ) + + assert result.returncode == 0, result.stdout + output = flatten(result.stdout) + assert "pytorch: 1 trace(s)" in output + assert "rocprof_json: 1 trace(s)" in output + assert "pftrace: 1 trace(s)" in output + # The .db is surfaced with a pointer at the presets that do work. + assert "unsupported" in output + assert "rocprofv3_lightweight" in output + + def test_discover_only_reports_empty_root(self, tmp_path): + """An empty search root is a warning, not a failure.""" + result = run_report_cli( + "tracelens", + "--root", + str(tmp_path), + "--output-dir", + str(tmp_path / "tracelens_output"), + "--discover-only", + ) + + assert result.returncode == 0, result.stdout + assert "No trace artifacts found" in flatten(result.stdout) + + def test_missing_root_fails_with_clear_error(self, tmp_path): + """A nonexistent search root fails before any analysis is attempted.""" + result = run_report_cli( + "tracelens", "--root", str(tmp_path / "nope"), "--discover-only" + ) + + assert result.returncode != 0 + assert "directory not found" in flatten(result.stdout) + + def test_invalid_mode_is_rejected(self, tmp_path): + """--mode is validated against the supported trace kinds.""" + result = run_report_cli("tracelens", "--root", str(tmp_path), "--mode", "bogus") + + assert result.returncode != 0 + output = flatten(result.stdout) + assert "invalid --mode" in output + assert "collective" in output + + @pytest.mark.skipif( + tracelens_installed(), reason="test covers the TraceLens-missing path" + ) + def test_analysis_without_tracelens_explains_how_to_install(self, tmp_path): + """Without TraceLens the command fails with install guidance, not a crash.""" + write_pytorch_trace(tmp_path / "torch_profiler_output" / "libkineto_trace.json") + + result = run_report_cli( + "tracelens", + "--root", + str(tmp_path), + "--output-dir", + str(tmp_path / "tracelens_output"), + ) + + assert result.returncode != 0 + output = flatten(result.stdout) + assert "TraceLens is not importable" in output + assert "madengine[tracelens]" in output + + def test_compare_rejects_missing_reports(self, tmp_path): + """tracelens-compare validates its inputs before invoking TraceLens.""" + existing = tmp_path / "baseline.xlsx" + existing.write_bytes(b"not really a workbook") + + result = run_report_cli( + "tracelens-compare", + str(existing), + str(tmp_path / "missing.xlsx"), + "--output", + str(tmp_path / "diff.xlsx"), + ) + + assert result.returncode != 0 + output = flatten(result.stdout) + assert "not found" in output + assert "missing.xlsx" in output + + @pytest.mark.skipif( + tracelens_installed(), reason="test covers the TraceLens-missing path" + ) + def test_compare_without_tracelens_explains_how_to_install(self, tmp_path): + """tracelens-compare surfaces the same install guidance as analysis.""" + first = tmp_path / "a.xlsx" + second = tmp_path / "b.xlsx" + for report in (first, second): + report.write_bytes(b"not really a workbook") + + result = run_report_cli( + "tracelens-compare", + str(first), + str(second), + "--output", + str(tmp_path / "diff.xlsx"), + ) + + assert result.returncode != 0 + assert "madengine[tracelens]" in flatten(result.stdout) + + @pytest.mark.skipif( + not tracelens_installed(), reason="requires pip install 'madengine[tracelens]'" + ) + def test_analysis_generates_reports_for_a_pytorch_trace(self, tmp_path): + """With TraceLens installed, a Kineto trace yields a summarised report.""" + write_pytorch_trace(tmp_path / "torch_profiler_output" / "libkineto_trace.json") + output_dir = tmp_path / "tracelens_output" + + result = run_report_cli( + "tracelens", + "--root", + str(tmp_path), + "--output-dir", + str(output_dir), + ) + + summary_csv = output_dir / "tracelens_summary.csv" + assert summary_csv.is_file(), result.stdout + rows = summary_rows(str(summary_csv)) + assert [r for r in rows if r["kind"] == "pytorch"], rows + assert (output_dir / "tracelens_summary.json").is_file() + + +@requires_gpu("in-container TraceLens tools require GPU hardware") +@pytest.mark.skipif(is_nvidia(), reason="TraceLens targets AMD GPU traces") +@pytest.mark.slow +class TestTraceLensContainerTools: + """The `tracelens` and `torch_profiler_dynolog` tools stacked onto a run.""" + + @pytest.mark.parametrize( + "clean_test_temp_files", + [DEFAULT_CLEAN_FILES + ["rocprof_output", "tracelens_output"]], + indirect=True, + ) + def test_tracelens_analyzes_rocprofv3_json_traces( + self, global_data, clean_test_temp_files + ): + """rocprofv3 JSON output is analyzed in-container and collected to cwd.""" + global_data["console"].sh( + build_run_command( + "dummy_prof", + additional_context=run_context( + [{"name": "rocprofv3_lightweight"}, {"name": "tracelens"}] + ), + ), + canFail=True, + ) + + summary_csv = os.path.join( + BASE_DIR, "tracelens_output", "tracelens_summary.csv" + ) + if not os.path.isfile(summary_csv): + pytest.fail( + "tracelens_output/tracelens_summary.csv not collected when stacking " + "tracelens onto rocprofv3_lightweight." + ) + rows = summary_rows(summary_csv) + if not [r for r in rows if r["kind"] == "rocprof_json"]: + pytest.fail(f"no rocprofv3 JSON trace was analyzed; summary rows: {rows}") + + @pytest.mark.parametrize( + "clean_test_temp_files", + [DEFAULT_CLEAN_FILES + ["rocprof_output", "tracelens_output"]], + indirect=True, + ) + def test_tracelens_explains_unreadable_rocprofv3_db_output( + self, global_data, clean_test_temp_files + ): + """The default .db output is reported as skipped with actionable guidance.""" + global_data["console"].sh( + build_run_command( + "dummy_prof", + additional_context=run_context( + [{"name": "rocprofv3"}, {"name": "tracelens"}] + ), + ), + canFail=True, + ) + + summary_csv = os.path.join( + BASE_DIR, "tracelens_output", "tracelens_summary.csv" + ) + if not os.path.isfile(summary_csv): + pytest.fail( + "tracelens_output/tracelens_summary.csv not collected when stacking " + "tracelens onto rocprofv3." + ) + skipped = [r for r in summary_rows(summary_csv) if r["status"] == "SKIPPED"] + if not skipped: + pytest.fail( + "rocprofv3 .db output should be reported as SKIPPED with guidance." + ) + if not any("rocprofv3_lightweight" in r["detail"] for r in skipped): + pytest.fail(f"skip reason does not point at a usable preset: {skipped}") + + @pytest.mark.parametrize( + "clean_test_temp_files", + [DEFAULT_CLEAN_FILES + ["torch_profiler_output"]], + indirect=True, + ) + def test_dynolog_collects_a_kineto_trace_from_pytorch( + self, global_data, clean_test_temp_files + ): + """torch_profiler_dynolog captures an on-demand trace from a PyTorch run. + + The warmup is shortened from the 60s default because the fixture workload + is far shorter lived than a real training job. + """ + global_data["console"].sh( + build_run_command( + "dummy_torchrun", + additional_context=run_context( + [ + { + "name": "torch_profiler_dynolog", + "env_vars": { + "TORCH_PROFILE_WARMUP_S": "20", + "TORCH_PROFILE_RETRY_INTERVAL_S": "5", + "TORCH_PROFILE_MAX_ATTEMPTS": "10", + }, + } + ] + ), + ), + canFail=True, + ) + + output_dir = os.path.join(BASE_DIR, "torch_profiler_output") + if not os.path.isdir(output_dir): + pytest.fail( + "torch_profiler_output/ not collected with the " + "torch_profiler_dynolog tool." + ) + collected = os.listdir(output_dir) + traces = [f for f in collected if f.endswith((".json", ".json.gz"))] + if not traces: + pytest.fail( + "no Kineto trace captured; dynolog never matched a PyTorch process " + f"(torch_profiler_output/ contains {collected})." + ) diff --git a/tests/integration/test_tracelens_tools_config.py b/tests/integration/test_tracelens_tools_config.py new file mode 100644 index 00000000..e55a8576 --- /dev/null +++ b/tests/integration/test_tracelens_tools_config.py @@ -0,0 +1,202 @@ +"""Integration tests for the TraceLens and dynolog tools: tools.json wiring. + +Verifies the shipped tools.json entries and that ContainerRunner.apply_tools +stacks them correctly, including the profiler-then-analysis ordering that makes +the tracelens tool useful. +""" + +import json +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +from madengine.execution.container_runner import ContainerRunner +from madengine.utils.path_utils import get_madengine_root + + +def _tools_json() -> Path: + return get_madengine_root() / "scripts" / "common" / "tools.json" + + +@pytest.fixture(scope="module") +def tools() -> dict: + with open(_tools_json(), encoding="utf-8") as f: + return json.load(f)["tools"] + + +def _apply(tool_names) -> tuple: + """Run apply_tools for the given tool names, returning (scripts, env).""" + ctx = MagicMock() + ctx.ctx = {"tools": [{"name": name} for name in tool_names]} + runner = ContainerRunner(context=ctx, console=MagicMock()) + scripts = { + "pre_scripts": [], + "encapsulate_script": "bash model_run.sh", + "post_scripts": [], + } + env: dict = {} + runner.apply_tools(scripts, env, str(_tools_json())) + return scripts, env + + +class TestDynologToolConfig: + """torch_profiler_dynolog drives torch.profiler on unmodified workloads.""" + + def test_config_sets_kineto_daemon_env(self, tools): + cfg = tools["torch_profiler_dynolog"] + assert cfg["env_vars"]["KINETO_USE_DAEMON"] == "1" + assert "KINETO_DAEMON_INIT_DELAY_S" in cfg["env_vars"] + assert cfg["env_vars"]["TORCH_PROFILE_OUTPUT_DIR"] == "torch_profiler_output" + + def test_requests_the_metadata_tracelens_needs(self, tools): + """Roofline and per-op analysis need input shapes; nn.Module view needs modules.""" + env = tools["torch_profiler_dynolog"]["env_vars"] + assert env["TORCH_PROFILE_RECORD_SHAPES"] == "1" + assert env["TORCH_PROFILE_WITH_STACKS"] == "1" + assert env["TORCH_PROFILE_WITH_MODULES"] == "1" + + def test_raises_upstream_process_limit_for_multi_gpu(self, tools): + """dyno gputrace defaults to 3 processes, which silently drops most ranks.""" + assert int(tools["torch_profiler_dynolog"]["env_vars"]["TORCH_PROFILE_PROCESS_LIMIT"]) >= 8 + + def test_does_not_wrap_the_model_command(self, tools): + """Tracing is triggered out-of-band, so the workload command is untouched.""" + assert tools["torch_profiler_dynolog"]["cmd"] == "" + + def test_installs_then_starts_then_stops_then_collects(self, tools): + cfg = tools["torch_profiler_dynolog"] + assert [Path(s["path"]).name for s in cfg["pre_scripts"]] == [ + "trace.sh", + "dynolog_start.sh", + ] + assert cfg["pre_scripts"][0]["args"] == "dynolog" + assert [Path(s["path"]).name for s in cfg["post_scripts"]] == [ + "dynolog_stop.sh", + "trace.sh", + ] + assert cfg["post_scripts"][1]["args"] == "torch_profiler" + + def test_referenced_scripts_exist(self, tools): + root = get_madengine_root() + cfg = tools["torch_profiler_dynolog"] + for script in cfg["pre_scripts"] + cfg["post_scripts"]: + assert (root / script["path"]).is_file(), script["path"] + # dynolog_start.sh launches the trigger, which tools.json does not name. + assert (root / "scripts/common/tools/dynolog_trigger.sh").is_file() + + def test_apply_tools_wires_env_and_scripts(self): + scripts, env = _apply(["torch_profiler_dynolog"]) + assert env["KINETO_USE_DAEMON"] == "1" + # An empty cmd must not disturb the model invocation. + assert scripts["encapsulate_script"].strip() == "bash model_run.sh" + assert any( + Path(s["path"]).name == "dynolog_start.sh" for s in scripts["pre_scripts"] + ) + assert any( + Path(s["path"]).name == "dynolog_stop.sh" for s in scripts["post_scripts"] + ) + + def test_not_in_rocprof_family(self): + """dynolog does not need rocprofv3, so multi-node runs must not drop it.""" + from madengine.deployment.common import tools_include_rocprof_family + + assert not tools_include_rocprof_family([{"name": "torch_profiler_dynolog"}]) + + +class TestTraceLensToolConfig: + """The tracelens tools run TraceLens over whatever traces the run produced.""" + + ALL_VARIANTS = ( + "tracelens", + "tracelens_pytorch", + "tracelens_rocprof", + "tracelens_pftrace", + "tracelens_collective", + ) + + def test_all_variants_are_defined(self, tools): + for name in self.ALL_VARIANTS: + assert name in tools, name + + def test_variants_differ_only_by_mode(self, tools): + modes = { + name: tools[name]["env_vars"]["TRACELENS_MODE"] for name in self.ALL_VARIANTS + } + assert modes == { + "tracelens": "auto", + "tracelens_pytorch": "pytorch", + "tracelens_rocprof": "rocprof", + "tracelens_pftrace": "pftrace", + "tracelens_collective": "collective", + } + + def test_installs_into_an_isolated_venv(self, tools): + """TraceLens pins protobuf/xprof; it must not touch the workload's env.""" + venv = tools["tracelens"]["env_vars"]["TRACELENS_VENV"] + assert venv.startswith("/") + assert "site-packages" not in venv + + def test_analysis_runs_after_the_model_not_around_it(self, tools): + cfg = tools["tracelens"] + assert cfg["cmd"] == "" + assert [Path(s["path"]).name for s in cfg["post_scripts"]] == [ + "tracelens.sh", + "trace.sh", + ] + assert cfg["post_scripts"][1]["args"] == "tracelens" + + def test_referenced_scripts_exist(self, tools): + root = get_madengine_root() + for name in self.ALL_VARIANTS: + cfg = tools[name] + for script in cfg["pre_scripts"] + cfg["post_scripts"]: + assert (root / script["path"]).is_file(), script["path"] + assert (root / "scripts/common/tools/tracelens_analyze.py").is_file() + + def test_stacks_after_a_profiler(self): + """Profiler setup must precede analysis, and analysis must run last.""" + scripts, env = _apply(["rocprofv3_perfetto", "tracelens"]) + + # The profiler still wraps the model command. + assert "rocprof_wrapper.sh" in scripts["encapsulate_script"] + assert "bash model_run.sh" in scripts["encapsulate_script"] + + post = [Path(s["path"]).name for s in scripts["post_scripts"]] + # rocprofv3_perfetto collects its trace before TraceLens reads it. + assert post.index("trace.sh") < post.index("tracelens.sh") + assert env["TRACELENS_MODE"] == "auto" + + def test_stacks_with_dynolog_for_the_full_pytorch_path(self): + scripts, env = _apply(["torch_profiler_dynolog", "tracelens"]) + + assert env["KINETO_USE_DAEMON"] == "1" + assert env["TRACELENS_MODE"] == "auto" + pre = [Path(s["path"]).name for s in scripts["pre_scripts"]] + assert "dynolog_start.sh" in pre + post = [Path(s["path"]).name for s in scripts["post_scripts"]] + # Kineto traces are collected before TraceLens analyses them. + assert post.index("dynolog_stop.sh") < post.index("tracelens.sh") + + def test_env_vars_can_be_overridden_per_run(self): + ctx = MagicMock() + ctx.ctx = { + "tools": [ + {"name": "tracelens", "env_vars": {"TRACELENS_GPU_ARCH": "MI300X"}} + ] + } + runner = ContainerRunner(context=ctx, console=MagicMock()) + scripts = { + "pre_scripts": [], + "encapsulate_script": "bash model_run.sh", + "post_scripts": [], + } + env: dict = {} + runner.apply_tools(scripts, env, str(_tools_json())) + assert env["TRACELENS_GPU_ARCH"] == "MI300X" + assert env["TRACELENS_MODE"] == "auto" + + def test_not_in_rocprof_family(self): + from madengine.deployment.common import tools_include_rocprof_family + + assert not tools_include_rocprof_family([{"name": name} for name in self.ALL_VARIANTS]) diff --git a/tests/unit/test_tracelens_analyze.py b/tests/unit/test_tracelens_analyze.py new file mode 100644 index 00000000..53d016e3 --- /dev/null +++ b/tests/unit/test_tracelens_analyze.py @@ -0,0 +1,335 @@ +"""Unit tests for the TraceLens trace analyzer script. + +The analyzer is shipped as a standalone stdlib-only script under +``scripts/common/tools/`` so it can run both inside a workload container and on +the host, so it is loaded here by path rather than imported as a module. +""" + +import csv +import importlib.util +import json +import sys +from pathlib import Path + +import pytest + +from madengine.utils.path_utils import get_madengine_root + +CHROME_TRACE = b'{"traceEvents": [], "schemaVersion": 1}' + + +def _load_analyzer(): + script = ( + get_madengine_root() / "scripts" / "common" / "tools" / "tracelens_analyze.py" + ) + assert script.is_file(), f"analyzer script missing at {script}" + spec = importlib.util.spec_from_file_location("tracelens_analyze", script) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def analyzer(): + return _load_analyzer() + + +def _write(root: Path, rel: str, content: bytes = b"{}") -> Path: + path = root / rel + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(content) + return path + + +@pytest.fixture +def trace_tree(tmp_path): + """A directory shaped like the working directory after a profiled run.""" + _write(tmp_path, "torch_profiler_output/libkineto_trace_rank0_1.json", CHROME_TRACE) + _write(tmp_path, "torch_profiler_output/libkineto_trace_rank1_2.json", CHROME_TRACE) + _write(tmp_path, "traces/model_rank2.pt.trace.json", CHROME_TRACE) + _write(tmp_path, "rocprof_output/9999_results.json", b"{}") + _write(tmp_path, "rocprof_output/model_trace.pftrace", b"\x00\x01") + _write(tmp_path, "rocprof_output/1e4d92661463/1234_results.db", b"sqlite") + _write(tmp_path, "rpd_output/trace.rpd", b"sqlite") + _write(tmp_path, "rpd_output/trace.json", CHROME_TRACE) + _write(tmp_path, "perf.csv", b"model,performance\n") + return tmp_path + + +class TestDiscovery: + """Trace classification must route each artifact to the right TraceLens reader.""" + + def test_classifies_each_trace_kind(self, analyzer, trace_tree): + traces, _ = analyzer.discover_traces(str(trace_tree)) + names = { + kind: sorted(Path(p).name for p in paths) for kind, paths in traces.items() + } + + assert names[analyzer.KIND_PYTORCH] == [ + "libkineto_trace_rank0_1.json", + "libkineto_trace_rank1_2.json", + "model_rank2.pt.trace.json", + "trace.json", + ] + assert names[analyzer.KIND_ROCPROF_JSON] == ["9999_results.json"] + assert names[analyzer.KIND_PFTRACE] == ["model_trace.pftrace"] + + def test_reports_unreadable_formats_with_guidance(self, analyzer, trace_tree): + _, unsupported = analyzer.discover_traces(str(trace_tree)) + by_name = {Path(p).name: reason for p, reason in unsupported} + + assert "1234_results.db" in by_name + assert "--output-format" in by_name["1234_results.db"] + assert "trace.rpd" in by_name + + def test_sniffs_ambiguous_trace_json(self, analyzer, tmp_path): + """trace.json is written by both rpd (Chrome trace) and rocm-trace-lite.""" + _write(tmp_path, "rpd_output/trace.json", CHROME_TRACE) + _write(tmp_path, "other_output/trace.json", b'{"not": "a trace"}') + + traces, _ = analyzer.discover_traces(str(tmp_path)) + claimed = [Path(p).parent.name for p in traces.get(analyzer.KIND_PYTORCH, [])] + assert claimed == ["rpd_output"] + + def test_excludes_report_output_directory(self, analyzer, tmp_path): + """Re-running analysis must not treat previous reports as new inputs.""" + _write(tmp_path, "torch_profiler_output/libkineto_trace_1.json", CHROME_TRACE) + _write(tmp_path, "tracelens_output/stale_results.json", b"{}") + + traces, unsupported = analyzer.discover_traces( + str(tmp_path), exclude_dirs=[str(tmp_path / "tracelens_output")] + ) + assert analyzer.KIND_ROCPROF_JSON not in traces + assert unsupported == [] + + def test_empty_tree_discovers_nothing(self, analyzer, tmp_path): + traces, unsupported = analyzer.discover_traces(str(tmp_path)) + assert traces == {} + assert unsupported == [] + + +class TestCommandConstruction: + """TraceLens must be invoked with the flags each report generator expects.""" + + def test_falls_back_to_module_when_console_script_absent(self, analyzer): + command = analyzer._build_command( + "/nonexistent/bin/python3", + "TraceLens_generate_perf_report_pytorch", + ["--profile_json_path", "trace.json"], + ) + assert command[0] == "/nonexistent/bin/python3" + assert command[1] == "-c" + assert "TraceLens.Reporting.generate_perf_report_pytorch" in command[2] + assert command[-2:] == ["--profile_json_path", "trace.json"] + + def test_prefers_console_script_in_interpreter_bindir(self, analyzer, tmp_path): + bindir = tmp_path / "bin" + bindir.mkdir() + python = bindir / "python3" + python.write_text("") + script = bindir / "TraceLens_generate_perf_report_rocprof" + script.write_text("") + script.chmod(0o755) + + command = analyzer._build_command( + str(python), "TraceLens_generate_perf_report_rocprof", ["--x"] + ) + assert command == [str(script), "--x"] + + def test_every_entry_point_has_a_module_fallback(self, analyzer): + for name, module in analyzer._ENTRY_POINTS.items(): + assert name.startswith("TraceLens_") + assert module.startswith("TraceLens.Reporting.") + + def test_pytorch_args_request_shapes_and_roofline(self, analyzer): + args = analyzer._pytorch_args("t.json", "/out/t", "MI300X", []) + assert args[:2] == ["--profile_json_path", "t.json"] + assert "--output_csvs_dir" in args + assert args[args.index("--gpu_arch_platform") + 1] == "MI300X" + + def test_pytorch_args_omit_roofline_without_arch(self, analyzer): + args = analyzer._pytorch_args("t.json", "/out/t", None, []) + assert "--gpu_arch_platform" not in args + + def test_pftrace_produces_three_complementary_reports(self, analyzer): + jobs = analyzer._pftrace_jobs("t.pftrace", "/out/t", []) + assert [tool for tool, _ in jobs] == [ + "TraceLens_generate_perf_report_pftrace_hip_activity", + "TraceLens_generate_perf_report_pftrace_hip_api", + "TraceLens_generate_perf_report_pftrace_memory_copy", + ] + for _, args in jobs: + assert args[:2] == ["--trace_path", "t.pftrace"] + + def test_collective_args_carry_rank_regex_and_world_size(self, analyzer): + args = analyzer._collective_args("/root", "/out/coll", 8, []) + assert args[args.index("--world_size") + 1] == "8" + assert "rank" in args[args.index("--rank_regex") + 1] + + def test_extra_args_are_forwarded(self, analyzer): + args = analyzer._pytorch_args("t.json", "/out/t", None, ["--detect_recompute"]) + assert args[-1] == "--detect_recompute" + + +class TestReportStem: + """Report names must be unique per trace and safe as filenames.""" + + def test_strips_known_trace_suffixes(self, analyzer, tmp_path): + stem = analyzer._report_stem( + str(tmp_path / "traces" / "model_rank0.pt.trace.json"), str(tmp_path) + ) + assert stem == "traces_model_rank0" + + def test_distinguishes_same_name_in_different_directories(self, analyzer, tmp_path): + a = analyzer._report_stem(str(tmp_path / "node_0" / "trace.json"), str(tmp_path)) + b = analyzer._report_stem(str(tmp_path / "node_1" / "trace.json"), str(tmp_path)) + assert a != b + + +class TestAnalyze: + """The analyze() driver must schedule one job per trace and record outcomes.""" + + def test_schedules_a_job_for_every_trace_and_writes_summary( + self, analyzer, trace_tree, monkeypatch + ): + calls = [] + + def fake_run(command, cwd=None): + calls.append(list(command)) + return 0, "" + + monkeypatch.setattr(analyzer, "_run", fake_run) + out = trace_tree / "tracelens_output" + summary = analyzer.analyze( + root=str(trace_tree), output_dir=str(out), python=sys.executable + ) + + # 4 pytorch + 1 rocprof + 3 pftrace + 1 collective + assert len(calls) == 9 + assert summary["succeeded"] == 9 + assert summary["failed"] == 0 + # The unreadable .db and .rpd artifacts are surfaced as skipped. + assert summary["skipped"] == 2 + + with open(summary["summary_csv"], newline="", encoding="utf-8") as handle: + rows = list(csv.DictReader(handle)) + assert len(rows) == 11 + assert set(rows[0]) == set(analyzer.SUMMARY_CSV_FIELDS) + + def test_records_failure_detail(self, analyzer, tmp_path, monkeypatch): + _write(tmp_path, "rocprof_output/1_results.json", b"{}") + monkeypatch.setattr( + analyzer, "_run", lambda command, cwd=None: (1, "boom\nNot a valid file") + ) + + summary = analyzer.analyze( + root=str(tmp_path), output_dir=str(tmp_path / "out"), python=sys.executable + ) + assert summary["failed"] == 1 + assert summary["results"][0]["detail"] == "Not a valid file" + + def test_mode_restricts_to_one_trace_kind(self, analyzer, trace_tree, monkeypatch): + calls = [] + monkeypatch.setattr( + analyzer, + "_run", + lambda command, cwd=None: (calls.append(list(command)), (0, ""))[1], + ) + + analyzer.analyze( + root=str(trace_tree), + output_dir=str(trace_tree / "out"), + mode="rocprof", + python=sys.executable, + ) + assert len(calls) == 1 + assert any("generate_perf_report_rocprof" in part for part in calls[0]) + + def test_collective_mode_emits_only_the_multi_rank_report( + self, analyzer, trace_tree, monkeypatch + ): + calls = [] + monkeypatch.setattr( + analyzer, + "_run", + lambda command, cwd=None: (calls.append(list(command)), (0, ""))[1], + ) + + analyzer.analyze( + root=str(trace_tree), + output_dir=str(trace_tree / "out"), + mode="collective", + python=sys.executable, + ) + assert len(calls) == 1 + assert any("multi_rank_collective_report" in part for part in calls[0]) + + def test_no_collective_report_for_a_single_rank(self, analyzer, tmp_path, monkeypatch): + _write(tmp_path, "torch_profiler_output/libkineto_trace_1.json", CHROME_TRACE) + calls = [] + monkeypatch.setattr( + analyzer, + "_run", + lambda command, cwd=None: (calls.append(list(command)), (0, ""))[1], + ) + + analyzer.analyze( + root=str(tmp_path), output_dir=str(tmp_path / "out"), python=sys.executable + ) + assert len(calls) == 1 + assert not any("multi_rank" in part for call in calls for part in call) + + def test_max_traces_caps_work_per_kind(self, analyzer, trace_tree, monkeypatch): + calls = [] + monkeypatch.setattr( + analyzer, + "_run", + lambda command, cwd=None: (calls.append(list(command)), (0, ""))[1], + ) + + analyzer.analyze( + root=str(trace_tree), + output_dir=str(trace_tree / "out"), + mode="pytorch", + max_traces=2, + python=sys.executable, + ) + assert len(calls) == 2 + + def test_empty_tree_reports_nothing_without_error(self, analyzer, tmp_path): + summary = analyzer.analyze( + root=str(tmp_path), output_dir=str(tmp_path / "out"), python=sys.executable + ) + assert summary["results"] == [] + assert summary["succeeded"] == 0 + + +class TestCli: + """The script's CLI is the contract used by tracelens.sh and the host wrapper.""" + + def test_discover_only_writes_json_summary_and_succeeds( + self, analyzer, trace_tree, capsys + ): + summary_path = trace_tree / "discovery.json" + code = analyzer.main( + ["--root", str(trace_tree), "--discover-only", "--json-summary", str(summary_path)] + ) + capsys.readouterr() + + assert code == 0 + summary = json.loads(summary_path.read_text(encoding="utf-8")) + assert summary["discovered"]["pytorch"] == 4 + assert len(summary["unsupported"]) == 2 + + def test_exit_code_reflects_failures(self, analyzer, tmp_path, monkeypatch, capsys): + _write(tmp_path, "rocprof_output/1_results.json", b"{}") + monkeypatch.setattr(analyzer, "_run", lambda command, cwd=None: (1, "failed")) + + code = analyzer.main(["--root", str(tmp_path), "--output-dir", str(tmp_path / "out")]) + capsys.readouterr() + assert code == 1 + + def test_rejects_unknown_mode(self, analyzer): + with pytest.raises(SystemExit): + analyzer.main(["--mode", "nonsense"]) diff --git a/tests/unit/test_tracelens_report.py b/tests/unit/test_tracelens_report.py new file mode 100644 index 00000000..c4440604 --- /dev/null +++ b/tests/unit/test_tracelens_report.py @@ -0,0 +1,282 @@ +"""Unit tests for host-side TraceLens report generation and its CLI commands.""" + +import json +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +from typer.testing import CliRunner + +from madengine.reporting import tracelens_report as tlr + + +@pytest.fixture +def runner(): + return CliRunner() + + +@pytest.fixture +def report_app(): + from madengine.cli.commands.report import report_app + + return report_app + + +class TestAnalyzerDiscovery: + """The host path drives the same packaged analyzer as the in-container tool.""" + + def test_finds_the_packaged_analyzer_script(self): + script = tlr.find_analyzer_script() + assert script.is_file() + assert script.name == "tracelens_analyze.py" + + def test_raises_when_the_script_is_missing(self, tmp_path): + with patch.object(tlr, "get_madengine_root", return_value=tmp_path): + with pytest.raises(FileNotFoundError, match="Trace analyzer not found"): + tlr.find_analyzer_script() + + +class TestResolvePython: + """TraceLens runs out-of-process so its pinned deps stay isolated.""" + + def test_explicit_python_wins(self): + assert tlr.resolve_python("/custom/python3") == "/custom/python3" + + def test_falls_back_to_tracelens_venv(self, tmp_path, monkeypatch): + bindir = tmp_path / "bin" + bindir.mkdir() + (bindir / "python3").write_text("") + monkeypatch.setenv("TRACELENS_VENV", str(tmp_path)) + + assert tlr.resolve_python() == str(bindir / "python3") + + def test_ignores_a_venv_without_an_interpreter(self, tmp_path, monkeypatch): + import sys + + monkeypatch.setenv("TRACELENS_VENV", str(tmp_path)) + assert tlr.resolve_python() == sys.executable + + +class TestGenerateReports: + """generate_tracelens_reports must fail fast and surface the analyzer summary.""" + + def test_requires_tracelens_to_be_installed(self, tmp_path): + with patch.object(tlr, "check_tracelens_available", return_value=False): + with pytest.raises(tlr.TraceLensNotInstalledError, match="madengine\\[tracelens\\]"): + tlr.generate_tracelens_reports(root=str(tmp_path)) + + def test_passes_options_through_to_the_analyzer(self, tmp_path): + captured = {} + + def fake_run(command): + captured["command"] = list(command) + summary_path = Path(command[command.index("--json-summary") + 1]) + summary_path.write_text(json.dumps({"succeeded": 2}), encoding="utf-8") + return MagicMock(returncode=0) + + with patch.object(tlr, "check_tracelens_available", return_value=True), patch.object( + tlr.subprocess, "run", side_effect=fake_run + ): + summary = tlr.generate_tracelens_reports( + root=str(tmp_path), + output_dir=str(tmp_path / "out"), + mode="pytorch", + python="/venv/bin/python3", + gpu_arch="MI300X", + world_size=8, + max_traces=4, + ) + + command = captured["command"] + assert command[command.index("--mode") + 1] == "pytorch" + assert command[command.index("--python") + 1] == "/venv/bin/python3" + assert command[command.index("--gpu-arch") + 1] == "MI300X" + assert command[command.index("--world-size") + 1] == "8" + assert command[command.index("--max-traces") + 1] == "4" + assert summary["succeeded"] == 2 + assert summary["exit_code"] == 0 + + def test_omits_unset_options(self, tmp_path): + captured = {} + + def fake_run(command): + captured["command"] = list(command) + return MagicMock(returncode=0) + + with patch.object(tlr, "check_tracelens_available", return_value=True), patch.object( + tlr.subprocess, "run", side_effect=fake_run + ): + tlr.generate_tracelens_reports( + root=str(tmp_path), output_dir=str(tmp_path / "out") + ) + + assert "--gpu-arch" not in captured["command"] + assert "--world-size" not in captured["command"] + + def test_creates_the_output_directory(self, tmp_path): + out = tmp_path / "nested" / "out" + with patch.object(tlr, "check_tracelens_available", return_value=True), patch.object( + tlr.subprocess, "run", return_value=MagicMock(returncode=0) + ): + tlr.generate_tracelens_reports(root=str(tmp_path), output_dir=str(out)) + assert out.is_dir() + + def test_discovery_does_not_require_tracelens(self, tmp_path): + """Users must be able to see what traces exist before installing TraceLens.""" + with patch.object(tlr, "check_tracelens_available", return_value=False), patch.object( + tlr.subprocess, "run", return_value=MagicMock(returncode=0) + ) as run: + tlr.discover_traces(root=str(tmp_path)) + assert "--discover-only" in run.call_args[0][0] + + +class TestCompareReports: + def test_rejects_a_single_report(self): + with pytest.raises(ValueError, match="at least two"): + tlr.compare_tracelens_reports(["only.xlsx"]) + + def test_rejects_mismatched_names(self): + with pytest.raises(ValueError, match="counts must match"): + tlr.compare_tracelens_reports(["a.xlsx", "b.xlsx"], names=["only-one"]) + + def test_requires_tracelens_to_be_installed(self): + with patch.object(tlr, "check_tracelens_available", return_value=False): + with pytest.raises(tlr.TraceLensNotInstalledError): + tlr.compare_tracelens_reports(["a.xlsx", "b.xlsx"]) + + def test_forwards_reports_and_names(self): + captured = {} + + def fake_run(command): + captured["command"] = list(command) + return MagicMock(returncode=0) + + with patch.object(tlr, "check_tracelens_available", return_value=True), patch.object( + tlr.subprocess, "run", side_effect=fake_run + ): + tlr.compare_tracelens_reports( + ["base.xlsx", "cand.xlsx"], output="diff.xlsx", names=["base", "cand"] + ) + + command = captured["command"] + assert command[command.index("--compare") + 1 : command.index("--compare") + 3] == [ + "base.xlsx", + "cand.xlsx", + ] + assert command[command.index("--compare-output") + 1] == "diff.xlsx" + assert command[command.index("--compare-names") + 1 :] == ["base", "cand"] + + +class TestReportTraceLensCli: + def test_rejects_an_invalid_mode(self, runner, report_app, tmp_path): + result = runner.invoke( + report_app, ["tracelens", "--root", str(tmp_path), "--mode", "nonsense"] + ) + assert result.exit_code != 0 + assert "invalid --mode" in result.output + + def test_rejects_a_missing_root(self, runner, report_app): + result = runner.invoke(report_app, ["tracelens", "--root", "does/not/exist"]) + assert result.exit_code != 0 + assert "directory not found" in result.output + + def test_reports_install_guidance_when_tracelens_is_absent( + self, runner, report_app, tmp_path + ): + with patch( + "madengine.cli.commands.report.generate_tracelens_reports", + side_effect=tlr.TraceLensNotInstalledError(tlr.INSTALL_HINT), + ): + result = runner.invoke(report_app, ["tracelens", "--root", str(tmp_path)]) + assert result.exit_code != 0 + assert "madengine[tracelens]" in result.output + + def test_succeeds_and_lists_generated_reports(self, runner, report_app, tmp_path): + summary = { + "succeeded": 1, + "failed": 0, + "skipped": 0, + "results": [ + { + "trace_file": "rocprof_output/1_results.json", + "kind": "rocprof_json", + "tracelens_tool": "TraceLens_generate_perf_report_rocprof", + "status": "SUCCESS", + "detail": "", + } + ], + } + with patch( + "madengine.cli.commands.report.generate_tracelens_reports", + return_value=summary, + ): + result = runner.invoke(report_app, ["tracelens", "--root", str(tmp_path)]) + assert result.exit_code == 0 + assert "1 report(s) generated" in result.output + + def test_fails_when_a_report_fails(self, runner, report_app, tmp_path): + summary = { + "succeeded": 0, + "failed": 1, + "skipped": 0, + "results": [ + { + "trace_file": "t.json", + "kind": "pytorch", + "tracelens_tool": "TraceLens_generate_perf_report_pytorch", + "status": "FAILURE", + "detail": "Not a valid trace", + } + ], + } + with patch( + "madengine.cli.commands.report.generate_tracelens_reports", + return_value=summary, + ): + result = runner.invoke(report_app, ["tracelens", "--root", str(tmp_path)]) + assert result.exit_code != 0 + assert "1 failed" in result.output + + def test_guides_the_user_when_no_traces_exist(self, runner, report_app, tmp_path): + with patch( + "madengine.cli.commands.report.generate_tracelens_reports", + return_value={"succeeded": 0, "failed": 0, "skipped": 0, "results": []}, + ): + result = runner.invoke(report_app, ["tracelens", "--root", str(tmp_path)]) + assert result.exit_code == 0 + assert "torch_profiler_dynolog" in result.output + + def test_discover_only_lists_kinds(self, runner, report_app, tmp_path): + with patch( + "madengine.cli.commands.report.discover_traces", + return_value={"discovered": {"pytorch": 3}, "unsupported": []}, + ): + result = runner.invoke( + report_app, ["tracelens", "--root", str(tmp_path), "--discover-only"] + ) + assert result.exit_code == 0 + assert "pytorch" in result.output + + +class TestReportTraceLensCompareCli: + def test_rejects_missing_report_files(self, runner, report_app): + result = runner.invoke(report_app, ["tracelens-compare", "a.xlsx", "b.xlsx"]) + assert result.exit_code != 0 + assert "not found" in result.output + + def test_succeeds_on_a_clean_comparison(self, runner, report_app, tmp_path): + a = tmp_path / "a.xlsx" + b = tmp_path / "b.xlsx" + a.write_text("") + b.write_text("") + out = str(tmp_path / "diff.xlsx") + + with patch( + "madengine.cli.commands.report.compare_tracelens_reports", + return_value={"status": "SUCCESS"}, + ): + result = runner.invoke( + report_app, ["tracelens-compare", str(a), str(b), "-o", out] + ) + assert result.exit_code == 0 + assert "Comparison written to" in result.output