mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
4602efd598
* feat(config): add profiling config settings - `profile_graphs` enables graph profiling with cProfile - `profiles_dir` sets the output for profiles * feat(nodes): add Profiler util Simple wrapper around cProfile. * feat(nodes): use Profiler in invocation processor * scripts: add generate_profile_graphs.sh script Helper to generate graphs for profiles. * pkg: add snakeviz and gprof2dot to dev deps These are useful for profiling. * tests: add tests for profiler util * fix(profiler): handle previous profile not stopped cleanly * feat(profiler): add profile_prefix config setting The prefix is used when writing profile output files. Useful to organise profiles into sessions. * tidy(profiler): add `_` to private API * feat(profiler): simplify API * feat(profiler): use child logger for profiler logs * chore(profiler): update docstrings * feat(profiler): stop() returns output path * chore(profiler): fix docstring * tests(profiler): update tests * chore: ruff
68 lines
2.1 KiB
Python
68 lines
2.1 KiB
Python
import cProfile
|
|
from logging import Logger
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
|
|
class Profiler:
|
|
"""
|
|
Simple wrapper around cProfile.
|
|
|
|
Usage
|
|
```
|
|
# Create a profiler
|
|
profiler = Profiler(logger, output_dir, "sql_query_perf")
|
|
# Start a new profile
|
|
profiler.start("my_profile")
|
|
# Do stuff
|
|
profiler.stop()
|
|
```
|
|
|
|
Visualize a profile as a flamegraph with [snakeviz](https://jiffyclub.github.io/snakeviz/)
|
|
```sh
|
|
snakeviz my_profile.prof
|
|
```
|
|
|
|
Visualize a profile as directed graph with [graphviz](https://graphviz.org/download/) & [gprof2dot](https://github.com/jrfonseca/gprof2dot)
|
|
```sh
|
|
gprof2dot -f pstats my_profile.prof | dot -Tpng -o my_profile.png
|
|
# SVG or PDF may be nicer - you can search for function names
|
|
gprof2dot -f pstats my_profile.prof | dot -Tsvg -o my_profile.svg
|
|
gprof2dot -f pstats my_profile.prof | dot -Tpdf -o my_profile.pdf
|
|
```
|
|
"""
|
|
|
|
def __init__(self, logger: Logger, output_dir: Path, prefix: Optional[str] = None) -> None:
|
|
self._logger = logger.getChild(f"profiler.{prefix}" if prefix else "profiler")
|
|
self._output_dir = output_dir
|
|
self._output_dir.mkdir(parents=True, exist_ok=True)
|
|
self._profiler: Optional[cProfile.Profile] = None
|
|
self._prefix = prefix
|
|
|
|
self.profile_id: Optional[str] = None
|
|
|
|
def start(self, profile_id: str) -> None:
|
|
if self._profiler:
|
|
self.stop()
|
|
|
|
self.profile_id = profile_id
|
|
|
|
self._profiler = cProfile.Profile()
|
|
self._profiler.enable()
|
|
self._logger.info(f"Started profiling {self.profile_id}.")
|
|
|
|
def stop(self) -> Path:
|
|
if not self._profiler:
|
|
raise RuntimeError("Profiler not initialized. Call start() first.")
|
|
self._profiler.disable()
|
|
|
|
filename = f"{self._prefix}_{self.profile_id}.prof" if self._prefix else f"{self.profile_id}.prof"
|
|
path = Path(self._output_dir, filename)
|
|
|
|
self._profiler.dump_stats(path)
|
|
self._logger.info(f"Stopped profiling, profile dumped to {path}.")
|
|
self._profiler = None
|
|
self.profile_id = None
|
|
|
|
return path
|