overhaul timeouts for Lighthouse, Manager, checkpoint server

Author: d4l3k

Cargo.toml

@@ -12,6 +12,7 @@ log = "0.4.22"
 prost = "0.13.3"
 prost-types = "0.13.3"
 pyo3 = {version="0.22.3", features = ["extension-module"]}
+rand = "0.8.5"
 slog = "2.7.0"
 slog-stdlog = "4.1.1"
 stderrlog = "0.6.0"

src/lib.rs

@@ -6,6 +6,9 @@
 
 pub mod lighthouse;
 pub mod manager;
+mod net;
+mod retry;
+mod timeout;
 
 use core::time::Duration;
 use std::env;
@@ -84,36 +87,33 @@ impl Manager {
 struct ManagerClient {
     runtime: Runtime,
     client: ManagerServiceClient<Channel>,
-    timeout: Duration,
 }
 
 #[pymethods]
 impl ManagerClient {
     #[new]
-    fn new(py: Python<'_>, addr: String, timeout: Duration) -> PyResult<Self> {
+    fn new(py: Python<'_>, addr: String) -> PyResult<Self> {
         py.allow_threads(move || {
             let runtime = Runtime::new()?;
             let client = runtime
-                .block_on(manager::manager_client_new(addr, timeout))
+                .block_on(manager::manager_client_new(addr, Duration::from_secs(60)))
                 .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
 
             Ok(Self {
                 runtime: runtime,
                 client: client,
-                timeout: timeout,
             })
         })
     }
 
-    #[pyo3(signature = (rank, step, checkpoint_server_addr, shrink_only, timeout=None))]
     fn quorum(
         &mut self,
         py: Python<'_>,
         rank: i64,
         step: i64,
         checkpoint_server_addr: String,
         shrink_only: bool,
-        timeout: Option<Duration>,
+        timeout: Duration,
     ) -> Result<(i64, i64, i64, String, String, i64, Option<i64>, i64, bool), StatusError> {
         py.allow_threads(move || {
             let mut request = tonic::Request::new(ManagerQuorumRequest {
@@ -122,9 +122,10 @@ impl ManagerClient {
                 checkpoint_server_addr: checkpoint_server_addr,
                 shrink_only: shrink_only,
             });
-            // This notifies the server about the timeout but doesn't affect the
-            // endpoint timeout which we set on client creation.
-            request.set_timeout(timeout.unwrap_or(self.timeout));
+
+            // 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.quorum(request))?;
             let resp = response.into_inner();
@@ -142,18 +143,18 @@ impl ManagerClient {
         })
     }
 
-    #[pyo3(signature = (rank, timeout=None))]
     fn checkpoint_address(
         &mut self,
         py: Python<'_>,
         rank: i64,
-        timeout: Option<Duration>,
+        timeout: Duration,
     ) -> Result<String, StatusError> {
         py.allow_threads(move || {
             let mut request = tonic::Request::new(CheckpointAddressRequest { rank: rank });
-            // This notifies the server about the timeout but doesn't affect the
-            // endpoint timeout which we set on client creation.
-            request.set_timeout(timeout.unwrap_or(self.timeout));
+
+            // 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
@@ -163,24 +164,24 @@ impl ManagerClient {
         })
     }
 
-    #[pyo3(signature = (rank, step, should_commit, timeout=None))]
     fn should_commit(
         &mut self,
         py: Python<'_>,
         rank: i64,
         step: i64,
         should_commit: bool,
-        timeout: Option<Duration>,
+        timeout: Duration,
     ) -> Result<bool, StatusError> {
         py.allow_threads(move || {
             let mut request = tonic::Request::new(ShouldCommitRequest {
                 rank: rank,
                 step: step,
                 should_commit: should_commit,
             });
+
             // This notifies the server about the timeout but doesn't affect the
             // endpoint timeout which we set on client creation.
-            request.set_timeout(timeout.unwrap_or(self.timeout));
+            request.set_timeout(timeout);
 
             let response = self.runtime.block_on(self.client.should_commit(request))?;
             let resp = response.into_inner();

src/lighthouse.rs

@@ -543,15 +543,13 @@ mod tests {
     use super::*;
     use std::ops::Sub;
 
-    use tonic::transport::{Channel, Endpoint};
+    use tonic::transport::Channel;
 
+    use crate::net::connect;
     use crate::torchftpb::lighthouse_service_client::LighthouseServiceClient;
 
     async fn lighthouse_client_new(addr: String) -> Result<LighthouseServiceClient<Channel>> {
-        let conn = Endpoint::new(addr)?
-            .connect_timeout(Duration::from_secs(10))
-            .connect()
-            .await?;
+        let conn = connect(addr, Duration::from_secs(10)).await?;
         Ok(LighthouseServiceClient::new(conn))
     }
 

src/manager.rs

@@ -16,10 +16,12 @@ use tokio::sync::Mutex;
 use tokio::task::JoinSet;
 use tokio::time::sleep;
 use tonic::transport::server::TcpIncoming;
+use tonic::transport::Channel;
 use tonic::transport::Server;
-use tonic::transport::{Channel, Endpoint};
 use tonic::{Request, Response, Status};
 
+use crate::net::connect;
+use crate::timeout::try_parse_grpc_timeout;
 use crate::torchftpb::lighthouse_service_client::LighthouseServiceClient;
 use crate::torchftpb::manager_service_client::ManagerServiceClient;
 use crate::torchftpb::{
@@ -59,16 +61,10 @@ pub struct Manager {
 
 pub async fn manager_client_new(
     addr: String,
-    timeout: Duration,
+    connect_timeout: Duration,
 ) -> Result<ManagerServiceClient<Channel>> {
-    // TODO add retries + backoff so other nodes can start before the rank0 comes up
-
     info!("ManagerClient: establishing connection to {}", &addr);
-    let conn = Endpoint::new(addr.clone())?
-        .timeout(timeout)
-        .connect_timeout(Duration::from_secs(60))
-        .connect()
-        .await?;
+    let conn = connect(addr, connect_timeout).await?;
     Ok(ManagerServiceClient::new(conn))
 }
 
@@ -163,12 +159,57 @@ impl Manager {
             &self.lighthouse_addr
         );
 
-        let conn = Endpoint::new(self.lighthouse_addr.clone())?
-            .connect_timeout(Duration::from_secs(60))
-            .connect()
-            .await?;
+        let conn = connect(self.lighthouse_addr.clone(), Duration::from_secs(60)).await?;
         Ok(LighthouseServiceClient::new(conn))
     }
+
+    async fn _run_quorum(
+        &self,
+        state: &mut ManagerState,
+        requester: QuorumMember,
+        timeout: Duration,
+    ) -> Result<(), Status> {
+        if (state.participants.len() as u64) < self.world_size {
+            return Ok(());
+        }
+
+        state.participants.clear();
+        info!("all workers joined -- starting quorum");
+
+        // TODO: don't hold the lock during quorum
+
+        let mut client = self
+            .lighthouse_client_new()
+            .await
+            .map_err(|e| Status::from_error(e.into()))?;
+
+        let mut lighthouse_request = tonic::Request::new(LighthouseQuorumRequest {
+            requester: Some(requester),
+        });
+        lighthouse_request.set_timeout(timeout);
+
+        let response = tokio::time::timeout(timeout, client.quorum(lighthouse_request))
+            .await
+            .unwrap_or_else(|e| {
+                Err(Status::cancelled(format!(
+                    "lighthouse quorum timed out: {}",
+                    e.to_string()
+                )))
+            })?;
+        let resp = response.into_inner();
+
+        info!("got lighthouse quorum {:?}", resp);
+
+        state
+            .channel
+            .send(
+                resp.quorum
+                    .ok_or_else(|| Status::internal("missing quorum"))?,
+            )
+            .map_err(|e| Status::from_error(e.into()))?;
+
+        Ok(())
+    }
 }
 
 #[tonic::async_trait]
@@ -182,6 +223,15 @@ impl ManagerService for Arc<Manager> {
 
         info!("got quorum request for rank {}", rank);
 
+        let timeout = try_parse_grpc_timeout(&request.metadata())
+            .map_err(|e| {
+                Status::invalid_argument(format!(
+                    "invalid timeout {}",
+                    e.to_str().unwrap_or("invalid")
+                ))
+            })?
+            .ok_or_else(|| Status::invalid_argument("missing timeout"))?;
+
         let mut rx = {
             let mut state = self.state.lock().await;
 
@@ -195,50 +245,19 @@ impl ManagerService for Arc<Manager> {
             state.participants.insert(rank);
             let rx = state.channel.subscribe();
 
-            if state.participants.len() as u64 >= self.world_size {
-                state.participants.clear();
-                info!("all workers joined -- starting quorum");
-
-                // TODO: don't hold the lock during quorum
-
-                let mut client = self
-                    .lighthouse_client_new()
-                    .await
-                    .map_err(|e| Status::from_error(e.into()))?;
-
-                let mut lighthouse_request = tonic::Request::new(LighthouseQuorumRequest {
-                    requester: Some(QuorumMember {
-                        replica_id: self.replica_id.clone(),
-                        address: self.address(),
-                        store_address: self.store_address.clone(),
-                        step: req.step,
-                        world_size: self.world_size,
-                        shrink_only: req.shrink_only,
-                    }),
-                });
-
-                // propagate timeout from request to lighthouse
-                let timeout = request
-                    .metadata()
-                    .get("grpc-timeout")
-                    .ok_or_else(|| Status::internal("grpc-timeout not set"))?;
-                lighthouse_request
-                    .metadata_mut()
-                    .insert("grpc-timeout", timeout.clone());
-
-                let response = client.quorum(lighthouse_request).await.unwrap();
-                let resp = response.into_inner();
-
-                info!("got lighthouse quorum {:?}", resp);
-
-                state
-                    .channel
-                    .send(
-                        resp.quorum
-                            .ok_or_else(|| Status::internal("missing quorum"))?,
-                    )
-                    .map_err(|e| Status::from_error(e.into()))?;
-            }
+            self._run_quorum(
+                &mut state,
+                QuorumMember {
+                    replica_id: self.replica_id.clone(),
+                    address: self.address(),
+                    store_address: self.store_address.clone(),
+                    step: req.step,
+                    world_size: self.world_size,
+                    shrink_only: req.shrink_only,
+                },
+                timeout,
+            )
+            .await?;
 
             rx
         };

src/net.rs

@@ -0,0 +1,34 @@
+use std::time::Duration;
+
+use anyhow::Result;
+use tonic::transport::{Channel, Endpoint};
+
+use crate::retry::{retry_backoff, ExponentialBackoff};
+
+pub async fn connect_once(addr: String, connect_timeout: Duration) -> Result<Channel> {
+    let conn = Endpoint::new(addr)?
+        .connect_timeout(connect_timeout)
+        // Enable HTTP2 keep alives
+        .http2_keep_alive_interval(Duration::from_secs(60))
+        // Time taken for server to respond. 20s is default for GRPC.
+        .keep_alive_timeout(Duration::from_secs(20))
+        // Enable alive for idle connections.
+        .keep_alive_while_idle(true)
+        .connect()
+        .await?;
+    Ok(conn)
+}
+
+pub async fn connect(addr: String, connect_timeout: Duration) -> Result<Channel> {
+    retry_backoff(
+        ExponentialBackoff {
+            initial_backoff: Duration::from_millis(100),
+            max_backoff: Duration::from_secs(10),
+            timeout: connect_timeout,
+            factor: 1.5,
+            jitter: Duration::from_millis(100),
+        },
+        || Box::pin(connect_once(addr.clone(), connect_timeout)),
+    )
+    .await
+}