Evict a worker that never acknowledges a kill

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

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<dyn WorkerStateManager>,
         platform_property_manager: Arc<PlatformPropertyManager>,
         allocation_strategy: WorkerAllocationStrategy,
         worker_change_notify: Arc<Notify>,
         worker_timeout_s: u64,
+        unacknowledged_kill_timeout_s: u64,
         worker_registry: SharedWorkerRegistry,
         maybe_origin_event_tx: Option<mpsc::Sender<OriginEvent>>,
     ) -> Arc<Self> {
@@ -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();

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

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<WorkerTimestamp>,
 }
 
 /// 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,
                     },
                 );
 

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(())
+}

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,
     );