In heartbeat handle, proactively check for missed min_time.

Author: spirali

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

@@ -141,7 +141,7 @@ pub enum WorkerTaskUpdate {
     RunningPrefilled(TaskRunningMsg),
     RejectRequest {
         task_id: TaskId,
-        rv_id: ResourceVariantId,
+        rv_id: Option<ResourceVariantId>,
     },
     EnableRequest {
         resource_rq_id: ResourceRqId,

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

@@ -357,7 +357,7 @@ fn task_reject(
     comm: &mut impl Comm,
     worker_id: WorkerId,
     task_id: TaskId,
-    resource_rq_variant: ResourceVariantId,
+    resource_rq_variant: Option<ResourceVariantId>,
 ) -> bool {
     let CoreSplitMut {
         task_map,
@@ -371,23 +371,27 @@ fn task_reject(
         log::debug!("Unknown task rejected id={task_id}");
         return false;
     };
-    log::debug!("Task id={task_id} (variant={resource_rq_variant}) rejected on worker={worker_id}");
+    log::debug!(
+        "Task id={task_id} (variant={resource_rq_variant:?}) rejected on worker={worker_id}"
+    );
     let worker = worker_map.get_worker_mut(worker_id);
     let resource_rq_id = task.resource_rq_id;
-    worker.block_request(resource_rq_id, resource_rq_variant);
+    if let Some(rv_id) = resource_rq_variant {
+        worker.block_request(resource_rq_id, rv_id);
+    }
     match &task.state {
         TaskRuntimeState::Assigned {
             worker_id: w_id,
             rv_id,
         } => {
             if worker_id != *w_id {
                 log::debug!("Rejection from invalid worker");
+            } else if resource_rq_variant != Some(*rv_id) {
+                log::debug!("Rejection invalid variant");
+            } else {
+                let rq = request_map.get(resource_rq_id).get(*rv_id);
+                worker.remove_sn_task(task_id, rq);
             }
-            if resource_rq_variant != *rv_id {
-                log::debug!("Rejection from invalid worker");
-            }
-            let rq = request_map.get(resource_rq_id).get(resource_rq_variant);
-            worker.remove_sn_task(task_id, rq);
         }
         TaskRuntimeState::Prefilled { worker_id: w_id } => {
             if worker_id != *w_id {

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

@@ -672,7 +672,7 @@ fn test_task_reject1() {
         w,
         smallvec![WorkerTaskUpdate::RejectRequest {
             task_id: t,
-            rv_id: ResourceVariantId::new(0)
+            rv_id: Some(ResourceVariantId::new(0))
         }],
     );
     comm.check_need_scheduling();
@@ -717,7 +717,7 @@ fn test_task_reject2() {
         w,
         smallvec![WorkerTaskUpdate::RejectRequest {
             task_id: t,
-            rv_id: ResourceVariantId::new(0)
+            rv_id: Some(ResourceVariantId::new(0))
         }],
     );
     comm.check_need_scheduling();
@@ -746,7 +746,7 @@ fn test_task_reject3() {
         w,
         smallvec![WorkerTaskUpdate::RejectRequest {
             task_id: t1,
-            rv_id: ResourceVariantId::new(0)
+            rv_id: Some(ResourceVariantId::new(0))
         }],
     );
     comm.check_need_scheduling();
@@ -760,7 +760,7 @@ fn test_task_reject3() {
         w,
         smallvec![WorkerTaskUpdate::RejectRequest {
             task_id: t2,
-            rv_id: ResourceVariantId::new(0)
+            rv_id: Some(ResourceVariantId::new(0))
         }],
     );
     comm.check_need_scheduling();
@@ -894,7 +894,7 @@ fn test_prefill_rejected() {
     let (w1, _t1, t2) = setup_prefill(&mut rt);
     let up = WorkerTaskUpdate::RejectRequest {
         task_id: t2,
-        rv_id: 0.into(),
+        rv_id: Some(0.into()),
     };
     let mut comm = TestComm::new();
     on_task_update(rt.core(), &mut comm, w1, smallvec![up]);
@@ -1080,7 +1080,7 @@ fn test_steal_rejected() {
     let (w1, w2, t) = setup_retracting(&mut rt);
     let up = WorkerTaskUpdate::RejectRequest {
         task_id: t,
-        rv_id: 0.into(),
+        rv_id: Some(0.into()),
     };
     let mut comm = TestComm::new();
     on_task_update(rt.core(), &mut comm, w1, smallvec![up]);

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

@@ -62,7 +62,7 @@ fn try_alloc_and_start_task(
         state.blocked_requests.insert((task.resource_rq_id, rv_id));
         task_updates.push(WorkerTaskUpdate::RejectRequest {
             task_id: task.id,
-            rv_id,
+            rv_id: Some(rv_id),
         });
         return;
     };
@@ -111,7 +111,7 @@ fn try_start_task(
         // Hard reject, we never unblock this rejection so we do not need to update blocked requests
         task_updates.push(WorkerTaskUpdate::RejectRequest {
             task_id: task.id,
-            rv_id,
+            rv_id: Some(rv_id),
         });
         return Some(allocation);
     }

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

@@ -19,8 +19,8 @@ use crate::internal::common::WrappedRcRefCell;
 use crate::internal::common::resources::Allocation;
 use crate::internal::common::resources::map::ResourceIdMap;
 use crate::internal::messages::worker::{
-    FromWorkerMessage, RetractResponseMsg, TaskResourceAllocation, ToWorkerMessage, WorkerOverview,
-    WorkerRegistrationResponse, WorkerStopReason,
+    FromWorkerMessage, RetractResponseMsg, TaskResourceAllocation, TaskUpdates, ToWorkerMessage,
+    WorkerOverview, WorkerRegistrationResponse, WorkerStopReason, WorkerTaskUpdate,
 };
 use crate::internal::server::rpc::ConnectionDescriptor;
 use crate::internal::transfer::auth::{
@@ -305,10 +305,39 @@ async fn heartbeat_process(heartbeat_interval: Duration, state_ref: WrappedRcRef
     let mut interval = tokio::time::interval(heartbeat_interval);
     loop {
         interval.tick().await;
-        state_ref
-            .get_mut()
-            .comm()
-            .send_message_to_server(FromWorkerMessage::Heartbeat);
+        {
+            let mut state = state_ref.get_mut();
+            state
+                .comm()
+                .send_message_to_server(FromWorkerMessage::Heartbeat);
+            if !state.prefilled_tasks.is_empty()
+                && let Some(remaining_time) = state.remaining_time()
+            {
+                let mut to_remove = Vec::new();
+                let mut updates = TaskUpdates::new();
+                for (rq_id, tasks) in &state.prefilled_tasks {
+                    let rqv = state.resource_rq_map.get(*rq_id);
+                    if remaining_time < rqv.min_time() {
+                        to_remove.push(*rq_id);
+                        for task in tasks {
+                            // Hard reject, we never unblock this rejection so we do not need to update blocked requests
+                            updates.push(WorkerTaskUpdate::RejectRequest {
+                                task_id: task.id,
+                                rv_id: None,
+                            });
+                        }
+                    }
+                }
+                if !updates.is_empty() {
+                    state
+                        .comm()
+                        .send_message_to_server(FromWorkerMessage::TaskUpdate(updates));
+                    for rq_id in to_remove {
+                        state.prefilled_tasks.remove(&rq_id);
+                    }
+                }
+            }
+        }
         log::debug!("Heartbeat sent");
     }
 }

tests/test_job.py

@@ -1427,3 +1427,17 @@ def test_resource_weight(hq_env: HqEnv):
     hq_env.start_worker(cpus=2, args=["--resource", "gpus=sum(2)"])
     wait_for_job_state(hq_env, 1, "RUNNING")
     wait_for_job_state(hq_env, 2, "WAITING")
+
+
+def test_fast_retracting(hq_env: HqEnv):
+    hq_env.start_server()
+    hq_env.command(["submit", "--array=1-350", "--time-request=78s", "--", "sleep", "200"])
+    hq_env.start_worker(cpus=4, args=["--time-limit=80s", "--heartbeat=500ms"])
+    wait_for_job_state(hq_env, 1, "RUNNING")
+    out = hq_env.command(["--output-mode=json", "job", "info", "1"], as_json=True)
+    assert sorted(t["id"] for t in out[0]["tasks"] if t["state"] == "running") == list(range(1, 5))
+    time.sleep(6)
+    hq_env.start_worker(cpus=4)
+    time.sleep(3)
+    out = hq_env.command(["--output-mode=json", "job", "info", "1"], as_json=True)
+    assert sorted(t["id"] for t in out[0]["tasks"] if t["state"] == "running") == list(range(1, 9))