diff --git a/nativelink-config/src/schedulers.rs b/nativelink-config/src/schedulers.rs index e0e66b04c..3f41741ea 100644 --- a/nativelink-config/src/schedulers.rs +++ b/nativelink-config/src/schedulers.rs @@ -151,6 +151,16 @@ pub struct SimpleSpec { #[serde(default, deserialize_with = "convert_duration_with_shellexpand")] pub max_action_executing_timeout_s: u64, + /// Evict a worker that has not reported back on an operation it was + /// told to kill within this many seconds. A healthy worker + /// acknowledges a kill in moments; one that cannot is wedged, and its + /// keepalives would otherwise keep `worker_timeout_s` from ever firing + /// while the dead operation holds its slot. Eviction requeues the + /// worker's other operations. + /// Default: 60 seconds + #[serde(default, deserialize_with = "convert_duration_with_shellexpand")] + pub unacknowledged_kill_timeout_s: u64, + /// If a job returns an internal error or times out this many times when /// attempting to run on a worker the scheduler will return the last error /// to the client. Jobs will be retried and this configuration is to help diff --git a/nativelink-scheduler/src/api_worker_scheduler.rs b/nativelink-scheduler/src/api_worker_scheduler.rs index ad1461ec3..45cb42cc0 100644 --- a/nativelink-scheduler/src/api_worker_scheduler.rs +++ b/nativelink-scheduler/src/api_worker_scheduler.rs @@ -601,6 +601,10 @@ pub struct ApiWorkerScheduler { help = "Timeout of how long to evict workers if no response in this given amount of time in seconds." )] worker_timeout_s: u64, + #[metric( + help = "How long a sent kill may go unacknowledged before the worker is evicted, in seconds." + )] + unacknowledged_kill_timeout_s: u64, /// Shared worker registry for checking worker liveness. worker_registry: SharedWorkerRegistry, @@ -613,12 +617,14 @@ pub struct ApiWorkerScheduler { } impl ApiWorkerScheduler { + #[expect(clippy::too_many_arguments)] pub fn new( worker_state_manager: Arc, platform_property_manager: Arc, allocation_strategy: WorkerAllocationStrategy, worker_change_notify: Arc, worker_timeout_s: u64, + unacknowledged_kill_timeout_s: u64, worker_registry: SharedWorkerRegistry, maybe_origin_event_tx: Option>, ) -> Arc { @@ -634,6 +640,7 @@ impl ApiWorkerScheduler { }), platform_property_manager, worker_timeout_s, + unacknowledged_kill_timeout_s, worker_registry, metrics: Arc::new(SchedulerMetrics::default()), maybe_origin_event_tx, @@ -905,20 +912,40 @@ impl WorkerScheduler for ApiWorkerScheduler { let now = UNIX_EPOCH + Duration::from_secs(now_timestamp); let timeout_threshold = now_timestamp.saturating_sub(self.worker_timeout_s); - let workers_to_check: Vec<(WorkerId, bool)> = { + let workers_to_check: Vec<(WorkerId, bool, bool)> = { let inner = self.inner.lock().await; inner .workers .iter() .map(|(worker_id, worker)| { let local_alive = worker.last_update_timestamp > timeout_threshold; - (worker_id.clone(), local_alive) + let kill_overdue = worker.running_action_infos.values().any(|info| { + info.kill_requested_at.is_some_and(|at| { + now_timestamp.saturating_sub(at) > self.unacknowledged_kill_timeout_s + }) + }); + (worker_id.clone(), local_alive, kill_overdue) }) .collect() }; let mut worker_ids_to_remove = Vec::new(); - for (worker_id, local_alive) in workers_to_check { + for (worker_id, local_alive, kill_overdue) in workers_to_check { + // A healthy worker acknowledges a kill in moments; one that + // cannot is wedged, and its keepalives keep the liveness checks + // below from ever firing while the dead operation holds its + // slot (the nativelink#2672 symptom). Evicting requeues its + // other operations. + if kill_overdue { + warn!( + ?worker_id, + unacknowledged_kill_timeout_s = self.unacknowledged_kill_timeout_s, + "Worker did not acknowledge a kill in time, removing from pool" + ); + worker_ids_to_remove.push((worker_id, true)); + continue; + } + if local_alive { continue; } @@ -936,7 +963,7 @@ impl WorkerScheduler for ApiWorkerScheduler { timeout_threshold, "Worker timed out - neither local nor registry shows alive" ); - worker_ids_to_remove.push(worker_id); + worker_ids_to_remove.push((worker_id, false)); } } @@ -947,20 +974,21 @@ impl WorkerScheduler for ApiWorkerScheduler { let mut inner = self.inner.lock().await; let mut result = Ok(()); - for worker_id in &worker_ids_to_remove { - warn!(?worker_id, "Worker timed out, removing from pool"); - result = result.merge( - inner - .immediate_evict_worker( - worker_id, - make_err!( - Code::Internal, - "Worker {worker_id} timed out, removing from pool" - ), - false, - ) - .await, - ); + for (worker_id, kill_overdue) in &worker_ids_to_remove { + let err = if *kill_overdue { + make_err!( + Code::Internal, + "Worker {worker_id} did not acknowledge a kill within {}s, removing from pool", + self.unacknowledged_kill_timeout_s + ) + } else { + warn!(?worker_id, "Worker timed out, removing from pool"); + make_err!( + Code::Internal, + "Worker {worker_id} timed out, removing from pool" + ) + }; + result = result.merge(inner.immediate_evict_worker(worker_id, err, false).await); } result @@ -982,13 +1010,15 @@ impl WorkerScheduler for ApiWorkerScheduler { .running_action_infos .iter() // An operation already told to die is never swept - // again, and nothing re-sends or times out the kill - // itself. A live worker that drops or ignores the - // kill request therefore holds the slot until - // keepalive-timeout eviction (remove_timedout_workers) - // reclaims the whole worker; that is the sole - // recovery path for an unresponsive-but-alive worker. - .filter(|(_, pending_action_info)| !pending_action_info.kill_requested) + // again; the kill itself is not re-sent. A worker + // that received the kill but never reports back is + // evicted by remove_timedout_workers once the kill + // has gone unacknowledged longer than + // `unacknowledged_kill_timeout_s`; a dead worker by + // the ordinary keepalive timeout. + .filter(|(_, pending_action_info)| { + pending_action_info.kill_requested_at.is_none() + }) .map(|(operation_id, _)| (worker_id.clone(), operation_id.clone())) }) .collect(); diff --git a/nativelink-scheduler/src/simple_scheduler.rs b/nativelink-scheduler/src/simple_scheduler.rs index 47586e46a..d2589cbb5 100644 --- a/nativelink-scheduler/src/simple_scheduler.rs +++ b/nativelink-scheduler/src/simple_scheduler.rs @@ -58,6 +58,10 @@ use crate::worker_scheduler::WorkerScheduler; /// If this changes, remember to change the documentation in the config. const DEFAULT_WORKER_TIMEOUT_S: u64 = 5; +/// Default timeout for a sent kill to be acknowledged in seconds. +/// If this changes, remember to change the documentation in the config. +const DEFAULT_UNACKNOWLEDGED_KILL_TIMEOUT_S: u64 = 60; + /// Mark operations as completed with error if no client has updated them /// within this duration. /// If this changes, remember to change the documentation in the config. @@ -513,6 +517,11 @@ impl SimpleScheduler { max_job_retries = DEFAULT_MAX_JOB_RETRIES; } + let mut unacknowledged_kill_timeout_s = spec.unacknowledged_kill_timeout_s; + if unacknowledged_kill_timeout_s == 0 { + unacknowledged_kill_timeout_s = DEFAULT_UNACKNOWLEDGED_KILL_TIMEOUT_S; + } + let worker_change_notify = Arc::new(Notify::new()); // Create shared worker registry for single heartbeat per worker. @@ -538,6 +547,7 @@ impl SimpleScheduler { spec.allocation_strategy, worker_change_notify.clone(), worker_timeout_s, + unacknowledged_kill_timeout_s, worker_registry, maybe_origin_event_tx.clone(), ); diff --git a/nativelink-scheduler/src/worker.rs b/nativelink-scheduler/src/worker.rs index e5bd8d3fc..5eb552957 100644 --- a/nativelink-scheduler/src/worker.rs +++ b/nativelink-scheduler/src/worker.rs @@ -65,10 +65,14 @@ pub enum WorkerUpdate { pub struct PendingActionInfoData { #[metric] pub action_info: ActionInfoWithProps, - /// Set once the worker has been told to kill this operation. Its later - /// report then only settles the worker's own bookkeeping. - #[metric(help = "If the worker has been asked to kill this operation.")] - pub kill_requested: bool, + /// When the worker was told to kill this operation, recorded as the + /// worker's last-seen timestamp at send time; keepalives keep that + /// within the worker timeout of now, so a deadline computed from it + /// fires at most one worker timeout late, never early. `None` until a + /// kill is sent. The operation's later report then only settles the + /// worker's own bookkeeping. + #[metric(help = "When the worker was asked to kill this operation.")] + pub kill_requested_at: Option, } /// Represents a connection to a worker and used as the medium to @@ -200,6 +204,7 @@ impl Worker { send_msg_to_worker(&self.tx, update_for_worker::Update::Disconnect(())) } WorkerUpdate::KillOperation(operation_id) => { + let last_seen = self.last_update_timestamp; let pending_action_info = self .running_action_infos .get_mut(&operation_id) @@ -211,7 +216,7 @@ impl Worker { })?; // Set before the send so a racing update_action cannot slip // through in between. - pending_action_info.kill_requested = true; + pending_action_info.kill_requested_at = Some(last_seen); self.metrics.kill_operation.inc(); send_msg_to_worker( &self.tx, @@ -227,7 +232,7 @@ impl Worker { pub(crate) fn is_kill_requested(&self, operation_id: &OperationId) -> bool { self.running_action_infos .get(operation_id) - .is_some_and(|pending_action_info| pending_action_info.kill_requested) + .is_some_and(|pending_action_info| pending_action_info.kill_requested_at.is_some()) } pub fn keep_alive(&mut self) -> Result<(), Error> { @@ -268,7 +273,7 @@ impl Worker { operation_id, PendingActionInfoData { action_info, - kill_requested: false, + kill_requested_at: None, }, ); diff --git a/nativelink-scheduler/tests/simple_scheduler_test.rs b/nativelink-scheduler/tests/simple_scheduler_test.rs index 8281dd393..f588cae73 100644 --- a/nativelink-scheduler/tests/simple_scheduler_test.rs +++ b/nativelink-scheduler/tests/simple_scheduler_test.rs @@ -3384,3 +3384,38 @@ async fn unreachable_worker_is_evicted_when_kill_cannot_be_sent() -> Result<(), Ok(()) } + +#[nativelink_test] +async fn live_worker_that_never_acknowledges_a_kill_is_evicted() -> Result<(), Error> { + let worker_id = WorkerId("worker_id".to_string()); + let (scheduler, mut rx_from_worker, _operation_id, _action1_listener) = + setup_worker_holding_a_finished_operation().await?; + + // The kill is delivered, but the worker is wedged (nativelink#2672): it + // keeps its keepalives up and never reports the operation. + scheduler.kill_revoked_operations().await?; + match rx_from_worker.try_recv().unwrap().update { + Some(update_for_worker::Update::KillOperationRequest(_)) => {} + v => panic!("Expected KillOperationRequest, got : {v:?}"), + } + + // Inside the acknowledgement window the fresh keepalives shield it from + // the ordinary worker timeout and nothing is evicted. + scheduler + .worker_keep_alive_received(&worker_id, NOW_TIME + 30) + .await?; + scheduler.remove_timedout_workers(NOW_TIME + 30).await?; + assert!(!logs_contain("Evicting worker from pool")); + + // Past the window the worker is evicted despite looking alive; its + // keepalives are exactly what would otherwise let the dead operation + // hold the slot forever. + scheduler + .worker_keep_alive_received(&worker_id, NOW_TIME + 61) + .await?; + drop(scheduler.remove_timedout_workers(NOW_TIME + 61).await); + assert!(logs_contain("did not acknowledge a kill in time")); + assert!(logs_contain("Evicting worker from pool")); + + Ok(()) +} diff --git a/nativelink-service/tests/worker_api_server_test.rs b/nativelink-service/tests/worker_api_server_test.rs index d9d40e117..108d76b1a 100644 --- a/nativelink-service/tests/worker_api_server_test.rs +++ b/nativelink-service/tests/worker_api_server_test.rs @@ -175,6 +175,7 @@ async fn setup_api_server_with_task_limit( WorkerAllocationStrategy::default(), tasks_or_worker_change_notify, worker_timeout, + 60, // unacknowledged_kill_timeout_s worker_registry, None, );