DOCSEARCHbyAlgolia
Submit Docs

AutoFaiss Documentation

repository·master·Indexed 21 days ago

https://github.com/criteo/autofaiss

AutoFaiss is a library for automatically creating optimal Faiss k-nearest neighbor (KNN) indices. It selects indexing parameters to maximize recall based on user-defined memory and query speed constraints. The library supports in-memory numpy arrays, .npy files on disk, and distributed index construction using PySpark for large datasets. It includes a Command Line Interface (CLI) for index generation and a Quantizer module for compressing indices while maintaining search accuracy.

Tokens
9.6K
Snippets
26
Records
33
Agent score
74%

What's inside AutoFaiss

  3. Generate memory-mapped indices

    master

    To minimize the memory footprint (at the cost of higher latency, typically >50ms), you can generate memory-mapped indices by setting should_be_memory_mappable=True in build_index. Note that this will always result in an IVF index, as only IVF indices support memory-mapping in Faiss.

    To load a memory-mapped index, use faiss.read_index with the appropriate flags:

    import faiss
faiss.read_index("my_index_folder/knn.index", faiss.IO_FLAG_MMAP | faiss.IO_FLAG_READ_ONLY)

  4. Install AutoFaiss from source

    master

    To install AutoFaiss, create a Python virtual environment, activate it, navigate to the deepr_python/autofaiss directory, and install the package in editable mode using pip install -e .. This allows you to use the CLI commands and run the provided notebooks.

    python3 -m venv .venvs/quantization_env
source .venvs/quantization_env/bin/activate
cd deepr_python/autofaiss
pip install -e .

  5. Install AutoFaiss

    master

    Install AutoFaiss using pip. It is recommended to use a virtual environment to manage dependencies.

    python -m venv .venv/autofaiss_env
source .venv/autofaiss_env/bin/activate
pip install -U pip
pip install autofaiss

  6. Use AutoFaiss with PySpark

    master

    AutoFaiss supports PySpark for two main distributed use cases:

    1. Building a large index in a distributed way: Useful when the index exceeds single-machine memory. You can control the number of splits using nb_indices_to_keep.
    2. Building partitioned indexes: If you have a partitioned dataset, you can build one index per partition in parallel using build_partitioned_indexes.

    Prerequisites:

    • Install pyspark: pip install pyspark.
    • Prepare your embeddings files.
    • Ensure a Spark session is created before calling AutoFaiss.

  7. Set up a distributed AutoFaiss cluster with PySpark

    master

    To generate an index from billions of embeddings, you can use PySpark to run AutoFaiss across multiple nodes. This requires setting up a Spark master node and multiple worker nodes (ideally on Ubuntu 20.04).

    1. Master Node Setup

    • Download and extract Spark (e.g., version 3.2.1).
    • Download the autofaiss.pex file, make it executable, and ensure it is available on all nodes.
    • If the master node is not directly accessible, use SSH tunneling to view the Spark UI at http://localhost:8080:
      ssh -L 8080:localhost:8080 -L 4040:localhost:4040 master_node

    2. Worker Node Setup

    • SSH Configuration: Ensure the master can SSH into workers without passwords. Use ssh-keyscan to add worker IPs to known_hosts and configure ~/.ssh/config for easy access.
    • Dependencies: Install required packages (OpenJDK, libgl1, etc.) using parallel-ssh.
    • AutoFaiss & Spark: Download autofaiss.pex and the Spark binaries onto all worker nodes.

    3. Cluster Lifecycle

    • Start Master: ./spark-3.2.1-bin-hadoop3.2/sbin/start-master.sh -p 7077
    • Start Workers: Use parallel-ssh to run ./spark-3.2.1-bin-hadoop3.2/sbin/start-worker.sh on all nodes, pointing them to the master's URI (e.g., spark://<master_ip>:7077).
    • Stop Cluster: Kill the Java processes (pkill java) on both master and workers to shut down the cluster.
    # Start master
./spark-3.2.1-bin-hadoop3.2/sbin/start-master.sh -p 7077

# Start workers via parallel-ssh
parallel-ssh -l $USER -i -h ips.txt './spark-3.2.1-bin-hadoop3.2/sbin/start-worker.sh -c 16 -m 28G "spark://172.31.35.188:7077"'

  10. Prepare embeddings for AutoFaiss

    master

    AutoFaiss reads embeddings from a directory of .npy files. If your dataset is large, you can split your embeddings into multiple parts. The files will be loaded in lexicographical order of their filenames.

    Workflow:

    1. Create a directory.
    2. Save your embeddings as np.float32 arrays using np.save into that directory.
    import os
import numpy as np

# Create a directory
embeddings_dir = "embeddings_folder"
os.makedirs(embeddings_dir, exist_ok=True)

# Create dummy embeddings
embeddings = np.float32(np.random.rand(4000, 100))

# Save in parts (lexicographical order: part1.npy, then part2.npy)
np.save(f"{embeddings_dir}/part1.npy", embeddings[:2000]) 
np.save(f"{embeddings_dir}/part2.npy", embeddings[2000:])

  11. Use build_index() for S3-based distributed build

    master

    You can also use build_index to build indices from multiple S3 paths. This is useful for large-scale datasets like LAION.

    index, index_infos = build_index(
    embeddings=[
        "s3://laion-us-east-1/embeddings/vit-l-14/laion2B-en/img_emb",
        "s3://laion-us-east-1/embeddings/vit-l-14/laion2B-multi/img_emb",
        "s3://laion-us-east-1/embeddings/vit-l-14/laion1B-nolang/img_emb"
    ],
    distributed="pyspark",
    max_index_memory_usage="200G",
    current_memory_available="24G",
    nb_indices_to_keep=10,
    file_format="npy",
    temporary_indices_folder="s3://laion-us-east-1/mytest/my_tmp_folder5",
    index_path="s3://laion-us-east-1/indices/vit-l-14/image/knn.index",
    index_infos_path="s3://laion-us-east-1/indices/vit-l-14/image/infos.json"
)
    index, index_infos = build_index(
    embeddings=["s3://laion-us-east-1/embeddings/vit-l-14/laion2B-en/img_emb","s3://laion-us-east-1/embeddings/vit-l-14/laion2B-multi/img_emb","s3://laion-us-east-1/embeddings/vit-l-14/laion1B-nolang/img_emb"],
    distributed="pyspark",
    max_index_memory_usage="200G",
    current_memory_available="24G",
    nb_indices_to_keep=10,
    file_format="npy",
    temporary_indices_folder="s3://laion-us-east-1/mytest/my_tmp_folder5",
    index_path="s3://laion-us-east-1/indices/vit-l-14/image/knn.index",
    index_infos_path="s3://laion-us-east-1/indices/vit-l-14/image/infos.json"
)

  12. Use build_index() for distributed index construction

    master

    To build an index across a Spark cluster, call build_index with distributed="pyspark". You must configure a SparkSession that correctly identifies the master node and sets the PYSPARK_PYTHON environment variable to the location of autofaiss.pex on your workers.

    Example: HDFS-based distributed build

    from autofaiss import build_index
from pyspark.sql import SparkSession
import os

def create_spark_session():
    # Path to autofaiss.pex must be identical on all workers
    os.environ['PYSPARK_PYTHON'] = "/home/ubuntu/autofaiss.pex"
    spark = (
        SparkSession.builder
        .config("spark.submit.deployMode", "client") \
        .config("spark.executorEnv.PEX_ROOT", "./.pex") \
        .config("spark.task.cpus", "16") \
        .config("spark.driver.port", "5678") \
        .config("spark.driver.blockManager.port", "6678") \
        .config("spark.driver.host", "172.31.35.188") \
        .config("spark.driver.bindAddress", "172.31.35.188") \
        .config("spark.executor.memory", "18G") \
        .config("spark.executor.memoryOverhead", "8G") \
        .config("spark.task.maxFailures", "100") \
        .master("spark://172.31.35.188:7077") \
        .appName("spark-stats") \
        .getOrCreate()
    )
    return spark

spark = create_spark_session()

index, index_infos = build_index(
    embeddings="hdfs://root/path/to/your/embeddings/folder",
    distributed="pyspark",
    file_format="parquet",
    max_index_memory_usage="16G",
    current_memory_available="24G",
    temporary_indices_folder="hdfs://root/tmp/distributed_autofaiss_indices",
    index_path="hdfs://root/path/to/your/index/knn.index",
    index_infos_path="hdfs://root/path/to/your/index/infos.json"
)
    from autofaiss import build_index
from pyspark.sql import SparkSession  # pylint: disable=import-outside-toplevel

from pyspark import SparkConf, SparkContext

def create_spark_session():
    # this must be a path that is available on all worker nodes
    
    os.environ['PYSPARK_PYTHON'] = "/home/ubuntu/autofaiss.pex"
    spark = (
        SparkSession.builder
        .config("spark.submit.deployMode", "client") \
        .config("spark.executorEnv.PEX_ROOT", "./.pex")
        #.config("spark.executor.cores", "16")
        #.config("spark.cores.max", "48") # you can reduce this number if you want to use only some cores ; if you're using yarn the option name is different, check spark doc
        .config("spark.task.cpus", "16")
        .config("spark.driver.port", "5678")
        .config("spark.driver.blockManager.port", "6678")
        .config("spark.driver.host", "172.31.35.188")
        .config("spark.driver.bindAddress", "172.31.35.188")
        .config("spark.executor.memory", "18G") # make sure to increase this if you're using more cores per executor
        .config("spark.executor.memoryOverhead", "8G")
        .config("spark.task.maxFailures", "100")
        .master("spark://172.31.35.188:7077") # this should point to your master node, if using the tunnelling version, keep this to localhost
        .appName("spark-stats")
        .getOrCreate()
    )
    return spark

spark = create_spark_session()

index, index_infos = build_index(
    embeddings="hdfs://root/path/to/your/embeddings/folder",
    distributed="pyspark",
    file_format="parquet",
    max_index_memory_usage="16G",
    current_memory_available="24G",
    temporary_indices_folder="hdfs://root/tmp/distributed_autofaiss_indices",
    index_path="hdfs://root/path/to/your/index/knn.index",
    index_infos_path="hdfs://root/path/to/your/index/infos.json"
)