#2074: Fix logic for exsiting offlineLB test

DARMA-tasking · Sep 20, 2024 · e2b3b9b · e2b3b9b
1 parent f24e2cc
commit e2b3b9b
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 15 deletions.
diff --git a/src/vt/vrt/collection/balance/lb_data_restart_reader.h b/src/vt/vrt/collection/balance/lb_data_restart_reader.h
@@ -132,10 +132,8 @@ struct LBDataRestartReader : runtime::component::Component<LBDataRestartReader>
    */
   std::shared_ptr<const std::set<ElementIDStruct>> getDistro(PhaseType phase) const {
     auto iter = history_.find(phase);
-    if (iter != history_.end()) {
-      return iter->second;
-    }
-    return nullptr;
+    vtAssert(iter != history_.end(), "Must have a valid, not skipped phase");
+    return iter->second;
   }
 
   /**
@@ -193,7 +191,7 @@ struct LBDataRestartReader : runtime::component::Component<LBDataRestartReader>
   std::vector<bool> changed_distro_;
 
   /// History of mapping that was read in from the data files
-  std::map<PhaseType, std::shared_ptr<std::set<ElementIDStruct>>> history_;
+  std::unordered_map<PhaseType, std::shared_ptr<std::set<ElementIDStruct>>> history_;
 
   struct DepartMsg : vt::Message {
     DepartMsg(NodeType in_depart_node, PhaseType in_phase, ElementIDStruct in_elm)

diff --git a/src/vt/vrt/collection/balance/offlinelb/offlinelb.cc b/src/vt/vrt/collection/balance/offlinelb/offlinelb.cc
@@ -56,12 +56,10 @@ void OfflineLB::init(objgroup::proxy::Proxy<OfflineLB> in_proxy) {
 
 void OfflineLB::runLB(LoadType) {
   auto const distro = theLBDataReader()->getDistro(phase_ + 1);
-  if (distro) {
-    for (auto&& elm : *distro) {
-      migrateObjectTo(elm, theContext()->getNode());
-    }
-    theLBDataReader()->clearDistro(phase_ + 1);
+  for (auto&& elm : *distro) {
+    migrateObjectTo(elm, theContext()->getNode());
   }
+  theLBDataReader()->clearDistro(phase_ + 1);
 }
 
 }}}} /* end namespace vt::vrt::collection::lb */
diff --git a/tests/unit/lb/test_offlinelb.cc b/tests/unit/lb/test_offlinelb.cc
@@ -232,17 +232,20 @@ TEST_F(TestOfflineLB, test_offlinelb_2) {
   // Preapre configuration file
   std::string file_name = getUniqueFilenameWithRanks(".txt");
   std::ofstream out(file_name);
+
+  // NoLB for phases on skipped list and one phase before them.
+  // Phases on the skipped list: 1, 2, 9
   out << ""
-    "0 OfflineLB\n"
-    "1 OfflineLB\n"
-    "2 OfflineLB\n"
+    "0 NoLB\n"
+    "1 NoLB\n"
+    "2 NoLB\n"
     "3 OfflineLB\n"
     "4 OfflineLB\n"
     "5 OfflineLB\n"
     "6 OfflineLB\n"
     "7 OfflineLB\n"
-    "8 OfflineLB\n"
-    "9 OfflineLB\n";
+    "8 NoLB\n"
+    "9 NoLB\n";
   out.close();
 
   theConfig()->vt_lb = true;