Fixed invalid semantics for finished tasks

Author: spirali

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

@@ -1,7 +1,7 @@
-use crate::TaskId;
 use crate::internal::common::{Map, Set};
 use crate::internal::server::task::Task;
 use crate::internal::server::taskmap::TaskMap;
+use crate::TaskId;
 
 pub fn compute_b_level_metric(tasks: &mut TaskMap) {
     crawl(tasks, |t| t.get_consumers());
@@ -30,26 +30,25 @@ fn crawl<F1: Fn(&Task) -> &Set<TaskId>>(tasks: &mut TaskMap, predecessor_fn: F1)
         task.set_scheduler_priority(level + 1);
 
         for t in task.task_deps.iter() {
-            let v: &mut u32 = neighbours
-                .get_mut(t)
-                .expect("Couldn't find task neighbour in level computation");
-            if *v <= 1 {
-                assert_eq!(*v, 1);
-                stack.push(*t);
-            } else {
-                *v -= 1;
+            if let Some(v) = neighbours.get_mut(t) {
+                if *v <= 1 {
+                    assert_eq!(*v, 1);
+                    stack.push(*t);
+                } else {
+                    *v -= 1;
+                }
             }
         }
     }
 }
 
 #[cfg(test)]
 mod tests {
-    use crate::TaskId;
     use crate::internal::common::index::ItemId;
     use crate::internal::scheduler::metrics::compute_b_level_metric;
     use crate::internal::server::core::Core;
     use crate::internal::tests::utils::workflows::submit_example_2;
+    use crate::TaskId;
 
     #[test]
     fn b_level_simple_graph() {

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

@@ -276,7 +276,7 @@ impl SchedulerState {
             }
             TaskRuntimeState::Running { .. }
             | TaskRuntimeState::RunningMultiNode(_)
-            | TaskRuntimeState::Finished(_) => {
+            | TaskRuntimeState::Finished => {
                 panic!("Invalid state {:?}", task.state);
             }
         };
@@ -379,14 +379,13 @@ impl SchedulerState {
                     //log::debug!("Task {} initially assigned to {}", task.id, worker_id);
                 };
                 if let Some(worker_id) = worker_id {
-                    debug_assert!(
-                        core.get_worker_map()
-                            .get_worker(worker_id)
-                            .is_capable_to_run_rqv(
-                                &core.get_task(task_id).configuration.resources,
-                                self.now
-                            )
-                    );
+                    debug_assert!(core
+                        .get_worker_map()
+                        .get_worker(worker_id)
+                        .is_capable_to_run_rqv(
+                            &core.get_task(task_id).configuration.resources,
+                            self.now
+                        ));
                     self.assign(core, task_id, worker_id);
                 } else {
                     core.add_sleeping_sn_task(task_id);

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

@@ -65,11 +65,9 @@ pub(crate) async fn process_client_message(
         ),
         FromGatewayMessage::ServerInfo => {
             let core = core_ref.get();
-            assert!(
-                client_sender
-                    .send(ToGatewayMessage::ServerInfo(core.get_server_info()))
-                    .is_ok()
-            );
+            assert!(client_sender
+                .send(ToGatewayMessage::ServerInfo(core.get_server_info()))
+                .is_ok());
             None
         }
         FromGatewayMessage::GetTaskInfo(request) => {
@@ -91,18 +89,16 @@ pub(crate) async fn process_client_message(
                             TaskRuntimeState::Stealing(_, _) => TaskState::Waiting,
                             TaskRuntimeState::Running { .. } => TaskState::Waiting,
                             TaskRuntimeState::RunningMultiNode(_) => TaskState::Waiting,
-                            TaskRuntimeState::Finished(_) => TaskState::Finished,
+                            TaskRuntimeState::Finished => TaskState::Finished,
                         },
                     }
                 })
                 .collect();
-            assert!(
-                client_sender
-                    .send(ToGatewayMessage::TaskInfo(TasksInfoResponse {
-                        tasks: task_infos
-                    }))
-                    .is_ok()
-            );
+            assert!(client_sender
+                .send(ToGatewayMessage::TaskInfo(TasksInfoResponse {
+                    tasks: task_infos
+                }))
+                .is_ok());
             None
         }
         FromGatewayMessage::CancelTasks(msg) => {
@@ -111,14 +107,12 @@ pub(crate) async fn process_client_message(
             let mut comm = comm_ref.get_mut();
             let (cancelled_tasks, already_finished) =
                 on_cancel_tasks(&mut core, &mut *comm, &msg.tasks);
-            assert!(
-                client_sender
-                    .send(ToGatewayMessage::CancelTasksResponse(CancelTasksResponse {
-                        cancelled_tasks,
-                        already_finished
-                    }))
-                    .is_ok()
-            );
+            assert!(client_sender
+                .send(ToGatewayMessage::CancelTasksResponse(CancelTasksResponse {
+                    cancelled_tasks,
+                    already_finished
+                }))
+                .is_ok());
             None
         }
         FromGatewayMessage::StopWorker(msg) => {
@@ -151,11 +145,9 @@ pub(crate) async fn process_client_message(
                 let core = core_ref.get();
                 compute_new_worker_query(&core, &msg.worker_queries)
             };
-            assert!(
-                client_sender
-                    .send(ToGatewayMessage::NewWorkerAllocationQueryResponse(response))
-                    .is_ok()
-            );
+            assert!(client_sender
+                .send(ToGatewayMessage::NewWorkerAllocationQueryResponse(response))
+                .is_ok());
             None
         }
         FromGatewayMessage::TryReleaseMemory => {
@@ -169,11 +161,9 @@ pub(crate) async fn process_client_message(
                 .get_worker_map()
                 .get(&worker_id)
                 .map(|w| w.worker_info(core.task_map()));
-            assert!(
-                client_sender
-                    .send(ToGatewayMessage::WorkerInfo(response))
-                    .is_ok()
-            );
+            assert!(client_sender
+                .send(ToGatewayMessage::WorkerInfo(response))
+                .is_ok());
             None
         }
     }
@@ -227,12 +217,10 @@ fn handle_new_tasks(
     }
     on_new_tasks(core, comm, tasks);
 
-    assert!(
-        client_sender
-            .send(ToGatewayMessage::NewTasksResponse(NewTasksResponse {
-                n_waiting_for_workers: 0 // TODO
-            }))
-            .is_ok()
-    );
+    assert!(client_sender
+        .send(ToGatewayMessage::NewTasksResponse(NewTasksResponse {
+            n_waiting_for_workers: 0 // TODO
+        }))
+        .is_ok());
     None
 }

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

@@ -313,7 +313,13 @@ impl Core {
             .tasks
             .remove(task_id)
             .expect("Trying to remove non-existent task");
-        assert!(!task.has_consumers());
+        if matches!(&task.state, TaskRuntimeState::Waiting(w) if w.unfinished_deps > 0) {
+            for input_id in task.task_deps {
+                if let Some(input) = self.find_task_mut(input_id) {
+                    assert!(input.remove_consumer(task_id));
+                }
+            }
+        }
         task.state
     }
 
@@ -364,7 +370,7 @@ impl Core {
     pub fn sanity_check(&self) {
         let fw_check = |task: &Task| {
             for task_dep in &task.task_deps {
-                assert!(self.tasks.get_task(*task_dep).is_finished());
+                assert!(self.tasks.find_task(*task_dep).is_none());
             }
             for &task_id in task.get_consumers() {
                 assert!(self.tasks.get_task(task_id).is_waiting());
@@ -399,7 +405,11 @@ impl Core {
                 TaskRuntimeState::Waiting(winfo) => {
                     let mut count = 0;
                     for task_dep in &task.task_deps {
-                        if !self.tasks.get_task(*task_dep).is_finished() {
+                        if !self
+                            .tasks
+                            .find_task(*task_dep)
+                            .is_none_or(|t| t.is_finished())
+                        {
                             count += 1;
                         }
                     }
@@ -424,9 +434,9 @@ impl Core {
                     worker_check_sn(self, task.id, target.unwrap_or(WorkerId::new(0)));
                 }
 
-                TaskRuntimeState::Finished(_) => {
+                TaskRuntimeState::Finished => {
                     for task_dep in &task.task_deps {
-                        assert!(self.tasks.get_task(*task_dep).is_finished());
+                        assert!(self.tasks.find_task(*task_dep).is_none());
                     }
                 }
                 TaskRuntimeState::RunningMultiNode(ws) => {

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

@@ -7,7 +7,7 @@ use crate::internal::messages::worker::{
 };
 use crate::internal::server::comm::Comm;
 use crate::internal::server::core::Core;
-use crate::internal::server::task::{FinishInfo, WaitingInfo};
+use crate::internal::server::task::WaitingInfo;
 use crate::internal::server::task::{Task, TaskRuntimeState};
 use crate::internal::server::worker::Worker;
 use crate::internal::server::workermap::WorkerMap;
@@ -84,7 +84,7 @@ pub(crate) fn on_remove_worker(
                     continue;
                 }
             }
-            TaskRuntimeState::Finished(_finfo) => {
+            TaskRuntimeState::Finished => {
                 continue;
             }
             TaskRuntimeState::RunningMultiNode(ws) => {
@@ -164,13 +164,17 @@ pub(crate) fn on_new_tasks(core: &mut Core, comm: &mut impl Comm, new_tasks: Vec
     assert!(!new_tasks.is_empty());
     for mut task in new_tasks.into_iter() {
         let mut count = 0;
-        for t in task.task_deps.iter() {
-            let task_dep = core.get_task_mut(*t);
-            task_dep.add_consumer(task.id);
-            if !task_dep.is_finished() {
-                count += 1
+        task.task_deps.retain(|t| {
+            if let Some(task_dep) = core.find_task_mut(*t) {
+                task_dep.add_consumer(task.id);
+                if !task_dep.is_finished() {
+                    count += 1
+                }
+                true
+            } else {
+                false
             }
-        }
+        });
         assert!(matches!(
             task.state,
             TaskRuntimeState::Waiting(WaitingInfo { unfinished_deps: 0 })
@@ -225,7 +229,7 @@ pub(crate) fn on_task_running(
             }
             TaskRuntimeState::Running { .. }
             | TaskRuntimeState::Waiting(_)
-            | TaskRuntimeState::Finished(_) => {
+            | TaskRuntimeState::Finished => {
                 unreachable!()
             }
         };
@@ -276,7 +280,7 @@ pub(crate) fn on_task_finished(
                     assert_eq!(*w_id, worker_id);
                     /* Do nothing */
                 }
-                TaskRuntimeState::Waiting(_) | TaskRuntimeState::Finished(_) => {
+                TaskRuntimeState::Waiting(_) | TaskRuntimeState::Finished => {
                     unreachable!();
                 }
             }
@@ -287,7 +291,7 @@ pub(crate) fn on_task_finished(
                 placement.insert(worker_id);
             }
 
-            task.state = TaskRuntimeState::Finished(FinishInfo {});
+            task.state = TaskRuntimeState::Finished;
             comm.ask_for_scheduling();
             comm.send_client_task_finished(task.id);
         } else {
@@ -315,8 +319,8 @@ pub(crate) fn on_task_finished(
         core.add_ready_to_assign(id);
         comm.ask_for_scheduling();
     }
-    unregister_as_consumer(core, comm, msg.id);
-    remove_task_if_possible(core, comm, msg.id);
+    let state = core.remove_task(msg.id);
+    assert!(matches!(state, TaskRuntimeState::Finished));
 }
 
 pub(crate) fn on_steal_response(
@@ -423,18 +427,13 @@ fn fail_task_helper(
         }
     };
 
-    // TODO: take taskmap in `unregister_as_consumer`
-    unregister_as_consumer(core, comm, task_id);
-
     for &consumer in &consumers {
-        {
-            let task = core.get_task(consumer);
-            log::debug!("Task={} canceled because of failed dependency", task.id);
-            assert!(task.is_waiting());
-        }
-        unregister_as_consumer(core, comm, consumer);
+        log::debug!("Task={} canceled because of failed dependency", consumer);
+        assert!(matches!(
+            core.remove_task(consumer),
+            TaskRuntimeState::Waiting(_)
+        ));
     }
-
     let state = core.remove_task(task_id);
     if worker_id.is_some() {
         assert!(matches!(
@@ -448,14 +447,6 @@ fn fail_task_helper(
         assert!(matches!(state, TaskRuntimeState::Waiting(_)));
     }
     drop(state);
-
-    for &consumer in &consumers {
-        // We can drop the resulting state as checks was done earlier
-        assert!(matches!(
-            core.remove_task(consumer),
-            TaskRuntimeState::Waiting(_)
-        ));
-    }
     comm.send_client_task_error(task_id, consumers, error_info);
 }
 
@@ -522,17 +513,13 @@ pub(crate) fn on_cancel_tasks(
                     }
                     running_ids.entry(from_id).or_default().push(task_id);
                 }
-                TaskRuntimeState::Finished(_) => {
+                TaskRuntimeState::Finished => {
                     already_finished.push(task_id);
                 }
             };
         }
     }
 
-    for &task_id in &to_unregister {
-        unregister_as_consumer(core, comm, task_id);
-    }
-
     core.remove_tasks_batched(&to_unregister);
 
     for (w_id, ids) in running_ids {
@@ -542,23 +529,3 @@ pub(crate) fn on_cancel_tasks(
     comm.ask_for_scheduling();
     (to_unregister.into_iter().collect(), already_finished)
 }
-
-fn unregister_as_consumer(core: &mut Core, comm: &mut impl Comm, task_id: TaskId) {
-    let inputs: Vec<TaskId> = core.get_task(task_id).task_deps.iter().copied().collect();
-    for input_id in inputs {
-        let input = core.get_task_mut(input_id);
-        assert!(input.remove_consumer(task_id));
-        remove_task_if_possible(core, comm, input_id);
-    }
-}
-
-fn remove_task_if_possible(core: &mut Core, _comm: &mut impl Comm, task_id: TaskId) {
-    if !core.get_task(task_id).is_removable() {
-        return;
-    }
-    match core.remove_task(task_id) {
-        TaskRuntimeState::Finished(_finfo) => { /* Ok */ }
-        _ => unreachable!(),
-    };
-    log::debug!("Task id={task_id} is no longer needed");
-}