Add continuous benchmark for tenants (#183)

qdrant · Aug 15, 2024 · a11ebc0 · a11ebc0
1 parent 4ffc3ce
commit a11ebc0
Show file tree

Hide file tree

Showing 9 changed files with 240 additions and 29 deletions.
diff --git a/.github/workflows/continuous-benchmark.yaml b/.github/workflows/continuous-benchmark.yaml
@@ -44,6 +44,20 @@ jobs:
               timeout 30m bash -x tools/run_ci.sh
             done
 
+            # Benchmark filtered search by tenants with mem limitation
+
+            export ENGINE_NAME="qdrant-all-on-disk-scalar-q"
+            export DATASETS="random-768-100-tenants"
+            export CONTAINER_MEM_LIMIT=150mb
+
+            # Benchmark the dev branch:
+            export QDRANT_VERSION=ghcr/dev
+            timeout 30m bash -x tools/run_ci.sh
+
+            # Benchmark the master branch:
+            export QDRANT_VERSION=docker/master
+            timeout 30m bash -x tools/run_ci.sh
+
             set -e
       - name: Fail job if any of the benches failed
         if: steps.benches.outputs.failed == 'error' || steps.benches.outputs.failed == 'timeout'
@@ -67,4 +81,4 @@ jobs:
             }
         env:
           SLACK_WEBHOOK_URL: ${{ secrets.CI_ALERTS_CHANNEL_WEBHOOK_URL }}
-          SLACK_WEBHOOK_TYPE: INCOMING_WEBHOOK
+          SLACK_WEBHOOK_TYPE: INCOMING_WEBHOOK
diff --git a/datasets/datasets.json b/datasets/datasets.json
@@ -336,6 +336,17 @@
       "b": "keyword"
     }
   },
+  {
+    "name": "random-768-100-tenants",
+    "vector_size": 768,
+    "distance": "cosine",
+    "type": "tar",
+    "link": "https://storage.googleapis.com/ann-filtered-benchmark/datasets/random_keywords_1m_768_vocab_100.tgz",
+    "path": "random-768-100-tenants/random_keywords_1m_768_vocab_100",
+    "schema": {
+      "a": "keyword"
+    }
+  },
   {
     "name": "random-100-match-kw-small-vocab-no-filters",
     "vector_size": 256,

diff --git a/engine/clients/qdrant/configure.py b/engine/clients/qdrant/configure.py
@@ -21,6 +21,13 @@ class QdrantConfigurator(BaseConfigurator):
         "float": rest.PayloadSchemaType.FLOAT,
         "geo": rest.PayloadSchemaType.GEO,
     }
+    INDEX_PARAMS_TYPE_MAPPING = {
+        "int": rest.IntegerIndexParams,
+        "keyword": rest.KeywordIndexParams,
+        "text": rest.TextIndexParams,
+        "float": rest.FloatIndexParams,
+        "geo": rest.GeoIndexParams,
+    }
 
     def __init__(self, host, collection_params: dict, connection_params: dict):
         super().__init__(host, collection_params, connection_params)
@@ -43,15 +50,25 @@ def recreate(self, dataset: Dataset, collection_params):
                 },
             }
         else:
+            is_vectors_on_disk = self.collection_params.get("vectors_config", {}).get(
+                "on_disk", False
+            )
+            self.collection_params.pop("vectors_config", None)
+
             vectors_config = {
                 "vectors_config": (
                     rest.VectorParams(
                         size=dataset.config.vector_size,
                         distance=self.DISTANCE_MAPPING.get(dataset.config.distance),
+                        on_disk=is_vectors_on_disk,
                     )
                 )
             }
 
+        payload_index_params = self.collection_params.pop("payload_index_params", {})
+        if not set(payload_index_params.keys()).issubset(dataset.config.schema.keys()):
+            raise ValueError("payload_index_params are not found in dataset schema")
+
         self.client.recreate_collection(
             collection_name=QDRANT_COLLECTION_NAME,
             **vectors_config,
@@ -65,8 +82,24 @@ def recreate(self, dataset: Dataset, collection_params):
             ),
         )
         for field_name, field_type in dataset.config.schema.items():
-            self.client.create_payload_index(
-                collection_name=QDRANT_COLLECTION_NAME,
-                field_name=field_name,
-                field_schema=self.INDEX_TYPE_MAPPING.get(field_type),
-            )
+            if field_type in ["keyword", "uuid"]:
+                is_tenant = payload_index_params.get(field_name, {}).get(
+                    "is_tenant", None
+                )
+                on_disk = payload_index_params.get(field_name, {}).get("on_disk", None)
+
+                self.client.create_payload_index(
+                    collection_name=QDRANT_COLLECTION_NAME,
+                    field_name=field_name,
+                    field_schema=self.INDEX_PARAMS_TYPE_MAPPING.get(field_type)(
+                        type=self.INDEX_TYPE_MAPPING.get(field_type),
+                        is_tenant=is_tenant,
+                        on_disk=on_disk,
+                    ),
+                )
+            else:
+                self.client.create_payload_index(
+                    collection_name=QDRANT_COLLECTION_NAME,
+                    field_name=field_name,
+                    field_schema=self.INDEX_TYPE_MAPPING.get(field_type),
+                )
diff --git a/engine/servers/qdrant-continuous-benchmarks-with-volume/docker-compose.yaml b/engine/servers/qdrant-continuous-benchmarks-with-volume/docker-compose.yaml
@@ -0,0 +1,29 @@
+version: '3.7'
+
+services:
+  qdrant_bench:
+    image: ${CONTAINER_REGISTRY:-docker.io}/qdrant/qdrant:${QDRANT_VERSION}
+    container_name: qdrant-continuous
+    ports:
+      - "6333:6333"
+      - "6334:6334"
+    volumes:
+      - qdrant_storage:/qdrant/storage
+    logging:
+      driver: "json-file"
+      options:
+        max-file: 1
+        max-size: 10m
+    deploy:
+      resources:
+        limits:
+          memory: ${CONTAINER_MEM_LIMIT:-25Gb}
+
+volumes:
+  qdrant_storage:
+    name: "qdrant_storage"
+    driver: local
+    driver_opts:
+      type: none
+      device: ${PWD}/qdrant_storage
+      o: bind
diff --git a/experiments/configurations/qdrant-on-disk.json b/experiments/configurations/qdrant-on-disk.json
@@ -11,5 +11,24 @@
       { "parallel": 8, "config": { "hnsw_ef": 128 } }
     ],
     "upload_params": { "parallel": 4 }
+  },
+  {
+    "name": "qdrant-all-on-disk-scalar-q",
+    "engine": "qdrant",
+    "connection_params": {},
+    "collection_params": {
+      "optimizers_config": { "default_segment_number": 17 },
+      "quantization_config": { "scalar": {"type": "int8", "quantile": 0.99, "always_ram": false} },
+      "vectors_config": { "on_disk": true },
+      "hnsw_config": { "on_disk": true, "m": 0, "payload_m": 16 },
+      "on_disk_payload": true,
+      "payload_index_params": {
+        "a": { "is_tenant": true, "on_disk": true }
+      }
+    },
+    "search_params": [
+      { "parallel": 8 }
+    ],
+    "upload_params": { "parallel": 4 }
   }
 ]
diff --git a/tools/run_client_script.sh b/tools/run_client_script.sh
@@ -3,6 +3,9 @@
 PS4='ts=$(date "+%Y-%m-%dT%H:%M:%SZ") level=DEBUG line=$LINENO file=$BASH_SOURCE '
 set -euo pipefail
 
+# Possible values are: full|upload|search
+EXPERIMENT_MODE=${1:-"full"}
+
 CLOUD_NAME=${CLOUD_NAME:-"hetzner"}
 SERVER_USERNAME=${SERVER_USERNAME:-"root"}
 
@@ -22,16 +25,26 @@ DATASETS=${DATASETS:-"laion-small-clip"}
 
 PRIVATE_IP_OF_THE_SERVER=$(bash "${SCRIPT_PATH}/${CLOUD_NAME}/get_private_ip.sh" "$BENCH_SERVER_NAME")
 
-RUN_EXPERIMENT="ENGINE_NAME=${ENGINE_NAME} DATASETS=${DATASETS} PRIVATE_IP_OF_THE_SERVER=${PRIVATE_IP_OF_THE_SERVER} bash ~/run_experiment.sh"
+RUN_EXPERIMENT="ENGINE_NAME=${ENGINE_NAME} DATASETS=${DATASETS} PRIVATE_IP_OF_THE_SERVER=${PRIVATE_IP_OF_THE_SERVER} EXPERIMENT_MODE=${EXPERIMENT_MODE} bash ~/run_experiment.sh"
 
 ssh -tt -o ServerAliveInterval=60 -o ServerAliveCountMax=3 "${SERVER_USERNAME}@${IP_OF_THE_CLIENT}" "${RUN_EXPERIMENT}"
 
-SEARCH_RESULT_FILE=$(ssh "${SERVER_USERNAME}@${IP_OF_THE_CLIENT}" "ls -t results/*-search-*.json | head -n 1")
-UPLOAD_RESULT_FILE=$(ssh "${SERVER_USERNAME}@${IP_OF_THE_CLIENT}" "ls -t results/*-upload-*.json | head -n 1")
+echo "Gather experiment results..."
+result_files_arr=()
+
+if [[ "$EXPERIMENT_MODE" == "full" ]] || [[ "$EXPERIMENT_MODE" == "upload" ]]; then
+  UPLOAD_RESULT_FILE=$(ssh "${SERVER_USERNAME}@${IP_OF_THE_CLIENT}" "ls -t results/*-upload-*.json | head -n 1")
+  result_files_arr+=("$UPLOAD_RESULT_FILE")
+fi
+
+if [[ "$EXPERIMENT_MODE" == "full" ]] || [[ "$EXPERIMENT_MODE" == "search" ]]; then
+  SEARCH_RESULT_FILE=$(ssh "${SERVER_USERNAME}@${IP_OF_THE_CLIENT}" "ls -t results/*-search-*.json | head -n 1")
+  result_files_arr+=("$SEARCH_RESULT_FILE")
+fi
 
 mkdir -p results
 
-for RESULT_FILE in $SEARCH_RESULT_FILE $UPLOAD_RESULT_FILE; do
+for RESULT_FILE in "${result_files_arr[@]}"; do
     # -p preseves modification time, access time, and modes (but not change time)
     scp -p "${SERVER_USERNAME}@${IP_OF_THE_CLIENT}:~/${RESULT_FILE}" "./results"
 done
diff --git a/tools/run_experiment.sh b/tools/run_experiment.sh
@@ -9,6 +9,8 @@ DATASETS=${DATASETS:-""}
 
 PRIVATE_IP_OF_THE_SERVER=${PRIVATE_IP_OF_THE_SERVER:-""}
 
+EXPERIMENT_MODE=${EXPERIMENT_MODE:-"full"}
+
 if [[ -z "$ENGINE_NAME" ]]; then
   echo "ENGINE_NAME is not set"
   exit 1
@@ -24,23 +26,40 @@ if [[ -z "$PRIVATE_IP_OF_THE_SERVER" ]]; then
   exit 1
 fi
 
+if [[ -z "$EXPERIMENT_MODE" ]]; then
+  echo "EXPERIMENT_MODE is not set, possible values are: full | upload | search"
+  exit 1
+fi
 docker container rm -f ci-benchmark-upload || true
 docker container rm -f ci-benchmark-search || true
 
 docker rmi --force qdrant/vector-db-benchmark:latest || true
 
-docker run \
-  --rm \
-  -it \
-  --name ci-benchmark-upload \
-  -v "$HOME/results:/code/results" \
-  qdrant/vector-db-benchmark:latest \
-  python run.py --engines "${ENGINE_NAME}" --datasets "${DATASETS}" --host "${PRIVATE_IP_OF_THE_SERVER}" --no-skip-if-exists --skip-search
-
-docker run \
-  --rm \
-  -it \
-  --name ci-benchmark-search \
-  -v "$HOME/results:/code/results" \
-  qdrant/vector-db-benchmark:latest \
-  python run.py --engines "${ENGINE_NAME}" --datasets "${DATASETS}" --host "${PRIVATE_IP_OF_THE_SERVER}" --no-skip-if-exists --skip-upload
+if [[ "$EXPERIMENT_MODE" == "full" ]] || [[ "$EXPERIMENT_MODE" == "upload" ]]; then
+  echo "EXPERIMENT_MODE=$EXPERIMENT_MODE"
+  docker run \
+    --rm \
+    -it \
+    --name ci-benchmark-upload \
+    -v "$HOME/results:/code/results" \
+    qdrant/vector-db-benchmark:latest \
+    python run.py --engines "${ENGINE_NAME}" --datasets "${DATASETS}" --host "${PRIVATE_IP_OF_THE_SERVER}" --no-skip-if-exists --skip-search
+fi
+
+
+if [[ "$EXPERIMENT_MODE" == "full" ]] || [[ "$EXPERIMENT_MODE" == "search" ]]; then
+  echo "EXPERIMENT_MODE=$EXPERIMENT_MODE"
+
+  if [[ "$EXPERIMENT_MODE" == "search" ]]; then
+    echo "Drop caches before running the experiment"
+    sudo bash -c 'sync; echo 1 > /proc/sys/vm/drop_caches'
+  fi
+
+  docker run \
+    --rm \
+    -it \
+    --name ci-benchmark-search \
+    -v "$HOME/results:/code/results" \
+    qdrant/vector-db-benchmark:latest \
+    python run.py --engines "${ENGINE_NAME}" --datasets "${DATASETS}" --host "${PRIVATE_IP_OF_THE_SERVER}" --no-skip-if-exists --skip-upload
+fi
diff --git a/tools/run_remote_benchmark.sh b/tools/run_remote_benchmark.sh
@@ -34,11 +34,31 @@ trap 'cleanup' EXIT
 SERVER_NAME=$BENCH_SERVER_NAME bash -x "${SCRIPT_PATH}/${CLOUD_NAME}/check_ssh_connection.sh"
 SERVER_NAME=$BENCH_CLIENT_NAME bash -x "${SCRIPT_PATH}/${CLOUD_NAME}/check_ssh_connection.sh"
 
+if [[ -z "${CONTAINER_MEM_LIMIT:-}" ]]; then
+  echo "CONTAINER_MEM_LIMIT is not set, run without memory limit"
 
-SERVER_CONTAINER_NAME=${SERVER_CONTAINER_NAME:-"qdrant-continuous-benchmarks"}
+  SERVER_CONTAINER_NAME=${SERVER_CONTAINER_NAME:-"qdrant-continuous-benchmarks"}
 
-bash -x "${SCRIPT_PATH}/run_server_container.sh" "$SERVER_CONTAINER_NAME"
+  bash -x "${SCRIPT_PATH}/run_server_container.sh" "$SERVER_CONTAINER_NAME"
 
-bash -x "${SCRIPT_PATH}/run_client_script.sh"
+  bash -x "${SCRIPT_PATH}/run_client_script.sh"
+
+  bash -x "${SCRIPT_PATH}/qdrant_collect_stats.sh" "$SERVER_CONTAINER_NAME"
+
+else
+  echo "CONTAINER_MEM_LIMIT is set, run search with memory limit: ${CONTAINER_MEM_LIMIT}"
+
+  SERVER_CONTAINER_NAME=${SERVER_CONTAINER_NAME:-"qdrant-continuous-benchmarks-with-volume"}
+
+  bash -x "${SCRIPT_PATH}/run_server_container_with_volume.sh" "$SERVER_CONTAINER_NAME"
+
+  bash -x "${SCRIPT_PATH}/run_client_script.sh" "upload"
+
+  bash -x "${SCRIPT_PATH}/run_server_container_with_volume.sh" "$SERVER_CONTAINER_NAME" "$CONTAINER_MEM_LIMIT" "continue"
+
+  bash -x "${SCRIPT_PATH}/run_client_script.sh" "search"
+
+  bash -x "${SCRIPT_PATH}/qdrant_collect_stats.sh" "$SERVER_CONTAINER_NAME"
+
+fi
 
-bash -x "${SCRIPT_PATH}/qdrant_collect_stats.sh" "$SERVER_CONTAINER_NAME"
diff --git a/tools/run_server_container_with_volume.sh b/tools/run_server_container_with_volume.sh
@@ -0,0 +1,53 @@
+#!/bin/bash
+
+set -e
+
+# Examples: qdrant-continuous-benchmarks-with-volume
+CONTAINER_NAME=$1
+CONTAINER_MEM_LIMIT=${2:-"25Gb"}
+EXECUTION_MODE=${3:-"init"}
+
+CLOUD_NAME=${CLOUD_NAME:-"hetzner"}
+SERVER_USERNAME=${SERVER_USERNAME:-"root"}
+
+
+SCRIPT=$(realpath "$0")
+SCRIPT_PATH=$(dirname "$SCRIPT")
+
+BENCH_SERVER_NAME=${SERVER_NAME:-"benchmark-server-1"}
+
+QDRANT_VERSION=${QDRANT_VERSION:-"dev"}
+
+IP_OF_THE_SERVER=$(bash "${SCRIPT_PATH}/${CLOUD_NAME}/get_public_ip.sh" "$BENCH_SERVER_NAME")
+
+bash -x "${SCRIPT_PATH}/sync_servers.sh" "root@$IP_OF_THE_SERVER"
+
+# if version is starts with "docker" or "ghcr", use container
+if [[ ${QDRANT_VERSION} == docker/* ]] || [[ ${QDRANT_VERSION} == ghcr/* ]]; then
+
+    if [[ ${QDRANT_VERSION} == docker/* ]]; then
+        # pull from docker hub
+        QDRANT_VERSION=${QDRANT_VERSION#docker/}
+        CONTAINER_REGISTRY='docker.io'
+    elif [[ ${QDRANT_VERSION} == ghcr/* ]]; then
+        # pull from github container registry
+        QDRANT_VERSION=${QDRANT_VERSION#ghcr/}
+        CONTAINER_REGISTRY='ghcr.io'
+    fi
+
+    if [[ "$EXECUTION_MODE" == "init" ]]; then
+      # create volume qdrant_storage
+      echo "Initialize qdrant from scratch"
+      DOCKER_VOLUME_SET_UP="docker volume rm -f qdrant_storage; sudo rm -rf qdrant_storage; mkdir qdrant_storage"
+      DOCKER_COMPOSE="export QDRANT_VERSION=${QDRANT_VERSION}; export CONTAINER_REGISTRY=${CONTAINER_REGISTRY}; export CONTAINER_MEM_LIMIT=${CONTAINER_MEM_LIMIT}; docker compose down; pkill qdrant; docker rm -f qdrant-continuous || true; docker rmi -f ${CONTAINER_REGISTRY}/qdrant/qdrant:${QDRANT_VERSION} || true ; ${DOCKER_VOLUME_SET_UP}; docker compose up -d; docker container ls -a"
+    else
+      # suggest that volume qdrant_storage exist and start qdrant
+      echo "Reload qdrant with existing data"
+      DOCKER_COMPOSE="export QDRANT_VERSION=${QDRANT_VERSION}; export CONTAINER_REGISTRY=${CONTAINER_REGISTRY}; export CONTAINER_MEM_LIMIT=${CONTAINER_MEM_LIMIT}; docker compose down; pkill qdrant; docker rm -f qdrant-continuous || true; docker rmi -f ${CONTAINER_REGISTRY}/qdrant/qdrant:${QDRANT_VERSION} || true ; sudo bash -c 'sync; echo 1 > /proc/sys/vm/drop_caches'; docker compose up -d; docker container ls -a"
+    fi
+
+    ssh -t  -o ServerAliveInterval=60 -o ServerAliveCountMax=3 "${SERVER_USERNAME}@${IP_OF_THE_SERVER}" "cd ./projects/vector-db-benchmark/engine/servers/${CONTAINER_NAME} ; $DOCKER_COMPOSE"
+else
+    echo "Error: unknown version ${QDRANT_VERSION}. Version name should start with 'docker/' or 'ghcr/'"
+    exit 1
+fi