#2299: Move logic for checking max load to BaseLB

DARMA-tasking · Sep 20, 2024 · 6aa4287 · 6aa4287
1 parent 664a3bc
commit 6aa4287
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 16 deletions.
diff --git a/src/vt/vrt/collection/balance/baselb/baselb.h b/src/vt/vrt/collection/balance/baselb/baselb.h
@@ -190,14 +190,31 @@ struct BaseLB {
   );
 
   /**
-   * \brief Get the estimated time needed for load balancing
+   * \brief Get the estimated time needed for load-balancing
    *
    * \return the estimated time
    */
   auto getCollectiveEpochCost() const {
     return std::chrono::nanoseconds(100);
   }
 
+  /**
+   * \brief Check if load-balancing should be done
+   *
+   * \return true when the maximum load exceeds the cost of load balancing; false otherwise
+   */
+  bool maxLoadExceedsLBCost() const {
+    auto const max = base_stats_->at(lb::Statistic::Rank_load_modeled).at(
+      lb::StatisticQuantity::max
+    );
+    auto max_in_ns = std::chrono::duration_cast<std::chrono::nanoseconds>(
+      std::chrono::duration<double>(max)
+    );
+
+    // Compare the maximum rank load to the estimated load-balancing cost
+    return max_in_ns > getCollectiveEpochCost();
+  }
+
 private:
   TransferVecType transfers_                      = {};
   TransferType off_node_migrate_                  = {};

diff --git a/src/vt/vrt/collection/balance/greedylb/greedylb.cc b/src/vt/vrt/collection/balance/greedylb/greedylb.cc
@@ -160,11 +160,7 @@ void GreedyLB::loadStats() {
   bool should_lb = false;
   this_load_begin = this_load;
 
-  // Use an estimated load-balancing cost on average rank load to load-balance
-  auto avg_ns = std::chrono::duration_cast<std::chrono::nanoseconds>(
-    std::chrono::duration<double>(avg_load)
-  );
-  if (avg_ns > getCollectiveEpochCost()) {
+  if (maxLoadExceedsLBCost()) {
     should_lb = I > greedy_tolerance;
   }
 

diff --git a/src/vt/vrt/collection/balance/hierarchicallb/hierlb.cc b/src/vt/vrt/collection/balance/hierarchicallb/hierlb.cc
@@ -263,11 +263,7 @@ void HierarchicalLB::loadStats() {
   bool should_lb = false;
   this_load_begin = this_load;
 
-  // Use an estimated load-balancing cost on average rank load to load-balance
-  auto avg_ns = std::chrono::duration_cast<std::chrono::nanoseconds>(
-    std::chrono::duration<double>(avg_load)
-  );
-  if (avg_ns > getCollectiveEpochCost()) {
+  if (maxLoadExceedsLBCost()) {
     should_lb = I > hierlb_tolerance;
   }
 

diff --git a/src/vt/vrt/collection/balance/temperedlb/temperedlb.cc b/src/vt/vrt/collection/balance/temperedlb/temperedlb.cc
@@ -449,11 +449,7 @@ void TemperedLB::runLB(LoadType total_load) {
     target_max_load_ = avg;
   }
 
-  // Use an estimated load-balancing cost on average rank load to load-balance
-  auto avg_ns = std::chrono::duration_cast<std::chrono::nanoseconds>(
-    std::chrono::duration<double>(avg)
-  );
-  if (avg_ns > getCollectiveEpochCost()) {
+  if (maxLoadExceedsLBCost()) {
     should_lb = max > (run_temperedlb_tolerance + 1.0) * target_max_load_;
   }