Tell workers to kill operations the scheduler no longer has executing on them

nativelink-config/src/cas_server.rs

@@ -389,6 +389,28 @@ pub struct WorkerApiConfig {
     /// The scheduler name referenced in the `schedulers` map in the main config.
     #[serde(deserialize_with = "convert_string_with_shellexpand")]
     pub scheduler: SchedulerRefName,
+
+    /// Disable the periodic sweep that tells workers to kill operations
+    /// the scheduler no longer has executing on them (for example because
+    /// every client disconnected before the action finished). With the
+    /// sweep disabled such orphaned actions run to completion and still
+    /// warm the action cache, which can be preferable for deployments
+    /// with long or expensive actions whose clients merely retry.
+    ///
+    /// Default: false (the sweep runs)
+    #[serde(default)]
+    pub disable_kill_revoked_operations: bool,
+
+    /// How often, in seconds, the scheduler checks for operations that
+    /// are still running on a worker but are no longer executing
+    /// according to the state manager, and tells the worker to kill
+    /// them. Each pass costs one state-manager lookup per running
+    /// operation, so store-backed deployments with many concurrent
+    /// actions may want a longer interval.
+    ///
+    /// Default: 5 (0 uses the default)
+    #[serde(default)]
+    pub kill_revoked_operations_interval_s: u64,
 }
 
 #[derive(Deserialize, Serialize, Debug, Default)]

nativelink-scheduler/src/api_worker_scheduler.rs

@@ -19,6 +19,7 @@ use std::sync::Arc;
 use std::time::{Instant, UNIX_EPOCH};
 
 use async_lock::Mutex;
+use futures::{StreamExt, future};
 use lru::LruCache;
 use nativelink_config::schedulers::WorkerAllocationStrategy;
 use nativelink_error::{Code, Error, ResultExt, error_if, make_err, make_input_err};
@@ -42,7 +43,11 @@ use nativelink_util::platform_properties::PlatformProperties;
 use nativelink_util::shutdown_guard::ShutdownGuard;
 use tokio::sync::{Notify, mpsc};
 use tonic::async_trait;
-use tracing::{error, info, trace, warn};
+use tracing::{debug, error, info, trace, warn};
+
+/// How many state-manager lookups `kill_revoked_operations` has in flight
+/// at once while checking which running operations were revoked.
+const MAX_CONCURRENT_REVOKED_CHECKS: usize = 32;
 use uuid::Uuid;
 
 /// Metrics for tracking scheduler performance.
@@ -374,19 +379,29 @@ impl ApiWorkerSchedulerImpl {
         };
 
         // Update the operation in the worker state manager.
-        let update_operation_res = self
-            .worker_state_manager
-            .update_operation(operation_id, worker_id, update)
-            .await
-            .err_tip(|| "in update_operation on SimpleScheduler::update_action");
-        if let Err(err) = &update_operation_res {
-            error!(
+        let update_operation_res = if worker.is_kill_requested(operation_id) {
+            debug!(
                 %operation_id,
                 ?worker_id,
-                ?err,
-                "Failed to update_operation on update_action"
+                "Ignoring update for operation the worker was told to kill"
             );
-        }
+            Ok(())
+        } else {
+            let update_operation_res = self
+                .worker_state_manager
+                .update_operation(operation_id, worker_id, update)
+                .await
+                .err_tip(|| "in update_operation on SimpleScheduler::update_action");
+            if let Err(err) = &update_operation_res {
+                error!(
+                    %operation_id,
+                    ?worker_id,
+                    ?err,
+                    "Failed to update_operation on update_action"
+                );
+            }
+            update_operation_res
+        };
 
         if !is_finished {
             return update_operation_res;
@@ -480,6 +495,49 @@ impl ApiWorkerSchedulerImpl {
         }
     }
 
+    /// Tells the worker to kill an operation it is still running but the
+    /// state manager no longer has executing on it. A worker that cannot be
+    /// reached is evicted, the same as for a failed run request.
+    async fn worker_notify_kill_operation(
+        &mut self,
+        worker_id: &WorkerId,
+        operation_id: OperationId,
+    ) -> Result<(), Error> {
+        let Some(worker) = self.workers.get_mut(worker_id) else {
+            // Gone between the snapshot and now; its actions were requeued.
+            return Ok(());
+        };
+        // Already told, or finished in the meantime; nothing more to send.
+        if !worker.running_action_infos.contains_key(&operation_id)
+            || worker.is_kill_requested(&operation_id)
+        {
+            return Ok(());
+        }
+        info!(
+            ?worker_id,
+            %operation_id,
+            "Killing operation the state manager no longer has executing on this worker"
+        );
+        if let Err(err) = worker
+            .notify_update(WorkerUpdate::KillOperation(operation_id.clone()))
+            .await
+        {
+            warn!(
+                ?worker_id,
+                %operation_id,
+                ?err,
+                "Worker command failed, removing worker"
+            );
+            let err = make_err!(
+                Code::Internal,
+                "Worker command failed, removing worker {worker_id} -- {err:?}",
+            );
+            return Result::<(), _>::Err(err.clone())
+                .merge(self.immediate_evict_worker(worker_id, err, true).await);
+        }
+        Ok(())
+    }
+
     /// Evicts the worker from the pool and puts items back into the queue if anything was being executed on it.
     async fn immediate_evict_worker(
         &mut self,
@@ -912,6 +970,98 @@ impl WorkerScheduler for ApiWorkerScheduler {
         let mut inner = self.inner.lock().await;
         inner.set_drain_worker(worker_id, is_draining).await
     }
+
+    async fn kill_revoked_operations(&self) -> Result<(), Error> {
+        let (worker_state_manager, running) = {
+            let inner = self.inner.lock().await;
+            let running: Vec<(WorkerId, OperationId)> = inner
+                .workers
+                .iter()
+                .flat_map(|(worker_id, worker)| {
+                    worker
+                        .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)
+                        .map(|(operation_id, _)| (worker_id.clone(), operation_id.clone()))
+                })
+                .collect();
+            (inner.worker_state_manager.clone(), running)
+        };
+
+        // On store-backed deployments each check is a network round-trip,
+        // so run them lock-free with bounded concurrency.
+        let revoked: Vec<(WorkerId, OperationId)> = futures::stream::iter(running)
+            .map(|(worker_id, operation_id)| {
+                let worker_state_manager = worker_state_manager.clone();
+                async move {
+                    match worker_state_manager
+                        .is_executing_on_worker(&operation_id, &worker_id)
+                        .await
+                    {
+                        Ok(true) => None,
+                        Ok(false) => Some((worker_id, operation_id)),
+                        // Only kill on positive evidence; try again next pass.
+                        Err(err) => {
+                            warn!(
+                                ?worker_id,
+                                %operation_id,
+                                ?err,
+                                "Could not check whether operation is still executing on worker"
+                            );
+                            None
+                        }
+                    }
+                }
+            })
+            .buffer_unordered(MAX_CONCURRENT_REVOKED_CHECKS)
+            .filter_map(future::ready)
+            .collect()
+            .await;
+
+        if revoked.is_empty() {
+            return Ok(());
+        }
+
+        // Re-check lock-free so the scheduler mutex is never held across
+        // store I/O; the checks above may be stale by the time we get here.
+        // The remaining TOCTOU window is fine: worker_notify_kill_operation
+        // re-guards with contains_key + is_kill_requested under the lock.
+        let mut confirmed = Vec::new();
+        for (worker_id, operation_id) in revoked {
+            // A result of Ok(true) or Err(_) means the operation was
+            // reassigned to this worker after the first check, or the state
+            // manager went quiet: nothing to kill on this pass.
+            let revoked = worker_state_manager
+                .is_executing_on_worker(&operation_id, &worker_id)
+                .await
+                .is_ok_and(|executing| !executing);
+            if revoked {
+                confirmed.push((worker_id, operation_id));
+            }
+        }
+
+        if confirmed.is_empty() {
+            return Ok(());
+        }
+
+        let mut inner = self.inner.lock().await;
+        let mut result = Ok(());
+        for (worker_id, operation_id) in confirmed {
+            result = result.merge(
+                inner
+                    .worker_notify_kill_operation(&worker_id, operation_id)
+                    .await,
+            );
+        }
+        result
+    }
 }
 
 impl RootMetricsComponent for ApiWorkerScheduler {}

nativelink-scheduler/src/simple_scheduler.rs

@@ -792,6 +792,10 @@ impl WorkerScheduler for SimpleScheduler {
             .set_drain_worker(worker_id, is_draining)
             .await
     }
+
+    async fn kill_revoked_operations(&self) -> Result<(), Error> {
+        self.worker_scheduler.kill_revoked_operations().await
+    }
 }
 
 impl RootMetricsComponent for SimpleScheduler {}

nativelink-scheduler/src/simple_scheduler_state_manager.rs

@@ -1183,6 +1183,35 @@ where
         self.inner_update_operation(operation_id, Some(worker_id), update)
             .await
     }
+
+    async fn is_executing_on_worker(
+        &self,
+        operation_id: &OperationId,
+        worker_id: &WorkerId,
+    ) -> Result<bool, Error> {
+        let Some(subscriber) = self
+            .action_db
+            .get_by_operation_id(operation_id)
+            .await
+            .err_tip(|| "In SimpleSchedulerStateManager::is_executing_on_worker")?
+        else {
+            return Ok(false);
+        };
+        let awaited_action = match subscriber.borrow().await {
+            Ok(awaited_action) => awaited_action,
+            // Store-backed dbs hand out a subscriber for any id and only
+            // discover the operation is gone on read.
+            Err(err) if err.code == Code::NotFound => return Ok(false),
+            Err(err) => {
+                return Err(err)
+                    .err_tip(|| "In SimpleSchedulerStateManager::is_executing_on_worker");
+            }
+        };
+        Ok(
+            matches!(awaited_action.state().stage, ActionStage::Executing)
+                && awaited_action.worker_id() == Some(worker_id),
+        )
+    }
 }
 
 #[async_trait]

nativelink-scheduler/src/worker.rs

@@ -20,7 +20,7 @@ use std::time::{SystemTime, UNIX_EPOCH};
 use nativelink_error::{Code, Error, ResultExt};
 use nativelink_metric::MetricsComponent;
 use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{
-    ConnectionResult, StartExecute, UpdateForWorker, update_for_worker,
+    ConnectionResult, KillOperationRequest, StartExecute, UpdateForWorker, update_for_worker,
 };
 use nativelink_util::action_messages::{ActionInfo, OperationId, WorkerId};
 use nativelink_util::metrics_utils::{AsyncCounterWrapper, CounterWithTime, FuncCounterWrapper};
@@ -56,12 +56,19 @@ pub enum WorkerUpdate {
 
     /// Request that the worker is no longer in the pool and may discard any jobs.
     Disconnect,
+
+    /// Requests that the worker stop executing this operation.
+    KillOperation(OperationId),
 }
 
 #[derive(Debug, MetricsComponent)]
 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,
 }
 
 /// Represents a connection to a worker and used as the medium to
@@ -164,6 +171,7 @@ impl Worker {
                 run_action: AsyncCounterWrapper::default(),
                 keep_alive: FuncCounterWrapper::default(),
                 notify_disconnect: CounterWithTime::default(),
+                kill_operation: CounterWithTime::default(),
             }),
         }
     }
@@ -191,9 +199,37 @@ impl Worker {
                 self.metrics.notify_disconnect.inc();
                 send_msg_to_worker(&self.tx, update_for_worker::Update::Disconnect(()))
             }
+            WorkerUpdate::KillOperation(operation_id) => {
+                let pending_action_info = self
+                    .running_action_infos
+                    .get_mut(&operation_id)
+                    .err_tip(|| {
+                        format!(
+                            "Worker {} asked to kill operation {operation_id} that is not running on it",
+                            self.id
+                        )
+                    })?;
+                // Set before the send so a racing update_action cannot slip
+                // through in between.
+                pending_action_info.kill_requested = true;
+                self.metrics.kill_operation.inc();
+                send_msg_to_worker(
+                    &self.tx,
+                    update_for_worker::Update::KillOperationRequest(KillOperationRequest {
+                        operation_id: operation_id.to_string(),
+                    }),
+                )
+            }
         }
     }
 
+    /// Whether the worker has been told to kill this operation.
+    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)
+    }
+
     pub fn keep_alive(&mut self) -> Result<(), Error> {
         let tx = &mut self.tx;
         let id = &self.id;
@@ -228,7 +264,13 @@ impl Worker {
                     worker_platform_properties,
                     &action_info.platform_properties,
                 );
-                running_action_infos.insert(operation_id, PendingActionInfoData { action_info });
+                running_action_infos.insert(
+                    operation_id,
+                    PendingActionInfoData {
+                        action_info,
+                        kill_requested: false,
+                    },
+                );
 
                 send_msg_to_worker(tx, update_for_worker::Update::StartAction(start_execute))
             })
@@ -317,4 +359,6 @@ struct Metrics {
     keep_alive: FuncCounterWrapper,
     #[metric(help = "The number of notify_disconnect sent to this worker.")]
     notify_disconnect: CounterWithTime,
+    #[metric(help = "The number of kill_operation sent to this worker.")]
+    kill_operation: CounterWithTime,
 }

nativelink-scheduler/src/worker_scheduler.rs

@@ -70,4 +70,9 @@ pub trait WorkerScheduler: Sync + Send + Unpin + RootMetricsComponent + 'static
 
     /// Sets if the worker is draining or not.
     async fn set_drain_worker(&self, worker_id: &WorkerId, is_draining: bool) -> Result<(), Error>;
+
+    /// Tells workers to kill operations they are still running but the
+    /// scheduler has finished, requeued or dropped without them (client
+    /// timeouts and cancellations, execution deadlines, retries elsewhere).
+    async fn kill_revoked_operations(&self) -> Result<(), Error>;
 }