456 lines
16 KiB
Rust
456 lines
16 KiB
Rust
#![cfg(feature = "disabled_tests")]
|
|
|
|
use rtx_preprocessing::{
|
|
FaultToleranceConfig,
|
|
// Loaders
|
|
IntegratedDataLoader,
|
|
IntegratedLoaderConfig,
|
|
LoaderConfig,
|
|
MinMaxScaler,
|
|
PrefetchStrategy,
|
|
PreprocessingError,
|
|
RebalancingStrategy,
|
|
// Core types
|
|
Result,
|
|
ShardingConfig,
|
|
ShardingStrategy,
|
|
// Transformers
|
|
StandardScaler,
|
|
Transformer,
|
|
WorkerInfo,
|
|
rtx_stubs::rtx_tensor::{Device, Tensor},
|
|
};
|
|
use std::fs::{File, create_dir_all};
|
|
use std::io::Write;
|
|
use std::path::{Path, PathBuf};
|
|
use tempfile::TempDir;
|
|
|
|
/// Helper function to create test dataset with known patterns
|
|
fn create_structured_test_dataset(
|
|
dir: &Path,
|
|
num_files: usize,
|
|
samples_per_file: usize,
|
|
) -> Vec<PathBuf> {
|
|
let mut file_paths = Vec::new();
|
|
|
|
for file_idx in 0..num_files {
|
|
let file_name = format!("dataset_{:03}.bin", file_idx);
|
|
let file_path = dir.join(&file_name);
|
|
let mut file = File::create(&file_path).unwrap();
|
|
|
|
// Create data with a predictable pattern for testing transformers
|
|
// Each file has a different base value to make sharding effects visible
|
|
let base_value = file_idx as f32 * 10.0;
|
|
|
|
for i in 0..samples_per_file {
|
|
// Create a simple linear pattern: base + i + some variance
|
|
let variance = (i % 7) as f32 * 0.1; // Small variance for testing
|
|
let value = base_value + (i as f32) + variance;
|
|
file.write_all(&value.to_le_bytes()).unwrap();
|
|
}
|
|
file.flush().unwrap();
|
|
|
|
file_paths.push(file_path);
|
|
}
|
|
|
|
file_paths
|
|
}
|
|
|
|
#[test]
|
|
fn test_integrated_loader_initialization() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let file_paths = create_structured_test_dataset(temp_dir.path(), 6, 200);
|
|
|
|
let config = IntegratedLoaderConfig {
|
|
mmap_config: LoaderConfig {
|
|
page_size: 4096,
|
|
prefetch_strategy: PrefetchStrategy::Sequential { window_size: 2 },
|
|
use_shared_memory: true,
|
|
enable_lazy_loading: true,
|
|
},
|
|
sharding_config: ShardingConfig {
|
|
num_shards: 3,
|
|
sharding_strategy: ShardingStrategy::RoundRobin,
|
|
rebalancing_strategy: RebalancingStrategy::Static,
|
|
fault_tolerance: FaultToleranceConfig::default(),
|
|
},
|
|
enable_cross_shard_prefetch: true,
|
|
batch_size: 100,
|
|
};
|
|
|
|
let workers = vec![
|
|
WorkerInfo::new(0, "worker0".to_string(), "127.0.0.1:8000".to_string()),
|
|
WorkerInfo::new(1, "worker1".to_string(), "127.0.0.1:8001".to_string()),
|
|
WorkerInfo::new(2, "worker2".to_string(), "127.0.0.1:8002".to_string()),
|
|
];
|
|
|
|
let mut loader = IntegratedDataLoader::new(config, workers);
|
|
|
|
// Test initialization
|
|
let result = loader.initialize(&file_paths);
|
|
assert!(result.is_ok());
|
|
assert!(loader.is_initialized());
|
|
|
|
// Verify distributed sharding worked
|
|
let distributed_stats = loader.distributed_loader().get_statistics();
|
|
assert_eq!(distributed_stats.total_shards, 3);
|
|
assert_eq!(distributed_stats.total_files, 6);
|
|
assert_eq!(distributed_stats.active_workers, 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_integrated_loader_with_transformers() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let file_paths = create_structured_test_dataset(temp_dir.path(), 4, 100);
|
|
|
|
let config = IntegratedLoaderConfig::default();
|
|
let workers = vec![
|
|
WorkerInfo::new(0, "worker0".to_string(), "127.0.0.1:8000".to_string()),
|
|
WorkerInfo::new(1, "worker1".to_string(), "127.0.0.1:8001".to_string()),
|
|
];
|
|
|
|
let mut loader = IntegratedDataLoader::new(config, workers);
|
|
loader.initialize(&file_paths).unwrap();
|
|
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
|
|
// Load data from worker 0
|
|
let worker_batch = loader.load_worker_batch(0, 0, &device).unwrap();
|
|
assert!(!worker_batch.is_empty());
|
|
|
|
// Test with StandardScaler
|
|
let tensor = &worker_batch[0];
|
|
let mut scaler = StandardScaler::new();
|
|
|
|
// Fit the scaler
|
|
let fit_result = scaler.fit(tensor);
|
|
assert!(fit_result.is_ok());
|
|
assert!(scaler.is_fitted());
|
|
|
|
// Transform the data
|
|
let scaled_tensor = scaler.transform(tensor).unwrap();
|
|
|
|
// Verify the transformation worked
|
|
assert_eq!(scaled_tensor.shape().dims(), tensor.shape().dims());
|
|
|
|
// The scaled data should have different values than original
|
|
let original_data: Vec<f32> = tensor.iter().map(|&x| x as f32).take(10).collect();
|
|
let scaled_data: Vec<f32> = scaled_tensor.iter().map(|&x| x as f32).take(10).collect();
|
|
|
|
// Values should be different after scaling (unless variance is exactly 1.0)
|
|
let has_different_values = original_data
|
|
.iter()
|
|
.zip(scaled_data.iter())
|
|
.any(|(&orig, &scaled)| (orig - scaled).abs() > 1e-6);
|
|
|
|
assert!(has_different_values);
|
|
}
|
|
|
|
#[test]
|
|
fn test_distributed_batch_loading_with_preprocessing() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let file_paths = create_structured_test_dataset(temp_dir.path(), 8, 50);
|
|
|
|
let config = IntegratedLoaderConfig {
|
|
sharding_config: ShardingConfig {
|
|
num_shards: 4,
|
|
sharding_strategy: ShardingStrategy::SizeAware,
|
|
rebalancing_strategy: RebalancingStrategy::Static,
|
|
fault_tolerance: FaultToleranceConfig::default(),
|
|
},
|
|
batch_size: 25,
|
|
..Default::default()
|
|
};
|
|
|
|
let workers = vec![
|
|
WorkerInfo::new(0, "worker0".to_string(), "127.0.0.1:8000".to_string()),
|
|
WorkerInfo::new(1, "worker1".to_string(), "127.0.0.1:8001".to_string()),
|
|
];
|
|
|
|
let mut loader = IntegratedDataLoader::new(config, workers.clone());
|
|
loader.initialize(&file_paths).unwrap();
|
|
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
|
|
// Load distributed batch
|
|
let distributed_batch = loader.load_distributed_batch(0, &device).unwrap();
|
|
|
|
// Should have data for both workers
|
|
assert!(distributed_batch.contains_key(&0));
|
|
assert!(distributed_batch.contains_key(&1));
|
|
|
|
// Process each worker's data with different scalers
|
|
let mut worker_scalers: std::collections::HashMap<u32, (StandardScaler, MinMaxScaler)> =
|
|
std::collections::HashMap::new();
|
|
|
|
for (&worker_id, tensors) in &distributed_batch {
|
|
let mut std_scaler = StandardScaler::new();
|
|
let mut minmax_scaler = MinMaxScaler::new();
|
|
|
|
if let Some(tensor) = tensors.first() {
|
|
// Fit both scalers
|
|
std_scaler.fit(tensor).unwrap();
|
|
minmax_scaler.fit(tensor).unwrap();
|
|
|
|
worker_scalers.insert(worker_id, (std_scaler, minmax_scaler));
|
|
}
|
|
}
|
|
|
|
// Verify that both workers have fitted scalers
|
|
assert_eq!(worker_scalers.len(), 2);
|
|
|
|
// Transform data for each worker
|
|
for (&worker_id, tensors) in &distributed_batch {
|
|
if let Some((std_scaler, minmax_scaler)) = worker_scalers.get(&worker_id) {
|
|
for tensor in tensors {
|
|
let std_transformed = std_scaler.transform(tensor).unwrap();
|
|
let minmax_transformed = minmax_scaler.transform(tensor).unwrap();
|
|
|
|
// Verify transformations
|
|
assert_eq!(std_transformed.shape().dims(), tensor.shape().dims());
|
|
assert_eq!(minmax_transformed.shape().dims(), tensor.shape().dims());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_fault_tolerance_with_preprocessing_pipeline() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let file_paths = create_structured_test_dataset(temp_dir.path(), 6, 100);
|
|
|
|
let config = IntegratedLoaderConfig {
|
|
sharding_config: ShardingConfig {
|
|
num_shards: 3,
|
|
sharding_strategy: ShardingStrategy::RoundRobin,
|
|
rebalancing_strategy: RebalancingStrategy::Dynamic {
|
|
imbalance_threshold: 0.3,
|
|
rebalance_interval_secs: 5,
|
|
},
|
|
fault_tolerance: FaultToleranceConfig {
|
|
enable_health_checks: true,
|
|
health_check_interval_secs: 1,
|
|
max_failures: 2,
|
|
recovery_timeout_secs: 10,
|
|
},
|
|
},
|
|
batch_size: 50,
|
|
..Default::default()
|
|
};
|
|
|
|
let workers = vec![
|
|
WorkerInfo::new(0, "worker0".to_string(), "127.0.0.1:8000".to_string()),
|
|
WorkerInfo::new(1, "worker1".to_string(), "127.0.0.1:8001".to_string()),
|
|
WorkerInfo::new(2, "worker2".to_string(), "127.0.0.1:8002".to_string()),
|
|
];
|
|
|
|
let mut loader = IntegratedDataLoader::new(config, workers);
|
|
loader.initialize(&file_paths).unwrap();
|
|
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
|
|
// Initial state: all workers healthy
|
|
assert!(loader.distributed_loader().is_worker_healthy(0));
|
|
assert!(loader.distributed_loader().is_worker_healthy(1));
|
|
assert!(loader.distributed_loader().is_worker_healthy(2));
|
|
|
|
// Simulate worker failure
|
|
loader.distributed_loader().report_worker_failure(1);
|
|
loader.distributed_loader().report_worker_failure(1); // Make worker 1 unhealthy
|
|
|
|
assert!(!loader.distributed_loader().is_worker_healthy(1));
|
|
|
|
// Load data should still work with remaining workers
|
|
let distributed_batch = loader.load_distributed_batch(0, &device).unwrap();
|
|
|
|
// Should have data from healthy workers only
|
|
assert!(distributed_batch.contains_key(&0));
|
|
assert!(distributed_batch.contains_key(&2));
|
|
assert!(!distributed_batch.contains_key(&1)); // Failed worker should have no data
|
|
|
|
// Test preprocessing still works
|
|
for tensors in distributed_batch.values() {
|
|
for tensor in tensors {
|
|
let mut scaler = StandardScaler::new();
|
|
scaler.fit(tensor).unwrap();
|
|
let transformed = scaler.transform(tensor).unwrap();
|
|
assert_eq!(transformed.shape().dims(), tensor.shape().dims());
|
|
}
|
|
}
|
|
|
|
// Test rebalancing
|
|
let rebalance_result = loader.rebalance_and_reload();
|
|
assert!(rebalance_result.is_ok());
|
|
|
|
// After rebalancing, healthy workers should have more shards
|
|
let stats = loader.get_comprehensive_statistics();
|
|
assert_eq!(stats.distributed_stats.active_workers, 2); // Only 2 healthy workers
|
|
assert!(stats.distributed_stats.rebalancing_events > 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_performance_metrics_and_statistics() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let file_paths = create_structured_test_dataset(temp_dir.path(), 10, 100);
|
|
|
|
let config = IntegratedLoaderConfig {
|
|
mmap_config: LoaderConfig {
|
|
page_size: 2048,
|
|
prefetch_strategy: PrefetchStrategy::Sequential { window_size: 3 },
|
|
use_shared_memory: true,
|
|
enable_lazy_loading: true,
|
|
},
|
|
sharding_config: ShardingConfig {
|
|
num_shards: 5,
|
|
sharding_strategy: ShardingStrategy::HashBased,
|
|
rebalancing_strategy: RebalancingStrategy::Static,
|
|
fault_tolerance: FaultToleranceConfig::default(),
|
|
},
|
|
enable_cross_shard_prefetch: true,
|
|
batch_size: 200,
|
|
};
|
|
|
|
let workers = vec![
|
|
WorkerInfo::new(0, "worker0".to_string(), "127.0.0.1:8000".to_string()),
|
|
WorkerInfo::new(1, "worker1".to_string(), "127.0.0.1:8001".to_string()),
|
|
];
|
|
|
|
let mut loader = IntegratedDataLoader::new(config, workers);
|
|
loader.initialize(&file_paths).unwrap();
|
|
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
|
|
// Perform multiple loads to generate statistics
|
|
for batch_idx in 0..3 {
|
|
let _batch = loader.load_distributed_batch(batch_idx, &device).unwrap();
|
|
}
|
|
|
|
// Report some performance metrics
|
|
loader
|
|
.distributed_loader()
|
|
.report_worker_performance(0, 100);
|
|
loader
|
|
.distributed_loader()
|
|
.report_worker_performance(1, 150);
|
|
|
|
// Get comprehensive statistics
|
|
let stats = loader.get_comprehensive_statistics();
|
|
|
|
// Verify statistics
|
|
assert_eq!(stats.distributed_stats.total_shards, 5);
|
|
assert_eq!(stats.distributed_stats.total_files, 10);
|
|
assert_eq!(stats.distributed_stats.active_workers, 2);
|
|
assert!(stats.distributed_stats.average_worker_performance > 0.0);
|
|
|
|
// Memory-mapped statistics
|
|
assert!(stats.mmap_stats.total_reads > 0);
|
|
assert!(stats.active_mmap_loaders > 0);
|
|
|
|
// Test derived metrics
|
|
let cache_hit_ratio = stats.cache_hit_ratio();
|
|
assert!(cache_hit_ratio >= 0.0 && cache_hit_ratio <= 1.0);
|
|
|
|
let efficiency_score = stats.efficiency_score();
|
|
assert!(efficiency_score >= 0.0 && efficiency_score <= 1.0);
|
|
|
|
// With no failures, efficiency should be high
|
|
assert!(efficiency_score > 0.5);
|
|
}
|
|
|
|
#[test]
|
|
fn test_concurrent_access_with_transformers() {
|
|
use std::sync::Arc;
|
|
use std::thread;
|
|
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let file_paths = create_structured_test_dataset(temp_dir.path(), 12, 50);
|
|
|
|
let config = IntegratedLoaderConfig {
|
|
sharding_config: ShardingConfig {
|
|
num_shards: 6,
|
|
sharding_strategy: ShardingStrategy::RoundRobin,
|
|
rebalancing_strategy: RebalancingStrategy::Static,
|
|
fault_tolerance: FaultToleranceConfig::default(),
|
|
},
|
|
batch_size: 25,
|
|
..Default::default()
|
|
};
|
|
|
|
let workers = vec![
|
|
WorkerInfo::new(0, "worker0".to_string(), "127.0.0.1:8000".to_string()),
|
|
WorkerInfo::new(1, "worker1".to_string(), "127.0.0.1:8001".to_string()),
|
|
WorkerInfo::new(2, "worker2".to_string(), "127.0.0.1:8002".to_string()),
|
|
];
|
|
|
|
let mut loader = IntegratedDataLoader::new(config, workers.clone());
|
|
loader.initialize(&file_paths).unwrap();
|
|
let loader = Arc::new(loader);
|
|
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let mut handles = vec![];
|
|
|
|
// Spawn threads for concurrent data loading and processing
|
|
for worker_id in 0..3 {
|
|
let loader_clone = Arc::clone(&loader);
|
|
let device_clone = device.clone();
|
|
|
|
let handle = thread::spawn(move || {
|
|
// Each thread loads data for its assigned worker and applies transformations
|
|
let batch = loader_clone
|
|
.load_worker_batch(worker_id, 0, &device_clone)
|
|
.unwrap();
|
|
|
|
let mut processed_count = 0;
|
|
for tensor in batch {
|
|
// Apply different transformers based on worker_id
|
|
match worker_id {
|
|
0 => {
|
|
let mut scaler = StandardScaler::new();
|
|
scaler.fit(&tensor).unwrap();
|
|
let _transformed = scaler.transform(&tensor).unwrap();
|
|
processed_count += 1;
|
|
}
|
|
1 => {
|
|
let mut scaler = MinMaxScaler::new();
|
|
scaler.fit(&tensor).unwrap();
|
|
let _transformed = scaler.transform(&tensor).unwrap();
|
|
processed_count += 1;
|
|
}
|
|
2 => {
|
|
// Worker 2 uses both transformers in sequence
|
|
let mut std_scaler = StandardScaler::new();
|
|
let mut minmax_scaler = MinMaxScaler::new();
|
|
|
|
std_scaler.fit(&tensor).unwrap();
|
|
let std_transformed = std_scaler.transform(&tensor).unwrap();
|
|
|
|
minmax_scaler.fit(&std_transformed).unwrap();
|
|
let _final_transformed = minmax_scaler.transform(&std_transformed).unwrap();
|
|
processed_count += 1;
|
|
}
|
|
_ => unreachable!(),
|
|
}
|
|
}
|
|
|
|
processed_count
|
|
});
|
|
|
|
handles.push(handle);
|
|
}
|
|
|
|
// Wait for all threads and collect results
|
|
let mut total_processed = 0;
|
|
for handle in handles {
|
|
let count = handle.join().unwrap();
|
|
total_processed += count;
|
|
}
|
|
|
|
// Verify that data was processed
|
|
assert!(total_processed > 0);
|
|
|
|
// Verify loader state is still consistent
|
|
let final_stats = loader.get_comprehensive_statistics();
|
|
assert_eq!(final_stats.distributed_stats.total_files, 12);
|
|
assert_eq!(final_stats.distributed_stats.active_workers, 3);
|
|
}
|