Fix restoring crash counters and instance ids

Author: spirali

CHANGELOG.md

@@ -10,6 +10,7 @@
 ### Fixes
 
 * Fixed a problem with journal loading when task dependencies are used
+* Fixed restoring crash counters and instance ids from journal
 
 ## 0.22.0
 

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

@@ -321,7 +321,7 @@ fn build_tasks_array(
             priority: task_desc.priority,
             crash_limit: task_desc.crash_limit,
         }],
-        adjust_instance_id: Default::default(),
+        adjust_instance_id_and_crash_counters: Default::default(),
     }
 }
 
@@ -391,7 +391,7 @@ fn build_tasks_graph(
     NewTasksMessage {
         tasks: task_configs,
         shared_data,
-        adjust_instance_id: Default::default(),
+        adjust_instance_id_and_crash_counters: Default::default(),
     }
 }
 

crates/hyperqueue/src/server/restore.rs

@@ -10,10 +10,12 @@ use crate::worker::start::RunningTaskContext;
 use crate::{JobId, JobTaskId, Map, make_tako_id, unwrap_tako_id};
 use std::path::Path;
 use tako::gateway::NewTasksMessage;
-use tako::{ItemId, WorkerId};
+use tako::{InstanceId, ItemId, WorkerId};
 
 struct RestorerTaskInfo {
     state: JobTaskState,
+    instance_id: Option<InstanceId>,
+    crash_counter: u32,
 }
 
 impl RestorerTaskInfo {
@@ -77,15 +79,17 @@ impl RestorerJob {
 
             for (task_id, job_task) in job.tasks.iter_mut() {
                 if let Some(task) = self.tasks.get_mut(task_id) {
+                    if task.crash_counter > 0 || task.instance_id.is_some() {
+                        new_tasks.adjust_instance_id_and_crash_counters.insert(
+                            make_tako_id(job_id, *task_id),
+                            (
+                                task.instance_id.map(|x| x.as_num() + 1).unwrap_or(0).into(),
+                                task.crash_counter,
+                            ),
+                        );
+                    }
                     match &task.state {
-                        JobTaskState::Waiting => continue,
-                        JobTaskState::Running { started_data } => {
-                            let instance_id = started_data.context.instance_id.as_num() + 1;
-                            new_tasks
-                                .adjust_instance_id
-                                .insert(make_tako_id(job_id, *task_id), instance_id.into());
-                            continue;
-                        }
+                        JobTaskState::Waiting | JobTaskState::Running { .. } => continue,
                         JobTaskState::Finished { .. } => job.counters.n_finished_tasks += 1,
                         JobTaskState::Failed { .. } => job.counters.n_failed_tasks += 1,
                         JobTaskState::Canceled { .. } => job.counters.n_canceled_tasks += 1,
@@ -112,6 +116,19 @@ impl RestorerJob {
     pub fn add_submit(&mut self, submit: SubmittedJobDescription) {
         self.submit_descs.push(submit)
     }
+
+    pub fn increase_crash_counters(&mut self, worker_id: WorkerId) {
+        for task in self.tasks.values_mut() {
+            match &task.state {
+                JobTaskState::Running { started_data }
+                    if started_data.worker_ids.contains(&worker_id) =>
+                {
+                    task.crash_counter += 1;
+                }
+                _ => {}
+            }
+        }
+    }
 }
 
 #[derive(Default)]
@@ -183,7 +200,13 @@ impl StateRestorer {
                     log::debug!("Replaying: WorkerConnected {worker_id}");
                     self.max_worker_id = self.max_worker_id.max(worker_id.as_num());
                 }
-                EventPayload::WorkerLost(_, _) => {}
+                EventPayload::WorkerLost(worker_id, reason) => {
+                    if reason.is_failure() {
+                        for job in self.jobs.values_mut() {
+                            job.increase_crash_counters(worker_id);
+                        }
+                    }
+                }
                 EventPayload::WorkerOverviewReceived(_) => {}
                 EventPayload::Submit {
                     job_id,
@@ -232,6 +255,8 @@ impl StateRestorer {
                                         worker_ids: workers,
                                     },
                                 },
+                                instance_id: Some(instance_id),
+                                crash_counter: 0,
                             },
                         );
                     }
@@ -300,6 +325,8 @@ impl StateRestorer {
                                         started_data: None,
                                         cancelled_date: event.time,
                                     },
+                                    instance_id: None,
+                                    crash_counter: 0,
                                 },
                             );
                         }

crates/tako/src/gateway.rs

@@ -129,7 +129,7 @@ pub struct TaskConfiguration {
 pub struct NewTasksMessage {
     pub tasks: Vec<TaskConfiguration>,
     pub shared_data: Vec<SharedTaskConfiguration>,
-    pub adjust_instance_id: Map<TaskId, InstanceId>,
+    pub adjust_instance_id_and_crash_counters: Map<TaskId, (InstanceId, u32)>,
 }
 
 #[derive(Deserialize, Serialize, Debug)]

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

@@ -218,10 +218,13 @@ fn handle_new_tasks(
         let task = Task::new(task.id, task.task_deps, conf.clone(), task.body);
         tasks.push(task);
     }
-    if !msg.adjust_instance_id.is_empty() {
+    if !msg.adjust_instance_id_and_crash_counters.is_empty() {
         for task in &mut tasks {
-            if let Some(instance_id) = msg.adjust_instance_id.get(&task.id) {
+            if let Some((instance_id, crash_counter)) =
+                msg.adjust_instance_id_and_crash_counters.get(&task.id)
+            {
                 task.instance_id = *instance_id;
+                task.crash_counter = *crash_counter;
             }
         }
     }

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

@@ -135,8 +135,11 @@ pub(crate) fn on_remove_worker(
 
     comm.broadcast_worker_message(&ToWorkerMessage::LostWorker(worker_id));
 
-    // IMPORTANT: We have to announce error BEFORE we announce lost worker (+ running tasks)
-    // because HQ does not recognize switch from waiting to failed stated.
+    // IMPORTANT: We need to announce lost worker before failing the jobs
+    // so in journal restoration we can detect what tasks were running
+    // without explicit logging
+    comm.send_client_worker_lost(worker_id, running_tasks, reason);
+
     for task_id in crashed_tasks {
         let count = core.get_task(task_id).crash_counter;
         log::debug!("Task {} reached crash limit {}", task_id, count);
@@ -154,9 +157,6 @@ pub(crate) fn on_remove_worker(
             },
         );
     }
-
-    comm.send_client_worker_lost(worker_id, running_tasks, reason);
-
     comm.ask_for_scheduling();
 }
 

crates/tako/src/internal/tests/integration/utils/server.rs

@@ -194,7 +194,7 @@ impl ServerHandle {
         let msg = NewTasksMessage {
             tasks,
             shared_data: configurations,
-            adjust_instance_id: Default::default(),
+            adjust_instance_id_and_crash_counters: Default::default(),
         };
         self.send(FromGatewayMessage::NewTasks(msg)).await;
         wait_for_msg!(self, ToGatewayMessage::NewTasksResponse(NewTasksResponse { .. }) => ());

tests/test_job.py

@@ -989,7 +989,7 @@ def test_zero_custom_error_message(hq_env: HqEnv):
     # print(table)
 
 
-@pytest.mark.parametrize("count", [None, 1, 7])
+@pytest.mark.parametrize("count", [None, 1, 2, 7])
 def test_crashing_job_status_default(count: Optional[int], hq_env: HqEnv):
     hq_env.start_server()
 

tests/test_journal.py

@@ -409,3 +409,19 @@ def test_restore_dependencies2(hq_env: HqEnv, tmp_path):
     hq_env.start_server(args=["--journal", journal_path])
     hq_env.start_worker(args=["--resource", "x=sum(2)"])
     wait_for_job_state(hq_env, 1, "FINISHED")
+
+
+def test_restore_crash_counters(hq_env: HqEnv, tmp_path):
+    journal_path = os.path.join(tmp_path, "my.journal")
+    hq_env.start_server(args=["--journal", journal_path])
+    hq_env.start_worker()
+    hq_env.command(["submit", "--crash-limit=2", "--", "sleep", "4"])
+    wait_for_job_state(hq_env, 1, "RUNNING")
+    hq_env.kill_worker(1)
+    wait_for_job_state(hq_env, 1, "WAITING")
+    hq_env.stop_server()
+    hq_env.start_server(args=["--journal", journal_path])
+    hq_env.start_worker()
+    wait_for_job_state(hq_env, 1, "RUNNING")
+    hq_env.kill_worker(3)
+    wait_for_job_state(hq_env, 1, "FAILED")