diff --git a/src/vt/vrt/collection/balance/lb_data_restart_reader.h b/src/vt/vrt/collection/balance/lb_data_restart_reader.h index 830ff4aa5e..6373a81153 100644 --- 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 */ std::shared_ptr> 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 std::vector changed_distro_; /// History of mapping that was read in from the data files - std::map>> history_; + std::unordered_map>> 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 index a7ddb3aad8..b9abd6d03f 100644 --- 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 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 index db83c49c68..161ee50a9c 100644 --- 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;