[FEATURE] Read a half-written source document as incomplete by noel-improv · Pull Request #560 · awslabs/graphrag-toolkit

noel-improv · 2026-09-16T20:11:02Z

Description

#532 writes a completion marker once every chunk for a document has been stored, but nothing reads it. A run killed part way through leaves a source document prefix holding some of its chunks and no marker, and the reader merges whatever it finds into one document. The build sees a short document rather than a missing one, and a restart has no way to tell the two apart.

This makes the reader skip a prefix its markers do not account for, so an interrupted document is re-extracted instead of being built from half its chunks.

The hard part is what to do about everything already staged. A strict reader would call all of it incomplete and re-stage the lot, which means paying for extraction again. Backfilling markers for prefixes that look whole does not work either: a truncated document and a short document are the same bytes, which is why the marker exists. So the collection answers the question. Staging writes a record under the collection before any document, and a reader that finds the record holds every prefix under it to its markers. A collection staged before markers existed carries no record and is read as it stands.

A date cutoff would do the same job with two weaknesses: someone has to configure the date, and a collection staged by older code after that date still reads wrong. The record has neither.

Changes

collection_record_key / writes_completion_markers: _staging.json under the collection, written before any document. A missing record answers no; a denied read or a throttle raises rather than quietly dropping the check.
is_complete: a prefix is whole when its markers cover exactly the nodes found under it. Chunks no marker covers mean a truncated document; a marker covering objects that are gone means the prefix lost something.
S3ChunkDownloader and S3DocDownloader skip a prefix that fails the check, logging the prefix at warning level.
S3DocUploader writes a marker for each object it stores. Markers started on the per-chunk path, so without this a recorded JSONL collection would read back as incomplete for good.
_write_completion_marker moved onto EncryptedPut, which both uploaders already share.

Problem

Nothing distinguishes a document whose chunks all landed from one whose run was killed half way through. Both read back as a document, one of them short, and the build indexes the short one.

Related issue (if any): #532

Testing

Unit tests added/updated
Integration tests added (as appropriate)
Existing tests pass (pytest)
Tested manually (describe below)

Unit tests cover a prefix with no marker, a marker covering a chunk that is gone, several markers accounting for one prefix between them, a collection with no record read as it stands, the record written before any document, a failed record stopping staging, and the JSONL uploader marking what it wrote. Full unit suite: 2,299 passed.

tests/integration/indexing/load/test_s3_staging_recovery_live.py stages a collection against real S3, damages it the way an interrupted run would, and reads it back. Twelve cases across both storage formats: the marker deleted, an object deleted from under a marker, re-staging repairing the damage, an undamaged collection reading whole, and a collection whose record was removed reading everything. All twelve pass against S3 in us-east-1. Objects are written under a unique prefix per run and deleted afterwards.

The completeness check was then forced to always pass, and the four bad-state cases failed while the other eight held.

Checklist

Code follows existing style and conventions
License headers present on new files
Documentation updated (if applicable)
No breaking changes (or clearly documented)

Two behaviour notes for the release:

An incomplete prefix is skipped, not raised on. The document is absent from the build and is re-extracted on the next run. Raising instead would stop a build over one interrupted document.
A collection staged before this change is read exactly as it is today. The check starts applying to collections staged from here on.

Not in this PR

Reporting which documents were skipped back to the caller. The warning log names each prefix; a restart that acts on the list needs more than that.
A failed marker write is still logged rather than raised, as it was in [FEATURE] Write a completion marker when every chunk for a document stored #532. It costs a re-stage, where raising would lose a stream the chunks themselves survived.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

A run killed part way through leaves a source document prefix holding some of its chunks and no completion marker. The reader merged whatever it found into one document, so the build saw a short document rather than a missing one, and a restart had no way to tell the two apart. The collection decides whether its prefixes must carry a marker. Staging writes a record under the collection before any document; a reader that finds the record holds every prefix under it to its markers and skips the ones they do not account for. A collection staged before markers existed carries no record and is read as it stands, so nothing already staged is re-extracted. Markers started on the per-chunk path only. Without one on the JSONL path a recorded collection in that format would read back as incomplete for good, so S3DocUploader writes a marker for each object it stores.

Stages a collection, damages it the way an interrupted run would - the marker missing, or an object gone from under a marker - and reads it back. Mocked tests can state that shape but not that S3 produces it: the listing, the reserved marker segment and the collection record have to line up against a real endpoint.

…isting Addresses review of the incomplete-document reader. Only an empty collection is recorded, so staging into one written before markers existed leaves its documents reading as they do today rather than turning incomplete and being extracted a second time. The record is written before the uploader is kept, so a failed record leaves the next call to try again instead of staging on unrecorded. Whether a collection is recorded now comes out of the listing the reader already performs, rather than a head_object. That drops a call per read and the ambiguity of a 403 standing in for a 404 under a prefix-limited policy. A marker name carries the whole digest, so two markers under one prefix cannot share a name and lose the first. A chunk id is taken by stripping what the writer added, so an id holding a '/' or a '.' still matches its marker. Markers are read during the prefetched listing rather than by a blocking GET per document on the way out. is_complete now states what a marker does not cover: extraction emits a source as several SourceDocuments when its chunks span rounds, and each is marked as it lands, so a source whose later parts were never stored reads as complete. Closing that needs the extraction pipeline to record when a source is finished.

acarbonetto · 2026-09-17T18:41:55Z

+        """
+        collection_path = join(key_prefix, collection_id, '')
+
+        existing = s3_client.list_objects_v2(


ensure that s3:ListBucket is included in all IAM roles and add this to the documentation - or else we will get an access denied here.

acarbonetto · 2026-09-17T18:46:46Z

                    collection_prefix=collection_prefix,
                    s3_encryption_key_id=self.s3_encryption_key_id,
                    num_threads=self.num_threads,
                    deterministic_document_key=self.deterministic_document_key


the document key continues to truncate to 5 characters. This could be a problem in some circumstances. We should consider removing this truncation: _doc_suffix (line 383–385)

…ion prefix The S3 uploaders joined a source id and a node id onto the collection prefix without checking either, so a separator in one wrote the object somewhere else. A node id under the reserved marker segment was worse than a stray write: the chunk read back as a completion marker, the declared set no longer matched what the prefix held, and the source document was skipped on every read, including after a re-stage that wrote the same key again. FileSystemTap and S3ChunkStore each carried their own id check, one a denylist and one an allowlist. Both now call one allowlist, which closes the encoded and Unicode separator variants a denylist misses.

How a separator behaves in a key segment is a property of the service, not of the string: botocore sends segments unencoded. These stage a hostile source id and a hostile node id against a real bucket and assert nothing of the document lands, alongside a generated-id control that still stages and reads back.

oussamahansal · 2026-09-18T19:01:23Z

+def chunk_id_from_key(chunk_key:str, source_doc_prefix:str) -> str:
+    """
+    The node id a chunk object is keyed on. Taken by stripping what the writer
+    added, because a node id can hold both a '/' and a '.'.


The stated reason no longer holds: source_document_prefix now runs every node id through validate_id_segment, which rejects /. Only the '.' half is still load-bearing (a node id ending in .json keys an object ending .json.json).

oussamahansal · 2026-09-18T19:02:00Z

+                for n in nodes
+            ])

            self._put(doc_output_path, s, 'text/plain', s3_client)


The marker write is guarded on if nodes but the object put isn't, so a document whose nodes all carry INDEX_KEY passes the source_document.nodes guard and writes an empty .jsonl with no marker — I staged one and got p/c/src1/src1-35df1.jsonl with a zero-byte body. The prefix then satisfies is_complete only by vacuous equality (set() == set()) and reads back as a SourceDocument with no nodes. Move the put under the same if nodes check.

noel-improv added 3 commits September 16, 2026 11:15

noel-improv self-assigned this Sep 17, 2026

noel-improv mentioned this pull request Sep 17, 2026

[FEATURE] End a source document when its last part lands #561

Open

8 tasks

noel-improv force-pushed the feat/detect-incomplete-staged-documents branch from d92c159 to d40c1ca Compare September 17, 2026 02:48

acarbonetto approved these changes Sep 17, 2026

View reviewed changes

noel-improv added 2 commits September 17, 2026 16:30

noel-improv requested a review from acarbonetto September 17, 2026 22:54

oussamahansal approved these changes Sep 18, 2026

View reviewed changes

Navigation Menu

[FEATURE] Read a half-written source document as incomplete - #560

[FEATURE] Read a half-written source document as incomplete#560
noel-improv wants to merge 5 commits into
awslabs:mainfrom
noel-improv:feat/detect-incomplete-staged-documents

noel-improv commented Sep 16, 2026 •

edited

Loading

Uh oh!

acarbonetto Sep 17, 2026

Uh oh!

acarbonetto Sep 17, 2026

Uh oh!

oussamahansal Sep 18, 2026

Uh oh!

oussamahansal Sep 18, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

noel-improv commented Sep 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Changes

Problem

Testing

Checklist

Not in this PR

Uh oh!

acarbonetto Sep 17, 2026

Choose a reason for hiding this comment

Uh oh!

acarbonetto Sep 17, 2026

Choose a reason for hiding this comment

Uh oh!

oussamahansal Sep 18, 2026

Choose a reason for hiding this comment

Uh oh!

oussamahansal Sep 18, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

noel-improv commented Sep 16, 2026 •

edited

Loading