InvokeAI/invokeai/app/util/profiler.py

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