diff --git a/src/vt/vrt/collection/balance/temperedlb/temperedlb.cc b/src/vt/vrt/collection/balance/temperedlb/temperedlb.cc index b575afbe1b..e93535bb8c 100644 --- a/src/vt/vrt/collection/balance/temperedlb/temperedlb.cc +++ b/src/vt/vrt/collection/balance/temperedlb/temperedlb.cc @@ -42,6 +42,8 @@ */ #include "vt/config.h" +#include "vt/configs/types/types_sentinels.h" +#include "vt/configs/types/types_type.h" #include "vt/timing/timing.h" #include "vt/vrt/collection/balance/baselb/baselb.h" #include "vt/vrt/collection/balance/model/load_model.h" @@ -474,6 +476,12 @@ void TemperedLB::runLB(TimeType total_load) { } } +void TemperedLB::clearDataStructures() { + underloaded_.clear(); + load_info_.clear(); + is_overloaded_ = is_underloaded_ = false; +} + void TemperedLB::doLBStages(TimeType start_imb) { decltype(this->cur_objs_) best_objs; LoadType best_load = 0; @@ -483,11 +491,7 @@ void TemperedLB::doLBStages(TimeType start_imb) { auto this_node = theContext()->getNode(); for (trial_ = 0; trial_ < num_trials_; ++trial_) { - // Clear out data structures - selected_.clear(); - underloaded_.clear(); - load_info_.clear(); - is_overloaded_ = is_underloaded_ = false; + clearDataStructures(); TimeType best_imb_this_trial = start_imb + 10; @@ -504,11 +508,7 @@ void TemperedLB::doLBStages(TimeType start_imb) { } this_new_load_ = this_load; } else { - // Clear out data structures from previous iteration - selected_.clear(); - underloaded_.clear(); - load_info_.clear(); - is_overloaded_ = is_underloaded_ = false; + clearDataStructures(); } vt_debug_print( @@ -667,7 +667,7 @@ void TemperedLB::informAsync() { vtAssert(k_max_ > 0, "Number of rounds (k) must be greater than zero"); auto const this_node = theContext()->getNode(); - if (is_underloaded_) { + if (canPropagate()) { underloaded_.insert(this_node); } @@ -682,7 +682,7 @@ void TemperedLB::informAsync() { auto propagate_epoch = theTerm()->makeEpochCollective("TemperedLB: informAsync"); // Underloaded start the round - if (is_underloaded_) { + if (canPropagate()) { uint8_t k_cur_async = 0; propagateRound(k_cur_async, false, propagate_epoch); } @@ -718,11 +718,11 @@ void TemperedLB::informSync() { vtAssert(k_max_ > 0, "Number of rounds (k) must be greater than zero"); auto const this_node = theContext()->getNode(); - if (is_underloaded_) { + if (canPropagate()) { underloaded_.insert(this_node); } - auto propagate_this_round = is_underloaded_; + auto propagate_this_round = canPropagate(); propagate_next_round_ = false; new_underloaded_ = underloaded_; new_load_info_ = load_info_; @@ -793,8 +793,7 @@ void TemperedLB::propagateRound(uint8_t k_cur, bool sync, EpochType epoch) { gen_propagate_.seed(seed_()); } - auto& selected = selected_; - selected = underloaded_; + auto& selected = underloaded_; if (selected.find(this_node) == selected.end()) { selected.insert(this_node); } @@ -1203,7 +1202,7 @@ void TemperedLB::decide() { int n_transfers = 0, n_rejected = 0; - if (is_overloaded_) { + if (canMigrate()) { std::vector under = makeUnderloaded(); std::unordered_map migrate_objs; diff --git a/src/vt/vrt/collection/balance/temperedlb/temperedlb.h b/src/vt/vrt/collection/balance/temperedlb/temperedlb.h index 8e6f68b5b6..f8f9046b33 100644 --- a/src/vt/vrt/collection/balance/temperedlb/temperedlb.h +++ b/src/vt/vrt/collection/balance/temperedlb/temperedlb.h @@ -94,6 +94,15 @@ struct TemperedLB : BaseLB { void informSync(); void decide(); void migrate(); + void clearDataStructures(); + + virtual bool canMigrate() const { return is_overloaded_; } + /** + * \brief Decides whether the rank can initiate information propagation stage + * + * TemperedLB restricts this to underloaded ranks + */ + virtual bool canPropagate() const { return is_underloaded_; } void propagateRound(uint8_t k_cur_async, bool sync, EpochType epoch = no_epoch); void propagateIncomingAsync(LoadMsgAsync* msg); @@ -164,7 +173,6 @@ struct TemperedLB : BaseLB { objgroup::proxy::Proxy proxy_ = {}; bool is_overloaded_ = false; bool is_underloaded_ = false; - std::unordered_set selected_ = {}; std::unordered_set underloaded_ = {}; std::unordered_set new_underloaded_ = {}; std::unordered_map cur_objs_ = {}; diff --git a/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.h b/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.h index c008741b4a..386e0c03d9 100644 --- a/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.h +++ b/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.h @@ -61,6 +61,11 @@ struct TemperedWMin : TemperedLB { void inputParams(balance::SpecEntry* spec) override; protected: + /** + * All ranks are allowed to initiate the information propagation stage + */ + bool canPropagate() const override { return true; } + TimeType getModeledWork(const elm::ElementIDStruct& obj) override; private: