#482: add flexibilty for multiple phases

DARMA-tasking · Dec 19, 2023 · b0d2c38 · b0d2c38
1 parent 72df2e2
commit b0d2c38
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/src/lbaf/Execution/lbsAlgorithmBase.py b/src/lbaf/Execution/lbsAlgorithmBase.py
@@ -225,8 +225,8 @@ def _initialize(self, p_id, phases, distributions, statistics):
             self._logger.error("Algorithm execution requires a dictionary of phases")
             raise SystemExit(1)
 
-        self._initial_communications = phases[p_id].get_communications()
-        self._metadata = phases[p_id].get_metadata()
+        self._initial_communications[p_id] = phases[p_id].get_communications()
+        self._metadata[p_id] = phases[p_id].get_metadata()
 
         # Create a new phase to preserve phase to be rebalanced
         self._logger.info(f"Creating new phase {p_id} for rebalancing")

diff --git a/src/lbaf/Execution/lbsRuntime.py b/src/lbaf/Execution/lbsRuntime.py
@@ -106,11 +106,11 @@ def execute(self, p_id: int, phase_increment=0):
             pp.set_id((pp_id := pp.get_id() + phase_increment))
 
             # Share initial communication with new phase
-            initial_communications = self.__algorithm.get_initial_communications()
+            initial_communications = self.__algorithm.get_initial_communications()[pp.get_id()]
             pp.set_communications(initial_communications)
 
             # Share original metadata with new phase
-            metadata = self.__algorithm.get_metadata()
+            metadata = self.__algorithm.get_metadata()[pp.get_id()]
             pp.set_metadata(metadata)
 
             self.__logger.info(f"Created rebalanced phase {pp_id}")