vLLM-PrefixCache-Monitor/visualize.py at main · Kingpython123/vLLM-PrefixCache-Monitor · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
"""Visualize vLLM Prefix Cache benchmark results."""

import json
import argparse
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np


def plot_hit_rate(results: list[dict], output_path: str):
    workloads = [r["workload"] for r in results]
    hit_rates = [
        r.get("cache_stats", {}).get("query_hit_rate", 0) for r in results
    ]
    labels = ["Short Chat", "Long-doc QA", "Same Prompt", "No Reuse"]

    fig, ax = plt.subplots(figsize=(8, 5))
    colors = ["#e74c3c", "#2ecc71", "#3498db", "#95a5a6"]

    bars = ax.bar(range(len(workloads)), hit_rates, color=colors, width=0.5)
    for bar, rate in zip(bars, hit_rates):
        ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.02,
                f"{rate*100:.1f}%", ha="center", fontweight="bold")

    ax.set_xticks(range(len(workloads)))
    ax.set_xticklabels(labels)
    ax.set_ylabel("Cache Hit Rate")
    ax.set_title("Prefix Cache Hit Rate by Workload")
    ax.set_ylim(0, 1.1)
    ax.axhline(y=0.5, color="gray", linestyle="--", alpha=0.3)
    plt.tight_layout()
    plt.savefig(output_path, dpi=150)


def plot_throughput(results: list[dict], output_path: str):
    workloads_disp = ["Short Chat", "Long-doc QA", "Same Prompt", "No Reuse"]
    throughputs = [r["throughput_rps"] for r in results]
    cache_hit = [
        r.get("cache_stats", {}).get("query_hit_rate", 0) for r in results
    ]

    fig, ax1 = plt.subplots(figsize=(10, 5))

    x = range(len(workloads_disp))
    colors = ["#e74c3c", "#2ecc71", "#3498db", "#95a5a6"]
    ax1.bar(x, throughputs, color=colors, width=0.4, label="Throughput (rps)")
    ax1.set_ylabel("Throughput (requests/s)")
    ax1.set_xticks(x)
    ax1.set_xticklabels(workloads_disp)

    ax2 = ax1.twinx()
    ax2.plot(x, [h * 100 for h in cache_hit], "ko-", linewidth=2, markersize=8, label="Hit Rate %")
    ax2.set_ylabel("Cache Hit Rate (%)")
    ax2.set_ylim(0, 100)

    lines1, labels1 = ax1.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    ax1.legend(lines1 + lines2, labels1 + labels2, loc="upper left")

    ax1.set_title("Throughput vs Cache Hit Rate by Workload")
    plt.tight_layout()
    plt.savefig(output_path, dpi=150)


def plot_latency(results: list[dict], output_path: str):
    workloads_disp = ["Short Chat", "Long-doc QA", "Same Prompt", "No Reuse"]
    p50 = [r["latency_p50_ms"] for r in results]
    p99 = [r["latency_p99_ms"] for r in results]

    fig, ax = plt.subplots(figsize=(8, 5))
    x = np.arange(len(workloads_disp))
    width = 0.35

    ax.bar(x - width/2, p50, width, color="#3498db", label="P50 Latency")
    ax.bar(x + width/2, p99, width, color="#e74c3c", label="P99 Latency")
    ax.set_xticks(x)
    ax.set_xticklabels(workloads_disp)
    ax.set_ylabel("Latency (ms)")
    ax.set_title("Latency Distribution by Workload")
    ax.legend()
    plt.tight_layout()
    plt.savefig(output_path, dpi=150)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--results", required=True)
    parser.add_argument("--output_dir", default="./figures")
    args = parser.parse_args()

    with open(args.results) as f:
        results = json.load(f)

    import os
    os.makedirs(args.output_dir, exist_ok=True)

    plot_hit_rate(results, f"{args.output_dir}/hit_rate.png")
    plot_throughput(results, f"{args.output_dir}/throughput.png")
    plot_latency(results, f"{args.output_dir}/latency.png")
    print(f"Figures saved to {args.output_dir}/")


if __name__ == "__main__":
    main()