From d2554c2d62ce036b3e983f5015da1312768f08ef Mon Sep 17 00:00:00 2001 From: noel-improv Date: Fri, 4 Sep 2026 10:10:53 -0600 Subject: [PATCH 1/2] fix(lexical-graph): surface a failed document upload instead of hanging staging _upload_batch polls until it has seen one item per submitted document, but three except blocks could each leave that count unreachable, so a failed upload became a run that never returned. What hid it was _upload_doc swallowing its own S3 error and returning None: the count advanced, and the document was yielded and counted as staged. The two are a pair, so removing the swallow alone turns a silent success into a hang. Every submitted document now puts exactly one item on the queue, a failure marker when the upload raised. The publisher puts its count in a finally so a producer that dies still releases the consumer, the consumer also breaks out when its producer is gone, and the first failure is raised once the batch drains. A document that failed to write is no longer yielded. Confirmed against a real boto3 client: a ClientError from S3 propagates, nothing is reported as staged, and the call returns. --- .../indexing/load/s3_based_docs.py | 78 +++++++++---- .../unit/indexing/load/test_s3_based_docs.py | 104 ++++++++++++++++++ 2 files changed, 162 insertions(+), 20 deletions(-) diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/load/s3_based_docs.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/load/s3_based_docs.py index dbb273bd..f5ebb99c 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/load/s3_based_docs.py +++ b/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,38 +216,50 @@ 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)}') + raise def _task_complete_callback(self, future): self._semaphore.release() 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: @@ -240,15 +267,15 @@ def _doc_publisher(self, queue:queue.Queue, source_documents:List[SourceDocument 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 +286,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: diff --git a/lexical-graph/tests/unit/indexing/load/test_s3_based_docs.py b/lexical-graph/tests/unit/indexing/load/test_s3_based_docs.py index d695401a..0d48c5b5 100644 --- a/lexical-graph/tests/unit/indexing/load/test_s3_based_docs.py +++ b/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' From ce141cf2626c11c453e2fb4ab1e7ad461e934e64 Mon Sep 17 00:00:00 2001 From: noel-improv Date: Tue, 8 Sep 2026 14:57:56 -0600 Subject: [PATCH 2/2] fix(lexical-graph): remove the unused upload callback that never signalled the queue The class carried a second _task_complete_callback that released the semaphore without putting anything on the queue. Nothing has ever wired it up, and connecting it would reproduce the staging hang this change fixes, on every document rather than only on a failed upload. --- .../lexical_graph/indexing/load/s3_based_docs.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/load/s3_based_docs.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/load/s3_based_docs.py index f5ebb99c..878fe06a 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/load/s3_based_docs.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/load/s3_based_docs.py @@ -218,9 +218,6 @@ def _upload_doc(self, root_path:str, doc:SourceDocument, s3_client): logger.error(f'Error while writing source document to S3: {str(e)}') raise - def _task_complete_callback(self, future): - self._semaphore.release() - def _get_callback_fn(self, queue:queue.Queue): def _task_complete_callback(future): # The consumer counts one item per submitted document, so a failure