use std::collections::{HashMap, HashSet};
use std::hash::Hash;
use std::net::SocketAddr;
use std::sync::{Arc, Mutex, RwLock};
use std::thread;
use std::time::{Duration, Instant};
pub type MessageHandlerFn =
Box<dyn Fn(&DistributedMessage) -> Result<MessageResponse, DistributedError> + Send + Sync>;
#[derive(Debug, Clone)]
pub struct ResourceReservation {
pub id: String,
pub node_id: String,
pub start_time: Instant,
pub duration: Duration,
pub resources: ResourceAllocation,
}
#[derive(Debug, Clone, PartialEq)]
pub enum ConsensusState {
Follower,
Candidate,
Leader,
}
#[derive(Debug, Clone)]
pub struct ClusterConfig {
pub heartbeat_interval: Duration,
pub node_timeout: Duration,
pub job_timeout: Duration,
pub max_retries: u32,
pub load_threshold: f64,
pub replication_factor: u32,
}
#[derive(Debug, Clone)]
pub enum MessagePriority {
Low,
Normal,
High,
Critical,
}
#[derive(Debug, Clone)]
pub struct MessageResponse {
pub message_id: String,
pub success: bool,
pub data: Vec<u8>,
pub error: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub enum JobPriority {
Low,
Normal,
High,
Critical,
}
#[derive(Debug, Clone, PartialEq)]
pub enum JobStatus {
Pending,
Running,
Completed,
Failed,
Cancelled,
Timeout,
}
#[derive(Debug, Clone)]
pub struct SchedulingDecision {
pub job_id: String,
pub node_id: String,
pub estimated_start_time: Instant,
pub resource_allocation: ResourceAllocation,
}
#[derive(Debug, thiserror::Error)]
pub enum DistributedError {
#[error("Node not found")]
NodeNotFound,
#[error("Job not found")]
JobNotFound,
#[error("Insufficient resources")]
InsufficientResources,
#[error("Node unreachable")]
NodeUnreachable,
#[error("Job timeout")]
JobTimeout,
#[error("Scheduling error: {0}")]
SchedulingError(String),
#[error("Communication error: {0}")]
CommunicationError(String),
}
#[derive(Debug, Clone)]
pub struct VoteRequest {
pub term: u64,
pub candidate_id: String,
pub last_log_index: usize,
pub last_log_term: u64,
}
#[derive(Debug, Clone)]
pub struct VoteResponse {
pub term: u64,
pub vote_granted: bool,
}
#[allow(dead_code)]
pub struct MessagePassingSystem {
node_id: String,
message_handlers: HashMap<String, MessageHandler>,
pending_messages: Arc<Mutex<Vec<DistributedMessage>>>,
pub(crate) message_queue: Arc<Mutex<Vec<DistributedMessage>>>,
pub(crate) routing_table: Arc<RwLock<HashMap<String, SocketAddr>>>,
}
impl MessagePassingSystem {
pub fn new(node_id: String) -> Self {
Self {
node_id,
message_handlers: HashMap::new(),
pending_messages: Arc::new(Mutex::new(Vec::new())),
message_queue: Arc::new(Mutex::new(Vec::new())),
routing_table: Arc::new(RwLock::new(HashMap::new())),
}
}
pub fn send_message(&self, message: DistributedMessage) -> Result<(), DistributedError> {
let routing_table = self.routing_table.read().expect("operation should succeed");
if let Some(_address) = routing_table.get(&message.destination) {
self.message_queue
.lock()
.expect("operation should succeed")
.push(message);
Ok(())
} else {
Err(DistributedError::NodeUnreachable)
}
}
pub fn broadcast_message(
&self,
message_type: MessageType,
data: Vec<u8>,
) -> Result<(), DistributedError> {
let routing_table = self.routing_table.read().expect("operation should succeed");
for node_id in routing_table.keys() {
if node_id != &self.node_id {
let message = DistributedMessage {
id: format!("{}_{}", self.node_id, Instant::now().elapsed().as_millis()),
source: self.node_id.clone(),
destination: node_id.clone(),
message_type: message_type.clone(),
data: data.clone(),
timestamp: Instant::now(),
priority: MessagePriority::Normal,
};
self.send_message(message)?;
}
}
Ok(())
}
pub fn register_handler(&mut self, message_type: String, handler: MessageHandler) {
self.message_handlers.insert(message_type, handler);
}
pub fn process_messages(&self) -> Result<Vec<MessageResponse>, DistributedError> {
let mut responses = Vec::new();
let mut queue = self.message_queue.lock().expect("operation should succeed");
for message in queue.drain(..) {
if let Some(handler) = self
.message_handlers
.get(&format!("{}", message.message_type))
{
let response = handler.handle(&message)?;
responses.push(response);
}
}
Ok(responses)
}
}
#[derive(Debug, Clone)]
pub struct PeerInfo {
pub id: String,
pub next_index: usize,
pub match_index: usize,
pub last_response: Instant,
}
#[allow(dead_code)]
pub struct AdvancedJobScheduler {
scheduling_policies: Vec<SchedulingPolicy>,
pub(crate) resource_reservations: HashMap<String, ResourceReservation>,
job_dependencies: HashMap<String, Vec<String>>,
priority_queues: HashMap<JobPriority, Vec<DistributedJob>>,
backfill_enabled: bool,
}
impl AdvancedJobScheduler {
pub fn new() -> Self {
Self {
scheduling_policies: vec![
SchedulingPolicy::FIFO,
SchedulingPolicy::ShortestJobFirst,
SchedulingPolicy::GangScheduling,
],
resource_reservations: HashMap::new(),
job_dependencies: HashMap::new(),
priority_queues: HashMap::new(),
backfill_enabled: true,
}
}
pub fn gang_schedule(
&mut self,
jobs: &[DistributedJob],
nodes: &HashMap<String, ClusterNode>,
) -> Result<Vec<SchedulingDecision>, DistributedError> {
let mut decisions = Vec::new();
let job_groups = self.group_related_jobs(jobs);
for group in job_groups {
if let Some(node_assignment) = self.find_gang_assignment(&group, nodes) {
for (job, node_id) in group.iter().zip(node_assignment.iter()) {
decisions.push(SchedulingDecision {
job_id: job.id.clone(),
node_id: node_id.clone(),
estimated_start_time: Instant::now(),
resource_allocation: self
.calculate_resource_allocation(job, node_id, nodes),
});
}
}
}
Ok(decisions)
}
pub fn backfill_schedule(
&mut self,
waiting_jobs: &[DistributedJob],
nodes: &HashMap<String, ClusterNode>,
) -> Result<Vec<SchedulingDecision>, DistributedError> {
let mut decisions = Vec::new();
if !self.backfill_enabled {
return Ok(decisions);
}
for job in waiting_jobs {
for (node_id, node) in nodes {
if self.can_backfill_job(job, node) {
decisions.push(SchedulingDecision {
job_id: job.id.clone(),
node_id: node_id.clone(),
estimated_start_time: Instant::now(),
resource_allocation: self
.calculate_resource_allocation(job, node_id, nodes),
});
break;
}
}
}
Ok(decisions)
}
pub fn reserve_resources(
&mut self,
reservation: ResourceReservation,
) -> Result<(), DistributedError> {
self.resource_reservations
.insert(reservation.id.clone(), reservation);
Ok(())
}
fn group_related_jobs(&self, jobs: &[DistributedJob]) -> Vec<Vec<DistributedJob>> {
vec![jobs.to_vec()]
}
fn find_gang_assignment(
&self,
job_group: &[DistributedJob],
nodes: &HashMap<String, ClusterNode>,
) -> Option<Vec<String>> {
if job_group.len() <= nodes.len() {
Some(nodes.keys().take(job_group.len()).cloned().collect())
} else {
None
}
}
fn can_backfill_job(&self, _job: &DistributedJob, _node: &ClusterNode) -> bool {
true
}
fn calculate_resource_allocation(
&self,
job: &DistributedJob,
_node_id: &str,
_nodes: &HashMap<String, ClusterNode>,
) -> ResourceAllocation {
ResourceAllocation {
cpu_cores: job.requirements.min_cpu_cores,
memory_gb: job.requirements.min_memory_gb,
gpu_count: job.requirements.min_gpu_count,
storage_gb: job.requirements.min_storage_gb,
network_bandwidth: 100,
}
}
}
#[derive(Debug, Clone)]
pub struct PartitionMove {
pub partition: usize,
pub from_node: String,
pub to_node: String,
}
#[derive(Debug, Clone)]
pub struct CheckpointStats {
pub total_checkpoints: usize,
pub total_size_bytes: u64,
pub compressed_checkpoints: usize,
pub compression_ratio: f64,
}
pub struct JobScheduler {
scheduling_strategy: SchedulingStrategy,
}
impl JobScheduler {
pub fn new() -> Self {
Self {
scheduling_strategy: SchedulingStrategy::LeastLoaded,
}
}
pub fn find_suitable_node(
&self,
job: &DistributedJob,
nodes: &HashMap<String, ClusterNode>,
) -> Option<String> {
let suitable_nodes: Vec<_> = nodes
.values()
.filter(|node| self.can_run_job(node, job))
.collect();
if suitable_nodes.is_empty() {
return None;
}
match self.scheduling_strategy {
SchedulingStrategy::LeastLoaded => suitable_nodes
.iter()
.min_by(|a, b| {
let load_a = a.load_metrics.cpu_usage + a.load_metrics.memory_usage;
let load_b = b.load_metrics.cpu_usage + b.load_metrics.memory_usage;
load_a
.partial_cmp(&load_b)
.unwrap_or(std::cmp::Ordering::Equal)
})
.map(|node| node.id.clone()),
SchedulingStrategy::RoundRobin => suitable_nodes.first().map(|node| node.id.clone()),
SchedulingStrategy::HighestCapacity => suitable_nodes
.iter()
.max_by_key(|node| node.capabilities.cpu_cores * node.capabilities.memory_gb)
.map(|node| node.id.clone()),
}
}
fn can_run_job(&self, node: &ClusterNode, job: &DistributedJob) -> bool {
node.status == NodeStatus::Available
&& node.capabilities.cpu_cores >= job.requirements.min_cpu_cores
&& node.capabilities.memory_gb >= job.requirements.min_memory_gb
&& node.capabilities.gpu_count >= job.requirements.min_gpu_count
&& node.capabilities.storage_gb >= job.requirements.min_storage_gb
}
}
pub struct CheckpointManager {
pub(crate) checkpoint_storage: HashMap<String, Checkpoint>,
#[allow(dead_code)]
checkpoint_interval: Duration,
compression_enabled: bool,
}
impl CheckpointManager {
pub fn new(checkpoint_interval: Duration) -> Self {
Self {
checkpoint_storage: HashMap::new(),
checkpoint_interval,
compression_enabled: true,
}
}
pub fn create_checkpoint(
&mut self,
job_id: &str,
state: JobState,
) -> Result<String, DistributedError> {
let checkpoint_id = format!("{job_id}_{}", Instant::now().elapsed().as_millis());
let checkpoint = Checkpoint {
id: checkpoint_id.clone(),
job_id: job_id.to_string(),
state,
created_at: Instant::now(),
size_bytes: 1024,
compressed: self.compression_enabled,
};
self.checkpoint_storage
.insert(checkpoint_id.clone(), checkpoint);
Ok(checkpoint_id)
}
pub fn restore_checkpoint(&self, checkpoint_id: &str) -> Result<JobState, DistributedError> {
self.checkpoint_storage
.get(checkpoint_id)
.map(|checkpoint| checkpoint.state.clone())
.ok_or(DistributedError::JobNotFound)
}
pub fn cleanup_old_checkpoints(&mut self, retention_period: Duration) {
let cutoff = Instant::now() - retention_period;
self.checkpoint_storage
.retain(|_, checkpoint| checkpoint.created_at > cutoff);
}
pub fn get_checkpoint_stats(&self) -> CheckpointStats {
let total_checkpoints = self.checkpoint_storage.len();
let total_size: u64 = self.checkpoint_storage.values().map(|c| c.size_bytes).sum();
let compressed_checkpoints = self
.checkpoint_storage
.values()
.filter(|c| c.compressed)
.count();
CheckpointStats {
total_checkpoints,
total_size_bytes: total_size,
compressed_checkpoints,
compression_ratio: if total_checkpoints > 0 {
compressed_checkpoints as f64 / total_checkpoints as f64
} else {
0.0
},
}
}
}
#[derive(Debug, Clone)]
pub struct ClusterNode {
pub id: String,
pub address: SocketAddr,
pub capabilities: NodeCapabilities,
pub status: NodeStatus,
pub last_heartbeat: Instant,
pub load_metrics: LoadMetrics,
pub job_history: Vec<JobExecution>,
}
#[derive(Debug, Clone)]
pub struct LogEntry {
pub term: u64,
pub index: usize,
pub command: String,
pub data: Vec<u8>,
}
#[derive(Debug, Clone)]
pub struct DistributedMessage {
pub id: String,
pub source: String,
pub destination: String,
pub message_type: MessageType,
pub data: Vec<u8>,
pub timestamp: Instant,
pub priority: MessagePriority,
}
#[derive(Debug, Clone)]
pub struct JobState {
pub progress: f64,
pub intermediate_results: HashMap<String, Vec<u8>>,
pub runtime_state: Vec<u8>,
}
pub struct DataPartitioner {
partitioning_strategy: PartitioningStrategy,
partition_count: usize,
node_assignments: HashMap<usize, String>,
replication_factor: usize,
}
impl DataPartitioner {
pub fn new(
strategy: PartitioningStrategy,
partition_count: usize,
replication_factor: usize,
) -> Self {
Self {
partitioning_strategy: strategy,
partition_count,
node_assignments: HashMap::new(),
replication_factor,
}
}
pub fn get_partition(&self, key: &str) -> usize {
match self.partitioning_strategy {
PartitioningStrategy::Hash => {
use std::collections::hash_map::DefaultHasher;
use std::hash::Hasher;
let mut hasher = DefaultHasher::new();
key.hash(&mut hasher);
(hasher.finish() as usize) % self.partition_count
}
PartitioningStrategy::Range => key.len() % self.partition_count,
PartitioningStrategy::Random => key.len() % self.partition_count,
}
}
pub fn get_partition_nodes(&self, partition: usize) -> Vec<String> {
let mut nodes = Vec::new();
if let Some(primary_node) = self.node_assignments.get(&partition) {
nodes.push(primary_node.clone());
for i in 1..self.replication_factor {
let replica_partition = (partition + i) % self.partition_count;
if let Some(replica_node) = self.node_assignments.get(&replica_partition) {
if !nodes.contains(replica_node) {
nodes.push(replica_node.clone());
}
}
}
}
nodes
}
pub fn assign_partition(&mut self, partition: usize, node_id: String) {
self.node_assignments.insert(partition, node_id);
}
pub fn rebalance_partitions(&mut self, available_nodes: &[String]) -> PartitioningResult {
let mut assignments_changed = 0;
let mut partitions_moved = Vec::new();
for partition in 0..self.partition_count {
let optimal_node = &available_nodes[partition % available_nodes.len()];
if let Some(current_node) = self.node_assignments.get(&partition) {
if current_node != optimal_node {
partitions_moved.push(PartitionMove {
partition,
from_node: current_node.clone(),
to_node: optimal_node.clone(),
});
self.node_assignments
.insert(partition, optimal_node.clone());
assignments_changed += 1;
}
} else {
self.node_assignments
.insert(partition, optimal_node.clone());
assignments_changed += 1;
}
}
PartitioningResult {
assignments_changed,
partitions_moved,
rebalance_time: Instant::now(),
}
}
}
#[derive(Debug, Clone)]
pub struct ClusterStats {
pub total_nodes: usize,
pub available_nodes: usize,
pub busy_nodes: usize,
pub overloaded_nodes: usize,
pub unreachable_nodes: usize,
pub total_jobs: usize,
pub running_jobs: usize,
pub completed_jobs: usize,
pub failed_jobs: usize,
pub queued_jobs: usize,
pub total_cpu_cores: u32,
pub total_memory_gb: u32,
pub total_gpu_count: u32,
pub avg_cpu_usage: f64,
pub avg_memory_usage: f64,
}
pub struct MessageHandler {
pub handler_fn: MessageHandlerFn,
}
impl MessageHandler {
pub fn new<F>(f: F) -> Self
where
F: Fn(&DistributedMessage) -> Result<MessageResponse, DistributedError>
+ Send
+ Sync
+ 'static,
{
Self {
handler_fn: Box::new(f),
}
}
pub fn handle(
&self,
message: &DistributedMessage,
) -> Result<MessageResponse, DistributedError> {
(self.handler_fn)(message)
}
}
pub struct FaultDetector {
failure_history: HashMap<String, Vec<Instant>>,
}
impl FaultDetector {
pub fn new() -> Self {
Self {
failure_history: HashMap::new(),
}
}
pub fn handle_failure(&mut self, node_id: &str) -> Result<(), DistributedError> {
let failures = self.failure_history.entry(node_id.to_string()).or_default();
failures.push(Instant::now());
let cutoff = Instant::now() - Duration::from_secs(3600);
failures.retain(|&failure_time| failure_time > cutoff);
if failures.len() > 3 {
println!("Node {node_id} has too many failures, marking as problematic");
}
Ok(())
}
pub fn is_problematic(&self, node_id: &str) -> bool {
self.failure_history
.get(node_id)
.map(|failures| failures.len() > 3)
.unwrap_or(false)
}
}
#[derive(Debug, Clone)]
pub enum MessageType {
JobSubmission,
JobResult,
Heartbeat,
ResourceUpdate,
ConsensusRequest,
DataPartition,
Custom(String),
}
pub struct LoadBalancer {
rebalance_threshold: f64,
}
impl LoadBalancer {
pub fn new() -> Self {
Self {
rebalance_threshold: 0.8,
}
}
pub fn rebalance(&self, nodes: &HashMap<String, ClusterNode>) -> Result<(), DistributedError> {
let overloaded_nodes: Vec<_> = nodes
.values()
.filter(|node| {
let total_load = node.load_metrics.cpu_usage + node.load_metrics.memory_usage;
total_load > self.rebalance_threshold * 2.0
})
.collect();
let underloaded_nodes: Vec<_> = nodes
.values()
.filter(|node| {
let total_load = node.load_metrics.cpu_usage + node.load_metrics.memory_usage;
total_load < self.rebalance_threshold
})
.collect();
println!(
"Rebalancing: {} overloaded nodes, {} underloaded nodes",
overloaded_nodes.len(),
underloaded_nodes.len()
);
Ok(())
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum NodeStatus {
Available,
Busy,
Overloaded,
Unreachable,
Maintenance,
}
#[derive(Debug, Clone, PartialEq)]
pub enum JobType {
Training,
Inference,
DataProcessing,
ModelEvaluation,
Hyperparameter,
Custom(String),
}
#[derive(Debug, Clone)]
pub struct Checkpoint {
pub id: String,
pub job_id: String,
pub state: JobState,
pub created_at: Instant,
pub size_bytes: u64,
pub compressed: bool,
}
#[derive(Debug, Clone)]
pub struct LoadMetrics {
pub cpu_usage: f64,
pub memory_usage: f64,
pub gpu_usage: f64,
pub network_io: f64,
pub disk_io: f64,
pub active_jobs: u32,
pub queue_size: u32,
}
#[derive(Debug, Clone)]
pub struct DistributedJob {
pub id: String,
pub name: String,
pub job_type: JobType,
pub priority: JobPriority,
pub requirements: ResourceRequirements,
pub created_at: Instant,
pub timeout: Duration,
pub retry_count: u32,
pub dependencies: Vec<String>,
pub metadata: HashMap<String, String>,
}
#[derive(Debug, Clone)]
pub enum SchedulingStrategy {
LeastLoaded,
RoundRobin,
HighestCapacity,
}
#[derive(Debug, Clone)]
pub struct ResourceUsage {
pub cpu_time: Duration,
pub memory_peak: u64,
pub gpu_time: Duration,
pub network_bytes: u64,
pub disk_bytes: u64,
}
#[derive(Debug, Clone)]
pub struct NodeCapabilities {
pub cpu_cores: u32,
pub memory_gb: u32,
pub gpu_count: u32,
pub storage_gb: u32,
pub network_bandwidth_mbps: u32,
pub supported_tasks: HashSet<String>,
}
#[allow(dead_code)]
pub struct ConsensusManager {
node_id: String,
pub(crate) state: ConsensusState,
term: u64,
voted_for: Option<String>,
pub(crate) log: Vec<LogEntry>,
commit_index: usize,
last_applied: usize,
peers: HashMap<String, PeerInfo>,
}
impl ConsensusManager {
pub fn new(node_id: String, peers: Vec<String>) -> Self {
let mut peer_map = HashMap::new();
for peer in peers {
peer_map.insert(
peer.clone(),
PeerInfo {
id: peer,
next_index: 0,
match_index: 0,
last_response: Instant::now(),
},
);
}
Self {
node_id,
state: ConsensusState::Follower,
term: 0,
voted_for: None,
log: Vec::new(),
commit_index: 0,
last_applied: 0,
peers: peer_map,
}
}
pub fn start_election(&mut self) -> Result<(), DistributedError> {
self.state = ConsensusState::Candidate;
self.term += 1;
self.voted_for = Some(self.node_id.clone());
println!(
"Node {} starting election for term {}",
self.node_id, self.term
);
Ok(())
}
pub fn handle_vote_request(&mut self, request: VoteRequest) -> VoteResponse {
let grant_vote = if request.term > self.term {
self.term = request.term;
self.voted_for = None;
self.state = ConsensusState::Follower;
true
} else if request.term == self.term
&& (self.voted_for.is_none() || self.voted_for.as_ref() == Some(&request.candidate_id))
{
self.voted_for = Some(request.candidate_id.clone());
true
} else {
false
};
VoteResponse {
term: self.term,
vote_granted: grant_vote,
}
}
pub fn append_entry(&mut self, entry: LogEntry) -> Result<(), DistributedError> {
if self.state != ConsensusState::Leader {
return Err(DistributedError::SchedulingError("Not leader".to_string()));
}
self.log.push(entry);
Ok(())
}
pub fn get_state(&self) -> (ConsensusState, u64) {
(self.state.clone(), self.term)
}
}
#[derive(Debug, Clone)]
pub struct PartitioningResult {
pub assignments_changed: usize,
pub partitions_moved: Vec<PartitionMove>,
pub rebalance_time: Instant,
}
#[derive(Debug, Clone)]
pub enum PartitioningStrategy {
Hash,
Range,
Random,
}
#[derive(Debug, Clone)]
pub enum SchedulingPolicy {
FIFO,
ShortestJobFirst,
GangScheduling,
Backfill,
PriorityBased,
}
pub struct DistributedCluster {
nodes: Arc<RwLock<HashMap<String, ClusterNode>>>,
jobs: Arc<RwLock<HashMap<String, DistributedJob>>>,
executions: Arc<RwLock<HashMap<String, JobExecution>>>,
pub(crate) job_queue: Arc<Mutex<Vec<DistributedJob>>>,
scheduler: Arc<Mutex<JobScheduler>>,
load_balancer: Arc<Mutex<LoadBalancer>>,
fault_detector: Arc<Mutex<FaultDetector>>,
config: ClusterConfig,
}
impl DistributedCluster {
pub fn new(config: ClusterConfig) -> Self {
Self {
nodes: Arc::new(RwLock::new(HashMap::new())),
jobs: Arc::new(RwLock::new(HashMap::new())),
executions: Arc::new(RwLock::new(HashMap::new())),
job_queue: Arc::new(Mutex::new(Vec::new())),
scheduler: Arc::new(Mutex::new(JobScheduler::new())),
load_balancer: Arc::new(Mutex::new(LoadBalancer::new())),
fault_detector: Arc::new(Mutex::new(FaultDetector::new())),
config,
}
}
pub fn register_node(&self, node: ClusterNode) -> Result<(), DistributedError> {
let mut nodes = self.nodes.write().expect("operation should succeed");
nodes.insert(node.id.clone(), node);
Ok(())
}
pub fn remove_node(&self, node_id: &str) -> Result<(), DistributedError> {
let mut nodes = self.nodes.write().expect("operation should succeed");
nodes
.remove(node_id)
.ok_or(DistributedError::NodeNotFound)?;
Ok(())
}
pub fn get_nodes(&self) -> Vec<ClusterNode> {
self.nodes
.read()
.expect("operation should succeed")
.values()
.cloned()
.collect()
}
pub fn get_available_nodes(&self) -> Vec<ClusterNode> {
self.nodes
.read()
.expect("operation should succeed")
.values()
.filter(|node| node.status == NodeStatus::Available)
.cloned()
.collect()
}
pub fn submit_job(&self, job: DistributedJob) -> Result<String, DistributedError> {
let job_id = job.id.clone();
self.jobs
.write()
.expect("operation should succeed")
.insert(job_id.clone(), job.clone());
self.job_queue
.lock()
.expect("operation should succeed")
.push(job);
self.schedule_jobs()?;
Ok(job_id)
}
pub fn schedule_jobs(&self) -> Result<(), DistributedError> {
let scheduler = self.scheduler.lock().expect("operation should succeed");
let mut queue = self.job_queue.lock().expect("operation should succeed");
let nodes = self.nodes.read().expect("operation should succeed");
queue.sort_by(|a, b| {
b.priority
.cmp(&a.priority)
.then_with(|| a.created_at.cmp(&b.created_at))
});
let mut scheduled_jobs = Vec::new();
for job in queue.iter() {
if let Some(node_id) = scheduler.find_suitable_node(job, &nodes) {
let execution = JobExecution {
job_id: job.id.clone(),
node_id: node_id.clone(),
start_time: Instant::now(),
end_time: None,
status: JobStatus::Running,
progress: 0.0,
result: None,
error: None,
resource_usage: None,
};
self.executions
.write()
.expect("operation should succeed")
.insert(job.id.clone(), execution);
scheduled_jobs.push(job.id.clone());
}
}
queue.retain(|job| !scheduled_jobs.contains(&job.id));
Ok(())
}
pub fn get_job_status(&self, job_id: &str) -> Option<JobStatus> {
self.executions
.read()
.expect("operation should succeed")
.get(job_id)
.map(|exec| exec.status.clone())
}
pub fn get_job_execution(&self, job_id: &str) -> Option<JobExecution> {
self.executions
.read()
.expect("operation should succeed")
.get(job_id)
.cloned()
}
pub fn cancel_job(&self, job_id: &str) -> Result<(), DistributedError> {
let mut executions = self.executions.write().expect("operation should succeed");
if let Some(execution) = executions.get_mut(job_id) {
execution.status = JobStatus::Cancelled;
execution.end_time = Some(Instant::now());
Ok(())
} else {
Err(DistributedError::JobNotFound)
}
}
pub fn update_heartbeat(
&self,
node_id: &str,
load_metrics: LoadMetrics,
) -> Result<(), DistributedError> {
let mut nodes = self.nodes.write().expect("operation should succeed");
if let Some(node) = nodes.get_mut(node_id) {
node.last_heartbeat = Instant::now();
node.load_metrics = load_metrics;
node.status = self.determine_node_status(&node.load_metrics);
Ok(())
} else {
Err(DistributedError::NodeNotFound)
}
}
fn determine_node_status(&self, metrics: &LoadMetrics) -> NodeStatus {
if metrics.cpu_usage > 0.9 || metrics.memory_usage > 0.9 {
NodeStatus::Overloaded
} else if metrics.cpu_usage > 0.7 || metrics.memory_usage > 0.7 {
NodeStatus::Busy
} else {
NodeStatus::Available
}
}
pub fn get_cluster_stats(&self) -> ClusterStats {
let nodes = self.nodes.read().expect("operation should succeed");
let jobs = self.jobs.read().expect("operation should succeed");
let executions = self.executions.read().expect("operation should succeed");
let total_nodes = nodes.len();
let available_nodes = nodes
.values()
.filter(|n| n.status == NodeStatus::Available)
.count();
let busy_nodes = nodes
.values()
.filter(|n| n.status == NodeStatus::Busy)
.count();
let overloaded_nodes = nodes
.values()
.filter(|n| n.status == NodeStatus::Overloaded)
.count();
let total_jobs = jobs.len();
let running_jobs = executions
.values()
.filter(|e| e.status == JobStatus::Running)
.count();
let completed_jobs = executions
.values()
.filter(|e| e.status == JobStatus::Completed)
.count();
let failed_jobs = executions
.values()
.filter(|e| e.status == JobStatus::Failed)
.count();
let total_cpu_cores: u32 = nodes.values().map(|n| n.capabilities.cpu_cores).sum();
let total_memory_gb: u32 = nodes.values().map(|n| n.capabilities.memory_gb).sum();
let total_gpu_count: u32 = nodes.values().map(|n| n.capabilities.gpu_count).sum();
let avg_cpu_usage = if !nodes.is_empty() {
nodes
.values()
.map(|n| n.load_metrics.cpu_usage)
.sum::<f64>()
/ nodes.len() as f64
} else {
0.0
};
let avg_memory_usage = if !nodes.is_empty() {
nodes
.values()
.map(|n| n.load_metrics.memory_usage)
.sum::<f64>()
/ nodes.len() as f64
} else {
0.0
};
ClusterStats {
total_nodes,
available_nodes,
busy_nodes,
overloaded_nodes,
unreachable_nodes: total_nodes - available_nodes - busy_nodes - overloaded_nodes,
total_jobs,
running_jobs,
completed_jobs,
failed_jobs,
queued_jobs: self
.job_queue
.lock()
.expect("operation should succeed")
.len(),
total_cpu_cores,
total_memory_gb,
total_gpu_count,
avg_cpu_usage,
avg_memory_usage,
}
}
pub fn start_monitoring(&self) -> Result<(), DistributedError> {
let nodes = Arc::clone(&self.nodes);
let executions = Arc::clone(&self.executions);
let config = self.config.clone();
thread::spawn(move || loop {
let now = Instant::now();
let mut nodes_guard = nodes.write().expect("operation should succeed");
for node in nodes_guard.values_mut() {
if now.duration_since(node.last_heartbeat) > config.node_timeout {
node.status = NodeStatus::Unreachable;
}
}
let mut executions_guard = executions.write().expect("operation should succeed");
for execution in executions_guard.values_mut() {
if execution.status == JobStatus::Running
&& now.duration_since(execution.start_time) > config.job_timeout
{
execution.status = JobStatus::Timeout;
execution.end_time = Some(now);
}
}
drop(nodes_guard);
drop(executions_guard);
thread::sleep(config.heartbeat_interval);
});
Ok(())
}
pub fn rebalance_workload(&self) -> Result<(), DistributedError> {
let load_balancer = self.load_balancer.lock().expect("operation should succeed");
let nodes = self.nodes.read().expect("operation should succeed");
load_balancer.rebalance(&nodes)?;
Ok(())
}
pub fn handle_node_failure(&self, node_id: &str) -> Result<(), DistributedError> {
let mut fault_detector = self
.fault_detector
.lock()
.expect("operation should succeed");
let mut executions = self.executions.write().expect("operation should succeed");
for execution in executions.values_mut() {
if execution.node_id == node_id && execution.status == JobStatus::Running {
execution.status = JobStatus::Failed;
execution.end_time = Some(Instant::now());
execution.error = Some("Node failure".to_string());
}
}
fault_detector.handle_failure(node_id)?;
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct ResourceAllocation {
pub cpu_cores: u32,
pub memory_gb: u32,
pub gpu_count: u32,
pub storage_gb: u32,
pub network_bandwidth: u32,
}
#[derive(Debug, Clone)]
pub struct JobExecution {
pub job_id: String,
pub node_id: String,
pub start_time: Instant,
pub end_time: Option<Instant>,
pub status: JobStatus,
pub progress: f64,
pub result: Option<String>,
pub error: Option<String>,
pub resource_usage: Option<ResourceUsage>,
}
#[derive(Debug, Clone)]
pub struct ResourceRequirements {
pub min_cpu_cores: u32,
pub min_memory_gb: u32,
pub min_gpu_count: u32,
pub min_storage_gb: u32,
pub preferred_node_tags: HashSet<String>,
pub exclusive_access: bool,
}