Capsule 02: Scaling Strategies

Capsule description

Scaling a RAG system isn't just "adding more machines". First you need to understand which component limits performance: ingestion, retrieval, or generation. In this capsule you'll see the three types of scaling (vertical, horizontal, sharding), ready-to-use Docker Compose configurations, Python code for load testing and bottleneck detection, and a clear process for deciding when to apply each strategy.

1. Components that can limit performance

Before scaling, identify the bottleneck:

If you haven't measured a baseline for at least a week, don't scale yet. The next capsule covers monitoring in detail; here we assume you already know where the bottleneck is.

2. Vertical scaling (more CPU/RAM on the same node)

When to use it

• You're not yet saturating a node (CPU < 70%, RAM < 85% sustained).
• The index fits on a single machine.
• Your load is stable and predictable.
• You're in the early-stage or MVP phase.

How it works

You increase the resources of the same instance: more vCPUs, more RAM, faster disk. The advantage is simplicity: you don't introduce coordination between nodes, you don't change code, you just bump up the machine tier.

Trade-offs

When to stop scaling vertical

• Your cloud's largest tier no longer gives more.
• Concurrency is the problem (many simultaneous queries), not the power of a single request.
• The cost per node spikes relative to the benefit.

3. Horizontal scaling (replicas + load balancer)

When to use it

• You already have concurrency bottlenecks: many simultaneous queries saturate a node.
• Retrieval p95 rises during traffic peaks.
• A single node can't absorb the load even though it has free resources on average.

Typical architecture

                    ┌─────────────────┐
                    │  Load Balancer  │
                    │  (Nginx/HAProxy)│
                    └────────┬────────┘
                             │
              ┌──────────────┼──────────────┐
              │              │              │
              ▼              ▼              ▼
        ┌──────────┐  ┌──────────┐  ┌──────────┐
        │  RAG #1  │  │  RAG #2  │  │  RAG #3  │
        │ (Chroma) │  │ (Chroma) │  │ (Chroma) │
        └──────────┘  └──────────┘  └──────────┘
              │              │              │
              └──────────────┼──────────────┘
                             │
                    ┌────────▼────────┐
                    │ Vector DB       │
                    │ (Chroma/Qdrant) │
                    └─────────────────┘

Each replica of the RAG API shares the same vector database (or a read replica). The load balancer distributes requests among the replicas.

Docker Compose: replicas + Nginx

# docker-compose.horizontal.yaml
version: '3.8'

services:
  nginx:
    image: nginx:alpine
    ports:
      - "8000:80"
    volumes:
      - ./nginx.conf:/etc/nginx/nginx.conf:ro
    depends_on:
      - rag-api-1
      - rag-api-2
      - rag-api-3

  rag-api-1:
    build: .
    environment:
      - CHROMA_HOST=chroma
      - REPLICA_ID=1
    deploy:
      replicas: 1

  rag-api-2:
    build: .
    environment:
      - CHROMA_HOST=chroma
      - REPLICA_ID=2
    deploy:
      replicas: 1

  rag-api-3:
    build: .
    environment:
      - CHROMA_HOST=chroma
      - REPLICA_ID=3
    deploy:
      replicas: 1

  chroma:
    image: chromadb/chroma:latest
    volumes:
      - chroma_data:/chroma/chroma
    environment:
      - IS_PERSISTENT=TRUE

volumes:
  chroma_data: {}

nginx.conf for round-robin:

events { worker_connections 1024; }

http {
    upstream rag_backend {
        least_conn;  # Alternative option: round_robin
        server rag-api-1:8000;
        server rag-api-2:8000;
        server rag-api-3:8000;
    }

    server {
        listen 80;
        location / {
            proxy_pass http://rag_backend;
            proxy_connect_timeout 5s;
            proxy_read_timeout 60s;
        }
    }
}

Autoscaling with Docker Swarm (optional)

If you use Swarm, you can scale dynamically:

docker service scale rag_api=5

And in docker-compose:

rag-api:
  image: my-rag-api:latest
  deploy:
    replicas: 3
    resources:
      limits:
        cpus: '2'
        memory: 4G

4. Sharding (by tenant, date, or domain)

When to use it

• The index grows so much that a single unit can't keep latency stable.
• You have multi-tenancy and each tenant has very different volumes.
• You want to isolate load by domain (for example, by region or product).

Sharding criteria

Practical rule

Don't shard until the index and traffic can't be sustained by a single reasonable unit. Sharding introduces complexity: query routing, rebalancing, maintaining multiple collections or clusters.

Routing in code

def get_shard_for_query(tenant_id: str, query_date: date | None) -> str:
    """Determine which shard to use based on tenant and optionally date."""
    if tenant_id:
        # Sharding by tenant (e.g. hash for distribution)
        h = hash(tenant_id) % 8
        return f"shard_tenant_{h}"
    if query_date:
        # Sharding by quarter
        quarter = (query_date.month - 1) // 3 + 1
        return f"shard_{query_date.year}_q{quarter}"
    return "shard_default"

Docker Compose: multiple ChromaDB shards

# docker-compose.sharding.yaml
version: '3.8'

services:
  chroma-shard-0:
    image: chromadb/chroma:latest
    volumes:
      - chroma_shard_0:/chroma/chroma
    environment:
      - IS_PERSISTENT=TRUE

  chroma-shard-1:
    image: chromadb/chroma:latest
    volumes:
      - chroma_shard_1:/chroma/chroma
    environment:
      - IS_PERSISTENT=TRUE

  chroma-shard-2:
    image: chromadb/chroma:latest
    volumes:
      - chroma_shard_2:/chroma/chroma
    environment:
      - IS_PERSISTENT=TRUE

  rag-api:
    build: .
    environment:
      - CHROMA_SHARD_0=chroma-shard-0:8000
      - CHROMA_SHARD_1=chroma-shard-1:8000
      - CHROMA_SHARD_2=chroma-shard-2:8000
    depends_on:
      - chroma-shard-0
      - chroma-shard-1
      - chroma-shard-2

volumes:
  chroma_shard_0: {}
  chroma_shard_1: {}
  chroma_shard_2: {}

5. Load testing with Python

Basic script with Locust or requests

# load_test_rag.py
"""
Load test for a RAG API. Usage: pip install locust && locust -f load_test_rag.py
Or run directly: python load_test_rag.py
"""
import asyncio
import time
import statistics
from concurrent.futures import ThreadPoolExecutor
import httpx

BASE_URL = "http://localhost:8000"  # Adjust to your API


def single_query(client: httpx.Client, query: str = "What is RAG?") -> float:
    """Run a query and return the latency in seconds."""
    start = time.perf_counter()
    resp = client.post(
        f"{BASE_URL}/v1/query",
        json={"query": query, "top_k": 5},
        timeout=30.0,
    )
    elapsed = time.perf_counter() - start
    resp.raise_for_status()
    return elapsed


def run_load_test(
    num_requests: int = 100,
    concurrency: int = 10,
    query: str = "What is RAG?",
) -> dict:
    """
    Run a load test with multiple workers.
    Returns latency percentiles and throughput.
    """
    latencies: list[float] = []

    def worker(_: int) -> float:
        with httpx.Client() as client:
            return single_query(client, query)

    with ThreadPoolExecutor(max_workers=concurrency) as ex:
        futures = [ex.submit(worker, i) for i in range(num_requests)]
        latencies = [f.result() for f in futures]

    latencies_sorted = sorted(latencies)
    n = len(latencies_sorted)

    return {
        "count": n,
        "p50_ms": statistics.median(latencies) * 1000,
        "p95_ms": latencies_sorted[int(n * 0.95)] * 1000 if n else 0,
        "p99_ms": latencies_sorted[int(n * 0.99)] * 1000 if n else 0,
        "mean_ms": statistics.mean(latencies) * 1000,
        "throughput_rps": n / (max(latencies) or 1),
    }


if __name__ == "__main__":
    result = run_load_test(num_requests=200, concurrency=20)
    print("Load test results:")
    for k, v in result.items():
        print(f"  {k}: {v}")

Usage

pip install httpx
python load_test_rag.py

Adjust BASE_URL, num_requests, and concurrency to your environment. Compare p50, p95, and p99 before and after scaling.

6. Bottleneck detection

Script to measure per component

# bottleneck_detection.py
"""
Measure latency per component: retrieval vs generation vs total.
Helps identify whether the bottleneck is in the vector DB or the LLM.
"""
import time
import httpx

BASE_URL = "http://localhost:8000"


def measure_retrieval_only(query: str) -> float:
    """Retrieval only (no LLM). Your API must expose a /retrieve endpoint."""
    start = time.perf_counter()
    resp = httpx.post(
        f"{BASE_URL}/v1/retrieve",
        json={"query": query, "top_k": 5},
        timeout=30.0,
    )
    resp.raise_for_status()
    return (time.perf_counter() - start) * 1000


def measure_full_rag(query: str) -> tuple[float, float]:
    """
    Full query. Assumes the API returns retrieval_time_ms and generation_time_ms.
    If not, you'll have to instrument your API to return these fields.
    """
    start = time.perf_counter()
    resp = httpx.post(
        f"{BASE_URL}/v1/query",
        json={"query": query, "top_k": 5},
        timeout=60.0,
    )
    resp.raise_for_status()
    total_ms = (time.perf_counter() - start) * 1000
    data = resp.json()
    retrieval_ms = data.get("retrieval_time_ms", 0)
    generation_ms = data.get("generation_time_ms", 0)
    return total_ms, retrieval_ms, generation_ms


def run_bottleneck_check(num_samples: int = 20, query: str = "What is RAG?"):
    """Run measurements and report where the bottleneck is."""
    retrieval_times = [measure_retrieval_only(query) for _ in range(num_samples)]
    retrieval_avg = sum(retrieval_times) / len(retrieval_times)

    full_times = [measure_full_rag(query) for _ in range(num_samples)]
    total_avg = sum(t[0] for t in full_times) / len(full_times)
    ret_avg = sum(t[1] for t in full_times) / len(full_times)
    gen_avg = sum(t[2] for t in full_times) / len(full_times)

    print("=== Bottleneck analysis ===\n")
    print(f"Pure retrieval (avg): {retrieval_avg:.0f} ms")
    print(f"Total RAG (avg):      {total_avg:.0f} ms")
    print(f"Retrieval in RAG:     {ret_avg:.0f} ms")
    print(f"Generation in RAG:    {gen_avg:.0f} ms")
    print()

    retrieval_pct = (ret_avg / total_avg * 100) if total_avg else 0
    generation_pct = (gen_avg / total_avg * 100) if total_avg else 0

    if retrieval_pct > 50:
        print("→ The bottleneck is in RETRIEVAL. Consider: more replicas, sharding, or index tuning.")
    elif generation_pct > 50:
        print("→ The bottleneck is in GENERATION (LLM). Consider: a faster model, cache, or more workers.")
    else:
        print("→ Latency is spread out. Check network, I/O, or other components.")

Example endpoint to instrument

Your RAG API should return per-stage times. Conceptual example:

# In your /v1/query endpoint
start_retrieval = time.perf_counter()
docs = vector_store.similarity_search(query, k=5)
retrieval_time_ms = (time.perf_counter() - start_retrieval) * 1000

start_generation = time.perf_counter()
response = llm.generate(context=docs, query=query)
generation_time_ms = (time.perf_counter() - start_generation) * 1000

return {
    "answer": response,
    "retrieval_time_ms": retrieval_time_ms,
    "generation_time_ms": generation_time_ms,
}

7. Practical decision rule

1. If you're not yet saturating a node: scale vertical first.
2. If you already have concurrency bottlenecks: move to horizontal (replicas + LB).
3. If the index grows a lot and latency degrades: apply sharding with clear criteria (tenant, date, domain).
4. Always: measure baseline → intervene once → measure again.

8. Signs that you need to scale

9. Suggested scaling process

1. Measure a baseline for at least a week (p50, p95, p99, memory, CPU).
2. Identify the bottleneck with the bottleneck detection script or APM metrics.
3. Run a single intervention (only vertical OR only horizontal OR only sharding).
4. Repeat the measurement and compare against the baseline.
5. If there isn't enough improvement, repeat the cycle with another scaling type or another component.

10. Practical exercises

Exercise 1: Vertical vs horizontal decision

Case: p95 went from 220 ms to 480 ms in two weeks; memory at 85% during peak hours.

Questions:

• Would you scale vertical or horizontal first?
• Which metric would you use to validate success?
• What condition would trigger a second intervention?

Exercise 2: Interpret load test results

You have these results before and after adding 2 replicas:

Question: Was the horizontal scaling effective? What next step would you consider?

Exercise 3: Choose a sharding criterion

You have a system with 50 tenants. 3 of them represent 70% of the document volume and the traffic.

Question: Would you shard by tenant? If so, how would you distribute the shards?

Exercise 4: Docker Compose for 5 replicas

Task: Modify this capsule's docker-compose.horizontal.yaml to use a single rag-api service with 5 replicas instead of 3 separate services, keeping Nginx as the load balancer.

version: '3.8'

services:
  nginx:
    image: nginx:alpine
    ports:
      - "8000:80"
    volumes:
      - ./nginx.conf:/etc/nginx/nginx.conf:ro
    depends_on:
      - rag-api

  rag-api:
    build: .
    environment:
      - CHROMA_HOST=chroma
    deploy:
      replicas: 5

  chroma:
    image: chromadb/chroma:latest
    volumes:
      - chroma_data:/chroma/chroma
    environment:
      - IS_PERSISTENT=TRUE

volumes:
  chroma_data: {}

upstream rag_backend {
    least_conn;
    server rag-api:8000;  # Docker resolves to all replicas
}

Exercise 5: Bottleneck script with timeout

Task: Modify measure_full_rag in the bottleneck script so that, if the request exceeds 45 seconds, it records a timeout and returns (45000, 0, 0) instead of failing. Add a timeout counter at the end of the report.

def measure_full_rag_with_timeout(
    query: str, timeout_sec: float = 45.0
) -> tuple[float, float, float, bool]:
    """
    Full query with a timeout. Returns (total_ms, retrieval_ms, generation_ms, timed_out).
    """
    start = time.perf_counter()
    try:
        resp = httpx.post(
            f"{BASE_URL}/v1/query",
            json={"query": query, "top_k": 5},
            timeout=timeout_sec,
        )
        resp.raise_for_status()
        total_ms = (time.perf_counter() - start) * 1000
        data = resp.json()
        return (
            total_ms,
            data.get("retrieval_time_ms", 0),
            data.get("generation_time_ms", 0),
            False,
        )
    except (httpx.TimeoutException, httpx.ConnectTimeout):
        return (timeout_sec * 1000, 0, 0, True)


def run_bottleneck_check(num_samples: int = 20, query: str = "What is RAG?"):
    # ... setup ...
    timeouts = 0
    for _ in range(num_samples):
        total_ms, ret_ms, gen_ms, timed_out = measure_full_rag_with_timeout(query)
        if timed_out:
            timeouts += 1
        # ... accumulate in lists ...
    print(f"\nTimeouts: {timeouts}/{num_samples}")

Exercise 6: When NOT to shard

Question: Give two concrete situations where sharding would be premature or counterproductive.

11. Scaling troubleshooting

"I scaled and it didn't improve"

You probably didn't attack the right bottleneck. If the problem is in retrieval and you only scaled the API (or vice versa), you won't see improvement. Use the bottleneck detection script to confirm whether the time goes to retrieval, generation, or network. Make sure you have per-component metrics before scaling.

"Latency improves but cost spikes"

Introduce autoscaling limits (min/max replicas) and review caches. Many repeated queries can be served from cache instead of going to the vector DB and the LLM. Also check whether you're over-provisioning: reducing the vertical tier or the number of replicas during off-peak hours can lower costs without affecting peaks.

"I don't know whether to shard yet"

Shard when the index and traffic can no longer be sustained by a single reasonable unit: retrieval p95 grows even though you have free resources, or the index is so large that a single machine can't keep it in RAM efficiently. If you can still scale vertical or horizontal with good ROI, wait.

"The replicas don't distribute load well"

Review the load balancer policy. round_robin treats all replicas equally; least_conn sends more traffic to the least busy ones. If some requests are much heavier than others, least_conn usually works better. Also verify there isn't a single contention point (for example, a single ChromaDB) saturating all the replicas.

"After sharding, some queries are slow"

You may be querying several shards and merging results. If the routing isn't precise (for example, a query needs data from 3 shards), the latency will be the sum of the 3 queries. Optimize the routing so most queries hit a single shard, or consider whether the sharding criterion is the right one.

12. Summary

• Scaling well starts with observation (baseline, per-component metrics), not intuition.
• Vertical: when you're not saturating a node; simple but with a ceiling.
• Horizontal: when concurrency is the bottleneck; replicas + load balancer.
• Sharding: when the index grows too much; by tenant, date, or domain.
• Run a single intervention at a time and measure again before the next.
• Use load testing and bottleneck detection scripts before deciding what to scale.
• The next capsule gives you the metrics and tools to observe correctly in production.

13. Additional resources

• Microservice Trade-offs (Martin Fowler) – When scaling and distribution add complexity.
• ChromaDB Production – ChromaDB deployment and scaling.
• Locust – Load Testing – Python load testing tool.
• Nginx Load Balancing – Load balancing configuration.
• Docker Compose Scale – Replicas with Compose.
• Scaling Vector Search (Pinecone) – Pinecone's approach to scaling indexes.
• Qdrant Sharding – Sharding in Qdrant.
• The Twelve-Factor App – Principles for scalable and portable apps.

Estimated time: 25–35 minutes
Next: 03-monitoring-and-observability-2.md