#2201: replay: set the actual LBData structure so user-defined fields…

… are passed to LB
DARMA-tasking · Nov 29, 2023 · 7f051f8 · 7f051f8
1 parent 9789a48
commit 7f051f8
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/src/vt/vrt/collection/balance/node_lb_data.h b/src/vt/vrt/collection/balance/node_lb_data.h
@@ -256,6 +256,18 @@ struct NodeLBData : runtime::component::Component<NodeLBData> {
    */
   LBDataHolder* getLBData() { return lb_data_.get(); }
 
+  /**
+   * \brief Set the LB data
+   *
+   * \note Used for replay
+   * \warning Squashes all the existing data!
+   *
+   * \param[in] lb_data LB data to override
+   */
+  void setLBData(LBDataHolder const& lb_data) {
+    lb_data_ = std::make_unique<LBDataHolder>(lb_data);
+  }
+
   template <typename SerializerT>
   void serialize(SerializerT& s) {
     s | proxy_

diff --git a/src/vt/vrt/collection/balance/workload_replay.cc b/src/vt/vrt/collection/balance/workload_replay.cc
@@ -222,6 +222,7 @@ void replayWorkloads(
       auto cb = theCB()->makeFunc<ReassignmentMsg>(
         vt::pipe::LifetimeEnum::Once, postLBWork
       );
+      theNodeLBData()->setLBData(*workloads.get());
       theLBManager()->selectStartLB(phase, cb);
     });
     runInEpochCollective("WorkloadReplayDriver -> destroyLB", [&] {