rustytorch/crates/training/rtx-distributed/src/coordination.rs

//! Distributed training coordination mechanisms

use crate::WorldInfo;
use crate::error::{DistributedError, Result};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Instant;
use tokio::sync::{Barrier, RwLock};

/// Distributed coordinator for training orchestration
#[derive(Debug)]
pub struct DistributedCoordinator {
    world: WorldInfo,
    is_initialized: bool,
    metrics: Arc<RwLock<HashMap<String, MetricAggregator>>>,
    lr_schedule: Arc<RwLock<Option<LRSchedule>>>,
    throughput_history: Arc<RwLock<Vec<f64>>>,
    memory_history: Arc<RwLock<Vec<f64>>>,
}

impl DistributedCoordinator {
    pub async fn new(world: WorldInfo) -> Result<Self> {
        Ok(Self {
            world,
            is_initialized: true,
            metrics: Arc::new(RwLock::new(HashMap::new())),
            lr_schedule: Arc::new(RwLock::new(None)),
            throughput_history: Arc::new(RwLock::new(Vec::new())),
            memory_history: Arc::new(RwLock::new(Vec::new())),
        })
    }

    pub fn is_initialized(&self) -> bool {
        self.is_initialized
    }

    pub fn world_size(&self) -> usize {
        self.world.world_size
    }

    pub fn rank(&self) -> usize {
        self.world.rank
    }

    pub fn is_master(&self) -> bool {
        self.world.rank == 0
    }

    pub async fn record_throughput(&self, samples_per_sec: f64) {
        let mut history = self.throughput_history.write().await;
        history.push(samples_per_sec);
        if history.len() > 100 {
            history.remove(0);
        }
    }

    pub async fn record_memory_usage(&self, usage: f64) {
        let mut history = self.memory_history.write().await;
        history.push(usage);
        if history.len() > 100 {
            history.remove(0);
        }
    }

    pub async fn optimize_batch_size(&self, current_batch: usize) -> usize {
        let memory = self.memory_history.read().await;
        let avg_memory = if memory.is_empty() {
            0.5
        } else {
            memory.iter().sum::<f64>() / memory.len() as f64
        };

        if avg_memory > 0.9 {
            (current_batch as f64 * 0.8) as usize
        } else if avg_memory < 0.6 {
            (current_batch as f64 * 1.2) as usize
        } else {
            current_batch
        }
    }

    pub async fn set_lr_schedule(&self, schedule: LRSchedule) {
        let mut lr = self.lr_schedule.write().await;
        *lr = Some(schedule);
    }

    pub async fn get_learning_rate(&self, step: usize) -> f64 {
        let lr_schedule = self.lr_schedule.read().await;

        match lr_schedule.as_ref() {
            Some(LRSchedule::CosineAnnealing {
                initial_lr,
                min_lr,
                warmup_steps,
                total_steps,
            }) => {
                if step < *warmup_steps {
                    initial_lr * (step as f64 / *warmup_steps as f64)
                } else if step >= *total_steps {
                    *min_lr
                } else {
                    let progress =
                        (step - warmup_steps) as f64 / (total_steps - warmup_steps) as f64;
                    min_lr
                        + (initial_lr - min_lr) * (1.0 + (progress * std::f64::consts::PI).cos())
                            / 2.0
                }
            }
            None => 0.001,
        }
    }

    pub async fn submit_metrics(&self, _rank: usize, name: &str, value: f64) {
        let mut metrics = self.metrics.write().await;
        let aggregator = metrics
            .entry(name.to_string())
            .or_insert_with(MetricAggregator::new);
        aggregator.add(value);
    }

    pub async fn aggregate_metrics(&self) -> HashMap<String, AggregatedMetric> {
        let metrics = self.metrics.read().await;
        let mut result = HashMap::new();

        for (name, aggregator) in metrics.iter() {
            result.insert(name.clone(), aggregator.aggregate());
        }

        result
    }

    pub async fn get_metric_history(&self, name: &str, _max_entries: usize) -> Vec<f64> {
        let metrics = self.metrics.read().await;
        metrics
            .get(name)
            .map(|agg| agg.values.clone())
            .unwrap_or_default()
    }
}

/// Training scheduler for work distribution
#[derive(Debug)]
pub struct TrainingScheduler {
    world_size: usize,
    total_samples: usize,
}

impl TrainingScheduler {
    pub fn new(world_size: usize, total_samples: usize) -> Self {
        Self {
            world_size,
            total_samples,
        }
    }

    pub async fn create_schedule(&self, batch_size: usize) -> Result<Schedule> {
        let samples_per_node = self.total_samples / self.world_size;
        let batches_per_epoch = self.total_samples / batch_size;

        Ok(Schedule {
            num_epochs: 1,
            batches_per_epoch,
            samples_per_node,
            batch_size,
            world_size: self.world_size,
        })
    }
}

/// Training schedule
#[derive(Debug)]
pub struct Schedule {
    num_epochs: usize,
    batches_per_epoch: usize,
    samples_per_node: usize,
    batch_size: usize,
    world_size: usize,
}

impl Schedule {
    pub fn num_epochs(&self) -> usize {
        self.num_epochs
    }

    pub fn batches_per_epoch(&self) -> usize {
        self.batches_per_epoch
    }

    pub fn samples_per_node(&self) -> usize {
        self.samples_per_node
    }

    pub fn get_work_for_rank(&self, rank: usize) -> WorkItem {
        let start = rank * self.samples_per_node;
        let end = start + self.samples_per_node;

        WorkItem {
            start_idx: start,
            end_idx: end,
        }
    }
}

/// Synchronization barrier for distributed training
#[derive(Debug, Clone)]
pub struct SynchronizationBarrier {
    barrier: Arc<Barrier>,
    world_size: usize,
    reached: Arc<RwLock<Vec<bool>>>,
}

impl SynchronizationBarrier {
    pub fn new(world_size: usize) -> Self {
        Self {
            barrier: Arc::new(Barrier::new(world_size)),
            world_size,
            reached: Arc::new(RwLock::new(vec![false; world_size])),
        }
    }

    pub async fn wait(&self, rank: usize) -> Result<()> {
        let mut reached = self.reached.write().await;
        reached[rank] = true;
        drop(reached);

        self.barrier.wait().await;
        Ok(())
    }

    pub async fn all_reached(&self) -> bool {
        self.reached.read().await.iter().all(|&r| r)
    }
}

/// Hyperparameter server for distributed settings
#[derive(Debug)]
pub struct HyperparameterServer {
    params: Arc<RwLock<HashMap<String, HyperParam>>>,
    updates: Arc<RwLock<Vec<ParamUpdate>>>,
}

impl Default for HyperparameterServer {
    fn default() -> Self {
        Self::new()
    }
}

impl HyperparameterServer {
    pub fn new() -> Self {
        Self {
            params: Arc::new(RwLock::new(HashMap::new())),
            updates: Arc::new(RwLock::new(Vec::new())),
        }
    }

    pub async fn register<T: Into<HyperParam>>(&mut self, name: &str, value: T) {
        let mut params = self.params.write().await;
        params.insert(name.to_string(), value.into());
    }

    pub async fn get_f64(&self, name: &str) -> Result<f64> {
        let params = self.params.read().await;
        params
            .get(name)
            .and_then(HyperParam::as_f64)
            .ok_or_else(|| DistributedError::runtime(name.to_string()))
    }

    pub async fn get_i32(&self, name: &str) -> Result<i32> {
        let params = self.params.read().await;
        params
            .get(name)
            .and_then(HyperParam::as_i32)
            .ok_or_else(|| DistributedError::runtime(name.to_string()))
    }

    pub async fn update<T: Into<HyperParam>>(&mut self, name: &str, value: T) -> Result<()> {
        let mut params = self.params.write().await;
        params.insert(name.to_string(), value.into());

        let mut updates = self.updates.write().await;
        updates.push(ParamUpdate {
            param_name: name.to_string(),
            timestamp: Instant::now(),
        });

        Ok(())
    }

    pub async fn get_updates_since(&self, _since: usize) -> Vec<ParamUpdate> {
        self.updates.read().await.clone()
    }
}

/// Global state manager
#[derive(Debug)]
pub struct GlobalStateManager {
    epoch: Arc<RwLock<usize>>,
    global_step: Arc<RwLock<usize>>,
    best_loss: Arc<RwLock<f64>>,
}

impl GlobalStateManager {
    pub fn new(_world_size: usize) -> Self {
        Self {
            epoch: Arc::new(RwLock::new(0)),
            global_step: Arc::new(RwLock::new(0)),
            best_loss: Arc::new(RwLock::new(f64::MAX)),
        }
    }

    pub async fn set_epoch(&self, epoch: usize) {
        let mut e = self.epoch.write().await;
        *e = epoch;
    }

    pub async fn get_epoch(&self) -> usize {
        *self.epoch.read().await
    }

    pub async fn increment_epoch(&self) {
        let mut e = self.epoch.write().await;
        *e += 1;
    }

    pub async fn set_global_step(&self, step: usize) {
        let mut s = self.global_step.write().await;
        *s = step;
    }

    pub async fn get_global_step(&self) -> usize {
        *self.global_step.read().await
    }

    pub async fn add_steps(&self, steps: usize) {
        let mut s = self.global_step.write().await;
        *s += steps;
    }

    pub async fn set_best_loss(&self, loss: f64) {
        let mut l = self.best_loss.write().await;
        *l = loss;
    }

    pub async fn get_best_loss(&self) -> f64 {
        *self.best_loss.read().await
    }

    pub async fn snapshot(&self) -> StateSnapshot {
        StateSnapshot {
            epoch: self.get_epoch().await,
            global_step: self.get_global_step().await,
            best_loss: self.get_best_loss().await,
        }
    }
}

/// Consensus protocol for distributed decisions
#[derive(Debug)]
pub struct ConsensusProtocol {
    world_size: usize,
    rank: usize,
    proposals: Arc<RwLock<HashMap<u64, Proposal>>>,
    votes: Arc<RwLock<HashMap<u64, Vec<bool>>>>,
}

impl ConsensusProtocol {
    pub fn new(world_size: usize, rank: usize) -> Self {
        Self {
            world_size,
            rank,
            proposals: Arc::new(RwLock::new(HashMap::new())),
            votes: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    pub async fn propose(&self, _key: &str, value: &str) -> Result<Proposal> {
        let proposal = Proposal {
            id: rand::random(),
            value: value.to_string(),
        };

        let mut proposals = self.proposals.write().await;
        proposals.insert(proposal.id, proposal.clone());

        let mut votes = self.votes.write().await;
        votes.insert(proposal.id, vec![false; self.world_size]);

        Ok(proposal)
    }

    pub async fn vote(&self, _rank: usize, proposal_id: u64, vote: bool) {
        let mut votes = self.votes.write().await;
        if let Some(proposal_votes) = votes.get_mut(&proposal_id)
            && _rank < proposal_votes.len()
        {
            proposal_votes[_rank] = vote;
        }
    }

    pub async fn has_consensus(&self, proposal_id: u64) -> bool {
        let votes = self.votes.read().await;
        votes
            .get(&proposal_id)
            .is_some_and(|v| v.iter().filter(|&&vote| vote).count() > self.world_size / 2)
    }

    pub async fn get_decision(&self, proposal_id: u64) -> Result<Decision> {
        let proposals = self.proposals.read().await;
        let proposal = proposals
            .get(&proposal_id)
            .ok_or_else(|| DistributedError::runtime("Proposal not found".to_string()))?;

        let accepted = self.has_consensus(proposal_id).await;

        Ok(Decision {
            value: proposal.value.clone(),
            accepted,
        })
    }
}

/// Workload distributor
#[derive(Debug)]
pub struct WorkloadDistributor {
    world_size: usize,
}

impl WorkloadDistributor {
    pub fn new(world_size: usize) -> Self {
        Self { world_size }
    }

    pub async fn distribute_dataset(&self, dataset_size: usize) -> Vec<DataDistribution> {
        let base_samples = dataset_size / self.world_size;
        let remainder = dataset_size % self.world_size;

        let mut distributions = Vec::new();
        let mut start_idx = 0;

        for rank in 0..self.world_size {
            let num_samples = if rank < remainder {
                base_samples + 1
            } else {
                base_samples
            };

            distributions.push(DataDistribution {
                rank,
                num_samples,
                start_idx,
                end_idx: start_idx + num_samples,
            });

            start_idx += num_samples;
        }

        distributions
    }
}

/// Gradient aggregator
#[derive(Debug)]
pub struct GradientAggregator {
    world_size: usize,
    gradients: Arc<RwLock<HashMap<usize, Vec<f32>>>>,
}

impl GradientAggregator {
    pub fn new(world_size: usize) -> Self {
        Self {
            world_size,
            gradients: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    pub async fn submit_gradients(&self, rank: usize, grads: Vec<f32>) {
        let mut gradients = self.gradients.write().await;
        gradients.insert(rank, grads);
    }

    pub async fn aggregate_mean(&self) -> Result<Vec<f32>> {
        let gradients = self.gradients.read().await;

        if gradients.is_empty() {
            return Err(DistributedError::runtime(
                "No gradients to aggregate".to_string(),
            ));
        }

        let grad_len = gradients.values().next().unwrap().len();
        let mut aggregated = vec![0.0; grad_len];

        for grad in gradients.values() {
            for (i, &g) in grad.iter().enumerate() {
                aggregated[i] += g;
            }
        }

        for val in &mut aggregated {
            *val /= gradients.len() as f32;
        }

        Ok(aggregated)
    }

    pub async fn aggregate_with_clipping(&self, max_norm: f32) -> Result<Vec<f32>> {
        let mut aggregated = self.aggregate_mean().await?;

        let norm: f32 = aggregated.iter().map(|g| g * g).sum::<f32>().sqrt();

        if norm > max_norm {
            let scale = max_norm / norm;
            for val in &mut aggregated {
                *val *= scale;
            }
        }

        Ok(aggregated)
    }
}

/// Model synchronizer
#[derive(Debug)]
pub struct ModelSynchronizer {
    world_size: usize,
    versions: Arc<RwLock<HashMap<usize, ModelVersion>>>,
    params: Arc<RwLock<HashMap<usize, Vec<f32>>>>,
}

impl ModelSynchronizer {
    pub fn new(world_size: usize) -> Self {
        Self {
            world_size,
            versions: Arc::new(RwLock::new(HashMap::new())),
            params: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    pub async fn register_model_version(
        &self,
        _name: &str,
        params: &[f32],
    ) -> Result<ModelVersion> {
        let checksum = format!("{:032}", params.len()); // Simplified checksum

        let version = ModelVersion {
            version: 1,
            checksum,
        };

        let mut versions = self.versions.write().await;
        for rank in 0..self.world_size {
            versions.insert(rank, version.clone());
        }

        Ok(version)
    }

    pub async fn check_sync_status(&self) -> SyncStatus {
        let versions = self.versions.read().await;

        if versions.is_empty() {
            return SyncStatus {
                all_synced: true,
                out_of_sync_nodes: vec![],
            };
        }

        let first_version = versions.values().next().unwrap();
        let out_of_sync: Vec<usize> = versions
            .iter()
            .filter(|(_, v)| v.checksum != first_version.checksum)
            .map(|(rank, _)| *rank)
            .collect();

        SyncStatus {
            all_synced: out_of_sync.is_empty(),
            out_of_sync_nodes: out_of_sync,
        }
    }

    pub async fn update_model_params(&self, rank: usize, params: &[f32]) -> Result<()> {
        let mut model_params = self.params.write().await;
        model_params.insert(rank, params.to_vec());

        let mut versions = self.versions.write().await;
        let checksum = format!("{:032}", params.len());
        versions.insert(
            rank,
            ModelVersion {
                version: 2,
                checksum,
            },
        );

        Ok(())
    }

    pub async fn broadcast_model_update(&self) -> Result<()> {
        // Synchronize all nodes to latest version
        let versions = self.versions.read().await;
        if let Some(latest) = versions.values().max_by_key(|v| v.version) {
            let mut versions = self.versions.write().await;
            for rank in 0..self.world_size {
                versions.insert(rank, latest.clone());
            }
        }
        Ok(())
    }
}

// Helper types
#[derive(Debug, Clone)]
pub struct WorkItem {
    pub start_idx: usize,
    pub end_idx: usize,
}

#[derive(Debug, Clone)]
pub enum HyperParam {
    Float(f64),
    Int(i32),
}

impl HyperParam {
    fn as_f64(&self) -> Option<f64> {
        match self {
            Self::Float(v) => Some(*v),
            _ => None,
        }
    }

    fn as_i32(&self) -> Option<i32> {
        match self {
            Self::Int(v) => Some(*v),
            _ => None,
        }
    }
}

impl From<f64> for HyperParam {
    fn from(v: f64) -> Self {
        Self::Float(v)
    }
}

impl From<i32> for HyperParam {
    fn from(v: i32) -> Self {
        Self::Int(v)
    }
}

#[derive(Debug, Clone)]
pub struct ParamUpdate {
    pub param_name: String,
    pub timestamp: Instant,
}

#[derive(Debug)]
pub struct StateSnapshot {
    pub epoch: usize,
    pub global_step: usize,
    pub best_loss: f64,
}

#[derive(Debug, Clone)]
pub struct Proposal {
    pub id: u64,
    pub value: String,
}

#[derive(Debug)]
pub struct Decision {
    pub value: String,
    pub accepted: bool,
}

#[derive(Debug)]
pub struct DataDistribution {
    pub rank: usize,
    pub num_samples: usize,
    pub start_idx: usize,
    pub end_idx: usize,
}

#[derive(Debug, Clone)]
pub struct ModelVersion {
    pub version: usize,
    pub checksum: String,
}

#[derive(Debug)]
pub struct SyncStatus {
    pub all_synced: bool,
    pub out_of_sync_nodes: Vec<usize>,
}

#[derive(Debug)]
struct MetricAggregator {
    values: Vec<f64>,
}

impl MetricAggregator {
    fn new() -> Self {
        Self { values: Vec::new() }
    }

    fn add(&mut self, value: f64) {
        self.values.push(value);
    }

    fn aggregate(&self) -> AggregatedMetric {
        let mean = if self.values.is_empty() {
            0.0
        } else {
            self.values.iter().sum::<f64>() / self.values.len() as f64
        };

        AggregatedMetric { mean }
    }
}

#[derive(Debug)]
pub struct AggregatedMetric {
    pub mean: f64,
}

#[derive(Debug, Clone)]
pub enum LRSchedule {
    CosineAnnealing {
        initial_lr: f64,
        min_lr: f64,
        warmup_steps: usize,
        total_steps: usize,
    },
}