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..174b998b 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,49 @@ 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) - except Exception as e: - logger.error(f'Error getting result from future: {str(e)}') - self._semaphore.release() + queue.put(future.result(timeout=1.0)) + except BaseException as e: + # BaseException, matching _doc_publisher: an interrupt in a worker + # strands the consumer as surely as a ClientError does. + 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 +266,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 +285,44 @@ def _upload_batch(self, source_docs_batch:List[SourceDocument]): logger.debug(f'About to start polling queue [count: {count}, target_count: {target_count}]') + failures:List[BaseException] = [] + 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): + failures.append(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. Reaching here + # means the batch is short, so it is a failure rather than an end. + if not thread.is_alive(): + failures.append(RuntimeError( + f'Upload producer stopped before reporting its document count ' + f'[count: {count}, target_count: {target_count}]' + )) + break logger.debug(f'Waiting on queue to empty [count: {count}, target_count: {target_count}]') thread.join() + if failures: + # The first is raised rather than aggregated: 3.10 has no ExceptionGroup. + if len(failures) > 1: + logger.error( + f'{len(failures)} source documents failed to upload, raising the first ' + f'[batch_size: {len(source_docs_batch)}]' + ) + raise failures[0] + 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..d5b52b3c 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,8 @@ # SPDX-License-Identifier: Apache-2.0 import contextlib +import logging +import queue import threading import time @@ -17,6 +19,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 +1065,127 @@ 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() + + error = [] + + def consume(): + try: + consumed.extend(uploader._upload_batch(self._docs(1))) + except BaseException as e: + error.append(e) + 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' + + assert error, 'a truncated batch returned without raising' + assert 'before reporting its document count' in str(error[0]) + + def test_every_failure_is_reported_not_only_the_first(self, caplog): + """Two documents fail. The first is raised, and the log says how many.""" + 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[0] or doc is docs[1]: + raise RuntimeError('S3 write failed') + return doc + + with caplog.at_level(logging.ERROR): + finished, _, error = self._consume(uploader, docs, upload_doc) + + assert finished and error + assert any('2 source documents failed to upload' in r.message for r in caplog.records)