#1934: Update unit tests to check all relevant fields in NodeLBData

Author: thearusable

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

@@ -103,7 +103,7 @@ NodeLBData::getUserData() const {
   return &lb_data_->user_defined_lb_info_;
 }
 
-std::unordered_map<PhaseType, DataMapType> const*
+DataMapBufferType const*
 NodeLBData::getPhaseAttributes() const {
   return &lb_data_->node_user_attributes_;
 }
@@ -147,6 +147,7 @@ void NodeLBData::resizeLBDataHistory(uint32_t new_hist_len) {
     lb_data_->node_subphase_comm_.resize(new_hist_len);
     lb_data_->user_defined_lb_info_.resize(new_hist_len);
     lb_data_->user_defined_json_.resize(new_hist_len);
+    lb_data_->node_user_attributes_.resize(new_hist_len);
   }
 
   NodeLBData::node_migrate_.clear();

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

@@ -194,7 +194,7 @@ struct NodeLBData : runtime::component::Component<NodeLBData> {
    *
    * \return an observer pointer to the user-defined attributes
    */
-  std::unordered_map<PhaseType, DataMapType> const* getPhaseAttributes() const;
+  DataMapBufferType const* getPhaseAttributes() const;
 
   /**
    * \internal \brief Get stored object comm data for a specific phase

tests/unit/collection/test_lb_data_holder.cc

@@ -271,7 +271,7 @@ TEST_F(TestLBDataHolder, test_lb_entity_attributes) {
   auto id = vt::vrt::collection::balance::ElementIDStruct{524291, 0};
 
   LBDataHolder testObj(json);
-  EXPECT_TRUE(testObj.node_user_attributes_.find(0) != testObj.node_user_attributes_.end());
+  EXPECT_TRUE(testObj.node_user_attributes_.contains(0));
   EXPECT_TRUE(testObj.node_user_attributes_[0].find(id) != testObj.node_user_attributes_[0].end());
   auto attributes = testObj.node_user_attributes_[0][id];
   EXPECT_EQ(123, std::get<int>(attributes["intSample"]));

tests/unit/collection/test_lb_data_retention.cc

@@ -65,14 +65,16 @@ void validatePersistedPhases(std::vector<PhaseType> expected_phases) {
     EXPECT_EQ(expected_phases.size(), theNodeLBData()->getLBData()->node_data_.size());
     EXPECT_EQ(expected_phases.size(), theNodeLBData()->getLBData()->node_subphase_comm_.size());
     EXPECT_EQ(expected_phases.size(), theNodeLBData()->getLBData()->user_defined_json_.size());
-    // EXPECT_EQ(expected_phases.size(), theNodeLBData()->getLBData()->user_defined_lb_info_.size());
+    EXPECT_EQ(expected_phases.size(), theNodeLBData()->getLBData()->user_defined_lb_info_.size());
+    EXPECT_EQ(expected_phases.size(), theNodeLBData()->getLBData()->node_user_attributes_.size());
     // Check if each phase is present
     for(auto&& phase : expected_phases) {
       EXPECT_TRUE(theNodeLBData()->getLBData()->node_comm_.contains(phase));
       EXPECT_TRUE(theNodeLBData()->getLBData()->node_data_.contains(phase));
       EXPECT_TRUE(theNodeLBData()->getLBData()->node_subphase_comm_.contains(phase));
       EXPECT_TRUE(theNodeLBData()->getLBData()->user_defined_json_.contains(phase));
-      // EXPECT_TRUE(theNodeLBData()->getLBData()->user_defined_lb_info_.contains(phase));
+      EXPECT_TRUE(theNodeLBData()->getLBData()->user_defined_lb_info_.contains(phase));
+      EXPECT_TRUE(theNodeLBData()->getLBData()->node_user_attributes_.contains(phase));
     }
   #else
     (void)expected_phases;
@@ -81,21 +83,20 @@ void validatePersistedPhases(std::vector<PhaseType> expected_phases) {
     EXPECT_EQ(0, theNodeLBData()->getLBData()->node_subphase_comm_.size());
     EXPECT_EQ(0, theNodeLBData()->getLBData()->user_defined_json_.size());
     EXPECT_EQ(0, theNodeLBData()->getLBData()->user_defined_lb_info_.size());
+    EXPECT_EQ(0, theNodeLBData()->getLBData()->node_user_attributes_.size());
   #endif
 }
 
 struct TestCol : vt::Collection<TestCol,vt::Index1D> {
   unsigned int prev_calls_ = thePhase()->getCurrentPhase();
 
-  TestCol() {
-    // Insert dummy lb info data
-    valInsert("foo", 10, true, true, true);
-  }
-
   unsigned int prevCalls() { return prev_calls_++; }
 
-  static void colHandler(TestCol* col) {
+  static void insertValue(TestCol* col) {
+    col->valInsert("foo", 10, true, true, true);
+  }
 
+  static void colHandler(TestCol* col) {
     auto& lb_data = col->lb_data_;
     auto load_phase_count = lb_data.getLoadPhaseCount();
     auto comm_phase_count = lb_data.getCommPhaseCount();
@@ -175,6 +176,7 @@ TEST_F(TestLBDataRetention, test_lbdata_retention_last1) {
   for (int i=0; i<num_phases; ++i) {
     runInEpochCollective([&]{
       // Do some work.
+      proxy.broadcastCollective<TestCol::insertValue>();
       proxy.broadcastCollective<TestCol::colHandler>();
     });
     // Go to the next phase.
@@ -214,6 +216,7 @@ TEST_F(TestLBDataRetention, test_lbdata_retention_last2) {
   for (int i=0; i<num_phases; ++i) {
     runInEpochCollective([&]{
       // Do some work.
+      proxy.broadcastCollective<TestCol::insertValue>();
       proxy.broadcastCollective<TestCol::colHandler>();
     });
     // Go to the next phase.
@@ -253,6 +256,7 @@ TEST_F(TestLBDataRetention, test_lbdata_retention_last4) {
   for (int i=0; i<num_phases; ++i) {
     runInEpochCollective([&]{
       // Do some work.
+      proxy.broadcastCollective<TestCol::insertValue>();
       proxy.broadcastCollective<TestCol::colHandler>();
     });
     // Go to the next phase.
@@ -295,6 +299,7 @@ TEST_F(TestLBDataRetention, test_lbdata_config_retention_higher) {
   for (uint32_t i=0; i<theConfig()->vt_lb_data_retention * 2; ++i) {
     runInEpochCollective([&]{
       // Do some work.
+      proxy.broadcastCollective<TestCol::insertValue>();
       proxy.broadcastCollective<TestCol::colHandler>();
     });
     // Go to the next phase.
@@ -336,6 +341,7 @@ TEST_F(TestLBDataRetention, test_lbdata_retention_model_switch_1) {
   for (uint32_t i=0; i<first_stage_num_phases; ++i) {
     runInEpochCollective([&]{
       // Do some work.
+      proxy.broadcastCollective<TestCol::insertValue>();
       proxy.broadcastCollective<TestCol::colHandler>();
     });
     // Go to the next phase.
@@ -354,6 +360,7 @@ TEST_F(TestLBDataRetention, test_lbdata_retention_model_switch_1) {
   for (uint32_t i=0; i<first_stage_num_phases; ++i) {
     runInEpochCollective([&]{
       // Do some work.
+      proxy.broadcastCollective<TestCol::insertValue>();
       proxy.broadcastCollective<TestCol::colHandler>();
     });
     // Go to the next phase.
@@ -393,6 +400,7 @@ TEST_F(TestLBDataRetention, test_lbdata_retention_model_switch_2) {
   for (uint32_t i=0; i<first_stage_num_phases; ++i) {
     runInEpochCollective([&]{
       // Do some work.
+      proxy.broadcastCollective<TestCol::insertValue>();
       proxy.broadcastCollective<TestCol::colHandler>();
     });
     // Go to the next phase.
@@ -416,6 +424,7 @@ TEST_F(TestLBDataRetention, test_lbdata_retention_model_switch_2) {
   for (uint32_t i=0; i<10; ++i) {
     runInEpochCollective([&]{
       // Do some work.
+      proxy.broadcastCollective<TestCol::insertValue>();
       proxy.broadcastCollective<TestCol::colHandler>();
     });
     // Go to the next phase.