MN schedulding limited to groups

spirali · spirali · commit d88434e4b19e · 2022-10-12T20:40:52.000+02:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -26,6 +26,9 @@ instead of removed.
   If the limit is reached, the task is marked as failed.
   The limit can be configured by `--crash-limit` in submit.
 
+* Groups of workers are introduced. A multi-node task is now started only on workers from the same group.
+  By default, workers are grouped by PBS/Slurm allocations, but it can be configured manually. 
+
 ## Changes
 
 ### Resource management
diff --git a/crates/pyhq/src/cluster/worker.rs b/crates/pyhq/src/cluster/worker.rs
@@ -42,6 +42,7 @@ impl RunningWorker {
                 }]),
                 listen_address: Default::default(),
                 hostname: get_hostname(None),
+                group: "default".to_string(),
                 work_dir,
                 log_dir,
                 heartbeat_interval: Duration::from_secs(10),
diff --git a/crates/tako/src/internal/scheduler/state.rs b/crates/tako/src/internal/scheduler/state.rs
@@ -333,16 +333,40 @@ impl SchedulerState {
     // }
 
     fn try_start_multinode_tasks(&mut self, core: &mut Core) {
+        let mut selected_workers = Vec::new();
         loop {
             // "while let" not used because of lifetime problems
-            let (mn_queue, task_map, worker_map) = core.multi_node_queue_split_mut();
+            let (mn_queue, task_map, worker_map, worker_groups) = core.multi_node_queue_split_mut();
             if let Some((task_id, _)) = mn_queue.queue.peek() {
                 let task_id = *task_id;
                 let task = task_map.get_task_mut(task_id);
                 let n_nodes = task.configuration.resources.n_nodes() as usize;
                 assert!(n_nodes > 0);
 
-                if worker_map.len() < n_nodes {
+                let mut found = false;
+                let mut big_enough = false;
+                'outer: for group in worker_groups.values() {
+                    if group.size() < n_nodes {
+                        continue;
+                    }
+                    big_enough = true;
+                    selected_workers.clear();
+                    for worker_id in group.worker_ids() {
+                        let worker = worker_map.get_worker(worker_id);
+                        if worker.is_free() {
+                            selected_workers.push(worker_id);
+                        }
+                        if selected_workers.len() == n_nodes {
+                            found = true;
+                            break 'outer;
+                        }
+                    }
+                }
+                if found {
+                    mn_queue.queue.pop();
+                    self.assign_multinode(worker_map, task, std::mem::take(&mut selected_workers));
+                    continue;
+                } else if !big_enough {
                     log::debug!(
                         "Multi-node task {} put into sleep. (n_nodes={}, workers={})",
                         task_id,
@@ -352,23 +376,6 @@ impl SchedulerState {
                     mn_queue.queue.pop();
                     core.add_sleeping_mn_task(task_id);
                     continue;
-                }
-
-                let mut selected_workers = Vec::new();
-                let mut found = false;
-                for worker in worker_map.values() {
-                    if worker.is_free() {
-                        selected_workers.push(worker.id);
-                    }
-                    if selected_workers.len() == n_nodes {
-                        found = true;
-                        break;
-                    }
-                }
-                if found {
-                    mn_queue.queue.pop();
-                    self.assign_multinode(worker_map, task, selected_workers);
-                    continue;
                 } else {
                     return;
                 }
@@ -391,13 +398,13 @@ impl SchedulerState {
                 compute_b_level_metric(core.task_map_mut())
             });
 
-            let (multi_node_queue, task_map, _) = core.multi_node_queue_split_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();
+            let (multi_node_queue, task_map, _, _) = core.multi_node_queue_split_mut();
             for task_id in multi_node_ready_tasks {
                 if let Some(task) = task_map.find_task(task_id) {
                     multi_node_queue.add_task(task)
diff --git a/crates/tako/src/internal/server/core.rs b/crates/tako/src/internal/server/core.rs
@@ -82,6 +82,10 @@ impl Core {
         WorkerId::new(self.worker_id_counter)
     }
 
+    pub fn worker_groups(&self) -> &Map<String, WorkerGroup> {
+        &self.worker_groups
+    }
+
     #[inline]
     pub fn is_used_task_id(&self, task_id: TaskId) -> bool {
         task_id <= self.maximal_task_id
@@ -93,11 +97,17 @@ impl Core {
 
     pub(crate) fn multi_node_queue_split_mut(
         &mut self,
-    ) -> (&mut MultiNodeQueue, &mut TaskMap, &mut WorkerMap) {
+    ) -> (
+        &mut MultiNodeQueue,
+        &mut TaskMap,
+        &mut WorkerMap,
+        &Map<String, WorkerGroup>,
+    ) {
         (
             &mut self.multi_node_queue,
             &mut self.tasks,
             &mut self.workers,
+            &self.worker_groups,
         )
     }
 
@@ -180,13 +190,13 @@ impl Core {
 
         let worker_id = worker.id;
         if let Some(g) = self.worker_groups.get_mut(&worker.configuration.group) {
-            g.worker_ids.insert(worker_id);
+            g.new_worker(worker_id);
         } else {
             let mut worker_ids = Set::new();
             worker_ids.insert(worker_id);
             self.worker_groups.insert(
                 worker.configuration.group.clone(),
-                WorkerGroup { worker_ids },
+                WorkerGroup::new(worker_ids),
             );
         }
         self.workers.insert(worker_id, worker);
@@ -198,8 +208,8 @@ impl Core {
             .worker_groups
             .get_mut(&worker.configuration.group)
             .unwrap();
-        assert!(group.worker_ids.remove(&worker_id));
-        if group.worker_ids.is_empty() {
+        group.remove_worker(worker_id);
+        if group.is_empty() {
             self.worker_groups.remove(&worker.configuration.group);
         }
         self.workers.remove(&worker_id).unwrap()
diff --git a/crates/tako/src/internal/server/workergroup.rs b/crates/tako/src/internal/server/workergroup.rs
@@ -1,15 +1,31 @@
 use crate::{Set, WorkerId};
 
 pub struct WorkerGroup {
-    pub worker_ids: Set<WorkerId>,
+    worker_ids: Set<WorkerId>,
 }
 
 impl WorkerGroup {
+    pub fn new(worker_ids: Set<WorkerId>) -> Self {
+        WorkerGroup { worker_ids }
+    }
+
+    pub fn worker_ids(&self) -> impl Iterator<Item = WorkerId> + '_ {
+        self.worker_ids.iter().copied()
+    }
+
     pub fn new_worker(&mut self, worker_id: WorkerId) {
         assert!(self.worker_ids.insert(worker_id));
     }
 
     pub fn remove_worker(&mut self, worker_id: WorkerId) {
         assert!(self.worker_ids.remove(&worker_id));
     }
+
+    pub fn size(&self) -> usize {
+        self.worker_ids.len()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.worker_ids.is_empty()
+    }
 }
diff --git a/crates/tako/src/internal/tests/test_reactor.rs b/crates/tako/src/internal/tests/test_reactor.rs
@@ -1145,7 +1145,7 @@ fn test_worker_groups() {
     create_test_workers(&mut core, &[1, 1]);
     let g = core.worker_group("default").unwrap();
     assert_eq!(
-        sorted_vec(g.worker_ids.iter().copied().collect()),
+        sorted_vec(g.worker_ids().collect()),
         vec![WorkerId::new(100), WorkerId::new(101)]
     );
     let mut comm = create_test_comm();
@@ -1157,7 +1157,7 @@ fn test_worker_groups() {
     );
     let g = core.worker_group("default").unwrap();
     assert_eq!(
-        sorted_vec(g.worker_ids.iter().copied().collect()),
+        sorted_vec(g.worker_ids().collect()),
         vec![WorkerId::new(100)]
     );
     let mut comm = create_test_comm();
diff --git a/crates/tako/src/internal/tests/test_scheduler_mn.rs b/crates/tako/src/internal/tests/test_scheduler_mn.rs
@@ -5,11 +5,12 @@ use crate::internal::server::core::Core;
 use crate::internal::server::task::Task;
 use crate::internal::tests::utils::env::{create_test_comm, TestComm};
 use crate::internal::tests::utils::schedule::{
-    create_test_scheduler, create_test_worker, create_test_workers, finish_on_worker,
-    submit_test_tasks,
+    create_test_scheduler, create_test_worker, create_test_worker_config, create_test_workers,
+    finish_on_worker, new_test_worker, submit_test_tasks,
 };
 use crate::internal::tests::utils::sorted_vec;
 use crate::internal::tests::utils::task::TaskBuilder;
+use crate::resources::{ResourceDescriptor, ResourceMap};
 use crate::{Priority, TaskId, WorkerId};
 
 /*fn get_mn_placement(task: &Task) -> Vec<WorkerId> {
@@ -300,3 +301,37 @@ fn test_mn_sleep_wakeup_at_once() {
     assert!(core.task_map().get_task(1.into()).is_waiting());
     assert!(core.task_map().get_task(2.into()).is_mn_running());
 }
+
+#[test]
+fn test_mn_schedule_on_groups() {
+    let mut core = Core::default();
+
+    let worker_id = WorkerId::new(100);
+    let mut wcfg1 = create_test_worker_config(worker_id, ResourceDescriptor::simple(1));
+    wcfg1.group = "group1".to_string();
+    new_test_worker(
+        &mut core,
+        worker_id,
+        wcfg1,
+        ResourceMap::from_vec(vec!["cpus".to_string()]),
+    );
+
+    let worker_id = WorkerId::new(101);
+    let mut wcfg2 = create_test_worker_config(worker_id, ResourceDescriptor::simple(1));
+    wcfg2.group = "group2".to_string();
+    new_test_worker(
+        &mut core,
+        worker_id,
+        wcfg2,
+        ResourceMap::from_vec(vec!["cpus".to_string()]),
+    );
+
+    let mut comm = create_test_comm();
+    let task1 = TaskBuilder::new(1).n_nodes(2).build();
+    submit_test_tasks(&mut core, vec![task1]);
+
+    let mut scheduler = create_test_scheduler();
+    scheduler.run_scheduling(&mut core, &mut comm);
+    core.sanity_check();
+    assert!(core.task_map().get_task(1.into()).is_waiting());
+}
diff --git a/crates/tako/src/internal/tests/utils/schedule.rs b/crates/tako/src/internal/tests/utils/schedule.rs
@@ -14,9 +14,12 @@ use crate::worker::{ServerLostPolicy, WorkerConfiguration};
 use crate::{TaskId, WorkerId};
 use std::time::{Duration, Instant};
 
-pub fn create_test_worker(core: &mut Core, worker_id: WorkerId, cpus: u64) {
-    let wcfg = WorkerConfiguration {
-        resources: ResourceDescriptor::simple(cpus),
+pub fn create_test_worker_config(
+    worker_id: WorkerId,
+    resources: ResourceDescriptor,
+) -> WorkerConfiguration {
+    WorkerConfiguration {
+        resources,
         listen_address: format!("1.1.1.{}:123", worker_id),
         hostname: format!("test{}", worker_id),
         group: "default".to_string(),
@@ -28,14 +31,27 @@ pub fn create_test_worker(core: &mut Core, worker_id: WorkerId, cpus: u64) {
         time_limit: None,
         on_server_lost: ServerLostPolicy::Stop,
         extra: Default::default(),
-    };
+    }
+}
 
-    let worker = Worker::new(
+pub fn new_test_worker(
+    core: &mut Core,
+    worker_id: WorkerId,
+    configuration: WorkerConfiguration,
+    resource_map: ResourceMap,
+) {
+    let worker = Worker::new(worker_id, configuration, resource_map);
+    on_new_worker(core, &mut TestComm::default(), worker);
+}
+
+pub fn create_test_worker(core: &mut Core, worker_id: WorkerId, cpus: u64) {
+    let wcfg = create_test_worker_config(worker_id, ResourceDescriptor::simple(cpus));
+    new_test_worker(
+        core,
         worker_id,
         wcfg,
         ResourceMap::from_vec(vec!["cpus".to_string()]),
     );
-    on_new_worker(core, &mut TestComm::default(), worker);
 }
 
 pub fn create_test_workers(core: &mut Core, cpus: &[u64]) {
diff --git a/docs/deployment/worker.md b/docs/deployment/worker.md
@@ -159,3 +159,7 @@ If you also want to include workers that are offline (i.e. that have crashed or
 ```bash
 $ hq worker info <worker-id>
 ```
+
+### Worker groups
+
+Each worker is a member exactly of one group. Groups are used when multi-node tasks are used. See more [here](../jobs/multinode.md#groups)
diff --git a/docs/jobs/multinode.md b/docs/jobs/multinode.md
@@ -28,9 +28,20 @@ Hostnames of all assigned nodes can be found in file which path is in
 environmental variable ``HQ_NODE_FILE``. Each line is now host name.
 The first line is always the root node.
 
-
 Note: Multi-node tasks always enables task directory (``--task-dir``).
 
+## Groups
+
+A multi-node task is started only on workers that belong to the same group.
+By default, workers are grouped by PBS/Slurm allocations and workers outside any allocation 
+are put in "default" group.
+
+A group of a worker can be specified at the start of the worker and it may be any string. Example:
+
+```commandline
+$ hq worker start --group my_group
+```
+
 ## Running MPI tasks
 
 A script that starts an MPI program in multi-node task may look like as follows:
diff --git a/tests/output/test_json.py b/tests/output/test_json.py
@@ -48,6 +48,7 @@ def test_print_worker_info(hq_env: HqEnv):
                 "resources": {"resources": RESOURCE_DESCRIPTOR_SCHEMA},
                 "time_limit": None,
                 "work_dir": str,
+                "group": str,
                 "on_server_lost": "stop",
             },
             "ended": None,
diff --git a/tests/pyapi/__init__.py b/tests/pyapi/__init__.py
@@ -1,10 +1,11 @@
 import os.path
 from typing import List, Tuple
 
-from hyperqueue import LocalCluster
 from hyperqueue.client import Client
 from hyperqueue.job import Job
 
+from hyperqueue import LocalCluster
+
 from ..conftest import HqEnv
 from ..utils.mock import ProgramMock
 
diff --git a/tests/pyapi/test_cluster.py b/tests/pyapi/test_cluster.py
@@ -1,5 +1,4 @@
 import pytest
-
 from hyperqueue.cluster import LocalCluster, WorkerConfig
 from hyperqueue.job import Job
 
diff --git a/tests/pyapi/test_function.py b/tests/pyapi/test_function.py
@@ -3,7 +3,6 @@
 from pathlib import Path
 
 import pytest
-
 from hyperqueue.client import Client, FailedJobsException, PythonEnv
 from hyperqueue.ffi.protocol import ResourceRequest
 from hyperqueue.job import Job
diff --git a/tests/pyapi/test_job.py b/tests/pyapi/test_job.py
@@ -4,7 +4,6 @@
 
 import iso8601
 import pytest
-
 from hyperqueue.client import FailedJobsException
 from hyperqueue.ffi.protocol import ResourceRequest
 from hyperqueue.job import Job
diff --git a/tests/test_job_mn.py b/tests/test_job_mn.py
@@ -92,3 +92,16 @@ def test_worker_lost_mn_task(hq_env: HqEnv, root: bool):
     wait_for_job_state(hq_env, 1, "WAITING")
     hq_env.start_workers(1, cpus=1)
     wait_for_job_state(hq_env, 1, "FINISHED")
+
+
+def test_submit_mn_different_groups(hq_env: HqEnv):
+    hq_env.start_server()
+    hq_env.start_workers(2, args=["--group=g1"])
+    hq_env.start_workers(2, args=["--group=g2"])
+
+    hq_env.command(["submit", "--nodes=3", "--", "/bin/hostname"])
+    time.sleep(0.5)
+    table = hq_env.command(["job", "info", "1"], as_table=True)
+    table.check_row_value("State", "WAITING")
+    hq_env.start_workers(1, args=["--group=g2"])
+    wait_for_job_state(hq_env, 1, "FINISHED")
