offloading_benchmarks_plotter.py

import os
import json
import matplotlib.pyplot as plt
# import seaborn as sns
from collections import defaultdict
import pandas as pd

plt.style.use("ggplot")

# Adjust this path to the folder where the JSON files are located
DATA_DIR = "offload_results/benchmark_10240_swap/"  # e.g., "./results/"

def plot_layer_comparison():
    # Store all data indexed by chunk_size -> block_size -> list of offload results
    results = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
    parameter_metadata = {}

    # Scan all JSON files
    for filename in os.listdir(DATA_DIR):
        if filename.endswith(".json") and filename.startswith("bs_"):
            path = os.path.join(DATA_DIR, filename)
            with open(path) as f:
                data = json.load(f)
            
            # Extract sizes
            parts = filename[:-5].split("_")
            block_size = int(parts[1])
            chunk_size = int(parts[3])
            num_layers = int(parts[5])

            parameter_metadata[(chunk_size, block_size, num_layers)] = data["parameters"]

            # Parse offload results
            transfer_times = data["data"]
            for entry in transfer_times:
                results[chunk_size][block_size][num_layers].append(
                    {k: float(v) for k, v in entry.items()}
                )
    # Plotting
    for chunk_size in sorted(results):
        for block_size in sorted(results[chunk_size]):
            plt.figure(figsize=(12, 6))
            offload_keys = set()
            for num_layers in sorted(results[chunk_size][block_size]):
                entries = results[chunk_size][block_size][num_layers]
                # Average over all entries
                averaged = defaultdict(float)
                for entry in entries:
                    for k, v in entry.items():
                        averaged[k] += v
                        offload_keys.add(k)
                for k in averaged:
                    averaged[k] /= len(entries)

                # Plot
                keys = sorted(offload_keys)
                values = [averaged[k] for k in keys]
                plt.plot(keys, values, marker='o', label=f"{num_layers} layers")

            # Metadata from first (chunk_size, block_size) combo
            # sample_param = parameter_metadata[(chunk_size, sorted(results[chunk_size])[0])]
            # meta_text = "\n".join([
            #     f"num_blocks: {sample_param['num_blocks']}",
            #     f"num_layers: {sample_param['num_layers']}",
            #     f"num_heads: {sample_param['num_heads']}",
            #     f"head_size: {sample_param['head_size']}",
            #     f"num_tokens: {sample_param['num_tokens']}",
            #     f"chunk_size: {sample_param['chunk_size']}"
            # ])

            plt.title(f"Offload Times Comparison\nBlock size {block_size} Chunk size {chunk_size}")
            plt.xlabel("Chunk iteration")
            plt.ylabel("Average Time (seconds)")
            plt.legend()
            plt.xticks(rotation=45)
            plt.tight_layout()

            # Add metadata as text box
            # plt.gcf().text(0.75, 0.5, meta_text, fontsize=10, bbox=dict(facecolor='white', alpha=0.5))

            # Show or save plot
            # plt.show()
            plt.savefig(f"plots/offload_plots/num_layers_comparisons/layers_comparison_bs_{block_size}_cs_{chunk_size}.png")

def plot_block_size_comparison():
    # Store all data indexed by chunk_size -> block_size -> list of offload results
    results = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
    parameter_metadata = {}

    # Scan all JSON files
    for filename in os.listdir(DATA_DIR):
        if filename.endswith(".json") and filename.startswith("bs_"):
            path = os.path.join(DATA_DIR, filename)
            with open(path) as f:
                data = json.load(f)
            
            # Extract sizes
            parts = filename[:-5].split("_")
            block_size = int(parts[1])
            chunk_size = int(parts[3])
            num_layers = int(parts[5])

            parameter_metadata[(chunk_size, block_size, num_layers)] = data["parameters"]

            # Parse offload results
            transfer_times = data["data"]["cpu_to_gpu"]
            for entry in transfer_times:
                results[chunk_size][num_layers][block_size].append(
                    {k: float(v) for k, v in entry.items()}
                )
    # Plotting
    for chunk_size in sorted(results):
        for num_layers in sorted(results[chunk_size]):
            plt.figure(figsize=(12, 6))
            offload_keys = set()
            for block_size in sorted(results[chunk_size][num_layers]):
                entries = results[chunk_size][num_layers][block_size]
                # Average over all entries
                averaged = defaultdict(float)
                for entry in entries:
                    for k, v in entry.items():
                        averaged[k] += v
                        offload_keys.add(k)
                for k in averaged:
                    averaged[k] /= len(entries)

                # Plot
                keys = sorted(offload_keys)
                values = [averaged[k] for k in keys]
                plt.plot(keys, values, marker='o', label=f"block size {block_size}")

            plt.title(f"Offload Times Comparison\nChunk size {chunk_size}, {num_layers} layers")
            plt.xlabel("Chunk iteration")
            plt.ylabel("Average Time (seconds)")
            plt.legend()
            plt.xticks(rotation=45)
            plt.tight_layout()

            plt.savefig(f"plots/offload_plots/block_size_comparisons/block_size_comparison_cs_{chunk_size}_l_{num_layers}.png")    

def plot_block_size_comparison_median():
    # Store all data indexed by chunk_size -> block_size -> list of offload results
    results = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
    parameter_metadata = {}

    # Scan all JSON files
    for filename in os.listdir(DATA_DIR):
        if filename.endswith(".json") and filename.startswith("bs_"):
            path = os.path.join(DATA_DIR, filename)
            with open(path) as f:
                data = json.load(f)
            
            # Extract sizes
            parts = filename[:-5].split("_")
            block_size = int(parts[1])
            chunk_size = int(parts[3])
            num_layers = int(parts[5])

            parameter_metadata[(chunk_size, block_size, num_layers)] = data["parameters"]

            # Parse offload results
            transfer_times = data["data"]["cpu_to_gpu"]
            for entry in transfer_times:
                results[chunk_size][num_layers][block_size].append(
                    {k: float(v) for k, v in entry.items()}
                )
    # Plotting
    for chunk_size in sorted(results):
        for num_layers in sorted(results[chunk_size]):
            plt.figure(figsize=(12, 6))
            offload_keys = set()
            for block_size in sorted(results[chunk_size][num_layers]):
                entries = results[chunk_size][num_layers][block_size]
                # Average over all entries
                lists_per_offload = defaultdict(list)
                for entry in entries:
                    for k, v in entry.items():
                        lists_per_offload[k].append(v)
                        offload_keys.add(k)
                medians = defaultdict(float)
                for k in lists_per_offload:
                    mid_index = len(lists_per_offload[k])//2
                    lists_per_offload[k].sort()
                    medians[k] = lists_per_offload[k][mid_index]

                # Plot
                keys = sorted(offload_keys)
                values = [medians[k] for k in keys]
                plt.plot(keys, values, marker='o', label=f"block size {block_size}")

            plt.title(f"Offload Times medians\nChunk size {chunk_size}, {num_layers} layers")
            plt.xlabel("Chunk iteration")
            plt.ylabel("Median Time (seconds)")
            plt.legend()
            plt.xticks(rotation=45)
            plt.tight_layout()

            plt.savefig(f"plots/offload_plots/block_size_comparisons_median/block_size_comparison_median_cs_{chunk_size}_l_{num_layers}.png")


def tps_tables():
    num_tokens = 800
    # Store all data indexed by chunk_size -> block_size -> list of offload results
    results = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
    parameter_metadata = {}

    # Scan all JSON files
    for filename in os.listdir(DATA_DIR):
        if filename.endswith(".json") and filename.startswith("bs_"):
            path = os.path.join(DATA_DIR, filename)
            with open(path) as f:
                data = json.load(f)
            
            # Extract sizes
            parts = filename[:-5].split("_")
            block_size = int(parts[1])
            chunk_size = int(parts[3])
            num_layers = int(parts[5])

            parameter_metadata[(chunk_size, block_size, num_layers)] = data["parameters"]

            # Parse offload results
            transfer_times = data["data"]["cpu_to_gpu"]
            sum = 0
            for entry in transfer_times:
                for v in entry.items():
                    print(v)
                    sum += float(v[1])
                print(sum)
            results[num_layers][chunk_size][block_size] = sum / (len(transfer_times) * num_tokens)
    
    # Step 1: Collect all unique block_sizes and chunk_sizes
    all_block_sizes = sorted({block_size for layer in results.values() for block_size in layer})
    all_chunk_sizes = sorted({chunk_size for layer in results.values()
                            for block in layer.values() for chunk_size in block})
    
    # Set up figure with 2 rows × 3 columns of subplots
    fig, axes = plt.subplots(2, 3, figsize=(18, 10))
    axes = axes.flatten()

    # Plot each table
    for i, (num_layers, block_dict) in enumerate(sorted(results.items())):
        # Build a full DataFrame with consistent row and column order
        df = pd.DataFrame(index=all_block_sizes, columns=all_chunk_sizes)
        for block_size, chunks in block_dict.items():
            for chunk_size, value in chunks.items():
                df.loc[block_size, chunk_size] = value

        # Convert to float and format
        df = df.astype(float)
        formatted_values = [[f"{val:.10f}" for val in row] for row in df.values]

        axes[i].axis('off')
        axes[i].set_title(f'num_layers = {num_layers}', fontsize=16)
        table = axes[i].table(
            cellText=formatted_values,
            rowLabels=df.index,
            colLabels=df.columns,
            loc='center'
        )
        table.scale(1.2, 1.5)

        # Set font size
        for key, cell in table.get_celld().items():
            cell.set_fontsize(12)

    # Hide the unused subplot if fewer than 6 tables
    for j in range(len(results), 6):
        axes[j].axis('off')

    plt.tight_layout()
    plt.savefig("plots/offload_plots/tps_by_layer.png", dpi=300)
    plt.show()

def tps_tables_2():
    num_tokens = 10240
    # Store all data indexed by chunk_size -> block_size -> list of offload results
    results = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
    parameter_metadata = {}

    # Scan all JSON files
    for filename in os.listdir(DATA_DIR):
        if filename.endswith(".json") and filename.startswith("bs_"):
            path = os.path.join(DATA_DIR, filename)
            with open(path) as f:
                data = json.load(f)
            
            # Extract sizes
            parts = filename[:-5].split("_")
            block_size = int(parts[1])
            chunk_size = int(parts[3])
            num_layers = int(parts[5])

            parameter_metadata[(chunk_size, block_size, num_layers)] = data["parameters"]

            # Parse offload results
            transfer_times = data["data"]
            bpt = 2 * num_layers * 8 * 128 * 2
            average_transfer_time = sum(transfer_times) / len(transfer_times)
            tps = num_tokens/average_transfer_time
            bps = tps * bpt
            results[num_layers][chunk_size][block_size] = (tps, bps)
    
    # Step 1: Collect all unique block_sizes and chunk_sizes
    all_block_sizes = sorted({block_size for layer in results.values() for block_size in layer})
    all_chunk_sizes = sorted({chunk_size for layer in results.values()
                            for block in layer.values() for chunk_size in block})
    
    # Set up figure with 2 rows × 3 columns of subplots
    fig, axes = plt.subplots(2, 3, figsize=(18, 10))
    axes = axes.flatten()

    # Plot each table
    for i, (num_layers, block_dict) in enumerate(sorted(results.items())):
        # Build a full DataFrame with consistent row and column order
        df = pd.DataFrame(index=all_block_sizes, columns=all_chunk_sizes)
        for block_size, chunks in block_dict.items():
            for chunk_size, value in chunks.items():
                df.loc[block_size, chunk_size] = value

        # Convert to float and format
        # df = df.astype(float, float)
        formatted_values = [[f"{tps:.2f}, {bps/1000000000:.2f}" for tps, bps in row] for row in df.values]

        axes[i].axis('off')
        axes[i].set_title(f'num_layers = {num_layers}', fontsize=16)
        table = axes[i].table(
            cellText=formatted_values,
            rowLabels=df.index,
            colLabels=df.columns,
            loc='center'
        )
        table.scale(1.2, 1.5)

        # Set font size
        for key, cell in table.get_celld().items():
            cell.set_fontsize(12)

    # Hide the unused subplot if fewer than 6 tables
    for j in range(len(results), 6):
        axes[j].axis('off')

    plt.tight_layout()
    plt.savefig(f"plots/offload_plots/tps_by_layer_swap_{num_tokens}_cpu2gpu.png", dpi=300)
    plt.show()

def tps_tables_swap():
    num_tokens = 10240
    # Store all data indexed by chunk_size -> block_size -> list of offload results
    results = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
    parameter_metadata = {}

    # Scan all JSON files
    for filename in os.listdir(DATA_DIR):
        if filename.endswith(".json") and filename.startswith("bs_"):
            path = os.path.join(DATA_DIR, filename)
            with open(path) as f:
                data = json.load(f)
            
            # Extract sizes
            parts = filename[:-5].split("_")
            block_size = int(parts[1])
            chunk_size = int(parts[3])
            num_layers = int(parts[5])

            if chunk_size == 0:
                parameter_metadata[(block_size, num_layers)] = data["parameters"]

                # Parse offload results
                transfer_times = data["data"]
                bpt = 2 * num_layers * 8 * 128 * 2
                average_transfer_time = sum(transfer_times) / len(transfer_times)
                tps = num_tokens/average_transfer_time
                bps = tps * bpt
                results[num_layers][block_size] = (tps, bps)
    
    # Step 1: Collect all unique block_sizes and chunk_sizes
    all_block_sizes = sorted({block_size for layer in results.values() for block_size in layer})
    
    # Set up figure with 2 rows × 3 columns of subplots
    fig, axes = plt.subplots(2, 3, figsize=(18, 10))
    axes = axes.flatten()

    # Plot each table
    for i, (num_layers, block_dict) in enumerate(sorted(results.items())):
        # Build a full DataFrame with consistent row and column order
        df = pd.DataFrame(index=all_block_sizes, columns=["TPS", "BPS"])
        for block_size, value in block_dict.items():
            df.loc[block_size, "TPS"] = block_dict[block_size][0]
            df.loc[block_size, "BPS"] = block_dict[block_size][1]
        
        formatted_values = [
            [f"{row['TPS']:.2f}, {row['BPS']/1e9:.2f}"] if pd.notnull(row['TPS']) else "" 
            for _, row in df.iterrows()
        ]
        formatted_values = [[val] for val in formatted_values]  # Because we only have one column

        axes[i].axis('off')
        axes[i].set_title(f'num_layers = {num_layers}', fontsize=16)
        table = axes[i].table(
            cellText=formatted_values,
            rowLabels=df.index,
            colLabels=["TPS, BPS (GB/s)"],
            loc='center'
        )
        table.scale(1.2, 1.5)

        # Set font size
        for key, cell in table.get_celld().items():
            cell.set_fontsize(12)

    # Hide the unused subplot if fewer than 6 tables
    for j in range(len(results), 6):
        axes[j].axis('off')

    plt.tight_layout()
    plt.savefig(f"plots/offload_plots/tps_by_layer_swap_{num_tokens}_cpu2gpu.png", dpi=300)
    plt.show()

if __name__ == "__main__":
    # plot_block_size_comparison_median()
    # plot_layer_comparison()
    tps_tables_swap()