System Design 2

Design a Generic Batch Processing System

Designing a batch processing system involves handling large volumes of work items with a focus on reliability, fault tolerance, and efficient resource utilization. This system enables users to submit collections of work items for distributed processing with comprehensive progress tracking and failure recovery.

Requirements

Functional Requirements

• Users should be able to submit batch jobs with multiple work items. Each batch job contains a collection of items to be processed by the same task type with different parameters.
• The system should distribute work items across multiple workers for parallel processing. Large batches should be split into smaller work units for efficient distribution.
• The system should track progress at both batch and item levels. Users need visibility into overall batch completion and individual item status.
• The system should handle partial failures gracefully. When some items in a batch fail, the system should retry failed items while continuing to process successful ones.
• Users should be able to monitor batch execution status and retrieve results. This includes real-time progress updates and final result aggregation.

Non-Functional Requirements

• Scalability: Handle up to 50,000 concurrent batches with 10 million total work items processing simultaneously
• Durability: No work items should be lost due to system failures; all progress must be recoverable
• At-least-once processing: Each work item is guaranteed to be processed at least once, with idempotency handling
• Fault tolerance: System continues operating with individual component failures
• Cost efficiency: Auto-scaling capabilities to minimize resource costs during low utilization

Capacity Estimation

• Peak load: 50,000 batches/hour, average 200 items per batch = 10M items/hour
• Processing rate: ~2,800 items/second sustained
• Storage: ~100GB metadata, 10TB temporary work data
• Worker fleet: 1,000-5,000 workers depending on item complexity

Core Entities

• BatchJob: A user-submitted collection of work items sharing the same task type
• WorkItem: An individual unit of work within a batch, containing input parameters and tracking state
• Task: A reusable processing template that defines how to execute work items
• BatchExecution: Runtime instance tracking the overall progress and status of a batch
• WorkerAssignment: Tracks which worker is processing which work items for failure recovery

System Interface

Batch Management API

POST /batches
{
  "task_id": "image_resize",
  "batch_name": "user_avatars_q4_2024",
  "work_items": [
    {"item_id": "img_1", "input_url": "s3://...", "params": {...}},
    {"item_id": "img_2", "input_url": "s3://...", "params": {...}}
  ],
  "batch_config": {
    "failure_threshold": 0.05,  // Allow 5% failure rate
    "retry_limit": 3,
    "timeout_seconds": 300
  }
}
Response: {"batch_id": "batch_123", "estimated_completion": "2024-01-15T10:30:00Z"}

GET /batches/{batch_id}/status
Response: {
  "batch_id": "batch_123",
  "status": "RUNNING",
  "progress": {
    "total_items": 1000,
    "completed": 750,
    "failed": 25,
    "retrying": 10,
    "pending": 215
  },
  "estimated_completion": "2024-01-15T10:25:00Z"
}

GET /batches/{batch_id}/items/{item_id}
Response: {
  "item_id": "img_1",
  "status": "COMPLETED",
  "worker_id": "worker_456",
  "start_time": "2024-01-15T10:00:00Z",
  "end_time": "2024-01-15T10:02:00Z",
  "output_location": "s3://results/img_1_thumb.jpg"
}

Data Flow

1. Batch Submission: User submits batch job through API Gateway → Batch Management Service
2. Work Item Creation: Batch Management Service creates individual WorkItem records and queues them
3. Work Distribution: Workers pull work items from queue and claim ownership
4. Progress Tracking: Workers update item status throughout processing lifecycle
5. Completion Detection: Batch Coordinator monitors progress and determines batch completion
6. Result Aggregation: Completed batch results are aggregated and made available to users

High-Level Design

Core Components

1. API Gateway: Entry point handling authentication, rate limiting, and request routing

2. Batch Management Service:

   • Creates and validates batch jobs
   • Splits batches into individual work items
   • Manages batch lifecycle and configuration

3. Work Item Database:

   • Choice: PostgreSQL with read replicas for metadata consistency
   • Schema:
     • batches: batch_id, task_id, status, config, created_at, completed_at
     • work_items: item_id, batch_id, status, worker_id, input_data, output_location, retry_count, last_error
     • worker_assignments: assignment_id, worker_id, item_id, claimed_at, heartbeat_at

4. Work Queue:

   • Choice: AWS SQS for reliability and auto-scaling
   • Separate queues for new work, retries, and dead letters
   • Message contains: item_id, batch_id, task_id, input_parameters

5. Worker Pool:

   • Choice: ECS with auto-scaling based on queue depth
   • Workers are stateless and task-agnostic
   • Implement heartbeat mechanism for failure detection

6. Batch Coordinator Service:

   • Monitors batch progress and determines completion
   • Handles retry logic for failed items
   • Manages batch-level timeouts and failure thresholds

7. Data Storage:

   • Input/Output: S3 for large data files with presigned URLs
   • Results: Structured results in PostgreSQL, large outputs in S3

8. Monitoring & Alerting:

   • CloudWatch for infrastructure metrics
   • Custom dashboards for batch processing KPIs
   • Alerts for stuck batches, high failure rates, queue backlogs

Workflow

1. Batch Creation:

   • Batch Management Service validates request and creates batch record
   • Individual work items are created in database with PENDING status
   • Work items are queued for processing

2. Work Processing:

   • Workers poll queue and claim work items (update worker_id and status to CLAIMED)
   • Workers send periodic heartbeats during processing
   • Upon completion, workers update status and output location

3. Progress Monitoring:

   • Batch Coordinator periodically scans active batches
   • Calculates completion percentage and estimates completion time
   • Detects stuck workers via missed heartbeats and reassigns work

4. Failure Handling:

   • Failed items are retried up to configured limit
   • Items exceeding retry limit are marked FAILED
   • Batch completion determined by failure threshold configuration

Deep Dives

1. Work Distribution and Load Balancing

Challenge: Efficiently distribute work items to prevent worker starvation and handle varying processing times.

Solution:

• Use SQS FIFO queues partitioned by batch_id to maintain order for dependent items
• Implement worker prefetch limits to prevent hoarding
• Dynamic batching: workers can claim multiple small items or single large items based on estimated processing time

Work Claiming Protocol:

-- Atomic work claiming
UPDATE work_items
SET status = 'CLAIMED', worker_id = $1, claimed_at = NOW()
WHERE item_id IN (
  SELECT item_id FROM work_items
  WHERE status = 'PENDING' AND batch_id = $2
  LIMIT $3 FOR UPDATE SKIP LOCKED
)
RETURNING item_id, input_data;

2. Failure Detection and Recovery

Worker Failure Detection:

• Workers send heartbeats every 30 seconds while processing
• Batch Coordinator detects missed heartbeats and marks items as ABANDONED
• Abandoned items are re-queued after updating retry_count

Partial Failure Handling:

• Individual item failures don't affect other items in the batch
• Failed items are retried with exponential backoff
• Persistent failures are sent to dead letter queue for manual inspection

Recovery Scenarios:

• Worker crash: Heartbeat timeout triggers work reassignment
• Network partition: Worker continues processing; coordinator detects timeout and may create duplicate work (handled by idempotency)
• Database failure: Work queue persists items; processing continues when database recovers

3. Batch Completion Logic

Completion Criteria:

• All items completed successfully, OR
• Failure rate exceeds configured threshold (batch marked FAILED), OR
• Batch timeout exceeded

Completion Detection:

-- Efficient batch status calculation
SELECT
  b.batch_id,
  COUNT(*) as total_items,
  COUNT(CASE WHEN wi.status = 'COMPLETED' THEN 1 END) as completed,
  COUNT(CASE WHEN wi.status = 'FAILED' THEN 1 END) as failed
FROM batches b
JOIN work_items wi ON b.batch_id = wi.batch_id
WHERE b.status = 'RUNNING'
GROUP BY b.batch_id
HAVING COUNT(*) = COUNT(CASE WHEN wi.status IN ('COMPLETED', 'FAILED') THEN 1 END);

4. Scaling and Performance

Database Scaling:

• Read replicas for status queries
• Partition work_items table by batch_id for large batches
• Connection pooling and prepared statements

Queue Scaling:

• SQS auto-scales to handle message volume
• Use multiple queues for different priority levels
• Implement message batching for high-throughput scenarios

Worker Auto-Scaling:

• Scale workers based on queue depth and message age
• Use spot instances for cost optimization
• Implement graceful shutdown for worker termination

Performance Optimizations:

• Batch database updates for progress tracking
• Use Redis for high-frequency heartbeat storage
• Implement result caching for repeated status queries

5. Monitoring and Observability

Key Metrics:

• Batch throughput (batches/hour, items/second)
• Processing latency (time from submission to completion)
• Failure rates by task type and error category
• Worker utilization and queue depth
• Cost per processed item

Alerting Thresholds:

• Queue depth > 10,000 items (scale up workers)
• Batch failure rate > 10% (investigate task issues)
• Worker heartbeat gaps > 5 minutes (check worker health)
• Database connection pool exhaustion

Dashboards:

• Real-time batch processing overview
• Historical performance trends
• Cost analysis and optimization opportunities
• Error analysis and retry patterns

This architecture provides a robust foundation for batch processing that handles the complexities of work distribution, failure recovery, and progress tracking at scale while maintaining cost efficiency and operational simplicity.

ChatGPT with Inference Batching

Understanding the Problem

What are we building? A ChatGPT-like system that efficiently utilizes a pool of GPU servers for text inference. Each GPU server has a fixed latency for processing batches of 1-100 strings, making request aggregation the key optimization strategy.

Functional Requirements

Core Requirements

1. Users should be able to submit text prompts and receive generated responses
2. The system should aggregate individual requests into batches for GPU processing
3. Responses should be correlated back to original users accurately
4. The system should handle variable request arrival patterns efficiently

Below the line (out of scope)

• Real-time streaming responses (token-by-token generation)
• Conversation history and context management
• User authentication and authorization
• Advanced prompt engineering or content filtering

Non-Functional Requirements

Core Requirements

1. High GPU Utilization: Maximize batch sizes (target: 90+ requests per batch) to optimize cost efficiency
2. Acceptable User Latency: Balance batching delays with user experience (target: <500ms total response time)
3. High Throughput: Handle thousands of concurrent requests efficiently
4. Scalability: System should scale horizontally as demand increases
5. Fault Tolerance: Basic resilience acceptable; memory-based components are fine per problem constraints

Below the line (out of scope)

• Strong consistency guarantees
• Persistent storage requirements
• Multi-region deployment
• Advanced security features

The Set Up

Planning the Approach

This system's core challenge is request aggregation - efficiently batching individual user requests to maximize GPU utilization while maintaining acceptable response times. The fixed GPU latency characteristic means sending 1 request takes the same time as sending 100 requests.

Defining the Core Entities

Request: An individual user prompt with associated metadata (requestId, timestamp, client connection info)

Batch: A collection of 1-100 requests sent to GPU servers as a single inference call

Response: Generated text corresponding to each request in a batch

System Interface

Simple HTTP REST API for request-response pattern:

POST /generate
{
  "prompt": "What is the capital of France?",
  "max_wait_ms": 200  // Optional: client timeout preference
}

Response:
{
  "response": "The capital of France is Paris.",
  "batch_size": 87,
  "processing_time_ms": 245
}

High-Level Design

Core Architecture

[Clients] --HTTP--> [Load Balancer] ---> [Frontend Services] ---> [GPU Pool Manager] ---> [GPU Servers]
                                              |                                              ^
                                              v                                              |
                                       [Batch Accumulator]                                  |
                                       (In-Memory Queue)                                    |
                                              |                                              |
                                              +----------------------------------------------+

Component Overview

Load Balancer: Distributes incoming HTTP requests across Frontend Service instances using round-robin or least-connections

Frontend Service: Stateless HTTP servers that handle request batching and response correlation

Batch Accumulator: In-memory data structure within each Frontend Service for collecting requests

GPU Pool Manager: Simple load balancer that routes batches to available GPU servers

GPU Servers: External inference servers with the fixed-latency API

Data Flow

Request Processing Flow

1. Request Arrival: Client sends HTTP POST to load balancer with text prompt
2. Load Distribution: Load balancer routes request to available Frontend Service instance
3. Batch Accumulation: Frontend Service adds request to in-memory batch accumulator
4. Batch Triggering: Batch is sent to GPU Pool Manager when either:
   • 100 requests accumulated (size trigger)
   • Timeout reached (75-100ms default)
5. GPU Routing: GPU Pool Manager selects least-loaded GPU server
6. Inference Processing: GPU server processes entire batch with fixed latency
7. Response Disaggregation: GPU Pool Manager returns batch response to Frontend Service
8. Response Correlation: Frontend Service maps individual responses back to waiting HTTP requests
9. Client Response: Each client receives their specific response over original HTTP connection

Key Data Transformations

Individual Request → Batch Request:

• Input: Single prompt string per HTTP request
• Transform: Collect into list of 1-100 prompt strings
• Output: Batch API call to GPU with list of strings

Batch Response → Individual Responses:

• Input: List of response strings from GPU
• Transform: Map each response to original request using stored correlation
• Output: Individual HTTP responses to waiting clients

Potential Deep Dives

1) How do we optimize the batching timeout strategy?

The timeout value creates a fundamental trade-off between user latency and GPU utilization. This is the most critical optimization lever in the system.

Dynamic Timeout Adjustment: During high traffic periods, shorter timeouts (20-30ms) work well because batches fill quickly. During low traffic, longer timeouts (100ms) are necessary to collect reasonable batch sizes. The system monitors request arrival rates and adjusts timeouts accordingly.

Traffic Pattern Analysis:

• Track requests per second over rolling windows
• High traffic (>50 req/sec): Use 25ms timeout
• Medium traffic (10-50 req/sec): Use 50ms timeout
• Low traffic (<10 req/sec): Use 100ms timeout

User-Specified Preferences: Allow clients to specify maximum wait time in their requests. Premium users might get priority treatment with faster timeouts, while batch processing jobs can accept longer delays.

2) How do we handle uneven load distribution across Frontend Services?

Each Frontend Service maintains its own batch accumulator, which can lead to suboptimal batching if traffic is unevenly distributed.

Load Balancer Strategies:

• Least-connections routing ensures requests go to Frontend Services with fewer active batches
• Health checks prevent routing to overloaded instances
• Sticky sessions are not needed since each request is independent

Cross-Service Batch Sharing: Frontend Services could share accumulated requests through a lightweight coordination mechanism. When one service has a partial batch approaching timeout, it could request additional items from other services' accumulators.

Capacity Planning: Monitor batch utilization metrics across all Frontend Services. If average batch sizes drop below targets, this indicates either uneven load distribution or insufficient traffic density requiring more aggressive timeout tuning.

3) How do we ensure GPU pool efficiency and handle failures?

GPU Load Balancing: The GPU Pool Manager tracks pending requests per GPU server and routes new batches to the least loaded server. This prevents GPU hotspots while maintaining utilization.

Failure Handling: When a GPU server fails or times out, the batch is immediately retried on a different server. The correlation mapping is preserved, so clients don't experience duplicate responses. Circuit breaker patterns prevent sending traffic to consistently failing GPUs.

Health Monitoring: Track GPU response times, error rates, and processing capacity. Servers showing degraded performance are temporarily removed from rotation. Automatic re-inclusion occurs after a recovery period with successful health checks.

4) How do we scale the system horizontally?

Frontend Service Scaling: Frontend Services are stateless and scale independently. Each instance maintains only its local batch accumulator and correlation mappings. Adding instances increases both batch accumulation capacity and HTTP request handling capacity.

Memory Management: Each Frontend Service has bounded memory usage since batch accumulators have maximum sizes and correlation mappings are short-lived. Monitor memory utilization and set limits to prevent resource exhaustion.

GPU Pool Scaling: Adding GPU servers increases overall processing capacity. The GPU Pool Manager automatically includes new servers in routing decisions. Removal requires draining existing batches before taking servers offline.

5) How do we handle different request priorities and patterns?

Priority Tiers: Implement separate batch accumulators for different priority levels. High-priority requests get smaller batch sizes and shorter timeouts, while standard requests wait for full batches. This balances responsiveness for critical users with efficiency for regular traffic.

Predictive Batching: Analyze historical traffic patterns to predict peak periods. Pre-warm batch accumulators during expected traffic surges and adjust timeout strategies based on time-of-day patterns.

Burst Handling: During traffic spikes, the system can temporarily reduce batch size requirements to maintain responsiveness. Monitoring request queue depths helps detect when burst mode should activate.

6) How do we monitor and optimize system performance?

Key Metrics:

• Average batch utilization percentage
• Request latency distribution (P50, P95, P99)
• GPU server utilization and response times
• Timeout frequency vs size-triggered batches
• Request correlation accuracy

Performance Optimization: Continuously analyze the relationship between timeout values and batch sizes. Optimal timeout tuning maximizes the product of batch efficiency and user satisfaction. A/B testing different timeout strategies helps identify optimal configurations.

Alerting Strategy: Monitor for degraded batch utilization, increased latency, or GPU failures. Set thresholds that trigger alerts before user experience degrades significantly.

Design ChatGPT Playground

Understanding the Problem

🎮 What is a ChatGPT Playground?

A ChatGPT Playground is an interactive environment where users can experiment with Large Language Models (LLMs) by sending prompts and receiving responses. Unlike standard chat interfaces, playgrounds provide advanced configuration options like adjusting model parameters (temperature, max tokens, top-p), switching between different models, and managing multiple conversation sessions for different experiments.

Functional Requirements

Core Requirements

1. Users should be able to send text prompts to an LLM and receive streaming responses in real-time.
2. Users should be able to view and persist conversation history within playground sessions.
3. Users should be able to configure LLM parameters (temperature, max tokens, model selection, system prompts).
4. Users should be able to create, name, and switch between multiple playground sessions.

Below the line (out of scope)

• Sharing playground sessions with other users
• Advanced prompt templating or batch processing
• Custom model fine-tuning or training
• Integration with external APIs or data sources

Non-Functional Requirements

Core Requirements

1. The system should provide low latency responses with first token arriving within 500ms of prompt submission.
2. The system should scale to support 50,000 concurrent users with 5,000 prompts per second at peak.
3. The system should be highly available (99.9% uptime) with eventual consistency acceptable for conversation history.
4. The system should guarantee data durability - no conversation data should be lost due to system failures.
5. The system should handle WebSocket connections gracefully with automatic reconnection and session recovery.

Below the line (out of scope)

• Strong consistency guarantees across all operations
• Sub-100ms response times
• Cross-region data replication

Here's how it might look on your whiteboard:

Playground Requirements

The Set Up

Planning the Approach

For this problem, we'll follow the delivery framework focusing on real-time streaming capabilities and session management. The core challenge is handling persistent WebSocket connections while maintaining conversation state and integrating with external LLM providers reliably.

Defining the Core Entities

User: Represents an individual using the playground platform, with authentication and usage tracking.

PlaygroundSession: A named workspace containing LLM configuration and associated conversation history. Users can create multiple sessions for different experiments.

Message: Individual prompts (user) or responses (assistant) within a session, with support for streaming token metadata.

LLMConfiguration: Parameters controlling LLM behavior (model, temperature, max_tokens, system_prompt) associated with each session.

StreamingState: Transient state tracking active streaming responses, token buffers, and connection status.

Core Entities

The API

We'll use REST for session management and WebSockets for real-time prompt/response interaction:

Session Management (REST)

POST /sessions
{
  "name": "My Experiment",
  "config": {
    "model": "gpt-4",
    "temperature": 0.7,
    "max_tokens": 1000
  }
}

GET /sessions/{sessionId}
Response: {
  "session_id": "uuid",
  "name": "My Experiment",
  "config": {...},
  "messages": [...] // paginated
}

PUT /sessions/{sessionId}/config
{
  "temperature": 0.9,
  "max_tokens": 500
}

Real-time Interaction (WebSocket)

// Connection
ws://api.playground.com/sessions/{sessionId}/chat

// Client → Server
{
  "type": "prompt",
  "content": "Explain quantum computing",
  "message_id": "uuid"
}

// Server → Client (streaming)
{
  "type": "token",
  "content": "Quantum",
  "message_id": "uuid",
  "position": 0
}
{
  "type": "token",
  "content": " computing",
  "message_id": "uuid",
  "position": 1
}
{
  "type": "complete",
  "message_id": "uuid",
  "full_response_id": "uuid"
}

High-Level Design

1) Users should be able to send prompts and receive streaming responses

The core challenge is maintaining persistent WebSocket connections while orchestrating calls to external LLM providers and streaming responses back to clients.

Pattern: Real-time Updates

Streaming LLM responses represent a classic real-time updates pattern where servers must push data to clients as it becomes available. This requires persistent connections and careful handling of connection lifecycle, backpressure, and failure recovery.

Our architecture centers around a WebSocket Gateway that manages persistent connections and a Chat Service that orchestrates LLM interactions:

Basic Streaming Architecture

WebSocket Gateway: Manages persistent connections, handles authentication, and routes messages between clients and backend services. Maintains connection state in memory with Redis backing for connection recovery.

Chat Service: Stateless service that processes prompts, calls the LLM Provider, and streams responses back through the WebSocket Gateway. Handles message persistence and session state management.

LLM Provider: External service (OpenAI, Anthropic, etc.) that processes prompts and returns streaming responses.

Flow for sending a prompt:

1. Client sends prompt via WebSocket connection to Gateway
2. Gateway forwards to Chat Service with connection identifier
3. Chat Service validates session, retrieves LLM config from database
4. Chat Service calls LLM Provider with streaming enabled
5. As tokens arrive from LLM Provider, Chat Service immediately forwards to Gateway
6. Gateway streams tokens to client via WebSocket
7. Chat Service persists complete message to database once streaming finishes

2) Users should be able to view and persist conversation history

Conversation persistence requires careful coordination between real-time streaming and durable storage, with support for session recovery and message replay.

We'll use PostgreSQL with a conversation-optimized schema:

CREATE TABLE playground_sessions (
  session_id UUID PRIMARY KEY,
  user_id UUID NOT NULL,
  name VARCHAR(255) NOT NULL,
  config JSONB NOT NULL,
  created_at TIMESTAMP DEFAULT NOW(),
  updated_at TIMESTAMP DEFAULT NOW()
);

CREATE TABLE messages (
  message_id UUID PRIMARY KEY,
  session_id UUID NOT NULL,
  role VARCHAR(20) NOT NULL, -- 'user' or 'assistant'
  content TEXT NOT NULL,
  tokens JSONB, -- streaming metadata, token timing
  sequence_number INTEGER NOT NULL,
  created_at TIMESTAMP DEFAULT NOW(),
  UNIQUE(session_id, sequence_number)
) PARTITION BY HASH(session_id);

-- Index for efficient conversation retrieval
CREATE INDEX idx_messages_session_sequence
ON messages (session_id, sequence_number);

Message Persistence Strategy:

1. Optimistic Writing: Start persisting user prompts immediately when received
2. Streaming Buffer: Accumulate assistant response tokens in memory during streaming
3. Atomic Completion: Persist complete assistant response as single transaction when streaming finishes
4. Sequence Guarantees: Use incrementing sequence numbers per session to ensure message ordering

3) Users should be able to configure LLM parameters

LLM configuration management requires handling parameter validation, hot-swapping configurations mid-session, and maintaining compatibility across different model providers.

Session Configuration

Configuration Service: Dedicated service managing LLM parameters with validation rules, default values, and provider-specific constraints.

Configuration validation examples:

• Temperature: 0.0-1.0 for OpenAI models, 0.0-2.0 for Anthropic
• Max tokens: Provider-specific limits (4096 for GPT-3.5, 100k for Claude)
• Model availability: Check provider APIs for supported models

Hot Configuration Updates:

-- Track configuration versions for rollback
CREATE TABLE config_versions (
  version_id UUID PRIMARY KEY,
  session_id UUID NOT NULL,
  config JSONB NOT NULL,
  created_at TIMESTAMP DEFAULT NOW()
);

When users update configurations mid-conversation, the Chat Service validates parameters against provider constraints and applies changes to subsequent prompts without affecting in-flight requests.

4) Users should be able to create and switch between multiple sessions

Session management requires efficient session listing, context switching, and maintaining session state across WebSocket reconnections.

Session Management Architecture

Session Router: Routes WebSocket connections to appropriate Chat Service instances based on session ID, enabling sticky session behavior for connection management.

Session State Recovery: When clients reconnect or switch sessions:

1. Authenticate WebSocket connection with session ID
2. Retrieve last N messages from database for context
3. Check for any incomplete streaming responses
4. Resume or restart as appropriate

Efficient session switching:

// Client-side session switching
websocket.send({
  type: "switch_session",
  new_session_id: "uuid",
  resume_from_message: "last_message_id",
});

Potential Deep Dives

1) How do we handle WebSocket connection reliability and session recovery?

WebSocket connections are inherently fragile - networks drop, browsers close, mobile apps background. Our system needs robust connection management and seamless recovery.

Approach: Heartbeat and Reconnection

Implement client-server heartbeating with exponential backoff reconnection. Client sends ping every 30 seconds, server responds with pong. If either side misses 3 consecutive heartbeats, assume connection is dead.

// Client-side connection management
class PlaygroundWebSocket {
  connect(sessionId) {
    this.ws = new WebSocket(`/sessions/${sessionId}/chat`);
    this.startHeartbeat();
    this.setupReconnection();
  }

  startHeartbeat() {
    this.heartbeat = setInterval(() => {
      if (this.ws.readyState === WebSocket.OPEN) {
        this.ws.send(JSON.stringify({ type: "ping" }));
      }
    }, 30000);
  }

  setupReconnection() {
    this.ws.onclose = () => {
      this.reconnectWithBackoff();
    };
  }
}

Challenges

Heartbeating adds bandwidth overhead and complexity. Mobile clients may not handle backgrounding gracefully, leading to unnecessary reconnections.

Approach: Stateless Recovery with Message Replay

Store connection state in Redis with TTL. When clients reconnect, they provide their last received message ID, and the server replays any missed messages.

Connection State Storage:

SET connection:{connectionId} {
  "session_id": "uuid",
  "user_id": "uuid",
  "last_message_sequence": 42,
  "streaming_message_id": "uuid" // if streaming was interrupted
} EX 3600

Recovery Protocol:

1. Client reconnects with last known message sequence
2. Server queries Redis for connection state
3. Server replays messages from database where sequence > last_known
4. If streaming was interrupted, restart from LLM provider or serve cached partial response

Challenges

Redis becomes a single point of failure. Message replay can be expensive for long-lived sessions with extensive history.

Approach: Connection Pooling with Sticky Sessions

Use consistent hashing to ensure users consistently connect to the same WebSocket Gateway instance. This enables in-memory connection state without Redis dependency.

Each Gateway instance maintains:

// In-memory connection registry
const connections = new Map(); // connectionId -> WebSocketConnection
const sessionConnections = new Map(); // sessionId -> Set<connectionId>

Load balancer uses consistent hashing on session ID to route WebSocket upgrade requests to the same Gateway instance.

Benefits: No external dependencies, faster connection recovery, simpler debugging.

Trade-offs: Gateway failures lose all connection state, requiring clients to reconnect and replay from database.

For our scale (50k concurrent users), consistent hashing with 10-20 Gateway instances provides good balance of reliability and simplicity.

2) How do we handle LLM provider failures and rate limiting?

External LLM providers have API limits, experience outages, and may return errors. Our system needs graceful degradation and retry strategies.

Pattern: External Service Integration

LLM provider integration demonstrates classic external service integration challenges including rate limiting, circuit breaking, and graceful degradation when dependencies fail.

Circuit Breaker Pattern:

class LLMProviderClient:
    def __init__(self):
        self.circuit_breaker = CircuitBreaker(
            failure_threshold=5,  # failures before opening
            recovery_timeout=60,  # seconds before trying again
            expected_exception=LLMProviderError
        )

    @circuit_breaker
    def stream_completion(self, prompt, config):
        # Call external LLM provider
        response = requests.post(
            f"{provider_url}/completions",
            json={"prompt": prompt, **config},
            stream=True,
            timeout=30
        )

        for chunk in response.iter_content():
            yield parse_token(chunk)

Provider Rate Limiting:

Implement token bucket rate limiting per API key:

class RateLimiter:
    def __init__(self, requests_per_minute=1000):
        self.bucket = TokenBucket(
            capacity=requests_per_minute,
            refill_rate=requests_per_minute / 60
        )

    async def acquire(self):
        if not self.bucket.consume(1):
            raise RateLimitExceeded("Provider rate limit exceeded")

Graceful Degradation Strategies:

1. Cached Responses: For identical prompts with deterministic settings (temperature=0), serve cached responses
2. Fallback Providers: Route to secondary LLM provider when primary fails
3. Queue with SLA: Buffer requests for up to 10 seconds during provider issues
4. User Communication: Send clear error messages explaining delays

3) How do we scale to handle 50,000 concurrent WebSocket connections?

WebSocket connections are stateful and memory-intensive. Scaling requires careful resource management and connection distribution.

Connection Capacity Planning:

Each WebSocket connection consumes approximately:

• 8KB TCP buffers (send/receive)
• 2KB application state (user ID, session ID, last message tracking)
• 1KB average message buffer for streaming

Total: ~11KB per connection × 50,000 = 550MB per Gateway instance.

With modern servers having 32GB+ RAM, a single Gateway instance can handle 50k connections, but we'll use multiple instances for fault tolerance.

Scaling Architecture:

Gateway Scaling

WebSocket Gateway Scaling:

• Deploy 5 Gateway instances behind L4 load balancer
• Each handles 10k connections (plenty of headroom)
• Use consistent hashing on session ID for sticky routing
• Auto-scale based on connection count metrics

Chat Service Scaling:

• Stateless horizontal scaling based on CPU/memory usage
• Queue depth monitoring for LLM provider calls
• Circuit breaker integration prevents cascade failures

Database Connection Pooling:

With 5k prompts/second peak, database connections become critical:

# Connection pool configuration
DATABASE_POOL = {
    "min_connections": 10,
    "max_connections": 100,
    "connection_timeout": 30,
    "idle_timeout": 300
}

PostgreSQL Scaling Strategy:

• Primary for writes (session creation, message persistence)
• Read replicas for conversation history retrieval
• Partition messages table by session_id hash for write distribution

4) How do we ensure data consistency and handle partial failures?

Distributed systems create opportunities for partial failures where some operations succeed while others fail, potentially leaving the system in inconsistent state.

Message Persistence Consistency:

The critical consistency challenge occurs during streaming: what happens if we successfully start streaming to the client but fail to persist the complete response?

Write-Ahead Logging with Compensation: Before starting LLM call, write message intent to database with PENDING status. Stream to client while buffering complete response. Then atomically update intent to COMPLETED and insert complete message in single transaction. Background cleanup job identifies stale PENDING intents and handles recovery.

Idempotency for Duplicate Prompts: Clients may retry prompts during network issues. Implement idempotency using client-generated message IDs. Check for existing messages with same ID before processing and return cached responses for duplicates.

Session State Consistency: Configuration updates during active streaming create race conditions. Use optimistic locking with version numbers on session configuration. Update operations include version checks and fail if concurrent modifications occurred.

Distribute ML models in a Data Center

P2P

1. So it seems like one machine should download the model and distribute it to other machines.
2. This looks suspiciously like p2p torrenting so I am going to approach it similarly.
3. Machine 1 has to download the ML model fully. Then it should chunk it in let's say 4 mb chunks. Each chunk is SHA1 hashed and this machine keeps tracks of the chunks and the hashes associated with it. Let's call this the seed machine.
4. The seed machine communicates with other machines telling there's a ML model to be downloaded. Since these machines are in the same DC, they know each other's IPs. We will have a dedicated gRPC server and client on each machine for the purpose of downloading ML models. So the seed machine starts by distributing the manifest to each of the machines.
5. The other machines then start downloading the chunks. The very first sets of chunks are downloaded from the seed machine.
6. Each machine periodically gossips with a few random peers to learn who has what chunks. When a machine needs a specific chunk, it picks a peer that has it (from its gossip knowledge) and downloads directly via gRPC.
7. Eventually at some time, all the machines have all the chunks, and they can validate they indeed have that because each has a copy of the manifest file.

The trade off between gossip versus having some central tracker:

1. network overhead in gossip, while with a central tracker, one can make efficient queries.
2. gossip protocol is self-healing, while a central tracker is a single point of failure.

Pros: potentially fast due to parallelism. inherent fault tolerance: multiple sources for each chunk. scalable, work from 10 to 10_000 machines.

Cons: Complex coordination (gossip, central tracker), unpredictable timing.

Tree

If we know there are exactly N machines, then scaling is not important and maybe controlled behaviour is. Each machine can pass the model to exactly one other machine in each round. So:

1. Machine 1 -> Machine 2
2. Machine 1 -> Machine 3, Machine 2 -> Machine 4
3. Machine 1 -> Machine 5, Machine 2 -> Machine 6, Machine 3 -> Machine 7, Machine 4 -> Machine 8

For N machines, the number of rounds we will have is log2(N). For 100, that's at least 7 rounds. Pros: There's no complex discovery here, and bandwith usage is optimal - each round, for a machine, it's committed to 100% upload or 100% download. Predictable timing and progress tracking.

Cons: half the machines are idle in each round. If one node fails, sub tree fails. Load imbalance: Early nodes (closer to root) work harder than leaf nodes.

Hybrid

Tree approach for first few nodes, then P2P file transfer for stragglers or failures.

Spotify Top-K

Functional Requirements

1. Real time Top K songs for all users.

Non-functional

• Real time - low latency.
• Availability over consistency over all

Data flow

Naive:

• Cassadandra for write heavy system, to store analytics.
  • Cassandra has CDC that writes to a commit log, files on disk.
• CDC into OLAP (a dataware house, maybe clickhouse).
  • You must build a service that parses the commit logs and push changes to OLAP.
  • Could even be Kafka.
• Followed by hourly map reduce workers to compute states.

It's naive because it's not realtime. At least you diffeentiated between OLTP and OLAP.

Better:

• User plays song -> Analytics service collects info. -> pushes this into Kafka.
  • Save the song id, the user id, the timestamp.
• Stream processor like flink or Spark Streams consumes it in windows.
  • Pull the above data from kafka.
  • Aggregate the counts every minute or 5 minutes.

High Level Data Flow

1. User has a post id. He uses that to post a comment.
2. The comment is processed by a CommentService, which makes an entry in a db. The db records the user id, comment id, post id, and timestamp.
3. This comment is then made available to all users via an SSE connection established by the clients
   • The client establishes the SSE GET request. The service keeps the connection open and continuously sends data to the client.
4. The service keeps track of a few things:
   • The user id/session id and the specific TCP connection object.
   • The user id and the post id the user is listening to.
   • Comments for a post id.
5. So when a comment comes in, a service should publish this comment to everyone who is subscribing to a post. So a pub/sub is the means of routing comments.
6. We can use a Redis pub/sub (since Redis will push, fire once and forget) and is fast. When a comment is posted, it's published to Redis, which maps a post id to the user id/connection. An SSE service manages the SSE connections, and listens to Redis channels.
   • The service subscribes to various channels based on the what the user connections it's managing are interested in. It receives corresponding messages, iterates though its in-memory list of active connections and sends messages to those users.