mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
feat: add profiler util (#5601)
* 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
This commit is contained in:
67
invokeai/app/util/profiler.py
Normal file
67
invokeai/app/util/profiler.py
Normal file
@ -0,0 +1,67 @@
|
||||
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
|
Reference in New Issue
Block a user