From 7a2ca750eb47d1d3bb7747665f20a0f328deaceb Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Fri, 10 Apr 2026 14:18:40 -0500 Subject: [PATCH 01/22] Add initial opensearch deployment Signed-off-by: James Bourbeau --- deploy/README.md | 220 +++++++++++++++++ deploy/bench/Dockerfile | 38 +++ deploy/bench/run.py | 213 +++++++++++++++++ deploy/docker-compose.yml | 155 ++++++++++++ deploy/opensearch/Dockerfile | 16 ++ deploy/opensearch/opensearch.yml | 21 ++ deploy/remote-index-build/Dockerfile | 8 + deploy/remote-index-build/run.py | 337 +++++++++++++++++++++++++++ 8 files changed, 1008 insertions(+) create mode 100644 deploy/README.md create mode 100644 deploy/bench/Dockerfile create mode 100644 deploy/bench/run.py create mode 100644 deploy/docker-compose.yml create mode 100644 deploy/opensearch/Dockerfile create mode 100644 deploy/opensearch/opensearch.yml create mode 100644 deploy/remote-index-build/Dockerfile create mode 100644 deploy/remote-index-build/run.py diff --git a/deploy/README.md b/deploy/README.md new file mode 100644 index 0000000000..f83bd82e48 --- /dev/null +++ b/deploy/README.md @@ -0,0 +1,220 @@ +# OpenSearch GPU Remote Index Build Benchmark + +Docker Compose benchmark of [OpenSearch's GPU-accelerated remote index build](https://docs.opensearch.org/latest/vector-search/remote-index-build/) feature using [cuvs-bench](https://github.com/jrbourbeau/cuvs/tree/main/python/cuvs_bench). Spins up all required services, builds a kNN index on a GPU, and runs a cuvs-bench search benchmark sweep. + +## How it works + +OpenSearch's kNN plugin can offload Faiss HNSW index construction to a dedicated GPU service. Rather than building the index in-process on the OpenSearch node, the workflow is: + +``` +OpenSearch flushes a segment + → uploads raw vectors + doc-IDs to S3 (MinIO in this demo) + → POSTs /_build to the remote-index-builder service + → service downloads vectors from MinIO + → builds Faiss HNSW index on GPU + → uploads finished index back to MinIO + → OpenSearch downloads the GPU-built index and merges it into the shard +``` + +## Services + +| Service | Image | Purpose | +|---|---|---| +| `minio` | `minio/minio` | S3-compatible object store — staging area for vectors and built indexes | +| `minio-init` | `minio/mc` | One-shot: creates the `opensearch-vectors` bucket | +| `opensearch` | custom build of `opensearchproject/opensearch` | OpenSearch node with kNN plugin and `repository-s3` plugin | +| `remote-index-builder` | `opensearchproject/remote-vector-index-builder:api-latest` | FastAPI service that builds Faiss indexes on the GPU | +| `bench` | custom Python | Registers repo + cluster settings, runs cuvs-bench build/search benchmark | + +## Requirements + +- **NVIDIA GPU** with CUDA support +- **NVIDIA Container Toolkit** — [installation guide](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html) +- **Docker Compose v2** +- **ANN benchmark dataset** in binary format (`.fbin`) — see [Dataset format](#dataset-format) + +## Usage + +Set the host kernel parameter required by OpenSearch (once per reboot): + +```bash +sudo sysctl -w vm.max_map_count=262144 +``` + +Set required environment variables: + +```bash +export DATASET_PATH=/path/to/ann-benchmark-datasets # directory containing dataset files +``` + +Optionally configure the benchmark: + +```bash +export DATASET=sift-128-euclidean # default +export BENCH_GROUPS=test # test | base | large (default: test) +export K=10 # number of neighbors (default: 10) +``` + +Download the dataset (one-time setup, skipped automatically if already present): + +```bash +docker compose build bench +docker compose run --rm --no-deps bench python -m cuvs_bench.get_dataset \ + --dataset ${DATASET:-sift-128-euclidean} \ + --dataset-path /data/datasets +``` + +Start all services: + +```bash +docker compose up --build +``` + +The `bench` container logs its progress through index build, GPU verification, and search. When complete you'll see a results table: + +``` +════════════════════════════════════════════════════════════ + Benchmark complete! +════════════════════════════════════════════════════════════ + + OpenSearch : http://opensearch:9200 + MinIO console : http://localhost:9001 (minioadmin / minioadmin) +``` + +To tear everything down: + +```bash +docker compose down -v +``` + +## What the bench script does + +1. Registers MinIO as an OpenSearch S3 snapshot repository +2. Applies cluster settings to enable remote index build and point OpenSearch at the builder service +3. Runs `cuvs-bench` build phase (handled entirely by the OpenSearch backend): + - Creates the kNN index and bulk-ingests dataset vectors + - Force-merges the index to trigger the GPU build + - Polls MinIO every 5 s for `.faiss` files confirming GPU build completion + - Records total build time (ingestion + GPU build) in the result +4. Runs `cuvs-bench` search phase and prints a recall/QPS/latency table + +## Dataset format + +cuvs-bench reads binary vector files with a simple header: + +``` +[4 bytes: n_rows as uint32] +[4 bytes: n_cols as uint32] +[n_rows × n_cols × itemsize bytes: vector data] +``` + +Supported extensions: `.fbin` (float32), `.f16bin` (float16), `.u8bin` (uint8), `.i8bin` (int8). + +`DATASET_PATH` should be a directory where each dataset lives in its own subdirectory named after the dataset, e.g.: + +``` +$DATASET_PATH/ + sift-128-euclidean/ + base.fbin + query.fbin + groundtruth.neighbors.ibin +``` + +## Key configuration + +**Cluster settings** (applied by `bench/run.py`): + +```json +{ + "persistent": { + "knn.remote_index_build.enabled": true, + "knn.remote_index_build.repository": "vector-repo", + "knn.remote_index_build.service.endpoint": "http://remote-index-builder:1025" + } +} +``` + +**Parameter groups** (`BENCH_GROUPS`): + +| Group | Build params | Search params | Use case | +|---|---|---|---| +| `test` | 1 combo (m=16, ef_construction=100) | ef_search: 50, 100 | Quick smoke test | +| `base` | 6 combos (m × ef_construction sweep) | ef_search: 50–512 | Standard benchmark | +| `large` | 2 combos (larger m, ef_construction) | ef_search: 100–1024 | High-recall benchmark | + +## GPU build verification + +The cuvs-bench OpenSearch backend polls MinIO every 5 seconds for `.faiss` files under `s3://opensearch-vectors/knn-indexes/`. The remote-index-builder is the only component that writes `.faiss` files back to the bucket, so their presence is definitive proof the GPU build completed. + +The build raises a `TimeoutError` (causing the `bench` container to exit with code 1) if the expected number of `.faiss` files does not appear within 600 seconds. + +## Running without a GPU + +This is intentionally not supported. The `bench` container will exit 1 if `.faiss` files do not appear in MinIO within 600 s. If you want to experiment without hardware, remove the `deploy.resources.reservations` block from `remote-index-builder` in `docker-compose.yml` and be aware the benchmark will fail at the verification step. + +## Running tests + +The cuvs-bench OpenSearch backend has three tiers of tests. All run inside the `bench` container so no local Python environment is needed. + +**Build the bench image first** (or after any code changes): + +```bash +docker compose build --no-cache bench +``` + +### Unit tests (no server required) + +```bash +docker compose run --rm --no-deps bench \ + pytest /opt/cuvs/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py -v +``` + +### Integration tests (live OpenSearch node) + +Requires a running OpenSearch node. Start it with: + +```bash +docker compose up -d --wait opensearch +``` + +Then run: + +```bash +docker compose run --rm --no-deps \ + -e OPENSEARCH_URL=http://opensearch:9200 \ + bench \ + pytest /opt/cuvs/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py -v -m integration +``` + +### Remote index build integration tests (full GPU stack) + +Requires the full stack (OpenSearch, MinIO, and the remote index builder). Start all services with: + +```bash +docker compose up -d --wait minio minio-init opensearch remote-index-builder +``` + +Then run: + +```bash +docker compose run --rm --no-deps \ + -e OPENSEARCH_URL=http://opensearch:9200 \ + -e BUILDER_URL=http://remote-index-builder:1025 \ + -e S3_ENDPOINT=http://minio:9000 \ + -e S3_BUCKET=opensearch-vectors \ + -e S3_ACCESS_KEY=minioadmin \ + -e S3_SECRET_KEY=minioadmin \ + bench \ + pytest /opt/cuvs/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py -v -m integration +``` + +This runs all integration tests including `TestOpenSearchRemoteIndexBuildIntegration`, which verifies the full GPU build flow end-to-end. + +## Ports + +| Port | Service | +|---|---| +| `9200` | OpenSearch REST API | +| `9000` | MinIO S3 API | +| `9001` | MinIO web console | +| `1025` | Remote index builder API | diff --git a/deploy/bench/Dockerfile b/deploy/bench/Dockerfile new file mode 100644 index 0000000000..db3403b0e8 --- /dev/null +++ b/deploy/bench/Dockerfile @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +FROM python:3.11-slim +WORKDIR /app + +RUN apt-get update && apt-get install -y --no-install-recommends git && rm -rf /var/lib/apt/lists/* + +# Sparse-clone just the cuvs_bench Python package. +# Installing via pip is not possible without CUDA (rapids-build-backend requires +# nvcc, and cuvs_bench declares a hard dep on the `cuvs` CUDA package). +# The opensearch backend is pure Python and needs neither, so we add the +# package directly to PYTHONPATH and install only the actual runtime deps. +RUN git clone --depth=1 --filter=blob:none --sparse --branch cuvs-bench-opensearch https://github.com/jrbourbeau/cuvs.git /opt/cuvs \ + && cd /opt/cuvs \ + && git sparse-checkout set python/cuvs_bench + +ENV PYTHONPATH=/opt/cuvs/python/cuvs_bench + +# Runtime dependencies from cuvs_bench/pyproject.toml (excluding `cuvs` itself) +# plus extras needed by the opensearch backend and this benchmark script. +RUN pip install --no-cache-dir \ + boto3 \ + pytest \ + botocore \ + click \ + h5py \ + matplotlib \ + "numpy<2" \ + "opensearch-py>=2.4.0,<3.0.0" \ + pandas \ + pyyaml \ + requests \ + scikit-learn \ + scipy + +COPY run.py . +CMD ["python", "-u", "run.py"] diff --git a/deploy/bench/run.py b/deploy/bench/run.py new file mode 100644 index 0000000000..645988d94c --- /dev/null +++ b/deploy/bench/run.py @@ -0,0 +1,213 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +OpenSearch GPU Remote Index Build Benchmark +=========================================== +Steps: + 1. Register MinIO S3 snapshot repository with OpenSearch + 2. Configure cluster settings for GPU remote index build + 3. Build kNN index via cuvs-bench: + a. Bulk-ingest dataset vectors + b. Trigger force-merge to kick off the GPU build + c. Poll MinIO for .faiss files confirming GPU build completion + 4. Run cuvs-bench search benchmarks and print results +""" + +import os +import sys + +import requests +from cuvs_bench.orchestrator import BenchmarkOrchestrator +from cuvs_bench.backends.base import BuildResult, SearchResult + + +OPENSEARCH_URL = os.environ.get("OPENSEARCH_URL", "http://opensearch:9200") +OPENSEARCH_HOST = os.environ.get("OPENSEARCH_HOST", "opensearch") +OPENSEARCH_PORT = int(os.environ.get("OPENSEARCH_PORT", "9200")) +BUILDER_URL = os.environ.get("BUILDER_URL", "http://remote-index-builder:1025") +MINIO_URL = os.environ.get("MINIO_URL", "http://minio:9000") + +DATASET = os.environ.get("DATASET", "sift-128-euclidean") +DATASET_PATH = os.environ.get("DATASET_PATH", "/data/datasets") +BENCH_GROUPS = os.environ.get("BENCH_GROUPS", "test") +K = int(os.environ.get("K", "10")) + +BUCKET = "opensearch-vectors" +REPO_NAME = "vector-repo" + +session = requests.Session() +session.headers.update({"Content-Type": "application/json"}) + + +# ── helpers ─────────────────────────────────────────────────────────────────── + + +def banner(msg: str) -> None: + print(f"\n{'─' * 60}\n {msg}\n{'─' * 60}") + + +# ── OpenSearch setup ────────────────────────────────────────────────────────── + + +def register_repository() -> None: + banner(f"Registering S3 repository '{REPO_NAME}' (backed by MinIO)") + r = session.put( + f"{OPENSEARCH_URL}/_snapshot/{REPO_NAME}", + json={ + "type": "s3", + "settings": { + # endpoint / protocol / path_style_access are client-level settings + # configured in opensearch/opensearch.yml, not per-repository settings. + "bucket": BUCKET, + "base_path": "knn-indexes", + }, + }, + ) + r.raise_for_status() + print(f" {r.json()}") + + +def configure_cluster() -> None: + banner("Enabling GPU remote index build (cluster settings)") + r = session.put( + f"{OPENSEARCH_URL}/_cluster/settings", + json={ + "persistent": { + "knn.remote_index_build.enabled": True, + "knn.remote_index_build.repository": REPO_NAME, + "knn.remote_index_build.service.endpoint": BUILDER_URL, + } + }, + ) + r.raise_for_status() + print(f" {r.json()}") + + +# ── results ─────────────────────────────────────────────────────────────────── + + +def _print_result_row( + params: dict, recall: float, qps: float, latency_ms: float +) -> None: + params_str = ", ".join(f"{k}={v}" for k, v in params.items()) + print( + f" {params_str:<40} {recall:<12.4f} {qps:>8.1f} {latency_ms:>12.2f}" + ) + + +def print_results(results: list) -> None: + banner("Benchmark Results") + search_results = [r for r in results if isinstance(r, SearchResult)] + if not search_results: + print(" No search results returned.") + return + + header = f" {'params':<40} {'recall@' + str(K):<12} {'QPS':>8} {'latency (ms)':>12}" + print(header) + print(" " + "─" * (len(header) - 2)) + for r in search_results: + per_param = (r.metadata or {}).get("per_search_param_results") + if per_param: + for entry in per_param: + _print_result_row( + entry["search_params"], + entry["recall"], + entry["queries_per_second"], + entry["search_time_ms"], + ) + else: + _print_result_row( + r.search_params[0] if r.search_params else {}, + r.recall, + r.queries_per_second, + r.search_time_ms, + ) + + +# ── entrypoint ──────────────────────────────────────────────────────────────── + + +def main() -> None: + print("\n" + "═" * 60) + print(" OpenSearch GPU Remote Index Build Benchmark") + print("═" * 60) + print(f" OpenSearch : {OPENSEARCH_URL}") + print(f" GPU builder: {BUILDER_URL}") + print(f" Dataset : {DATASET} (path: {DATASET_PATH})") + print(f" Groups : {BENCH_GROUPS} k={K}") + + register_repository() + configure_cluster() + + orchestrator = BenchmarkOrchestrator(backend_type="opensearch") + + # Shared kwargs for both build and search phases + bench_kwargs = dict( + dataset=DATASET, + dataset_path=DATASET_PATH, + algorithms="opensearch_faiss_hnsw", + groups=BENCH_GROUPS, + host=OPENSEARCH_HOST, + port=OPENSEARCH_PORT, + use_ssl=False, + verify_certs=False, + remote_index_build=True, + # S3/MinIO config for GPU build verification (used by the backend) + remote_build_s3_endpoint=MINIO_URL, + remote_build_s3_bucket=BUCKET, + remote_build_s3_prefix="knn-indexes/", + remote_build_s3_access_key="minioadmin", + remote_build_s3_secret_key="minioadmin", + ) + + # ── Build phase ─────────────────────────────────────────────────────────── + # The backend handles the full GPU build flow: ingest vectors → force merge + # → poll MinIO for .faiss files confirming GPU build completion. + banner("Building index (GPU remote build via cuvs-bench)") + build_results = orchestrator.run_benchmark( + build=True, + search=False, + force=True, + bulk_batch_size=500, + **bench_kwargs, + ) + + index_names = [ + r.index_path + for r in build_results + if isinstance(r, BuildResult) and r.success and r.index_path + ] + if not index_names: + print(" ERROR: no indexes were successfully built") + sys.exit(1) + build_times = { + r.index_path: r.build_time_seconds + for r in build_results + if isinstance(r, BuildResult) and r.success and r.index_path + } + for name, t in build_times.items(): + print(f" {name}: built in {t:.1f}s") + + # ── Search phase ────────────────────────────────────────────────────────── + banner("Running search benchmarks (via cuvs-bench)") + search_results = orchestrator.run_benchmark( + build=False, + search=True, + count=K, + **bench_kwargs, + ) + + print_results(search_results) + + print("\n" + "═" * 60) + print(" Benchmark complete!") + print("═" * 60) + print(f"\n OpenSearch : {OPENSEARCH_URL}") + print(" MinIO console : http://localhost:9001 (minioadmin / minioadmin)") + print() + + +if __name__ == "__main__": + main() diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml new file mode 100644 index 0000000000..9f10f0aaca --- /dev/null +++ b/deploy/docker-compose.yml @@ -0,0 +1,155 @@ +# OpenSearch GPU Remote Index Build — Docker Compose Demo +# +# Architecture: +# minio S3-compatible object store used to exchange vector data +# minio-init One-shot container that creates the MinIO bucket +# opensearch OpenSearch node with kNN plugin (requires 2.17+) +# remote-index-builder GPU-accelerated Faiss index builder (FastAPI service) +# bench Registers repo + cluster settings, runs cuvs-bench build/search +# +# Requirements: +# - NVIDIA GPU with CUDA support +# - NVIDIA Container Toolkit https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html +# - Docker Compose v2 +# - vm.max_map_count >= 262144 on the host: +# sudo sysctl -w vm.max_map_count=262144 +# +# Usage: +# docker compose up --build +# +# Data flow: +# OpenSearch flushes a segment → uploads vectors + doc-IDs to MinIO +# OpenSearch POSTs /_build to the remote-index-builder with the S3 paths +# remote-index-builder downloads from MinIO, builds index on GPU, uploads result +# OpenSearch downloads the finished index from MinIO and merges it into the shard + +services: + + # ── Object Store ──────────────────────────────────────────────────────────── + minio: + image: minio/minio:latest + environment: + MINIO_ROOT_USER: minioadmin + MINIO_ROOT_PASSWORD: minioadmin + command: server /data --console-address ":9001" + ports: + - "9000:9000" # S3 API + - "9001:9001" # MinIO web console + volumes: + - minio-data:/data + healthcheck: + test: ["CMD", "curl", "-f", "http://localhost:9000/minio/health/live"] + interval: 5s + timeout: 5s + retries: 12 + start_period: 10s + + minio-init: + image: minio/mc:latest + depends_on: + minio: + condition: service_healthy + entrypoint: > + /bin/sh -c " + mc alias set local http://minio:9000 minioadmin minioadmin && + mc mb --ignore-existing local/opensearch-vectors && + echo 'Bucket opensearch-vectors ready' + " + restart: "no" + + # ── OpenSearch ─────────────────────────────────────────────────────────────── + opensearch: + build: + context: ./opensearch + environment: + - OPENSEARCH_JAVA_OPTS=-Xms4g -Xmx4g + # Bypass the Docker entrypoint script — all config lives in opensearch.yml. + # Keystore credentials are baked into the image (see opensearch/Dockerfile). + command: ["/usr/share/opensearch/bin/opensearch"] + ulimits: + nofile: + soft: 65536 + hard: 65536 + volumes: + - opensearch-data:/usr/share/opensearch/data + # Mounts S3 client config (endpoint, protocol, path_style_access). + # These are client-level settings and cannot be set in the snapshot API. + - ./opensearch/opensearch.yml:/usr/share/opensearch/config/opensearch.yml:ro + ports: + - "9200:9200" + depends_on: + minio-init: + condition: service_completed_successfully + healthcheck: + # wait_for_status=yellow blocks until the cluster is at least yellow + test: ["CMD-SHELL", "curl -sf 'http://localhost:9200/_cluster/health?wait_for_status=yellow&timeout=5s'"] + interval: 15s + timeout: 10s + retries: 20 + start_period: 30s + + # ── GPU Index Builder ──────────────────────────────────────────────────────── + remote-index-builder: + image: opensearchproject/remote-vector-index-builder:api-latest + environment: + - AWS_ACCESS_KEY_ID=minioadmin + - AWS_SECRET_ACCESS_KEY=minioadmin + - AWS_DEFAULT_REGION=us-east-1 + # Route S3 calls to MinIO instead of AWS (requires boto3 >= 1.28) + - AWS_ENDPOINT_URL=http://minio:9000 + ports: + - "1025:1025" + depends_on: + minio-init: + condition: service_completed_successfully + healthcheck: + test: ["CMD-SHELL", "python3 -c 'import socket; socket.create_connection((\"localhost\", 1025), 2).close()'"] + interval: 5s + timeout: 5s + retries: 24 + start_period: 10s + restart: on-failure + # GPU device reservation + deploy: + resources: + reservations: + devices: + - driver: nvidia + count: 1 + capabilities: [gpu] + + # ── Benchmark ──────────────────────────────────────────────────────────────── + # Required environment variables (set in shell or a .env file): + # DATASET_PATH Absolute path to the directory containing dataset files + # e.g. export DATASET_PATH=/data/ann-benchmarks + # Optional: + # DATASET Dataset name (default: sift-128-euclidean) + # BENCH_GROUPS Parameter sweep group: test | base | large (default: test) + # K Number of neighbors to search for (default: 10) + bench: + build: + context: ./bench + depends_on: + opensearch: + condition: service_healthy + remote-index-builder: + condition: service_healthy + minio: + condition: service_healthy + environment: + - OPENSEARCH_URL=http://opensearch:9200 + - OPENSEARCH_HOST=opensearch + - OPENSEARCH_PORT=9200 + - BUILDER_URL=http://remote-index-builder:1025 + - MINIO_URL=http://minio:9000 + - DATASET=${DATASET:-sift-128-euclidean} + - DATASET_PATH=/data/datasets + - BENCH_GROUPS=${BENCH_GROUPS:-test} + - K=${K:-10} + volumes: + - ${DATASET_PATH:?DATASET_PATH must be set to the directory containing dataset files}:/data/datasets + restart: "no" + +volumes: + minio-data: + opensearch-data: diff --git a/deploy/opensearch/Dockerfile b/deploy/opensearch/Dockerfile new file mode 100644 index 0000000000..af5c9b5629 --- /dev/null +++ b/deploy/opensearch/Dockerfile @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +FROM opensearchproject/opensearch:latest + +# The repository-s3 plugin is not bundled in the default OpenSearch image but +# is required for the remote vector index build feature: OpenSearch uses it to +# upload raw vectors and download GPU-built indexes via the MinIO bucket. +RUN /usr/share/opensearch/bin/opensearch-plugin install --batch repository-s3 + +# Pre-populate the keystore with MinIO credentials. The keystore lives in +# /usr/share/opensearch/config (not the data volume) so this survives restarts. +# The repository-s3 plugin reads credentials exclusively from the keystore. +RUN /usr/share/opensearch/bin/opensearch-keystore create && \ + echo "minioadmin" | /usr/share/opensearch/bin/opensearch-keystore add --stdin s3.client.default.access_key && \ + echo "minioadmin" | /usr/share/opensearch/bin/opensearch-keystore add --stdin s3.client.default.secret_key diff --git a/deploy/opensearch/opensearch.yml b/deploy/opensearch/opensearch.yml new file mode 100644 index 0000000000..ebb17cced3 --- /dev/null +++ b/deploy/opensearch/opensearch.yml @@ -0,0 +1,21 @@ +# Bind to all interfaces so other containers can reach OpenSearch +network.host: 0.0.0.0 + +# Single-node cluster — suppresses the production bootstrap checks that +# require seed_hosts / initial_cluster_manager_nodes to be configured. +# Must be in opensearch.yml (not an env var) because we exec the opensearch +# binary directly rather than going through the Docker entrypoint script. +discovery.type: single-node + +# Disable the security plugin — no SSL certs needed for this demo +plugins.security.disabled: true + +# S3 client used by the repository-s3 plugin. +# Credentials are injected into the keystore at container startup +# (see the `command` block in docker-compose.yml). +s3.client.default.endpoint: "minio:9000" +s3.client.default.protocol: "http" +s3.client.default.path_style_access: "true" +# The async S3 client requires a non-empty region even for non-AWS endpoints. +# "us-east-1" is a placeholder — MinIO ignores the region value entirely. +s3.client.default.region: "us-east-1" diff --git a/deploy/remote-index-build/Dockerfile b/deploy/remote-index-build/Dockerfile new file mode 100644 index 0000000000..e67d495429 --- /dev/null +++ b/deploy/remote-index-build/Dockerfile @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +FROM python:3.11-slim +WORKDIR /app +RUN pip install --no-cache-dir requests boto3 numpy +COPY run.py . +CMD ["python", "-u", "run.py"] diff --git a/deploy/remote-index-build/run.py b/deploy/remote-index-build/run.py new file mode 100644 index 0000000000..63ba918ff2 --- /dev/null +++ b/deploy/remote-index-build/run.py @@ -0,0 +1,337 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +OpenSearch GPU Remote Index Build — End-to-End Demo +==================================================== +Steps: + 1. Register a MinIO-backed S3 snapshot repository with OpenSearch + 2. Configure cluster settings to enable GPU-based remote index building + 3. Create a kNN index (Faiss HNSW / L2) with remote build enabled + 4. Ingest 100,000 random 256-dimensional float vectors via the bulk API (8 parallel workers) + 5. Flush + force-merge to consolidate segments and trigger the GPU build + 6. Poll MinIO for a .faiss file — hard-fail if the GPU build never completes + 7. Execute a kNN search and print the top-10 nearest neighbors +""" + +import json +import os +import sys +import time +from concurrent.futures import ThreadPoolExecutor, as_completed + +import boto3 +from botocore.config import Config +import numpy as np +import requests + +OPENSEARCH_URL = os.environ.get("OPENSEARCH_URL", "http://opensearch:9200") +BUILDER_URL = os.environ.get("BUILDER_URL", "http://remote-index-builder:1025") +MINIO_URL = os.environ.get("MINIO_URL", "http://minio:9000") + +INDEX_NAME = "gpu-demo" +DIMENSION = 256 # matches common embedding model output sizes +NUM_DOCS = 200_000 +BUCKET = "opensearch-vectors" +REPO_NAME = "vector-repo" + +session = requests.Session() +session.headers.update({"Content-Type": "application/json"}) + + +def banner(msg: str) -> None: + print(f"\n{'─' * 60}") + print(f" {msg}") + print(f"{'─' * 60}") + + +# ── configuration ───────────────────────────────────────────────────────────── + + +def register_repository() -> None: + banner(f"Registering S3 repository '{REPO_NAME}' (backed by MinIO)") + r = session.put( + f"{OPENSEARCH_URL}/_snapshot/{REPO_NAME}", + json={ + "type": "s3", + "settings": { + # endpoint / protocol / path_style_access are client-level settings + # configured in opensearch/opensearch.yml, not per-repository settings. + "bucket": BUCKET, + "base_path": "knn-indexes", + }, + }, + ) + r.raise_for_status() + print(f" {r.json()}") + + +def configure_cluster() -> None: + banner("Enabling remote GPU index build (cluster settings)") + r = session.put( + f"{OPENSEARCH_URL}/_cluster/settings", + json={ + "persistent": { + "knn.remote_index_build.enabled": True, + "knn.remote_index_build.repository": REPO_NAME, + "knn.remote_index_build.service.endpoint": BUILDER_URL, + } + }, + ) + r.raise_for_status() + print(f" {r.json()}") + + +# ── index ───────────────────────────────────────────────────────────────────── + + +def create_index() -> None: + banner(f"Creating kNN index '{INDEX_NAME}'") + + resp = session.delete(f"{OPENSEARCH_URL}/{INDEX_NAME}") + if resp.status_code == 200: + print(" Deleted existing index") + + r = session.put( + f"{OPENSEARCH_URL}/{INDEX_NAME}", + json={ + "settings": { + "index.knn": True, + "index.knn.remote_index_build.enabled": True, + # Trigger GPU build for segments >= 1 KB (covers any real dataset) + "index.knn.remote_index_build.size.min": "1kb", + "number_of_shards": 1, + "number_of_replicas": 0, + }, + "mappings": { + "properties": { + "vector": { + "type": "knn_vector", + "dimension": DIMENSION, + "method": { + "name": "hnsw", + "engine": "faiss", + "space_type": "l2", + "parameters": {"m": 32, "ef_construction": 512}, + }, + }, + "doc_id": {"type": "integer"}, + "label": {"type": "keyword"}, + } + }, + }, + ) + r.raise_for_status() + print(f" {r.json()}") + + +# ── ingest ──────────────────────────────────────────────────────────────────── + + +def ingest_vectors() -> None: + batch_size = 500 + banner( + f"Ingesting {NUM_DOCS:,} random {DIMENSION}-dim vectors (bulk API, 8 workers)" + ) + + def send_batch(start: int) -> int: + end = min(start + batch_size, NUM_DOCS) + vecs = np.random.randn(end - start, DIMENSION).astype(np.float32) + lines = [] + for i, vec in enumerate(vecs, start): + lines.append( + json.dumps({"index": {"_index": INDEX_NAME, "_id": str(i)}}) + ) + lines.append( + json.dumps( + { + "vector": vec.tolist(), + "doc_id": i, + "label": f"item-{i:04d}", + } + ) + ) + payload = ("\n".join(lines) + "\n").encode("utf-8") + r = session.post( + f"{OPENSEARCH_URL}/_bulk", + data=payload, + headers={"Content-Type": "application/x-ndjson"}, + ) + r.raise_for_status() + body = r.json() + if body.get("errors"): + failed = [ + item["index"]["error"] + for item in body["items"] + if "error" in item.get("index", {}) + ] + print( + f" Warning: {len(failed)} error(s) in batch {start}–{end}: {failed[0]}" + ) + return (end - start) - len(failed) + return end - start + + ingested = 0 + starts = list(range(0, NUM_DOCS, batch_size)) + with ThreadPoolExecutor(max_workers=8) as executor: + futures = {executor.submit(send_batch, s): s for s in starts} + for future in as_completed(futures): + ingested += future.result() + if ingested % 10_000 == 0 or ingested >= NUM_DOCS: + print(f" Ingested {ingested:,}/{NUM_DOCS:,}") + + session.post(f"{OPENSEARCH_URL}/{INDEX_NAME}/_flush") + r = session.get(f"{OPENSEARCH_URL}/{INDEX_NAME}/_count") + print(f" Document count after flush: {r.json()['count']:,}") + + +# ── GPU build ───────────────────────────────────────────────────────────────── + + +def trigger_gpu_build() -> None: + banner("Triggering GPU index build via force merge") + print( + " OpenSearch will upload vectors to MinIO, then call the GPU builder." + ) + print( + " force_merge max_num_segments=1 consolidates all segments into one." + ) + r = session.post( + f"{OPENSEARCH_URL}/{INDEX_NAME}/_forcemerge?max_num_segments=1", + timeout=300, + ) + print(f" Force merge HTTP {r.status_code}") + + +def verify_gpu_build(timeout: int = 600) -> None: + """Confirm the GPU builder uploaded a .faiss index file to MinIO. + + The remote-index-builder is the *only* component that writes .faiss files + back to the S3 bucket, so their presence is definitive proof that the GPU + build completed. The kNN stats API does not expose remote build counters + in OpenSearch 3.x, so we poll MinIO directly via boto3 instead. + + Exits with code 1 if no .faiss file appears within `timeout` seconds. + """ + banner("Verifying GPU index build (polling MinIO for .faiss files)") + print(f" Bucket : {BUCKET}/knn-indexes/") + print(f" Timeout : {timeout}s (poll interval: 5s)\n") + + s3 = boto3.client( + "s3", + endpoint_url=MINIO_URL, + aws_access_key_id="minioadmin", + aws_secret_access_key="minioadmin", + region_name="us-east-1", + config=Config(signature_version="s3v4"), + ) + + deadline = time.time() + timeout + while time.time() < deadline: + try: + resp = s3.list_objects_v2(Bucket=BUCKET, Prefix="knn-indexes/") + faiss_files = [ + obj["Key"] + for obj in resp.get("Contents", []) + if obj["Key"].endswith(".faiss") + ] + if faiss_files: + print( + f" PASS: GPU build confirmed — {len(faiss_files)} .faiss file(s) in MinIO:" + ) + for f in faiss_files: + print(f" s3://{BUCKET}/{f}") + return + + remaining = int(deadline - time.time()) + all_keys = [obj["Key"] for obj in resp.get("Contents", [])] + print( + f" Waiting for .faiss file... objects={all_keys} ({remaining}s left)" + ) + except Exception as e: + print(f" MinIO check error: {e}") + time.sleep(5) + + print( + f"\n FAIL: no GPU-built .faiss index appeared in MinIO after {timeout}s" + ) + print("\n Possible causes:") + print( + " 1. remote-index-builder is unreachable from the OpenSearch container." + ) + print( + f" Verify the container is running and BUILDER_URL={BUILDER_URL} is correct." + ) + print( + " 2. Segment size never exceeded index.knn.remote_index_build.size.min." + ) + print(" Try increasing NUM_DOCS or lowering the size.min threshold.") + print( + " 3. No GPU is available inside the remote-index-builder container." + ) + print(" Check: docker compose logs remote-index-builder") + print(" Ensure the NVIDIA Container Toolkit is installed on the host.") + sys.exit(1) + + +# ── search ──────────────────────────────────────────────────────────────────── + + +def search_vectors() -> None: + banner("kNN test search (top-5 nearest neighbors)") + query_vec = np.random.randn(DIMENSION).astype(np.float32).tolist() + + r = session.post( + f"{OPENSEARCH_URL}/{INDEX_NAME}/_search", + json={ + "size": 10, + "query": {"knn": {"vector": {"vector": query_vec, "k": 10}}}, + "_source": ["doc_id", "label"], + }, + ) + r.raise_for_status() + hits = r.json()["hits"]["hits"] + total = r.json()["hits"]["total"]["value"] + + print(f" Index contains {total} documents") + print(f" Top {len(hits)} results:") + for rank, hit in enumerate(hits, 1): + src = hit["_source"] + print( + f" #{rank:>2} id={hit['_id']:>6} score={hit['_score']:.6f} label={src['label']}" + ) + + +# ── entrypoint ──────────────────────────────────────────────────────────────── + + +def main() -> None: + print("\n" + "═" * 60) + print(" OpenSearch GPU Remote Index Build — End-to-End Demo") + print("═" * 60) + print(f" OpenSearch : {OPENSEARCH_URL}") + print(f" GPU builder: {BUILDER_URL}") + print( + f" Vectors : {NUM_DOCS} × dim={DIMENSION} engine=faiss method=hnsw space=l2" + ) + + register_repository() + configure_cluster() + create_index() + ingest_vectors() + trigger_gpu_build() + verify_gpu_build() + search_vectors() + + print("\n" + "═" * 60) + print(" Demo complete!") + print("═" * 60) + print(f"\n OpenSearch is still running at {OPENSEARCH_URL}") + print(" MinIO console : http://localhost:9001 (minioadmin / minioadmin)") + print(f" GPU builder : {BUILDER_URL}") + print() + + +if __name__ == "__main__": + main() From d8bf286bd664a6b4fe565c7d21c046a372fa5f9a Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Tue, 28 Apr 2026 17:56:40 -0500 Subject: [PATCH 02/22] Update Signed-off-by: James Bourbeau --- deploy/DEPLOYMENT.md | 191 ++++++++++++++++++++++++ deploy/README.md | 118 +++++++-------- deploy/bench/Dockerfile | 7 +- deploy/bench/entrypoint.sh | 48 ++++++ deploy/bench/run.py | 212 ++++++++++++++++++--------- deploy/docker-compose.yml | 114 ++++++-------- deploy/opensearch/Dockerfile | 18 +-- deploy/opensearch/entrypoint.sh | 31 ++++ deploy/opensearch/opensearch.yml | 14 +- deploy/remote-index-build/Dockerfile | 3 - deploy/remote-index-build/run.py | 162 +++++++------------- 11 files changed, 576 insertions(+), 342 deletions(-) create mode 100644 deploy/DEPLOYMENT.md create mode 100644 deploy/bench/entrypoint.sh create mode 100644 deploy/opensearch/entrypoint.sh diff --git a/deploy/DEPLOYMENT.md b/deploy/DEPLOYMENT.md new file mode 100644 index 0000000000..181132e309 --- /dev/null +++ b/deploy/DEPLOYMENT.md @@ -0,0 +1,191 @@ +# OpenSearch GPU Remote Index Build — Deployment Guide + +This guide walks through running OpenSearch with GPU-accelerated vector index construction using the [remote index build service](https://docs.opensearch.org/latest/vector-search/remote-index-build/). When enabled, OpenSearch offloads Faiss HNSW index building to a dedicated GPU service rather than building indexes in-process on the data nodes. + +## How it works + +The GPU build is triggered automatically during normal ingest — no changes to your indexing workflow are required beyond the one-time cluster and index configuration described below. + +```mermaid +sequenceDiagram + participant Client + participant OpenSearch + participant S3 + participant Builder as Remote Index Builder (GPU) + + Client->>OpenSearch: Bulk ingest vectors + note over OpenSearch: segment flush + OpenSearch->>S3: Upload raw vectors + doc IDs + OpenSearch->>Builder: POST /_build (S3 paths) + Builder->>S3: Download vectors + note over Builder: Build Faiss HNSW on GPU + Builder->>S3: Upload .faiss index + OpenSearch->>S3: Download .faiss index + note over OpenSearch: Merge into shard + Client->>OpenSearch: kNN search query + OpenSearch->>Client: Top-k results +``` + +## Services + +| Service | Image | Purpose | +|---|---|---| +| `opensearch` | custom build of `opensearchproject/opensearch:3.6.0` | OpenSearch node with kNN plugin and `repository-s3` plugin | +| `remote-index-builder` | `opensearchproject/remote-vector-index-builder:api-latest` | GPU-accelerated Faiss HNSW index builder | + +The custom OpenSearch image adds the `repository-s3` plugin (required for S3-backed vector staging) and populates the S3 keystore from environment variables at startup so credentials are never baked into image layers. + +## Requirements + +- **Docker Compose v2** +- **NVIDIA GPU** with CUDA support +- **NVIDIA Container Toolkit** — [installation guide](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html) +- **AWS S3 bucket** for staging vectors during the build + +## Setup + +Set the host kernel parameter required by OpenSearch (once per reboot): + +```bash +sudo sysctl -w vm.max_map_count=262144 +``` + +Set required environment variables: + +```bash +export S3_BUCKET= +export AWS_ACCESS_KEY_ID= +export AWS_SECRET_ACCESS_KEY= +``` + +Optionally configure the region and session token for temporary credentials: + +```bash +export AWS_DEFAULT_REGION=us-east-1 # default: us-east-1 +export AWS_SESSION_TOKEN= # required for temporary (STS) credentials +``` + +Start OpenSearch and the GPU builder: + +```bash +docker compose --profile gpu up --build -d --wait opensearch remote-index-builder +``` + +## Connecting OpenSearch to the GPU builder + +Before any index can use GPU builds, you need to register your S3 bucket as a snapshot repository and apply the cluster settings that point OpenSearch at the builder service. Run these once against a live cluster. + +**Register S3 repository:** + +```bash +curl -X PUT http://localhost:9200/_snapshot/ \ + -H "Content-Type: application/json" \ + -d '{ + "type": "s3", + "settings": { + "bucket": "", + "base_path": "knn-indexes", + "region": "us-east-1" + } + }' +``` + +**Apply cluster settings:** + +```bash +curl -X PUT http://localhost:9200/_cluster/settings \ + -H "Content-Type: application/json" \ + -d '{ + "persistent": { + "knn.remote_index_build.enabled": true, + "knn.remote_index_build.repository": "", + "knn.remote_index_build.service.endpoint": "http://remote-index-builder:1025" + } + }' +``` + +> **Note:** `remote-index-builder` resolves inside the Docker network. If OpenSearch and the builder are not on the same Docker network, replace this with a reachable hostname or IP. + +## Creating an index with GPU builds enabled + +Add `"index.knn.remote_index_build.enabled": true` to your index settings alongside the standard kNN configuration: + +```bash +curl -X PUT http://localhost:9200/my-vectors \ + -H "Content-Type: application/json" \ + -d '{ + "settings": { + "index.knn": true, + "index.knn.remote_index_build.enabled": true, + "number_of_shards": 1, + "number_of_replicas": 1 + }, + "mappings": { + "properties": { + "vector": { + "type": "knn_vector", + "dimension": 256, + "method": { + "name": "hnsw", + "engine": "faiss", + "space_type": "l2", + "parameters": { + "m": 32, + "ef_construction": 512 + } + } + } + } + } + }' +``` + +GPU builds are only available with the `faiss` engine. The `lucene` engine always builds locally. + +## Verifying the GPU build + +The `remote-index-build/` directory contains an end-to-end demo script that ingests 200,000 random vectors, triggers a force-merge, and confirms the GPU build completed by polling S3 for the resulting `.faiss` file. + +Run it inside a temporary container on the same Docker network: + +```bash +docker compose run --rm \ + -e OPENSEARCH_URL=http://opensearch:9200 \ + -e BUILDER_URL=http://remote-index-builder:1025 \ + -e S3_BUCKET=${S3_BUCKET} \ + -e AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID} \ + -e AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY} \ + -e AWS_SESSION_TOKEN=${AWS_SESSION_TOKEN} \ + -e AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION:-us-east-1} \ + -v $(pwd)/remote-index-build:/app/remote-index-build \ + --no-deps bench \ + python remote-index-build/run.py +``` + +Or run it directly if you have Python and the dependencies installed locally (`boto3`, `numpy`, `requests`), pointing `OPENSEARCH_URL` at `http://localhost:9200`. + +A successful run prints a `.faiss` file path in S3 and returns top-10 nearest-neighbor results. + +## Tearing down + +```bash +docker compose --profile gpu down -v +``` + +The `-v` flag removes the OpenSearch data volume. Omit it to preserve indexed data across restarts. + +## Production considerations + +This setup is a working demonstration, not a production-hardened deployment. Key differences to address before running in production: + +- **Security plugin**: `opensearch.yml` has `plugins.security.disabled: true`. Re-enable it and configure TLS and authentication for any non-local deployment. +- **Single-node cluster**: `discovery.type: single-node` bypasses multi-node bootstrap checks. Replace with a properly configured multi-node cluster for production. +- **Replicas**: The demo uses `number_of_replicas: 0`. Set this to at least `1` for production workloads. +- **S3 permissions**: The IAM credentials need `s3:GetObject`, `s3:PutObject`, `s3:ListBucket`, and `s3:DeleteObject` on the staging bucket. + +## Ports + +| Port | Service | +|---|---| +| `9200` | OpenSearch REST API | +| `1025` | Remote index builder API | diff --git a/deploy/README.md b/deploy/README.md index f83bd82e48..d5678d55a1 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -1,6 +1,6 @@ -# OpenSearch GPU Remote Index Build Benchmark +# OpenSearch kNN Benchmark -Docker Compose benchmark of [OpenSearch's GPU-accelerated remote index build](https://docs.opensearch.org/latest/vector-search/remote-index-build/) feature using [cuvs-bench](https://github.com/jrbourbeau/cuvs/tree/main/python/cuvs_bench). Spins up all required services, builds a kNN index on a GPU, and runs a cuvs-bench search benchmark sweep. +Docker Compose benchmark comparing CPU and GPU kNN index builds in OpenSearch using [cuvs-bench](https://github.com/jrbourbeau/cuvs/tree/main/python/cuvs_bench). Supports both local CPU builds and [GPU-accelerated remote index builds](https://docs.opensearch.org/latest/vector-search/remote-index-build/) via the `REMOTE_INDEX_BUILD` environment variable. ## How it works @@ -8,11 +8,11 @@ OpenSearch's kNN plugin can offload Faiss HNSW index construction to a dedicated ``` OpenSearch flushes a segment - → uploads raw vectors + doc-IDs to S3 (MinIO in this demo) + → uploads raw vectors + doc-IDs to S3 → POSTs /_build to the remote-index-builder service - → service downloads vectors from MinIO + → service downloads vectors from S3 → builds Faiss HNSW index on GPU - → uploads finished index back to MinIO + → uploads finished index back to S3 → OpenSearch downloads the GPU-built index and merges it into the shard ``` @@ -20,18 +20,18 @@ OpenSearch flushes a segment | Service | Image | Purpose | |---|---|---| -| `minio` | `minio/minio` | S3-compatible object store — staging area for vectors and built indexes | -| `minio-init` | `minio/mc` | One-shot: creates the `opensearch-vectors` bucket | | `opensearch` | custom build of `opensearchproject/opensearch` | OpenSearch node with kNN plugin and `repository-s3` plugin | | `remote-index-builder` | `opensearchproject/remote-vector-index-builder:api-latest` | FastAPI service that builds Faiss indexes on the GPU | -| `bench` | custom Python | Registers repo + cluster settings, runs cuvs-bench build/search benchmark | +| `bench` | custom Python | Downloads dataset, registers repo + cluster settings, runs cuvs-bench build/search benchmark, exports results, generates plots | ## Requirements -- **NVIDIA GPU** with CUDA support -- **NVIDIA Container Toolkit** — [installation guide](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html) - **Docker Compose v2** - **ANN benchmark dataset** in binary format (`.fbin`) — see [Dataset format](#dataset-format) +- **GPU mode only** (`--profile gpu`, `REMOTE_INDEX_BUILD=true`): + - NVIDIA GPU with CUDA support + - NVIDIA Container Toolkit — [installation guide](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html) + - AWS S3 bucket (or S3-compatible store) for staging vectors and built indexes ## Usage @@ -47,39 +47,35 @@ Set required environment variables: export DATASET_PATH=/path/to/ann-benchmark-datasets # directory containing dataset files ``` -Optionally configure the benchmark: +GPU mode also requires S3 credentials: ```bash -export DATASET=sift-128-euclidean # default -export BENCH_GROUPS=test # test | base | large (default: test) -export K=10 # number of neighbors (default: 10) +export S3_BUCKET=my-opensearch-vectors # S3 bucket name +export AWS_ACCESS_KEY_ID= +export AWS_SECRET_ACCESS_KEY= ``` -Download the dataset (one-time setup, skipped automatically if already present): +Optionally configure the benchmark: ```bash -docker compose build bench -docker compose run --rm --no-deps bench python -m cuvs_bench.get_dataset \ - --dataset ${DATASET:-sift-128-euclidean} \ - --dataset-path /data/datasets +export AWS_SESSION_TOKEN= # required when using temporary (STS) credentials +export AWS_DEFAULT_REGION=us-east-1 # AWS region for the S3 bucket (default: us-east-1) +export DATASET=sift-128-euclidean # default +export BENCH_GROUPS=test # test | base (default: test) +export K=10 # number of neighbors (default: 10) ``` Start all services: ```bash +# CPU build (no GPU required) docker compose up --build -``` - -The `bench` container logs its progress through index build, GPU verification, and search. When complete you'll see a results table: +# GPU build +docker compose --profile gpu up --build ``` -════════════════════════════════════════════════════════════ - Benchmark complete! -════════════════════════════════════════════════════════════ - OpenSearch : http://opensearch:9200 - MinIO console : http://localhost:9001 (minioadmin / minioadmin) -``` +The `bench` container logs its progress through each phase. When complete you'll see a results table followed by the paths to the generated plot PNGs under `$DATASET_PATH`. To tear everything down: @@ -87,16 +83,19 @@ To tear everything down: docker compose down -v ``` -## What the bench script does +## What the bench container does -1. Registers MinIO as an OpenSearch S3 snapshot repository -2. Applies cluster settings to enable remote index build and point OpenSearch at the builder service -3. Runs `cuvs-bench` build phase (handled entirely by the OpenSearch backend): +1. Downloads the dataset (skipped if already present in `$DATASET_PATH`) +2. **GPU mode only**: Registers the S3 bucket as an OpenSearch snapshot repository +3. **GPU mode only**: Applies cluster settings to enable remote index build and point OpenSearch at the builder service +4. Runs `cuvs-bench` build phase (handled entirely by the OpenSearch backend): - Creates the kNN index and bulk-ingests dataset vectors - - Force-merges the index to trigger the GPU build - - Polls MinIO every 5 s for `.faiss` files confirming GPU build completion - - Records total build time (ingestion + GPU build) in the result -4. Runs `cuvs-bench` search phase and prints a recall/QPS/latency table + - **GPU mode**: Waits for all ingestion-time GPU builds to complete, then force-merges to one segment to trigger the final GPU build and polls the kNN stats API every 5 s until the build is confirmed complete + - **CPU mode**: Force-merges the index to one segment + - Records total build time in the result +5. Runs `cuvs-bench` search phase and prints a recall/QPS/latency table +6. Exports benchmark JSON results to CSV (`cuvs_bench.run --data-export`) +7. Generates recall vs. latency/throughput plots as PNGs in `$DATASET_PATH` (`cuvs_bench.plot`) ## Dataset format @@ -139,18 +138,27 @@ $DATASET_PATH/ | Group | Build params | Search params | Use case | |---|---|---|---| | `test` | 1 combo (m=16, ef_construction=100) | ef_search: 50, 100 | Quick smoke test | -| `base` | 6 combos (m × ef_construction sweep) | ef_search: 50–512 | Standard benchmark | -| `large` | 2 combos (larger m, ef_construction) | ef_search: 100–1024 | High-recall benchmark | +| `base` | 16 combos (m=[32,64,96,128] × ef_construction=[64,128,256,512]) | ef_search: 10–800 | Standard benchmark | ## GPU build verification -The cuvs-bench OpenSearch backend polls MinIO every 5 seconds for `.faiss` files under `s3://opensearch-vectors/knn-indexes/`. The remote-index-builder is the only component that writes `.faiss` files back to the bucket, so their presence is definitive proof the GPU build completed. +The cuvs-bench OpenSearch backend polls the kNN stats API every 5 seconds, waiting for `index_build_success_count` to increment by the expected number of new builds and for all in-flight flush and merge operations to reach zero. -The build raises a `TimeoutError` (causing the `bench` container to exit with code 1) if the expected number of `.faiss` files does not appear within 600 seconds. +The build raises a `TimeoutError` (causing the `bench` container to exit with code 1) if the expected number of successful builds is not confirmed within 600 seconds. -## Running without a GPU +## CPU vs GPU comparison -This is intentionally not supported. The `bench` container will exit 1 if `.faiss` files do not appear in MinIO within 600 s. If you want to experiment without hardware, remove the `deploy.resources.reservations` block from `remote-index-builder` in `docker-compose.yml` and be aware the benchmark will fail at the verification step. +To compare CPU and GPU builds on the same dataset, run the benchmark twice — once in each mode — clearing the OpenSearch volume between runs so the index is rebuilt from scratch each time: + +```bash +# GPU build (starts the remote-index-builder via --profile gpu) +docker compose --profile gpu up --build +docker compose --profile gpu down -v + +# CPU build (no GPU or S3 required) +docker compose up --build +docker compose down -v +``` ## Running tests @@ -171,15 +179,10 @@ docker compose run --rm --no-deps bench \ ### Integration tests (live OpenSearch node) -Requires a running OpenSearch node. Start it with: +Requires a running OpenSearch node. S3 credentials are not required for these tests. ```bash docker compose up -d --wait opensearch -``` - -Then run: - -```bash docker compose run --rm --no-deps \ -e OPENSEARCH_URL=http://opensearch:9200 \ bench \ @@ -188,22 +191,17 @@ docker compose run --rm --no-deps \ ### Remote index build integration tests (full GPU stack) -Requires the full stack (OpenSearch, MinIO, and the remote index builder). Start all services with: - -```bash -docker compose up -d --wait minio minio-init opensearch remote-index-builder -``` - -Then run: +Requires the full stack (OpenSearch and the remote index builder) and S3 credentials. ```bash +docker compose --profile gpu up -d --wait opensearch remote-index-builder docker compose run --rm --no-deps \ -e OPENSEARCH_URL=http://opensearch:9200 \ -e BUILDER_URL=http://remote-index-builder:1025 \ - -e S3_ENDPOINT=http://minio:9000 \ - -e S3_BUCKET=opensearch-vectors \ - -e S3_ACCESS_KEY=minioadmin \ - -e S3_SECRET_KEY=minioadmin \ + -e S3_BUCKET=${S3_BUCKET} \ + -e AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID} \ + -e AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY} \ + -e AWS_SESSION_TOKEN=${AWS_SESSION_TOKEN} \ bench \ pytest /opt/cuvs/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py -v -m integration ``` @@ -215,6 +213,4 @@ This runs all integration tests including `TestOpenSearchRemoteIndexBuildIntegra | Port | Service | |---|---| | `9200` | OpenSearch REST API | -| `9000` | MinIO S3 API | -| `9001` | MinIO web console | | `1025` | Remote index builder API | diff --git a/deploy/bench/Dockerfile b/deploy/bench/Dockerfile index db3403b0e8..b4eb2af023 100644 --- a/deploy/bench/Dockerfile +++ b/deploy/bench/Dockerfile @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - FROM python:3.11-slim WORKDIR /app @@ -34,5 +31,5 @@ RUN pip install --no-cache-dir \ scikit-learn \ scipy -COPY run.py . -CMD ["python", "-u", "run.py"] +COPY --chmod=755 run.py entrypoint.sh . +CMD ["/app/entrypoint.sh"] diff --git a/deploy/bench/entrypoint.sh b/deploy/bench/entrypoint.sh new file mode 100644 index 0000000000..4f48da9c4c --- /dev/null +++ b/deploy/bench/entrypoint.sh @@ -0,0 +1,48 @@ +#!/bin/bash +set -e + +DATASET="${DATASET:-sift-128-euclidean}" +BENCH_GROUPS="${BENCH_GROUPS:-test}" +K="${K:-10}" +# Auto-detect GPU mode: remote-index-builder only appears in Docker DNS when +# started via --profile gpu. DNS entries are registered at network setup time +# (before containers run), so this check is reliable by the time entrypoint +# executes (OpenSearch healthy check alone takes 30+ seconds). +if getent hosts remote-index-builder > /dev/null 2>&1; then + echo "remote-index-builder detected — waiting for it to be ready..." + until python3 -c 'import socket; socket.create_connection(("remote-index-builder", 1025), 2).close()' 2>/dev/null; do + sleep 5 + done + echo "remote-index-builder is ready." + export REMOTE_INDEX_BUILD=true +else + echo "remote-index-builder not available — using CPU build mode." + export REMOTE_INDEX_BUILD=false +fi + +# Step 1: Download dataset (skipped automatically if already present) +python -m cuvs_bench.get_dataset \ + --dataset "$DATASET" \ + --dataset-path /data/datasets + +# Step 2: Run benchmark (build + search + writes result JSON files) +python -u run.py + +# Step 3: Export JSON → CSV (required by cuvs_bench.plot) +python -m cuvs_bench.run --data-export \ + --dataset "$DATASET" \ + --dataset-path /data/datasets \ + --algorithms opensearch_faiss_hnsw \ + --groups "$BENCH_GROUPS" \ + --count "$K" \ + --batch-size 10000 \ + --search-mode latency + +# Step 4: Plot — PNGs written to /data/datasets (mounted from host $DATASET_PATH) +python -m cuvs_bench.plot \ + --dataset "$DATASET" \ + --dataset-path /data/datasets \ + --algorithms opensearch_faiss_hnsw \ + --groups "$BENCH_GROUPS" \ + --count "$K" \ + --output-filepath /data/datasets diff --git a/deploy/bench/run.py b/deploy/bench/run.py index 645988d94c..84b894438a 100644 --- a/deploy/bench/run.py +++ b/deploy/bench/run.py @@ -1,20 +1,20 @@ #!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """ OpenSearch GPU Remote Index Build Benchmark =========================================== Steps: - 1. Register MinIO S3 snapshot repository with OpenSearch + 1. Register S3 snapshot repository with OpenSearch 2. Configure cluster settings for GPU remote index build 3. Build kNN index via cuvs-bench: a. Bulk-ingest dataset vectors b. Trigger force-merge to kick off the GPU build - c. Poll MinIO for .faiss files confirming GPU build completion + c. Poll S3 for .faiss files confirming GPU build completion 4. Run cuvs-bench search benchmarks and print results + 5. Write gbench-compatible JSON result files so cuvs_bench.run --data-export + and cuvs_bench.plot can be used for CSV export and plotting """ +import json import os import sys @@ -23,18 +23,21 @@ from cuvs_bench.backends.base import BuildResult, SearchResult -OPENSEARCH_URL = os.environ.get("OPENSEARCH_URL", "http://opensearch:9200") +OPENSEARCH_URL = os.environ.get("OPENSEARCH_URL", "http://opensearch:9200") OPENSEARCH_HOST = os.environ.get("OPENSEARCH_HOST", "opensearch") OPENSEARCH_PORT = int(os.environ.get("OPENSEARCH_PORT", "9200")) -BUILDER_URL = os.environ.get("BUILDER_URL", "http://remote-index-builder:1025") -MINIO_URL = os.environ.get("MINIO_URL", "http://minio:9000") +BUILDER_URL = os.environ.get("BUILDER_URL", "http://remote-index-builder:1025") + +REMOTE_INDEX_BUILD = os.environ.get("REMOTE_INDEX_BUILD", "false").lower() == "true" -DATASET = os.environ.get("DATASET", "sift-128-euclidean") -DATASET_PATH = os.environ.get("DATASET_PATH", "/data/datasets") -BENCH_GROUPS = os.environ.get("BENCH_GROUPS", "test") -K = int(os.environ.get("K", "10")) +S3_BUCKET = os.environ.get("S3_BUCKET", "") +S3_REGION = os.environ.get("AWS_DEFAULT_REGION", "us-east-1") + +DATASET = os.environ.get("DATASET", "sift-128-euclidean") +DATASET_PATH = os.environ.get("DATASET_PATH", "/data/datasets") +BENCH_GROUPS = os.environ.get("BENCH_GROUPS", "test") +K = int(os.environ.get("K", "10")) -BUCKET = "opensearch-vectors" REPO_NAME = "vector-repo" session = requests.Session() @@ -43,25 +46,22 @@ # ── helpers ─────────────────────────────────────────────────────────────────── - def banner(msg: str) -> None: - print(f"\n{'─' * 60}\n {msg}\n{'─' * 60}") + print(f"\n{'─'*60}\n {msg}\n{'─'*60}") # ── OpenSearch setup ────────────────────────────────────────────────────────── - def register_repository() -> None: - banner(f"Registering S3 repository '{REPO_NAME}' (backed by MinIO)") + banner(f"Registering S3 repository '{REPO_NAME}'") r = session.put( f"{OPENSEARCH_URL}/_snapshot/{REPO_NAME}", json={ "type": "s3", "settings": { - # endpoint / protocol / path_style_access are client-level settings - # configured in opensearch/opensearch.yml, not per-repository settings. - "bucket": BUCKET, + "bucket": S3_BUCKET, "base_path": "knn-indexes", + "region": S3_REGION, }, }, ) @@ -70,31 +70,103 @@ def register_repository() -> None: def configure_cluster() -> None: - banner("Enabling GPU remote index build (cluster settings)") + if REMOTE_INDEX_BUILD: + banner("Enabling GPU remote index build (cluster settings)") + settings = { + "knn.remote_index_build.enabled": True, + "knn.remote_index_build.repository": REPO_NAME, + "knn.remote_index_build.service.endpoint": BUILDER_URL, + } + else: + banner("Disabling GPU remote index build (cluster settings)") + settings = { + "knn.remote_index_build.enabled": False, + } r = session.put( f"{OPENSEARCH_URL}/_cluster/settings", - json={ - "persistent": { - "knn.remote_index_build.enabled": True, - "knn.remote_index_build.repository": REPO_NAME, - "knn.remote_index_build.service.endpoint": BUILDER_URL, - } - }, + json={"persistent": settings}, ) r.raise_for_status() print(f" {r.json()}") -# ── results ─────────────────────────────────────────────────────────────────── - +# ── result files ───────────────────────────────────────────────────────────── -def _print_result_row( - params: dict, recall: float, qps: float, latency_ms: float +def write_result_files( + build_results: list, + search_results: list, + dataset: str, + dataset_path: str, + algo: str, + groups: str, + k: int, + batch_size: int = 10000, ) -> None: + """Write gbench-compatible JSON result files. + + Creates files under //result/{build,search}/ in the + same format the C++ backend produces, so ``cuvs_bench.run --data-export`` + and ``cuvs_bench.plot`` work without modification. + """ + build_dir = os.path.join(dataset_path, dataset, "result", "build") + search_dir = os.path.join(dataset_path, dataset, "result", "search") + os.makedirs(build_dir, exist_ok=True) + os.makedirs(search_dir, exist_ok=True) + + # Build JSON – one record per successfully built index. + # data_export.py assumes the build CSV has columns: + # [algo_name, index_name, time, threads, cpu_time, ...] + # so "threads" and "cpu_time" must be present in the JSON (they are not in + # skip_build_cols and therefore get included as columns 3 and 4). + build_benchmarks = [ + { + "name": r.index_path, + "real_time": r.build_time_seconds, + "time_unit": "s", + "threads": 1, + "cpu_time": r.build_time_seconds, + } + for r in build_results + if isinstance(r, BuildResult) and r.success and r.index_path + ] + + # Search JSON – zip build + search results to recover the index name, then + # expand per-search-param entries so each (index, ef_search) is one record. + build_list = [r for r in build_results if isinstance(r, BuildResult)] + search_list = [r for r in search_results if isinstance(r, SearchResult)] + search_benchmarks = [] + for build_r, search_r in zip(build_list, search_list): + if not search_r.success or not build_r.index_path: + continue + for entry in (search_r.metadata or {}).get("per_search_param_results", []): + search_benchmarks.append({ + "name": build_r.index_path, + "real_time": entry["search_time_ms"], + "time_unit": "ms", + "Recall": entry["recall"], + "items_per_second": entry["queries_per_second"], + # Latency field expected by data_export in seconds + "Latency": entry["search_time_ms"] / 1000.0, + }) + + build_file = os.path.join(build_dir, f"{algo},{groups}.json") + search_file = os.path.join(search_dir, f"{algo},{groups},k{k},bs{batch_size}.json") + + with open(build_file, "w") as fh: + json.dump({"benchmarks": build_benchmarks}, fh, indent=2) + with open(search_file, "w") as fh: + json.dump({"benchmarks": search_benchmarks}, fh, indent=2) + + print(f"\n Result files written:") + print(f" {build_file}") + print(f" {search_file}") + + +# ── results ─────────────────────────────────────────────────────────────────── + +def _print_result_row(params: dict, recall: float, qps: float, latency_ms: float) -> None: params_str = ", ".join(f"{k}={v}" for k, v in params.items()) - print( - f" {params_str:<40} {recall:<12.4f} {qps:>8.1f} {latency_ms:>12.2f}" - ) + print(f" {params_str:<40} {recall:<12.4f} {qps:>8.1f} {latency_ms:>12.2f}") def print_results(results: list) -> None: @@ -104,41 +176,38 @@ def print_results(results: list) -> None: print(" No search results returned.") return - header = f" {'params':<40} {'recall@' + str(K):<12} {'QPS':>8} {'latency (ms)':>12}" + header = f" {'params':<40} {'recall@'+str(K):<12} {'QPS':>8} {'latency (ms)':>12}" print(header) print(" " + "─" * (len(header) - 2)) for r in search_results: per_param = (r.metadata or {}).get("per_search_param_results") if per_param: for entry in per_param: - _print_result_row( - entry["search_params"], - entry["recall"], - entry["queries_per_second"], - entry["search_time_ms"], - ) + _print_result_row(entry["search_params"], entry["recall"], entry["queries_per_second"], entry["search_time_ms"]) else: - _print_result_row( - r.search_params[0] if r.search_params else {}, - r.recall, - r.queries_per_second, - r.search_time_ms, - ) + _print_result_row(r.search_params[0] if r.search_params else {}, r.recall, r.queries_per_second, r.search_time_ms) # ── entrypoint ──────────────────────────────────────────────────────────────── - def main() -> None: + if REMOTE_INDEX_BUILD and not S3_BUCKET: + print("ERROR: S3_BUCKET must be set when REMOTE_INDEX_BUILD=true") + sys.exit(1) + print("\n" + "═" * 60) - print(" OpenSearch GPU Remote Index Build Benchmark") + print(" OpenSearch kNN Benchmark") print("═" * 60) - print(f" OpenSearch : {OPENSEARCH_URL}") - print(f" GPU builder: {BUILDER_URL}") - print(f" Dataset : {DATASET} (path: {DATASET_PATH})") - print(f" Groups : {BENCH_GROUPS} k={K}") - - register_repository() + print(f" OpenSearch : {OPENSEARCH_URL}") + print(f" Remote index build : {REMOTE_INDEX_BUILD}") + if REMOTE_INDEX_BUILD: + print(f" GPU builder : {BUILDER_URL}") + print(f" S3 bucket : s3://{S3_BUCKET}/knn-indexes/ (region: {S3_REGION})") + print(f" Dataset : {DATASET} (path: {DATASET_PATH})") + print(f" Groups : {BENCH_GROUPS} k={K}") + + if REMOTE_INDEX_BUILD: + register_repository() configure_cluster() orchestrator = BenchmarkOrchestrator(backend_type="opensearch") @@ -153,24 +222,16 @@ def main() -> None: port=OPENSEARCH_PORT, use_ssl=False, verify_certs=False, - remote_index_build=True, - # S3/MinIO config for GPU build verification (used by the backend) - remote_build_s3_endpoint=MINIO_URL, - remote_build_s3_bucket=BUCKET, - remote_build_s3_prefix="knn-indexes/", - remote_build_s3_access_key="minioadmin", - remote_build_s3_secret_key="minioadmin", + remote_index_build=REMOTE_INDEX_BUILD, ) - # ── Build phase ─────────────────────────────────────────────────────────── - # The backend handles the full GPU build flow: ingest vectors → force merge - # → poll MinIO for .faiss files confirming GPU build completion. - banner("Building index (GPU remote build via cuvs-bench)") + mode = "GPU remote build" if REMOTE_INDEX_BUILD else "CPU" + banner(f"Building index ({mode} via cuvs-bench)") build_results = orchestrator.run_benchmark( build=True, search=False, force=True, - bulk_batch_size=500, + bulk_batch_size=10_000, **bench_kwargs, ) @@ -201,11 +262,22 @@ def main() -> None: print_results(search_results) + write_result_files( + build_results=build_results, + search_results=search_results, + dataset=DATASET, + dataset_path=DATASET_PATH, + algo=bench_kwargs["algorithms"], + groups=BENCH_GROUPS, + k=K, + ) + print("\n" + "═" * 60) print(" Benchmark complete!") print("═" * 60) - print(f"\n OpenSearch : {OPENSEARCH_URL}") - print(" MinIO console : http://localhost:9001 (minioadmin / minioadmin)") + print(f"\n OpenSearch : {OPENSEARCH_URL}") + if REMOTE_INDEX_BUILD: + print(f" GPU builder : {BUILDER_URL}") print() diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml index 9f10f0aaca..3aea398496 100644 --- a/deploy/docker-compose.yml +++ b/deploy/docker-compose.yml @@ -1,85 +1,68 @@ # OpenSearch GPU Remote Index Build — Docker Compose Demo # # Architecture: -# minio S3-compatible object store used to exchange vector data -# minio-init One-shot container that creates the MinIO bucket # opensearch OpenSearch node with kNN plugin (requires 2.17+) # remote-index-builder GPU-accelerated Faiss index builder (FastAPI service) # bench Registers repo + cluster settings, runs cuvs-bench build/search # # Requirements: -# - NVIDIA GPU with CUDA support -# - NVIDIA Container Toolkit https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html # - Docker Compose v2 # - vm.max_map_count >= 262144 on the host: # sudo sysctl -w vm.max_map_count=262144 # +# GPU mode only (--profile gpu): +# - NVIDIA GPU with CUDA support +# - NVIDIA Container Toolkit https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html +# - An S3 bucket and AWS credentials (set via environment variables below) +# +# Required environment variables (set in shell or a .env file): +# DATASET_PATH Absolute path to the directory containing dataset files +# +# GPU mode only: +# S3_BUCKET S3 bucket name for staging vectors and built indexes +# AWS_ACCESS_KEY_ID AWS access key ID +# AWS_SECRET_ACCESS_KEY AWS secret access key +# +# Optional environment variables: +# AWS_SESSION_TOKEN STS session token (required for temporary credentials) +# AWS_DEFAULT_REGION AWS region for the S3 bucket (default: us-east-1) +# REMOTE_INDEX_BUILD Override GPU/CPU mode detection (true/false); normally auto-detected +# DATASET Dataset name (default: sift-128-euclidean) +# BENCH_GROUPS Parameter sweep group: test | base | large (default: test) +# K Number of neighbors to search for (default: 10) +# # Usage: -# docker compose up --build +# CPU: docker compose up --build +# GPU: docker compose --profile gpu up --build # # Data flow: -# OpenSearch flushes a segment → uploads vectors + doc-IDs to MinIO +# OpenSearch flushes a segment → uploads vectors + doc-IDs to S3 # OpenSearch POSTs /_build to the remote-index-builder with the S3 paths -# remote-index-builder downloads from MinIO, builds index on GPU, uploads result -# OpenSearch downloads the finished index from MinIO and merges it into the shard +# remote-index-builder downloads from S3, builds index on GPU, uploads result +# OpenSearch downloads the finished index from S3 and merges it into the shard services: - # ── Object Store ──────────────────────────────────────────────────────────── - minio: - image: minio/minio:latest - environment: - MINIO_ROOT_USER: minioadmin - MINIO_ROOT_PASSWORD: minioadmin - command: server /data --console-address ":9001" - ports: - - "9000:9000" # S3 API - - "9001:9001" # MinIO web console - volumes: - - minio-data:/data - healthcheck: - test: ["CMD", "curl", "-f", "http://localhost:9000/minio/health/live"] - interval: 5s - timeout: 5s - retries: 12 - start_period: 10s - - minio-init: - image: minio/mc:latest - depends_on: - minio: - condition: service_healthy - entrypoint: > - /bin/sh -c " - mc alias set local http://minio:9000 minioadmin minioadmin && - mc mb --ignore-existing local/opensearch-vectors && - echo 'Bucket opensearch-vectors ready' - " - restart: "no" - # ── OpenSearch ─────────────────────────────────────────────────────────────── opensearch: build: context: ./opensearch environment: - - OPENSEARCH_JAVA_OPTS=-Xms4g -Xmx4g - # Bypass the Docker entrypoint script — all config lives in opensearch.yml. - # Keystore credentials are baked into the image (see opensearch/Dockerfile). - command: ["/usr/share/opensearch/bin/opensearch"] + - OPENSEARCH_JAVA_OPTS=-Xms16g -Xmx16g + # S3 credentials — entrypoint.sh writes these into the keystore at startup. + # If unset, OpenSearch starts without S3 configured (remote index build unavailable). + - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID:-} + - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY:-} + - AWS_SESSION_TOKEN=${AWS_SESSION_TOKEN:-} ulimits: nofile: soft: 65536 hard: 65536 volumes: - opensearch-data:/usr/share/opensearch/data - # Mounts S3 client config (endpoint, protocol, path_style_access). - # These are client-level settings and cannot be set in the snapshot API. - ./opensearch/opensearch.yml:/usr/share/opensearch/config/opensearch.yml:ro ports: - "9200:9200" - depends_on: - minio-init: - condition: service_completed_successfully healthcheck: # wait_for_status=yellow blocks until the cluster is at least yellow test: ["CMD-SHELL", "curl -sf 'http://localhost:9200/_cluster/health?wait_for_status=yellow&timeout=5s'"] @@ -90,18 +73,15 @@ services: # ── GPU Index Builder ──────────────────────────────────────────────────────── remote-index-builder: + profiles: [gpu] image: opensearchproject/remote-vector-index-builder:api-latest environment: - - AWS_ACCESS_KEY_ID=minioadmin - - AWS_SECRET_ACCESS_KEY=minioadmin - - AWS_DEFAULT_REGION=us-east-1 - # Route S3 calls to MinIO instead of AWS (requires boto3 >= 1.28) - - AWS_ENDPOINT_URL=http://minio:9000 + - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID:-} + - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY:-} + - AWS_SESSION_TOKEN=${AWS_SESSION_TOKEN:-} + - AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION:-us-east-1} ports: - "1025:1025" - depends_on: - minio-init: - condition: service_completed_successfully healthcheck: test: ["CMD-SHELL", "python3 -c 'import socket; socket.create_connection((\"localhost\", 1025), 2).close()'"] interval: 5s @@ -119,37 +99,29 @@ services: capabilities: [gpu] # ── Benchmark ──────────────────────────────────────────────────────────────── - # Required environment variables (set in shell or a .env file): - # DATASET_PATH Absolute path to the directory containing dataset files - # e.g. export DATASET_PATH=/data/ann-benchmarks - # Optional: - # DATASET Dataset name (default: sift-128-euclidean) - # BENCH_GROUPS Parameter sweep group: test | base | large (default: test) - # K Number of neighbors to search for (default: 10) bench: build: context: ./bench depends_on: opensearch: condition: service_healthy - remote-index-builder: - condition: service_healthy - minio: - condition: service_healthy environment: - OPENSEARCH_URL=http://opensearch:9200 - OPENSEARCH_HOST=opensearch - OPENSEARCH_PORT=9200 - BUILDER_URL=http://remote-index-builder:1025 - - MINIO_URL=http://minio:9000 + - S3_BUCKET=${S3_BUCKET:-} + - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID:-} + - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY:-} + - AWS_SESSION_TOKEN=${AWS_SESSION_TOKEN:-} + - AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION:-us-east-1} - DATASET=${DATASET:-sift-128-euclidean} - DATASET_PATH=/data/datasets - BENCH_GROUPS=${BENCH_GROUPS:-test} - K=${K:-10} volumes: - - ${DATASET_PATH:?DATASET_PATH must be set to the directory containing dataset files}:/data/datasets + - ${DATASET_PATH:-/tmp/datasets}:/data/datasets restart: "no" volumes: - minio-data: opensearch-data: diff --git a/deploy/opensearch/Dockerfile b/deploy/opensearch/Dockerfile index af5c9b5629..84192b53ce 100644 --- a/deploy/opensearch/Dockerfile +++ b/deploy/opensearch/Dockerfile @@ -1,16 +1,12 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -FROM opensearchproject/opensearch:latest +FROM opensearchproject/opensearch:3.6.0 # The repository-s3 plugin is not bundled in the default OpenSearch image but # is required for the remote vector index build feature: OpenSearch uses it to -# upload raw vectors and download GPU-built indexes via the MinIO bucket. +# upload raw vectors and download GPU-built indexes via S3. RUN /usr/share/opensearch/bin/opensearch-plugin install --batch repository-s3 -# Pre-populate the keystore with MinIO credentials. The keystore lives in -# /usr/share/opensearch/config (not the data volume) so this survives restarts. -# The repository-s3 plugin reads credentials exclusively from the keystore. -RUN /usr/share/opensearch/bin/opensearch-keystore create && \ - echo "minioadmin" | /usr/share/opensearch/bin/opensearch-keystore add --stdin s3.client.default.access_key && \ - echo "minioadmin" | /usr/share/opensearch/bin/opensearch-keystore add --stdin s3.client.default.secret_key +# entrypoint.sh populates the keystore from S3_ACCESS_KEY / S3_SECRET_KEY +# environment variables at container startup, then execs opensearch. +COPY --chmod=755 entrypoint.sh /entrypoint.sh + +ENTRYPOINT ["/entrypoint.sh"] diff --git a/deploy/opensearch/entrypoint.sh b/deploy/opensearch/entrypoint.sh new file mode 100644 index 0000000000..eff9812fdb --- /dev/null +++ b/deploy/opensearch/entrypoint.sh @@ -0,0 +1,31 @@ +#!/bin/bash +# +# Populate the OpenSearch keystore with S3 credentials from environment variables, +# then start OpenSearch. Doing this at runtime (not image build time) avoids +# baking credentials into image layers. +# +# Required environment variables: +# AWS_ACCESS_KEY_ID AWS access key ID +# AWS_SECRET_ACCESS_KEY AWS secret access key +# +# Optional environment variables: +# AWS_SESSION_TOKEN STS session token (required for temporary credentials) +# +set -e + +# The repository-s3 plugin reads credentials exclusively from the keystore. +# If credentials are not set, skip keystore setup — S3 and remote index build +# will be unavailable, but OpenSearch itself will start normally. +if [ -n "${AWS_ACCESS_KEY_ID}" ] && [ -n "${AWS_SECRET_ACCESS_KEY}" ]; then + rm -f /usr/share/opensearch/config/opensearch.keystore + /usr/share/opensearch/bin/opensearch-keystore create + printf '%s' "${AWS_ACCESS_KEY_ID}" | /usr/share/opensearch/bin/opensearch-keystore add --stdin s3.client.default.access_key + printf '%s' "${AWS_SECRET_ACCESS_KEY}" | /usr/share/opensearch/bin/opensearch-keystore add --stdin s3.client.default.secret_key + if [ -n "${AWS_SESSION_TOKEN}" ]; then + printf '%s' "${AWS_SESSION_TOKEN}" | /usr/share/opensearch/bin/opensearch-keystore add --stdin s3.client.default.session_token + fi +else + echo "Warning: AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY not set — S3 repository and remote index build will not be available" >&2 +fi + +exec /usr/share/opensearch/bin/opensearch diff --git a/deploy/opensearch/opensearch.yml b/deploy/opensearch/opensearch.yml index ebb17cced3..96a68a6bc0 100644 --- a/deploy/opensearch/opensearch.yml +++ b/deploy/opensearch/opensearch.yml @@ -3,19 +3,9 @@ network.host: 0.0.0.0 # Single-node cluster — suppresses the production bootstrap checks that # require seed_hosts / initial_cluster_manager_nodes to be configured. -# Must be in opensearch.yml (not an env var) because we exec the opensearch -# binary directly rather than going through the Docker entrypoint script. +# Must be in opensearch.yml (not an env var) because our entrypoint.sh +# execs the opensearch binary directly. discovery.type: single-node # Disable the security plugin — no SSL certs needed for this demo plugins.security.disabled: true - -# S3 client used by the repository-s3 plugin. -# Credentials are injected into the keystore at container startup -# (see the `command` block in docker-compose.yml). -s3.client.default.endpoint: "minio:9000" -s3.client.default.protocol: "http" -s3.client.default.path_style_access: "true" -# The async S3 client requires a non-empty region even for non-AWS endpoints. -# "us-east-1" is a placeholder — MinIO ignores the region value entirely. -s3.client.default.region: "us-east-1" diff --git a/deploy/remote-index-build/Dockerfile b/deploy/remote-index-build/Dockerfile index e67d495429..d8b2485a7e 100644 --- a/deploy/remote-index-build/Dockerfile +++ b/deploy/remote-index-build/Dockerfile @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - FROM python:3.11-slim WORKDIR /app RUN pip install --no-cache-dir requests boto3 numpy diff --git a/deploy/remote-index-build/run.py b/deploy/remote-index-build/run.py index 63ba918ff2..a9f6621689 100644 --- a/deploy/remote-index-build/run.py +++ b/deploy/remote-index-build/run.py @@ -1,17 +1,14 @@ #!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """ OpenSearch GPU Remote Index Build — End-to-End Demo ==================================================== Steps: - 1. Register a MinIO-backed S3 snapshot repository with OpenSearch + 1. Register an S3 snapshot repository with OpenSearch 2. Configure cluster settings to enable GPU-based remote index building 3. Create a kNN index (Faiss HNSW / L2) with remote build enabled 4. Ingest 100,000 random 256-dimensional float vectors via the bulk API (8 parallel workers) 5. Flush + force-merge to consolidate segments and trigger the GPU build - 6. Poll MinIO for a .faiss file — hard-fail if the GPU build never completes + 6. Poll S3 for a .faiss file — hard-fail if the GPU build never completes 7. Execute a kNN search and print the top-10 nearest neighbors """ @@ -22,44 +19,43 @@ from concurrent.futures import ThreadPoolExecutor, as_completed import boto3 -from botocore.config import Config import numpy as np import requests OPENSEARCH_URL = os.environ.get("OPENSEARCH_URL", "http://opensearch:9200") -BUILDER_URL = os.environ.get("BUILDER_URL", "http://remote-index-builder:1025") -MINIO_URL = os.environ.get("MINIO_URL", "http://minio:9000") +BUILDER_URL = os.environ.get("BUILDER_URL", "http://remote-index-builder:1025") + +S3_BUCKET = os.environ["S3_BUCKET"] +S3_REGION = os.environ.get("AWS_DEFAULT_REGION", "us-east-1") +# boto3 reads AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY automatically INDEX_NAME = "gpu-demo" -DIMENSION = 256 # matches common embedding model output sizes -NUM_DOCS = 200_000 -BUCKET = "opensearch-vectors" -REPO_NAME = "vector-repo" +DIMENSION = 256 # matches common embedding model output sizes +NUM_DOCS = 200_000 +REPO_NAME = "vector-repo" session = requests.Session() session.headers.update({"Content-Type": "application/json"}) def banner(msg: str) -> None: - print(f"\n{'─' * 60}") + print(f"\n{'─'*60}") print(f" {msg}") - print(f"{'─' * 60}") + print(f"{'─'*60}") # ── configuration ───────────────────────────────────────────────────────────── - def register_repository() -> None: - banner(f"Registering S3 repository '{REPO_NAME}' (backed by MinIO)") + banner(f"Registering S3 repository '{REPO_NAME}'") r = session.put( f"{OPENSEARCH_URL}/_snapshot/{REPO_NAME}", json={ "type": "s3", "settings": { - # endpoint / protocol / path_style_access are client-level settings - # configured in opensearch/opensearch.yml, not per-repository settings. - "bucket": BUCKET, + "bucket": S3_BUCKET, "base_path": "knn-indexes", + "region": S3_REGION, }, }, ) @@ -73,8 +69,8 @@ def configure_cluster() -> None: f"{OPENSEARCH_URL}/_cluster/settings", json={ "persistent": { - "knn.remote_index_build.enabled": True, - "knn.remote_index_build.repository": REPO_NAME, + "knn.remote_index_build.enabled": True, + "knn.remote_index_build.repository": REPO_NAME, "knn.remote_index_build.service.endpoint": BUILDER_URL, } }, @@ -85,7 +81,6 @@ def configure_cluster() -> None: # ── index ───────────────────────────────────────────────────────────────────── - def create_index() -> None: banner(f"Creating kNN index '{INDEX_NAME}'") @@ -97,27 +92,25 @@ def create_index() -> None: f"{OPENSEARCH_URL}/{INDEX_NAME}", json={ "settings": { - "index.knn": True, - "index.knn.remote_index_build.enabled": True, - # Trigger GPU build for segments >= 1 KB (covers any real dataset) - "index.knn.remote_index_build.size.min": "1kb", - "number_of_shards": 1, + "index.knn": True, + "index.knn.remote_index_build.enabled": True, + "number_of_shards": 1, "number_of_replicas": 0, }, "mappings": { "properties": { "vector": { - "type": "knn_vector", + "type": "knn_vector", "dimension": DIMENSION, "method": { - "name": "hnsw", - "engine": "faiss", + "name": "hnsw", + "engine": "faiss", "space_type": "l2", "parameters": {"m": 32, "ef_construction": 512}, }, }, "doc_id": {"type": "integer"}, - "label": {"type": "keyword"}, + "label": {"type": "keyword"}, } }, }, @@ -128,30 +121,17 @@ def create_index() -> None: # ── ingest ──────────────────────────────────────────────────────────────────── - def ingest_vectors() -> None: batch_size = 500 - banner( - f"Ingesting {NUM_DOCS:,} random {DIMENSION}-dim vectors (bulk API, 8 workers)" - ) + banner(f"Ingesting {NUM_DOCS:,} random {DIMENSION}-dim vectors (bulk API, 8 workers)") def send_batch(start: int) -> int: - end = min(start + batch_size, NUM_DOCS) + end = min(start + batch_size, NUM_DOCS) vecs = np.random.randn(end - start, DIMENSION).astype(np.float32) lines = [] for i, vec in enumerate(vecs, start): - lines.append( - json.dumps({"index": {"_index": INDEX_NAME, "_id": str(i)}}) - ) - lines.append( - json.dumps( - { - "vector": vec.tolist(), - "doc_id": i, - "label": f"item-{i:04d}", - } - ) - ) + lines.append(json.dumps({"index": {"_index": INDEX_NAME, "_id": str(i)}})) + lines.append(json.dumps({"vector": vec.tolist(), "doc_id": i, "label": f"item-{i:04d}"})) payload = ("\n".join(lines) + "\n").encode("utf-8") r = session.post( f"{OPENSEARCH_URL}/_bulk", @@ -161,14 +141,8 @@ def send_batch(start: int) -> int: r.raise_for_status() body = r.json() if body.get("errors"): - failed = [ - item["index"]["error"] - for item in body["items"] - if "error" in item.get("index", {}) - ] - print( - f" Warning: {len(failed)} error(s) in batch {start}–{end}: {failed[0]}" - ) + failed = [item["index"]["error"] for item in body["items"] if "error" in item.get("index", {})] + print(f" Warning: {len(failed)} error(s) in batch {start}–{end}: {failed[0]}") return (end - start) - len(failed) return end - start @@ -188,15 +162,10 @@ def send_batch(start: int) -> int: # ── GPU build ───────────────────────────────────────────────────────────────── - def trigger_gpu_build() -> None: banner("Triggering GPU index build via force merge") - print( - " OpenSearch will upload vectors to MinIO, then call the GPU builder." - ) - print( - " force_merge max_num_segments=1 consolidates all segments into one." - ) + print(" OpenSearch will upload vectors to S3, then call the GPU builder.") + print(" force_merge max_num_segments=1 consolidates all segments into one.") r = session.post( f"{OPENSEARCH_URL}/{INDEX_NAME}/_forcemerge?max_num_segments=1", timeout=300, @@ -205,71 +174,52 @@ def trigger_gpu_build() -> None: def verify_gpu_build(timeout: int = 600) -> None: - """Confirm the GPU builder uploaded a .faiss index file to MinIO. + """Confirm the GPU builder uploaded a .faiss index file to S3. The remote-index-builder is the *only* component that writes .faiss files back to the S3 bucket, so their presence is definitive proof that the GPU build completed. The kNN stats API does not expose remote build counters - in OpenSearch 3.x, so we poll MinIO directly via boto3 instead. + in OpenSearch 3.x, so we poll S3 directly via boto3 instead. Exits with code 1 if no .faiss file appears within `timeout` seconds. """ - banner("Verifying GPU index build (polling MinIO for .faiss files)") - print(f" Bucket : {BUCKET}/knn-indexes/") + banner("Verifying GPU index build (polling S3 for .faiss files)") + print(f" Bucket : s3://{S3_BUCKET}/knn-indexes/") print(f" Timeout : {timeout}s (poll interval: 5s)\n") - s3 = boto3.client( - "s3", - endpoint_url=MINIO_URL, - aws_access_key_id="minioadmin", - aws_secret_access_key="minioadmin", - region_name="us-east-1", - config=Config(signature_version="s3v4"), - ) + # boto3 picks up AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_DEFAULT_REGION + # from the environment automatically. + s3 = boto3.client("s3", region_name=S3_REGION) deadline = time.time() + timeout while time.time() < deadline: try: - resp = s3.list_objects_v2(Bucket=BUCKET, Prefix="knn-indexes/") + resp = s3.list_objects_v2(Bucket=S3_BUCKET, Prefix="knn-indexes/") faiss_files = [ obj["Key"] for obj in resp.get("Contents", []) if obj["Key"].endswith(".faiss") ] if faiss_files: - print( - f" PASS: GPU build confirmed — {len(faiss_files)} .faiss file(s) in MinIO:" - ) + print(f" PASS: GPU build confirmed — {len(faiss_files)} .faiss file(s) in S3:") for f in faiss_files: - print(f" s3://{BUCKET}/{f}") + print(f" s3://{S3_BUCKET}/{f}") return remaining = int(deadline - time.time()) all_keys = [obj["Key"] for obj in resp.get("Contents", [])] - print( - f" Waiting for .faiss file... objects={all_keys} ({remaining}s left)" - ) + print(f" Waiting for .faiss file... objects={all_keys} ({remaining}s left)") except Exception as e: - print(f" MinIO check error: {e}") + print(f" S3 check error: {e}") time.sleep(5) - print( - f"\n FAIL: no GPU-built .faiss index appeared in MinIO after {timeout}s" - ) + print(f"\n FAIL: no GPU-built .faiss index appeared in S3 after {timeout}s") print("\n Possible causes:") - print( - " 1. remote-index-builder is unreachable from the OpenSearch container." - ) - print( - f" Verify the container is running and BUILDER_URL={BUILDER_URL} is correct." - ) - print( - " 2. Segment size never exceeded index.knn.remote_index_build.size.min." - ) + print(" 1. remote-index-builder is unreachable from the OpenSearch container.") + print(f" Verify the container is running and BUILDER_URL={BUILDER_URL} is correct.") + print(" 2. Segment size never exceeded index.knn.remote_index_build.size.min.") print(" Try increasing NUM_DOCS or lowering the size.min threshold.") - print( - " 3. No GPU is available inside the remote-index-builder container." - ) + print(" 3. No GPU is available inside the remote-index-builder container.") print(" Check: docker compose logs remote-index-builder") print(" Ensure the NVIDIA Container Toolkit is installed on the host.") sys.exit(1) @@ -277,7 +227,6 @@ def verify_gpu_build(timeout: int = 600) -> None: # ── search ──────────────────────────────────────────────────────────────────── - def search_vectors() -> None: banner("kNN test search (top-5 nearest neighbors)") query_vec = np.random.randn(DIMENSION).astype(np.float32).tolist() @@ -291,30 +240,26 @@ def search_vectors() -> None: }, ) r.raise_for_status() - hits = r.json()["hits"]["hits"] + hits = r.json()["hits"]["hits"] total = r.json()["hits"]["total"]["value"] print(f" Index contains {total} documents") print(f" Top {len(hits)} results:") for rank, hit in enumerate(hits, 1): src = hit["_source"] - print( - f" #{rank:>2} id={hit['_id']:>6} score={hit['_score']:.6f} label={src['label']}" - ) + print(f" #{rank:>2} id={hit['_id']:>6} score={hit['_score']:.6f} label={src['label']}") # ── entrypoint ──────────────────────────────────────────────────────────────── - def main() -> None: print("\n" + "═" * 60) print(" OpenSearch GPU Remote Index Build — End-to-End Demo") print("═" * 60) print(f" OpenSearch : {OPENSEARCH_URL}") print(f" GPU builder: {BUILDER_URL}") - print( - f" Vectors : {NUM_DOCS} × dim={DIMENSION} engine=faiss method=hnsw space=l2" - ) + print(f" S3 bucket : s3://{S3_BUCKET}/knn-indexes/ (region: {S3_REGION})") + print(f" Vectors : {NUM_DOCS} × dim={DIMENSION} engine=faiss method=hnsw space=l2") register_repository() configure_cluster() @@ -328,7 +273,6 @@ def main() -> None: print(" Demo complete!") print("═" * 60) print(f"\n OpenSearch is still running at {OPENSEARCH_URL}") - print(" MinIO console : http://localhost:9001 (minioadmin / minioadmin)") print(f" GPU builder : {BUILDER_URL}") print() From d2627c95e2002736a498ee651469322c3bd49b77 Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Fri, 15 May 2026 16:48:53 -0500 Subject: [PATCH 03/22] Updates Signed-off-by: James Bourbeau --- deploy/DEPLOYMENT.md | 7 ++- deploy/README.md | 34 ++++++----- deploy/bench/Dockerfile | 4 +- deploy/bench/entrypoint.sh | 46 ++++++++++---- deploy/bench/run.py | 101 ++++++++++++++++++++++++------- deploy/docker-compose.yml | 54 +++++++++-------- deploy/opensearch/Dockerfile | 4 +- deploy/remote-index-build/run.py | 100 +++++++++++++++++++++--------- 8 files changed, 243 insertions(+), 107 deletions(-) diff --git a/deploy/DEPLOYMENT.md b/deploy/DEPLOYMENT.md index 181132e309..1a87f3098a 100644 --- a/deploy/DEPLOYMENT.md +++ b/deploy/DEPLOYMENT.md @@ -117,6 +117,7 @@ curl -X PUT http://localhost:9200/my-vectors \ "settings": { "index.knn": true, "index.knn.remote_index_build.enabled": true, + "index.knn.remote_index_build.size.min": "1kb", "number_of_shards": 1, "number_of_replicas": 1 }, @@ -140,11 +141,11 @@ curl -X PUT http://localhost:9200/my-vectors \ }' ``` -GPU builds are only available with the `faiss` engine. The `lucene` engine always builds locally. +GPU builds are only available with the `faiss` engine. The `lucene` engine always builds locally. The low `size.min` value above is useful for demos because it forces small flushed segments onto the remote path; use OpenSearch's default or a larger production threshold for real workloads. ## Verifying the GPU build -The `remote-index-build/` directory contains an end-to-end demo script that ingests 200,000 random vectors, triggers a force-merge, and confirms the GPU build completed by polling S3 for the resulting `.faiss` file. +The `remote-index-build/` directory contains an end-to-end demo script that ingests 200,000 random vectors, flushes the index to trigger remote builds, and confirms the GPU build completed by polling S3 for the resulting `.faiss` file. Run it inside a temporary container on the same Docker network: @@ -157,6 +158,8 @@ docker compose run --rm \ -e AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY} \ -e AWS_SESSION_TOKEN=${AWS_SESSION_TOKEN} \ -e AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION:-us-east-1} \ + -e REMOTE_BUILD_SIZE_MIN=${REMOTE_BUILD_SIZE_MIN:-} \ + -e REMOTE_BUILD_TIMEOUT=${REMOTE_BUILD_TIMEOUT:-1800} \ -v $(pwd)/remote-index-build:/app/remote-index-build \ --no-deps bench \ python remote-index-build/run.py diff --git a/deploy/README.md b/deploy/README.md index d5678d55a1..3ff8b8cab3 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -1,6 +1,6 @@ # OpenSearch kNN Benchmark -Docker Compose benchmark comparing CPU and GPU kNN index builds in OpenSearch using [cuvs-bench](https://github.com/jrbourbeau/cuvs/tree/main/python/cuvs_bench). Supports both local CPU builds and [GPU-accelerated remote index builds](https://docs.opensearch.org/latest/vector-search/remote-index-build/) via the `REMOTE_INDEX_BUILD` environment variable. +Docker Compose benchmark comparing CPU and GPU kNN index builds in OpenSearch using `cuvs-bench`. Supports both local CPU builds and [GPU-accelerated remote index builds](https://docs.opensearch.org/latest/vector-search/remote-index-build/) via the `REMOTE_INDEX_BUILD` environment variable. ## How it works @@ -31,7 +31,7 @@ OpenSearch flushes a segment - **GPU mode only** (`--profile gpu`, `REMOTE_INDEX_BUILD=true`): - NVIDIA GPU with CUDA support - NVIDIA Container Toolkit — [installation guide](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html) - - AWS S3 bucket (or S3-compatible store) for staging vectors and built indexes + - AWS S3 bucket for staging vectors and built indexes ## Usage @@ -50,7 +50,7 @@ export DATASET_PATH=/path/to/ann-benchmark-datasets # directory containing dat GPU mode also requires S3 credentials: ```bash -export S3_BUCKET=my-opensearch-vectors # S3 bucket name +export S3_BUCKET=opepsearch-s3-bucket # S3 bucket name export AWS_ACCESS_KEY_ID= export AWS_SECRET_ACCESS_KEY= ``` @@ -63,6 +63,8 @@ export AWS_DEFAULT_REGION=us-east-1 # AWS region for the S3 bucket (defau export DATASET=sift-128-euclidean # default export BENCH_GROUPS=test # test | base (default: test) export K=10 # number of neighbors (default: 10) +export BATCH_SIZE=10000 # query/bulk batch size (default: 10000) +export REMOTE_BUILD_TIMEOUT=1800 # seconds to wait for remote builds (default: 1800) ``` Start all services: @@ -90,8 +92,8 @@ docker compose down -v 3. **GPU mode only**: Applies cluster settings to enable remote index build and point OpenSearch at the builder service 4. Runs `cuvs-bench` build phase (handled entirely by the OpenSearch backend): - Creates the kNN index and bulk-ingests dataset vectors - - **GPU mode**: Waits for all ingestion-time GPU builds to complete, then force-merges to one segment to trigger the final GPU build and polls the kNN stats API every 5 s until the build is confirmed complete - - **CPU mode**: Force-merges the index to one segment + - **GPU mode**: Flushes segments, waits for all submitted remote GPU builds to complete, and polls the kNN stats API every 5 s until the build is confirmed complete + - **CPU mode**: Flushes and refreshes the local OpenSearch index - Records total build time in the result 5. Runs `cuvs-bench` search phase and prints a recall/QPS/latency table 6. Exports benchmark JSON results to CSV (`cuvs_bench.run --data-export`) @@ -138,13 +140,13 @@ $DATASET_PATH/ | Group | Build params | Search params | Use case | |---|---|---|---| | `test` | 1 combo (m=16, ef_construction=100) | ef_search: 50, 100 | Quick smoke test | -| `base` | 16 combos (m=[32,64,96,128] × ef_construction=[64,128,256,512]) | ef_search: 10–800 | Standard benchmark | +| `base` | 9 combos (m=[32,64,96] × ef_construction=[64,128,256]) | ef_search: 10–800 | Standard benchmark | ## GPU build verification -The cuvs-bench OpenSearch backend polls the kNN stats API every 5 seconds, waiting for `index_build_success_count` to increment by the expected number of new builds and for all in-flight flush and merge operations to reach zero. +The cuvs-bench OpenSearch backend snapshots remote-build stats before ingest, then polls the kNN stats API every 5 seconds until `index_build_success_count` catches up with `build_request_success_count` and all in-flight flush and merge operations reach zero. -The build raises a `TimeoutError` (causing the `bench` container to exit with code 1) if the expected number of successful builds is not confirmed within 600 seconds. +The build raises a `TimeoutError` (causing the `bench` container to exit with code 1) if the expected successful builds are not confirmed within `REMOTE_BUILD_TIMEOUT` seconds. If no remote build is observed shortly after ingest, the backend raises an error that suggests lowering `REMOTE_BUILD_SIZE_MIN`; leave it unset to use OpenSearch's default threshold, or set it explicitly to override that value. ## CPU vs GPU comparison @@ -177,9 +179,9 @@ docker compose run --rm --no-deps bench \ pytest /opt/cuvs/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py -v ``` -### Integration tests (live OpenSearch node) +### Integration tests (live OpenSearch node only) -Requires a running OpenSearch node. S3 credentials are not required for these tests. +Requires a running OpenSearch node. S3 credentials and the GPU profile are not required for these tests. ```bash docker compose up -d --wait opensearch @@ -191,22 +193,24 @@ docker compose run --rm --no-deps \ ### Remote index build integration tests (full GPU stack) -Requires the full stack (OpenSearch and the remote index builder) and S3 credentials. +Requires the full stack (OpenSearch and the remote index builder), S3 credentials, and a GPU-capable host. Export the AWS credentials and region before starting OpenSearch; its S3 keystore is populated at container startup. Use `--profile gpu` when starting the services so Docker Compose includes `remote-index-builder`. The pytest command itself does not need the profile flag because it runs against the already-started services. ```bash +export AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION:-us-east-1} docker compose --profile gpu up -d --wait opensearch remote-index-builder docker compose run --rm --no-deps \ -e OPENSEARCH_URL=http://opensearch:9200 \ -e BUILDER_URL=http://remote-index-builder:1025 \ -e S3_BUCKET=${S3_BUCKET} \ - -e AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID} \ - -e AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY} \ - -e AWS_SESSION_TOKEN=${AWS_SESSION_TOKEN} \ + -e AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION} \ + -e S3_ACCESS_KEY=${AWS_ACCESS_KEY_ID} \ + -e S3_SECRET_KEY=${AWS_SECRET_ACCESS_KEY} \ + -e S3_SESSION_TOKEN=${AWS_SESSION_TOKEN} \ bench \ pytest /opt/cuvs/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py -v -m integration ``` -This runs all integration tests including `TestOpenSearchRemoteIndexBuildIntegration`, which verifies the full GPU build flow end-to-end. +This lets the pytest `integration` marker decide which tests run. With only OpenSearch running, remote-build tests skip because the GPU builder and S3 environment are unavailable. With the GPU stack running, the same marker includes the remote-build coverage. ## Ports diff --git a/deploy/bench/Dockerfile b/deploy/bench/Dockerfile index b4eb2af023..429534d9e2 100644 --- a/deploy/bench/Dockerfile +++ b/deploy/bench/Dockerfile @@ -1,7 +1,9 @@ FROM python:3.11-slim WORKDIR /app -RUN apt-get update && apt-get install -y --no-install-recommends git && rm -rf /var/lib/apt/lists/* +RUN apt-get update \ + && apt-get install -y --no-install-recommends git \ + && rm -rf /var/lib/apt/lists/* # Sparse-clone just the cuvs_bench Python package. # Installing via pip is not possible without CUDA (rapids-build-backend requires diff --git a/deploy/bench/entrypoint.sh b/deploy/bench/entrypoint.sh index 4f48da9c4c..c7ed3cb585 100644 --- a/deploy/bench/entrypoint.sh +++ b/deploy/bench/entrypoint.sh @@ -4,20 +4,44 @@ set -e DATASET="${DATASET:-sift-128-euclidean}" BENCH_GROUPS="${BENCH_GROUPS:-test}" K="${K:-10}" -# Auto-detect GPU mode: remote-index-builder only appears in Docker DNS when -# started via --profile gpu. DNS entries are registered at network setup time -# (before containers run), so this check is reliable by the time entrypoint -# executes (OpenSearch healthy check alone takes 30+ seconds). -if getent hosts remote-index-builder > /dev/null 2>&1; then +BATCH_SIZE="${BATCH_SIZE:-10000}" +ALGORITHM="opensearch_faiss_hnsw" + +wait_for_builder() { echo "remote-index-builder detected — waiting for it to be ready..." until python3 -c 'import socket; socket.create_connection(("remote-index-builder", 1025), 2).close()' 2>/dev/null; do sleep 5 done echo "remote-index-builder is ready." - export REMOTE_INDEX_BUILD=true +} + +if [ -n "${REMOTE_INDEX_BUILD:-}" ]; then + case "${REMOTE_INDEX_BUILD,,}" in + true|1|yes) + wait_for_builder + export REMOTE_INDEX_BUILD=true + ;; + false|0|no) + echo "REMOTE_INDEX_BUILD=false — using CPU build mode." + export REMOTE_INDEX_BUILD=false + ;; + *) + echo "ERROR: REMOTE_INDEX_BUILD must be true or false when set (got '${REMOTE_INDEX_BUILD}')" >&2 + exit 1 + ;; + esac else - echo "remote-index-builder not available — using CPU build mode." - export REMOTE_INDEX_BUILD=false + # Auto-detect GPU mode: remote-index-builder only appears in Docker DNS when + # started via --profile gpu. DNS entries are registered at network setup time + # (before containers run), so this check is reliable by the time entrypoint + # executes (OpenSearch healthy check alone takes 30+ seconds). + if getent hosts remote-index-builder > /dev/null 2>&1; then + wait_for_builder + export REMOTE_INDEX_BUILD=true + else + echo "remote-index-builder not available — using CPU build mode." + export REMOTE_INDEX_BUILD=false + fi fi # Step 1: Download dataset (skipped automatically if already present) @@ -32,17 +56,17 @@ python -u run.py python -m cuvs_bench.run --data-export \ --dataset "$DATASET" \ --dataset-path /data/datasets \ - --algorithms opensearch_faiss_hnsw \ + --algorithms "$ALGORITHM" \ --groups "$BENCH_GROUPS" \ --count "$K" \ - --batch-size 10000 \ + --batch-size "$BATCH_SIZE" \ --search-mode latency # Step 4: Plot — PNGs written to /data/datasets (mounted from host $DATASET_PATH) python -m cuvs_bench.plot \ --dataset "$DATASET" \ --dataset-path /data/datasets \ - --algorithms opensearch_faiss_hnsw \ + --algorithms "$ALGORITHM" \ --groups "$BENCH_GROUPS" \ --count "$K" \ --output-filepath /data/datasets diff --git a/deploy/bench/run.py b/deploy/bench/run.py index 84b894438a..a66b852a83 100644 --- a/deploy/bench/run.py +++ b/deploy/bench/run.py @@ -7,8 +7,8 @@ 2. Configure cluster settings for GPU remote index build 3. Build kNN index via cuvs-bench: a. Bulk-ingest dataset vectors - b. Trigger force-merge to kick off the GPU build - c. Poll S3 for .faiss files confirming GPU build completion + b. Flush segments to kick off remote GPU builds when enabled + c. Poll kNN stats until every submitted remote build completes 4. Run cuvs-bench search benchmarks and print results 5. Write gbench-compatible JSON result files so cuvs_bench.run --data-export and cuvs_bench.plot can be used for CSV export and plotting @@ -29,6 +29,8 @@ BUILDER_URL = os.environ.get("BUILDER_URL", "http://remote-index-builder:1025") REMOTE_INDEX_BUILD = os.environ.get("REMOTE_INDEX_BUILD", "false").lower() == "true" +REMOTE_BUILD_SIZE_MIN = os.environ.get("REMOTE_BUILD_SIZE_MIN", "").strip() +REMOTE_BUILD_TIMEOUT = int(os.environ.get("REMOTE_BUILD_TIMEOUT", "1800")) S3_BUCKET = os.environ.get("S3_BUCKET", "") S3_REGION = os.environ.get("AWS_DEFAULT_REGION", "us-east-1") @@ -37,7 +39,9 @@ DATASET_PATH = os.environ.get("DATASET_PATH", "/data/datasets") BENCH_GROUPS = os.environ.get("BENCH_GROUPS", "test") K = int(os.environ.get("K", "10")) +BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "10000")) +ALGORITHM = "opensearch_faiss_hnsw" REPO_NAME = "vector-repo" session = requests.Session() @@ -50,6 +54,16 @@ def banner(msg: str) -> None: print(f"\n{'─'*60}\n {msg}\n{'─'*60}") +def _recall_for_entry( + result: SearchResult, entry_index: int, entry_count: int +) -> float | None: + # cuVS computes recall in the orchestrator from SearchResult.neighbors. The + # OpenSearch backend returns neighbors for the final search-parameter run. + if entry_index == entry_count - 1: + return float(result.recall) + return None + + # ── OpenSearch setup ────────────────────────────────────────────────────────── def register_repository() -> None: @@ -135,22 +149,33 @@ def write_result_files( build_list = [r for r in build_results if isinstance(r, BuildResult)] search_list = [r for r in search_results if isinstance(r, SearchResult)] search_benchmarks = [] + skipped_without_recall = 0 for build_r, search_r in zip(build_list, search_list): if not search_r.success or not build_r.index_path: continue - for entry in (search_r.metadata or {}).get("per_search_param_results", []): - search_benchmarks.append({ - "name": build_r.index_path, - "real_time": entry["search_time_ms"], - "time_unit": "ms", - "Recall": entry["recall"], - "items_per_second": entry["queries_per_second"], - # Latency field expected by data_export in seconds - "Latency": entry["search_time_ms"] / 1000.0, - }) + per_param = (search_r.metadata or {}).get("per_search_param_results", []) + for entry_index, entry in enumerate(per_param): + recall = _recall_for_entry(search_r, entry_index, len(per_param)) + if recall is None: + skipped_without_recall += 1 + continue + latency_ms = float(entry["search_time_ms"]) + search_benchmarks.append( + { + "name": build_r.index_path, + "real_time": latency_ms, + "time_unit": "ms", + "Recall": recall, + "items_per_second": float(entry["queries_per_second"]), + # Latency field expected by data_export in seconds + "Latency": latency_ms / 1000.0, + } + ) build_file = os.path.join(build_dir, f"{algo},{groups}.json") - search_file = os.path.join(search_dir, f"{algo},{groups},k{k},bs{batch_size}.json") + search_file = os.path.join( + search_dir, f"{algo},{groups},k{k},bs{batch_size}.json" + ) with open(build_file, "w") as fh: json.dump({"benchmarks": build_benchmarks}, fh, indent=2) @@ -160,13 +185,23 @@ def write_result_files( print(f"\n Result files written:") print(f" {build_file}") print(f" {search_file}") + if skipped_without_recall: + print( + " skipped " + f"{skipped_without_recall} search rows without per-parameter recall" + ) # ── results ─────────────────────────────────────────────────────────────────── -def _print_result_row(params: dict, recall: float, qps: float, latency_ms: float) -> None: +def _print_result_row( + params: dict, recall: float | None, qps: float | None, latency_ms: float | None +) -> None: params_str = ", ".join(f"{k}={v}" for k, v in params.items()) - print(f" {params_str:<40} {recall:<12.4f} {qps:>8.1f} {latency_ms:>12.2f}") + recall_str = "n/a" if recall is None else f"{recall:.4f}" + qps_str = "n/a" if qps is None else f"{qps:.1f}" + latency_str = "n/a" if latency_ms is None else f"{latency_ms:.2f}" + print(f" {params_str:<40} {recall_str:<12} {qps_str:>8} {latency_str:>12}") def print_results(results: list) -> None: @@ -179,13 +214,25 @@ def print_results(results: list) -> None: header = f" {'params':<40} {'recall@'+str(K):<12} {'QPS':>8} {'latency (ms)':>12}" print(header) print(" " + "─" * (len(header) - 2)) + missing_recall_rows = 0 for r in search_results: - per_param = (r.metadata or {}).get("per_search_param_results") - if per_param: - for entry in per_param: - _print_result_row(entry["search_params"], entry["recall"], entry["queries_per_second"], entry["search_time_ms"]) - else: - _print_result_row(r.search_params[0] if r.search_params else {}, r.recall, r.queries_per_second, r.search_time_ms) + per_param = (r.metadata or {}).get("per_search_param_results", []) + entry_count = len(per_param) + for entry_index, entry in enumerate(per_param): + recall = _recall_for_entry(r, entry_index, entry_count) + if recall is None: + missing_recall_rows += 1 + _print_result_row( + entry["search_params"], + recall, + float(entry["queries_per_second"]), + float(entry["search_time_ms"]), + ) + if missing_recall_rows: + print( + "\n Note: this cuVS version does not report recall for every " + "OpenSearch search parameter row." + ) # ── entrypoint ──────────────────────────────────────────────────────────────── @@ -203,7 +250,10 @@ def main() -> None: if REMOTE_INDEX_BUILD: print(f" GPU builder : {BUILDER_URL}") print(f" S3 bucket : s3://{S3_BUCKET}/knn-indexes/ (region: {S3_REGION})") + print(f" Build size minimum : {REMOTE_BUILD_SIZE_MIN or 'OpenSearch default'}") + print(f" Build timeout : {REMOTE_BUILD_TIMEOUT}s") print(f" Dataset : {DATASET} (path: {DATASET_PATH})") + print(f" Algorithm : {ALGORITHM}") print(f" Groups : {BENCH_GROUPS} k={K}") if REMOTE_INDEX_BUILD: @@ -216,7 +266,7 @@ def main() -> None: bench_kwargs = dict( dataset=DATASET, dataset_path=DATASET_PATH, - algorithms="opensearch_faiss_hnsw", + algorithms=ALGORITHM, groups=BENCH_GROUPS, host=OPENSEARCH_HOST, port=OPENSEARCH_PORT, @@ -224,6 +274,10 @@ def main() -> None: verify_certs=False, remote_index_build=REMOTE_INDEX_BUILD, ) + if REMOTE_INDEX_BUILD: + bench_kwargs["remote_build_timeout"] = REMOTE_BUILD_TIMEOUT + if REMOTE_BUILD_SIZE_MIN: + bench_kwargs["remote_build_size_min"] = REMOTE_BUILD_SIZE_MIN # ── Build phase ─────────────────────────────────────────────────────────── mode = "GPU remote build" if REMOTE_INDEX_BUILD else "CPU" banner(f"Building index ({mode} via cuvs-bench)") @@ -231,7 +285,7 @@ def main() -> None: build=True, search=False, force=True, - bulk_batch_size=10_000, + bulk_batch_size=BATCH_SIZE, **bench_kwargs, ) @@ -270,6 +324,7 @@ def main() -> None: algo=bench_kwargs["algorithms"], groups=BENCH_GROUPS, k=K, + batch_size=BATCH_SIZE, ) print("\n" + "═" * 60) diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml index 3aea398496..e0da2e1adc 100644 --- a/deploy/docker-compose.yml +++ b/deploy/docker-compose.yml @@ -27,9 +27,12 @@ # AWS_SESSION_TOKEN STS session token (required for temporary credentials) # AWS_DEFAULT_REGION AWS region for the S3 bucket (default: us-east-1) # REMOTE_INDEX_BUILD Override GPU/CPU mode detection (true/false); normally auto-detected +# REMOTE_BUILD_SIZE_MIN Optional minimum segment size override for remote builds +# REMOTE_BUILD_TIMEOUT Remote build wait timeout in seconds (default: 1800) # DATASET Dataset name (default: sift-128-euclidean) -# BENCH_GROUPS Parameter sweep group: test | base | large (default: test) +# BENCH_GROUPS Parameter sweep group: test | base (default: test) # K Number of neighbors to search for (default: 10) +# BATCH_SIZE cuvs-bench query/bulk batch size (default: 10000) # # Usage: # CPU: docker compose up --build @@ -41,6 +44,12 @@ # remote-index-builder downloads from S3, builds index on GPU, uploads result # OpenSearch downloads the finished index from S3 and merges it into the shard +x-aws-env: &aws-env + AWS_ACCESS_KEY_ID: ${AWS_ACCESS_KEY_ID:-} + AWS_SECRET_ACCESS_KEY: ${AWS_SECRET_ACCESS_KEY:-} + AWS_SESSION_TOKEN: ${AWS_SESSION_TOKEN:-} + AWS_DEFAULT_REGION: ${AWS_DEFAULT_REGION:-us-east-1} + services: # ── OpenSearch ─────────────────────────────────────────────────────────────── @@ -48,12 +57,8 @@ services: build: context: ./opensearch environment: - - OPENSEARCH_JAVA_OPTS=-Xms16g -Xmx16g - # S3 credentials — entrypoint.sh writes these into the keystore at startup. - # If unset, OpenSearch starts without S3 configured (remote index build unavailable). - - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID:-} - - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY:-} - - AWS_SESSION_TOKEN=${AWS_SESSION_TOKEN:-} + <<: *aws-env + OPENSEARCH_JAVA_OPTS: -Xms16g -Xmx16g ulimits: nofile: soft: 65536 @@ -76,10 +81,7 @@ services: profiles: [gpu] image: opensearchproject/remote-vector-index-builder:api-latest environment: - - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID:-} - - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY:-} - - AWS_SESSION_TOKEN=${AWS_SESSION_TOKEN:-} - - AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION:-us-east-1} + <<: *aws-env ports: - "1025:1025" healthcheck: @@ -106,19 +108,23 @@ services: opensearch: condition: service_healthy environment: - - OPENSEARCH_URL=http://opensearch:9200 - - OPENSEARCH_HOST=opensearch - - OPENSEARCH_PORT=9200 - - BUILDER_URL=http://remote-index-builder:1025 - - S3_BUCKET=${S3_BUCKET:-} - - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID:-} - - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY:-} - - AWS_SESSION_TOKEN=${AWS_SESSION_TOKEN:-} - - AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION:-us-east-1} - - DATASET=${DATASET:-sift-128-euclidean} - - DATASET_PATH=/data/datasets - - BENCH_GROUPS=${BENCH_GROUPS:-test} - - K=${K:-10} + <<: *aws-env + OPENSEARCH_URL: http://opensearch:9200 + OPENSEARCH_HOST: opensearch + OPENSEARCH_PORT: "9200" + BUILDER_URL: http://remote-index-builder:1025 + REMOTE_INDEX_BUILD: ${REMOTE_INDEX_BUILD:-} + REMOTE_BUILD_SIZE_MIN: ${REMOTE_BUILD_SIZE_MIN:-} + REMOTE_BUILD_TIMEOUT: ${REMOTE_BUILD_TIMEOUT:-1800} + S3_BUCKET: ${S3_BUCKET:-} + S3_ACCESS_KEY: ${AWS_ACCESS_KEY_ID:-} + S3_SECRET_KEY: ${AWS_SECRET_ACCESS_KEY:-} + S3_SESSION_TOKEN: ${AWS_SESSION_TOKEN:-} + DATASET: ${DATASET:-sift-128-euclidean} + DATASET_PATH: /data/datasets + BENCH_GROUPS: ${BENCH_GROUPS:-test} + K: ${K:-10} + BATCH_SIZE: ${BATCH_SIZE:-10000} volumes: - ${DATASET_PATH:-/tmp/datasets}:/data/datasets restart: "no" diff --git a/deploy/opensearch/Dockerfile b/deploy/opensearch/Dockerfile index 84192b53ce..b90fe9a073 100644 --- a/deploy/opensearch/Dockerfile +++ b/deploy/opensearch/Dockerfile @@ -5,8 +5,8 @@ FROM opensearchproject/opensearch:3.6.0 # upload raw vectors and download GPU-built indexes via S3. RUN /usr/share/opensearch/bin/opensearch-plugin install --batch repository-s3 -# entrypoint.sh populates the keystore from S3_ACCESS_KEY / S3_SECRET_KEY -# environment variables at container startup, then execs opensearch. +# entrypoint.sh populates the keystore from AWS credential environment +# variables at container startup, then execs opensearch. COPY --chmod=755 entrypoint.sh /entrypoint.sh ENTRYPOINT ["/entrypoint.sh"] diff --git a/deploy/remote-index-build/run.py b/deploy/remote-index-build/run.py index a9f6621689..6e3af60d88 100644 --- a/deploy/remote-index-build/run.py +++ b/deploy/remote-index-build/run.py @@ -6,8 +6,8 @@ 1. Register an S3 snapshot repository with OpenSearch 2. Configure cluster settings to enable GPU-based remote index building 3. Create a kNN index (Faiss HNSW / L2) with remote build enabled - 4. Ingest 100,000 random 256-dimensional float vectors via the bulk API (8 parallel workers) - 5. Flush + force-merge to consolidate segments and trigger the GPU build + 4. Ingest 200,000 random 256-dimensional float vectors via the bulk API (8 parallel workers) + 5. Flush segments to trigger the GPU build 6. Poll S3 for a .faiss file — hard-fail if the GPU build never completes 7. Execute a kNN search and print the top-10 nearest neighbors """ @@ -33,6 +33,8 @@ DIMENSION = 256 # matches common embedding model output sizes NUM_DOCS = 200_000 REPO_NAME = "vector-repo" +REMOTE_BUILD_SIZE_MIN = os.environ.get("REMOTE_BUILD_SIZE_MIN", "").strip() +REMOTE_BUILD_TIMEOUT = int(os.environ.get("REMOTE_BUILD_TIMEOUT", "1800")) session = requests.Session() session.headers.update({"Content-Type": "application/json"}) @@ -88,15 +90,21 @@ def create_index() -> None: if resp.status_code == 200: print(" Deleted existing index") + index_settings = { + "index.knn": True, + "index.knn.remote_index_build.enabled": True, + "number_of_shards": 1, + "number_of_replicas": 0, + } + if REMOTE_BUILD_SIZE_MIN: + index_settings["index.knn.remote_index_build.size.min"] = ( + REMOTE_BUILD_SIZE_MIN + ) + r = session.put( f"{OPENSEARCH_URL}/{INDEX_NAME}", json={ - "settings": { - "index.knn": True, - "index.knn.remote_index_build.enabled": True, - "number_of_shards": 1, - "number_of_replicas": 0, - }, + "settings": index_settings, "mappings": { "properties": { "vector": { @@ -123,15 +131,28 @@ def create_index() -> None: def ingest_vectors() -> None: batch_size = 500 - banner(f"Ingesting {NUM_DOCS:,} random {DIMENSION}-dim vectors (bulk API, 8 workers)") + banner( + f"Ingesting {NUM_DOCS:,} random {DIMENSION}-dim vectors " + "(bulk API, 8 workers)" + ) def send_batch(start: int) -> int: end = min(start + batch_size, NUM_DOCS) vecs = np.random.randn(end - start, DIMENSION).astype(np.float32) lines = [] for i, vec in enumerate(vecs, start): - lines.append(json.dumps({"index": {"_index": INDEX_NAME, "_id": str(i)}})) - lines.append(json.dumps({"vector": vec.tolist(), "doc_id": i, "label": f"item-{i:04d}"})) + lines.append( + json.dumps({"index": {"_index": INDEX_NAME, "_id": str(i)}}) + ) + lines.append( + json.dumps( + { + "vector": vec.tolist(), + "doc_id": i, + "label": f"item-{i:04d}", + } + ) + ) payload = ("\n".join(lines) + "\n").encode("utf-8") r = session.post( f"{OPENSEARCH_URL}/_bulk", @@ -141,8 +162,15 @@ def send_batch(start: int) -> int: r.raise_for_status() body = r.json() if body.get("errors"): - failed = [item["index"]["error"] for item in body["items"] if "error" in item.get("index", {})] - print(f" Warning: {len(failed)} error(s) in batch {start}–{end}: {failed[0]}") + failed = [ + item["index"]["error"] + for item in body["items"] + if "error" in item.get("index", {}) + ] + print( + f" Warning: {len(failed)} error(s) in batch " + f"{start}–{end}: {failed[0]}" + ) return (end - start) - len(failed) return end - start @@ -155,25 +183,25 @@ def send_batch(start: int) -> int: if ingested % 10_000 == 0 or ingested >= NUM_DOCS: print(f" Ingested {ingested:,}/{NUM_DOCS:,}") - session.post(f"{OPENSEARCH_URL}/{INDEX_NAME}/_flush") + session.post(f"{OPENSEARCH_URL}/{INDEX_NAME}/_refresh") r = session.get(f"{OPENSEARCH_URL}/{INDEX_NAME}/_count") - print(f" Document count after flush: {r.json()['count']:,}") + print(f" Document count after ingest: {r.json()['count']:,}") # ── GPU build ───────────────────────────────────────────────────────────────── def trigger_gpu_build() -> None: - banner("Triggering GPU index build via force merge") - print(" OpenSearch will upload vectors to S3, then call the GPU builder.") - print(" force_merge max_num_segments=1 consolidates all segments into one.") - r = session.post( - f"{OPENSEARCH_URL}/{INDEX_NAME}/_forcemerge?max_num_segments=1", - timeout=300, + banner("Triggering GPU index build via flush") + print( + " OpenSearch will upload eligible flushed segments to S3, " + "then call the GPU builder." ) - print(f" Force merge HTTP {r.status_code}") + r = session.post(f"{OPENSEARCH_URL}/{INDEX_NAME}/_flush", timeout=300) + r.raise_for_status() + print(f" Flush complete: {r.json()}") -def verify_gpu_build(timeout: int = 600) -> None: +def verify_gpu_build(timeout: int = REMOTE_BUILD_TIMEOUT) -> None: """Confirm the GPU builder uploaded a .faiss index file to S3. The remote-index-builder is the *only* component that writes .faiss files @@ -201,14 +229,20 @@ def verify_gpu_build(timeout: int = 600) -> None: if obj["Key"].endswith(".faiss") ] if faiss_files: - print(f" PASS: GPU build confirmed — {len(faiss_files)} .faiss file(s) in S3:") + print( + " PASS: GPU build confirmed — " + f"{len(faiss_files)} .faiss file(s) in S3:" + ) for f in faiss_files: print(f" s3://{S3_BUCKET}/{f}") return remaining = int(deadline - time.time()) all_keys = [obj["Key"] for obj in resp.get("Contents", [])] - print(f" Waiting for .faiss file... objects={all_keys} ({remaining}s left)") + print( + f" Waiting for .faiss file... objects={all_keys} " + f"({remaining}s left)" + ) except Exception as e: print(f" S3 check error: {e}") time.sleep(5) @@ -216,7 +250,10 @@ def verify_gpu_build(timeout: int = 600) -> None: print(f"\n FAIL: no GPU-built .faiss index appeared in S3 after {timeout}s") print("\n Possible causes:") print(" 1. remote-index-builder is unreachable from the OpenSearch container.") - print(f" Verify the container is running and BUILDER_URL={BUILDER_URL} is correct.") + print( + " Verify the container is running and " + f"BUILDER_URL={BUILDER_URL} is correct." + ) print(" 2. Segment size never exceeded index.knn.remote_index_build.size.min.") print(" Try increasing NUM_DOCS or lowering the size.min threshold.") print(" 3. No GPU is available inside the remote-index-builder container.") @@ -228,7 +265,7 @@ def verify_gpu_build(timeout: int = 600) -> None: # ── search ──────────────────────────────────────────────────────────────────── def search_vectors() -> None: - banner("kNN test search (top-5 nearest neighbors)") + banner("kNN test search (top-10 nearest neighbors)") query_vec = np.random.randn(DIMENSION).astype(np.float32).tolist() r = session.post( @@ -259,14 +296,19 @@ def main() -> None: print(f" OpenSearch : {OPENSEARCH_URL}") print(f" GPU builder: {BUILDER_URL}") print(f" S3 bucket : s3://{S3_BUCKET}/knn-indexes/ (region: {S3_REGION})") - print(f" Vectors : {NUM_DOCS} × dim={DIMENSION} engine=faiss method=hnsw space=l2") + print( + f" Vectors : {NUM_DOCS} × dim={DIMENSION} " + "engine=faiss method=hnsw space=l2" + ) + print(f" Build size minimum: {REMOTE_BUILD_SIZE_MIN or 'OpenSearch default'}") + print(f" Build timeout: {REMOTE_BUILD_TIMEOUT}s") register_repository() configure_cluster() create_index() ingest_vectors() trigger_gpu_build() - verify_gpu_build() + verify_gpu_build(timeout=REMOTE_BUILD_TIMEOUT) search_vectors() print("\n" + "═" * 60) From 80b8251b7aecc7d2ceefa0b389884c89969955dd Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Mon, 18 May 2026 10:45:53 -0500 Subject: [PATCH 04/22] Update Signed-off-by: James Bourbeau --- deploy/DEPLOYMENT.md | 20 ++++++++++++-------- deploy/README.md | 17 +++++++++++++---- deploy/docker-compose.yml | 18 ++++++++---------- deploy/opensearch/Dockerfile | 2 +- deploy/opensearch/entrypoint.sh | 18 +++++++++--------- deploy/remote-index-build/run.py | 7 ++++--- 6 files changed, 47 insertions(+), 35 deletions(-) diff --git a/deploy/DEPLOYMENT.md b/deploy/DEPLOYMENT.md index 1a87f3098a..28b184d206 100644 --- a/deploy/DEPLOYMENT.md +++ b/deploy/DEPLOYMENT.md @@ -33,7 +33,7 @@ sequenceDiagram | `opensearch` | custom build of `opensearchproject/opensearch:3.6.0` | OpenSearch node with kNN plugin and `repository-s3` plugin | | `remote-index-builder` | `opensearchproject/remote-vector-index-builder:api-latest` | GPU-accelerated Faiss HNSW index builder | -The custom OpenSearch image adds the `repository-s3` plugin (required for S3-backed vector staging) and populates the S3 keystore from environment variables at startup so credentials are never baked into image layers. +The custom OpenSearch image adds the `repository-s3` plugin (required for S3-backed vector staging). When static AWS keys are provided, the image populates the S3 keystore at startup so credentials are never baked into image layers. Without static keys, OpenSearch can fall back to the AWS default credential provider chain, such as an EC2 instance role. ## Requirements @@ -50,19 +50,24 @@ Set the host kernel parameter required by OpenSearch (once per reboot): sudo sysctl -w vm.max_map_count=262144 ``` -Set required environment variables: +Set the required bucket name: ```bash export S3_BUCKET= +``` + +If you are using static credentials instead of a default AWS credential provider, also export: + +```bash export AWS_ACCESS_KEY_ID= export AWS_SECRET_ACCESS_KEY= +export AWS_SESSION_TOKEN= # required for temporary (STS) credentials ``` -Optionally configure the region and session token for temporary credentials: +Optionally configure the region: ```bash export AWS_DEFAULT_REGION=us-east-1 # default: us-east-1 -export AWS_SESSION_TOKEN= # required for temporary (STS) credentials ``` Start OpenSearch and the GPU builder: @@ -154,9 +159,6 @@ docker compose run --rm \ -e OPENSEARCH_URL=http://opensearch:9200 \ -e BUILDER_URL=http://remote-index-builder:1025 \ -e S3_BUCKET=${S3_BUCKET} \ - -e AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID} \ - -e AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY} \ - -e AWS_SESSION_TOKEN=${AWS_SESSION_TOKEN} \ -e AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION:-us-east-1} \ -e REMOTE_BUILD_SIZE_MIN=${REMOTE_BUILD_SIZE_MIN:-} \ -e REMOTE_BUILD_TIMEOUT=${REMOTE_BUILD_TIMEOUT:-1800} \ @@ -165,6 +167,8 @@ docker compose run --rm \ python remote-index-build/run.py ``` +Static AWS credential environment variables are passed through by the `bench` service when they are exported on the host. + Or run it directly if you have Python and the dependencies installed locally (`boto3`, `numpy`, `requests`), pointing `OPENSEARCH_URL` at `http://localhost:9200`. A successful run prints a `.faiss` file path in S3 and returns top-10 nearest-neighbor results. @@ -184,7 +188,7 @@ This setup is a working demonstration, not a production-hardened deployment. Key - **Security plugin**: `opensearch.yml` has `plugins.security.disabled: true`. Re-enable it and configure TLS and authentication for any non-local deployment. - **Single-node cluster**: `discovery.type: single-node` bypasses multi-node bootstrap checks. Replace with a properly configured multi-node cluster for production. - **Replicas**: The demo uses `number_of_replicas: 0`. Set this to at least `1` for production workloads. -- **S3 permissions**: The IAM credentials need `s3:GetObject`, `s3:PutObject`, `s3:ListBucket`, and `s3:DeleteObject` on the staging bucket. +- **S3 permissions**: The IAM principal used by OpenSearch and the builder needs `s3:GetObject`, `s3:PutObject`, `s3:ListBucket`, and `s3:DeleteObject` on the staging bucket. ## Ports diff --git a/deploy/README.md b/deploy/README.md index 3ff8b8cab3..a6181a0cf2 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -47,18 +47,25 @@ Set required environment variables: export DATASET_PATH=/path/to/ann-benchmark-datasets # directory containing dataset files ``` -GPU mode also requires S3 credentials: +GPU mode also requires an S3 bucket. Static AWS keys are supported, but optional +when the containers can use another AWS default credential provider such as an +EC2 instance role: + +```bash +export S3_BUCKET=opensearch-s3-bucket # S3 bucket name +``` + +If you are using static credentials instead of a default provider, also export: ```bash -export S3_BUCKET=opepsearch-s3-bucket # S3 bucket name export AWS_ACCESS_KEY_ID= export AWS_SECRET_ACCESS_KEY= +export AWS_SESSION_TOKEN= # required when using temporary (STS) credentials ``` Optionally configure the benchmark: ```bash -export AWS_SESSION_TOKEN= # required when using temporary (STS) credentials export AWS_DEFAULT_REGION=us-east-1 # AWS region for the S3 bucket (default: us-east-1) export DATASET=sift-128-euclidean # default export BENCH_GROUPS=test # test | base (default: test) @@ -193,7 +200,9 @@ docker compose run --rm --no-deps \ ### Remote index build integration tests (full GPU stack) -Requires the full stack (OpenSearch and the remote index builder), S3 credentials, and a GPU-capable host. Export the AWS credentials and region before starting OpenSearch; its S3 keystore is populated at container startup. Use `--profile gpu` when starting the services so Docker Compose includes `remote-index-builder`. The pytest command itself does not need the profile flag because it runs against the already-started services. +Requires the full stack (OpenSearch and the remote index builder), S3 access, and a GPU-capable host. Export the S3 bucket and region before starting OpenSearch. If you are using static AWS keys, export them before startup so OpenSearch can populate its S3 keystore; otherwise the containers can use the AWS default credential provider chain. Use `--profile gpu` when starting the services so Docker Compose includes `remote-index-builder`. The pytest command itself does not need the profile flag because it runs against the already-started services. + +When using static AWS keys, map them into the test fixture's `S3_*` variable names: ```bash export AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION:-us-east-1} diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml index e0da2e1adc..4680107006 100644 --- a/deploy/docker-compose.yml +++ b/deploy/docker-compose.yml @@ -13,18 +13,19 @@ # GPU mode only (--profile gpu): # - NVIDIA GPU with CUDA support # - NVIDIA Container Toolkit https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html -# - An S3 bucket and AWS credentials (set via environment variables below) +# - An S3 bucket and credentials from AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY +# or another AWS default credential provider, such as an EC2 instance role # # Required environment variables (set in shell or a .env file): # DATASET_PATH Absolute path to the directory containing dataset files # # GPU mode only: # S3_BUCKET S3 bucket name for staging vectors and built indexes -# AWS_ACCESS_KEY_ID AWS access key ID -# AWS_SECRET_ACCESS_KEY AWS secret access key # # Optional environment variables: -# AWS_SESSION_TOKEN STS session token (required for temporary credentials) +# AWS_ACCESS_KEY_ID AWS access key ID (optional if using default credentials) +# AWS_SECRET_ACCESS_KEY AWS secret access key (optional if using default credentials) +# AWS_SESSION_TOKEN STS session token (required for temporary static credentials) # AWS_DEFAULT_REGION AWS region for the S3 bucket (default: us-east-1) # REMOTE_INDEX_BUILD Override GPU/CPU mode detection (true/false); normally auto-detected # REMOTE_BUILD_SIZE_MIN Optional minimum segment size override for remote builds @@ -45,9 +46,9 @@ # OpenSearch downloads the finished index from S3 and merges it into the shard x-aws-env: &aws-env - AWS_ACCESS_KEY_ID: ${AWS_ACCESS_KEY_ID:-} - AWS_SECRET_ACCESS_KEY: ${AWS_SECRET_ACCESS_KEY:-} - AWS_SESSION_TOKEN: ${AWS_SESSION_TOKEN:-} + AWS_ACCESS_KEY_ID: + AWS_SECRET_ACCESS_KEY: + AWS_SESSION_TOKEN: AWS_DEFAULT_REGION: ${AWS_DEFAULT_REGION:-us-east-1} services: @@ -117,9 +118,6 @@ services: REMOTE_BUILD_SIZE_MIN: ${REMOTE_BUILD_SIZE_MIN:-} REMOTE_BUILD_TIMEOUT: ${REMOTE_BUILD_TIMEOUT:-1800} S3_BUCKET: ${S3_BUCKET:-} - S3_ACCESS_KEY: ${AWS_ACCESS_KEY_ID:-} - S3_SECRET_KEY: ${AWS_SECRET_ACCESS_KEY:-} - S3_SESSION_TOKEN: ${AWS_SESSION_TOKEN:-} DATASET: ${DATASET:-sift-128-euclidean} DATASET_PATH: /data/datasets BENCH_GROUPS: ${BENCH_GROUPS:-test} diff --git a/deploy/opensearch/Dockerfile b/deploy/opensearch/Dockerfile index b90fe9a073..06be41a26c 100644 --- a/deploy/opensearch/Dockerfile +++ b/deploy/opensearch/Dockerfile @@ -6,7 +6,7 @@ FROM opensearchproject/opensearch:3.6.0 RUN /usr/share/opensearch/bin/opensearch-plugin install --batch repository-s3 # entrypoint.sh populates the keystore from AWS credential environment -# variables at container startup, then execs opensearch. +# variables when provided, then execs opensearch. COPY --chmod=755 entrypoint.sh /entrypoint.sh ENTRYPOINT ["/entrypoint.sh"] diff --git a/deploy/opensearch/entrypoint.sh b/deploy/opensearch/entrypoint.sh index eff9812fdb..4ea84342a8 100644 --- a/deploy/opensearch/entrypoint.sh +++ b/deploy/opensearch/entrypoint.sh @@ -1,21 +1,18 @@ #!/bin/bash # -# Populate the OpenSearch keystore with S3 credentials from environment variables, +# If static S3 credentials are provided, write them to the OpenSearch keystore, # then start OpenSearch. Doing this at runtime (not image build time) avoids -# baking credentials into image layers. +# baking credentials into image layers. If static credentials are not provided, +# repository-s3 can fall back to the AWS default credential provider chain, such +# as an EC2 instance role. # -# Required environment variables: +# Static credential environment variables: # AWS_ACCESS_KEY_ID AWS access key ID # AWS_SECRET_ACCESS_KEY AWS secret access key -# -# Optional environment variables: # AWS_SESSION_TOKEN STS session token (required for temporary credentials) # set -e -# The repository-s3 plugin reads credentials exclusively from the keystore. -# If credentials are not set, skip keystore setup — S3 and remote index build -# will be unavailable, but OpenSearch itself will start normally. if [ -n "${AWS_ACCESS_KEY_ID}" ] && [ -n "${AWS_SECRET_ACCESS_KEY}" ]; then rm -f /usr/share/opensearch/config/opensearch.keystore /usr/share/opensearch/bin/opensearch-keystore create @@ -24,8 +21,11 @@ if [ -n "${AWS_ACCESS_KEY_ID}" ] && [ -n "${AWS_SECRET_ACCESS_KEY}" ]; then if [ -n "${AWS_SESSION_TOKEN}" ]; then printf '%s' "${AWS_SESSION_TOKEN}" | /usr/share/opensearch/bin/opensearch-keystore add --stdin s3.client.default.session_token fi +elif [ -n "${AWS_ACCESS_KEY_ID}" ] || [ -n "${AWS_SECRET_ACCESS_KEY}" ]; then + echo "ERROR: AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY must be set together" >&2 + exit 1 else - echo "Warning: AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY not set — S3 repository and remote index build will not be available" >&2 + echo "AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY not set; using the AWS default credential provider chain for S3" >&2 fi exec /usr/share/opensearch/bin/opensearch diff --git a/deploy/remote-index-build/run.py b/deploy/remote-index-build/run.py index 6e3af60d88..de6ab477a0 100644 --- a/deploy/remote-index-build/run.py +++ b/deploy/remote-index-build/run.py @@ -27,7 +27,8 @@ S3_BUCKET = os.environ["S3_BUCKET"] S3_REGION = os.environ.get("AWS_DEFAULT_REGION", "us-east-1") -# boto3 reads AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY automatically +# boto3 uses AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY when set, otherwise it +# falls through to the default credential provider chain. INDEX_NAME = "gpu-demo" DIMENSION = 256 # matches common embedding model output sizes @@ -215,8 +216,8 @@ def verify_gpu_build(timeout: int = REMOTE_BUILD_TIMEOUT) -> None: print(f" Bucket : s3://{S3_BUCKET}/knn-indexes/") print(f" Timeout : {timeout}s (poll interval: 5s)\n") - # boto3 picks up AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_DEFAULT_REGION - # from the environment automatically. + # boto3 uses static AWS env vars when set, otherwise it falls through to + # the default credential provider chain. s3 = boto3.client("s3", region_name=S3_REGION) deadline = time.time() + timeout From 4d21712b70110e5a18bd24404bcffc0f60714748 Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Mon, 18 May 2026 13:07:57 -0500 Subject: [PATCH 05/22] Update Signed-off-by: James Bourbeau --- deploy/DEPLOYMENT.md | 6 +++--- deploy/README.md | 8 ++++---- deploy/bench/run.py | 16 +++++++++++----- deploy/docker-compose.yml | 4 ++-- deploy/remote-index-build/run.py | 13 +++++++++++-- 5 files changed, 31 insertions(+), 16 deletions(-) diff --git a/deploy/DEPLOYMENT.md b/deploy/DEPLOYMENT.md index 28b184d206..a04bc259d7 100644 --- a/deploy/DEPLOYMENT.md +++ b/deploy/DEPLOYMENT.md @@ -67,7 +67,7 @@ export AWS_SESSION_TOKEN= # required for temporary (STS) creden Optionally configure the region: ```bash -export AWS_DEFAULT_REGION=us-east-1 # default: us-east-1 +export AWS_DEFAULT_REGION=us-west-2 # default: us-west-2 ``` Start OpenSearch and the GPU builder: @@ -90,7 +90,7 @@ curl -X PUT http://localhost:9200/_snapshot/ \ "settings": { "bucket": "", "base_path": "knn-indexes", - "region": "us-east-1" + "region": "us-west-2" } }' ``` @@ -159,7 +159,7 @@ docker compose run --rm \ -e OPENSEARCH_URL=http://opensearch:9200 \ -e BUILDER_URL=http://remote-index-builder:1025 \ -e S3_BUCKET=${S3_BUCKET} \ - -e AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION:-us-east-1} \ + -e AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION:-us-west-2} \ -e REMOTE_BUILD_SIZE_MIN=${REMOTE_BUILD_SIZE_MIN:-} \ -e REMOTE_BUILD_TIMEOUT=${REMOTE_BUILD_TIMEOUT:-1800} \ -v $(pwd)/remote-index-build:/app/remote-index-build \ diff --git a/deploy/README.md b/deploy/README.md index a6181a0cf2..0d60367e65 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -44,7 +44,7 @@ sudo sysctl -w vm.max_map_count=262144 Set required environment variables: ```bash -export DATASET_PATH=/path/to/ann-benchmark-datasets # directory containing dataset files +export DATASET_PATH="$(pwd)/ann-benchmark-datasets" # directory containing dataset files ``` GPU mode also requires an S3 bucket. Static AWS keys are supported, but optional @@ -52,7 +52,7 @@ when the containers can use another AWS default credential provider such as an EC2 instance role: ```bash -export S3_BUCKET=opensearch-s3-bucket # S3 bucket name +export S3_BUCKET=opensearch-cuvs-bench # S3 bucket name ``` If you are using static credentials instead of a default provider, also export: @@ -66,7 +66,7 @@ export AWS_SESSION_TOKEN= # required when using temporary (STS) Optionally configure the benchmark: ```bash -export AWS_DEFAULT_REGION=us-east-1 # AWS region for the S3 bucket (default: us-east-1) +export AWS_DEFAULT_REGION=us-west-2 # AWS region for the S3 bucket (default: us-west-2) export DATASET=sift-128-euclidean # default export BENCH_GROUPS=test # test | base (default: test) export K=10 # number of neighbors (default: 10) @@ -205,7 +205,7 @@ Requires the full stack (OpenSearch and the remote index builder), S3 access, an When using static AWS keys, map them into the test fixture's `S3_*` variable names: ```bash -export AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION:-us-east-1} +export AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION:-us-west-2} docker compose --profile gpu up -d --wait opensearch remote-index-builder docker compose run --rm --no-deps \ -e OPENSEARCH_URL=http://opensearch:9200 \ diff --git a/deploy/bench/run.py b/deploy/bench/run.py index a66b852a83..25424bf033 100644 --- a/deploy/bench/run.py +++ b/deploy/bench/run.py @@ -32,8 +32,8 @@ REMOTE_BUILD_SIZE_MIN = os.environ.get("REMOTE_BUILD_SIZE_MIN", "").strip() REMOTE_BUILD_TIMEOUT = int(os.environ.get("REMOTE_BUILD_TIMEOUT", "1800")) -S3_BUCKET = os.environ.get("S3_BUCKET", "") -S3_REGION = os.environ.get("AWS_DEFAULT_REGION", "us-east-1") +S3_BUCKET = os.environ.get("S3_BUCKET", "").strip() +S3_REGION = os.environ.get("AWS_DEFAULT_REGION", "us-west-2") DATASET = os.environ.get("DATASET", "sift-128-euclidean") DATASET_PATH = os.environ.get("DATASET_PATH", "/data/datasets") @@ -222,6 +222,7 @@ def print_results(results: list) -> None: recall = _recall_for_entry(r, entry_index, entry_count) if recall is None: missing_recall_rows += 1 + continue _print_result_row( entry["search_params"], recall, @@ -230,8 +231,8 @@ def print_results(results: list) -> None: ) if missing_recall_rows: print( - "\n Note: this cuVS version does not report recall for every " - "OpenSearch search parameter row." + "\n Omitted " + f"{missing_recall_rows} timing-only search rows without recall." ) @@ -239,7 +240,12 @@ def print_results(results: list) -> None: def main() -> None: if REMOTE_INDEX_BUILD and not S3_BUCKET: - print("ERROR: S3_BUCKET must be set when REMOTE_INDEX_BUILD=true") + print( + "ERROR: S3_BUCKET is not set. Remote index build requires an S3 " + "bucket for vector and index staging. Set it before starting the " + "stack, for example: export S3_BUCKET=", + file=sys.stderr, + ) sys.exit(1) print("\n" + "═" * 60) diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml index 4680107006..f089d2e894 100644 --- a/deploy/docker-compose.yml +++ b/deploy/docker-compose.yml @@ -26,7 +26,7 @@ # AWS_ACCESS_KEY_ID AWS access key ID (optional if using default credentials) # AWS_SECRET_ACCESS_KEY AWS secret access key (optional if using default credentials) # AWS_SESSION_TOKEN STS session token (required for temporary static credentials) -# AWS_DEFAULT_REGION AWS region for the S3 bucket (default: us-east-1) +# AWS_DEFAULT_REGION AWS region for the S3 bucket (default: us-west-2) # REMOTE_INDEX_BUILD Override GPU/CPU mode detection (true/false); normally auto-detected # REMOTE_BUILD_SIZE_MIN Optional minimum segment size override for remote builds # REMOTE_BUILD_TIMEOUT Remote build wait timeout in seconds (default: 1800) @@ -49,7 +49,7 @@ x-aws-env: &aws-env AWS_ACCESS_KEY_ID: AWS_SECRET_ACCESS_KEY: AWS_SESSION_TOKEN: - AWS_DEFAULT_REGION: ${AWS_DEFAULT_REGION:-us-east-1} + AWS_DEFAULT_REGION: ${AWS_DEFAULT_REGION:-us-west-2} services: diff --git a/deploy/remote-index-build/run.py b/deploy/remote-index-build/run.py index de6ab477a0..a7548ded04 100644 --- a/deploy/remote-index-build/run.py +++ b/deploy/remote-index-build/run.py @@ -25,8 +25,8 @@ OPENSEARCH_URL = os.environ.get("OPENSEARCH_URL", "http://opensearch:9200") BUILDER_URL = os.environ.get("BUILDER_URL", "http://remote-index-builder:1025") -S3_BUCKET = os.environ["S3_BUCKET"] -S3_REGION = os.environ.get("AWS_DEFAULT_REGION", "us-east-1") +S3_BUCKET = os.environ.get("S3_BUCKET", "").strip() +S3_REGION = os.environ.get("AWS_DEFAULT_REGION", "us-west-2") # boto3 uses AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY when set, otherwise it # falls through to the default credential provider chain. @@ -291,6 +291,15 @@ def search_vectors() -> None: # ── entrypoint ──────────────────────────────────────────────────────────────── def main() -> None: + if not S3_BUCKET: + print( + "ERROR: S3_BUCKET is not set. The remote index build demo requires " + "an S3 bucket for vector and index staging. Set it before running, " + "for example: export S3_BUCKET=", + file=sys.stderr, + ) + sys.exit(1) + print("\n" + "═" * 60) print(" OpenSearch GPU Remote Index Build — End-to-End Demo") print("═" * 60) From b5ecb4b8b4a150d3a1beb294108ce3b43da38c92 Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Wed, 20 May 2026 10:36:48 -0500 Subject: [PATCH 06/22] Update Signed-off-by: James Bourbeau --- deploy/README.md | 8 +++-- deploy/bench/Dockerfile | 2 +- deploy/bench/entrypoint.sh | 8 +++-- deploy/bench/run.py | 62 ++++++++++++++++++++++++++++++-------- deploy/docker-compose.yml | 6 ++-- 5 files changed, 66 insertions(+), 20 deletions(-) diff --git a/deploy/README.md b/deploy/README.md index 0d60367e65..64f644de4b 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -70,7 +70,8 @@ export AWS_DEFAULT_REGION=us-west-2 # AWS region for the S3 bucket (defau export DATASET=sift-128-euclidean # default export BENCH_GROUPS=test # test | base (default: test) export K=10 # number of neighbors (default: 10) -export BATCH_SIZE=10000 # query/bulk batch size (default: 10000) +export BATCH_SIZE= # optional query batch size override +export BUILD_BATCH_SIZE= # optional bulk ingest batch size override export REMOTE_BUILD_TIMEOUT=1800 # seconds to wait for remote builds (default: 1800) ``` @@ -101,6 +102,7 @@ docker compose down -v - Creates the kNN index and bulk-ingests dataset vectors - **GPU mode**: Flushes segments, waits for all submitted remote GPU builds to complete, and polls the kNN stats API every 5 s until the build is confirmed complete - **CPU mode**: Flushes and refreshes the local OpenSearch index + - Uses the backend's automatic OpenSearch bulk-ingest batch sizing by default; set `BUILD_BATCH_SIZE` to override it - Records total build time in the result 5. Runs `cuvs-bench` search phase and prints a recall/QPS/latency table 6. Exports benchmark JSON results to CSV (`cuvs_bench.run --data-export`) @@ -146,8 +148,8 @@ $DATASET_PATH/ | Group | Build params | Search params | Use case | |---|---|---|---| -| `test` | 1 combo (m=16, ef_construction=100) | ef_search: 50, 100 | Quick smoke test | -| `base` | 9 combos (m=[32,64,96] × ef_construction=[64,128,256]) | ef_search: 10–800 | Standard benchmark | +| `test` | 1 combo (m=16) | ef_search: 10, 20 | Quick smoke test | +| `base` | 4 combos (m=[16,32,48,64]) | ef_search: 10, 20, 40, 60, 80, 120, 200, 400, 600, 800 | Standard benchmark | ## GPU build verification diff --git a/deploy/bench/Dockerfile b/deploy/bench/Dockerfile index 429534d9e2..c709e5af20 100644 --- a/deploy/bench/Dockerfile +++ b/deploy/bench/Dockerfile @@ -26,7 +26,7 @@ RUN pip install --no-cache-dir \ h5py \ matplotlib \ "numpy<2" \ - "opensearch-py>=2.4.0,<3.0.0" \ + "opensearch-py>=2.4.0" \ pandas \ pyyaml \ requests \ diff --git a/deploy/bench/entrypoint.sh b/deploy/bench/entrypoint.sh index c7ed3cb585..ad7ccda117 100644 --- a/deploy/bench/entrypoint.sh +++ b/deploy/bench/entrypoint.sh @@ -4,8 +4,12 @@ set -e DATASET="${DATASET:-sift-128-euclidean}" BENCH_GROUPS="${BENCH_GROUPS:-test}" K="${K:-10}" -BATCH_SIZE="${BATCH_SIZE:-10000}" ALGORITHM="opensearch_faiss_hnsw" +BATCH_SIZE="${BATCH_SIZE:-}" +DATA_EXPORT_BATCH_ARGS=() +if [ -n "$BATCH_SIZE" ]; then + DATA_EXPORT_BATCH_ARGS=(--batch-size "$BATCH_SIZE") +fi wait_for_builder() { echo "remote-index-builder detected — waiting for it to be ready..." @@ -59,7 +63,7 @@ python -m cuvs_bench.run --data-export \ --algorithms "$ALGORITHM" \ --groups "$BENCH_GROUPS" \ --count "$K" \ - --batch-size "$BATCH_SIZE" \ + "${DATA_EXPORT_BATCH_ARGS[@]}" \ --search-mode latency # Step 4: Plot — PNGs written to /data/datasets (mounted from host $DATASET_PATH) diff --git a/deploy/bench/run.py b/deploy/bench/run.py index 25424bf033..cc4bc15c16 100644 --- a/deploy/bench/run.py +++ b/deploy/bench/run.py @@ -39,7 +39,10 @@ DATASET_PATH = os.environ.get("DATASET_PATH", "/data/datasets") BENCH_GROUPS = os.environ.get("BENCH_GROUPS", "test") K = int(os.environ.get("K", "10")) -BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "10000")) +BATCH_SIZE = os.environ.get("BATCH_SIZE", "").strip() +BATCH_SIZE = int(BATCH_SIZE) if BATCH_SIZE else None +BUILD_BATCH_SIZE = os.environ.get("BUILD_BATCH_SIZE", "").strip() +BUILD_BATCH_SIZE = int(BUILD_BATCH_SIZE) if BUILD_BATCH_SIZE else None ALGORITHM = "opensearch_faiss_hnsw" REPO_NAME = "vector-repo" @@ -64,6 +67,18 @@ def _recall_for_entry( return None +def _get_search_batch_size(search_results: list) -> int | None: + if BATCH_SIZE is not None: + return BATCH_SIZE + for result in search_results: + if not isinstance(result, SearchResult): + continue + batch_size = (result.metadata or {}).get("batch_size") + if batch_size is not None: + return int(batch_size) + return None + + # ── OpenSearch setup ────────────────────────────────────────────────────────── def register_repository() -> None: @@ -114,7 +129,7 @@ def write_result_files( algo: str, groups: str, k: int, - batch_size: int = 10000, + batch_size: int, ) -> None: """Write gbench-compatible JSON result files. @@ -261,6 +276,14 @@ def main() -> None: print(f" Dataset : {DATASET} (path: {DATASET_PATH})") print(f" Algorithm : {ALGORITHM}") print(f" Groups : {BENCH_GROUPS} k={K}") + print( + " Search batch size : " + f"{BATCH_SIZE if BATCH_SIZE is not None else 'backend default'}" + ) + print( + " Build batch size : " + f"{BUILD_BATCH_SIZE if BUILD_BATCH_SIZE is not None else 'backend auto'}" + ) if REMOTE_INDEX_BUILD: register_repository() @@ -268,8 +291,8 @@ def main() -> None: orchestrator = BenchmarkOrchestrator(backend_type="opensearch") - # Shared kwargs for both build and search phases - bench_kwargs = dict( + # Shared kwargs for both build and search phases. + common_kwargs = dict( dataset=DATASET, dataset_path=DATASET_PATH, algorithms=ALGORITHM, @@ -278,12 +301,20 @@ def main() -> None: port=OPENSEARCH_PORT, use_ssl=False, verify_certs=False, + ) + + build_kwargs = dict( + common_kwargs, remote_index_build=REMOTE_INDEX_BUILD, ) + if BUILD_BATCH_SIZE is not None: + build_kwargs["build_batch_size"] = BUILD_BATCH_SIZE if REMOTE_INDEX_BUILD: - bench_kwargs["remote_build_timeout"] = REMOTE_BUILD_TIMEOUT + build_kwargs["remote_build_timeout"] = REMOTE_BUILD_TIMEOUT if REMOTE_BUILD_SIZE_MIN: - bench_kwargs["remote_build_size_min"] = REMOTE_BUILD_SIZE_MIN + build_kwargs["remote_build_size_min"] = REMOTE_BUILD_SIZE_MIN + + # ── Build phase ─────────────────────────────────────────────────────────── mode = "GPU remote build" if REMOTE_INDEX_BUILD else "CPU" banner(f"Building index ({mode} via cuvs-bench)") @@ -291,8 +322,7 @@ def main() -> None: build=True, search=False, force=True, - bulk_batch_size=BATCH_SIZE, - **bench_kwargs, + **build_kwargs, ) index_names = [ @@ -313,24 +343,32 @@ def main() -> None: # ── Search phase ────────────────────────────────────────────────────────── banner("Running search benchmarks (via cuvs-bench)") - search_results = orchestrator.run_benchmark( + search_run_kwargs = dict( build=False, search=True, count=K, - **bench_kwargs, + **common_kwargs, ) + if BATCH_SIZE is not None: + search_run_kwargs["batch_size"] = BATCH_SIZE + search_results = orchestrator.run_benchmark(**search_run_kwargs) print_results(search_results) + batch_size = _get_search_batch_size(search_results) + if batch_size is None: + print(" ERROR: no successful search result reported a batch size") + sys.exit(1) + write_result_files( build_results=build_results, search_results=search_results, dataset=DATASET, dataset_path=DATASET_PATH, - algo=bench_kwargs["algorithms"], + algo=common_kwargs["algorithms"], groups=BENCH_GROUPS, k=K, - batch_size=BATCH_SIZE, + batch_size=batch_size, ) print("\n" + "═" * 60) diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml index f089d2e894..bbba494ace 100644 --- a/deploy/docker-compose.yml +++ b/deploy/docker-compose.yml @@ -33,7 +33,8 @@ # DATASET Dataset name (default: sift-128-euclidean) # BENCH_GROUPS Parameter sweep group: test | base (default: test) # K Number of neighbors to search for (default: 10) -# BATCH_SIZE cuvs-bench query/bulk batch size (default: 10000) +# BATCH_SIZE Optional cuvs-bench query batch size override +# BUILD_BATCH_SIZE Optional OpenSearch bulk ingest batch size override # # Usage: # CPU: docker compose up --build @@ -122,7 +123,8 @@ services: DATASET_PATH: /data/datasets BENCH_GROUPS: ${BENCH_GROUPS:-test} K: ${K:-10} - BATCH_SIZE: ${BATCH_SIZE:-10000} + BATCH_SIZE: ${BATCH_SIZE:-} + BUILD_BATCH_SIZE: ${BUILD_BATCH_SIZE:-} volumes: - ${DATASET_PATH:-/tmp/datasets}:/data/datasets restart: "no" From b1573ab1374dbbaf5dc4ce72197df5589e4f8f34 Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Wed, 20 May 2026 11:29:48 -0500 Subject: [PATCH 07/22] Update Signed-off-by: James Bourbeau --- deploy/README.md | 2 +- deploy/bench/entrypoint.sh | 9 +++------ deploy/bench/run.py | 7 +++---- 3 files changed, 7 insertions(+), 11 deletions(-) diff --git a/deploy/README.md b/deploy/README.md index 64f644de4b..337b2866d3 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -105,7 +105,7 @@ docker compose down -v - Uses the backend's automatic OpenSearch bulk-ingest batch sizing by default; set `BUILD_BATCH_SIZE` to override it - Records total build time in the result 5. Runs `cuvs-bench` search phase and prints a recall/QPS/latency table -6. Exports benchmark JSON results to CSV (`cuvs_bench.run --data-export`) +6. Exports benchmark JSON results to CSV 7. Generates recall vs. latency/throughput plots as PNGs in `$DATASET_PATH` (`cuvs_bench.plot`) ## Dataset format diff --git a/deploy/bench/entrypoint.sh b/deploy/bench/entrypoint.sh index ad7ccda117..64a65a41b5 100644 --- a/deploy/bench/entrypoint.sh +++ b/deploy/bench/entrypoint.sh @@ -5,11 +5,6 @@ DATASET="${DATASET:-sift-128-euclidean}" BENCH_GROUPS="${BENCH_GROUPS:-test}" K="${K:-10}" ALGORITHM="opensearch_faiss_hnsw" -BATCH_SIZE="${BATCH_SIZE:-}" -DATA_EXPORT_BATCH_ARGS=() -if [ -n "$BATCH_SIZE" ]; then - DATA_EXPORT_BATCH_ARGS=(--batch-size "$BATCH_SIZE") -fi wait_for_builder() { echo "remote-index-builder detected — waiting for it to be ready..." @@ -57,13 +52,15 @@ python -m cuvs_bench.get_dataset \ python -u run.py # Step 3: Export JSON → CSV (required by cuvs_bench.plot) +# --batch-size is ignored when --data-export is set, but Click prompts for it +# before entering main(), so pass a dummy value to keep the container non-interactive. python -m cuvs_bench.run --data-export \ --dataset "$DATASET" \ --dataset-path /data/datasets \ --algorithms "$ALGORITHM" \ --groups "$BENCH_GROUPS" \ --count "$K" \ - "${DATA_EXPORT_BATCH_ARGS[@]}" \ + --batch-size 1 \ --search-mode latency # Step 4: Plot — PNGs written to /data/datasets (mounted from host $DATASET_PATH) diff --git a/deploy/bench/run.py b/deploy/bench/run.py index cc4bc15c16..9a29b832a4 100644 --- a/deploy/bench/run.py +++ b/deploy/bench/run.py @@ -10,8 +10,7 @@ b. Flush segments to kick off remote GPU builds when enabled c. Poll kNN stats until every submitted remote build completes 4. Run cuvs-bench search benchmarks and print results - 5. Write gbench-compatible JSON result files so cuvs_bench.run --data-export - and cuvs_bench.plot can be used for CSV export and plotting + 5. Write gbench-compatible JSON result files for CSV export and plotting """ import json @@ -134,8 +133,8 @@ def write_result_files( """Write gbench-compatible JSON result files. Creates files under //result/{build,search}/ in the - same format the C++ backend produces, so ``cuvs_bench.run --data-export`` - and ``cuvs_bench.plot`` work without modification. + same format the C++ backend produces, so the cuvs-bench CSV exporters and + ``cuvs_bench.plot`` work without modification. """ build_dir = os.path.join(dataset_path, dataset, "result", "build") search_dir = os.path.join(dataset_path, dataset, "result", "search") From 26e8640c7eec6145f8692578b8181c9fd23fe3e0 Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Wed, 20 May 2026 12:45:02 -0500 Subject: [PATCH 08/22] Multi-node Signed-off-by: James Bourbeau --- deploy/bench/entrypoint.sh | 7 +- deploy/bench/run.py | 10 +- deploy/docker-compose.multinode.yml | 147 ++++++++++++++++++++++++++++ deploy/docker-compose.yml | 4 + deploy/opensearch/Dockerfile | 1 + 5 files changed, 162 insertions(+), 7 deletions(-) create mode 100644 deploy/docker-compose.multinode.yml diff --git a/deploy/bench/entrypoint.sh b/deploy/bench/entrypoint.sh index 64a65a41b5..78f050f9e4 100644 --- a/deploy/bench/entrypoint.sh +++ b/deploy/bench/entrypoint.sh @@ -7,11 +7,12 @@ K="${K:-10}" ALGORITHM="opensearch_faiss_hnsw" wait_for_builder() { - echo "remote-index-builder detected — waiting for it to be ready..." - until python3 -c 'import socket; socket.create_connection(("remote-index-builder", 1025), 2).close()' 2>/dev/null; do + builder_url="${BUILDER_URL:-http://remote-index-builder:1025}" + echo "Remote index build enabled — waiting for builder at ${builder_url}..." + until BUILDER_URL="${builder_url}" python3 -c 'import os, socket; from urllib.parse import urlparse; url = urlparse(os.environ["BUILDER_URL"]); socket.create_connection((url.hostname, url.port or 1025), 2).close()' 2>/dev/null; do sleep 5 done - echo "remote-index-builder is ready." + echo "Remote index builder is ready." } if [ -n "${REMOTE_INDEX_BUILD:-}" ]; then diff --git a/deploy/bench/run.py b/deploy/bench/run.py index 9a29b832a4..55bdc1777f 100644 --- a/deploy/bench/run.py +++ b/deploy/bench/run.py @@ -32,6 +32,7 @@ REMOTE_BUILD_TIMEOUT = int(os.environ.get("REMOTE_BUILD_TIMEOUT", "1800")) S3_BUCKET = os.environ.get("S3_BUCKET", "").strip() +S3_PREFIX = os.environ.get("S3_PREFIX", "knn-indexes").strip() or "knn-indexes" S3_REGION = os.environ.get("AWS_DEFAULT_REGION", "us-west-2") DATASET = os.environ.get("DATASET", "sift-128-euclidean") @@ -44,7 +45,8 @@ BUILD_BATCH_SIZE = int(BUILD_BATCH_SIZE) if BUILD_BATCH_SIZE else None ALGORITHM = "opensearch_faiss_hnsw" -REPO_NAME = "vector-repo" +REPO_NAME = os.environ.get("REMOTE_VECTOR_REPOSITORY", "vector-repo").strip() +REPO_NAME = REPO_NAME or "vector-repo" session = requests.Session() session.headers.update({"Content-Type": "application/json"}) @@ -88,7 +90,7 @@ def register_repository() -> None: "type": "s3", "settings": { "bucket": S3_BUCKET, - "base_path": "knn-indexes", + "base_path": S3_PREFIX, "region": S3_REGION, }, }, @@ -269,7 +271,8 @@ def main() -> None: print(f" Remote index build : {REMOTE_INDEX_BUILD}") if REMOTE_INDEX_BUILD: print(f" GPU builder : {BUILDER_URL}") - print(f" S3 bucket : s3://{S3_BUCKET}/knn-indexes/ (region: {S3_REGION})") + print(f" S3 bucket : s3://{S3_BUCKET}/{S3_PREFIX}/ (region: {S3_REGION})") + print(f" Repository : {REPO_NAME}") print(f" Build size minimum : {REMOTE_BUILD_SIZE_MIN or 'OpenSearch default'}") print(f" Build timeout : {REMOTE_BUILD_TIMEOUT}s") print(f" Dataset : {DATASET} (path: {DATASET_PATH})") @@ -313,7 +316,6 @@ def main() -> None: if REMOTE_BUILD_SIZE_MIN: build_kwargs["remote_build_size_min"] = REMOTE_BUILD_SIZE_MIN - # ── Build phase ─────────────────────────────────────────────────────────── mode = "GPU remote build" if REMOTE_INDEX_BUILD else "CPU" banner(f"Building index ({mode} via cuvs-bench)") diff --git a/deploy/docker-compose.multinode.yml b/deploy/docker-compose.multinode.yml new file mode 100644 index 0000000000..c88f95d9e5 --- /dev/null +++ b/deploy/docker-compose.multinode.yml @@ -0,0 +1,147 @@ +# Multi-node split of the benchmark stack. +# +# Copy this file to all three EC2 instances from the deploy-opensearch-tmp +# branch of https://github.com/jrbourbeau/cuvs. The default OpenSearch and +# bench images build from that checkout, so keep the opensearch/ directory on +# the OpenSearch node and the bench/ directory on the client node unless you +# set OPENSEARCH_IMAGE and BENCH_IMAGE to prebuilt images. The OpenSearch image +# contains opensearch.yml, so a prebuilt image does not need the opensearch/ +# directory at runtime. +# +# OpenSearch node: +# docker compose -f docker-compose.multinode.yml --profile opensearch up -d +# +# GPU remote index builder node: +# docker compose -f docker-compose.multinode.yml --profile builder up -d +# +# Client/benchmark node: +# docker compose -f docker-compose.multinode.yml --profile client build bench +# docker compose -f docker-compose.multinode.yml --profile client run --rm bench +# +# Docker Compose networks are host-local. Across EC2 instances, use private +# EC2 DNS names or private IPv4 addresses in OPENSEARCH_URL, +# OPENSEARCH_HOST, and REMOTE_INDEX_BUILDER_URL. + +name: opensearch-cuvs-bench-multinode + +x-aws-env: &aws-env + AWS_ACCESS_KEY_ID: ${AWS_ACCESS_KEY_ID:-} + AWS_SECRET_ACCESS_KEY: ${AWS_SECRET_ACCESS_KEY:-} + AWS_SESSION_TOKEN: ${AWS_SESSION_TOKEN:-} + AWS_REGION: ${AWS_REGION:-us-west-2} + AWS_DEFAULT_REGION: ${AWS_DEFAULT_REGION:-us-west-2} + S3_BUCKET: ${S3_BUCKET:-} + S3_PREFIX: ${S3_PREFIX:-knn-indexes} + +x-bench-env: &bench-env + <<: *aws-env + OPENSEARCH_URL: ${OPENSEARCH_URL:-http://opensearch:9200} + OPENSEARCH_HOST: ${OPENSEARCH_HOST:-opensearch} + OPENSEARCH_PORT: ${OPENSEARCH_PORT:-9200} + BUILDER_URL: ${REMOTE_INDEX_BUILDER_URL:-http://remote-index-builder:1025} + REMOTE_INDEX_BUILDER_URL: ${REMOTE_INDEX_BUILDER_URL:-http://remote-index-builder:1025} + REMOTE_INDEX_BUILD: ${REMOTE_INDEX_BUILD:-true} + REMOTE_BUILD_SIZE_MIN: ${REMOTE_BUILD_SIZE_MIN:-} + REMOTE_BUILD_TIMEOUT: ${REMOTE_BUILD_TIMEOUT:-1800} + REMOTE_VECTOR_REPOSITORY: ${REMOTE_VECTOR_REPOSITORY:-vector-repo} + DATASET: ${DATASET:-sift-128-euclidean} + DATASET_PATH: /data/datasets + BENCH_GROUPS: ${BENCH_GROUPS:-test} + K: ${K:-10} + BATCH_SIZE: ${BATCH_SIZE:-} + BUILD_BATCH_SIZE: ${BUILD_BATCH_SIZE:-} + +services: + opensearch: + profiles: ["opensearch"] + image: ${OPENSEARCH_IMAGE:-opensearch-cuvs-bench-opensearch} + build: + context: ${OPENSEARCH_BUILD_CONTEXT:-./opensearch} + dockerfile: ${OPENSEARCH_DOCKERFILE:-Dockerfile} + container_name: opensearch + restart: unless-stopped + ports: + - "${OPENSEARCH_BIND_ADDR:-0.0.0.0}:${OPENSEARCH_PORT:-9200}:9200" + - "${OPENSEARCH_PERF_BIND_ADDR:-127.0.0.1}:${OPENSEARCH_PERF_PORT:-9600}:9600" + environment: + <<: *aws-env + cluster.name: ${OPENSEARCH_CLUSTER_NAME:-opensearch-cuvs-bench} + node.name: ${OPENSEARCH_NODE_NAME:-opensearch-node-1} + discovery.type: single-node + bootstrap.memory_lock: "true" + OPENSEARCH_JAVA_OPTS: ${OPENSEARCH_JAVA_OPTS:--Xms16g -Xmx16g} + DISABLE_SECURITY_PLUGIN: ${DISABLE_SECURITY_PLUGIN:-true} + OPENSEARCH_INITIAL_ADMIN_PASSWORD: ${OPENSEARCH_INITIAL_ADMIN_PASSWORD:-} + ulimits: + memlock: + soft: -1 + hard: -1 + nofile: + soft: 65536 + hard: 65536 + volumes: + - opensearch-data:/usr/share/opensearch/data + healthcheck: + test: ["CMD-SHELL", "curl -fsS http://localhost:9200/_cluster/health >/dev/null || exit 1"] + interval: 20s + timeout: 5s + retries: 30 + + remote-index-builder: + profiles: ["builder", "gpu"] + image: ${REMOTE_INDEX_BUILDER_IMAGE:-opensearchproject/remote-vector-index-builder:api-latest} + container_name: remote-index-builder + restart: unless-stopped + ports: + - "${REMOTE_INDEX_BUILDER_BIND_ADDR:-0.0.0.0}:${REMOTE_INDEX_BUILDER_HOST_PORT:-1025}:1025" + environment: + <<: *aws-env + healthcheck: + test: ["CMD-SHELL", "python3 -c 'import socket; socket.create_connection((\"localhost\", 1025), 2).close()'"] + interval: 5s + timeout: 5s + retries: 24 + start_period: 10s + deploy: + resources: + reservations: + devices: + - driver: nvidia + count: ${GPU_COUNT:-1} + capabilities: ["gpu"] + + bench: + profiles: ["client"] + image: ${BENCH_IMAGE:-opensearch-cuvs-bench-bench} + build: + context: ${BENCH_BUILD_CONTEXT:-./bench} + dockerfile: ${BENCH_DOCKERFILE:-Dockerfile} + environment: + <<: *bench-env + extra_hosts: + - "host.docker.internal:host-gateway" + volumes: + - ${DATASET_PATH:-/tmp/datasets}:/data/datasets + + configure-remote-index-build: + profiles: ["configure"] + image: curlimages/curl:latest + environment: + <<: *bench-env + command: + - sh + - -ec + - | + : "$${S3_BUCKET:?S3_BUCKET must be set}" + : "$${REMOTE_INDEX_BUILDER_URL:?REMOTE_INDEX_BUILDER_URL must be set}" + + curl -fsS -X PUT "$${OPENSEARCH_URL}/_snapshot/$${REMOTE_VECTOR_REPOSITORY}" \ + -H 'Content-Type: application/json' \ + -d "{\"type\":\"s3\",\"settings\":{\"bucket\":\"$${S3_BUCKET}\",\"base_path\":\"$${S3_PREFIX}\",\"region\":\"$${AWS_REGION}\"}}" + + curl -fsS -X PUT "$${OPENSEARCH_URL}/_cluster/settings" \ + -H 'Content-Type: application/json' \ + -d "{\"persistent\":{\"knn.remote_index_build.enabled\":\"true\",\"knn.remote_index_build.repository\":\"$${REMOTE_VECTOR_REPOSITORY}\",\"knn.remote_index_build.service.endpoint\":\"$${REMOTE_INDEX_BUILDER_URL}\"}}" + +volumes: + opensearch-data: diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml index bbba494ace..4b6300ecb8 100644 --- a/deploy/docker-compose.yml +++ b/deploy/docker-compose.yml @@ -30,6 +30,8 @@ # REMOTE_INDEX_BUILD Override GPU/CPU mode detection (true/false); normally auto-detected # REMOTE_BUILD_SIZE_MIN Optional minimum segment size override for remote builds # REMOTE_BUILD_TIMEOUT Remote build wait timeout in seconds (default: 1800) +# REMOTE_VECTOR_REPOSITORY Optional snapshot repository name (default: vector-repo) +# S3_PREFIX Optional S3 prefix for staged vectors/indexes (default: knn-indexes) # DATASET Dataset name (default: sift-128-euclidean) # BENCH_GROUPS Parameter sweep group: test | base (default: test) # K Number of neighbors to search for (default: 10) @@ -118,7 +120,9 @@ services: REMOTE_INDEX_BUILD: ${REMOTE_INDEX_BUILD:-} REMOTE_BUILD_SIZE_MIN: ${REMOTE_BUILD_SIZE_MIN:-} REMOTE_BUILD_TIMEOUT: ${REMOTE_BUILD_TIMEOUT:-1800} + REMOTE_VECTOR_REPOSITORY: ${REMOTE_VECTOR_REPOSITORY:-vector-repo} S3_BUCKET: ${S3_BUCKET:-} + S3_PREFIX: ${S3_PREFIX:-knn-indexes} DATASET: ${DATASET:-sift-128-euclidean} DATASET_PATH: /data/datasets BENCH_GROUPS: ${BENCH_GROUPS:-test} diff --git a/deploy/opensearch/Dockerfile b/deploy/opensearch/Dockerfile index 06be41a26c..9ab8bfcdd8 100644 --- a/deploy/opensearch/Dockerfile +++ b/deploy/opensearch/Dockerfile @@ -7,6 +7,7 @@ RUN /usr/share/opensearch/bin/opensearch-plugin install --batch repository-s3 # entrypoint.sh populates the keystore from AWS credential environment # variables when provided, then execs opensearch. +COPY opensearch.yml /usr/share/opensearch/config/opensearch.yml COPY --chmod=755 entrypoint.sh /entrypoint.sh ENTRYPOINT ["/entrypoint.sh"] From 9cf5c45ea07d884e07242f4727c027354e5232b0 Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Mon, 22 Jun 2026 16:22:03 -0500 Subject: [PATCH 09/22] Update Signed-off-by: James Bourbeau --- deploy/AWS_MULTINODE_SETUP.md | 364 ++++++++++++++++++++++ deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md | 354 +++++++++++++++++++++ deploy/DEPLOYMENT.md | 10 +- deploy/README.md | 3 - deploy/bench/Dockerfile | 2 +- 5 files changed, 724 insertions(+), 9 deletions(-) create mode 100644 deploy/AWS_MULTINODE_SETUP.md create mode 100644 deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md diff --git a/deploy/AWS_MULTINODE_SETUP.md b/deploy/AWS_MULTINODE_SETUP.md new file mode 100644 index 0000000000..cff5d973e4 --- /dev/null +++ b/deploy/AWS_MULTINODE_SETUP.md @@ -0,0 +1,364 @@ +# Opensearch + cuVS multi-node benchmarking setup + +This guide sets up the three-node version of the OpenSearch GPU benchmark stack. + +```text +client node --> runs the benchmark submitter +opensearch node --> runs OpenSearch +builder GPU node --> runs the remote index build service +``` + +The examples use `us-west-2`, but that region is not required. Choose any AWS region where your needed instance types, GPU AMI, and quotas are available, then keep all resources in that same region. + +The goal is to keep the deployment simple while separating the three major components onto their own EC2 instances. Docker Compose still runs locally on each instance; it does not create a cross-host network. Cross-node communication uses EC2 private DNS names or private IPv4 addresses. + +## 1. Choose one AWS region + +In the AWS Console, set the region selector to your chosen region. The examples below use: + +```text +US West (Oregon) us-west-2 +``` + +Use this same region for S3, EC2, security groups, and the instances' IAM roles. The EC2 instances and security groups should also be in the same VPC so private DNS, private IPs, and security-group source rules work as expected. + +When running commands, set both AWS region variables from one value: + +```bash +export AWS_DEFAULT_REGION=us-west-2 +export AWS_REGION="$AWS_DEFAULT_REGION" +``` + +`AWS_REGION` is used when registering the OpenSearch S3 repository region. `AWS_DEFAULT_REGION` is used by AWS CLI and boto-style tooling. Setting both from the same value avoids accidental mismatches. + +## 2. Create the S3 bucket + +Go to **S3 > Create bucket**. + +Use: + +```text +Bucket name: globally unique name +Region: your chosen AWS region, for example US West (Oregon) us-west-2 +Object Ownership: ACLs disabled +Block Public Access: block all public access +Encryption: SSE-S3 is fine +Versioning: optional +``` + +This bucket stores the remote-build staging objects, including vectors and +generated index artifacts. Benchmark datasets and result plots stay on the +client node under `DATASET_PATH`. + +## 3. Create an IAM policy for S3 + +Go to **IAM > Policies > Create policy > JSON**. + +Use this policy: + +```json +{ + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": ["s3:ListBucket"], + "Resource": "arn:aws:s3:::opensearch-cuvs-bench" + }, + { + "Effect": "Allow", + "Action": ["s3:GetObject", "s3:PutObject", "s3:DeleteObject"], + "Resource": "arn:aws:s3:::opensearch-cuvs-bench/*" + } + ] +} +``` + +Name it: + +```text +opensearch-cuvs-bench-s3-policy +``` + +The delete permission lets the snapshot repository and staging workflow clean up +temporary objects when needed. + +## 4. Create the EC2 IAM role + +Go to **IAM > Roles > Create role**. + +Choose: + +```text +Trusted entity: AWS service +Use case: EC2 +``` + +Attach: + +```text +opensearch-cuvs-bench-s3-policy +AmazonSSMManagedInstanceCore +``` + +Name it: + +```text +opensearch-cuvs-bench-ec2-role +``` + +This role gives the instances refreshable S3 credentials and enables Session Manager access. Prefer this over fixed `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, and `AWS_SESSION_TOKEN` values. + +## 5. Create security groups + +Go to **EC2 > Security Groups > Create security group**. + +Create these three security groups in the same VPC: + +```text +sg-cuvs-client +sg-cuvs-opensearch +sg-cuvs-builder +``` + +Inbound rules for `sg-cuvs-client`: + +```text +No inbound rules needed +``` + +Inbound rules for `sg-cuvs-opensearch`: + +```text +TCP 9200 from sg-cuvs-client +TCP 9200 from sg-cuvs-builder +``` + +Inbound rules for `sg-cuvs-builder`: + +```text +TCP 1025 from sg-cuvs-opensearch +TCP 1025 from sg-cuvs-client +``` + +Leave outbound as the default `allow all`. Use Session Manager instead of SSH if possible. If you need SSH, add TCP `22` only from your own IP. + +## 6. Launch the OpenSearch node + +Go to **EC2 > Instances > Launch instance**. + +Use: + +```text +Name: opensearch-cuvs-db +AMI: Ubuntu 22.04 or Amazon Linux 2023 +Instance type: r7i.xlarge, r7i.2xlarge, or m7i.2xlarge +Security group: sg-cuvs-opensearch +IAM role: opensearch-cuvs-bench-ec2-role +Storage: 100+ GiB gp3, larger if needed +``` + +Under **Advanced details > Metadata options**, use: + +```text +IMDS endpoint: Enabled +IMDSv2: Required +Hop limit: 2 +``` + +After launch, copy the instance's **Private IPv4 DNS** or **Private IPv4 address**. You will use it as `OPENSEARCH_URL`. + +## 7. Launch the GPU builder node + +Launch a second EC2 instance: + +```text +Name: opensearch-cuvs-builder +AMI: AWS Deep Learning Base AMI with CUDA, Ubuntu 22.04 +Instance type: g5.xlarge or g6.xlarge +Security group: sg-cuvs-builder +IAM role: opensearch-cuvs-bench-ec2-role +Storage: 100+ GiB gp3 +``` + +Use the same metadata options: + +```text +IMDS endpoint: Enabled +IMDSv2: Required +Hop limit: 2 +``` + +After launch, copy the instance's **Private IPv4 DNS** or **Private IPv4 address**. You will use it as `REMOTE_INDEX_BUILDER_URL`. + +## 8. Launch the client node + +Launch the benchmark submitter instance: + +```text +Name: opensearch-cuvs-client +AMI: Ubuntu 22.04 or Amazon Linux 2023 +Instance type: c7i.xlarge or m7i.xlarge +Security group: sg-cuvs-client +IAM role: opensearch-cuvs-bench-ec2-role +Storage: 50-100 GiB gp3 +``` + +Use the same metadata options: + +```text +IMDS endpoint: Enabled +IMDSv2: Required +Hop limit: 2 +``` + +## 9. Install Docker and Compose on each node + +Connect to each instance with **EC2 > Instances > Connect > Session Manager**. + +Install Docker and Docker Compose if they are not already installed. Then verify: + +```bash +docker compose version +``` + +On the GPU builder node, also verify GPU access: + +```bash +nvidia-smi +docker run --rm --gpus all nvidia/cuda:12.9.0-base-ubuntu22.04 nvidia-smi +``` + +If you want to avoid `sudo docker`, add your login user to the `docker` group and reconnect: + +```bash +sudo groupadd docker 2>/dev/null || true +sudo usermod -aG docker "$(whoami)" +``` + +## 10. Copy the deployment checkout to each node + +On all three nodes, clone or copy the `deploy-opensearch-tmp` branch of the +`jrbourbeau/cuvs` fork: + +```bash +git clone --branch deploy-opensearch-tmp https://github.com/jrbourbeau/cuvs.git +cd cuvs +``` + +Make sure this file is present: + +```text +docker-compose.multinode.yml +``` + +## 11. Start OpenSearch + +On the OpenSearch node: + +```bash +export S3_BUCKET=opensearch-cuvs-bench +export AWS_DEFAULT_REGION=us-west-2 +export AWS_REGION="$AWS_DEFAULT_REGION" +docker compose -f docker-compose.multinode.yml --profile opensearch up -d +``` + +Verify: + +```bash +curl http://localhost:9200 +``` + +## 12. Start the remote index builder + +On the GPU builder node: + +```bash +export S3_BUCKET=opensearch-cuvs-bench +export AWS_DEFAULT_REGION=us-west-2 +export AWS_REGION="$AWS_DEFAULT_REGION" +docker compose -f docker-compose.multinode.yml --profile builder up -d +``` + +Verify the container is running: + +```bash +docker ps +``` + +From the OpenSearch node, verify that the builder is reachable. OpenSearch is +the service that calls the remote builder: + +```bash +python3 -c 'import socket; socket.create_connection(("BUILDER_PRIVATE_DNS_OR_IP", 1025), 5).close(); print("builder reachable")' +``` + +From the client node, run the same check. The benchmark container also waits +for the builder before starting a remote-build run: + +```bash +python3 -c 'import socket; socket.create_connection(("BUILDER_PRIVATE_DNS_OR_IP", 1025), 5).close(); print("builder reachable")' +``` + +## 13. Configure OpenSearch from the client node + +On the client node: + +```bash +export OPENSEARCH_HOST=OPENSEARCH_PRIVATE_DNS_OR_IP +export OPENSEARCH_URL=http://${OPENSEARCH_HOST}:9200 +export REMOTE_INDEX_BUILDER_URL=http://BUILDER_PRIVATE_DNS_OR_IP:1025 +export S3_BUCKET=opensearch-cuvs-bench +export AWS_DEFAULT_REGION=us-west-2 +export AWS_REGION="$AWS_DEFAULT_REGION" +export S3_PREFIX=knn-indexes +export REMOTE_VECTOR_REPOSITORY=vector-repo +``` + +Then run the one-shot configure profile: + +```bash +docker compose -f docker-compose.multinode.yml --profile configure run --rm configure-remote-index-build +``` + +This registers the S3-backed remote vector repository and tells OpenSearch where the remote index builder service lives. + +## 14. Run the benchmark + +Still on the client node: + +```bash +export REMOTE_INDEX_BUILD=true +export DATASET_PATH="$(pwd)/opensearch-cuvs-datasets" +export DATASET=sift-128-euclidean +export BENCH_GROUPS=test +export K=10 +export BATCH_SIZE= +export BUILD_BATCH_SIZE= + +mkdir -p "${DATASET_PATH}" + +docker compose -f docker-compose.multinode.yml --profile client build bench +docker compose -f docker-compose.multinode.yml --profile client run --rm bench +``` + +## 15. Debug checklist + +From the client node, these should all work: + +```bash +curl "$OPENSEARCH_URL" +python3 -c 'import os, socket; from urllib.parse import urlparse; url = urlparse(os.environ["REMOTE_INDEX_BUILDER_URL"]); socket.create_connection((url.hostname, url.port or 1025), 5).close(); print("builder reachable")' +aws s3 ls s3://$S3_BUCKET --region "$AWS_DEFAULT_REGION" +``` + +If S3 fails, check the IAM role and S3 policy. If OpenSearch or builder +connectivity fails, check private DNS/IP values and security group source rules. + +## Notes + +- Keep the S3 bucket and EC2 resources in the same AWS region. The examples use `us-west-2`, but the setup is not region-specific. +- Within that region, keep all three instances and security groups in the same VPC. Prefer the same Availability Zone for the first benchmark run. +- Use private DNS or private IPs, not public IPs, for cross-node service traffic. +- Docker Compose networks are local to one host; Compose service names do not resolve across EC2 instances. +- The EC2 IAM role credentials rotate automatically. Avoid freezing temporary credentials into `.env` unless the application absolutely requires literal `AWS_*` variables. diff --git a/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md b/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md new file mode 100644 index 0000000000..c9d555bde1 --- /dev/null +++ b/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md @@ -0,0 +1,354 @@ +# AWS single-role multinode setup + +This guide sets up the three-node version of the OpenSearch GPU benchmark stack in `us-west-2`: + +```text +client node runs the benchmark submitter +opensearch node runs OpenSearch +builder GPU node runs the remote index build service +``` + +The goal is to keep the deployment simple while separating the three major components onto their own EC2 instances. Docker Compose still runs locally on each instance; it does not create a cross-host network. Cross-node communication uses EC2 private DNS names or private IPv4 addresses. + +## 1. Set the AWS region + +In the AWS Console, set the region selector to: + +```text +US West (Oregon) us-west-2 +``` + +Use this same region for S3, EC2, security groups, and the instances' IAM roles. + +## 2. Create the S3 bucket + +Go to **S3 > Create bucket**. + +Use: + +```text +Bucket name: globally unique name +Region: US West (Oregon) us-west-2 +Object Ownership: ACLs disabled +Block Public Access: block all public access +Encryption: SSE-S3 is fine +Versioning: optional +``` + +This bucket stores the remote-build staging objects, including vectors and +generated index artifacts. Benchmark datasets and result plots stay on the +client node under `DATASET_PATH`. + +## 3. Create an IAM policy for S3 + +Go to **IAM > Policies > Create policy > JSON**. + +Use this policy: + +```json +{ + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": ["s3:ListBucket"], + "Resource": "arn:aws:s3:::opensearch-cuvs-bench" + }, + { + "Effect": "Allow", + "Action": ["s3:GetObject", "s3:PutObject", "s3:DeleteObject"], + "Resource": "arn:aws:s3:::opensearch-cuvs-bench/*" + } + ] +} +``` + +Name it: + +```text +opensearch-cuvs-bench-s3-policy +``` + +The delete permission lets the snapshot repository and staging workflow clean up +temporary objects when needed. + +## 4. Create the EC2 IAM role + +Go to **IAM > Roles > Create role**. + +Choose: + +```text +Trusted entity: AWS service +Use case: EC2 +``` + +Attach: + +```text +opensearch-cuvs-bench-s3-policy +AmazonSSMManagedInstanceCore +``` + +Name it: + +```text +opensearch-cuvs-bench-ec2-role +``` + +This role gives the instances refreshable S3 credentials and enables Session Manager access. Prefer this over fixed `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, and `AWS_SESSION_TOKEN` values. + +## 5. Create security groups + +Go to **EC2 > Security Groups > Create security group**. + +Create these three security groups in the same VPC: + +```text +sg-cuvs-client +sg-cuvs-opensearch +sg-cuvs-builder +``` + +Inbound rules for `sg-cuvs-client`: + +```text +No inbound rules needed +``` + +Inbound rules for `sg-cuvs-opensearch`: + +```text +TCP 9200 from sg-cuvs-client +TCP 9200 from sg-cuvs-builder +``` + +Inbound rules for `sg-cuvs-builder`: + +```text +TCP 1025 from sg-cuvs-opensearch +TCP 1025 from sg-cuvs-client +``` + +Leave outbound as the default `allow all`. Use Session Manager instead of SSH if possible. If you need SSH, add TCP `22` only from your own IP. + +## 6. Launch the OpenSearch node + +Go to **EC2 > Instances > Launch instance**. + +Use: + +```text +Name: opensearch-cuvs-db +AMI: Ubuntu 22.04 or Amazon Linux 2023 +Instance type: r7i.xlarge, r7i.2xlarge, or m7i.2xlarge +Security group: sg-cuvs-opensearch +IAM role: opensearch-cuvs-bench-ec2-role +Storage: 100+ GiB gp3, larger if needed +``` + +Under **Advanced details > Metadata options**, use: + +```text +IMDS endpoint: Enabled +IMDSv2: Required +Hop limit: 2 +``` + +After launch, copy the instance's **Private IPv4 DNS** or **Private IPv4 address**. You will use it as `OPENSEARCH_URL`. + +## 7. Launch the GPU builder node + +Launch a second EC2 instance: + +```text +Name: opensearch-cuvs-builder +AMI: AWS Deep Learning Base AMI with CUDA, Ubuntu 22.04 +Instance type: g5.xlarge or g6.xlarge +Security group: sg-cuvs-builder +IAM role: opensearch-cuvs-bench-ec2-role +Storage: 100+ GiB gp3 +``` + +Use the same metadata options: + +```text +IMDS endpoint: Enabled +IMDSv2: Required +Hop limit: 2 +``` + +After launch, copy the instance's **Private IPv4 DNS** or **Private IPv4 address**. You will use it as `REMOTE_INDEX_BUILDER_URL`. + +## 8. Launch the client node + +Launch the benchmark submitter instance: + +```text +Name: opensearch-cuvs-client +AMI: Ubuntu 22.04 or Amazon Linux 2023 +Instance type: c7i.xlarge or m7i.xlarge +Security group: sg-cuvs-client +IAM role: opensearch-cuvs-bench-ec2-role +Storage: 50-100 GiB gp3 +``` + +Use the same metadata options: + +```text +IMDS endpoint: Enabled +IMDSv2: Required +Hop limit: 2 +``` + +## 9. Install Docker and Compose on each node + +Connect to each instance with **EC2 > Instances > Connect > Session Manager**. + +Install Docker and Docker Compose if they are not already installed. Then verify: + +```bash +docker compose version +``` + +On the GPU builder node, also verify GPU access: + +```bash +nvidia-smi +docker run --rm --gpus all nvidia/cuda:12.9.0-base-ubuntu22.04 nvidia-smi +``` + +If you want to avoid `sudo docker`, add your login user to the `docker` group and reconnect: + +```bash +sudo groupadd docker 2>/dev/null || true +sudo usermod -aG docker "$(whoami)" +``` + +## 10. Copy the deployment checkout to each node + +On all three nodes, clone or copy the `deploy-opensearch-tmp` branch of the +`jrbourbeau/cuvs` fork: + +```bash +git clone --branch deploy-opensearch-tmp https://github.com/jrbourbeau/cuvs.git +cd cuvs +``` + +Make sure this file is present: + +```text +docker-compose.multinode.yml +``` + +## 11. Start OpenSearch + +On the OpenSearch node: + +```bash +export S3_BUCKET=opensearch-cuvs-bench +export AWS_REGION=us-west-2 +export AWS_DEFAULT_REGION=us-west-2 +docker compose -f docker-compose.multinode.yml --profile opensearch up -d +``` + +Verify: + +```bash +curl http://localhost:9200 +``` + +## 12. Start the remote index builder + +On the GPU builder node: + +```bash +export S3_BUCKET=opensearch-cuvs-bench +export AWS_REGION=us-west-2 +export AWS_DEFAULT_REGION=us-west-2 +docker compose -f docker-compose.multinode.yml --profile builder up -d +``` + +Verify the container is running: + +```bash +docker ps +``` + +From the OpenSearch node, verify that the builder is reachable. OpenSearch is +the service that calls the remote builder: + +```bash +python3 -c 'import socket; socket.create_connection(("BUILDER_PRIVATE_DNS_OR_IP", 1025), 5).close(); print("builder reachable")' +``` + +From the client node, run the same check. The benchmark container also waits +for the builder before starting a remote-build run: + +```bash +python3 -c 'import socket; socket.create_connection(("BUILDER_PRIVATE_DNS_OR_IP", 1025), 5).close(); print("builder reachable")' +``` + +## 13. Configure OpenSearch from the client node + +On the client node: + +```bash +export S3_BUCKET=opensearch-cuvs-bench +export OPENSEARCH_HOST=OPENSEARCH_PRIVATE_DNS_OR_IP +export OPENSEARCH_PORT=9200 +export OPENSEARCH_URL=http://${OPENSEARCH_HOST}:${OPENSEARCH_PORT} +export REMOTE_INDEX_BUILDER_URL=http://BUILDER_PRIVATE_DNS_OR_IP:1025 +export AWS_REGION=us-west-2 +export AWS_DEFAULT_REGION=us-west-2 +export S3_PREFIX=knn-indexes +export REMOTE_VECTOR_REPOSITORY=vector-repo +``` + +Then run the one-shot configure profile: + +```bash +docker compose -f docker-compose.multinode.yml --profile configure run --rm configure-remote-index-build +``` + +This registers the S3-backed remote vector repository and tells OpenSearch where the remote index builder service lives. + +## 14. Run the benchmark + +Still on the client node: + +```bash +export REMOTE_INDEX_BUILD=true +export DATASET_PATH=/data/opensearch-cuvs-datasets +export DATASET=sift-128-euclidean +export BENCH_GROUPS=test +export K=10 +export BATCH_SIZE= +export BUILD_BATCH_SIZE= + +mkdir -p "${DATASET_PATH}" + +docker compose -f docker-compose.multinode.yml --profile client build bench +docker compose -f docker-compose.multinode.yml --profile client run --rm bench +``` + +## 15. Debug checklist + +From the client node, these should all work: + +```bash +curl "$OPENSEARCH_URL" +python3 -c 'import os, socket; from urllib.parse import urlparse; url = urlparse(os.environ["REMOTE_INDEX_BUILDER_URL"]); socket.create_connection((url.hostname, url.port or 1025), 5).close(); print("builder reachable")' +aws s3 ls s3://$S3_BUCKET --region us-west-2 +``` + +If S3 fails, check the IAM role and S3 policy. If OpenSearch or builder +connectivity fails, check private DNS/IP values and security group source rules. + +## Notes + +- Keep all three instances in `us-west-2`. +- Prefer the same VPC and same Availability Zone for the first benchmark run. +- Use private DNS or private IPs, not public IPs, for cross-node service traffic. +- Docker Compose networks are local to one host; Compose service names do not resolve across EC2 instances. +- The EC2 IAM role credentials rotate automatically. Avoid freezing temporary credentials into `.env` unless the application absolutely requires literal `AWS_*` variables. diff --git a/deploy/DEPLOYMENT.md b/deploy/DEPLOYMENT.md index a04bc259d7..4084eb4cbf 100644 --- a/deploy/DEPLOYMENT.md +++ b/deploy/DEPLOYMENT.md @@ -53,7 +53,7 @@ sudo sysctl -w vm.max_map_count=262144 Set the required bucket name: ```bash -export S3_BUCKET= +export S3_BUCKET=opensearch-cuvs-bench ``` If you are using static credentials instead of a default AWS credential provider, also export: @@ -78,17 +78,17 @@ docker compose --profile gpu up --build -d --wait opensearch remote-index-builde ## Connecting OpenSearch to the GPU builder -Before any index can use GPU builds, you need to register your S3 bucket as a snapshot repository and apply the cluster settings that point OpenSearch at the builder service. Run these once against a live cluster. +Before any index can use GPU builds, you need to register an S3-backed snapshot repository and apply the cluster settings that point OpenSearch at the builder service. Run these once against a live cluster. **Register S3 repository:** ```bash -curl -X PUT http://localhost:9200/_snapshot/ \ +curl -X PUT http://localhost:9200/_snapshot/vector-repo \ -H "Content-Type: application/json" \ -d '{ "type": "s3", "settings": { - "bucket": "", + "bucket": "opensearch-cuvs-bench", "base_path": "knn-indexes", "region": "us-west-2" } @@ -103,7 +103,7 @@ curl -X PUT http://localhost:9200/_cluster/settings \ -d '{ "persistent": { "knn.remote_index_build.enabled": true, - "knn.remote_index_build.repository": "", + "knn.remote_index_build.repository": "vector-repo", "knn.remote_index_build.service.endpoint": "http://remote-index-builder:1025" } }' diff --git a/deploy/README.md b/deploy/README.md index 337b2866d3..ae516a8e03 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -214,9 +214,6 @@ docker compose run --rm --no-deps \ -e BUILDER_URL=http://remote-index-builder:1025 \ -e S3_BUCKET=${S3_BUCKET} \ -e AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION} \ - -e S3_ACCESS_KEY=${AWS_ACCESS_KEY_ID} \ - -e S3_SECRET_KEY=${AWS_SECRET_ACCESS_KEY} \ - -e S3_SESSION_TOKEN=${AWS_SESSION_TOKEN} \ bench \ pytest /opt/cuvs/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py -v -m integration ``` diff --git a/deploy/bench/Dockerfile b/deploy/bench/Dockerfile index c709e5af20..20adbca588 100644 --- a/deploy/bench/Dockerfile +++ b/deploy/bench/Dockerfile @@ -10,7 +10,7 @@ RUN apt-get update \ # nvcc, and cuvs_bench declares a hard dep on the `cuvs` CUDA package). # The opensearch backend is pure Python and needs neither, so we add the # package directly to PYTHONPATH and install only the actual runtime deps. -RUN git clone --depth=1 --filter=blob:none --sparse --branch cuvs-bench-opensearch https://github.com/jrbourbeau/cuvs.git /opt/cuvs \ +RUN git clone --depth=1 --filter=blob:none --sparse https://github.com/rapidsai/cuvs.git /opt/cuvs \ && cd /opt/cuvs \ && git sparse-checkout set python/cuvs_bench From a926b29b44b4a43c90331152cb403cdcf37aba99 Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Tue, 21 Jul 2026 11:56:07 -0500 Subject: [PATCH 10/22] Update Signed-off-by: James Bourbeau --- deploy/AWS_MULTINODE_SETUP.md | 2 +- deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md | 2 +- deploy/DEPLOYMENT.md | 2 + deploy/README.md | 36 +++++++- deploy/bench/Dockerfile | 8 +- deploy/bench/entrypoint.sh | 20 ++++- deploy/bench/prepare_custom_dataset.py | 85 +++++++++++++++++++ deploy/bench/run.py | 6 ++ deploy/docker-compose.multinode.yml | 4 + deploy/docker-compose.yml | 9 +- deploy/remote-index-build/run.py | 3 + .../cuvs_bench/backends/opensearch.py | 15 +++- .../cuvs_bench/tests/test_opensearch.py | 28 ++++++ 13 files changed, 207 insertions(+), 13 deletions(-) create mode 100644 deploy/bench/prepare_custom_dataset.py diff --git a/deploy/AWS_MULTINODE_SETUP.md b/deploy/AWS_MULTINODE_SETUP.md index cff5d973e4..c181c04d28 100644 --- a/deploy/AWS_MULTINODE_SETUP.md +++ b/deploy/AWS_MULTINODE_SETUP.md @@ -330,7 +330,7 @@ Still on the client node: ```bash export REMOTE_INDEX_BUILD=true export DATASET_PATH="$(pwd)/opensearch-cuvs-datasets" -export DATASET=sift-128-euclidean +export DATASET=miracl-en-5m-1024d-fp32 export BENCH_GROUPS=test export K=10 export BATCH_SIZE= diff --git a/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md b/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md index c9d555bde1..d33cb71bcb 100644 --- a/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md +++ b/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md @@ -320,7 +320,7 @@ Still on the client node: ```bash export REMOTE_INDEX_BUILD=true export DATASET_PATH=/data/opensearch-cuvs-datasets -export DATASET=sift-128-euclidean +export DATASET=miracl-en-5m-1024d-fp32 export BENCH_GROUPS=test export K=10 export BATCH_SIZE= diff --git a/deploy/DEPLOYMENT.md b/deploy/DEPLOYMENT.md index 4084eb4cbf..8701eabaa8 100644 --- a/deploy/DEPLOYMENT.md +++ b/deploy/DEPLOYMENT.md @@ -121,6 +121,7 @@ curl -X PUT http://localhost:9200/my-vectors \ -d '{ "settings": { "index.knn": true, + "index.knn.advanced.approximate_threshold": 10000, "index.knn.remote_index_build.enabled": true, "index.knn.remote_index_build.size.min": "1kb", "number_of_shards": 1, @@ -161,6 +162,7 @@ docker compose run --rm \ -e S3_BUCKET=${S3_BUCKET} \ -e AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION:-us-west-2} \ -e REMOTE_BUILD_SIZE_MIN=${REMOTE_BUILD_SIZE_MIN:-} \ + -e APPROXIMATE_THRESHOLD=${APPROXIMATE_THRESHOLD:-10000} \ -e REMOTE_BUILD_TIMEOUT=${REMOTE_BUILD_TIMEOUT:-1800} \ -v $(pwd)/remote-index-build:/app/remote-index-build \ --no-deps bench \ diff --git a/deploy/README.md b/deploy/README.md index ae516a8e03..2a352f98bb 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -72,6 +72,7 @@ export BENCH_GROUPS=test # test | base (default: test) export K=10 # number of neighbors (default: 10) export BATCH_SIZE= # optional query batch size override export BUILD_BATCH_SIZE= # optional bulk ingest batch size override +export APPROXIMATE_THRESHOLD=10000 # AWS-recommended GPU indexing starting value export REMOTE_BUILD_TIMEOUT=1800 # seconds to wait for remote builds (default: 1800) ``` @@ -87,6 +88,27 @@ docker compose --profile gpu up --build The `bench` container logs its progress through each phase. When complete you'll see a results table followed by the paths to the generated plot PNGs under `$DATASET_PATH`. +### MIRACL 5M custom dataset + +Set `DATASET=miracl-en-5m-1024d-fp32` to use the custom dataset from +`s3://opensearch-cuvs-bench/miracl-en-5m-1024d-fp32/miracl-en-5m-1024d-fp32/` +instead of a built-in cuvs-bench dataset: + +```bash +export DATASET=miracl-en-5m-1024d-fp32 +export AWS_DEFAULT_REGION=us-west-2 +docker compose up --build # CPU build +# or: docker compose --profile gpu up --build +``` + +The bench container downloads the dataset's `config.yaml` and every file it +references—including `base.fbin`, `query.fbin`, and +`groundtruth.neighbors.ibin`—into `$DATASET_PATH`, skipping files that are +already present. AWS credentials are optional when the container can use an +AWS default credential provider such as an EC2 instance role; otherwise set +the AWS credential variables shown above. The uploaded ground-truth neighbors +are used for the benchmark's recall calculation. + To tear everything down: ```bash @@ -99,7 +121,7 @@ docker compose down -v 2. **GPU mode only**: Registers the S3 bucket as an OpenSearch snapshot repository 3. **GPU mode only**: Applies cluster settings to enable remote index build and point OpenSearch at the builder service 4. Runs `cuvs-bench` build phase (handled entirely by the OpenSearch backend): - - Creates the kNN index and bulk-ingests dataset vectors + - Creates the kNN index with `index.knn.advanced.approximate_threshold=10000` and bulk-ingests dataset vectors - **GPU mode**: Flushes segments, waits for all submitted remote GPU builds to complete, and polls the kNN stats API every 5 s until the build is confirmed complete - **CPU mode**: Flushes and refreshes the local OpenSearch index - Uses the backend's automatic OpenSearch bulk-ingest batch sizing by default; set `BUILD_BATCH_SIZE` to override it @@ -124,7 +146,7 @@ Supported extensions: `.fbin` (float32), `.f16bin` (float16), `.u8bin` (uint8), ``` $DATASET_PATH/ - sift-128-euclidean/ + miracl-en-5m-1024d-fp32/ base.fbin query.fbin groundtruth.neighbors.ibin @@ -144,6 +166,16 @@ $DATASET_PATH/ } ``` +For GPU runs, `bench/run.py` passes `APPROXIMATE_THRESHOLD` to the OpenSearch +backend, which sets `index.knn.advanced.approximate_threshold` in the index +creation request. The default is `10000`, the starting value recommended by +AWS to avoid smaller segment index builds. + +The benchmark image clones `CUVS_BRANCH=deploy-opensearch-tmp` from +`CUVS_REPOSITORY=https://github.com/jrbourbeau/cuvs.git` so it includes the +matching OpenSearch backend. Both values are Docker build arguments and can be +overridden with environment variables before running Docker Compose. + **Parameter groups** (`BENCH_GROUPS`): | Group | Build params | Search params | Use case | diff --git a/deploy/bench/Dockerfile b/deploy/bench/Dockerfile index 20adbca588..f133d2a52e 100644 --- a/deploy/bench/Dockerfile +++ b/deploy/bench/Dockerfile @@ -1,6 +1,9 @@ FROM python:3.11-slim WORKDIR /app +ARG CUVS_REPOSITORY=https://github.com/jrbourbeau/cuvs.git +ARG CUVS_BRANCH=deploy-opensearch-tmp + RUN apt-get update \ && apt-get install -y --no-install-recommends git \ && rm -rf /var/lib/apt/lists/* @@ -10,7 +13,8 @@ RUN apt-get update \ # nvcc, and cuvs_bench declares a hard dep on the `cuvs` CUDA package). # The opensearch backend is pure Python and needs neither, so we add the # package directly to PYTHONPATH and install only the actual runtime deps. -RUN git clone --depth=1 --filter=blob:none --sparse https://github.com/rapidsai/cuvs.git /opt/cuvs \ +RUN git clone --depth=1 --filter=blob:none --sparse \ + --branch "${CUVS_BRANCH}" "${CUVS_REPOSITORY}" /opt/cuvs \ && cd /opt/cuvs \ && git sparse-checkout set python/cuvs_bench @@ -33,5 +37,5 @@ RUN pip install --no-cache-dir \ scikit-learn \ scipy -COPY --chmod=755 run.py entrypoint.sh . +COPY --chmod=755 run.py entrypoint.sh prepare_custom_dataset.py . CMD ["/app/entrypoint.sh"] diff --git a/deploy/bench/entrypoint.sh b/deploy/bench/entrypoint.sh index 78f050f9e4..b897dcc398 100644 --- a/deploy/bench/entrypoint.sh +++ b/deploy/bench/entrypoint.sh @@ -2,10 +2,18 @@ set -e DATASET="${DATASET:-sift-128-euclidean}" +CUSTOM_DATASET="miracl-en-5m-1024d-fp32" BENCH_GROUPS="${BENCH_GROUPS:-test}" K="${K:-10}" ALGORITHM="opensearch_faiss_hnsw" +export DATASET +if [ "$DATASET" = "$CUSTOM_DATASET" ]; then + export DATASET_CONFIGURATION="/data/datasets/${DATASET}/config.yaml" +else + unset DATASET_CONFIGURATION +fi + wait_for_builder() { builder_url="${BUILDER_URL:-http://remote-index-builder:1025}" echo "Remote index build enabled — waiting for builder at ${builder_url}..." @@ -44,10 +52,14 @@ else fi fi -# Step 1: Download dataset (skipped automatically if already present) -python -m cuvs_bench.get_dataset \ - --dataset "$DATASET" \ - --dataset-path /data/datasets +# Step 1: Prepare the selected dataset (skipped when files already exist). +if [ "$DATASET" = "$CUSTOM_DATASET" ]; then + python -u prepare_custom_dataset.py +else + python -m cuvs_bench.get_dataset \ + --dataset "$DATASET" \ + --dataset-path /data/datasets +fi # Step 2: Run benchmark (build + search + writes result JSON files) python -u run.py diff --git a/deploy/bench/prepare_custom_dataset.py b/deploy/bench/prepare_custom_dataset.py new file mode 100644 index 0000000000..4b13e7be3c --- /dev/null +++ b/deploy/bench/prepare_custom_dataset.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +"""Download the custom cuvs-bench dataset described by its S3 YAML file.""" + +import os +from pathlib import Path, PurePosixPath + +import boto3 +import yaml +from botocore.exceptions import ClientError + + +DATASET_NAME = "miracl-en-5m-1024d-fp32" +S3_BUCKET = "opensearch-cuvs-bench" +S3_PREFIX = "miracl-en-5m-1024d-fp32/miracl-en-5m-1024d-fp32" + + +def download_if_needed(s3, bucket: str, key: str, destination: Path) -> None: + if destination.is_file() and destination.stat().st_size > 0: + print(f"Using existing {destination}") + return + + destination.parent.mkdir(parents=True, exist_ok=True) + temporary = destination.with_suffix(destination.suffix + ".part") + print(f"Downloading s3://{bucket}/{key} -> {destination}") + try: + s3.download_file(bucket, key, str(temporary)) + except ClientError as error: + temporary.unlink(missing_ok=True) + raise RuntimeError( + f"Could not download s3://{bucket}/{key}. The dataset YAML " + "references this file, so it is required by the benchmark." + ) from error + temporary.replace(destination) + + +def main() -> None: + dataset_root = Path(os.environ.get("DATASET_PATH", "/data/datasets")) + region = os.environ.get("AWS_DEFAULT_REGION", "us-west-2") + s3 = boto3.client("s3", region_name=region) + + config_path = dataset_root / DATASET_NAME / "config.yaml" + download_if_needed( + s3, S3_BUCKET, f"{S3_PREFIX}/config.yaml", config_path + ) + + with config_path.open() as config_file: + configs = yaml.safe_load(config_file) + if not isinstance(configs, list): + raise ValueError(f"Expected a list in {config_path}") + + try: + config = next(item for item in configs if item["name"] == DATASET_NAME) + except (KeyError, StopIteration) as error: + raise ValueError( + f"Dataset {DATASET_NAME!r} was not found in {config_path}" + ) from error + + file_fields = ( + "base_file", + "query_file", + "groundtruth_neighbors_file", + "groundtruth_distances_file", + ) + for field in file_fields: + relative_path = config.get(field) + if not relative_path: + continue + relative_path = PurePosixPath(relative_path) + if relative_path.is_absolute() or ".." in relative_path.parts: + raise ValueError( + f"Unsafe {field} path in {config_path}: {relative_path}" + ) + download_if_needed( + s3, + S3_BUCKET, + f"{S3_PREFIX}/{relative_path.name}", + dataset_root.joinpath(*relative_path.parts), + ) + + print(f"Custom dataset is ready: {DATASET_NAME}") + print(f"Dataset configuration: {config_path}") + + +if __name__ == "__main__": + main() diff --git a/deploy/bench/run.py b/deploy/bench/run.py index 55bdc1777f..d7bb00bdee 100644 --- a/deploy/bench/run.py +++ b/deploy/bench/run.py @@ -30,6 +30,7 @@ REMOTE_INDEX_BUILD = os.environ.get("REMOTE_INDEX_BUILD", "false").lower() == "true" REMOTE_BUILD_SIZE_MIN = os.environ.get("REMOTE_BUILD_SIZE_MIN", "").strip() REMOTE_BUILD_TIMEOUT = int(os.environ.get("REMOTE_BUILD_TIMEOUT", "1800")) +APPROXIMATE_THRESHOLD = int(os.environ.get("APPROXIMATE_THRESHOLD", "10000")) S3_BUCKET = os.environ.get("S3_BUCKET", "").strip() S3_PREFIX = os.environ.get("S3_PREFIX", "knn-indexes").strip() or "knn-indexes" @@ -37,6 +38,7 @@ DATASET = os.environ.get("DATASET", "sift-128-euclidean") DATASET_PATH = os.environ.get("DATASET_PATH", "/data/datasets") +DATASET_CONFIGURATION = os.environ.get("DATASET_CONFIGURATION", "").strip() BENCH_GROUPS = os.environ.get("BENCH_GROUPS", "test") K = int(os.environ.get("K", "10")) BATCH_SIZE = os.environ.get("BATCH_SIZE", "").strip() @@ -274,6 +276,7 @@ def main() -> None: print(f" S3 bucket : s3://{S3_BUCKET}/{S3_PREFIX}/ (region: {S3_REGION})") print(f" Repository : {REPO_NAME}") print(f" Build size minimum : {REMOTE_BUILD_SIZE_MIN or 'OpenSearch default'}") + print(f" Approx. threshold : {APPROXIMATE_THRESHOLD}") print(f" Build timeout : {REMOTE_BUILD_TIMEOUT}s") print(f" Dataset : {DATASET} (path: {DATASET_PATH})") print(f" Algorithm : {ALGORITHM}") @@ -304,6 +307,8 @@ def main() -> None: use_ssl=False, verify_certs=False, ) + if DATASET_CONFIGURATION: + common_kwargs["dataset_configuration"] = DATASET_CONFIGURATION build_kwargs = dict( common_kwargs, @@ -313,6 +318,7 @@ def main() -> None: build_kwargs["build_batch_size"] = BUILD_BATCH_SIZE if REMOTE_INDEX_BUILD: build_kwargs["remote_build_timeout"] = REMOTE_BUILD_TIMEOUT + build_kwargs["approximate_threshold"] = APPROXIMATE_THRESHOLD if REMOTE_BUILD_SIZE_MIN: build_kwargs["remote_build_size_min"] = REMOTE_BUILD_SIZE_MIN diff --git a/deploy/docker-compose.multinode.yml b/deploy/docker-compose.multinode.yml index c88f95d9e5..ef7b533cf8 100644 --- a/deploy/docker-compose.multinode.yml +++ b/deploy/docker-compose.multinode.yml @@ -42,6 +42,7 @@ x-bench-env: &bench-env REMOTE_INDEX_BUILDER_URL: ${REMOTE_INDEX_BUILDER_URL:-http://remote-index-builder:1025} REMOTE_INDEX_BUILD: ${REMOTE_INDEX_BUILD:-true} REMOTE_BUILD_SIZE_MIN: ${REMOTE_BUILD_SIZE_MIN:-} + APPROXIMATE_THRESHOLD: ${APPROXIMATE_THRESHOLD:-10000} REMOTE_BUILD_TIMEOUT: ${REMOTE_BUILD_TIMEOUT:-1800} REMOTE_VECTOR_REPOSITORY: ${REMOTE_VECTOR_REPOSITORY:-vector-repo} DATASET: ${DATASET:-sift-128-euclidean} @@ -116,6 +117,9 @@ services: build: context: ${BENCH_BUILD_CONTEXT:-./bench} dockerfile: ${BENCH_DOCKERFILE:-Dockerfile} + args: + CUVS_REPOSITORY: ${CUVS_REPOSITORY:-https://github.com/jrbourbeau/cuvs.git} + CUVS_BRANCH: ${CUVS_BRANCH:-deploy-opensearch-tmp} environment: <<: *bench-env extra_hosts: diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml index 4b6300ecb8..cdaa95908a 100644 --- a/deploy/docker-compose.yml +++ b/deploy/docker-compose.yml @@ -29,14 +29,17 @@ # AWS_DEFAULT_REGION AWS region for the S3 bucket (default: us-west-2) # REMOTE_INDEX_BUILD Override GPU/CPU mode detection (true/false); normally auto-detected # REMOTE_BUILD_SIZE_MIN Optional minimum segment size override for remote builds +# APPROXIMATE_THRESHOLD Minimum vectors before approximate indexing (default: 10000) # REMOTE_BUILD_TIMEOUT Remote build wait timeout in seconds (default: 1800) # REMOTE_VECTOR_REPOSITORY Optional snapshot repository name (default: vector-repo) # S3_PREFIX Optional S3 prefix for staged vectors/indexes (default: knn-indexes) -# DATASET Dataset name (default: sift-128-euclidean) +# DATASET Dataset name (default: sift-128-euclidean); MIRACL uses custom S3 # BENCH_GROUPS Parameter sweep group: test | base (default: test) # K Number of neighbors to search for (default: 10) # BATCH_SIZE Optional cuvs-bench query batch size override # BUILD_BATCH_SIZE Optional OpenSearch bulk ingest batch size override +# CUVS_REPOSITORY Repository cloned into the benchmark image +# CUVS_BRANCH Repository branch cloned into the benchmark image # # Usage: # CPU: docker compose up --build @@ -108,6 +111,9 @@ services: bench: build: context: ./bench + args: + CUVS_REPOSITORY: ${CUVS_REPOSITORY:-https://github.com/jrbourbeau/cuvs.git} + CUVS_BRANCH: ${CUVS_BRANCH:-deploy-opensearch-tmp} depends_on: opensearch: condition: service_healthy @@ -119,6 +125,7 @@ services: BUILDER_URL: http://remote-index-builder:1025 REMOTE_INDEX_BUILD: ${REMOTE_INDEX_BUILD:-} REMOTE_BUILD_SIZE_MIN: ${REMOTE_BUILD_SIZE_MIN:-} + APPROXIMATE_THRESHOLD: ${APPROXIMATE_THRESHOLD:-10000} REMOTE_BUILD_TIMEOUT: ${REMOTE_BUILD_TIMEOUT:-1800} REMOTE_VECTOR_REPOSITORY: ${REMOTE_VECTOR_REPOSITORY:-vector-repo} S3_BUCKET: ${S3_BUCKET:-} diff --git a/deploy/remote-index-build/run.py b/deploy/remote-index-build/run.py index a7548ded04..1a21b65c62 100644 --- a/deploy/remote-index-build/run.py +++ b/deploy/remote-index-build/run.py @@ -36,6 +36,7 @@ REPO_NAME = "vector-repo" REMOTE_BUILD_SIZE_MIN = os.environ.get("REMOTE_BUILD_SIZE_MIN", "").strip() REMOTE_BUILD_TIMEOUT = int(os.environ.get("REMOTE_BUILD_TIMEOUT", "1800")) +APPROXIMATE_THRESHOLD = int(os.environ.get("APPROXIMATE_THRESHOLD", "10000")) session = requests.Session() session.headers.update({"Content-Type": "application/json"}) @@ -93,6 +94,7 @@ def create_index() -> None: index_settings = { "index.knn": True, + "index.knn.advanced.approximate_threshold": APPROXIMATE_THRESHOLD, "index.knn.remote_index_build.enabled": True, "number_of_shards": 1, "number_of_replicas": 0, @@ -311,6 +313,7 @@ def main() -> None: "engine=faiss method=hnsw space=l2" ) print(f" Build size minimum: {REMOTE_BUILD_SIZE_MIN or 'OpenSearch default'}") + print(f" Approx. threshold : {APPROXIMATE_THRESHOLD}") print(f" Build timeout: {REMOTE_BUILD_TIMEOUT}s") register_repository() diff --git a/python/cuvs_bench/cuvs_bench/backends/opensearch.py b/python/cuvs_bench/cuvs_bench/backends/opensearch.py index 026fe98fd6..c197418204 100644 --- a/python/cuvs_bench/cuvs_bench/backends/opensearch.py +++ b/python/cuvs_bench/cuvs_bench/backends/opensearch.py @@ -134,6 +134,7 @@ def _build_benchmark_configs( # Remote Index Build (OpenSearch 3.0+, faiss engine only) "remote_index_build", "remote_build_size_min", + "approximate_threshold", "remote_build_timeout", ) conn_kwargs = {k: kwargs[k] for k in _conn_keys if k in kwargs} @@ -282,6 +283,8 @@ class OpenSearchBackend(BenchmarkBackend): on the index at creation time, opting it into the GPU build path (default: ``False``). - ``remote_build_size_min`` – minimum segment size to trigger GPU build, e.g. ``"1kb"`` (default: OpenSearch's default) + - ``approximate_threshold`` – minimum number of vectors in a segment before + building an approximate index (default: OpenSearch's default) - ``remote_build_timeout`` – seconds to wait for GPU build (default: ``1800``) """ @@ -348,6 +351,7 @@ def _build_index_mapping( build_param: Dict[str, Any], remote_index_build: bool = False, remote_build_size_min: Optional[str] = None, + approximate_threshold: Optional[int] = None, ) -> Dict[str, Any]: """ Construct the OpenSearch index mapping dict for k-NN. @@ -359,8 +363,9 @@ def _build_index_mapping( is set to ``true`` in the index settings, opting qualifying segments into the GPU build path. If ``remote_build_size_min`` is provided it overrides ``index.knn.remote_index_build.size.min``; otherwise OpenSearch's default - applies. The cluster-level infrastructure is assumed to be pre-configured - externally. + applies. If ``approximate_threshold`` is provided, it sets + ``index.knn.advanced.approximate_threshold`` when the index is created. + The cluster-level infrastructure is assumed to be pre-configured externally. """ m = build_param.get("m", 16) ef_construction = build_param.get("ef_construction", 100) @@ -397,6 +402,10 @@ def _build_index_mapping( "number_of_shards": build_param.get("number_of_shards", 1), "number_of_replicas": build_param.get("number_of_replicas", 0), } + if approximate_threshold is not None: + index_settings["knn.advanced.approximate_threshold"] = ( + approximate_threshold + ) if remote_index_build: if engine != "faiss": raise ValueError( @@ -702,6 +711,7 @@ def build( build_batch_size = self.config.get("build_batch_size") remote_index_build = bool(self.config.get("remote_index_build", False)) remote_build_size_min = self.config.get("remote_build_size_min") + approximate_threshold = self.config.get("approximate_threshold") if dry_run: print( @@ -753,6 +763,7 @@ def build( build_param, remote_index_build, remote_build_size_min, + approximate_threshold, ) self._client.indices.create(index=index_name, body=mapping) diff --git a/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py b/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py index d16eedf25d..8bc366b642 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py @@ -155,6 +155,7 @@ def test_load_forwards_remote_build_kwargs(self, config_dir): dataset_path="/data", remote_index_build=True, remote_build_size_min="2kb", + approximate_threshold=10_000, remote_build_timeout=123, remote_build_s3_endpoint="http://s3:9000", ) @@ -162,6 +163,7 @@ def test_load_forwards_remote_build_kwargs(self, config_dir): bc = configs[0].backend_config assert bc["remote_index_build"] is True assert bc["remote_build_size_min"] == "2kb" + assert bc["approximate_threshold"] == 10_000 assert bc["remote_build_timeout"] == 123 assert "remote_build_s3_endpoint" not in bc @@ -221,6 +223,32 @@ def test_remote_build_size_min_overrides_default(self): settings = mapping["settings"]["index"] assert settings["knn.remote_index_build.size.min"] == "2kb" + def test_approximate_threshold_is_set_when_specified(self): + backend = _make_backend() + mapping = backend._build_index_mapping( + dims=4, + engine="faiss", + space_type="l2", + build_param={}, + remote_index_build=True, + approximate_threshold=10_000, + ) + + settings = mapping["settings"]["index"] + assert settings["knn.advanced.approximate_threshold"] == 10_000 + + def test_approximate_threshold_uses_opensearch_default_when_unspecified(self): + backend = _make_backend() + mapping = backend._build_index_mapping( + dims=4, + engine="faiss", + space_type="l2", + build_param={}, + ) + + settings = mapping["settings"]["index"] + assert "knn.advanced.approximate_threshold" not in settings + def test_wait_for_remote_build_raises_on_failure_count(self): backend = _make_backend() initial_stats = { From 5b6ba4f44a2b45a17eae312f5a1294e72c802e8e Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Tue, 21 Jul 2026 14:03:48 -0500 Subject: [PATCH 11/22] Revert approximate_threshold additions Signed-off-by: James Bourbeau --- deploy/DEPLOYMENT.md | 2 -- deploy/README.md | 8 +----- deploy/bench/run.py | 3 -- deploy/docker-compose.multinode.yml | 1 - deploy/docker-compose.yml | 2 -- deploy/remote-index-build/run.py | 3 -- .../cuvs_bench/backends/opensearch.py | 15 ++-------- .../cuvs_bench/tests/test_opensearch.py | 28 ------------------- 8 files changed, 3 insertions(+), 59 deletions(-) diff --git a/deploy/DEPLOYMENT.md b/deploy/DEPLOYMENT.md index 8701eabaa8..4084eb4cbf 100644 --- a/deploy/DEPLOYMENT.md +++ b/deploy/DEPLOYMENT.md @@ -121,7 +121,6 @@ curl -X PUT http://localhost:9200/my-vectors \ -d '{ "settings": { "index.knn": true, - "index.knn.advanced.approximate_threshold": 10000, "index.knn.remote_index_build.enabled": true, "index.knn.remote_index_build.size.min": "1kb", "number_of_shards": 1, @@ -162,7 +161,6 @@ docker compose run --rm \ -e S3_BUCKET=${S3_BUCKET} \ -e AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION:-us-west-2} \ -e REMOTE_BUILD_SIZE_MIN=${REMOTE_BUILD_SIZE_MIN:-} \ - -e APPROXIMATE_THRESHOLD=${APPROXIMATE_THRESHOLD:-10000} \ -e REMOTE_BUILD_TIMEOUT=${REMOTE_BUILD_TIMEOUT:-1800} \ -v $(pwd)/remote-index-build:/app/remote-index-build \ --no-deps bench \ diff --git a/deploy/README.md b/deploy/README.md index 2a352f98bb..c185e6b742 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -72,7 +72,6 @@ export BENCH_GROUPS=test # test | base (default: test) export K=10 # number of neighbors (default: 10) export BATCH_SIZE= # optional query batch size override export BUILD_BATCH_SIZE= # optional bulk ingest batch size override -export APPROXIMATE_THRESHOLD=10000 # AWS-recommended GPU indexing starting value export REMOTE_BUILD_TIMEOUT=1800 # seconds to wait for remote builds (default: 1800) ``` @@ -121,7 +120,7 @@ docker compose down -v 2. **GPU mode only**: Registers the S3 bucket as an OpenSearch snapshot repository 3. **GPU mode only**: Applies cluster settings to enable remote index build and point OpenSearch at the builder service 4. Runs `cuvs-bench` build phase (handled entirely by the OpenSearch backend): - - Creates the kNN index with `index.knn.advanced.approximate_threshold=10000` and bulk-ingests dataset vectors + - Creates the kNN index and bulk-ingests dataset vectors - **GPU mode**: Flushes segments, waits for all submitted remote GPU builds to complete, and polls the kNN stats API every 5 s until the build is confirmed complete - **CPU mode**: Flushes and refreshes the local OpenSearch index - Uses the backend's automatic OpenSearch bulk-ingest batch sizing by default; set `BUILD_BATCH_SIZE` to override it @@ -166,11 +165,6 @@ $DATASET_PATH/ } ``` -For GPU runs, `bench/run.py` passes `APPROXIMATE_THRESHOLD` to the OpenSearch -backend, which sets `index.knn.advanced.approximate_threshold` in the index -creation request. The default is `10000`, the starting value recommended by -AWS to avoid smaller segment index builds. - The benchmark image clones `CUVS_BRANCH=deploy-opensearch-tmp` from `CUVS_REPOSITORY=https://github.com/jrbourbeau/cuvs.git` so it includes the matching OpenSearch backend. Both values are Docker build arguments and can be diff --git a/deploy/bench/run.py b/deploy/bench/run.py index d7bb00bdee..8f05419865 100644 --- a/deploy/bench/run.py +++ b/deploy/bench/run.py @@ -30,7 +30,6 @@ REMOTE_INDEX_BUILD = os.environ.get("REMOTE_INDEX_BUILD", "false").lower() == "true" REMOTE_BUILD_SIZE_MIN = os.environ.get("REMOTE_BUILD_SIZE_MIN", "").strip() REMOTE_BUILD_TIMEOUT = int(os.environ.get("REMOTE_BUILD_TIMEOUT", "1800")) -APPROXIMATE_THRESHOLD = int(os.environ.get("APPROXIMATE_THRESHOLD", "10000")) S3_BUCKET = os.environ.get("S3_BUCKET", "").strip() S3_PREFIX = os.environ.get("S3_PREFIX", "knn-indexes").strip() or "knn-indexes" @@ -276,7 +275,6 @@ def main() -> None: print(f" S3 bucket : s3://{S3_BUCKET}/{S3_PREFIX}/ (region: {S3_REGION})") print(f" Repository : {REPO_NAME}") print(f" Build size minimum : {REMOTE_BUILD_SIZE_MIN or 'OpenSearch default'}") - print(f" Approx. threshold : {APPROXIMATE_THRESHOLD}") print(f" Build timeout : {REMOTE_BUILD_TIMEOUT}s") print(f" Dataset : {DATASET} (path: {DATASET_PATH})") print(f" Algorithm : {ALGORITHM}") @@ -318,7 +316,6 @@ def main() -> None: build_kwargs["build_batch_size"] = BUILD_BATCH_SIZE if REMOTE_INDEX_BUILD: build_kwargs["remote_build_timeout"] = REMOTE_BUILD_TIMEOUT - build_kwargs["approximate_threshold"] = APPROXIMATE_THRESHOLD if REMOTE_BUILD_SIZE_MIN: build_kwargs["remote_build_size_min"] = REMOTE_BUILD_SIZE_MIN diff --git a/deploy/docker-compose.multinode.yml b/deploy/docker-compose.multinode.yml index ef7b533cf8..cadd95bd2a 100644 --- a/deploy/docker-compose.multinode.yml +++ b/deploy/docker-compose.multinode.yml @@ -42,7 +42,6 @@ x-bench-env: &bench-env REMOTE_INDEX_BUILDER_URL: ${REMOTE_INDEX_BUILDER_URL:-http://remote-index-builder:1025} REMOTE_INDEX_BUILD: ${REMOTE_INDEX_BUILD:-true} REMOTE_BUILD_SIZE_MIN: ${REMOTE_BUILD_SIZE_MIN:-} - APPROXIMATE_THRESHOLD: ${APPROXIMATE_THRESHOLD:-10000} REMOTE_BUILD_TIMEOUT: ${REMOTE_BUILD_TIMEOUT:-1800} REMOTE_VECTOR_REPOSITORY: ${REMOTE_VECTOR_REPOSITORY:-vector-repo} DATASET: ${DATASET:-sift-128-euclidean} diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml index cdaa95908a..6572571968 100644 --- a/deploy/docker-compose.yml +++ b/deploy/docker-compose.yml @@ -29,7 +29,6 @@ # AWS_DEFAULT_REGION AWS region for the S3 bucket (default: us-west-2) # REMOTE_INDEX_BUILD Override GPU/CPU mode detection (true/false); normally auto-detected # REMOTE_BUILD_SIZE_MIN Optional minimum segment size override for remote builds -# APPROXIMATE_THRESHOLD Minimum vectors before approximate indexing (default: 10000) # REMOTE_BUILD_TIMEOUT Remote build wait timeout in seconds (default: 1800) # REMOTE_VECTOR_REPOSITORY Optional snapshot repository name (default: vector-repo) # S3_PREFIX Optional S3 prefix for staged vectors/indexes (default: knn-indexes) @@ -125,7 +124,6 @@ services: BUILDER_URL: http://remote-index-builder:1025 REMOTE_INDEX_BUILD: ${REMOTE_INDEX_BUILD:-} REMOTE_BUILD_SIZE_MIN: ${REMOTE_BUILD_SIZE_MIN:-} - APPROXIMATE_THRESHOLD: ${APPROXIMATE_THRESHOLD:-10000} REMOTE_BUILD_TIMEOUT: ${REMOTE_BUILD_TIMEOUT:-1800} REMOTE_VECTOR_REPOSITORY: ${REMOTE_VECTOR_REPOSITORY:-vector-repo} S3_BUCKET: ${S3_BUCKET:-} diff --git a/deploy/remote-index-build/run.py b/deploy/remote-index-build/run.py index 1a21b65c62..a7548ded04 100644 --- a/deploy/remote-index-build/run.py +++ b/deploy/remote-index-build/run.py @@ -36,7 +36,6 @@ REPO_NAME = "vector-repo" REMOTE_BUILD_SIZE_MIN = os.environ.get("REMOTE_BUILD_SIZE_MIN", "").strip() REMOTE_BUILD_TIMEOUT = int(os.environ.get("REMOTE_BUILD_TIMEOUT", "1800")) -APPROXIMATE_THRESHOLD = int(os.environ.get("APPROXIMATE_THRESHOLD", "10000")) session = requests.Session() session.headers.update({"Content-Type": "application/json"}) @@ -94,7 +93,6 @@ def create_index() -> None: index_settings = { "index.knn": True, - "index.knn.advanced.approximate_threshold": APPROXIMATE_THRESHOLD, "index.knn.remote_index_build.enabled": True, "number_of_shards": 1, "number_of_replicas": 0, @@ -313,7 +311,6 @@ def main() -> None: "engine=faiss method=hnsw space=l2" ) print(f" Build size minimum: {REMOTE_BUILD_SIZE_MIN or 'OpenSearch default'}") - print(f" Approx. threshold : {APPROXIMATE_THRESHOLD}") print(f" Build timeout: {REMOTE_BUILD_TIMEOUT}s") register_repository() diff --git a/python/cuvs_bench/cuvs_bench/backends/opensearch.py b/python/cuvs_bench/cuvs_bench/backends/opensearch.py index c197418204..026fe98fd6 100644 --- a/python/cuvs_bench/cuvs_bench/backends/opensearch.py +++ b/python/cuvs_bench/cuvs_bench/backends/opensearch.py @@ -134,7 +134,6 @@ def _build_benchmark_configs( # Remote Index Build (OpenSearch 3.0+, faiss engine only) "remote_index_build", "remote_build_size_min", - "approximate_threshold", "remote_build_timeout", ) conn_kwargs = {k: kwargs[k] for k in _conn_keys if k in kwargs} @@ -283,8 +282,6 @@ class OpenSearchBackend(BenchmarkBackend): on the index at creation time, opting it into the GPU build path (default: ``False``). - ``remote_build_size_min`` – minimum segment size to trigger GPU build, e.g. ``"1kb"`` (default: OpenSearch's default) - - ``approximate_threshold`` – minimum number of vectors in a segment before - building an approximate index (default: OpenSearch's default) - ``remote_build_timeout`` – seconds to wait for GPU build (default: ``1800``) """ @@ -351,7 +348,6 @@ def _build_index_mapping( build_param: Dict[str, Any], remote_index_build: bool = False, remote_build_size_min: Optional[str] = None, - approximate_threshold: Optional[int] = None, ) -> Dict[str, Any]: """ Construct the OpenSearch index mapping dict for k-NN. @@ -363,9 +359,8 @@ def _build_index_mapping( is set to ``true`` in the index settings, opting qualifying segments into the GPU build path. If ``remote_build_size_min`` is provided it overrides ``index.knn.remote_index_build.size.min``; otherwise OpenSearch's default - applies. If ``approximate_threshold`` is provided, it sets - ``index.knn.advanced.approximate_threshold`` when the index is created. - The cluster-level infrastructure is assumed to be pre-configured externally. + applies. The cluster-level infrastructure is assumed to be pre-configured + externally. """ m = build_param.get("m", 16) ef_construction = build_param.get("ef_construction", 100) @@ -402,10 +397,6 @@ def _build_index_mapping( "number_of_shards": build_param.get("number_of_shards", 1), "number_of_replicas": build_param.get("number_of_replicas", 0), } - if approximate_threshold is not None: - index_settings["knn.advanced.approximate_threshold"] = ( - approximate_threshold - ) if remote_index_build: if engine != "faiss": raise ValueError( @@ -711,7 +702,6 @@ def build( build_batch_size = self.config.get("build_batch_size") remote_index_build = bool(self.config.get("remote_index_build", False)) remote_build_size_min = self.config.get("remote_build_size_min") - approximate_threshold = self.config.get("approximate_threshold") if dry_run: print( @@ -763,7 +753,6 @@ def build( build_param, remote_index_build, remote_build_size_min, - approximate_threshold, ) self._client.indices.create(index=index_name, body=mapping) diff --git a/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py b/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py index 8bc366b642..d16eedf25d 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py @@ -155,7 +155,6 @@ def test_load_forwards_remote_build_kwargs(self, config_dir): dataset_path="/data", remote_index_build=True, remote_build_size_min="2kb", - approximate_threshold=10_000, remote_build_timeout=123, remote_build_s3_endpoint="http://s3:9000", ) @@ -163,7 +162,6 @@ def test_load_forwards_remote_build_kwargs(self, config_dir): bc = configs[0].backend_config assert bc["remote_index_build"] is True assert bc["remote_build_size_min"] == "2kb" - assert bc["approximate_threshold"] == 10_000 assert bc["remote_build_timeout"] == 123 assert "remote_build_s3_endpoint" not in bc @@ -223,32 +221,6 @@ def test_remote_build_size_min_overrides_default(self): settings = mapping["settings"]["index"] assert settings["knn.remote_index_build.size.min"] == "2kb" - def test_approximate_threshold_is_set_when_specified(self): - backend = _make_backend() - mapping = backend._build_index_mapping( - dims=4, - engine="faiss", - space_type="l2", - build_param={}, - remote_index_build=True, - approximate_threshold=10_000, - ) - - settings = mapping["settings"]["index"] - assert settings["knn.advanced.approximate_threshold"] == 10_000 - - def test_approximate_threshold_uses_opensearch_default_when_unspecified(self): - backend = _make_backend() - mapping = backend._build_index_mapping( - dims=4, - engine="faiss", - space_type="l2", - build_param={}, - ) - - settings = mapping["settings"]["index"] - assert "knn.advanced.approximate_threshold" not in settings - def test_wait_for_remote_build_raises_on_failure_count(self): backend = _make_backend() initial_stats = { From 4916070f6136d2f3c3e5abccabb8dcf33d317456 Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Tue, 21 Jul 2026 15:31:53 -0500 Subject: [PATCH 12/22] Try multiple shards Signed-off-by: James Bourbeau --- .../cuvs_bench/config/algos/opensearch_faiss_hnsw.yaml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/cuvs_bench/cuvs_bench/config/algos/opensearch_faiss_hnsw.yaml b/python/cuvs_bench/cuvs_bench/config/algos/opensearch_faiss_hnsw.yaml index 6c7439be46..056fce5839 100644 --- a/python/cuvs_bench/cuvs_bench/config/algos/opensearch_faiss_hnsw.yaml +++ b/python/cuvs_bench/cuvs_bench/config/algos/opensearch_faiss_hnsw.yaml @@ -15,10 +15,12 @@ groups: base: build: m: [16, 32, 48, 64] + number_of_shards: [6] search: ef_search: [10, 20, 40, 60, 80, 120, 200, 400, 600, 800] test: build: m: [16] + number_of_shards: [6] search: ef_search: [10, 20] From 36d49882096f6484e7eb37bb27496aeedaf2d00a Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Wed, 22 Jul 2026 15:51:49 -0500 Subject: [PATCH 13/22] Update bucket prefix Signed-off-by: James Bourbeau --- deploy/README.md | 2 +- deploy/bench/prepare_custom_dataset.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/deploy/README.md b/deploy/README.md index c185e6b742..e7251c7002 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -90,7 +90,7 @@ The `bench` container logs its progress through each phase. When complete you'll ### MIRACL 5M custom dataset Set `DATASET=miracl-en-5m-1024d-fp32` to use the custom dataset from -`s3://opensearch-cuvs-bench/miracl-en-5m-1024d-fp32/miracl-en-5m-1024d-fp32/` +`s3://opensearch-cuvs-bench/miracl-en-5m-1024d-fp32/` instead of a built-in cuvs-bench dataset: ```bash diff --git a/deploy/bench/prepare_custom_dataset.py b/deploy/bench/prepare_custom_dataset.py index 4b13e7be3c..4b57eaf488 100644 --- a/deploy/bench/prepare_custom_dataset.py +++ b/deploy/bench/prepare_custom_dataset.py @@ -11,7 +11,7 @@ DATASET_NAME = "miracl-en-5m-1024d-fp32" S3_BUCKET = "opensearch-cuvs-bench" -S3_PREFIX = "miracl-en-5m-1024d-fp32/miracl-en-5m-1024d-fp32" +S3_PREFIX = "miracl-en-5m-1024d-fp32" def download_if_needed(s3, bucket: str, key: str, destination: Path) -> None: From e98d0b33c406854b1140ab44c6a36322b691f604 Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Wed, 22 Jul 2026 16:11:39 -0500 Subject: [PATCH 14/22] Make number of shards configurable Signed-off-by: James Bourbeau --- deploy/AWS_MULTINODE_SETUP.md | 1 + deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md | 1 + deploy/DEPLOYMENT.md | 1 + deploy/README.md | 1 + deploy/bench/run.py | 5 ++++ deploy/docker-compose.multinode.yml | 1 + deploy/docker-compose.yml | 2 ++ deploy/remote-index-build/run.py | 5 +++- .../cuvs_bench/backends/opensearch.py | 10 ++++++- .../cuvs_bench/tests/test_opensearch.py | 27 +++++++++++++++++++ 10 files changed, 52 insertions(+), 2 deletions(-) diff --git a/deploy/AWS_MULTINODE_SETUP.md b/deploy/AWS_MULTINODE_SETUP.md index c181c04d28..df82c12a8b 100644 --- a/deploy/AWS_MULTINODE_SETUP.md +++ b/deploy/AWS_MULTINODE_SETUP.md @@ -335,6 +335,7 @@ export BENCH_GROUPS=test export K=10 export BATCH_SIZE= export BUILD_BATCH_SIZE= +export NUMBER_OF_SHARDS=1 mkdir -p "${DATASET_PATH}" diff --git a/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md b/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md index d33cb71bcb..88f83fa70f 100644 --- a/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md +++ b/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md @@ -325,6 +325,7 @@ export BENCH_GROUPS=test export K=10 export BATCH_SIZE= export BUILD_BATCH_SIZE= +export NUMBER_OF_SHARDS=1 mkdir -p "${DATASET_PATH}" diff --git a/deploy/DEPLOYMENT.md b/deploy/DEPLOYMENT.md index 4084eb4cbf..d4d4ef1c7c 100644 --- a/deploy/DEPLOYMENT.md +++ b/deploy/DEPLOYMENT.md @@ -162,6 +162,7 @@ docker compose run --rm \ -e AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION:-us-west-2} \ -e REMOTE_BUILD_SIZE_MIN=${REMOTE_BUILD_SIZE_MIN:-} \ -e REMOTE_BUILD_TIMEOUT=${REMOTE_BUILD_TIMEOUT:-1800} \ + -e NUMBER_OF_SHARDS=${NUMBER_OF_SHARDS:-1} \ -v $(pwd)/remote-index-build:/app/remote-index-build \ --no-deps bench \ python remote-index-build/run.py diff --git a/deploy/README.md b/deploy/README.md index e7251c7002..cd9b5ecc59 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -72,6 +72,7 @@ export BENCH_GROUPS=test # test | base (default: test) export K=10 # number of neighbors (default: 10) export BATCH_SIZE= # optional query batch size override export BUILD_BATCH_SIZE= # optional bulk ingest batch size override +export NUMBER_OF_SHARDS=1 # number of primary index shards (default: 1) export REMOTE_BUILD_TIMEOUT=1800 # seconds to wait for remote builds (default: 1800) ``` diff --git a/deploy/bench/run.py b/deploy/bench/run.py index 8f05419865..379a1b2719 100644 --- a/deploy/bench/run.py +++ b/deploy/bench/run.py @@ -44,6 +44,9 @@ BATCH_SIZE = int(BATCH_SIZE) if BATCH_SIZE else None BUILD_BATCH_SIZE = os.environ.get("BUILD_BATCH_SIZE", "").strip() BUILD_BATCH_SIZE = int(BUILD_BATCH_SIZE) if BUILD_BATCH_SIZE else None +NUMBER_OF_SHARDS = int(os.environ.get("NUMBER_OF_SHARDS", "1")) +if NUMBER_OF_SHARDS < 1: + raise ValueError("NUMBER_OF_SHARDS must be at least 1") ALGORITHM = "opensearch_faiss_hnsw" REPO_NAME = os.environ.get("REMOTE_VECTOR_REPOSITORY", "vector-repo").strip() @@ -287,6 +290,7 @@ def main() -> None: " Build batch size : " f"{BUILD_BATCH_SIZE if BUILD_BATCH_SIZE is not None else 'backend auto'}" ) + print(f" Index shards : {NUMBER_OF_SHARDS}") if REMOTE_INDEX_BUILD: register_repository() @@ -302,6 +306,7 @@ def main() -> None: groups=BENCH_GROUPS, host=OPENSEARCH_HOST, port=OPENSEARCH_PORT, + number_of_shards=NUMBER_OF_SHARDS, use_ssl=False, verify_certs=False, ) diff --git a/deploy/docker-compose.multinode.yml b/deploy/docker-compose.multinode.yml index cadd95bd2a..1e832e8917 100644 --- a/deploy/docker-compose.multinode.yml +++ b/deploy/docker-compose.multinode.yml @@ -50,6 +50,7 @@ x-bench-env: &bench-env K: ${K:-10} BATCH_SIZE: ${BATCH_SIZE:-} BUILD_BATCH_SIZE: ${BUILD_BATCH_SIZE:-} + NUMBER_OF_SHARDS: ${NUMBER_OF_SHARDS:-1} services: opensearch: diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml index 6572571968..04da4c87ef 100644 --- a/deploy/docker-compose.yml +++ b/deploy/docker-compose.yml @@ -37,6 +37,7 @@ # K Number of neighbors to search for (default: 10) # BATCH_SIZE Optional cuvs-bench query batch size override # BUILD_BATCH_SIZE Optional OpenSearch bulk ingest batch size override +# NUMBER_OF_SHARDS Number of primary index shards (default: 1) # CUVS_REPOSITORY Repository cloned into the benchmark image # CUVS_BRANCH Repository branch cloned into the benchmark image # @@ -134,6 +135,7 @@ services: K: ${K:-10} BATCH_SIZE: ${BATCH_SIZE:-} BUILD_BATCH_SIZE: ${BUILD_BATCH_SIZE:-} + NUMBER_OF_SHARDS: ${NUMBER_OF_SHARDS:-1} volumes: - ${DATASET_PATH:-/tmp/datasets}:/data/datasets restart: "no" diff --git a/deploy/remote-index-build/run.py b/deploy/remote-index-build/run.py index a7548ded04..9f8a4c2bfc 100644 --- a/deploy/remote-index-build/run.py +++ b/deploy/remote-index-build/run.py @@ -36,6 +36,9 @@ REPO_NAME = "vector-repo" REMOTE_BUILD_SIZE_MIN = os.environ.get("REMOTE_BUILD_SIZE_MIN", "").strip() REMOTE_BUILD_TIMEOUT = int(os.environ.get("REMOTE_BUILD_TIMEOUT", "1800")) +NUMBER_OF_SHARDS = int(os.environ.get("NUMBER_OF_SHARDS", "1")) +if NUMBER_OF_SHARDS < 1: + raise ValueError("NUMBER_OF_SHARDS must be at least 1") session = requests.Session() session.headers.update({"Content-Type": "application/json"}) @@ -94,7 +97,7 @@ def create_index() -> None: index_settings = { "index.knn": True, "index.knn.remote_index_build.enabled": True, - "number_of_shards": 1, + "number_of_shards": NUMBER_OF_SHARDS, "number_of_replicas": 0, } if REMOTE_BUILD_SIZE_MIN: diff --git a/python/cuvs_bench/cuvs_bench/backends/opensearch.py b/python/cuvs_bench/cuvs_bench/backends/opensearch.py index 026fe98fd6..b01331f1ee 100644 --- a/python/cuvs_bench/cuvs_bench/backends/opensearch.py +++ b/python/cuvs_bench/cuvs_bench/backends/opensearch.py @@ -141,6 +141,11 @@ def _build_benchmark_configs( tune_mode = kwargs.get("_tune_mode", False) tune_build_params = kwargs.get("_tune_build_params") tune_search_params = kwargs.get("_tune_search_params") + number_of_shards = kwargs.get("number_of_shards") + if number_of_shards is not None: + number_of_shards = int(number_of_shards) + if number_of_shards < 1: + raise ValueError("number_of_shards must be at least 1") benchmark_configs: List[BenchmarkConfig] = [] @@ -161,7 +166,10 @@ def _build_benchmark_configs( actual_build = build_combos actual_search = search_combos - for build_param in actual_build: + for raw_build_param in actual_build: + build_param = raw_build_param.copy() + if number_of_shards is not None: + build_param["number_of_shards"] = number_of_shards prefix = ( algo_name if group_name == "base" diff --git a/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py b/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py index d16eedf25d..74e83a7e12 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py @@ -165,6 +165,33 @@ def test_load_forwards_remote_build_kwargs(self, config_dir): assert bc["remote_build_timeout"] == 123 assert "remote_build_s3_endpoint" not in bc + def test_load_overrides_number_of_shards(self, config_dir): + loader = OpenSearchConfigLoader(config_path=config_dir) + _, configs = loader.load( + dataset="test-ds", + dataset_path="/data", + groups="test", + number_of_shards=4, + ) + + assert all( + config.indexes[0].build_param["number_of_shards"] == 4 + for config in configs + ) + assert all( + "number_of_shards4" in config.indexes[0].name + for config in configs + ) + + def test_load_rejects_invalid_number_of_shards(self, config_dir): + loader = OpenSearchConfigLoader(config_path=config_dir) + with pytest.raises(ValueError, match="number_of_shards must be at least 1"): + loader.load( + dataset="test-ds", + dataset_path="/data", + number_of_shards=0, + ) + class TestOpenSearchBackend: def test_build_dry_run(self): From 174379bfa1ccf23c423efb1d541769e655ece37d Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Thu, 23 Jul 2026 10:28:50 -0500 Subject: [PATCH 15/22] Configure approx thresh Signed-off-by: James Bourbeau --- deploy/AWS_MULTINODE_SETUP.md | 1 + deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md | 1 + deploy/DEPLOYMENT.md | 2 + deploy/README.md | 7 ++++ deploy/bench/run.py | 5 +++ deploy/docker-compose.multinode.yml | 1 + deploy/docker-compose.yml | 2 + deploy/remote-index-build/run.py | 4 ++ .../cuvs_bench/backends/opensearch.py | 16 +++++++ .../cuvs_bench/tests/test_opensearch.py | 42 +++++++++++++++++++ 10 files changed, 81 insertions(+) diff --git a/deploy/AWS_MULTINODE_SETUP.md b/deploy/AWS_MULTINODE_SETUP.md index df82c12a8b..e7bc69d135 100644 --- a/deploy/AWS_MULTINODE_SETUP.md +++ b/deploy/AWS_MULTINODE_SETUP.md @@ -336,6 +336,7 @@ export K=10 export BATCH_SIZE= export BUILD_BATCH_SIZE= export NUMBER_OF_SHARDS=1 +export APPROXIMATE_THRESHOLD=10000 mkdir -p "${DATASET_PATH}" diff --git a/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md b/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md index 88f83fa70f..870b56e2b3 100644 --- a/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md +++ b/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md @@ -326,6 +326,7 @@ export K=10 export BATCH_SIZE= export BUILD_BATCH_SIZE= export NUMBER_OF_SHARDS=1 +export APPROXIMATE_THRESHOLD=10000 mkdir -p "${DATASET_PATH}" diff --git a/deploy/DEPLOYMENT.md b/deploy/DEPLOYMENT.md index d4d4ef1c7c..423c1e8b1c 100644 --- a/deploy/DEPLOYMENT.md +++ b/deploy/DEPLOYMENT.md @@ -123,6 +123,7 @@ curl -X PUT http://localhost:9200/my-vectors \ "index.knn": true, "index.knn.remote_index_build.enabled": true, "index.knn.remote_index_build.size.min": "1kb", + "index.knn.advanced.approximate_threshold": 10000, "number_of_shards": 1, "number_of_replicas": 1 }, @@ -163,6 +164,7 @@ docker compose run --rm \ -e REMOTE_BUILD_SIZE_MIN=${REMOTE_BUILD_SIZE_MIN:-} \ -e REMOTE_BUILD_TIMEOUT=${REMOTE_BUILD_TIMEOUT:-1800} \ -e NUMBER_OF_SHARDS=${NUMBER_OF_SHARDS:-1} \ + -e APPROXIMATE_THRESHOLD=${APPROXIMATE_THRESHOLD:-10000} \ -v $(pwd)/remote-index-build:/app/remote-index-build \ --no-deps bench \ python remote-index-build/run.py diff --git a/deploy/README.md b/deploy/README.md index cd9b5ecc59..923b04f21b 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -73,9 +73,16 @@ export K=10 # number of neighbors (default: 10) export BATCH_SIZE= # optional query batch size override export BUILD_BATCH_SIZE= # optional bulk ingest batch size override export NUMBER_OF_SHARDS=1 # number of primary index shards (default: 1) +export APPROXIMATE_THRESHOLD=10000 # vectors per segment before ANN build (default: 10000) export REMOTE_BUILD_TIMEOUT=1800 # seconds to wait for remote builds (default: 1800) ``` +`APPROXIMATE_THRESHOLD` sets +`index.knn.advanced.approximate_threshold` on each benchmark index. AWS +recommends `10000` as a starting point for GPU-accelerated indexing so smaller +segments do not build ANN structures prematurely. Set it to `0` to always +build ANN structures or `-1` to disable them. + Start all services: ```bash diff --git a/deploy/bench/run.py b/deploy/bench/run.py index 379a1b2719..3065256710 100644 --- a/deploy/bench/run.py +++ b/deploy/bench/run.py @@ -47,6 +47,9 @@ NUMBER_OF_SHARDS = int(os.environ.get("NUMBER_OF_SHARDS", "1")) if NUMBER_OF_SHARDS < 1: raise ValueError("NUMBER_OF_SHARDS must be at least 1") +APPROXIMATE_THRESHOLD = int(os.environ.get("APPROXIMATE_THRESHOLD", "10000")) +if APPROXIMATE_THRESHOLD < -1: + raise ValueError("APPROXIMATE_THRESHOLD must be -1 or greater") ALGORITHM = "opensearch_faiss_hnsw" REPO_NAME = os.environ.get("REMOTE_VECTOR_REPOSITORY", "vector-repo").strip() @@ -291,6 +294,7 @@ def main() -> None: f"{BUILD_BATCH_SIZE if BUILD_BATCH_SIZE is not None else 'backend auto'}" ) print(f" Index shards : {NUMBER_OF_SHARDS}") + print(f" Approx. threshold : {APPROXIMATE_THRESHOLD}") if REMOTE_INDEX_BUILD: register_repository() @@ -316,6 +320,7 @@ def main() -> None: build_kwargs = dict( common_kwargs, remote_index_build=REMOTE_INDEX_BUILD, + approximate_threshold=APPROXIMATE_THRESHOLD, ) if BUILD_BATCH_SIZE is not None: build_kwargs["build_batch_size"] = BUILD_BATCH_SIZE diff --git a/deploy/docker-compose.multinode.yml b/deploy/docker-compose.multinode.yml index 1e832e8917..62fa512ab3 100644 --- a/deploy/docker-compose.multinode.yml +++ b/deploy/docker-compose.multinode.yml @@ -51,6 +51,7 @@ x-bench-env: &bench-env BATCH_SIZE: ${BATCH_SIZE:-} BUILD_BATCH_SIZE: ${BUILD_BATCH_SIZE:-} NUMBER_OF_SHARDS: ${NUMBER_OF_SHARDS:-1} + APPROXIMATE_THRESHOLD: ${APPROXIMATE_THRESHOLD:-10000} services: opensearch: diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml index 04da4c87ef..a08149445b 100644 --- a/deploy/docker-compose.yml +++ b/deploy/docker-compose.yml @@ -38,6 +38,7 @@ # BATCH_SIZE Optional cuvs-bench query batch size override # BUILD_BATCH_SIZE Optional OpenSearch bulk ingest batch size override # NUMBER_OF_SHARDS Number of primary index shards (default: 1) +# APPROXIMATE_THRESHOLD Vectors per segment before ANN build (default: 10000) # CUVS_REPOSITORY Repository cloned into the benchmark image # CUVS_BRANCH Repository branch cloned into the benchmark image # @@ -136,6 +137,7 @@ services: BATCH_SIZE: ${BATCH_SIZE:-} BUILD_BATCH_SIZE: ${BUILD_BATCH_SIZE:-} NUMBER_OF_SHARDS: ${NUMBER_OF_SHARDS:-1} + APPROXIMATE_THRESHOLD: ${APPROXIMATE_THRESHOLD:-10000} volumes: - ${DATASET_PATH:-/tmp/datasets}:/data/datasets restart: "no" diff --git a/deploy/remote-index-build/run.py b/deploy/remote-index-build/run.py index 9f8a4c2bfc..9dbde482c3 100644 --- a/deploy/remote-index-build/run.py +++ b/deploy/remote-index-build/run.py @@ -39,6 +39,9 @@ NUMBER_OF_SHARDS = int(os.environ.get("NUMBER_OF_SHARDS", "1")) if NUMBER_OF_SHARDS < 1: raise ValueError("NUMBER_OF_SHARDS must be at least 1") +APPROXIMATE_THRESHOLD = int(os.environ.get("APPROXIMATE_THRESHOLD", "10000")) +if APPROXIMATE_THRESHOLD < -1: + raise ValueError("APPROXIMATE_THRESHOLD must be -1 or greater") session = requests.Session() session.headers.update({"Content-Type": "application/json"}) @@ -99,6 +102,7 @@ def create_index() -> None: "index.knn.remote_index_build.enabled": True, "number_of_shards": NUMBER_OF_SHARDS, "number_of_replicas": 0, + "index.knn.advanced.approximate_threshold": APPROXIMATE_THRESHOLD, } if REMOTE_BUILD_SIZE_MIN: index_settings["index.knn.remote_index_build.size.min"] = ( diff --git a/python/cuvs_bench/cuvs_bench/backends/opensearch.py b/python/cuvs_bench/cuvs_bench/backends/opensearch.py index b01331f1ee..bbce4d94a9 100644 --- a/python/cuvs_bench/cuvs_bench/backends/opensearch.py +++ b/python/cuvs_bench/cuvs_bench/backends/opensearch.py @@ -131,6 +131,7 @@ def _build_benchmark_configs( "use_ssl", "verify_certs", "build_batch_size", + "approximate_threshold", # Remote Index Build (OpenSearch 3.0+, faiss engine only) "remote_index_build", "remote_build_size_min", @@ -285,6 +286,8 @@ class OpenSearchBackend(BenchmarkBackend): - ``verify_certs`` – verify SSL certs (default: ``False``) - ``build_batch_size`` – vectors per bulk request. If omitted, choose a batch size with roughly 1 MiB of raw vector data. + - ``approximate_threshold`` – minimum vectors per segment before + building ANN data structures (default: OpenSearch's default). - ``requires_network`` – trigger network pre-flight check (default: ``True``) - ``remote_index_build`` – set ``index.knn.remote_index_build.enabled=true`` on the index at creation time, opting it into the GPU build path (default: ``False``). @@ -356,6 +359,7 @@ def _build_index_mapping( build_param: Dict[str, Any], remote_index_build: bool = False, remote_build_size_min: Optional[str] = None, + approximate_threshold: Optional[int] = None, ) -> Dict[str, Any]: """ Construct the OpenSearch index mapping dict for k-NN. @@ -405,6 +409,15 @@ def _build_index_mapping( "number_of_shards": build_param.get("number_of_shards", 1), "number_of_replicas": build_param.get("number_of_replicas", 0), } + if approximate_threshold is not None: + approximate_threshold = int(approximate_threshold) + if approximate_threshold < -1: + raise ValueError( + "approximate_threshold must be -1 or greater" + ) + index_settings["knn.advanced.approximate_threshold"] = ( + approximate_threshold + ) if remote_index_build: if engine != "faiss": raise ValueError( @@ -710,6 +723,7 @@ def build( build_batch_size = self.config.get("build_batch_size") remote_index_build = bool(self.config.get("remote_index_build", False)) remote_build_size_min = self.config.get("remote_build_size_min") + approximate_threshold = self.config.get("approximate_threshold") if dry_run: print( @@ -761,6 +775,7 @@ def build( build_param, remote_index_build, remote_build_size_min, + approximate_threshold, ) self._client.indices.create(index=index_name, body=mapping) @@ -801,6 +816,7 @@ def build( "engine": engine, "space_type": space_type, "remote_index_build": remote_index_build, + "approximate_threshold": approximate_threshold, }, success=True, ) diff --git a/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py b/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py index 74e83a7e12..802752984d 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py @@ -157,12 +157,14 @@ def test_load_forwards_remote_build_kwargs(self, config_dir): remote_build_size_min="2kb", remote_build_timeout=123, remote_build_s3_endpoint="http://s3:9000", + approximate_threshold=10_000, ) bc = configs[0].backend_config assert bc["remote_index_build"] is True assert bc["remote_build_size_min"] == "2kb" assert bc["remote_build_timeout"] == 123 + assert bc["approximate_threshold"] == 10_000 assert "remote_build_s3_endpoint" not in bc def test_load_overrides_number_of_shards(self, config_dir): @@ -248,6 +250,46 @@ def test_remote_build_size_min_overrides_default(self): settings = mapping["settings"]["index"] assert settings["knn.remote_index_build.size.min"] == "2kb" + def test_approximate_threshold_is_added_to_index_settings(self): + backend = _make_backend() + mapping = backend._build_index_mapping( + dims=4, + engine="faiss", + space_type="l2", + build_param={}, + approximate_threshold=10_000, + ) + + settings = mapping["settings"]["index"] + assert settings["knn.advanced.approximate_threshold"] == 10_000 + + def test_approximate_threshold_accepts_disable_value(self): + backend = _make_backend() + mapping = backend._build_index_mapping( + dims=4, + engine="faiss", + space_type="l2", + build_param={}, + approximate_threshold=-1, + ) + + settings = mapping["settings"]["index"] + assert settings["knn.advanced.approximate_threshold"] == -1 + + def test_approximate_threshold_rejects_values_below_minus_one(self): + backend = _make_backend() + with pytest.raises( + ValueError, + match="approximate_threshold must be -1 or greater", + ): + backend._build_index_mapping( + dims=4, + engine="faiss", + space_type="l2", + build_param={}, + approximate_threshold=-2, + ) + def test_wait_for_remote_build_raises_on_failure_count(self): backend = _make_backend() initial_stats = { From 037e513a10d5b7cd16e3f5a6e9e991610a216d19 Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Tue, 28 Jul 2026 11:54:02 -0500 Subject: [PATCH 16/22] Initial commit Signed-off-by: James Bourbeau --- python/cuvs_bench/cuvs_bench/backends/base.py | 15 +- .../cuvs_bench/backends/cpp_gbench.py | 187 ++++++----- .../cuvs_bench/backends/opensearch.py | 142 ++++---- .../cuvs_bench/orchestrator/orchestrator.py | 129 ++++---- python/cuvs_bench/cuvs_bench/run/__main__.py | 124 ++++--- .../cuvs_bench/cuvs_bench/run/data_export.py | 166 +++++++++- .../cuvs_bench/tests/test_cpp_gbench.py | 6 +- .../cuvs_bench/tests/test_data_export.py | 304 ++++++++++++++++++ .../cuvs_bench/tests/test_opensearch.py | 67 +++- .../cuvs_bench/tests/test_registry.py | 48 +-- 10 files changed, 883 insertions(+), 305 deletions(-) create mode 100644 python/cuvs_bench/cuvs_bench/tests/test_data_export.py diff --git a/python/cuvs_bench/cuvs_bench/backends/base.py b/python/cuvs_bench/cuvs_bench/backends/base.py index bf802d2128..15209decb6 100644 --- a/python/cuvs_bench/cuvs_bench/backends/base.py +++ b/python/cuvs_bench/cuvs_bench/backends/base.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ @@ -274,8 +274,8 @@ class SearchResult: algorithm : str Algorithm name search_params : List[Dict[str, Any]] - List of search parameter combinations used (e.g., [{"nprobe": 1}, {"nprobe": 5}]) - All are batched into one C++ command (matches runners.py behavior) + Search parameter combinations represented by this result. Backends + normally return one result per combination. latency_percentiles : Optional[Dict[str, float]] Latency percentiles in milliseconds (p50, p95, p99) gpu_time_seconds : Optional[float] @@ -421,7 +421,7 @@ def search( force: bool = False, search_threads: Optional[int] = None, dry_run: bool = False, - ) -> SearchResult: + ) -> List[SearchResult]: """ Search for nearest neighbors using the built indexes. @@ -452,8 +452,11 @@ def search( Returns ------- - SearchResult - Search timing, results, and recall metrics + List[SearchResult] + One or more search result objects. Backends normally return one + result per independently measurable search point, but may return + an aggregate result when their native output remains authoritative + (for example, the C++ Google Benchmark backend). Raises ------ diff --git a/python/cuvs_bench/cuvs_bench/backends/cpp_gbench.py b/python/cuvs_bench/cuvs_bench/backends/cpp_gbench.py index a888e957e5..1e2d861456 100644 --- a/python/cuvs_bench/cuvs_bench/backends/cpp_gbench.py +++ b/python/cuvs_bench/cuvs_bench/backends/cpp_gbench.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ @@ -309,7 +309,7 @@ def search( force: bool = False, search_threads: Optional[int] = None, dry_run: bool = False, - ) -> SearchResult: + ) -> List[SearchResult]: """ Search using C++ Google Benchmark executable. @@ -339,40 +339,45 @@ def search( Returns ------- - SearchResult - Search timing, recall, and QPS (aggregated across all indexes) + List[SearchResult] + A single aggregate result containing search timing, recall, and + QPS across all indexes. """ if not indexes: - return SearchResult( - neighbors=np.array([]), - distances=np.array([]), - search_time_ms=0.0, - queries_per_second=0.0, - recall=0.0, - algorithm="", - search_params=[], - metadata={"skipped": True, "reason": "no_indexes"}, - success=True, - ) + return [ + SearchResult( + neighbors=np.array([]), + distances=np.array([]), + search_time_ms=0.0, + queries_per_second=0.0, + recall=0.0, + algorithm="", + search_params=[], + metadata={"skipped": True, "reason": "no_indexes"}, + success=True, + ) + ] first_index = indexes[0] # Pre-flight check (GPU, network, etc.) skip_reason = self._pre_flight_check() if skip_reason: - return SearchResult( - neighbors=np.array([]), - distances=np.array([]), - search_time_ms=0.0, - queries_per_second=0.0, - recall=0.0, - algorithm=first_index.algo, - search_params=first_index.search_params - if first_index.search_params - else [], - metadata={"skipped": True, "reason": skip_reason}, - success=True, - ) + return [ + SearchResult( + neighbors=np.array([]), + distances=np.array([]), + search_time_ms=0.0, + queries_per_second=0.0, + recall=0.0, + algorithm=first_index.algo, + search_params=first_index.search_params + if first_index.search_params + else [], + metadata={"skipped": True, "reason": skip_reason}, + success=True, + ) + ] # Note: runners.py doesn't validate and lets C++ fail. We validate here for # better Python-side error messages. @@ -470,21 +475,23 @@ def search( f"Benchmark command for {self.output_filename[1]}:\n{' '.join(cmd)}\n" ) Path(temp_config_path).unlink(missing_ok=True) - return SearchResult( - neighbors=np.array([]), - distances=np.array([]), - search_time_ms=0.0, - queries_per_second=0.0, - recall=0.0, - algorithm=first_index.algo, - search_params=first_index.search_params, - metadata={ - "dry_run": True, - "num_indexes": len(indexes), - "total_search_configs": total_search_configs, - }, - success=True, - ) + return [ + SearchResult( + neighbors=np.array([]), + distances=np.array([]), + search_time_ms=0.0, + queries_per_second=0.0, + recall=0.0, + algorithm=first_index.algo, + search_params=first_index.search_params, + metadata={ + "dry_run": True, + "num_indexes": len(indexes), + "total_search_configs": total_search_configs, + }, + success=True, + ) + ] # Execute subprocess start_time = time.perf_counter() @@ -544,54 +551,60 @@ def search( # Note: C++ Google Benchmark doesn't return actual neighbors/distances # This is a limitation of the current system - return SearchResult( - neighbors=np.array([]), # Not available from C++ benchmark - distances=np.array([]), # Not available from C++ benchmark - search_time_ms=total_search_time, - queries_per_second=avg_qps, - recall=avg_recall, - algorithm=first_index.algo, - search_params=first_index.search_params, - metadata={ - "num_indexes": len(indexes), - "num_benchmarks": len(benchmarks), - "elapsed_time": elapsed_time, - "latency_us": benchmarks[0].get("Latency") - if benchmarks - else None, - "end_to_end": benchmarks[0].get("end_to_end") - if benchmarks - else None, - "context": gbench_results.get("context", {}), - }, - success=True, - ) + return [ + SearchResult( + neighbors=np.array([]), # Not available from C++ benchmark + distances=np.array([]), # Not available from C++ benchmark + search_time_ms=total_search_time, + queries_per_second=avg_qps, + recall=avg_recall, + algorithm=first_index.algo, + search_params=first_index.search_params, + metadata={ + "num_indexes": len(indexes), + "num_benchmarks": len(benchmarks), + "elapsed_time": elapsed_time, + "latency_us": benchmarks[0].get("Latency") + if benchmarks + else None, + "end_to_end": benchmarks[0].get("end_to_end") + if benchmarks + else None, + "context": gbench_results.get("context", {}), + }, + success=True, + ) + ] except subprocess.CalledProcessError as e: - return SearchResult( - neighbors=np.array([]), - distances=np.array([]), - search_time_ms=time.perf_counter() - start_time, - queries_per_second=0.0, - recall=0.0, - algorithm=first_index.algo, - search_params=first_index.search_params, - success=False, - error_message=f"Search failed: {e.stderr}", - ) + return [ + SearchResult( + neighbors=np.array([]), + distances=np.array([]), + search_time_ms=time.perf_counter() - start_time, + queries_per_second=0.0, + recall=0.0, + algorithm=first_index.algo, + search_params=first_index.search_params, + success=False, + error_message=f"Search failed: {e.stderr}", + ) + ] except Exception as e: - return SearchResult( - neighbors=np.array([]), - distances=np.array([]), - search_time_ms=time.perf_counter() - start_time, - queries_per_second=0.0, - recall=0.0, - algorithm=first_index.algo, - search_params=first_index.search_params, - success=False, - error_message=f"Search error: {str(e)}", - ) + return [ + SearchResult( + neighbors=np.array([]), + distances=np.array([]), + search_time_ms=time.perf_counter() - start_time, + queries_per_second=0.0, + recall=0.0, + algorithm=first_index.algo, + search_params=first_index.search_params, + success=False, + error_message=f"Search error: {str(e)}", + ) + ] finally: # Cleanup temporary config diff --git a/python/cuvs_bench/cuvs_bench/backends/opensearch.py b/python/cuvs_bench/cuvs_bench/backends/opensearch.py index bbce4d94a9..6c78645f1d 100644 --- a/python/cuvs_bench/cuvs_bench/backends/opensearch.py +++ b/python/cuvs_bench/cuvs_bench/backends/opensearch.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -199,6 +199,7 @@ def _build_benchmark_configs( backend_cfg: Dict[str, Any] = { "name": index_label, + "group": group_name, "host": host, "port": port, "index_name": os_index_name, @@ -273,6 +274,7 @@ class OpenSearchBackend(BenchmarkBackend): Required: - ``name`` – index label (e.g. ``"opensearch_faiss_hnsw.m16.ef_construction100"``) + - ``group`` – algorithm configuration group selected from YAML - ``index_name`` – OpenSearch index name (lowercase, no dots) - ``engine`` – ``"faiss"`` or ``"lucene"`` - ``algo`` – algorithm name (e.g. ``"opensearch_faiss_hnsw"``) @@ -564,11 +566,12 @@ def _failed_build_result( self, error_message: str, build_params: Optional[Dict[str, Any]] = None ) -> BuildResult: return BuildResult( - index_path="", + index_path=self.config.get("index_name", ""), build_time_seconds=0.0, index_size_bytes=0, algorithm=self.algo, build_params=build_params or {}, + metadata={"group": self.config["group"]}, success=False, error_message=error_message, ) @@ -587,6 +590,10 @@ def _failed_search_result( recall=0.0, algorithm=self.algo, search_params=search_params or [], + metadata={ + "group": self.config["group"], + "index_name": self.config.get("index_name", ""), + }, success=False, error_message=error_message, ) @@ -737,6 +744,7 @@ def build( index_size_bytes=0, algorithm=self.algo, build_params=build_param, + metadata={"group": self.config["group"]}, success=True, ) @@ -751,6 +759,10 @@ def build( index_size_bytes=0, algorithm=self.algo, build_params=build_param, + metadata={ + "group": self.config["group"], + "skipped": True, + }, success=True, ) @@ -813,6 +825,7 @@ def build( algorithm=self.algo, build_params=build_param, metadata={ + "group": self.config["group"], "engine": engine, "space_type": space_type, "remote_index_build": remote_index_build, @@ -831,21 +844,14 @@ def search( force: bool = False, search_threads: Optional[int] = None, dry_run: bool = False, - ) -> SearchResult: + ) -> List[SearchResult]: """ Search the OpenSearch k-NN index for nearest neighbors. Iterates over every search-parameter combination defined in the index config, updating the index-level ``ef_search`` setting between runs. - Metrics (QPS, latency) are collected per parameter set and stored in - ``SearchResult.metadata["per_search_param_results"]``. - - The *neighbors* and *distances* arrays in the returned result reflect - the **last** search-parameter combination (highest ef_search by - convention), while *queries_per_second* is the average across all - parameter combinations. This backend returns ``recall=0.0``; the - shared orchestrator path computes recall from the returned neighbors - and dataset ground truth. + Returns one result per parameter set so the orchestrator can compute + recall for each set independently. Parameters ---------- @@ -871,16 +877,18 @@ def search( Returns ------- - SearchResult + List[SearchResult] """ skip = self._pre_flight_check() if skip: - return self._failed_search_result( - k, f"pre-flight check failed: {skip}" - ) + return [ + self._failed_search_result( + k, f"pre-flight check failed: {skip}" + ) + ] if not indexes: - return self._failed_search_result(k, "No indexes provided") + return [self._failed_search_result(k, "No indexes provided")] index_cfg = indexes[0] index_name = self._resolve_index_name(index_cfg) @@ -894,35 +902,43 @@ def search( f"(k={k}, batch_size={batch_size})" ) - return SearchResult( - neighbors=np.zeros((0, k), dtype=np.int64), - distances=np.zeros((0, k), dtype=np.float32), - search_time_ms=0.0, - queries_per_second=0.0, - recall=0.0, - algorithm=self.algo, - search_params=search_params_list, - success=True, - ) + return [ + SearchResult( + neighbors=np.zeros((0, k), dtype=np.int64), + distances=np.zeros((0, k), dtype=np.float32), + search_time_ms=0.0, + queries_per_second=0.0, + recall=0.0, + algorithm=self.algo, + search_params=[search_params], + metadata={ + "group": self.config["group"], + "index_name": index_name, + "dry_run": True, + }, + success=True, + ) + for search_params in search_params_list + ] # Dataset handles lazy loading from query files when needed. query_vectors = dataset.query_vectors if query_vectors.size == 0: - return self._failed_search_result( - k, - "No query vectors available. Provide dataset.query_vectors " - "or a valid dataset.query_file path.", - search_params=search_params_list, - ) + return [ + self._failed_search_result( + k, + "No query vectors available. Provide " + "dataset.query_vectors or a valid dataset.query_file path.", + search_params=search_params_list, + ) + ] n_queries = query_vectors.shape[0] n_batches = (n_queries + batch_size - 1) // batch_size # Run search for each search-parameter combination - per_param_results: List[Dict[str, Any]] = [] - last_neighbors = np.full((n_queries, k), -1, dtype=np.int64) - last_distances = np.zeros((n_queries, k), dtype=np.float32) + results: List[SearchResult] = [] for sp in search_params_list: ef_search = sp.get("ef_search", 100) @@ -981,39 +997,25 @@ def search( elapsed = time.perf_counter() - t0 qps = n_queries / elapsed if elapsed > 0 else 0.0 - per_param_results.append( - { - "search_params": sp, - "search_time_ms": elapsed * 1000.0, - "queries_per_second": qps, - "batch_size": batch_size, - "num_batches": n_batches, - } + results.append( + SearchResult( + neighbors=neighbors, + distances=distances, + search_time_ms=elapsed * 1000.0, + queries_per_second=qps, + recall=0.0, + algorithm=self.algo, + search_params=[sp], + metadata={ + "group": self.config["group"], + "index_name": index_name, + "engine": engine, + "batch_size": batch_size, + "num_batches": n_batches, + "latency_seconds": elapsed / n_batches, + }, + success=True, + ) ) - last_neighbors = neighbors - last_distances = distances - # Aggregate across all search-param combinations - avg_qps = float( - np.mean([r["queries_per_second"] for r in per_param_results]) - ) - total_search_time_ms = float( - sum(r["search_time_ms"] for r in per_param_results) - ) - - return SearchResult( - neighbors=last_neighbors, - distances=last_distances, - search_time_ms=total_search_time_ms, - queries_per_second=avg_qps, - recall=0.0, - algorithm=self.algo, - search_params=search_params_list, - metadata={ - "engine": engine, - "batch_size": batch_size, - "num_batches": n_batches, - "per_search_param_results": per_param_results, - }, - success=True, - ) + return results diff --git a/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py b/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py index 580291c2f0..0a50f1d8c9 100644 --- a/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py +++ b/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -143,9 +143,12 @@ def run_benchmark( Returns ------- List[Union[BuildResult, SearchResult]] - List of result objects, one per benchmark run: - - sweep mode: One result per IndexConfig (Cartesian product of params) - - tune mode: One result per Optuna trial (n_trials total) + All measurements produced by the benchmark: + - sweep mode: build results plus the search results returned for + each benchmark configuration. + - tune mode: build and search results produced by each trial. A + successful build-and-search trial normally contributes two + objects. Each SearchResult contains: recall, search_time_ms, queries_per_second, success, metadata, etc. @@ -246,7 +249,7 @@ def _run_sweep( # Pass ALL indexes at once - ONE C++ command searches all # Each index has its own search_params list # Total benchmarks = sum(len(idx.search_params) for idx in indexes) - search_result = backend.search( + search_results = backend.search( dataset=bench_dataset, indexes=config.indexes, k=count, @@ -257,28 +260,17 @@ def _run_sweep( dry_run=dry_run, ) - # Compute recall for backends that return actual neighbors. - # The C++ backend computes recall in the subprocess and returns - # empty neighbors, so this is skipped for it. - # Empty neighbors or nonzero recall indicate that the backend - # already handled recall itself. - if ( - search_result.success - and search_result.neighbors.size > 0 - and search_result.recall == 0.0 - ): - gt = bench_dataset.groundtruth_neighbors - if gt is not None: - search_result.recall = compute_recall( - search_result.neighbors, gt, count - ) - - results.append(search_result) - - if not search_result.success: - print( - f"Search failed for {config.index_name}: {search_result.error_message}" + for search_result in search_results: + self._finalize_search_result( + search_result, bench_dataset, count ) + results.append(search_result) + + if not search_result.success: + print( + f"Search failed for {config.index_name}: " + f"{search_result.error_message}" + ) finally: backend.cleanup() @@ -430,7 +422,7 @@ def objective(trial) -> float: # Run single trial with these specific parameters # First trial (trial.number=0) overwrites, subsequent trials append - result = self._run_trial( + trial_results = self._run_trial( algorithm=algorithm, build_params=build_params, search_params=search_params_dict, @@ -446,13 +438,20 @@ def objective(trial) -> float: **loader_kwargs, ) - # Store result for pareto plot - all_results.append(result) + # Retain every measurement for export. The last result is the + # search result used as the Optuna objective on successful trials. + all_results.extend(trial_results) + result = trial_results[-1] # Check if trial failed if not result.success: raise optuna.TrialPruned() + if not isinstance(result, SearchResult): + raise RuntimeError( + "Successful tune trial did not produce a search result" + ) + # Build metrics dict from SearchResult attributes # No fallbacks - if metrics are missing, let it fail loudly so we can fix the root cause metrics = { @@ -525,7 +524,7 @@ def _run_trial( search_threads: Optional[int], append_results: bool = False, **loader_kwargs, - ) -> Union[BuildResult, SearchResult]: + ) -> List[Union[BuildResult, SearchResult]]: """ Run a single benchmark trial with specific parameters. @@ -558,12 +557,19 @@ def _run_trial( # Should have exactly one config for single trial if not benchmark_configs: - return SearchResult( - success=False, - error_message="No config generated for trial", - metrics={}, - search_params=[], - ) + return [ + SearchResult( + neighbors=np.empty((0, count), dtype=np.int64), + distances=np.empty((0, count), dtype=np.float32), + search_time_ms=0.0, + queries_per_second=0.0, + recall=0.0, + algorithm=algorithm, + search_params=[], + success=False, + error_message="No config generated for trial", + ) + ] config = benchmark_configs[0] # Pass append_results via config (backend-specific, not in base class) @@ -575,20 +581,21 @@ def _run_trial( try: backend.initialize() - result = None + trial_results: List[Union[BuildResult, SearchResult]] = [] if build: - result = backend.build( + build_result = backend.build( dataset=bench_dataset, indexes=config.indexes, force=force, dry_run=dry_run, ) - if not result.success: - return result + trial_results.append(build_result) + if not build_result.success: + return trial_results if search: - result = backend.search( + search_results = backend.search( dataset=bench_dataset, indexes=config.indexes, k=count, @@ -599,24 +606,36 @@ def _run_trial( dry_run=dry_run, ) - # Compute recall for backends that return actual neighbors. - # Empty neighbors or nonzero recall indicate that the backend - # already handled recall itself. - if ( - result.success - and result.neighbors.size > 0 - and result.recall == 0.0 - ): - gt = bench_dataset.groundtruth_neighbors - if gt is not None: - result.recall = compute_recall( - result.neighbors, gt, count - ) + if len(search_results) != 1: + raise RuntimeError( + "Tune mode expected one search-parameter result" + ) + search_result = search_results[0] + self._finalize_search_result( + search_result, bench_dataset, count + ) + trial_results.append(search_result) - return result + return trial_results finally: backend.cleanup() + @staticmethod + def _finalize_search_result( + result: SearchResult, dataset: Dataset, k: int + ) -> None: + """Compute recall for backends that return neighbor arrays.""" + if ( + result.success + and result.neighbors.size > 0 + and result.recall == 0.0 + ): + groundtruth = dataset.groundtruth_neighbors + if groundtruth is not None: + result.recall = compute_recall( + result.neighbors, groundtruth, k + ) + def _create_dataset(self, dataset_config: DatasetConfig) -> Dataset: """ Create a Dataset object from DatasetConfig. diff --git a/python/cuvs_bench/cuvs_bench/run/__main__.py b/python/cuvs_bench/cuvs_bench/run/__main__.py index 6950ff7202..225a3d3f9c 100644 --- a/python/cuvs_bench/cuvs_bench/run/__main__.py +++ b/python/cuvs_bench/cuvs_bench/run/__main__.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -11,7 +11,11 @@ import click import yaml -from .data_export import convert_json_to_csv_build, convert_json_to_csv_search +from .data_export import ( + convert_json_to_csv_build, + convert_json_to_csv_search, + write_results_to_csv, +) from ..orchestrator import BenchmarkOrchestrator @@ -144,12 +148,8 @@ @click.option( "--data-export", is_flag=True, - help="By default, the intermediate JSON outputs produced by " - "cuvs_bench.run to more easily readable CSV files is done " - "automatically, which are needed to build charts made by " - "cuvs_bench.plot. But if some of the benchmark runs failed or " - "were interrupted, use this option to convert those intermediate " - "files manually.", + help="Deprecated: convert existing C++ benchmark JSON files to CSV. " + "Benchmark runs now export CSV automatically.", ) @click.option( "--mode", @@ -243,7 +243,7 @@ def main( dry_run : bool Whether to perform a dry run without actual execution. data_export : bool - Whether to export intermediate JSON results to CSV. + Deprecated option for converting existing C++ JSON results to CSV. mode : str Benchmark mode: 'sweep' (exhaustive) or 'tune' (Optuna-based). constraints : Optional[str] @@ -257,52 +257,70 @@ def main( and any backend-specific connection parameters (host, port, etc.). """ - if not data_export: - # Determine backend type and extra kwargs from --backend-config - backend_type = "cpp_gbench" - backend_kwargs = {} - if backend_config: - with open(backend_config, "r") as f: - cfg = yaml.safe_load(f) - if not isinstance(cfg, dict): - raise ValueError( - f"--backend-config must parse to a mapping, " - f"got {type(cfg).__name__}" - ) - if "backend" not in cfg: - raise ValueError( - "--backend-config must include a 'backend' field" - ) - backend_type = cfg.pop("backend") - backend_kwargs = cfg - - orchestrator = BenchmarkOrchestrator(backend_type=backend_type) - orchestrator.run_benchmark( - mode=mode, - constraints=json.loads(constraints) if constraints else None, - n_trials=n_trials, - dataset=dataset, - dataset_path=dataset_path, - build=build, - search=search, - force=force, - dry_run=dry_run, - count=count, - batch_size=batch_size, - search_mode=search_mode, - search_threads=search_threads, - dataset_configuration=dataset_configuration, - algorithm_configuration=configuration, - algorithms=algorithms, - groups=groups, - algo_groups=algo_groups, - subset_size=subset_size, - executable_dir=executable_dir, - **backend_kwargs, + if data_export: + click.echo( + "Warning: --data-export is deprecated because benchmark runs now " + "export CSV automatically. Converting existing C++ JSON results.", + err=True, ) + convert_json_to_csv_build(dataset, dataset_path) + convert_json_to_csv_search(dataset, dataset_path) + return + + # The CLI historically runs both phases when neither flag is specified. + if not build and not search: + build = search = True + + backend_type = "cpp_gbench" + backend_kwargs = {} + if backend_config: + with open(backend_config, "r") as f: + cfg = yaml.safe_load(f) + if not isinstance(cfg, dict): + raise ValueError( + f"--backend-config must parse to a mapping, " + f"got {type(cfg).__name__}" + ) + if "backend" not in cfg: + raise ValueError("--backend-config must include a 'backend' field") + backend_type = cfg.pop("backend") + backend_kwargs = cfg + + orchestrator = BenchmarkOrchestrator(backend_type=backend_type) + results = orchestrator.run_benchmark( + mode=mode, + constraints=json.loads(constraints) if constraints else None, + n_trials=n_trials, + dataset=dataset, + dataset_path=dataset_path, + build=build, + search=search, + force=force, + dry_run=dry_run, + count=count, + batch_size=batch_size, + search_mode=search_mode, + search_threads=search_threads, + dataset_configuration=dataset_configuration, + algorithm_configuration=configuration, + algorithms=algorithms, + groups=groups, + algo_groups=algo_groups, + subset_size=subset_size, + executable_dir=executable_dir, + **backend_kwargs, + ) + + if dry_run: + return - convert_json_to_csv_build(dataset, dataset_path) - convert_json_to_csv_search(dataset, dataset_path) + if backend_type == "cpp_gbench": + if build: + convert_json_to_csv_build(dataset, dataset_path) + if search: + convert_json_to_csv_search(dataset, dataset_path) + else: + write_results_to_csv(results, dataset, dataset_path, count, batch_size) if __name__ == "__main__": diff --git a/python/cuvs_bench/cuvs_bench/run/data_export.py b/python/cuvs_bench/cuvs_bench/run/data_export.py index 707677a083..50afa7a5ea 100644 --- a/python/cuvs_bench/cuvs_bench/run/data_export.py +++ b/python/cuvs_bench/cuvs_bench/run/data_export.py @@ -1,14 +1,17 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # import json import os import traceback +from collections import defaultdict import pandas as pd +from ..backends.base import BuildResult, SearchResult + skip_build_cols = set( [ "algo_name", @@ -50,6 +53,165 @@ } +def write_results_to_csv(results, dataset, dataset_path, count, batch_size): + """Write Python-backend results using the existing plotting CSV schema.""" + grouped = defaultdict(list) + for result in results: + group = result.metadata.get("group") + index_name = result.metadata.get("index_name") + if group is None or ( + isinstance(result, SearchResult) and index_name is None + ): + continue + method = "build" if isinstance(result, BuildResult) else "search" + grouped[(method, result.algorithm, group)].append(result) + + for (method, algorithm, group), group_results in grouped.items(): + if method == "build": + _write_build_results( + group_results, algorithm, group, dataset, dataset_path + ) + else: + _write_search_results( + group_results, + algorithm, + group, + dataset, + dataset_path, + count, + batch_size, + ) + + +def _write_build_results(results, algorithm, group, dataset, dataset_path): + output_dir = os.path.join(dataset_path, dataset, "result", "build") + os.makedirs(output_dir, exist_ok=True) + algo_name = algorithm if group == "base" else f"{algorithm}_{group}" + + rows = [] + for result in results: + if not result.success or result.metadata.get("skipped"): + continue + metadata = _scalar_metadata(result.metadata) + rows.append( + { + **result.build_params, + **metadata, + "algo_name": algo_name, + "index_name": result.index_path, + "time": result.build_time_seconds, + } + ) + + columns = ["algo_name", "index_name", "time"] + dataframe = pd.DataFrame(rows) + build_file = os.path.join(output_dir, f"{algorithm},{group}.csv") + + complete_run = all( + result.success and not result.metadata.get("skipped") + for result in results + ) + if not complete_run and os.path.exists(build_file): + dataframe = pd.concat( + [pd.read_csv(build_file), dataframe], + ignore_index=True, + sort=False, + ) + + if dataframe.empty: + # Do not replace an existing measurement with a skipped or failed + # build, and do not create an empty result file. + return + + dataframe = dataframe.drop_duplicates(subset=["index_name"], keep="last") + dataframe = dataframe[ + columns + [name for name in dataframe if name not in columns] + ] + dataframe.to_csv(build_file, index=False) + + +def _write_search_results( + results, algorithm, group, dataset, dataset_path, count, batch_size +): + output_dir = os.path.join(dataset_path, dataset, "result", "search") + os.makedirs(output_dir, exist_ok=True) + algo_name = algorithm if group == "base" else f"{algorithm}_{group}" + + rows = [] + for result in results: + if not result.success: + continue + metadata = _scalar_metadata(result.metadata) + search_params = ( + result.search_params[0] if len(result.search_params) == 1 else {} + ) + rows.append( + { + **search_params, + **metadata, + "algo_name": algo_name, + "index_name": result.metadata["index_name"], + "recall": result.recall, + "throughput": result.queries_per_second, + "latency": result.metadata.get( + "latency_seconds", result.search_time_ms / 1000.0 + ), + } + ) + + columns = [ + "algo_name", + "index_name", + "recall", + "throughput", + "latency", + ] + dataframe = pd.DataFrame(rows) + if dataframe.empty: + dataframe = pd.DataFrame(columns=columns) + else: + dataframe = dataframe[ + columns + [name for name in dataframe if name not in columns] + ] + + build_file = os.path.join( + dataset_path, + dataset, + "result", + "build", + f"{algorithm},{group}.csv", + ) + if os.path.exists(build_file): + build = pd.read_csv(build_file).drop_duplicates( + subset=["index_name"], keep="last" + ) + if "time" in build: + dataframe = dataframe.merge( + build[["index_name", "time"]].rename( + columns={"time": "build time"} + ), + on="index_name", + how="left", + ) + + stem = f"{algorithm},{group},k{count},bs{batch_size}" + raw_file = os.path.join(output_dir, f"{stem},raw.csv") + dataframe.to_csv(raw_file, index=False) + frontier_file = os.path.join(output_dir, f"{stem}.json") + write_frontier(frontier_file, dataframe, "throughput") + write_frontier(frontier_file, dataframe, "latency") + + +def _scalar_metadata(metadata): + reserved = {"group", "index_name", "latency_seconds"} + return { + key: value + for key, value in metadata.items() + if key not in reserved + and isinstance(value, (str, int, float, bool, type(None))) + } + + def read_json_files(dataset, dataset_path, method): """ Yield file paths, algo names, and loaded JSON data as pandas DataFrames. @@ -70,6 +232,8 @@ def read_json_files(dataset, dataset_path, method): DataFrame of JSON content. """ dir_path = os.path.join(dataset_path, dataset, "result", method) + if not os.path.isdir(dir_path): + return for file in os.listdir(dir_path): if file.endswith(".json"): file_path = os.path.join(dir_path, file) diff --git a/python/cuvs_bench/cuvs_bench/tests/test_cpp_gbench.py b/python/cuvs_bench/cuvs_bench/tests/test_cpp_gbench.py index a25b4fbe04..02711c46c0 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_cpp_gbench.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_cpp_gbench.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ @@ -243,7 +243,7 @@ def test_search_with_no_indexes(self): # Search with empty indexes list result = backend.search( dataset=dataset, indexes=[], k=10, batch_size=1000 - ) + )[0] assert result.success is True assert result.metadata.get("skipped") is True @@ -351,7 +351,7 @@ def test_search_dry_run(self): k=10, batch_size=1000, dry_run=True, - ) + )[0] assert result.success is True assert result.metadata.get("dry_run") is True diff --git a/python/cuvs_bench/cuvs_bench/tests/test_data_export.py b/python/cuvs_bench/cuvs_bench/tests/test_data_export.py new file mode 100644 index 0000000000..2595d32095 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_data_export.py @@ -0,0 +1,304 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +import numpy as np +import pandas as pd +from click.testing import CliRunner + +from cuvs_bench.backends.base import BuildResult, SearchResult +from cuvs_bench.orchestrator.config_loaders import ( + BenchmarkConfig, + DatasetConfig, + IndexConfig, +) +from cuvs_bench.orchestrator.orchestrator import BenchmarkOrchestrator +from cuvs_bench.plot.__main__ import load_all_results +from cuvs_bench.run.__main__ import main +from cuvs_bench.run.data_export import write_results_to_csv + + +def test_python_backend_csv_is_plot_compatible(tmp_path): + dataset = "test-dataset" + algorithm = "opensearch_faiss_hnsw" + index_name = "test-index" + results = [ + BuildResult( + index_path=index_name, + build_time_seconds=1.5, + index_size_bytes=1024, + algorithm=algorithm, + build_params={"m": 16}, + metadata={"group": "base"}, + ), + SearchResult( + neighbors=np.empty((0, 2), dtype=np.int64), + distances=np.empty((0, 2), dtype=np.float32), + search_time_ms=20.0, + queries_per_second=100.0, + recall=0.5, + algorithm=algorithm, + search_params=[{"ef_search": 50}], + metadata={ + "group": "base", + "index_name": index_name, + "latency_seconds": 0.01, + }, + ), + SearchResult( + neighbors=np.empty((0, 2), dtype=np.int64), + distances=np.empty((0, 2), dtype=np.float32), + search_time_ms=10.0, + queries_per_second=200.0, + recall=1.0, + algorithm=algorithm, + search_params=[{"ef_search": 100}], + metadata={ + "group": "base", + "index_name": index_name, + "latency_seconds": 0.005, + }, + ), + ] + + write_results_to_csv( + results, dataset, str(tmp_path), count=2, batch_size=2 + ) + + result_path = tmp_path / dataset / "result" + raw_file = result_path / "search" / f"{algorithm},base,k2,bs2,raw.csv" + raw = pd.read_csv(raw_file) + assert raw.columns[:5].tolist() == [ + "algo_name", + "index_name", + "recall", + "throughput", + "latency", + ] + assert raw["ef_search"].tolist() == [50, 100] + assert raw["build time"].tolist() == [1.5, 1.5] + assert not list(result_path.rglob("*.json")) + + write_results_to_csv( + [ + BuildResult( + index_path=index_name, + build_time_seconds=0.0, + index_size_bytes=0, + algorithm=algorithm, + build_params={"m": 16}, + metadata={"group": "base", "skipped": True}, + ) + ], + dataset, + str(tmp_path), + count=2, + batch_size=2, + ) + build = pd.read_csv(result_path / "build" / f"{algorithm},base.csv") + assert build["time"].tolist() == [1.5] + + for mode in ("throughput", "latency"): + plotted = load_all_results( + str(result_path.parent), + algorithms=[algorithm], + groups=["base"], + algo_groups=[], + k=2, + batch_size=2, + method="search", + index_key="algo", + raw=False, + mode=mode, + time_unit="s", + ) + assert algorithm in plotted + assert plotted[algorithm] + + +def test_run_command_always_writes_python_backend_csv(tmp_path, monkeypatch): + algorithm = "opensearch_faiss_hnsw" + result = SearchResult( + neighbors=np.empty((0, 2), dtype=np.int64), + distances=np.empty((0, 2), dtype=np.float32), + search_time_ms=10.0, + queries_per_second=200.0, + recall=1.0, + algorithm=algorithm, + search_params=[{"ef_search": 100}], + metadata={ + "group": "base", + "index_name": "test-index", + "latency_seconds": 0.01, + }, + ) + + class FakeOrchestrator: + def __init__(self, backend_type): + assert backend_type == "opensearch" + + def run_benchmark(self, **kwargs): + return [result] + + monkeypatch.setattr( + "cuvs_bench.run.__main__.BenchmarkOrchestrator", FakeOrchestrator + ) + backend_config = tmp_path / "backend.yaml" + backend_config.write_text("backend: opensearch\n") + + cli_result = CliRunner().invoke( + main, + [ + "--dataset", + "test-dataset", + "--dataset-path", + str(tmp_path), + "--algorithms", + algorithm, + "--groups", + "base", + "--count", + "2", + "--batch-size", + "2", + "--search-mode", + "latency", + "--search", + "--backend-config", + str(backend_config), + ], + ) + + assert cli_result.exit_code == 0, cli_result.output + assert ( + tmp_path + / "test-dataset" + / "result" + / "search" + / f"{algorithm},base,k2,bs2,raw.csv" + ).exists() + help_output = CliRunner().invoke(main, ["--help"]).output + assert "--data-export" in help_output + assert "Deprecated" in help_output + + +def test_data_export_option_is_deprecated(tmp_path, monkeypatch): + converted = [] + monkeypatch.setattr( + "cuvs_bench.run.__main__.convert_json_to_csv_build", + lambda dataset, dataset_path: converted.append("build"), + ) + monkeypatch.setattr( + "cuvs_bench.run.__main__.convert_json_to_csv_search", + lambda dataset, dataset_path: converted.append("search"), + ) + + cli_result = CliRunner().invoke( + main, + [ + "--dataset", + "test-dataset", + "--dataset-path", + str(tmp_path), + "--algorithms", + "cuvs_cagra", + "--groups", + "base", + "--count", + "2", + "--batch-size", + "2", + "--search-mode", + "latency", + "--data-export", + ], + ) + + assert cli_result.exit_code == 0, cli_result.output + assert "--data-export is deprecated" in cli_result.output + assert converted == ["build", "search"] + + +def test_tune_trial_retains_build_and_search_results(): + algorithm = "opensearch_faiss_hnsw" + index = IndexConfig( + name="test-index", + algo=algorithm, + build_param={"m": 16}, + search_params=[{"ef_search": 100}], + file="test-index", + ) + dataset = DatasetConfig(name="test-dataset") + + class FakeLoader: + def load(self, **kwargs): + return dataset, [ + BenchmarkConfig( + indexes=[index], + backend_config={"name": index.name, "group": "base"}, + ) + ] + + class FakeBackend: + def __init__(self, config): + pass + + def initialize(self): + pass + + def cleanup(self): + pass + + def build(self, dataset, indexes, force, dry_run): + return BuildResult( + index_path=index.name, + build_time_seconds=1.5, + index_size_bytes=1024, + algorithm=algorithm, + build_params=index.build_param, + metadata={"group": "base"}, + ) + + def search(self, dataset, indexes, k, **kwargs): + return [ + SearchResult( + neighbors=np.array([[0, 1]], dtype=np.int64), + distances=np.zeros((1, k), dtype=np.float32), + search_time_ms=10.0, + queries_per_second=100.0, + recall=0.0, + algorithm=algorithm, + search_params=index.search_params, + metadata={ + "group": "base", + "index_name": index.name, + }, + ) + ] + + orchestrator = object.__new__(BenchmarkOrchestrator) + orchestrator.config_loader = FakeLoader() + orchestrator.backend_class = FakeBackend + orchestrator._create_dataset = lambda config: type( + "Dataset", + (), + {"groundtruth_neighbors": np.array([[0, 1]], dtype=np.int64)}, + )() + + results = orchestrator._run_trial( + algorithm=algorithm, + build_params=index.build_param, + search_params=index.search_params[0], + build=True, + search=True, + force=False, + dry_run=False, + count=2, + batch_size=1, + search_mode="latency", + search_threads=None, + ) + + assert [type(result) for result in results] == [BuildResult, SearchResult] + assert results[1].recall == 1.0 diff --git a/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py b/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py index 802752984d..884e80bb15 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # """ @@ -21,12 +21,14 @@ OpenSearchConfigLoader, ) from cuvs_bench.orchestrator.config_loaders import IndexConfig +from cuvs_bench.orchestrator.orchestrator import BenchmarkOrchestrator def _make_backend(config_overrides: dict = None) -> OpenSearchBackend: """Backend with no network requirement so pre-flight passes without a server.""" config = { "name": "test_index", + "group": "base", "index_name": "test_index", "engine": "faiss", "algo": "opensearch_faiss_hnsw", @@ -90,6 +92,7 @@ def live_backend(opensearch_url): backend = OpenSearchBackend( { "name": index_name, + "group": "base", "index_name": index_name, "engine": "faiss", "algo": "opensearch_faiss_hnsw", @@ -146,6 +149,7 @@ def test_load_produces_correct_configs(self, config_dir): assert len(benchmark_configs) == 4 bc = benchmark_configs[0] assert bc.backend_config["engine"] == "faiss" + assert bc.backend_config["group"] == "test" assert len(bc.indexes[0].search_params) == 2 # ef_search: [50, 100] def test_load_forwards_remote_build_kwargs(self, config_dir): @@ -205,11 +209,56 @@ def test_build_dry_run(self): assert result.index_path == backend.config["index_name"] def test_search_dry_run(self): - result = _make_backend().search( + results = _make_backend().search( _make_dataset(), [_make_index_cfg()], k=3, dry_run=True ) - assert result.success - assert len(result.search_params) == 2 + assert len(results) == 2 + assert all(result.success for result in results) + assert [result.search_params for result in results] == [ + [{"ef_search": 50}], + [{"ef_search": 100}], + ] + + def test_recall_is_computed_for_each_search_parameter(self): + class FakeIndices: + def __init__(self): + self.ef_search = None + + def put_settings(self, index, body): + self.ef_search = body["index.knn.algo_param.ef_search"] + + class FakeClient: + def __init__(self): + self.indices = FakeIndices() + + def msearch(self, index, body): + ids = [2, 3] if self.indices.ef_search == 50 else [0, 1] + response = { + "hits": { + "hits": [ + {"_id": str(neighbor), "_score": 1.0} + for neighbor in ids + ] + } + } + return {"responses": [response for _ in body[::2]]} + + dataset = Dataset( + name="test", + query_vectors=np.zeros((2, 4), dtype=np.float32), + groundtruth_neighbors=np.array([[0, 1], [0, 1]]), + ) + backend = _make_backend() + backend._OpenSearchBackend__client = FakeClient() + + results = backend.search( + dataset, [_make_index_cfg()], k=2, batch_size=2 + ) + for result in results: + BenchmarkOrchestrator._finalize_search_result(result, dataset, 2) + + assert [result.recall for result in results] == [0.0, 1.0] + assert all(result.neighbors.shape == (2, 2) for result in results) def test_remote_build_requires_faiss_engine(self): backend = _make_backend({"engine": "lucene"}) @@ -395,7 +444,7 @@ def test_search_fails_without_query_vectors(self): training_vectors=np.empty((0, 4), dtype=np.float32), query_vectors=np.empty((0, 4), dtype=np.float32), ) - result = _make_backend().search(dataset, [_make_index_cfg()], k=3) + result = _make_backend().search(dataset, [_make_index_cfg()], k=3)[0] assert not result.success assert "No query vectors" in result.error_message @@ -541,6 +590,7 @@ def live_remote_build_backend(opensearch_url, remote_build_env): backend = OpenSearchBackend( { "name": index_name, + "group": "base", "index_name": index_name, "engine": "faiss", "algo": "opensearch_faiss_hnsw", @@ -572,12 +622,11 @@ def test_build_and_search(self, live_backend): assert build_result.build_time_seconds > 0 assert build_result.index_size_bytes > 0 - search_result = live_backend.search(dataset, [idx], k=k) + search_result = live_backend.search(dataset, [idx], k=k)[0] assert search_result.success assert search_result.recall == 0.0 assert search_result.queries_per_second > 0 assert search_result.neighbors.shape == (10, k) - assert len(search_result.metadata["per_search_param_results"]) == 1 @pytest.mark.opensearch @@ -594,7 +643,9 @@ def test_remote_build_and_search(self, live_remote_build_backend): assert build_result.build_time_seconds > 0 assert build_result.metadata["remote_index_build"] is True - search_result = live_remote_build_backend.search(dataset, [idx], k=k) + search_result = live_remote_build_backend.search(dataset, [idx], k=k)[ + 0 + ] assert search_result.success assert search_result.recall == 0.0 assert search_result.queries_per_second > 0 diff --git a/python/cuvs_bench/cuvs_bench/tests/test_registry.py b/python/cuvs_bench/cuvs_bench/tests/test_registry.py index 960ea9cb25..eadf804db3 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_registry.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_registry.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ @@ -59,16 +59,18 @@ def search( distances = np.random.rand(n_queries, k) first = indexes[0] - return SearchResult( - neighbors=neighbors, - distances=distances, - search_time_ms=0.1, - queries_per_second=n_queries / 0.1, - recall=0.95, - algorithm=self.algo, - search_params=first.search_params, - success=True, - ) + return [ + SearchResult( + neighbors=neighbors, + distances=distances, + search_time_ms=0.1, + queries_per_second=n_queries / 0.1, + recall=0.95, + algorithm=self.algo, + search_params=first.search_params, + success=True, + ) + ] class AnotherDummyBackend(BenchmarkBackend): @@ -109,16 +111,18 @@ def search( distances = np.random.rand(n_queries, k) first = indexes[0] - return SearchResult( - neighbors=neighbors, - distances=distances, - search_time_ms=0.2, - queries_per_second=n_queries / 0.2, - recall=0.90, - algorithm=self.algo, - search_params=first.search_params if first else [], - success=True, - ) + return [ + SearchResult( + neighbors=neighbors, + distances=distances, + search_time_ms=0.2, + queries_per_second=n_queries / 0.2, + recall=0.90, + algorithm=self.algo, + search_params=first.search_params if first else [], + success=True, + ) + ] class TestDataset: @@ -386,7 +390,7 @@ def test_dummy_backend_search(self, tmp_path): ) ] - result = backend.search(dataset=dataset, indexes=indexes, k=10) + result = backend.search(dataset=dataset, indexes=indexes, k=10)[0] assert result.success assert result.recall == 0.95 From a82ff0fa30734047d45f593baa43004bc24220a1 Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Tue, 28 Jul 2026 13:29:47 -0500 Subject: [PATCH 17/22] Update Signed-off-by: James Bourbeau --- .../cuvs_bench/tests/test_data_export.py | 107 +----------------- 1 file changed, 1 insertion(+), 106 deletions(-) diff --git a/python/cuvs_bench/cuvs_bench/tests/test_data_export.py b/python/cuvs_bench/cuvs_bench/tests/test_data_export.py index 2595d32095..721d03c5ac 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_data_export.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_data_export.py @@ -5,7 +5,6 @@ import numpy as np import pandas as pd -from click.testing import CliRunner from cuvs_bench.backends.base import BuildResult, SearchResult from cuvs_bench.orchestrator.config_loaders import ( @@ -15,7 +14,6 @@ ) from cuvs_bench.orchestrator.orchestrator import BenchmarkOrchestrator from cuvs_bench.plot.__main__ import load_all_results -from cuvs_bench.run.__main__ import main from cuvs_bench.run.data_export import write_results_to_csv @@ -117,109 +115,6 @@ def test_python_backend_csv_is_plot_compatible(tmp_path): assert plotted[algorithm] -def test_run_command_always_writes_python_backend_csv(tmp_path, monkeypatch): - algorithm = "opensearch_faiss_hnsw" - result = SearchResult( - neighbors=np.empty((0, 2), dtype=np.int64), - distances=np.empty((0, 2), dtype=np.float32), - search_time_ms=10.0, - queries_per_second=200.0, - recall=1.0, - algorithm=algorithm, - search_params=[{"ef_search": 100}], - metadata={ - "group": "base", - "index_name": "test-index", - "latency_seconds": 0.01, - }, - ) - - class FakeOrchestrator: - def __init__(self, backend_type): - assert backend_type == "opensearch" - - def run_benchmark(self, **kwargs): - return [result] - - monkeypatch.setattr( - "cuvs_bench.run.__main__.BenchmarkOrchestrator", FakeOrchestrator - ) - backend_config = tmp_path / "backend.yaml" - backend_config.write_text("backend: opensearch\n") - - cli_result = CliRunner().invoke( - main, - [ - "--dataset", - "test-dataset", - "--dataset-path", - str(tmp_path), - "--algorithms", - algorithm, - "--groups", - "base", - "--count", - "2", - "--batch-size", - "2", - "--search-mode", - "latency", - "--search", - "--backend-config", - str(backend_config), - ], - ) - - assert cli_result.exit_code == 0, cli_result.output - assert ( - tmp_path - / "test-dataset" - / "result" - / "search" - / f"{algorithm},base,k2,bs2,raw.csv" - ).exists() - help_output = CliRunner().invoke(main, ["--help"]).output - assert "--data-export" in help_output - assert "Deprecated" in help_output - - -def test_data_export_option_is_deprecated(tmp_path, monkeypatch): - converted = [] - monkeypatch.setattr( - "cuvs_bench.run.__main__.convert_json_to_csv_build", - lambda dataset, dataset_path: converted.append("build"), - ) - monkeypatch.setattr( - "cuvs_bench.run.__main__.convert_json_to_csv_search", - lambda dataset, dataset_path: converted.append("search"), - ) - - cli_result = CliRunner().invoke( - main, - [ - "--dataset", - "test-dataset", - "--dataset-path", - str(tmp_path), - "--algorithms", - "cuvs_cagra", - "--groups", - "base", - "--count", - "2", - "--batch-size", - "2", - "--search-mode", - "latency", - "--data-export", - ], - ) - - assert cli_result.exit_code == 0, cli_result.output - assert "--data-export is deprecated" in cli_result.output - assert converted == ["build", "search"] - - def test_tune_trial_retains_build_and_search_results(): algorithm = "opensearch_faiss_hnsw" index = IndexConfig( @@ -277,7 +172,7 @@ def search(self, dataset, indexes, k, **kwargs): ) ] - orchestrator = object.__new__(BenchmarkOrchestrator) + orchestrator = BenchmarkOrchestrator(backend_type="opensearch") orchestrator.config_loader = FakeLoader() orchestrator.backend_class = FakeBackend orchestrator._create_dataset = lambda config: type( From 85d4ac23fa8660110e9aee4da6e5535515745767 Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Wed, 29 Jul 2026 10:45:43 -0500 Subject: [PATCH 18/22] Remove custom run.py module Signed-off-by: James Bourbeau --- deploy/README.md | 17 +- deploy/bench/Dockerfile | 2 +- deploy/bench/configure_opensearch.py | 123 +++++++++ deploy/bench/entrypoint.sh | 36 ++- deploy/bench/run.py | 398 --------------------------- 5 files changed, 156 insertions(+), 420 deletions(-) create mode 100644 deploy/bench/configure_opensearch.py delete mode 100644 deploy/bench/run.py diff --git a/deploy/README.md b/deploy/README.md index 923b04f21b..2795f5f81e 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -22,7 +22,7 @@ OpenSearch flushes a segment |---|---|---| | `opensearch` | custom build of `opensearchproject/opensearch` | OpenSearch node with kNN plugin and `repository-s3` plugin | | `remote-index-builder` | `opensearchproject/remote-vector-index-builder:api-latest` | FastAPI service that builds Faiss indexes on the GPU | -| `bench` | custom Python | Downloads dataset, registers repo + cluster settings, runs cuvs-bench build/search benchmark, exports results, generates plots | +| `bench` | `python:3.11-slim` | Downloads the dataset, configures OpenSearch, runs the standard cuvs-bench CLI, and generates plots | ## Requirements @@ -127,15 +127,14 @@ docker compose down -v 1. Downloads the dataset (skipped if already present in `$DATASET_PATH`) 2. **GPU mode only**: Registers the S3 bucket as an OpenSearch snapshot repository 3. **GPU mode only**: Applies cluster settings to enable remote index build and point OpenSearch at the builder service -4. Runs `cuvs-bench` build phase (handled entirely by the OpenSearch backend): +4. Runs the standard `python -m cuvs_bench.run` build and search workflow: - Creates the kNN index and bulk-ingests dataset vectors - **GPU mode**: Flushes segments, waits for all submitted remote GPU builds to complete, and polls the kNN stats API every 5 s until the build is confirmed complete - **CPU mode**: Flushes and refreshes the local OpenSearch index - Uses the backend's automatic OpenSearch bulk-ingest batch sizing by default; set `BUILD_BATCH_SIZE` to override it - Records total build time in the result -5. Runs `cuvs-bench` search phase and prints a recall/QPS/latency table -6. Exports benchmark JSON results to CSV -7. Generates recall vs. latency/throughput plots as PNGs in `$DATASET_PATH` (`cuvs_bench.plot`) + - Computes recall for each search-parameter set and writes the Python-backend results directly to the plotting CSV schema +5. Generates recall vs. latency/throughput plots as PNGs in `$DATASET_PATH` (`cuvs_bench.plot`) ## Dataset format @@ -161,7 +160,7 @@ $DATASET_PATH/ ## Key configuration -**Cluster settings** (applied by `bench/run.py`): +**Cluster settings** (applied by `bench/configure_opensearch.py`): ```json { @@ -231,7 +230,7 @@ docker compose up -d --wait opensearch docker compose run --rm --no-deps \ -e OPENSEARCH_URL=http://opensearch:9200 \ bench \ - pytest /opt/cuvs/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py -v -m integration + pytest /opt/cuvs/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py -v -m opensearch ``` ### Remote index build integration tests (full GPU stack) @@ -249,10 +248,10 @@ docker compose run --rm --no-deps \ -e S3_BUCKET=${S3_BUCKET} \ -e AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION} \ bench \ - pytest /opt/cuvs/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py -v -m integration + pytest /opt/cuvs/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py -v -m opensearch ``` -This lets the pytest `integration` marker decide which tests run. With only OpenSearch running, remote-build tests skip because the GPU builder and S3 environment are unavailable. With the GPU stack running, the same marker includes the remote-build coverage. +This lets the pytest `opensearch` marker decide which tests run. With only OpenSearch running, remote-build tests skip because the GPU builder and S3 environment are unavailable. With the GPU stack running, the same marker includes the remote-build coverage. ## Ports diff --git a/deploy/bench/Dockerfile b/deploy/bench/Dockerfile index f133d2a52e..60df983989 100644 --- a/deploy/bench/Dockerfile +++ b/deploy/bench/Dockerfile @@ -37,5 +37,5 @@ RUN pip install --no-cache-dir \ scikit-learn \ scipy -COPY --chmod=755 run.py entrypoint.sh prepare_custom_dataset.py . +COPY --chmod=755 configure_opensearch.py entrypoint.sh prepare_custom_dataset.py . CMD ["/app/entrypoint.sh"] diff --git a/deploy/bench/configure_opensearch.py b/deploy/bench/configure_opensearch.py new file mode 100644 index 0000000000..948158a82e --- /dev/null +++ b/deploy/bench/configure_opensearch.py @@ -0,0 +1,123 @@ +#!/usr/bin/env python3 +"""Configure OpenSearch and write the cuvs-bench backend configuration.""" + +import argparse +import os + +import requests +import yaml + + +def _optional_int(name: str) -> int | None: + value = os.environ.get(name, "").strip() + return int(value) if value else None + + +def create_backend_config() -> dict: + remote_index_build = ( + os.environ.get("REMOTE_INDEX_BUILD", "false").lower() == "true" + ) + number_of_shards = int(os.environ.get("NUMBER_OF_SHARDS", "1")) + if number_of_shards < 1: + raise ValueError("NUMBER_OF_SHARDS must be at least 1") + + approximate_threshold = int( + os.environ.get("APPROXIMATE_THRESHOLD", "10000") + ) + if approximate_threshold < -1: + raise ValueError("APPROXIMATE_THRESHOLD must be -1 or greater") + + config = { + "backend": "opensearch", + "host": os.environ.get("OPENSEARCH_HOST", "opensearch"), + "port": int(os.environ.get("OPENSEARCH_PORT", "9200")), + "use_ssl": False, + "verify_certs": False, + "number_of_shards": number_of_shards, + "approximate_threshold": approximate_threshold, + "remote_index_build": remote_index_build, + } + + build_batch_size = _optional_int("BUILD_BATCH_SIZE") + if build_batch_size is not None: + config["build_batch_size"] = build_batch_size + + if remote_index_build: + config["remote_build_timeout"] = int( + os.environ.get("REMOTE_BUILD_TIMEOUT", "1800") + ) + remote_build_size_min = os.environ.get( + "REMOTE_BUILD_SIZE_MIN", "" + ).strip() + if remote_build_size_min: + config["remote_build_size_min"] = remote_build_size_min + + return config + + +def configure_cluster() -> None: + opensearch_url = os.environ.get( + "OPENSEARCH_URL", "http://opensearch:9200" + ) + remote_index_build = ( + os.environ.get("REMOTE_INDEX_BUILD", "false").lower() == "true" + ) + session = requests.Session() + session.headers.update({"Content-Type": "application/json"}) + + if remote_index_build: + bucket = os.environ.get("S3_BUCKET", "").strip() + if not bucket: + raise ValueError( + "S3_BUCKET must be set when REMOTE_INDEX_BUILD=true" + ) + repository = ( + os.environ.get("REMOTE_VECTOR_REPOSITORY", "vector-repo").strip() + or "vector-repo" + ) + response = session.put( + f"{opensearch_url}/_snapshot/{repository}", + json={ + "type": "s3", + "settings": { + "bucket": bucket, + "base_path": ( + os.environ.get("S3_PREFIX", "knn-indexes").strip() + or "knn-indexes" + ), + "region": os.environ.get( + "AWS_DEFAULT_REGION", "us-west-2" + ), + }, + }, + ) + response.raise_for_status() + settings = { + "knn.remote_index_build.enabled": True, + "knn.remote_index_build.repository": repository, + "knn.remote_index_build.service.endpoint": os.environ.get( + "BUILDER_URL", "http://remote-index-builder:1025" + ), + } + else: + settings = {"knn.remote_index_build.enabled": False} + + response = session.put( + f"{opensearch_url}/_cluster/settings", + json={"persistent": settings}, + ) + response.raise_for_status() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("output") + args = parser.parse_args() + + configure_cluster() + with open(args.output, "w") as file: + yaml.safe_dump(create_backend_config(), file, sort_keys=False) + + +if __name__ == "__main__": + main() diff --git a/deploy/bench/entrypoint.sh b/deploy/bench/entrypoint.sh index b897dcc398..1bc717ed45 100644 --- a/deploy/bench/entrypoint.sh +++ b/deploy/bench/entrypoint.sh @@ -5,7 +5,9 @@ DATASET="${DATASET:-sift-128-euclidean}" CUSTOM_DATASET="miracl-en-5m-1024d-fp32" BENCH_GROUPS="${BENCH_GROUPS:-test}" K="${K:-10}" +BATCH_SIZE="${BATCH_SIZE:-10000}" ALGORITHM="opensearch_faiss_hnsw" +BACKEND_CONFIG="/tmp/opensearch-backend.yaml" export DATASET if [ "$DATASET" = "$CUSTOM_DATASET" ]; then @@ -61,20 +63,29 @@ else --dataset-path /data/datasets fi -# Step 2: Run benchmark (build + search + writes result JSON files) -python -u run.py +# Step 2: Configure OpenSearch and write the backend configuration. +python -u configure_opensearch.py "$BACKEND_CONFIG" -# Step 3: Export JSON → CSV (required by cuvs_bench.plot) -# --batch-size is ignored when --data-export is set, but Click prompts for it -# before entering main(), so pass a dummy value to keep the container non-interactive. -python -m cuvs_bench.run --data-export \ - --dataset "$DATASET" \ - --dataset-path /data/datasets \ - --algorithms "$ALGORITHM" \ - --groups "$BENCH_GROUPS" \ - --count "$K" \ - --batch-size 1 \ +# Step 3: Run the standard cuvs-bench CLI. Python backends write plotting CSV +# files automatically. +run_args=( + python -m cuvs_bench.run + --backend-config "$BACKEND_CONFIG" + --dataset "$DATASET" + --dataset-path /data/datasets + --algorithms "$ALGORITHM" + --groups "$BENCH_GROUPS" + --count "$K" + --batch-size "$BATCH_SIZE" --search-mode latency + --build + --search + --force +) +if [ -n "${DATASET_CONFIGURATION:-}" ]; then + run_args+=(--dataset-configuration "$DATASET_CONFIGURATION") +fi +"${run_args[@]}" # Step 4: Plot — PNGs written to /data/datasets (mounted from host $DATASET_PATH) python -m cuvs_bench.plot \ @@ -83,4 +94,5 @@ python -m cuvs_bench.plot \ --algorithms "$ALGORITHM" \ --groups "$BENCH_GROUPS" \ --count "$K" \ + --batch-size "$BATCH_SIZE" \ --output-filepath /data/datasets diff --git a/deploy/bench/run.py b/deploy/bench/run.py deleted file mode 100644 index 3065256710..0000000000 --- a/deploy/bench/run.py +++ /dev/null @@ -1,398 +0,0 @@ -#!/usr/bin/env python3 -""" -OpenSearch GPU Remote Index Build Benchmark -=========================================== -Steps: - 1. Register S3 snapshot repository with OpenSearch - 2. Configure cluster settings for GPU remote index build - 3. Build kNN index via cuvs-bench: - a. Bulk-ingest dataset vectors - b. Flush segments to kick off remote GPU builds when enabled - c. Poll kNN stats until every submitted remote build completes - 4. Run cuvs-bench search benchmarks and print results - 5. Write gbench-compatible JSON result files for CSV export and plotting -""" - -import json -import os -import sys - -import requests -from cuvs_bench.orchestrator import BenchmarkOrchestrator -from cuvs_bench.backends.base import BuildResult, SearchResult - - -OPENSEARCH_URL = os.environ.get("OPENSEARCH_URL", "http://opensearch:9200") -OPENSEARCH_HOST = os.environ.get("OPENSEARCH_HOST", "opensearch") -OPENSEARCH_PORT = int(os.environ.get("OPENSEARCH_PORT", "9200")) -BUILDER_URL = os.environ.get("BUILDER_URL", "http://remote-index-builder:1025") - -REMOTE_INDEX_BUILD = os.environ.get("REMOTE_INDEX_BUILD", "false").lower() == "true" -REMOTE_BUILD_SIZE_MIN = os.environ.get("REMOTE_BUILD_SIZE_MIN", "").strip() -REMOTE_BUILD_TIMEOUT = int(os.environ.get("REMOTE_BUILD_TIMEOUT", "1800")) - -S3_BUCKET = os.environ.get("S3_BUCKET", "").strip() -S3_PREFIX = os.environ.get("S3_PREFIX", "knn-indexes").strip() or "knn-indexes" -S3_REGION = os.environ.get("AWS_DEFAULT_REGION", "us-west-2") - -DATASET = os.environ.get("DATASET", "sift-128-euclidean") -DATASET_PATH = os.environ.get("DATASET_PATH", "/data/datasets") -DATASET_CONFIGURATION = os.environ.get("DATASET_CONFIGURATION", "").strip() -BENCH_GROUPS = os.environ.get("BENCH_GROUPS", "test") -K = int(os.environ.get("K", "10")) -BATCH_SIZE = os.environ.get("BATCH_SIZE", "").strip() -BATCH_SIZE = int(BATCH_SIZE) if BATCH_SIZE else None -BUILD_BATCH_SIZE = os.environ.get("BUILD_BATCH_SIZE", "").strip() -BUILD_BATCH_SIZE = int(BUILD_BATCH_SIZE) if BUILD_BATCH_SIZE else None -NUMBER_OF_SHARDS = int(os.environ.get("NUMBER_OF_SHARDS", "1")) -if NUMBER_OF_SHARDS < 1: - raise ValueError("NUMBER_OF_SHARDS must be at least 1") -APPROXIMATE_THRESHOLD = int(os.environ.get("APPROXIMATE_THRESHOLD", "10000")) -if APPROXIMATE_THRESHOLD < -1: - raise ValueError("APPROXIMATE_THRESHOLD must be -1 or greater") - -ALGORITHM = "opensearch_faiss_hnsw" -REPO_NAME = os.environ.get("REMOTE_VECTOR_REPOSITORY", "vector-repo").strip() -REPO_NAME = REPO_NAME or "vector-repo" - -session = requests.Session() -session.headers.update({"Content-Type": "application/json"}) - - -# ── helpers ─────────────────────────────────────────────────────────────────── - -def banner(msg: str) -> None: - print(f"\n{'─'*60}\n {msg}\n{'─'*60}") - - -def _recall_for_entry( - result: SearchResult, entry_index: int, entry_count: int -) -> float | None: - # cuVS computes recall in the orchestrator from SearchResult.neighbors. The - # OpenSearch backend returns neighbors for the final search-parameter run. - if entry_index == entry_count - 1: - return float(result.recall) - return None - - -def _get_search_batch_size(search_results: list) -> int | None: - if BATCH_SIZE is not None: - return BATCH_SIZE - for result in search_results: - if not isinstance(result, SearchResult): - continue - batch_size = (result.metadata or {}).get("batch_size") - if batch_size is not None: - return int(batch_size) - return None - - -# ── OpenSearch setup ────────────────────────────────────────────────────────── - -def register_repository() -> None: - banner(f"Registering S3 repository '{REPO_NAME}'") - r = session.put( - f"{OPENSEARCH_URL}/_snapshot/{REPO_NAME}", - json={ - "type": "s3", - "settings": { - "bucket": S3_BUCKET, - "base_path": S3_PREFIX, - "region": S3_REGION, - }, - }, - ) - r.raise_for_status() - print(f" {r.json()}") - - -def configure_cluster() -> None: - if REMOTE_INDEX_BUILD: - banner("Enabling GPU remote index build (cluster settings)") - settings = { - "knn.remote_index_build.enabled": True, - "knn.remote_index_build.repository": REPO_NAME, - "knn.remote_index_build.service.endpoint": BUILDER_URL, - } - else: - banner("Disabling GPU remote index build (cluster settings)") - settings = { - "knn.remote_index_build.enabled": False, - } - r = session.put( - f"{OPENSEARCH_URL}/_cluster/settings", - json={"persistent": settings}, - ) - r.raise_for_status() - print(f" {r.json()}") - - -# ── result files ───────────────────────────────────────────────────────────── - -def write_result_files( - build_results: list, - search_results: list, - dataset: str, - dataset_path: str, - algo: str, - groups: str, - k: int, - batch_size: int, -) -> None: - """Write gbench-compatible JSON result files. - - Creates files under //result/{build,search}/ in the - same format the C++ backend produces, so the cuvs-bench CSV exporters and - ``cuvs_bench.plot`` work without modification. - """ - build_dir = os.path.join(dataset_path, dataset, "result", "build") - search_dir = os.path.join(dataset_path, dataset, "result", "search") - os.makedirs(build_dir, exist_ok=True) - os.makedirs(search_dir, exist_ok=True) - - # Build JSON – one record per successfully built index. - # data_export.py assumes the build CSV has columns: - # [algo_name, index_name, time, threads, cpu_time, ...] - # so "threads" and "cpu_time" must be present in the JSON (they are not in - # skip_build_cols and therefore get included as columns 3 and 4). - build_benchmarks = [ - { - "name": r.index_path, - "real_time": r.build_time_seconds, - "time_unit": "s", - "threads": 1, - "cpu_time": r.build_time_seconds, - } - for r in build_results - if isinstance(r, BuildResult) and r.success and r.index_path - ] - - # Search JSON – zip build + search results to recover the index name, then - # expand per-search-param entries so each (index, ef_search) is one record. - build_list = [r for r in build_results if isinstance(r, BuildResult)] - search_list = [r for r in search_results if isinstance(r, SearchResult)] - search_benchmarks = [] - skipped_without_recall = 0 - for build_r, search_r in zip(build_list, search_list): - if not search_r.success or not build_r.index_path: - continue - per_param = (search_r.metadata or {}).get("per_search_param_results", []) - for entry_index, entry in enumerate(per_param): - recall = _recall_for_entry(search_r, entry_index, len(per_param)) - if recall is None: - skipped_without_recall += 1 - continue - latency_ms = float(entry["search_time_ms"]) - search_benchmarks.append( - { - "name": build_r.index_path, - "real_time": latency_ms, - "time_unit": "ms", - "Recall": recall, - "items_per_second": float(entry["queries_per_second"]), - # Latency field expected by data_export in seconds - "Latency": latency_ms / 1000.0, - } - ) - - build_file = os.path.join(build_dir, f"{algo},{groups}.json") - search_file = os.path.join( - search_dir, f"{algo},{groups},k{k},bs{batch_size}.json" - ) - - with open(build_file, "w") as fh: - json.dump({"benchmarks": build_benchmarks}, fh, indent=2) - with open(search_file, "w") as fh: - json.dump({"benchmarks": search_benchmarks}, fh, indent=2) - - print(f"\n Result files written:") - print(f" {build_file}") - print(f" {search_file}") - if skipped_without_recall: - print( - " skipped " - f"{skipped_without_recall} search rows without per-parameter recall" - ) - - -# ── results ─────────────────────────────────────────────────────────────────── - -def _print_result_row( - params: dict, recall: float | None, qps: float | None, latency_ms: float | None -) -> None: - params_str = ", ".join(f"{k}={v}" for k, v in params.items()) - recall_str = "n/a" if recall is None else f"{recall:.4f}" - qps_str = "n/a" if qps is None else f"{qps:.1f}" - latency_str = "n/a" if latency_ms is None else f"{latency_ms:.2f}" - print(f" {params_str:<40} {recall_str:<12} {qps_str:>8} {latency_str:>12}") - - -def print_results(results: list) -> None: - banner("Benchmark Results") - search_results = [r for r in results if isinstance(r, SearchResult)] - if not search_results: - print(" No search results returned.") - return - - header = f" {'params':<40} {'recall@'+str(K):<12} {'QPS':>8} {'latency (ms)':>12}" - print(header) - print(" " + "─" * (len(header) - 2)) - missing_recall_rows = 0 - for r in search_results: - per_param = (r.metadata or {}).get("per_search_param_results", []) - entry_count = len(per_param) - for entry_index, entry in enumerate(per_param): - recall = _recall_for_entry(r, entry_index, entry_count) - if recall is None: - missing_recall_rows += 1 - continue - _print_result_row( - entry["search_params"], - recall, - float(entry["queries_per_second"]), - float(entry["search_time_ms"]), - ) - if missing_recall_rows: - print( - "\n Omitted " - f"{missing_recall_rows} timing-only search rows without recall." - ) - - -# ── entrypoint ──────────────────────────────────────────────────────────────── - -def main() -> None: - if REMOTE_INDEX_BUILD and not S3_BUCKET: - print( - "ERROR: S3_BUCKET is not set. Remote index build requires an S3 " - "bucket for vector and index staging. Set it before starting the " - "stack, for example: export S3_BUCKET=", - file=sys.stderr, - ) - sys.exit(1) - - print("\n" + "═" * 60) - print(" OpenSearch kNN Benchmark") - print("═" * 60) - print(f" OpenSearch : {OPENSEARCH_URL}") - print(f" Remote index build : {REMOTE_INDEX_BUILD}") - if REMOTE_INDEX_BUILD: - print(f" GPU builder : {BUILDER_URL}") - print(f" S3 bucket : s3://{S3_BUCKET}/{S3_PREFIX}/ (region: {S3_REGION})") - print(f" Repository : {REPO_NAME}") - print(f" Build size minimum : {REMOTE_BUILD_SIZE_MIN or 'OpenSearch default'}") - print(f" Build timeout : {REMOTE_BUILD_TIMEOUT}s") - print(f" Dataset : {DATASET} (path: {DATASET_PATH})") - print(f" Algorithm : {ALGORITHM}") - print(f" Groups : {BENCH_GROUPS} k={K}") - print( - " Search batch size : " - f"{BATCH_SIZE if BATCH_SIZE is not None else 'backend default'}" - ) - print( - " Build batch size : " - f"{BUILD_BATCH_SIZE if BUILD_BATCH_SIZE is not None else 'backend auto'}" - ) - print(f" Index shards : {NUMBER_OF_SHARDS}") - print(f" Approx. threshold : {APPROXIMATE_THRESHOLD}") - - if REMOTE_INDEX_BUILD: - register_repository() - configure_cluster() - - orchestrator = BenchmarkOrchestrator(backend_type="opensearch") - - # Shared kwargs for both build and search phases. - common_kwargs = dict( - dataset=DATASET, - dataset_path=DATASET_PATH, - algorithms=ALGORITHM, - groups=BENCH_GROUPS, - host=OPENSEARCH_HOST, - port=OPENSEARCH_PORT, - number_of_shards=NUMBER_OF_SHARDS, - use_ssl=False, - verify_certs=False, - ) - if DATASET_CONFIGURATION: - common_kwargs["dataset_configuration"] = DATASET_CONFIGURATION - - build_kwargs = dict( - common_kwargs, - remote_index_build=REMOTE_INDEX_BUILD, - approximate_threshold=APPROXIMATE_THRESHOLD, - ) - if BUILD_BATCH_SIZE is not None: - build_kwargs["build_batch_size"] = BUILD_BATCH_SIZE - if REMOTE_INDEX_BUILD: - build_kwargs["remote_build_timeout"] = REMOTE_BUILD_TIMEOUT - if REMOTE_BUILD_SIZE_MIN: - build_kwargs["remote_build_size_min"] = REMOTE_BUILD_SIZE_MIN - - # ── Build phase ─────────────────────────────────────────────────────────── - mode = "GPU remote build" if REMOTE_INDEX_BUILD else "CPU" - banner(f"Building index ({mode} via cuvs-bench)") - build_results = orchestrator.run_benchmark( - build=True, - search=False, - force=True, - **build_kwargs, - ) - - index_names = [ - r.index_path - for r in build_results - if isinstance(r, BuildResult) and r.success and r.index_path - ] - if not index_names: - print(" ERROR: no indexes were successfully built") - sys.exit(1) - build_times = { - r.index_path: r.build_time_seconds - for r in build_results - if isinstance(r, BuildResult) and r.success and r.index_path - } - for name, t in build_times.items(): - print(f" {name}: built in {t:.1f}s") - - # ── Search phase ────────────────────────────────────────────────────────── - banner("Running search benchmarks (via cuvs-bench)") - search_run_kwargs = dict( - build=False, - search=True, - count=K, - **common_kwargs, - ) - if BATCH_SIZE is not None: - search_run_kwargs["batch_size"] = BATCH_SIZE - search_results = orchestrator.run_benchmark(**search_run_kwargs) - - print_results(search_results) - - batch_size = _get_search_batch_size(search_results) - if batch_size is None: - print(" ERROR: no successful search result reported a batch size") - sys.exit(1) - - write_result_files( - build_results=build_results, - search_results=search_results, - dataset=DATASET, - dataset_path=DATASET_PATH, - algo=common_kwargs["algorithms"], - groups=BENCH_GROUPS, - k=K, - batch_size=batch_size, - ) - - print("\n" + "═" * 60) - print(" Benchmark complete!") - print("═" * 60) - print(f"\n OpenSearch : {OPENSEARCH_URL}") - if REMOTE_INDEX_BUILD: - print(f" GPU builder : {BUILDER_URL}") - print() - - -if __name__ == "__main__": - main() From edca7e812a53454dfa3a7f78e1a9b0d08b72f5a3 Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Wed, 29 Jul 2026 21:03:51 -0500 Subject: [PATCH 19/22] Refresh interval Signed-off-by: James Bourbeau --- deploy/AWS_MULTINODE_SETUP.md | 1 + deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md | 1 + deploy/README.md | 9 +- deploy/bench/Dockerfile | 2 +- deploy/bench/configure_opensearch.py | 4 + deploy/bench/entrypoint.sh | 11 +- deploy/bench/print_results.py | 102 ++++++++++++++++++ deploy/docker-compose.multinode.yml | 1 + deploy/docker-compose.yml | 2 + .../cuvs_bench/backends/opensearch.py | 24 +++-- .../cuvs_bench/tests/test_opensearch.py | 25 +++++ 11 files changed, 172 insertions(+), 10 deletions(-) create mode 100644 deploy/bench/print_results.py diff --git a/deploy/AWS_MULTINODE_SETUP.md b/deploy/AWS_MULTINODE_SETUP.md index e7bc69d135..1dcd543cf3 100644 --- a/deploy/AWS_MULTINODE_SETUP.md +++ b/deploy/AWS_MULTINODE_SETUP.md @@ -337,6 +337,7 @@ export BATCH_SIZE= export BUILD_BATCH_SIZE= export NUMBER_OF_SHARDS=1 export APPROXIMATE_THRESHOLD=10000 +export REFRESH_INTERVAL= mkdir -p "${DATASET_PATH}" diff --git a/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md b/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md index 870b56e2b3..927b9149ba 100644 --- a/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md +++ b/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md @@ -327,6 +327,7 @@ export BATCH_SIZE= export BUILD_BATCH_SIZE= export NUMBER_OF_SHARDS=1 export APPROXIMATE_THRESHOLD=10000 +export REFRESH_INTERVAL= mkdir -p "${DATASET_PATH}" diff --git a/deploy/README.md b/deploy/README.md index 2795f5f81e..dba9f081d7 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -74,6 +74,7 @@ export BATCH_SIZE= # optional query batch size override export BUILD_BATCH_SIZE= # optional bulk ingest batch size override export NUMBER_OF_SHARDS=1 # number of primary index shards (default: 1) export APPROXIMATE_THRESHOLD=10000 # vectors per segment before ANN build (default: 10000) +export REFRESH_INTERVAL= # optional index refresh interval (for example: 30s or -1) export REMOTE_BUILD_TIMEOUT=1800 # seconds to wait for remote builds (default: 1800) ``` @@ -83,6 +84,11 @@ recommends `10000` as a starting point for GPU-accelerated indexing so smaller segments do not build ANN structures prematurely. Set it to `0` to always build ANN structures or `-1` to disable them. +`REFRESH_INTERVAL` sets `index.refresh_interval` on each benchmark index. +Leave it empty to use the OpenSearch default, or set it to `-1` to disable +automatic refreshes during ingestion. cuvs-bench performs an explicit refresh +before searching. + Start all services: ```bash @@ -134,7 +140,8 @@ docker compose down -v - Uses the backend's automatic OpenSearch bulk-ingest batch sizing by default; set `BUILD_BATCH_SIZE` to override it - Records total build time in the result - Computes recall for each search-parameter set and writes the Python-backend results directly to the plotting CSV schema -5. Generates recall vs. latency/throughput plots as PNGs in `$DATASET_PATH` (`cuvs_bench.plot`) +5. Prints a compact build-time and search recall/QPS/latency overview +6. Generates recall vs. latency/throughput plots as PNGs in `$DATASET_PATH` (`cuvs_bench.plot`) ## Dataset format diff --git a/deploy/bench/Dockerfile b/deploy/bench/Dockerfile index 60df983989..e457fb2f15 100644 --- a/deploy/bench/Dockerfile +++ b/deploy/bench/Dockerfile @@ -37,5 +37,5 @@ RUN pip install --no-cache-dir \ scikit-learn \ scipy -COPY --chmod=755 configure_opensearch.py entrypoint.sh prepare_custom_dataset.py . +COPY --chmod=755 configure_opensearch.py entrypoint.sh prepare_custom_dataset.py print_results.py . CMD ["/app/entrypoint.sh"] diff --git a/deploy/bench/configure_opensearch.py b/deploy/bench/configure_opensearch.py index 948158a82e..664faf8c5e 100644 --- a/deploy/bench/configure_opensearch.py +++ b/deploy/bench/configure_opensearch.py @@ -42,6 +42,10 @@ def create_backend_config() -> dict: if build_batch_size is not None: config["build_batch_size"] = build_batch_size + refresh_interval = os.environ.get("REFRESH_INTERVAL", "").strip() + if refresh_interval: + config["refresh_interval"] = refresh_interval + if remote_index_build: config["remote_build_timeout"] = int( os.environ.get("REMOTE_BUILD_TIMEOUT", "1800") diff --git a/deploy/bench/entrypoint.sh b/deploy/bench/entrypoint.sh index 1bc717ed45..e23a5c7b2f 100644 --- a/deploy/bench/entrypoint.sh +++ b/deploy/bench/entrypoint.sh @@ -87,7 +87,16 @@ if [ -n "${DATASET_CONFIGURATION:-}" ]; then fi "${run_args[@]}" -# Step 4: Plot — PNGs written to /data/datasets (mounted from host $DATASET_PATH) +# Step 4: Print a compact overview of the generated results. +python -u print_results.py \ + --dataset-path /data/datasets \ + --dataset "$DATASET" \ + --algorithm "$ALGORITHM" \ + --groups "$BENCH_GROUPS" \ + --count "$K" \ + --batch-size "$BATCH_SIZE" + +# Step 5: Plot — PNGs written to /data/datasets (mounted from host $DATASET_PATH) python -m cuvs_bench.plot \ --dataset "$DATASET" \ --dataset-path /data/datasets \ diff --git a/deploy/bench/print_results.py b/deploy/bench/print_results.py new file mode 100644 index 0000000000..b9ab30d075 --- /dev/null +++ b/deploy/bench/print_results.py @@ -0,0 +1,102 @@ +#!/usr/bin/env python3 +"""Print a compact overview of cuvs-bench CSV results.""" + +import argparse +from pathlib import Path + +import pandas as pd + +_BUILD_COLUMNS = {"algo_name", "index_name", "time"} +_SEARCH_COLUMNS = { + "algo_name", + "index_name", + "recall", + "throughput", + "latency", +} +_METADATA_COLUMNS = { + "batch_size", + "build time", + "engine", + "num_batches", + "space_type", +} + + +def _format_params(row, excluded: set[str]) -> str: + values = [] + for name, value in row.items(): + if name in excluded or pd.isna(value): + continue + values.append(f"{name}={value}") + return ", ".join(values) or "default" + + +def print_results( + dataset_path: str, + dataset: str, + algorithm: str, + groups: str, + count: int, + batch_size: int, +) -> None: + result_dir = Path(dataset_path) / dataset / "result" + group_names = [ + group.strip() for group in groups.split(",") if group.strip() + ] + + print("\nBuild results:") + for group in group_names: + build_file = result_dir / "build" / f"{algorithm},{group}.csv" + if not build_file.exists(): + print(f" [{group}] no build results") + continue + for _, row in pd.read_csv(build_file).iterrows(): + params = _format_params( + row, _BUILD_COLUMNS | _METADATA_COLUMNS + ) + print( + f" {row['algo_name']} index={row['index_name']} " + f"time={float(row['time']):.2f}s params={params}" + ) + + print("\nSearch results:") + for group in group_names: + stem = f"{algorithm},{group},k{count},bs{batch_size},raw.csv" + search_file = result_dir / "search" / stem + if not search_file.exists(): + print(f" [{group}] no search results") + continue + for _, row in pd.read_csv(search_file).iterrows(): + params = _format_params( + row, _SEARCH_COLUMNS | _METADATA_COLUMNS + ) + print( + f" {row['algo_name']} index={row['index_name']} " + f"params={params} recall={float(row['recall']):.4f} " + f"qps={float(row['throughput']):.1f} " + f"latency={float(row['latency']) * 1000.0:.2f}ms" + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--dataset-path", required=True) + parser.add_argument("--dataset", required=True) + parser.add_argument("--algorithm", required=True) + parser.add_argument("--groups", required=True) + parser.add_argument("--count", required=True, type=int) + parser.add_argument("--batch-size", required=True, type=int) + args = parser.parse_args() + print_results( + args.dataset_path, + args.dataset, + args.algorithm, + args.groups, + args.count, + args.batch_size, + ) + + +if __name__ == "__main__": + main() diff --git a/deploy/docker-compose.multinode.yml b/deploy/docker-compose.multinode.yml index 62fa512ab3..68aefcb7af 100644 --- a/deploy/docker-compose.multinode.yml +++ b/deploy/docker-compose.multinode.yml @@ -52,6 +52,7 @@ x-bench-env: &bench-env BUILD_BATCH_SIZE: ${BUILD_BATCH_SIZE:-} NUMBER_OF_SHARDS: ${NUMBER_OF_SHARDS:-1} APPROXIMATE_THRESHOLD: ${APPROXIMATE_THRESHOLD:-10000} + REFRESH_INTERVAL: ${REFRESH_INTERVAL:-} services: opensearch: diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml index a08149445b..0b9ce5407e 100644 --- a/deploy/docker-compose.yml +++ b/deploy/docker-compose.yml @@ -39,6 +39,7 @@ # BUILD_BATCH_SIZE Optional OpenSearch bulk ingest batch size override # NUMBER_OF_SHARDS Number of primary index shards (default: 1) # APPROXIMATE_THRESHOLD Vectors per segment before ANN build (default: 10000) +# REFRESH_INTERVAL Optional index refresh interval (for example: 30s or -1) # CUVS_REPOSITORY Repository cloned into the benchmark image # CUVS_BRANCH Repository branch cloned into the benchmark image # @@ -138,6 +139,7 @@ services: BUILD_BATCH_SIZE: ${BUILD_BATCH_SIZE:-} NUMBER_OF_SHARDS: ${NUMBER_OF_SHARDS:-1} APPROXIMATE_THRESHOLD: ${APPROXIMATE_THRESHOLD:-10000} + REFRESH_INTERVAL: ${REFRESH_INTERVAL:-} volumes: - ${DATASET_PATH:-/tmp/datasets}:/data/datasets restart: "no" diff --git a/python/cuvs_bench/cuvs_bench/backends/opensearch.py b/python/cuvs_bench/cuvs_bench/backends/opensearch.py index 6c78645f1d..1ae341fd57 100644 --- a/python/cuvs_bench/cuvs_bench/backends/opensearch.py +++ b/python/cuvs_bench/cuvs_bench/backends/opensearch.py @@ -132,6 +132,7 @@ def _build_benchmark_configs( "verify_certs", "build_batch_size", "approximate_threshold", + "refresh_interval", # Remote Index Build (OpenSearch 3.0+, faiss engine only) "remote_index_build", "remote_build_size_min", @@ -290,6 +291,9 @@ class OpenSearchBackend(BenchmarkBackend): a batch size with roughly 1 MiB of raw vector data. - ``approximate_threshold`` – minimum vectors per segment before building ANN data structures (default: OpenSearch's default). + - ``refresh_interval`` – how often OpenSearch refreshes the index, + e.g. ``"1s"`` or ``"-1"`` to disable automatic refreshes during + ingestion (default: OpenSearch's default). - ``requires_network`` – trigger network pre-flight check (default: ``True``) - ``remote_index_build`` – set ``index.knn.remote_index_build.enabled=true`` on the index at creation time, opting it into the GPU build path (default: ``False``). @@ -362,6 +366,7 @@ def _build_index_mapping( remote_index_build: bool = False, remote_build_size_min: Optional[str] = None, approximate_threshold: Optional[int] = None, + refresh_interval: Optional[str] = None, ) -> Dict[str, Any]: """ Construct the OpenSearch index mapping dict for k-NN. @@ -420,6 +425,8 @@ def _build_index_mapping( index_settings["knn.advanced.approximate_threshold"] = ( approximate_threshold ) + if refresh_interval is not None: + index_settings["refresh_interval"] = str(refresh_interval) if remote_index_build: if engine != "faiss": raise ValueError( @@ -731,6 +738,7 @@ def build( remote_index_build = bool(self.config.get("remote_index_build", False)) remote_build_size_min = self.config.get("remote_build_size_min") approximate_threshold = self.config.get("approximate_threshold") + refresh_interval = self.config.get("refresh_interval") if dry_run: print( @@ -781,13 +789,14 @@ def build( # Create index mapping = self._build_index_mapping( - dims, - engine, - space_type, - build_param, - remote_index_build, - remote_build_size_min, - approximate_threshold, + dims=dims, + engine=engine, + space_type=space_type, + build_param=build_param, + remote_index_build=remote_index_build, + remote_build_size_min=remote_build_size_min, + approximate_threshold=approximate_threshold, + refresh_interval=refresh_interval, ) self._client.indices.create(index=index_name, body=mapping) @@ -830,6 +839,7 @@ def build( "space_type": space_type, "remote_index_build": remote_index_build, "approximate_threshold": approximate_threshold, + "refresh_interval": refresh_interval, }, success=True, ) diff --git a/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py b/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py index 884e80bb15..43f11dd701 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py @@ -162,6 +162,7 @@ def test_load_forwards_remote_build_kwargs(self, config_dir): remote_build_timeout=123, remote_build_s3_endpoint="http://s3:9000", approximate_threshold=10_000, + refresh_interval="-1", ) bc = configs[0].backend_config @@ -169,6 +170,7 @@ def test_load_forwards_remote_build_kwargs(self, config_dir): assert bc["remote_build_size_min"] == "2kb" assert bc["remote_build_timeout"] == 123 assert bc["approximate_threshold"] == 10_000 + assert bc["refresh_interval"] == "-1" assert "remote_build_s3_endpoint" not in bc def test_load_overrides_number_of_shards(self, config_dir): @@ -339,6 +341,29 @@ def test_approximate_threshold_rejects_values_below_minus_one(self): approximate_threshold=-2, ) + def test_refresh_interval_is_added_to_index_settings(self): + backend = _make_backend() + mapping = backend._build_index_mapping( + dims=4, + engine="faiss", + space_type="l2", + build_param={}, + refresh_interval="-1", + ) + + assert mapping["settings"]["index"]["refresh_interval"] == "-1" + + def test_refresh_interval_uses_opensearch_default_when_unspecified(self): + backend = _make_backend() + mapping = backend._build_index_mapping( + dims=4, + engine="faiss", + space_type="l2", + build_param={}, + ) + + assert "refresh_interval" not in mapping["settings"]["index"] + def test_wait_for_remote_build_raises_on_failure_count(self): backend = _make_backend() initial_stats = { From 1dcf66e5f985fd0f7ad6b24214e642a24bbffa43 Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Thu, 30 Jul 2026 12:08:22 -0500 Subject: [PATCH 20/22] Make approximate_threshold optional Signed-off-by: James Bourbeau --- deploy/AWS_MULTINODE_SETUP.md | 2 +- deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md | 2 +- deploy/DEPLOYMENT.md | 2 +- deploy/README.md | 13 +++++++------ deploy/bench/configure_opensearch.py | 9 ++++----- deploy/bench/entrypoint.sh | 1 + deploy/docker-compose.multinode.yml | 2 +- deploy/docker-compose.yml | 4 ++-- deploy/remote-index-build/run.py | 12 +++++++++--- 9 files changed, 27 insertions(+), 20 deletions(-) diff --git a/deploy/AWS_MULTINODE_SETUP.md b/deploy/AWS_MULTINODE_SETUP.md index 1dcd543cf3..f249b2c966 100644 --- a/deploy/AWS_MULTINODE_SETUP.md +++ b/deploy/AWS_MULTINODE_SETUP.md @@ -336,7 +336,7 @@ export K=10 export BATCH_SIZE= export BUILD_BATCH_SIZE= export NUMBER_OF_SHARDS=1 -export APPROXIMATE_THRESHOLD=10000 +export APPROXIMATE_THRESHOLD= export REFRESH_INTERVAL= mkdir -p "${DATASET_PATH}" diff --git a/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md b/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md index 927b9149ba..0fac09b54b 100644 --- a/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md +++ b/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md @@ -326,7 +326,7 @@ export K=10 export BATCH_SIZE= export BUILD_BATCH_SIZE= export NUMBER_OF_SHARDS=1 -export APPROXIMATE_THRESHOLD=10000 +export APPROXIMATE_THRESHOLD= export REFRESH_INTERVAL= mkdir -p "${DATASET_PATH}" diff --git a/deploy/DEPLOYMENT.md b/deploy/DEPLOYMENT.md index 423c1e8b1c..9ec929aa22 100644 --- a/deploy/DEPLOYMENT.md +++ b/deploy/DEPLOYMENT.md @@ -164,7 +164,7 @@ docker compose run --rm \ -e REMOTE_BUILD_SIZE_MIN=${REMOTE_BUILD_SIZE_MIN:-} \ -e REMOTE_BUILD_TIMEOUT=${REMOTE_BUILD_TIMEOUT:-1800} \ -e NUMBER_OF_SHARDS=${NUMBER_OF_SHARDS:-1} \ - -e APPROXIMATE_THRESHOLD=${APPROXIMATE_THRESHOLD:-10000} \ + -e APPROXIMATE_THRESHOLD=${APPROXIMATE_THRESHOLD:-} \ -v $(pwd)/remote-index-build:/app/remote-index-build \ --no-deps bench \ python remote-index-build/run.py diff --git a/deploy/README.md b/deploy/README.md index dba9f081d7..4f875eaf89 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -73,16 +73,17 @@ export K=10 # number of neighbors (default: 10) export BATCH_SIZE= # optional query batch size override export BUILD_BATCH_SIZE= # optional bulk ingest batch size override export NUMBER_OF_SHARDS=1 # number of primary index shards (default: 1) -export APPROXIMATE_THRESHOLD=10000 # vectors per segment before ANN build (default: 10000) +export APPROXIMATE_THRESHOLD= # optional vectors per segment before ANN build export REFRESH_INTERVAL= # optional index refresh interval (for example: 30s or -1) export REMOTE_BUILD_TIMEOUT=1800 # seconds to wait for remote builds (default: 1800) ``` -`APPROXIMATE_THRESHOLD` sets -`index.knn.advanced.approximate_threshold` on each benchmark index. AWS -recommends `10000` as a starting point for GPU-accelerated indexing so smaller -segments do not build ANN structures prematurely. Set it to `0` to always -build ANN structures or `-1` to disable them. +When set, `APPROXIMATE_THRESHOLD` sets +`index.knn.advanced.approximate_threshold` on each benchmark index. Leave it +empty to use OpenSearch's default. AWS recommends `10000` as a starting point +for GPU-accelerated indexing so smaller segments do not build ANN structures +prematurely. Set it to `0` to always build ANN structures or `-1` to disable +them. `REFRESH_INTERVAL` sets `index.refresh_interval` on each benchmark index. Leave it empty to use the OpenSearch default, or set it to `-1` to disable diff --git a/deploy/bench/configure_opensearch.py b/deploy/bench/configure_opensearch.py index 664faf8c5e..cff14a938e 100644 --- a/deploy/bench/configure_opensearch.py +++ b/deploy/bench/configure_opensearch.py @@ -21,10 +21,8 @@ def create_backend_config() -> dict: if number_of_shards < 1: raise ValueError("NUMBER_OF_SHARDS must be at least 1") - approximate_threshold = int( - os.environ.get("APPROXIMATE_THRESHOLD", "10000") - ) - if approximate_threshold < -1: + approximate_threshold = _optional_int("APPROXIMATE_THRESHOLD") + if approximate_threshold is not None and approximate_threshold < -1: raise ValueError("APPROXIMATE_THRESHOLD must be -1 or greater") config = { @@ -34,9 +32,10 @@ def create_backend_config() -> dict: "use_ssl": False, "verify_certs": False, "number_of_shards": number_of_shards, - "approximate_threshold": approximate_threshold, "remote_index_build": remote_index_build, } + if approximate_threshold is not None: + config["approximate_threshold"] = approximate_threshold build_batch_size = _optional_int("BUILD_BATCH_SIZE") if build_batch_size is not None: diff --git a/deploy/bench/entrypoint.sh b/deploy/bench/entrypoint.sh index e23a5c7b2f..3513e51c06 100644 --- a/deploy/bench/entrypoint.sh +++ b/deploy/bench/entrypoint.sh @@ -104,4 +104,5 @@ python -m cuvs_bench.plot \ --groups "$BENCH_GROUPS" \ --count "$K" \ --batch-size "$BATCH_SIZE" \ + --raw \ --output-filepath /data/datasets diff --git a/deploy/docker-compose.multinode.yml b/deploy/docker-compose.multinode.yml index 68aefcb7af..c4ff34b5c2 100644 --- a/deploy/docker-compose.multinode.yml +++ b/deploy/docker-compose.multinode.yml @@ -51,7 +51,7 @@ x-bench-env: &bench-env BATCH_SIZE: ${BATCH_SIZE:-} BUILD_BATCH_SIZE: ${BUILD_BATCH_SIZE:-} NUMBER_OF_SHARDS: ${NUMBER_OF_SHARDS:-1} - APPROXIMATE_THRESHOLD: ${APPROXIMATE_THRESHOLD:-10000} + APPROXIMATE_THRESHOLD: ${APPROXIMATE_THRESHOLD:-} REFRESH_INTERVAL: ${REFRESH_INTERVAL:-} services: diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml index 0b9ce5407e..c5abf465b8 100644 --- a/deploy/docker-compose.yml +++ b/deploy/docker-compose.yml @@ -38,7 +38,7 @@ # BATCH_SIZE Optional cuvs-bench query batch size override # BUILD_BATCH_SIZE Optional OpenSearch bulk ingest batch size override # NUMBER_OF_SHARDS Number of primary index shards (default: 1) -# APPROXIMATE_THRESHOLD Vectors per segment before ANN build (default: 10000) +# APPROXIMATE_THRESHOLD Optional vectors per segment before ANN build # REFRESH_INTERVAL Optional index refresh interval (for example: 30s or -1) # CUVS_REPOSITORY Repository cloned into the benchmark image # CUVS_BRANCH Repository branch cloned into the benchmark image @@ -138,7 +138,7 @@ services: BATCH_SIZE: ${BATCH_SIZE:-} BUILD_BATCH_SIZE: ${BUILD_BATCH_SIZE:-} NUMBER_OF_SHARDS: ${NUMBER_OF_SHARDS:-1} - APPROXIMATE_THRESHOLD: ${APPROXIMATE_THRESHOLD:-10000} + APPROXIMATE_THRESHOLD: ${APPROXIMATE_THRESHOLD:-} REFRESH_INTERVAL: ${REFRESH_INTERVAL:-} volumes: - ${DATASET_PATH:-/tmp/datasets}:/data/datasets diff --git a/deploy/remote-index-build/run.py b/deploy/remote-index-build/run.py index 9dbde482c3..e9bab54c51 100644 --- a/deploy/remote-index-build/run.py +++ b/deploy/remote-index-build/run.py @@ -39,8 +39,11 @@ NUMBER_OF_SHARDS = int(os.environ.get("NUMBER_OF_SHARDS", "1")) if NUMBER_OF_SHARDS < 1: raise ValueError("NUMBER_OF_SHARDS must be at least 1") -APPROXIMATE_THRESHOLD = int(os.environ.get("APPROXIMATE_THRESHOLD", "10000")) -if APPROXIMATE_THRESHOLD < -1: +_approximate_threshold = os.environ.get("APPROXIMATE_THRESHOLD", "").strip() +APPROXIMATE_THRESHOLD = ( + int(_approximate_threshold) if _approximate_threshold else None +) +if APPROXIMATE_THRESHOLD is not None and APPROXIMATE_THRESHOLD < -1: raise ValueError("APPROXIMATE_THRESHOLD must be -1 or greater") session = requests.Session() @@ -102,8 +105,11 @@ def create_index() -> None: "index.knn.remote_index_build.enabled": True, "number_of_shards": NUMBER_OF_SHARDS, "number_of_replicas": 0, - "index.knn.advanced.approximate_threshold": APPROXIMATE_THRESHOLD, } + if APPROXIMATE_THRESHOLD is not None: + index_settings["index.knn.advanced.approximate_threshold"] = ( + APPROXIMATE_THRESHOLD + ) if REMOTE_BUILD_SIZE_MIN: index_settings["index.knn.remote_index_build.size.min"] = ( REMOTE_BUILD_SIZE_MIN From 855467d4523d2c7a2722cea457494705f04018bc Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Thu, 30 Jul 2026 12:39:52 -0500 Subject: [PATCH 21/22] Add force merge option Signed-off-by: James Bourbeau --- deploy/AWS_MULTINODE_SETUP.md | 1 + deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md | 1 + deploy/README.md | 5 +++ deploy/bench/configure_opensearch.py | 13 +++++++ deploy/docker-compose.multinode.yml | 1 + deploy/docker-compose.yml | 2 ++ .../cuvs_bench/backends/opensearch.py | 34 ++++++++++++++++--- 7 files changed, 53 insertions(+), 4 deletions(-) diff --git a/deploy/AWS_MULTINODE_SETUP.md b/deploy/AWS_MULTINODE_SETUP.md index f249b2c966..ae88805b99 100644 --- a/deploy/AWS_MULTINODE_SETUP.md +++ b/deploy/AWS_MULTINODE_SETUP.md @@ -338,6 +338,7 @@ export BUILD_BATCH_SIZE= export NUMBER_OF_SHARDS=1 export APPROXIMATE_THRESHOLD= export REFRESH_INTERVAL= +export FORCE_MERGE=false mkdir -p "${DATASET_PATH}" diff --git a/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md b/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md index 0fac09b54b..aa377f3fb3 100644 --- a/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md +++ b/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md @@ -328,6 +328,7 @@ export BUILD_BATCH_SIZE= export NUMBER_OF_SHARDS=1 export APPROXIMATE_THRESHOLD= export REFRESH_INTERVAL= +export FORCE_MERGE=false mkdir -p "${DATASET_PATH}" diff --git a/deploy/README.md b/deploy/README.md index 4f875eaf89..097f77f073 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -75,6 +75,7 @@ export BUILD_BATCH_SIZE= # optional bulk ingest batch size ove export NUMBER_OF_SHARDS=1 # number of primary index shards (default: 1) export APPROXIMATE_THRESHOLD= # optional vectors per segment before ANN build export REFRESH_INTERVAL= # optional index refresh interval (for example: 30s or -1) +export FORCE_MERGE=false # optionally merge each shard to one segment export REMOTE_BUILD_TIMEOUT=1800 # seconds to wait for remote builds (default: 1800) ``` @@ -90,6 +91,10 @@ Leave it empty to use the OpenSearch default, or set it to `-1` to disable automatic refreshes during ingestion. cuvs-bench performs an explicit refresh before searching. +Set `FORCE_MERGE=true` to merge every primary shard down to one segment after +ingestion and flush complete. The force-merge time is included in the reported +build time. It is disabled by default. + Start all services: ```bash diff --git a/deploy/bench/configure_opensearch.py b/deploy/bench/configure_opensearch.py index cff14a938e..99e4bcad27 100644 --- a/deploy/bench/configure_opensearch.py +++ b/deploy/bench/configure_opensearch.py @@ -13,6 +13,18 @@ def _optional_int(name: str) -> int | None: return int(value) if value else None +def _bool(name: str, default: bool = False) -> bool: + value = os.environ.get(name) + if value is None or not value.strip(): + return default + normalized = value.strip().lower() + if normalized in {"true", "1", "yes"}: + return True + if normalized in {"false", "0", "no"}: + return False + raise ValueError(f"{name} must be true or false") + + def create_backend_config() -> dict: remote_index_build = ( os.environ.get("REMOTE_INDEX_BUILD", "false").lower() == "true" @@ -33,6 +45,7 @@ def create_backend_config() -> dict: "verify_certs": False, "number_of_shards": number_of_shards, "remote_index_build": remote_index_build, + "force_merge": _bool("FORCE_MERGE"), } if approximate_threshold is not None: config["approximate_threshold"] = approximate_threshold diff --git a/deploy/docker-compose.multinode.yml b/deploy/docker-compose.multinode.yml index c4ff34b5c2..d6140acd19 100644 --- a/deploy/docker-compose.multinode.yml +++ b/deploy/docker-compose.multinode.yml @@ -53,6 +53,7 @@ x-bench-env: &bench-env NUMBER_OF_SHARDS: ${NUMBER_OF_SHARDS:-1} APPROXIMATE_THRESHOLD: ${APPROXIMATE_THRESHOLD:-} REFRESH_INTERVAL: ${REFRESH_INTERVAL:-} + FORCE_MERGE: ${FORCE_MERGE:-false} services: opensearch: diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml index c5abf465b8..49f4b63e10 100644 --- a/deploy/docker-compose.yml +++ b/deploy/docker-compose.yml @@ -40,6 +40,7 @@ # NUMBER_OF_SHARDS Number of primary index shards (default: 1) # APPROXIMATE_THRESHOLD Optional vectors per segment before ANN build # REFRESH_INTERVAL Optional index refresh interval (for example: 30s or -1) +# FORCE_MERGE Merge each shard to one segment after ingest (default: false) # CUVS_REPOSITORY Repository cloned into the benchmark image # CUVS_BRANCH Repository branch cloned into the benchmark image # @@ -140,6 +141,7 @@ services: NUMBER_OF_SHARDS: ${NUMBER_OF_SHARDS:-1} APPROXIMATE_THRESHOLD: ${APPROXIMATE_THRESHOLD:-} REFRESH_INTERVAL: ${REFRESH_INTERVAL:-} + FORCE_MERGE: ${FORCE_MERGE:-false} volumes: - ${DATASET_PATH:-/tmp/datasets}:/data/datasets restart: "no" diff --git a/python/cuvs_bench/cuvs_bench/backends/opensearch.py b/python/cuvs_bench/cuvs_bench/backends/opensearch.py index 1ae341fd57..49c1bc7be0 100644 --- a/python/cuvs_bench/cuvs_bench/backends/opensearch.py +++ b/python/cuvs_bench/cuvs_bench/backends/opensearch.py @@ -133,6 +133,7 @@ def _build_benchmark_configs( "build_batch_size", "approximate_threshold", "refresh_interval", + "force_merge", # Remote Index Build (OpenSearch 3.0+, faiss engine only) "remote_index_build", "remote_build_size_min", @@ -294,6 +295,8 @@ class OpenSearchBackend(BenchmarkBackend): - ``refresh_interval`` – how often OpenSearch refreshes the index, e.g. ``"1s"`` or ``"-1"`` to disable automatic refreshes during ingestion (default: OpenSearch's default). + - ``force_merge`` – merge each shard down to one segment after + ingestion and flush complete (default: ``False``). - ``requires_network`` – trigger network pre-flight check (default: ``True``) - ``remote_index_build`` – set ``index.knn.remote_index_build.enabled=true`` on the index at creation time, opting it into the GPU build path (default: ``False``). @@ -564,6 +567,23 @@ def _flush_index(self, index_name: str) -> None: f"Flush did not complete on all shards for {index_name}: {resp}" ) + def _force_merge_index(self, index_name: str) -> None: + resp = self._client.indices.forcemerge( + index=index_name, + max_num_segments=1, + request_timeout=None, + ) + shards = resp.get("_shards", {}) + total = shards.get("total", 0) + successful = shards.get("successful", 0) + failed = shards.get("failed", 0) + + if failed or successful != total: + raise RuntimeError( + f"Force merge did not complete on all shards for " + f"{index_name}: {resp}" + ) + def _resolve_index_name(self, index_cfg: IndexConfig) -> str: return self.config.get( "index_name", index_cfg.name.replace(".", "_").lower() @@ -701,9 +721,10 @@ def build( vectors. If the index already exists and ``force=False`` the build is skipped. - Build time measures ingest and flush. When ``remote_index_build=True`` - it also includes waiting for GPU build confirmation via the kNN stats - API. The final index refresh runs after build timing is recorded. + Build time measures ingest, flush, and the optional force merge. When + ``remote_index_build=True`` it also includes waiting for GPU build + confirmation via the kNN stats API. The final index refresh runs after + build timing is recorded. Parameters ---------- @@ -739,11 +760,13 @@ def build( remote_build_size_min = self.config.get("remote_build_size_min") approximate_threshold = self.config.get("approximate_threshold") refresh_interval = self.config.get("refresh_interval") + force_merge = bool(self.config.get("force_merge", False)) if dry_run: print( f"[dry_run] Would build OpenSearch index '{index_name}' " - f"(engine={engine}, remote_index_build={remote_index_build}, build_param={build_param})" + f"(engine={engine}, remote_index_build={remote_index_build}, " + f"force_merge={force_merge}, build_param={build_param})" ) return BuildResult( @@ -817,6 +840,8 @@ def build( initial_stats=pre_ingest_stats, timeout=remote_timeout, ) + if force_merge: + self._force_merge_index(index_name) build_time = time.perf_counter() - t0 self._client.indices.refresh(index=index_name, request_timeout=120) @@ -840,6 +865,7 @@ def build( "remote_index_build": remote_index_build, "approximate_threshold": approximate_threshold, "refresh_interval": refresh_interval, + "force_merge": force_merge, }, success=True, ) From 1c8684f9832ae9d5af4dd6e196e0dfee4d07e901 Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Thu, 30 Jul 2026 15:16:02 -0500 Subject: [PATCH 22/22] Use msearch-level request params Signed-off-by: James Bourbeau --- .../cuvs_bench/cuvs_bench/backends/opensearch.py | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/python/cuvs_bench/cuvs_bench/backends/opensearch.py b/python/cuvs_bench/cuvs_bench/backends/opensearch.py index 49c1bc7be0..9f44c821d2 100644 --- a/python/cuvs_bench/cuvs_bench/backends/opensearch.py +++ b/python/cuvs_bench/cuvs_bench/backends/opensearch.py @@ -885,9 +885,9 @@ def search( Search the OpenSearch k-NN index for nearest neighbors. Iterates over every search-parameter combination defined in the index - config, updating the index-level ``ef_search`` setting between runs. - Returns one result per parameter set so the orchestrator can compute - recall for each set independently. + config, passing ``ef_search`` directly in every k-NN query. Returns one + result per parameter set so the orchestrator can compute recall for + each set independently. Parameters ---------- @@ -979,12 +979,6 @@ def search( for sp in search_params_list: ef_search = sp.get("ef_search", 100) - if engine == "faiss": - self._client.indices.put_settings( - index=index_name, - body={"index.knn.algo_param.ef_search": ef_search}, - ) - neighbors = np.full((n_queries, k), -1, dtype=np.int64) distances = np.zeros((n_queries, k), dtype=np.float32) @@ -1002,6 +996,9 @@ def search( "vector": { "vector": q_vec.tolist(), "k": k, + "method_parameters": { + "ef_search": ef_search, + }, } } },