system_monitor.py

import time
import threading
from config import logger

import psutil
from pynvml import nvmlInit, nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo, nvmlDeviceGetUtilizationRates # resources utilisation system monitor package
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.backends.backend_pdf import PdfPages

nvmlInit()

class SystemMonitor:
    """
    Polls CPU, RAM, GPU utilisation, and VRAM on a background thread.
    Stage boundaries are marked and included in the benchmark chart.
    """

    POLL_INTERVAL = 2

    def __init__(self):
        self.handle = nvmlDeviceGetHandleByIndex(0)
        self.timestamps = []
        self.cpu = []
        self.ram_used = []
        self.gpu_util = []
        self.vram_used = []
        self.stages = []
        self._running = False
        self._thread = None
        self._start_time = None
        self._stage_start = None

    def start(self):
        self._running = True
        self._start_time = time.time()
        self._stage_start = time.time()
        self._thread = threading.Thread(target=self._poll, daemon=True)
        self._thread.start()
        logger.info("System monitor started.")

    def stop(self):
        self._running = False
        self._thread.join()
        logger.info("System monitor stopped.")

    def mark_stage(self, label: str):
        now = time.time()
        elapsed = now - self._start_time
        if self.stages:
            prev_duration = now - self._stage_start
            logger.info(f"Stage '{self.stages[-1][0]}' completed in {prev_duration:.1f}s")
            self.stages[-1] = (self.stages[-1][0], self.stages[-1][1], prev_duration)
        self._stage_start = now
        self.stages.append((label, elapsed, None))
        logger.info(f"Stage started: '{label}' at {elapsed:.1f}s into run")

    def get_stage_summary(self) -> str:
        lines = ["Pipeline stage wall times:"]
        for label, _, duration in self.stages:
            lines.append(f"  {label}: {duration:.1f}s" if duration else f"  {label}: in progress")
        return "\n".join(lines)

    def _poll(self):
        while self._running:
            elapsed = time.time() - self._start_time
            self.cpu.append(psutil.cpu_percent(interval=None))
            self.ram_used.append(psutil.virtual_memory().used / 1024 ** 3)
            mem = nvmlDeviceGetMemoryInfo(self.handle)
            util = nvmlDeviceGetUtilizationRates(self.handle)
            self.gpu_util.append(util.gpu)
            self.vram_used.append(mem.used / 1024 ** 3)
            self.timestamps.append(elapsed)
            time.sleep(self.POLL_INTERVAL)

    def build_benchmark_figure(self) -> plt.Figure:
        fig, axes = plt.subplots(4, 1, figsize=(16, 14), sharex=True)
        fig.suptitle("Sentibot - System Resource Usage During Pipeline", fontsize=14, fontweight="bold")
        t = self.timestamps
        stage_colours = ["#e6194b", "#3cb44b", "#4363d8", "#f58231", "#911eb4", "#42d4f4", "#f032e6"]

        def add_markers(ax):
            for i, (label, ts, duration) in enumerate(self.stages):
                ax.axvline(x=ts, color=stage_colours[i % len(stage_colours)], linestyle="--", linewidth=1.2, alpha=0.8)

        axes[0].plot(t, self.cpu, color="#4363d8", linewidth=1.5)
        axes[0].set_ylabel("CPU (%)"); axes[0].set_ylim(0, 100); axes[0].grid(True, alpha=0.3); add_markers(axes[0])

        total_ram = psutil.virtual_memory().total / 1024 ** 3
        axes[1].plot(t, self.ram_used, color="#911eb4", linewidth=1.5)
        axes[1].set_ylabel("RAM (GB)"); axes[1].set_ylim(0, total_ram); axes[1].grid(True, alpha=0.3); add_markers(axes[1])

        axes[2].plot(t, self.gpu_util, color="#3cb44b", linewidth=1.5)
        axes[2].set_ylabel("GPU (%)"); axes[2].set_ylim(0, 100); axes[2].grid(True, alpha=0.3); add_markers(axes[2])

        axes[3].plot(t, self.vram_used, color="#f58231", linewidth=1.5)
        axes[3].set_ylabel("VRAM (GB)"); axes[3].set_ylim(0, 4); axes[3].set_xlabel("Time (seconds)")
        axes[3].grid(True, alpha=0.3); add_markers(axes[3])

        handles = [ mpatches.Patch(color=stage_colours[i % len(stage_colours)], label=f"{label}" + (f" ({d:.1f}s)" if (d := duration) else "")) for i, (label, _, duration) in enumerate(self.stages) ]
        fig.legend(handles=handles, loc="lower center", ncol=4, fontsize=8, title="Pipeline Stages", bbox_to_anchor=(0.5, 0.01))
        plt.tight_layout(rect=[0, 0.06, 1, 1])
        return fig

    def save_benchmark_pdf(self, output_path: str) -> bool:
        """Save the benchmark figure to its own PDF"""
        if len(self.timestamps) <= 3:
            logger.warning("Not enough monitor data for benchmark page.")
            return False
        with PdfPages(output_path) as pdf:
            fig = self.build_benchmark_figure()
            pdf.savefig(fig, bbox_inches="tight")
            plt.close(fig)
        logger.info(f"Benchmark PDF saved locally: {output_path}")
        return True