From 31325d3a526a6f6a6bc6a798f0ee961cb91d985d Mon Sep 17 00:00:00 2001
From: Phil Miller <pbmille@sandia.gov>
Date: Tue, 28 Jul 2020 17:29:21 -0400
Subject: [PATCH] #649: LBManager: use epoch/scheduler helper routines instead
 of explicit creation/management

---
 .../balance/lb_invoke/lb_manager.cc           | 25 +++++--------------
 1 file changed, 6 insertions(+), 19 deletions(-)

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 7fd46503d4..6e490c8b8f 100644
--- a/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc
+++ b/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc
@@ -151,38 +151,27 @@ LBManager::makeLB(MsgSharedPtr<StartLBMsg> msg) {
 
   destroy_lb_ = [proxy]{ proxy.destroyCollective(); };
 
-  EpochType model_epoch = theTerm()->makeEpochCollective("LBManager::model_epoch");
-  EpochType balance_epoch = theTerm()->makeEpochCollective("LBManager::balance_epoch");
-  EpochType migrate_epoch = theTerm()->makeEpochCollective("LBManager::migrate_epoch");
-
-  theMsg()->pushEpoch(model_epoch);
-  model_->updateLoads(phase);
-  theMsg()->popEpoch(model_epoch);
-  theTerm()->finishedEpoch(model_epoch);
+  runInEpochCollective([=] {
+    model_->updateLoads(phase);
+  });
 
-  theTerm()->addAction(model_epoch, [=] {
+  runInEpochCollective([=] {
     vt_debug_print(
       lb, node,
       "LBManager: running strategy\n"
     );
-    theMsg()->pushEpoch(balance_epoch);
     strat->startLB(phase, base_proxy, model_.get(), theProcStats()->getProcComm()->back());
-    theMsg()->popEpoch(balance_epoch);
-    theTerm()->finishedEpoch(balance_epoch);
   });
 
-  theTerm()->addAction(balance_epoch, [=] {
+  runInEpochCollective([=] {
     vt_debug_print(
       lb, node,
       "LBManager: starting migrations\n"
     );
-    theMsg()->pushEpoch(migrate_epoch);
     strat->applyMigrations(strat->getTransfers());
-    theMsg()->popEpoch(migrate_epoch);
-    theTerm()->finishedEpoch(migrate_epoch);
   });
 
-  theTerm()->addAction(migrate_epoch, [=] {
+  runInEpochCollective([=] {
     vt_debug_print(
       lb, node,
       "LBManager: finished migrations\n"
@@ -190,8 +179,6 @@ LBManager::makeLB(MsgSharedPtr<StartLBMsg> msg) {
     theProcStats()->startIterCleanup();
     this->finishedRunningLB(phase);
   });
-
-  runSchedulerThrough(migrate_epoch);
 }
 
 void LBManager::collectiveImpl(