Idle timeout checked in server

Author: spirali

crates/tako/src/internal/messages/worker.rs

@@ -71,7 +71,6 @@ pub enum ToWorkerMessage {
     CancelTasks(TaskIdsMsg),
     NewWorker(NewWorkerMsg),
     LostWorker(WorkerId),
-    SetReservation(bool),
     /// Override the internally set overview interval with a new duration
     /// if it is **disabled** on the worker.
     /// If the worker has already enabled overview interval, then this does nothing.

crates/tako/src/internal/scheduler/state.rs

@@ -189,7 +189,7 @@ impl SchedulerState {
                     );
                     for worker_id in &worker_ids[1..] {
                         let worker = worker_map.get_worker_mut(*worker_id);
-                        worker.set_reservation(true, comm);
+                        worker.set_reservation(true);
                     }
                 }
                 (
@@ -236,7 +236,7 @@ impl SchedulerState {
                     })
                     .unwrap_or(true);
                 if unreserve {
-                    worker.set_reservation(false, comm);
+                    worker.set_reservation(false);
                 }
             }
         }

crates/tako/src/internal/server/rpc.rs

@@ -188,7 +188,7 @@ async fn worker_rpc_loop(
                 // Idle timeout might be smaller than 500ms in tests.
                 #[allow(clippy::manual_clamp)]
                 heartbeat_interval
-                    .min(idle_timeout.unwrap_or(heartbeat_interval))
+                    .min(idle_timeout.map(|t| t / 16).unwrap_or(heartbeat_interval))
                     // Sanity check that interval is not too short
                     .max(Duration::from_millis(500)),
             )
@@ -204,6 +204,15 @@ async fn worker_rpc_loop(
                 log::debug!("Heartbeat not arrived, worker={}", worker.id);
                 break LostWorkerReason::HeartbeatLost;
             }
+
+            if let Some(timeout) = idle_timeout {
+                if worker.is_free()
+                    && !worker.is_reserved()
+                    && worker.idle_timestamp + timeout < now
+                {
+                    break LostWorkerReason::IdleTimeout;
+                }
+            }
         }
     };
 

crates/tako/src/internal/server/worker.rs

@@ -5,14 +5,12 @@ use crate::internal::common::Set;
 use crate::internal::common::resources::TimeRequest;
 use crate::internal::common::resources::map::ResourceMap;
 use crate::internal::common::resources::{ResourceRequest, ResourceRequestVariants};
-use crate::internal::messages::worker::ToWorkerMessage;
-use crate::internal::server::comm::Comm;
 use crate::internal::server::task::Task;
 use crate::internal::server::taskmap::TaskMap;
 use crate::internal::server::workerload::{ResourceRequestLowerBound, WorkerLoad, WorkerResources};
 use crate::internal::worker::configuration::WorkerConfiguration;
 use crate::{TaskId, WorkerId};
-use std::time::Duration;
+use std::time::{Duration, Instant};
 
 bitflags::bitflags! {
     pub(crate) struct WorkerFlags: u32 {
@@ -55,12 +53,15 @@ pub struct Worker {
     pub(crate) resources: WorkerResources,
     pub(crate) flags: WorkerFlags,
     // When the worker will be terminated
-    pub(crate) termination_time: Option<std::time::Instant>,
+    pub(crate) termination_time: Option<Instant>,
 
     pub(crate) mn_task: Option<MultiNodeTaskAssignment>,
 
+    // Saved timestamp when a worker is put into an idle state
+    pub(crate) idle_timestamp: Instant,
+
     // COLD DATA move it into a box (?)
-    pub(crate) last_heartbeat: std::time::Instant,
+    pub(crate) last_heartbeat: Instant,
     pub(crate) configuration: WorkerConfiguration,
 }
 
@@ -123,12 +124,12 @@ impl Worker {
 
     pub fn reset_mn_task(&mut self) {
         self.mn_task = None;
+        self.idle_timestamp = Instant::now();
     }
 
-    pub fn set_reservation(&mut self, value: bool, comm: &mut impl Comm) {
+    pub fn set_reservation(&mut self, value: bool) {
         if self.is_reserved() != value {
             self.flags.set(WorkerFlags::RESERVED, value);
-            comm.send_worker_message(self.id, &ToWorkerMessage::SetReservation(value))
         }
     }
 
@@ -148,6 +149,9 @@ impl Worker {
 
     pub fn remove_sn_task(&mut self, task: &Task) {
         assert!(self.sn_tasks.remove(&task.id));
+        if self.sn_tasks.is_empty() {
+            self.idle_timestamp = Instant::now();
+        }
         self.sn_load
             .remove_request(task.id, &task.configuration.resources, &self.resources);
     }
@@ -205,7 +209,7 @@ impl Worker {
         self.flags.contains(WorkerFlags::PARKED)
     }
 
-    pub fn is_capable_to_run(&self, request: &ResourceRequest, now: std::time::Instant) -> bool {
+    pub fn is_capable_to_run(&self, request: &ResourceRequest, now: Instant) -> bool {
         self.has_time_to_run(request.min_time(), now)
             && self.resources.is_capable_to_run_request(request)
     }
@@ -221,7 +225,7 @@ impl Worker {
     }
 
     // Returns None if there is no time limit for a worker or time limit was passed
-    pub fn remaining_time(&self, now: std::time::Instant) -> Option<Duration> {
+    pub fn remaining_time(&self, now: Instant) -> Option<Duration> {
         self.termination_time.map(|time| {
             if time < now {
                 Duration::default()
@@ -231,19 +235,15 @@ impl Worker {
         })
     }
 
-    pub fn has_time_to_run(&self, time_request: TimeRequest, now: std::time::Instant) -> bool {
+    pub fn has_time_to_run(&self, time_request: TimeRequest, now: Instant) -> bool {
         if let Some(time) = self.termination_time {
             now + time_request < time
         } else {
             true
         }
     }
 
-    pub fn has_time_to_run_for_rqv(
-        &self,
-        rqv: &ResourceRequestVariants,
-        now: std::time::Instant,
-    ) -> bool {
+    pub fn has_time_to_run_for_rqv(&self, rqv: &ResourceRequestVariants, now: Instant) -> bool {
         if self.termination_time.is_none() {
             return true;
         }
@@ -280,6 +280,7 @@ impl Worker {
             flags: WorkerFlags::empty(),
             last_heartbeat: now,
             mn_task: None,
+            idle_timestamp: now,
         }
     }
 }

crates/tako/src/internal/tests/test_reactor.rs

@@ -749,14 +749,6 @@ fn test_task_mn_cancel() {
     let mut scheduler = create_test_scheduler();
     scheduler.run_scheduling(&mut core, &mut comm);
     core.sanity_check();
-    assert!(matches!(
-        &comm.take_worker_msgs(100, 1)[0],
-        &ToWorkerMessage::SetReservation(false)
-    ));
-    assert!(matches!(
-        &comm.take_worker_msgs(101, 1)[0],
-        &ToWorkerMessage::SetReservation(false)
-    ));
     comm.emptiness_check();
 
     assert!(core.find_task(1.into()).is_none());

crates/tako/src/internal/tests/test_scheduler_mn.rs

@@ -44,19 +44,15 @@ fn check_worker_status_change(s1: WorkerStatus, s2: WorkerStatus, ms: &[ToWorker
         (WorkerStatus::Root, WorkerStatus::Root) | (WorkerStatus::None, WorkerStatus::Root) => {
             assert!(matches!(ms, &[ToWorkerMessage::ComputeTask(_)]))
         }
-        (WorkerStatus::NonRoot, WorkerStatus::Root) => assert!(matches!(
-            ms,
-            &[
-                ToWorkerMessage::ComputeTask(_),
-                ToWorkerMessage::SetReservation(false)
-            ]
-        )),
+        (WorkerStatus::NonRoot, WorkerStatus::Root) => {
+            assert!(matches!(ms, &[ToWorkerMessage::ComputeTask(_),]))
+        }
         (WorkerStatus::NonRoot, WorkerStatus::None) => {
-            assert!(matches!(ms, &[ToWorkerMessage::SetReservation(false)]))
+            assert!(matches!(ms, &[]))
         }
         (WorkerStatus::None, WorkerStatus::NonRoot)
         | (WorkerStatus::Root, WorkerStatus::NonRoot) => {
-            assert!(matches!(ms, &[ToWorkerMessage::SetReservation(true)]))
+            assert!(matches!(ms, &[]))
         }
         (WorkerStatus::NonRoot, WorkerStatus::NonRoot)
         | (WorkerStatus::Root, WorkerStatus::None)
@@ -86,25 +82,21 @@ fn test_schedule_mn_simple() {
     scheduler.run_scheduling(&mut core, &mut comm);
     core.sanity_check();
 
-    let test_mn_task = |task: &Task, comm: &mut TestComm, reservation: bool| -> Vec<WorkerId> {
+    let test_mn_task = |task: &Task, comm: &mut TestComm| -> Vec<WorkerId> {
         let ws = task.mn_placement().unwrap().to_vec();
         assert_eq!(ws.len(), 2);
         if let ToWorkerMessage::ComputeTask(m) = &comm.take_worker_msgs(ws[0], 1)[0] {
             assert_eq!(&m.node_list, &ws);
         } else {
             unreachable!()
         }
-        if reservation {
-            let msgs = comm.take_worker_msgs(ws[1], 1);
-            assert!(matches!(&msgs[0], ToWorkerMessage::SetReservation(true)));
-        }
         ws
     };
 
     let task3 = core.get_task(3.into());
-    let ws3 = test_mn_task(task3, &mut comm, true);
+    let ws3 = test_mn_task(task3, &mut comm);
     let task4 = core.get_task(4.into());
-    let ws4 = test_mn_task(task4, &mut comm, true);
+    let ws4 = test_mn_task(task4, &mut comm);
     for w in &ws4 {
         assert!(!ws3.contains(w));
     }
@@ -124,7 +116,7 @@ fn test_schedule_mn_simple() {
     core.sanity_check();
 
     let task2 = core.get_task(2.into());
-    let ws2 = test_mn_task(task2, &mut comm, false);
+    let ws2 = test_mn_task(task2, &mut comm);
     comm.emptiness_check();
 
     finish_on_worker(&mut core, 3, ws2[0]);
@@ -152,14 +144,9 @@ fn test_schedule_mn_reserve() {
         comm.take_worker_msgs(ws1[0], 1)[0],
         ToWorkerMessage::ComputeTask(_)
     ));
-    assert!(matches!(
-        comm.take_worker_msgs(ws1[1], 1)[0],
-        ToWorkerMessage::SetReservation(true)
-    ));
-    assert!(matches!(
-        comm.take_worker_msgs(ws1[2], 1)[0],
-        ToWorkerMessage::SetReservation(true)
-    ));
+    assert!(!core.get_worker_by_id_or_panic(ws1[0]).is_reserved());
+    assert!(core.get_worker_by_id_or_panic(ws1[1]).is_reserved());
+    assert!(core.get_worker_by_id_or_panic(ws1[2]).is_reserved());
     comm.emptiness_check();
     finish_on_worker(&mut core, 1, ws1[0]);
     scheduler.run_scheduling(&mut core, &mut comm);

crates/tako/src/internal/tests/test_worker.rs

@@ -140,24 +140,6 @@ fn test_worker_start_task_resource_variants() {
     assert_eq!(requests[0], rq);
 }
 
-#[test]
-fn test_worker_reservation() {
-    let state_ref = create_test_worker_state(create_test_worker_config());
-    let mut state = state_ref.get_mut();
-    let finish_time = state.last_task_finish_time;
-    assert!(!state.reservation);
-    process_worker_message(&mut state, ToWorkerMessage::SetReservation(true));
-    let comm = state.comm().test();
-    comm.check_emptiness();
-    assert!(state.reservation);
-    assert_eq!(state.last_task_finish_time, finish_time);
-    process_worker_message(&mut state, ToWorkerMessage::SetReservation(false));
-    let comm = state.comm().test();
-    comm.check_emptiness();
-    assert!(!state.reservation);
-    assert_ne!(state.last_task_finish_time, finish_time);
-}
-
 #[test]
 fn test_worker_other_workers() {
     let state_ref = create_test_worker_state(create_test_worker_config());

crates/tako/src/internal/worker/rpc.rs

@@ -177,10 +177,6 @@ pub async fn run_worker(
     let local_comm_fut = handle_local_comm(local_conn_listener, state_ref.clone());
 
     let heartbeat_fut = heartbeat_process(heartbeat_interval, state_ref.clone());
-    let idle_timeout_fut = match configuration.idle_timeout {
-        Some(timeout) => Either::Left(idle_timeout_process(timeout, state_ref.clone())),
-        None => Either::Right(futures::future::pending()),
-    };
     let overview_fut = send_overview_loop(state_ref.clone());
 
     let time_limit_fut = match time_limit {
@@ -220,10 +216,6 @@ pub async fn run_worker(
                 log::info!("Time limit reached");
                 Ok(Some(FromWorkerMessage::Stop(WorkerStopReason::TimeLimitReached)))
             }
-            _ = idle_timeout_fut => {
-                log::info!("Idle timeout reached");
-                Ok(Some(FromWorkerMessage::Stop(WorkerStopReason::IdleTimeout)))
-            }
             _ = stop_flag.notified() => {
                 log::info!("Worker received an external stop notification");
                 Ok(Some(FromWorkerMessage::Stop(WorkerStopReason::Interrupted)))
@@ -379,24 +371,6 @@ async fn heartbeat_process(heartbeat_interval: Duration, state_ref: WrappedRcRef
     }
 }
 
-/// Runs until an idle timeout happens.
-/// Idle timeout occurs when the worker doesn't have anything to do for the specified duration.
-async fn idle_timeout_process(idle_timeout: Duration, state_ref: WrappedRcRefCell<WorkerState>) {
-    let mut interval = tokio::time::interval(Duration::from_secs(1));
-
-    loop {
-        interval.tick().await;
-
-        let state = state_ref.get();
-        if !state.has_tasks() && !state.reservation {
-            let elapsed = state.last_task_finish_time.elapsed();
-            if elapsed > idle_timeout {
-                break;
-            }
-        }
-    }
-}
-
 pub(crate) fn process_worker_message(state: &mut WorkerState, message: ToWorkerMessage) -> bool {
     match message {
         ToWorkerMessage::ComputeTask(msg) => {
@@ -449,12 +423,6 @@ pub(crate) fn process_worker_message(state: &mut WorkerState, message: ToWorkerM
         ToWorkerMessage::LostWorker(worker_id) => {
             state.remove_worker(worker_id);
         }
-        ToWorkerMessage::SetReservation(on_off) => {
-            state.reservation = on_off;
-            if !on_off {
-                state.reset_idle_timer();
-            }
-        }
         ToWorkerMessage::PlacementResponse(data_id, placement) => {
             state.process_resolved_placement(data_id, placement);
         }

crates/tako/src/internal/worker/state.rs

@@ -55,9 +55,6 @@ pub struct WorkerState {
     //pub(crate) secret_key: Option<Arc<SecretKey>>,
     pub(crate) start_time: Instant,
 
-    pub(crate) reservation: bool, // If true, idle timeout is blocked
-    pub(crate) last_task_finish_time: Instant,
-
     pub(crate) lc_state: RefCell<LocalCommState>,
     pub(crate) data_storage: DataStorage,
     download_manager: Option<WorkerDownloadManagerRef>,
@@ -166,10 +163,6 @@ impl WorkerState {
         !self.tasks.is_empty()
     }
 
-    pub fn reset_idle_timer(&mut self) {
-        self.last_task_finish_time = Instant::now();
-    }
-
     fn remove_task(
         &mut self,
         task_id: TaskId,
@@ -214,7 +207,6 @@ impl WorkerState {
         if self.tasks.is_empty() {
             self.comm.notify_worker_is_empty();
         }
-        self.reset_idle_timer();
         outputs
     }
 
@@ -466,8 +458,6 @@ impl WorkerStateRef {
             start_time: now,
             resource_map,
             resource_label_map,
-            last_task_finish_time: now,
-            reservation: false,
             worker_addresses: Default::default(),
             lc_state: RefCell::new(LocalCommState::new()),
             data_storage: DataStorage::new(),