[FIX] Overlap chunk uploads across documents and carry the thread count into workers - #24
Closed
noel-improv wants to merge 3 commits into
Closed
[FIX] Overlap chunk uploads across documents and carry the thread count into workers#24noel-improv wants to merge 3 commits into
noel-improv wants to merge 3 commits into
Conversation
S3ChunkUploader waited for one source document's chunk uploads before submitting the next document's. In-flight uploads were therefore bounded by a single document's chunk count rather than by the thread pool - 3.31 chunks per document on the WikiHow benchmark corpus - so most of the pool sat idle whatever it was sized to. A thread sweep on that corpus measured the write phase at ~13 minutes for 16, 32 and 64 threads alike, while the extraction phase either side of it halved with each doubling. Keep a bounded window of documents in flight instead. End to end on the same LLM work that measures 1.15x, with byte-identical output. Replaying the corpus against a stand-in S3 client shows a far larger gain, but it sizes the pool in-process and so assumes away the pool question this change does not address. Documents are still yielded in order, and still only once their own chunk uploads have been attempted. A failed chunk is logged rather than raised, as before.
S3BasedDocs sized its upload pool from GraphRAGConfig at upload time. That read happens in a worker spawned by run_pipeline, which inherits no parent memory, so a thread count set programmatically on the config was absent and read back as the default of 4. A benchmark asking for 64 uploaded with 4. Resolve the count in __init__, which runs in the process that configured GraphRAGConfig, and carry it as a field so it pickles with the handler. This is how an extractor's num_workers already survives the same trip. Callers can pass num_threads explicitly. Uploaders built directly still fall back to reading the config, so nothing outside this path changes. The chunk downloader reads the config the same way and is left alone; the write path is the one with measurements behind it.
|
Lexical Graph Coverage Report: The coverage is at 65.02% (target: 80%). Download the HTML report here. |
|
Lexical Graph Coverage Report: The coverage is at 65.02% (target: 80%). Download the HTML report here. |
… readback The uploaders each had their own copy of the num_threads fallback, and the two downloaders still read GraphRAGConfig at use time - the read this branch calls broken in a spawned worker. A ConfiguredThreadCount mixin holds the one definition and all four components take it, so read back is sized by whatever configured the writer. Also drops the benchmark figures from S3ChunkUploader.upload's docstring and corrects what it promises: _drain logs a failed chunk rather than raising, so a yielded document is one whose uploads were attempted, not one whose chunks are all in S3.
noel-improv
force-pushed
the
fix/s3-doc-store-upload-concurrency
branch
from
August 25, 2026 16:11
77ee254 to
a8883b8
Compare
|
Lexical Graph Coverage Report: The coverage is at 65.02% (target: 80%). Download the HTML report here. |
Owner
Author
|
Superseded by awslabs#507, which tracks the same branch. Closing this one so review happens in a single place. |
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
1 participant
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
S3ChunkUploaderwaited for one source document's chunk uploads before submitting the next document's, so in-flight uploads were bounded by a single document's chunk count rather than by the thread pool. On the WikiHow benchmark corpus that is 3.31 chunks per document, so most of the pool sat idle whatever it was sized to. A thread sweep measured the write phase at ~13 minutes for 16, 32 and 64 threads alike, while the extraction phase either side of it halved with each doubling.The pool was also getting the wrong size.
S3BasedDocsreadGraphRAGConfig.extraction_num_threads_per_workerat upload time, and that read happens in a worker spawned byrun_pipeline, which inherits no parent memory. A run asking for 64 uploaded with the default of 4.End to end this measures 1.15x on the same LLM work, with byte-identical output. A replay against a stand-in S3 client shows a far larger gain, but it sizes the pool in-process and so assumes away the pool question, which this PR does not address.
Changes
S3ChunkUploader.uploadkeeps a bounded window of documents in flight instead of draining one document before submitting the next. Documents are still yielded in order, and still only once their own chunk uploads have been attempted.S3BasedDocs.__init__, which runs in the process that configuredGraphRAGConfig, and carried as a field so it pickles with the handler. This is the trip an extractor'snum_workersalready survives.ConfiguredThreadCountmixin holds the single definition of the count-or-config fallback. Both uploaders and both downloaders use it, so read back is sized by whatever configured the writer rather than by a config read in the wrong process.Problem
Uploads to the S3 doc store did not scale with
extraction_num_threads_per_worker, for two independent reasons: the uploader never had enough work queued to use the pool, and the pool was sized from a config read that returned the default in the worker process.Related issue (if any): #
Testing
pytest)pytest tests/uniton this branch: 1981 passed. The one error,test_integ_dependency_compatibility::test_boto3_botocore_version_alignment, is a pre-existing pip-resolution failure unrelated to this change.New tests cover overlap across documents (a
threading.Barrierthat only releases on genuinely simultaneous uploads, so it fails against the serial version), yield order, every chunk uploaded exactly once, theINDEX_KEYskip, a failed chunk not stopping the stream, and the thread count surviving into an uploader and a downloader whose process sees the default.Measured against the WikiHow benchmark corpus on real S3, output byte-identical to the previous implementation.
Checklist
One behaviour change worth naming:
num_threadsis now resolved whenS3BasedDocsis constructed, so settingGraphRAGConfig.extraction_num_threads_per_workerafterwards no longer affects it. The call sites inbenchmarks/scripts/benchmark_extract.pyand the integration-test scripts take the value from the environment, which the config reads lazily on first access, so none of them are affected.A failed chunk upload is logged and not raised, and its document is still yielded. That is unchanged behaviour, but it means a yielded document does not guarantee every one of its chunks reached S3. The docstring now says so rather than claiming durability.
By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.