feat: expose lighthouse join timeout (#61)

* feat: expose lighthouse join timeout

* lint

* add join timeout test

* expose quorum tick ms

* blacken

Author: Jackmin801

src/lib.rs

@@ -225,16 +225,25 @@ struct Lighthouse {
 #[pymethods]
 impl Lighthouse {
     #[new]
-    fn new(py: Python<'_>, bind: String, min_replicas: u64) -> PyResult<Self> {
+    fn new(
+        py: Python<'_>,
+        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()?;
 
             let lighthouse = rt
                 .block_on(lighthouse::Lighthouse::new(lighthouse::LighthouseOpt {
                     bind: bind,
                     min_replicas: min_replicas,
-                    join_timeout_ms: 100,
-                    quorum_tick_ms: 100,
+                    join_timeout_ms: join_timeout_ms,
+                    quorum_tick_ms: quorum_tick_ms,
                 }))
                 .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
 

torchft/lighthouse_test.py

@@ -0,0 +1,93 @@
+import time
+from unittest import TestCase
+
+import torch.distributed as dist
+
+from torchft import Manager, ProcessGroupGloo
+from torchft.torchft import Lighthouse
+
+
+class TestLighthouse(TestCase):
+    def test_join_timeout_behavior(self) -> None:
+        """Test that join_timeout_ms affects joining behavior"""
+        # To test, we create a lighthouse with 100ms and 400ms join timeouts
+        # and measure the time taken to validate the quorum.
+        lighthouse = Lighthouse(
+            bind="[::]:0",
+            min_replicas=1,
+            join_timeout_ms=100,
+        )
+
+        # Create a manager that tries to join
+        try:
+            store = dist.TCPStore(
+                host_name="localhost",
+                port=0,
+                is_master=True,
+                wait_for_workers=False,
+            )
+            pg = ProcessGroupGloo()
+            manager = Manager(
+                pg=pg,
+                min_replica_size=1,
+                load_state_dict=lambda x: None,
+                state_dict=lambda: None,
+                replica_id=f"lighthouse_test",
+                store_addr="localhost",
+                store_port=store.port,
+                rank=0,
+                world_size=1,
+                use_async_quorum=False,
+                lighthouse_addr=lighthouse.address(),
+            )
+
+            start_time = time.time()
+            manager.start_quorum()
+            time_taken = time.time() - start_time
+            assert time_taken < 0.4, f"Time taken to join: {time_taken} > 0.4s"
+
+        finally:
+            # Cleanup
+            lighthouse.shutdown()
+            if "manager" in locals():
+                manager.shutdown()
+
+        lighthouse = Lighthouse(
+            bind="[::]:0",
+            min_replicas=1,
+            join_timeout_ms=400,
+        )
+
+        # Create a manager that tries to join
+        try:
+            store = dist.TCPStore(
+                host_name="localhost",
+                port=0,
+                is_master=True,
+                wait_for_workers=False,
+            )
+            pg = ProcessGroupGloo()
+            manager = Manager(
+                pg=pg,
+                min_replica_size=1,
+                load_state_dict=lambda x: None,
+                state_dict=lambda: None,
+                replica_id=f"lighthouse_test",
+                store_addr="localhost",
+                store_port=store.port,
+                rank=0,
+                world_size=1,
+                use_async_quorum=False,
+                lighthouse_addr=lighthouse.address(),
+            )
+
+            start_time = time.time()
+            manager.start_quorum()
+            time_taken = time.time() - start_time
+            assert time_taken > 0.4, f"Time taken to join: {time_taken} < 0.4s"
+
+        finally:
+            # Cleanup
+            lighthouse.shutdown()
+            if "manager" in locals():
+                manager.shutdown()

torchft/torchft.pyi

@@ -23,6 +23,6 @@ class Manager:
     def shutdown(self) -> None: ...
 
 class Lighthouse:
-    def __init__(self, bind: str, min_replicas: int) -> 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: ...