diff --git a/src/vt/vrt/collection/balance/lb_data_holder.cc b/src/vt/vrt/collection/balance/lb_data_holder.cc index 15ffd34827..fdc1487606 100644 --- a/src/vt/vrt/collection/balance/lb_data_holder.cc +++ b/src/vt/vrt/collection/balance/lb_data_holder.cc @@ -144,6 +144,15 @@ std::unique_ptr LBDataHolder::rankAttributesToJson() const { return std::make_unique(std::move(j)); } +void LBDataHolder::addInitialTask(nlohmann::json& j, std::size_t n) const { + j["tasks"][n]["resource"] = "cpu"; + j["tasks"][n]["node"] = vt::theContext()->getNode(); + j["tasks"][n]["time"] = 0; + outputEntity( + j["tasks"][n]["entity"], ElementIDStruct() + ); +} + std::unique_ptr LBDataHolder::toJson(PhaseType phase) const { using json = nlohmann::json; @@ -194,6 +203,10 @@ std::unique_ptr LBDataHolder::toJson(PhaseType phase) const { i++; } + + if ((phase == 0) and (i > 0)) { + addInitialTask(j, i); + } } i = 0; diff --git a/src/vt/vrt/collection/balance/lb_data_holder.h b/src/vt/vrt/collection/balance/lb_data_holder.h index 39af9eded7..a3badad166 100644 --- a/src/vt/vrt/collection/balance/lb_data_holder.h +++ b/src/vt/vrt/collection/balance/lb_data_holder.h @@ -125,6 +125,8 @@ struct LBDataHolder { */ void outputEntity(nlohmann::json& j, ElementIDStruct const& elm_id) const; + void addInitialTask(nlohmann::json& j, std::size_t n) const; + /** * \brief Read the LB phase's metadata * diff --git a/tests/unit/collection/test_lb.extended.cc b/tests/unit/collection/test_lb.extended.cc index a865ef5113..2c6379cc38 100644 --- a/tests/unit/collection/test_lb.extended.cc +++ b/tests/unit/collection/test_lb.extended.cc @@ -438,6 +438,10 @@ TEST_P(TestNodeLBDataDumper, test_node_lb_data_dumping_with_interval) { EXPECT_EQ(phase["user_defined"]["new_time"], static_cast(phase["id"])); } + auto num_tasks = json["phases"][0]["tasks"].size(); + auto entity = json["phases"][0]["tasks"][num_tasks - 1]["entity"]; + EXPECT_EQ(entity["home"], 0); + EXPECT_EQ(entity["id"], 0); }); if (vt::theContext()->getNode() == 0) { @@ -615,6 +619,10 @@ TEST_F(TestRestoreLBData, test_restore_lb_data_data_1) { auto &orig_load_map = phase_data.second; for (auto &entry : read_load_map) { auto read_elm_id = entry.first; + if ((read_elm_id.id == vt::elm::no_element_id) + and (read_elm_id.getHomeNode() == 0)) { + continue; + } EXPECT_FALSE(orig_load_map.find(read_elm_id) == orig_load_map.end()); if (orig_load_map.find(read_elm_id) == orig_load_map.end()) { fmt::print(