gpu-energy-profiler 0.1.0__tar.gz

gpu_energy_profiler-0.1.0/PKG-INFO

@@ -0,0 +1,14 @@
+Metadata-Version: 2.4
+Name: gpu-energy-profiler
+Version: 0.1.0
+Summary: 
+Author: Vincent Eichhorn
+Author-email: v.eichhorn@posteo.de
+Requires-Python: >=3.14
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Programming Language :: Python :: 3.15
+Requires-Dist: nvidia-ml-py (>=12.0,<13.0)
+Requires-Dist: pandas (>=3.0.6,<4.0.0)
+Requires-Dist: plotly (>=7.1.0,<8.0.0)
+Requires-Dist: torch (>=2.14.1,<3.0.0)

gpu_energy_profiler-0.1.0/pyproject.toml

@@ -0,0 +1,25 @@
+[project]
+name = "gpu-energy-profiler"
+version = "0.1.0"
+description = ""
+authors = [
+    {name = "Vincent Eichhorn",email = "v.eichhorn@posteo.de"}
+]
+requires-python = ">=3.14"
+dependencies = [
+    "pandas (>=3.0.6,<4.0.0)",
+    "torch (>=2.14.1,<3.0.0)",
+    "plotly (>=7.1.0,<8.0.0)",
+    "nvidia-ml-py (>=12.0,<13.0)"
+]
+
+[tool.poetry.group.dev.dependencies]
+pytest = "^8.0"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+
+[build-system]
+requires = ["poetry-core>=2.0.0,<3.0.0"]
+build-backend = "poetry.core.masonry.api"

gpu_energy_profiler-0.1.0/src/gpu_energy_profiler/__init__.py

@@ -0,0 +1,7 @@
+"""Tools for profiling PyTorch workloads and GPU power usage."""
+
+from .base import Profiler
+from .nvidia_profiler import NvidiaProfiler
+from .torch_profiler import TorchProfiler
+
+__all__ = ["NvidiaProfiler", "Profiler", "TorchProfiler"]

gpu_energy_profiler-0.1.0/src/gpu_energy_profiler/base.py

@@ -0,0 +1,35 @@
+"""Shared functionality for profiler implementations."""
+
+from contextlib import contextmanager
+from datetime import datetime
+from typing import Generator
+
+
+class Profiler:
+    """Record named timestamps while a profiling session is running."""
+
+    def __init__(self) -> None:
+        """Initialize an empty step timeline."""
+        self.record_steps: list[tuple[datetime, str]] = []
+        self.record_step("__init__")
+
+    def record_step(self, name: str) -> None:
+        """Record a named timestamp.
+
+        Args:
+            name: Label for the recorded step.
+        """
+        self.record_steps.append((datetime.now(), name))
+
+    @contextmanager
+    def record_context(self, name: str) -> Generator[None, None, None]:
+        """Record a context's start and end markers.
+
+        Args:
+            name: Label recorded when entering the context.
+        """
+        try:
+            self.record_step(name)
+            yield
+        finally:
+            self.record_step("__other__")

gpu_energy_profiler-0.1.0/src/gpu_energy_profiler/multiprocessing_util.py

@@ -0,0 +1,236 @@
+"""Result handlers and helpers for profiler worker processes."""
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from itertools import zip_longest
+from multiprocessing import Manager, Process, Queue
+from pathlib import Path
+from typing import Any
+import warnings
+
+Result = tuple[Any, ...]
+ResultOrSentinel = Result | None
+
+
+class ResultHandler(ABC):
+    """Store worker results behind a queue-like interface."""
+
+    def __init__(self) -> None:
+        """Initialize an unconfigured result handler."""
+        self.column_names: tuple[str, ...] = ()
+        self.dtypes: tuple[type, ...] = ()
+
+    def __enter__(self) -> "ResultHandler":
+        """Enter the handler context.
+
+        Returns:
+            This result handler.
+        """
+        return self
+
+    def __exit__(self, *context_info: Any) -> None:
+        """Exit the handler context.
+
+        Args:
+            *context_info: Context-manager exception information.
+        """
+        del context_info
+
+    def set_columns(
+        self,
+        names: tuple[str, ...],
+        dtypes: tuple[type, ...] | None = None,
+    ) -> None:
+        """Set column names and conversion types used by file-backed results.
+
+        Args:
+            names: Column names in the serialized result rows.
+            dtypes: Conversion types for serialized values. Defaults to ``str``.
+
+        Raises:
+            ValueError: If the number of names and types differs.
+        """
+        dtypes = dtypes or (str,) * len(names)
+        if len(names) != len(dtypes):
+            raise ValueError("names and dtypes must have the same length")
+        self.column_names = names
+        self.dtypes = dtypes
+
+    @abstractmethod
+    def put(self, data: ResultOrSentinel) -> None:
+        """Store one result or ``None`` as the end-of-results sentinel.
+
+        Args:
+            data: Result tuple or ``None`` to mark the end of the stream.
+        """
+
+    @abstractmethod
+    def latest(self) -> ResultOrSentinel:
+        """Return the latest result, or ``None`` when no result is available."""
+
+    @abstractmethod
+    def all(self) -> list[Result]:
+        """Return all results up to the end-of-results sentinel."""
+
+    def get(self) -> ResultOrSentinel:
+        """Return the next result using the legacy method name.
+
+        Returns:
+            The next result or ``None``.
+        """
+        return self.latest()
+
+    def get_all(self) -> list[Result]:
+        """Return all results using the legacy method name.
+
+        Returns:
+            All results up to the end-of-results sentinel.
+        """
+        return self.all()
+
+
+class MPQueueResultHandler(ResultHandler):
+    """Store results in a multiprocessing queue."""
+
+    def __init__(self) -> None:
+        """Initialize an in-memory queue result handler."""
+        super().__init__()
+        self.queue: Queue[ResultOrSentinel] = Queue()
+
+    def put(self, data: ResultOrSentinel) -> None:
+        """Put a result or sentinel into the queue.
+
+        Args:
+            data: Result tuple or ``None`` as the stream sentinel.
+        """
+        self.queue.put(data)
+
+    def latest(self) -> ResultOrSentinel:
+        """Block until and return the next queued result.
+
+        Returns:
+            The next queued result or ``None``.
+        """
+        return self.queue.get()
+
+    def all(self) -> list[Result]:
+        """Read results until the sentinel is encountered.
+
+        Returns:
+            All queued result tuples.
+        """
+        return [result for result in iter(self.queue.get, None)]
+
+
+class FileCacheResultHandler(ResultHandler):
+    """Append typed result rows to a CSV cache file."""
+
+    def __init__(self, file_path: str, force: bool = False) -> None:
+        """Initialize a CSV-backed result handler.
+
+        Args:
+            file_path: Path to the CSV cache.
+            force: Whether to replace an existing cache.
+        """
+        super().__init__()
+        self.file_path = Path(file_path)
+        self.file_path.parent.mkdir(parents=True, exist_ok=True)
+        if force and self.file_path.is_file():
+            self.file_path.unlink()
+        if self.file_path.is_file() and self.file_path.stat().st_size > 0:
+            warnings.warn(
+                f"File {self.file_path} already exists and is not empty",
+                UserWarning,
+                stacklevel=2,
+            )
+
+    def __enter__(self) -> "FileCacheResultHandler":
+        """Enter the file-cache context.
+
+        Returns:
+            This file-cache result handler.
+        """
+        return self
+
+    def set_columns(
+        self,
+        names: tuple[str, ...],
+        dtypes: tuple[type, ...] | None = None,
+    ) -> None:
+        """Write the CSV header when the cache is empty."""
+        super().set_columns(names, dtypes)
+        if self.file_path.exists() and self.file_path.stat().st_size > 0:
+            return
+        self.file_path.write_text(",".join(names) + "\n", encoding="utf-8")
+
+    def put(self, data: ResultOrSentinel) -> None:
+        """Append a tuple result and ignore the worker sentinel.
+
+        Args:
+            data: Tuple to append, or ``None`` to end a worker stream.
+        """
+        if not isinstance(data, tuple):
+            return
+        with self.file_path.open("a", encoding="utf-8") as file:
+            file.write(",".join(str(value) for value in data) + "\n")
+
+    def _parse_line(self, line: str) -> Result:
+        """Convert one serialized CSV row to a typed result tuple.
+
+        Args:
+            line: Comma-separated result row.
+
+        Returns:
+            The converted result tuple.
+        """
+        values = zip_longest(
+            line.strip().split(","),
+            self.dtypes,
+            fillvalue=str,
+        )
+        return tuple(converter(value) for value, converter in values)
+
+    def latest(self) -> ResultOrSentinel:
+        """Return the last cached row, or ``None`` for an empty cache."""
+        rows = self._rows()
+        return self._parse_line(rows[-1]) if rows else None
+
+    def all(self) -> list[Result]:
+        """Return every cached row in file order.
+
+        Returns:
+            Typed result tuples in cache order.
+        """
+        return [self._parse_line(row) for row in self._rows()]
+
+    def _rows(self) -> list[str]:
+        """Read data rows from the cache file.
+
+        Returns:
+            Raw CSV data rows without the optional header.
+        """
+        with self.file_path.open("r", encoding="utf-8") as file:
+            rows = file.readlines()
+        return rows[1:] if self.column_names else rows
+
+
+def start_separate_process(target: Callable[..., Any], arguments: list[Any]) -> Any:
+    """Run ``target`` in a spawned process and return its first result.
+
+    Args:
+        target: Worker function receiving a managed queue first.
+        arguments: Additional arguments passed to the worker.
+
+    Returns:
+        The first result placed in the managed queue.
+    """
+    with Manager() as manager:
+        queue = manager.Queue()
+        process = Process(target=target, args=[queue, *arguments])
+        process.start()
+        process.join()
+        return queue.get()
+
+
+# Compatibility alias for the original misspelled function name.
+start_seprate_process = start_separate_process

gpu_energy_profiler-0.1.0/src/gpu_energy_profiler/nvidia_profiler.py

@@ -0,0 +1,458 @@
+"""GPU power and memory profiling through NVML or ``nvidia-smi``."""
+
+from datetime import datetime
+import subprocess
+from multiprocessing import Array, Event, Process, Value
+import time
+from typing import Any, Literal
+
+import pandas as pd
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
+
+from .base import Profiler
+from .multiprocessing_util import FileCacheResultHandler, MPQueueResultHandler, ResultHandler
+from .plotting_util import sample_colors
+
+
+def _load_pynvml() -> Any:
+    """Import and return the optional ``pynvml`` module.
+
+    Returns:
+        The imported NVML Python bindings.
+
+    Raises:
+        RuntimeError: If the optional dependency is not installed.
+    """
+    try:
+        import pynvml
+    except ImportError as error:
+        raise RuntimeError(
+            "The pynvml backend requires the optional 'nvidia-ml-py' dependency."
+        ) from error
+    return pynvml
+
+
+class NvidiaProfiler(Profiler):
+    """Sample GPU power and memory usage with a selectable backend."""
+
+    _COLUMNS = ("gpu_id", "timestamp", "power", "memory", "record_step")
+    _BACKENDS = ("nvidia-smi", "pynvml")
+
+    def __init__(
+        self,
+        interval: int = 1,
+        cache_file: str | None = None,
+        force_cache: bool = False,
+        backend: Literal["nvidia-smi", "pynvml"] = "nvidia-smi",
+    ) -> None:
+        """Initialize an NVIDIA power and memory profiler.
+
+        Args:
+            interval: Sampling interval in milliseconds.
+            cache_file: Optional CSV path for persisted samples.
+            force_cache: Whether to replace an existing cache file.
+            backend: Sampling implementation, either ``"nvidia-smi"`` or
+                ``"pynvml"``.
+
+        Raises:
+            ValueError: If ``backend`` is not supported.
+        """
+        if backend not in self._BACKENDS:
+            raise ValueError(f"Unknown backend {backend!r}; choose from {self._BACKENDS}.")
+        self.current_record_step: Any = Array("c", 1000)
+        self.interval = interval
+        self.backend = backend
+        self.data: list[tuple[Any, ...]] = []
+        self.should_profiling_run = Value("i", 1)
+        self.profiling_started = Event()
+        self.profiling_stopped = Event()
+        self.result_handler: ResultHandler = (
+            MPQueueResultHandler()
+            if cache_file is None
+            else FileCacheResultHandler(cache_file, force_cache)
+        )
+        self.result_handler.set_columns(self._COLUMNS, (int, str, float, float, str))
+        self.process = Process(
+            target=self._profiling_process,
+            args=(
+                self.should_profiling_run,
+                self.profiling_started,
+                self.profiling_stopped,
+                self.result_handler,
+                self.current_record_step,
+                self.interval,
+                self.backend,
+            ),
+        )
+        super().__init__()
+
+    def record_step(self, name: str) -> None:
+        """Record a step and share its name with the sampling process.
+
+        Args:
+            name: Label for the recorded step.
+        """
+        super().record_step(name)
+        self.current_record_step.value = name.encode("utf-8")
+
+    @staticmethod
+    def _parse_nvidia_smi_row(line: str, current_record_step: Any) -> tuple[Any, ...]:
+        """Parse one CSV row emitted by ``nvidia-smi``.
+
+        Args:
+            line: Raw CSV sample line.
+            current_record_step: Shared character array containing the step name.
+
+        Returns:
+            A typed sample tuple containing GPU, timestamp, power, memory, and step.
+        """
+        values = line.strip().split(", ")
+        return (
+            int(values[0]),
+            datetime.strptime(values[1], "%Y/%m/%d %H:%M:%S.%f"),
+            float(values[2].split(" ")[0]),
+            float(values[3].split(" ")[0]),
+            current_record_step.value.decode("utf-8"),
+        )
+
+    @staticmethod
+    def _profiling_process(
+        should_run: Any,
+        started: Any,
+        stopped: Any,
+        result_handler: ResultHandler,
+        current_record_step: Any,
+        interval: int,
+        backend: Literal["nvidia-smi", "pynvml"],
+    ) -> None:
+        """Run the selected sampling backend in a worker process.
+
+        Args:
+            should_run: Shared flag controlling the worker lifecycle.
+            started: Event set after the sampler starts.
+            stopped: Event set after sampling ends.
+            result_handler: Destination for sampled rows.
+            current_record_step: Shared current step name.
+            interval: Sampling interval in milliseconds.
+            backend: Backend to run.
+        """
+        if backend == "nvidia-smi":
+            NvidiaProfiler._nvidia_smi_profiling_process(
+                should_run,
+                started,
+                stopped,
+                result_handler,
+                current_record_step,
+                interval,
+            )
+            return
+        NvidiaProfiler._pynvml_profiling_process(
+            should_run,
+            started,
+            stopped,
+            result_handler,
+            current_record_step,
+            interval,
+        )
+
+    @staticmethod
+    def _nvidia_smi_profiling_process(
+        should_run: Any,
+        started: Any,
+        stopped: Any,
+        result_handler: ResultHandler,
+        current_record_step: Any,
+        interval: int,
+    ) -> None:
+        """Read samples until ``should_run`` is cleared.
+
+        Args:
+            should_run: Shared flag controlling the worker lifecycle.
+            started: Event set after the sampler starts.
+            stopped: Event set after sampling ends.
+            result_handler: Destination for sampled rows.
+            current_record_step: Shared current step name.
+            interval: Sampling interval in milliseconds.
+        """
+        command = (
+            "nvidia-smi --query-gpu=index,timestamp,power.draw,memory.used "
+            f"--format=csv -lms {interval}"
+        )
+        with (
+            subprocess.Popen(
+                command,
+                shell=True,
+                text=True,
+                stdout=subprocess.PIPE,
+            ) as nvidia_smi_process,
+            result_handler as result,
+        ):
+            assert nvidia_smi_process.stdout is not None
+            with nvidia_smi_process.stdout as output:
+                output.readline()
+                if should_run.value:
+                    started.set()
+                while should_run.value:
+                    result.put(
+                        NvidiaProfiler._parse_nvidia_smi_row(output.readline(), current_record_step)
+                    )
+        result.put(None)  # type: ignore[arg-type]
+        stopped.set()
+
+    @staticmethod
+    def _pynvml_profiling_process(
+        should_run: Any,
+        started: Any,
+        stopped: Any,
+        result_handler: ResultHandler,
+        current_record_step: Any,
+        interval: int,
+    ) -> None:
+        """Read power and memory samples through the optional ``pynvml`` package.
+
+        Args:
+            should_run: Shared flag controlling the worker lifecycle.
+            started: Event set after the sampler starts.
+            stopped: Event set after sampling ends.
+            result_handler: Destination for sampled rows.
+            current_record_step: Shared current step name.
+            interval: Sampling interval in milliseconds.
+        """
+        pynvml = _load_pynvml()
+        pynvml.nvmlInit()
+        try:
+            with result_handler as result:
+                if should_run.value:
+                    started.set()
+                while should_run.value:
+                    for gpu_id in range(pynvml.nvmlDeviceGetCount()):
+                        handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id)
+                        memory = pynvml.nvmlDeviceGetMemoryInfo(handle)
+                        result.put(
+                            (
+                                gpu_id,
+                                datetime.now(),
+                                pynvml.nvmlDeviceGetPowerUsage(handle) / 1000.0,
+                                memory.used / (1024 * 1024),
+                                current_record_step.value.decode("utf-8"),
+                            )
+                        )
+                    time.sleep(interval / 1000.0)
+                result.put(None)  # type: ignore[arg-type]
+        finally:
+            pynvml.nvmlShutdown()
+            stopped.set()
+
+    def __enter__(self) -> "NvidiaProfiler":
+        """Start sampling and return this profiler.
+
+        Returns:
+            This active NVIDIA profiler.
+
+        Raises:
+            RuntimeError: If the selected backend is unavailable.
+        """
+        if self.backend == "nvidia-smi":
+            status, output = subprocess.getstatusoutput(
+                "nvidia-smi --query-gpu=index --format=csv,noheader"
+            )
+            if status != 0 or not output.strip():
+                raise RuntimeError(
+                    "nvidia-smi is unavailable or could not query an NVIDIA GPU. "
+                    f"Command output: {output.strip() or '<empty>'}"
+                )
+        else:
+            pynvml = _load_pynvml()
+            try:
+                pynvml.nvmlInit()
+                if pynvml.nvmlDeviceGetCount() == 0:
+                    raise RuntimeError("pynvml found no NVIDIA GPU devices.")
+            except pynvml.NVMLError as error:
+                raise RuntimeError(
+                    "pynvml could not initialize or query an NVIDIA GPU."
+                ) from error
+            finally:
+                pynvml.nvmlShutdown()
+        self.process.start()
+        if not self.profiling_started.wait(timeout=5):
+            if self.process.is_alive():
+                self.process.terminate()
+                self.process.join(timeout=5)
+            raise RuntimeError(
+                f"The {self.backend} profiling worker did not start. "
+                "Check the NVIDIA driver and backend installation."
+            )
+        return self
+
+    def __exit__(self, *context_info: Any) -> None:
+        """Stop sampling and collect all recorded rows.
+
+        Args:
+            *context_info: Context-manager exception information.
+        """
+        del context_info
+        self.should_profiling_run.value = 0
+        self.profiling_stopped.wait()
+        self.data = self.result_handler.all()
+        self.process.join()
+        self.process.terminate()
+
+    @staticmethod
+    def from_cache(cache_file: str) -> "NvidiaProfiler":
+        """Create a profiler populated from an existing cache file.
+
+        Args:
+            cache_file: Path to a previously written sample cache.
+
+        Returns:
+            A profiler containing the cached samples.
+        """
+        profiler = NvidiaProfiler(cache_file=cache_file)
+        profiler.data = profiler.result_handler.all()
+        return profiler
+
+    def to_pandas(self) -> pd.DataFrame:
+        """Return samples with parsed timestamps and standard column names.
+
+        Returns:
+            A dataframe containing GPU sample data.
+        """
+        df = pd.DataFrame(self.data, columns=self._COLUMNS)
+        df["timestamp"] = pd.to_datetime(
+            df["timestamp"], format="%Y-%m-%d %H:%M:%S.%f", errors="coerce"
+        )
+        return df
+
+    def profiled_gpus(self) -> list[int]:
+        """Return the IDs of GPUs present in the samples.
+
+        Returns:
+            Unique GPU IDs found in the samples.
+        """
+        return [int(gpu_id) for gpu_id in self.to_pandas()["gpu_id"].unique()]
+
+    def total_energy(
+        self,
+        gpu_ids: list[int] | None = None,
+        record_steps: list[str] | None = None,
+        return_data: bool = False,
+    ) -> float | list[float]:
+        """Return sampled energy in watt-seconds.
+
+        Args:
+            gpu_ids: GPUs to include. Defaults to the first sampled GPU.
+            record_steps: Step labels to include. Defaults to all steps.
+            return_data: Return one energy value per record-step group.
+
+        Returns:
+            Total energy, or grouped energy values when ``return_data`` is true.
+        """
+        if not self.data:
+            return 0.0
+        df = self.to_pandas()
+        df["record_step_id"] = df["record_step"].ne(df["record_step"].shift()).cumsum()
+        gpu_ids = gpu_ids or [df["gpu_id"].unique()[0]]
+        df = df[df["gpu_id"].isin(gpu_ids)].copy()
+        df["time_interval"] = df["timestamp"].diff().dt.total_seconds().fillna(0)
+        df["energy_interval"] = df["power"] * df["time_interval"]
+        record_steps = record_steps or list(df["record_step"].unique())
+        df = df[df["record_step"].isin(record_steps)]
+        if return_data:
+            return list(df.groupby("record_step_id")["energy_interval"].sum())
+        return df["energy_interval"].sum()
+
+    def total_time(self) -> float:
+        """Return the time between the first and last sample in seconds.
+
+        Returns:
+            Elapsed profiling time in seconds.
+        """
+        if not self.data:
+            return 0.0
+        df = self.to_pandas()
+        return (df["timestamp"].max() - df["timestamp"].min()).total_seconds()
+
+    def avg_memory_usage(self, gpu_id: int | None = None) -> float:
+        """Return average memory usage in MiB for a GPU.
+
+        Args:
+            gpu_id: GPU to measure. Defaults to the first sampled GPU.
+
+        Returns:
+            Average memory usage in MiB.
+        """
+        if not self.data:
+            return 0.0
+        df = self.to_pandas()
+        gpu_id = gpu_id or df["gpu_id"].unique()[0]
+        return df.loc[df["gpu_id"] == gpu_id, "memory"].mean()
+
+    def time_series_plot(self) -> go.Figure | None:
+        """Return a Plotly chart showing power, memory, and recorded steps.
+
+        Returns:
+            A Plotly figure, or ``None`` when no samples are available.
+        """
+        if not self.data:
+            return None
+        df = self.to_pandas()
+        figure = make_subplots(specs=[[{"secondary_y": True}]])
+        gpu_ids = self.profiled_gpus()
+        gpu_colors = sample_colors("Rainbow", len(gpu_ids))
+        for index, gpu_id in enumerate(gpu_ids):
+            gpu_df = df[df["gpu_id"] == gpu_id]
+            for metric, unit in (("power", "W"), ("memory", "MiB")):
+                figure.add_trace(
+                    go.Scatter(
+                        x=gpu_df["timestamp"],
+                        y=gpu_df[metric],
+                        name=f"{metric.capitalize()} ({unit})",
+                        mode="lines+markers",
+                        legendgroup=str(gpu_id),
+                        legendgrouptitle_text=f"GPU #{gpu_id}",
+                        line=dict(
+                            color=gpu_colors[index],
+                            width=4,
+                            dash="dot" if metric == "memory" else "solid",
+                        ),
+                    ),
+                    secondary_y=metric == "memory",
+                )
+        _add_record_step_regions(figure, df, self.record_steps)
+        figure.update_layout(
+            title="GPU Memory and Power Usage", legend=dict(groupclick="toggleitem")
+        )
+        figure.update_xaxes(title_text="Time")
+        figure.update_yaxes(title_text="Power (W)", secondary_y=False)
+        figure.update_yaxes(title_text="Memory (MiB)", secondary_y=True)
+        return figure
+
+def _add_record_step_regions(
+    figure: go.Figure,
+    samples: pd.DataFrame,
+    record_steps: list[tuple[datetime, str]],
+) -> None:
+    """Add shaded regions for recorded steps to a Plotly figure.
+
+    Args:
+        figure: Figure to modify.
+        samples: Sample dataframe containing timestamps.
+        record_steps: Timestamp and label pairs defining the regions.
+    """
+    max_timestamp: datetime = samples["timestamp"].max()
+    steps = record_steps + [(max_timestamp, ".")]
+    step_names = list(dict.fromkeys(name for _, name in steps))
+    colors = dict(zip(step_names, sample_colors("viridis", len(step_names))))
+    previous_timestamp, previous_name = steps[0]
+    for timestamp, name in steps[1:]:
+        figure.add_vrect(
+            x0=previous_timestamp,
+            x1=timestamp,
+            annotation_text=previous_name,
+            annotation_position="top left",
+            line_width=0,
+            opacity=0.25,
+            fillcolor=colors[previous_name],
+        )
+        previous_timestamp, previous_name = timestamp, name

gpu_energy_profiler-0.1.0/src/gpu_energy_profiler/plotting_util.py

@@ -0,0 +1,22 @@
+"""Shared plotting helpers for profiler visualizations."""
+
+import plotly.express as px
+
+
+def sample_colors(scale: str, item_count: int) -> list[str]:
+    """Return evenly spaced colors for a collection.
+
+    Args:
+        scale: Plotly color scale name.
+        item_count: Number of colors required.
+
+    Returns:
+        At least two colors sampled from the requested scale.
+    """
+    count = max(item_count, 2)
+    return [
+        str(color)
+        for color in px.colors.sample_colorscale(
+            scale, [index / (count - 1) for index in range(count)]
+        )
+    ]

gpu_energy_profiler-0.1.0/src/gpu_energy_profiler/torch_profiler.py

@@ -0,0 +1,206 @@
+"""PyTorch profiler integration and event analysis."""
+
+from typing import Any
+
+import pandas as pd
+from torch.autograd.profiler_util import EventList, FunctionEvent
+from torch.profiler import ProfilerActivity, profile
+
+from .base import Profiler
+
+
+class TorchProfiler(profile, Profiler):
+    """Collect and analyse CPU and CUDA events from PyTorch."""
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        """Initialize a PyTorch profiler with useful profiling defaults.
+
+        Args:
+            *args: Positional arguments forwarded to ``torch.profiler.profile``.
+            **kwargs: Keyword arguments overriding the profiler defaults.
+        """
+        defaults = {
+            "with_flops": True,
+            "profile_memory": True,
+            "activities": [ProfilerActivity.CPU, ProfilerActivity.CUDA],
+        }
+        profile.__init__(self, *args, **{**defaults, **kwargs})
+        Profiler.__init__(self)
+        self.numeric_columns = [
+            "flops",
+            "count",
+            "self_device_time_total",
+            "self_cpu_time_total",
+            "device_time_total",
+            "cpu_time_total",
+            "self_device_memory_usage",
+            "self_cpu_memory_usage",
+            "device_memory_usage",
+            "cpu_memory_usage",
+        ]
+
+    def _get_profiler_events(self) -> EventList:
+        """Return the collected PyTorch events.
+
+        Returns:
+            The finalized PyTorch event list.
+
+        Raises:
+            AssertionError: If profiling has not been stopped correctly.
+        """
+        profiler = self.profiler
+        assert profiler is not None, "Profiling not stopped correctly"
+        profiler._ensure_function_events()
+        events = profiler._function_events
+        assert events is not None, "Profiler events are unavailable"
+        return events
+
+    def _get_profiler_events_by_record_step(self) -> dict[str, list[FunctionEvent]]:
+        """Group events by the most recent recorded step.
+
+        Returns:
+            A mapping from recorded step names to their events.
+        """
+        events = self._get_profiler_events()
+        matched_events = {step: [] for _, step in self.record_steps}
+        profiler = self.profiler
+        assert profiler is not None, "Profiling not stopped correctly"
+        kineto_results = profiler.kineto_results
+        assert kineto_results is not None, "Profiler trace is unavailable"
+        base_timestamp = kineto_results.trace_start_ns() * 1e-3
+
+        for event in events:
+            event_timestamp = base_timestamp + event.time_range.start
+            previous_steps = [
+                (event_timestamp - timestamp.timestamp() * 1e6, name)
+                for timestamp, name in self.record_steps
+                if event_timestamp >= timestamp.timestamp() * 1e6
+            ]
+            matched_events[min(previous_steps)[1]].append(event)
+        return matched_events
+
+    def _event_rows(self) -> list[dict[str, Any]]:
+        """Convert grouped PyTorch events into dataframe rows.
+
+        Returns:
+            Rows containing event metrics and their recorded step.
+        """
+        rows = []
+        for step, events in self._get_profiler_events_by_record_step().items():
+            for event in events:
+                row = {column: getattr(event, column, None) for column in self.numeric_columns}
+                row.update(
+                    name=getattr(event, "name", None),
+                    is_annotation=getattr(event, "is_user_annotation", None),
+                    device=getattr(event, "device_type", None).name,  # type: ignore[union-attr]
+                    record_step=step,
+                )
+                rows.append(row)
+        return rows
+
+    def to_pandas(self) -> pd.DataFrame:
+        """Return one row per PyTorch event with timing and memory metrics.
+
+        Returns:
+            A dataframe containing event metrics, devices, and percentages.
+        """
+        df = pd.DataFrame(self._event_rows())
+        df.loc[df["device"] == "CPU", ["self_device_time_total", "device_time_total"]] = 0
+        df.loc[df["device"] == "CUDA", ["self_cpu_time_total", "cpu_time_total"]] = 0
+        df["self_cpu_time_total_percentage"] = (
+            df["self_cpu_time_total"] / df["self_cpu_time_total"].sum() * 100
+        )
+        df["cpu_time_total_percentage"] = df["cpu_time_total"] / df["cpu_time_total"].sum() * 100
+        return df
+
+    def summary(self) -> pd.DataFrame:
+        """Return metrics summed by event name, excluding annotations.
+
+        Returns:
+            A dataframe with metrics grouped and sorted by event name.
+        """
+        df = self.to_pandas()
+        return (
+            df[~df["is_annotation"]][["name"] + self.numeric_columns]
+            .groupby("name")
+            .sum()
+            .sort_values(by=["flops", "count"])
+        )
+
+    def totals(self) -> pd.Series:
+        """Return total metrics for non-annotation events.
+
+        Returns:
+            A series containing the sum of each numeric metric.
+        """
+        df = self.to_pandas()
+        return df[~df["is_annotation"]][self.numeric_columns].sum(axis=0)
+
+    def total_time(self, device: str = "CUDA") -> float:
+        """Return total self time in microseconds for one device.
+
+        Args:
+            device: Device to measure, either ``"CPU"`` or ``"CUDA"``.
+
+        Returns:
+            Total self time in microseconds.
+
+        Raises:
+            AssertionError: If ``device`` is not ``"CPU"`` or ``"CUDA"``.
+        """
+        assert device in {"CPU", "CUDA"}, "device must be either CPU or CUDA"
+        time_field = "self_cpu_time_total" if device == "CPU" else "self_device_time_total"
+        return sum(
+            getattr(event, time_field, 0.0)
+            for event in self._get_profiler_events()
+            if event.device_type.name == device and not event.is_user_annotation
+        )
+
+    def total_flops(self) -> int:
+        """Return the total FLOPs recorded by PyTorch.
+
+        Returns:
+            The total number of floating-point operations.
+        """
+        return int(sum((getattr(event, "flops", 0.0) or 0.0) for event in self._get_profiler_events()))
+
+    def flops_by_step(self) -> pd.DataFrame:
+        """Return FLOPs summed for each recorded step.
+
+        Returns:
+            A dataframe indexed by step name with a ``flops`` column.
+        """
+        flops_by_step = {
+            name: sum((event.flops or 0.0) for event in events)
+            for name, events in self._get_profiler_events_by_record_step().items()
+        }
+        return pd.DataFrame.from_dict(flops_by_step, orient="index", columns=["flops"])
+
+    def time_by_step(self) -> pd.DataFrame:
+        """Return CPU and CUDA self time in microseconds for each step.
+
+        Returns:
+            A dataframe indexed by step with ``cpu_time`` and ``gpu_time``.
+        """
+        time_by_step = {}
+        for step, events in self._get_profiler_events_by_record_step().items():
+            cpu_time = sum(
+                getattr(event, "self_cpu_time_total", 0.0)
+                for event in events
+                if event.device_type.name == "CPU" and not event.is_user_annotation
+            )
+            gpu_time = sum(
+                getattr(event, "self_device_time_total", 0.0)
+                for event in events
+                if event.device_type.name == "CUDA" and not event.is_user_annotation
+            )
+            time_by_step[step] = (cpu_time, gpu_time)
+        return pd.DataFrame.from_dict(
+            time_by_step, orient="index", columns=["cpu_time", "gpu_time"]
+        )
+
+    # Compatibility aliases for versions before the API cleanup.
+    get_total_time = total_time
+    get_total_flops = total_flops
+    get_flops_by_step = flops_by_step
+    get_time_by_step = time_by_step