awslabs / awslabs/graphrag-toolkit

[FEATURE] Cut `lexical.long` integration-test runtime dominated by the batch-extract steps

Open
#485 1 comment 0 reactions 1 assignee Claimed by @oussamahansal View on GitHub
enhancement
Dominant language
Python
Stars
442
Forks
106
Avg merge
2d 3h
Merged PRs (30d)
41

Description

### Package

lexical-graph

### Problem statement

The `lexical.long` integration-test suite takes **over 2 hours** end-to-end, and the
overwhelming majority of that time is spent in the two batch-extract steps:

- `batch_extract.BatchExtractToS3`
- `batch_extract.BatchExtractAutoTuneToS3`

Together these two steps take roughly **1 hour 20 minutes** (for ONE of the extract processes) for a **70-document** corpus
(`integration-tests/source-data/corpus-modified.json`). This is too slow for a
70-document test set and makes the long suite painful to iterate on.

### Why it is slow (root cause)

Both tests call `graph_index.extract(...)` over the **same** 70-doc corpus, and they run
**sequentially** in a single SageMaker notebook (`test_suite.py` runs the test list in a
`for` loop). So we pay the batch-extract cost **twice, back-to-back**.

Wall-clock for a batch-extract step is **not** driven by the number of documents. From the
lexical-graph batch-inference implementation:

- **Fixed per-job Bedrock overhead dominates.** Each Bedrock `create_model_invocation_job`
carries a large fixed queue/startup cost (job scheduling + model provisioning) that is
independent of record count.
(`lexical-graph/.../indexing/utils/batch_inference_utils.py`)
- **Completion is discovered on a 60-second poll cadence, and the poll sleeps *before* the
first check** — so even an instantly-completing job costs ≥60s of polling latency:
```python
# batch_inference_utils.py — wait_for_job_completion
while status not in ['Completed', 'Failed', 'Stopped', 'PartiallyCompleted', 'Expired']:
time.sleep(60)
response = bedrock_client.get_model_invocation_job(jobIdentifier=job_arn)
status = response['status']
```
- **Bedrock has a 100-record-per-job minimum** (`BEDROCK_MIN_BATCH_SIZE = 100`) and the code
merges any sub-100 tail into the previous job. So a small corpus does **not** produce
smaller/faster jobs — it produces the *same single fully-overhead job*.
- Rounds are **sequential**: with `BatchExtractToS3`'s `extraction_batch_size=100` (documents)
the 70-doc corpus is one round; auto-tune uses `round_capacity = num_workers * max_batch_size`
(`2 * 250 = 500`), also one round. Within a round, jobs run concurrently, so total wall-clock
≈ *number of sequential rounds × (Bedrock startup overhead + 60s poll granularity)*.

**Net:** each of the two extract tests is ≈ one Bedrock batch-job lifetime, and we run two of
them one after another. Total runtime is therefore driven by (a) doing the extract twice
serially and (b) fixed Bedrock/polling overhead — **not** by the 70-document count.

### Proposed solution

Ranked by impact-to-effort. (1) is the immediate, mechanical win and is being raised as a
companion PR.

### ~~1. Split the suite so the two extract steps run in parallel (highest impact, low effort)~~ — ✅ Addressed in #486 (see comment below)

`batch_extract.BatchExtractAutoTuneToS3` is **standalone**: it writes
`params['auto_tune_batch_collection_id']`, which is **never consumed** by any downstream test
(only `batch_extract.BatchExtractToS3`'s `batch_collection_id` /
`multihop_expected_num_batch_docs` feed `batch_build.BuildFromS3` and the queries). So the two
extract runs are fully independent.

Split `lexical.long` into two suite files that can run as **two parallel CloudFormation
stacks / CI jobs**:

- `lexical.long` → `BatchExtractToS3` + `BuildFromS3` + queries + progress monitor
- `lexical.autotune.long` → `BatchExtractAutoTuneToS3` (standalone)

Because both extracts are on the critical path today and each is ≈ one batch-job lifetime,
running them in parallel roughly **halves** the wall-clock of the extract phase
(~80 min serial → ~40 min, bounded by the slower of the two).

### 2. Make the batch-inference poll interval configurable and short in tests (low effort)

The hard-coded `time.sleep(60)` (sleep-*before*-first-check) adds a full minute of latency per
poll and at least one minute even for a job that is already done. For the small test corpus this
is pure overhead. Expose the poll interval on `BatchConfig` (e.g. `poll_interval_seconds`,
default 60) and set it low (e.g. 10–15s) in the integration tests. Requires a small
`lexical-graph` change.

### 3. Decouple `build`/`query` tests from batch extract with a cached "dry-run" fixture (medium effort)

Run the batch extraction **once**, save the extracted S3 docs, and commit/stage them as a
fixture. Point `batch_build.BuildFromS3` (and the multi-hop queries) at the fixture so they no
longer wait on Bedrock batch inference on every run. Keep a **small smoke corpus** (a handful of
docs, still ≥ the 100-record floor) to validate the batch-extract *wiring* end-to-end. This
removes the slow, non-deterministic Bedrock step from routine build/query regression runs and is
the biggest structural lever for day-to-day CI cycle time.

### 4. Reduce `InferClassificationsConfig` iterations in the tests (low effort, needs measurement)

Both extract tests use `InferClassificationsConfig(num_samples=5, num_iterations=10)`. If
classification inference issues real-time (non-batch) LLM calls on top of the batch job, cutting
`num_iterations` (e.g. to 3) for the test config would trim time without changing what the batch
pipeline exercises. Measure first to confirm the contribution.

### Alternatives considered

## Alternatives considered

- **Halve the test set (70 → 35 docs).** This is intuitive but yields **little batch-extract
savings**: because of the 100-record-per-job minimum and the sub-100 tail-merge, 35 docs still
produce a single fully-overhead Bedrock job, so the fixed per-job wall-clock is unchanged.
Fewer docs *does* speed up the downstream **build** and **query** steps, and reduces cost, so
it is worth doing for those reasons — but it is not the fix for the extract bottleneck, and it
slightly weakens the "one extracted doc per source URL" assertion's coverage.

- **Increase parallelism knobs (`max_num_concurrent_batches`, `extraction_num_workers`,
`max_batch_size`).** These help **large** corpora by packing/parallelizing jobs, but for the
70-doc test corpus everything already fits in a **single round / job per worker**, so the
bottleneck is fixed per-job overhead + the 60s poll — not throughput. Raising these values
will not meaningfully speed up the current test corpus.

- **Drop the auto-tune extract test entirely.** Rejected — it validates the auto-tuning code
path (job packing / round sizing) added in #469. Splitting it into `lexical.autotune.long`
(proposal 1) keeps the coverage while removing it from the critical path. Its corpus could
additionally be shrunk since its output is never built downstream — its only assertion is a
doc-count round-trip.

## Appendix — key references

| Concern | Location |
|---|---|
| Suite definition | `integration-tests/lexical.long` |
| Extract tests | `integration-tests/test-scripts/graphrag_toolkit_tests/batch_extract.py` |
| Downstream consumer of `batch_collection_id` | `integration-tests/test-scripts/graphrag_toolkit_tests/batch_build.py` |
| Sequential test loop | `integration-tests/test-scripts/test_suite.py` |
| 60s poll / min-records | `lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/utils/batch_inference_utils.py` |
| `BatchConfig` defaults | `lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/batch_config.py` |
| Auto-tune round packing | `lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/bucket_filler.py` |

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.