mlmf 0.2.0

//! Distributed model loading and management for scalable inference.
//!
//! This module provides comprehensive support for distributed model deployment,
//! including model sharding, device placement, cross-node communication, and
//! load balancing for large-scale inference scenarios.

use candle_core::{Device, Tensor};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::net::SocketAddr;
use std::path::Path;
use std::sync::Arc;
use tokio::sync::RwLock;

use crate::cached_loader::CachedModelLoader;
use crate::LoadedModel;
use anyhow::{anyhow, Result};

/// Configuration for distributed model deployment
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DistributedConfig {
    /// List of nodes participating in the distributed deployment
    pub nodes: Vec<NodeConfig>,

    /// Strategy for distributing model shards across nodes
    pub sharding_strategy: ShardingStrategy,

    /// Device placement configuration
    pub device_placement: DevicePlacementConfig,

    /// Communication settings between nodes
    pub communication: CommunicationConfig,

    /// Load balancing configuration
    pub load_balancing: LoadBalancingConfig,

    /// Fault tolerance settings
    pub fault_tolerance: FaultToleranceConfig,
}

/// Configuration for a single node in the distributed system
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeConfig {
    /// Unique identifier for this node
    pub node_id: String,

    /// Network address for this node
    pub address: SocketAddr,

    /// Available devices on this node
    pub devices: Vec<DeviceInfo>,

    /// Node capabilities and resources
    pub capabilities: NodeCapabilities,

    /// Role of this node in the distributed system
    pub role: NodeRole,
}

/// Information about available devices on a node
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DeviceInfo {
    /// Device identifier (e.g., "gpu:0", "cpu")
    pub device_id: String,

    /// Device type
    pub device_type: DeviceType,

    /// Available memory in bytes
    pub memory_bytes: u64,

    /// Compute capability or performance score
    pub compute_score: f32,

    /// Current utilization (0.0 to 1.0)
    pub utilization: f32,
}

/// Type of compute device
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum DeviceType {
    /// CPU device
    Cpu,
    /// CUDA GPU device
    Cuda,
    /// Metal GPU device (Apple Silicon)
    Metal,
    /// Custom accelerator device
    Custom(String),
}

/// Node capabilities and resource information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeCapabilities {
    /// Total available memory in bytes
    pub total_memory: u64,

    /// Network bandwidth in bytes per second
    pub network_bandwidth: u64,

    /// Storage capacity in bytes
    pub storage_capacity: u64,

    /// Supported data types
    pub supported_dtypes: Vec<String>,

    /// Special capabilities (e.g., "quantization", "flash_attention")
    pub special_capabilities: Vec<String>,
}

/// Role of a node in the distributed system
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum NodeRole {
    /// Coordinator node that manages the cluster
    Coordinator,
    /// Worker node that performs inference
    Worker,
    /// Storage node that provides model data
    Storage,
    /// Hybrid node with multiple roles
    Hybrid(Vec<NodeRole>),
}

/// Strategy for distributing model shards
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ShardingStrategy {
    /// No sharding - full model on each node
    NoSharding,

    /// Shard by layers
    LayerSharding {
        /// Number of layers per shard
        layers_per_shard: usize,
    },

    /// Shard by model dimensions
    DimensionSharding {
        /// Dimension along which to shard
        dimension: ShardingDimension,
        /// Size of each shard
        shard_size: usize,
    },

    /// Pipeline parallelism - different stages on different nodes
    PipelineSharding {
        /// Number of pipeline stages
        num_stages: usize,
    },

    /// Tensor parallelism - split tensors across nodes
    TensorSharding {
        /// Parallelism degree
        degree: usize,
    },

    /// Custom sharding strategy
    Custom {
        /// Function to determine shard placement
        placement_fn: String, // Serialized function identifier
    },

    /// Modality-specific sharding for multi-modal models
    ModalitySpecific {
        /// Assignment of modalities to node groups
        modality_assignments: HashMap<crate::multimodal::Modality, Vec<String>>,
    },
}

/// Dimension for sharding operations
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ShardingDimension {
    /// Shard along rows
    Rows,
    /// Shard along columns  
    Columns,
    /// Shard along batch dimension
    Batch,
    /// Shard along sequence length
    Sequence,
}

/// Device placement configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DevicePlacementConfig {
    /// Strategy for placing model shards on devices
    pub strategy: PlacementStrategy,

    /// Constraints for device placement
    pub constraints: PlacementConstraints,

    /// Memory allocation settings
    pub memory_allocation: MemoryAllocationConfig,
}

/// Strategy for device placement
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum PlacementStrategy {
    /// Round-robin placement across available devices
    RoundRobin,

    /// Place based on available memory
    MemoryBased,

    /// Place based on compute capability
    ComputeBased,

    /// Load-balanced placement considering current utilization
    LoadBalanced,

    /// Locality-aware placement to minimize communication
    LocalityAware,

    /// Custom placement strategy
    Custom(String),
}

/// Constraints for device placement
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PlacementConstraints {
    /// Minimum memory required per device
    pub min_memory_per_device: u64,

    /// Maximum devices to use
    pub max_devices: Option<usize>,

    /// Preferred device types
    pub preferred_device_types: Vec<DeviceType>,

    /// Co-location requirements (shards that should be on same device)
    pub colocation_groups: Vec<Vec<String>>,

    /// Anti-affinity requirements (shards that should not be on same device)
    pub anti_affinity_groups: Vec<Vec<String>>,
}

/// Memory allocation configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryAllocationConfig {
    /// Memory allocation strategy
    pub strategy: MemoryAllocationStrategy,

    /// Reserve memory percentage for system operations
    pub system_reserve_percent: f32,

    /// Enable memory pooling across shards
    pub enable_memory_pooling: bool,

    /// Memory fragmentation threshold
    pub fragmentation_threshold: f32,
}

/// Strategy for memory allocation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum MemoryAllocationStrategy {
    /// Eager allocation - allocate all memory upfront
    Eager,

    /// Lazy allocation - allocate memory as needed
    Lazy,

    /// Pooled allocation - use memory pools
    Pooled,

    /// Streaming allocation - for very large models
    Streaming,
}

/// Communication configuration between nodes
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CommunicationConfig {
    /// Communication protocol
    pub protocol: CommunicationProtocol,

    /// Timeout settings
    pub timeouts: TimeoutConfig,

    /// Compression settings
    pub compression: CompressionConfig,

    /// Security settings
    pub security: SecurityConfig,
}

/// Communication protocol options
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum CommunicationProtocol {
    /// HTTP/REST protocol
    Http,

    /// gRPC protocol
    Grpc,

    /// TCP sockets
    Tcp,

    /// UDP sockets (for low-latency scenarios)
    Udp,

    /// NCCL for GPU communication
    Nccl,

    /// MPI for HPC environments
    Mpi,
}

/// Timeout configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TimeoutConfig {
    /// Connection timeout in milliseconds
    pub connection_timeout_ms: u64,

    /// Request timeout in milliseconds
    pub request_timeout_ms: u64,

    /// Heartbeat interval in milliseconds
    pub heartbeat_interval_ms: u64,

    /// Node failure detection timeout in milliseconds
    pub failure_detection_timeout_ms: u64,
}

/// Compression configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompressionConfig {
    /// Enable compression for model data transfer
    pub enable_compression: bool,

    /// Compression algorithm
    pub algorithm: CompressionAlgorithm,

    /// Compression level (algorithm-specific)
    pub level: u32,

    /// Minimum data size to trigger compression
    pub min_compress_size: usize,
}

/// Compression algorithms
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum CompressionAlgorithm {
    /// No compression
    None,
    /// LZ4 compression (fast)
    Lz4,
    /// Zstd compression (balanced)
    Zstd,
    /// Gzip compression (high ratio)
    Gzip,
    /// Custom compression
    Custom(String),
}

/// Security configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SecurityConfig {
    /// Enable TLS encryption
    pub enable_tls: bool,

    /// TLS certificate path
    pub cert_path: Option<String>,

    /// TLS private key path
    pub key_path: Option<String>,

    /// Authentication method
    pub auth_method: AuthMethod,

    /// API key for authentication
    pub api_key: Option<String>,
}

/// Authentication methods
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum AuthMethod {
    /// No authentication
    None,
    /// API key authentication
    ApiKey,
    /// Token-based authentication
    Token,
    /// mTLS authentication
    MutualTls,
}

/// Load balancing configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LoadBalancingConfig {
    /// Load balancing strategy
    pub strategy: LoadBalancingStrategy,

    /// Health check configuration
    pub health_check: HealthCheckConfig,

    /// Request routing configuration
    pub routing: RoutingConfig,
}

/// Load balancing strategies
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum LoadBalancingStrategy {
    /// Round-robin distribution
    RoundRobin,

    /// Least connections
    LeastConnections,

    /// Weighted round-robin
    WeightedRoundRobin,

    /// Based on response time
    ResponseTime,

    /// Based on resource utilization
    ResourceBased,

    /// Custom load balancing
    Custom(String),
}

/// Health check configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HealthCheckConfig {
    /// Enable health checks
    pub enabled: bool,

    /// Health check interval in milliseconds
    pub interval_ms: u64,

    /// Health check timeout in milliseconds
    pub timeout_ms: u64,

    /// Number of consecutive failures before marking unhealthy
    pub failure_threshold: u32,

    /// Number of consecutive successes before marking healthy
    pub success_threshold: u32,
}

/// Request routing configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RoutingConfig {
    /// Enable sticky sessions
    pub sticky_sessions: bool,

    /// Session affinity method
    pub affinity_method: AffinityMethod,

    /// Retry configuration
    pub retry: RetryConfig,
}

/// Session affinity methods
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum AffinityMethod {
    /// No affinity
    None,

    /// Client IP-based affinity
    ClientIp,

    /// Session token-based affinity
    SessionToken,

    /// Model-specific affinity
    ModelBased,
}

/// Retry configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RetryConfig {
    /// Maximum number of retries
    pub max_retries: u32,

    /// Initial retry delay in milliseconds
    pub initial_delay_ms: u64,

    /// Backoff multiplier
    pub backoff_multiplier: f32,

    /// Maximum retry delay in milliseconds
    pub max_delay_ms: u64,
}

/// Fault tolerance configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FaultToleranceConfig {
    /// Replication configuration
    pub replication: ReplicationConfig,

    /// Failure handling configuration
    pub failure_handling: FailureHandlingConfig,

    /// Recovery configuration
    pub recovery: RecoveryConfig,
}

/// Model replication configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ReplicationConfig {
    /// Enable model replication
    pub enabled: bool,

    /// Replication factor (number of copies)
    pub factor: u32,

    /// Replication strategy
    pub strategy: ReplicationStrategy,
}

/// Replication strategies
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ReplicationStrategy {
    /// Synchronous replication
    Synchronous,

    /// Asynchronous replication
    Asynchronous,

    /// Quorum-based replication
    Quorum { min_replicas: u32 },
}

/// Failure handling configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FailureHandlingConfig {
    /// Enable automatic failover
    pub auto_failover: bool,

    /// Failover timeout in milliseconds
    pub failover_timeout_ms: u64,

    /// Circuit breaker configuration
    pub circuit_breaker: CircuitBreakerConfig,
}

/// Circuit breaker configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CircuitBreakerConfig {
    /// Enable circuit breaker
    pub enabled: bool,

    /// Failure threshold to open circuit
    pub failure_threshold: u32,

    /// Recovery timeout in milliseconds
    pub recovery_timeout_ms: u64,

    /// Half-open request count
    pub half_open_requests: u32,
}

/// Recovery configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RecoveryConfig {
    /// Enable automatic recovery
    pub auto_recovery: bool,

    /// Recovery strategy
    pub strategy: RecoveryStrategy,

    /// Recovery timeout in milliseconds
    pub recovery_timeout_ms: u64,
}

/// Recovery strategies
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum RecoveryStrategy {
    /// Restart failed nodes
    RestartNodes,

    /// Redistribute shards
    RedistributeShards,

    /// Scale out with new nodes
    ScaleOut,

    /// Manual intervention required
    Manual,
}

/// Distributed model information
#[derive(Debug, Clone)]
pub struct DistributedModel {
    /// Model metadata
    pub metadata: DistributedModelMetadata,

    /// Shard information
    pub shards: Vec<ModelShard>,

    /// Active nodes participating in this model
    pub active_nodes: HashMap<String, NodeConfig>,

    /// Current deployment status
    pub status: DeploymentStatus,
}

/// Metadata for distributed models
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DistributedModelMetadata {
    /// Model identifier
    pub model_id: String,

    /// Total model size in bytes
    pub total_size: u64,

    /// Number of shards
    pub num_shards: usize,

    /// Sharding strategy used
    pub sharding_strategy: ShardingStrategy,

    /// Creation timestamp
    pub created_at: chrono::DateTime<chrono::Utc>,
}

/// Information about a model shard
#[derive(Debug, Clone)]
pub struct ModelShard {
    /// Shard identifier
    pub shard_id: String,

    /// Node where this shard is located
    pub node_id: String,

    /// Device where this shard is loaded
    pub device_id: String,

    /// Shard size in bytes
    pub size_bytes: u64,

    /// Model loaded status
    pub loaded: bool,

    /// Shard status
    pub status: ShardStatus,
}

/// Status of a model shard
#[derive(Debug, Clone, PartialEq)]
pub enum ShardStatus {
    /// Shard is not loaded
    NotLoaded,

    /// Shard is currently loading
    Loading,

    /// Shard is loaded and ready
    Ready,

    /// Shard failed to load
    Failed(String),

    /// Shard is being migrated
    Migrating,
}

/// Deployment status of distributed model
#[derive(Debug, Clone, PartialEq)]
pub enum DeploymentStatus {
    /// Deployment is being planned
    Planning,

    /// Model is being deployed
    Deploying,

    /// Model is fully deployed and ready
    Ready,

    /// Deployment failed
    Failed(String),

    /// Model is being updated
    Updating,

    /// Model is being shutdown
    ShuttingDown,
}

impl Default for DistributedConfig {
    fn default() -> Self {
        Self {
            nodes: Vec::new(),
            sharding_strategy: ShardingStrategy::NoSharding,
            device_placement: DevicePlacementConfig::default(),
            communication: CommunicationConfig::default(),
            load_balancing: LoadBalancingConfig::default(),
            fault_tolerance: FaultToleranceConfig::default(),
        }
    }
}

impl Default for DevicePlacementConfig {
    fn default() -> Self {
        Self {
            strategy: PlacementStrategy::LoadBalanced,
            constraints: PlacementConstraints::default(),
            memory_allocation: MemoryAllocationConfig::default(),
        }
    }
}

impl Default for PlacementConstraints {
    fn default() -> Self {
        Self {
            min_memory_per_device: 1024 * 1024 * 1024, // 1GB
            max_devices: None,
            preferred_device_types: vec![DeviceType::Cuda, DeviceType::Cpu],
            colocation_groups: Vec::new(),
            anti_affinity_groups: Vec::new(),
        }
    }
}

impl Default for MemoryAllocationConfig {
    fn default() -> Self {
        Self {
            strategy: MemoryAllocationStrategy::Lazy,
            system_reserve_percent: 0.1, // Reserve 10% for system
            enable_memory_pooling: true,
            fragmentation_threshold: 0.2,
        }
    }
}

impl Default for CommunicationConfig {
    fn default() -> Self {
        Self {
            protocol: CommunicationProtocol::Http,
            timeouts: TimeoutConfig::default(),
            compression: CompressionConfig::default(),
            security: SecurityConfig::default(),
        }
    }
}

impl Default for TimeoutConfig {
    fn default() -> Self {
        Self {
            connection_timeout_ms: 30_000,
            request_timeout_ms: 60_000,
            heartbeat_interval_ms: 10_000,
            failure_detection_timeout_ms: 30_000,
        }
    }
}

impl Default for CompressionConfig {
    fn default() -> Self {
        Self {
            enable_compression: true,
            algorithm: CompressionAlgorithm::Lz4,
            level: 1,
            min_compress_size: 1024, // 1KB
        }
    }
}

impl Default for SecurityConfig {
    fn default() -> Self {
        Self {
            enable_tls: false,
            cert_path: None,
            key_path: None,
            auth_method: AuthMethod::None,
            api_key: None,
        }
    }
}

impl Default for LoadBalancingConfig {
    fn default() -> Self {
        Self {
            strategy: LoadBalancingStrategy::RoundRobin,
            health_check: HealthCheckConfig::default(),
            routing: RoutingConfig::default(),
        }
    }
}

impl Default for HealthCheckConfig {
    fn default() -> Self {
        Self {
            enabled: true,
            interval_ms: 30_000,
            timeout_ms: 5_000,
            failure_threshold: 3,
            success_threshold: 2,
        }
    }
}

impl Default for RoutingConfig {
    fn default() -> Self {
        Self {
            sticky_sessions: false,
            affinity_method: AffinityMethod::None,
            retry: RetryConfig::default(),
        }
    }
}

impl Default for RetryConfig {
    fn default() -> Self {
        Self {
            max_retries: 3,
            initial_delay_ms: 1000,
            backoff_multiplier: 2.0,
            max_delay_ms: 30_000,
        }
    }
}

impl Default for FaultToleranceConfig {
    fn default() -> Self {
        Self {
            replication: ReplicationConfig::default(),
            failure_handling: FailureHandlingConfig::default(),
            recovery: RecoveryConfig::default(),
        }
    }
}

impl Default for ReplicationConfig {
    fn default() -> Self {
        Self {
            enabled: false,
            factor: 1,
            strategy: ReplicationStrategy::Asynchronous,
        }
    }
}

impl Default for FailureHandlingConfig {
    fn default() -> Self {
        Self {
            auto_failover: true,
            failover_timeout_ms: 60_000,
            circuit_breaker: CircuitBreakerConfig::default(),
        }
    }
}

impl Default for CircuitBreakerConfig {
    fn default() -> Self {
        Self {
            enabled: true,
            failure_threshold: 5,
            recovery_timeout_ms: 60_000,
            half_open_requests: 3,
        }
    }
}

impl Default for RecoveryConfig {
    fn default() -> Self {
        Self {
            auto_recovery: true,
            strategy: RecoveryStrategy::RedistributeShards,
            recovery_timeout_ms: 300_000, // 5 minutes
        }
    }
}

/// Builder for distributed configuration
pub struct DistributedConfigBuilder {
    config: DistributedConfig,
}

impl DistributedConfigBuilder {
    /// Create a new distributed configuration builder
    pub fn new() -> Self {
        Self {
            config: DistributedConfig::default(),
        }
    }

    /// Add a node to the cluster
    pub fn add_node(mut self, node: NodeConfig) -> Self {
        self.config.nodes.push(node);
        self
    }

    /// Set the sharding strategy
    pub fn sharding_strategy(mut self, strategy: ShardingStrategy) -> Self {
        self.config.sharding_strategy = strategy;
        self
    }

    /// Set device placement configuration
    pub fn device_placement(mut self, placement: DevicePlacementConfig) -> Self {
        self.config.device_placement = placement;
        self
    }

    /// Set communication configuration
    pub fn communication(mut self, communication: CommunicationConfig) -> Self {
        self.config.communication = communication;
        self
    }

    /// Set load balancing configuration
    pub fn load_balancing(mut self, load_balancing: LoadBalancingConfig) -> Self {
        self.config.load_balancing = load_balancing;
        self
    }

    /// Set fault tolerance configuration
    pub fn fault_tolerance(mut self, fault_tolerance: FaultToleranceConfig) -> Self {
        self.config.fault_tolerance = fault_tolerance;
        self
    }

    /// Build the configuration
    pub fn build(self) -> DistributedConfig {
        self.config
    }
}

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