Benchmark Framework V2

benchmark_v2/.python-version

@@ -0,0 +1 @@
+3.12

benchmark_v2/Justfile

@@ -0,0 +1,41 @@
+set shell := ["bash", "-lc"]
+export CACHE_ROOT := `echo ~/.cache/datasketch/benchmark_v2`
+
+# Default target: list available recipes.
+default:
+	@just --list
+
+install:
+	uv sync
+
+format:
+	uv run ruff format .
+
+lint:
+	uv run ruff check .
+
+test:
+	uv run python -m pytest
+
+coverage:
+	uv run coverage run -m pytest && uv run coverage report
+
+# Pass-through CLI: `just cli -- subcommand --flag value`
+cli *args:
+	uv run python -m benchmark_v2 {{args}}
+
+list-datasets:
+	uv run python -m benchmark_v2 datasets list
+
+# Example: just sync my-dataset
+#          just sync my-dataset with_ground_truth=""   (to omit the flag)
+#          just sync my-dataset extra_flags="--limit 100"
+sync dataset with_ground_truth="--with-ground-truth" memmap="--memmap" extra_flags="":
+	uv run python -m benchmark_v2 datasets sync {{dataset}} --cache-root {{CACHE_ROOT}} {{with_ground_truth}} {{memmap}} {{extra_flags}}
+
+# Example: just sync-all
+#          just sync-all workers="--workers 8"
+#          just sync-all with_ground_truth="" memmap=""
+sync-all workers="" with_ground_truth="--with-ground-truth" memmap="--memmap" extra_flags="":
+	uv run python -m benchmark_v2 datasets sync-all --cache-root {{CACHE_ROOT}} {{with_ground_truth}} {{memmap}} {{workers}} {{extra_flags}}
+

benchmark_v2/benchmark_v2/__init__.py

@@ -0,0 +1,11 @@
+"""Benchmark v2 package."""
+
+from importlib import import_module
+
+__all__ = ["get_cli_app"]
+
+
+def get_cli_app():
+    """Import the Typer CLI lazily to avoid import-time side effects."""
+    module = import_module("benchmark_v2.cli")
+    return module.app

benchmark_v2/benchmark_v2/__main__.py

@@ -0,0 +1,10 @@
+from . import get_cli_app
+
+
+def main() -> None:
+    app = get_cli_app()
+    app()
+
+
+if __name__ == "__main__":
+    main()

benchmark_v2/benchmark_v2/cli.py

@@ -0,0 +1,183 @@
+"""Typer-based CLI for benchmark v2."""
+
+from __future__ import annotations
+
+import os
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+from typing import Callable, Optional
+
+import typer
+
+from .datasets import cache, memmap, registry
+
+app = typer.Typer(help="Benchmark v2 orchestration")
+datasets_app = typer.Typer(help="Manage benchmark datasets")
+app.add_typer(datasets_app, name="datasets")
+
+
+@datasets_app.command("list")
+def list_datasets() -> None:
+    """Display the datasets available for benchmarking."""
+
+    entries = sorted(registry.list_datasets(), key=lambda entry: entry.name)
+    for entry in entries:
+        typer.echo(f"- {entry.name}: corpus={entry.corpus.url}")
+        if entry.queries:
+            query_names = ", ".join(entry.queries.keys())
+            typer.echo(f"    queries: {query_names}")
+        if entry.ground_truth:
+            gt_names = ", ".join(entry.ground_truth.keys())
+            typer.echo(f"    ground truth: {gt_names}")
+        if entry.tags:
+            typer.echo(f"    tags: {', '.join(entry.tags)}")
+        if entry.notes:
+            typer.echo(f"    notes: {entry.notes}")
+
+
+def _sync_single_dataset(
+    name: str,
+    cache_root: Optional[Path],
+    include_queries: bool,
+    include_ground_truth: bool,
+    dtype: str,
+    build_memmap: bool,
+    force_download: bool,
+    force_memmap: bool,
+    write: Callable[[str], None],
+) -> None:
+    """Ensure dataset assets are cached locally and optionally memmapped."""
+
+    result = cache.ensure_dataset(
+        name,
+        cache_root=cache_root,
+        include_queries=include_queries,
+        include_ground_truth=include_ground_truth,
+        force_download=force_download,
+    )
+
+    write(f"Corpus stored at {result.corpus_path}")
+    if build_memmap:
+        corpus_memmap = memmap.ensure_memmap(
+            result.corpus_path, dtype=dtype, force=force_memmap
+        )
+        write(f"Corpus memmap ready at {corpus_memmap.tokens_path}")
+
+    for query_name, path in result.query_paths.items():
+        write(f"Query '{query_name}' stored at {path}")
+        if build_memmap:
+            q_memmap = memmap.ensure_memmap(path, dtype=dtype, force=force_memmap)
+            write(f"  memmap: {q_memmap.tokens_path}")
+
+    for gt_name, path in result.ground_truth_paths.items():
+        write(f"Ground truth '{gt_name}' stored at {path}")
+
+
+@datasets_app.command("sync")
+def sync_dataset(
+    name: str = typer.Argument(..., help="Dataset identifier"),
+    cache_root: Optional[Path] = typer.Option(
+        None, "--cache-root", help="Override cache directory"
+    ),
+    include_queries: bool = typer.Option(
+        True, "--queries/--no-queries", help="Include query assets"
+    ),
+    include_ground_truth: bool = typer.Option(
+        False,
+        "--with-ground-truth/--without-ground-truth",
+        help="Download ground-truth artifacts when available",
+    ),
+    dtype: str = typer.Option("int32", help="dtype for memmapped integer tokens"),
+    build_memmap: bool = typer.Option(
+        True, "--memmap/--no-memmap", help="Construct memmaps"
+    ),
+    force_download: bool = typer.Option(
+        False, "--force", help="Force re-download of assets"
+    ),
+    force_memmap: bool = typer.Option(
+        False, "--force-memmap", help="Rebuild memmaps from scratch"
+    ),
+) -> None:
+    """Ensure dataset assets are cached locally and optionally memmapped."""
+
+    _sync_single_dataset(
+        name=name,
+        cache_root=cache_root,
+        include_queries=include_queries,
+        include_ground_truth=include_ground_truth,
+        dtype=dtype,
+        build_memmap=build_memmap,
+        force_download=force_download,
+        force_memmap=force_memmap,
+        write=typer.echo,
+    )
+
+
+@datasets_app.command("sync-all")
+def sync_all_datasets(
+    cache_root: Optional[Path] = typer.Option(
+        None, "--cache-root", help="Override cache directory"
+    ),
+    include_queries: bool = typer.Option(
+        True, "--queries/--no-queries", help="Include query assets"
+    ),
+    include_ground_truth: bool = typer.Option(
+        False,
+        "--with-ground-truth/--without-ground-truth",
+        help="Download ground-truth artifacts when available",
+    ),
+    dtype: str = typer.Option("int32", help="dtype for memmapped integer tokens"),
+    build_memmap: bool = typer.Option(
+        True, "--memmap/--no-memmap", help="Construct memmaps"
+    ),
+    force_download: bool = typer.Option(
+        False, "--force", help="Force re-download of assets"
+    ),
+    force_memmap: bool = typer.Option(
+        False, "--force-memmap", help="Rebuild memmaps from scratch"
+    ),
+    max_workers: Optional[int] = typer.Option(
+        None,
+        "--max-workers",
+        min=1,
+        help="Maximum concurrent dataset downloads (default: number of datasets)",
+    ),
+) -> None:
+    """Ensure every registered dataset is cached locally."""
+
+    entries = sorted(registry.list_datasets(), key=lambda entry: entry.name.lower())
+    if not entries:
+        typer.echo("No datasets registered.")
+        return
+
+    default_workers = os.cpu_count() or 4
+    worker_count = max_workers or min(len(entries), default_workers)
+    results: dict[str, list[str]] = {}
+
+    def _run(entry_name: str) -> tuple[str, list[str]]:
+        buffer: list[str] = [f"=== {entry_name} ==="]
+        _sync_single_dataset(
+            name=entry_name,
+            cache_root=cache_root,
+            include_queries=include_queries,
+            include_ground_truth=include_ground_truth,
+            dtype=dtype,
+            build_memmap=build_memmap,
+            force_download=force_download,
+            force_memmap=force_memmap,
+            write=buffer.append,
+        )
+        return entry_name, buffer
+
+    executor = ThreadPoolExecutor(max_workers=worker_count)
+    future_map = {executor.submit(_run, entry.name): entry.name for entry in entries}
+    try:
+        for future in as_completed(future_map):
+            entry_name, buffer = future.result()
+            results[entry_name] = buffer
+    finally:
+        executor.shutdown(wait=False, cancel_futures=True)
+
+    for entry in entries:
+        for line in results.get(entry.name, []):
+            typer.echo(line)

benchmark_v2/benchmark_v2/datasets/__init__.py

@@ -0,0 +1,35 @@
+"""Dataset utilities for benchmark v2."""
+
+from . import cache, memmap, registry
+from .cache import DatasetCacheResult, ensure_dataset, get_cache_root
+from .ground_truth import GroundTruthRecord, ensure_ground_truth, iter_ground_truth
+from .memmap import CSRMemmap, build_memmap, ensure_memmap
+from .registry import (
+    DatasetAsset,
+    DatasetEntry,
+    datasets_dict,
+    get_dataset,
+    list_datasets,
+    register_dataset,
+)
+
+__all__ = [
+    "cache",
+    "memmap",
+    "registry",
+    "DatasetAsset",
+    "DatasetEntry",
+    "DatasetCacheResult",
+    "CSRMemmap",
+    "ensure_dataset",
+    "ensure_memmap",
+    "build_memmap",
+    "get_cache_root",
+    "list_datasets",
+    "get_dataset",
+    "register_dataset",
+    "datasets_dict",
+    "GroundTruthRecord",
+    "ensure_ground_truth",
+    "iter_ground_truth",
+]