#1830: TemperedWMin: provide canMigrate override

src/vt/vrt/collection/balance/temperedlb/temperedlb.cc

@@ -475,7 +475,7 @@ void TemperedLB::runLB(TimeType total_load) {
 }
 
 void TemperedLB::clearDataStructures() {
-  underloaded_.clear();
+  potential_recipients_.clear();
   load_info_.clear();
   is_overloaded_ = is_underloaded_ = false;
 }
@@ -666,7 +666,7 @@ void TemperedLB::informAsync() {
 
   auto const this_node = theContext()->getNode();
   if (canPropagate()) {
-    underloaded_.insert(this_node);
+    potential_recipients_.insert(this_node);
   }
 
   setup_done_ = false;
@@ -693,7 +693,7 @@ void TemperedLB::informAsync() {
     vt_debug_print(
       terse, temperedlb,
       "TemperedLB::informAsync: trial={}, iter={}, known underloaded={}\n",
-      trial_, iter_, underloaded_.size()
+      trial_, iter_, potential_recipients_.size()
     );
   }
 
@@ -717,12 +717,12 @@ void TemperedLB::informSync() {
 
   auto const this_node = theContext()->getNode();
   if (canPropagate()) {
-    underloaded_.insert(this_node);
+    potential_recipients_.insert(this_node);
   }
 
   auto propagate_this_round = canPropagate();
   propagate_next_round_ = false;
-  new_underloaded_ = underloaded_;
+  new_potential_recipients_ = potential_recipients_;
   new_load_info_ = load_info_;
 
   setup_done_ = false;
@@ -752,15 +752,15 @@ void TemperedLB::informSync() {
 
     propagate_this_round = propagate_next_round_;
     propagate_next_round_ = false;
-    underloaded_ = new_underloaded_;
+    potential_recipients_ = new_potential_recipients_;
     load_info_ = new_load_info_;
   }
 
   if (is_overloaded_) {
     vt_debug_print(
       terse, temperedlb,
       "TemperedLB::informSync: trial={}, iter={}, known underloaded={}\n",
-      trial_, iter_, underloaded_.size()
+      trial_, iter_, potential_recipients_.size()
     );
   }
 
@@ -791,7 +791,7 @@ void TemperedLB::propagateRound(uint8_t k_cur, bool sync, EpochType epoch) {
     gen_propagate_.seed(seed_());
   }
 
-  auto& selected = underloaded_;
+  auto& selected = potential_recipients_;
   if (selected.find(this_node) == selected.end()) {
     selected.insert(this_node);
   }
@@ -868,7 +868,7 @@ void TemperedLB::propagateIncomingAsync(LoadMsgAsync* msg) {
       load_info_[elm.first] = elm.second;
 
       if (isUnderloaded(elm.second)) {
-        underloaded_.insert(elm.first);
+        potential_recipients_.insert(elm.first);
       }
     }
   }
@@ -902,7 +902,7 @@ void TemperedLB::propagateIncomingSync(LoadMsgSync* msg) {
       new_load_info_[elm.first] = elm.second;
 
       if (isUnderloaded(elm.second)) {
-        new_underloaded_.insert(elm.first);
+        new_potential_recipients_.insert(elm.first);
       }
     }
   }
@@ -1204,7 +1204,7 @@ void TemperedLB::decide() {
     auto potential_recipients = getPotentialRecipients();
     std::unordered_map<NodeType, ObjsType> migrate_objs;
 
-    if (potential_recipients.size() > 0) {
+    if (not potential_recipients.empty()) {
       std::vector<ObjIDType> ordered_obj_ids = orderObjects(
         obj_ordering_, cur_objs_, this_new_load_, target_max_load_
       );

src/vt/vrt/collection/balance/temperedlb/temperedlb.h

@@ -96,6 +96,9 @@ struct TemperedLB : BaseLB {
   void migrate();
   void clearDataStructures();
 
+  /**
+   * \brief Decides whether the rank can perform the migration
+   */
   virtual bool canMigrate() const {  return is_overloaded_; }
   /**
    * \brief Decides whether the rank can initiate information propagation stage
@@ -132,6 +135,7 @@ struct TemperedLB : BaseLB {
   void setupDone(ReduceMsgType* msg);
 
   std::unordered_map<NodeType, LoadType> load_info_ = {};
+  std::unordered_map<ObjIDType, TimeType> cur_objs_ = {};
 
 private:
   uint16_t f_                                       = 0;
@@ -175,9 +179,8 @@ struct TemperedLB : BaseLB {
   objgroup::proxy::Proxy<TemperedLB> proxy_         = {};
   bool is_overloaded_                               = false;
   bool is_underloaded_                              = false;
-  std::unordered_set<NodeType> underloaded_         = {};
-  std::unordered_set<NodeType> new_underloaded_     = {};
-  std::unordered_map<ObjIDType, TimeType> cur_objs_ = {};
+  std::unordered_set<NodeType> potential_recipients_      = {};
+  std::unordered_set<NodeType> new_potential_recipients_  = {};
   LoadType this_new_load_                           = 0.0;
   TimeType new_imbalance_                           = 0.0;
   TimeType target_max_load_                         = 0.0;

src/vt/vrt/collection/balance/temperedwmin/temperedwmin.cc

@@ -46,6 +46,8 @@
 #include "vt/vrt/collection/balance/lb_common.h"
 #include "vt/vrt/collection/balance/model/load_model.h"
 
+#include <algorithm>
+
 namespace vt { namespace vrt { namespace collection { namespace lb {
 
 void TemperedWMin::init(objgroup::proxy::Proxy<TemperedWMin> in_proxy) {
@@ -122,4 +124,13 @@ TimeType TemperedWMin::getModeledWork(const elm::ElementIDStruct& obj) const {
     beta_ * load_model_->getModeledComm(obj, when) + gamma_;
 }
 
+bool TemperedWMin::canMigrate() const {
+  auto const this_node    = theContext()->getNode();
+  auto const another_rank = std::find_if(
+    load_info_.begin(), load_info_.end(),
+    [this_node](auto const& elm) { return elm.first != this_node; }
+  );
+  return (not cur_objs_.empty()) and (another_rank != load_info_.end());
+}
+
 }}}} // namespace vt::vrt::collection::lb

src/vt/vrt/collection/balance/temperedwmin/temperedwmin.h

@@ -61,6 +61,10 @@ struct TemperedWMin : TemperedLB {
   void inputParams(balance::SpecEntry* spec) override;
 
 protected:
+  /**
+   * Allow migration when there are objects to migrate and other ranks are known
+   */
+  bool canMigrate() const override;
   /**
    * All ranks are allowed to initiate the information propagation stage
    */