Mastering Testing and Optimization of AI Pipelines 

Below you will find some suggested materials on these topics, which we believe may help you complete the final activity in the lesson. Please note that using these resources is entirely optional, and you are welcome to explore any other sources or courses at your own discretion.

Introduction (2 Minutes)

Opening Statement:

• Welcome to this lecture on Performance Testing and Optimization of Vector Searches.

• In today’s session, we’ll cover:

  1. Methods for testing the performance of vector searches.

  2. Tools for measuring latency and throughput in vector databases.

  3. Techniques for optimizing vector search for large-scale datasets.

Why is this important?

• Vector search underpins many modern AI systems, including recommendation engines, semantic search, and Retrieval-Augmented Generation (RAG).

• Ensuring these searches are fast and accurate is critical for user satisfaction and system reliability.

Part 1: Methods for Testing Vector Search Performance (5 Minutes)

Key Performance Metrics:

1. Latency:

   • Measures the time it takes to execute a query.

   • Importance: Users expect near-instantaneous results, especially in real-time systems.

2. Throughput:

   • Measures the number of queries processed per second.

   • Importance: Determines how well the system handles high traffic.

3. Recall and Precision:

   • Recall: Measures the proportion of relevant items retrieved.

   • Precision: Measures the proportion of retrieved items that are relevant.

   • Importance: Balances performance with search quality.

Common Testing Scenarios:

• Single Query Latency: Test how quickly a single query executes.

• Batch Query Performance: Measure how the system performs under bulk requests.

• Scaling Tests: Evaluate performance as data size or query volume increases.

Example Setup:

• Use a vector database like FAISS or Pinecone.

• Perform queries on a dataset of 1M+ embeddings.

• Collect metrics for both single and batch queries.

Part 2: Tools for Measuring Latency and Throughput (5 Minutes)

1. Python Libraries:

time:
python
Copy code
import time

start = time.time()

# Perform vector search

end = time.time()

print(f"Latency: {end - start} seconds")

• perf_counter: High-resolution timer for more accurate measurements.

2. Built-In Database Metrics:

• FAISS:

  • Provides logs for query execution times.

  • Use search method and log timing manually.

• Pinecone:

  • Includes detailed latency and throughput metrics in its dashboard.

3. Profiling and Benchmarking Tools:

• Locust: Load testing tool to simulate high query volumes.

  • Simulate thousands of concurrent queries to measure throughput.

• Apache JMeter: Useful for evaluating query latency under different conditions.

4. Monitoring Tools:

• Prometheus + Grafana:

  • Collect query performance metrics.

  • Visualize latency and throughput over time.

  • Example Dashboard: Query Latency (ms), Query Count per Second.

Part 3: Optimizing Vector Search for Large-Scale Datasets (8 Minutes)

Challenges with Large-Scale Datasets:

• High Latency: As the dataset grows, search times can increase significantly.

• Memory Constraints: Embedding vectors require significant storage, which can strain resources.

• Accuracy vs. Speed Trade-offs: Fast searches may compromise recall and precision.

Optimization Techniques:

1. Index Selection:

   • Choose the right index for your use case.

   • Example: Use FAISS’s IVF (Inverted File Index) for large datasets.

2. Quantization:

   • Reduce memory usage by approximating embeddings.

   • Example: Use FAISS’s PQ (Product Quantization).

3. Hybrid Approaches:

   • Combine vector search with traditional filtering.

   • Example: Use metadata filters before executing a vector search.

4. Sharding:

   • Split the dataset across multiple nodes to distribute the load.

   • Improves both memory usage and query speed.

5. Batching Queries:

   • Aggregate multiple queries into a single request.

   • Reduces overhead from repeated operations.

Example: FAISS Optimization

python

Copy code

import faiss

# Build an IVF index

dimension = 128

nlist = 100  # Number of clusters

quantizer = faiss.IndexFlatL2(dimension)

index = faiss.IndexIVFFlat(quantizer, dimension, nlist)

# Train the index with your dataset

index.train(embedding_dataset)

index.add(embedding_dataset)

Best Practices:

• Always benchmark changes to ensure optimizations improve performance.

• Monitor query patterns to identify bottlenecks and adjust index configurations.

• Regularly update and re-train indices to reflect changes in the dataset.

Conclusion and Key Takeaways (2 Minutes)

Recap:

1. Testing Performance: Measure latency, throughput, recall, and precision to ensure your system meets requirements.

2. Using Tools: Utilize Python libraries, monitoring tools, and benchmarking frameworks for detailed performance insights.

3. Optimizing Searches: Apply techniques like quantization, sharding, and batching to improve performance for large datasets.

Call to Action:

• Start testing and optimizing your vector search pipeline today.

• Experiment with different indices and configurations to find the best fit for your use case.

Objective

By the end of this activity, you will:

1. Measure the latency of embedding retrieval and vector searches.

2. Optimize the vector database configuration to achieve faster query times.

Step-by-Step Instructions

Step 1: Set Up the Environment

Install Required Libraries:
bash
Copy code
pip install faiss-cpu numpy

1. Prepare the Dataset:

   • Use a sample dataset of embeddings.

   • Example: 1M embeddings of 128 dimensions generated randomly.

python
Copy code

import numpy as np

# Generate random embeddings

dimension = 128

num_embeddings = 1000000

embeddings = np.random.random((num_embeddings, dimension)).astype('float32')

# Save embeddings for reuse

np.save("embeddings.npy", embeddings)

Step 2: Measure Latency of Embedding Retrieval

Load the Embeddings and Initialize FAISS:
python
Copy code

import faiss

import time

import numpy as np

# Load embeddings

embeddings = np.load("embeddings.npy")

# Create a FlatL2 index

index = faiss.IndexFlatL2(embeddings.shape[1])

index.add(embeddings)

1. Define a Function to Measure Latency:
   python
   Copy code

def measure_latency(query_vector, k=10):

    start = time.perf_counter()

    distances, indices = index.search(query_vector, k)

    end = time.perf_counter()

    latency = (end - start) * 1000  # Convert to milliseconds

    return latency, indices

# Create a random query vector

query_vector = np.random.random((1, embeddings.shape[1])).astype('float32')

# Measure latency

latency, results = measure_latency(query_vector)

print(f"Latency: {latency:.2f} ms")

2. Record Results for Different Query Sizes:

   • Test for 1, 10, and 100 query vectors.

python
Copy code

for n_queries in [1, 10, 100]:

    query_vectors = np.random.random((n_queries, embeddings.shape[1])).astype('float32')

    latency, _ = measure_latency(query_vectors)

    print(f"Latency for {n_queries} queries: {latency:.2f} ms")

Step 3: Optimize Vector Search

1. Switch to an Inverted Index (IVF):

   • Use FAISS’s IVF (Inverted File Index) for faster searches on large datasets.

python
Copy code

# Create an IVF index

num_clusters = 100  # Number of clusters

quantizer = faiss.IndexFlatL2(dimension)  # Quantizer for clustering

ivf_index = faiss.IndexIVFFlat(quantizer, dimension, num_clusters)

# Train and add embeddings to the IVF index

ivf_index.train(embeddings)

ivf_index.add(embeddings)

# Enable clustering for faster searches

ivf_index.nprobe = 10  # Number of clusters to search

2. Compare Performance with FlatL2:
   python
   Copy code

def compare_performance(query_vector):

    # FlatL2 search

    start_flat = time.perf_counter()

    _, flat_results = index.search(query_vector, 10)

    end_flat = time.perf_counter()

    # IVF search

    start_ivf = time.perf_counter()

    _, ivf_results = ivf_index.search(query_vector, 10)

    end_ivf = time.perf_counter()

    print(f"FlatL2 Latency: {(end_flat - start_flat) * 1000:.2f} ms")

    print(f"IVF Latency: {(end_ivf - start_ivf) * 1000:.2f} ms")

# Test performance

compare_performance(query_vector)

Step 4: Additional Optimizations

1. Apply Product Quantization (PQ):

   • Use PQ to reduce memory usage and speed up searches.

python
Copy code

# Create a PQ index

pq_index = faiss.IndexPQ(dimension, 8, 8)  # 8 sub-quantizers

pq_index.train(embeddings)

pq_index.add(embeddings)

# Measure PQ performance

def measure_pq_latency(query_vector):

    start = time.perf_counter()

    _, pq_results = pq_index.search(query_vector, 10)

    end = time.perf_counter()

    print(f"PQ Latency: {(end - start) * 1000:.2f} ms")

measure_pq_latency(query_vector)

2. Optimize Cluster Parameters:

   • Experiment with the nprobe parameter to balance speed and recall.

python
Copy code

for nprobe in [1, 5, 10, 20]:

    ivf_index.nprobe = nprobe

    latency, _ = measure_latency(query_vector)

    print(f"nprobe={nprobe}, Latency: {latency:.2f} ms")

Expected Outcomes

1. Latency Measurements:

   • Report on the latency for FlatL2, IVF, and PQ indices under various query sizes.

2. Performance Comparison:

   • Document the trade-offs between latency, recall, and memory usage for each index type.

3. Optimized Configuration:

   • Identify the optimal parameters (e.g., nprobe) for your dataset and workload.

Objective

• Write a Python script to benchmark the performance of FAISS on a large dataset, measuring key metrics like latency and recall, and document the findings.

Step-by-Step Instructions

Step 1: Prepare the Dataset

1. Generate or Load a Large Dataset:

   • Use a dataset with 1M+ embeddings of 128 dimensions. You can either generate random embeddings or load a real-world dataset.

python
Copy code

import numpy as np

# Generate random embeddings

dimension = 128

num_embeddings = 1000000

embeddings = np.random.random((num_embeddings, dimension)).astype('float32')

# Save the dataset for reuse

np.save("embeddings.npy", embeddings)

2. Generate Query Embeddings:

   • Create a smaller set of query vectors to test the search performance.

python
Copy code

num_queries = 100

query_embeddings = np.random.random((num_queries, dimension)).astype('float32')

np.save("query_embeddings.npy", query_embeddings)

Step 2: Initialize FAISS and Benchmark

Set Up FAISS:
python
Copy code

import faiss

import time

# Load the dataset

embeddings = np.load("embeddings.npy")

# Create a FlatL2 index

dimension = embeddings.shape[1]

index = faiss.IndexFlatL2(dimension)

index.add(embeddings)

1. Benchmark Latency and Recall:

   • Define functions to measure latency and recall for search queries.

python
Copy code

def measure_latency(index, queries, k=10):

    start = time.perf_counter()

    _, indices = index.search(queries, k)

    end = time.perf_counter()

    latency = (end - start) * 1000  # Convert to milliseconds

    return latency

def measure_recall(index, queries, ground_truth, k=10):

    _, indices = index.search(queries, k)

    recall = 0

    for query_idx, gt in enumerate(ground_truth):

        recall += len(set(gt).intersection(set(indices[query_idx]))) / len(gt)

    return recall / len(queries)

# Generate ground truth using brute-force search

ground_truth = []

for query in query_embeddings:

    distances = ((embeddings - query) ** 2).sum(axis=1)

    ground_truth.append(np.argsort(distances)[:10])

3. Run Benchmarks:

   • Measure the latency and recall of the FlatL2 index.

python
Copy code

queries = np.load("query_embeddings.npy")

# Measure latency

flat_latency = measure_latency(index, queries)

print(f"FlatL2 Latency: {flat_latency:.2f} ms")

# Measure recall

flat_recall = measure_recall(index, queries, ground_truth)

print(f"FlatL2 Recall: {flat_recall:.2%}")

Step 3: Optimize with IVF Index

Set Up an Inverted File Index (IVF):
python
Copy code

# Create an IVF index

num_clusters = 100  # Number of clusters

quantizer = faiss.IndexFlatL2(dimension)

ivf_index = faiss.IndexIVFFlat(quantizer, dimension, num_clusters)

ivf_index.train(embeddings)

ivf_index.add(embeddings)

1. Benchmark the IVF Index:

   • Compare the latency and recall for different nprobe values.

python
Copy code

for nprobe in [1, 5, 10, 20]:

    ivf_index.nprobe = nprobe

    latency = measure_latency(ivf_index, queries)

    recall = measure_recall(ivf_index, queries, ground_truth)

    print(f"IVF Latency (nprobe={nprobe}): {latency:.2f} ms")

    print(f"IVF Recall (nprobe={nprobe}): {recall:.2%}")

Step 4: Document Findings

Template for Results Documentation:

Metric

FlatL2 Index

IVF (nprobe=5)

IVF (nprobe=10)

Latency (ms)

X ms

Y ms

Z ms

Recall (%)

X%

Y%

Z%

Memory Usage (MB)

X MB

Y MB

Z MB

Example Insights:

1. FlatL2 Performance:

   • Latency: 150ms per query.

   • Recall: 100%.

   • Memory Usage: High.

2. IVF Performance:

   • nprobe=5: Balanced latency (30ms) and recall (95%).

   • nprobe=10: Improved recall (98%) at the cost of higher latency (60ms).

Recommendations:

• Use FlatL2 for small datasets or when high recall is critical.

• Use IVF with nprobe=10 for large datasets to balance speed and accuracy.

Step 5: Save and Share Results

Save Results to CSV:
python
Copy code

import pandas as pd

results = [

    {"Index Type": "FlatL2", "Latency (ms)": flat_latency, "Recall (%)": flat_recall * 100},

    {"Index Type": "IVF (nprobe=5)", "Latency (ms)": 30, "Recall (%)": 95},

    {"Index Type": "IVF (nprobe=10)", "Latency (ms)": 60, "Recall (%)": 98},

]

results_df = pd.DataFrame(results)

results_df.to_csv("faiss_benchmark_results.csv", index=False)

print("Results saved to faiss_benchmark_results.csv")

1. Share Findings:

   • Include the CSV file and insights in your project documentation.

Expected Outcomes

1. Performance Metrics:

   • A clear comparison of latency and recall for different FAISS index types.

2. Optimized Configuration:

   • Recommendations on index and parameter configurations for your dataset.

3. Documentation:

   • Comprehensive findings in a CSV or Markdown format for future reference.