First version of multi node tasks

Author: spirali

crates/hyperqueue/src/client/commands/submit/command.rs

@@ -6,7 +6,7 @@ use std::{fs, io};
 use anyhow::anyhow;
 use bstr::BString;
 use clap::Parser;
-use tako::common::resources::{CpuRequest, GenericResourceAmount};
+use tako::common::resources::{CpuRequest, GenericResourceAmount, NumOfNodes};
 use tako::messages::common::{ProgramDefinition, StdioDef};
 use tako::messages::gateway::{GenericResourceRequest, ResourceRequest};
 
@@ -134,6 +134,18 @@ impl From<PinModeArg> for PinMode {
 */
 #[derive(Parser)]
 pub struct SubmitJobConfOpts {
+    /// Number of nodes; 0
+    /// [default: 0]
+    #[clap(
+        long,
+        conflicts_with("pin"),
+        conflicts_with("cpus"),
+        conflicts_with("time-request")
+    )]
+    nodes: Option<NumOfNodes>,
+    /* Other resource configurations is not yet supported in combination of nodes,
+      remove conflict_with as support is done
+    */
     /// Number and placement of CPUs for each job
     /// [default: 1]
     #[clap(long)]
@@ -241,6 +253,7 @@ impl SubmitJobConfOpts {
             };
 
         SubmitJobConfOpts {
+            nodes: self.nodes.or(other.nodes),
             cpus: self.cpus.or(other.cpus),
             resource,
             time_request: self.time_request.or(other.time_request),
@@ -330,6 +343,7 @@ impl JobSubmitOpts {
             .collect();
 
         ResourceRequest {
+            n_nodes: self.conf.nodes.unwrap_or(0),
             cpus: self
                 .conf
                 .cpus
@@ -407,6 +421,7 @@ pub async fn submit_computation(
         directives: _,
         conf:
             SubmitJobConfOpts {
+                nodes: _,
                 cpus: _,
                 resource: _,
                 time_request: _,
@@ -462,6 +477,13 @@ pub async fn submit_computation(
         stdin,
     };
 
+    // Force task_dir for multi node tasks (for a place where to create node file)
+    let task_dir = if resources.n_nodes > 0 {
+        true
+    } else {
+        task_dir
+    };
+
     let task_desc = TaskDescription {
         program: program_def,
         resources,

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

@@ -23,6 +23,7 @@ use crate::{JobId, JobTaskCount, WorkerId};
 use chrono::{DateTime, Local, SubsecRound, Utc};
 use core::time::Duration;
 use humantime::format_duration;
+use std::borrow::Cow;
 
 use std::path::Path;
 use std::time::SystemTime;
@@ -169,12 +170,12 @@ impl CliOutput {
             .iter()
             .filter_map(|t: &JobTaskInfo| match &t.state {
                 JobTaskState::Failed {
-                    started_data: StartedTaskData { worker_id, .. },
+                    started_data: StartedTaskData { worker_ids, .. },
                     error,
                     ..
                 } => Some(vec![
                     t.task_id.cell(),
-                    format_worker(*worker_id, worker_map).cell(),
+                    format_workers(worker_ids, worker_map).cell(),
                     error.to_owned().cell().foreground_color(Some(Color::Red)),
                 ]),
                 _ => None,
@@ -573,9 +574,9 @@ impl Output for CliOutput {
                         job_rows.append(&mut vec![
                             task.task_id.cell().justify(Justify::Right),
                             task_status_to_cell(get_task_status(&task.state)),
-                            match task.state.get_worker() {
-                                Some(worker) => format_worker(worker, &worker_map),
-                                _ => "",
+                            match task.state.get_workers() {
+                                Some(workers) => format_workers(workers, &worker_map),
+                                _ => "".into(),
                             }
                             .cell(),
                             multiline_cell(vec![
@@ -960,6 +961,9 @@ fn task_status_to_cell(status: Status) -> CellStruct {
 }
 
 fn format_resource_request(rq: &ResourceRequest) -> String {
+    if rq.n_nodes > 0 {
+        return format!("nodes: {}", rq.n_nodes);
+    }
     let mut result = format_cpu_request(&rq.cpus);
     for grq in &rq.generic {
         result.push_str(&format!("\n{}: {}", grq.resource, grq.amount))
@@ -987,14 +991,15 @@ pub fn format_job_workers(tasks: &[JobTaskInfo], worker_map: &WorkerMap) -> Stri
     // BTreeSet is used to both filter duplicates and keep a stable order
     let worker_set: BTreeSet<_> = tasks
         .iter()
-        .filter_map(|task| task.state.get_worker())
+        .filter_map(|task| task.state.get_workers())
+        .flatten()
         .collect();
     let worker_count = worker_set.len();
 
     let mut result = worker_set
         .into_iter()
         .take(MAX_DISPLAYED_WORKERS)
-        .map(|id| format_worker(id, worker_map))
+        .map(|id| format_worker(*id, worker_map))
         .collect::<Vec<_>>()
         .join(", ");
 
@@ -1012,6 +1017,21 @@ fn format_worker(id: WorkerId, worker_map: &WorkerMap) -> &str {
         .unwrap_or_else(|| "N/A")
 }
 
+fn format_workers<'a>(ids: &[WorkerId], worker_map: &'a WorkerMap) -> Cow<'a, str> {
+    if ids.len() == 1 {
+        format_worker(ids[0], worker_map).into()
+    } else {
+        assert!(!ids.is_empty());
+        let mut result = String::new();
+        //result.push_str(format_worker(ids[0], worker_map));
+        for id in ids {
+            result.push_str(format_worker(*id, worker_map));
+            result.push('\n');
+        }
+        result.into()
+    }
+}
+
 fn get_task_time(state: &JobTaskState) -> (Option<DateTime<Utc>>, Option<DateTime<Utc>>) {
     match state {
         JobTaskState::Canceled {

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

@@ -115,6 +115,7 @@ impl Output for JsonOutput {
                         },
                     resources:
                         ResourceRequest {
+                            n_nodes,
                             cpus,
                             generic,
                             min_time,
@@ -135,6 +136,7 @@ impl Output for JsonOutput {
                 "stdout": format_stdio_def(&stdout),
             });
             json["resources"] = json!({
+                "n_nodes": n_nodes,
                 "cpus": format_cpu_request(cpus),
                 "generic": generic,
                 "min_time": format_duration(min_time)
@@ -225,7 +227,16 @@ impl Output for JsonOutput {
 
 fn fill_task_started_data(dict: &mut Value, data: StartedTaskData) {
     dict["started_at"] = format_datetime(data.start_date);
-    dict["worker"] = data.worker_id.as_num().into();
+    if data.worker_ids.len() == 1 {
+        dict["worker"] = data.worker_ids[0].as_num().into();
+    } else {
+        dict["workers"] = data
+            .worker_ids
+            .iter()
+            .map(|worker_id| worker_id.as_num().into())
+            .collect::<Vec<Value>>()
+            .into();
+    }
 }
 
 fn fill_task_paths(dict: &mut Value, map: &TaskToPathsMap, task_id: JobTaskId) {

crates/hyperqueue/src/common/env.rs

@@ -22,3 +22,4 @@ pub const HQ_PIN: &str = create_hq_env!("PIN");
 pub const HQ_TASK_DIR: &str = create_hq_env!("TASK_DIR");
 pub const HQ_ERROR_FILENAME: &str = create_hq_env!("ERROR_FILENAME");
 pub const HQ_CPUS: &str = create_hq_env!("CPUS");
+pub const HQ_NODE_FILE: &str = create_hq_env!("NODE_FILE");

crates/hyperqueue/src/server/autoalloc/process.rs

@@ -1584,6 +1584,7 @@ mod tests {
             cwd: Default::default(),
         };
         let resources = ResourceRequest {
+            n_nodes: 0,
             cpus: Default::default(),
             generic: vec![],
             min_time,

crates/hyperqueue/src/server/client/submit.rs

@@ -478,6 +478,7 @@ mod tests {
                 cwd: Default::default(),
             },
             resources: ResourceRequest {
+                n_nodes: 0,
                 cpus: CpuRequest::Compact(cpu_count),
                 generic: vec![GenericResourceRequest {
                     resource: "a".to_string(),

crates/hyperqueue/src/server/job.rs

@@ -9,6 +9,7 @@ use crate::transfer::messages::{
 use crate::worker::start::RunningTaskContext;
 use crate::{JobId, JobTaskCount, JobTaskId, Map, TakoTaskId, WorkerId};
 use chrono::{DateTime, Utc};
+use smallvec::SmallVec;
 use std::path::PathBuf;
 use tako::common::index::ItemId;
 use tako::common::Set;
@@ -23,7 +24,7 @@ use tokio::sync::oneshot;
 pub struct StartedTaskData {
     pub start_date: DateTime<Utc>,
     pub context: RunningTaskContext,
-    pub worker_id: WorkerId,
+    pub worker_ids: SmallVec<[WorkerId; 1]>,
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -58,8 +59,8 @@ impl JobTaskState {
         }
     }
 
-    pub fn get_worker(&self) -> Option<WorkerId> {
-        self.started_data().map(|data| data.worker_id)
+    pub fn get_workers(&self) -> Option<&[WorkerId]> {
+        self.started_data().map(|data| data.worker_ids.as_slice())
     }
 }
 
@@ -274,7 +275,7 @@ impl Job {
     pub fn set_running_state(
         &mut self,
         tako_task_id: TakoTaskId,
-        worker: WorkerId,
+        workers: SmallVec<[WorkerId; 1]>,
         context: SerializedTaskContext,
     ) {
         let (_, state) = self.get_task_state_mut(tako_task_id);
@@ -287,7 +288,7 @@ impl Job {
                 started_data: StartedTaskData {
                     start_date: Utc::now(),
                     context,
-                    worker_id: worker,
+                    worker_ids: workers,
                 },
             };
             self.counters.n_running_tasks += 1;

crates/hyperqueue/src/server/state.rs

@@ -206,10 +206,16 @@ impl State {
         log::debug!("Task id={} updated {:?}", msg.id, msg.state);
         let (mut job_id, mut is_job_terminated): (Option<JobId>, bool) = (None, false);
         match msg.state {
-            TaskState::Running { worker_id, context } => {
+            TaskState::Running {
+                worker_ids,
+                context,
+            } => {
                 let job = self.get_job_mut_by_tako_task_id(msg.id).unwrap();
-                job.set_running_state(msg.id, worker_id, context);
-                self.event_storage.on_task_started(msg.id, worker_id);
+                job.set_running_state(msg.id, worker_ids.clone(), context);
+
+                // TODO: Prepare it for multi-node tasks
+                // This (incomplete) version just takes the first worker as "the worker" for task
+                self.event_storage.on_task_started(msg.id, worker_ids[0]);
             }
             TaskState::Finished => {
                 let job = self.get_job_mut_by_tako_task_id(msg.id).unwrap();

crates/hyperqueue/src/worker/start.rs

@@ -1,7 +1,6 @@
-use std::fmt::Write;
 use std::fs::File;
 use std::io;
-use std::io::{ErrorKind, Read};
+use std::io::{BufWriter, ErrorKind, Read, Write};
 use std::path::{Path, PathBuf};
 use std::process::ExitStatus;
 use std::rc::Rc;
@@ -21,10 +20,10 @@ use tako::worker::launcher::{command_from_definitions, TaskLaunchData, TaskLaunc
 use tako::worker::state::WorkerState;
 use tako::worker::task::Task;
 use tako::worker::taskenv::{StopReason, TaskResult};
-use tako::{InstanceId, TaskId};
+use tako::{InstanceId, TaskId, WorkerId};
 
 use crate::common::env::{
-    HQ_CPUS, HQ_ERROR_FILENAME, HQ_INSTANCE_ID, HQ_PIN, HQ_SUBMIT_DIR, HQ_TASK_DIR,
+    HQ_CPUS, HQ_ERROR_FILENAME, HQ_INSTANCE_ID, HQ_NODE_FILE, HQ_PIN, HQ_SUBMIT_DIR, HQ_TASK_DIR,
 };
 use crate::common::placeholders::{
     fill_placeholders_in_paths, CompletePlaceholderCtx, ResolvablePaths,
@@ -107,12 +106,23 @@ impl TaskLauncher for HqTaskLauncher {
                         .to_string()
                         .into(),
                 );
+                if !task.node_list.is_empty() {
+                    let filename = task_dir.path().join("hq-nodelist");
+                    write_node_file(state, &task.node_list, &filename)?;
+                    program.env.insert(
+                        HQ_NODE_FILE.into(),
+                        filename.to_string_lossy().to_string().into(),
+                    );
+                }
                 Some(task_dir)
             } else {
                 None
             };
 
-            insert_resources_into_env(state, task, allocation, &mut program);
+            // Do not insert resources for multi-node tasks, semantics has to be cleared
+            if task.node_list.is_empty() {
+                insert_resources_into_env(state, task, allocation, &mut program);
+            }
 
             let submit_dir: PathBuf = program.env[<&BStr>::from(HQ_SUBMIT_DIR)].to_string().into();
             program
@@ -195,6 +205,21 @@ fn get_custom_error_filename(task_dir: &TempDir) -> PathBuf {
     task_dir.path().join("hq-error")
 }
 
+fn write_node_file(
+    state: &WorkerState,
+    node_list: &[WorkerId],
+    path: &Path,
+) -> std::io::Result<()> {
+    let file = File::create(path)?;
+    let mut file = BufWriter::new(file);
+    for worker_id in node_list {
+        file.write_all(state.worker_hostname(*worker_id).unwrap().as_bytes())?;
+        file.write_all(b"\n")?;
+    }
+    file.flush()?;
+    Ok(())
+}
+
 fn insert_resources_into_env(
     state: &WorkerState,
     task: &Task,
@@ -326,6 +351,7 @@ async fn run_task(
 
 /// Provide a more detailed error message when a process fails to be spawned.
 fn map_spawn_error(error: std::io::Error, program: &ProgramDefinition) -> tako::Error {
+    use std::fmt::Write;
     let context = match &error.kind() {
         ErrorKind::NotFound => {
             let file = &program.args[0];

crates/tako/benches/benchmarks/scheduler.rs

@@ -91,7 +91,7 @@ impl Comm for NullComm {
     fn send_client_task_started(
         &mut self,
         _task_id: TaskId,
-        _worker_id: WorkerId,
+        _worker_id: &[WorkerId],
         _context: SerializedTaskContext,
     ) {
     }