From fb87c9c61e805d375222b8574c027a014b5a2ef2 Mon Sep 17 00:00:00 2001 From: MauriceDHanisch Date: Thu, 2 Jul 2026 21:35:28 -0700 Subject: [PATCH] fix(worker): retry the connect+auth+registration handshake on transient stalls A busy server (mid MILP scheduler solve on its single-threaded executor, or handling a burst of other workers registering) could fail to poll a new worker's connection in time for any leg of the auth handshake or the registration round-trip to complete within its own hardcoded 15s timeout. None of those legs were retried (only the raw TCP connect already was), so one slow attempt killed the whole hq worker start process -- on an autoalloc-managed SLURM pilot, that means the entire allocation is wasted, not just a delayed worker. Reproduced live: a combined registration-burst + real scheduler-load stress test that previously left a second batch of workers permanently unable to register now has all of them connect cleanly (tested up to 30 workers). Retries the whole connect+authenticate+register+await-response sequence (connect_and_register_with_retry, 8 attempts / 10s backoff) rather than raising the fixed 15s timeouts themselves, since there's no fixed stall duration to size a bigger timeout against. Also fixes two latent panics on the registration-response wait (a timeout or early connection close used to abort the process outright instead of returning a proper error). Added DsError::AuthenticationRejected, distinct from the existing transient GenericError, so a deterministic rejection (wrong secret key, incompatible protocol/role) fails fast instead of retrying 8 times for an outcome that will never change -- caught by the existing test_secret.rs tests, which went from ~0.1s to ~70s combined before this fix. Scoped to the worker's own connection to the server only (worker/rpc.rs); client CLI commands and server-side code are untouched. --- crates/tako/src/internal/common/error.rs | 6 + crates/tako/src/internal/transfer/auth.rs | 33 +++-- crates/tako/src/internal/worker/rpc.rs | 170 +++++++++++++++------- 3 files changed, 148 insertions(+), 61 deletions(-) diff --git a/crates/tako/src/internal/common/error.rs b/crates/tako/src/internal/common/error.rs index b93f2b04f..1fff4313e 100644 --- a/crates/tako/src/internal/common/error.rs +++ b/crates/tako/src/internal/common/error.rs @@ -11,6 +11,12 @@ pub enum DsError { SchedulerError(String), #[error("{0}")] GenericError(String), + /// The peer deterministically rejected the authentication handshake (wrong + /// secret key, incompatible protocol/role). Unlike a timed-out or dropped + /// connection, retrying reproduces the exact same rejection every time, so + /// callers should treat this as immediately fatal rather than retryable. + #[error("{0}")] + AuthenticationRejected(String), } impl From for DsError { diff --git a/crates/tako/src/internal/transfer/auth.rs b/crates/tako/src/internal/transfer/auth.rs index fb2c808a1..c29144fd8 100644 --- a/crates/tako/src/internal/transfer/auth.rs +++ b/crates/tako/src/internal/transfer/auth.rs @@ -136,12 +136,17 @@ impl Authenticator { message: AuthenticationResponse, ) -> crate::Result<(Option, Option)> { if let Some(error) = std::mem::take(&mut self.error) { - return Err(format!("Authentication failed: {error}").into()); + return Err(DsError::AuthenticationRejected(format!( + "Authentication failed: {error}" + ))); } let opener = match (message, &self.secret_key) { (AuthenticationResponse::Error(error), _) => { - return Err(format!("Received authentication error: {}", error.message).into()); + return Err(DsError::AuthenticationRejected(format!( + "Received authentication error: {}", + error.message + ))); } (AuthenticationResponse::NoAuth, None) => { log::trace!("Empty authentication finished"); @@ -149,26 +154,32 @@ impl Authenticator { } (AuthenticationResponse::Encryption(response), Some(key)) => { log::trace!("Challenge verification started"); - let remote_nonce = - &Nonce::from_slice(&response.nonce).map_err(|_| "Invalid nonce")?; - let mut opener = - StreamOpener::new(key, remote_nonce).map_err(|_| "Failed to create opener")?; - let (opened_challenge, tag) = opener - .open_chunk(&response.response) - .map_err(|_| DsError::GenericError("Cannot verify challenge".to_string()))?; + let remote_nonce = &Nonce::from_slice(&response.nonce) + .map_err(|_| DsError::AuthenticationRejected("Invalid nonce".to_string()))?; + let mut opener = StreamOpener::new(key, remote_nonce).map_err(|_| { + DsError::AuthenticationRejected("Failed to create opener".to_string()) + })?; + let (opened_challenge, tag) = + opener.open_chunk(&response.response).map_err(|_| { + DsError::AuthenticationRejected("Cannot verify challenge".to_string()) + })?; let mut expected_response = Vec::new(); expected_response.extend_from_slice(self.peer_role.as_bytes()); expected_response.extend_from_slice(&self.challenge); if tag != StreamTag::Message || opened_challenge != expected_response { - return Err("Received challenge does not match.".into()); + return Err(DsError::AuthenticationRejected( + "Received challenge does not match.".to_string(), + )); } log::trace!("Challenge verification finished"); Some(opener) } (_, _) => { - return Err("Invalid authentication state".into()); + return Err(DsError::AuthenticationRejected( + "Invalid authentication state".to_string(), + )); } }; Ok((self.sealer, opener)) diff --git a/crates/tako/src/internal/worker/rpc.rs b/crates/tako/src/internal/worker/rpc.rs index 872ba2802..0bf799a0b 100644 --- a/crates/tako/src/internal/worker/rpc.rs +++ b/crates/tako/src/internal/worker/rpc.rs @@ -89,6 +89,87 @@ pub async fn connect_to_server_and_authenticate( }) } +// A busy server (e.g. mid scheduler solve on its single-threaded executor, or +// handling a burst of other workers registering) may not poll a new worker's +// connection in time for any single attempt at the auth handshake or the +// registration round-trip below to complete within their own timeouts. Unlike +// connect_to_server's raw TCP connect (already retried), neither of those had any +// retry at all: one slow attempt used to kill the worker process outright, which +// on an autoalloc-managed SLURM pilot means the whole allocation is wasted, not +// just a delayed worker. Retried here, not by raising the handshake's own +// timeouts, since there is no fixed stall duration to size a bigger timeout +// against. +const REGISTRATION_MAX_ATTEMPTS: u32 = 8; +const REGISTRATION_RETRY_DELAY: Duration = Duration::from_secs(10); + +async fn connect_and_register( + scheduler_addresses: &[SocketAddr], + secret_key: Option>, + configuration: &WorkerConfiguration, +) -> crate::Result<(ConnectionDescriptor, WorkerRegistrationResponse)> { + let ConnectionDescriptor { + address, + mut sender, + mut receiver, + mut opener, + mut sealer, + } = connect_to_server_and_authenticate(scheduler_addresses, secret_key).await?; + + let message = ConnectionRegistration::Worker(RegisterWorker { + configuration: configuration.clone(), + }); + let data = serialize(&message)?.into(); + sender.send(seal_message(&mut sealer, data)).await?; + + let response = timeout(Duration::from_secs(15), receiver.next()) + .await + .map_err(|_| { + crate::Error::GenericError("Did not receive worker registration response".into()) + })? + .ok_or_else(|| { + crate::Error::GenericError( + "Connection closed without receiving registration response".into(), + ) + })??; + let response: WorkerRegistrationResponse = open_message(&mut opener, &response)?; + + Ok(( + ConnectionDescriptor { + address, + sender, + receiver, + opener, + sealer, + }, + response, + )) +} + +async fn connect_and_register_with_retry( + scheduler_addresses: &[SocketAddr], + secret_key: Option>, + configuration: &WorkerConfiguration, +) -> crate::Result<(ConnectionDescriptor, WorkerRegistrationResponse)> { + for attempt in 1..=REGISTRATION_MAX_ATTEMPTS { + match connect_and_register(scheduler_addresses, secret_key.clone(), configuration).await { + Ok(result) => return Ok(result), + // A deterministic rejection (wrong secret key, incompatible protocol/role) will + // reproduce identically on every attempt; retrying only delays reporting a + // misconfiguration that will never fix itself. + Err(e @ crate::Error::AuthenticationRejected(_)) => return Err(e), + Err(e) if attempt < REGISTRATION_MAX_ATTEMPTS => { + log::warn!( + "Worker registration attempt {attempt}/{REGISTRATION_MAX_ATTEMPTS} failed: \ + {e}; retrying in {REGISTRATION_RETRY_DELAY:?}" + ); + sleep(REGISTRATION_RETRY_DELAY).await; + } + Err(e) => return Err(e), + } + } + unreachable!("loop always returns on the last attempt") +} + // Maximum time to wait for running tasks to be shutdown when worker ends. const MAX_WAIT_FOR_RUNNING_TASKS_SHUTDOWN: Duration = Duration::from_secs(5); @@ -106,20 +187,25 @@ pub async fn run_worker( )> { let (_listener, port) = start_listener().await?; configuration.listen_address = format!("{}:{}", configuration.hostname, port); - let ConnectionDescriptor { - mut sender, - mut receiver, - mut opener, - mut sealer, - .. - } = connect_to_server_and_authenticate(&scheduler_addresses, secret_key.clone()).await?; - { - let message = ConnectionRegistration::Worker(RegisterWorker { - configuration: configuration.clone(), - }); - let data = serialize(&message)?.into(); - sender.send(seal_message(&mut sealer, data)).await?; - } + let ( + ConnectionDescriptor { + sender, + receiver, + opener, + sealer, + .. + }, + WorkerRegistrationResponse { + worker_id, + other_workers, + resource_names, + resource_rq_map, + server_idle_timeout, + server_uid, + worker_overview_interval_override, + }, + ) = connect_and_register_with_retry(&scheduler_addresses, secret_key.clone(), &configuration) + .await?; let (queue_sender, queue_receiver) = tokio::sync::mpsc::unbounded_channel::(); let heartbeat_interval = configuration.heartbeat_interval; @@ -127,45 +213,29 @@ pub async fn run_worker( let time_limit = configuration.time_limit; let (worker_id, state_ref) = { - match timeout(Duration::from_secs(15), receiver.next()).await { - Ok(Some(data)) => { - let WorkerRegistrationResponse { - worker_id, - other_workers, - resource_names, - resource_rq_map, - server_idle_timeout, - server_uid, - worker_overview_interval_override, - } = open_message(&mut opener, &data?)?; - - sync_worker_configuration(&mut configuration, server_idle_timeout); - - let comm = WorkerComm::new(queue_sender); - let launcher = launcher_setup(&server_uid, worker_id); - let state_ref = WorkerStateRef::new( - comm, - worker_id, - configuration.clone(), - ResourceIdMap::from_vec(resource_names), - resource_rq_map, - launcher, - server_uid, - ); - - { - let mut state = state_ref.get_mut(); - state.worker_overview_interval_override = worker_overview_interval_override; - for worker_info in other_workers { - state.new_worker(worker_info); - } - } + sync_worker_configuration(&mut configuration, server_idle_timeout); + + let comm = WorkerComm::new(queue_sender); + let launcher = launcher_setup(&server_uid, worker_id); + let state_ref = WorkerStateRef::new( + comm, + worker_id, + configuration.clone(), + ResourceIdMap::from_vec(resource_names), + resource_rq_map, + launcher, + server_uid, + ); - (worker_id, state_ref) + { + let mut state = state_ref.get_mut(); + state.worker_overview_interval_override = worker_overview_interval_override; + for worker_info in other_workers { + state.new_worker(worker_info); } - Ok(None) => panic!("Connection closed without receiving registration response"), - Err(_) => panic!("Did not receive worker registration response"), } + + (worker_id, state_ref) }; let local_conn_listener = state_ref.get().lc_state.borrow().create_listener()?;