From 4602efd59877e1dcce6cd08097ee735364bcc5e1 Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Wed, 31 Jan 2024 21:51:57 +1100 Subject: [PATCH] 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 --- .../app/services/config/config_default.py | 9 +++ .../invocation_processor_default.py | 19 +++++- invokeai/app/util/profiler.py | 67 +++++++++++++++++++ pyproject.toml | 2 +- scripts/generate_profile_graphs.sh | 27 ++++++++ tests/test_profiler.py | 53 +++++++++++++++ 6 files changed, 175 insertions(+), 2 deletions(-) create mode 100644 invokeai/app/util/profiler.py create mode 100755 scripts/generate_profile_graphs.sh create mode 100644 tests/test_profiler.py diff --git a/invokeai/app/services/config/config_default.py b/invokeai/app/services/config/config_default.py index 83a831524d..3af906bd04 100644 --- a/invokeai/app/services/config/config_default.py +++ b/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.""" diff --git a/invokeai/app/services/invocation_processor/invocation_processor_default.py b/invokeai/app/services/invocation_processor/invocation_processor_default.py index 09608dca2b..1cfb0c1822 100644 --- a/invokeai/app/services/invocation_processor/invocation_processor_default.py +++ b/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 diff --git a/invokeai/app/util/profiler.py b/invokeai/app/util/profiler.py new file mode 100644 index 0000000000..d1ce126b04 --- /dev/null +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 11c0e60dad..b96a6756f3 100644 --- a/pyproject.toml +++ b/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", diff --git a/scripts/generate_profile_graphs.sh b/scripts/generate_profile_graphs.sh new file mode 100755 index 0000000000..fa6cd9c103 --- /dev/null +++ b/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 +# is the path to the directory containing the .prof files +# 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 diff --git a/tests/test_profiler.py b/tests/test_profiler.py new file mode 100644 index 0000000000..ab152e1a2b --- /dev/null +++ b/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()