[FEATURE] End a source document when its last part lands - #561
Open
noel-improv wants to merge 12 commits into
Open
noel-improv wants to merge 12 commits into
noel-improv wants to merge 12 commits into
Conversation
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.
noel-improv
force-pushed
the
feat/end-a-source-when-its-last-part-lands
branch
from
September 17, 2026 02:48
68c338c to
81af972
Compare
…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.
Extraction emits a source as several source documents when its chunks span rounds, and staging marked each part as it arrived. The markers under the prefix then accounted for exactly the chunks present, so a source that lost its later parts read back as a whole document. The pipeline now names the sources a round leaves with nothing left to extract, and only the last document emitted for one of those ends its source. Chunks are counted on the way in rather than on the way out, because a checkpointed run extracts fewer nodes than it was handed and counting output would leave a resumed source open for good. A source can also split within one round: the batch extractor sorts its output by node id, and documents are cut on contiguous runs of source id. The uploader accumulates the chunk ids it stored for each source, and the marker on the part that ends a source declares the whole set. That is what the reader compares against, so a prefix missing an earlier part fails whatever order the parts were written in, and neither the unordered JSONL upload pool nor a logged-and-swallowed chunk failure can pass a short prefix off as whole. A round can finish a source without emitting a document for it, since a resumed run drops chunks it has already extracted, so the part that would have ended the source never arrives. Running out of stream ends those, declaring what was stored. A run that dies never gets there, which leaves the prefix incomplete. A source whose chunk upload failed is dropped instead, so it gets no marker either way.
Three ways a completion marker said more than it knew. A closing marker declaring an empty chunk set was read as declaring nothing at all, because the fallback to the part's own ids treated an empty list the same as a missing field. Such a marker then certified only the part it sat on. A source that lost a chunk was dropped from the open set, but a later part put it back, so a surviving final part closed a source whose prefix is missing a chunk. The loss is recorded now and checked before any marker is written. The doc uploader ended its open sources per batch rather than per stream, so a source whose parts straddled a batch cut was closed declaring only the parts staged up to the cut. If the run then died, the truncated prefix read complete. A collection that is not recorded because it already holds an object is read without completeness checking for the rest of its life, which is worth more than a debug line.
_extract_fixed_batch held its own copy of the tail _emit_extracted runs, and the copy never set final_part. So when a caller passes pre-chunked nodes and node_batcher slices one source across worker batches, every part marked itself as ending its source and declared only its own chunks. A prefix missing the rest then read as complete. The copy is gone. _emit_extracted now decides final_part for every caller, and finished_sources of None means every source it was handed is finished, which is what the fixed-batch path has: each batch extracts every node fed into it.
EncryptedPut had collected the collection record, the open-source tracking and the marker write, none of which are encryption. Those move to CompletionMarkers, which both uploaders inherit alongside it, and EncryptedPut is back to _put. The tracking itself was written once per uploader with different failure handling, which is where the closed-a-poisoned-source bug lived. Both now call _open_source, _end_source and _poison_source on the mixin, so there is one place to get it right. Ending a source that was never opened returns None rather than an empty list, so it stays open instead of declaring that it stores nothing. The JSONL reader downloaded and parsed every object under a prefix before deciding the prefix was incomplete and discarding the result. A prefix with no marker cannot read complete, and the listing already says so, so that case returns before any download. Its chunk ids live in the object bodies rather than the keys, so the rest of the check still needs them read. written_nodes ran twice for the same document in the publisher and twice again in the worker; it runs once in each now.
noel-improv
force-pushed
the
feat/end-a-source-when-its-last-part-lands
branch
from
September 17, 2026 22:53
81af972 to
501c251
Compare
noel-improv
requested review from
acarbonetto and
oussamahansal
and removed request for
oussamahansal
September 17, 2026 23:02
oussamahansal
requested changes
Sep 18, 2026
…its own Every node a part carries can be a vector store artefact, which written_nodes filters out. The doc publisher still ended the source for such a part, but the worker writes no marker when there is nothing to store, so nothing closed the source and the end of the stream had nothing left to close. The prefix then read incomplete on every later run. The chunk uploader never had this, because it keys its submit loop on written_nodes. A part with nothing to store now leaves its source open, and the end of the stream ends it declaring what the earlier parts stored.
A run that dies mid-source leaves a part marker behind. The resumed run's checkpoint drops the chunks that part already stored, so the part that ends the source declares only what this run staged, and the completeness check compared that against every object under the prefix. The source read incomplete with all of its chunks present, and the marker key is a digest of the chunk ids, so a re-run reproduced the same declaration and the document never built. A part that leaves its source open now counts towards the declaration. This also closes the other direction: a prefix whose earlier part's objects are gone used to read complete, because nothing in the declaration accounted for them.
CompletionMarkers annotates _open_sources and _poisoned_sources as supplied by the host, but pydantic does not collect PrivateAttr from a plain mixin, so a host that leaves them out fails on first use rather than at construction. sources_finished_by reads as a query and consumes from outstanding, so it answers for each round once.
oussamahansal
approved these changes
Sep 18, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Follows #560, which skips a source document prefix its markers do not account for. A marker covers one
SourceDocument, and extraction emits a source as several of those when its chunks span rounds. Each part was marked as it landed, so the markers under the prefix accounted for exactly the chunks present and a source that lost its later parts read back as a whole document.The part that ends a source now declares every chunk id stored for that source, and the reader compares the objects present against that declaration.
Reproduced against S3 before the change: a source staged as two rounds read back whole at 4 chunks, and with round two removed it read back as a valid 2-chunk document rather than being skipped.
Changes
SourceDocument.final_partsays whether a document holds the last of its source's nodes. Defaults to true, so every producer that emits a source as one document is unchanged.ExtractionPipelinenames the sources a round leaves with nothing left to extract, and marks only the last document emitted for one of those. Chunks are counted on the way in: a checkpointed run extracts fewer nodes than it was handed, so counting output would leave a resumed source open for good.is_completecompares the prefix against that.Problem
Three things make a bare per-document marker insufficient:
A declared set survives all three. A flag on its own does not.
Related issue (if any): #560
Testing
pytest)Unit tests cover a source split across rounds, a source split within one round by sorted output, a round that drops chunks, each source ending on its own last document, the declared set on both uploaders, a source left open by the end of the stream, and a source whose chunk failed getting no marker. Full unit suite: 2,320 passed.
tests/integration/indexing/load/test_s3_staging_recovery_live.pyruns against real S3 and now covers a split source losing its opening part, in both storage formats. 18 cases pass in us-east-1.Each half was reverted in turn: forcing every part to be final failed the two split tests, and dropping the declared set from the marker failed the three declaration tests.
End to end against Bedrock, S3 and Neo4j: 3 documents extracted with 3 markers and the collection record, an intact build giving 3
__Source__nodes, and a rebuild after removing one marker giving 2.Checklist
One behaviour worth naming for review: a run that completes normally closes every source still open, declaring what was stored. That is what makes a source whose closing part never arrived readable at all, and it means the end of a run is treated as evidence. A run that dies is not, and leaves the prefix incomplete.
Not in this PR
FileBasedDocswrites no markers and is unchanged. Thefinal_partdefault keeps it behaving as it does today.By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.