Add timeout capability to TaskCollection

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Current develop
 
 ### Added (new features/APIs/variables/...)
+- [[PR 1244]](https://github.com/parthenon-hpc-lab/parthenon/pull/1244) Add TaskCollection timeout capability
 - [[PR 1142]](https://github.com/parthenon-hpc-lab/parthenon/pull/1142) Unify par_dispatch, par_for_outer & par_for_inner overloads
 - [[PR 1255]](https://github.com/parthenon-hpc-lab/parthenon/pull/1255) RK34 low storage 3rd order 4 stage SSP integrator with CFL <= 2 from Spiteri & Ruuth 2002, SIAM Journal on Numerical Analysis, 40(2):469–491
 - [[PR 1283]](https://github.com/parthenon-hpc-lab/parthenon/pull/1283) Ability to automatically document ParameterInputs

doc/sphinx/src/driver.rst

@@ -47,7 +47,9 @@ More details on integrators can be found on the :ref:`integrators` page.
 Second, it defines a ``Step()`` function, which is reponsible for taking
 a timestep by looping over stages and calling the
 ``ConstructAndExecuteTaskLists`` function which builds and executes the
-tasks for each stage. Applications that derive from ``MultiStageDriver``
+tasks for each stage. The task list called here has a timeout set by the
+parameter ``parthenon/mesh/task_collection_timeout_in_seconds``, which
+defaults to five minutes. Applications that derive from ``MultiStageDriver``
 are responsible for defining a ``MakeTaskCollection`` function that
 makes a ``TaskCollection`` given a ``BlockList_t &`` and an integer
 stage. The advection example

doc/sphinx/src/generated/diffusion-parth-table.csv

@@ -23,6 +23,7 @@ block,parameters,type,default,description
 "parthenon/mesh","packs_per_rank","int","1","number of meshblockpacks per rank, overrides pack_size"
 "parthenon/mesh","refinement","string","none","mesh refinement mode"
 "parthenon/mesh","refinement_in_one_min_nbufs","int","64",""
+"parthenon/mesh","task_collection_timeout_in_seconds","int","300",""
 "parthenon/mesh","x1max","Real","","maximum x1 value of domain"
 "parthenon/mesh","x1min","Real","","minimum x1 value of domain"
 "parthenon/mesh","x1rat","Real","1","unused"

src/driver/driver.hpp

@@ -161,15 +161,15 @@ TaskListStatus ConstructAndExecuteBlockTasks(T *driver, Args... args) {
   for (auto &pmb : driver->pmesh->block_list) {
     tr[i++] = driver->MakeTaskList(pmb.get(), std::forward<Args>(args)...);
   }
-  TaskListStatus status = tc.Execute();
+  TaskListStatus status = tc.Execute(driver->pmesh->task_collection_timeout_in_seconds);
   return status;
 }
 
 template <typename T, class... Args>
 TaskListStatus ConstructAndExecuteTaskLists(T *driver, Args... args) {
   TaskCollection tc =
       driver->MakeTaskCollection(driver->pmesh->block_list, std::forward<Args>(args)...);
-  TaskListStatus status = tc.Execute();
+  TaskListStatus status = tc.Execute(driver->pmesh->task_collection_timeout_in_seconds);
   return status;
 }
 

src/mesh/mesh.cpp

@@ -78,6 +78,8 @@ Mesh::Mesh(ParameterInput *pin, ApplicationInput *app_in, Packages_t &packages,
                                      "enable a multigrid mesh")),
       nbnew(), nbdel(), step_since_lb(), gflag(), packages(packages),
       resolved_packages(ResolvePackages(packages)),
+      task_collection_timeout_in_seconds(pin->GetOrAddInteger(
+          "parthenon/mesh", "task_collection_timeout_in_seconds", 60 * 5)),
       // private members:
       num_mesh_threads_(
           pin->GetOrAddInteger("parthenon/mesh", "num_threads", 1,
@@ -653,70 +655,17 @@ void Mesh::BuildTagMapAndBoundaryBuffers() {
 void Mesh::CommunicateBoundaries(std::string md_name,
                                  const std::vector<std::string> &fields) {
   const int num_partitions = DefaultNumPartitions();
-  const int nmb = GetNumMeshBlocksThisRank(Globals::my_rank);
-  constexpr std::int64_t max_it = 1e10;
-  std::vector<bool> sent(num_partitions, false);
-  bool all_sent;
-  std::int64_t send_iters = 0;
 
+  TaskCollection tc;
+  TaskRegion &region = tc.AddRegion(num_partitions);
   auto partitions = GetDefaultBlockPartitions();
-  do {
-    all_sent = true;
-    for (int i = 0; i < partitions.size(); ++i) {
-      auto &md = mesh_data.Add(md_name, partitions[i], fields);
-      if (!sent[i]) {
-        if (SendBoundaryBuffers(md) != TaskStatus::complete) {
-          all_sent = false;
-        } else {
-          sent[i] = true;
-        }
-      }
-    }
-    send_iters++;
-  } while (!all_sent && send_iters < max_it);
-  PARTHENON_REQUIRE(
-      send_iters < max_it,
-      "Too many iterations waiting to send boundary communication buffers.");
-
-  // wait to receive FillGhost variables
-  // TODO(someone) evaluate if ReceiveWithWait kind of logic is better, also related to
-  // https://github.com/lanl/parthenon/issues/418
-  std::vector<bool> received(num_partitions, false);
-  bool all_received;
-  std::int64_t receive_iters = 0;
-  do {
-    all_received = true;
-    for (int i = 0; i < partitions.size(); ++i) {
-      auto &md = mesh_data.Add(md_name, partitions[i], fields);
-      if (!received[i]) {
-        if (ReceiveBoundaryBuffers(md) != TaskStatus::complete) {
-          all_received = false;
-        } else {
-          received[i] = true;
-        }
-      }
-    }
-    receive_iters++;
-  } while (!all_received && receive_iters < max_it);
-  PARTHENON_REQUIRE(
-      receive_iters < max_it,
-      "Too many iterations waiting to receive boundary communication buffers.");
-
-  for (auto &partition : partitions) {
-    auto &md = mesh_data.Add(md_name, partition, fields);
-    // unpack FillGhost variables
-    SetBoundaries(md);
-  }
-
-  //  Now do prolongation, compute primitives, apply BCs
-  for (auto &partition : partitions) {
-    auto &md = mesh_data.Add(md_name, partition, fields);
-    if (multilevel) {
-      ApplyBoundaryConditionsOnCoarseOrFineMD(md, true);
-      ProlongateBoundaries(md);
-    }
-    ApplyBoundaryConditionsOnCoarseOrFineMD(md, false);
+  for (int i = 0; i < num_partitions; i++) {
+    auto &md = mesh_data.Add(md_name, partitions[i], fields);
+    auto bound = AddBoundaryExchangeTasks(TaskID(0), region[i], md, multilevel);
   }
+  TaskListStatus status = tc.Execute(task_collection_timeout_in_seconds);
+  PARTHENON_REQUIRE(status == TaskListStatus::complete,
+                    "Boundary communication called internal by mesh failed.");
 }
 
 void Mesh::PreCommFillDerived() {

src/mesh/mesh.hpp

@@ -136,6 +136,7 @@ class Mesh {
 
   int step_since_lb;
   int gflag;
+  int task_collection_timeout_in_seconds;
 
   BlockList_t block_list;
   Packages_t packages;

src/tasks/tasks.cpp

@@ -156,12 +156,11 @@ TaskListStatus TaskRegion::Execute(Pool_t &pool) {
   }
 
   // then wait until everything is done
-  pool.wait();
-  pool.check_task_returns();
+  const TaskStatus status = pool.wait();
 
   // Check the results, so as to fire any exceptions from threads
   // Return failure if a task failed
-  if (task_failed.load(std::memory_order_acquire)) {
+  if (status != TaskStatus::complete || task_failed.load(std::memory_order_acquire)) {
     return TaskListStatus::fail;
   } else {
     return TaskListStatus::complete;

src/tasks/tasks.hpp

@@ -493,8 +493,8 @@ class TaskCollection {
     regions.emplace_back(num_lists);
     return regions.back();
   }
-  TaskListStatus Execute() {
-    static Pool_t pool(1);
+  TaskListStatus Execute(std::size_t timeout_in_seconds = 60 * 5) {
+    static Pool_t pool(timeout_in_seconds, 1);
     return Execute(pool);
   }
   TaskListStatus Execute(Pool_t &pool) {

src/tasks/thread_pool.hpp

@@ -14,6 +14,7 @@
 #ifndef TASKS_THREAD_POOL_HPP_
 #define TASKS_THREAD_POOL_HPP_
 
+#include <chrono>
 #include <condition_variable>
 #include <functional>
 #include <future>
@@ -66,17 +67,22 @@ class ThreadQueue {
     complete = true;
     cv.notify_all();
   }
-  void wait_for_complete() {
-    std::unique_lock<std::mutex> lock(mutex);
-    waiting = true;
-    if (queue.empty() && nwaiting == nworkers) {
+  bool wait_for_complete(std::chrono::seconds max_time) {
+    bool timeout{false};
+    {
+      std::unique_lock<std::mutex> lock(mutex);
+      waiting = true;
+      if (queue.empty() && nwaiting == nworkers) {
+        complete = false;
+        waiting = false;
+        return timeout;
+      }
+      timeout = !complete_cv.wait_for(lock, max_time, [this]() { return complete; });
       complete = false;
       waiting = false;
-      return;
     }
-    complete_cv.wait(lock, [this]() { return complete; });
-    complete = false;
-    waiting = false;
+    if (!timeout) signal_kill();
+    return timeout;
   }
 
  private:
@@ -124,8 +130,9 @@ class ThreadVector {
 
 class ThreadPool {
  public:
-  explicit ThreadPool(const int numthreads = std::thread::hardware_concurrency())
-      : nthreads(numthreads), queue(nthreads) {
+  explicit ThreadPool(const std::size_t timeout_in_seconds,
+                      const int numthreads = std::thread::hardware_concurrency())
+      : nthreads(numthreads), queue(nthreads), timeout_duration(timeout_in_seconds) {
     for (int i = 0; i < nthreads; i++) {
       auto worker = [&]() {
         while (true) {
@@ -145,8 +152,6 @@ class ThreadPool {
     }
   }
 
-  void wait() { queue.wait_for_complete(); }
-
   void kill() { queue.signal_kill(); }
 
   template <typename F, class... Args>
@@ -164,29 +169,32 @@ class ThreadPool {
   // Mostly this exists to throw any exceptions,
   // but we can check returns too.
   // Would need changes for >1 failure mode
-  TaskStatus check_task_returns() {
-    queue.wait_for_complete();
+  TaskStatus wait() {
+    const bool timeout = queue.wait_for_complete(timeout_duration);
     TaskStatus overall = TaskStatus::complete;
     for (auto &task : run_tasks) {
       TaskStatus task_return = task->get_future().get();
       if (task_return == TaskStatus::fail) overall = TaskStatus::fail;
     }
     run_tasks.clear();
 
-    return overall;
+    return timeout ? TaskStatus::fail : overall;
   }
 
  private:
   const int nthreads;
   std::vector<std::thread> threads;
   ThreadQueue<std::function<void()>> queue;
   ThreadVector<std::shared_ptr<std::packaged_task<TaskStatus()>>> run_tasks;
+  std::chrono::seconds timeout_duration;
 };
 
 template <typename return_t = TaskStatus>
 class SerialPool {
  public:
-  explicit SerialPool([[maybe_unused]] const int numthreads = 1) {}
+  explicit SerialPool(const std::size_t timeout_in_seconds,
+                      [[maybe_unused]] const int numthreads = 1)
+      : timeout_duration(timeout_in_seconds) {}
 
   template <typename F, class... Args>
   void enqueue(F &&f, Args &&...args) {
@@ -197,23 +205,33 @@ class SerialPool {
   }
 
   int size() const { return 1; }
-  void wait() {}
 
-  TaskStatus check_task_returns() {
+  TaskStatus wait() {
     TaskStatus overall = TaskStatus::complete;
+
+    const auto start_time = std::chrono::high_resolution_clock::now();
+    auto timeout_check = [start_time, timeout_duration = timeout_duration]() {
+      const auto end_time = std::chrono::high_resolution_clock::now();
+      const auto duration =
+          std::chrono::duration_cast<std::chrono::seconds>(end_time - start_time);
+      return duration >= timeout_duration;
+    };
+
     while (!queue.empty()) {
       auto f = queue.front();
       auto ret = f();
       if constexpr (std::is_same<return_t, TaskStatus>::value) {
         if (ret == TaskStatus::fail) overall = TaskStatus::fail;
       }
       queue.pop();
+      if (timeout_check()) return TaskStatus::fail;
     }
     return overall;
   }
 
  private:
   std::queue<std::function<return_t()>> queue;
+  std::chrono::seconds timeout_duration;
 };
 
 #ifdef PARTHENON_USE_SERIAL_POOL

tst/unit/test_tasklist.cpp

@@ -48,3 +48,14 @@ TEST_CASE("Task Object Lifecycle", "[TaskList][AddTask]") {
     REQUIRE(track_destruction.expired());
   }
 }
+
+TEST_CASE("TaskCollection timeout", "[TaskList][AddTask]") {
+  GIVEN("A TaskCollection") {
+    parthenon::TaskCollection tc;
+    parthenon::TaskRegion &region = tc.AddRegion(1);
+    region[0].AddTask(TaskID(0), []() { return TaskStatus::incomplete; });
+    const std::size_t timeout_in_seconds = 4;
+    tc.Execute(timeout_in_seconds);
+    REQUIRE(true);
+  }
+}