chore: import upstream snapshot with attribution

invokeai/app/services/invocation_stats/invocation_stats_base.py

@@ -0,0 +1,93 @@
+# Copyright 2023 Lincoln D. Stein <lincoln.stein@gmail.com>
+"""Utility to collect execution time and GPU usage stats on invocations in flight
+
+Usage:
+
+statistics = InvocationStatsService()
+with statistics.collect_stats(invocation, graph_execution_state.id):
+      ... execute graphs...
+statistics.log_stats()
+
+Typical output:
+[2023-08-02 18:03:04,507]::[InvokeAI]::INFO --> Graph stats: c7764585-9c68-4d9d-a199-55e8186790f3
+[2023-08-02 18:03:04,507]::[InvokeAI]::INFO --> Node                 Calls  Seconds  VRAM Used
+[2023-08-02 18:03:04,507]::[InvokeAI]::INFO --> main_model_loader        1   0.005s     0.01G
+[2023-08-02 18:03:04,508]::[InvokeAI]::INFO --> clip_skip                1   0.004s     0.01G
+[2023-08-02 18:03:04,508]::[InvokeAI]::INFO --> compel                   2   0.512s     0.26G
+[2023-08-02 18:03:04,508]::[InvokeAI]::INFO --> rand_int                 1   0.001s     0.01G
+[2023-08-02 18:03:04,508]::[InvokeAI]::INFO --> range_of_size            1   0.001s     0.01G
+[2023-08-02 18:03:04,508]::[InvokeAI]::INFO --> iterate                  1   0.001s     0.01G
+[2023-08-02 18:03:04,508]::[InvokeAI]::INFO --> metadata_accumulator     1   0.002s     0.01G
+[2023-08-02 18:03:04,508]::[InvokeAI]::INFO --> noise                    1   0.002s     0.01G
+[2023-08-02 18:03:04,508]::[InvokeAI]::INFO --> t2l                      1   3.541s     1.93G
+[2023-08-02 18:03:04,508]::[InvokeAI]::INFO --> l2i                      1   0.679s     0.58G
+[2023-08-02 18:03:04,508]::[InvokeAI]::INFO --> TOTAL GRAPH EXECUTION TIME:  4.749s
+[2023-08-02 18:03:04,508]::[InvokeAI]::INFO --> Current VRAM utilization 0.01G
+
+The abstract base class for this class is InvocationStatsServiceBase. An implementing class which
+writes to the system log is stored in InvocationServices.performance_statistics.
+"""
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import ContextManager
+
+from invokeai.app.invocations.baseinvocation import BaseInvocation
+from invokeai.app.services.invocation_stats.invocation_stats_common import InvocationStatsSummary
+
+
+class InvocationStatsServiceBase(ABC):
+    "Abstract base class for recording node memory/time performance statistics"
+
+    @abstractmethod
+    def __init__(self) -> None:
+        """
+        Initialize the InvocationStatsService and reset counters to zero
+        """
+
+    @abstractmethod
+    def collect_stats(
+        self,
+        invocation: BaseInvocation,
+        graph_execution_state_id: str,
+    ) -> ContextManager[None]:
+        """
+        Return a context object that will capture the statistics on the execution
+        of invocaation. Use with: to place around the part of the code that executes the invocation.
+        :param invocation: BaseInvocation object from the current graph.
+        :param graph_execution_state_id: The id of the current session.
+        """
+        pass
+
+    @abstractmethod
+    def reset_stats(self, graph_execution_state_id: str) -> None:
+        """Reset all stored statistics."""
+        pass
+
+    @abstractmethod
+    def log_stats(self, graph_execution_state_id: str) -> None:
+        """
+        Write out the accumulated statistics to the log or somewhere else.
+        :param graph_execution_state_id: The id of the session whose stats to log.
+        :raises GESStatsNotFoundError: if the graph isn't tracked in the stats.
+        """
+        pass
+
+    @abstractmethod
+    def get_stats(self, graph_execution_state_id: str) -> InvocationStatsSummary:
+        """
+        Gets the accumulated statistics for the indicated graph.
+        :param graph_execution_state_id: The id of the session whose stats to get.
+        :raises GESStatsNotFoundError: if the graph isn't tracked in the stats.
+        """
+        pass
+
+    @abstractmethod
+    def dump_stats(self, graph_execution_state_id: str, output_path: Path) -> None:
+        """
+        Write out the accumulated statistics to the indicated path as JSON.
+        :param graph_execution_state_id: The id of the session whose stats to dump.
+        :param output_path: The file to write the stats to.
+        :raises GESStatsNotFoundError: if the graph isn't tracked in the stats.
+        """
+        pass

invokeai/app/services/invocation_stats/invocation_stats_common.py

@@ -0,0 +1,183 @@
+from collections import defaultdict
+from dataclasses import asdict, dataclass
+from typing import Any, Optional
+
+
+class GESStatsNotFoundError(Exception):
+    """Raised when execution stats are not found for a given Graph Execution State."""
+
+
+@dataclass
+class NodeExecutionStatsSummary:
+    """The stats for a specific type of node."""
+
+    node_type: str
+    num_calls: int
+    time_used_seconds: float
+    delta_vram_gb: float
+
+
+@dataclass
+class ModelCacheStatsSummary:
+    """The stats for the model cache."""
+
+    high_water_mark_gb: float
+    cache_size_gb: float
+    total_usage_gb: float
+    cache_hits: int
+    cache_misses: int
+    models_cached: int
+    models_cleared: int
+
+
+@dataclass
+class GraphExecutionStatsSummary:
+    """The stats for the graph execution state."""
+
+    graph_execution_state_id: str
+    execution_time_seconds: float
+    # `wall_time_seconds`, `ram_usage_gb` and `ram_change_gb` are derived from the node execution stats.
+    # In some situations, there are no node stats, so these values are optional.
+    wall_time_seconds: Optional[float]
+    ram_usage_gb: Optional[float]
+    ram_change_gb: Optional[float]
+
+
+@dataclass
+class InvocationStatsSummary:
+    """
+    The accumulated stats for a graph execution.
+    Its `__str__` method returns a human-readable stats summary.
+    """
+
+    vram_usage_gb: Optional[float]
+    graph_stats: GraphExecutionStatsSummary
+    model_cache_stats: ModelCacheStatsSummary
+    node_stats: list[NodeExecutionStatsSummary]
+
+    def __str__(self) -> str:
+        _str = ""
+        _str = f"Graph stats: {self.graph_stats.graph_execution_state_id}\n"
+        _str += f"{'Node':>30} {'Calls':>7} {'Seconds':>9} {'VRAM Change':+>10}\n"
+
+        for summary in self.node_stats:
+            _str += f"{summary.node_type:>30} {summary.num_calls:>7} {summary.time_used_seconds:>8.3f}s {summary.delta_vram_gb:+10.3f}G\n"
+
+        _str += f"TOTAL GRAPH EXECUTION TIME: {self.graph_stats.execution_time_seconds:7.3f}s\n"
+
+        if self.graph_stats.wall_time_seconds is not None:
+            _str += f"TOTAL GRAPH WALL TIME: {self.graph_stats.wall_time_seconds:7.3f}s\n"
+
+        if self.graph_stats.ram_usage_gb is not None and self.graph_stats.ram_change_gb is not None:
+            _str += f"RAM used by InvokeAI process: {self.graph_stats.ram_usage_gb:4.2f}G ({self.graph_stats.ram_change_gb:+5.3f}G)\n"
+
+        _str += f"RAM used to load models: {self.model_cache_stats.total_usage_gb:4.2f}G\n"
+        if self.vram_usage_gb:
+            _str += f"VRAM in use: {self.vram_usage_gb:4.3f}G\n"
+        _str += "RAM cache statistics:\n"
+        _str += f"   Model cache hits: {self.model_cache_stats.cache_hits}\n"
+        _str += f"   Model cache misses: {self.model_cache_stats.cache_misses}\n"
+        _str += f"   Models cached: {self.model_cache_stats.models_cached}\n"
+        _str += f"   Models cleared from cache: {self.model_cache_stats.models_cleared}\n"
+        _str += f"   Cache high water mark: {self.model_cache_stats.high_water_mark_gb:4.2f}/{self.model_cache_stats.cache_size_gb:4.2f}G\n"
+
+        return _str
+
+    def as_dict(self) -> dict[str, Any]:
+        """Returns the stats as a dictionary."""
+        return asdict(self)
+
+
+@dataclass
+class NodeExecutionStats:
+    """Class for tracking execution stats of an invocation node."""
+
+    invocation_type: str
+
+    start_time: float  # Seconds since the epoch.
+    end_time: float  # Seconds since the epoch.
+
+    start_ram_gb: float  # GB
+    end_ram_gb: float  # GB
+
+    delta_vram_gb: float  # GB
+
+    def total_time(self) -> float:
+        return self.end_time - self.start_time
+
+
+class GraphExecutionStats:
+    """Class for tracking execution stats of a graph."""
+
+    def __init__(self):
+        self._node_stats_list: list[NodeExecutionStats] = []
+
+    def add_node_execution_stats(self, node_stats: NodeExecutionStats):
+        self._node_stats_list.append(node_stats)
+
+    def get_total_run_time(self) -> float:
+        """Get the total time spent executing nodes in the graph."""
+        total = 0.0
+        for node_stats in self._node_stats_list:
+            total += node_stats.total_time()
+        return total
+
+    def get_first_node_stats(self) -> NodeExecutionStats | None:
+        """Get the stats of the first node in the graph (by start_time)."""
+        first_node = None
+        for node_stats in self._node_stats_list:
+            if first_node is None or node_stats.start_time < first_node.start_time:
+                first_node = node_stats
+
+        assert first_node is not None
+        return first_node
+
+    def get_last_node_stats(self) -> NodeExecutionStats | None:
+        """Get the stats of the last node in the graph (by end_time)."""
+        last_node = None
+        for node_stats in self._node_stats_list:
+            if last_node is None or node_stats.end_time > last_node.end_time:
+                last_node = node_stats
+
+        return last_node
+
+    def get_graph_stats_summary(self, graph_execution_state_id: str) -> GraphExecutionStatsSummary:
+        """Get a summary of the graph stats."""
+        first_node = self.get_first_node_stats()
+        last_node = self.get_last_node_stats()
+
+        wall_time_seconds: Optional[float] = None
+        ram_usage_gb: Optional[float] = None
+        ram_change_gb: Optional[float] = None
+
+        if last_node and first_node:
+            wall_time_seconds = last_node.end_time - first_node.start_time
+            ram_usage_gb = last_node.end_ram_gb
+            ram_change_gb = last_node.end_ram_gb - first_node.start_ram_gb
+
+        return GraphExecutionStatsSummary(
+            graph_execution_state_id=graph_execution_state_id,
+            execution_time_seconds=self.get_total_run_time(),
+            wall_time_seconds=wall_time_seconds,
+            ram_usage_gb=ram_usage_gb,
+            ram_change_gb=ram_change_gb,
+        )
+
+    def get_node_stats_summaries(self) -> list[NodeExecutionStatsSummary]:
+        """Get a summary of the node stats."""
+        summaries: list[NodeExecutionStatsSummary] = []
+        node_stats_by_type: dict[str, list[NodeExecutionStats]] = defaultdict(list)
+
+        for node_stats in self._node_stats_list:
+            node_stats_by_type[node_stats.invocation_type].append(node_stats)
+
+        for node_type, node_type_stats_list in node_stats_by_type.items():
+            num_calls = len(node_type_stats_list)
+            time_used = sum([n.total_time() for n in node_type_stats_list])
+            delta_vram = max([n.delta_vram_gb for n in node_type_stats_list])
+            summary = NodeExecutionStatsSummary(
+                node_type=node_type, num_calls=num_calls, time_used_seconds=time_used, delta_vram_gb=delta_vram
+            )
+            summaries.append(summary)
+
+        return summaries

invokeai/app/services/invocation_stats/invocation_stats_default.py

@@ -0,0 +1,143 @@
+import json
+import time
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Generator
+
+import psutil
+import torch
+
+import invokeai.backend.util.logging as logger
+from invokeai.app.invocations.baseinvocation import BaseInvocation
+from invokeai.app.services.invocation_stats.invocation_stats_base import InvocationStatsServiceBase
+from invokeai.app.services.invocation_stats.invocation_stats_common import (
+    GESStatsNotFoundError,
+    GraphExecutionStats,
+    GraphExecutionStatsSummary,
+    InvocationStatsSummary,
+    ModelCacheStatsSummary,
+    NodeExecutionStats,
+    NodeExecutionStatsSummary,
+)
+from invokeai.app.services.invoker import Invoker
+from invokeai.backend.model_manager.load.model_cache.cache_stats import CacheStats
+
+# Size of 1GB in bytes.
+GB = 2**30
+
+
+class InvocationStatsService(InvocationStatsServiceBase):
+    """Accumulate performance information about a running graph. Collects time spent in each node,
+    as well as the maximum and current VRAM utilisation for CUDA systems"""
+
+    def __init__(self):
+        # Maps graph_execution_state_id to GraphExecutionStats.
+        self._stats: dict[str, GraphExecutionStats] = {}
+        # Maps graph_execution_state_id to model manager CacheStats.
+        self._cache_stats: dict[str, CacheStats] = {}
+
+    def start(self, invoker: Invoker) -> None:
+        self._invoker = invoker
+
+    @contextmanager
+    def collect_stats(self, invocation: BaseInvocation, graph_execution_state_id: str) -> Generator[None, None, None]:
+        # This is to handle case of the model manager not being initialized, which happens
+        # during some tests.
+        services = self._invoker.services
+        if not self._stats.get(graph_execution_state_id):
+            # First time we're seeing this graph_execution_state_id.
+            self._stats[graph_execution_state_id] = GraphExecutionStats()
+            self._cache_stats[graph_execution_state_id] = CacheStats()
+
+        # Record state before the invocation.
+        start_time = time.time()
+        start_ram = psutil.Process().memory_info().rss
+
+        # Remember current VRAM usage
+        vram_in_use = torch.cuda.memory_allocated() if torch.cuda.is_available() else 0.0
+
+        assert services.model_manager.load is not None
+        services.model_manager.load.ram_cache.stats = self._cache_stats[graph_execution_state_id]
+
+        try:
+            # Let the invocation run.
+            yield None
+        finally:
+            # Record delta VRAM
+            delta_vram_gb = ((torch.cuda.memory_allocated() - vram_in_use) / GB) if torch.cuda.is_available() else 0.0
+
+            node_stats = NodeExecutionStats(
+                invocation_type=invocation.get_type(),
+                start_time=start_time,
+                end_time=time.time(),
+                start_ram_gb=start_ram / GB,
+                end_ram_gb=psutil.Process().memory_info().rss / GB,
+                delta_vram_gb=delta_vram_gb,
+            )
+            self._stats[graph_execution_state_id].add_node_execution_stats(node_stats)
+
+    def reset_stats(self, graph_execution_state_id: str) -> None:
+        self._stats.pop(graph_execution_state_id, None)
+        self._cache_stats.pop(graph_execution_state_id, None)
+
+    def get_stats(self, graph_execution_state_id: str) -> InvocationStatsSummary:
+        graph_stats_summary = self._get_graph_summary(graph_execution_state_id)
+        node_stats_summaries = self._get_node_summaries(graph_execution_state_id)
+        model_cache_stats_summary = self._get_model_cache_summary(graph_execution_state_id)
+        # Note: We use memory_allocated() here (not memory_reserved()) because we want to show
+        # the current actively-used VRAM, not the total reserved memory including PyTorch's cache.
+        vram_usage_gb = torch.cuda.memory_allocated() / GB if torch.cuda.is_available() else None
+
+        return InvocationStatsSummary(
+            graph_stats=graph_stats_summary,
+            model_cache_stats=model_cache_stats_summary,
+            node_stats=node_stats_summaries,
+            vram_usage_gb=vram_usage_gb,
+        )
+
+    def log_stats(self, graph_execution_state_id: str) -> None:
+        stats = self.get_stats(graph_execution_state_id)
+        logger.info(str(stats))
+
+    def dump_stats(self, graph_execution_state_id: str, output_path: Path) -> None:
+        stats = self.get_stats(graph_execution_state_id)
+        with open(output_path, "w") as f:
+            f.write(json.dumps(stats.as_dict(), indent=2))
+
+    def _get_model_cache_summary(self, graph_execution_state_id: str) -> ModelCacheStatsSummary:
+        try:
+            cache_stats = self._cache_stats[graph_execution_state_id]
+        except KeyError as e:
+            raise GESStatsNotFoundError(
+                f"Attempted to get model cache statistics for unknown graph {graph_execution_state_id}: {e}."
+            ) from e
+
+        return ModelCacheStatsSummary(
+            cache_hits=cache_stats.hits,
+            cache_misses=cache_stats.misses,
+            high_water_mark_gb=cache_stats.high_watermark / GB,
+            cache_size_gb=cache_stats.cache_size / GB,
+            total_usage_gb=sum(list(cache_stats.loaded_model_sizes.values())) / GB,
+            models_cached=cache_stats.in_cache,
+            models_cleared=cache_stats.cleared,
+        )
+
+    def _get_graph_summary(self, graph_execution_state_id: str) -> GraphExecutionStatsSummary:
+        try:
+            graph_stats = self._stats[graph_execution_state_id]
+        except KeyError as e:
+            raise GESStatsNotFoundError(
+                f"Attempted to get graph statistics for unknown graph {graph_execution_state_id}: {e}."
+            ) from e
+
+        return graph_stats.get_graph_stats_summary(graph_execution_state_id)
+
+    def _get_node_summaries(self, graph_execution_state_id: str) -> list[NodeExecutionStatsSummary]:
+        try:
+            graph_stats = self._stats[graph_execution_state_id]
+        except KeyError as e:
+            raise GESStatsNotFoundError(
+                f"Attempted to get node statistics for unknown graph {graph_execution_state_id}: {e}."
+            ) from e
+
+        return graph_stats.get_node_stats_summaries()