automatically register send fields based on class attributes

Author: bknueven

mpisppy/cylinders/cross_scen_spoke.py

@@ -11,12 +11,16 @@
 from mpisppy import MPI
 from mpisppy.utils.lshaped_cuts import LShapedCutGenerator
 from mpisppy.cylinders.spwindow import Field
+from mpisppy.cylinders.spoke import Spoke
 
 import numpy as np
 import pyomo.environ as pyo
-import mpisppy.cylinders.spoke as spoke
 
-class CrossScenarioCutSpoke(spoke.Spoke):
+class CrossScenarioCutSpoke(Spoke):
+
+    send_fields = (*Spoke.send_fields, Field.CROSS_SCENARIO_CUT)
+    receive_fields = (*Spoke.receive_fields, Field.NONANT, Field.CROSS_SCENARIO_COST)
+    optional_receive_fields = (*Spoke.optional_receive_fields, )
 
     def register_send_fields(self) -> None:
 
@@ -35,15 +39,13 @@ def register_send_fields(self) -> None:
         (self.nonant_per_scen, remainder) = divmod(vbuflen, local_scen_count)
         assert(remainder == 0)
 
-        ## the _locals will also have the kill signal
         self.all_nonant_len = vbuflen
         self.all_eta_len = nscen*local_scen_count
 
         self.all_nonants = self.register_recv_field(Field.NONANT, 0, vbuflen)
         self.all_etas = self.register_recv_field(Field.CROSS_SCENARIO_COST, 0, nscen * nscen)
 
-        self.all_coefs = self.register_send_field(Field.CROSS_SCENARIO_CUT,
-                                                  nscen*(self.nonant_per_scen + 1 + 1))
+        self.all_coefs = self.send_buffers[Field.CROSS_SCENARIO_CUT]
 
         return
 
@@ -301,7 +303,6 @@ def main(self):
 
         # main loop
         while not (self.got_kill_signal()):
-            # if self._new_locals:
             if self.all_nonants.is_new() and self.all_etas.is_new():
                 self.make_cut()
             ## End if

mpisppy/cylinders/hub.py

@@ -30,6 +30,10 @@
 
 class Hub(SPCommunicator):
 
+    send_fields = (*SPCommunicator.send_fields, Field.SHUTDOWN, Field.BEST_OBJECTIVE_BOUNDS,)
+    receive_fields = (*SPCommunicator.receive_fields, )
+    optional_receive_fields = (*SPCommunicator.optional_receive_fields, Field.OBJECTIVE_INNER_BOUND, Field.OBJECTIVE_OUTER_BOUND, )
+
     _hub_algo_best_bound_provider = False
 
     def __init__(self, spbase_object, fullcomm, strata_comm, cylinder_comm, communicators, options=None):
@@ -85,7 +89,7 @@ def register_extension_recv_field(self, field: Field, strata_rank: int, buf_len:
         to the extension sync_with_spokes function.
         """
         key = self._make_key(field, strata_rank)
-        if key not in self._locals:
+        if key not in self.receive_buffers:
             # if it is not already registered, we need to update the local buffer
             self.extension_recv.add(key)
         ## End if
@@ -103,7 +107,7 @@ def register_extension_send_field(self, field: Field, buf_len: int) -> SendArray
         return self.register_send_field(field, buf_len)
 
     def is_send_field_registered(self, field: Field) -> bool:
-        return field in self._sends
+        return field in self.send_buffers
 
     def extension_send_field(self, field: Field, buf: SendArray):
         """
@@ -117,7 +121,7 @@ def sync_extension_fields(self):
         Update all registered extension fields. Safe to call even when there are no extension fields.
         """
         for key in self.extension_recv:
-            ext_buf = self._locals[key]
+            ext_buf = self.receive_buffers[key]
             (field, srank) = self._split_key(key)
             ext_buf._is_new = self.hub_from_spoke(ext_buf, srank, field)
         ## End for
@@ -233,7 +237,7 @@ def receive_innerbounds(self):
         logging.debug("Hub is trying to receive from InnerBounds")
         for idx in self.innerbound_spoke_indices:
             key = self._make_key(Field.OBJECTIVE_INNER_BOUND, idx)
-            recv_buf = self._locals[key]
+            recv_buf = self.receive_buffers[key]
             is_new = self.hub_from_spoke(recv_buf, idx, Field.OBJECTIVE_INNER_BOUND)
             if is_new:
                 bound = recv_buf[0]
@@ -249,7 +253,7 @@ def receive_outerbounds(self):
         logging.debug("Hub is trying to receive from OuterBounds")
         for idx in self.outerbound_spoke_indices:
             key = self._make_key(Field.OBJECTIVE_OUTER_BOUND, idx)
-            recv_buf = self._locals[key]
+            recv_buf = self.receive_buffers[key]
             is_new = self.hub_from_spoke(recv_buf, idx, Field.OBJECTIVE_OUTER_BOUND)
             if is_new:
                 bound = recv_buf[0]
@@ -320,18 +324,18 @@ def initialize_inner_bound_buffers(self):
     def _populate_boundsout_cache(self, buf):
         """ Populate a given buffer with the current bounds
         """
-        buf[-3] = self.BestOuterBound
-        buf[-2] = self.BestInnerBound
+        buf[0] = self.BestOuterBound
+        buf[1] = self.BestInnerBound
 
     def send_boundsout(self):
         """ Send bounds to the appropriate spokes
         This is called only for spokes which are bounds only.
         w and nonant spokes are passed bounds through the w and nonant buffers
         """
-        my_bounds = self.boundsout_send_buffer
+        my_bounds = self.send_buffers[Field.BEST_OBJECTIVE_BOUNDS]
         self._populate_boundsout_cache(my_bounds.array())
         logging.debug("hub is sending bounds={}".format(my_bounds))
-        self.hub_to_spoke(my_bounds, Field.OBJECTIVE_BOUNDS)
+        self.hub_to_spoke(my_bounds, Field.BEST_OBJECTIVE_BOUNDS)
         return
 
     def initialize_spoke_indices(self):
@@ -392,45 +396,7 @@ def initialize_spoke_indices(self):
 
 
     def register_send_fields(self):
-
-        self.shutdown = self.register_send_field(Field.SHUTDOWN, 1)
-
-        required_fields = set()
-        for i, spoke in enumerate(self.communicators):
-            if i == self.strata_rank:
-                continue
-            spoke_class = spoke["spcomm_class"]
-            if hasattr(spoke_class, "converger_spoke_types"):
-                for cst in spoke_class.converger_spoke_types:
-                    if cst == ConvergerSpokeType.W_GETTER:
-                        required_fields.add(Field.DUALS)
-                    elif cst == ConvergerSpokeType.NONANT_GETTER:
-                        required_fields.add(Field.NONANT)
-                    elif cst == ConvergerSpokeType.INNER_BOUND or cst == ConvergerSpokeType.OUTER_BOUND:
-                        required_fields.add(Field.OBJECTIVE_BOUNDS)
-                    else:
-                        pass # Intentional no-op
-                    ## End if
-                ## End for
-            else:
-                # Intentional no-op. Non-converger spokes need to register any needed
-                # fields separately. See the functions `register_extension_recv_field`
-                # and `register_extension_send_field`.
-                pass
-            ## End if
-        ## End for
-
-        n_nonants = 0
-        for s in self.opt.local_scenarios.values():
-            n_nonants += len(s._mpisppy_data.nonant_indices)
-        ## End for
-
-        if Field.DUALS in required_fields:
-            self.w_send_buffer = self.register_send_field(Field.DUALS, n_nonants)
-        if Field.NONANT in required_fields:
-            self.nonant_send_buffer = self.register_send_field(Field.NONANT, n_nonants)
-        if Field.OBJECTIVE_BOUNDS in required_fields:
-            self.boundsout_send_buffer = self.register_send_field(Field.OBJECTIVE_BOUNDS, 2)
+        super().register_send_fields()
 
         # Not all opt classes may have extensions
         if getattr(self.opt, "extensions", None) is not None:
@@ -439,7 +405,6 @@ def register_send_fields(self):
         return
 
 
-
     def hub_to_spoke(self, buf: SendArray, field: Field):
         """ Put the specified values into the specified locally-owned buffer
             for the spoke to pick up.
@@ -534,13 +499,17 @@ def send_terminate(self):
             buffer, so every spoke will see it simultaneously.
             processes (don't need to call them one at a time).
         """
-        shutdown = self.shutdown
-        shutdown[0] = 1.0
-        self.hub_to_spoke(shutdown, Field.SHUTDOWN)
+        self.send_buffers[Field.SHUTDOWN][0] = 1.0
+        self.hub_to_spoke(self.send_buffers[Field.SHUTDOWN], Field.SHUTDOWN)
         return
 
 
 class PHHub(Hub):
+
+    send_fields = (*Hub.send_fields, Field.NONANT, Field.DUALS)
+    receive_fields = (*Hub.receive_fields,)
+    optional_receive_fields = (*Hub.optional_receive_fields,)
+
     def setup_hub(self):
         """ Must be called after make_windows(), so that
             the hub knows the sizes of all the spokes windows
@@ -673,8 +642,7 @@ def send_nonants(self):
         """
         self.opt._save_nonants()
         ci = 0  ## index to self.nonant_send_buffer
-        # my_nonants = self._sends[Field.NONANT]
-        nonant_send_buffer = self.nonant_send_buffer
+        nonant_send_buffer = self.send_buffers[Field.NONANT]
         for k, s in self.opt.local_scenarios.items():
             for xvar in s._mpisppy_data.nonant_indices.values():
                 nonant_send_buffer[ci] = xvar._value
@@ -690,7 +658,7 @@ def send_ws(self):
         """ Send dual weights to the appropriate spokes
         """
         # NOTE: my_ws.array() and self.w_send_buffer should be the same array.
-        my_ws = self._sends[Field.DUALS]
+        my_ws = self.send_buffers[Field.DUALS]
         self.opt._populate_W_cache(my_ws.array(), padding=1)
         logging.debug("hub is sending Ws={}".format(my_ws.array()))
 
@@ -701,6 +669,10 @@ def send_ws(self):
 
 class LShapedHub(Hub):
 
+    send_fields = (*Hub.send_fields, Field.NONANT,)
+    receive_fields = (*Hub.receive_fields,)
+    optional_receive_fields = (*Hub.optional_receive_fields,)
+
     def setup_hub(self):
         """ Must be called after make_windows(), so that
             the hub knows the sizes of all the spokes windows
@@ -781,7 +753,7 @@ def send_nonants(self):
             TODO: Will likely fail with bundling
         """
         ci = 0  ## index to self.nonant_send_buffer
-        nonant_send_buffer = self.nonant_send_buffer
+        nonant_send_buffer = self.send_buffers[Field.NONANT]
         for k, s in self.opt.local_scenarios.items():
             nonant_to_root_var_map = s._mpisppy_model.subproblem_to_root_vars_map
             for xvar in s._mpisppy_data.nonant_indices.values():
@@ -797,6 +769,8 @@ def send_nonants(self):
 
 class SubgradientHub(PHHub):
 
+    # send / receive fields are same as PHHub
+
     _hub_algo_best_bound_provider = True
 
     def main(self):
@@ -806,6 +780,8 @@ def main(self):
 
 class APHHub(PHHub):
 
+    # send / receive fields are same as PHHub
+
     def main(self):
         """ SPComm gets attached by self.__init___; holding APH harmless """
         logger.critical("aph debug main in hub.py")

mpisppy/cylinders/reduced_costs_spoke.py

@@ -15,6 +15,10 @@
 
 class ReducedCostsSpoke(LagrangianOuterBound):
 
+    send_fields = (*LagrangianOuterBound.send_fields, Field.EXPECTED_REDUCED_COST, Field.SCENARIO_REDUCED_COST ,)
+    receive_fields = (*LagrangianOuterBound.receive_fields,)
+    optional_receive_fields = (*LagrangianOuterBound.optional_receive_fields,)
+
     converger_spoke_char = 'R'
 
     def __init__(self, *args, **kwargs):
@@ -54,28 +58,25 @@ def register_send_fields(self) -> None:
             scenario_buffer_len += len(s._mpisppy_data.nonant_indices)
         self._scenario_rc_buffer = np.zeros(scenario_buffer_len)
 
-        self.register_send_field(Field.EXPECTED_REDUCED_COST, self.nonant_length)
-        self.register_send_field(Field.SCENARIO_REDUCED_COST, scenario_buffer_len)
-
         return
 
     @property
     def rc_global(self):
-        return self._sends[Field.EXPECTED_REDUCED_COST].value_array()
+        return self.send_buffers[Field.EXPECTED_REDUCED_COST].value_array()
 
     @rc_global.setter
     def rc_global(self, vals):
-        arr = self._sends[Field.EXPECTED_REDUCED_COST].value_array()
+        arr = self.send_buffers[Field.EXPECTED_REDUCED_COST].value_array()
         arr[:] = vals
         return
 
     @property
     def rc_scenario(self):
-        return self._sends[Field.SCENARIO_REDUCED_COST].value_array()
+        return self.send_buffers[Field.SCENARIO_REDUCED_COST].value_array()
 
     @rc_scenario.setter
     def rc_scenario(self, vals):
-        arr = self._sends[Field.SCENARIO_REDUCED_COST].value_array()
+        arr = self.send_buffers[Field.SCENARIO_REDUCED_COST].value_array()
         arr[:] = vals
         return
 

mpisppy/cylinders/spcommunicator.py

@@ -23,7 +23,7 @@
 import abc
 import time
 
-from mpisppy.cylinders.spwindow import Field, SPWindow
+from mpisppy.cylinders.spwindow import Field, FieldLengths, SPWindow
 
 def communicator_array(size):
     arr = np.empty(size+1, dtype='d')
@@ -138,26 +138,26 @@ def __init__(self, spbase_object, fullcomm, strata_comm, cylinder_comm, communic
             self.options = options
 
         # Common fields for spokes and hubs
-        self._locals = dict()
-        self._sends = dict()
+        self.receive_buffers = dict()
+        self.send_buffers = dict()
 
         # setup FieldLengths which calculates
         # the length of each buffer type based
         # on the problem data
-        # self._field_lengths = FieldLengths(self.opt)
+        self._field_lengths = FieldLengths(self.opt)
 
         # attach the SPCommunicator to
         # the SPBase object
         self.opt.spcomm = self
 
-        # self.register_send_fields()
+        self.register_send_fields()
 
         return
 
     def _make_key(self, field: Field, origin: int):
         """
         Given a field and an origin (i.e. a strata_rank), generate a key for indexing
-        into the self._locals dictionary and getting the corresponding RecvArray.
+        into the self.receive_buffers dictionary and getting the corresponding RecvArray.
 
         Undone by `_split_key`. Currently, the key is simply a Tuple[field, origin].
         """
@@ -175,9 +175,8 @@ def _split_key(self, key) -> tuple[Field, int]:
     def _build_window_spec(self) -> dict[Field, int]:
         """ Build dict with fields and lengths needed for local MPI window
         """
-        self.register_send_fields()
         window_spec = dict()
-        for (field,buf) in self._sends.items():
+        for (field,buf) in self.send_buffers.items():
             window_spec[field] = np.size(buf.array())
         ## End for
         return window_spec
@@ -186,28 +185,28 @@ def register_recv_field(self, field: Field, origin: int, length: int = -1) -> Re
         key = self._make_key(field, origin)
         if length == -1:
             length = self._field_lengths[field]
-        if key in self._locals:
-            my_fa = self._locals[key]
+        if key in self.receive_buffers:
+            my_fa = self.receive_buffers[key]
             assert(length + 1 == np.size(my_fa.array()))
         else:
             my_fa = RecvArray(length)
-            self._locals[key] = my_fa
+            self.receive_buffers[key] = my_fa
         ## End if
         return my_fa
 
     def register_send_field(self, field: Field, length: int = -1) -> SendArray:
-        assert field not in self._sends, "Field {} is already registered".format(field)
+        assert field not in self.send_buffers, "Field {} is already registered".format(field)
         if length == -1:
             length = self._field_lengths[field]
-        # if field in self._sends:
-        #     my_fa = self._sends[field]
+        # if field in self.send_buffers:
+        #     my_fa = self.send_buffers[field]
         #     assert(length + 1 == np.size(my_fa.array()))
         # else:
         #     my_fa = SendArray(length)
-        #     self._sends[field] = my_fa
+        #     self.send_buffers[field] = my_fa
         # ## End if else
         my_fa = SendArray(length)
-        self._sends[field] = my_fa
+        self.send_buffers[field] = my_fa
         return my_fa
 
     @abc.abstractmethod
@@ -259,6 +258,7 @@ def make_windows(self) -> None:
 
         return
 
-    @abc.abstractmethod
     def register_send_fields(self) -> None:
-        pass
+        self.send_buffers = {}
+        for field in self.send_fields:
+            self.send_buffers[field] = self.register_send_field(field)