#2201: temperedlb: add cluster summary to messages

DARMA-tasking · Nov 29, 2023 · 4e217ac · 4e217ac
1 parent ec201d9
commit 4e217ac
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 5 deletions.
diff --git a/src/vt/vrt/collection/balance/temperedlb/tempered_msgs.h b/src/vt/vrt/collection/balance/temperedlb/tempered_msgs.h
@@ -49,13 +49,24 @@
 #include <vector>
 #include <unordered_map>
 
+namespace vt::vrt::collection::lb {
+
+using SharedIDType     = int;
+using BytesType        = double;
+using ClusterSummaryType =
+  std::unordered_map<SharedIDType, std::tuple<BytesType, LoadType>>;
+
+} /* end namespace vt::vrt::collection::lb */
+
 namespace vt { namespace vrt { namespace collection { namespace balance {
 
 struct LoadMsg : vt::Message {
   using MessageParentType = vt::Message;
   vt_msg_serialize_required(); // node_load_
 
   using NodeLoadType = std::unordered_map<NodeType, LoadType>;
+  using NodeClusterSummaryType =
+    std::unordered_map<NodeType, lb::ClusterSummaryType>;
 
   LoadMsg() = default;
   LoadMsg(NodeType in_from_node, NodeLoadType const& in_node_load)
@@ -66,22 +77,32 @@ struct LoadMsg : vt::Message {
     return node_load_;
   }
 
+  NodeClusterSummaryType const& getNodeClusterSummary() const {
+    return node_cluster_summary_;
+  }
+
   void addNodeLoad(NodeType node, LoadType load) {
     node_load_[node] = load;
   }
 
+  void addNodeClusters(NodeType node, lb::ClusterSummaryType summary) {
+    node_cluster_summary_[node] = summary;
+  }
+
   NodeType getFromNode() const { return from_node_; }
 
   template <typename SerializerT>
   void serialize(SerializerT& s) {
     MessageParentType::serialize(s);
     s | from_node_;
     s | node_load_;
+    s | node_cluster_summary_;
   }
 
 private:
   NodeType from_node_     = uninitialized_destination;
   NodeLoadType node_load_ = {};
+  NodeClusterSummaryType node_cluster_summary_ = {};
 };
 
 struct LoadMsgAsync : LoadMsg {

diff --git a/src/vt/vrt/collection/balance/temperedlb/temperedlb.cc b/src/vt/vrt/collection/balance/temperedlb/temperedlb.cc
@@ -591,7 +591,7 @@ void TemperedLB::computeClusterSummary() {
   }
 }
 
-TemperedLB::BytesType TemperedLB::computeMemoryUsage() const {
+BytesType TemperedLB::computeMemoryUsage() const {
   // Compute bytes used by shared blocks mapped here based on object mapping
   auto const blocks_here = getSharedBlocksHere();
 
@@ -618,7 +618,7 @@ TemperedLB::BytesType TemperedLB::computeMemoryUsage() const {
   return rank_bytes_ + total_shared_bytes + max_object_working_bytes;
 }
 
-std::set<TemperedLB::SharedIDType> TemperedLB::getSharedBlocksHere() const {
+std::set<SharedIDType> TemperedLB::getSharedBlocksHere() const {
   std::set<SharedIDType> blocks_here;
   for (auto const& [obj, _] : cur_objs_) {
     if (obj_shared_block_.find(obj) != obj_shared_block_.end()) {
@@ -1015,6 +1015,9 @@ void TemperedLB::propagateRound(uint8_t k_cur, bool sync, EpochType epoch) {
         envelopeSetEpoch(msg->env, epoch);
       }
       msg->addNodeLoad(this_node, this_new_load_);
+      if (has_memory_data_) {
+        msg->addNodeClusters(this_node, cur_blocks_);
+      }
       proxy_[random_node].sendMsg<
         LoadMsgSync, &TemperedLB::propagateIncomingSync
       >(msg.get());
@@ -1025,6 +1028,9 @@ void TemperedLB::propagateRound(uint8_t k_cur, bool sync, EpochType epoch) {
         envelopeSetEpoch(msg->env, epoch);
       }
       msg->addNodeLoad(this_node, this_new_load_);
+      if (has_memory_data_) {
+        msg->addNodeClusters(this_node, cur_blocks_);
+      }
       proxy_[random_node].sendMsg<
         LoadMsgAsync, &TemperedLB::propagateIncomingAsync
       >(msg.get());

diff --git a/src/vt/vrt/collection/balance/temperedlb/temperedlb.h b/src/vt/vrt/collection/balance/temperedlb/temperedlb.h
@@ -68,8 +68,6 @@ struct TemperedLB : BaseLB {
   using ReduceMsgType  = vt::collective::ReduceNoneMsg;
   using QuantityType     = std::map<lb::StatisticQuantity, double>;
   using StatisticMapType = std::unordered_map<lb::Statistic, QuantityType>;
-  using SharedIDType     = int;
-  using BytesType        = double;
 
   TemperedLB() = default;
   TemperedLB(TemperedLB const&) = delete;
@@ -229,7 +227,7 @@ struct TemperedLB : BaseLB {
   /// Working bytes for each object
   std::unordered_map<ObjIDType, BytesType> obj_working_bytes_;
   /// Current assignment memory/load summary
-  std::unordered_map<SharedIDType, std::tuple<BytesType, LoadType>> cur_blocks_;
+  ClusterSummaryType cur_blocks_;
   /// User-defined memory threshold
   BytesType mem_thresh_ = 0;
 };