New scheduler priority assignment

Author: spirali

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ### New features
 
+* The server scheduler now slightly prioritizes tasks from older jobs and finishing partially-computed task graphs
 * `hq worker info` contains more information
 * `hq job forget` tries to free more memory
 * You can now configure Job name in the Python API.

crates/tako/benches/benchmarks/scheduler.rs

@@ -6,36 +6,13 @@ use crate::{add_tasks, create_worker};
 use tako::gateway::LostWorkerReason;
 use tako::internal::messages::common::TaskFailInfo;
 use tako::internal::messages::worker::ToWorkerMessage;
-use tako::internal::scheduler::metrics::compute_b_level_metric;
 use tako::internal::scheduler::state::SchedulerState;
 use tako::internal::server::comm::Comm;
 use tako::internal::server::core::Core;
 use tako::task::SerializedTaskContext;
 use tako::worker::{WorkerConfiguration, WorkerOverview};
 use tako::{InstanceId, TaskId, WorkerId};
 
-fn bench_b_level(c: &mut BenchmarkGroup<WallTime>) {
-    for task_count in [10, 1_000, 100_000] {
-        c.bench_with_input(
-            BenchmarkId::new("compute b-level", task_count),
-            &task_count,
-            |b, &task_count| {
-                b.iter_batched_ref(
-                    || {
-                        let mut core = Core::default();
-                        add_tasks(&mut core, task_count);
-                        core
-                    },
-                    |core| {
-                        compute_b_level_metric(core.task_map_mut());
-                    },
-                    BatchSize::SmallInput,
-                );
-            },
-        );
-    }
-}
-
 fn bench_schedule(c: &mut BenchmarkGroup<WallTime>) {
     for task_count in [10, 1_000, 100_000] {
         for worker_count in [1, 8, 16, 32] {
@@ -72,7 +49,6 @@ fn bench_schedule(c: &mut BenchmarkGroup<WallTime>) {
 pub fn benchmark(c: &mut Criterion) {
     let mut group = c.benchmark_group("scheduler");
 
-    bench_b_level(&mut group);
     bench_schedule(&mut group);
 }
 

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

@@ -1,4 +1,3 @@
-pub mod metrics;
 pub mod multinode;
 pub(crate) mod query;
 pub mod state;

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

@@ -37,13 +37,6 @@ fn task_priority_tuple(task: &Task) -> PriorityTuple {
 }
 
 impl MultiNodeQueue {
-    pub fn recompute_priorities(&mut self, _task_map: &TaskMap) {
-        /*if self.queues.is_empty() {
-            return;
-        }*/
-        // TODO: Not priority now, but it should be implemented also for multi node tasks
-    }
-
     pub fn shrink_to_fit(&mut self) {
         self.queues.shrink_to_fit();
         self.requests.shrink_to_fit();

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

@@ -5,7 +5,6 @@ use std::time::{Duration, Instant};
 use tokio::sync::Notify;
 use tokio::time::sleep;
 
-use super::metrics::compute_b_level_metric;
 use crate::internal::common::Map;
 use crate::internal::messages::worker::{TaskIdsMsg, ToWorkerMessage};
 use crate::internal::scheduler::multinode::MultiNodeAllocator;
@@ -346,16 +345,6 @@ impl SchedulerState {
         }
         log::debug!("Scheduling started");
 
-        if core.check_has_new_tasks_and_reset() {
-            // TODO: utilize information and do not recompute all b-levels
-            trace_time!("scheduler", "blevel", {
-                compute_b_level_metric(core.task_map_mut())
-            });
-
-            let (multi_node_queue, task_map, _, _) = core.multi_node_queue_split_mut();
-            multi_node_queue.recompute_priorities(task_map);
-        }
-
         let multi_node_ready_tasks = core.take_multi_node_ready_to_assign();
         if !multi_node_ready_tasks.is_empty() {
             let (multi_node_queue, task_map, _, _) = core.multi_node_queue_split_mut();

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

@@ -251,13 +251,14 @@ fn handle_new_tasks(
             return Some(format!("Invalid configuration index {idx}"));
         }
         let conf = &configurations[idx];
-        let task = Task::new(
+        let mut task = Task::new(
             task.id,
             task.task_deps,
             task.dataobj_deps,
             conf.clone(),
             task.body,
         );
+        task.scheduler_priority = -(task.id.job_id().as_num() as i32);
         tasks.push(task);
     }
     if !msg.adjust_instance_id_and_crash_counters.is_empty() {

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

@@ -16,6 +16,9 @@ use crate::internal::server::workermap::WorkerMap;
 use crate::{TaskId, WorkerId};
 use std::fmt::Write;
 
+// Scheduler priority increase for each t-level
+pub(crate) const T_LEVEL_WEIGHT: i32 = 256;
+
 pub(crate) fn on_new_worker(core: &mut Core, comm: &mut impl Comm, worker: Worker) {
     comm.broadcast_worker_message(&ToWorkerMessage::NewWorker(NewWorkerMsg {
         worker_id: worker.id,
@@ -169,17 +172,26 @@ 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;
+        // We assign scheduler priority here, the goal is to set scheduler_priority as follows = t-level * T_LEVEL_WEIGHT - job_id
+        // where t-level is the length of the maximal path from root tasks
+        // Goal is to prioritize task graph components that were partially computed + prioritize older tasks (according job_id)
+        // T-level is T_LEVEL_WEIGHT-times more important than job_id difference,
+        // but large job_id difference will overweight t-level which is usually bounded, that is done by design.
+        let mut priority = -(task.id.job_id().as_num() as i32);
         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() {
+                    priority =
+                        std::cmp::max(priority, task_dep.scheduler_priority + T_LEVEL_WEIGHT);
                     count += 1
                 }
                 true
             } else {
                 false
             }
         });
+        task.set_scheduler_priority(priority);
         assert!(matches!(
             task.state,
             TaskRuntimeState::Waiting(WaitingInfo { unfinished_deps: 0 })

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

@@ -11,8 +11,8 @@ use crate::internal::messages::worker::{StealResponse, StealResponseMsg};
 use crate::internal::scheduler::state::SchedulerState;
 use crate::internal::server::core::Core;
 use crate::internal::server::reactor::{
-    on_cancel_tasks, on_new_tasks, on_new_worker, on_remove_worker, on_steal_response,
-    on_task_error, on_task_finished, on_task_running,
+    T_LEVEL_WEIGHT, on_cancel_tasks, on_new_tasks, on_new_worker, on_remove_worker,
+    on_steal_response, on_task_error, on_task_finished, on_task_running,
 };
 use crate::internal::server::task::{Task, TaskRuntimeState};
 use crate::internal::server::worker::Worker;
@@ -144,6 +144,48 @@ fn test_worker_add() {
     assert_eq!(core.get_workers().count(), 2);
 }
 
+#[test]
+fn test_scheduler_priority() {
+    let mut core = Core::default();
+    let mut comm = create_test_comm();
+    //new_workers(&mut core, &mut comm, vec![1]);
+
+    let t1 = task(501);
+    let t2 = task_with_deps(502, &[&t1]);
+    let t3 = task(503);
+    let t4 = task_with_deps(504, &[&t2]);
+
+    let task_id5 = TaskId::new(123.into(), 1.into());
+    let t5 = TaskBuilder::new(task_id5).build();
+    let task_id6 = TaskId::new(122.into(), 0.into());
+    let t6 = TaskBuilder::new(task_id6).build();
+    let task_id7 = TaskId::new(123.into(), 2.into());
+    let t7 = TaskBuilder::new(task_id7).task_deps(&[&t5]).build();
+    let task_id8 = TaskId::new(123.into(), 4.into());
+    let t8 = TaskBuilder::new(task_id8).build();
+
+    on_new_tasks(&mut core, &mut comm, vec![t1, t2, t3, t4, t5, t6, t7, t8]);
+
+    assert_eq!(core.get_task(TaskId::new_test(501)).scheduler_priority, 0);
+    assert_eq!(
+        core.get_task(TaskId::new_test(502)).scheduler_priority,
+        T_LEVEL_WEIGHT
+    );
+    assert_eq!(core.get_task(TaskId::new_test(503)).scheduler_priority, 0);
+    assert_eq!(
+        core.get_task(TaskId::new_test(504)).scheduler_priority,
+        2 * T_LEVEL_WEIGHT
+    );
+
+    assert_eq!(core.get_task(task_id5).scheduler_priority, -123);
+    assert_eq!(core.get_task(task_id6).scheduler_priority, -122);
+    assert_eq!(
+        core.get_task(task_id7).scheduler_priority,
+        -123 + T_LEVEL_WEIGHT
+    );
+    assert_eq!(core.get_task(task_id8).scheduler_priority, -123);
+}
+
 #[test]
 fn test_submit_jobs() {
     let mut core = Core::default();