diff --git a/deploy/AWS_MULTINODE_SETUP.md b/deploy/AWS_MULTINODE_SETUP.md new file mode 100644 index 0000000000..ae88805b99 --- /dev/null +++ b/deploy/AWS_MULTINODE_SETUP.md @@ -0,0 +1,368 @@ +# 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=miracl-en-5m-1024d-fp32 +export BENCH_GROUPS=test +export K=10 +export BATCH_SIZE= +export BUILD_BATCH_SIZE= +export NUMBER_OF_SHARDS=1 +export APPROXIMATE_THRESHOLD= +export REFRESH_INTERVAL= +export FORCE_MERGE=false + +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..aa377f3fb3 --- /dev/null +++ b/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md @@ -0,0 +1,358 @@ +# 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=miracl-en-5m-1024d-fp32 +export BENCH_GROUPS=test +export K=10 +export BATCH_SIZE= +export BUILD_BATCH_SIZE= +export NUMBER_OF_SHARDS=1 +export APPROXIMATE_THRESHOLD= +export REFRESH_INTERVAL= +export FORCE_MERGE=false + +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 new file mode 100644 index 0000000000..9ec929aa22 --- /dev/null +++ b/deploy/DEPLOYMENT.md @@ -0,0 +1,201 @@ +# 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). 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 + +- **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 the required bucket name: + +```bash +export S3_BUCKET=opensearch-cuvs-bench +``` + +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: + +```bash +export AWS_DEFAULT_REGION=us-west-2 # default: us-west-2 +``` + +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 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/vector-repo \ + -H "Content-Type: application/json" \ + -d '{ + "type": "s3", + "settings": { + "bucket": "opensearch-cuvs-bench", + "base_path": "knn-indexes", + "region": "us-west-2" + } + }' +``` + +**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": "vector-repo", + "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, + "index.knn.remote_index_build.size.min": "1kb", + "index.knn.advanced.approximate_threshold": 10000, + "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. 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, 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: + +```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_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} \ + -e APPROXIMATE_THRESHOLD=${APPROXIMATE_THRESHOLD:-} \ + -v $(pwd)/remote-index-build:/app/remote-index-build \ + --no-deps bench \ + 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. + +## 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 principal used by OpenSearch and the builder needs `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 new file mode 100644 index 0000000000..097f77f073 --- /dev/null +++ b/deploy/README.md @@ -0,0 +1,274 @@ +# OpenSearch kNN Benchmark + +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 + +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 + → POSTs /_build to the remote-index-builder service + → service downloads vectors from S3 + → builds Faiss HNSW index on GPU + → uploads finished index back to S3 + → OpenSearch downloads the GPU-built index and merges it into the shard +``` + +## Services + +| Service | Image | Purpose | +|---|---|---| +| `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` | `python:3.11-slim` | Downloads the dataset, configures OpenSearch, runs the standard cuvs-bench CLI, and generates plots | + +## Requirements + +- **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 for staging vectors and built indexes + +## 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="$(pwd)/ann-benchmark-datasets" # directory containing dataset files +``` + +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-cuvs-bench # S3 bucket name +``` + +If you are using static credentials instead of a default provider, also export: + +```bash +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_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) +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= # 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) +``` + +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 +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 +# CPU build (no GPU required) +docker compose up --build + +# GPU build +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/` +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 +docker compose down -v +``` + +## What the bench container does + +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 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 + - Computes recall for each search-parameter set and writes the Python-backend results directly to the plotting CSV schema +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 + +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/ + miracl-en-5m-1024d-fp32/ + base.fbin + query.fbin + groundtruth.neighbors.ibin +``` + +## Key configuration + +**Cluster settings** (applied by `bench/configure_opensearch.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" + } +} +``` + +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 | +|---|---|---|---| +| `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 + +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 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 + +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 + +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 only) + +Requires a running OpenSearch node. S3 credentials and the GPU profile are not required for these tests. + +```bash +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 opensearch +``` + +### Remote index build integration tests (full GPU stack) + +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-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 \ + -e BUILDER_URL=http://remote-index-builder:1025 \ + -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 opensearch +``` + +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 + +| Port | Service | +|---|---| +| `9200` | OpenSearch REST API | +| `1025` | Remote index builder API | diff --git a/deploy/bench/Dockerfile b/deploy/bench/Dockerfile new file mode 100644 index 0000000000..e457fb2f15 --- /dev/null +++ b/deploy/bench/Dockerfile @@ -0,0 +1,41 @@ +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/* + +# 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_BRANCH}" "${CUVS_REPOSITORY}" /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" \ + pandas \ + pyyaml \ + requests \ + scikit-learn \ + scipy + +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 new file mode 100644 index 0000000000..99e4bcad27 --- /dev/null +++ b/deploy/bench/configure_opensearch.py @@ -0,0 +1,139 @@ +#!/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 _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" + ) + 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 = _optional_int("APPROXIMATE_THRESHOLD") + if approximate_threshold is not None and 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, + "remote_index_build": remote_index_build, + "force_merge": _bool("FORCE_MERGE"), + } + 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: + 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") + ) + 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 new file mode 100644 index 0000000000..3513e51c06 --- /dev/null +++ b/deploy/bench/entrypoint.sh @@ -0,0 +1,108 @@ +#!/bin/bash +set -e + +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 + 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}..." + 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." +} + +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 + # 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: 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: Configure OpenSearch and write the backend configuration. +python -u configure_opensearch.py "$BACKEND_CONFIG" + +# 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: 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 \ + --algorithms "$ALGORITHM" \ + --groups "$BENCH_GROUPS" \ + --count "$K" \ + --batch-size "$BATCH_SIZE" \ + --raw \ + --output-filepath /data/datasets diff --git a/deploy/bench/prepare_custom_dataset.py b/deploy/bench/prepare_custom_dataset.py new file mode 100644 index 0000000000..4b57eaf488 --- /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" + + +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/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 new file mode 100644 index 0000000000..d6140acd19 --- /dev/null +++ b/deploy/docker-compose.multinode.yml @@ -0,0 +1,154 @@ +# 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:-} + NUMBER_OF_SHARDS: ${NUMBER_OF_SHARDS:-1} + APPROXIMATE_THRESHOLD: ${APPROXIMATE_THRESHOLD:-} + REFRESH_INTERVAL: ${REFRESH_INTERVAL:-} + FORCE_MERGE: ${FORCE_MERGE:-false} + +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} + args: + CUVS_REPOSITORY: ${CUVS_REPOSITORY:-https://github.com/jrbourbeau/cuvs.git} + CUVS_BRANCH: ${CUVS_BRANCH:-deploy-opensearch-tmp} + 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 new file mode 100644 index 0000000000..49f4b63e10 --- /dev/null +++ b/deploy/docker-compose.yml @@ -0,0 +1,150 @@ +# OpenSearch GPU Remote Index Build — Docker Compose Demo +# +# Architecture: +# 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: +# - 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 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 +# +# Optional environment variables: +# 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-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) +# 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); 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 +# 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 +# +# Usage: +# CPU: docker compose up --build +# GPU: docker compose --profile gpu up --build +# +# Data flow: +# 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 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_SECRET_ACCESS_KEY: + AWS_SESSION_TOKEN: + AWS_DEFAULT_REGION: ${AWS_DEFAULT_REGION:-us-west-2} + +services: + + # ── OpenSearch ─────────────────────────────────────────────────────────────── + opensearch: + build: + context: ./opensearch + environment: + <<: *aws-env + OPENSEARCH_JAVA_OPTS: -Xms16g -Xmx16g + ulimits: + nofile: + soft: 65536 + hard: 65536 + volumes: + - opensearch-data:/usr/share/opensearch/data + - ./opensearch/opensearch.yml:/usr/share/opensearch/config/opensearch.yml:ro + ports: + - "9200:9200" + 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: + profiles: [gpu] + image: opensearchproject/remote-vector-index-builder:api-latest + environment: + <<: *aws-env + ports: + - "1025:1025" + 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 ──────────────────────────────────────────────────────────────── + 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 + environment: + <<: *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} + 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} + K: ${K:-10} + BATCH_SIZE: ${BATCH_SIZE:-} + BUILD_BATCH_SIZE: ${BUILD_BATCH_SIZE:-} + 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" + +volumes: + opensearch-data: diff --git a/deploy/opensearch/Dockerfile b/deploy/opensearch/Dockerfile new file mode 100644 index 0000000000..9ab8bfcdd8 --- /dev/null +++ b/deploy/opensearch/Dockerfile @@ -0,0 +1,13 @@ +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 S3. +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"] diff --git a/deploy/opensearch/entrypoint.sh b/deploy/opensearch/entrypoint.sh new file mode 100644 index 0000000000..4ea84342a8 --- /dev/null +++ b/deploy/opensearch/entrypoint.sh @@ -0,0 +1,31 @@ +#!/bin/bash +# +# 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. If static credentials are not provided, +# repository-s3 can fall back to the AWS default credential provider chain, such +# as an EC2 instance role. +# +# Static credential environment variables: +# AWS_ACCESS_KEY_ID AWS access key ID +# AWS_SECRET_ACCESS_KEY AWS secret access key +# AWS_SESSION_TOKEN STS session token (required for temporary credentials) +# +set -e + +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 +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 "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/opensearch/opensearch.yml b/deploy/opensearch/opensearch.yml new file mode 100644 index 0000000000..96a68a6bc0 --- /dev/null +++ b/deploy/opensearch/opensearch.yml @@ -0,0 +1,11 @@ +# 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 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 diff --git a/deploy/remote-index-build/Dockerfile b/deploy/remote-index-build/Dockerfile new file mode 100644 index 0000000000..d8b2485a7e --- /dev/null +++ b/deploy/remote-index-build/Dockerfile @@ -0,0 +1,5 @@ +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..e9bab54c51 --- /dev/null +++ b/deploy/remote-index-build/run.py @@ -0,0 +1,346 @@ +#!/usr/bin/env python3 +""" +OpenSearch GPU Remote Index Build — End-to-End Demo +==================================================== +Steps: + 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 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 +""" + +import json +import os +import sys +import time +from concurrent.futures import ThreadPoolExecutor, as_completed + +import boto3 +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") + +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. + +INDEX_NAME = "gpu-demo" +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")) +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 = 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() +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}'") + r = session.put( + f"{OPENSEARCH_URL}/_snapshot/{REPO_NAME}", + json={ + "type": "s3", + "settings": { + "bucket": S3_BUCKET, + "base_path": "knn-indexes", + "region": S3_REGION, + }, + }, + ) + 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") + + index_settings = { + "index.knn": True, + "index.knn.remote_index_build.enabled": True, + "number_of_shards": NUMBER_OF_SHARDS, + "number_of_replicas": 0, + } + 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 + ) + + r = session.put( + f"{OPENSEARCH_URL}/{INDEX_NAME}", + json={ + "settings": index_settings, + "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 " + f"{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}/_refresh") + r = session.get(f"{OPENSEARCH_URL}/{INDEX_NAME}/_count") + print(f" Document count after ingest: {r.json()['count']:,}") + + +# ── GPU build ───────────────────────────────────────────────────────────────── + +def trigger_gpu_build() -> None: + banner("Triggering GPU index build via flush") + print( + " OpenSearch will upload eligible flushed segments to S3, " + "then call the GPU builder." + ) + 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 = 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 + 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 S3 directly via boto3 instead. + + Exits with code 1 if no .faiss file appears within `timeout` seconds. + """ + 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") + + # 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 + while time.time() < deadline: + try: + 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( + " 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} " + f"({remaining}s left)" + ) + except Exception as e: + print(f" S3 check error: {e}") + time.sleep(5) + + 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( + " 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.") + 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-10 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: + 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) + 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" 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(timeout=REMOTE_BUILD_TIMEOUT) + search_vectors() + + print("\n" + "═" * 60) + print(" Demo complete!") + print("═" * 60) + print(f"\n OpenSearch is still running at {OPENSEARCH_URL}") + print(f" GPU builder : {BUILDER_URL}") + print() + + +if __name__ == "__main__": + main() 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 026fe98fd6..9f44c821d2 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 # @@ -131,6 +131,9 @@ def _build_benchmark_configs( "use_ssl", "verify_certs", "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", @@ -141,6 +144,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 +169,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" @@ -190,6 +201,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, @@ -264,6 +276,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"``) @@ -277,6 +290,13 @@ 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). + - ``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``). @@ -348,6 +368,8 @@ 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, + refresh_interval: Optional[str] = None, ) -> Dict[str, Any]: """ Construct the OpenSearch index mapping dict for k-NN. @@ -397,6 +419,17 @@ 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 refresh_interval is not None: + index_settings["refresh_interval"] = str(refresh_interval) if remote_index_build: if engine != "faiss": raise ValueError( @@ -534,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() @@ -543,11 +593,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, ) @@ -566,6 +617,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, ) @@ -666,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 ---------- @@ -702,11 +758,15 @@ 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") + 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( @@ -715,6 +775,7 @@ def build( index_size_bytes=0, algorithm=self.algo, build_params=build_param, + metadata={"group": self.config["group"]}, success=True, ) @@ -729,6 +790,10 @@ def build( index_size_bytes=0, algorithm=self.algo, build_params=build_param, + metadata={ + "group": self.config["group"], + "skipped": True, + }, success=True, ) @@ -747,12 +812,14 @@ def build( # Create index mapping = self._build_index_mapping( - dims, - engine, - space_type, - build_param, - remote_index_build, - remote_build_size_min, + 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) @@ -773,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) @@ -790,9 +859,13 @@ 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, + "approximate_threshold": approximate_threshold, + "refresh_interval": refresh_interval, + "force_merge": force_merge, }, success=True, ) @@ -807,21 +880,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. + 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 ---------- @@ -847,16 +913,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) @@ -870,45 +938,47 @@ 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) - 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) @@ -926,6 +996,9 @@ def search( "vector": { "vector": q_vec.tolist(), "k": k, + "method_parameters": { + "ef_search": ef_search, + }, } } }, @@ -957,39 +1030,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/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] 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..721d03c5ac --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_data_export.py @@ -0,0 +1,199 @@ +# +# 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 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.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_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 = BenchmarkOrchestrator(backend_type="opensearch") + 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 d16eedf25d..43f11dd701 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): @@ -157,14 +161,45 @@ 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, + refresh_interval="-1", ) 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 bc["refresh_interval"] == "-1" 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): @@ -176,11 +211,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"}) @@ -221,6 +301,69 @@ 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_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 = { @@ -326,7 +469,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 @@ -472,6 +615,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", @@ -503,12 +647,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 @@ -525,7 +668,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