Expose heartbeat method (src/lib.rs LighthouseClient) on pyo3 bindings (#176)

* Expose heartbeat method (src/lib.rs LighthouseClient) on pyo3 bindings

* fix lint

---------

Co-authored-by: Matt Kotzbauer <matt@dhcp-10-250-15-170.harvard.edu>
Co-authored-by: Tristan Rice <rice@fn.lc>

Author: 3 people

src/lib.rs

@@ -35,7 +35,8 @@ pub mod torchftpb {
 use crate::torchftpb::lighthouse_service_client::LighthouseServiceClient;
 use crate::torchftpb::manager_service_client::ManagerServiceClient;
 use crate::torchftpb::{
-    CheckpointMetadataRequest, LighthouseQuorumRequest, ManagerQuorumRequest, ShouldCommitRequest,
+    CheckpointMetadataRequest, LighthouseHeartbeatRequest, LighthouseQuorumRequest,
+    ManagerQuorumRequest, ShouldCommitRequest,
 };
 use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyString};
@@ -562,6 +563,26 @@ impl LighthouseClient {
         });
         Ok(convert_quorum(py, &quorum?)?)
     }
+
+    /// Send a single heartbeat to the lighthouse.
+    ///
+    /// Args:
+    ///     replica_id (str):  The replica_id you registered with.
+    ///     timeout      (timedelta, optional):  Per-RPC deadline.  Default = 5 s.
+    #[pyo3(signature = (replica_id, timeout = Duration::from_secs(5)))]
+    fn heartbeat(
+        &self,
+        py: Python<'_>,
+        replica_id: String,
+        timeout: Duration,
+    ) -> Result<(), StatusError> {
+        py.allow_threads(move || {
+            let mut req = tonic::Request::new(LighthouseHeartbeatRequest { replica_id });
+            req.set_timeout(timeout);
+            self.runtime.block_on(self.client.clone().heartbeat(req))?;
+            Ok(())
+        })
+    }
 }
 
 /// LighthouseServer is a GRPC server for the lighthouse service.

torchft/_torchft.pyi

@@ -99,3 +99,8 @@ class LighthouseClient:
         shrink_only: Optional[bool] = None,
         data: Optional[dict[Hashable, object]] = None,
     ) -> Quorum: ...
+    def heartbeat(
+        self,
+        replica_id: str,
+        timeout: timedelta = timedelta(seconds=5),
+    ) -> None: ...

torchft/lighthouse_test.py

@@ -121,3 +121,37 @@ def test_lighthouse_client_behavior(self) -> None:
         finally:
             # Cleanup
             lighthouse.shutdown()
+
+    def test_heartbeat_round_trip(self) -> None:
+        lighthouse = LighthouseServer(
+            bind="[::]:0",
+            min_replicas=1,
+            heartbeat_timeout_ms=200,
+        )
+        try:
+            client = LighthouseClient(
+                addr=lighthouse.address(),
+                connect_timeout=timedelta(seconds=1),
+            )
+
+            client.heartbeat("rep0")
+
+            # (Should still be alive, as sleep time is less than timeout)
+            time.sleep(0.15)
+            q = client.quorum(
+                replica_id="rep0",
+                timeout=timedelta(milliseconds=500),
+            )
+            assert any(m.replica_id == "rep0" for m in q.participants)
+
+            # (Wait long enough for timeout to trigger)
+            time.sleep(0.25)
+            # "Probe" with different replica so we don't revive rep0
+            probe = client.quorum(
+                replica_id="probe",
+                timeout=timedelta(milliseconds=500),
+            )
+            assert all(m.replica_id != "rep0" for m in probe.participants)
+
+        finally:
+            lighthouse.shutdown()