Files
rustytorch/crates/training/rtx-distributed/src/elastic_training/health.rs
T
2026-03-04 00:08:42 +00:00

75 lines
1.9 KiB
Rust

//! 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
}
}