Function main

perfs/plot_memory_metrics.py:112–164

()

def main() -> None:
    args = parse_args()
    if not args.input_file.is_file():
        raise SystemExit(f"Input file does not exist: {args.input_file}")

    try:
        import matplotlib.pyplot as plt
    except ImportError as exc:
        raise SystemExit(
            "Plotting requires matplotlib. Install it in the environment first."
        ) from exc

    loaded_runs = []
    for label, metrics_file in _resolve_metrics_files(args.input_file):
        if not metrics_file.is_file():
            raise SystemExit(f"Metrics file does not exist: {metrics_file}")
        total_pages, loaded_metrics = load_points(metrics_file)
        if total_pages:
            loaded_runs.append((label, total_pages, loaded_metrics))

    if not loaded_runs:
        raise SystemExit(f"No plotable data found in {args.input_file}")

    fig, ax = plt.subplots(1, 1, figsize=(12, 6))

    for label, total_pages, loaded_metrics in loaded_runs:
        point_count = min(
            [len(total_pages)] + [len(values) for values in loaded_metrics.values()]
        )
        x_values = total_pages[:point_count]

        for key in MEMORY_KEYS:
            series = loaded_metrics[key][:point_count]
            if all(value != value for value in series):
                continue
            metric_label = key.replace("_mb", "").upper()
            ax.plot(
                x_values,
                series,
                label=f"{label} {metric_label}",
                linewidth=1.5,
            )

    ax.set_xlabel("Total page number")
    ax.set_ylabel("Memory (MiB)")
    ax.set_title("Memory at page loaded")
    ax.grid(True, alpha=0.3)
    ax.legend()
    fig.tight_layout()

    args.output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(args.output, dpi=150)
    print(f"Wrote plot to {args.output}")


if __name__ == "__main__":