unifying register_receive_fields

Author: bknueven

mpisppy/cylinders/cross_scen_spoke.py

@@ -41,13 +41,16 @@ def register_send_fields(self) -> None:
         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.send_buffers[Field.CROSS_SCENARIO_CUT]
 
         return
 
+    def register_receive_fields(self):
+        super().register_receive_fields()
+        self.all_nonants = self.register_recv_field(Field.NONANT, 0)
+        self.all_etas = self.register_recv_field(Field.CROSS_SCENARIO_COST, 0)
+
     def prep_cs_cuts(self):
         # create a map scenario -> index, this index is used for various lists containing scenario dependent info.
         self.scenario_to_index = { scen : indx for indx, scen in enumerate(self.opt.all_scenario_names) }

mpisppy/cylinders/hub.py

@@ -35,7 +35,11 @@ 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()}")
         # for logging
         self.print_init = True
@@ -47,8 +51,6 @@ def __init__(self, spbase_object, fullcomm, strata_comm, cylinder_comm, communic
         self.stalled_iter_cnt = 0
         self.last_gap = float('inf')  # abs_gap tracker
 
-        self.extension_recv = set()
-
         self.initialize_bound_values()
 
         return

mpisppy/cylinders/spcommunicator.py

@@ -22,7 +22,6 @@
 import numpy as np
 import abc
 import time
-import itertools
 
 from mpisppy.cylinders.spwindow import Field, FieldLengths, SPWindow
 
@@ -191,25 +190,40 @@ def _build_window_spec(self) -> dict[Field, int]:
 
     def _exchange_send_fields(self) -> None:
         """ Do an all-to-all so we know what the other communicators are sending """
-        self.send_fields_by_rank = self.strata_comm.allgather(tuple(self.send_buffers.keys()))
+        send_buffers = tuple((k, buff._length) for k, buff in self.send_buffers.items())
+        self.send_fields_lengths_by_rank = self.strata_comm.allgather(send_buffers)
+
+        self.send_fields_by_rank = {}
 
         self.available_receive_fields = {}
-        for rank, fields in enumerate(self.send_fields_by_rank):
+        for rank, fields_lengths in enumerate(self.send_fields_lengths_by_rank):
             if rank == self.strata_rank:
                 continue
-            for f in fields:
+            self.send_fields_by_rank[rank] = []
+            for f, length in fields_lengths:
                 if f not in self.available_receive_fields:
                     self.available_receive_fields[f] = []
                 self.available_receive_fields[f].append(rank)
+                self.send_fields_by_rank[rank].append(f)
+
+        # print(f"{self.__class__.__name__}: {self.available_receive_fields=}")
 
     def register_recv_field(self, field: Field, origin: int, length: int = -1) -> RecvArray:
+        # print(f"{self.__class__.__name__}.register_recv_field, {field=}, {origin=}")
         key = self._make_key(field, origin)
         if length == -1:
             length = self._field_lengths[field]
         if key in self.receive_buffers:
             my_fa = self.receive_buffers[key]
             assert(length + 1 == np.size(my_fa.array()))
         else:
+            available_fields_from_origin = self.send_fields_lengths_by_rank[origin]
+            for _field, _length in available_fields_from_origin:
+                if field == _field:
+                    assert length == _length
+                    break
+            else: # couldn't find field!
+                raise RuntimeError(f"Couldn't find {field=} from {origin=}")
             my_fa = RecvArray(length)
             self.receive_buffers[key] = my_fa
         ## End if
@@ -284,6 +298,7 @@ def register_send_fields(self) -> None:
             self.register_send_field(field)
 
     def register_receive_fields(self) -> None:
+        # print(f"{self.__class__.__name__}: {self.receive_fields=}")
         for field in self.receive_fields:
             self.receive_field_spcomms[field] = []
             for strata_rank, comm in enumerate(self.communicators):

mpisppy/cylinders/spoke.py

@@ -29,10 +29,6 @@ def __init__(self, spbase_object, fullcomm, strata_comm, cylinder_comm, communic
 
         self.last_call_to_got_kill_signal = time.time()
 
-        # All spokes need the SHUTDOWN field to know when to terminate. Just
-        # register that here.
-        self.shutdown = self.register_recv_field(Field.SHUTDOWN, 0, 1)
-
         return
 
     def spoke_to_hub(self, buf: SendArray, field: Field):
@@ -95,7 +91,7 @@ def _spoke_from_hub(self,
     def _got_kill_signal(self):
         shutdown_buf = self.receive_buffers[self._make_key(Field.SHUTDOWN, 0)]
         if shutdown_buf.is_new():
-            shutdown = (self.shutdown[0] == 1.0)
+            shutdown = (shutdown_buf[0] == 1.0)
         else:
             shutdown = False
         ## End if
@@ -158,9 +154,12 @@ def __init__(self, spbase_object, fullcomm, strata_comm, cylinder_comm, options=
     def register_send_fields(self) -> None:
         super().register_send_fields()
         self._bound = self.send_buffers[self.bound_type()]
-        self._hub_bounds = self.register_recv_field(Field.BEST_OBJECTIVE_BOUNDS, 0, 2)
         return
 
+    def register_receive_fields(self) -> None:
+        super().register_receive_fields()
+        self._hub_bounds = self.register_recv_field(Field.BEST_OBJECTIVE_BOUNDS, 0, 2)
+
     @abc.abstractmethod
     def bound_type(self) -> Field:
         pass
@@ -203,19 +202,6 @@ def nonant_len_type(self) -> Field:
         # TODO: Make this a static method?
         pass
 
-    def register_send_fields(self) -> None:
-
-        super().register_send_fields()
-
-        vbuflen = 0
-        for s in self.opt.local_scenarios.values():
-            vbuflen += len(s._mpisppy_data.nonant_indices)
-        ## End for
-
-        self.register_recv_field(self.nonant_len_type(), 0, vbuflen)
-
-        return
-
 
 class InnerBoundSpoke(_BoundSpoke):
     """ For Spokes that provide an inner bound through self.bound to the