diff --git a/scripts/JSON_data_files_validator.py b/scripts/JSON_data_files_validator.py index bacf268dde..d3e3a30ef1 100644 --- a/scripts/JSON_data_files_validator.py +++ b/scripts/JSON_data_files_validator.py @@ -64,6 +64,7 @@ def _get_valid_schema(self) -> Schema: 'range': [[int]], }, }, + Optional('attributes'): dict }, 'phases': [ { @@ -88,7 +89,8 @@ def _get_valid_schema(self) -> Schema: } ], 'time': float, - Optional('user_defined'): dict + Optional('user_defined'): dict, + Optional('attributes'): dict }, ], Optional('communications'): [ diff --git a/src/vt/vrt/collection/balance/lb_data_holder.cc b/src/vt/vrt/collection/balance/lb_data_holder.cc index 9ab01a380c..43b1e61802 100644 --- a/src/vt/vrt/collection/balance/lb_data_holder.cc +++ b/src/vt/vrt/collection/balance/lb_data_holder.cc @@ -125,6 +125,26 @@ std::unique_ptr LBDataHolder::metadataToJson() const { return std::make_unique(std::move(j)); } +std::unique_ptr LBDataHolder::rankAttributesToJson() const { + if (rank_attributes_.empty()) { + return nullptr; + } + + nlohmann::json j; + + for (auto const& [key, value] : rank_attributes_) { + if (std::holds_alternative(value)) { + j["attributes"][key] = std::get(value); + } else if (std::holds_alternative(value)) { + j["attributes"][key] = std::get(value); + } else if (std::holds_alternative(value)) { + j["attributes"][key] = std::get(value); + } + } + + return std::make_unique(std::move(j)); +} + std::unique_ptr LBDataHolder::toJson(PhaseType phase) const { using json = nlohmann::json; @@ -150,6 +170,20 @@ std::unique_ptr LBDataHolder::toJson(PhaseType phase) const { } outputEntity(j["tasks"][i]["entity"], id); + if (node_user_attributes_.find(phase) != node_user_attributes_.end()) { + if (node_user_attributes_.at(phase).find(id) != node_user_attributes_.at(phase).end()) { + for (auto const& [key, value] : node_user_attributes_.at(phase).at(id)) { + if (std::holds_alternative(value)) { + j["tasks"][i]["attributes"][key] = std::get(value); + } else if (std::holds_alternative(value)) { + j["tasks"][i]["attributes"][key] = std::get(value); + } else if (std::holds_alternative(value)) { + j["tasks"][i]["attributes"][key] = std::get(value); + } + } + } + } + auto const& subphase_times = elm.second.subphase_loads; std::size_t const subphases = subphase_times.size(); if (subphases != 0) { @@ -303,6 +337,18 @@ LBDataHolder::LBDataHolder(nlohmann::json const& j) } } } + + if (task.find("attributes") != task.end()) { + for (auto const& [key, value] : task["attributes"].items()) { + if (value.is_number_integer()) { + node_user_attributes_[id][elm][key] = value.get(); + } else if (value.is_number_float()) { + node_user_attributes_[id][elm][key] = value.get(); + } else if (value.is_string()) { + node_user_attributes_[id][elm][key] = value.get(); + } + } + } } } } @@ -441,6 +487,18 @@ void LBDataHolder::readMetadata(nlohmann::json const& j) { } } } + // load rank user atrributes + if (metadata.find("attributes") != metadata.end()) { + for (auto const& [key, value] : metadata["attributes"].items()) { + if (value.is_number_integer()) { + rank_attributes_[key] = value.get(); + } else if (value.is_number_float()) { + rank_attributes_[key] = value.get(); + } else if (value.is_string()) { + rank_attributes_[key] = value.get(); + } + } + } } } diff --git a/src/vt/vrt/collection/balance/lb_data_holder.h b/src/vt/vrt/collection/balance/lb_data_holder.h index 3a69999d98..caeb4e4cea 100644 --- a/src/vt/vrt/collection/balance/lb_data_holder.h +++ b/src/vt/vrt/collection/balance/lb_data_holder.h @@ -85,6 +85,8 @@ struct LBDataHolder { s | skipped_phases_; s | identical_phases_; s | user_defined_lb_info_; + s | node_user_attributes_; + s | rank_attributes_; } /** @@ -97,12 +99,19 @@ struct LBDataHolder { std::unique_ptr toJson(PhaseType phase) const; /** - * \brief Output a LB phase's metdadata to JSON + * \brief Output a LB phase's metadata to JSON * * \return the json data structure */ std::unique_ptr metadataToJson() const; + /** + * \brief Output a LB rank attributes metadata to JSON + * + * \return the json data structure + */ + std::unique_ptr rankAttributesToJson() const; + /** * \brief Clear all LB data */ @@ -125,6 +134,8 @@ struct LBDataHolder { void readMetadata(nlohmann::json const& j); public: + /// Node attributes for the current rank + ElmUserDataType rank_attributes_; /// Node timings for each local object std::unordered_map node_data_; /// Node communication graph for each local object @@ -139,6 +150,8 @@ struct LBDataHolder { std::unordered_map> user_per_phase_json_; /// User-defined data from each phase for LB std::unordered_map user_defined_lb_info_; + /// User-defined attributes from each phase + std::unordered_map node_user_attributes_; /// Node indices for each ID along with the proxy ID std::unordered_map>> node_idx_; /// Map from id to objgroup proxy diff --git a/src/vt/vrt/collection/balance/node_lb_data.cc b/src/vt/vrt/collection/balance/node_lb_data.cc index 18828426cb..6d65cad1b2 100644 --- a/src/vt/vrt/collection/balance/node_lb_data.cc +++ b/src/vt/vrt/collection/balance/node_lb_data.cc @@ -102,6 +102,11 @@ NodeLBData::getUserData() const { return &lb_data_->user_defined_lb_info_; } +std::unordered_map const* +NodeLBData::getPhaseAttributes() const { + return &lb_data_->node_user_attributes_; +} + std::unordered_map const* NodeLBData::getNodeComm() const { return &lb_data_->node_comm_; } @@ -110,6 +115,10 @@ std::unordered_map> con return &lb_data_->node_subphase_comm_; } +ElmUserDataType const* NodeLBData::getNodeAttributes() const { + return &lb_data_->rank_attributes_; +} + CommMapType* NodeLBData::getNodeComm(PhaseType phase) { auto iter = lb_data_->node_comm_.find(phase); return (iter != lb_data_->node_comm_.end()) ? &iter->second : nullptr; @@ -127,6 +136,7 @@ void NodeLBData::startIterCleanup(PhaseType phase, unsigned int look_back) { lb_data_->node_comm_.erase(phase - look_back); lb_data_->node_subphase_comm_.erase(phase - look_back); lb_data_->user_defined_lb_info_.erase(phase - look_back); + lb_data_->node_user_attributes_.erase(phase - look_back); } // Clear migrate lambdas and proxy lookup since LB is complete @@ -215,6 +225,10 @@ void NodeLBData::createLBDataFile() { if(phasesMetadata) { metadata["phases"] = *phasesMetadata; } + auto attributesMetadata = lb_data_->rankAttributesToJson(); + if(attributesMetadata) { + metadata["attributes"] = *attributesMetadata; + } lb_data_writer_ = std::make_unique( "phases", metadata, file_name, compress ); @@ -352,6 +366,11 @@ void NodeLBData::addNodeLBData( lb_data_->user_defined_lb_info_[phase][id][key] = val->toVariant(); } ); + storable->collectAttributes( + [&](std::string const& key, auto val) { + lb_data_->node_user_attributes_[phase][id][key] = val->toVariant(); + } + ); } in->updatePhase(1); diff --git a/src/vt/vrt/collection/balance/node_lb_data.h b/src/vt/vrt/collection/balance/node_lb_data.h index 3060f944c8..686ac258b3 100644 --- a/src/vt/vrt/collection/balance/node_lb_data.h +++ b/src/vt/vrt/collection/balance/node_lb_data.h @@ -188,6 +188,13 @@ struct NodeLBData : runtime::component::Component { */ std::unordered_map const* getUserData() const; + /** + * \internal \brief Get the user-defined attributes + * + * \return an observer pointer to the user-defined attributes + */ + std::unordered_map const* getPhaseAttributes() const; + /** * \internal \brief Get stored object comm data for a specific phase * @@ -204,6 +211,13 @@ struct NodeLBData : runtime::component::Component { */ std::unordered_map> const* getNodeSubphaseComm() const; + /** + * \internal \brief Get stored node attributes + * + * \return an observer pointer to the node attributes + */ + ElmUserDataType const* getNodeAttributes() const; + /** * \internal \brief Test if this node has an object to migrate * diff --git a/src/vt/vrt/collection/types/storage/storable.h b/src/vt/vrt/collection/types/storage/storable.h index 444828f3b8..c23fcc03a3 100644 --- a/src/vt/vrt/collection/types/storage/storable.h +++ b/src/vt/vrt/collection/types/storage/storable.h @@ -94,13 +94,14 @@ struct Storable { * \param[in] u the value * \param[in] dump_to_json whether to include in json file * \param[in] provide_to_lb whether to provide the data to the load balancers + * \param[in] dump_to_attributes whether to dump to attributes in JSON output * * \note If \c provide_to_lb is enabled, the data must be convertible to an * integer, double, or string */ template void valInsert( - std::string const& str, U&& u, bool dump_to_json, bool provide_to_lb + std::string const& str, U&& u, bool dump_to_json, bool provide_to_lb, bool dump_to_attributes ); /** @@ -160,6 +161,20 @@ struct Storable { } } + /** + * \brief Traverse the map and calls Callable on each attribute + * + * \param[in] c the callable + */ + template + void collectAttributes(Callable&& c) { + for (auto const& [key, value] : map_) { + if (value->isAttribute()) { + c(key, value.get()); + } + } + } + private: /// Map of type-erased, stored values std::unordered_map> map_; diff --git a/src/vt/vrt/collection/types/storage/storable.impl.h b/src/vt/vrt/collection/types/storage/storable.impl.h index 1ecf9b4734..b03b04378c 100644 --- a/src/vt/vrt/collection/types/storage/storable.impl.h +++ b/src/vt/vrt/collection/types/storage/storable.impl.h @@ -60,7 +60,7 @@ void Storable::valInsert(std::string const& str, U&& u) { std::forward_as_tuple(str), std::forward_as_tuple( std::make_unique::type>>( - std::forward(u), false, false + std::forward(u), false, false, false ) ) ); @@ -68,14 +68,14 @@ void Storable::valInsert(std::string const& str, U&& u) { template void Storable::valInsert( - std::string const& str, U&& u, bool dump_to_json, bool provide_to_lb + std::string const& str, U&& u, bool dump_to_json, bool provide_to_lb, bool dump_to_attributes ) { map_.emplace( std::piecewise_construct, std::forward_as_tuple(str), std::forward_as_tuple( std::make_unique::type>>( - std::forward(u), dump_to_json, provide_to_lb + std::forward(u), dump_to_json, provide_to_lb, dump_to_attributes ) ) ); diff --git a/src/vt/vrt/collection/types/storage/store_elm.h b/src/vt/vrt/collection/types/storage/store_elm.h index 0f9524bdb4..1937179243 100644 --- a/src/vt/vrt/collection/types/storage/store_elm.h +++ b/src/vt/vrt/collection/types/storage/store_elm.h @@ -74,10 +74,12 @@ struct StoreElmBase { * * \param[in] dump_to_json whether to dump to JSON output * \param[in] provide_to_lb whether to provide to the LB + * \param[in] dump_to_attributes whether to dump to attributes in JSON output */ - StoreElmBase(bool dump_to_json, bool provide_to_lb) + StoreElmBase(bool dump_to_json, bool provide_to_lb, bool dump_to_attributes) : dump_to_json_(dump_to_json), - provide_to_lb_(provide_to_lb) + provide_to_lb_(provide_to_lb), + dump_to_attributes_(dump_to_attributes) {} virtual ~StoreElmBase() {} @@ -121,6 +123,7 @@ struct StoreElmBase { void serialize(SerializerT& s) { s | dump_to_json_; s | provide_to_lb_; + s | dump_to_attributes_; } /** @@ -139,6 +142,15 @@ struct StoreElmBase { return provide_to_lb_; } + /** + * \brief Whether the value should be dumped to the json attributes field + * + * \return whether it is an attribute + */ + bool isAttribute() const { + return dump_to_attributes_; + } + /** * \brief Generate the json because it is jsonable * @@ -177,6 +189,7 @@ struct StoreElmBase { protected: bool dump_to_json_ = false; bool provide_to_lb_ = false; + bool dump_to_attributes_ = false; }; /** @@ -209,10 +222,11 @@ struct StoreElm< * \param[in] u the value * \param[in] dump_to_json whether to dump to json * \param[in] provide_to_lb whether to provide to LB + * \param[in] dump_to_attributes whether to dump to attributes in JSON output */ template - explicit StoreElm(U&& u, bool dump_to_json, bool provide_to_lb) - : StoreElmBase(dump_to_json, provide_to_lb), + explicit StoreElm(U&& u, bool dump_to_json, bool provide_to_lb, bool dump_to_attributes) + : StoreElmBase(dump_to_json, provide_to_lb, dump_to_attributes), elm_(std::forward(u)) { } @@ -328,10 +342,11 @@ struct StoreElm< * \param[in] u the value * \param[in] dump_to_json whether to dump to json * \param[in] provide_to_lb whether to provide to LB + * \param[in] dump_to_attributes whether to dump to attributes in JSON output */ template - explicit StoreElm(U&& u, bool dump_to_json, bool provide_to_lb) - : StoreElmBase(dump_to_json, provide_to_lb), + explicit StoreElm(U&& u, bool dump_to_json, bool provide_to_lb, bool dump_to_attributes) + : StoreElmBase(dump_to_json, provide_to_lb, dump_to_attributes), wrapper_(detail::ByteWrapper{std::forward(u)}) { } diff --git a/tests/unit/collection/test_lb.extended.cc b/tests/unit/collection/test_lb.extended.cc index d9d0dd199c..a865ef5113 100644 --- a/tests/unit/collection/test_lb.extended.cc +++ b/tests/unit/collection/test_lb.extended.cc @@ -745,8 +745,8 @@ TEST_P(TestDumpUserdefinedData, test_dump_userdefined_json) { EXPECT_NE(elm_ptr, nullptr); if (elm_ptr != nullptr) { auto elm_id = elm_ptr->getElmID(); - elm_ptr->valInsert("hello", std::string("world"), should_dump, false); - elm_ptr->valInsert("elephant", 123456789, should_dump, false); + elm_ptr->valInsert("hello", std::string("world"), should_dump, false, false); + elm_ptr->valInsert("elephant", 123456789, should_dump, false, false); lbdh.user_defined_json_[phase][elm_id] = std::make_shared( elm_ptr->toJson() ); @@ -777,6 +777,68 @@ INSTANTIATE_TEST_SUITE_P( DumpUserdefinedDataExplode, TestDumpUserdefinedData, booleans ); +struct TestDumpAttributesFieldData : TestParallelHarnessParam { }; + +TEST_P(TestDumpAttributesFieldData, test_dump_attributes_json) { + bool is_attribute = GetParam(); + + auto this_node = vt::theContext()->getNode(); + auto num_nodes = vt::theContext()->getNumNodes(); + + vt::vrt::collection::CollectionProxy proxy; + auto const range = vt::Index1D(num_nodes * 1); + + // Construct a collection + runInEpochCollective([&] { + proxy = vt::theCollection()->constructCollective( + range, "test_dump_attributes_json" + ); + }); + + vt::vrt::collection::balance::LBDataHolder lbdh; + PhaseType phase = 0; + + { + lbdh.node_data_[phase]; + lbdh.node_comm_[phase]; + + vt::Index1D idx(this_node * 1); + auto elm_ptr = proxy(idx).tryGetLocalPtr(); + EXPECT_NE(elm_ptr, nullptr); + if (elm_ptr != nullptr) { + auto elm_id = elm_ptr->getElmID(); + elm_ptr->valInsert("intSample", 123, false, false, is_attribute); + elm_ptr->valInsert("doubleSample", 123.456, false, false, is_attribute); + elm_ptr->valInsert("stringSample", std::string("abc"), false, false, is_attribute); + + elm_ptr->collectAttributes( + [&](std::string const& key, auto val) { + lbdh.node_user_attributes_[phase][elm_id][key] = val->toVariant(); + } + ); + lbdh.node_data_[phase][elm_id].whole_phase_load = 1.0; + } + } + + auto json_str = getJsonStringForPhase(phase, lbdh); + fmt::print("{}\n", json_str); + if (is_attribute) { + EXPECT_NE(json_str.find("attributes"), std::string::npos); + EXPECT_NE(json_str.find("intSample"), std::string::npos); + EXPECT_NE(json_str.find("doubleSample"), std::string::npos); + EXPECT_NE(json_str.find("stringSample"), std::string::npos); + } else { + EXPECT_EQ(json_str.find("attributes"), std::string::npos); + EXPECT_EQ(json_str.find("intSample"), std::string::npos); + EXPECT_EQ(json_str.find("doubleSample"), std::string::npos); + EXPECT_EQ(json_str.find("stringSample"), std::string::npos); + } +} + +INSTANTIATE_TEST_SUITE_P( + DumpAttributesFieldDataExplode, TestDumpAttributesFieldData, booleans +); + struct SerializationTestCol : vt::Collection { template void serialize(SerializerT &s) { vt::Collection::serialize(s); diff --git a/tests/unit/collection/test_lb_data_holder.cc b/tests/unit/collection/test_lb_data_holder.cc index a5f9cb2c88..af95937b92 100644 --- a/tests/unit/collection/test_lb_data_holder.cc +++ b/tests/unit/collection/test_lb_data_holder.cc @@ -58,7 +58,7 @@ namespace vt { namespace tests { namespace unit { namespace lb { static constexpr int const num_phases = 10; -struct TestOptionalPhasesMetadata : TestParallelHarnessParam { }; +struct TestLBDataHolder : TestParallelHarnessParam { }; void addPhasesDataToJson(nlohmann::json& json, PhaseType amountOfPhasesToAdd, std::vector toSkip) { using ElementIDStruct = vt::vrt::collection::balance::ElementIDStruct; @@ -90,7 +90,7 @@ void addPhasesDataToJson(nlohmann::json& json, PhaseType amountOfPhasesToAdd, st json["phases"] = phases; } -TEST_F(TestOptionalPhasesMetadata, test_no_metadata) { +TEST_F(TestLBDataHolder, test_no_metadata) { using LBDataHolder = vt::vrt::collection::balance::LBDataHolder; nlohmann::json json; @@ -117,7 +117,7 @@ TEST_F(TestOptionalPhasesMetadata, test_no_metadata) { EXPECT_EQ((*outJsonPtr)["identical_to_previous"]["range"], expectedIdenticalRanges); } -TEST_F(TestOptionalPhasesMetadata, test_no_lb_phases_metadata) { +TEST_F(TestLBDataHolder, test_no_lb_phases_metadata) { using LBDataHolder = vt::vrt::collection::balance::LBDataHolder; nlohmann::json json; @@ -146,7 +146,7 @@ TEST_F(TestOptionalPhasesMetadata, test_no_lb_phases_metadata) { EXPECT_EQ((*outJsonPtr)["identical_to_previous"]["range"], expectedIdenticalRanges); } -TEST_F(TestOptionalPhasesMetadata, test_lb_phases_metadata_empty) { +TEST_F(TestLBDataHolder, test_lb_phases_metadata_empty) { using LBDataHolder = vt::vrt::collection::balance::LBDataHolder; nlohmann::json metadata, phasesMetadata, json; @@ -182,7 +182,7 @@ TEST_F(TestOptionalPhasesMetadata, test_lb_phases_metadata_empty) { EXPECT_EQ((*outJsonPtr)["identical_to_previous"]["range"], expectedIdenticalRanges); } -TEST_F(TestOptionalPhasesMetadata, test_lb_phases_metadata_filled) { +TEST_F(TestLBDataHolder, test_lb_phases_metadata_filled) { using LBDataHolder = vt::vrt::collection::balance::LBDataHolder; nlohmann::json metadata, phasesMetadata, json; @@ -218,6 +218,83 @@ TEST_F(TestOptionalPhasesMetadata, test_lb_phases_metadata_filled) { EXPECT_EQ((*outJsonPtr)["identical_to_previous"]["range"], expectedIdenticalRanges); } +TEST_F(TestLBDataHolder, test_lb_rank_attributes) { + using LBDataHolder = vt::vrt::collection::balance::LBDataHolder; + + nlohmann::json json = R"( + { + "type": "LBDatafile", + "metadata" : { + "attributes": { + "intSample": 123, + "doubleSample": 1.99, + "stringSample": "abc" + } + }, + "phases" : [] + } + )"_json; + + LBDataHolder testObj(json); + EXPECT_EQ(123, std::get(testObj.rank_attributes_["intSample"])); + EXPECT_EQ(1.99, std::get(testObj.rank_attributes_["doubleSample"])); + EXPECT_EQ("abc", std::get(testObj.rank_attributes_["stringSample"])); + + auto outAttributesJsonPtr = testObj.rankAttributesToJson(); + ASSERT_TRUE(outAttributesJsonPtr != nullptr); + EXPECT_EQ(123, (*outAttributesJsonPtr)["attributes"]["intSample"]); + EXPECT_EQ(1.99, (*outAttributesJsonPtr)["attributes"]["doubleSample"]); + EXPECT_EQ("abc", (*outAttributesJsonPtr)["attributes"]["stringSample"]); +} + +TEST_F(TestLBDataHolder, test_lb_entity_attributes) { + using LBDataHolder = vt::vrt::collection::balance::LBDataHolder; + + nlohmann::json json = R"( + { + "type": "LBDatafile", + "phases": [ + { + "id": 0, + "tasks": [ + { + "entity": { + "home": 0, + "id": 524291, + "type": "object", + "migratable": true + }, + "node": 0, + "resource": "cpu", + "time": 3.0, + "attributes": { + "intSample": 123, + "doubleSample": 1.99, + "stringSample": "abc" + } + } + ] + } + ] + } + )"_json; + auto id = vt::vrt::collection::balance::ElementIDStruct{524291, 0}; + + LBDataHolder testObj(json); + EXPECT_TRUE(testObj.node_user_attributes_.find(0) != testObj.node_user_attributes_.end()); + EXPECT_TRUE(testObj.node_user_attributes_[0].find(id) != testObj.node_user_attributes_[0].end()); + auto attributes = testObj.node_user_attributes_[0][id]; + EXPECT_EQ(123, std::get(attributes["intSample"])); + EXPECT_EQ(1.99, std::get(attributes["doubleSample"])); + EXPECT_EQ("abc", std::get(attributes["stringSample"])); + + auto outJsonPtr = testObj.toJson(0); + ASSERT_TRUE(outJsonPtr != nullptr); + EXPECT_EQ(123, (*outJsonPtr)["tasks"][0]["attributes"]["intSample"]); + EXPECT_EQ(1.99, (*outJsonPtr)["tasks"][0]["attributes"]["doubleSample"]); + EXPECT_EQ("abc", (*outJsonPtr)["tasks"][0]["attributes"]["stringSample"]); +} + }}}} // end namespace vt::tests::unit::lb #endif /*vt_check_enabled(lblite)*/