gpu_plot.py

#!/usr/bin/env python3
"""
GPU Data Plotter - runs on local machine, converts CSV to line chart PNG.
Requires: pip install matplotlib

Usage:
    python3 gpu_plot.py gpu_data.csv                    # generates gpu_data.png
    python3 gpu_plot.py gpu_data.csv -o chart.png       # specify output filename
    python3 gpu_plot.py gpu_data.csv --dpi 200          # high resolution
    python3 gpu_plot.py gpu_data.csv --metrics sm mem   # only plot SM and mem usage
"""

import argparse
import csv
import sys
from collections import defaultdict
from datetime import datetime

try:
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    import matplotlib.dates as mdates
    from matplotlib.font_manager import FontProperties
except ImportError:
    print("Error: matplotlib required. Run: pip install matplotlib")
    sys.exit(1)


METRIC_INFO = {
    "pwr":   {"label": "Power (W)",            "color": "#e74c3c"},
    "gtemp": {"label": "GPU Temp (C)",         "color": "#e67e22"},
    "sm":    {"label": "SM Utilization (%)",    "color": "#2ecc71"},
    "mem":   {"label": "Memory Util (%)",       "color": "#3498db"},
    "enc":   {"label": "Encoder (%)",           "color": "#9b59b6"},
    "dec":   {"label": "Decoder (%)",           "color": "#1abc9c"},
    "mclk":  {"label": "Memory Clock (MHz)",    "color": "#f39c12"},
    "pclk":  {"label": "SM Clock (MHz)",        "color": "#e91e63"},
}

DEFAULT_METRICS = ["pwr", "gtemp", "sm", "mem", "pclk"]


def parse_args():
    p = argparse.ArgumentParser(description="Convert GPU monitor CSV to line chart PNG")
    p.add_argument("csv_file", help="CSV file generated by gpu_collect.py")
    p.add_argument("-o", "--output", type=str, default=None,
                   help="Output PNG path (default: same name as CSV with .png)")
    p.add_argument("--dpi", type=int, default=150, help="Output DPI (default: 150)")
    p.add_argument("--metrics", nargs="+", default=DEFAULT_METRICS,
                   choices=list(METRIC_INFO.keys()),
                   help=f"Metrics to plot (default: {' '.join(DEFAULT_METRICS)})")
    return p.parse_args()


def load_csv(path):
    """Return (data, uuid_map).
    data: {gpu_idx: {metric: [(datetime, value), ...]}}
    uuid_map: {gpu_idx: uuid_str}
    """
    data = defaultdict(lambda: defaultdict(list))
    uuid_map = {}

    with open(path, "r") as f:
        reader = csv.DictReader(f)
        for row in reader:
            try:
                ts = datetime.strptime(row["timestamp"], "%Y-%m-%d %H:%M:%S")
            except (ValueError, KeyError):
                continue

            gpu = row.get("gpu", "0")

            uuid = row.get("uuid", "").strip()
            if uuid and uuid != "N/A" and gpu not in uuid_map:
                uuid_map[gpu] = uuid

            for metric in METRIC_INFO:
                raw = row.get(metric, "-")
                if raw == "-" or raw == "":
                    continue
                try:
                    val = float(raw)
                    data[gpu][metric].append((ts, val))
                except ValueError:
                    pass

    return data, uuid_map


def compute_stats(data, metrics):
    """Return {gpu_idx: {metric: {"max": v, "min": v, "avg": v}}}"""
    stats = {}
    for gpu, metric_data in data.items():
        stats[gpu] = {}
        for metric in metrics:
            points = metric_data.get(metric, [])
            vals = [v for _, v in points if v is not None]
            if vals:
                stats[gpu][metric] = {
                    "max": max(vals),
                    "min": min(vals),
                    "avg": sum(vals) / len(vals),
                }
            else:
                stats[gpu][metric] = {"max": "-", "min": "-", "avg": "-"}
    return stats


def plot(data, uuid_map, metrics, output_path, dpi):
    gpu_indices = sorted(data.keys(), key=lambda x: int(x) if x.isdigit() else x)
    if not gpu_indices:
        print("Error: no valid data found in CSV")
        sys.exit(1)

    metrics = [m for m in metrics if m in METRIC_INFO]
    n = len(metrics)
    stats = compute_stats(data, metrics)

    # Build summary table: one row per GPU, columns grouped by metric
    table_columns = ["GPU", "UUID"]
    for m in metrics:
        short = METRIC_INFO[m]["label"]
        table_columns.extend([f"{short}\nMax", f"{short}\nMin", f"{short}\nAvg"])

    table_rows = []
    for gpu in gpu_indices:
        uuid = uuid_map.get(gpu, "N/A")
        # Truncate UUID to last 12 chars for readability
        uuid_short = uuid if uuid == "N/A" else "..." + uuid[-12:]
        row = [f"GPU {gpu}", uuid_short]
        for m in metrics:
            s = stats[gpu].get(m, {})
            for key in ("max", "min", "avg"):
                v = s.get(key, "-")
                if isinstance(v, float):
                    row.append(f"{v:.1f}")
                else:
                    row.append(str(v))
        table_rows.append(row)

    # Calculate table height
    n_table_rows = len(table_rows) + 1  # +1 for header
    table_height = max(2.0, 0.45 * n_table_rows + 1.2)

    # Wide figure to accommodate the table
    fig_width = max(24, 2.0 * len(table_columns))
    fig_height = 3 * n + table_height + 1.5
    fig = plt.figure(figsize=(fig_width, fig_height))

    chart_top = 1.0 - 0.4 / fig_height
    chart_bottom = (table_height + 1.0) / fig_height
    gs_charts = fig.add_gridspec(n, 1, top=chart_top, bottom=chart_bottom,
                                  left=0.07, right=0.97, hspace=0.15)

    axes = [fig.add_subplot(gs_charts[i]) for i in range(n)]

    total_points = 0
    for ax, metric in zip(axes, metrics):
        info = METRIC_INFO[metric]
        for gpu in gpu_indices:
            points = data[gpu].get(metric, [])
            if not points:
                continue
            ts_list, val_list = zip(*points)
            total_points = max(total_points, len(ts_list))
            ax.plot(ts_list, val_list, label=f"GPU {gpu}", linewidth=1.2, alpha=0.85)

        ax.set_ylabel(info["label"], fontsize=10)
        ax.legend(loc="upper left", fontsize=8)
        ax.grid(True, alpha=0.3)

        if metric in ("sm", "mem", "enc", "dec"):
            ax.set_ylim(-2, 105)

    axes[-1].set_xlabel("Time", fontsize=10)
    axes[-1].xaxis.set_major_formatter(mdates.DateFormatter("%H:%M:%S"))
    if total_points > 500:
        axes[-1].xaxis.set_major_locator(mdates.AutoDateLocator())

    for ax in axes[:-1]:
        plt.setp(ax.get_xticklabels(), visible=False)
    fig.autofmt_xdate(rotation=30)

    # Summary table at the bottom
    table_bottom_pos = 0.3 / fig_height
    table_top_pos = chart_bottom - 0.4 / fig_height
    ax_table = fig.add_axes([0.03, table_bottom_pos, 0.94,
                              table_top_pos - table_bottom_pos])
    ax_table.axis("off")

    cell_colors = []
    header_color = "#34495e"
    for _ in table_rows:
        cell_colors.append(["#f8f9fa"] * len(table_columns))

    table = ax_table.table(
        cellText=table_rows,
        colLabels=table_columns,
        cellLoc="center",
        loc="center",
    )
    table.auto_set_font_size(False)
    table.set_fontsize(9)
    table.scale(1.0, 1.6)

    for (row, col), cell in table.get_celld().items():
        cell.set_edgecolor("#bdc3c7")
        if row == 0:
            cell.set_facecolor(header_color)
            cell.set_text_props(color="white", fontweight="bold")
        else:
            cell.set_facecolor("#f8f9fa" if row % 2 == 1 else "#ffffff")

    fig.suptitle(f"GPU Monitor ({total_points} data points)", fontsize=13,
                 fontweight="bold", y=1.0 - 0.15 / fig_height)
    fig.savefig(output_path, dpi=dpi, bbox_inches="tight")
    plt.close(fig)
    print(f"Chart saved -> {output_path}")


def main():
    args = parse_args()
    output = args.output
    if output is None:
        output = args.csv_file.rsplit(".", 1)[0] + ".png"

    data, uuid_map = load_csv(args.csv_file)
    plot(data, uuid_map, args.metrics, output, args.dpi)


if __name__ == "__main__":
    main()