Merge pull request #2228 from DARMA-tasking/2227-add-user-lb-info-to-…

…load-models 2227 Add user LB info to load models
DARMA-tasking · Dec 4, 2023 · e307da5 · e307da5
2 parents ab2cdbe + 437d23e
commit e307da5
Show file tree

Hide file tree

Showing 32 changed files with 395 additions and 54 deletions.
diff --git a/src/vt/vrt/collection/balance/baselb/baselb.cc b/src/vt/vrt/collection/balance/baselb/baselb.cc
@@ -103,7 +103,6 @@ void BaseLB::importProcessorData(
 
   comm_data = &comm_in;
   base_stats_ = &in_stats;
-  user_data_ = &in_data_map;
 }
 
 void BaseLB::getArgs(PhaseType phase) {
@@ -191,8 +190,12 @@ std::shared_ptr<const balance::Reassignment> BaseLB::normalizeReassignments() {
       auto const raw_load_summary = getObjectRawLoads(
         load_model_, obj_id, {PhaseOffset::NEXT_PHASE, PhaseOffset::WHOLE_PHASE}
       );
+      auto const obj_user_data = getObjectUserData(
+        load_model_, obj_id, {PhaseOffset::NEXT_PHASE, PhaseOffset::WHOLE_PHASE}
+      );
+
       depart_map[dest].push_back(
-        std::make_tuple(obj_id, load_summary, raw_load_summary)
+        std::make_tuple(obj_id, load_summary, raw_load_summary, obj_user_data)
       );
     }
 
@@ -230,8 +233,9 @@ void BaseLB::notifyNewHostNodeOfObjectsArriving(
     auto const obj_id = std::get<0>(arrival);
     auto const& load_summary = std::get<1>(arrival);
     auto const& raw_load_summary = std::get<2>(arrival);
+    auto const& obj_user_data = std::get<3>(arrival);
     pending_reassignment_->arrive_[obj_id] = std::make_tuple(
-      load_summary, raw_load_summary
+      load_summary, raw_load_summary, obj_user_data
     );
   }
 }

diff --git a/src/vt/vrt/collection/balance/baselb/baselb.h b/src/vt/vrt/collection/balance/baselb/baselb.h
@@ -78,7 +78,7 @@ struct BaseLB {
   using StatisticMapType = std::unordered_map<lb::Statistic, QuantityType>;
   using LoadSummary      = balance::LoadSummary;
   using ObjLoadListType  = std::vector<
-    std::tuple<ObjIDType, LoadSummary, LoadSummary>
+    std::tuple<ObjIDType, LoadSummary, LoadSummary, balance::ElmUserDataType>
   >;
   using ObjDestinationListType = std::vector<std::tuple<ObjIDType, NodeType>>;
 
@@ -169,7 +169,6 @@ struct BaseLB {
 
   TimeType start_time_                                = TimeType{0.0};
   ElementCommType const* comm_data                    = nullptr;
-  balance::DataMapType const* user_data_              = nullptr;
   objgroup::proxy::Proxy<BaseLB> proxy_               = {};
   PhaseType phase_                                    = 0;
   std::unique_ptr<balance::ConfigEntry> config_entry_ = nullptr;

diff --git a/src/vt/vrt/collection/balance/lb_common.cc b/src/vt/vrt/collection/balance/lb_common.cc
@@ -97,6 +97,26 @@ LoadSummary getObjectRawLoads(
   return ret;
 }
 
+ElmUserDataType getObjectUserData(
+  std::shared_ptr<LoadModel> model, ElementIDStruct object, PhaseOffset when
+) {
+  return getObjectUserData(model.get(), object, when);
+}
+
+ElmUserDataType getObjectUserData(
+  LoadModel* model, ElementIDStruct object, PhaseOffset when
+) {
+  ElmUserDataType ret;
+
+  if (model->hasUserData()) {
+    ret = model->getUserData(
+      object, {when.phases, PhaseOffset::WHOLE_PHASE}
+    );
+  }
+
+  return ret;
+}
+
 LoadSummary getNodeLoads(std::shared_ptr<LoadModel> model, PhaseOffset when)
 {
   LoadSummary ret;

diff --git a/src/vt/vrt/collection/balance/lb_common.h b/src/vt/vrt/collection/balance/lb_common.h
@@ -112,14 +112,14 @@ struct LoadSummary {
   }
 };
 
+/// User-defined LB values
+using UserDataValueType = std::variant<int, double, std::string>;
+using ElmUserDataType = std::unordered_map<std::string, UserDataValueType>;
+
 using LoadMapType         = std::unordered_map<ElementIDStruct, LoadSummary>;
 using SubphaseLoadMapType = std::unordered_map<ElementIDStruct, std::vector<LoadType>>;
-
 /// User-defined LB values map
-using DataMapType         = std::unordered_map<
-  ElementIDStruct,
-  std::unordered_map<std::string, std::variant<int, double, std::string>>
->;
+using DataMapType         = std::unordered_map<ElementIDStruct, ElmUserDataType>;
 
 struct Reassignment {
   // Include the subject node so that these structures can be formed
@@ -130,7 +130,7 @@ struct Reassignment {
   int32_t global_migration_count;
   std::unordered_map<ElementIDStruct, NodeType> depart_;
   std::unordered_map<
-    ElementIDStruct, std::tuple<LoadSummary, LoadSummary>
+    ElementIDStruct, std::tuple<LoadSummary, LoadSummary, ElmUserDataType>
   > arrive_;
 };
 
@@ -175,6 +175,14 @@ LoadSummary getObjectRawLoads(
   LoadModel* model, ElementIDStruct object, PhaseOffset when
 );
 
+ElmUserDataType getObjectUserData(
+  std::shared_ptr<LoadModel> model, ElementIDStruct object, PhaseOffset when
+);
+
+ElmUserDataType getObjectUserData(
+  LoadModel* model, ElementIDStruct object, PhaseOffset when
+);
+
 LoadSummary getNodeLoads(std::shared_ptr<LoadModel> model, PhaseOffset when);
 
 } /* end namespace balance */

diff --git a/src/vt/vrt/collection/balance/lb_data_holder.cc b/src/vt/vrt/collection/balance/lb_data_holder.cc
@@ -290,6 +290,19 @@ LBDataHolder::LBDataHolder(nlohmann::json const& j)
                 }
               }
             }
+
+            if (task.find("user_defined") != task.end()) {
+              for (auto const& [key, value] : task["user_defined"].items()) {
+                if (value.is_string()) {
+                  user_defined_lb_info_[id][elm][key] =
+                    value.template get<std::string>();
+                }
+                if (value.is_number()) {
+                  user_defined_lb_info_[id][elm][key] =
+                    value.template get<double>();
+                }
+              }
+            }
           }
         }
       }

diff --git a/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc b/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc
@@ -164,7 +164,8 @@ void LBManager::setLoadModel(std::shared_ptr<LoadModel> model) {
   model_ = model;
   auto nlb_data = theNodeLBData();
   model_->setLoads(nlb_data->getNodeLoad(),
-                   nlb_data->getNodeComm());
+                   nlb_data->getNodeComm(),
+                   nlb_data->getUserData());
 }
 
 template <typename LB>

diff --git a/src/vt/vrt/collection/balance/model/comm_overhead.cc b/src/vt/vrt/collection/balance/model/comm_overhead.cc
@@ -55,9 +55,10 @@ CommOverhead::CommOverhead(
 { }
 
 void CommOverhead::setLoads(std::unordered_map<PhaseType, LoadMapType> const* proc_load,
-                            std::unordered_map<PhaseType, CommMapType> const* proc_comm) {
+                            std::unordered_map<PhaseType, CommMapType> const* proc_comm,
+                            std::unordered_map<PhaseType, DataMapType> const* user_data) {
   proc_comm_ = proc_comm;
-  ComposedModel::setLoads(proc_load, proc_comm);
+  ComposedModel::setLoads(proc_load, proc_comm, user_data);
 }
 
 LoadType

diff --git a/src/vt/vrt/collection/balance/model/comm_overhead.h b/src/vt/vrt/collection/balance/model/comm_overhead.h
@@ -66,7 +66,8 @@ struct CommOverhead : public ComposedModel {
   );
 
   void setLoads(std::unordered_map<PhaseType, LoadMapType> const* proc_load,
-                std::unordered_map<PhaseType, CommMapType> const* proc_comm) override;
+                std::unordered_map<PhaseType, CommMapType> const* proc_comm,
+                std::unordered_map<PhaseType, DataMapType> const* user_data) override;
 
   LoadType getModeledLoad(ElementIDStruct object, PhaseOffset when) const override;
 

diff --git a/src/vt/vrt/collection/balance/model/composed_model.cc b/src/vt/vrt/collection/balance/model/composed_model.cc
@@ -46,8 +46,9 @@
 namespace vt { namespace vrt { namespace collection { namespace balance {
 
 void ComposedModel::setLoads(std::unordered_map<PhaseType, LoadMapType> const* proc_load,
-                             std::unordered_map<PhaseType, CommMapType> const* proc_comm) {
-  base_->setLoads(proc_load, proc_comm);
+                             std::unordered_map<PhaseType, CommMapType> const* proc_comm,
+                             std::unordered_map<PhaseType, DataMapType> const* user_data) {
+  base_->setLoads(proc_load, proc_comm, user_data);
 }
 
 void ComposedModel::updateLoads(PhaseType last_completed_phase) {
@@ -72,6 +73,14 @@ LoadType ComposedModel::getRawLoad(ElementIDStruct object, PhaseOffset when) con
   return base_->getRawLoad(object, when);
 }
 
+bool ComposedModel::hasUserData() const {
+  return base_->hasUserData();
+}
+
+ElmUserDataType ComposedModel::getUserData(ElementIDStruct object, PhaseOffset when) const {
+  return base_->getUserData(object, when);
+}
+
 unsigned int ComposedModel::getNumPastPhasesNeeded(unsigned int look_back) const
 {
   return base_->getNumPastPhasesNeeded(look_back);

diff --git a/src/vt/vrt/collection/balance/model/composed_model.h b/src/vt/vrt/collection/balance/model/composed_model.h
@@ -65,14 +65,17 @@ class ComposedModel : public LoadModel
   explicit ComposedModel(std::shared_ptr<LoadModel> base) : base_(base) {}
 
   void setLoads(std::unordered_map<PhaseType, LoadMapType> const* proc_load,
-                std::unordered_map<PhaseType, CommMapType> const* proc_comm) override;
+                std::unordered_map<PhaseType, CommMapType> const* proc_comm,
+                std::unordered_map<PhaseType, DataMapType> const* user_data) override;
 
   void updateLoads(PhaseType last_completed_phase) override;
 
   LoadType getModeledLoad(ElementIDStruct object, PhaseOffset when) const override;
   LoadType getModeledComm(ElementIDStruct object, PhaseOffset when) const override;
   bool hasRawLoad() const override;
   LoadType getRawLoad(ElementIDStruct object, PhaseOffset when) const override;
+  bool hasUserData() const override;
+  ElmUserDataType getUserData(ElementIDStruct object, PhaseOffset when) const override;
   unsigned int getNumPastPhasesNeeded(unsigned int look_back) const override;
 
   ObjectIterator begin() const override;

diff --git a/src/vt/vrt/collection/balance/model/load_model.h b/src/vt/vrt/collection/balance/model/load_model.h
@@ -102,7 +102,7 @@ struct LoadMapObjectIterator : public ObjectIteratorImpl {
 
 struct DualLoadMapObjectIterator : public ObjectIteratorImpl {
   using DualLoadMapType = std::unordered_map<
-    ElementIDStruct, std::tuple<LoadSummary, LoadSummary>
+    ElementIDStruct, std::tuple<LoadSummary, LoadSummary, ElmUserDataType>
   >;
   using map_iterator_type = DualLoadMapType::const_iterator;
   using iterator_category = std::iterator_traits<map_iterator_type>::iterator_category;
@@ -196,7 +196,8 @@ struct LoadModel
    */
   virtual void setLoads(
     std::unordered_map<PhaseType, LoadMapType> const* proc_load,
-    std::unordered_map<PhaseType, CommMapType> const* proc_comm
+    std::unordered_map<PhaseType, CommMapType> const* proc_comm,
+    std::unordered_map<PhaseType, DataMapType> const* user_data
   ) = 0;
 
   /**
@@ -247,6 +248,26 @@ struct LoadModel
     return 0.0;
   };
 
+  /**
+   * \brief Whether or not the model is based on the RawData model
+   */
+  virtual bool hasUserData() const { return false; }
+
+  /**
+   * \brief Provide the given object's user data during a specified interval
+   *
+   * \param[in] object The object whose user data is desired
+   * \param[in] when The interval in which the user data is desired
+   *
+   * \return The user data associated with the object
+   */
+  virtual ElmUserDataType getUserData(ElementIDStruct object, PhaseOffset when) const {
+    vtAbort(
+      "LoadModel::getUserData() called on a model that does not implement it"
+    );
+    return ElmUserDataType{};
+  };
+
   /**
    * \brief Provide an estimate of the communication cost for a given object
    * during a specified interval

diff --git a/src/vt/vrt/collection/balance/model/naive_persistence.cc b/src/vt/vrt/collection/balance/model/naive_persistence.cc
@@ -66,6 +66,14 @@ LoadType NaivePersistence::getRawLoad(ElementIDStruct object, PhaseOffset offset
   return ComposedModel::getRawLoad(object, offset);
 }
 
+ElmUserDataType NaivePersistence::getUserData(ElementIDStruct object, PhaseOffset offset) const
+{
+  if (offset.phases >= 0)
+    offset.phases = -1;
+
+  return ComposedModel::getUserData(object, offset);
+}
+
 unsigned int NaivePersistence::getNumPastPhasesNeeded(unsigned int look_back) const
 {
   return ComposedModel::getNumPastPhasesNeeded(std::max(1u, look_back));

diff --git a/src/vt/vrt/collection/balance/model/naive_persistence.h b/src/vt/vrt/collection/balance/model/naive_persistence.h
@@ -62,6 +62,7 @@ struct NaivePersistence : public ComposedModel {
   explicit NaivePersistence(std::shared_ptr<balance::LoadModel> base);
   LoadType getModeledLoad(ElementIDStruct object, PhaseOffset when) const override;
   LoadType getRawLoad(ElementIDStruct object, PhaseOffset offset) const override;
+  ElmUserDataType getUserData(ElementIDStruct object, PhaseOffset when) const override;
   unsigned int getNumPastPhasesNeeded(unsigned int look_back) const override;
 }; // class NaivePersistence
 

diff --git a/src/vt/vrt/collection/balance/model/per_collection.cc b/src/vt/vrt/collection/balance/model/per_collection.cc
@@ -56,10 +56,11 @@ void PerCollection::addModel(CollectionID proxy, std::shared_ptr<LoadModel> mode
 }
 
 void PerCollection::setLoads(std::unordered_map<PhaseType, LoadMapType> const* proc_load,
-                             std::unordered_map<PhaseType, CommMapType> const* proc_comm) {
+                             std::unordered_map<PhaseType, CommMapType> const* proc_comm,
+                             std::unordered_map<PhaseType, DataMapType> const* user_data) {
   for (auto& m : models_)
-    m.second->setLoads(proc_load, proc_comm);
-  ComposedModel::setLoads(proc_load, proc_comm);
+    m.second->setLoads(proc_load, proc_comm, user_data);
+  ComposedModel::setLoads(proc_load, proc_comm, user_data);
 }
 
 void PerCollection::updateLoads(PhaseType last_completed_phase) {
@@ -100,6 +101,27 @@ LoadType PerCollection::getRawLoad(ElementIDStruct object, PhaseOffset when) con
   return ComposedModel::getRawLoad(object, when);
 }
 
+bool PerCollection::hasUserData() const {
+  // Only return true if all possible paths lead to true
+  bool has_user_data = true;
+
+  for (auto it = models_.begin(); it != models_.end(); ++it) {
+    has_user_data = has_user_data and it->second->hasUserData();
+  }
+
+  return has_user_data and ComposedModel::hasUserData();
+}
+
+ElmUserDataType PerCollection::getUserData(ElementIDStruct object, PhaseOffset when) const {
+  // See if some specific model has been given for the object in question
+  auto mi = models_.find(theNodeLBData()->getCollectionProxyForElement(object));
+  if (mi != models_.end())
+    return mi->second->getUserData(object, when);
+
+  // Otherwise, default to the given base model
+  return ComposedModel::getUserData(object, when);
+}
+
 unsigned int PerCollection::getNumPastPhasesNeeded(unsigned int look_back) const
 {
   unsigned int needed = ComposedModel::getNumPastPhasesNeeded(look_back);

diff --git a/src/vt/vrt/collection/balance/model/per_collection.h b/src/vt/vrt/collection/balance/model/per_collection.h
@@ -74,13 +74,16 @@ struct PerCollection : public ComposedModel
   void addModel(CollectionID proxy, std::shared_ptr<LoadModel> model);
 
   void setLoads(std::unordered_map<PhaseType, LoadMapType> const* proc_load,
-                std::unordered_map<PhaseType, CommMapType> const* proc_comm) override;
+                std::unordered_map<PhaseType, CommMapType> const* proc_comm,
+                std::unordered_map<PhaseType, DataMapType> const* user_data) override;
 
   void updateLoads(PhaseType last_completed_phase) override;
 
   LoadType getModeledLoad(ElementIDStruct object, PhaseOffset when) const override;
   bool hasRawLoad() const override;
   LoadType getRawLoad(ElementIDStruct object, PhaseOffset when) const override;
+  bool hasUserData() const override;
+  ElmUserDataType getUserData(ElementIDStruct object, PhaseOffset when) const override;
   unsigned int getNumPastPhasesNeeded(unsigned int look_back) const override;
 
 private:

diff --git a/src/vt/vrt/collection/balance/model/proposed_reassignment.cc b/src/vt/vrt/collection/balance/model/proposed_reassignment.cc
@@ -121,4 +121,18 @@ LoadType ProposedReassignment::getRawLoad(ElementIDStruct object, PhaseOffset wh
   return ComposedModel::getRawLoad(object, when);
 }
 
+ElmUserDataType ProposedReassignment::getUserData(ElementIDStruct object, PhaseOffset when) const
+{
+  auto a = reassignment_->arrive_.find(object);
+  if (a != reassignment_->arrive_.end()) {
+    return std::get<2>(a->second);
+  }
+
+  // Check this *after* arrivals to handle hypothetical self-migration
+  vtAssert(reassignment_->depart_.find(object) == reassignment_->depart_.end(),
+           "Departing object should not appear as a user data query subject");
+
+  return ComposedModel::getUserData(object, when);
+}
+
 }}}}
diff --git a/src/vt/vrt/collection/balance/model/proposed_reassignment.h b/src/vt/vrt/collection/balance/model/proposed_reassignment.h
@@ -59,6 +59,7 @@ struct ProposedReassignment : public ComposedModel {
   int getNumObjects() const override;
   LoadType getModeledLoad(ElementIDStruct object, PhaseOffset when) const override;
   LoadType getRawLoad(ElementIDStruct object, PhaseOffset when) const override;
+  ElmUserDataType getUserData(ElementIDStruct object, PhaseOffset when) const override;
 
 private:
   std::shared_ptr<const Reassignment> reassignment_;