diff --git a/deploy/AWS_MULTINODE_SETUP.md b/deploy/AWS_MULTINODE_SETUP.md new file mode 100644 index 0000000000..cff5d973e4 --- /dev/null +++ b/deploy/AWS_MULTINODE_SETUP.md @@ -0,0 +1,364 @@ +# Opensearch + cuVS multi-node benchmarking setup + +This guide sets up the three-node version of the OpenSearch GPU benchmark stack. + +```text +client node --> runs the benchmark submitter +opensearch node --> runs OpenSearch +builder GPU node --> runs the remote index build service +``` + +The examples use `us-west-2`, but that region is not required. Choose any AWS region where your needed instance types, GPU AMI, and quotas are available, then keep all resources in that same region. + +The goal is to keep the deployment simple while separating the three major components onto their own EC2 instances. Docker Compose still runs locally on each instance; it does not create a cross-host network. Cross-node communication uses EC2 private DNS names or private IPv4 addresses. + +## 1. Choose one AWS region + +In the AWS Console, set the region selector to your chosen region. The examples below use: + +```text +US West (Oregon) us-west-2 +``` + +Use this same region for S3, EC2, security groups, and the instances' IAM roles. The EC2 instances and security groups should also be in the same VPC so private DNS, private IPs, and security-group source rules work as expected. + +When running commands, set both AWS region variables from one value: + +```bash +export AWS_DEFAULT_REGION=us-west-2 +export AWS_REGION="$AWS_DEFAULT_REGION" +``` + +`AWS_REGION` is used when registering the OpenSearch S3 repository region. `AWS_DEFAULT_REGION` is used by AWS CLI and boto-style tooling. Setting both from the same value avoids accidental mismatches. + +## 2. Create the S3 bucket + +Go to **S3 > Create bucket**. + +Use: + +```text +Bucket name: globally unique name +Region: your chosen AWS region, for example US West (Oregon) us-west-2 +Object Ownership: ACLs disabled +Block Public Access: block all public access +Encryption: SSE-S3 is fine +Versioning: optional +``` + +This bucket stores the remote-build staging objects, including vectors and +generated index artifacts. Benchmark datasets and result plots stay on the +client node under `DATASET_PATH`. + +## 3. Create an IAM policy for S3 + +Go to **IAM > Policies > Create policy > JSON**. + +Use this policy: + +```json +{ + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": ["s3:ListBucket"], + "Resource": "arn:aws:s3:::opensearch-cuvs-bench" + }, + { + "Effect": "Allow", + "Action": ["s3:GetObject", "s3:PutObject", "s3:DeleteObject"], + "Resource": "arn:aws:s3:::opensearch-cuvs-bench/*" + } + ] +} +``` + +Name it: + +```text +opensearch-cuvs-bench-s3-policy +``` + +The delete permission lets the snapshot repository and staging workflow clean up +temporary objects when needed. + +## 4. Create the EC2 IAM role + +Go to **IAM > Roles > Create role**. + +Choose: + +```text +Trusted entity: AWS service +Use case: EC2 +``` + +Attach: + +```text +opensearch-cuvs-bench-s3-policy +AmazonSSMManagedInstanceCore +``` + +Name it: + +```text +opensearch-cuvs-bench-ec2-role +``` + +This role gives the instances refreshable S3 credentials and enables Session Manager access. Prefer this over fixed `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, and `AWS_SESSION_TOKEN` values. + +## 5. Create security groups + +Go to **EC2 > Security Groups > Create security group**. + +Create these three security groups in the same VPC: + +```text +sg-cuvs-client +sg-cuvs-opensearch +sg-cuvs-builder +``` + +Inbound rules for `sg-cuvs-client`: + +```text +No inbound rules needed +``` + +Inbound rules for `sg-cuvs-opensearch`: + +```text +TCP 9200 from sg-cuvs-client +TCP 9200 from sg-cuvs-builder +``` + +Inbound rules for `sg-cuvs-builder`: + +```text +TCP 1025 from sg-cuvs-opensearch +TCP 1025 from sg-cuvs-client +``` + +Leave outbound as the default `allow all`. Use Session Manager instead of SSH if possible. If you need SSH, add TCP `22` only from your own IP. + +## 6. Launch the OpenSearch node + +Go to **EC2 > Instances > Launch instance**. + +Use: + +```text +Name: opensearch-cuvs-db +AMI: Ubuntu 22.04 or Amazon Linux 2023 +Instance type: r7i.xlarge, r7i.2xlarge, or m7i.2xlarge +Security group: sg-cuvs-opensearch +IAM role: opensearch-cuvs-bench-ec2-role +Storage: 100+ GiB gp3, larger if needed +``` + +Under **Advanced details > Metadata options**, use: + +```text +IMDS endpoint: Enabled +IMDSv2: Required +Hop limit: 2 +``` + +After launch, copy the instance's **Private IPv4 DNS** or **Private IPv4 address**. You will use it as `OPENSEARCH_URL`. + +## 7. Launch the GPU builder node + +Launch a second EC2 instance: + +```text +Name: opensearch-cuvs-builder +AMI: AWS Deep Learning Base AMI with CUDA, Ubuntu 22.04 +Instance type: g5.xlarge or g6.xlarge +Security group: sg-cuvs-builder +IAM role: opensearch-cuvs-bench-ec2-role +Storage: 100+ GiB gp3 +``` + +Use the same metadata options: + +```text +IMDS endpoint: Enabled +IMDSv2: Required +Hop limit: 2 +``` + +After launch, copy the instance's **Private IPv4 DNS** or **Private IPv4 address**. You will use it as `REMOTE_INDEX_BUILDER_URL`. + +## 8. Launch the client node + +Launch the benchmark submitter instance: + +```text +Name: opensearch-cuvs-client +AMI: Ubuntu 22.04 or Amazon Linux 2023 +Instance type: c7i.xlarge or m7i.xlarge +Security group: sg-cuvs-client +IAM role: opensearch-cuvs-bench-ec2-role +Storage: 50-100 GiB gp3 +``` + +Use the same metadata options: + +```text +IMDS endpoint: Enabled +IMDSv2: Required +Hop limit: 2 +``` + +## 9. Install Docker and Compose on each node + +Connect to each instance with **EC2 > Instances > Connect > Session Manager**. + +Install Docker and Docker Compose if they are not already installed. Then verify: + +```bash +docker compose version +``` + +On the GPU builder node, also verify GPU access: + +```bash +nvidia-smi +docker run --rm --gpus all nvidia/cuda:12.9.0-base-ubuntu22.04 nvidia-smi +``` + +If you want to avoid `sudo docker`, add your login user to the `docker` group and reconnect: + +```bash +sudo groupadd docker 2>/dev/null || true +sudo usermod -aG docker "$(whoami)" +``` + +## 10. Copy the deployment checkout to each node + +On all three nodes, clone or copy the `deploy-opensearch-tmp` branch of the +`jrbourbeau/cuvs` fork: + +```bash +git clone --branch deploy-opensearch-tmp https://github.com/jrbourbeau/cuvs.git +cd cuvs +``` + +Make sure this file is present: + +```text +docker-compose.multinode.yml +``` + +## 11. Start OpenSearch + +On the OpenSearch node: + +```bash +export S3_BUCKET=opensearch-cuvs-bench +export AWS_DEFAULT_REGION=us-west-2 +export AWS_REGION="$AWS_DEFAULT_REGION" +docker compose -f docker-compose.multinode.yml --profile opensearch up -d +``` + +Verify: + +```bash +curl http://localhost:9200 +``` + +## 12. Start the remote index builder + +On the GPU builder node: + +```bash +export S3_BUCKET=opensearch-cuvs-bench +export AWS_DEFAULT_REGION=us-west-2 +export AWS_REGION="$AWS_DEFAULT_REGION" +docker compose -f docker-compose.multinode.yml --profile builder up -d +``` + +Verify the container is running: + +```bash +docker ps +``` + +From the OpenSearch node, verify that the builder is reachable. OpenSearch is +the service that calls the remote builder: + +```bash +python3 -c 'import socket; socket.create_connection(("BUILDER_PRIVATE_DNS_OR_IP", 1025), 5).close(); print("builder reachable")' +``` + +From the client node, run the same check. The benchmark container also waits +for the builder before starting a remote-build run: + +```bash +python3 -c 'import socket; socket.create_connection(("BUILDER_PRIVATE_DNS_OR_IP", 1025), 5).close(); print("builder reachable")' +``` + +## 13. Configure OpenSearch from the client node + +On the client node: + +```bash +export OPENSEARCH_HOST=OPENSEARCH_PRIVATE_DNS_OR_IP +export OPENSEARCH_URL=http://${OPENSEARCH_HOST}:9200 +export REMOTE_INDEX_BUILDER_URL=http://BUILDER_PRIVATE_DNS_OR_IP:1025 +export S3_BUCKET=opensearch-cuvs-bench +export AWS_DEFAULT_REGION=us-west-2 +export AWS_REGION="$AWS_DEFAULT_REGION" +export S3_PREFIX=knn-indexes +export REMOTE_VECTOR_REPOSITORY=vector-repo +``` + +Then run the one-shot configure profile: + +```bash +docker compose -f docker-compose.multinode.yml --profile configure run --rm configure-remote-index-build +``` + +This registers the S3-backed remote vector repository and tells OpenSearch where the remote index builder service lives. + +## 14. Run the benchmark + +Still on the client node: + +```bash +export REMOTE_INDEX_BUILD=true +export DATASET_PATH="$(pwd)/opensearch-cuvs-datasets" +export DATASET=sift-128-euclidean +export BENCH_GROUPS=test +export K=10 +export BATCH_SIZE= +export BUILD_BATCH_SIZE= + +mkdir -p "${DATASET_PATH}" + +docker compose -f docker-compose.multinode.yml --profile client build bench +docker compose -f docker-compose.multinode.yml --profile client run --rm bench +``` + +## 15. Debug checklist + +From the client node, these should all work: + +```bash +curl "$OPENSEARCH_URL" +python3 -c 'import os, socket; from urllib.parse import urlparse; url = urlparse(os.environ["REMOTE_INDEX_BUILDER_URL"]); socket.create_connection((url.hostname, url.port or 1025), 5).close(); print("builder reachable")' +aws s3 ls s3://$S3_BUCKET --region "$AWS_DEFAULT_REGION" +``` + +If S3 fails, check the IAM role and S3 policy. If OpenSearch or builder +connectivity fails, check private DNS/IP values and security group source rules. + +## Notes + +- Keep the S3 bucket and EC2 resources in the same AWS region. The examples use `us-west-2`, but the setup is not region-specific. +- Within that region, keep all three instances and security groups in the same VPC. Prefer the same Availability Zone for the first benchmark run. +- Use private DNS or private IPs, not public IPs, for cross-node service traffic. +- Docker Compose networks are local to one host; Compose service names do not resolve across EC2 instances. +- The EC2 IAM role credentials rotate automatically. Avoid freezing temporary credentials into `.env` unless the application absolutely requires literal `AWS_*` variables. diff --git a/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md b/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md new file mode 100644 index 0000000000..c9d555bde1 --- /dev/null +++ b/deploy/AWS_MULTINODE_SINGLE_ROLE_SETUP.md @@ -0,0 +1,354 @@ +# AWS single-role multinode setup + +This guide sets up the three-node version of the OpenSearch GPU benchmark stack in `us-west-2`: + +```text +client node runs the benchmark submitter +opensearch node runs OpenSearch +builder GPU node runs the remote index build service +``` + +The goal is to keep the deployment simple while separating the three major components onto their own EC2 instances. Docker Compose still runs locally on each instance; it does not create a cross-host network. Cross-node communication uses EC2 private DNS names or private IPv4 addresses. + +## 1. Set the AWS region + +In the AWS Console, set the region selector to: + +```text +US West (Oregon) us-west-2 +``` + +Use this same region for S3, EC2, security groups, and the instances' IAM roles. + +## 2. Create the S3 bucket + +Go to **S3 > Create bucket**. + +Use: + +```text +Bucket name: globally unique name +Region: US West (Oregon) us-west-2 +Object Ownership: ACLs disabled +Block Public Access: block all public access +Encryption: SSE-S3 is fine +Versioning: optional +``` + +This bucket stores the remote-build staging objects, including vectors and +generated index artifacts. Benchmark datasets and result plots stay on the +client node under `DATASET_PATH`. + +## 3. Create an IAM policy for S3 + +Go to **IAM > Policies > Create policy > JSON**. + +Use this policy: + +```json +{ + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": ["s3:ListBucket"], + "Resource": "arn:aws:s3:::opensearch-cuvs-bench" + }, + { + "Effect": "Allow", + "Action": ["s3:GetObject", "s3:PutObject", "s3:DeleteObject"], + "Resource": "arn:aws:s3:::opensearch-cuvs-bench/*" + } + ] +} +``` + +Name it: + +```text +opensearch-cuvs-bench-s3-policy +``` + +The delete permission lets the snapshot repository and staging workflow clean up +temporary objects when needed. + +## 4. Create the EC2 IAM role + +Go to **IAM > Roles > Create role**. + +Choose: + +```text +Trusted entity: AWS service +Use case: EC2 +``` + +Attach: + +```text +opensearch-cuvs-bench-s3-policy +AmazonSSMManagedInstanceCore +``` + +Name it: + +```text +opensearch-cuvs-bench-ec2-role +``` + +This role gives the instances refreshable S3 credentials and enables Session Manager access. Prefer this over fixed `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, and `AWS_SESSION_TOKEN` values. + +## 5. Create security groups + +Go to **EC2 > Security Groups > Create security group**. + +Create these three security groups in the same VPC: + +```text +sg-cuvs-client +sg-cuvs-opensearch +sg-cuvs-builder +``` + +Inbound rules for `sg-cuvs-client`: + +```text +No inbound rules needed +``` + +Inbound rules for `sg-cuvs-opensearch`: + +```text +TCP 9200 from sg-cuvs-client +TCP 9200 from sg-cuvs-builder +``` + +Inbound rules for `sg-cuvs-builder`: + +```text +TCP 1025 from sg-cuvs-opensearch +TCP 1025 from sg-cuvs-client +``` + +Leave outbound as the default `allow all`. Use Session Manager instead of SSH if possible. If you need SSH, add TCP `22` only from your own IP. + +## 6. Launch the OpenSearch node + +Go to **EC2 > Instances > Launch instance**. + +Use: + +```text +Name: opensearch-cuvs-db +AMI: Ubuntu 22.04 or Amazon Linux 2023 +Instance type: r7i.xlarge, r7i.2xlarge, or m7i.2xlarge +Security group: sg-cuvs-opensearch +IAM role: opensearch-cuvs-bench-ec2-role +Storage: 100+ GiB gp3, larger if needed +``` + +Under **Advanced details > Metadata options**, use: + +```text +IMDS endpoint: Enabled +IMDSv2: Required +Hop limit: 2 +``` + +After launch, copy the instance's **Private IPv4 DNS** or **Private IPv4 address**. You will use it as `OPENSEARCH_URL`. + +## 7. Launch the GPU builder node + +Launch a second EC2 instance: + +```text +Name: opensearch-cuvs-builder +AMI: AWS Deep Learning Base AMI with CUDA, Ubuntu 22.04 +Instance type: g5.xlarge or g6.xlarge +Security group: sg-cuvs-builder +IAM role: opensearch-cuvs-bench-ec2-role +Storage: 100+ GiB gp3 +``` + +Use the same metadata options: + +```text +IMDS endpoint: Enabled +IMDSv2: Required +Hop limit: 2 +``` + +After launch, copy the instance's **Private IPv4 DNS** or **Private IPv4 address**. You will use it as `REMOTE_INDEX_BUILDER_URL`. + +## 8. Launch the client node + +Launch the benchmark submitter instance: + +```text +Name: opensearch-cuvs-client +AMI: Ubuntu 22.04 or Amazon Linux 2023 +Instance type: c7i.xlarge or m7i.xlarge +Security group: sg-cuvs-client +IAM role: opensearch-cuvs-bench-ec2-role +Storage: 50-100 GiB gp3 +``` + +Use the same metadata options: + +```text +IMDS endpoint: Enabled +IMDSv2: Required +Hop limit: 2 +``` + +## 9. Install Docker and Compose on each node + +Connect to each instance with **EC2 > Instances > Connect > Session Manager**. + +Install Docker and Docker Compose if they are not already installed. Then verify: + +```bash +docker compose version +``` + +On the GPU builder node, also verify GPU access: + +```bash +nvidia-smi +docker run --rm --gpus all nvidia/cuda:12.9.0-base-ubuntu22.04 nvidia-smi +``` + +If you want to avoid `sudo docker`, add your login user to the `docker` group and reconnect: + +```bash +sudo groupadd docker 2>/dev/null || true +sudo usermod -aG docker "$(whoami)" +``` + +## 10. Copy the deployment checkout to each node + +On all three nodes, clone or copy the `deploy-opensearch-tmp` branch of the +`jrbourbeau/cuvs` fork: + +```bash +git clone --branch deploy-opensearch-tmp https://github.com/jrbourbeau/cuvs.git +cd cuvs +``` + +Make sure this file is present: + +```text +docker-compose.multinode.yml +``` + +## 11. Start OpenSearch + +On the OpenSearch node: + +```bash +export S3_BUCKET=opensearch-cuvs-bench +export AWS_REGION=us-west-2 +export AWS_DEFAULT_REGION=us-west-2 +docker compose -f docker-compose.multinode.yml --profile opensearch up -d +``` + +Verify: + +```bash +curl http://localhost:9200 +``` + +## 12. Start the remote index builder + +On the GPU builder node: + +```bash +export S3_BUCKET=opensearch-cuvs-bench +export AWS_REGION=us-west-2 +export AWS_DEFAULT_REGION=us-west-2 +docker compose -f docker-compose.multinode.yml --profile builder up -d +``` + +Verify the container is running: + +```bash +docker ps +``` + +From the OpenSearch node, verify that the builder is reachable. OpenSearch is +the service that calls the remote builder: + +```bash +python3 -c 'import socket; socket.create_connection(("BUILDER_PRIVATE_DNS_OR_IP", 1025), 5).close(); print("builder reachable")' +``` + +From the client node, run the same check. The benchmark container also waits +for the builder before starting a remote-build run: + +```bash +python3 -c 'import socket; socket.create_connection(("BUILDER_PRIVATE_DNS_OR_IP", 1025), 5).close(); print("builder reachable")' +``` + +## 13. Configure OpenSearch from the client node + +On the client node: + +```bash +export S3_BUCKET=opensearch-cuvs-bench +export OPENSEARCH_HOST=OPENSEARCH_PRIVATE_DNS_OR_IP +export OPENSEARCH_PORT=9200 +export OPENSEARCH_URL=http://${OPENSEARCH_HOST}:${OPENSEARCH_PORT} +export REMOTE_INDEX_BUILDER_URL=http://BUILDER_PRIVATE_DNS_OR_IP:1025 +export AWS_REGION=us-west-2 +export AWS_DEFAULT_REGION=us-west-2 +export S3_PREFIX=knn-indexes +export REMOTE_VECTOR_REPOSITORY=vector-repo +``` + +Then run the one-shot configure profile: + +```bash +docker compose -f docker-compose.multinode.yml --profile configure run --rm configure-remote-index-build +``` + +This registers the S3-backed remote vector repository and tells OpenSearch where the remote index builder service lives. + +## 14. Run the benchmark + +Still on the client node: + +```bash +export REMOTE_INDEX_BUILD=true +export DATASET_PATH=/data/opensearch-cuvs-datasets +export DATASET=sift-128-euclidean +export BENCH_GROUPS=test +export K=10 +export BATCH_SIZE= +export BUILD_BATCH_SIZE= + +mkdir -p "${DATASET_PATH}" + +docker compose -f docker-compose.multinode.yml --profile client build bench +docker compose -f docker-compose.multinode.yml --profile client run --rm bench +``` + +## 15. Debug checklist + +From the client node, these should all work: + +```bash +curl "$OPENSEARCH_URL" +python3 -c 'import os, socket; from urllib.parse import urlparse; url = urlparse(os.environ["REMOTE_INDEX_BUILDER_URL"]); socket.create_connection((url.hostname, url.port or 1025), 5).close(); print("builder reachable")' +aws s3 ls s3://$S3_BUCKET --region us-west-2 +``` + +If S3 fails, check the IAM role and S3 policy. If OpenSearch or builder +connectivity fails, check private DNS/IP values and security group source rules. + +## Notes + +- Keep all three instances in `us-west-2`. +- Prefer the same VPC and same Availability Zone for the first benchmark run. +- Use private DNS or private IPs, not public IPs, for cross-node service traffic. +- Docker Compose networks are local to one host; Compose service names do not resolve across EC2 instances. +- The EC2 IAM role credentials rotate automatically. Avoid freezing temporary credentials into `.env` unless the application absolutely requires literal `AWS_*` variables. diff --git a/deploy/DEPLOYMENT.md b/deploy/DEPLOYMENT.md new file mode 100644 index 0000000000..4084eb4cbf --- /dev/null +++ b/deploy/DEPLOYMENT.md @@ -0,0 +1,198 @@ +# 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", + "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} \ + -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..ae516a8e03 --- /dev/null +++ b/deploy/README.md @@ -0,0 +1,228 @@ +# 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` | custom Python | Downloads dataset, registers repo + cluster settings, runs cuvs-bench build/search benchmark, exports results, 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 REMOTE_BUILD_TIMEOUT=1800 # seconds to wait for remote builds (default: 1800) +``` + +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`. + +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 `cuvs-bench` build phase (handled entirely by the OpenSearch backend): + - Creates the kNN index and bulk-ingests dataset vectors + - **GPU mode**: Flushes segments, waits for all submitted remote GPU builds to complete, and polls the kNN stats API every 5 s until the build is confirmed complete + - **CPU mode**: Flushes and refreshes the local OpenSearch index + - Uses the backend's automatic OpenSearch bulk-ingest batch sizing by default; set `BUILD_BATCH_SIZE` to override it + - Records total build time in the result +5. Runs `cuvs-bench` search phase and prints a recall/QPS/latency table +6. Exports benchmark JSON results to CSV +7. Generates recall vs. latency/throughput plots as PNGs in `$DATASET_PATH` (`cuvs_bench.plot`) + +## Dataset format + +cuvs-bench reads binary vector files with a simple header: + +``` +[4 bytes: n_rows as uint32] +[4 bytes: n_cols as uint32] +[n_rows × n_cols × itemsize bytes: vector data] +``` + +Supported extensions: `.fbin` (float32), `.f16bin` (float16), `.u8bin` (uint8), `.i8bin` (int8). + +`DATASET_PATH` should be a directory where each dataset lives in its own subdirectory named after the dataset, e.g.: + +``` +$DATASET_PATH/ + sift-128-euclidean/ + base.fbin + query.fbin + groundtruth.neighbors.ibin +``` + +## Key configuration + +**Cluster settings** (applied by `bench/run.py`): + +```json +{ + "persistent": { + "knn.remote_index_build.enabled": true, + "knn.remote_index_build.repository": "vector-repo", + "knn.remote_index_build.service.endpoint": "http://remote-index-builder:1025" + } +} +``` + +**Parameter groups** (`BENCH_GROUPS`): + +| Group | Build params | Search params | Use case | +|---|---|---|---| +| `test` | 1 combo (m=16) | ef_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 integration +``` + +### 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 integration +``` + +This lets the pytest `integration` marker decide which tests run. With only OpenSearch running, remote-build tests skip because the GPU builder and S3 environment are unavailable. With the GPU stack running, the same marker includes the remote-build coverage. + +## Ports + +| 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..20adbca588 --- /dev/null +++ b/deploy/bench/Dockerfile @@ -0,0 +1,37 @@ +FROM python:3.11-slim +WORKDIR /app + +RUN apt-get update \ + && apt-get install -y --no-install-recommends git \ + && rm -rf /var/lib/apt/lists/* + +# Sparse-clone just the cuvs_bench Python package. +# Installing via pip is not possible without CUDA (rapids-build-backend requires +# nvcc, and cuvs_bench declares a hard dep on the `cuvs` CUDA package). +# The opensearch backend is pure Python and needs neither, so we add the +# package directly to PYTHONPATH and install only the actual runtime deps. +RUN git clone --depth=1 --filter=blob:none --sparse https://github.com/rapidsai/cuvs.git /opt/cuvs \ + && cd /opt/cuvs \ + && git sparse-checkout set python/cuvs_bench + +ENV PYTHONPATH=/opt/cuvs/python/cuvs_bench + +# Runtime dependencies from cuvs_bench/pyproject.toml (excluding `cuvs` itself) +# plus extras needed by the opensearch backend and this benchmark script. +RUN pip install --no-cache-dir \ + boto3 \ + pytest \ + botocore \ + click \ + h5py \ + matplotlib \ + "numpy<2" \ + "opensearch-py>=2.4.0" \ + pandas \ + pyyaml \ + requests \ + scikit-learn \ + scipy + +COPY --chmod=755 run.py entrypoint.sh . +CMD ["/app/entrypoint.sh"] diff --git a/deploy/bench/entrypoint.sh b/deploy/bench/entrypoint.sh new file mode 100644 index 0000000000..78f050f9e4 --- /dev/null +++ b/deploy/bench/entrypoint.sh @@ -0,0 +1,74 @@ +#!/bin/bash +set -e + +DATASET="${DATASET:-sift-128-euclidean}" +BENCH_GROUPS="${BENCH_GROUPS:-test}" +K="${K:-10}" +ALGORITHM="opensearch_faiss_hnsw" + +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: Download dataset (skipped automatically if already present) +python -m cuvs_bench.get_dataset \ + --dataset "$DATASET" \ + --dataset-path /data/datasets + +# Step 2: Run benchmark (build + search + writes result JSON files) +python -u run.py + +# Step 3: Export JSON → CSV (required by cuvs_bench.plot) +# --batch-size is ignored when --data-export is set, but Click prompts for it +# before entering main(), so pass a dummy value to keep the container non-interactive. +python -m cuvs_bench.run --data-export \ + --dataset "$DATASET" \ + --dataset-path /data/datasets \ + --algorithms "$ALGORITHM" \ + --groups "$BENCH_GROUPS" \ + --count "$K" \ + --batch-size 1 \ + --search-mode latency + +# Step 4: Plot — PNGs written to /data/datasets (mounted from host $DATASET_PATH) +python -m cuvs_bench.plot \ + --dataset "$DATASET" \ + --dataset-path /data/datasets \ + --algorithms "$ALGORITHM" \ + --groups "$BENCH_GROUPS" \ + --count "$K" \ + --output-filepath /data/datasets diff --git a/deploy/bench/run.py b/deploy/bench/run.py new file mode 100644 index 0000000000..55bdc1777f --- /dev/null +++ b/deploy/bench/run.py @@ -0,0 +1,385 @@ +#!/usr/bin/env python3 +""" +OpenSearch GPU Remote Index Build Benchmark +=========================================== +Steps: + 1. Register S3 snapshot repository with OpenSearch + 2. Configure cluster settings for GPU remote index build + 3. Build kNN index via cuvs-bench: + a. Bulk-ingest dataset vectors + b. Flush segments to kick off remote GPU builds when enabled + c. Poll kNN stats until every submitted remote build completes + 4. Run cuvs-bench search benchmarks and print results + 5. Write gbench-compatible JSON result files for CSV export and plotting +""" + +import json +import os +import sys + +import requests +from cuvs_bench.orchestrator import BenchmarkOrchestrator +from cuvs_bench.backends.base import BuildResult, SearchResult + + +OPENSEARCH_URL = os.environ.get("OPENSEARCH_URL", "http://opensearch:9200") +OPENSEARCH_HOST = os.environ.get("OPENSEARCH_HOST", "opensearch") +OPENSEARCH_PORT = int(os.environ.get("OPENSEARCH_PORT", "9200")) +BUILDER_URL = os.environ.get("BUILDER_URL", "http://remote-index-builder:1025") + +REMOTE_INDEX_BUILD = os.environ.get("REMOTE_INDEX_BUILD", "false").lower() == "true" +REMOTE_BUILD_SIZE_MIN = os.environ.get("REMOTE_BUILD_SIZE_MIN", "").strip() +REMOTE_BUILD_TIMEOUT = int(os.environ.get("REMOTE_BUILD_TIMEOUT", "1800")) + +S3_BUCKET = os.environ.get("S3_BUCKET", "").strip() +S3_PREFIX = os.environ.get("S3_PREFIX", "knn-indexes").strip() or "knn-indexes" +S3_REGION = os.environ.get("AWS_DEFAULT_REGION", "us-west-2") + +DATASET = os.environ.get("DATASET", "sift-128-euclidean") +DATASET_PATH = os.environ.get("DATASET_PATH", "/data/datasets") +BENCH_GROUPS = os.environ.get("BENCH_GROUPS", "test") +K = int(os.environ.get("K", "10")) +BATCH_SIZE = os.environ.get("BATCH_SIZE", "").strip() +BATCH_SIZE = int(BATCH_SIZE) if BATCH_SIZE else None +BUILD_BATCH_SIZE = os.environ.get("BUILD_BATCH_SIZE", "").strip() +BUILD_BATCH_SIZE = int(BUILD_BATCH_SIZE) if BUILD_BATCH_SIZE else None + +ALGORITHM = "opensearch_faiss_hnsw" +REPO_NAME = os.environ.get("REMOTE_VECTOR_REPOSITORY", "vector-repo").strip() +REPO_NAME = REPO_NAME or "vector-repo" + +session = requests.Session() +session.headers.update({"Content-Type": "application/json"}) + + +# ── helpers ─────────────────────────────────────────────────────────────────── + +def banner(msg: str) -> None: + print(f"\n{'─'*60}\n {msg}\n{'─'*60}") + + +def _recall_for_entry( + result: SearchResult, entry_index: int, entry_count: int +) -> float | None: + # cuVS computes recall in the orchestrator from SearchResult.neighbors. The + # OpenSearch backend returns neighbors for the final search-parameter run. + if entry_index == entry_count - 1: + return float(result.recall) + return None + + +def _get_search_batch_size(search_results: list) -> int | None: + if BATCH_SIZE is not None: + return BATCH_SIZE + for result in search_results: + if not isinstance(result, SearchResult): + continue + batch_size = (result.metadata or {}).get("batch_size") + if batch_size is not None: + return int(batch_size) + return None + + +# ── OpenSearch setup ────────────────────────────────────────────────────────── + +def register_repository() -> None: + banner(f"Registering S3 repository '{REPO_NAME}'") + r = session.put( + f"{OPENSEARCH_URL}/_snapshot/{REPO_NAME}", + json={ + "type": "s3", + "settings": { + "bucket": S3_BUCKET, + "base_path": S3_PREFIX, + "region": S3_REGION, + }, + }, + ) + r.raise_for_status() + print(f" {r.json()}") + + +def configure_cluster() -> None: + if REMOTE_INDEX_BUILD: + banner("Enabling GPU remote index build (cluster settings)") + settings = { + "knn.remote_index_build.enabled": True, + "knn.remote_index_build.repository": REPO_NAME, + "knn.remote_index_build.service.endpoint": BUILDER_URL, + } + else: + banner("Disabling GPU remote index build (cluster settings)") + settings = { + "knn.remote_index_build.enabled": False, + } + r = session.put( + f"{OPENSEARCH_URL}/_cluster/settings", + json={"persistent": settings}, + ) + r.raise_for_status() + print(f" {r.json()}") + + +# ── result files ───────────────────────────────────────────────────────────── + +def write_result_files( + build_results: list, + search_results: list, + dataset: str, + dataset_path: str, + algo: str, + groups: str, + k: int, + batch_size: int, +) -> None: + """Write gbench-compatible JSON result files. + + Creates files under //result/{build,search}/ in the + same format the C++ backend produces, so the cuvs-bench CSV exporters and + ``cuvs_bench.plot`` work without modification. + """ + build_dir = os.path.join(dataset_path, dataset, "result", "build") + search_dir = os.path.join(dataset_path, dataset, "result", "search") + os.makedirs(build_dir, exist_ok=True) + os.makedirs(search_dir, exist_ok=True) + + # Build JSON – one record per successfully built index. + # data_export.py assumes the build CSV has columns: + # [algo_name, index_name, time, threads, cpu_time, ...] + # so "threads" and "cpu_time" must be present in the JSON (they are not in + # skip_build_cols and therefore get included as columns 3 and 4). + build_benchmarks = [ + { + "name": r.index_path, + "real_time": r.build_time_seconds, + "time_unit": "s", + "threads": 1, + "cpu_time": r.build_time_seconds, + } + for r in build_results + if isinstance(r, BuildResult) and r.success and r.index_path + ] + + # Search JSON – zip build + search results to recover the index name, then + # expand per-search-param entries so each (index, ef_search) is one record. + build_list = [r for r in build_results if isinstance(r, BuildResult)] + search_list = [r for r in search_results if isinstance(r, SearchResult)] + search_benchmarks = [] + skipped_without_recall = 0 + for build_r, search_r in zip(build_list, search_list): + if not search_r.success or not build_r.index_path: + continue + per_param = (search_r.metadata or {}).get("per_search_param_results", []) + for entry_index, entry in enumerate(per_param): + recall = _recall_for_entry(search_r, entry_index, len(per_param)) + if recall is None: + skipped_without_recall += 1 + continue + latency_ms = float(entry["search_time_ms"]) + search_benchmarks.append( + { + "name": build_r.index_path, + "real_time": latency_ms, + "time_unit": "ms", + "Recall": recall, + "items_per_second": float(entry["queries_per_second"]), + # Latency field expected by data_export in seconds + "Latency": latency_ms / 1000.0, + } + ) + + build_file = os.path.join(build_dir, f"{algo},{groups}.json") + search_file = os.path.join( + search_dir, f"{algo},{groups},k{k},bs{batch_size}.json" + ) + + with open(build_file, "w") as fh: + json.dump({"benchmarks": build_benchmarks}, fh, indent=2) + with open(search_file, "w") as fh: + json.dump({"benchmarks": search_benchmarks}, fh, indent=2) + + print(f"\n Result files written:") + print(f" {build_file}") + print(f" {search_file}") + if skipped_without_recall: + print( + " skipped " + f"{skipped_without_recall} search rows without per-parameter recall" + ) + + +# ── results ─────────────────────────────────────────────────────────────────── + +def _print_result_row( + params: dict, recall: float | None, qps: float | None, latency_ms: float | None +) -> None: + params_str = ", ".join(f"{k}={v}" for k, v in params.items()) + recall_str = "n/a" if recall is None else f"{recall:.4f}" + qps_str = "n/a" if qps is None else f"{qps:.1f}" + latency_str = "n/a" if latency_ms is None else f"{latency_ms:.2f}" + print(f" {params_str:<40} {recall_str:<12} {qps_str:>8} {latency_str:>12}") + + +def print_results(results: list) -> None: + banner("Benchmark Results") + search_results = [r for r in results if isinstance(r, SearchResult)] + if not search_results: + print(" No search results returned.") + return + + header = f" {'params':<40} {'recall@'+str(K):<12} {'QPS':>8} {'latency (ms)':>12}" + print(header) + print(" " + "─" * (len(header) - 2)) + missing_recall_rows = 0 + for r in search_results: + per_param = (r.metadata or {}).get("per_search_param_results", []) + entry_count = len(per_param) + for entry_index, entry in enumerate(per_param): + recall = _recall_for_entry(r, entry_index, entry_count) + if recall is None: + missing_recall_rows += 1 + continue + _print_result_row( + entry["search_params"], + recall, + float(entry["queries_per_second"]), + float(entry["search_time_ms"]), + ) + if missing_recall_rows: + print( + "\n Omitted " + f"{missing_recall_rows} timing-only search rows without recall." + ) + + +# ── entrypoint ──────────────────────────────────────────────────────────────── + +def main() -> None: + if REMOTE_INDEX_BUILD and not S3_BUCKET: + print( + "ERROR: S3_BUCKET is not set. Remote index build requires an S3 " + "bucket for vector and index staging. Set it before starting the " + "stack, for example: export S3_BUCKET=", + file=sys.stderr, + ) + sys.exit(1) + + print("\n" + "═" * 60) + print(" OpenSearch kNN Benchmark") + print("═" * 60) + print(f" OpenSearch : {OPENSEARCH_URL}") + print(f" Remote index build : {REMOTE_INDEX_BUILD}") + if REMOTE_INDEX_BUILD: + print(f" GPU builder : {BUILDER_URL}") + print(f" S3 bucket : s3://{S3_BUCKET}/{S3_PREFIX}/ (region: {S3_REGION})") + print(f" Repository : {REPO_NAME}") + print(f" Build size minimum : {REMOTE_BUILD_SIZE_MIN or 'OpenSearch default'}") + print(f" Build timeout : {REMOTE_BUILD_TIMEOUT}s") + print(f" Dataset : {DATASET} (path: {DATASET_PATH})") + print(f" Algorithm : {ALGORITHM}") + print(f" Groups : {BENCH_GROUPS} k={K}") + print( + " Search batch size : " + f"{BATCH_SIZE if BATCH_SIZE is not None else 'backend default'}" + ) + print( + " Build batch size : " + f"{BUILD_BATCH_SIZE if BUILD_BATCH_SIZE is not None else 'backend auto'}" + ) + + if REMOTE_INDEX_BUILD: + register_repository() + configure_cluster() + + orchestrator = BenchmarkOrchestrator(backend_type="opensearch") + + # Shared kwargs for both build and search phases. + common_kwargs = dict( + dataset=DATASET, + dataset_path=DATASET_PATH, + algorithms=ALGORITHM, + groups=BENCH_GROUPS, + host=OPENSEARCH_HOST, + port=OPENSEARCH_PORT, + use_ssl=False, + verify_certs=False, + ) + + build_kwargs = dict( + common_kwargs, + remote_index_build=REMOTE_INDEX_BUILD, + ) + if BUILD_BATCH_SIZE is not None: + build_kwargs["build_batch_size"] = BUILD_BATCH_SIZE + if REMOTE_INDEX_BUILD: + build_kwargs["remote_build_timeout"] = REMOTE_BUILD_TIMEOUT + if REMOTE_BUILD_SIZE_MIN: + build_kwargs["remote_build_size_min"] = REMOTE_BUILD_SIZE_MIN + + # ── Build phase ─────────────────────────────────────────────────────────── + mode = "GPU remote build" if REMOTE_INDEX_BUILD else "CPU" + banner(f"Building index ({mode} via cuvs-bench)") + build_results = orchestrator.run_benchmark( + build=True, + search=False, + force=True, + **build_kwargs, + ) + + index_names = [ + r.index_path + for r in build_results + if isinstance(r, BuildResult) and r.success and r.index_path + ] + if not index_names: + print(" ERROR: no indexes were successfully built") + sys.exit(1) + build_times = { + r.index_path: r.build_time_seconds + for r in build_results + if isinstance(r, BuildResult) and r.success and r.index_path + } + for name, t in build_times.items(): + print(f" {name}: built in {t:.1f}s") + + # ── Search phase ────────────────────────────────────────────────────────── + banner("Running search benchmarks (via cuvs-bench)") + search_run_kwargs = dict( + build=False, + search=True, + count=K, + **common_kwargs, + ) + if BATCH_SIZE is not None: + search_run_kwargs["batch_size"] = BATCH_SIZE + search_results = orchestrator.run_benchmark(**search_run_kwargs) + + print_results(search_results) + + batch_size = _get_search_batch_size(search_results) + if batch_size is None: + print(" ERROR: no successful search result reported a batch size") + sys.exit(1) + + write_result_files( + build_results=build_results, + search_results=search_results, + dataset=DATASET, + dataset_path=DATASET_PATH, + algo=common_kwargs["algorithms"], + groups=BENCH_GROUPS, + k=K, + batch_size=batch_size, + ) + + print("\n" + "═" * 60) + print(" Benchmark complete!") + print("═" * 60) + print(f"\n OpenSearch : {OPENSEARCH_URL}") + if REMOTE_INDEX_BUILD: + print(f" GPU builder : {BUILDER_URL}") + print() + + +if __name__ == "__main__": + main() diff --git a/deploy/docker-compose.multinode.yml b/deploy/docker-compose.multinode.yml new file mode 100644 index 0000000000..c88f95d9e5 --- /dev/null +++ b/deploy/docker-compose.multinode.yml @@ -0,0 +1,147 @@ +# Multi-node split of the benchmark stack. +# +# Copy this file to all three EC2 instances from the deploy-opensearch-tmp +# branch of https://github.com/jrbourbeau/cuvs. The default OpenSearch and +# bench images build from that checkout, so keep the opensearch/ directory on +# the OpenSearch node and the bench/ directory on the client node unless you +# set OPENSEARCH_IMAGE and BENCH_IMAGE to prebuilt images. The OpenSearch image +# contains opensearch.yml, so a prebuilt image does not need the opensearch/ +# directory at runtime. +# +# OpenSearch node: +# docker compose -f docker-compose.multinode.yml --profile opensearch up -d +# +# GPU remote index builder node: +# docker compose -f docker-compose.multinode.yml --profile builder up -d +# +# Client/benchmark node: +# docker compose -f docker-compose.multinode.yml --profile client build bench +# docker compose -f docker-compose.multinode.yml --profile client run --rm bench +# +# Docker Compose networks are host-local. Across EC2 instances, use private +# EC2 DNS names or private IPv4 addresses in OPENSEARCH_URL, +# OPENSEARCH_HOST, and REMOTE_INDEX_BUILDER_URL. + +name: opensearch-cuvs-bench-multinode + +x-aws-env: &aws-env + AWS_ACCESS_KEY_ID: ${AWS_ACCESS_KEY_ID:-} + AWS_SECRET_ACCESS_KEY: ${AWS_SECRET_ACCESS_KEY:-} + AWS_SESSION_TOKEN: ${AWS_SESSION_TOKEN:-} + AWS_REGION: ${AWS_REGION:-us-west-2} + AWS_DEFAULT_REGION: ${AWS_DEFAULT_REGION:-us-west-2} + S3_BUCKET: ${S3_BUCKET:-} + S3_PREFIX: ${S3_PREFIX:-knn-indexes} + +x-bench-env: &bench-env + <<: *aws-env + OPENSEARCH_URL: ${OPENSEARCH_URL:-http://opensearch:9200} + OPENSEARCH_HOST: ${OPENSEARCH_HOST:-opensearch} + OPENSEARCH_PORT: ${OPENSEARCH_PORT:-9200} + BUILDER_URL: ${REMOTE_INDEX_BUILDER_URL:-http://remote-index-builder:1025} + REMOTE_INDEX_BUILDER_URL: ${REMOTE_INDEX_BUILDER_URL:-http://remote-index-builder:1025} + REMOTE_INDEX_BUILD: ${REMOTE_INDEX_BUILD:-true} + REMOTE_BUILD_SIZE_MIN: ${REMOTE_BUILD_SIZE_MIN:-} + REMOTE_BUILD_TIMEOUT: ${REMOTE_BUILD_TIMEOUT:-1800} + REMOTE_VECTOR_REPOSITORY: ${REMOTE_VECTOR_REPOSITORY:-vector-repo} + DATASET: ${DATASET:-sift-128-euclidean} + DATASET_PATH: /data/datasets + BENCH_GROUPS: ${BENCH_GROUPS:-test} + K: ${K:-10} + BATCH_SIZE: ${BATCH_SIZE:-} + BUILD_BATCH_SIZE: ${BUILD_BATCH_SIZE:-} + +services: + opensearch: + profiles: ["opensearch"] + image: ${OPENSEARCH_IMAGE:-opensearch-cuvs-bench-opensearch} + build: + context: ${OPENSEARCH_BUILD_CONTEXT:-./opensearch} + dockerfile: ${OPENSEARCH_DOCKERFILE:-Dockerfile} + container_name: opensearch + restart: unless-stopped + ports: + - "${OPENSEARCH_BIND_ADDR:-0.0.0.0}:${OPENSEARCH_PORT:-9200}:9200" + - "${OPENSEARCH_PERF_BIND_ADDR:-127.0.0.1}:${OPENSEARCH_PERF_PORT:-9600}:9600" + environment: + <<: *aws-env + cluster.name: ${OPENSEARCH_CLUSTER_NAME:-opensearch-cuvs-bench} + node.name: ${OPENSEARCH_NODE_NAME:-opensearch-node-1} + discovery.type: single-node + bootstrap.memory_lock: "true" + OPENSEARCH_JAVA_OPTS: ${OPENSEARCH_JAVA_OPTS:--Xms16g -Xmx16g} + DISABLE_SECURITY_PLUGIN: ${DISABLE_SECURITY_PLUGIN:-true} + OPENSEARCH_INITIAL_ADMIN_PASSWORD: ${OPENSEARCH_INITIAL_ADMIN_PASSWORD:-} + ulimits: + memlock: + soft: -1 + hard: -1 + nofile: + soft: 65536 + hard: 65536 + volumes: + - opensearch-data:/usr/share/opensearch/data + healthcheck: + test: ["CMD-SHELL", "curl -fsS http://localhost:9200/_cluster/health >/dev/null || exit 1"] + interval: 20s + timeout: 5s + retries: 30 + + remote-index-builder: + profiles: ["builder", "gpu"] + image: ${REMOTE_INDEX_BUILDER_IMAGE:-opensearchproject/remote-vector-index-builder:api-latest} + container_name: remote-index-builder + restart: unless-stopped + ports: + - "${REMOTE_INDEX_BUILDER_BIND_ADDR:-0.0.0.0}:${REMOTE_INDEX_BUILDER_HOST_PORT:-1025}:1025" + environment: + <<: *aws-env + healthcheck: + test: ["CMD-SHELL", "python3 -c 'import socket; socket.create_connection((\"localhost\", 1025), 2).close()'"] + interval: 5s + timeout: 5s + retries: 24 + start_period: 10s + deploy: + resources: + reservations: + devices: + - driver: nvidia + count: ${GPU_COUNT:-1} + capabilities: ["gpu"] + + bench: + profiles: ["client"] + image: ${BENCH_IMAGE:-opensearch-cuvs-bench-bench} + build: + context: ${BENCH_BUILD_CONTEXT:-./bench} + dockerfile: ${BENCH_DOCKERFILE:-Dockerfile} + environment: + <<: *bench-env + extra_hosts: + - "host.docker.internal:host-gateway" + volumes: + - ${DATASET_PATH:-/tmp/datasets}:/data/datasets + + configure-remote-index-build: + profiles: ["configure"] + image: curlimages/curl:latest + environment: + <<: *bench-env + command: + - sh + - -ec + - | + : "$${S3_BUCKET:?S3_BUCKET must be set}" + : "$${REMOTE_INDEX_BUILDER_URL:?REMOTE_INDEX_BUILDER_URL must be set}" + + curl -fsS -X PUT "$${OPENSEARCH_URL}/_snapshot/$${REMOTE_VECTOR_REPOSITORY}" \ + -H 'Content-Type: application/json' \ + -d "{\"type\":\"s3\",\"settings\":{\"bucket\":\"$${S3_BUCKET}\",\"base_path\":\"$${S3_PREFIX}\",\"region\":\"$${AWS_REGION}\"}}" + + curl -fsS -X PUT "$${OPENSEARCH_URL}/_cluster/settings" \ + -H 'Content-Type: application/json' \ + -d "{\"persistent\":{\"knn.remote_index_build.enabled\":\"true\",\"knn.remote_index_build.repository\":\"$${REMOTE_VECTOR_REPOSITORY}\",\"knn.remote_index_build.service.endpoint\":\"$${REMOTE_INDEX_BUILDER_URL}\"}}" + +volumes: + opensearch-data: diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml new file mode 100644 index 0000000000..4b6300ecb8 --- /dev/null +++ b/deploy/docker-compose.yml @@ -0,0 +1,137 @@ +# 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) +# 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 +# +# 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 + 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:-} + 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..a7548ded04 --- /dev/null +++ b/deploy/remote-index-build/run.py @@ -0,0 +1,333 @@ +#!/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")) + +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": 1, + "number_of_replicas": 0, + } + if REMOTE_BUILD_SIZE_MIN: + index_settings["index.knn.remote_index_build.size.min"] = ( + REMOTE_BUILD_SIZE_MIN + ) + + r = session.put( + f"{OPENSEARCH_URL}/{INDEX_NAME}", + json={ + "settings": index_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()