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

[acarbonetto]

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:

  # 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 [FEATURE] (lexical-graph): Auto-tuning batch extraction #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

[mykola-pereyma]

✅ Addressed by #486 (merged)

Proposal (1) — removing the redundant back-to-back extract from lexical.long — is now done (struck through in the description above). What #486 shipped:

• lexical.long now runs BatchExtractAutoTuneToS3 + BuildFromS3 (replacing BatchExtractToS3). Rather than splitting into two parallel suites, it consolidates to a single auto-tune extract that also feeds the build, which removes the second, back-to-back extract that dominated the runtime.
• The auto-tune extract was fixed to match the build contract: writes to the …/extracted prefix and publishes params['batch_collection_id'] / params['multihop_expected_num_batch_docs'] (previously the orphaned auto_tune_batch_collection_id), so BuildFromS3 and its assertion consume it.
• Added opt-in BENCHMARK_EXTRACT_DOC_LIMIT (env / build-tests.sh --benchmark-extraction-doc-limit N) to cap extracted docs for dev/smoke runs (unset = unchanged).
• Removed BatchExtractSmallBatchToS3.

Note: the implementation differs from the proposal as written — there is no lexical.autotune.long, and standalone BatchExtractToS3 is no longer referenced by lexical.long (the class still exists). If we want the plain (non-auto-tune) extract path exercised in CI, that needs a separate decision.

Still open on this issue (not addressed by #486):

• (2) configurable / shorter batch-poll interval — wait_for_job_completion still hard-codes time.sleep(60) before the first check and BatchConfig has no poll_interval_seconds.
• (3) cached dry-run fixture to decouple BuildFromS3/queries from live Bedrock.
• (4) reduce InferClassificationsConfig(num_iterations=10) in the test config (measure first).