Support generic quorum api on LighthouseClient (#150)

* support generic quorum api on LighthouseClient with document updated.

Author: fduwjj

Cargo.toml

@@ -14,7 +14,7 @@ gethostname = "0.5.0"
 log = "0.4.22"
 prost = "0.13.3"
 prost-types = "0.13.3"
-pyo3 = {version = "0.22.3", features = ["extension-module"]}
+pyo3 = {version = "0.24", features = ["extension-module"]}
 rand = "0.8.5"
 slog = "2.7.0"
 slog-stdlog = "4.1.1"

proto/torchft.proto

@@ -42,6 +42,8 @@ message QuorumMember {
     int64 step = 4;
     uint64 world_size = 5;
     bool shrink_only = 6;
+    // User passing in data stored as JSON string.
+    string data = 7;
 }
 
 message Quorum {

src/lib.rs

@@ -32,9 +32,13 @@ pub mod torchftpb {
     tonic::include_proto!("torchft");
 }
 
+use crate::torchftpb::lighthouse_service_client::LighthouseServiceClient;
 use crate::torchftpb::manager_service_client::ManagerServiceClient;
-use crate::torchftpb::{CheckpointMetadataRequest, ManagerQuorumRequest, ShouldCommitRequest};
+use crate::torchftpb::{
+    CheckpointMetadataRequest, LighthouseQuorumRequest, ManagerQuorumRequest, ShouldCommitRequest,
+};
 use pyo3::prelude::*;
+use pyo3::types::{PyDict, PyString};
 
 // Get the number of threads to use for the tokio runtime
 fn num_threads() -> usize {
@@ -304,7 +308,7 @@ impl QuorumResult {
 fn reset_python_signals(py: Python<'_>) -> PyResult<()> {
     // clear python signal handlers
     // signal.signal(signal.SIGINT, signal.SIG_DFL)
-    let signal = py.import_bound("signal")?;
+    let signal = py.import("signal")?;
     let set_signal = signal.getattr("signal")?;
     let args = (signal.getattr("SIGINT")?, signal.getattr("SIG_DFL")?);
     set_signal.call1(args)?;
@@ -337,6 +341,217 @@ async fn lighthouse_main_async(opt: lighthouse::LighthouseOpt) -> Result<()> {
     Ok(())
 }
 
+/// quorum member of one quorum.
+///
+/// Args:
+///     replica_id (str): The string id of the replica calling quorum.
+///     address (str): The address of the replica calling quorum.
+///     store_address (str): The address of the store.
+///     step (int): The step of the replica calling quorum.
+///     world_size (int): The world size of the replica calling quorum.
+///     shrink_only (bool): Whether the quorum is for shrinking only.
+///     timeout (timedelta): The timeout for quorum.
+///     data (dict or None): The data to be passed with quorum.
+#[pyclass(get_all, set_all)]
+pub struct QuorumMember {
+    replica_id: String,
+    address: String,
+    store_address: String,
+    step: i64,
+    world_size: u64,
+    shrink_only: bool,
+    data: Option<Py<PyDict>>,
+}
+
+impl QuorumMember {
+    // PyDict has not implemeted Clone, so we need to implement it manually
+    pub fn clone_with_py(&self, py: Python) -> Self {
+        QuorumMember {
+            replica_id: self.replica_id.clone(),
+            address: self.address.clone(),
+            store_address: self.store_address.clone(),
+            step: self.step,
+            world_size: self.world_size,
+            shrink_only: self.shrink_only,
+            data: self.data.as_ref().map(|d| d.clone_ref(py)),
+        }
+    }
+}
+
+impl Clone for QuorumMember {
+    fn clone(&self) -> Self {
+        Python::with_gil(|py| self.clone_with_py(py))
+    }
+}
+
+#[pyclass(get_all, set_all)]
+#[derive(Clone)]
+pub struct Timestamp {
+    pub seconds: i64,
+    pub nanos: i32,
+}
+
+/// quorum result.
+///
+/// Args:
+///     quorum_id (int): The id of current quorum.
+///     participants (list[QuorumMember]): All members within the quorum.
+///     created (timedelta): Time of quorum created in server.
+#[pyclass(get_all, set_all)]
+struct Quorum {
+    quorum_id: i64,
+    participants: Vec<QuorumMember>,
+    created: Timestamp,
+}
+
+impl From<prost_types::Timestamp> for Timestamp {
+    fn from(ts: prost_types::Timestamp) -> Self {
+        Timestamp {
+            seconds: ts.seconds,
+            nanos: ts.nanos,
+        }
+    }
+}
+
+// Util functions to convert between python dict and rust string using json.
+fn pydict_to_string<'py>(py: Python, data: Option<&Bound<'_, PyDict>>) -> PyResult<String> {
+    match data {
+        Some(d) => {
+            let json = py.import("json")?;
+            let json_obj = json.call_method1("dumps", (d,))?;
+            let py_str: &Bound<PyString> = json_obj.downcast()?;
+            Ok(py_str.to_str()?.to_owned())
+        }
+        None => Ok(String::new()),
+    }
+}
+
+fn string_to_pydict(py: Python, s: &str) -> PyResult<Option<Py<PyDict>>> {
+    if s.is_empty() {
+        return Ok(None); // Treat empty string as None
+    }
+
+    let json = py.import("json")?;
+    let obj = json.call_method1("loads", (s,))?;
+    let dict: &Bound<PyDict> = obj.downcast()?;
+    Ok(Some(dict.to_owned().into())) // convert Bound<PyDict> -> Py<PyDict>
+}
+
+fn convert_quorum_member(py: Python, m: &torchftpb::QuorumMember) -> PyResult<QuorumMember> {
+    Ok(QuorumMember {
+        replica_id: m.replica_id.clone(),
+        address: m.address.clone(),
+        store_address: m.store_address.clone(),
+        step: m.step.clone(),
+        world_size: m.world_size.clone(),
+        shrink_only: m.shrink_only.clone(),
+        data: string_to_pydict(py, &m.data)?,
+    })
+}
+
+fn convert_quorum(py: Python, q: &torchftpb::Quorum) -> PyResult<Quorum> {
+    let participants: Vec<QuorumMember> = q
+        .participants
+        .iter()
+        .map(|m| convert_quorum_member(py, m)) // this expects &m
+        .collect::<Result<Vec<_>, _>>()?;
+
+    Ok(Quorum {
+        quorum_id: q.quorum_id,
+        participants: participants,
+        created: Timestamp::from(q.created.unwrap()),
+    })
+}
+
+/// LighthouseClient is a GRPC client to the lighthouse service.
+///
+/// It is used to directly communicate with the lighthouse Server.
+///
+/// Args:
+///     addr (str): The HTTP address of the lighthouse server.
+///     connect_timeout (timedelta): The timeout for connecting to the lighthouse server.
+#[pyclass]
+struct LighthouseClient {
+    client: LighthouseServiceClient<Channel>,
+    runtime: Runtime,
+}
+
+#[pymethods]
+impl LighthouseClient {
+    #[pyo3(signature = (addr, connect_timeout))]
+    #[new]
+    fn new(py: Python<'_>, addr: String, connect_timeout: Duration) -> PyResult<Self> {
+        py.allow_threads(move || {
+            let runtime = tokio::runtime::Builder::new_multi_thread()
+                .worker_threads(num_threads())
+                .thread_name("torchft-lhclnt")
+                .enable_all()
+                .build()?;
+            let client = runtime
+                .block_on(manager::lighthouse_client_new(addr, connect_timeout))
+                .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+            Ok(Self {
+                client: client,
+                runtime: runtime,
+            })
+        })
+    }
+
+    /// quorum sends a request to the lighthouse server to form a quorum.
+    ///
+    /// Args:
+    ///     replica_id (str): The string id of the replica calling quorum.
+    ///     address (str): The address of the replica calling quorum.
+    ///     store_address (str): The address of the store.
+    ///     step (int): The step of the replica calling quorum.
+    ///     world_size (int): The world size of the replica calling quorum.
+    ///     shrink_only (bool): Whether the quorum is for shrinking only.
+    ///     timeout (timedelta): The timeout for quorum.
+    ///     data (Optional[dict]): The data to be passed with quorum.
+    ///
+    /// Returns:
+    ///     Quorum: Current quorum if successful.
+    fn quorum<'py>(
+        &self,
+        py: Python<'_>,
+        replica_id: String,
+        address: String,
+        store_address: String,
+        step: i64,
+        world_size: u64,
+        shrink_only: bool,
+        timeout: Duration,
+        data: Option<&Bound<'_, PyDict>>,
+    ) -> Result<Quorum, StatusError> {
+        let data_string = pydict_to_string(py, data)?;
+        let quorum: Result<torchftpb::Quorum, StatusError> = py.allow_threads(move || {
+            let mut request = tonic::Request::new(LighthouseQuorumRequest {
+                requester: Some(torchftpb::QuorumMember {
+                    replica_id: replica_id,
+                    address: address,
+                    store_address: store_address,
+                    step: step,
+                    world_size: world_size,
+                    shrink_only: shrink_only,
+                    data: data_string,
+                }),
+            });
+
+            // This timeout is processed on the server side so we also enable
+            // keep alives to detect server health.
+            request.set_timeout(timeout);
+
+            let response = self.runtime.block_on(self.client.clone().quorum(request))?;
+            let resp = response.into_inner();
+            let quorum = resp
+                .quorum
+                .ok_or_else(|| Status::internal("missing quorum"))?;
+            Ok(quorum)
+        });
+        Ok(convert_quorum(py, &quorum?)?)
+    }
+}
+
 /// LighthouseServer is a GRPC server for the lighthouse service.
 ///
 /// It is used to coordinate the ManagerServer for each replica group.
@@ -428,6 +643,12 @@ impl From<StatusError> for PyErr {
     }
 }
 
+impl From<pyo3::PyErr> for StatusError {
+    fn from(err: pyo3::PyErr) -> Self {
+        StatusError(Status::internal(err.to_string()))
+    }
+}
+
 impl From<Status> for StatusError {
     fn from(other: Status) -> Self {
         Self(other)
@@ -479,9 +700,13 @@ fn _torchft(m: &Bound<'_, PyModule>) -> PyResult<()> {
     // setup logging on import
     setup_logging().map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
 
+    m.add_class::<Timestamp>()?;
+    m.add_class::<QuorumMember>()?;
+    m.add_class::<Quorum>()?;
     m.add_class::<ManagerServer>()?;
     m.add_class::<ManagerClient>()?;
     m.add_class::<LighthouseServer>()?;
+    m.add_class::<LighthouseClient>()?;
     m.add_class::<QuorumResult>()?;
     m.add_function(wrap_pyfunction!(lighthouse_main, m)?)?;