[FIX] Surface a failed document upload instead of hanging S3 staging

lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/load/s3_based_docs.py

@@ -31,12 +31,27 @@
 QUEUE_SIZE = 1000
 BATCH_SIZE = 100
 
+# How often the staging consumer wakes to check whether its producer is still
+# alive. Uploads are slow, so this paces the liveness check, not the work.
+QUEUE_POLL_SECONDS = 1.0
+
 logger = logging.getLogger(__name__)
 
 # Joins node ids before hashing them. Without a separator ['ab', 'c'] and
 # ['a', 'bc'] hash alike.
 _DOC_SUFFIX_DELIMITER = '\x00'
 
+class _UploadFailed:
+    """
+    Stands in for a document that did not upload.
+
+    The consumer counts one item per submitted document, so a failure has to
+    occupy a slot rather than being dropped.
+    """
+
+    def __init__(self, cause:BaseException):
+        self.cause = cause
+
 class ConfiguredThreadCount:
     """
     Sizes a thread pool from a caller-supplied count, falling back to the config
@@ -201,54 +216,63 @@ def _upload_doc(self, root_path:str, doc:SourceDocument, s3_client):
 
         except Exception as e:
             logger.error(f'Error while writing source document to S3: {str(e)}')
-
-    def _task_complete_callback(self, future):
-        self._semaphore.release()
+            raise
 
     def _get_callback_fn(self, queue:queue.Queue):
         def _task_complete_callback(future):
+            # The consumer counts one item per submitted document, so a failure
+            # has to put something too. Putting nothing leaves it polling for a
+            # count it can never reach.
             try:
-                doc = future.result(timeout=1.0)
-                queue.put(doc)
+                queue.put(future.result(timeout=1.0))
             except Exception as e:
-                logger.error(f'Error getting result from future: {str(e)}')
-            self._semaphore.release()
+                logger.error(f'Error uploading source document: {str(e)}')
+                queue.put(_UploadFailed(e))
+            finally:
+                self._semaphore.release()
         return _task_complete_callback
 
-    def _submit_proxy(self, function, executor, queue:queue.Queue, *args, **kwargs):
+    def _submit_proxy(self, function, executor, queue:queue.Queue, *args, **kwargs) -> bool:
+        """Submit one upload. Returns whether the consumer should expect an item."""
+        self._semaphore.acquire()
         try:
-            self._semaphore.acquire()
             future = executor.submit(function, *args, **kwargs)
-            future.add_done_callback(self._get_callback_fn(queue))
         except Exception as e:
             logger.exception(f'Error in submit proxy: {str(e)}')
+            self._semaphore.release()
+            return False
+
+        future.add_done_callback(self._get_callback_fn(queue))
+        return True
 
     def _doc_publisher(self, queue:queue.Queue, source_documents:List[SourceDocument]=[]):
 
         s3_client = GraphRAGConfig.s3
-
+
+        count = 0
+
+        # The count goes on the queue whatever happens above, including a
+        # KeyboardInterrupt, or the consumer waits on a producer that has gone.
         try:
 
             with concurrent.futures.ThreadPoolExecutor(max_workers=self._num_threads()) as executor:
 
-                count = 0
-
                 for source_document in source_documents:
 
                     if not source_document.nodes:
                         continue
 
                     root_path = join(self.collection_prefix, source_document.source_id())
 
-                    self._submit_proxy(self._upload_doc, executor, queue, root_path, source_document, s3_client)
-
-                    count += 1
-
-            self._queue.put(count)
+                    if self._submit_proxy(self._upload_doc, executor, queue, root_path, source_document, s3_client):
+                        count += 1
 
-        except Exception as e:
+        except BaseException as e:
             logger.exception(f'Error in doc publisher: {str(e)}')
 
+        finally:
+            self._queue.put(count)
+
     def _upload_batch(self, source_docs_batch:List[SourceDocument]):
 
         thread = threading.Thread(target=self._doc_publisher, daemon=True, kwargs={'source_documents': source_docs_batch, 'queue': self._queue})
@@ -259,22 +283,33 @@ def _upload_batch(self, source_docs_batch:List[SourceDocument]):
 
         logger.debug(f'About to start polling queue [count: {count}, target_count: {target_count}]')
 
+        failure = None
+
         while target_count is None or count < target_count:
             try:
-                item = self._queue.get(timeout=60.0)
+                item = self._queue.get(timeout=QUEUE_POLL_SECONDS)
                 if isinstance(item, int):
                     target_count = item
                 else:
                     count += 1
-                    yield item
+                    if isinstance(item, _UploadFailed):
+                        failure = failure or item.cause
+                    else:
+                        yield item
                 self._queue.task_done()
-            except queue.Empty as e:
-                continue
+            except queue.Empty:
+                # A producer that died before putting its count would otherwise
+                # keep the consumer here for the rest of the run.
+                if not thread.is_alive():
+                    break
 
         logger.debug(f'Waiting on queue to empty [count: {count}, target_count: {target_count}]')
 
         thread.join()
 
+        if failure is not None:
+            raise failure
+
     def upload(self, source_documents: List[SourceDocument]):
 
         if not self._queue:

lexical-graph/tests/unit/indexing/load/test_s3_based_docs.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import contextlib
+import queue
 import threading
 import time
 
@@ -17,6 +18,9 @@
 )
 from graphrag_toolkit.lexical_graph.indexing.model import SourceDocument
 from graphrag_toolkit.lexical_graph.storage.constants import INDEX_KEY
+from threading import Semaphore
+
+S3_BASED_DOCS = 'graphrag_toolkit.lexical_graph.indexing.load.s3_based_docs'
 
 
 class TestS3BasedDocsInitialization:
@@ -1060,3 +1064,103 @@ def test_distinct_objects_are_all_kept(self):
         })
 
         assert len(doc.nodes) == 2
+
+
+class TestStagingSurfacesUploadFailures:
+    """
+    Every submitted document has to put exactly one item on the queue. Without
+    that, _upload_batch polls for a count it can never reach, and a failed
+    upload becomes a hung run rather than an error.
+    """
+
+    def _docs(self, n):
+        docs = []
+        for i in range(n):
+            node = TextNode(text='chunk text', id_=f'c{i}')
+            node.relationships[NodeRelationship.SOURCE] = RelatedNodeInfo(
+                node_id=f'aws::src{i}:d41d'
+            )
+            docs.append(SourceDocument(nodes=[node]))
+        return docs
+
+    def _consume(self, uploader, docs, upload_doc, timeout=8):
+        """Run upload() on a thread so a hang shows up as a timeout, not a stall."""
+        yielded, finished, error = [], threading.Event(), []
+
+        def consume():
+            try:
+                with patch.object(S3DocUploader, '_upload_doc', side_effect=upload_doc), \
+                     patch(f'{S3_BASED_DOCS}.GraphRAGConfig'):
+                    yielded.extend(uploader.upload(list(docs)))
+            except Exception as e:
+                error.append(e)
+            finished.set()
+
+        threading.Thread(target=consume, daemon=True).start()
+        return finished.wait(timeout=timeout), yielded, error
+
+    def test_a_failed_upload_does_not_hang_the_run(self):
+        docs = self._docs(3)
+        uploader = S3DocUploader(bucket_name='b', collection_prefix='p', num_threads=2)
+
+        def upload_doc(root_path, doc, s3_client):
+            if doc is docs[1]:
+                raise RuntimeError('S3 write failed')
+            return doc
+
+        finished, _, _ = self._consume(uploader, docs, upload_doc)
+
+        assert finished, 'upload() never returned after a document failed'
+
+    def test_a_failed_upload_raises(self):
+        docs = self._docs(3)
+        uploader = S3DocUploader(bucket_name='b', collection_prefix='p', num_threads=2)
+
+        def upload_doc(root_path, doc, s3_client):
+            if doc is docs[1]:
+                raise RuntimeError('S3 write failed')
+            return doc
+
+        _, _, error = self._consume(uploader, docs, upload_doc)
+
+        assert error, 'a failed upload was not surfaced to the caller'
+
+    def test_a_failed_document_is_not_yielded_as_staged(self):
+        docs = self._docs(3)
+        uploader = S3DocUploader(bucket_name='b', collection_prefix='p', num_threads=2)
+
+        def upload_doc(root_path, doc, s3_client):
+            if doc is docs[1]:
+                raise RuntimeError('S3 write failed')
+            return doc
+
+        _, yielded, _ = self._consume(uploader, docs, upload_doc)
+
+        assert docs[1] not in yielded
+        assert None not in yielded
+
+    def test_every_document_is_yielded_when_all_succeed(self):
+        docs = self._docs(3)
+        uploader = S3DocUploader(bucket_name='b', collection_prefix='p', num_threads=2)
+
+        finished, yielded, error = self._consume(
+            uploader, docs, lambda root_path, doc, s3_client: doc
+        )
+
+        assert finished and not error
+        assert len(yielded) == 3
+
+    def test_a_dead_producer_does_not_hang_the_consumer(self):
+        uploader = S3DocUploader(bucket_name='b', collection_prefix='p', num_threads=2)
+        uploader._semaphore = Semaphore(2)
+        uploader._queue = queue.Queue()
+        consumed, finished = [], threading.Event()
+
+        def consume():
+            consumed.extend(uploader._upload_batch(self._docs(1)))
+            finished.set()
+
+        with patch.object(S3DocUploader, '_doc_publisher', side_effect=RuntimeError('dead')):
+            threading.Thread(target=consume, daemon=True).start()
+
+            assert finished.wait(timeout=8), 'consumer waited on a producer that had died'