Initial commit

crates/training/rtx-distributed/src/elastic_training/health.rs

@@ -0,0 +1,74 @@
+//! Health monitoring for elastic cluster.
+
+use parking_lot::RwLock;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::time::Duration;
+
+use super::cluster::ElasticCluster;
+use super::worker::WorkerId;
+
+/// Health monitoring for elastic cluster
+pub struct HealthMonitor {
+    /// Cluster reference
+    cluster: Arc<ElasticCluster>,
+    /// Check interval
+    check_interval: Duration,
+    /// Is monitoring
+    running: AtomicBool,
+    /// Failure callbacks
+    failure_handlers: RwLock<Vec<Box<dyn Fn(WorkerId) + Send + Sync>>>,
+}
+
+impl HealthMonitor {
+    /// Create a new health monitor
+    pub fn new(cluster: Arc<ElasticCluster>, check_interval_ms: u64) -> Self {
+        Self {
+            cluster,
+            check_interval: Duration::from_millis(check_interval_ms),
+            running: AtomicBool::new(false),
+            failure_handlers: RwLock::new(Vec::new()),
+        }
+    }
+
+    /// Start monitoring
+    pub fn start(&self) {
+        self.running.store(true, Ordering::SeqCst);
+    }
+
+    /// Stop monitoring
+    pub fn stop(&self) {
+        self.running.store(false, Ordering::SeqCst);
+    }
+
+    /// Check for failures once
+    pub fn check_once(&self) -> Vec<WorkerId> {
+        let failed = self.cluster.detect_failures();
+
+        for &worker_id in &failed {
+            let _ = self.cluster.mark_failed(worker_id);
+
+            // Call handlers
+            let handlers = self.failure_handlers.read();
+            for handler in handlers.iter() {
+                handler(worker_id);
+            }
+        }
+
+        failed
+    }
+
+    /// Add failure handler
+    pub fn on_failure<F>(&self, handler: F)
+    where
+        F: Fn(WorkerId) + Send + Sync + 'static,
+    {
+        let mut handlers = self.failure_handlers.write();
+        handlers.push(Box::new(handler));
+    }
+
+    /// Get check interval
+    pub fn check_interval(&self) -> Duration {
+        self.check_interval
+    }
+}