feat(profiling): integrate TraceLens trace analysis - #167
Open
Cemberk wants to merge 1 commit into
Open
Conversation
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 <cursoragent@cursor.com>
Cemberk
requested review from
Rohan138,
coketaste,
gargrahul and
leconcio
as code owners
August 11, 2026 15:19
There was a problem hiding this comment.
Pull request overview
Integrates TraceLens-based trace analysis into madengine, providing both host-side reporting (madengine report tracelens / tracelens-compare) and in-container tooling to capture PyTorch Kineto traces on-demand (dynolog) and generate TraceLens reports. This fits into the reporting + tools pipeline by turning collected profiler artifacts into actionable operator/kernel/roofline/collective outputs and ensuring distributed collectors gather the needed directories.
Changes:
- Add a stdlib-only TraceLens analyzer script plus a host-side wrapper module and new report CLI commands (
tracelens,tracelens-compare). - Add in-container tools (
torch_profiler_dynolog,tracelensand mode-specific variants) with supporting pre/post scripts and TraceLens-venv isolation. - Extend SLURM/Kubernetes artifact collection, add examples/docs, and add unit/integration/e2e coverage.
Reviewed changes
Copilot reviewed 24 out of 25 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit/test_tracelens_report.py | Unit tests for host-side TraceLens wrapper functions and report CLI behavior. |
| tests/unit/test_tracelens_analyze.py | Unit tests for the standalone analyzer script (discovery, command construction, summaries). |
| tests/integration/test_tracelens_tools_config.py | Integration tests validating tools.json wiring and tool stacking behavior. |
| tests/e2e/test_tracelens_workflows.py | E2E coverage for host reporting and (GPU) container tool workflows. |
| src/madengine/scripts/common/tools/tracelens_analyze.py | Standalone (stdlib-only) analyzer that discovers traces and runs TraceLens entry points out-of-process. |
| src/madengine/scripts/common/tools/dynolog_trigger.sh | Background trigger script to request torch.profiler traces via dynolog. |
| src/madengine/scripts/common/tools.json | Registers torch_profiler_dynolog and TraceLens tool variants and their scripts/env. |
| src/madengine/scripts/common/pre_scripts/trace.sh | Adds installers for dynolog and TraceLens (isolated venv + optional traceconv staging). |
| src/madengine/scripts/common/pre_scripts/dynolog_start.sh | Starts dynolog daemon and arms the trigger alongside the workload. |
| src/madengine/scripts/common/post_scripts/tracelens.sh | Runs the analyzer post-workload to generate TraceLens reports without failing the run. |
| src/madengine/scripts/common/post_scripts/trace.sh | Collects new output directories (torch_profiler_output, tracelens_output) as artifacts. |
| src/madengine/scripts/common/post_scripts/dynolog_stop.sh | Stops dynolog/trigger and reports captured trace count + tail logs. |
| src/madengine/reporting/tracelens_report.py | Host-side driver for running the analyzer and surfacing install guidance/errors. |
| src/madengine/deployment/templates/slurm/job.sh.j2 | Collects additional per-node tool output directories on successful SLURM runs. |
| src/madengine/deployment/templates/kubernetes/job.yaml.j2 | Copies additional per-pod tool output directories into results for K8s runs. |
| src/madengine/deployment/k8s_scripts.py | Ensures helper scripts referenced by pre/post scripts under scripts/common/tools are bundled for K8s. |
| src/madengine/deployment/k8s_results.py | Collects .pftrace and additional known output directories via kubectl cp. |
| src/madengine/cli/commands/report.py | Adds report tracelens and report tracelens-compare commands with Rich output. |
| pyproject.toml | Adds optional dependency extra madengine[tracelens] pinned to a TraceLens git ref. |
| examples/profiling-configs/tracelens_rocprofv3.json | Example config for rocprofv3_lightweight + TraceLens analysis. |
| examples/profiling-configs/torch_profiler_tracelens.json | Example config for dynolog capture + TraceLens PyTorch analysis. |
| examples/profiling-configs/README.md | Documents TraceLens/dynolog example configs and produced artifacts. |
| docs/profiling.md | Adds documentation for dynolog-based capture and TraceLens host/container analysis workflows. |
| docs/cli-reference.md | Adds CLI reference entries for report tracelens and report tracelens-compare. |
| .gitignore | Anchors scripts/ and *.json ignore rules to repo root to avoid ignoring packaged runtime scripts. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+187
to
+190
| 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 |
| return _run_analyzer(args, summary_path) | ||
|
|
||
|
|
||
| def discover_traces(root: str = ".", output_dir: str = "tracelens_output") -> Dict[str, object]: |
Comment on lines
236
to
239
| 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: |
5 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 tracelenskeeps TraceLens' pinned protobuf and xprof out of workload images entirely. In-container, thetracelenstool 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*.jsonignore 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.