393 lines
13 KiB
Rust
393 lines
13 KiB
Rust
//! Tests for continuous batching scheduler with SLA lanes
|
|
//!
|
|
//! This module contains comprehensive tests for the continuous batching
|
|
//! system that are designed to FAIL FIRST, then drive implementation.
|
|
|
|
use rtx_inference::{
|
|
error::InferenceError,
|
|
request::{InferenceRequest, RequestId, RequestPriority},
|
|
scheduler::{BatchScheduler, BatchSchedulerConfig, SlaLane},
|
|
};
|
|
use std::time::{Duration, Instant};
|
|
use tokio::time::sleep;
|
|
|
|
/// Test that scheduler can be created with valid configuration
|
|
#[tokio::test]
|
|
async fn test_scheduler_creation() {
|
|
let config = BatchSchedulerConfig {
|
|
max_batch_size: 32,
|
|
max_wait_time: Duration::from_millis(100),
|
|
sla_lanes: vec![
|
|
SlaLane {
|
|
name: "premium".to_string(),
|
|
priority: RequestPriority::High,
|
|
max_latency: Duration::from_millis(50),
|
|
max_batch_size: 16,
|
|
memory_limit: Some(1024 * 1024 * 100), // 100MB
|
|
min_wait_time: Some(Duration::from_millis(10)),
|
|
},
|
|
SlaLane {
|
|
name: "standard".to_string(),
|
|
priority: RequestPriority::Normal,
|
|
max_latency: Duration::from_millis(150),
|
|
max_batch_size: 32,
|
|
memory_limit: Some(1024 * 1024 * 200), // 200MB
|
|
min_wait_time: Some(Duration::from_millis(20)),
|
|
},
|
|
],
|
|
..Default::default()
|
|
};
|
|
|
|
let scheduler = BatchScheduler::new(config).await;
|
|
assert!(scheduler.is_ok());
|
|
|
|
let scheduler = scheduler.unwrap();
|
|
assert_eq!(scheduler.config().max_batch_size, 32);
|
|
assert_eq!(scheduler.config().sla_lanes.len(), 2);
|
|
}
|
|
|
|
/// Test request submission and queue management
|
|
#[tokio::test]
|
|
async fn test_request_submission() {
|
|
let config = BatchSchedulerConfig::default();
|
|
let mut scheduler = BatchScheduler::new(config).await.unwrap();
|
|
|
|
let request = InferenceRequest {
|
|
id: RequestId::new(),
|
|
model_name: "llama-7b".to_string(),
|
|
input_tokens: vec![1, 2, 3, 4, 5],
|
|
max_new_tokens: 100,
|
|
temperature: 0.8,
|
|
priority: RequestPriority::Normal,
|
|
deadline: Some(Instant::now() + Duration::from_millis(1000)),
|
|
..Default::default()
|
|
};
|
|
|
|
let result = scheduler.submit_request(request.clone()).await;
|
|
assert!(result.is_ok());
|
|
|
|
let queue_stats = scheduler.queue_stats().await;
|
|
assert_eq!(queue_stats.total_pending, 1);
|
|
assert_eq!(queue_stats.lanes.get("standard").unwrap().pending_count, 1);
|
|
}
|
|
|
|
/// Test SLA lane assignment based on request priority
|
|
#[tokio::test]
|
|
async fn test_sla_lane_assignment() {
|
|
let config = BatchSchedulerConfig {
|
|
sla_lanes: vec![
|
|
SlaLane {
|
|
name: "premium".to_string(),
|
|
priority: RequestPriority::High,
|
|
max_latency: Duration::from_millis(50),
|
|
max_batch_size: 8,
|
|
memory_limit: Some(1024 * 1024 * 50), // 50MB
|
|
min_wait_time: Some(Duration::from_millis(5)),
|
|
},
|
|
SlaLane {
|
|
name: "standard".to_string(),
|
|
priority: RequestPriority::Normal,
|
|
max_latency: Duration::from_millis(150),
|
|
max_batch_size: 32,
|
|
memory_limit: Some(1024 * 1024 * 200), // 200MB
|
|
min_wait_time: Some(Duration::from_millis(20)),
|
|
},
|
|
],
|
|
..Default::default()
|
|
};
|
|
|
|
let mut scheduler = BatchScheduler::new(config).await.unwrap();
|
|
|
|
// Submit high priority request
|
|
let high_priority_request = InferenceRequest {
|
|
id: RequestId::new(),
|
|
priority: RequestPriority::High,
|
|
input_tokens: vec![1, 2, 3],
|
|
..Default::default()
|
|
};
|
|
|
|
scheduler
|
|
.submit_request(high_priority_request)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Submit normal priority request
|
|
let normal_priority_request = InferenceRequest {
|
|
id: RequestId::new(),
|
|
priority: RequestPriority::Normal,
|
|
input_tokens: vec![4, 5, 6],
|
|
..Default::default()
|
|
};
|
|
|
|
scheduler
|
|
.submit_request(normal_priority_request)
|
|
.await
|
|
.unwrap();
|
|
|
|
let queue_stats = scheduler.queue_stats().await;
|
|
assert_eq!(queue_stats.lanes.get("premium").unwrap().pending_count, 1);
|
|
assert_eq!(queue_stats.lanes.get("standard").unwrap().pending_count, 1);
|
|
}
|
|
|
|
/// Test dynamic batch formation with different batch sizes
|
|
#[tokio::test]
|
|
async fn test_dynamic_batch_formation() {
|
|
let config = BatchSchedulerConfig {
|
|
max_batch_size: 4,
|
|
max_wait_time: Duration::from_millis(50),
|
|
..Default::default()
|
|
};
|
|
|
|
let mut scheduler = BatchScheduler::new(config).await.unwrap();
|
|
|
|
// Submit multiple requests rapidly
|
|
let mut request_ids = Vec::new();
|
|
for i in 0..6 {
|
|
let request = InferenceRequest {
|
|
id: RequestId::new(),
|
|
input_tokens: vec![i; 10],
|
|
max_new_tokens: 50,
|
|
..Default::default()
|
|
};
|
|
request_ids.push(request.id);
|
|
scheduler.submit_request(request).await.unwrap();
|
|
}
|
|
|
|
// Should form batches of size 4, 2
|
|
let batch1 = scheduler.get_next_batch().await.unwrap();
|
|
assert_eq!(batch1.requests.len(), 4);
|
|
|
|
let batch2 = scheduler.get_next_batch().await.unwrap();
|
|
assert_eq!(batch2.requests.len(), 2);
|
|
}
|
|
|
|
/// Test batch scheduling with SLA constraints
|
|
#[tokio::test]
|
|
async fn test_sla_constrained_batching() {
|
|
let config = BatchSchedulerConfig {
|
|
sla_lanes: vec![
|
|
SlaLane {
|
|
name: "premium".to_string(),
|
|
priority: RequestPriority::High,
|
|
max_latency: Duration::from_millis(30),
|
|
max_batch_size: 2, // Small batch size for low latency
|
|
memory_limit: Some(1024 * 1024 * 10), // 10MB
|
|
min_wait_time: Some(Duration::from_millis(1)),
|
|
},
|
|
SlaLane {
|
|
name: "standard".to_string(),
|
|
priority: RequestPriority::Normal,
|
|
max_latency: Duration::from_millis(100),
|
|
max_batch_size: 8,
|
|
memory_limit: Some(1024 * 1024 * 50), // 50MB
|
|
min_wait_time: Some(Duration::from_millis(5)),
|
|
},
|
|
],
|
|
..Default::default()
|
|
};
|
|
|
|
let mut scheduler = BatchScheduler::new(config).await.unwrap();
|
|
|
|
// Submit premium requests - should batch quickly due to SLA
|
|
for i in 0..3 {
|
|
let request = InferenceRequest {
|
|
id: RequestId::new(),
|
|
priority: RequestPriority::High,
|
|
input_tokens: vec![i; 5],
|
|
..Default::default()
|
|
};
|
|
scheduler.submit_request(request).await.unwrap();
|
|
}
|
|
|
|
// Premium batch should form with max 2 requests
|
|
let premium_batch = scheduler.get_next_batch().await.unwrap();
|
|
assert_eq!(premium_batch.requests.len(), 2);
|
|
assert_eq!(premium_batch.lane_name, "premium");
|
|
|
|
// Remaining premium request should form another batch
|
|
let remaining_batch = scheduler.get_next_batch().await.unwrap();
|
|
assert_eq!(remaining_batch.requests.len(), 1);
|
|
assert_eq!(remaining_batch.lane_name, "premium");
|
|
}
|
|
|
|
/// Test preemption of lower priority batches
|
|
#[tokio::test]
|
|
async fn test_batch_preemption() {
|
|
let config = BatchSchedulerConfig {
|
|
enable_preemption: true,
|
|
sla_lanes: vec![
|
|
SlaLane {
|
|
name: "premium".to_string(),
|
|
priority: RequestPriority::High,
|
|
max_latency: Duration::from_millis(50),
|
|
max_batch_size: 4,
|
|
memory_limit: Some(1024 * 1024 * 25), // 25MB
|
|
min_wait_time: Some(Duration::from_millis(2)),
|
|
},
|
|
SlaLane {
|
|
name: "standard".to_string(),
|
|
priority: RequestPriority::Normal,
|
|
max_latency: Duration::from_millis(200),
|
|
max_batch_size: 8,
|
|
memory_limit: Some(1024 * 1024 * 50), // 50MB
|
|
min_wait_time: Some(Duration::from_millis(5)),
|
|
},
|
|
],
|
|
..Default::default()
|
|
};
|
|
|
|
let mut scheduler = BatchScheduler::new(config).await.unwrap();
|
|
|
|
// Submit normal priority requests
|
|
for i in 0..4 {
|
|
let request = InferenceRequest {
|
|
id: RequestId::new(),
|
|
priority: RequestPriority::Normal,
|
|
input_tokens: vec![i; 20],
|
|
..Default::default()
|
|
};
|
|
scheduler.submit_request(request).await.unwrap();
|
|
}
|
|
|
|
// Get standard batch for processing
|
|
let standard_batch = scheduler.get_next_batch().await.unwrap();
|
|
assert_eq!(standard_batch.lane_name, "standard");
|
|
|
|
// Mark batch as executing
|
|
scheduler
|
|
.mark_batch_executing(standard_batch.id)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Submit high priority request - should trigger preemption consideration
|
|
let high_priority_request = InferenceRequest {
|
|
id: RequestId::new(),
|
|
priority: RequestPriority::High,
|
|
input_tokens: vec![99; 10],
|
|
deadline: Some(Instant::now() + Duration::from_millis(30)),
|
|
..Default::default()
|
|
};
|
|
|
|
scheduler
|
|
.submit_request(high_priority_request)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Should suggest preemption
|
|
let preemption_decision = scheduler.evaluate_preemption().await.unwrap();
|
|
assert!(preemption_decision.should_preempt);
|
|
assert_eq!(preemption_decision.target_batch_id, Some(standard_batch.id));
|
|
}
|
|
|
|
/// Test SLA violation detection and handling
|
|
#[tokio::test]
|
|
async fn test_sla_violation_detection() {
|
|
let config = BatchSchedulerConfig {
|
|
sla_lanes: vec![SlaLane {
|
|
name: "premium".to_string(),
|
|
priority: RequestPriority::High,
|
|
max_latency: Duration::from_millis(100),
|
|
max_batch_size: 4,
|
|
memory_limit: Some(1024 * 1024 * 25), // 25MB
|
|
min_wait_time: Some(Duration::from_millis(2)),
|
|
}],
|
|
..Default::default()
|
|
};
|
|
|
|
let mut scheduler = BatchScheduler::new(config).await.unwrap();
|
|
|
|
let request = InferenceRequest {
|
|
id: RequestId::new(),
|
|
priority: RequestPriority::High,
|
|
input_tokens: vec![1, 2, 3],
|
|
deadline: Some(Instant::now() + Duration::from_millis(50)),
|
|
..Default::default()
|
|
};
|
|
|
|
let request_id = request.id;
|
|
scheduler.submit_request(request).await.unwrap();
|
|
|
|
// Wait longer than deadline
|
|
sleep(Duration::from_millis(60)).await;
|
|
|
|
// Check for SLA violations
|
|
let violations = scheduler.check_sla_violations().await.unwrap();
|
|
assert!(!violations.is_empty());
|
|
assert!(violations.iter().any(|v| v.request_id == request_id));
|
|
}
|
|
|
|
/// Test memory pressure handling in batch formation
|
|
#[tokio::test]
|
|
async fn test_memory_pressure_batching() {
|
|
let config = BatchSchedulerConfig {
|
|
max_batch_size: 16,
|
|
memory_pressure_threshold: 0.8, // 80% memory usage
|
|
..Default::default()
|
|
};
|
|
|
|
let mut scheduler = BatchScheduler::new(config).await.unwrap();
|
|
|
|
// Simulate high memory pressure
|
|
scheduler.set_memory_pressure(0.85).await.unwrap();
|
|
|
|
// Submit requests with varying memory requirements
|
|
let large_request = InferenceRequest {
|
|
id: RequestId::new(),
|
|
input_tokens: vec![1; 2048], // Large context
|
|
max_new_tokens: 1024,
|
|
..Default::default()
|
|
};
|
|
|
|
scheduler.submit_request(large_request).await.unwrap();
|
|
|
|
let small_request = InferenceRequest {
|
|
id: RequestId::new(),
|
|
input_tokens: vec![1; 128], // Small context
|
|
max_new_tokens: 64,
|
|
..Default::default()
|
|
};
|
|
|
|
scheduler.submit_request(small_request).await.unwrap();
|
|
|
|
// Under memory pressure, should create smaller batches
|
|
let batch = scheduler.get_next_batch().await.unwrap();
|
|
assert!(batch.requests.len() <= 2); // Should limit batch size due to memory pressure
|
|
assert!(batch.estimated_memory_usage <= scheduler.config().memory_limit);
|
|
}
|
|
|
|
/// Test graceful degradation under extreme load
|
|
#[tokio::test]
|
|
async fn test_graceful_degradation() {
|
|
let config = BatchSchedulerConfig {
|
|
max_batch_size: 4,
|
|
max_queue_size: 10, // Limited queue size
|
|
degradation_mode_threshold: 8,
|
|
..Default::default()
|
|
};
|
|
|
|
let mut scheduler = BatchScheduler::new(config).await.unwrap();
|
|
|
|
// Submit more requests than queue can handle
|
|
let mut rejected_count = 0;
|
|
for i in 0..15 {
|
|
let request = InferenceRequest {
|
|
id: RequestId::new(),
|
|
input_tokens: vec![i; 10],
|
|
..Default::default()
|
|
};
|
|
|
|
match scheduler.submit_request(request).await {
|
|
Ok(_) => {}
|
|
Err(InferenceError::QueueFull { .. }) => rejected_count += 1,
|
|
Err(e) => panic!("Unexpected error: {}", e),
|
|
}
|
|
}
|
|
|
|
// Should reject some requests to maintain stability
|
|
assert!(rejected_count > 0);
|
|
assert!(rejected_count >= 5); // Should reject excess requests
|
|
|
|
// Queue should be in degradation mode
|
|
assert!(scheduler.is_degradation_mode().await);
|
|
}
|