Merge pull request #1728 from DARMA-tasking/1723-include-user-defined…

…-in-json

1723 Include user-defined fields in json LB data files

src/vt/messaging/active.cc

@@ -161,7 +161,7 @@ void ActiveMessenger::startup() {
   // Hook to collect LB data about objgroups
   thePhase()->registerHookCollective(phase::PhaseHook::End, [this]{
     theNodeLBData()->addNodeLBData(
-      bare_handler_dummy_elm_id_for_lb_data_, &bare_handler_lb_data_
+      bare_handler_dummy_elm_id_for_lb_data_, &bare_handler_lb_data_, nullptr
     );
   });
 #endif

src/vt/objgroup/manager.cc

@@ -68,7 +68,7 @@ void ObjGroupManager::startup() {
         auto proxy = elm::ElmIDBits::getObjGroupProxy(elm_id.id, false);
         vtAssertExpr(proxy == obj.first);
         theNodeLBData()->registerObjGroupInfo(elm_id, obj.first);
-        theNodeLBData()->addNodeLBData(elm_id, &holder->getLBData());
+        theNodeLBData()->addNodeLBData(elm_id, &holder->getLBData(), nullptr);
       }
     }
   });

src/vt/vrt/collection/balance/col_lb_data.impl.h

@@ -74,7 +74,7 @@ void CollectionLBData::syncNextPhase(CollectStatsMsg<ColT>* msg, ColT* col) {
 
   auto const proxy = col->getProxy();
   auto const subphase = getFocusedSubPhase(proxy);
-  theNodeLBData()->addNodeLBData(col->elm_id_, &col->lb_data_, subphase);
+  theNodeLBData()->addNodeLBData(col->elm_id_, &col->lb_data_, col, subphase);
 
   std::vector<uint64_t> idx;
   for (index::NumDimensionsType i = 0; i < col->getIndex().ndims(); i++) {

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

@@ -82,6 +82,15 @@ std::unique_ptr<nlohmann::json> LBDataHolder::toJson(PhaseType phase) const {
       j["tasks"][i]["resource"] = "cpu";
       j["tasks"][i]["node"] = theContext()->getNode();
       j["tasks"][i]["time"] = time;
+      if (user_defined_json_.find(phase) != user_defined_json_.end()) {
+        auto &user_def_this_phase = user_defined_json_.at(phase);
+        if (user_def_this_phase.find(id) != user_def_this_phase.end()) {
+          auto &user_def = user_def_this_phase.at(id);
+          if (!user_def->empty()) {
+            j["tasks"][i]["user_defined"] = *user_def;
+          }
+        }
+      }
       outputEntity(j["tasks"][i]["entity"], id);
 
       auto const& subphase_times = elm.second.subphase_loads;

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

@@ -76,6 +76,7 @@ struct LBDataHolder {
     s | node_data_;
     s | node_comm_;
     s | node_subphase_comm_;
+    s | user_defined_json_;
     s | node_idx_;
   }
 
@@ -109,6 +110,10 @@ struct LBDataHolder {
   std::unordered_map<PhaseType, CommMapType> node_comm_;
   /// Node communication graph for each subphase
   std::unordered_map<PhaseType, std::unordered_map<SubphaseType, CommMapType>> node_subphase_comm_;
+  /// User-defined data from each phase
+  std::unordered_map<PhaseType, std::unordered_map<
+    ElementIDStruct, std::shared_ptr<nlohmann::json>
+  >> user_defined_json_;
   /// Node indices for each ID along with the proxy ID
   std::unordered_map<ElementIDStruct, std::tuple<VirtualProxyType, std::vector<uint64_t>>> node_idx_;
   /// Map from id to objgroup proxy

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

@@ -56,6 +56,7 @@
 #include <unordered_map>
 #include <cstdio>
 #include <sys/stat.h>
+#include <memory>
 
 #include <fmt/core.h>
 
@@ -245,7 +246,8 @@ void NodeLBData::registerObjGroupInfo(
 }
 
 void NodeLBData::addNodeLBData(
-  ElementIDStruct id, elm::ElementLBData* in, SubphaseType focused_subphase
+  ElementIDStruct id, elm::ElementLBData* in, StorableType *storable,
+  SubphaseType focused_subphase
 ) {
   vt_debug_print(
     normal, lb,
@@ -281,6 +283,12 @@ void NodeLBData::addNodeLBData(
     }
   }
 
+  if (storable) {
+    lb_data_->user_defined_json_[phase][id] = std::make_shared<nlohmann::json>(
+      storable->toJson()
+    );
+  }
+
   in->updatePhase(1);
 
   auto model = theLBManager()->getLoadModel();

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

@@ -55,6 +55,7 @@
 #include "vt/objgroup/proxy/proxy_objgroup.h"
 #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 <string>
 #include <unordered_map>
@@ -78,6 +79,7 @@ namespace vt { namespace vrt { namespace collection { namespace balance {
  */
 struct NodeLBData : runtime::component::Component<NodeLBData> {
   using MigrateFnType       = std::function<void(NodeType)>;
+  using StorableType = vt::vrt::collection::storage::Storable;
 
   /**
    * \internal \brief System call to construct \c NodeLBData
@@ -131,7 +133,7 @@ struct NodeLBData : runtime::component::Component<NodeLBData> {
    * \param[in] focused_subphase the focused subphase (optional)
    */
   void addNodeLBData(
-    ElementIDStruct id, elm::ElementLBData* in,
+    ElementIDStruct id, elm::ElementLBData* in, StorableType *storable,
     SubphaseType focused_subphase = elm::ElementLBData::no_subphase
   );
 

src/vt/vrt/collection/types/storage/storable.cc

@@ -56,4 +56,14 @@ void Storable::valRemove(std::string const& str) {
   }
 }
 
+nlohmann::json Storable::toJson() {
+  nlohmann::json j;
+  for (auto iter = map_.begin(); iter != map_.end(); ++iter) {
+    if (iter->second->shouldJson()) {
+      j[iter->first] = iter->second->toJson();
+    }
+  }
+  return j;
+}
+
 }}}} /* end namespace vt::vrt::collection::storage */

src/vt/vrt/collection/types/storage/storable.h

@@ -78,6 +78,16 @@ struct Storable {
   template <typename U>
   void valInsert(std::string const& str, U&& u);
 
+  /**
+   * \brief Insert a new key/value pair
+   *
+   * \param[in] str the key
+   * \param[in] u the value
+   * \param[in] dump_to_json whether to include in json file
+   */
+  template <typename U>
+  void valInsert(std::string const& str, U&& u, bool dump_to_json);
+
   /**
    * \brief Get the value from a key
    *
@@ -114,6 +124,13 @@ struct Storable {
    */
   void valRemove(std::string const& str);
 
+  /**
+   * \brief Generate the json if applicable
+   *
+   * \return the json
+   */
+  nlohmann::json toJson();
+
 private:
   /// Map of type-erased, stored values
   std::unordered_map<std::string, std::unique_ptr<StoreElmBase>> map_;

src/vt/vrt/collection/types/storage/storable.impl.h

@@ -66,6 +66,19 @@ void Storable::valInsert(std::string const& str, U&& u) {
   );
 }
 
+template <typename U>
+void Storable::valInsert(std::string const& str, U&& u, bool dump_to_json) {
+  map_.emplace(
+    std::piecewise_construct,
+    std::forward_as_tuple(str),
+    std::forward_as_tuple(
+      std::make_unique<StoreElm<typename std::decay<U>::type>>(
+        std::forward<U>(u), dump_to_json
+      )
+    )
+  );
+}
+
 template <typename U>
 U& Storable::valGet(std::string const& str) {
   vtAssert(valExists(str), "Key must exist in map");

src/vt/vrt/collection/types/storage/store_elm.h

@@ -46,6 +46,8 @@
 
 #include "vt/config.h"
 
+#include <nlohmann/json.hpp>
+
 #include <checkpoint/checkpoint.h>
 
 #include <type_traits>
@@ -62,10 +64,23 @@ struct StoreElmBase {
   /// uses polymorphic serialization
   checkpoint_virtual_serialize_root()
 
+  using json = nlohmann::json;
+
   StoreElmBase() = default;
 
+  StoreElmBase(bool dump_to_json)
+    : dump_to_json_(dump_to_json)
+  {}
+
   virtual ~StoreElmBase() {}
 
+  /**
+   * \brief Generate the json if applicable
+   *
+   * \return the json
+   */
+  virtual nlohmann::json toJson() = 0;
+
   /**
    * \brief Get the value as \c T
    *
@@ -88,7 +103,49 @@ struct StoreElmBase {
    * \param[in] s the serializer
    */
   template <typename SerializerT>
-  void serialize(SerializerT& s) { }
+  void serialize(SerializerT& s) {
+    s | dump_to_json_;
+  }
+
+  /**
+   * \brief Whether the value should be dumped to the json LB data file
+   *
+   * \return whether to dump
+   */
+  bool shouldJson() const { return dump_to_json_; }
+
+  /**
+   * \brief Generate the json because it is jsonable
+   *
+   * \param[in] u the data to convert to json
+   */
+  template <typename U>
+  static json maybeGenerateJson(
+    const U& u, typename std::enable_if<
+      nlohmann::detail::has_to_json<json,U>::value
+    >::type* = nullptr
+  ) {
+    json j(u);
+    return j;
+  }
+
+  /**
+   * \brief Abort because it is not jsonable
+   *
+   * \param[in] u the data that cannot be converted to json
+   */
+  template <typename U>
+  static json maybeGenerateJson(
+    const U& u, typename std::enable_if<
+      not nlohmann::detail::has_to_json<json,U>::value
+    >::type* = nullptr
+  ) {
+    vtAbort("Instantiated maybeGenerateJson on non-jsonable type");
+    return json{};
+  }
+
+protected:
+  bool dump_to_json_ = false;
 };
 
 /**
@@ -125,6 +182,22 @@ struct StoreElm<
     : elm_(std::forward<U>(u))
   { }
 
+  /**
+   * \brief Construct with value
+   *
+   * \param[in] u the value
+   * \param[in] dump_to_json whether to dump this to json LB data file
+   */
+  template <typename U>
+  StoreElm(
+    U&& u, bool dump_to_json, typename std::enable_if<
+      nlohmann::detail::has_to_json<nlohmann::json,U>::value
+    >::type* = nullptr
+  )
+    : StoreElmBase(dump_to_json),
+      elm_(std::forward<U>(u))
+  { }
+
   /**
    * \brief Serialization re-constructor
    *
@@ -156,6 +229,16 @@ struct StoreElm<
    */
   T const& get() const { return elm_; }
 
+  /**
+   * \brief Generate the json if jsonable
+   *
+   * \return the json
+   */
+  nlohmann::json toJson() override
+  {
+    return StoreElm::maybeGenerateJson<T>(elm_);
+  }
+
 private:
   T elm_ = {};                  /**< The stored value */
 };
@@ -222,6 +305,22 @@ struct StoreElm<
     : wrapper_(detail::ByteWrapper<T>{std::forward<U>(u)})
   { }
 
+  /**
+   * \brief Construct with value
+   *
+   * \param[in] u the value
+   * \param[in] dump_to_json whether to dump this to json LB data file
+   */
+  template <typename U>
+  StoreElm(
+    U&& u, bool dump_to_json, typename std::enable_if<
+      nlohmann::detail::has_to_json<nlohmann::json,U>::value
+    >::type* = nullptr
+  )
+    : StoreElmBase(dump_to_json),
+      wrapper_(detail::ByteWrapper<T>{std::forward<U>(u)})
+  { }
+
   /**
    * \brief Serialization re-constructor
    *
@@ -253,6 +352,16 @@ struct StoreElm<
    */
   T const& get() const { return wrapper_.elm_; }
 
+  /**
+   * \brief Generate the json if jsonable
+   *
+   * \return the json
+   */
+  nlohmann::json toJson() override
+  {
+    return StoreElm::maybeGenerateJson<T>(wrapper_.elm_);
+  }
+
 private:
   detail::ByteWrapper<T> wrapper_ = {}; /**< The byte-copyable value wrapper */
 };