#2229: Add attributes field for rank and object

scripts/JSON_data_files_validator.py

@@ -64,6 +64,7 @@ def _get_valid_schema(self) -> Schema:
                             'range': [[int]],
                         },
                     },
+                    Optional('attributes'): dict
                 },
                 'phases': [
                     {
@@ -88,7 +89,8 @@ def _get_valid_schema(self) -> Schema:
                                     }
                                 ],
                                 'time': float,
-                                Optional('user_defined'): dict
+                                Optional('user_defined'): dict,
+                                Optional('attributes'): dict
                             },
                         ],
                         Optional('communications'): [

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

@@ -125,6 +125,22 @@ std::unique_ptr<nlohmann::json> LBDataHolder::metadataToJson() const {
   return std::make_unique<nlohmann::json>(std::move(j));
 }
 
+std::unique_ptr<nlohmann::json> LBDataHolder::rankAttributesToJson() const {
+  nlohmann::json j;
+
+  for (auto const& [key, value] : rank_attributes_) {
+    if (std::holds_alternative<int>(value)) {
+      j["attributes"][key] = std::get<int>(value);
+    } else if (std::holds_alternative<double>(value)) {
+      j["attributes"][key] = std::get<double>(value);
+    } else if (std::holds_alternative<std::string>(value)) {
+      j["attributes"][key] = std::get<std::string>(value);
+    }
+  }
+
+  return std::make_unique<nlohmann::json>(std::move(j));
+}
+
 std::unique_ptr<nlohmann::json> LBDataHolder::toJson(PhaseType phase) const {
   using json = nlohmann::json;
 
@@ -150,6 +166,20 @@ std::unique_ptr<nlohmann::json> LBDataHolder::toJson(PhaseType phase) const {
       }
       outputEntity(j["tasks"][i]["entity"], id);
 
+      if (node_user_attributes_.find(phase) != node_user_attributes_.end()) {
+        if (node_user_attributes_.at(phase).find(id) != node_user_attributes_.at(phase).end()) {
+          for (auto const& [key, value] : node_user_attributes_.at(phase).at(id)) {
+            if (std::holds_alternative<int>(value)) {
+              j["tasks"][i]["attributes"][key] = std::get<int>(value);
+            } else if (std::holds_alternative<double>(value)) {
+              j["tasks"][i]["attributes"][key] = std::get<double>(value);
+            } else if (std::holds_alternative<std::string>(value)) {
+              j["tasks"][i]["attributes"][key] = std::get<std::string>(value);
+            }
+          }
+        }
+      }
+
       auto const& subphase_times = elm.second.subphase_loads;
       std::size_t const subphases = subphase_times.size();
       if (subphases != 0) {
@@ -303,6 +333,18 @@ LBDataHolder::LBDataHolder(nlohmann::json const& j)
                 }
               }
             }
+
+            if (task.find("attributes") != task.end()) {
+              for (auto const& [key, value] : task["attributes"].items()) {
+                if (value.is_number_integer()) {
+                  node_user_attributes_[id][elm][key] = value.get<int>();
+                } else if (value.is_number_float()) {
+                  node_user_attributes_[id][elm][key] = value.get<double>();
+                } else if (value.is_string()) {
+                  node_user_attributes_[id][elm][key] = value.get<std::string>();
+                }
+              }
+            }
           }
         }
       }
@@ -441,6 +483,18 @@ void LBDataHolder::readMetadata(nlohmann::json const& j) {
         }
       }
     }
+    // load rank user atrributes
+    if (metadata.find("attributes") != metadata.end()) {
+      for (auto const& [key, value] : metadata["attributes"].items()) {
+        if (value.is_number_integer()) {
+          rank_attributes_[key] = value.get<int>();
+        } else if (value.is_number_float()) {
+          rank_attributes_[key] = value.get<double>();
+        } else if (value.is_string()) {
+          rank_attributes_[key] = value.get<std::string>();
+        }
+      }
+    }
   }
 }
 

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

@@ -85,6 +85,8 @@ struct LBDataHolder {
     s | skipped_phases_;
     s | identical_phases_;
     s | user_defined_lb_info_;
+    s | node_user_attributes_;
+    s | rank_attributes_;
   }
 
   /**
@@ -97,12 +99,19 @@ struct LBDataHolder {
   std::unique_ptr<nlohmann::json> toJson(PhaseType phase) const;
 
   /**
-   * \brief Output a LB phase's metdadata to JSON
+   * \brief Output a LB phase's metadata to JSON
    *
    * \return the json data structure
    */
   std::unique_ptr<nlohmann::json> metadataToJson() const;
 
+  /**
+   * \brief Output a LB rank attributes metadata to JSON
+   *
+   * \return the json data structure
+   */
+  std::unique_ptr<nlohmann::json> rankAttributesToJson() const;
+
   /**
    * \brief Clear all LB data
    */
@@ -125,6 +134,8 @@ struct LBDataHolder {
   void readMetadata(nlohmann::json const& j);
 
 public:
+  /// Node attributes for the current rank
+  ElmUserDataType rank_attributes_;
   /// Node timings for each local object
   std::unordered_map<PhaseType, LoadMapType> node_data_;
   /// Node communication graph for each local object
@@ -139,6 +150,8 @@ struct LBDataHolder {
   std::unordered_map<PhaseType, std::shared_ptr<nlohmann::json>> user_per_phase_json_;
   /// User-defined data from each phase for LB
   std::unordered_map<PhaseType, DataMapType> user_defined_lb_info_;
+  /// User-defined attributes from each phase
+  std::unordered_map<PhaseType, DataMapType> node_user_attributes_;
   /// Node indices for each ID along with the proxy ID
   std::unordered_map<ElementIDStruct, std::tuple<VirtualProxyType, std::vector<uint64_t>>> node_idx_;
   /// Map from id to objgroup proxy

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

@@ -102,6 +102,11 @@ NodeLBData::getUserData() const {
   return &lb_data_->user_defined_lb_info_;
 }
 
+std::unordered_map<PhaseType, DataMapType> const*
+NodeLBData::getPhaseAttributes() const {
+  return &lb_data_->node_user_attributes_;
+}
+
 std::unordered_map<PhaseType, CommMapType> const* NodeLBData::getNodeComm() const {
   return &lb_data_->node_comm_;
 }
@@ -110,6 +115,10 @@ std::unordered_map<PhaseType, std::unordered_map<SubphaseType, CommMapType>> con
   return &lb_data_->node_subphase_comm_;
 }
 
+ElmUserDataType const* NodeLBData::getNodeAttributes() const {
+  return &lb_data_->rank_attributes_;
+}
+
 CommMapType* NodeLBData::getNodeComm(PhaseType phase) {
   auto iter = lb_data_->node_comm_.find(phase);
   return (iter != lb_data_->node_comm_.end()) ? &iter->second : nullptr;
@@ -127,6 +136,7 @@ void NodeLBData::startIterCleanup(PhaseType phase, unsigned int look_back) {
     lb_data_->node_comm_.erase(phase - look_back);
     lb_data_->node_subphase_comm_.erase(phase - look_back);
     lb_data_->user_defined_lb_info_.erase(phase - look_back);
+    lb_data_->node_user_attributes_.erase(phase - look_back);
   }
 
   // Clear migrate lambdas and proxy lookup since LB is complete
@@ -215,6 +225,10 @@ void NodeLBData::createLBDataFile() {
     if(phasesMetadata) {
        metadata["phases"] = *phasesMetadata;
     }
+    auto attributesMetadata = lb_data_->rankAttributesToJson();
+    if(attributesMetadata) {
+      metadata["attributes"] = *attributesMetadata;
+    }
     lb_data_writer_ = std::make_unique<JSONAppender>(
       "phases", metadata, file_name, compress
     );
@@ -352,6 +366,11 @@ void NodeLBData::addNodeLBData(
         lb_data_->user_defined_lb_info_[phase][id][key] = val->toVariant();
       }
     );
+    storable->collectAttributes(
+      [&](std::string const& key, auto val) {
+        lb_data_->node_user_attributes_[phase][id][key] = val->toVariant();
+      }
+    );
   }
 
   in->updatePhase(1);

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

@@ -188,6 +188,13 @@ struct NodeLBData : runtime::component::Component<NodeLBData> {
    */
   std::unordered_map<PhaseType, DataMapType> const* getUserData() const;
 
+  /**
+   * \internal \brief Get the user-defined attributes
+   *
+   * \return an observer pointer to the user-defined attributes
+   */
+  std::unordered_map<PhaseType, DataMapType> const* getPhaseAttributes() const;
+
   /**
    * \internal \brief Get stored object comm data for a specific phase
    *
@@ -204,6 +211,13 @@ struct NodeLBData : runtime::component::Component<NodeLBData> {
    */
   std::unordered_map<PhaseType, std::unordered_map<SubphaseType, CommMapType>> const* getNodeSubphaseComm() const;
 
+  /**
+   * \internal \brief Get stored node attributes
+   *
+   * \return an observer pointer to the node attributes
+   */
+  ElmUserDataType const* getNodeAttributes() const;
+
   /**
    * \internal \brief Test if this node has an object to migrate
    *

src/vt/vrt/collection/types/storage/storable.h

@@ -94,13 +94,14 @@ struct Storable {
    * \param[in] u the value
    * \param[in] dump_to_json whether to include in json file
    * \param[in] provide_to_lb whether to provide the data to the load balancers
+   * \param[in] dump_to_attributes whether to dump to attributes in JSON output
    *
    * \note If \c provide_to_lb is enabled, the data must be convertible to an
    * integer, double, or string
    */
   template <typename U>
   void valInsert(
-    std::string const& str, U&& u, bool dump_to_json, bool provide_to_lb
+    std::string const& str, U&& u, bool dump_to_json, bool provide_to_lb, bool dump_to_attributes
   );
 
   /**
@@ -160,6 +161,20 @@ struct Storable {
     }
   }
 
+  /**
+   * \brief Traverse the map and calls Callable on each attribute
+   *
+   * \param[in] c the callable
+   */
+  template <typename Callable>
+  void collectAttributes(Callable&& c) {
+    for (auto const& [key, value] : map_) {
+      if (value->isAttribute()) {
+        c(key, value.get());
+      }
+    }
+  }
+
 private:
   /// Map of type-erased, stored values
   std::unordered_map<std::string, std::unique_ptr<StoreElmBase>> map_;

src/vt/vrt/collection/types/storage/storable.impl.h

@@ -60,22 +60,22 @@ void Storable::valInsert(std::string const& str, U&& u) {
     std::forward_as_tuple(str),
     std::forward_as_tuple(
       std::make_unique<StoreElm<typename std::decay<U>::type>>(
-        std::forward<U>(u), false, false
+        std::forward<U>(u), false, false, false
       )
     )
   );
 }
 
 template <typename U>
 void Storable::valInsert(
-  std::string const& str, U&& u, bool dump_to_json, bool provide_to_lb
+  std::string const& str, U&& u, bool dump_to_json, bool provide_to_lb, bool dump_to_attributes
 ) {
   map_.emplace(
     std::piecewise_construct,
     std::forward_as_tuple(str),
     std::forward_as_tuple(
       std::make_unique<StoreElm<typename std::decay<U>::type>>(
-        std::forward<U>(u), dump_to_json, provide_to_lb
+        std::forward<U>(u), dump_to_json, provide_to_lb, dump_to_attributes
       )
     )
   );

src/vt/vrt/collection/types/storage/store_elm.h

@@ -74,10 +74,12 @@ struct StoreElmBase {
    *
    * \param[in] dump_to_json whether to dump to JSON output
    * \param[in] provide_to_lb whether to provide to the LB
+   * \param[in] dump_to_attributes whether to dump to attributes in JSON output
    */
-  StoreElmBase(bool dump_to_json, bool provide_to_lb)
+  StoreElmBase(bool dump_to_json, bool provide_to_lb, bool dump_to_attributes)
     : dump_to_json_(dump_to_json),
-      provide_to_lb_(provide_to_lb)
+      provide_to_lb_(provide_to_lb),
+      dump_to_attributes_(dump_to_attributes)
   {}
 
   virtual ~StoreElmBase() {}
@@ -121,6 +123,7 @@ struct StoreElmBase {
   void serialize(SerializerT& s) {
     s | dump_to_json_;
     s | provide_to_lb_;
+    s | dump_to_attributes_;
   }
 
   /**
@@ -139,6 +142,15 @@ struct StoreElmBase {
     return provide_to_lb_;
   }
 
+  /**
+   * \brief Whether the value should be dumped to the json attributes field
+   *
+   * \return whether it is an attribute
+   */
+  bool isAttribute() const {
+    return dump_to_attributes_;
+  }
+
   /**
    * \brief Generate the json because it is jsonable
    *
@@ -177,6 +189,7 @@ struct StoreElmBase {
 protected:
   bool dump_to_json_ = false;
   bool provide_to_lb_ = false;
+  bool dump_to_attributes_ = false;
 };
 
 /**
@@ -209,10 +222,11 @@ struct StoreElm<
    * \param[in] u the value
    * \param[in] dump_to_json whether to dump to json
    * \param[in] provide_to_lb whether to provide to LB
+   * \param[in] dump_to_attributes whether to dump to attributes in JSON output
    */
   template <typename U>
-  explicit StoreElm(U&& u, bool dump_to_json, bool provide_to_lb)
-    : StoreElmBase(dump_to_json, provide_to_lb),
+  explicit StoreElm(U&& u, bool dump_to_json, bool provide_to_lb, bool dump_to_attributes)
+    : StoreElmBase(dump_to_json, provide_to_lb, dump_to_attributes),
       elm_(std::forward<U>(u))
   { }
 
@@ -328,10 +342,11 @@ struct StoreElm<
    * \param[in] u the value
    * \param[in] dump_to_json whether to dump to json
    * \param[in] provide_to_lb whether to provide to LB
+   * \param[in] dump_to_attributes whether to dump to attributes in JSON output
    */
   template <typename U>
-  explicit StoreElm(U&& u, bool dump_to_json, bool provide_to_lb)
-    : StoreElmBase(dump_to_json, provide_to_lb),
+  explicit StoreElm(U&& u, bool dump_to_json, bool provide_to_lb, bool dump_to_attributes)
+    : StoreElmBase(dump_to_json, provide_to_lb, dump_to_attributes),
       wrapper_(detail::ByteWrapper<T>{std::forward<U>(u)})
   { }