Design a Task Scheduling Service

Question

Design a distributed task scheduling system that handles millions of requests. Discuss trade-offs in consistency, availability, and performance.

Codemia · Accepted Answer

### Functional Requirements
1. **Task Queue Management**: Ability to queue millions of tasks for various ML operations like data preprocessing, model training, and evaluation.
2. **Scheduling**: Dynamic scheduling of tasks based on priority and resource availability.
3. **Monitoring and Logging**: Real-time monitoring of task status and logging for debugging and auditing purposes.
4. **A/B Testing Framework**: Support for running A/B tests on models and serving different versions to users.

### Non-Functional Requirements
1. **Scalability**: The system must scale horizontally to handle increasing load without performance degradation.
2. **High Availability**: The system should be available 99.99% of the time to ensure continuous ML operations.
3. **Low Latency**: Task scheduling and execution should have a response time of less than 100ms.
4. **Data Quality Assurance**: Mechanisms to ensure the integrity and quality of data used for training and evaluation.

Assuming Two Sigma handles an average of 5 million tasks per day and expects a 20% growth in task volume annually:

1. **Daily Tasks**: 5,000,000 tasks/day
2. **Per Second Load**: 5,000,000 tasks / (24 * 60 * 60) = ~57.87 tasks/second.
3. **Peak Load**: Assuming peak traffic is 3x the average, we need to handle ~173 tasks/second.
4. **Task Size**: If each task takes 200KB of data, we would require bandwidth to handle ~34.6 MB/sec on average and ~104 MB/sec at peak.

### Component Diagram
- **Load Balancer**: Distributes incoming requests to task scheduler instances.
- **Task Scheduler**: Manages queuing, prioritization, and scheduling of tasks. Implemented using **Redis** for in-memory data structure and fast access.
- **Worker Nodes**: Execute tasks. Each worker can run multiple instances of ML model training or evaluation using **Docker** containers for isolation.
- **Database**: A **PostgreSQL** database for persistent storage of task metadata and logs.
- **Monitoring**: Use **Prometheus** and **Grafana** for real-time monitoring of task execution and system health.
- **A/B Testing Service**: Dedicated service to manage A/B testing logic and routing of requests based on experiment versions.

### Schema Design
- **Tasks Table**: 
  - `task_id` (UUID, Primary Key)
  - `task_type` (VARCHAR)
  - `status` (ENUM: queued, running, completed, failed)
  - `created_at` (TIMESTAMP)
  - `updated_at` (TIMESTAMP)
  - `payload` (JSONB)
- **Results Table**: 
  - `result_id` (UUID, Primary Key)
  - `task_id` (UUID, Foreign Key)
  - `output` (JSONB)
  - `metrics` (JSONB)
  - `created_at` (TIMESTAMP)
- **A/B Tests Table**:
  - `experiment_id` (UUID, Primary Key)
  - `version_a` (UUID, Foreign Key)
  - `version_b` (UUID, Foreign Key)
  - `metrics` (JSONB)
  - `created_at` (TIMESTAMP)

This schema supports effective querying for monitoring task progress and A/B test results.

1. **Consistency vs. Availability**: Opted for eventual consistency in task processing to maximize availability. This means tasks may not be processed in the exact order they are received, but the system remains responsive.
2. **Complexity vs. Performance**: Chose a microservices architecture for the A/B testing framework to allow independent scaling and maintenance, but this adds complexity in service orchestration.
3. **Batch vs. Real-time Processing**: For tasks that require immediate feedback (like A/B testing), we chose real-time processing with a focus on low latency, while other less time-sensitive tasks can be handled in batch mode.
4. **Data Quality vs. Speed**: Prioritized speed in the task execution pipeline, but implemented a robust data validation step to ensure that only high-quality data is used in model training.

Design a Task Scheduling Service

Two Sigma

What the Interviewer Expects

Key Topics to Cover

How to Approach This

Possible Follow-up Questions

Practice a Similar Problem on Codemia

Sample Answer

Requirements

Functional Requirements

Capacity Estimation

Submit Your Answer

Two Sigma Software Engineer Interview Guide

Related Questions

Design a high-throughput Inventory Management System

Design a low-latency Rate Limiting System

Design a fault-tolerant Payment System

Design a fault-tolerant Messaging System

Design Walmart Product Search