From f8bb8bdf2233ada35d77431211262838745b1633 Mon Sep 17 00:00:00 2001
From: BergurDavidsen <bergurdavidsen@gmail.com>
Date: Thu, 9 Apr 2026 18:49:36 +0200
Subject: [PATCH] Fix retrieve_results to drain all output queues before
 exiting

Track termination per queue and wait for all output queues to finish before returning, instead of exiting on the first None. This ensures complete result draining
---
 pipeline/pipeline.py | 32 ++++++++++++--------------------
 1 file changed, 12 insertions(+), 20 deletions(-)

diff --git a/pipeline/pipeline.py b/pipeline/pipeline.py
index 1e8ee0b..0970186 100644
--- a/pipeline/pipeline.py
+++ b/pipeline/pipeline.py
@@ -122,28 +122,30 @@ def retrieve_results(self, event: Event) -> None:
         """
         This method continuously checks all output queues for new queries. It performs
         non-blocking retrieval from the queues and processes the queries if available.
-        If a terminating character (None) is received, the method returns and stops
-        further processing. Additionally, it logs the end of the pipeline execution
+        If a terminating character (None) is received, the method increments the finished
+        outputs counter and keeps processing if there are any remaining outputs. Additionally,
+        it logs the end of the pipeline execution
         and sets the provided event to signal completion.
 
         Args:
             event (Event): An event object used to signal the completion of the pipeline.
         """
 
-        while True:
+        finished_outputs = 0
+        total_outputs = len(self._output_queues)
+
+        while finished_outputs < total_outputs:
             for output_queue in self._output_queues:
-                # non-blocking retrieval from the queues
-                # if the queue is empty, simply move on
                 try:
-                    new_query: Query | None = output_queue.get(timeout=0.1)
+                    new_query = output_queue.get(timeout=0.1)
                 except Empty:
                     continue
 
-                # if terminating character is received, return
-                if not new_query:
-                    return
+                if new_query is None:
+                    finished_outputs += 1
+                    print(f"[Pipeline] Output queue finished ({finished_outputs}/{total_outputs})")
+                    continue
 
-                # log the end of pipeline execution
                 self._logger.info(
                     "%s, pipeline - %s, run, end, %d, %.6f, %d, %d",
                     self.name,
@@ -185,16 +187,6 @@ def run(self, query_queue: Queue, event: Event) -> None:
 
             # check if done
             if query is None:
-                # only terminate once all of the queries have been processed in order to protect the graph from circular dependencies during termination
-                while queries_sent > self.queries_processed:
-                    # print(
-                    #     "Waiting for all queries to be processed, sent: %d, processed: %d",
-                    #     queries_sent,
-                    #     self.queries_processed,
-                    # )
-                    # wait for a query to be processed
-                    event.wait(0.1)
-
                 # if so, send the termination element to the following stages and join the threads
                 for input_queue in self._input_queues:
                     input_queue.put(None)