awslabs / awslabs/graphrag-toolkit
[BUG] S3BasedDocs staging hangs, or reports a failed upload as written
- Dominant language
- Python
- Stars
- 442
- Forks
- 106
- Avg merge
- 2d 17h
- Merged PRs (30d)
- 52
Description
### Package version
3.19.1
### Package
lexical-graph
### Python version
3.12.13
### Operating System
macOS
### Description
A document that fails to upload during S3 staging either hangs the run or is reported as staged. Neither raises.
`_upload_batch` polls until `count == target_count` and treats `queue.Empty` as "keep waiting", without checking whether the producer thread is alive. Three `except Exception: log` blocks can put that count out of reach:
- `_get_callback_fn` puts nothing on the queue when `future.result()` raises
- `_submit_proxy` returns without releasing the semaphore, while `_doc_publisher` still runs `count += 1` for a document it never submitted
- `_doc_publisher` keeps `queue.put(count)` inside the `try`, so a failure above it skips the put
What keeps this from firing today is `_upload_doc` swallowing its own S3 exception and returning `None`. The callback puts the `None`, the count advances, and `accept()` yields it and counts it in the "Finished writing N source documents" total. So the failed document is reported as written, and a retry has no idea it needs to be redone.
That makes the two a pair: stop swallowing in `_upload_doc` on its own and the silent success becomes a hang.
Expected: an upload failure raises, and a document that failed to write is not counted as staged.
This is not #418. Same symptom and same path, but that one was a forked worker inheriting a held `BufferedWriter` lock, and it is fixed. This came in with `4b8743c0`.
### Steps to reproduce
```python
# 1. Producer dies -> the consumer never returns.
with patch.object(S3DocUploader, '_doc_publisher', side_effect=RuntimeError('boom')):
list(uploader.upload([doc])) # blocks indefinitely, no exception
# 2. The coupling. Against a working S3 stub, patch _upload_doc two ways:
# returns None (today) -> finished=True yielded=[None] <- failure reported as success
# raises -> finished=False yielded=[] <- hangs
```
### Other information
A fix would be one invariant: every submitted document puts exactly one item on the queue.
- `_get_callback_fn` always puts, a failure marker on exception
- `_submit_proxy` releases the semaphore in a `finally` and reports whether it submitted
- `_doc_publisher` puts the count in a `finally`, catching `BaseException`
- `_upload_batch` breaks on `queue.Empty` when the producer is dead, and re-raises after `join()`
- `_upload_doc` stops swallowing
- `accept()` excludes failures from its staged count
No test fails a document mid-batch, which is the gap that hid this.
Contributor guide
Research direction
Start by running the two S3DocUploader reproductions in the issue, then trace _upload_batch, _get_callback_fn, _submit_proxy, _doc_publisher, _upload_doc, and accept. Verify producer failures do not leave consumers waiting, upload failures raise, and failed documents are excluded from the staged count; add coverage for a mid-batch failure.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- aws, python
- Domain
- backend, cloud
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100