⏳Rework the batching system to speedup sklearn vectors

ADR-011: Dynamic Batching System for Vector Operations

Status

Proposed

Context

The current saan-multivectors architecture requires vectors to implement a singular embed() method that processes one value at a time:

async def embed(self, value: str, query: bool = False) -> np.ndarray:
    # Process single value
    return embedding_vector

This design creates significant performance bottlenecks for vectors that could benefit from vectorization, particularly:

1. API-based vectors (OpenAI, Cohere, Anthropic): Network overhead dominates when making individual requests instead of batch API calls
2. GPU-accelerated vectors: GPU utilization is poor with single-item processing; batch processing achieves 10-50x speedups
3. Statistical vectors: Many algorithms (like transformers) are inherently designed for batch processing
4. I/O bound operations: Database lookups, file reads, and other I/O operations benefit from batching

Precedent: Async Context Manager Pattern

Since ADR-010, we've successfully used async context managers to extend vector capabilities without breaking API compatibility. The metadata provider system demonstrates that complex functionality can be added transparently:

# Complex corpus statistics access through simple context
async with with_metadata_provider(provider) as ctx:
    # Vector.embed() automatically accesses corpus metadata
    embedding = await vector.embed("query text")

This pattern proved that we can enhance vector behavior while maintaining API simplicity and backward compatibility.

Performance Impact Analysis

Current measurements from production workloads show:

• API vectors: 300ms average latency for single requests vs 50ms for 32-item batches (6x improvement)
• Transformer vectors: Single embedding ~100ms vs batch of 32 ~120ms total (26x per-item improvement)
• Memory usage: Batch processing reduces memory allocation overhead by ~40%
• GPU utilization: Increases from 15% to 85% for transformer-based vectors

Decision

We will implement a Context-Aware Transparent Batching system that automatically batches vector operations for vectors that declare batching capability through an embed_batch() method. The system uses the presence of this method to determine batching eligibility, ensuring complete backward compatibility.

Core Architecture

The system introduces a DynamicVectorBatcher accessible through an async context manager that transparently intercepts and batches embed() calls for capable vectors:

async with with_dynamic_batching(max_batch_size=32, max_wait_time=0.1) as batcher:
    # Only vectors with embed_batch() methods are automatically batched
    embeddings = await asyncio.gather(*[
        vector.embed(text) for text in documents
    ])

Implementation Pattern

Vectors declare batching capability by implementing embed_batch() and checking for active batchers:

class EmbeddingVector:
    async def embed(self, value: str, query: bool = False, **kwargs) -> np.ndarray:
        # Check if there's an active batcher
        batcher = get_current_batcher()
        if batcher:
            return await batcher.process_embed(self, value, query=query, **kwargs)
        
        # No batcher - direct implementation
        return np.array([0, 0, 0, 0, 0])
    
    # Presence of this method signals batching capability
    async def embed_batch(self, values: List[str], query: bool = False, **kwargs) -> List[np.ndarray]:
        return [np.array([0, 0, 0, 0, 0]) for _ in values]

Legacy vectors without embed_batch() or batcher checks continue working unchanged:

class LegacyVector:
    async def embed(self, value: str, query: bool = False, **kwargs) -> np.ndarray:
        # Simple implementation - no batching, no changes needed
        return np.array([1, 1, 1, 1, 1])

Key Design Elements

1. Capability Detection: The presence of embed_batch() method determines batching eligibility
2. Opt-in Batching: Only vectors that implement both batcher checking and embed_batch() participate in batching
3. Complete Backward Compatibility: Vectors without embed_batch() work exactly as before
4. Transparent Operation: Batching happens automatically when context is active
5. Context Isolation: Different application components can use different batching configurations

Rationale

Technical Justification

Clear Capability Signaling: Using embed_batch() presence as a batching capability indicator provides clear, explicit opt-in semantics without requiring additional metadata or configuration.

Zero Breaking Changes: Vectors without embed_batch() continue working exactly as before, ensuring complete backward compatibility across all existing codebases.

Performance Optimization: Vectors with native batching achieve maximum performance benefits, while the system degrades gracefully for vectors without batch capabilities.

Simple Mental Model: Developers can easily understand that implementing embed_batch() enables batching, with no hidden behavior or complex configuration.

Systems Thinking Considerations

This design reflects systems thinking principles by:

• Emergent Performance: Batch efficiency emerges from the interaction of multiple embed calls within a context
• Clear Boundaries: Explicit capability signaling creates clear boundaries between batching-capable and legacy vectors
• Adaptive Behavior: System adapts behavior based on vector capabilities without requiring external configuration
• Minimal Coupling: Vectors and batching system interact through well-defined, minimal interfaces

Consequences

Positive

• Zero Migration Risk: All existing vectors continue working without any changes
• Dramatic Performance Gains: 6-26x speedups for vectors implementing embed_batch()
• Simple Adoption Model: Clear path to batching benefits through embed_batch() implementation
• Transparent Operation: Performance improvements happen automatically when batching context is active
• Clean Architecture: Batching concerns completely separated from embedding logic

Negative

• Dual Implementation: Batching-capable vectors must implement both embed() and embed_batch() methods
• Timing Unpredictability: Batch processing introduces variable latency due to wait times
• Memory Overhead: Batching queues consume additional memory for pending requests

Neutral

• Optional Feature: Batching is completely optional - vectors work fine without it
• Configuration Tuning: Optimal batch sizes and timeouts may require experimentation
• Context Management: Developers must understand when and how to use batching contexts

Alternatives Considered

Option 1: Automatic Vector Detection

• Pros: No vector changes needed, completely transparent
• Cons: Complex detection logic, potential for false positives, unclear capability signaling
• Why rejected: Detection heuristics are unreliable and make the system unpredictable

Option 2: Configuration-Based Batching

• Pros: Explicit control, no vector changes
• Cons: External configuration complexity, poor developer experience, maintenance burden
• Why rejected: Requires maintaining vector capability lists external to the code

Option 3: Decorator Pattern

• Pros: Explicit opt-in, clean separation
• Cons: Requires vector decoration, not transparent, complex setup
• Why rejected: Requires manual decoration of each vector instance, poor developer experience

Implementation Notes

Async Batching Flow

The following sequence diagram illustrates how the dynamic batching system coordinates multiple concurrent embed requests:

sequenceDiagram
    participant App as Application
    participant V1 as Vector Instance
    participant API as External API/GPU
    participant f as future results
    participant B as DynamicVectorBatcher
    participant Q as Batch Queue
    participant T as timer
    Note over App,T: Batching Context Active
    alt no batch embed method
        %% First embed call
        App->>+V1: await embed("text1")
        V1->>+B: process_embed(self, "text1")
        B-->>+V1: embed("text1)
        V1->>+API: atomic API call
        Note over f: No batching, pass through behavior
        API-->>-V1: result1
        V1 ->>-B: result1
        B -->> App: result1
    else batch embed method detected
        %% First embed call
        App->>+V1: await embed("text1")
        V1->>+B: process_embed(self, "text1")
        B->>+Q: Add to queue [batch_key]
        Note over Q: Queue: ["text1"]
        Q-->>-B: Queue size: 1
        B-->>+f: Future<result1>
        B->>+T: Start timer (max_wait_time)
        Note over API: No immediate processing, we wait for <br>the opportunity to batch with future calls <br>for a set time (or until the queue is full)
        %% Second embed call (concurrent)
        App->>+V1: await embed("text2")
        V1->>+B: process_embed(self, "text2")
        B->>+Q: Add to queue [same batch_key]
        Note over Q: Queue: ["text1", "text2"]
        Q-->>-B: Queue size: 2
        B-->>f: Future<result2>

        
        T->>-B: Cancel timer (times up!)
        B->>+Q: Extract batch items
        Q-->>-B: items[1..]
        Note over Q: Queue: []
        B->>+V1: embed_batch(["text1", ...])
        V1->>+API: Batch API call
        API-->>-V1: [result1,...]
        V1-->>-B: [result1,...]
        %% Resolve all futures
        Note over App: All futures resolve
        loop over (future, result) pairs
            B->>f: future.set_result(result)
            f-->>App: result
        end
    end

Context Management

from contextvars import ContextVar
from typing import Optional
import asyncio

_batching_context: ContextVar[Optional['DynamicVectorBatcher']] = ContextVar(
    'vector_batcher', default=None
)

def get_current_batcher() -> Optional['DynamicVectorBatcher']:
    """Get the current batching context."""
    return _batching_context.get()

@asynccontextmanager
async def with_dynamic_batching(
    max_batch_size: int = 32,
    max_wait_time: float = 0.1
):
    """Context manager for transparent vector batching."""
    batcher = DynamicVectorBatcher(max_batch_size, max_wait_time)
    
    token = _batching_context.set(batcher)
    try:
        yield batcher
    finally:
        await batcher.shutdown()
        _batching_context.reset(token)

Core Batcher Implementation

from collections import defaultdict
from dataclasses import dataclass
from typing import Dict, List, Any
import time

@dataclass
class BatchItem:
    value: str
    kwargs: dict
    future: asyncio.Future

class DynamicVectorBatcher:
    def __init__(self, max_batch_size: int = 32, max_wait_time: float = 0.1):
        self.max_batch_size = max_batch_size
        self.max_wait_time = max_wait_time
        self._queues: Dict[int, List[BatchItem]] = defaultdict(list)
        self._locks: Dict[int, asyncio.Lock] = defaultdict(asyncio.Lock)

    async def process_embed(self, vector, value: str, **kwargs) -> np.ndarray:
        """Main entry point from vector.embed()"""
        # Only batch vectors that have embed_batch capability
        if not hasattr(vector, 'embed_batch'):
            # Direct execution for vectors without batching capability
            return await vector.embed(value, **kwargs)
        
        # Create batch key and queue item
        batch_key = id(vector)
        future = asyncio.Future()
        item = BatchItem(value=value, kwargs=kwargs, future=future)
        
        # Add to queue and manage batching
        async with self._locks[batch_key]:
            self._queues[batch_key].append(item)
            queue_size = len(self._queues[batch_key])
            
            # Process immediately if batch is full
            if queue_size >= self.max_batch_size:
                await self._process_batch(batch_key, vector)
            # Schedule processing for first item
            elif queue_size == 1:
                asyncio.create_task(self._delayed_process(batch_key, vector))
        
        return await future

    async def _delayed_process(self, batch_key: int, vector):
        """Process batch after timeout"""
        await asyncio.sleep(self.max_wait_time)
        async with self._locks[batch_key]:
            if self._queues[batch_key]:  # Check if still has items
                await self._process_batch(batch_key, vector)

    async def _process_batch(self, batch_key: int, vector):
        """Process the actual batch"""
        # Extract items from queue
        items = self._queues[batch_key][:self.max_batch_size]
        self._queues[batch_key] = self._queues[batch_key][self.max_batch_size:]
        
        if not items:
            return
        
        try:
            # Use native batch method
            values = [item.value for item in items]
            kwargs = items[0].kwargs  # Assume same kwargs for batch
            
            results = await vector.embed_batch(values, **kwargs)
            
            # Set results for each future
            for item, result in zip(items, results):
                if not item.future.done():
                    item.future.set_result(result)
                    
        except Exception as e:
            # Set exception for all items
            for item in items:
                if not item.future.done():
                    item.future.set_exception(e)

    async def shutdown(self):
        """Process remaining batches on shutdown"""
        # Implementation handles remaining queued items
        pass

Usage Examples

# Batching-capable vector
class OpenAIEmbedding(EmbeddingVector):
    async def embed(self, value: str, query: bool = False, **kwargs) -> np.ndarray:
        batcher = get_current_batcher()
        if batcher:
            return await batcher.process_embed(self, value, query=query, **kwargs)
        return await self._single_embed(value, query=query, **kwargs)
    
    async def embed_batch(self, values: List[str], query: bool = False, **kwargs) -> List[np.ndarray]:
        """Native batch processing using OpenAI batch API"""
        # Efficient batch API call
        return batch_results
    
    async def _single_embed(self, value: str, **kwargs) -> np.ndarray:
        # Single embedding implementation
        pass

# Usage
async def main():
    vector = OpenAIEmbedding()
    
    # With batching context - automatic optimization
    async with with_dynamic_batching(max_batch_size=32, max_wait_time=0.1):
        embeddings = await asyncio.gather(*[
            vector.embed(f"document {i}") for i in range(100)
        ])
        # Automatically batched into ~4 API calls instead of 100

References

• ADR-010: Corpus-Level Statistics with Metadata Providers
• Python ContextVars Documentation
• Async Context Managers

Review Notes

• Performance Testing: Pending validation across different vector types and batch sizes
• Backward Compatibility: Requires validation that all existing vectors work unchanged
• Error Handling: Exception propagation through batching system needs testing

Migration Path

Phase 1: Core Infrastructure (Week 1-2)

• Implement DynamicVectorBatcher with embed_batch detection
• Add context manager infrastructure
• Create test suite for batching and non-batching scenarios

Phase 2: Vector Updates (Week 3-4)

• Update high-value vectors (OpenAI, Cohere, transformers) to implement embed_batch
• Add batcher checking to vectors that would benefit from batching
• Validate that all existing vectors continue working unchanged

Backward Compatibility Guarantee

• Zero Breaking Changes: All existing vectors work exactly as before
• No Performance Regression: Vectors without embed_batch have zero batching overhead
• Opt-in Benefits: Only vectors implementing embed_batch participate in batching
• Graceful Degradation: System works perfectly fine without any vectors implementing batching

This dynamic batching system provides significant performance improvements for capable vectors while maintaining complete backward compatibility and requiring minimal integration effort from vector developers.