bring back explicit window creation / distruction

Author: bknueven

mpisppy/cylinders/spcommunicator.py

@@ -146,17 +146,12 @@ def __init__(self, spbase_object, fullcomm, strata_comm, cylinder_comm, communic
         # on the problem data
         self._field_lengths = FieldLengths(self.opt)
 
+        self.window = None
+
         # attach the SPCommunicator to
         # the SPBase object
         self.opt.spcomm = self
 
-        self.register_send_fields()
-
-        self._make_windows()
-        self._create_field_rank_mappings()
-
-        self.register_receive_fields()
-
         return
 
     def _make_key(self, field: Field, origin: int):
@@ -273,13 +268,38 @@ def hub_finalize(self):
     def allreduce_or(self, val):
         return self.opt.allreduce_or(val)
 
-    def _make_windows(self) -> None:
+    def make_windows(self) -> None:
+        """ Make MPI windows: blocking call for all ranks in `strata_comm`.
+        """
+
+        if self.window is not None:
+            return
+
+        self.register_send_fields()
 
         window_spec = self._build_window_spec()
         self.window = SPWindow(window_spec, self.strata_comm)
 
+        self._create_field_rank_mappings()
+        self.register_receive_fields()
+
         return
 
+    def free_windows(self) -> None:
+        """ Free MPI windows: blocking call for all ranks in `strata_comm`.
+        """
+
+        if self.window is None:
+            return
+
+        self.receive_buffers = {}
+        self.send_buffers = {}
+        self.receive_field_spcomms = {}
+
+        self.window.free()
+
+        self.window = None
+
     def is_send_field_registered(self, field: Field) -> bool:
         return field in self.send_buffers
 

mpisppy/cylinders/spwindow.py

@@ -13,7 +13,6 @@
 import numpy.typing as nptyping
 
 import enum
-import weakref
 
 import pyomo.environ as pyo
 
@@ -112,7 +111,6 @@ def __init__(self, my_fields: dict, strata_comm: MPI.Comm, field_order=None):
         self.buffer_length = total_buffer_length
         self.window = MPI.Win.Allocate(window_size_bytes, MPI.DOUBLE.size, comm=strata_comm)
         # ensure the memory allocated for the window is freed
-        self._window_finalizer = weakref.finalize(self, self.window.Free)
         self.buff = np.ndarray(dtype="d", shape=(total_buffer_length,), buffer=self.window.tomemory())
         self.buff[:] = np.nan
 
@@ -126,6 +124,17 @@ def __init__(self, my_fields: dict, strata_comm: MPI.Comm, field_order=None):
 
         return
 
+    def free(self):
+        if self.window is not None:
+            self.window.Free()
+            self.buff = None
+            self.buffer_layout = None
+            self.buffer_length = 0
+            self.window = None
+            self.strata_buffer_layouts = None
+            self.window = None
+        return
+
     #### Functions ####
     def get(self, dest: nptyping.ArrayLike, strata_rank: int, field: Field):
 

mpisppy/spin_the_wheel.py

@@ -127,6 +127,8 @@ def run(self, comm_world=None):
         spcomm = sp_class(opt, fullcomm, strata_comm, cylinder_comm,
                           communicator_list, **sp_kwargs)
 
+        spcomm.make_windows()
+
         # Run main()
         if strata_rank == 0:
             spcomm.setup_hub()
@@ -148,6 +150,7 @@ def run(self, comm_world=None):
 
         ## give the hub the chance to catch new values
         spcomm.hub_finalize()
+        spcomm.free_windows()
 
         fullcomm.Barrier()
         global_toc("Finalize Complete")