#1938: update runLB and startLB interface for base and derived classes

DARMA-tasking · Oct 24, 2022 · 6ef8a04 · 6ef8a04
1 parent 0a2036f
commit 6ef8a04
Show file tree

Hide file tree

Showing 20 changed files with 36 additions and 34 deletions.
diff --git a/src/vt/vrt/collection/balance/baselb/baselb.cc b/src/vt/vrt/collection/balance/baselb/baselb.cc
@@ -66,19 +66,20 @@ std::shared_ptr<const balance::Reassignment> BaseLB::startLB(
   balance::LoadModel* model,
   StatisticMapType const& in_stats,
   ElementCommType const& in_comm_lb_data,
-  TimeType total_load
+  TimeType total_load,
+  TimeType current_time
 ) {
-  start_time_ = timing::getCurrentTime();
+  start_time_ = current_time;
   phase_ = phase;
   proxy_ = proxy;
   load_model_ = model;
 
   importProcessorData(in_stats, in_comm_lb_data);
 
-  runInEpochCollective("BaseLB::startLB -> runLB", [this,total_load]{
+  runInEpochCollective("BaseLB::startLB -> runLB", [this,total_load, current_time]{
     getArgs(phase_);
     inputParams(config_entry_.get());
-    runLB(total_load);
+    runLB(total_load, current_time);
   });
 
   return normalizeReassignments();

diff --git a/src/vt/vrt/collection/balance/baselb/baselb.h b/src/vt/vrt/collection/balance/baselb/baselb.h
@@ -111,7 +111,8 @@ struct BaseLB {
     balance::LoadModel *model,
     StatisticMapType const& in_stats,
     ElementCommType const& in_comm_lb_data,
-    TimeType total_load
+    TimeType total_load,
+    TimeType current_time
   );
 
   void importProcessorData(
@@ -137,7 +138,7 @@ struct BaseLB {
   void finalize(CountMsg* msg);
 
   virtual void inputParams(balance::ConfigEntry* config) = 0;
-  virtual void runLB(TimeType total_load) = 0;
+  virtual void runLB(TimeType total_load, TimeType current_time) = 0;
 
   StatisticMapType const* getStats() const {
     return base_stats_;

diff --git a/src/vt/vrt/collection/balance/greedylb/greedylb.cc b/src/vt/vrt/collection/balance/greedylb/greedylb.cc
@@ -143,7 +143,7 @@ void GreedyLB::inputParams(balance::ConfigEntry* config) {
   strat_ = strategy_converter_.getFromConfig(config, strat_);
 }
 
-void GreedyLB::runLB(TimeType total_load) {
+void GreedyLB::runLB(TimeType total_load, TimeType current_time) {
   this_load = loadMilli(total_load);
   buildHistogram();
   loadStats();

diff --git a/src/vt/vrt/collection/balance/greedylb/greedylb.h b/src/vt/vrt/collection/balance/greedylb/greedylb.h
@@ -81,7 +81,7 @@ struct GreedyLB : LoadSamplerBaseLB {
   virtual ~GreedyLB() {}
 
   void init(objgroup::proxy::Proxy<GreedyLB> in_proxy);
-  void runLB(TimeType total_load) override;
+  void runLB(TimeType total_load, TimeType current_time) override;
   void inputParams(balance::ConfigEntry* config) override;
 
   static std::unordered_map<std::string, std::string> getInputKeysWithHelp();

diff --git a/src/vt/vrt/collection/balance/hierarchicallb/hierlb.cc b/src/vt/vrt/collection/balance/hierarchicallb/hierlb.cc
@@ -788,7 +788,7 @@ void HierarchicalLB::clearObj(ObjSampleType& objs) {
   }
 }
 
-void HierarchicalLB::runLB(TimeType total_load) {
+void HierarchicalLB::runLB(TimeType total_load, TimeType current_time) {
   this_load = loadMilli(total_load);
   buildHistogram();
   setupTree(min_threshold);

diff --git a/src/vt/vrt/collection/balance/hierarchicallb/hierlb.h b/src/vt/vrt/collection/balance/hierarchicallb/hierlb.h
@@ -73,7 +73,7 @@ struct HierarchicalLB : LoadSamplerBaseLB {
   virtual ~HierarchicalLB() {}
 
   void init(objgroup::proxy::Proxy<HierarchicalLB> in_proxy);
-  void runLB(TimeType total_load) override;
+  void runLB(TimeType total_load, TimeType current_time) override;
   void inputParams(balance::ConfigEntry* config) override;
 
   static std::unordered_map<std::string, std::string> getInputKeysWithHelp();

diff --git a/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc b/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc
@@ -216,7 +216,7 @@ void LBManager::defaultPostLBWork(ReassignmentMsg* msg) {
 
 void
 LBManager::runLB(
-  LBProxyType base_proxy, PhaseType phase, vt::Callback<ReassignmentMsg> cb
+  LBProxyType base_proxy, PhaseType phase, vt::Callback<ReassignmentMsg> cb, TimeType current_time
 ) {
   runInEpochCollective("LBManager::runLB -> updateLoads", [=] {
     model_->updateLoads(phase);
@@ -252,30 +252,30 @@ LBManager::runLB(
   vt_debug_print(terse, lb, "LBManager: running strategy\n");
 
   auto reassignment = strat->startLB(
-    phase, base_proxy, model_.get(), stats, *comm, total_load_from_model
+    phase, base_proxy, model_.get(), stats, *comm, total_load_from_model, current_time
   );
   cb.send(reassignment, phase);
 }
 
-void LBManager::selectStartLB(PhaseType phase) {
+void LBManager::selectStartLB(PhaseType phase, TimeType current_time) {
   namespace ph = std::placeholders;
   auto post_lb_ptr = std::mem_fn(&LBManager::defaultPostLBWork);
   auto post_lb_fn = std::bind(post_lb_ptr, this, ph::_1);
   auto cb = theCB()->makeFunc<ReassignmentMsg>(
     vt::pipe::LifetimeEnum::Once, post_lb_fn
   );
-  selectStartLB(phase, cb);
+  selectStartLB(phase, cb, current_time);
 }
 
 void LBManager::selectStartLB(
-  PhaseType phase, vt::Callback<ReassignmentMsg> cb
+  PhaseType phase, vt::Callback<ReassignmentMsg> cb, TimeType current_time
 ) {
   LBType lb = decideLBToRun(phase, true);
-  startLB(phase, lb, cb);
+  startLB(phase, lb, cb, current_time);
 }
 
 void LBManager::startLB(
-  PhaseType phase, LBType lb, vt::Callback<ReassignmentMsg> cb
+  PhaseType phase, LBType lb, vt::Callback<ReassignmentMsg> cb, TimeType current_time
 ) {
   vt_debug_print(
     normal, lb,
@@ -343,7 +343,7 @@ void LBManager::startLB(
   }
 
   LBProxyType base_proxy = lb_instances_["chosen"];
-  runLB(base_proxy, phase, cb);
+  runLB(base_proxy, phase, cb, current_time);
 }
 
 /*static*/

diff --git a/src/vt/vrt/collection/balance/lb_invoke/lb_manager.h b/src/vt/vrt/collection/balance/lb_invoke/lb_manager.h
@@ -134,15 +134,15 @@ struct LBManager : runtime::component::Component<LBManager> {
    * \param[in] phase the phase
    * \param[in] cb the callback for delivering the reassignment
    */
-  void selectStartLB(PhaseType phase, vt::Callback<ReassignmentMsg> cb);
+  void selectStartLB(PhaseType phase, vt::Callback<ReassignmentMsg> cb, TimeType current_time);
 
   /**
    * \internal
    * \brief Collectively start load balancing after deciding which to run
    *
    * \param[in] phase the phase
    */
-  void selectStartLB(PhaseType phase);
+  void selectStartLB(PhaseType phase, TimeType current_time);
 
   /**
    * \internal
@@ -152,7 +152,7 @@ struct LBManager : runtime::component::Component<LBManager> {
    * \param[in] lb the load balancer to run
    * \param[in] cb the callback for delivering the reassignment
    */
-  void startLB(PhaseType phase, LBType lb, vt::Callback<ReassignmentMsg> cb);
+  void startLB(PhaseType phase, LBType lb, vt::Callback<ReassignmentMsg> cb, TimeType current_time);
 
   /**
    * \internal
@@ -242,7 +242,7 @@ struct LBManager : runtime::component::Component<LBManager> {
    * \param[in] cb the callback for delivering the reassignment
    */
   void runLB(
-    LBProxyType base_proxy, PhaseType phase, vt::Callback<ReassignmentMsg> cb
+    LBProxyType base_proxy, PhaseType phase, vt::Callback<ReassignmentMsg> cb, TimeType current_time
   );
 
   void defaultPostLBWork(ReassignmentMsg* r);

diff --git a/src/vt/vrt/collection/balance/offlinelb/offlinelb.cc b/src/vt/vrt/collection/balance/offlinelb/offlinelb.cc
@@ -54,7 +54,7 @@ void OfflineLB::init(objgroup::proxy::Proxy<OfflineLB> in_proxy) {
   proxy_ = in_proxy;
 }
 
-void OfflineLB::runLB(TimeType) {
+void OfflineLB::runLB(TimeType, TimeType current_time) {
   auto const& myNewList = theLBDataReader()->getMoveList(phase_);
   for (size_t in = 0; in < myNewList.size(); in += 2) {
     auto this_node = theContext()->getNode();

diff --git a/src/vt/vrt/collection/balance/offlinelb/offlinelb.h b/src/vt/vrt/collection/balance/offlinelb/offlinelb.h
@@ -61,7 +61,7 @@ struct OfflineLB : BaseLB {
   virtual ~OfflineLB() = default;
 
   void init(objgroup::proxy::Proxy<OfflineLB> in_proxy);
-  void runLB(TimeType) override;
+  void runLB(TimeType, TimeType current_time) override;
   void inputParams(balance::ConfigEntry* config) override { }
 
   static std::unordered_map<std::string, std::string> getInputKeysWithHelp() {

diff --git a/src/vt/vrt/collection/balance/randomlb/randomlb.cc b/src/vt/vrt/collection/balance/randomlb/randomlb.cc
@@ -95,7 +95,7 @@ void RandomLB::inputParams(balance::ConfigEntry* config) {
   randomize_seed_ = config->getOrDefault<bool>("randomize_seed", randomize_seed_);
 }
 
-void RandomLB::runLB(TimeType) {
+void RandomLB::runLB(TimeType, TimeType current_time) {
   auto const this_node = theContext()->getNode();
   auto const num_nodes = static_cast<int32_t>(theContext()->getNumNodes());
 

diff --git a/src/vt/vrt/collection/balance/randomlb/randomlb.h b/src/vt/vrt/collection/balance/randomlb/randomlb.h
@@ -55,7 +55,7 @@ struct RandomLB : BaseLB {
   RandomLB() = default;
 
   void init(objgroup::proxy::Proxy<RandomLB> in_proxy);
-  void runLB(TimeType) override;
+  void runLB(TimeType, TimeType current_time) override;
   void inputParams(balance::ConfigEntry* config) override;
 
   static std::unordered_map<std::string, std::string> getInputKeysWithHelp();

diff --git a/src/vt/vrt/collection/balance/rotatelb/rotatelb.cc b/src/vt/vrt/collection/balance/rotatelb/rotatelb.cc
@@ -61,7 +61,7 @@ RotateLB::getInputKeysWithHelp() {
 
 void RotateLB::inputParams(balance::ConfigEntry* config) { }
 
-void RotateLB::runLB(TimeType) {
+void RotateLB::runLB(TimeType, TimeType current_time) {
   auto const& this_node = theContext()->getNode();
   auto const& num_nodes = theContext()->getNumNodes();
   auto const next_node = this_node + 1 > num_nodes-1 ? 0 : this_node + 1;

diff --git a/src/vt/vrt/collection/balance/rotatelb/rotatelb.h b/src/vt/vrt/collection/balance/rotatelb/rotatelb.h
@@ -66,7 +66,7 @@ struct RotateLB : BaseLB {
   virtual ~RotateLB() {}
 
   void init(objgroup::proxy::Proxy<RotateLB> in_proxy);
-  void runLB(TimeType) override;
+  void runLB(TimeType, TimeType current_time) override;
   void inputParams(balance::ConfigEntry* config) override;
 
   static std::unordered_map<std::string, std::string> getInputKeysWithHelp();

diff --git a/src/vt/vrt/collection/balance/temperedlb/temperedlb.cc b/src/vt/vrt/collection/balance/temperedlb/temperedlb.cc
@@ -420,7 +420,7 @@ void TemperedLB::inputParams(balance::ConfigEntry* config) {
   }
 }
 
-void TemperedLB::runLB(TimeType total_load) {
+void TemperedLB::runLB(TimeType total_load, TimeType current_time) {
   bool should_lb = false;
 
   this_load = total_load;

diff --git a/src/vt/vrt/collection/balance/temperedlb/temperedlb.h b/src/vt/vrt/collection/balance/temperedlb/temperedlb.h
@@ -77,7 +77,7 @@ struct TemperedLB : BaseLB {
 
 public:
   void init(objgroup::proxy::Proxy<TemperedLB> in_proxy);
-  void runLB(TimeType total_load) override;
+  void runLB(TimeType total_load, TimeType current_time) override;
   void inputParams(balance::ConfigEntry* config) override;
 
   static std::unordered_map<std::string, std::string> getInputKeysWithHelp();

diff --git a/src/vt/vrt/collection/balance/testserializationlb/testserializationlb.cc b/src/vt/vrt/collection/balance/testserializationlb/testserializationlb.cc
@@ -60,7 +60,7 @@ void TestSerializationLB::init(objgroup::proxy::Proxy<TestSerializationLB>) {
 
 void TestSerializationLB::inputParams(balance::ConfigEntry*) { }
 
-void TestSerializationLB::runLB(TimeType) {
+void TestSerializationLB::runLB(TimeType, TimeType current_time) {
   auto const this_node = theContext()->getNode();
   for (auto obj : *load_model_) {
     auto const load = load_model_->getModeledLoad(

diff --git a/src/vt/vrt/collection/balance/testserializationlb/testserializationlb.h b/src/vt/vrt/collection/balance/testserializationlb/testserializationlb.h
@@ -60,7 +60,7 @@ struct TestSerializationLB : BaseLB {
   virtual ~TestSerializationLB() {}
 
   void init(objgroup::proxy::Proxy<TestSerializationLB> in_proxy);
-  void runLB(TimeType) override;
+  void runLB(TimeType, TimeType current_time) override;
   void inputParams(balance::ConfigEntry* config) override;
 
   static std::unordered_map<std::string, std::string> getInputKeysWithHelp();

diff --git a/src/vt/vrt/collection/balance/zoltanlb/zoltanlb.cc b/src/vt/vrt/collection/balance/zoltanlb/zoltanlb.cc
@@ -129,7 +129,7 @@ void ZoltanLB::inputParams(balance::ConfigEntry* config) {
   }
 }
 
-void ZoltanLB::runLB(TimeType total_load) {
+void ZoltanLB::runLB(TimeType total_load, TimeType current_time) {
   auto const& this_node = theContext()->getNode();
   this_load = loadMilli(total_load);
 

diff --git a/src/vt/vrt/collection/balance/zoltanlb/zoltanlb.h b/src/vt/vrt/collection/balance/zoltanlb/zoltanlb.h
@@ -64,7 +64,7 @@ struct ZoltanLB : BaseLB {
   ZoltanLB();
 
   void init(objgroup::proxy::Proxy<ZoltanLB> in_proxy);
-  void runLB(TimeType total_load) override;
+  void runLB(TimeType total_load, TimeType current_time) override;
   void inputParams(balance::ConfigEntry* config) override;
 
   static std::unordered_map<std::string, std::string> getInputKeysWithHelp();