Refactor stacking for reuse and overlappability

Currently, the stacking feature adds a lot of code to various parts of ShmemTrainMatcher. It would be nice if this could be cleaner, for:

• more maintainable ShmemTrainMatcher
• reusability: the DetectorAssembler also does stacking except in its own less good way
• reusability: we will at some point probably split ShmemTrainMatcher and some of the resulting matching-and-processing devices will want to stack
• improvability: the current implementation is pretty naive and though the handle_source is intended to be run in a thread pool, only one train can be processed at a time; this requires rethinking of buffer handling to change
  • bonus concern: the stacking buffers, if we start overlapping processing and decoupling this from sending, would not satisfy the safeNDArray flag in the current version

Plan for this MR:

• Move stacking functionality to "friend" class
• Get rid of global stacking buffers; idea: create, per train, a stacking context with fresh buffers
• Add some testing of stacking (not just utils, but the friend class working on tables)
• (Out of scope) Test that it works in a overlapping trains setting
• (Out of scope) Start overlapping train processing
• Refactor DetectorAssembler to use StackingFriend to stack before assembly
• Change the configuration that the manager passes to DetectorAssembler correspondingly
• Nevermind, DetectorAssembler is not a great fit actually; maybe use xarray and exploit extra-geom's partial data assembly functionality (Out of scope; see !74 (closed))
• Test everything with actual Karabo (pending node allocation)

src/calng/ShmemTrainMatcher.py

@@ -69,7 +69,9 @@ class ShmemTrainMatcher(TrainMatcher.TrainMatcher):
     def initialization(self):
         super().initialization()
         self._shmem_handler = shmem_utils.ShmemCircularBufferReceiver()
-        self._stacking_friend = StackingFriend(self.get("merge"), self.get("sources"))
+        self._stacking_friend = StackingFriend(
+            self, self.get("merge"), self.get("sources")
+        )
         self._frameselection_friend = FrameselectionFriend(self.get("frameSelector"))
         self._thread_pool = concurrent.futures.ThreadPoolExecutor(
             max_workers=self.get("processingThreads")
@@ -102,22 +104,21 @@ class ShmemTrainMatcher(TrainMatcher.TrainMatcher):
     def on_matched_data(self, train_id, sources):
         frame_selection_mask = self._frameselection_friend.get_mask(sources)
         # note: should not do stacking and frame selection for now!
-        self._stacking_friend.prepare_stacking_for_train(sources)
-
-        concurrent.futures.wait(
-            [
-                self._thread_pool.submit(
-                    self._handle_source,
-                    source,
-                    data,
-                    timestamp,
-                    new_sources_map,
-                    frame_selection_mask,
-                )
-                for source, (data, timestamp) in sources.items()
-            ]
-        )
-        sources.update(new_sources_map)
+        with self._stacking_friend.stacking_context as stacker:
+            concurrent.futures.wait(
+                [
+                    self._thread_pool.submit(
+                        self._handle_source,
+                        source,
+                        data,
+                        timestamp,
+                        stacker,
+                        frame_selection_mask,
+                    )
+                    for source, (data, timestamp) in sources.items()
+                ]
+            )
+        sources.update(stacker.new_source_map)
 
         # karabo output
         if self.output is not None:
@@ -141,10 +142,9 @@ class ShmemTrainMatcher(TrainMatcher.TrainMatcher):
         source,
         data_hash,
         timestamp,
-        new_sources_map,
+        stacker,
         frame_selection_mask,
-        ignore_stacking,
     ):
         self._shmem_handler.dereference_shmem_handles(data_hash)
         self._frameselection_friend.apply_mask(source, data_hash, frame_selection_mask)
-        self._stacking_friend.handle_source(...)
+        stacker.process(source, data_hash)