Mini-Nav/tests/test_sw_retrieval_benchmark.py

from __future__ import annotations

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

import numpy as np
import pytest


def load_sw_benchmark():
    path = Path(__file__).resolve().parent.parent / "scripts" / "sw_retrieval_benchmark.py"
    spec = importlib.util.spec_from_file_location("sw_retrieval_benchmark", path)
    assert spec is not None, f"Could not find spec for {path}"
    assert spec.loader is not None, f"No loader available for {path}"
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def _write_dataset(path: Path) -> None:
    np.savez_compressed(
        path,
        rows_words=np.array(
            [
                [0b11110000],
                [0b11100000],
                [0b00001111],
            ],
            dtype=np.uint64,
        ),
        row_labels=np.array([0, 0, 1], dtype=np.int64),
        queries_words=np.array(
            [
                [0b11110000],
                [0b00001111],
            ],
            dtype=np.uint64,
        ),
        query_labels=np.array([0, 1], dtype=np.int64),
    )


def test_load_dataset_preserves_little_endian_words(tmp_path):
    bench = load_sw_benchmark()
    dataset_path = tmp_path / "dataset.npz"
    np.savez_compressed(
        dataset_path,
        rows_words=np.array([[0x2, 0x1]], dtype=np.uint64),
        row_labels=np.array([7], dtype=np.int64),
        queries_words=np.array([[0x4, 0x3]], dtype=np.uint64),
        query_labels=np.array([8], dtype=np.int64),
    )

    dataset = bench.load_retrieval_dataset_npz(dataset_path)

    assert dataset.rows == [0x1_0000_0000_0000_0002]
    assert dataset.queries == [0x3_0000_0000_0000_0004]
    assert dataset.row_labels == [7]
    assert dataset.query_labels == [8]


def test_run_benchmark_reports_quality_and_query_speed(tmp_path):
    bench = load_sw_benchmark()
    dataset_path = tmp_path / "dataset.npz"
    _write_dataset(dataset_path)
    timer_values = iter([1_000, 3_000])

    result = bench.run_benchmark(
        dataset_path,
        hash_bits=64,
        topk_values=(1, 2),
        run_id="unit-test",
        timer_ns=lambda: next(timer_values),
    )

    assert result["mode"] == "software-numpy"
    assert result["status"] == "pass"
    assert result["params"]["engine"] == "numpy"
    assert result["params"]["workers"] == 1
    assert result["dataset"]["num_queries"] == 2
    assert result["params"]["num_rows"] == 3
    assert result["params"]["topk_k"] == 2
    assert result["metrics"]["1"]["exact_match_rate"] == 1.0
    assert result["metrics"]["1"]["recall@k"] == 1.0
    assert result["performance"]["total_elapsed_ns"] == 2_000
    assert result["performance"]["ns_per_query"] == 1_000.0
    assert result["performance"]["queries_per_second"] == 1_000_000.0


def test_numpy_topk_matches_reference_with_tiebreak(tmp_path):
    bench = load_sw_benchmark()
    dataset_path = tmp_path / "dataset.npz"
    np.savez_compressed(
        dataset_path,
        rows_words=np.array(
            [
                [0b11110000],
                [0b11110000],
                [0b11100000],
                [0b00001111],
            ],
            dtype=np.uint64,
        ),
        row_labels=np.array([0, 0, 0, 1], dtype=np.int64),
        queries_words=np.array([[0b11110000]], dtype=np.uint64),
        query_labels=np.array([0], dtype=np.int64),
    )
    dataset = bench.load_retrieval_dataset_npz(dataset_path)

    assert bench.match_topk_numpy(dataset.queries_words[0], dataset.rows_words, width=64, k=3) == [0, 1, 2]


def test_numpy_batch_topk_vectorizes_across_queries(monkeypatch):
    bench = load_sw_benchmark()
    rows_words = np.array(
        [
            [0b11110000],
            [0b11100000],
            [0b00001111],
        ],
        dtype=np.uint64,
    )
    queries_words = np.array(
        [
            [0b11110000],
            [0b00001111],
        ],
        dtype=np.uint64,
    )

    monkeypatch.setattr(
        bench,
        "match_topk_numpy",
        lambda *_args, **_kwargs: (_ for _ in ()).throw(AssertionError("scalar path called")),
    )

    assert bench._match_topk_numpy_batch(queries_words, rows_words, width=64, k=2) == [
        [0, 1],
        [2, 1],
    ]


def test_run_benchmark_exact_match_compares_against_reference(tmp_path, monkeypatch):
    bench = load_sw_benchmark()
    dataset_path = tmp_path / "dataset.npz"
    _write_dataset(dataset_path)

    monkeypatch.setattr(bench, "ref_match_topk", lambda query, rows, width, k: ([1, 0], []))

    result = bench.run_benchmark(
        dataset_path,
        hash_bits=64,
        topk_values=(1,),
        timer_ns=lambda: 0,
    )

    assert result["metrics"]["1"]["exact_match_rate"] == 0.0


def test_run_benchmark_threaded_numpy_matches_single_worker(tmp_path):
    bench = load_sw_benchmark()
    dataset_path = tmp_path / "dataset.npz"
    _write_dataset(dataset_path)

    single = bench.run_benchmark(
        dataset_path,
        hash_bits=64,
        topk_values=(1, 2),
        run_id="single",
        workers=1,
        timer_ns=lambda: 0,
    )
    threaded = bench.run_benchmark(
        dataset_path,
        hash_bits=64,
        topk_values=(1, 2),
        run_id="threaded",
        workers=2,
        timer_ns=lambda: 0,
    )

    assert single["mode"] == "software-numpy"
    assert threaded["mode"] == "software-numpy"
    assert single["params"]["engine"] == "numpy"
    assert threaded["params"]["workers"] == 2
    assert threaded["metrics"] == single["metrics"]


def test_run_benchmark_rejects_hash_bits_that_do_not_match_npz_width(tmp_path):
    bench = load_sw_benchmark()
    dataset_path = tmp_path / "dataset.npz"
    _write_dataset(dataset_path)

    with pytest.raises(ValueError, match="hash_bits"):
        bench.run_benchmark(dataset_path, hash_bits=8)


def test_cli_missing_dataset_prints_generation_hint_without_traceback(monkeypatch, capsys):
    bench = load_sw_benchmark()
    missing_path = Path("outputs/cam_retrieval_benchmark/datasets/cifar100_hash512_rows512_queries8192.npz")
    monkeypatch.setattr(
        sys,
        "argv",
        [
            "sw_retrieval_benchmark.py",
            "--dataset",
            str(missing_path),
            "--hash-bits",
            "512",
            "--topk-k",
            "5",
            "--workers",
            "1",
        ],
    )

    with pytest.raises(SystemExit) as exc_info:
        bench.main()

    assert exc_info.value.code == 2
    captured = capsys.readouterr()
    assert "retrieval dataset not found" in captured.err
    assert "python scripts/prepare_cam_retrieval_dataset.py" in captured.err
    assert "--dataset cifar100" in captured.err
    assert "--num-rows 512" in captured.err
    assert "--max-queries 8192" in captured.err
    assert "--hash-bits 512" in captured.err
    assert "Traceback" not in captured.err


def test_write_outputs_includes_quality_and_performance_csv(tmp_path):
    bench = load_sw_benchmark()
    result = {
        "run_id": "unit-test",
        "mode": "software-numpy",
        "status": "pass",
        "params": {"num_rows": 3, "hash_bits": 64, "topk_k": 2, "workers": 2, "engine": "numpy"},
        "dataset": {"num_classes": 2, "num_queries": 2},
        "metrics": {
            "1": {
                "macro_precision": 1.0,
                "retrieval_recall": 0.75,
                "macro_f1": 0.85,
                "exact_match_rate": 1.0,
                "recall@k": 1.0,
            }
        },
        "performance": {
            "total_elapsed_ns": 2_000,
            "total_elapsed_sec": 0.000002,
            "ns_per_query": 1_000.0,
            "queries_per_second": 1_000_000.0,
        },
    }

    bench.write_outputs(tmp_path, result)

    metrics_json = json.loads((tmp_path / "metrics.json").read_text(encoding="utf-8"))
    metrics_csv = (tmp_path / "metrics.csv").read_text(encoding="utf-8")
    summary_md = (tmp_path / "summary.md").read_text(encoding="utf-8")

    assert metrics_json["performance"]["queries_per_second"] == 1_000_000.0
    assert "queries_per_second" in metrics_csv
    assert "1000000.0" in metrics_csv
    assert "Software CAM Retrieval Benchmark Summary" in summary_md
    assert "software-numpy" in metrics_csv
    assert "workers" in metrics_csv
    assert "workers: `2`" in summary_md
    assert "queries_per_second" in summary_md