meta_analysis.py

import matplotlib as mpl
mpl.use('Agg')  # noqa
import matplotlib.pyplot as plt
import numpy as np
import argparse
import os
import pickle
import re
import json


from benchmark.datasets import DATASETS
from benchmark.algorithms.definitions import get_definitions
from benchmark.plotting.metrics import all_metrics as metrics
from benchmark.plotting.metrics import get_all_recall_values, get_count_at_certain_recall
from benchmark.plotting.utils import (get_plot_label, compute_metrics,
        create_linestyles, create_pointset)
from benchmark.results import (store_results, load_all_results, load_all_results_without_read, 
        get_result_filename, get_unique_algorithms)
from benchmark.dataset_io import knn_result_read
import benchmark.streaming.compute_gt
from benchmark.streaming.load_runbook import load_runbook
from benchmark.utils import read_gt_fromdir, read_end2end_gt_fromdir


def read_int_arg_from_filename(filename, argname, default):
    pattern = f"{argname}_(\d+)"
    match = re.search(pattern, filename)

    if match:
        number = int(match.group(1))
        return number
    else:
        return default

 
def read_float_arg_from_filename(filename, argname, default):
    pattern = f"{argname}_0_(\d+)"
    match = re.search(pattern, filename)

    if match:
        digits = len(match.group(1))
        number = int(match.group(1)) / (10 ** digits)
        return number
    else:
        return default


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--dataset',
        metavar="DATASET",
        required=True)
    parser.add_argument(
        '--count',
        default=-1,
        type=int)
    parser.add_argument(
        '--definitions',
        metavar='FILE',
        help='load algorithm definitions from FILE',
        default='algos-2021.yaml')
    parser.add_argument(
        '--neurips23track',
        choices=['filter', 'ood', 'sparse', 'streaming', 'none'],
        default='none'
    )
    parser.add_argument(
        '--runbook_path',
        metavar='FILE',
        help='paths to runbooks',
    )
    parser.add_argument(
        '--private-query',
        help='Use the private queries and ground truth',
        action='store_true')
    parser.add_argument(
        '--filtered',
        help='Use filtered queries.',
        action='store_true'
    )
    parser.add_argument(
        '--label_file', 
        type=str, 
        default=None,
        help='Path to the label file.'
    )
    parser.add_argument(
        '--filter_label_file', 
        type=str, 
        default=None,
        help='Path to the filter file.'
    )
    parser.add_argument(
        '--mode',
        type=str,
        default=''
    )
    parser.add_argument(
        '--results_base_path',
        type=str,
        default='results'
    )
    parser.add_argument(
        '--dump_json',
        help='Dump results to json.',
        action='store_true'
    )
    parser.add_argument(
        '--dump_distances',
        help='Dump average distances to a log file.',
        action='store_true'
    )
    parser.add_argument(
        '--dump_timeline',
        help='Dump timeline data for visualization.',
        action='store_true'
    )
    args = parser.parse_args()

    dataset = DATASETS[args.dataset]()
    dim = dataset.d

    if args.count == -1:
        root = os.path.join(args.results_base_path, "neurips23", args.neurips23track, os.path.split(args.runbook_path)[1], args.dataset)
        counts = os.listdir(root)
        counts = [int(count) if count != 'end2end' else 'end2end' for count in counts]
    else:
        counts = [int(args.count)]

    max_pts, runbook = load_runbook(args.dataset, dataset.nb, args.runbook_path)

    result_list = []
    valid_counts = set()
    recall_targets = set()

    for count in counts:
        optimal_dumped_for_this_count = False

        if count == 'end2end':
            results = load_all_results(args.dataset, 0, neurips23track=args.neurips23track, runbook_path=args.runbook_path, \
                                       filtered=args.filtered, label_file=args.label_file, filter_label_file=args.filter_label_file, 
                                       end2end=True, base_path=args.results_base_path)
        else:
            results = load_all_results(args.dataset, count, neurips23track=args.neurips23track, runbook_path=args.runbook_path, \
                                       filtered=args.filtered, label_file=args.label_file, filter_label_file=args.filter_label_file, 
                                       end2end=False, base_path=args.results_base_path)

        for i, (fileroot, filename, properties, run) in enumerate(results):
            if args.mode != '':
                if not filename.startswith(args.mode):
                    continue

            if filename.startswith("normal"):
                mode = "normal"
                train_mode = None
                train_counts = []
                train_queries = 0
            elif filename.startswith("early_stop_normal"):
                mode = "laet"
                train_mode = "train_normal"
                train_counts = properties['model_counts'].tolist()
                train_queries = properties['num_train_queries'] if properties['num_train_queries'] > 0 else dataset.ntrain
            elif filename.startswith("early_stop_opt_normal"):
                mode = "ours"
                train_mode = "train_opt_normal"
                train_counts = properties['model_counts'].tolist()
                train_queries = properties['num_train_queries'] if properties['num_train_queries'] > 0 else dataset.ntrain
            elif filename.startswith("early_stop_darth_normal"):
                mode = "darth"
                train_mode = "train_darth_normal"
                train_counts = properties['model_counts'].tolist()
                train_queries = properties['num_train_queries'] if properties['num_train_queries'] > 0 else dataset.ntrain
            elif filename.startswith("early_stop_darth_opt_normal"):
                mode = "darth_opt"
                train_mode = "train_darth_opt_normal"
                train_counts = properties['model_counts'].tolist()
                train_queries = properties['num_train_queries'] if properties['num_train_queries'] > 0 else dataset.ntrain
            elif filename.startswith("rem_normal"):
                mode = "rem"
                train_mode = "train_rem_normal"
                train_counts = properties['model_counts'].tolist()
                train_queries = properties['num_train_queries'] if properties['num_train_queries'] > 0 else dataset.ntrain
            else:
                mode = "train"
                train_mode = None
                train_counts = []
                train_queries = 0

            print(f"from {fileroot}/{filename}")
            if fileroot.split("/")[-1].endswith('.txt'):
                label_file, filter_label_file = fileroot.split("/")[-2], fileroot.split("/")[-1]
            else:
                label_file, filter_label_file = "", ""
            gt_dir = benchmark.streaming.compute_gt.gt_dir(dataset, args.runbook_path, label_file, filter_label_file)

            threads = read_int_arg_from_filename(filename, 'threadsSearch', 1)
            repeat = read_int_arg_from_filename(filename, 'repeat', 1)
            recall_target = read_float_arg_from_filename(filename, 'recallTarget', 1)
            traversal_window_size = read_int_arg_from_filename(filename, 'traversalWindow', 0)
            do_not_use_dynamic_recall_threshold = read_int_arg_from_filename(filename, 'disable1', 0)
            do_not_predict_average_recall = read_int_arg_from_filename(filename, 'disable2', 0)
            do_not_use_prediction_intervals = read_int_arg_from_filename(filename, 'disable3', 0)
            if count != 'end2end':
                valid_counts.add(count)
                recall_targets.add(recall_target)

            for i in range(0, properties['num_searches']):
                search_step_id = properties['step_' + str(i)]
                step_suffix = str(search_step_id)

                neighbors = np.array(run['neighbors_step' +  step_suffix])
                neighbor_distances = np.array(run['neighbor_distances_step' + step_suffix])
                try:
                    start_times = np.array(run['start_times_step' + step_suffix])
                    end_times = np.array(run['end_times_step' + step_suffix])
                except:
                    start_times = None
                    end_times = None
                total_cmps = np.array(run['total_cmps_step' + step_suffix])
                total_latency = np.array(run['total_latency_step' + step_suffix])
                cmps = np.array(run['cmps_step' + step_suffix])
                lats = np.array(run['lats_step' + step_suffix])
                mean_prediction_times = np.mean(np.array(run['prediction_times_step' + step_suffix]))
                mean_prediction_overhead = np.mean(np.array(run['prediction_overhead_step' + step_suffix]))

                if filename.startswith('train'):
                    if count == 'end2end':
                        Q, ks, recall_targets_ = dataset.get_query_sequence()
                        groundtruths, groundtruth_distances = read_end2end_gt_fromdir(gt_dir, step_suffix, ks, train=True)
                    else:
                        Q = dataset.get_queries()
                        groundtruths, groundtruth_distances = read_gt_fromdir(gt_dir, step_suffix, count, train=True)
                else:
                    if count == 'end2end':
                        Q, ks, recall_targets_ = dataset.get_query_sequence()
                        groundtruths, groundtruth_distances = read_end2end_gt_fromdir(gt_dir, step_suffix, ks, train=False)
                    else:
                        Q = dataset.get_queries()
                        groundtruths, groundtruth_distances = read_gt_fromdir(gt_dir, step_suffix, count, train=False)

                if args.dump_distances:
                    print(f"Shape of neighbor_distances: {neighbor_distances.shape}")
                    print(f"Shape of groundtruth_distances: {groundtruth_distances.shape}")
                    avg_neighbor_distances = np.mean(neighbor_distances, axis=0)
                    avg_groundtruth_distances = np.mean(groundtruth_distances, axis=0)
                    print(f"Avg neighbor_distances: {avg_neighbor_distances}")
                    print(f"Avg groundtruth_distances: {avg_groundtruth_distances}")
                    # dump avg_neighbor_distances and avg_groundtruth_distances to avg_distances_{count}.log
                    with open(f"avg_distances_{count}.log", "w") as f:
                        f.write(f"Avg neighbor_distances: \n{','.join([str(dist) for dist in avg_neighbor_distances])}\n")
                        f.write(f"Avg groundtruth_distances: \n{','.join([str(dist) for dist in avg_groundtruth_distances])}\n")

                if args.dump_timeline:
                    try:
                        dists_visited = np.array(run['dists_visited_step' + step_suffix])
                        cmps_visited = np.array(run['cmps_visited_step' + step_suffix])
                        hops_visited = np.array(run['hops_visited_step' + step_suffix])

                        with open(f"cmps_timeline_{count}.log", "w") as cmps_timeline:
                            with open(f"dists_timeline_{count}.log", "w") as dists_timeline:
                                with open(f"gt_cmps_timeline_{count}.log", "w") as gt_cmps_timeline:
                                    # for each query, group dists and cmps by hops
                                    for dists_visited_per_query, cmps_visited_per_query, hops_visited_per_query in zip(dists_visited, cmps_visited, hops_visited):
                                        dists_grouped_by_hops = {}
                                        cmps_grouped_by_hops = {}
                                        for curr_dist, curr_cmps, curr_hops in zip(dists_visited_per_query, cmps_visited_per_query, hops_visited_per_query):
                                            if curr_hops not in dists_grouped_by_hops or curr_dist < dists_grouped_by_hops[curr_hops]:
                                                dists_grouped_by_hops[curr_hops] = curr_dist
                                                cmps_grouped_by_hops[curr_hops] = curr_cmps
                                        dists_list = [dists_grouped_by_hops[hop] for hop in sorted(dists_grouped_by_hops.keys())]
                                        cmps_list = [cmps_grouped_by_hops[hop] for hop in sorted(cmps_grouped_by_hops.keys())]
                                        dists_cmps_pairs = list(zip(dists_list, cmps_list))
                                        dists_cmps_pairs.sort(key=lambda x: x[0])  # sort by distance
                                        gt_cmps = [cmps for _, cmps in dists_cmps_pairs[:count]]
                                        # write to file
                                        dists_timeline.write(','.join([str(dist) for dist in dists_list]) + '\n')
                                        cmps_timeline.write(','.join(str(cmps) for cmps in cmps_list) + '\n')
                                        gt_cmps_timeline.write(','.join(str(cmps) for cmps in gt_cmps) + '\n')
                    except:
                        print(f"No timeline data found in {fileroot}/{filename}.")

                mean_recall = 0
                mean_cmps = 0
                mean_lats = 0
                recall_list_for_all_results = []
                cmps_list_for_all_results = []
                lats_list_for_all_results = []
                query_id_for_all_results = []
                cmps_list_for_all_queries = []
                lats_list_for_all_queries = []
                recalls_list_for_all_queries = []
                num_queries_reached_recall_target = 0
                num_all_results = 0

                # 新增：用于存储按查询数量分组的召回率和延迟
                recall_by_count = {}  # {count: [recall_values]}
                latency_by_count = {}  # {count: [latency_values]}

                if start_times is not None and end_times is not None:
                    start_time = min(start_times)
                    end_time = max(end_times)
                    thpt = repeat * len(neighbors) / (end_time - start_time) * (10**6)
                else:
                    thpt = None

                # 合并两个循环以减少重复计算
                for query_id, (neighbors_per_query, groundtruths_per_query, cmps_per_query, lats_per_query, total_cmps_per_query, total_latency_per_query) in enumerate(zip(neighbors, groundtruths, cmps, lats, total_cmps, total_latency)):
                    neighbors_array = np.array(neighbors_per_query)
                    groundtruths_array = np.array(groundtruths_per_query)
                    cmps_array = np.array(cmps_per_query)
                    lats_array = np.array(lats_per_query)

                    is_in_groundtruth = np.isin(neighbors_array, groundtruths_array)
                    recall_per_query = np.sum(is_in_groundtruth)

                    recall = recall_per_query / len(neighbors_per_query)
                    mean_recall += recall
                    if recall >= (recall_targets_[query_id] if count == 'end2end' else recall_target):
                        num_queries_reached_recall_target += 1

                    num_all_results += len(neighbors_per_query)

                    # 第二个循环中的处理逻辑
                    true_cmps_per_query = cmps_array[is_in_groundtruth]
                    true_lats_per_query = lats_array[is_in_groundtruth]
                    
                    true_cmps_per_query.sort()
                    true_lats_per_query.sort()

                    if recall_per_query == 0:
                        continue

                    mean_cmps += total_cmps_per_query
                    mean_lats += total_latency_per_query

                    recall_list_for_all_results.extend([1] * recall_per_query)
                    cmps_list_for_all_results.extend(true_cmps_per_query)
                    lats_list_for_all_results.extend(true_lats_per_query)
                    query_id_for_all_results.extend([query_id] * recall_per_query)

                    cmps_list_for_all_queries.append(total_cmps_per_query)
                    lats_list_for_all_queries.append(total_latency_per_query)
                    recalls_list_for_all_queries.append(recall)

                    # 新增：记录每个查询的召回率和延迟，按键值（count）分组
                    if len(neighbors_per_query) not in recall_by_count:
                        recall_by_count[len(neighbors_per_query)] = []
                        latency_by_count[len(neighbors_per_query)] = []
                    
                    recall_by_count[len(neighbors_per_query)].append(recall)
                    latency_by_count[len(neighbors_per_query)].append(total_latency_per_query)

                mean_recall = mean_recall / len(neighbors)
                if thpt is not None:
                    goodput = thpt * num_queries_reached_recall_target / len(neighbors)

                mean_cmps /= len(neighbors)
                mean_lats /= len(neighbors)

                # 新增：计算每个count的平均召回率和延迟
                avg_recall_by_count = {k: float(np.mean(v)) for k, v in recall_by_count.items()}
                avg_latency_by_count = {k: float(np.mean(v)) for k, v in latency_by_count.items()}

                # 新增：计算latency CDF（前1%、前2%...前100%查询的latency值）
                recalls_list_for_all_queries_sorted = sorted(recalls_list_for_all_queries)
                recall_cdf = []
                num_queries = len(recalls_list_for_all_queries_sorted)
                if num_queries > 0:
                    for i in range(1, 101):
                        index = max(0, int(num_queries * i / 100) - 1)
                        recall_cdf.append(float(recalls_list_for_all_queries_sorted[index]))
                
                lats_list_for_all_queries_sorted = sorted(lats_list_for_all_queries)
                latency_cdf = []
                num_queries = len(lats_list_for_all_queries_sorted)
                if num_queries > 0:
                    for i in range(1, 101):
                        index = max(0, int(num_queries * i / 100) - 1)
                        latency_cdf.append(float(lats_list_for_all_queries_sorted[index]) / 1000) # ms

                if not optimal_dumped_for_this_count:

                    threshold = int(num_all_results * (np.mean(recall_targets_) if count == 'end2end' else recall_target))

                    # sort recall_list_for_all_results, cmps_list_for_all_results, lats_list_for_all_results, query_id_for_all_results by cmps
                    recall_list_for_all_results, cmps_list_for_all_results, lats_list_for_all_results, query_id_for_all_results = \
                        zip(*sorted(zip(recall_list_for_all_results, cmps_list_for_all_results, lats_list_for_all_results, query_id_for_all_results), key=lambda x: x[1]))

                    recall_list_for_all_results = recall_list_for_all_results[:threshold]
                    cmps_list_for_all_results = cmps_list_for_all_results[:threshold]
                    lats_list_for_all_results = lats_list_for_all_results[:threshold]

                    # Create dictionaries to store maximum cmps and lats for each query_id
                    max_cmps_by_query = {}
                    max_lats_by_query = {}
                    
                    # Populate dictionaries with maximum values for each query_id
                    for cmps_val, lats_val, query_id in zip(cmps_list_for_all_results, lats_list_for_all_results, query_id_for_all_results):
                        # For cmps, keep the maximum value for each query_id
                        if query_id not in max_cmps_by_query or cmps_val > max_cmps_by_query[query_id]:
                            max_cmps_by_query[query_id] = cmps_val
                        
                        # For lats, keep the maximum value for each query_id
                        if query_id not in max_lats_by_query or lats_val > max_lats_by_query[query_id]:
                            max_lats_by_query[query_id] = lats_val
                    
                    # Calculate mean of maximum values for each query_id
                    mean_cmps_optimized = np.mean(list(max_cmps_by_query.values())) if max_cmps_by_query else 0
                    mean_lats_optimized = np.mean(list(max_lats_by_query.values())) if max_lats_by_query else 0

                    mean_recall_optimized = np.sum(recall_list_for_all_results) / num_all_results

                    print(f"recall: {mean_recall:.4f} -> {mean_recall_optimized:.4f} (-{mean_recall-mean_recall_optimized:.4f})")
                    print(f"number of distance calculations: {mean_cmps:.2f} -> {mean_cmps_optimized:.2f} (-{(mean_cmps-mean_cmps_optimized)/mean_cmps*100:.2f}%)")
                    print(f"latency (us): {mean_lats:.2f} -> {mean_lats_optimized:.2f} (-{(mean_lats-mean_lats_optimized)/mean_lats*100:.2f}%)")
                    print(f"latency per distance calculation (us): {mean_lats/mean_cmps:.4f} -> {mean_lats_optimized/mean_cmps_optimized:.4f}")
                    print(f"prediction times: {mean_prediction_times:.2f}, prediction overhead (us): {mean_prediction_overhead:.2f}")
                    if thpt is not None:
                        print(f"QPS: {thpt:.2f}, Goodput: {goodput:.2f}")

                else:
                    print(f"recall: {mean_recall:.4f}")
                    print(f"number of distance calculations: {mean_cmps:.2f}")
                    print(f"latency (us): {mean_lats:.2f}")
                    print(f"latency per distance calculation (us): {mean_lats/mean_cmps:.4f}")
                    print(f"prediction times: {mean_prediction_times:.2f}, prediction overhead (us): {mean_prediction_overhead:.2f}")
                    if thpt is not None:
                        print(f"QPS: {thpt:.2f}, Goodput: {goodput:.2f}")

                preparation_run_latency = 0
                preparation_generate_latency = 0
                preparation_train_latency = 0
                if train_mode is not None:
                    for train_count in train_counts:
                        if not os.path.exists(os.path.join("models", args.dataset)):
                            continue
                        with open(os.path.join("models", args.dataset, os.path.split(args.runbook_path)[1], train_mode, str(train_count), "latency.json"), "r") as f:
                            preparation_latencies = json.load(f)
                            preparation_run_latency += preparation_latencies["run_latency"]
                            preparation_generate_latency += preparation_latencies["generate_latency"]
                            preparation_train_latency += preparation_latencies["train_latency"]

                if mode != "train":
                    result_list.append({
                        # search settings
                        "threads": threads,
                        "mode": mode,
                        "count": count,
                        "train_counts": train_counts,
                        "train_queries": int(train_queries),
                        "num_queries": Q.shape[0],
                        "traversal_window_size": traversal_window_size,
                        "do_not_use_dynamic_recall_threshold": do_not_use_dynamic_recall_threshold,
                        "do_not_predict_average_recall": do_not_predict_average_recall,
                        "do_not_use_prediction_intervals": do_not_use_prediction_intervals,
                        # recall
                        "recall": mean_recall,
                        "recall_target": recall_target,
                        "recall_groupby_count": avg_recall_by_count,
                        "recall_cdf": recall_cdf,
                        # performance metrics
                        "mean_cmps": mean_cmps,
                        # "p50_cmps": np.percentile(cmps_list_for_all_queries, 50),
                        # "p90_cmps": np.percentile(cmps_list_for_all_queries, 90),
                        # "p95_cmps": np.percentile(cmps_list_for_all_queries, 95),
                        # "p99_cmps": np.percentile(cmps_list_for_all_queries, 99),
                        "mean_latency": mean_lats,
                        # "p50_latency": np.percentile(lats_list_for_all_queries, 50),
                        # "p90_latency": np.percentile(lats_list_for_all_queries, 90),
                        # "p95_latency": np.percentile(lats_list_for_all_queries, 95),
                        # "p99_latency": np.percentile(lats_list_for_all_queries, 99),
                        "latency_cdf": latency_cdf,
                        "latency_groupby_count": avg_latency_by_count,
                        "thpt": thpt,
                        "goodput": goodput,
                        # prediction overhead metrics
                        "prediction_times": mean_prediction_times,
                        "prediction_overhead": mean_prediction_overhead,
                        # model preparation time
                        "run_latency": preparation_run_latency,
                        "generate_latency": preparation_generate_latency,
                        "train_latency": preparation_train_latency,
                        "total_preparation_latency": preparation_run_latency + preparation_generate_latency + preparation_train_latency,
                    })

                    if not optimal_dumped_for_this_count:
                        optimal_dumped_for_this_count = True
                        result_list.append({
                            # search settings
                            "threads": threads,
                            "mode": "optimal",
                            "count": count,
                            "train_counts": train_counts,
                            "train_queries": int(train_queries),
                            "num_queries": Q.shape[0],
                            "traversal_window_size": traversal_window_size,
                            "do_not_use_dynamic_recall_threshold": do_not_use_dynamic_recall_threshold,
                            "do_not_predict_average_recall": do_not_predict_average_recall,
                            "do_not_use_prediction_intervals": do_not_use_prediction_intervals,
                            # recall
                            "recall": mean_recall_optimized,
                            "recall_target": recall_target,
                            "recall_groupby_count": avg_recall_by_count,
                            "recall_cdf": recall_cdf,
                            # performance metrics
                            "mean_cmps": mean_cmps_optimized,
                            "mean_latency": mean_lats_optimized,
                            "latency_cdf": latency_cdf,
                            "latency_groupby_count": avg_latency_by_count,
                            "thpt": thpt,
                            "goodput": goodput,
                            # prediction overhead metrics
                            "prediction_times": mean_prediction_times,
                            "prediction_overhead": mean_prediction_overhead,
                            # model preparation time
                            "run_latency": preparation_run_latency,
                            "generate_latency": preparation_generate_latency,
                            "train_latency": preparation_train_latency,
                            "total_preparation_latency": preparation_run_latency + preparation_generate_latency + preparation_train_latency,
                        })
            
            print("")

    if args.dump_json:
        if not os.path.exists(os.path.join("plot", args.results_base_path)):
            os.makedirs(os.path.join("plot", args.results_base_path))
        with open(os.path.join("plot", args.results_base_path, f"{args.dataset}.json"), "w") as f:
            valid_counts = sorted(list(valid_counts))
            recall_targets = sorted(list(recall_targets))
            json.dump({
                "counts": valid_counts,
                "recall_targets": recall_targets,
                "microbench_results": [result for result in result_list if result["count"] != "end2end"],
                "end2end_results": [result for result in result_list if result["count"] == "end2end"],
            }, f, indent=4)