torsh-distributed 0.1.2

//! Horovod compatibility layer for ToRSh distributed training
//!
//! This module provides compatibility with Horovod's distributed training API,
//! allowing users to migrate from Horovod to ToRSh more easily.
//!
//! Horovod is a distributed training framework that provides:
//! - AllReduce-based distributed training
//! - Gradient compression
//! - Timeline profiling
//! - Elastic training
//! - Spark integration

use crate::{TorshDistributedError, TorshResult};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::Path;

/// Horovod configuration compatible with ToRSh
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HorovodConfig {
    /// Gradient compression configuration
    pub gradient_compression: Option<HorovodCompressionConfig>,
    /// Timeline profiling configuration
    pub timeline: Option<HorovodTimelineConfig>,
    /// Elastic training configuration
    pub elastic: Option<HorovodElasticConfig>,
    /// Optimizer fusion configuration
    pub optimizer_fusion: Option<HorovodOptimizerFusionConfig>,
    /// Backward passes aggregation
    pub backward_passes_per_step: Option<u32>,
    /// Gradient accumulation steps
    pub gradient_accumulation_steps: Option<u32>,
    /// Custom communication operations
    pub custom_ops: Option<HashMap<String, String>>,
}

/// Horovod gradient compression configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HorovodCompressionConfig {
    /// Compression algorithm type
    pub compression_type: HorovodCompressionType,
    /// Compression parameters
    pub compression_params: HashMap<String, f32>,
    /// Memory optimization enabled
    pub memory_optimization: Option<bool>,
    /// Compression period
    pub compression_period: Option<u32>,
}

/// Horovod compression algorithm types
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum HorovodCompressionType {
    /// No compression
    None,
    /// Top-K sparsification
    TopK,
    /// Quantization
    Quantization,
    /// Random-K sparsification
    RandomK,
    /// Threshold-based sparsification
    Threshold,
    /// Bernoulli sampling
    Bernoulli,
    /// Gaussian sampling
    Gaussian,
}

/// Horovod timeline profiling configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HorovodTimelineConfig {
    /// Timeline filename
    pub filename: String,
    /// Mark cycles in timeline
    pub mark_cycles: Option<bool>,
    /// Timeline sampling rate
    pub sampling_rate: Option<f32>,
    /// Include memory usage
    pub include_memory: Option<bool>,
}

/// Horovod elastic training configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HorovodElasticConfig {
    /// Minimum number of workers
    pub min_workers: u32,
    /// Maximum number of workers
    pub max_workers: u32,
    /// Initial number of workers
    pub initial_workers: Option<u32>,
    /// Discovery server address
    pub discovery_server: Option<String>,
    /// Health check interval (seconds)
    pub health_check_interval: Option<u64>,
    /// State broadcast timeout (seconds)
    pub state_broadcast_timeout: Option<u64>,
}

/// Horovod optimizer fusion configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HorovodOptimizerFusionConfig {
    /// Enable optimizer fusion
    pub enabled: bool,
    /// Fusion threshold (bytes)
    pub fusion_threshold: Option<u64>,
    /// Cycle time (ms)
    pub cycle_time_ms: Option<u64>,
    /// Fusion buffer size
    pub fusion_buffer_size: Option<u64>,
}

/// Horovod integration statistics
#[derive(Debug, Clone, Default)]
pub struct HorovodStats {
    /// Number of AllReduce operations
    pub allreduce_ops: u64,
    /// Total AllReduce time (seconds)
    pub allreduce_time_sec: f64,
    /// Total compressed bytes
    pub compressed_bytes: u64,
    /// Total uncompressed bytes
    pub uncompressed_bytes: u64,
    /// Number of elastic events
    pub elastic_events: u64,
    /// Average compression ratio
    pub compression_ratio: f64,
    /// Timeline events recorded
    pub timeline_events: u64,
}

/// Horovod compatibility integration
pub struct HorovodIntegration {
    /// Configuration
    config: HorovodConfig,
    /// Statistics
    stats: HorovodStats,
    /// Initialization status
    initialized: bool,
    /// Process rank
    rank: u32,
    /// World size
    world_size: u32,
    /// Local rank
    local_rank: u32,
    /// Local size
    local_size: u32,
}

impl HorovodIntegration {
    /// Create a new Horovod integration
    pub fn new(config: HorovodConfig) -> Self {
        Self {
            config,
            stats: HorovodStats::default(),
            initialized: false,
            rank: 0,
            world_size: 1,
            local_rank: 0,
            local_size: 1,
        }
    }

    /// Load configuration from JSON file
    pub fn from_file<P: AsRef<Path>>(path: P) -> TorshResult<Self> {
        let content = std::fs::read_to_string(path).map_err(|e| {
            TorshDistributedError::configuration_error(format!(
                "Failed to read Horovod config file: {}",
                e
            ))
        })?;

        let config: HorovodConfig = serde_json::from_str(&content).map_err(|e| {
            TorshDistributedError::configuration_error(format!(
                "Failed to parse Horovod config: {}",
                e
            ))
        })?;

        Ok(Self::new(config))
    }

    /// Initialize Horovod integration
    pub fn initialize(
        &mut self,
        rank: u32,
        world_size: u32,
        local_rank: u32,
        local_size: u32,
    ) -> TorshResult<()> {
        if self.initialized {
            return Err(TorshDistributedError::configuration_error(
                "Horovod integration already initialized",
            ));
        }

        self.rank = rank;
        self.world_size = world_size;
        self.local_rank = local_rank;
        self.local_size = local_size;

        self.validate_config()?;
        self.setup_compression()?;
        self.setup_timeline()?;
        self.setup_elastic()?;
        self.setup_optimizer_fusion()?;

        self.initialized = true;
        tracing::info!(
            "Horovod integration initialized - rank: {}, world_size: {}, local_rank: {}",
            self.rank,
            self.world_size,
            self.local_rank
        );

        Ok(())
    }

    /// Validate Horovod configuration
    fn validate_config(&self) -> TorshResult<()> {
        // Validate elastic configuration
        if let Some(ref elastic) = self.config.elastic {
            if elastic.min_workers == 0 {
                return Err(TorshDistributedError::configuration_error(
                    "Elastic training min_workers must be greater than 0",
                ));
            }
            if elastic.max_workers < elastic.min_workers {
                return Err(TorshDistributedError::configuration_error(
                    "Elastic training max_workers must be >= min_workers",
                ));
            }
            if let Some(initial) = elastic.initial_workers {
                if initial < elastic.min_workers || initial > elastic.max_workers {
                    return Err(TorshDistributedError::configuration_error(
                        "Elastic training initial_workers must be between min_workers and max_workers"
                    ));
                }
            }
        }

        // Validate compression configuration
        if let Some(ref compression) = self.config.gradient_compression {
            if matches!(compression.compression_type, HorovodCompressionType::TopK)
                && !compression.compression_params.contains_key("k")
            {
                return Err(TorshDistributedError::configuration_error(
                    "TopK compression requires 'k' parameter",
                ));
            }
            if matches!(
                compression.compression_type,
                HorovodCompressionType::Threshold
            ) && !compression.compression_params.contains_key("threshold")
            {
                return Err(TorshDistributedError::configuration_error(
                    "Threshold compression requires 'threshold' parameter",
                ));
            }
        }

        Ok(())
    }

    /// Setup gradient compression
    fn setup_compression(&self) -> TorshResult<()> {
        if let Some(ref compression) = self.config.gradient_compression {
            tracing::info!(
                "Setting up Horovod gradient compression: {:?}",
                compression.compression_type
            );

            match compression.compression_type {
                HorovodCompressionType::None => {
                    tracing::debug!("No gradient compression configured");
                }
                HorovodCompressionType::TopK => {
                    let k = compression.compression_params.get("k").unwrap_or(&0.01);
                    tracing::debug!("TopK compression configured with k={}", k);
                }
                HorovodCompressionType::Quantization => {
                    let bits = compression.compression_params.get("bits").unwrap_or(&8.0);
                    tracing::debug!("Quantization compression configured with bits={}", bits);
                }
                HorovodCompressionType::RandomK => {
                    let k = compression.compression_params.get("k").unwrap_or(&0.01);
                    tracing::debug!("RandomK compression configured with k={}", k);
                }
                HorovodCompressionType::Threshold => {
                    let threshold = compression
                        .compression_params
                        .get("threshold")
                        .unwrap_or(&0.01);
                    tracing::debug!(
                        "Threshold compression configured with threshold={}",
                        threshold
                    );
                }
                HorovodCompressionType::Bernoulli => {
                    let probability = compression
                        .compression_params
                        .get("probability")
                        .unwrap_or(&0.01);
                    tracing::debug!(
                        "Bernoulli compression configured with probability={}",
                        probability
                    );
                }
                HorovodCompressionType::Gaussian => {
                    let sigma = compression.compression_params.get("sigma").unwrap_or(&1.0);
                    tracing::debug!("Gaussian compression configured with sigma={}", sigma);
                }
            }
        }
        Ok(())
    }

    /// Setup timeline profiling
    fn setup_timeline(&self) -> TorshResult<()> {
        if let Some(ref timeline) = self.config.timeline {
            tracing::info!(
                "Setting up Horovod timeline profiling: {}",
                timeline.filename
            );

            if timeline.mark_cycles.unwrap_or(false) {
                tracing::debug!("Timeline cycle marking enabled");
            }

            if let Some(rate) = timeline.sampling_rate {
                tracing::debug!("Timeline sampling rate: {}", rate);
            }

            if timeline.include_memory.unwrap_or(false) {
                tracing::debug!("Timeline memory tracking enabled");
            }
        }
        Ok(())
    }

    /// Setup elastic training
    fn setup_elastic(&self) -> TorshResult<()> {
        if let Some(ref elastic) = self.config.elastic {
            tracing::info!(
                "Setting up Horovod elastic training: min={}, max={}",
                elastic.min_workers,
                elastic.max_workers
            );

            if let Some(ref discovery) = elastic.discovery_server {
                tracing::debug!("Elastic discovery server: {}", discovery);
            }

            let health_interval = elastic.health_check_interval.unwrap_or(30);
            tracing::debug!("Elastic health check interval: {}s", health_interval);

            let broadcast_timeout = elastic.state_broadcast_timeout.unwrap_or(300);
            tracing::debug!("Elastic state broadcast timeout: {}s", broadcast_timeout);
        }
        Ok(())
    }

    /// Setup optimizer fusion
    fn setup_optimizer_fusion(&self) -> TorshResult<()> {
        if let Some(ref fusion) = self.config.optimizer_fusion {
            if fusion.enabled {
                tracing::info!("Setting up Horovod optimizer fusion");

                let threshold = fusion.fusion_threshold.unwrap_or(64 * 1024 * 1024); // 64MB
                tracing::debug!("Optimizer fusion threshold: {} bytes", threshold);

                let cycle_time = fusion.cycle_time_ms.unwrap_or(1);
                tracing::debug!("Optimizer fusion cycle time: {} ms", cycle_time);

                let buffer_size = fusion.fusion_buffer_size.unwrap_or(128 * 1024 * 1024); // 128MB
                tracing::debug!("Optimizer fusion buffer size: {} bytes", buffer_size);
            }
        }
        Ok(())
    }

    /// Convert Horovod config to ToRSh DDP config
    pub fn to_ddp_config(&self) -> TorshResult<crate::ddp::BucketConfig> {
        use crate::ddp::BucketConfig;

        let _gradient_accumulation_steps = self.config.gradient_accumulation_steps.unwrap_or(1);
        let _backward_passes_per_step = self.config.backward_passes_per_step.unwrap_or(1);

        let bucket_config = BucketConfig {
            max_bucket_size_mb: 25.0, // Default Horovod bucket size
            enabled: true,
            min_bucket_size_mb: 1.0,
        };

        Ok(bucket_config)
    }

    /// Convert Horovod config to ToRSh gradient compression config
    pub fn to_compression_config(
        &self,
    ) -> TorshResult<Option<crate::gradient_compression::CompressionConfig>> {
        use crate::gradient_compression::{CompressionConfig, CompressionMethod};

        if let Some(ref compression) = self.config.gradient_compression {
            let method = match compression.compression_type {
                HorovodCompressionType::None => return Ok(None),
                HorovodCompressionType::TopK => {
                    let k = compression.compression_params.get("k").unwrap_or(&0.01);
                    CompressionMethod::TopK { k: *k }
                }
                HorovodCompressionType::Quantization => {
                    let bits = *compression.compression_params.get("bits").unwrap_or(&8.0) as u32;
                    CompressionMethod::Quantization {
                        bits: bits.try_into().unwrap_or(8),
                    }
                }
                HorovodCompressionType::RandomK => {
                    let k = compression.compression_params.get("k").unwrap_or(&0.01);
                    CompressionMethod::RandomK { k: *k }
                }
                HorovodCompressionType::Threshold => {
                    let threshold = compression
                        .compression_params
                        .get("threshold")
                        .unwrap_or(&0.01);
                    CompressionMethod::Threshold {
                        threshold: *threshold,
                    }
                }
                HorovodCompressionType::Bernoulli => {
                    // Map Bernoulli to RandomK with probability parameter
                    let k = compression
                        .compression_params
                        .get("probability")
                        .unwrap_or(&0.01);
                    CompressionMethod::RandomK { k: *k }
                }
                HorovodCompressionType::Gaussian => {
                    // Map Gaussian to NaturalCompression
                    let sigma = compression.compression_params.get("sigma").unwrap_or(&1.0);
                    CompressionMethod::NaturalCompression {
                        compression_factor: *sigma,
                    }
                }
            };

            let config = CompressionConfig {
                method,
                compression_ratio: *compression
                    .compression_params
                    .get("compression_ratio")
                    .unwrap_or(&0.1),
                error_feedback: true,
                error_feedback_momentum: 0.9,
                memory_efficient: compression.memory_optimization.unwrap_or(false),
                warmup_steps: 100,
            };

            Ok(Some(config))
        } else {
            Ok(None)
        }
    }

    /// Convert Horovod config to ToRSh elastic config
    pub fn to_elastic_config(&self) -> TorshResult<Option<crate::fault_tolerance::ElasticConfig>> {
        use crate::fault_tolerance::ElasticConfig;
        use std::time::Duration;

        if let Some(ref elastic) = self.config.elastic {
            let config = ElasticConfig {
                min_workers: elastic.min_workers as usize,
                max_workers: elastic.max_workers as usize,
                scaling_timeout: Duration::from_secs(
                    elastic.state_broadcast_timeout.unwrap_or(300),
                ),
                scaling_check_interval: Duration::from_secs(
                    elastic.health_check_interval.unwrap_or(30),
                ),
                enable_elastic_scheduling: true,
                rendezvous_backend: "etcd".to_string(),
                rendezvous_endpoint: elastic
                    .discovery_server
                    .clone()
                    .unwrap_or_else(|| "localhost:2379".to_string()),
            };

            Ok(Some(config))
        } else {
            Ok(None)
        }
    }

    /// Get the current configuration
    pub fn config(&self) -> &HorovodConfig {
        &self.config
    }

    /// Get current statistics
    pub fn stats(&self) -> &HorovodStats {
        &self.stats
    }

    /// Check if Horovod integration is initialized
    pub fn is_initialized(&self) -> bool {
        self.initialized
    }

    /// Get current rank (equivalent to hvd.rank())
    pub fn rank(&self) -> u32 {
        self.rank
    }

    /// Get world size (equivalent to hvd.size())
    pub fn size(&self) -> u32 {
        self.world_size
    }

    /// Get local rank (equivalent to hvd.local_rank())
    pub fn local_rank(&self) -> u32 {
        self.local_rank
    }

    /// Get local size (equivalent to hvd.local_size())
    pub fn local_size(&self) -> u32 {
        self.local_size
    }

    /// Simulate Horovod allreduce operation
    pub fn allreduce(&mut self, tensor_name: &str, tensor_size: usize) -> TorshResult<()> {
        if !self.initialized {
            return Err(TorshDistributedError::BackendNotInitialized);
        }

        let start_time = std::time::Instant::now();

        // Simulate allreduce operation
        tracing::debug!("Horovod allreduce: {} ({} bytes)", tensor_name, tensor_size);

        // Update statistics
        self.stats.allreduce_ops += 1;
        self.stats.allreduce_time_sec += start_time.elapsed().as_secs_f64();

        // Handle compression if enabled
        if let Some(ref compression) = self.config.gradient_compression {
            if !matches!(compression.compression_type, HorovodCompressionType::None) {
                let compression_ratio =
                    self.estimate_compression_ratio(compression.compression_type, tensor_size);
                self.stats.uncompressed_bytes += tensor_size as u64;
                self.stats.compressed_bytes += (tensor_size as f64 * compression_ratio) as u64;
                self.stats.compression_ratio =
                    self.stats.compressed_bytes as f64 / self.stats.uncompressed_bytes as f64;
            }
        }

        Ok(())
    }

    /// Estimate compression ratio for different algorithms
    fn estimate_compression_ratio(
        &self,
        compression_type: HorovodCompressionType,
        _tensor_size: usize,
    ) -> f64 {
        match compression_type {
            HorovodCompressionType::None => 1.0,
            HorovodCompressionType::TopK => {
                let k = self
                    .config
                    .gradient_compression
                    .as_ref()
                    .and_then(|c| c.compression_params.get("k"))
                    .unwrap_or(&0.01);
                *k as f64
            }
            HorovodCompressionType::Quantization => {
                let bits = self
                    .config
                    .gradient_compression
                    .as_ref()
                    .and_then(|c| c.compression_params.get("bits"))
                    .unwrap_or(&8.0);
                (*bits as f64) / 32.0 // Assuming 32-bit floats
            }
            HorovodCompressionType::RandomK => {
                let k = self
                    .config
                    .gradient_compression
                    .as_ref()
                    .and_then(|c| c.compression_params.get("k"))
                    .unwrap_or(&0.01);
                *k as f64
            }
            HorovodCompressionType::Threshold => 0.1, // Rough estimate
            HorovodCompressionType::Bernoulli => {
                let probability = self
                    .config
                    .gradient_compression
                    .as_ref()
                    .and_then(|c| c.compression_params.get("probability"))
                    .unwrap_or(&0.01);
                *probability as f64
            }
            HorovodCompressionType::Gaussian => 0.5, // Rough estimate
        }
    }

    /// Record timeline event
    pub fn record_timeline_event(&mut self, event_name: &str, event_type: &str) {
        if self.config.timeline.is_some() {
            tracing::debug!("Timeline event: {} ({})", event_name, event_type);
            self.stats.timeline_events += 1;
        }
    }

    /// Handle elastic event
    pub fn handle_elastic_event(
        &mut self,
        event_type: &str,
        new_world_size: u32,
    ) -> TorshResult<()> {
        if self.config.elastic.is_some() {
            tracing::info!(
                "Elastic event: {} - new world size: {}",
                event_type,
                new_world_size
            );
            self.world_size = new_world_size;
            self.stats.elastic_events += 1;
        }
        Ok(())
    }

    /// Create a default Horovod configuration
    pub fn default_config() -> HorovodConfig {
        HorovodConfig {
            gradient_compression: None,
            timeline: None,
            elastic: None,
            optimizer_fusion: Some(HorovodOptimizerFusionConfig {
                enabled: true,
                fusion_threshold: Some(64 * 1024 * 1024), // 64MB
                cycle_time_ms: Some(1),
                fusion_buffer_size: Some(128 * 1024 * 1024), // 128MB
            }),
            backward_passes_per_step: Some(1),
            gradient_accumulation_steps: Some(1),
            custom_ops: None,
        }
    }

    /// Create a configuration with TopK compression
    pub fn config_with_topk_compression(k: f32) -> HorovodConfig {
        let mut compression_params = HashMap::new();
        compression_params.insert("k".to_string(), k);

        HorovodConfig {
            gradient_compression: Some(HorovodCompressionConfig {
                compression_type: HorovodCompressionType::TopK,
                compression_params,
                memory_optimization: Some(true),
                compression_period: Some(1),
            }),
            timeline: None,
            elastic: None,
            optimizer_fusion: Some(HorovodOptimizerFusionConfig {
                enabled: true,
                fusion_threshold: Some(64 * 1024 * 1024),
                cycle_time_ms: Some(1),
                fusion_buffer_size: Some(128 * 1024 * 1024),
            }),
            backward_passes_per_step: Some(1),
            gradient_accumulation_steps: Some(1),
            custom_ops: None,
        }
    }

    /// Create a configuration with elastic training
    pub fn config_with_elastic(min_workers: u32, max_workers: u32) -> HorovodConfig {
        HorovodConfig {
            gradient_compression: None,
            timeline: None,
            elastic: Some(HorovodElasticConfig {
                min_workers,
                max_workers,
                initial_workers: Some(min_workers),
                discovery_server: None,
                health_check_interval: Some(30),
                state_broadcast_timeout: Some(300),
            }),
            optimizer_fusion: Some(HorovodOptimizerFusionConfig {
                enabled: true,
                fusion_threshold: Some(64 * 1024 * 1024),
                cycle_time_ms: Some(1),
                fusion_buffer_size: Some(128 * 1024 * 1024),
            }),
            backward_passes_per_step: Some(1),
            gradient_accumulation_steps: Some(1),
            custom_ops: None,
        }
    }
}

impl Default for HorovodConfig {
    fn default() -> Self {
        HorovodIntegration::default_config()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_horovod_config_validation() {
        let config = HorovodIntegration::default_config();
        let mut integration = HorovodIntegration::new(config);

        // Should succeed with valid parameters
        assert!(integration.initialize(0, 4, 0, 2).is_ok());
        assert!(integration.is_initialized());
        assert_eq!(integration.rank(), 0);
        assert_eq!(integration.size(), 4);
        assert_eq!(integration.local_rank(), 0);
    }

    #[test]
    fn test_horovod_topk_compression() {
        let config = HorovodIntegration::config_with_topk_compression(0.1);
        let mut integration = HorovodIntegration::new(config);

        assert!(integration.initialize(0, 4, 0, 2).is_ok());

        // Test compression config conversion
        let compression_config = integration.to_compression_config().unwrap();
        assert!(compression_config.is_some());

        if let Some(config) = compression_config {
            assert!(
                matches!(config.method, crate::gradient_compression::CompressionMethod::TopK { k } if k == 0.1)
            );
        }
    }

    #[test]
    fn test_horovod_elastic_config() {
        let config = HorovodIntegration::config_with_elastic(2, 8);
        let mut integration = HorovodIntegration::new(config);

        assert!(integration.initialize(0, 4, 0, 2).is_ok());

        // Test elastic config conversion
        let elastic_config = integration.to_elastic_config().unwrap();
        assert!(elastic_config.is_some());

        if let Some(config) = elastic_config {
            assert_eq!(config.min_workers, 2);
            assert_eq!(config.max_workers, 8);
        }
    }

    #[test]
    fn test_horovod_allreduce_simulation() {
        let config = HorovodIntegration::default_config();
        let mut integration = HorovodIntegration::new(config);

        assert!(integration.initialize(0, 4, 0, 2).is_ok());

        // Simulate allreduce operations
        assert!(integration.allreduce("layer1.weight", 1024).is_ok());
        assert!(integration.allreduce("layer1.bias", 256).is_ok());

        let stats = integration.stats();
        assert_eq!(stats.allreduce_ops, 2);
        assert!(stats.allreduce_time_sec >= 0.0);
    }

    #[test]
    fn test_horovod_invalid_elastic_config() {
        let config = HorovodIntegration::config_with_elastic(8, 2); // max < min
        let mut integration = HorovodIntegration::new(config);

        // Should fail validation
        assert!(integration.initialize(0, 4, 0, 2).is_err());
    }

    #[test]
    fn test_horovod_config_serialization() {
        let config = HorovodIntegration::config_with_topk_compression(0.05);

        // Test JSON serialization
        let json = serde_json::to_string(&config).unwrap();
        assert!(json.contains("TopK"));
        assert!(json.contains("0.05"));

        // Test deserialization
        let deserialized: HorovodConfig = serde_json::from_str(&json).unwrap();
        assert!(deserialized.gradient_compression.is_some());
    }
}