[gateway] RFC: Dedicated embeddings gateway plugin architecture for /v1/embeddings

[dittops]

Summary

This RFC proposes implementing a dedicated gateway plugin for the /v1/embeddings endpoint (Option 4) as a long-term architectural solution, building upon the analysis in #4.

Context

As discussed in #4, there are multiple approaches to implementing embeddings support:

1. Option 1: Extend current chat completions gateway (not suitable - different API patterns)
2. Option 2: Route through metadata service (MVP approach - limited functionality)
3. Option 3: Direct routing with extra headers (poor UX)
4. Option 4: Dedicated embeddings gateway plugin (this proposal)

While Option 2 provides a quick MVP, embeddings workloads have fundamentally different characteristics from chat completions that justify a dedicated gateway plugin.

Motivation for Dedicated Gateway Plugin

Key Differences from Chat Completions

< /dev/null | Aspect | Chat Completions | Embeddings |
|--------|-----------------|------------|
| Input Pattern | Sequential conversation | Batch text processing |
| Output | Streaming text generation | Fixed-size vectors |
| Latency Profile | Variable (token generation) | Predictable |
| Memory Usage | Dynamic | Deterministic |
| Caching Strategy | Prefix-based | Input hash-based |
| Routing Needs | SLO, load, prefix cache | Batch-aware, memory-based |

Benefits of Dedicated Plugin

1. Optimized Routing Algorithms

   • Batch-size aware routing
   • Memory-based pod selection
   • Simplified algorithms (no streaming/SLO complexity)

2. Clean Architecture

   • Separate concerns between generation and embedding
   • Independent scaling policies
   • Cleaner codebase without conditional logic

3. Performance Optimization

   • Embedding-specific request batching
   • Optimized memory allocation
   • Dedicated connection pooling

4. Future Extensibility

   • Multi-modal embeddings support
   • Custom embedding strategies
   • Advanced caching mechanisms

Proposed Architecture

Component Structure

pkg/plugins/gateway/embeddings/
├── gateway.go          # Main server implementation
├── req_body.go         # Request handling
├── types.go            # Embeddings-specific types
├── router.go           # Routing algorithms
├── metrics.go          # Prometheus metrics
└── cache.go            # Embedding cache interface

Routing Algorithms

• Least Memory: Route to pod with most available memory
• Batch Affinity: Route similar batch sizes to same pods
• Round Robin: Simple distribution for uniform workloads

Deployment Architecture

# Separate deployment for embeddings gateway
apiVersion: apps/v1
kind: Deployment
metadata:
  name: aibrix-embeddings-gateway
spec:
  replicas: 3  # Independent scaling
  template:
    spec:
      containers:
      - name: embeddings-gateway
        ports:
        - containerPort: 50053  # Different port

Implementation Plan

Phase 1: Core Gateway Structure

• Create embeddings gateway base structure
• Implement request/response handling
• Add batch validation logic
• Basic routing (round-robin)

Phase 2: Advanced Routing

• Implement memory-based routing
• Add batch affinity routing
• Create routing benchmarks

Phase 3: Integration

• Update Envoy configuration
• Add Kubernetes manifests
• Integrate with existing cache

Phase 4: Production Features

• Prometheus metrics
• Request batching optimization
• Circuit breaker patterns
• Comprehensive testing

Testing Requirements

• Unit tests for all routing algorithms
• Integration tests with multiple embedding models
• Chaos testing for failure scenarios

Success Metrics

• Independent scaling without affecting chat completions
• Clean separation of concerns in codebase

References

• Original embeddings issue: [API] Implement /v1/embeddings endpoint for OpenAI-compatible embeddings support #4
• OpenAI Embeddings API: https://platform.openai.com/docs/api-reference/embeddings
• Current gateway architecture: docs/source/gateways.md

Questions for Discussion

1. Should we support embedding-specific caching in Phase 1?
2. What batch size limits should we enforce?
3. How should we handle multi-modal embeddings in the future?
4. Should routing algorithms be configurable per model?

cc: @maintainers

[dittops]

Update: Proposed Routing Strategies for Embeddings Gateway

Based on further analysis of the existing gateway routing algorithms and embedding workload characteristics, here are the specific routing strategies that should be implemented for the embeddings gateway:

Core Routing Strategies

1. Random (random)

   • Already implemented in existing gateway
   • Simple baseline for uniform distribution
   • Useful when all pods have similar capacity

2. Least-Request (least-request)

   • Already implemented in existing gateway
   • Routes to pod with fewest active embedding requests
   • Effective for balancing concurrent request load

3. Least-Latency (least-latency)

   • Exists but needs adaptation for embeddings
   • Remove generation phase metrics
   • Focus on embedding computation latency only
   • Track average time per token for embeddings

4. Least-Load (least-load)

   • Similar to existing least-busy-time but focused on overall resource utilization
   • Considers: GPU memory, GPU utilization, active batch sizes
   • Formula: load_score = gpu_util * 0.4 + memory_util * 0.4 + (active_batch_size / max_batch_size) * 0.2

5. Least-Busy (least-busy)

   • Maps to existing least-busy-time implementation
   • Routes based on GPU busy time percentage
   • Already suitable for embeddings without modification

Additional Embedding-Specific Strategies

6. Batch-Affinity (batch-affinity)

   • New strategy specific to embeddings
   • Routes similar batch sizes to same pods
   • Optimizes memory allocation patterns
   • Example: batch_size < 10 → small_batch_pods, batch_size > 100 → large_batch_pods

7. Memory-Aware (memory-aware)

   • Extension of least-load focusing on memory
   • Critical for large batch embeddings
   • Prevents OOM by routing large batches to high-memory pods

Implementation Priority

Phase 1 (MVP):

• Random (reuse existing)
• Least-Request (reuse existing)
• Least-Load (adapt from least-busy-time)

Phase 2 (Optimization):

• Least-Latency (adapt existing)
• Least-Busy (reuse existing)
• Batch-Affinity (new implementation)

Phase 3 (Advanced):

• Memory-Aware routing
• Hybrid strategies combining multiple metrics

Configuration Example

apiVersion: v1
kind: ConfigMap
metadata:
  name: embeddings-gateway-config
data:
  routing_strategies:  < /dev/null | 
    default: "least-request"
    models:
      text-embedding-3-small:
        strategy: "least-request"
        fallback: "random"
      text-embedding-3-large:
        strategy: "memory-aware"
        fallback: "least-load"
    batch_thresholds:
      small: 10
      medium: 100
      large: 500

Router Selection Logic

func SelectRouter(model string, batchSize int) Router {
    // For large batches, prioritize memory-aware routing
    if batchSize > config.BatchThresholds.Large {
        return NewMemoryAwareRouter()
    }
    
    // For medium batches, use least-load
    if batchSize > config.BatchThresholds.Medium {
        return NewLeastLoadRouter()
    }
    
    // For small batches, use model-specific config
    if strategy, ok := config.Models[model]; ok {
        return RouterFactory.Create(strategy.Strategy)
    }
    
    // Default fallback
    return NewLeastRequestRouter()
}

This approach leverages existing routing infrastructure while adding embedding-specific optimizations where needed.