#2229: Change rank_attributes to std::variant

DARMA-tasking · Feb 16, 2024 · a621055 · a621055
1 parent e4323d4
commit a621055
Show file tree

Hide file tree

Showing 5 changed files with 52 additions and 12 deletions.
diff --git a/src/vt/vrt/collection/balance/lb_data_holder.cc b/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;
 
@@ -469,8 +485,15 @@ void LBDataHolder::readMetadata(nlohmann::json const& j) {
     }
     // load rank user atrributes
     if (metadata.find("attributes") != metadata.end()) {
-      rank_attributes_ = std::make_shared<nlohmann::json>();
-      *(rank_attributes_) = metadata["attributes"];
+      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>();
+        }
+      }
     }
   }
 }

diff --git a/src/vt/vrt/collection/balance/lb_data_holder.h b/src/vt/vrt/collection/balance/lb_data_holder.h
@@ -105,6 +105,13 @@ struct LBDataHolder {
    */
   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
    */
@@ -128,7 +135,7 @@ struct LBDataHolder {
 
 public:
   /// Node attributes for the current rank
-  std::shared_ptr<nlohmann::json> rank_attributes_;
+  ElmUserDataType rank_attributes_;
   /// Node timings for each local object
   std::unordered_map<PhaseType, LoadMapType> node_data_;
   /// Node communication graph for each local object

diff --git a/src/vt/vrt/collection/balance/node_lb_data.cc b/src/vt/vrt/collection/balance/node_lb_data.cc
@@ -115,8 +115,8 @@ std::unordered_map<PhaseType, std::unordered_map<SubphaseType, CommMapType>> con
   return &lb_data_->node_subphase_comm_;
 }
 
-std::shared_ptr<nlohmann::json> const NodeLBData::getNodeAttributes() const {
-  return lb_data_->rank_attributes_;
+ElmUserDataType const* NodeLBData::getNodeAttributes() const {
+  return &lb_data_->rank_attributes_;
 }
 
 CommMapType* NodeLBData::getNodeComm(PhaseType phase) {
@@ -225,8 +225,9 @@ void NodeLBData::createLBDataFile() {
     if(phasesMetadata) {
        metadata["phases"] = *phasesMetadata;
     }
-    if(lb_data_->rank_attributes_) {
-      metadata["attributes"] = *lb_data_->rank_attributes_;
+    auto attributesMetadata = lb_data_->rankAttributesToJson();
+    if(attributesMetadata) {
+      metadata["attributes"] = *attributesMetadata;
     }
     lb_data_writer_ = std::make_unique<JSONAppender>(
       "phases", metadata, file_name, compress

diff --git a/src/vt/vrt/collection/balance/node_lb_data.h b/src/vt/vrt/collection/balance/node_lb_data.h
@@ -214,9 +214,9 @@ struct NodeLBData : runtime::component::Component<NodeLBData> {
   /**
    * \internal \brief Get stored node attributes
    *
-   * \return an observer shared pointer to the node attributes
+   * \return an observer pointer to the node attributes
    */
-  std::shared_ptr<nlohmann::json> const getNodeAttributes() const;
+  ElmUserDataType const* getNodeAttributes() const;
 
   /**
    * \internal \brief Test if this node has an object to migrate

diff --git a/tests/unit/collection/test_lb_data_holder.cc b/tests/unit/collection/test_lb_data_holder.cc
@@ -226,16 +226,25 @@ TEST_F(TestLBDataHolder, test_lb_rank_attributes) {
       "type": "LBDatafile",
       "metadata" : {
         "attributes": {
-            "some_val": 123
+            "intSample": 123,
+            "doubleSample": 1.99,
+            "stringSample": "abc"
         }
       },
       "phases" : []
     }
   )"_json;
 
   LBDataHolder testObj(json);
-  EXPECT_TRUE(nullptr != testObj.rank_attributes_);
-  EXPECT_EQ(123, (*testObj.rank_attributes_)["some_val"]);
+  EXPECT_EQ(123, std::get<int>(testObj.rank_attributes_["intSample"]));
+  EXPECT_EQ(1.99, std::get<double>(testObj.rank_attributes_["doubleSample"]));
+  EXPECT_EQ("abc", std::get<std::string>(testObj.rank_attributes_["stringSample"]));
+
+  auto outAttributesJsonPtr = testObj.rankAttributesToJson();
+  ASSERT_TRUE(outAttributesJsonPtr != nullptr);
+  EXPECT_EQ(123, (*outAttributesJsonPtr)["attributes"]["intSample"]);
+  EXPECT_EQ(1.99, (*outAttributesJsonPtr)["attributes"]["doubleSample"]);
+  EXPECT_EQ("abc", (*outAttributesJsonPtr)["attributes"]["stringSample"]);
 }
 
 TEST_F(TestLBDataHolder, test_lb_entity_attributes) {