diff --git a/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc b/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc index 679ea0cc21..900149749f 100644 --- a/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc +++ b/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc @@ -643,13 +643,25 @@ void LBManager::computeStatistics( // FIXME: this is only TemperedWMin specific and requires composed // WeightedCommunicationVolume model (TemperedWMin::total_work_model_) - auto work = model->getModeledLoad( - elm, {balance::PhaseOffset::NEXT_PHASE, balance::PhaseOffset::WHOLE_PHASE} - ); - obj_work_model.emplace_back( - LoadData{lb::Statistic::Object_work_modeled, work} - ); - total_work_from_model += work; + auto proxy_iter = lb_instances_.find("chosen"); + if (proxy_iter == lb_instances_.end()) { + vt_debug_print(verbose, lb,"No load balancer instance chosen\n"); + continue; + } + + auto strat = proxy_iter->second.get(); + if (strat) { + auto tempered = reinterpret_cast(strat); + if (not tempered) { + continue; + } + + auto work = tempered->getModeledValue(elm); + obj_work_model.emplace_back( + LoadData{lb::Statistic::Object_work_modeled, work} + ); + total_work_from_model += work; + } } TimeType total_load_raw = 0.; diff --git a/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.cc b/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.cc index ce237742f1..4da4390daf 100644 --- a/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.cc +++ b/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.cc @@ -47,6 +47,7 @@ #include "vt/vrt/collection/balance/lb_invoke/lb_manager.h" #include "vt/vrt/collection/balance/model/load_model.h" #include "vt/vrt/collection/balance/model/weighted_communication_volume.h" +#include namespace vt { namespace vrt { namespace collection { namespace lb { diff --git a/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.h b/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.h index 2125e456e2..b3124738b3 100644 --- a/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.h +++ b/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.h @@ -63,7 +63,6 @@ struct TemperedWMin : TemperedLB { void inputParams(balance::SpecEntry* spec) override; -protected: TimeType getModeledValue(const elm::ElementIDStruct& obj) override; private: