#1014 dump node stats before trimming

src/vt/runtime/runtime.cc

@@ -1027,7 +1027,12 @@ bool Runtime::finalize(bool const force_now) {
 }
 
 void Runtime::sync() {
-  MPI_Barrier(theContext->getComm());
+  MPI_Comm comm = theContext->getComm();
+  if (comm == MPI_COMM_NULL) {
+    vtAbort("Trying to sync runtime while the communicator is not available");
+  } else {
+    MPI_Barrier(comm);
+  }
 }
 
 void Runtime::runScheduler() {
@@ -1216,6 +1221,7 @@ void Runtime::initializeComponents() {
     theTrace->initialize();
   #endif
   theEvent->initialize();
+  vrt::collection::balance::ProcStats::initialize();
 
   debug_print(runtime, node, "end: initializeComponents\n");
 }
@@ -1362,6 +1368,7 @@ void Runtime::finalizeComponents() {
 
   // Finalize memory usage component
   util::memory::MemoryUsage::finalize();
+  vrt::collection::balance::ProcStats::finalize();
 
   debug_print(runtime, node, "end: finalizeComponents\n");
 }

src/vt/vrt/collection/balance/lb_invoke/invoke.cc

@@ -240,6 +240,8 @@ void LBManager::releaseNow(PhaseType phase) {
   balance::ProcStats::startIterCleanup();
   balance::ProcStats::clearStats();
 
+  balance::ProcStats::outputStatsForPhase(phase);
+
   auto msg = makeMessage<CollectionPhaseMsg>();
 
   // Destruct the objgroup that was used for LB

src/vt/vrt/collection/balance/proc_stats.cc

@@ -61,6 +61,8 @@
 
 namespace vt { namespace vrt { namespace collection { namespace balance {
 
+using ArgType = vt::arguments::ArgConfig;
+
 /*static*/
 typename ProcStats::SparseMapType<std::unordered_map<ElementIDType,TimeType>>
   ProcStats::proc_data_ = {};
@@ -146,7 +148,6 @@ ProcStats::getProcSubphaseLoad(PhaseType phase) {
 }
 
 /*static*/ void ProcStats::createStatsFile() {
-  using ArgType = vt::arguments::ArgConfig;
   auto const node = theContext()->getNode();
   auto const base_file = std::string(ArgType::vt_lb_stats_file);
   auto const dir = std::string(ArgType::vt_lb_stats_dir);
@@ -155,7 +156,7 @@ ProcStats::getProcSubphaseLoad(PhaseType phase) {
 
   debug_print(
     lb, node,
-    "ProcStats: createStatsFile file={}\n", file_name
+    "ProcStats:: createStatsFile file={}\n", file_name
   );
 
   // Node 0 creates the directory
@@ -174,6 +175,7 @@ ProcStats::getProcSubphaseLoad(PhaseType phase) {
   }
 
   stats_file_ = fopen(file_name.c_str(), "w+");
+  vtAssertExpr(stats_file_ != nullptr);
 }
 
 /*static*/ void ProcStats::closeStatsFile() {
@@ -183,74 +185,88 @@ ProcStats::getProcSubphaseLoad(PhaseType phase) {
   }
 }
 
-/*static*/ void ProcStats::outputStatsFile() {
-  if (stats_file_ == nullptr) {
-    createStatsFile();
+std::pair<ElementIDType, ElementIDType>
+getRecvSendDirection(CommKeyType const& comm) {
+  switch (comm.cat_) {
+  case CommCategory::SendRecv:
+  case CommCategory::Broadcast:
+    return std::make_pair(comm.toObj(), comm.fromObj());
+
+  case CommCategory::NodeToCollection:
+  case CommCategory::NodeToCollectionBcast:
+    return std::make_pair(comm.toObj(), comm.fromNode());
+
+  case CommCategory::CollectionToNode:
+  case CommCategory::CollectionToNodeBcast:
+    return std::make_pair(comm.toNode(), comm.fromObj());
   }
 
-  vtAssertExpr(stats_file_ != nullptr);
+  vtAssert(false, "Invalid balance::CommCategory enum value");
+  return std::make_pair(ElementIDType{}, ElementIDType{});
+}
 
-  auto const num_iters = proc_data_.size();
+/*static*/ void ProcStats::initialize() {
+#if backend_check_enabled(lblite)
+  if (ArgType::vt_lb_stats and stats_file_ == nullptr) {
+    createStatsFile();
+  }
+#endif
+}
+
+/*static*/ void ProcStats::finalize() {
+  // If statistics are enabled, close output file and clear stats
+#if backend_check_enabled(lblite)
+  if (ArgType::vt_lb_stats) {
+    closeStatsFile();
+    clearStats();
+  }
+#endif
+}
 
-  vt_print(lb, "ProcStats::outputStatsFile: file={}, iter={}\n", print_ptr(stats_file_), num_iters);
+/*static*/ void ProcStats::outputStatsForPhase(PhaseType phase) {
+  // Statistics output when LB is enabled and appropriate flag is enabled
+  if (!ArgType::vt_lb_stats) {
+    return;
+  }
 
-  for (size_t i = 0; i < num_iters; i++) {
-    for (auto&& elm : proc_data_[i]) {
-      ElementIDType id = elm.first;
-      TimeType time = elm.second;
-      const auto& subphase_times = proc_subphase_data_[i][id];
-      size_t subphases = subphase_times.size();
+  vtAssertExpr(stats_file_ != nullptr);
+  debug_print(lb, node, "ProcStats::outputStatsForPhase: phase={}\n", phase);
 
-      auto obj_str = fmt::format("{},{},{},{},[", i, id, time, subphases);
-      for (size_t s = 0; s < subphases; s++) {
-        obj_str += std::to_string(subphase_times[s]);
-        if (s != subphases - 1)
-          obj_str += ",";
-      }
+  for (auto&& elm : proc_data_.at(phase)) {
+    ElementIDType id = elm.first;
+    TimeType time = elm.second;
+    const auto& subphase_times = proc_subphase_data_.at(phase)[id];
+    size_t subphases = subphase_times.size();
 
-      obj_str += "]\n";
+    auto obj_str = fmt::format("{},{},{},{},[", phase, id, time, subphases);
 
-      fprintf(stats_file_, "%s", obj_str.c_str());
-    }
-    for (auto&& elm : proc_comm_[i]) {
-      using E = typename std::underlying_type<CommCategory>::type;
-
-      auto const& key = elm.first;
-      auto const& val = elm.second;
-      auto const cat = static_cast<E>(key.cat_);
-
-      if (
-        key.cat_ == CommCategory::SendRecv or
-        key.cat_ == CommCategory::Broadcast
-      ) {
-        auto const to   = key.toObj();
-        auto const from = key.fromObj();
-        auto obj_str = fmt::format("{},{},{},{},{}\n", i, to, from, val, cat);
-        fprintf(stats_file_, "%s", obj_str.c_str());
-      } else if (
-        key.cat_ == CommCategory::NodeToCollection or
-        key.cat_ == CommCategory::NodeToCollectionBcast
-      ) {
-        auto const to   = key.toObj();
-        auto const from = key.fromNode();
-        auto obj_str = fmt::format("{},{},{},{},{}\n", i, to, from, val, cat);
-        fprintf(stats_file_, "%s", obj_str.c_str());
-      } else if (
-        key.cat_ == CommCategory::CollectionToNode or
-        key.cat_ == CommCategory::CollectionToNodeBcast
-      ) {
-        auto const to   = key.toNode();
-        auto const from = key.fromObj();
-        auto obj_str = fmt::format("{},{},{},{},{}\n", i, to, from, val, cat);
-        fprintf(stats_file_, "%s", obj_str.c_str());
-      } else {
-        vtAssert(false, "Invalid balance::CommCategory enum value");
+    for (size_t s = 0; s < subphases; s++) {
+      if (s > 0) {
+        obj_str += ",";
       }
+
+      obj_str += std::to_string(subphase_times[s]);
     }
+
+    obj_str += "]\n";
+
+    fprintf(stats_file_, "%s", obj_str.c_str());
+  }
+
+  for (auto&& elm : proc_comm_.at(phase)) {
+    using E = typename std::underlying_type<CommCategory>::type;
+
+    auto const& comm = elm.first;
+    auto const recvSend = getRecvSendDirection(comm);
+    auto const cat = static_cast<E>(comm.cat_);
+    auto obj_str = fmt::format(
+      "{},{},{},{},{}\n", phase, recvSend.first, recvSend.second,
+      elm.second, cat
+    );
+    fprintf(stats_file_, "%s", obj_str.c_str());
   }
-  fflush(stats_file_);
 
-  closeStatsFile();
+  fflush(stats_file_);
 }
 
 ElementIDType ProcStats::addProcStats(

src/vt/vrt/collection/balance/proc_stats.h

@@ -74,8 +74,10 @@ struct ProcStats {
   static void clearStats();
   static void startIterCleanup();
   static void releaseLB();
+  static void initialize();
+  static void finalize();
 
-  static void outputStatsFile();
+  static void outputStatsForPhase(PhaseType phase);
 
   static SubphaseLoadMapType const& getProcSubphaseLoad(PhaseType phase);
 

src/vt/vrt/collection/manager.cc

@@ -58,14 +58,6 @@ CollectionManager::CollectionManager() {
 /*virtual*/ CollectionManager::~CollectionManager() {
   cleanupAll<>();
 
-  // Statistics output when LB is enabled and appropriate flag is enabled
-#if backend_check_enabled(lblite)
-  if (ArgType::vt_lb_stats) {
-    balance::ProcStats::outputStatsFile();
-    balance::ProcStats::clearStats();
-  }
-#endif
-
   // Destroy the LBManager
   balance::LBManager::destroy();
 }