#1933: LB: Add loadAndBroadcastSpec method to NodeLBData class

DARMA-tasking · Sep 22, 2022 · 3bc0d02 · 3bc0d02
1 parent d061713
commit 3bc0d02
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 1 deletion.
diff --git a/src/vt/runtime/runtime.cc b/src/vt/runtime/runtime.cc
@@ -643,6 +643,7 @@ void Runtime::setup() {
 # if vt_check_enabled(trace_enabled)
   theTrace->loadAndBroadcastSpec();
 # endif
+  theNodeLBData->loadAndBroadcastSpec();
 
   if (theConfig()->vt_pause) {
     pauseForDebugger();

diff --git a/src/vt/vrt/collection/balance/node_lb_data.cc b/src/vt/vrt/collection/balance/node_lb_data.cc
@@ -143,6 +143,26 @@ void NodeLBData::initialize() {
 #endif
 }
 
+void NodeLBData::loadAndBroadcastSpec() {
+  using namespace ::vt::utils::file_spec;
+
+  if (theConfig()->vt_lb_spec) {
+    auto spec_proxy = FileSpec::construct(FileSpecType::LB);
+
+    theTerm()->produce();
+    if (theContext()->getNode() == 0) {
+      auto spec_ptr = spec_proxy.get();
+      spec_ptr->parse();
+      spec_ptr->broadcastSpec();
+    }
+    theSched()->runSchedulerWhile(
+      [&spec_proxy] { return not spec_proxy.get()->specReceived(); });
+    theTerm()->consume();
+
+    spec_proxy_ = spec_proxy.getProxy();
+  }
+}
+
 void NodeLBData::createLBDataFile() {
   auto const file_name = theConfig()->getLBDataFileOut();
   auto const compress = theConfig()->vt_lb_data_compress;
@@ -228,7 +248,10 @@ getRecvSendDirection(elm::CommKeyType const& comm) {
 
 void NodeLBData::outputLBDataForPhase(PhaseType phase) {
   // LB data output when LB is enabled and appropriate flag is enabled
-  if (!theConfig()->vt_lb_data) {
+  vt::objgroup::proxy::Proxy<utils::file_spec::FileSpec> proxy(spec_proxy_);
+
+  if (!theConfig()->vt_lb_data or not proxy.get()->checkEnabled(phase) ) {
+    fmt::print("skipping lb data on {} phase\n", phase);
     return;
   }
 

diff --git a/src/vt/vrt/collection/balance/node_lb_data.h b/src/vt/vrt/collection/balance/node_lb_data.h
@@ -56,6 +56,7 @@
 #include "vt/utils/json/base_appender.h"
 #include "vt/vrt/collection/balance/lb_data_holder.h"
 #include "vt/vrt/collection/types/storage/storable.h"
+#include "vt/utils/file_spec/spec.h"
 
 #include <string>
 #include <unordered_map>
@@ -147,6 +148,11 @@ struct NodeLBData : runtime::component::Component<NodeLBData> {
    */
   void startIterCleanup(PhaseType phase, unsigned int look_back);
 
+  /**
+   * \internal \brief Load and broadcast the LB specification file
+   */
+  void loadAndBroadcastSpec();
+
   /**
    * \internal \brief Output LB data file for given phase based on instrumented
    * data
@@ -246,11 +252,13 @@ struct NodeLBData : runtime::component::Component<NodeLBData> {
   template <typename SerializerT>
   void serialize(SerializerT& s) {
     s | proxy_
+      | spec_proxy_
       | node_migrate_
       | node_collection_lookup_
       | node_objgroup_lookup_
       | next_elm_
       | created_dir_
+      | lb_data_writer_
       | lb_data_;
   }
 
@@ -268,6 +276,8 @@ struct NodeLBData : runtime::component::Component<NodeLBData> {
 private:
   /// Local proxy to objgroup
   objgroup::proxy::Proxy<NodeLBData> proxy_;
+  /// Proxy to objgroup for sending spec file
+  ObjGroupProxyType spec_proxy_ = vt::no_obj_group;
   /// Local migration type-free lambdas for each object
   std::unordered_map<ElementIDStruct,MigrateFnType> node_migrate_;
   /// Map from element ID to the collection's virtual proxy (untyped)