Add Custom Faiss Base Image Build Code

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "custom-faiss-installed-image/faiss"]
+	path = custom-faiss-installed-image/faiss
+	url = https://github.com/facebookresearch/faiss.git

custom-faiss-installed-image/.dockerignore

@@ -0,0 +1 @@
+*.npy

custom-faiss-installed-image/0001-added-commit-to-enable-the-add_with_ids-function-on-.patch

@@ -0,0 +1,48 @@
+From 16c4099a97cc8c69e1e049b37df692ded346ed8f Mon Sep 17 00:00:00 2001
+From: Navneet Verma <[email protected]>
+Date: Fri, 27 Dec 2024 19:49:46 +0000
+Subject: [PATCH] added commit to enable the add_with_ids function on cagra
+ index
+
+Signed-off-by: Navneet Verma <[email protected]>
+---
+ faiss/gpu/GpuIndexCagra.cu | 8 ++++++++
+ faiss/gpu/GpuIndexCagra.h  | 2 ++
+ 2 files changed, 10 insertions(+)
+
+diff --git a/faiss/gpu/GpuIndexCagra.cu b/faiss/gpu/GpuIndexCagra.cu
+index fe0c82b..5a1f090 100644
+--- a/faiss/gpu/GpuIndexCagra.cu
++++ b/faiss/gpu/GpuIndexCagra.cu
+@@ -103,6 +103,14 @@ void GpuIndexCagra::train(idx_t n, const float* x) {
+     this->ntotal = n;
+ }
+
++/*
++Adding add function to support add_with_ids functionality from IndexIDMap index type.
++*/
++void GpuIndexCagra::add(idx_t n, const float* x) {
++    train(n, x);
++}
++
++
+ bool GpuIndexCagra::addImplRequiresIDs_() const {
+     return false;
+ };
+diff --git a/faiss/gpu/GpuIndexCagra.h b/faiss/gpu/GpuIndexCagra.h
+index d6fae29..d28bc6d 100644
+--- a/faiss/gpu/GpuIndexCagra.h
++++ b/faiss/gpu/GpuIndexCagra.h
+@@ -248,6 +248,8 @@ struct GpuIndexCagra : public GpuIndex {
+     /// Trains CAGRA based on the given vector data
+     void train(idx_t n, const float* x) override;
+
++    void add(idx_t n, const float* x) override;
++
+     /// Initialize ourselves from the given CPU index; will overwrite
+     /// all data in ourselves
+     void copyFrom(const faiss::IndexHNSWCagra* index);
+-- 
+2.25.1
+
+

custom-faiss-installed-image/Dockerfile

@@ -0,0 +1,45 @@
+FROM nvidia/cuda:12.1.1-base-ubuntu22.04
+
+# Install base utilities and also the dependencies for faiss
+RUN apt-get update \
+    && apt-get install -y build-essential gfortran libblas-dev libopenblas-base zlib1g zlib1g-dev git tmux libssl-dev \
+    && apt-get install -y wget \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN useradd -m appuser 
+
+# Install miniconda
+ENV CONDA_DIR /opt/conda
+RUN wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh && \
+    /bin/bash ~/miniconda.sh -b -p /opt/conda
+
+RUN chown -R appuser:appuser /opt
+COPY faiss-test.py /tmp
+RUN mkdir /tmp/faiss
+COPY faiss /tmp/faiss
+COPY build-faiss.sh /tmp/build-faiss.sh
+
+RUN chown -R appuser:appuser /tmp
+#RUN chmod -R 755 /opt
+
+USER appuser
+# Put conda in path so we can use conda activate
+ENV PATH=$CONDA_DIR/bin:$PATH
+#install some necessary dependencies
+RUN conda install -y -q python=3.11 cmake=3.26 make=4.2 swig=4.0 "numpy<2" scipy=1.14 pytest=7.4 gflags=2.2
+RUN cmake --version
+
+# special for 12.1.1
+ENV CUDA_ARCHS "70-real;72-real;75-real;80;86-real"
+
+# install base packages for X86_64
+RUN conda install -y -q -c conda-forge gxx_linux-64=14.2 sysroot_linux-64=2.17
+
+# install all the dependencies that we think we will need for installing faiss
+# https://github.com/facebookresearch/faiss/blob/3c8dc4194907e9b911551d5a009468106f8b9c7f/.github/actions/build_cmake/action.yml#L57C31-L57C44
+RUN conda install -y -q libcuvs=24.12 cuda-nvcc gxx_linux-64=12.4 -c rapidsai -c conda-forge
+
+RUN /tmp/build-faiss.sh
+
+CMD ["sh", "-c", "nvidia-smi && sleep 36000"]

custom-faiss-installed-image/README.md

@@ -0,0 +1,61 @@
+## Creating the image
+Below are the steps to create the image that can have a custom faiss python code build.
+1. The basic docker file is present in this folder named Dockerfile.
+2. Ensure that faiss is checked out before you trigger the build. The commit of faiss I have tested with is: 3c8dc4194907e9b911551d5a009468106f8b9c7f
+3. Build the docker image using this command
+```
+rm -rf faiss ; git restore faiss ; ./build-docker-image.sh
+```
+4. Once image is build use the below command to run the container.
+```
+docker run --gpus all -d -v ./:/tmp/ custom-faiss:latest
+```
+Make sure to run the command from the same folder where this readme is present otherwise wrong files get copied.
+5. Now ssh into the container
+```
+docker container exec -it <container-id>  /bin/bash
+```
+6. Once you ssh in the container make sure that you are sshed as `appuser`.
+7. Now go to `tmp/faiss` using `cd /tmp/faiss`
+8. Now run below command to trigger the cmake build.
+```
+cmake -B build \
+    -DBUILD_SHARED_LIBS=ON \
+    -DFAISS_ENABLE_GPU=ON \
+    -DFAISS_OPT_LEVEL=generic \
+    -DFAISS_ENABLE_C_API=ON \
+    -DFAISS_ENABLE_PYTHON=ON \
+    -DPYTHON_EXECUTABLE=$CONDA/bin/python \
+    -DCMAKE_BUILD_TYPE=Release \
+    -DCMAKE_CUDA_ARCHITECTURES="${CUDA_ARCHS}" \
+    -DFAISS_ENABLE_CUVS=ON \
+    -DCUDAToolkit_ROOT="/usr/local/cuda/lib64" \
+    .
+```
+9. Now run the below command to make faiss. Here I am using 4 cores to build faiss, you can use more if you have more cores. Just make sure that you are not using all cores otherwise the machine will get struck. This command takes a lot of time to run and build faiss
+```
+make -C build -j6 faiss swigfaiss
+```
+10. Once the build is complete now run the below command to create faiss bindings.
+```
+cd build/faiss/python && python3 setup.py build
+```
+11. One possible way to use faiss which you have built is like this, that can be added in the code.
+```
+import sys
+sys.path.append('/tmp/faiss/build/faiss/python')
+import faiss
+```
+Another approach can be setting the PYTHONPATH variable like this
+```
+export PYTHONPATH="$(ls -d `pwd`/tmp/faiss/build/faiss/python/build/lib*/):`pwd`/"
+```
+12. Now to test if everthing is build correctly or not run the below command
+```
+python faiss-test.py
+```
+and see an output like this
+```
+Creating GPU Index.. with IVF_PQ
+Indexing done
+```

custom-faiss-installed-image/build-docker-image.sh

@@ -0,0 +1,53 @@
+#!/bin/bash
+
+# Exit on any error
+set -xe
+
+# Function to check if command was successful
+check_status() {
+    if [ $? -ne 0 ]; then
+        echo "Error: $1"
+        exit 1
+    fi
+}
+
+# Function to check if patch file exists
+check_patch_file() {
+    if [ ! -f "$1" ]; then
+        echo "Error: Patch file '$1' not found!"
+        exit 1
+    fi
+}
+
+apply_patch() {
+    echo "Updating FAISS submodule..."
+    git submodule update --init -- faiss
+    check_status "Failed to update submodule"
+
+    # Check if patch file is provided as argument
+    if [ $# -ne 1 ]; then
+        echo "Usage: $0 "
+        exit 1
+    fi
+
+    PATCH_FILE="$1"
+    check_patch_file "$PATCH_FILE"
+
+    # Navigate into the faiss directory
+    echo "Entering FAISS directory..."
+    cd faiss || exit 1
+
+    # Apply the patch
+    echo "Applying patch..."
+    git apply ../"$PATCH_FILE"
+    check_status "Failed to apply patch"
+
+    chmod -R 777 .
+
+    echo "Successfully applied patch to FAISS!"
+    cd ..
+}
+
+apply_patch "0001-added-commit-to-enable-the-add_with_ids-function-on-.patch"
+
+docker build -t custom-faiss:latest .

custom-faiss-installed-image/build-faiss.sh

@@ -0,0 +1,35 @@
+#!/bin/bash
+
+# Exit on any error
+set -xe
+
+cd /tmp/faiss
+
+printenv
+
+cmake --version
+
+echo "Running cmake build"
+pwd
+cmake -B build \
+    -DBUILD_SHARED_LIBS=ON \
+    -DFAISS_ENABLE_GPU=ON \
+    -DFAISS_OPT_LEVEL=generic \
+    -DFAISS_ENABLE_C_API=ON \
+    -DFAISS_ENABLE_PYTHON=ON \
+    -DPYTHON_EXECUTABLE=$CONDA/bin/python \
+    -DCMAKE_BUILD_TYPE=Release \
+    -DCMAKE_CUDA_ARCHITECTURES="${CUDA_ARCHS}" \
+    -DFAISS_ENABLE_CUVS=ON \
+    -DCUDAToolkit_ROOT="/usr/local/cuda/lib64" \
+    .
+
+echo "Running make command"
+
+make -C build -j6 faiss swigfaiss
+
+# Setting up the python code
+cd build/faiss/python && python3 setup.py build
+
+# make sure that faiss python code is available for the use
+export PYTHONPATH="$(ls -d `pwd`/tmp/faiss/build/faiss/python/build/lib*/):`pwd`/"

custom-faiss-installed-image/create-dataset.py

@@ -0,0 +1,14 @@
+import numpy as np
+
+if __name__ == "__main__":
+    d = 768
+    np.random.seed(1234)
+    # Create a new memory-mapped array
+    shape = (7_000_000, 768)  # Example shape
+    print(f"Creating dataset of {shape}")
+    filename = 'mmap-7m.npy'
+
+    mmap_array = np.lib.format.open_memmap(filename, mode='w+', dtype='float32', shape=shape)
+    mmap_array[:] = np.random.rand(*shape)
+    del mmap_array
+    print("file is written")

custom-faiss-installed-image/faiss-test.py

@@ -0,0 +1,105 @@
+import sys
+import os.path
+# If you faiss python installed you should skip this line
+sys.path.append('/tmp/faiss/build/faiss/python')
+import faiss
+import logging
+from timeit import default_timer as timer
+import math
+import numpy as np
+import time
+
+from numpy import dtype
+
+
+def indexData(d:int, xb:np.ndarray, ids:np.ndarray, indexingParams:dict, space_type:str, file_to_write:str="gpuIndex.cagra.graph"):
+    num_of_parallel_threads = 8
+    logging.info(f"Setting number of parallel threads for graph build: {num_of_parallel_threads}")
+    faiss.omp_set_num_threads(num_of_parallel_threads)
+    res = faiss.StandardGpuResources()
+    metric = faiss.METRIC_L2
+    if space_type == "innerproduct":
+        metric = faiss.METRIC_INNER_PRODUCT
+    cagraIndexConfig = faiss.GpuIndexCagraConfig()
+    cagraIndexConfig.intermediate_graph_degree = 64 if indexingParams.get('intermediate_graph_degree') is None else indexingParams['intermediate_graph_degree']
+    cagraIndexConfig.graph_degree = 32 if indexingParams.get('graph_degree') == None else indexingParams['graph_degree']
+    cagraIndexConfig.device = faiss.get_num_gpus() - 1
+    #cagraIndexConfig.conservative_memory_allocation = True
+    # This was available earlier now this parameter is now available
+    #cagraIndexConfig.store_dataset = False
+
+    cagraIndexConfig.build_algo = faiss.graph_build_algo_IVF_PQ
+    cagraIndexIVFPQConfig = faiss.IVFPQBuildCagraConfig()
+    cagraIndexIVFPQConfig.kmeans_n_iters = 10 if indexingParams.get('kmeans_n_iters') == None else indexingParams['kmeans_n_iters']
+    cagraIndexIVFPQConfig.pq_bits = 8 if indexingParams.get('pq_bits') == None else indexingParams['pq_bits']
+    cagraIndexIVFPQConfig.pq_dim = 32 if indexingParams.get('pq_dim') == None else indexingParams['pq_dim']
+    cagraIndexIVFPQConfig.n_lists = int(math.sqrt(len(xb))) if indexingParams.get('n_lists') == None else indexingParams['n_lists']
+    cagraIndexIVFPQConfig.kmeans_trainset_fraction = 10 if indexingParams.get('kmeans_trainset_fraction') == None else indexingParams['kmeans_trainset_fraction']
+    cagraIndexIVFPQConfig.conservative_memory_allocation = True
+    cagraIndexConfig.ivf_pq_params = cagraIndexIVFPQConfig
+
+    cagraIndexSearchIVFPQConfig = faiss.IVFPQSearchCagraConfig()
+    cagraIndexSearchIVFPQConfig.n_probes = 30 if indexingParams.get('n_probes') == None else indexingParams['n_probes']
+    cagraIndexConfig.ivf_pq_search_params = cagraIndexSearchIVFPQConfig
+
+    print("Creating GPU Index.. with IVF_PQ")
+    t1 = timer()
+    cagraIVFPQIndex = faiss.GpuIndexCagra(res, d, metric, cagraIndexConfig)
+    idMapIVFPQIndex = faiss.IndexIDMap(cagraIVFPQIndex)
+    indexDataInIndex(idMapIVFPQIndex, ids, xb)
+    t2 = timer()
+    print(f"Indexing time: {t2 - t1} sec")
+
+def indexDataInIndex(index: faiss.Index, ids, xb):
+    if ids is None:
+        index.train(xb)
+    else:
+        index.add_with_ids(xb, ids)
+
+def runIndicesSearch(xq, graphFile:str, param:dict, gt) -> dict:
+    index:faiss.IndexIDMap = loadGraphFromFile(graphFile)
+    index.index.base_level_only = True
+    hnswParameters = faiss.SearchParametersHNSW()
+    hnswParameters.efSearch = 100 if param.get('ef_search') is None else param['ef_search']
+    logging.info(f"Ef search is : {hnswParameters.efSearch}")
+    k = 100 if param.get('K') is None else param['K']
+
+    def search(xq, k, params):
+        D, ids = index.search(xq, k, params=params)
+        return ids
+    # K is always set to 100
+    total_time = 0
+    I = []
+    query = xq[0]
+    t1 = timer()
+    result = search(np.array([query]), 100, hnswParameters)
+    t2 = timer()
+    I.append(result[0])
+    total_time = total_time + (t2-t1)
+    print(f"Total time for search: {total_time}")
+
+
+
+def loadGraphFromFile(graphFile: str) -> faiss.Index:
+    if os.path.isfile(graphFile) is False:
+        logging.error(f"The path provided: {graphFile} is not a file")
+        sys.exit(0)
+
+    return faiss.read_index(graphFile)
+
+
+
+if __name__ == "__main__":
+    d = 768
+    filename = 'mmap-7m.npy'
+    print("file is written, loading file now..")
+    xb = np.load(file = filename)
+    print(xb.shape)
+
+    #ids = [i for i in range(len(xb))]
+    ids = None
+    indexData(d, xb, ids, {}, "l2", "testgpuIndex.cagra.graph")
+    print("Indexing done")
+    del xb
+    time.sleep(10)
+    print("Deleted xb")