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()?;