extensions have to do nearly everything anyways; might as well be explicit about it

Author: bknueven

mpisppy/cylinders/hub.py

@@ -11,7 +11,7 @@
 import logging
 import mpisppy.log
 
-from mpisppy.cylinders.spcommunicator import RecvArray, SendArray, SPCommunicator
+from mpisppy.cylinders.spcommunicator import RecvArray, SPCommunicator
 from math import inf
 
 from mpisppy import global_toc
@@ -32,9 +32,6 @@ class Hub(SPCommunicator):
     _hub_algo_best_bound_provider = False
 
     def __init__(self, spbase_object, fullcomm, strata_comm, cylinder_comm, communicators, options=None):
-        # The extensions will be registered in SPCommunicator.__init__
-        self.extension_recv = set()
-
         super().__init__(spbase_object, fullcomm, strata_comm, cylinder_comm, communicators, options=options)
 
         logger.debug(f"Built the hub object on global rank {fullcomm.Get_rank()}")
@@ -80,51 +77,6 @@ def current_iteration(self):
     def main(self):
         pass
 
-    def register_extension_recv_field(self, field: Field, strata_rank: int, buf_len: int = -1) -> RecvArray:
-        """
-        Register an extensions interest in the given field from the given spoke. The hub
-        is then responsible for updating this field into a local buffer prior to the call
-        to the extension sync_with_spokes function.
-        """
-        key = self._make_key(field, strata_rank)
-        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
-        ra = self.register_recv_field(field, strata_rank, buf_len)
-        return ra
-
-    def register_extension_send_field(self, field: Field, buf_len: int) -> SendArray:
-        """
-        Register a field with the hub that an extension will be making available to spokes. Returns a
-        buffer that is usable for sending the desired values. The extension is responsible for calling
-        the hub publish_extension_field when ready to send the values. Returns a SendArray to use
-        to publish values to spokes. Meant to be called within the extension function
-        `register_send_fields`.
-        """
-        return self.register_send_field(field, buf_len)
-
-    def is_send_field_registered(self, field: Field) -> bool:
-        return field in self.send_buffers
-
-    def extension_send_field(self, field: Field, buf: SendArray):
-        """
-        Send the data in the SendArray `buf` which stores the Field `field`. This will make
-        the data available to the spokes in this strata.
-        """
-        return self.put_send_buffer(buf, field)
-
-    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.receive_buffers[key]
-            (field, srank) = self._split_key(key)
-            ext_buf._is_new = self.get_receive_buffer(ext_buf, field, srank)
-        ## End for
-        return
-
     def clear_latest_chars(self):
         self.latest_ib_char = None
         self.latest_ob_char = None
@@ -371,7 +323,6 @@ def sync(self):
         self.receive_outerbounds()
         self.receive_innerbounds()
         if self.opt.extensions is not None:
-            self.sync_extension_fields()
             self.opt.extobject.sync_with_spokes()
 
     def sync_with_spokes(self):
@@ -384,7 +335,6 @@ def sync_bounds(self):
 
     def sync_extensions(self):
         if self.opt.extensions is not None:
-            self.sync_extension_fields()
             self.opt.extobject.sync_with_spokes()
 
     def sync_nonants(self):
@@ -485,7 +435,6 @@ def sync(self, send_nonants=True):
         self.receive_innerbounds()
         # in case LShaped ever gets extensions
         if getattr(self.opt, "extensions", None) is not None:
-            self.sync_extension_fields()
             self.opt.extobject.sync_with_spokes()
 
     def is_converged(self):

mpisppy/cylinders/spcommunicator.py

@@ -287,6 +287,9 @@ def _make_windows(self) -> None:
 
         return
 
+    def is_send_field_registered(self, field: Field) -> bool:
+        return field in self.send_buffers
+
     def register_send_fields(self) -> None:
         for field in self.send_fields:
             self.register_send_field(field)

mpisppy/extensions/coeff_rho.py

@@ -7,6 +7,7 @@
 # full copyright and license information.
 ###############################################################################
 
+from mpisppy import global_toc
 import mpisppy.extensions.extension
 
 from mpisppy.utils.sputils import nonant_cost_coeffs
@@ -36,5 +37,4 @@ def post_iter0(self):
                 #     nv = s._mpisppy_data.nonant_indices[ndn_i] # var_data object
                 #     print(ndn_i,nv.getname(),cc[ndn_i],rho._value)
 
-        if self.ph.cylinder_rank == 0:
-            print("Rho values updated by CoeffRho Extension")
+        global_toc("Rho values updated by CoeffRho Extension", self.ph.cylinder_rank == 0)

mpisppy/extensions/cross_scen_extension.py

@@ -126,6 +126,11 @@ def _check_bound(self):
             cached_ph_obj[k].activate()
 
     def get_from_cross_cuts(self):
+        self.opt.spcomm.get_receive_buffer(
+            self.cuts,
+            Field.CROSS_SCENARIO_CUT,
+            self.cross_scenario_index,
+        )
         if self.cuts.is_new():
             self.make_cuts(self.cuts.array())
 
@@ -144,7 +149,7 @@ def send_to_cross_cuts(self):
                 ## End for
             ## End for
 
-            self.opt.spcomm.extension_send_field(Field.NONANT, all_nonants)
+            self.opt.spcomm.put_send_buffer(all_nonants, Field.NONANT)
 
         ## End if
 
@@ -156,7 +161,7 @@ def send_to_cross_cuts(self):
                 all_etas[ci] = s._mpisppy_model.eta[sn]._value
                 ci += 1
 
-        self.opt.spcomm.extension_send_field(Field.CROSS_SCENARIO_COST, all_etas)
+        self.opt.spcomm.put_send_buffer(all_etas, Field.CROSS_SCENARIO_COST)
 
         return
 
@@ -255,13 +260,13 @@ def register_send_fields(self):
         if spcomm.is_send_field_registered(Field.NONANT):
             self.send_nonants = False
         else:
-            self.all_nonants = spcomm.register_extension_send_field(
+            self.all_nonants = spcomm.register_send_field(
                 Field.NONANT,
                 local_scen_count * self.opt.nonant_length
             )
             self.send_nonants = True
         ## End if-else
-        self.all_etas = spcomm.register_extension_send_field(
+        self.all_etas = spcomm.register_send_field(
             Field.CROSS_SCENARIO_COST,
             nscen * nscen,
         )
@@ -282,11 +287,11 @@ def register_receive_fields(self):
         spcomm = self.opt.spcomm
         cross_scenario_cut_ranks = spcomm.fields_to_ranks[Field.CROSS_SCENARIO_CUT]
         assert len(cross_scenario_cut_ranks) == 1
-        index = cross_scenario_cut_ranks[0]
+        self.cross_scenario_index = cross_scenario_cut_ranks[0]
 
-        self.cuts = spcomm.register_extension_recv_field(
+        self.cuts = spcomm.register_recv_field(
             Field.CROSS_SCENARIO_CUT,
-            index,
+            self.cross_scenario_index,
         )
 
     def sync_with_spokes(self):

mpisppy/extensions/extension.py

@@ -34,8 +34,7 @@ def setup_hub(self):
     def register_send_fields(self):
         '''
         Method called by the Hub SPCommunicator to get any fields that the extension
-        will make available to spokes. Use hub function `register_extension_send_field`
-        to register a field.
+        will make available to spokes.
         '''
         return
 

mpisppy/extensions/reduced_costs_fixer.py

@@ -101,11 +101,11 @@ def register_receive_fields(self):
 
         self.reduced_costs_spoke_index = index
 
-        self.reduced_cost_buf = spcomm.register_extension_recv_field(
+        self.reduced_cost_buf = spcomm.register_recv_field(
             Field.EXPECTED_REDUCED_COST,
             self.reduced_costs_spoke_index,
         )
-        self.outer_bound_buf = spcomm.register_extension_recv_field(
+        self.outer_bound_buf = spcomm.register_recv_field(
             Field.OBJECTIVE_OUTER_BOUND,
             self.reduced_costs_spoke_index,
         )
@@ -116,6 +116,16 @@ def register_receive_fields(self):
     def sync_with_spokes(self, pre_iter0 = False):
         # TODO: if we calculate the new bounds in the spoke we don't need to check if the buffers
         #       have the same ID
+        self.opt.spcomm.get_receive_buffer(
+            self.reduced_cost_buf,
+            Field.EXPECTED_REDUCED_COST,
+            self.reduced_costs_spoke_index,
+        )
+        self.opt.spcomm.get_receive_buffer(
+            self.outer_bound_buf,
+            Field.OBJECTIVE_OUTER_BOUND,
+            self.reduced_costs_spoke_index,
+        )
         if self.reduced_cost_buf.is_new() and self.reduced_cost_buf.id() == self.outer_bound_buf.id():
             reduced_costs = self.reduced_cost_buf.value_array()
             this_outer_bound = self.outer_bound_buf.value_array()[0]

mpisppy/extensions/reduced_costs_rho.py

@@ -48,12 +48,17 @@ def register_receive_fields(self):
         assert len(reduced_cost_ranks) == 1
         self.reduced_costs_spoke_index = reduced_cost_ranks[0]
 
-        self.scenario_reduced_cost_buf = spcomm.register_extension_recv_field(
+        self.scenario_reduced_cost_buf = spcomm.register_recv_field(
             Field.SCENARIO_REDUCED_COST,
             self.reduced_costs_spoke_index,
         )
 
     def sync_with_spokes(self):
+        self.opt.spcomm.get_receive_buffer(
+            self.scenario_reduced_cost_buf,
+            Field.SCENARIO_REDUCED_COST,
+            self.reduced_costs_spoke_index,
+        )
         if self.scenario_reduced_cost_buf.is_new():
             self._scenario_rc_buffer[:] = self.scenario_reduced_cost_buf.value_array()
             # print(f"In ReducedCostsRho; {self._scenario_rc_buffer=}")

mpisppy/extensions/sensi_rho.py

@@ -54,8 +54,7 @@ def compute_and_update_rho(self):
                 #     nv = s._mpisppy_data.nonant_indices[ndn_i]  # var_data object
                 #     print(f"{s.name=}, {nv.name=}, {rho.value=}")
 
-        if ph.cylinder_rank == 0:
-            print(f"Rho values updated by {self.__class__.__name__} Extension")
+        global_toc(f"Rho values updated by {self.__class__.__name__} Extension", ph.cylinder_rank == 0)
 
     def miditer(self):
         self.update_caches()