From a66f5cf39f44c49733eb485bca207abfc0fb46cc Mon Sep 17 00:00:00 2001 From: JHopeCollins Date: Tue, 7 May 2024 17:40:18 +0100 Subject: [PATCH 1/5] ManualEnsemble class for specifying all comms in Ensemble --- asQ/ensemble.py | 76 ++++++++++++++++++++++++++++++++++++------------- 1 file changed, 56 insertions(+), 20 deletions(-) diff --git a/asQ/ensemble.py b/asQ/ensemble.py index cd774a63..15ad255e 100644 --- a/asQ/ensemble.py +++ b/asQ/ensemble.py @@ -1,7 +1,7 @@ from firedrake import COMM_WORLD, Ensemble -from pyop2.mpi import internal_comm +from pyop2.mpi import MPI, internal_comm, is_pyop2_comm, PyOP2CommError -__all__ = ['create_ensemble', 'split_ensemble', 'EnsembleConnector'] +__all__ = ['create_ensemble', 'split_ensemble'] def create_ensemble(time_partition, comm=COMM_WORLD): @@ -42,35 +42,71 @@ def split_ensemble(ensemble, split_size, **kwargs): split_rank = ensemble.ensemble_comm.rank // split_size # create split_ensemble.global_comm - split_comm = ensemble.global_comm.Split(color=split_rank, - key=ensemble.global_comm.rank) + split_global_comm = ensemble.global_comm.Split(color=split_rank, + key=ensemble.global_comm.rank) - return EnsembleConnector(split_comm, ensemble.comm, split_size, **kwargs) + # create split_ensemble.ensemble_comm + split_ensemble_comm = ensemble.ensemble_comm.Split(color=split_rank, + key=ensemble.global_comm.rank) + return ManualEnsemble(split_global_comm, ensemble.comm, split_ensemble_comm, **kwargs) -class EnsembleConnector(Ensemble): - def __init__(self, global_comm, local_comm, nmembers, **kwargs): + +class ManualEnsemble(Ensemble): + def __init__(self, global_comm, spatial_comm, ensemble_comm, **kwargs): """ - An Ensemble created from provided spatial communicators (ensemble.comm). + An Ensemble created from provided comms. :arg global_comm: global communicator the Ensemble is defined over. - :arg local_comm: communicator to use for the Ensemble.comm member. - :arg nmembers: number of Ensemble members (ensemble.ensemble_comm.size). + :arg spatial_comm: communicator to use for the Ensemble.comm member. + :arg ensemble_comm: communicator to use for the Ensemble.ensemble_comm member. + + The global_comm, spatial_comm, and ensemble_comm must have the same logical meaning + as they do in firedrake.Ensemble. i.e. the global_comm is the union of a cartesian + product of multiple spatial_comms and ensemble_comms. + - ManualEnsemble is logically defined over all ranks in global_comm. + - Each rank in global_comm belongs to only one spatial_comm and one ensemble_comm. + - The size of the intersection of any (spatial_comm, ensemble_comm) pair is 1. + + WARNING: Not meeting these requirements may produce in errors, hangs, and nonsensical results. """ - if nmembers*local_comm.size != global_comm.size: - msg = "The global ensemble must have the same number of ranks as the sum of the local comms" - raise ValueError(msg) + # are we handed user comms? + + for comm in (global_comm, spatial_comm, ensemble_comm): + if is_pyop2_comm(comm): + raise PyOP2CommError("Cannot construct Ensemble from PyOP2 internal comm") + + # check cartesian product consistency + + if spatial_comm.size*ensemble_comm.size != global_comm.size: + msg = "The global comm must have the same number of ranks as the product of spatial and ensemble comms" + raise PyOP2CommError(msg) + + global_group = global_comm.Get_group() + spatial_group = spatial_comm.Get_group() + ensemble_group = ensemble_comm.Get_group() + + if MPI.Group.Intersection(spatial_group, ensemble_group).size != 1: + raise PyOP2CommError("spatial and ensemble comms must be cartesian product in global_comm") + if MPI.Group.Intersection(global_group, spatial_group).size != spatial_group.size: + raise PyOP2CommError("spatial_comm must be subgroup of global_comm") + if MPI.Group.Intersection(global_group, ensemble_group).size != ensemble_group.size: + raise PyOP2CommError("ensemble_comm must be subgroup of global_comm") + + # create internal duplicates and name comms for debugging + ensemble_name = kwargs.get("name", "Ensemble") - ensemble_name = kwargs.get("ensemble_name", "Ensemble") self.global_comm = global_comm + if not hasattr(self.global_comm, "name"): + self.global_comm.name = f"{ensemble_name} global comm" self._comm = internal_comm(self.global_comm, self) - self.comm = local_comm - self.comm.name = f"{ensemble_name} spatial comm" + self.comm = spatial_comm + if not hasattr(self.comm, "name"): + self.comm.name = f"{ensemble_name} spatial comm" self._spatial_comm = internal_comm(self.comm, self) - self.ensemble_comm = self.global_comm.Split(color=self.comm.rank, - key=global_comm.rank) - self.ensemble_comm.name = f"{ensemble_name} ensemble comm" - + self.ensemble_comm = ensemble_comm + if not hasattr(self.ensemble_comm, "name"): + self.ensemble_comm.name = f"{ensemble_name} ensemble comm" self._ensemble_comm = internal_comm(self.ensemble_comm, self) From b5ccbfe5cbac162cdcfa7f734b0e94c9b9b84b81 Mon Sep 17 00:00:00 2001 From: JHopeCollins Date: Tue, 7 May 2024 18:31:12 +0100 Subject: [PATCH 2/5] use MPI.Group.Compare for the comm tests in ManualEnsemble --- asQ/ensemble.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/asQ/ensemble.py b/asQ/ensemble.py index 15ad255e..d5c3a134 100644 --- a/asQ/ensemble.py +++ b/asQ/ensemble.py @@ -88,9 +88,13 @@ def __init__(self, global_comm, spatial_comm, ensemble_comm, **kwargs): if MPI.Group.Intersection(spatial_group, ensemble_group).size != 1: raise PyOP2CommError("spatial and ensemble comms must be cartesian product in global_comm") - if MPI.Group.Intersection(global_group, spatial_group).size != spatial_group.size: + + spatial_intersection = MPI.Group.Intersection(global_group, spatial_group) + ensemble_intersection = MPI.Group.Intersection(global_group, ensemble_group) + + if MPI.Group.Compare(spatial_intersection, spatial_group) not in {MPI.IDENT, MPI.CONGRUENT}: raise PyOP2CommError("spatial_comm must be subgroup of global_comm") - if MPI.Group.Intersection(global_group, ensemble_group).size != ensemble_group.size: + if MPI.Group.Compare(ensemble_intersection, ensemble_group) not in {MPI.IDENT, MPI.CONGRUENT}: raise PyOP2CommError("ensemble_comm must be subgroup of global_comm") # create internal duplicates and name comms for debugging From 9113af8b8fe01e23d63f1a047170691fa2167268 Mon Sep 17 00:00:00 2001 From: JHopeCollins Date: Tue, 7 May 2024 18:37:15 +0100 Subject: [PATCH 3/5] use MPI.Group.Compare for the comm tests in ManualEnsemble --- asQ/ensemble.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/asQ/ensemble.py b/asQ/ensemble.py index d5c3a134..a05a2632 100644 --- a/asQ/ensemble.py +++ b/asQ/ensemble.py @@ -89,12 +89,11 @@ def __init__(self, global_comm, spatial_comm, ensemble_comm, **kwargs): if MPI.Group.Intersection(spatial_group, ensemble_group).size != 1: raise PyOP2CommError("spatial and ensemble comms must be cartesian product in global_comm") - spatial_intersection = MPI.Group.Intersection(global_group, spatial_group) - ensemble_intersection = MPI.Group.Intersection(global_group, ensemble_group) + is_subgroup = lambda sub, group: MPI.Group.Compare(sub, MPI.Group.Intersection(sub, group)) in {MPI.IDENT, MPI.CONGRUENT} - if MPI.Group.Compare(spatial_intersection, spatial_group) not in {MPI.IDENT, MPI.CONGRUENT}: + if not is_subgroup(spatial_group, global_group): raise PyOP2CommError("spatial_comm must be subgroup of global_comm") - if MPI.Group.Compare(ensemble_intersection, ensemble_group) not in {MPI.IDENT, MPI.CONGRUENT}: + if not is_subgroup(ensemble_group, global_group): raise PyOP2CommError("ensemble_comm must be subgroup of global_comm") # create internal duplicates and name comms for debugging From 714d9fbad0a9849f859be7b648b66119dbc2cf62 Mon Sep 17 00:00:00 2001 From: JHopeCollins Date: Tue, 14 May 2024 11:13:14 +0100 Subject: [PATCH 4/5] split_ensemble - use ensemble_comm.rank as key to split ensemble_comm --- asQ/ensemble.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/asQ/ensemble.py b/asQ/ensemble.py index a05a2632..1446b7d4 100644 --- a/asQ/ensemble.py +++ b/asQ/ensemble.py @@ -47,7 +47,7 @@ def split_ensemble(ensemble, split_size, **kwargs): # create split_ensemble.ensemble_comm split_ensemble_comm = ensemble.ensemble_comm.Split(color=split_rank, - key=ensemble.global_comm.rank) + key=ensemble.ensemble_comm.rank) return ManualEnsemble(split_global_comm, ensemble.comm, split_ensemble_comm, **kwargs) From b7fd79b40b9b8bc84ef27582c1102db3b33591fd Mon Sep 17 00:00:00 2001 From: JHopeCollins Date: Tue, 14 May 2024 13:54:17 +0100 Subject: [PATCH 5/5] make sure the new comms in split_comm are cleaned up --- asQ/ensemble.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/asQ/ensemble.py b/asQ/ensemble.py index 1446b7d4..fd50b06c 100644 --- a/asQ/ensemble.py +++ b/asQ/ensemble.py @@ -1,3 +1,4 @@ +import weakref from firedrake import COMM_WORLD, Ensemble from pyop2.mpi import MPI, internal_comm, is_pyop2_comm, PyOP2CommError @@ -49,7 +50,13 @@ def split_ensemble(ensemble, split_size, **kwargs): split_ensemble_comm = ensemble.ensemble_comm.Split(color=split_rank, key=ensemble.ensemble_comm.rank) - return ManualEnsemble(split_global_comm, ensemble.comm, split_ensemble_comm, **kwargs) + new_ensemble = ManualEnsemble(split_global_comm, ensemble.comm, split_ensemble_comm, **kwargs) + + # make sure the new comms are cleaned up when the split ensemble goes out of scope + weakref.finalize(new_ensemble, split_global_comm.Free) + weakref.finalize(new_ensemble, split_ensemble_comm.Free) + + return new_ensemble class ManualEnsemble(Ensemble): @@ -69,6 +76,8 @@ def __init__(self, global_comm, spatial_comm, ensemble_comm, **kwargs): - The size of the intersection of any (spatial_comm, ensemble_comm) pair is 1. WARNING: Not meeting these requirements may produce in errors, hangs, and nonsensical results. + + ManualEnsemble will not Free any of the comms. This is the responsibility of the user. """ # are we handed user comms?