Allow depending on tasks from an earlier submit in the jobfile

Author: Kobzol

CHANGELOG.md

@@ -26,6 +26,7 @@
 * `hq job progress` stopped using periodic polling.
   In the current version, the client passively waits for events.
   It saves some resources as the server is not queried every second.
+* You can now depend on tasks from an earlier submit when submitting a jobfile.
 
 ## v0.24.0
 

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

@@ -8,18 +8,20 @@ use crate::client::commands::submit::defs::{PinMode as PinModeDef, TaskConfigDef
 use crate::client::globalsettings::GlobalSettings;
 use crate::common::arraydef::IntArray;
 use crate::common::utils::fs::get_current_dir;
+use crate::rpc_call;
 use crate::transfer::connection::ClientSession;
 use crate::transfer::messages::{
-    JobDescription, JobSubmitDescription, JobTaskDescription, PinMode, SubmitRequest,
-    TaskDescription, TaskKind, TaskKindProgram, TaskWithDependencies,
+    FromClientMessage, IdSelector, JobDescription, JobDetailRequest, JobSubmitDescription,
+    JobTaskDescription, PinMode, SubmitRequest, TaskDescription, TaskIdSelector, TaskKind,
+    TaskKindProgram, TaskSelector, TaskStatusSelector, TaskWithDependencies, ToClientMessage,
 };
 use clap::Parser;
 use smallvec::smallvec;
 use std::path::PathBuf;
-use tako::Map;
 use tako::gateway::{EntryType, ResourceRequest, ResourceRequestVariants, TaskDataFlags};
 use tako::program::{FileOnCloseBehavior, ProgramDefinition, StdioDef};
 use tako::{JobId, JobTaskCount, JobTaskId};
+use tako::{Map, Set};
 
 #[derive(Parser)]
 pub struct JobSubmitFileOpts {
@@ -130,22 +132,33 @@ fn build_job_desc_individual_tasks(
     tasks: Vec<TaskDef>,
     data_flags: TaskDataFlags,
     has_streaming: bool,
+    existing_tasks: &[JobTaskId],
 ) -> crate::Result<JobTaskDescription> {
     let mut max_id: JobTaskId = tasks
         .iter()
         .map(|t| t.id)
+        .chain(existing_tasks.iter().copied().map(Some))
         .max()
         .flatten()
         .unwrap_or(JobTaskId::new(0));
 
+    let existing_tasks: Set<JobTaskId> = existing_tasks.iter().copied().collect();
+
     /* Topological sort */
     let original_len = tasks.len();
     let mut new_tasks = Vec::with_capacity(original_len);
     let mut unprocessed_tasks = Map::new();
     let mut in_degrees = Map::new();
     let mut consumers: Map<JobTaskId, Vec<_>> = Map::new();
+    let mut deps_from_previous_submit: Map<JobTaskId, Vec<_>> = Map::new();
     for task in tasks {
         let t = build_task(task, &mut max_id, data_flags, has_streaming);
+        if existing_tasks.contains(&t.id) {
+            return Err(crate::Error::GenericError(format!(
+                "Task {} has already been defined in this job",
+                t.id
+            )));
+        }
         if in_degrees.insert(t.id, t.task_deps.len()).is_some() {
             return Err(crate::Error::GenericError(format!(
                 "Task {} is defined multiple times",
@@ -157,14 +170,21 @@ fn build_job_desc_individual_tasks(
             new_tasks.push(t);
         } else {
             for dep in &t.task_deps {
-                consumers.entry(*dep).or_default().push(t.id);
+                if existing_tasks.contains(dep) {
+                    deps_from_previous_submit
+                        .entry(*dep)
+                        .or_default()
+                        .push(t.id);
+                } else {
+                    consumers.entry(*dep).or_default().push(t.id);
+                }
             }
             unprocessed_tasks.insert(t.id, t);
         }
     }
-    let mut idx = 0;
-    while idx < new_tasks.len() {
-        if let Some(consumers) = consumers.get(&new_tasks[idx].id) {
+
+    let mut handle_consumers =
+        |consumers: &[JobTaskId], new_tasks: &mut Vec<TaskWithDependencies>| {
             for c in consumers {
                 let d = in_degrees.get_mut(c).unwrap();
                 assert!(*d > 0);
@@ -173,6 +193,16 @@ fn build_job_desc_individual_tasks(
                     new_tasks.push(unprocessed_tasks.remove(c).unwrap())
                 }
             }
+        };
+
+    for consumers in deps_from_previous_submit.values() {
+        handle_consumers(consumers, &mut new_tasks);
+    }
+
+    let mut idx = 0;
+    while idx < new_tasks.len() {
+        if let Some(consumers) = consumers.get(&new_tasks[idx].id) {
+            handle_consumers(consumers, &mut new_tasks);
         }
         idx += 1;
     }
@@ -182,23 +212,35 @@ fn build_job_desc_individual_tasks(
     } else {
         let t = unprocessed_tasks.values().next().unwrap();
         return Err(crate::Error::GenericError(format!(
-            "Task {} is part of dependency cycle or has an invalid dependencies",
+            "Task {} is part of dependency cycle or has invalid dependencies",
             t.id
         )));
     }
 
     Ok(JobTaskDescription::Graph { tasks: new_tasks })
 }
 
-fn build_job_submit(jdef: JobDef, job_id: Option<JobId>) -> crate::Result<SubmitRequest> {
+fn build_job_submit(
+    jdef: JobDef,
+    job_info: Option<(JobId, Vec<JobTaskId>)>,
+) -> crate::Result<SubmitRequest> {
     let task_desc = if let Some(array) = jdef.array {
         build_job_desc_array(array, jdef.stream.is_some())
     } else {
         let mut data_flags = TaskDataFlags::empty();
         if jdef.data_layer {
             data_flags.insert(TaskDataFlags::ENABLE_DATA_LAYER);
         }
-        build_job_desc_individual_tasks(jdef.tasks, data_flags, jdef.stream.is_some())?
+        let existing_tasks = job_info
+            .as_ref()
+            .map(|(_, tasks)| tasks.as_slice())
+            .unwrap_or_default();
+        build_job_desc_individual_tasks(
+            jdef.tasks,
+            data_flags,
+            jdef.stream.is_some(),
+            existing_tasks,
+        )?
     };
     Ok(SubmitRequest {
         job_desc: JobDescription {
@@ -210,7 +252,7 @@ fn build_job_submit(jdef: JobDef, job_id: Option<JobId>) -> crate::Result<Submit
             submit_dir: get_current_dir(),
             stream_path: jdef.stream,
         },
-        job_id,
+        job_id: job_info.map(|j| j.0),
     })
 }
 
@@ -225,6 +267,25 @@ pub async fn submit_computation_from_job_file(
                 anyhow::anyhow!(format!("Cannot read {}: {}", opts.path.display(), e))
             })?)?
         };
-    let request = build_job_submit(jdef, opts.job)?;
+
+    let job_info = if let Some(job_id) = opts.job {
+        let mut response =
+            rpc_call!(session.connection(), FromClientMessage::JobDetail(JobDetailRequest {
+            job_id_selector: IdSelector::Specific(IntArray::from_id(job_id.as_num())),
+            task_selector: Some(TaskSelector {
+                id_selector: TaskIdSelector::All,
+                status_selector: TaskStatusSelector::All
+            })
+        }), ToClientMessage::JobDetailResponse(r) => r)
+            .await?;
+        let Some(job) = response.details.pop().and_then(|(_, detail)| detail) else {
+            return Err(anyhow::anyhow!("Job {job_id} not found"));
+        };
+        Some((job_id, job.tasks.into_iter().map(|(id, _)| id).collect()))
+    } else {
+        None
+    };
+
+    let request = build_job_submit(jdef, job_info)?;
     send_submit_request(gsettings, session, request, false, false, None).await
 }

docs/jobs/jobfile.md

@@ -120,11 +120,11 @@ command = ["sleep", "1"]
 
 ## Task dependencies
 
-Job Definition File allows to define a dependencies between tasks. In other words,
+Job Definition File allows to define dependencies between tasks. In other words,
 it means that the task may be executed only if the previous tasks are already finished.
 
-The task's option `deps` defines on which tasks the given task dependents.
-The task is addressed by their IDs.
+The task's option `deps` defines on which tasks the given task depends.
+Each dependency is identified by a task ID.
 
 The following example creates three tasks where the third task depends on the first two tasks.
 
@@ -143,6 +143,8 @@ command = [...]
 deps = [1, 3] # <---- Dependency on tasks 1 and 3
 ```
 
+Dependencies can also refer to tasks from [earlier submits](openjobs.md) in the same job.
+
 ## Resource variants
 
 More resource configurations may be defined for a task.

tests/test_jobfile.py

@@ -391,6 +391,46 @@ def test_job_file_dependencies_with_cycle(hq_env: HqEnv, tmp_path):
     hq_env.command(["job", "submit-file", "job.toml"], expect_fail="cycle")
 
 
+def test_job_file_dependencies_from_previous_submit(hq_env: HqEnv, tmp_path):
+    hq_env.start_server()
+    tmp_path.joinpath("job1.toml").write_text(
+        """
+[[task]]
+id = 1
+command = ["sleep", "2"]
+"""
+    )
+
+    hq_env.command(["job", "open"])
+    hq_env.command(["job", "submit-file", "--job", "1", "job1.toml"])
+
+    tmp_path.joinpath("job2.toml").write_text(
+        """
+[[task]]
+id = 3
+command = ["sleep", "1"]
+deps = [1]
+
+[[task]]
+id = 5
+command = ["sleep", "1"]
+deps = [1, 3]
+"""
+    )
+    hq_env.command(["job", "submit-file", "--job", "1", "job2.toml"])
+    hq_env.command(["job", "close", "1"])
+
+    table = hq_env.command(["task", "info", "1", "1"], as_table=True)
+    table.check_row_value("Dependencies", "")
+    table = hq_env.command(["task", "info", "1", "3"], as_table=True)
+    table.check_row_value("Dependencies", "1")
+    table = hq_env.command(["task", "info", "1", "5"], as_table=True)
+    table.check_row_value("Dependencies", "1,3")
+
+    hq_env.start_worker()
+    wait_for_job_state(hq_env, 1, "FINISHED")
+
+
 def test_job_file_attach(hq_env: HqEnv, tmp_path):
     hq_env.start_server()
     hq_env.command(["job", "open"])