awslabs / awslabs/graphrag-toolkit

[BUG] S3BasedDocs staging hangs, or reports a failed upload as written

Closed
#519 0 comments 0 reactions 0 assignees View on GitHub
bug triage_needed
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.