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

invokeai/app/services/config/config_default.py

@@ -251,7 +251,11 @@ class InvokeAIAppConfig(InvokeAISettings):
     log_level           : Literal["debug", "info", "warning", "error", "critical"] = Field(default="info", description="Emit logging messages at this level or  higher", json_schema_extra=Categories.Logging)
     log_sql             : bool = Field(default=False, description="Log SQL queries", json_schema_extra=Categories.Logging)
 
+    # Development
     dev_reload          : bool = Field(default=False, description="Automatically reload when Python sources are changed.", json_schema_extra=Categories.Development)
+    profile_graphs      : bool = Field(default=False, description="Enable graph profiling", json_schema_extra=Categories.Development)
+    profile_prefix      : Optional[str] = Field(default=None, description="An optional prefix for profile output files.", json_schema_extra=Categories.Development)
+    profiles_dir        : Path = Field(default=Path('profiles'), description="Directory for graph profiles", json_schema_extra=Categories.Development)
 
     version             : bool = Field(default=False, description="Show InvokeAI version and exit", json_schema_extra=Categories.Other)
 
@@ -449,6 +453,11 @@ class InvokeAIAppConfig(InvokeAISettings):
         disabled_in_config = not self.xformers_enabled
         return disabled_in_config and self.attention_type != "xformers"
 
+    @property
+    def profiles_path(self) -> Path:
+        """Path to the graph profiles directory."""
+        return self._resolve(self.profiles_dir)
+
     @staticmethod
     def find_root() -> Path:
         """Choose the runtime root directory when not specified on command line or init file."""

invokeai/app/services/invocation_processor/invocation_processor_default.py

@@ -6,6 +6,7 @@ from typing import Optional
 import invokeai.backend.util.logging as logger
 from invokeai.app.invocations.baseinvocation import InvocationContext
 from invokeai.app.services.invocation_queue.invocation_queue_common import InvocationQueueItem
+from invokeai.app.util.profiler import Profiler
 
 from ..invoker import Invoker
 from .invocation_processor_base import InvocationProcessorABC
@@ -18,7 +19,7 @@ class DefaultInvocationProcessor(InvocationProcessorABC):
     __invoker: Invoker
     __threadLimit: BoundedSemaphore
 
-    def start(self, invoker) -> None:
+    def start(self, invoker: Invoker) -> None:
         # if we do want multithreading at some point, we could make this configurable
         self.__threadLimit = BoundedSemaphore(1)
         self.__invoker = invoker
@@ -39,6 +40,16 @@ class DefaultInvocationProcessor(InvocationProcessorABC):
             self.__threadLimit.acquire()
             queue_item: Optional[InvocationQueueItem] = None
 
+            profiler = (
+                Profiler(
+                    logger=self.__invoker.services.logger,
+                    output_dir=self.__invoker.services.configuration.profiles_path,
+                    prefix=self.__invoker.services.configuration.profile_prefix,
+                )
+                if self.__invoker.services.configuration.profile_graphs
+                else None
+            )
+
             while not stop_event.is_set():
                 try:
                     queue_item = self.__invoker.services.queue.get()
@@ -49,6 +60,10 @@ class DefaultInvocationProcessor(InvocationProcessorABC):
                     # do not hammer the queue
                     time.sleep(0.5)
                     continue
+
+                if profiler and profiler.profile_id != queue_item.graph_execution_state_id:
+                    profiler.start(profile_id=queue_item.graph_execution_state_id)
+
                 try:
                     graph_execution_state = self.__invoker.services.graph_execution_manager.get(
                         queue_item.graph_execution_state_id
@@ -201,6 +216,8 @@ class DefaultInvocationProcessor(InvocationProcessorABC):
                         queue_id=queue_item.session_queue_id,
                         graph_execution_state_id=graph_execution_state.id,
                     )
+                    if profiler:
+                        profiler.stop()
 
         except KeyboardInterrupt:
             pass  # Log something? KeyboardInterrupt is probably not going to be seen by the processor

invokeai/app/util/profiler.py

@@ -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

pyproject.toml

@@ -109,7 +109,7 @@ dependencies = [
   "mkdocs-git-revision-date-localized-plugin",
   "mkdocs-redirects==1.2.0",
 ]
-"dev" = ["jurigged", "pudb"]
+"dev" = ["jurigged", "pudb", "snakeviz", "gprof2dot"]
 "test" = [
   "ruff==0.1.11",
   "ruff-lsp",

scripts/generate_profile_graphs.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+
+# Accepts a path to a directory containing .prof files and generates a graphs
+# for each of them. The default output format is pdf, but can be changed by
+# providing a second argument.
+
+# Usage: ./generate_profile_graphs.sh <path_to_profiles> <type>
+# <path_to_profiles> is the path to the directory containing the .prof files
+# <type> is the type of graph to generate. Defaults to 'pdf' if not provided.
+# Valid types are: 'svg', 'png' and 'pdf'.
+
+# Requires:
+# - graphviz: https://graphviz.org/download/
+# - gprof2dot: https://github.com/jrfonseca/gprof2dot
+
+if [ -z "$1" ]; then
+  echo "Missing path to profiles directory"
+  exit 1
+fi
+
+type=${2:-pdf}
+
+for file in $1/*.prof; do
+  base_name=$(basename "$file" .prof)
+  gprof2dot -f pstats "$file" | dot -T$type -Glabel="Session ID ${base_name}" -Glabelloc="t" -o "$1/$base_name.$type"
+  echo "Generated $1/$base_name.$type"
+done

tests/test_profiler.py

@@ -0,0 +1,53 @@
+import re
+from logging import Logger
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import pytest
+
+from invokeai.app.util.profiler import Profiler
+
+
+def test_profiler_starts():
+    with TemporaryDirectory() as tempdir:
+        profiler = Profiler(logger=Logger("test_profiler"), output_dir=Path(tempdir))
+        assert not profiler._profiler
+        assert not profiler.profile_id
+        profiler.start("test")
+        assert profiler._profiler
+        assert profiler.profile_id == "test"
+        profiler.stop()
+        assert not profiler._profiler
+        assert not profiler.profile_id
+        profiler.start("test2")
+        assert profiler._profiler
+        assert profiler.profile_id == "test2"
+        profiler.stop()
+
+
+def test_profiler_profiles():
+    with TemporaryDirectory() as tempdir:
+        profiler = Profiler(logger=Logger("test_profiler"), output_dir=Path(tempdir))
+        profiler.start("test")
+        for _ in range(1000000):
+            pass
+        profiler.stop()
+        assert (Path(tempdir) / "test.prof").exists()
+
+
+def test_profiler_profiles_with_prefix():
+    with TemporaryDirectory() as tempdir:
+        profiler = Profiler(logger=Logger("test_profiler"), output_dir=Path(tempdir), prefix="prefix")
+        profiler.start("test")
+        for _ in range(1000000):
+            pass
+        profiler.stop()
+        assert (Path(tempdir) / "prefix_test.prof").exists()
+
+
+def test_profile_fails_if_not_set_up():
+    with TemporaryDirectory() as tempdir:
+        profiler = Profiler(logger=Logger("test_profiler"), output_dir=Path(tempdir))
+        match = re.escape("Profiler not initialized. Call start() first.")
+        with pytest.raises(RuntimeError, match=match):
+            profiler.stop()