[FIX] Size the LLM call pool and the Bedrock connection pool to the thread count - #23
Open
noel-improv wants to merge 54 commits into
Open
noel-improv wants to merge 54 commits into
noel-improv wants to merge 54 commits into
Conversation
Bumps [astro](https://github.com/withastro/astro/tree/HEAD/packages/astro) from 7.1.1 to 7.1.3. - [Release notes](https://github.com/withastro/astro/releases) - [Changelog](https://github.com/withastro/astro/blob/main/packages/astro/CHANGELOG.md) - [Commits](https://github.com/withastro/astro/commits/astro@7.1.3/packages/astro) --- updated-dependencies: - dependency-name: astro dependency-version: 7.1.3 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
…s#434) Bumps [postcss](https://github.com/postcss/postcss) from 8.5.12 to 8.5.23. - [Release notes](https://github.com/postcss/postcss/releases) - [Changelog](https://github.com/postcss/postcss/blob/main/CHANGELOG.md) - [Commits](postcss/postcss@8.5.12...8.5.23) --- updated-dependencies: - dependency-name: postcss dependency-version: 8.5.23 dependency-type: indirect ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
…ite (awslabs#431) Bumps [@astrojs/starlight](https://github.com/withastro/starlight/tree/HEAD/packages/starlight) from 0.41.3 to 0.41.4. - [Release notes](https://github.com/withastro/starlight/releases) - [Changelog](https://github.com/withastro/starlight/blob/main/packages/starlight/CHANGELOG.md) - [Commits](https://github.com/withastro/starlight/commits/@astrojs/starlight@0.41.4/packages/starlight) --- updated-dependencies: - dependency-name: "@astrojs/starlight" dependency-version: 0.41.4 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Bumps [svgo](https://github.com/svg/svgo) from 4.0.1 to 4.0.2. - [Release notes](https://github.com/svg/svgo/releases) - [Commits](svg/svgo@v4.0.1...v4.0.2) --- updated-dependencies: - dependency-name: svgo dependency-version: 4.0.2 dependency-type: indirect ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Bumps [astral-sh/setup-uv](https://github.com/astral-sh/setup-uv) from 8.3.2 to 9.0.0. - [Release notes](https://github.com/astral-sh/setup-uv/releases) - [Commits](astral-sh/setup-uv@11f9893...c771a70) --- updated-dependencies: - dependency-name: astral-sh/setup-uv dependency-version: 9.0.0 dependency-type: direct:production update-type: version-update:semver-major ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Bumps [react](https://github.com/react/react/tree/HEAD/packages/react) and [@types/react](https://github.com/DefinitelyTyped/DefinitelyTyped/tree/HEAD/types/react). These dependencies needed to be updated together. Updates `react` from 18.3.1 to 19.2.8 - [Release notes](https://github.com/react/react/releases) - [Changelog](https://github.com/react/react/blob/main/CHANGELOG.md) - [Commits](https://github.com/react/react/commits/v19.2.8/packages/react) Updates `@types/react` from 18.3.28 to 19.2.17 - [Release notes](https://github.com/DefinitelyTyped/DefinitelyTyped/releases) - [Commits](https://github.com/DefinitelyTyped/DefinitelyTyped/commits/HEAD/types/react) --- updated-dependencies: - dependency-name: "@types/react" dependency-version: 19.2.17 dependency-type: direct:production update-type: version-update:semver-major - dependency-name: react dependency-version: 19.2.8 dependency-type: direct:production update-type: version-update:semver-major ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
…bs#430) Bumps [react-dom](https://github.com/react/react/tree/HEAD/packages/react-dom) and [@types/react-dom](https://github.com/DefinitelyTyped/DefinitelyTyped/tree/HEAD/types/react-dom). These dependencies needed to be updated together. Updates `react-dom` from 18.3.1 to 19.2.8 - [Release notes](https://github.com/react/react/releases) - [Changelog](https://github.com/react/react/blob/main/CHANGELOG.md) - [Commits](https://github.com/react/react/commits/v19.2.8/packages/react-dom) Updates `@types/react-dom` from 18.3.7 to 19.2.3 - [Release notes](https://github.com/DefinitelyTyped/DefinitelyTyped/releases) - [Commits](https://github.com/DefinitelyTyped/DefinitelyTyped/commits/HEAD/types/react-dom) --- updated-dependencies: - dependency-name: "@types/react-dom" dependency-version: 19.2.3 dependency-type: direct:production update-type: version-update:semver-major - dependency-name: react-dom dependency-version: 19.2.8 dependency-type: direct:production update-type: version-update:semver-major ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
…wslabs#448) Bumps [@types/react](https://github.com/DefinitelyTyped/DefinitelyTyped/tree/HEAD/types/react) from 19.2.17 to 19.2.18. - [Release notes](https://github.com/DefinitelyTyped/DefinitelyTyped/releases) - [Commits](https://github.com/DefinitelyTyped/DefinitelyTyped/commits/HEAD/types/react) --- updated-dependencies: - dependency-name: "@types/react" dependency-version: 19.2.18 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
…ers_recursive (awslabs#415) When MetadataFilters has an empty filters list, the function previously returned '()' (from joining an empty list within parentheses). This truthy string bypassed the guard in VersionManager._get_existing_source_nodes(), causing invalid Cypher: WHERE () AND <filter> Neptune Database rejects this with MalformedQueryException: Invalid input 'A': expected whitespace, comment or a relationship pattern Fix: return '' (empty string) when filter_strs is empty, before reaching the condition-based join logic. Fixes awslabs#408
Bumps [actions/stale](https://github.com/actions/stale) from 10 to 11. - [Release notes](https://github.com/actions/stale/releases) - [Changelog](https://github.com/actions/stale/blob/main/CHANGELOG.md) - [Commits](actions/stale@v10...v11) --- updated-dependencies: - dependency-name: actions/stale dependency-version: '11' dependency-type: direct:production update-type: version-update:semver-major ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
…slabs#443) Bumps [pypa/gh-action-pypi-publish](https://github.com/pypa/gh-action-pypi-publish) from 1.14.1 to 1.14.2. - [Release notes](https://github.com/pypa/gh-action-pypi-publish/releases) - [Commits](pypa/gh-action-pypi-publish@ba38be9...dc37677) --- updated-dependencies: - dependency-name: pypa/gh-action-pypi-publish dependency-version: 1.14.2 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
…wslabs#444) Bumps [@astrojs/react](https://github.com/withastro/astro/tree/HEAD/packages/integrations/react) from 6.0.1 to 6.0.2. - [Release notes](https://github.com/withastro/astro/releases) - [Changelog](https://github.com/withastro/astro/blob/main/packages/integrations/react/CHANGELOG.md) - [Commits](https://github.com/withastro/astro/commits/@astrojs/react@6.0.2/packages/integrations/react) --- updated-dependencies: - dependency-name: "@astrojs/react" dependency-version: 6.0.2 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
awslabs#445) Bumps [@types/react-dom](https://github.com/DefinitelyTyped/DefinitelyTyped/tree/HEAD/types/react-dom) from 19.2.3 to 19.2.4. - [Release notes](https://github.com/DefinitelyTyped/DefinitelyTyped/releases) - [Commits](https://github.com/DefinitelyTyped/DefinitelyTyped/commits/HEAD/types/react-dom) --- updated-dependencies: - dependency-name: "@types/react-dom" dependency-version: 19.2.4 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
…ite (awslabs#447) Bumps [@astrojs/starlight](https://github.com/withastro/starlight/tree/HEAD/packages/starlight) from 0.41.4 to 0.41.5. - [Release notes](https://github.com/withastro/starlight/releases) - [Changelog](https://github.com/withastro/starlight/blob/main/packages/starlight/CHANGELOG.md) - [Commits](https://github.com/withastro/starlight/commits/@astrojs/starlight@0.41.5/packages/starlight) --- updated-dependencies: - dependency-name: "@astrojs/starlight" dependency-version: 0.41.5 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Bumps [astro](https://github.com/withastro/astro/tree/HEAD/packages/astro) from 7.1.3 to 7.1.6. - [Release notes](https://github.com/withastro/astro/releases) - [Changelog](https://github.com/withastro/astro/blob/main/packages/astro/CHANGELOG.md) - [Commits](https://github.com/withastro/astro/commits/astro@7.1.6/packages/astro) --- updated-dependencies: - dependency-name: astro dependency-version: 7.1.6 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
…bs#417) * Separate benchmark pipeline into dedicated top-level directory * update readme * Suite file path resolution and Missing test dependencies --------- Co-authored-by: Oussama Hansal <haooussa@amazon.com>
…slabs#420) * perf(lexical-graph): parallelise per-document chunk listing in S3ChunkDownloader S3ChunkDownloader.download listed each source document's chunks with a separate list_objects_v2 call issued one at a time on the main thread. On the 5,000-doc WikiHow benchmark this listing took 255.7s of the 578s S3 readback; the other ~322s is the per-chunk downloads, which were already threaded and are unchanged here. Run the per-document list calls on a thread pool sized by extraction_num_threads_per_worker (the setting the downloads already use). executor.map preserves input order, so documents still yield in prefix order and the listings prefetch while each document's chunks download. Listing and downloading use separate pools. Also read Contents with .get('Contents', []) so an empty prefix returns no keys instead of raising KeyError. Measured against the same collection, listing only: 255.7s serial -> 16.5s at 16 threads, 8.3s at 32, 4.4s at 64. Listing is I/O-bound (no CPU, trivial memory), so raising extraction_num_threads_per_worker above its default of 4 compounds the gain. This addresses the listing half of readback; the download half is a separate change (store an S3 reference in the graph instead of the chunk). Tests: order preservation across documents, and a Barrier-based concurrency test that passes on the parallel implementation and fails (BrokenBarrierError) on the serial one. * perf(lexical-graph): bound listing window and overlap chunk downloads in S3ChunkDownloader The parallel listing used executor.map, which submits a listing task for every document up front and buffers each completed chunk-key list until the serial download consumer reaches it, so peak memory scaled with the whole collection instead of one document. Replace map with an explicit sliding window: at most extraction_num_threads_per_worker documents' listings are in flight, refilled one-per-consumed, so memory tracks the window. Each document's chunk downloads are dispatched onto a shared download executor as soon as its listing completes, so downloads for consecutive documents overlap. Reuse the collection-level paginator instead of rebuilding it per prefix. Document order is unchanged. Add a regression test asserting the window stays bounded when the consumer stalls on the first document. * fix(lexical-graph): dispatch S3 chunk downloads lazily per document Eager download dispatch in _list_and_dispatch submitted every chunk's download as soon as a document was listed, so the sliding window bounded listings but not downloaded payloads: up to num_threads look-ahead documents' chunk data was resident at once (OOM risk), and abandoning the generator early still drained those downloads. Split listing from downloading: list chunk keys on the bounded window, download only the current document's chunks before yielding. Restores the one-document memory bound and no-work-on-early-stop, keeps listing parallelism. Rework the window test to drive the generator from a consumer thread (removes the dead 10s-timeout handshake) and add a guard that fails on eager dispatch: no later-document download starts while stalled on doc 0. * fix(lexical-graph): size the S3 connection pool to the configured thread count Listing and downloading now run concurrently, so peak S3 connections reach twice the configured thread count. botocore defaults the pool to 10, well under the thread counts the extract stage uses, and once it's exhausted botocore discards and reopens connections, giving back the concurrency the threads were meant to buy. Size the S3 client's pool to 2x extraction_num_threads_per_worker, floored at botocore's own 10 so the default 4-thread case doesn't drop to 8. Other services keep botocore's defaults. Measured on a 50-object batch: 64 threads against a pool of 10 takes 321.6ms, and the same 64 threads with a pool of 64 takes 99.1ms. Below 32 threads the pool size makes no measurable difference. * fix(lexical-graph): close the chunk downloader deterministically on early exit download() holds two thread pools open across its yields, so they shut down when the generator is closed rather than when its loop ends. A consumer that stopped early left that to the garbage collector: prompt under CPython refcounting, but unbounded if a reference cycle keeps the generator alive, or on a runtime that doesn't refcount. Wrap the inner generator in contextlib.closing at the one caller, so its GeneratorExit propagates and shutdown(wait=True) runs on the way out. Document the requirement on download() for anyone calling it directly. Adds a test that a listing failure surfaces to the consumer instead of being swallowed while the executors unwind. The caller fix ships without a direct test of its own: CPython's refcounting cleans up either way, so a test written against it passes with the fix reverted.
Adds a generic LlamaIndex reader plugin provider that allows using any LlamaIndex reader package as an extraction source. Files: - providers/__init__.py (registration) - providers/llama_index_plugin_reader_provider.py (implementation) - reader_provider_config.py (config updates) - tests/test_llama_index_plugin_reader_provider.py (test) This is a clean PR as requested — contains only the plugin files, no unrelated changes.
* fix(lexical-graph): restore extract_tables on PDFReaderConfig awslabs#449 removed the field while AdvancedPDFReaderProvider still reads config.extract_tables in three places, including __init__, so constructing the provider raised AttributeError. The config tests asserted return_full_document and metadata_fn but never this field, which is why nothing upstream caught the removal. They now cover it. * ci(lexical-graph): install pymupdf so the PDF provider tests run The AdvancedPDFReaderProvider tests open with importorskip("pymupdf"), which requirements.txt does not carry because the dependency is optional. All ten skipped in CI and the job still reported success, so the provider shipped unguarded. Installing pymupdf for tests only leaves it optional for users. * fix(lexical-graph): declare pymupdf as a pinned test extra Address review on awslabs#459. Move pymupdf into the existing [project.optional-dependencies] test group, pinned, and install it through -e '.[test]' rather than listing test packages in the workflow. The workflow was duplicating the group by hand. Drop test_extract_tables_field_is_present: the two updated config tests already fail if the field is removed, which the reviewer pointed out and I confirmed. * fix(ci): install requirements and the editable package separately The combined command failed on every Python version: requirements.txt starts with `--only-binary :all:`, which applies to the whole invocation and blocks building the local package, so `-e '.[test]'` could not resolve. Splitting them keeps dependencies binary-only while the package still builds. Verified by reproducing the CI command locally rather than by dry run.
…wslabs#461) Bumps [pymupdf](https://github.com/pymupdf/pymupdf) from 1.27.2.3 to 1.28.2. - [Release notes](https://github.com/pymupdf/pymupdf/releases) - [Changelog](https://github.com/pymupdf/PyMuPDF/blob/main/changes.txt) - [Commits](pymupdf/PyMuPDF@1.27.2.3...1.28.2) --- updated-dependencies: - dependency-name: pymupdf dependency-version: 1.28.2 dependency-type: direct:development update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
…wslabs#462) Bumps [actions/attest-build-provenance](https://github.com/actions/attest-build-provenance) from 4.1.1 to 4.2.2. - [Release notes](https://github.com/actions/attest-build-provenance/releases) - [Changelog](https://github.com/actions/attest-build-provenance/blob/main/RELEASE.md) - [Commits](actions/attest-build-provenance@0f67c3f...4d10147) --- updated-dependencies: - dependency-name: actions/attest-build-provenance dependency-version: 4.2.2 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Bumps [astro](https://github.com/withastro/astro/tree/HEAD/packages/astro) from 7.1.6 to 7.2.0. - [Release notes](https://github.com/withastro/astro/releases) - [Changelog](https://github.com/withastro/astro/blob/main/packages/astro/CHANGELOG.md) - [Commits](https://github.com/withastro/astro/commits/astro@7.2.0/packages/astro) --- updated-dependencies: - dependency-name: astro dependency-version: 7.2.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
…ite (awslabs#464) Bumps [@astrojs/starlight](https://github.com/withastro/starlight/tree/HEAD/packages/starlight) from 0.41.5 to 0.41.7. - [Release notes](https://github.com/withastro/starlight/releases) - [Changelog](https://github.com/withastro/starlight/blob/main/packages/starlight/CHANGELOG.md) - [Commits](https://github.com/withastro/starlight/commits/@astrojs/starlight@0.41.7/packages/starlight) --- updated-dependencies: - dependency-name: "@astrojs/starlight" dependency-version: 0.41.7 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
…te (awslabs#454) Nova 2 Lite returns None for some documents during batch extraction. Add null checks in _update_node() for both BatchTopicExtractorSync and BatchLLMPropositionExtractorSync so they gracefully default to empty results instead of crashing with a TypeError. Also fix the else branch in BatchTopicExtractorSync to use {'topics': []} for consistency. Co-authored-by: Oussama Hansal <haooussa@amazon.com>
…am-search retrievers (awslabs#455) * feat(benchmarks): Add multi-retriever benchmark harness with beam-search retrievers Add support for running all retrievers in a single benchmark pass: - Add topic-beam-chunk_only and topic_beam_search to retriever factory - Import ChunkCosineSimilaritySearch, SemanticChunkBeamGraphSearch, SemanticGuidedChunkRetriever, TopicBeamSearch - Tune shared sub-retriever params: max_search_results=10, max_context_tokens=3000, token_truncation_mode=per_topic_cap - Create run_all_retrievers.sh loop script (extract+build once, query+evaluate per retriever) - Create benchmark.concurrentqa.all-retrievers dataset file - Forward BENCHMARK_ALL_RETRIEVERS and BENCHMARK_DATASET env vars through build-tests.sh to .env.testing - Conditionally invoke run_all_retrievers.sh from run_test_suite.sh * docs: Add BENCHMARK_ALL_RETRIEVERS and BENCHMARK_DATASET to env.template files * Address PR review comments for multi-retriever harness --------- Co-authored-by: Oussama Hansal <haooussa@amazon.com>
…on (awslabs#457) * feat(benchmarks): Upload benchmark results to S3 after evaluation Add upload_benchmark_results_to_s3() and upload_all_benchmark_results_to_s3() to benchmarks/utils/s3_utils.py using boto3. The per-retriever upload is called at the end of run_benchmark_evaluate() after scores are written but before assertions run, ensuring results persist even if assertions fail. Uses S3_RESULTS_BUCKET and S3_RESULTS_PREFIX env vars (same as integration_test_handler.py). No-ops with a warning when bucket is unset. * adress PR comments * upload_all_benchmark_results_to_s3 --------- Co-authored-by: Oussama Hansal <haooussa@amazon.com>
…slabs#456) * feat(benchmarks): Add PGA bio/stat split evaluation support Allow running benchmarks against just the PGA bio or PGA stat subset by setting BENCHMARK_DATASET=pga_bio or BENCHMARK_DATASET=pga_stat. - Add 'pga_bio' and 'pga_stat' entries to QA_FILE_MAP - Resolve data subdirectory to 'pga' for all pga* datasets in load_qa_pairs() and sync_benchmark_data_from_s3() - Read dataset_name from BENCHMARK_DATASET env var in PgaBenchmarkQuery * address PR comments --------- Co-authored-by: Oussama Hansal <haooussa@amazon.com>
…slabs#421) * feat(integration-tests): add S3 doc-store option to WikiHow extraction benchmark benchmark_extract hardcoded FileBasedDocs (local disk), so the benchmark could not measure S3-backed document storage. Add a BENCHMARK_DOC_STORE env switch: "file" (default, unchanged behaviour) keeps FileBasedDocs; "s3" uses S3BasedDocs, with BENCHMARK_S3_JSONL selecting the per-chunk or JSONL write path. S3BasedDocs uses collection_id=None (timestamp per run) so repeated runs do not accumulate into one collection. A log line records the resolved store, flag, and collection_id. build-tests.sh propagates the two new variables into .env.testing so they reach the notebook; without this they stay local to the caller's shell and the run silently falls back to FileBasedDocs. * fix(integration-tests): validate S3 env vars before S3 doc-store benchmark BENCHMARK_DOC_STORE=s3 read AWS_REGION_NAME/S3_RESULTS_BUCKET/S3_RESULTS_PREFIX via plain os.environ[...]. In non-batch (prototype) runs those vars are not otherwise required, so a missing one raised a bare KeyError and aborted the run before extraction. Check them up front and raise a clear ValueError naming the missing variables. * fix(benchmarks): address review on the S3 doc-store benchmark Set GraphRAGConfig.aws_region from AWS_REGION_NAME. The region argument on BatchConfig and S3BasedDocs is stored and never read, so every boto3 client was resolving region from AWS_REGION or the boto default instead, which can differ from the bucket's region. Check both paths' env vars before the BatchConfig block dereferences them, so a missing variable reports what is missing rather than raising a bare KeyError. Count extracted source documents with a delimiter listing instead of iterating the doc store, which downloaded every object back for a count (~16.5k GETs on the WikiHow run). Document BENCHMARK_DOC_STORE and BENCHMARK_S3_JSONL in both env templates and the benchmarks README.
…wslabs#436) * feat(lexical-graph): add pluggable ChunkStore interface Chunk text lives inline as a chunk.value graph property today, which Neptune keeps in memory for every node whether or not a query reads it. This adds a ChunkStore interface (get/put/get_batch) with a ChunkStoreFactory registration pattern mirroring GraphStoreFactory, and an InGraphChunkStore implementing today's behavior as the default. No existing write or read path is wired to it yet — this is the interface and default backend only, laying the groundwork for an external (e.g. S3) backend in a follow-up. * refactor(lexical-graph): give ChunkStore.get() a default implementation Every backend implementing ChunkStore ends up writing the same get_batch([chunk_id]).get(chunk_id) one-liner for get(). Make get() a concrete method on the ABC instead of forcing each implementation to repeat it; only put()/get_batch() stay abstract. * fix(lexical-graph): don't reject duck-typed graph clients in ChunkStoreFactory InGraphChunkStoreFactory checked isinstance(graph_store, GraphStore), but GraphBatchClient - what real graph builders actually pass as graph_client - duck-types GraphStore without subclassing it, matching how the rest of the codebase already treats this parameter. The isinstance check would have rejected it once ChunkGraphBuilder wires up to ChunkStore. Check for None instead of type. * fix(lexical-graph): let InGraphChunkStore read back what it writes get_batch() required a chunk to have an __EXTRACTED_FROM__ link to a __Source__, but put() only writes chunk.value and never creates that link. Text written through the store was therefore invisible to a read from the same store, and silently so - the query is well-formed, it just matches nothing. Matching on chunk id alone is the only coherent fix: put(chunk_id, text) has no source id in its signature, so it cannot create the link. Adds a live Neo4j round-trip test under tests/integration/, skipped unless NEO4J_TEST_URI is set. The mocked tests assert on the shape of the query string, which cannot catch a query that is well-formed but matches the wrong thing; four of the five live cases fail against the old query and pass against the new one. * docs(lexical-graph): correct ChunkStoreFactory docstrings, isolate registry in tests The docstrings claimed ChunkStoreFactory mirrors GraphStoreFactory and that InGraphChunkStoreFactory is a last-resort fallback in the way DummyGraphStoreFactory is. Neither holds. GraphStoreFactory has no default at all and raises for empty graph_info as well as unrecognized graph_info, and DummyGraphStoreFactory matches a literal dummy:// prefix rather than acting as a fallback. Describes what the code actually does and names the difference from GraphStoreFactory. Also pins the behavior the corrected text describes: an unrecognized non-empty chunk_info raises rather than falling back to the in-graph store, so a typo'd backend URI fails loudly. Registrations live in a module-level dict with no reset between tests, so a test registering a factory that matches any falsy chunk_info changed what later tests saw - reordering the file made two tests fail. An autouse fixture now snapshots and restores the registry. * refactor(lexical-graph): annotate try_create Optional, note the batched-write gap try_create returns None when a factory doesn't recognise chunk_info, so the annotation should say Optional[ChunkStore] rather than ChunkStore. Covers the ChunkStoreFactoryMethod ABC and InGraphChunkStoreFactory; S3ChunkStoreFactory carries the same change on the branch that adds it. GraphStoreFactoryMethod and VectorIndexFactoryMethod document returning None while annotated otherwise, but fixing those means touching code this PR doesn't, so they are left alone. Add a TODO on ChunkStore recording that it has no batched write. ChunkGraphBuilder calls put() once per chunk, and reads already batch and thread, so the write path is where the remaining ingestion cost sits. * test(lexical-graph): require an empty database instead of wiping it The live Neo4j fixture ran MATCH (n) DETACH DELETE n before and after every test, so pointing NEO4J_TEST_URI at a database with data in it destroyed that data. Count the nodes first and fail the test if any exist. Teardown now deletes only __Chunk__ nodes, which is all InGraphChunkStore.put() creates.
…nsert (downmerge) (awslabs#478) * fix KeyError 'params' in domain-entity batch insert * add regression tests --------- Co-authored-by: Oussama Hansal <haooussa@amazon.com>
Bumps [astro](https://github.com/withastro/astro/tree/HEAD/packages/astro) from 7.2.0 to 7.2.2. - [Release notes](https://github.com/withastro/astro/releases) - [Changelog](https://github.com/withastro/astro/blob/main/packages/astro/CHANGELOG.md) - [Commits](https://github.com/withastro/astro/commits/astro@7.2.2/packages/astro) --- updated-dependencies: - dependency-name: astro dependency-version: 7.2.2 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Bumps [astral-sh/setup-uv](https://github.com/astral-sh/setup-uv) from 9.0.0 to 10.0.1. - [Release notes](https://github.com/astral-sh/setup-uv/releases) - [Commits](astral-sh/setup-uv@c771a70...20cfd1b) --- updated-dependencies: - dependency-name: astral-sh/setup-uv dependency-version: 10.0.1 dependency-type: direct:production update-type: version-update:semver-major ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
…ount Extraction awaits its blocking LLM calls through asyncio.to_thread, which uses the event loop's default executor. CPython sizes that at min(32, cpu_count + 4), so past that point extraction_num_threads_per_worker buys nothing: the extra jobs queue for a thread instead of running. Route those calls through a pool sized to the caller's own worker count. The pool never shrinks below CPython's default, so a low thread setting cannot make things worse than before. Callers pass their worker count rather than reading GraphRAGConfig, because extraction runs in a spawned process and a thread count set programmatically on the config is absent there. The extractor's num_workers is pickled with the component and survives.
llm_cache built its bedrock-runtime client without max_pool_connections, so it took botocore's default of 10 while the thread fix drives up to 64 concurrent calls through it. Past the pool botocore discards and reopens connections, giving back the concurrency the thread fix bought. The pool takes the largest of the botocore default, the configured thread count and CPython's executor floor. No single source is reliable: extraction runs in a spawned process, so a thread count set on GraphRAGConfig in the parent reads back as the default, while the environment variable survives. Also collapses the three duplicated client-construction blocks into one helper, and the repeated patch and assertion blocks in the cross-region tests into two shared helpers.
|
Lexical Graph Coverage Report: The coverage is at 64.31% (target: 80%). Download the HTML report here. |
|
BYOKG-RAG Coverage Report: The coverage is at 94.88% (target: 80%). Download the HTML report here. |
…l heuristic Drop the machine-specific timing from the module docstring and the worked example from the pool-size helper, leaving the contract. Size the bedrock client pool the same way config.py already does. Cap the pool so a thread count set far too high cannot multiply across spawned processes, and move the module next to the cache that reads it.
|
BYOKG-RAG Coverage Report: The coverage is at 94.88% (target: 80%). Download the HTML report here. |
|
Lexical Graph Coverage Report: The coverage is at 64.34% (target: 80%). Download the HTML report here. |
The client sits behind the LLM call pool, and pool_size() reads 0 until something creates it, so a caller arriving first sized the pool from the config default while the executor still ran at its own floor.
|
BYOKG-RAG Coverage Report: The coverage is at 94.88% (target: 80%). Download the HTML report here. |
|
Lexical Graph Coverage Report: The coverage is at unknown% (target: unknown%). Download the HTML report here. |
Fold the client's sizing comment into its docstring so the rule is stated in one place. Treat an explicit thread count of zero as a value rather than as unset. Warn whenever a request exceeds the maximum, not only when that request happens to grow the pool, and warn once rather than per call.
|
BYOKG-RAG Coverage Report: The coverage is at 94.88% (target: 80%). Download the HTML report here. |
|
Lexical Graph Coverage Report: The coverage is at unknown% (target: unknown%). Download the HTML report here. |
The pool belongs to the process that created it. Pipeline extraction runs in workers that are spawned and torn down per batch, so it goes with them; the threads only outlive anything when a caller drives the extractors in-process.
…core count MIN_POOL_SIZE tracks cpu_count, so asserting the client pool equals it only held on a host with seven or more cores; CI runs on four and got 10 against an expected 8. Assert the floors as a property instead, cover the executor floor separately by making it the binding term, and cover the explicit-zero case that the switch away from truthiness introduced.
|
BYOKG-RAG Coverage Report: The coverage is at 94.88% (target: 80%). Download the HTML report here. |
|
Lexical Graph Coverage Report: The coverage is at 64.35% (target: 80%). Download the HTML report here. |
…ximum A caller-supplied thread count reached the bedrock client unclamped, so a request far above the pool's own maximum sized the socket pool from a number the executor will never reach. Cap it, reject a negative count rather than absorbing it into the floors, and say that the ceiling bounds one pool rather than the process.
|
BYOKG-RAG Coverage Report: The coverage is at 94.88% (target: 80%). Download the HTML report here. |
|
Lexical Graph Coverage Report: The coverage is at 64.35% (target: 80%). Download the HTML report here. |
… the process Growing the pool replaced it and shut the old one down without waiting, so the superseded threads ran until their work drained and the maximum bounded a single pool rather than the process. Size it from the first caller instead, and log once when a later caller asks for more than it has. That also removes the replacement race the lock was widened for. The lock still spans the submit, now to stop a concurrent shutdown clearing the executor, and the test that covered the old race is repointed at that one.
…code cannot say The earlier pass trimmed only the two blocks named in review and left the rest, and folding an inline comment into the client docstring grew it further. Keep the spawn rationale, the boto3 thread-safety reason, the context copy and why the lock spans the submit, and drop the elaboration around them.
|
BYOKG-RAG Coverage Report: The coverage is at 94.88% (target: 80%). Download the HTML report here. |
|
Lexical Graph Coverage Report: The coverage is at 64.35% (target: 80%). Download the HTML report here. |
The comment read as though the cap prevented pool sizes multiplying across processes. It does not: each spawned worker gets its own pool, so the fleet ceiling is workers x MAX_POOL_SIZE. State the per-process scope instead.
|
Lexical Graph Coverage Report: The coverage is at 64.35% (target: 80%). Download the HTML report here. |
|
BYOKG-RAG Coverage Report: The coverage is at 94.88% (target: 80%). Download the HTML report here. |
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
8 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.
Two commits, one concern in two layers: the executor that runs the LLM calls, and the client those calls go through.
The thread ceiling
Extraction runs one asyncio job per node, each awaiting a synchronous LLM call.
asyncio.to_threadhands those to the default executor, which CPython sizes atmin(32, cpu_count + 4). That cap, notextraction_num_threads_per_worker, decided how many calls were in flight.500 docs, one worker, on-demand Sonnet 4.6,
ml.m5.4xlarge, model-call phase only:16 is flat because both settings sit under the cap of 20 on that instance.
The connection pool
llm_cache.pybuilt its bedrock-runtime client with nomax_pool_connections, so it took botocore's default of 10 while the change above drives up to 64 concurrent calls through it. Past the pool botocore discards and reopens connections, costing back the concurrency.The pool takes the largest of botocore's default,
extraction_num_threads_per_worker, and CPython's executor size. Reading the config alone would be a no-op on the path that needs it, because extraction runs in a spawned process and a thread count set onGraphRAGConfigin the parent reads back in the worker as the default:Known limitation: a count set programmatically above the CPython floor stays invisible to the client, so the pool can still sit under peak in-flight calls. Closing that means threading the extractor's
num_workersthrough, asllm_concurrency.pydoes.Also collapses three duplicated client-construction blocks, and the cross-region tests' repeated patch and assert blocks, into helpers.
1977 tests pass. The one error in
test_integ_dependency_compatibility.pyis pre-existing on main.