diff --git a/nativelink-config/src/cas_server.rs b/nativelink-config/src/cas_server.rs index 727830b59..1131699a2 100644 --- a/nativelink-config/src/cas_server.rs +++ b/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)] diff --git a/nativelink-scheduler/src/api_worker_scheduler.rs b/nativelink-scheduler/src/api_worker_scheduler.rs index 2a79d0582..ad1461ec3 100644 --- a/nativelink-scheduler/src/api_worker_scheduler.rs +++ b/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 {} diff --git a/nativelink-scheduler/src/simple_scheduler.rs b/nativelink-scheduler/src/simple_scheduler.rs index 5a72b9f50..47586e46a 100644 --- a/nativelink-scheduler/src/simple_scheduler.rs +++ b/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 {} diff --git a/nativelink-scheduler/src/simple_scheduler_state_manager.rs b/nativelink-scheduler/src/simple_scheduler_state_manager.rs index 66fbcb466..492fa53bd 100644 --- a/nativelink-scheduler/src/simple_scheduler_state_manager.rs +++ b/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 { + 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] diff --git a/nativelink-scheduler/src/worker.rs b/nativelink-scheduler/src/worker.rs index 4e38daf0f..e5bd8d3fc 100644 --- a/nativelink-scheduler/src/worker.rs +++ b/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, } diff --git a/nativelink-scheduler/src/worker_scheduler.rs b/nativelink-scheduler/src/worker_scheduler.rs index f3f24ce60..5a506a89e 100644 --- a/nativelink-scheduler/src/worker_scheduler.rs +++ b/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>; } diff --git a/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs b/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs index 1a3807008..7c09794a1 100644 --- a/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs +++ b/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs @@ -3,21 +3,25 @@ use std::collections::HashMap; use std::sync::Arc; use std::time::SystemTime; +use futures::StreamExt; use mock_instant::thread_local::MockClock; -use nativelink_error::Error; +use nativelink_error::{Code, Error, make_err}; use nativelink_macro::nativelink_test; use nativelink_scheduler::awaited_action_db::AwaitedAction; use nativelink_scheduler::default_scheduler_factory::memory_awaited_action_db_factory; use nativelink_scheduler::simple_scheduler_state_manager::SimpleSchedulerStateManager; use nativelink_scheduler::worker_registry::WorkerRegistry; use nativelink_util::action_messages::{ - ActionInfo, ActionStage, ActionState, ActionUniqueKey, ActionUniqueQualifier, OperationId, - WorkerId, + ActionInfo, ActionResult, ActionStage, ActionState, ActionUniqueKey, ActionUniqueQualifier, + OperationId, WorkerId, }; use nativelink_util::common::DigestInfo; use nativelink_util::digest_hasher::DigestHasherFunc; use nativelink_util::instant_wrapper::MockInstantWrapped; -use nativelink_util::operation_state_manager::{UpdateOperationType, WorkerStateManager}; +use nativelink_util::operation_state_manager::{ + ClientStateManager, MatchingEngineStateManager, OperationFilter, OperationStageFlags, + UpdateOperationType, WorkerStateManager, +}; use tokio::sync::Notify; #[nativelink_test] @@ -64,10 +68,9 @@ fn make_system_time(add_time: u64) -> SystemTime { .unwrap() } -/// An action already assigned to `worker_id` and executing since `started`. -fn executing_action(worker_id: &WorkerId, started: SystemTime) -> AwaitedAction { +fn action_info(started: SystemTime) -> ActionInfo { let action_digest = DigestInfo::zero_digest(); - let action_info = ActionInfo { + ActionInfo { command_digest: action_digest, input_root_digest: action_digest, timeout: Duration::ZERO, @@ -80,7 +83,13 @@ fn executing_action(worker_id: &WorkerId, started: SystemTime) -> AwaitedAction digest_function: DigestHasherFunc::Sha256, digest: action_digest, }), - }; + } +} + +/// An action already assigned to `worker_id` and executing since `started`. +fn executing_action(worker_id: &WorkerId, started: SystemTime) -> AwaitedAction { + let action_info = action_info(started); + let action_digest = action_info.digest(); let operation_id = OperationId::default(); let mut action = AwaitedAction::new(operation_id.clone(), Arc::new(action_info), started); action.worker_set_state( @@ -280,3 +289,108 @@ async fn eventually_times_out_an_orphan() -> Result<(), Error> { ); Ok(()) } + +/// The worker scheduler asks this to decide whether a worker should be told +/// to kill an operation, so it must be false for every way an operation can +/// leave a worker: requeued, reassigned, finished, or gone. +#[nativelink_test] +async fn is_executing_on_worker_follows_the_assignment() -> Result<(), Error> { + MockClock::set_time(Duration::from_secs(NOW_TIME)); + let state_mgr = state_manager(Arc::new(WorkerRegistry::new())); + let worker_id = WorkerId::from(String::from("worker")); + let other_worker_id = WorkerId::from(String::from("other-worker")); + + let _client_listener = state_mgr + .add_action( + OperationId::default(), + Arc::new(action_info(make_system_time(0))), + ) + .await?; + let operation_id = MatchingEngineStateManager::filter_operations( + state_mgr.as_ref(), + OperationFilter { + stages: OperationStageFlags::Queued, + ..Default::default() + }, + ) + .await? + .next() + .await + .expect("the queued operation") + .as_state() + .await? + .0 + .client_operation_id + .clone(); + + // Queued: on nobody. + assert!( + !state_mgr + .is_executing_on_worker(&operation_id, &worker_id) + .await? + ); + + state_mgr + .assign_operation(&operation_id, Ok(&worker_id)) + .await?; + assert!( + state_mgr + .is_executing_on_worker(&operation_id, &worker_id) + .await? + ); + assert!( + !state_mgr + .is_executing_on_worker(&operation_id, &other_worker_id) + .await? + ); + + // Requeued (a timeout), then picked up by another worker. + state_mgr + .assign_operation( + &operation_id, + Err(make_err!(Code::DeadlineExceeded, "timed out")), + ) + .await?; + assert!( + !state_mgr + .is_executing_on_worker(&operation_id, &worker_id) + .await? + ); + state_mgr + .assign_operation(&operation_id, Ok(&other_worker_id)) + .await?; + assert!( + !state_mgr + .is_executing_on_worker(&operation_id, &worker_id) + .await? + ); + assert!( + state_mgr + .is_executing_on_worker(&operation_id, &other_worker_id) + .await? + ); + + // Finished. + state_mgr + .update_operation( + &operation_id, + &other_worker_id, + UpdateOperationType::UpdateWithActionStage(ActionStage::Completed( + ActionResult::default(), + )), + ) + .await?; + assert!( + !state_mgr + .is_executing_on_worker(&operation_id, &other_worker_id) + .await? + ); + + // Never existed. + assert!( + !state_mgr + .is_executing_on_worker(&OperationId::default(), &worker_id) + .await? + ); + Ok(()) +} diff --git a/nativelink-scheduler/tests/simple_scheduler_test.rs b/nativelink-scheduler/tests/simple_scheduler_test.rs index 6425f7e58..8281dd393 100644 --- a/nativelink-scheduler/tests/simple_scheduler_test.rs +++ b/nativelink-scheduler/tests/simple_scheduler_test.rs @@ -36,7 +36,8 @@ use nativelink_proto::com::github::trace_machina::nativelink::events::{ event, request_event, response_event, }; use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{ - ActionResourceUsage, ConnectionResult, StartExecute, UpdateForWorker, update_for_worker, + ActionResourceUsage, ConnectionResult, KillOperationRequest, StartExecute, UpdateForWorker, + update_for_worker, }; use nativelink_scheduler::awaited_action_db::{ AwaitedAction, AwaitedActionDb, AwaitedActionSubscriber, SortedAwaitedAction, @@ -3203,3 +3204,183 @@ async fn failed_final_update_does_not_leak_worker_capacity() -> Result<(), Error Ok(()) } + +/// Setup shared by the kill tests: a single-slot worker running action 1, +/// which a client-timeout sweep then finishes server-side while the worker +/// still holds it. Returns the scheduler, the worker's channel and the +/// operation id the worker was given. +async fn setup_worker_holding_a_finished_operation() -> Result< + ( + Arc, + mpsc::UnboundedReceiver, + OperationId, + Box, + ), + Error, +> { + let worker_id = WorkerId("worker_id".to_string()); + + let task_change_notify = Arc::new(Notify::new()); + let (scheduler, _worker_scheduler) = SimpleScheduler::new_with_callback( + &SimpleSpec { + supported_platform_properties: Some(HashMap::from([( + "cpu_count".to_string(), + PropertyType::Minimum, + )])), + ..Default::default() + }, + memory_awaited_action_db_factory( + // Keep the finished operation around, so its state can be + // inspected after the worker reports. + 100_000, + &task_change_notify.clone(), + MockInstantWrapped::default, + ), + || async move {}, + task_change_notify, + MockInstantWrapped::default, + None, + ); + + let mut rx_from_worker = setup_new_worker( + &scheduler, + worker_id.clone(), + PlatformProperties::new(HashMap::from([( + "cpu_count".to_string(), + PlatformPropertyValue::Minimum(1), + )])), + ) + .await?; + + let action1_listener = setup_action( + &scheduler, + DigestInfo::new([1u8; 32], 512), + HashMap::from([("cpu_count".to_string(), "1".to_string())]), + make_system_time(1), + ) + .await?; + let operation_id = match rx_from_worker.recv().await.unwrap().update { + Some(update_for_worker::Update::StartAction(start_execute)) => start_execute.operation_id, + v => panic!("Expected StartAction, got : {v:?}"), + }; + + // Still executing: nothing to kill. + scheduler.kill_revoked_operations().await?; + assert_eq!( + rx_from_worker.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + ); + + // The client stops sending keepalives past client_action_timeout_s, then + // a sweep over all operations times the executing operation out, marking + // it Completed(DeadlineExceeded) while the worker still runs it. + MockClock::advance(Duration::from_mins(2)); + drop( + scheduler + .filter_operations(OperationFilter::default()) + .await? + .collect::>() + .await, + ); + + Ok(( + scheduler, + rx_from_worker, + OperationId::from(operation_id), + action1_listener, + )) +} + +#[nativelink_test] +async fn revoked_operation_is_killed_once_and_its_report_frees_the_slot() -> 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?; + assert!(logs_contain( + "Operation timed out having no more clients listening" + )); + + // The worker is told to kill it, exactly once. + scheduler.kill_revoked_operations().await?; + assert_eq!( + rx_from_worker.try_recv().unwrap(), + UpdateForWorker { + update: Some(update_for_worker::Update::KillOperationRequest( + KillOperationRequest { + operation_id: operation_id.to_string(), + } + )), + } + ); + scheduler.kill_revoked_operations().await?; + assert_eq!( + rx_from_worker.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + ); + + // Action 2 queues: the worker's only slot is still held by action 1 + // until the worker confirms it stopped. + let _action2_listener = setup_action( + &scheduler, + DigestInfo::new([2u8; 32], 512), + HashMap::from([("cpu_count".to_string(), "1".to_string())]), + make_system_time(2), + ) + .await?; + assert_eq!( + rx_from_worker.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + ); + + // The killed action reports back with the kill's error. That is not + // forwarded to the state manager (it would be "already completed" here, + // or a burnt retry had the operation been requeued), so it is not an + // error, and the slot is freed. + scheduler + .update_action( + &worker_id, + &operation_id, + UpdateOperationType::UpdateWithError(make_err!( + Code::Aborted, + "Command was killed by scheduler" + )), + ) + .await?; + let (action1_state, _) = action1_listener.as_state().await?; + match &action1_state.stage { + ActionStage::Completed(result) => assert_eq!( + result.error.as_ref().map(|err| err.code), + Some(Code::DeadlineExceeded) + ), + stage => panic!("Expected the client timeout result to stand, got : {stage:?}"), + } + + scheduler.do_try_match_for_test().await?; + match rx_from_worker + .try_recv() + .expect("worker should have been sent action 2") + .update + { + Some(update_for_worker::Update::StartAction(_)) => {} + v => panic!("Expected StartAction for the second action, got : {v:?}"), + } + + Ok(()) +} + +#[nativelink_test] +async fn unreachable_worker_is_evicted_when_kill_cannot_be_sent() -> Result<(), Error> { + let (scheduler, rx_from_worker, _operation_id, _action1_listener) = + setup_worker_holding_a_finished_operation().await?; + + // The worker's connection is gone, so the kill cannot be delivered. + drop(rx_from_worker); + let err = scheduler + .kill_revoked_operations() + .await + .expect_err("an undeliverable kill should evict the worker"); + assert_eq!(err.code, Code::Internal); + assert!(logs_contain("Evicting worker from pool")); + + Ok(()) +} diff --git a/nativelink-service/src/worker_api_server.rs b/nativelink-service/src/worker_api_server.rs index ca24d4409..f9dfd6573 100644 --- a/nativelink-service/src/worker_api_server.rs +++ b/nativelink-service/src/worker_api_server.rs @@ -48,6 +48,11 @@ pub type ConnectWorkerStream = pub type NowFn = Box Result + Send + Sync>; +/// How often workers are told to kill operations the scheduler no longer +/// has executing on them, unless overridden by +/// `kill_revoked_operations_interval_s` in the worker API config. +const DEFAULT_KILL_REVOKED_OPERATIONS_INTERVAL_S: u64 = 5; + pub struct WorkerApiServer { scheduler: Arc, now_fn: Arc, @@ -72,28 +77,48 @@ impl WorkerApiServer { rand::rng().fill_bytes(&mut out); out }; + let kill_revoked_enabled = !config.disable_kill_revoked_operations; + let kill_revoked_interval_s = if config.kill_revoked_operations_interval_s == 0 { + DEFAULT_KILL_REVOKED_OPERATIONS_INTERVAL_S + } else { + config.kill_revoked_operations_interval_s + }; for scheduler in schedulers.values() { // This will protect us from holding a reference to the scheduler forever in the // event our ExecutionServer dies. Our scheduler is a weak ref, so the spawn will // eventually see the Arc went away and return. let weak_scheduler = Arc::downgrade(scheduler); background_spawn!("worker_api_server", async move { - let mut ticker = interval(Duration::from_secs(1)); + let mut timeout_ticker = interval(Duration::from_secs(1)); + let mut kill_ticker = interval(Duration::from_secs(kill_revoked_interval_s)); loop { - ticker.tick().await; - let timestamp = SystemTime::now() - .duration_since(UNIX_EPOCH) - .expect("Error: system time is now behind unix epoch"); - match weak_scheduler.upgrade() { - Some(scheduler) => { - if let Err(err) = - scheduler.remove_timedout_workers(timestamp.as_secs()).await - { - error!(?err, "Failed to remove_timedout_workers",); + tokio::select! { + _ = timeout_ticker.tick() => { + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("Error: system time is now behind unix epoch"); + match weak_scheduler.upgrade() { + Some(scheduler) => { + if let Err(err) = + scheduler.remove_timedout_workers(timestamp.as_secs()).await + { + error!(?err, "Failed to remove_timedout_workers",); + } + } + // If we fail to upgrade, our service is probably destroyed, so return. + None => return, + } + } + _ = kill_ticker.tick(), if kill_revoked_enabled => { + match weak_scheduler.upgrade() { + Some(scheduler) => { + if let Err(err) = scheduler.kill_revoked_operations().await { + error!(?err, "Failed to kill_revoked_operations"); + } + } + None => return, } } - // If we fail to upgrade, our service is probably destroyed, so return. - None => return, } } }); diff --git a/nativelink-service/tests/worker_api_server_test.rs b/nativelink-service/tests/worker_api_server_test.rs index 2b047777f..d9d40e117 100644 --- a/nativelink-service/tests/worker_api_server_test.rs +++ b/nativelink-service/tests/worker_api_server_test.rs @@ -125,6 +125,14 @@ impl WorkerStateManager for MockWorkerStateManager { WorkerStateManagerReturns::UpdateOperation(result) => result, } } + + async fn is_executing_on_worker( + &self, + _operation_id: &OperationId, + _worker_id: &WorkerId, + ) -> Result { + Ok(true) + } } struct TestContext { @@ -176,6 +184,8 @@ async fn setup_api_server_with_task_limit( let worker_api_server = WorkerApiServer::new_with_now_fn( &WorkerApiConfig { scheduler: SCHEDULER_NAME.to_string(), + disable_kill_revoked_operations: false, + kill_revoked_operations_interval_s: 0, }, &schedulers, now_fn, diff --git a/nativelink-util/src/operation_state_manager.rs b/nativelink-util/src/operation_state_manager.rs index 6dd140a63..ea6661ead 100644 --- a/nativelink-util/src/operation_state_manager.rs +++ b/nativelink-util/src/operation_state_manager.rs @@ -148,6 +148,14 @@ pub trait WorkerStateManager: Sync + Send + MetricsComponent { worker_id: &WorkerId, update: UpdateOperationType, ) -> Result<(), Error>; + + /// Whether the operation is still executing on this worker. False once + /// it has finished, been requeued or reassigned, or no longer exists. + async fn is_executing_on_worker( + &self, + operation_id: &OperationId, + worker_id: &WorkerId, + ) -> Result; } #[async_trait]