[lighthouse] detect unhealthy participants via heartbeats

pytorch · Jan 10, 2025 · 529233a · 529233a
1 parent 5dd6f38
commit 529233a
Show file tree

Hide file tree

Showing 6 changed files with 192 additions and 37 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -44,6 +44,7 @@ impl Manager {
         bind: String,
         store_addr: String,
         world_size: u64,
+        heartbeat_interval: Duration,
     ) -> PyResult<Self> {
         py.allow_threads(move || {
             let runtime = Runtime::new()?;
@@ -55,6 +56,7 @@ impl Manager {
                     bind,
                     store_addr,
                     world_size,
+                    heartbeat_interval,
                 ))
                 .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
             let handle = runtime.spawn(manager.clone().run());
@@ -224,16 +226,19 @@ struct Lighthouse {
 
 #[pymethods]
 impl Lighthouse {
+    #[pyo3(signature = (bind, min_replicas, join_timeout_ms=None, quorum_tick_ms=None, heartbeat_timeout_ms=None))]
     #[new]
     fn new(
         py: Python<'_>,
         bind: String,
         min_replicas: u64,
         join_timeout_ms: Option<u64>,
         quorum_tick_ms: Option<u64>,
+        heartbeat_timeout_ms: Option<u64>,
     ) -> PyResult<Self> {
         let join_timeout_ms = join_timeout_ms.unwrap_or(100);
         let quorum_tick_ms = quorum_tick_ms.unwrap_or(100);
+        let heartbeat_timeout_ms = heartbeat_timeout_ms.unwrap_or(5000);
 
         py.allow_threads(move || {
             let rt = Runtime::new()?;
@@ -244,6 +249,7 @@ impl Lighthouse {
                     min_replicas: min_replicas,
                     join_timeout_ms: join_timeout_ms,
                     quorum_tick_ms: quorum_tick_ms,
+                    heartbeat_timeout_ms: heartbeat_timeout_ms,
                 }))
                 .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;