Initial commit

2026-03-04 00:08:42 +00:00
commit 4d88dc0584
4449 changed files with 1556714 additions and 0 deletions
@@ -0,0 +1,585 @@
+//! Elastic Training
+//!
+//! This module provides elastic training capabilities for distributed training:
+//! - Dynamic worker scaling (add/remove workers during training)
+//! - Graceful scaling without stopping training
+//! - Automatic state redistribution on scale events
+//! - Checkpoint-based recovery for failed workers
+//! - Health monitoring and failure detection
+//!
+//! Elastic training enables efficient resource utilization by allowing
+//! the training cluster to grow or shrink based on availability.
+
+mod cluster;
+mod config;
+mod health;
+mod redistributor;
+mod scale_ops;
+mod stats;
+mod worker;
+
+// Re-export all public types
+pub use cluster::ElasticCluster;
+pub use config::{ElasticConfig, ScalingPolicy, WorkerState};
+pub use health::HealthMonitor;
+pub use redistributor::StateRedistributor;
+pub use scale_ops::{ScaleEventType, ScaleOpState, ScaleOpType, ScaleOperation, ShardMoveOp};
+pub use stats::{ElasticStats, ExecutionReport, RedistributionStats, ShardMove};
+pub use worker::{WorkerId, WorkerInfo};
+
+use std::sync::Arc;
+
+// =============================================================================
+// Thread-Safe Wrappers
+// =============================================================================
+
+/// Thread-safe shared elastic cluster
+pub type SharedElasticCluster = Arc<ElasticCluster>;
+
+/// Create a shared elastic cluster
+pub fn shared_elastic_cluster(config: ElasticConfig) -> SharedElasticCluster {
+    Arc::new(ElasticCluster::new(config))
+}
+
+// =============================================================================
+// Tests
+// =============================================================================
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::time::Duration;
+
+    #[test]
+    fn test_elastic_config_default() {
+        let config = ElasticConfig::default();
+        assert_eq!(config.min_workers, 1);
+        assert_eq!(config.max_workers, 64);
+        assert_eq!(config.initial_workers, 4);
+        assert_eq!(config.scaling_policy, ScalingPolicy::Manual);
+    }
+
+    #[test]
+    fn test_worker_info_creation() {
+        let worker = WorkerInfo::new(1, 0, "host1".to_string(), vec![0, 1]);
+        assert_eq!(worker.id, 1);
+        assert_eq!(worker.rank, 0);
+        assert_eq!(worker.num_devices(), 2);
+        assert_eq!(worker.state, WorkerState::Initializing);
+    }
+
+    #[test]
+    fn test_worker_heartbeat() {
+        let mut worker = WorkerInfo::new(1, 0, "host1".to_string(), vec![0]);
+        let old_hb = worker.last_heartbeat;
+
+        std::thread::sleep(Duration::from_millis(10));
+        worker.heartbeat();
+
+        assert!(worker.last_heartbeat > old_hb);
+    }
+
+    #[test]
+    fn test_worker_health_check() {
+        use std::time::Instant;
+
+        let worker = WorkerInfo {
+            id: 1,
+            rank: 0,
+            hostname: "host1".to_string(),
+            device_ids: vec![0],
+            state: WorkerState::Active,
+            last_heartbeat: Instant::now(),
+            joined_at: Instant::now(),
+            batches_processed: 0,
+            throughput: 0.0,
+        };
+
+        assert!(worker.is_healthy(Duration::from_secs(1)));
+        assert!(!worker.is_healthy(Duration::from_nanos(1)));
+    }
+
+    #[test]
+    fn test_elastic_cluster_creation() {
+        let config = ElasticConfig::default();
+        let cluster = ElasticCluster::new(config);
+
+        assert_eq!(cluster.world_size(), 0);
+        assert!(!cluster.is_running());
+    }
+
+    #[test]
+    fn test_cluster_start_stop() {
+        let config = ElasticConfig::default();
+        let cluster = ElasticCluster::new(config);
+
+        cluster.start().unwrap();
+        assert!(cluster.is_running());
+
+        cluster.stop();
+        assert!(!cluster.is_running());
+    }
+
+    #[test]
+    fn test_register_worker() {
+        let config = ElasticConfig::default();
+        let cluster = ElasticCluster::new(config);
+
+        let worker = cluster
+            .register_worker("host1".to_string(), vec![0, 1])
+            .unwrap();
+
+        assert_eq!(worker.rank, 0);
+        assert_eq!(cluster.world_size(), 1);
+    }
+
+    #[test]
+    fn test_register_multiple_workers() {
+        let config = ElasticConfig::default();
+        let cluster = ElasticCluster::new(config);
+
+        let w1 = cluster
+            .register_worker("host1".to_string(), vec![0])
+            .unwrap();
+        let w2 = cluster
+            .register_worker("host2".to_string(), vec![0])
+            .unwrap();
+
+        assert_eq!(w1.rank, 0);
+        assert_eq!(w2.rank, 1);
+        assert_eq!(cluster.world_size(), 2);
+    }
+
+    #[test]
+    fn test_activate_worker() {
+        let config = ElasticConfig::default();
+        let cluster = ElasticCluster::new(config);
+
+        let worker = cluster
+            .register_worker("host1".to_string(), vec![0])
+            .unwrap();
+        cluster.activate_worker(worker.id).unwrap();
+
+        let info = cluster.worker(worker.id).unwrap();
+        assert_eq!(info.state, WorkerState::Active);
+    }
+
+    #[test]
+    fn test_max_workers_limit() {
+        let config = ElasticConfig {
+            max_workers: 2,
+            ..Default::default()
+        };
+        let cluster = ElasticCluster::new(config);
+
+        cluster
+            .register_worker("host1".to_string(), vec![0])
+            .unwrap();
+        cluster
+            .register_worker("host2".to_string(), vec![0])
+            .unwrap();
+
+        // Third should fail
+        let result = cluster.register_worker("host3".to_string(), vec![0]);
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_heartbeat() {
+        let config = ElasticConfig::default();
+        let cluster = ElasticCluster::new(config);
+
+        let worker = cluster
+            .register_worker("host1".to_string(), vec![0])
+            .unwrap();
+        cluster.heartbeat(worker.id, 100, 50.0).unwrap();
+
+        let info = cluster.worker(worker.id).unwrap();
+        assert_eq!(info.batches_processed, 100);
+        assert!((info.throughput - 50.0).abs() < 0.01);
+    }
+
+    #[test]
+    fn test_total_throughput() {
+        let config = ElasticConfig::default();
+        let cluster = ElasticCluster::new(config);
+
+        let w1 = cluster
+            .register_worker("host1".to_string(), vec![0])
+            .unwrap();
+        let w2 = cluster
+            .register_worker("host2".to_string(), vec![0])
+            .unwrap();
+
+        cluster.activate_worker(w1.id).unwrap();
+        cluster.activate_worker(w2.id).unwrap();
+
+        cluster.heartbeat(w1.id, 100, 50.0).unwrap();
+        cluster.heartbeat(w2.id, 100, 30.0).unwrap();
+
+        let total = cluster.total_throughput();
+        assert!((total - 80.0).abs() < 0.01);
+    }
+
+    #[test]
+    fn test_scale_up_request() {
+        let config = ElasticConfig {
+            scale_cooldown_ms: 0, // Disable cooldown for test
+            ..Default::default()
+        };
+        let cluster = ElasticCluster::new(config);
+
+        // Register initial workers
+        cluster
+            .register_worker("host1".to_string(), vec![0])
+            .unwrap();
+
+        let op_id = cluster.request_scale_up(2, "test".to_string()).unwrap();
+        assert!(op_id > 0);
+        assert_eq!(cluster.pending_ops_count(), 1);
+    }
+
+    #[test]
+    fn test_scale_up_exceeds_max() {
+        let config = ElasticConfig {
+            max_workers: 4,
+            scale_cooldown_ms: 0,
+            ..Default::default()
+        };
+        let cluster = ElasticCluster::new(config);
+
+        // Register 3 workers
+        for i in 0..3 {
+            cluster
+                .register_worker(format!("host{}", i), vec![0])
+                .unwrap();
+        }
+
+        // Try to scale up by 2 (would exceed max of 4)
+        let result = cluster.request_scale_up(2, "test".to_string());
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_scale_down_request() {
+        let config = ElasticConfig {
+            min_workers: 1,
+            scale_cooldown_ms: 0,
+            ..Default::default()
+        };
+        let cluster = ElasticCluster::new(config);
+
+        // Register workers
+        let w1 = cluster
+            .register_worker("host1".to_string(), vec![0])
+            .unwrap();
+        let w2 = cluster
+            .register_worker("host2".to_string(), vec![0])
+            .unwrap();
+
+        cluster.activate_worker(w1.id).unwrap();
+        cluster.activate_worker(w2.id).unwrap();
+
+        let op_id = cluster.request_scale_down(1, "test".to_string()).unwrap();
+        assert!(op_id > 0);
+    }
+
+    #[test]
+    fn test_scale_down_below_min() {
+        let config = ElasticConfig {
+            min_workers: 2,
+            scale_cooldown_ms: 0,
+            ..Default::default()
+        };
+        let cluster = ElasticCluster::new(config);
+
+        // Register 2 workers
+        for i in 0..2 {
+            cluster
+                .register_worker(format!("host{}", i), vec![0])
+                .unwrap();
+        }
+
+        // Try to scale down by 1 (would go below min of 2)
+        let result = cluster.request_scale_down(1, "test".to_string());
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_failure_detection() {
+        let config = ElasticConfig {
+            failure_timeout_ms: 10,
+            ..Default::default()
+        };
+        let cluster = ElasticCluster::new(config);
+
+        let worker = cluster
+            .register_worker("host1".to_string(), vec![0])
+            .unwrap();
+        cluster.activate_worker(worker.id).unwrap();
+
+        // Wait for timeout
+        std::thread::sleep(Duration::from_millis(20));
+
+        let failed = cluster.detect_failures();
+        assert_eq!(failed.len(), 1);
+        assert_eq!(failed[0], worker.id);
+    }
+
+    #[test]
+    fn test_mark_failed() {
+        let config = ElasticConfig::default();
+        let cluster = ElasticCluster::new(config);
+
+        let worker = cluster
+            .register_worker("host1".to_string(), vec![0])
+            .unwrap();
+        cluster.mark_failed(worker.id).unwrap();
+
+        let info = cluster.worker(worker.id).unwrap();
+        assert_eq!(info.state, WorkerState::Failed);
+    }
+
+    #[test]
+    fn test_request_replacement() {
+        let config = ElasticConfig::default();
+        let cluster = ElasticCluster::new(config);
+
+        let worker = cluster
+            .register_worker("host1".to_string(), vec![0])
+            .unwrap();
+        let op_id = cluster.request_replacement(worker.id).unwrap();
+
+        assert!(op_id > 0);
+        assert_eq!(cluster.pending_ops_count(), 1);
+    }
+
+    #[test]
+    fn test_health_monitor() {
+        let config = ElasticConfig {
+            failure_timeout_ms: 10,
+            ..Default::default()
+        };
+        let cluster = Arc::new(ElasticCluster::new(config));
+
+        let worker = cluster
+            .register_worker("host1".to_string(), vec![0])
+            .unwrap();
+        cluster.activate_worker(worker.id).unwrap();
+
+        let monitor = HealthMonitor::new(cluster.clone(), 100);
+        monitor.start();
+
+        // Wait for timeout
+        std::thread::sleep(Duration::from_millis(20));
+
+        let failed = monitor.check_once();
+        assert_eq!(failed.len(), 1);
+    }
+
+    #[test]
+    fn test_state_redistributor() {
+        let config = ElasticConfig::default();
+        let cluster = Arc::new(ElasticCluster::new(config));
+        let redistributor = StateRedistributor::new(cluster);
+
+        // Redistribute from 2 to 4 workers
+        redistributor.redistribute(2, 4).unwrap();
+
+        let stats = redistributor.stats();
+        assert_eq!(stats.redistributions, 1);
+    }
+
+    #[test]
+    fn test_no_redistribution_same_size() {
+        let config = ElasticConfig::default();
+        let cluster = Arc::new(ElasticCluster::new(config));
+        let redistributor = StateRedistributor::new(cluster);
+
+        redistributor.redistribute(4, 4).unwrap();
+
+        let stats = redistributor.stats();
+        assert_eq!(stats.redistributions, 0);
+    }
+
+    #[test]
+    fn test_shared_elastic_cluster() {
+        let config = ElasticConfig::default();
+        let cluster = shared_elastic_cluster(config);
+
+        cluster
+            .register_worker("host1".to_string(), vec![0])
+            .unwrap();
+        assert_eq!(cluster.world_size(), 1);
+    }
+
+    #[test]
+    fn test_elastic_stats() {
+        let config = ElasticConfig::default();
+        let cluster = ElasticCluster::new(config);
+
+        cluster.start().unwrap();
+        cluster
+            .register_worker("host1".to_string(), vec![0])
+            .unwrap();
+        cluster
+            .register_worker("host2".to_string(), vec![0])
+            .unwrap();
+
+        let stats = cluster.stats();
+        assert_eq!(stats.total_workers_joined, 2);
+        assert!(stats.start_time.is_some());
+    }
+
+    #[test]
+    fn test_active_worker_count() {
+        let config = ElasticConfig::default();
+        let cluster = ElasticCluster::new(config);
+
+        let w1 = cluster
+            .register_worker("host1".to_string(), vec![0])
+            .unwrap();
+        let w2 = cluster
+            .register_worker("host2".to_string(), vec![0])
+            .unwrap();
+
+        assert_eq!(cluster.active_worker_count(), 0);
+
+        cluster.activate_worker(w1.id).unwrap();
+        assert_eq!(cluster.active_worker_count(), 1);
+
+        cluster.activate_worker(w2.id).unwrap();
+        assert_eq!(cluster.active_worker_count(), 2);
+    }
+
+    #[test]
+    fn test_all_workers() {
+        let config = ElasticConfig::default();
+        let cluster = ElasticCluster::new(config);
+
+        cluster
+            .register_worker("host1".to_string(), vec![0])
+            .unwrap();
+        cluster
+            .register_worker("host2".to_string(), vec![0])
+            .unwrap();
+
+        let workers = cluster.all_workers();
+        assert_eq!(workers.len(), 2);
+    }
+
+    #[test]
+    fn test_scaling_policy() {
+        assert_eq!(ScalingPolicy::Manual, ScalingPolicy::Manual);
+        assert_ne!(ScalingPolicy::Manual, ScalingPolicy::ResourceBased);
+    }
+
+    #[test]
+    fn test_execution_report() {
+        let mut report = ExecutionReport::new();
+        assert_eq!(report.operations_started, 0);
+        assert_eq!(report.operations_completed, 0);
+        assert!(report.all_succeeded());
+
+        report.operations_started = 2;
+        report.operations_completed = 1;
+        report.operations_failed = 1;
+        report.add_success(1, Duration::from_millis(100));
+        report.add_failure(2, "Test error".to_string());
+        report.add_warning(1, "Test warning".to_string());
+
+        assert!(!report.all_succeeded());
+        assert_eq!(report.successes.len(), 1);
+        assert_eq!(report.failures.len(), 1);
+        assert_eq!(report.warnings.len(), 1);
+    }
+
+    #[test]
+    fn test_shard_move_op() {
+        let move_op = ShardMoveOp {
+            shard_id: 1,
+            from_rank: 0,
+            to_rank: 2,
+            size_bytes: 1024,
+        };
+        assert_eq!(move_op.shard_id, 1);
+        assert_eq!(move_op.from_rank, 0);
+        assert_eq!(move_op.to_rank, 2);
+    }
+
+    #[test]
+    fn test_scale_event_type() {
+        assert_eq!(
+            ScaleEventType::ScaleUpStarting,
+            ScaleEventType::ScaleUpStarting
+        );
+        assert_ne!(
+            ScaleEventType::ScaleUpStarting,
+            ScaleEventType::ScaleDownStarting
+        );
+    }
+
+    #[test]
+    fn test_calculate_redistribution_plan_scale_up() {
+        let config = ElasticConfig::default();
+        let cluster = ElasticCluster::new(config);
+
+        // Scale from 4 to 8 workers - larger numbers to ensure moves
+        let plan = cluster.calculate_redistribution_plan(4, 8);
+
+        // Some shards should move to new workers
+        assert!(!plan.is_empty());
+        for shard_move in &plan {
+            // Moves should be to new workers (ranks 4-7)
+            assert!(shard_move.to_rank >= 4);
+        }
+    }
+
+    #[test]
+    fn test_calculate_redistribution_plan_scale_down() {
+        let config = ElasticConfig::default();
+        let cluster = ElasticCluster::new(config);
+
+        // Scale from 4 to 2 workers
+        let plan = cluster.calculate_redistribution_plan(4, 2);
+
+        // Shards from workers 2 and 3 should move to workers 0 and 1
+        assert!(!plan.is_empty());
+        for shard_move in &plan {
+            assert!(shard_move.from_rank >= 2);
+            assert!(shard_move.to_rank < 2);
+        }
+    }
+
+    #[test]
+    fn test_calculate_redistribution_plan_no_change() {
+        let config = ElasticConfig::default();
+        let cluster = ElasticCluster::new(config);
+
+        // Same size - no moves
+        let plan = cluster.calculate_redistribution_plan(4, 4);
+        assert!(plan.is_empty());
+    }
+
+    #[tokio::test]
+    async fn test_execute_pending_operations_empty() {
+        let config = ElasticConfig::default();
+        let cluster = ElasticCluster::new(config);
+
+        // No pending operations
+        let report = cluster.execute_pending_operations_async().await.unwrap();
+
+        assert_eq!(report.operations_started, 0);
+        assert!(report.all_succeeded());
+    }
+
+    #[tokio::test]
+    async fn test_redistribute_state_same_size() {
+        let config = ElasticConfig::default();
+        let cluster = ElasticCluster::new(config);
+
+        // Same size - should be no-op
+        let result = cluster.redistribute_state_async(4, 4).await;
+        assert!(result.is_ok());
+    }
+}