run.py

"""
run.py — VECTOR-ADIS simplified benchmark runner

Usage:
  python run.py --engines "qdrant-single-node,milvus-single-node" --datasets "glove-100-angular,h-and-m-2048-angular-filters"

This script orchestrates local benchmark runs for Qdrant and Milvus engines,
measuring ingestion and search performance under consistent dataset configurations.
Results are stored as structured JSON files under /results/.
"""

import glob
import hashlib
import json
import os
import signal
import traceback
from contextlib import contextmanager
from datetime import datetime, timezone
from pathlib import Path
from typing import Generator, List

import typer


class TimeoutException(Exception):
    """Raised when a timeout occurs."""

    pass


@contextmanager
def run_with_timeout(seconds: int) -> Generator[None, None, None]:
    """Context manager for timing out operations using SIGALRM (Unix only)."""

    def handler(signum, frame):
        raise TimeoutException(f"Operation timed out after {seconds} seconds")

    old_handler = signal.signal(signal.SIGALRM, handler)
    signal.alarm(seconds)
    try:
        yield
    finally:
        signal.alarm(0)
        signal.signal(signal.SIGALRM, old_handler)


from benchmark.config_read import read_dataset_config, read_engine_configs
from benchmark.dataset import Dataset
from engine.base_client import IncompatibilityError
from engine.clients.client_factory import ClientFactory

DEFAULT_ENGINES = ",".join(["qdrant-single-node", "milvus-single-node"])
DEFAULT_DATASETS = ",".join(
    [
        "random-100",
        "random-match-keyword-100-angular",
        "glove-100-angular",
        "deep-image-96-angular",
        "arxiv-titles-384-angular",
        "cohere-wiki-100k",
        "laion-small-clip",
        "dbpedia-openai-100K-1536-angular",
        "dbpedia-openai-1M-1536-angular",
        "h-and-m-2048-angular",
        "msmarco-sparse",
    ]
)

app = typer.Typer(help="Run local Qdrant vs Milvus benchmarks with JSON outputs")


def _split_csv(items: str) -> List[str]:
    """Convert a comma-separated option string into a list of identifiers."""

    return [entry.strip() for entry in items.split(",") if entry.strip()]


def _get_upload_key(engine_config: dict, dataset_name: str) -> str:
    """
    Generate a unique key for upload configuration.

    Uploads with the same engine, dataset, collection_params, and upload_params
    (index configuration) can share data - only search_params differ.
    """
    key_parts = {
        "engine": engine_config.get("engine"),
        "dataset": dataset_name,
        "collection_params": engine_config.get("collection_params", {}),
        "upload_params": engine_config.get("upload_params", {}),
    }
    key_str = json.dumps(key_parts, sort_keys=True)
    return hashlib.md5(key_str.encode()).hexdigest()[:12]


def _execute_run(
    *,
    engines: List[str],
    datasets: List[str],
    host: str,
    timeout: float,
    skip_upload: bool,
    skip_search: bool,
    force: bool,
):
    """Shared implementation for the legacy default command and the explicit `run` command."""

    configs_by_name, engine_manifests = read_engine_configs()
    all_datasets = read_dataset_config()

    # Prepare result directory
    os.makedirs("results", exist_ok=True)

    # Track which (engine, dataset, index_config) combos have been uploaded
    # to avoid redundant uploads when only search params differ
    uploaded_configs: set[str] = set()

    for engine_name in engines:
        if engine_name in engine_manifests:
            engine_configs = engine_manifests[engine_name]
        elif engine_name in configs_by_name:
            engine_configs = [configs_by_name[engine_name]]
        else:
            print(f"[WARN] Engine {engine_name} not found in configurations, skipping.")
            continue

        for engine_config in engine_configs:
            experiment_name = engine_config["name"]

            for dataset_name in datasets:
                if dataset_name not in all_datasets:
                    print(f"[WARN] Dataset {dataset_name} not found in registry, skipping.")
                    continue

                dataset_config = all_datasets[dataset_name]

                upload_key = _get_upload_key(engine_config, dataset_name)
                should_skip_upload = skip_upload or (upload_key in uploaded_configs)

                if should_skip_upload and upload_key in uploaded_configs:
                    print(
                        f"\n🔄 Running experiment: {experiment_name} × {dataset_name} (reusing existing upload)"
                    )
                else:
                    print(f"\n🚀 Running experiment: {experiment_name} × {dataset_name}")

                client = ClientFactory(host).build_client(engine_config)

                try:
                    dataset = Dataset(dataset_config)
                    dataset.download()

                    # Skip if incompatible (e.g. sparse dataset not supported)
                    if getattr(dataset.config, "type", "") == "sparse" and not getattr(
                        client, "sparse_vector_support", False
                    ):
                        raise IncompatibilityError(f"{client.name} does not support sparse vectors")

                    with run_with_timeout(int(timeout)):
                        metrics = client.run_experiment(
                            dataset=dataset,
                            skip_upload=should_skip_upload,
                            skip_search=skip_search,
                            skip_if_exists=not force,
                        )

                    uploaded_configs.add(upload_key)

                except IncompatibilityError as e:
                    print(
                        f"[SKIP] {experiment_name} × {dataset_name}: incompatible parameters → {e}"
                    )
                    continue
                except TimeoutException:
                    print(f"[TIMEOUT] {experiment_name} × {dataset_name} exceeded {timeout}s")
                    client.delete_client()
                    continue
                except Exception as e:
                    print(f"[ERROR] Experiment {experiment_name} × {dataset_name} failed → {e}")
                    traceback.print_exc()
                    continue

                client.delete_client()

                result_path = os.path.join(
                    "results",
                    f"{experiment_name}_{dataset_name}.json",
                )

                results = {
                    "manifest": engine_name,
                    "experiment": experiment_name,
                    "engine": engine_config.get("engine"),
                    "dataset": dataset_name,
                    "timestamp": datetime.now(timezone.utc).isoformat(),
                    **(metrics or {}),
                }

                with open(result_path, "w") as f:
                    json.dump(results, f, indent=2)

                print(f"✅ Results saved → {result_path}")


@app.callback(invoke_without_command=True)
def main(
    ctx: typer.Context,
    engines: str = typer.Option(
        DEFAULT_ENGINES,
        help="Comma-separated engine identifiers to benchmark",
    ),
    datasets: str = typer.Option(
        DEFAULT_DATASETS,
        help="Comma-separated dataset identifiers to benchmark",
    ),
    host: str = typer.Option("localhost", help="Target host running Docker containers"),
    timeout: float = typer.Option(3600.0, help="Timeout per benchmark (seconds)"),
    skip_upload: bool = typer.Option(False, help="Skip data upload"),
    skip_search: bool = typer.Option(False, help="Skip search phase"),
    force: bool = typer.Option(
        False,
        "--force",
        "--rerun",
        help="Rerun experiments even when previous JSON outputs exist",
    ),
):
    """Legacy entry point so `python run.py --engines ...` keeps working."""

    if ctx.invoked_subcommand is not None:
        return

    _execute_run(
        engines=_split_csv(engines),
        datasets=_split_csv(datasets),
        host=host,
        timeout=timeout,
        skip_upload=skip_upload,
        skip_search=skip_search,
        force=force,
    )


@app.command()
def run(
    engines: str = typer.Option(
        DEFAULT_ENGINES,
        help="Comma-separated engine identifiers to benchmark",
    ),
    datasets: str = typer.Option(
        DEFAULT_DATASETS,
        help="Comma-separated dataset identifiers to benchmark",
    ),
    host: str = typer.Option("localhost", help="Target host running Docker containers"),
    timeout: float = typer.Option(3600.0, help="Timeout per benchmark (seconds)"),
    skip_upload: bool = typer.Option(False, help="Skip data upload"),
    skip_search: bool = typer.Option(False, help="Skip search phase"),
    force: bool = typer.Option(
        False,
        "--force",
        "--rerun",
        help="Rerun experiments even when previous JSON outputs exist",
    ),
):
    """
    Run local vector database benchmarks for Qdrant and Milvus.

    Example:
        python run.py --engines "qdrant-single-node,milvus-single-node" --datasets "glove-100-angular,random-100,dbpedia-openai-1M-1536-angular,h-and-m-2048-angular-filters,msmarco-sparse-1M"
    """

    _execute_run(
        engines=_split_csv(engines),
        datasets=_split_csv(datasets),
        host=host,
        timeout=timeout,
        skip_upload=skip_upload,
        skip_search=skip_search,
        force=force,
    )


@app.command("list")
def list_configs(
    show_engines: bool = typer.Option(
        True,
        help="Show available engine identifiers",
        rich_help_panel="Sections",
    ),
    show_datasets: bool = typer.Option(
        True,
        help="Show available dataset identifiers",
        rich_help_panel="Sections",
    ),
):
    """Print registered engine and dataset identifiers."""

    if not show_engines and not show_datasets:
        raise typer.BadParameter("At least one of --show-engines/--show-datasets must be true")

    if show_engines:
        configs_by_name, engine_manifests = read_engine_configs()
        typer.echo("Engine manifests:")
        for manifest_name in sorted(engine_manifests.keys()):
            typer.echo(
                f"  - {manifest_name} ({len(engine_manifests[manifest_name])} configurations)"
            )

        typer.echo("")
        typer.echo("Experiment profiles:")
        for name in sorted(configs_by_name.keys()):
            typer.echo(f"  - {name}")

    if show_datasets:
        datasets = read_dataset_config()
        if show_engines:
            typer.echo("")
        typer.echo("Registered datasets:")
        for name in sorted(datasets.keys()):
            typer.echo(f"  - {name}")


@app.command()
def convert(
    input_dir: Path = typer.Argument(..., exists=True, file_okay=False, dir_okay=True),
    output_file: Path = typer.Argument(..., dir_okay=False),
):
    """Merge upload/search JSON outputs into a consolidated file."""

    search_files = sorted(glob.glob(str(input_dir / "*-search-*.json")))
    upload_files = sorted(glob.glob(str(input_dir / "*-upload-*.json")))

    typer.echo(f"Found {len(search_files)} search files and {len(upload_files)} upload files.")

    upload_data = {}
    for upload_file in upload_files:
        with open(upload_file) as fh:
            data = json.load(fh)
        experiment_name = data["params"]["experiment"]
        upload_data[experiment_name] = data

    consolidated = []

    for search_file in search_files:
        with open(search_file) as fh:
            data = json.load(fh)
        experiment_name = data["params"]["experiment"]
        dataset_name = data["params"]["dataset"]
        engine_params = data["params"]["config"]
        parallel = data["params"]["parallel"]
        engine_name = data["params"]["engine"]

        if experiment_name not in upload_data:
            typer.echo(
                f"[WARN] Upload results missing for experiment '{experiment_name}', skipping.",
                err=True,
            )
            continue

        upload_time = upload_data[experiment_name]["results"].get("upload_time")
        total_upload_time = upload_data[experiment_name]["results"].get("total_time")

        search_results = dict(data["results"])
        search_results.pop("total_time", None)

        consolidated.append(
            {
                "engine_name": engine_name,
                "setup_name": experiment_name,
                "dataset_name": dataset_name,
                "upload_time": upload_time,
                "total_upload_time": total_upload_time,
                "parallel": parallel,
                "engine_params": engine_params,
                **search_results,
            }
        )

    output_file.parent.mkdir(parents=True, exist_ok=True)

    with open(output_file, "w") as f:
        json.dump(consolidated, f, indent=2)

    typer.echo(f"Saved consolidated results to {output_file}")


@app.command("list")
def list_configs(
    show_engines: bool = typer.Option(
        True,
        help="Show available engine identifiers",
        rich_help_panel="Sections",
    ),
    show_datasets: bool = typer.Option(
        True,
        help="Show available dataset identifiers",
        rich_help_panel="Sections",
    ),
):
    """Print registered engine and dataset identifiers."""

    if not show_engines and not show_datasets:
        raise typer.BadParameter("At least one of --show-engines/--show-datasets must be true")

    if show_engines:
        engines = read_engine_configs()
        typer.echo("Registered engines:")
        for name in sorted(engines.keys()):
            typer.echo(f"  - {name}")

    if show_datasets:
        datasets = read_dataset_config()
        if show_engines:
            typer.echo("")
        typer.echo("Registered datasets:")
        for name in sorted(datasets.keys()):
            typer.echo(f"  - {name}")


@app.command()
def convert(
    input_dir: Path = typer.Argument(..., exists=True, file_okay=False, dir_okay=True),
    output_file: Path = typer.Argument(..., dir_okay=False),
):
    """Merge upload/search JSON outputs into a consolidated file.
    Gracefully handles missing config, params, or result fields.
    """

    search_files = sorted(glob.glob(str(input_dir / "*-search-*.json")))
    upload_files = sorted(glob.glob(str(input_dir / "*-upload-*.json")))

    typer.echo(f"Found {len(search_files)} search files and {len(upload_files)} upload files.")

    upload_data: dict[str, dict] = {}
    for upload_file in upload_files:
        try:
            with open(upload_file) as fh:
                data = json.load(fh)
            params = data.get("params", {})
            experiment_name = params.get("experiment", Path(upload_file).stem)
            upload_data[experiment_name] = data
        except Exception as e:
            typer.echo(f"[WARN] Failed to read {upload_file}: {e}", err=True)

    consolidated: list[dict] = []

    for search_file in search_files:
        try:
            with open(search_file) as fh:
                data = json.load(fh)
        except Exception as e:
            typer.echo(f"[WARN] Failed to read {search_file}: {e}", err=True)
            continue

        params = data.get("params", {})
        results = data.get("results", {})

        experiment_name = params.get("experiment", Path(search_file).stem)
        dataset_name = params.get("dataset", "unknown")
        engine_name = params.get("engine", "unknown")
        parallel = params.get("parallel", False)
        engine_params = params.get("config", {})  # ✅ safe access

        if experiment_name not in upload_data:
            typer.echo(
                f"[WARN] Upload results missing for experiment '{experiment_name}', skipping.",
                err=True,
            )
            continue

        upload_entry = upload_data[experiment_name]
        upload_results = upload_entry.get("results", {})

        upload_time = upload_results.get("upload_time")
        total_upload_time = upload_results.get("total_time")

        search_results = dict(results)
        search_results.pop("total_time", None)

        consolidated.append(
            {
                "engine_name": engine_name,
                "setup_name": experiment_name,
                "dataset_name": dataset_name,
                "upload_time": upload_time,
                "total_upload_time": total_upload_time,
                "parallel": parallel,
                "engine_params": engine_params,
                **search_results,
            }
        )

    output_file.parent.mkdir(parents=True, exist_ok=True)

    with open(output_file, "w") as f:
        json.dump(consolidated, f, indent=2)

    typer.echo(f"✅ Saved consolidated results to {output_file}")
    typer.echo(f"📊 Consolidated {len(consolidated)} benchmark entries.")


if __name__ == "__main__":
    app()