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
54 lines
1.7 KiB
Python
54 lines
1.7 KiB
Python
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()
|