Mini-Nav/tests/test_retrieval_benchmark_performance.py

from __future__ import annotations

import csv
import importlib.util
import sys
from pathlib import Path


HW_SIM_DIR = Path(__file__).resolve().parents[1] / "hw" / "sim"
if str(HW_SIM_DIR) in sys.path:
    sys.path.remove(str(HW_SIM_DIR))
sys.path.insert(0, str(HW_SIM_DIR))
benchmark_path = HW_SIM_DIR / "benchmarks" / "retrieval" / "test_retrieval_benchmark.py"
spec = importlib.util.spec_from_file_location("hw_retrieval_benchmark", benchmark_path)
assert spec is not None
assert spec.loader is not None
hw_retrieval_benchmark = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = hw_retrieval_benchmark
spec.loader.exec_module(hw_retrieval_benchmark)

QueryTiming = hw_retrieval_benchmark.QueryTiming
build_hardware_performance = hw_retrieval_benchmark.build_hardware_performance
summarize_query_timings = hw_retrieval_benchmark.summarize_query_timings
write_outputs = hw_retrieval_benchmark.write_outputs


def test_summarize_query_timings_uses_query_only_accept_to_last_cycles() -> None:
    summary = summarize_query_timings([
        QueryTiming(
            accept_to_first_result_cycles=10,
            accept_to_last_result_cycles=14,
            total_query_cycles=14,
        ),
        QueryTiming(
            accept_to_first_result_cycles=12,
            accept_to_last_result_cycles=18,
            total_query_cycles=18,
        ),
    ])

    assert summary == {
        "num_queries": 2,
        "query_only_total_cycles": 32,
        "query_only_cycles_per_query": 16.0,
        "query_only_min_cycles": 14,
        "query_only_max_cycles": 18,
        "query_only_queries_per_cycle": 2 / 32,
        "mean_accept_to_first_result_cycles": 11.0,
        "mean_accept_to_last_result_cycles": 16.0,
        "total_query_cycles": 32,
        "mean_total_query_cycles": 16.0,
        "min_total_query_cycles": 14,
        "max_total_query_cycles": 18,
        "cycles_per_query": 16.0,
        "queries_per_cycle": 2 / 32,
    }


def test_build_hardware_performance_separates_query_and_end_to_end_cycles() -> None:
    performance = build_hardware_performance(
        [
            QueryTiming(
                accept_to_first_result_cycles=10,
                accept_to_last_result_cycles=14,
                total_query_cycles=14,
            ),
            QueryTiming(
                accept_to_first_result_cycles=12,
                accept_to_last_result_cycles=18,
                total_query_cycles=18,
            ),
        ],
        load_write_noise_cycles=100,
        end_to_end_cycles=140,
    )

    assert performance["query_only_total_cycles"] == 32
    assert performance["query_only_cycles_per_query"] == 16.0
    assert performance["query_only_queries_per_cycle"] == 2 / 32
    assert performance["load_write_noise_cycles"] == 100
    assert performance["end_to_end_cycles"] == 140
    assert performance["end_to_end_cycles_per_query"] == 70.0
    assert performance["end_to_end_queries_per_cycle"] == 2 / 140
    assert performance["cycles_per_query"] == performance["query_only_cycles_per_query"]
    assert performance["queries_per_cycle"] == performance["query_only_queries_per_cycle"]


def test_write_outputs_includes_hardware_performance_fields(tmp_path: Path) -> None:
    result = {
        "run_id": "test-run",
        "mode": "no_noise",
        "status": "pass",
        "params": {
            "num_rows": 512,
            "hash_bits": 512,
            "lanes": 8,
            "topk_k": 5,
            "write_noise_en": 0,
            "write_noise_rate_num": 0,
            "write_noise_rate_den": 100,
        },
        "dataset": {
            "num_classes": 10,
            "positives_per_class": 0,
            "queries_per_class": 0,
            "num_queries": 2,
            "seed": 0,
        },
        "metrics": {
            "1": {
                "macro_precision": 1.0,
                "retrieval_recall": 0.5,
                "macro_f1": 2 / 3,
                "recall@k": 1.0,
                "exact_match_rate": 1.0,
            }
        },
        "performance": {
            "num_queries": 2,
            "query_only_total_cycles": 32,
            "query_only_cycles_per_query": 16.0,
            "query_only_min_cycles": 14,
            "query_only_max_cycles": 18,
            "query_only_queries_per_cycle": 2 / 32,
            "mean_accept_to_first_result_cycles": 11.0,
            "mean_accept_to_last_result_cycles": 16.0,
            "load_write_noise_cycles": 100,
            "end_to_end_cycles": 140,
            "end_to_end_cycles_per_query": 70.0,
            "end_to_end_queries_per_cycle": 2 / 140,
            "total_query_cycles": 32,
            "mean_total_query_cycles": 16.0,
            "min_total_query_cycles": 14,
            "max_total_query_cycles": 18,
            "cycles_per_query": 16.0,
            "queries_per_cycle": 2 / 32,
        },
    }

    write_outputs(tmp_path, result)

    with (tmp_path / "metrics.csv").open(newline="", encoding="utf-8") as f:
        row = next(csv.DictReader(f))

    assert row["query_only_cycles_per_query"] == "16.0"
    assert row["query_only_total_cycles"] == "32"
    assert row["query_only_queries_per_cycle"] == str(2 / 32)
    assert row["load_write_noise_cycles"] == "100"
    assert row["end_to_end_cycles"] == "140"
    assert row["end_to_end_queries_per_cycle"] == str(2 / 140)
    assert row["cycles_per_query"] == "16.0"
    assert row["mean_accept_to_first_result_cycles"] == "11.0"
    assert row["mean_accept_to_last_result_cycles"] == "16.0"
    assert row["queries_per_cycle"] == str(2 / 32)

    summary = (tmp_path / "summary.md").read_text(encoding="utf-8")
    assert "## Hardware performance" in summary
    assert "query-only cycles/query" in summary
    assert "load/write/noise cycles" in summary
    assert "end-to-end cycles/query" in summary
    assert "accept_to_last_result_cycles" in summary