Micro-Benchmark of KMeans for FAISS and cuVS

kmeans_benchmark.py

#!/usr/bin/env python3
"""
Benchmark script for comparing faiss kmeans (CPU/GPU) and cuvs kmeans
on a 1Mx1024 dataset.
"""

import argparse
from pathlib import Path
import sys
import time
from collections import namedtuple
from typing import Optional, Dict, Any

# Try importing required libraries with graceful error handling
try:
    import numpy as np
except ImportError:
    print("Error: numpy is required. Please install it: pip install numpy")
    sys.exit(1)

try:
    import cupy as cp
    CUPY_AVAILABLE = True
except ImportError:
    print("Warning: cupy not available. cuvs and faiss-gpu benchmarks will be skipped.")
    CUPY_AVAILABLE = False

try:
    import faiss
    FAISS_AVAILABLE = True
    # Check if GPU support is available
    try:
        FAISS_GPU_AVAILABLE = faiss.get_num_gpus() > 0
    except:
        FAISS_GPU_AVAILABLE = False
except ImportError:
    print("Warning: faiss not available. faiss benchmarks will be skipped.")
    FAISS_AVAILABLE = False
    FAISS_GPU_AVAILABLE = False

try:
    from cuvs.cluster import kmeans
    CUVS_AVAILABLE = True
except ImportError:
    print("Warning: cuvs not available. cuvs benchmarks will be skipped.")
    CUVS_AVAILABLE = False

try:
    import h5py
    H5PY_AVAILABLE = True
except ImportError:
    print("Warning: h5py not available. GIST dataset loading will fail.")
    H5PY_AVAILABLE = False

# Benchmark result structure
BenchmarkResult = namedtuple(
    'BenchmarkResult',
    [
        'name',
        'time_ms',
        'inertia',
        'n_iter',
        'success',
        'init_method',
        'oversampling_factor'
    ]
)


def generate_dataset(dataset_path: str = "gist-960-euclidean.hdf5",
                     dtype: np.dtype = np.float32) -> np.ndarray:
    """
    Load the GIST dataset for benchmarking.

    Parameters
    ----------
    dataset_path : str
        Path to the local GIST HDF5 file. The benchmark uses the "train" split.
    dtype : np.dtype
        Data type (default: float32)

    Returns
    -------
    np.ndarray
        Dataset loaded from the HDF5 "train" key
    """
    if not H5PY_AVAILABLE:
        raise RuntimeError(
            "h5py is required to load the GIST dataset. "
            "Please install it: pip install h5py"
        )

    dataset_file = Path(dataset_path)
    if not dataset_file.is_file():
        dataset_file = Path(__file__).resolve().parent / dataset_path

    if not dataset_file.is_file():
        raise FileNotFoundError(
            f"GIST dataset file not found: {dataset_path}. "
            "Expected a local HDF5 file containing a 'train' dataset."
        )

    with h5py.File(dataset_file, "r") as hdf5_file:
        if "train" not in hdf5_file:
            raise KeyError(
                f"Dataset key 'train' not found in {dataset_file}."
            )
        data = hdf5_file["train"][:]

    return np.ascontiguousarray(data, dtype=dtype)


def benchmark_faiss_cpu(data: np.ndarray,
                        n_clusters: int = 100,
                        max_iter: int = 300,
                        n_init: int = 1,
                        seed: int = 42) -> BenchmarkResult:
    """
    Benchmark faiss kmeans on CPU.

    Parameters
    ----------
    data : np.ndarray
        Input data of shape (n_samples, n_features)
    n_clusters : int
        Number of clusters
    max_iter : int
        Maximum number of iterations
    n_init : int
        Number of initializations
    seed : int
        Random seed

    Returns
    -------
    BenchmarkResult
        Benchmark results
    """
    if not FAISS_AVAILABLE:
        return BenchmarkResult("faiss-cpu", 0.0, 0.0, 0, False, None, None)

    try:
        # Prepare data for faiss (must be contiguous and float32)
        data_faiss = np.ascontiguousarray(data, dtype=np.float32)
        n_samples, n_features = data_faiss.shape

        # Create kmeans object
        kmeans = faiss.Kmeans(
            n_features,
            n_clusters,
            niter=max_iter,
            nredo=n_init,
            seed=seed,
            verbose=False,
            spherical=False,
            update_index=True,
            gpu=False
        )

        # Run kmeans
        start_time = time.perf_counter()
        kmeans.train(data_faiss)
        end_time = time.perf_counter()

        # Get results
        execution_time_ms = (end_time - start_time) * 1000.0
        inertia = float(kmeans.obj[-1]) if len(kmeans.obj) > 0 else 0.0
        n_iter = len(kmeans.obj) if len(kmeans.obj) > 0 else 0

        return BenchmarkResult(
            name="faiss-cpu",
            time_ms=execution_time_ms,
            inertia=inertia,
            n_iter=n_iter,
            success=True,
            init_method=None,
            oversampling_factor=None
        )
    except Exception as e:
        print(f"Error in faiss-cpu benchmark: {e}")
        return BenchmarkResult("faiss-cpu", 0.0, 0.0, 0, False, None, None)


def benchmark_faiss_gpu(data: np.ndarray,
                         n_clusters: int = 100,
                         max_iter: int = 300,
                         n_init: int = 1,
                         seed: int = 42) -> BenchmarkResult:
    """
    Benchmark faiss kmeans on GPU.

    Parameters
    ----------
    data : np.ndarray
        Input data of shape (n_samples, n_features)
    n_clusters : int
        Number of clusters
    max_iter : int
        Maximum number of iterations
    n_init : int
        Number of initializations
    seed : int
        Random seed

    Returns
    -------
    BenchmarkResult
        Benchmark results
    """
    if not FAISS_AVAILABLE or not FAISS_GPU_AVAILABLE:
        return BenchmarkResult("faiss-gpu", 0.0, 0.0, 0, False, None, None)

    try:
        # Prepare data for faiss (must be contiguous and float32)
        data_faiss = np.ascontiguousarray(data, dtype=np.float32)
        n_samples, n_features = data_faiss.shape

        # Create kmeans object for GPU
        kmeans = faiss.Kmeans(
            n_features,
            n_clusters,
            niter=max_iter,
            nredo=n_init,
            seed=seed,
            verbose=False,
            spherical=False,
            update_index=True,
            gpu=True
        )

        # Run kmeans
        start_time = time.perf_counter()
        kmeans.train(data_faiss)
        end_time = time.perf_counter()

        # Get results
        execution_time_ms = (end_time - start_time) * 1000.0
        inertia = float(kmeans.obj[-1]) if len(kmeans.obj) > 0 else 0.0
        n_iter = len(kmeans.obj) if len(kmeans.obj) > 0 else 0

        return BenchmarkResult(
            name="faiss-gpu",
            time_ms=execution_time_ms,
            inertia=inertia,
            n_iter=n_iter,
            success=True,
            init_method=None,
            oversampling_factor=None
        )
    except Exception as e:
        print(f"Error in faiss-gpu benchmark: {e}")
        return BenchmarkResult("faiss-gpu", 0.0, 0.0, 0, False, None, None)


def benchmark_cuvs(data: np.ndarray,
                   n_clusters: int = 100,
                   max_iter: int = 300,
                   seed: int = 42,
                   init_method: str = "KMeansPlusPlus",
                   oversampling_factor: Optional[float] = None) -> BenchmarkResult:
    """
    Benchmark cuvs kmeans on GPU.

    Parameters
    ----------
    data : np.ndarray
        Input data of shape (n_samples, n_features)
    n_clusters : int
        Number of clusters
    max_iter : int
        Maximum number of iterations
    seed : int
        Random seed

    Returns
    -------
    BenchmarkResult
        Benchmark results
    """
    run_name = (
        f"cuvs-{init_method}-of{oversampling_factor:g}"
        if oversampling_factor is not None
        else f"cuvs-{init_method}"
    )

    if not CUVS_AVAILABLE or not CUPY_AVAILABLE:
        return BenchmarkResult(
            run_name, 0.0, 0.0, 0, False, init_method, oversampling_factor
        )

    try:
        # Convert to cupy array (GPU)
        data_gpu = cp.asarray(data, dtype=cp.float32)

        # Create kmeans parameters
        params_kwargs: Dict[str, Any] = {
            "n_clusters": n_clusters,
            "max_iter": max_iter,
            "init_method": init_method,
            "tol": 1e-4
        }
        if oversampling_factor is not None:
            params_kwargs["oversampling_factor"] = oversampling_factor
        params = kmeans.KMeansParams(**params_kwargs)

        start_time = time.perf_counter()
        centroids, inertia, n_iter = kmeans.fit(params, data_gpu)
        end_time = time.perf_counter()

        # Get results
        execution_time_ms = (end_time - start_time) * 1000.0
        inertia_value = float(inertia) if inertia is not None else 0.0
        n_iter_value = int(n_iter) if n_iter is not None else 0

        return BenchmarkResult(
            name=run_name,
            time_ms=execution_time_ms,
            inertia=inertia_value,
            n_iter=n_iter_value,
            success=True,
            init_method=init_method,
            oversampling_factor=oversampling_factor
        )
    except Exception as e:
        print(f"Error in cuvs benchmark: {e}")
        return BenchmarkResult(
            run_name, 0.0, 0.0, 0, False, init_method, oversampling_factor
        )


def print_results(results: list[BenchmarkResult]):
    """
    Print benchmark results in a formatted table.

    Parameters
    ----------
    results : list[BenchmarkResult]
        List of benchmark results
    """
    print("\n" + "=" * 110)
    print("KMeans Benchmark Results")
    print("=" * 110)
    print(f"{'Implementation':<28} {'Init Method':<16} {'Over-factor':<12} "
          f"{'Time (ms)':<12} {'Inertia':<15} {'Iterations':<12} {'Status':<10}")
    print("-" * 110)

    for result in results:
        init_method_str = result.init_method if result.init_method is not None else "N/A"
        oversampling_str = (
            f"{result.oversampling_factor:g}"
            if result.oversampling_factor is not None
            else "N/A"
        )

        if result.success:
            status = "Success"
            time_str = f"{result.time_ms:.2f}"
            inertia_str = f"{result.inertia:.2e}" if result.inertia > 0 else "N/A"
            iter_str = str(result.n_iter)
        else:
            status = "Failed"
            time_str = "N/A"
            inertia_str = "N/A"
            iter_str = "N/A"

        print(f"{result.name:<28} {init_method_str:<16} {oversampling_str:<12} "
              f"{time_str:<12} {inertia_str:<15} {iter_str:<12} {status:<10}")

    print("=" * 110)

    # Print speedup comparison
    successful_results = [r for r in results if r.success and r.time_ms > 0]
    if len(successful_results) > 1:
        print("\nSpeedup Comparison (relative to slowest):")
        print("-" * 110)
        slowest_time = max(r.time_ms for r in successful_results)
        for result in successful_results:
            speedup = slowest_time / result.time_ms
            print(f"{result.name:<20} {speedup:.2f}x")
        print("=" * 110)


def main():
    """Main function to run all benchmarks."""
    parser = argparse.ArgumentParser(
        description="Benchmark faiss and cuvs kmeans on local GIST dataset"
    )
    parser.add_argument(
        "--dataset-path",
        type=str,
        default="gist-960-euclidean.hdf5",
        help="Path to local GIST HDF5 dataset (default: gist-960-euclidean.hdf5)"
    )
    parser.add_argument(
        "--n-clusters",
        type=int,
        default=100,
        help="Number of clusters (default: 100)"
    )
    parser.add_argument(
        "--max-iter",
        type=int,
        default=300,
        help="Maximum number of iterations (default: 300)"
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed (default: 42)"
    )
    parser.add_argument(
        "--skip-faiss-cpu",
        action="store_true",
        help="Skip faiss CPU benchmark"
    )
    parser.add_argument(
        "--skip-faiss-gpu",
        action="store_true",
        help="Skip faiss GPU benchmark"
    )
    parser.add_argument(
        "--skip-cuvs",
        action="store_true",
        help="Skip cuvs benchmark"
    )

    args = parser.parse_args()

    print("KMeans Benchmark Script")
    print("=" * 80)
    print(f"Dataset file: {args.dataset_path}")
    print(f"Clusters: {args.n_clusters}")
    print(f"Max iterations: {args.max_iter}")
    print(f"Random seed: {args.seed}")
    print("=" * 80)

    # Load dataset
    print("\nLoading dataset...")
    data = generate_dataset(
        dataset_path=args.dataset_path
    )
    print(f"Dataset shape: {data.shape}, dtype: {data.dtype}")

    results = []

    # Run benchmarks
    if not args.skip_faiss_cpu and FAISS_AVAILABLE:
        print("\nRunning faiss-cpu benchmark...")
        result = benchmark_faiss_cpu(
            data,
            n_clusters=args.n_clusters,
            max_iter=args.max_iter,
            seed=args.seed
        )
        results.append(result)

    if not args.skip_faiss_gpu and FAISS_AVAILABLE and FAISS_GPU_AVAILABLE:
        print("\nRunning faiss-gpu benchmark...")
        result = benchmark_faiss_gpu(
            data,
            n_clusters=args.n_clusters,
            max_iter=args.max_iter,
            seed=args.seed
        )
        results.append(result)

    if not args.skip_cuvs and CUVS_AVAILABLE and CUPY_AVAILABLE:
        cuvs_configs = [
            ("Random", None),
            ("KMeansPlusPlus", 0.0),
            ("KMeansPlusPlus", 0.5),
            ("KMeansPlusPlus", 0.75),
            ("KMeansPlusPlus", 1.0),
            ("KMeansPlusPlus", 1.25),
            ("KMeansPlusPlus", 2.0),
            ("KMeansPlusPlus", 6.0)
        ]
        for init_method, oversampling_factor in cuvs_configs:
            oversampling_msg = (
                f", oversampling_factor={oversampling_factor:g}"
                if oversampling_factor is not None
                else ""
            )
            print(
                f"\nRunning cuvs benchmark "
                f"(init_method={init_method}{oversampling_msg})..."
            )
            result = benchmark_cuvs(
                data,
                n_clusters=args.n_clusters,
                max_iter=args.max_iter,
                seed=args.seed,
                init_method=init_method,
                oversampling_factor=oversampling_factor
            )
            results.append(result)

    # Print results
    print_results(results)

    # Check if any benchmarks succeeded
    if not any(r.success for r in results):
        print("\nError: No benchmarks completed successfully.")
        sys.exit(1)


if __name__ == "__main__":
    main()