Retract check has own timed loop and configuration

Author: spirali

crates/hyperqueue/src/client/commands/worker.rs

@@ -289,6 +289,14 @@ pub struct WorkerStartOpts {
     /// It should *NOT* be placed on a network filesystem.
     #[arg(long)]
     pub work_dir: Option<PathBuf>,
+
+    #[arg(
+        long,
+        default_value = "30s",
+        value_parser = parse_hms_or_human_time,
+        help = duration_doc!("Defines the interval at which the worker checks for overdue tasks, based on the min_time setting")
+    )]
+    pub retract_check_interval: Duration,
 }
 
 pub async fn start_hq_worker(
@@ -335,6 +343,7 @@ fn gather_configuration(opts: WorkerStartOpts) -> anyhow::Result<WorkerConfigura
         hostname,
         on_server_lost,
         work_dir,
+        retract_check_interval,
     } = opts;
 
     let detect_resources = detect_resources_cli
@@ -461,6 +470,7 @@ fn gather_configuration(opts: WorkerStartOpts) -> anyhow::Result<WorkerConfigura
         overview_configuration,
         extra,
         min_utilization,
+        retract_check_interval,
     })
 }
 

crates/hyperqueue/src/client/output/json.rs

@@ -554,6 +554,7 @@ fn format_worker_info(worker_info: WorkerInfo) -> serde_json::Value {
                 overview_configuration: _,
                 idle_timeout,
                 time_limit,
+                retract_check_interval: _,
                 on_server_lost,
                 group,
                 min_utilization,

crates/pyhq/src/cluster/worker.rs

@@ -51,6 +51,7 @@ impl RunningWorker {
                 overview_configuration: OverviewConfiguration::disabled(),
                 idle_timeout: None,
                 time_limit: None,
+                retract_check_interval: Duration::from_secs(30),
                 on_server_lost: ServerLostPolicy::Stop,
                 extra: Default::default(),
                 min_utilization: 0.0,

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

@@ -6,6 +6,7 @@ use crate::internal::server::worker::Worker;
 use crate::resources::{ResourceAmount, ResourceDescriptorItem, ResourceDescriptorKind};
 use crate::worker::{ServerLostPolicy, WorkerConfiguration};
 use crate::{Set, WorkerId};
+use std::time::Duration;
 
 /// Read the documentation of `new_worker_query`` in control.rs
 pub(crate) fn compute_new_worker_query(
@@ -59,6 +60,7 @@ pub(crate) fn compute_new_worker_query(
                 on_server_lost: ServerLostPolicy::Stop,
                 min_utilization: query.min_utilization,
                 extra: Default::default(),
+                retract_check_interval: Duration::from_secs(30),
             };
             let worker = Worker::new(worker_id, configuration, &resource_map, now);
             fake_workers.push(worker);

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

@@ -120,7 +120,9 @@ async fn worker_rpc_loop(
     );
 
     let heartbeat_interval = msg.configuration.heartbeat_interval;
+
     log::debug!("Worker heartbeat interval: {heartbeat_interval:?}");
+
     // Sanity that interval is not too small
     assert!(heartbeat_interval.as_millis() > 150);
 

crates/tako/src/internal/tests/integration/utils/worker.rs

@@ -81,6 +81,7 @@ pub fn create_worker_configuration(
             time_limit: None,
             extra,
             min_utilization: 0.0,
+            retract_check_interval: Duration::from_secs(30),
         },
         secret_key,
     )

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

@@ -44,6 +44,7 @@ fn test_worker_add() {
         },
         idle_timeout: None,
         time_limit: None,
+        retract_check_interval: Duration::from_secs(30),
         on_server_lost: ServerLostPolicy::Stop,
         min_utilization: 0.0,
         extra: Default::default(),
@@ -102,6 +103,7 @@ fn test_worker_add() {
         },
         idle_timeout: None,
         time_limit: None,
+        retract_check_interval: Default::default(),
         on_server_lost: ServerLostPolicy::Stop,
         min_utilization: 0.0,
         extra: Default::default(),

crates/tako/src/internal/tests/utils/worker.rs

@@ -78,6 +78,7 @@ impl WorkerBuilder {
             },
             idle_timeout: None,
             time_limit: self.time_limit,
+            retract_check_interval: Duration::from_secs(30),
             on_server_lost: ServerLostPolicy::Stop,
             min_utilization: self.min_utilization,
             extra: Default::default(),

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

@@ -47,6 +47,7 @@ pub struct WorkerConfiguration {
     pub overview_configuration: OverviewConfiguration,
     pub idle_timeout: Option<Duration>,
     pub time_limit: Option<Duration>,
+    pub retract_check_interval: Duration,
     pub on_server_lost: ServerLostPolicy,
     pub min_utilization: f32,
     pub extra: Map<String, String>,

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

@@ -123,6 +123,7 @@ pub async fn run_worker(
 
     let (queue_sender, queue_receiver) = tokio::sync::mpsc::unbounded_channel::<Bytes>();
     let heartbeat_interval = configuration.heartbeat_interval;
+    let retract_check_interval = configuration.retract_check_interval;
     let time_limit = configuration.time_limit;
 
     let (worker_id, state_ref) = {
@@ -171,6 +172,7 @@ 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 retract_check_fut = retract_check_process(retract_check_interval, state_ref.clone());
     let overview_fut = send_overview_loop(state_ref.clone());
 
     let time_limit_fut = match time_limit {
@@ -201,6 +203,7 @@ pub async fn run_worker(
                 Ok(Some(FromWorkerMessage::Stop(WorkerStopReason::Interrupted)))
             }
             _ = heartbeat_fut => { unreachable!() }
+            _ = retract_check_fut => { unreachable!() }
             _ = overview_fut => { unreachable!() }
             _ = local_comm_fut => { unreachable!() }
         };
@@ -310,9 +313,22 @@ async fn heartbeat_process(heartbeat_interval: Duration, state_ref: WrappedRcRef
             state
                 .comm()
                 .send_message_to_server(FromWorkerMessage::Heartbeat);
+        }
+        log::debug!("Heartbeat sent");
+    }
+}
+
+/// Repeatedly check overdue task
+async fn retract_check_process(check_interval: Duration, state_ref: WrappedRcRefCell<WorkerState>) {
+    let mut interval = tokio::time::interval(check_interval);
+    loop {
+        interval.tick().await;
+        {
+            let mut state = state_ref.get_mut();
             if !state.prefilled_tasks.is_empty()
                 && let Some(remaining_time) = state.remaining_time()
             {
+                log::debug!("Checking tasks for retract");
                 let mut to_remove = Vec::new();
                 let mut updates = TaskUpdates::new();
                 for (rq_id, tasks) in &state.prefilled_tasks {
@@ -338,7 +354,6 @@ async fn heartbeat_process(heartbeat_interval: Duration, state_ref: WrappedRcRef
                 }
             }
         }
-        log::debug!("Heartbeat sent");
     }
 }