Maintenance/checkpointing

.pre-commit-config.yaml

@@ -5,6 +5,7 @@ repos:
     hooks:
       - id: trailing-whitespace
       - id: end-of-file-fixer
+        exclude: 'coverage'
       - id: check-yaml
       - id: check-added-large-files
   - repo: https://github.com/astral-sh/ruff-pre-commit

README.md

@@ -93,6 +93,8 @@ discussions](https://github.com/Helmholtz-AI-Energy/propulate/discussions) :octo
 
 ## Installation
 
+- Installations of MPI and hdf5 are required.
+- ``Propulate`` depends on parallel hdf5. Try ``CC="mpicc" HDF5_MPI="ON" pip install --no-binary=h5py h5py`` to install the parallel enabled Python bindings or follow instructions [here](https://docs.h5py.org/en/stable/build.html#install).
 - You can install the **latest stable release** from PyPI: ``pip install propulate``
 - If you need the **latest updates**, you can also install ``Propulate`` directly from the master branch.
 Pull and run ``pip install .``.

propulate/_globals.py

@@ -2,5 +2,4 @@
 SYNCHRONIZATION_TAG = 2
 INIT_TAG = 3
 POPULATION_TAG = 4
-DUMP_TAG = 5
 MIGRATION_TAG = 6

propulate/migrator.py

@@ -4,6 +4,7 @@
 from pathlib import Path
 from typing import Callable, Generator, List, Optional, Type, Union
 
+import h5py
 import numpy as np
 from mpi4py import MPI
 
@@ -123,7 +124,7 @@ def __init__(
             Individual
         ] = []  # Emigrated individuals to be deactivated on sending island
 
-    def _send_emigrants(self) -> None:
+    def _send_emigrants(self, hdf5_checkpoint: h5py.File) -> None:
         """Perform migration, i.e. island sends individuals out to other islands."""
         log_string = (
             f"Island {self.island_idx} Worker {self.island_comm.rank} "
@@ -188,6 +189,10 @@ def _send_emigrants(self) -> None:
 
                 # Send emigrants to target island.
                 departing = copy.deepcopy(emigrants)
+                for ind in departing:
+                    hdf5_checkpoint[f"{ind.island}"][f"{ind.island_rank}"][
+                        "active_on_island"
+                    ][ind.generation, self.island_idx] = False
                 # Determine new responsible worker on target island.
                 for ind in departing:
                     ind.current = self.rng.randrange(0, count)
@@ -236,7 +241,7 @@ def _send_emigrants(self) -> None:
                 f"to select {num_emigrants} migrants."
             )
 
-    def _receive_immigrants(self) -> None:
+    def _receive_immigrants(self, hdf5_checkpoint: h5py.File) -> None:
         """
         Check for and possibly receive immigrants send by other islands.
 
@@ -284,6 +289,9 @@ def _receive_immigrants(self) -> None:
                             log_string
                             + f"Identical immigrant {immigrant} already active on target  island {self.island_idx}."
                         )
+                    hdf5_checkpoint[f"{immigrant.island}"][f"{immigrant.island_rank}"][
+                        "active_on_island"
+                    ][immigrant.generation] = True
                     self.population.append(
                         copy.deepcopy(immigrant)
                     )  # Append immigrant to population.
@@ -421,53 +429,51 @@ def propulate(self, logging_interval: int = 10, debug: int = 1) -> None:
         if self.propulate_comm is None:
             while self.generations <= -1 or self.generation < self.generations:
                 # Breed and evaluate individual.
-                self._evaluate_individual()
+                # TODO this should be refactored, the subworkers don't need the logfile
+                # TODO this needs to be addressed before merge, since multirank workers should fail with this
+                self._evaluate_individual(None)
                 self.generation += 1
             return
 
         if self.island_comm.rank == 0:
             log.info(f"Island {self.island_idx} has {self.island_comm.size} workers.")
 
-        dump = True if self.island_comm.rank == 0 else False
         migration = True if self.migration_prob > 0 else False
         self.propulate_comm.barrier()
 
         # Loop over generations.
-        while self.generations <= -1 or self.generation < self.generations:
-            if self.generation % int(logging_interval) == 0:
-                log.info(
-                    f"Island {self.island_idx} Worker {self.island_comm.rank}: In generation {self.generation}..."
-                )
-
-            # Breed and evaluate individual.
-            self._evaluate_individual()
-
-            # Check for and possibly receive incoming individuals from other intra-island workers.
-            self._receive_intra_island_individuals()
+        # TODO this should probably be refactored, checkpointing can probably be handled in one place
+        with h5py.File(
+            self.checkpoint_path, "a", driver="mpio", comm=self.propulate_comm
+        ) as f:
+            while self.generation < self.generations:
+                if self.generation % int(logging_interval) == 0:
+                    log.info(
+                        f"Island {self.island_idx} Worker {self.island_comm.rank}: In generation {self.generation}..."
+                    )
 
-            # Migration.
-            if migration:
-                # Emigration: Island sends individuals out.
-                # Happens on per-worker basis with certain probability.
-                if self.rng.random() < self.migration_prob:
-                    self._send_emigrants()
+                # Breed and evaluate individual.
+                self._evaluate_individual(f)
 
-                # Immigration: Check for incoming individuals from other islands.
-                self._receive_immigrants()
+                # Check for and possibly receive incoming individuals from other intra-island workers.
+                self._receive_intra_island_individuals()
 
-                # Emigration: Check for emigrants from other intra-island workers to be deactivated.
-                self._deactivate_emigrants()
-                if debug == 2:
-                    check = self._check_emigrants_to_deactivate()
-                    assert check is False
+                # Migration.
+                if migration:
+                    # Emigration: Island sends individuals out.
+                    # Happens on per-worker basis with certain probability.
+                    if self.rng.random() < self.migration_prob:
+                        self._send_emigrants(f)
 
-            if dump:  # Dump checkpoint.
-                self._dump_checkpoint()
+                    # Immigration: Check for incoming individuals from other islands.
+                    self._receive_immigrants(f)
 
-            dump = (
-                self._determine_worker_dumping_next()
-            )  # Determine worker dumping checkpoint in the next generation.
-            self.generation += 1  # Go to next generation.
+                    # Emigration: Check for emigrants from other intra-island workers to be deactivated.
+                    self._deactivate_emigrants()
+                    if debug == 2:
+                        check = self._check_emigrants_to_deactivate()
+                        assert check is False
+                self.generation += 1  # Go to next generation.
 
         # Having completed all generations, the workers have to wait for each other.
         # Once all workers are done, they should check for incoming messages once again
@@ -486,7 +492,10 @@ def propulate(self, logging_interval: int = 10, debug: int = 1) -> None:
 
         if migration:
             # Final check for incoming individuals from other islands.
-            self._receive_immigrants()
+            with h5py.File(
+                self.checkpoint_path, "a", driver="mpio", comm=self.propulate_comm
+            ) as f:
+                self._receive_immigrants(f)
             self.propulate_comm.barrier()
 
             # Emigration: Final check for emigrants from other intra-island workers to be deactivated.
@@ -507,12 +516,3 @@ def propulate(self, logging_interval: int = 10, debug: int = 1) -> None:
                         raise ValueError(
                             "There should not be any individuals left that need to be deactivated."
                         )
-
-            self.propulate_comm.barrier()
-
-        # Final checkpointing on rank 0.
-        if self.island_comm.rank == 0:
-            self._dump_final_checkpoint()  # Dump checkpoint.
-        self.propulate_comm.barrier()
-        _ = self._determine_worker_dumping_next()
-        self.propulate_comm.barrier()

propulate/pollinator.py

@@ -1,9 +1,11 @@
 import copy
 import logging
 import random
+import time
 from pathlib import Path
 from typing import Callable, Generator, List, Optional, Tuple, Type, Union
 
+import h5py
 import numpy as np
 from mpi4py import MPI
 
@@ -204,7 +206,8 @@ def _send_emigrants(self) -> None:
                 f"to select {num_emigrants} migrants."
             )
 
-    def _receive_immigrants(self) -> None:
+    # TODO implement checkpoint update
+    def _receive_immigrants(self, hdf5_checkpoint: h5py.File) -> None:
         """Check for and possibly receive immigrants send by other islands."""
         replace_num = 0
         log_string = (
@@ -400,53 +403,54 @@ def propulate(self, logging_interval: int = 10, debug: int = 1) -> None:
         debug : int, optional
             The debug level; 0 - silent; 1 - moderate, 2 - noisy (debug mode). Default is 1.
         """
+        self.start_time = time.time_ns()
         if self.worker_sub_comm != MPI.COMM_SELF:
             self.generation = self.worker_sub_comm.bcast(self.generation, root=0)
         if self.propulate_comm is None:
             while self.generations <= -1 or self.generation < self.generations:
                 # Breed and evaluate individual.
-                self._evaluate_individual()
+                # TODO this should be refactored, the subworkers don't need the logfile
+                # TODO this needs to be addressed before merge, since multirank workers should fail with this
+                self._evaluate_individual(None)
                 self.generation += 1
             return
         if self.island_comm.rank == 0:
             log.info(f"Island {self.island_idx} has {self.island_comm.size} workers.")
 
-        dump = True if self.island_comm.rank == 0 else False
         migration = True if self.migration_prob > 0 else False
         self.propulate_comm.barrier()
 
         # Loop over generations.
-        while self.generations <= -1 or self.generation < self.generations:
-            if debug == 1 and self.generation % int(logging_interval) == 0:
-                log.info(
-                    f"Island {self.island_idx} Worker {self.island_comm.rank}: In generation {self.generation}..."
-                )
-
-            # Breed and evaluate individual.
-            self._evaluate_individual()
+        # TODO this should probably be refactored, checkpointing can probably be handled in one place
+        with h5py.File(
+            self.checkpoint_path, "a", driver="mpio", comm=self.propulate_comm
+        ) as f:
+            while self.generation < self.generations:
+                if self.generation % int(logging_interval) == 0:
+                    log.info(
+                        f"Island {self.island_idx} Worker {self.island_comm.rank}: In generation {self.generation}..."
+                    )
 
-            # Check for and possibly receive incoming individuals from other intra-island workers.
-            self._receive_intra_island_individuals()
+                # Breed and evaluate individual.
+                self._evaluate_individual(f)
 
-            if migration:
-                # Emigration: Island sends individuals out.
-                # Happens on per-worker basis with certain probability.
-                if self.rng.random() < self.migration_prob:
-                    self._send_emigrants()
+                # Check for and possibly receive incoming individuals from other intra-island workers.
+                self._receive_intra_island_individuals()
 
-                # Immigration: Island checks for incoming individuals from other islands.
-                self._receive_immigrants()
+                if migration:
+                    # Emigration: Island sends individuals out.
+                    # Happens on per-worker basis with certain probability.
+                    if self.rng.random() < self.migration_prob:
+                        self._send_emigrants()
 
-                # Immigration: Check for individuals replaced by other intra-island workers to be deactivated.
-                self._deactivate_replaced_individuals()
+                    # Immigration: Island checks for incoming individuals from other islands.
+                    # TODO this should probably update the checkpoint so it needs to pass the handle
+                    self._receive_immigrants(None)
 
-            if dump:  # Dump checkpoint.
-                self._dump_checkpoint()
+                    # Immigration: Check for individuals replaced by other intra-island workers to be deactivated.
+                    self._deactivate_replaced_individuals()
 
-            dump = (
-                self._determine_worker_dumping_next()
-            )  # Determine worker dumping checkpoint in the next generation.
-            self.generation += 1  # Go to next generation.
+                self.generation += 1  # Go to next generation.
 
         # Having completed all generations, the workers have to wait for each other.
         # Once all workers are done, they should check for incoming messages once again
@@ -464,7 +468,8 @@ def propulate(self, logging_interval: int = 10, debug: int = 1) -> None:
 
         if migration:
             # Final check for incoming individuals from other islands.
-            self._receive_immigrants()
+            # TODO this needs to update the checkpoint
+            self._receive_immigrants(None)
             self.propulate_comm.barrier()
 
             # Immigration: Final check for individuals replaced by other intra-island workers to be deactivated.
@@ -477,11 +482,3 @@ def propulate(self, logging_interval: int = 10, debug: int = 1) -> None:
                     f"Finally {len(self.replaced)} individual(s) in replaced: {self.replaced}:\n{self.population}"
                 )
                 self._deactivate_replaced_individuals()
-            self.propulate_comm.barrier()
-
-        # Final checkpointing on rank 0.
-        if self.island_comm.rank == 0:
-            self._dump_final_checkpoint()  # Dump checkpoint.
-        self.propulate_comm.barrier()
-        _ = self._determine_worker_dumping_next()
-        self.propulate_comm.barrier()

propulate/population.py

@@ -85,6 +85,7 @@ def __init__(
         self.migration_history: str = ""  # migration history
         self.evaltime = float("inf")  # evaluation time
         self.evalperiod = 0.0  # evaluation duration
+        self.island_rank = 0  # rank in the island comm
 
         # NOTE needed for PSO type propagators
         self.velocity = velocity
@@ -99,11 +100,11 @@ def __getitem__(self, key: str) -> Union[float, int, str]:
             return self.mapping[key]
         else:
             # continuous variable
-            if self.types[key] == float:
+            if self.types[key] is float:
                 return float(self.position[self.offsets[key]].item())
-            elif self.types[key] == int:
+            elif self.types[key] is int:
                 return int(np.rint(self.position[self.offsets[key]]).item())
-            elif self.types[key] == str:
+            elif self.types[key] is str:
                 offset = self.offsets[key]
                 upper = self.offsets[key] + len(self.limits[key])
                 return str(
@@ -120,13 +121,13 @@ def __setitem__(self, key: str, newvalue: Union[float, int, str, Any]) -> None:
         else:
             if key not in self.limits:
                 raise ValueError("Unknown gene.")
-            if self.types[key] == float:
+            if self.types[key] is float:
                 assert isinstance(newvalue, float)
                 self.position[self.offsets[key]] = newvalue
-            elif self.types[key] == int:
+            elif self.types[key] is int:
                 assert isinstance(newvalue, int)
                 self.position[self.offsets[key]] = float(newvalue)
-            elif self.types[key] == str:
+            elif self.types[key] is str:
                 assert newvalue in self.limits[key]
                 offset = self.offsets[key]
                 upper = len(self.limits[key])

propulate/propagators/base.py

@@ -73,6 +73,10 @@ def __init__(
         if rng is None:
             rng = random.Random()
         self.rng = rng  # Random number generator
+        # TODO
+        self.limits: Mapping[
+            str, Union[Tuple[float, float], Tuple[int, int], Tuple[str, ...]]
+        ] = dict()
 
     def __call__(self, inds: List[Individual]) -> Union[List[Individual], Individual]:
         """
@@ -177,6 +181,9 @@ class Conditional(Propagator):
 
     def __init__(
         self,
+        limits: Mapping[
+            str, Union[Tuple[float, float], Tuple[int, int], Tuple[str, ...]]
+        ],
         pop_size: int,
         true_prop: Propagator,
         false_prop: Propagator,
@@ -203,6 +210,7 @@ def __init__(
         self.pop_size = pop_size
         self.true_prop = true_prop
         self.false_prop = false_prop
+        self.limits = limits
 
     def __call__(self, inds: List[Individual]) -> Union[List[Individual], Individual]:
         """

propulate/propagators/cmaes.py

@@ -751,6 +751,7 @@ def __init__(
             The separate random number generator for the Propulate optimization.
         """
         self.adapter = adapter
+        self.limits = limits
         problem_dimension = len(limits)
         # Number of individuals considered for each generation
         lambd = (

propulate/propagators/pso.py

@@ -270,7 +270,7 @@ def __call__(self, individuals: List[Individual]) -> Individual:
         individuals : List[propulate.population.Individual]
             The list of individuals that must at least contain one individual that belongs to the propagator.
             This list is used to calculate personal and global best of the particle and the swarm,
-            respectively, and then to update the particle based on the retrieved results. 
+            respectively, and then to update the particle based on the retrieved results.
             cannot be used as ``Individual`` objects are converted to particles first.
 
         Returns