Basic integration with the py-spy profiler #356
Conversation
jarlsondre
changed the title
jarlsondre
changed the title
jarlsondre
changed the title
jarlsondre
changed the title
|
Printing multiple lines for identical function calls is not a good idea if one cannot relate them back to the call stack in the file. I would suggest aggregating identical function calls for the printed overview. Otherwise if there are identifiers per call stack, one can also add those as a column to the overview (but that may be a bit confusing too). |
Added this feature now, also updated the screenshot at the top |
jarlsondre
changed the title
There was a problem hiding this comment.
Pull Request Overview
This PR adds basic integration with the py‑spy profiler into the project. Key changes include:
- Adding py‑spy and tabulate to the project dependencies.
- Implementing new CLI commands for generating a py‑spy report and a flamegraph.
- Integrating py‑spy profiling options into the SLURM script generation flow for multi‑GPU training.
Reviewed Changes
Copilot reviewed 11 out of 12 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| use-cases/mnist/torch/config.yaml | Updated configuration flags for GPU data and communication overhead measurement. |
| src/itwinai/torch/profiling/py_spy_aggregation.py | Added functions to parse and aggregate profiling stack traces. |
| src/itwinai/slurm/utils.py | Introduced new arguments (profiling rate and py‑spy flag) for the SLURM job parser. |
| src/itwinai/slurm/slurm_script_builder.py | Extended script builder to optionally prepend the py‑spy record command to the training command. |
| src/itwinai/slurm/slurm_constants.py | Updated directives to include GPU resource options. |
| src/itwinai/cli.py | Added new CLI commands for flamegraph and py‑spy report generation, along with replacing print with typer.echo. |
| pyproject.toml | Added dependencies for py‑spy and tabulate. |
| licenses/CDDL-1.0.txt | Included license information for flamegraph.pl. |
| .github/workflows/lint.yml | Updated linter configuration to ignore Perl files. |
Files not reviewed (1)
- THIRD_PARTY_LICENSES: Language not supported
Comments suppressed due to low confidence (1)
src/itwinai/cli.py:68
- Ensure that the new py‑spy report generation functionality is adequately covered by unit tests, especially to handle potential edge cases in parsing and data aggregation.
def generate_py_spy_report(
| # Prepending the py-spy profiling command | ||
| py_spy_profiling_output = "py_spy_profiling_$SLURM_NODEID.txt" | ||
| bash_command = ( | ||
| f"py-spy record -r {self.profiling_sampling_rate} -s " | ||
| f"-o {py_spy_profiling_output} -f raw -- " + bash_command | ||
| ) |
There was a problem hiding this comment.
[nitpick] Consider extracting the construction and prepending of the py‑spy profiling command into a helper function to improve readability and ease future modifications.
| # Prepending the py-spy profiling command | |
| py_spy_profiling_output = "py_spy_profiling_$SLURM_NODEID.txt" | |
| bash_command = ( | |
| f"py-spy record -r {self.profiling_sampling_rate} -s " | |
| f"-o {py_spy_profiling_output} -f raw -- " + bash_command | |
| ) | |
| bash_command = self._prepend_py_spy_profiling(bash_command) |
Copilot uses AI. Check for mistakes.
* add aggregation script to cli * add py-spy dependency and some param hints * format with ruff * formatting * add flamegraph and cli for flamegraph * add more advanced error handling for cli application * remove perl linting * format cli.py * use typer.exit instead of return for error * add slurm option for using py-spy profiling * add missing arguments to cli.py for slurm generator * replace print with typer.echo in cli applications * add docstrings to new CLI functionality * update error handling in cli * change profiler output to raw * add licensing and comments from PR * aggregate functions * add relevant itwinai calls * add percentage to data aggregation * add aggregation wrt itwinai info as well * clean up code and add table library * add tabulate to project dependencies
Basic Integration with the py-spy Profiler
The goal of this PR is to add a basic level of integration with the py-spy profiler. This includes:
Basic Data Aggregation
The basic data aggregation consists of finding all the final function calls in each stack and sorting them by number of samples, which is a proxy of how much time was spent in them. This is printed as a table and the idea is that you expect most of the time to be spent in the backward and forward calls from PyTorch. If you see something else at the top, then this might suggest a bottleneck.
Example of table output (with and without aggregation):