expose quorum tick ms

pytorch · Jan 8, 2025 · 13e72f1 · 13e72f1
1 parent a41eb1c
commit 13e72f1
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -230,8 +230,10 @@ impl Lighthouse {
         bind: String,
         min_replicas: u64,
         join_timeout_ms: Option<u64>,
+        quorum_tick_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);
 
         py.allow_threads(move || {
             let rt = Runtime::new()?;
@@ -241,7 +243,7 @@ impl Lighthouse {
                     bind: bind,
                     min_replicas: min_replicas,
                     join_timeout_ms: join_timeout_ms,
-                    quorum_tick_ms: 100,
+                    quorum_tick_ms: quorum_tick_ms,
                 }))
                 .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
 

diff --git a/torchft/torchft.pyi b/torchft/torchft.pyi
@@ -23,6 +23,6 @@ class Manager:
     def shutdown(self) -> None: ...
 
 class Lighthouse:
-    def __init__(self, bind: str, min_replicas: int, join_timeout_ms: Optional[int] = None) -> None: ...
+    def __init__(self, bind: str, min_replicas: int, join_timeout_ms: Optional[int] = None, quorum_tick_ms: Optional[int] = None) -> None: ...
     def address(self) -> str: ...
     def shutdown(self) -> None: ...