Configuring Batch Extraction

Topics

• Overview
• BatchConfig parameters
  • Required parameters
    • bucket_name
    • region
    • role_arn
  • Optional parameters
    • key_prefix
    • max_batch_size
    • max_num_concurrent_batches
    • s3_encryption_key_id
  • VPC security parameters (optional)
    • subnet_ids
    • security_group_ids
  • File management
    • delete_on_success
  • Auto-tuning
    • auto_tune
• Optimizing batch extraction performance
• Auto-tuning batch extraction

Overview

BatchConfig parameters

The BatchConfig object manages the configuration settings for Amazon Bedrock batch inference jobs. Here’s a detailed explanation of each parameter:

Required parameters

bucket_name

You must specify the name of an Amazon S3 bucket where your batch processing files (both input and output) will be stored.

region

You need to provide the AWS Region name (such as “us-east-1”) where both your S3 bucket is located and where the Amazon Bedrock batch inference job will run.

role_arn

This is the Amazon Resource Name (ARN) for the service role that handles batch inference operations. You can either create a default service role through the console or follow the instructions in the Create a service role for batch inference documentation.

Optional parameters

key_prefix

If desired, you can specify an S3 key prefix for organizing your input and output files.

max_batch_size

Controls how many records (chunks) can be included in each batch inference job. The default value is 25000 records.

max_num_concurrent_batches

Determines how many batch inference jobs can run simultaneously per worker. This setting works in conjunction with GraphRAGConfig.extraction_num_workers. The default is 3 concurrent batches per worker.

s3_encryption_key_id

You can provide the unique identifier for an encryption key to secure the output data in S3.

VPC security parameters (optional)

For more information about VPC protection, see Protect batch inference jobs using a VPC.

subnet_ids

An array of subnet IDs within your Virtual Private Cloud (VPC) for protecting batch inference jobs.

security_group_ids

An array of security group IDs within your VPC for protecting batch inference jobs.

File management

delete_on_success

Controls whether input and output JSON files are automatically deleted from the local filesystem after successful batch job completion. By default, this is set to True. Note that this setting does not affect files stored in S3, which are preserved regardless.

Auto-tuning

auto_tune

When set to True, enables auto-tuning batch extraction. Instead of pulling a fixed number of source documents per iteration, the framework streams and chunks documents incrementally, filling GraphRAGConfig.extraction_num_workers buckets up to max_batch_size chunks each before submitting a round of batch jobs. For small inputs it consolidates chunks into the fewest jobs needed. With auto-tuning enabled you only need to specify extraction_num_workers and max_batch_size; extraction_batch_size becomes a derived value (if you set it explicitly, it acts as an optional upper bound on the number of documents consumed per round). The default is False.

Coordination with parallelization: under auto-tuning, extraction_num_workers sets how many batch jobs a round submits concurrently. max_num_concurrent_batches (the per-worker job concurrency of the fixed-batch path) does not apply, because each round job is a single batch. Document chunking runs in the main process so chunk counts are known before a round is packed; only extraction runs concurrently. Keep extraction_num_workers within Bedrock’s quota of 20 combined in-progress and submitted batch jobs per region.

Optimizing batch extraction performance

The most important settings for controlling batch extraction performance are:

• GraphRAGConfig.extraction_batch_size: Sets how many source documents go to the extraction pipeline. When calculating this value, consider that the total number of chunks (source documents × average chunks per document) should be sufficient to fill your planned simultaneous batch jobs. If you enable auto-tuning, this value is determined dynamically and you no longer need to calculate it; setting it explicitly only caps the number of documents consumed per round.
• GraphRAGConfig.extraction_num_workers: Sets how many CPUs run batch jobs simultaneously.
• BatchConfig.max_num_concurrent_batches: Sets how many concurrent batch jobs each worker runs.
• BatchConfig.max_batch_size: Sets the maximum number of chunks per batch job.

To maximize the efficiency of batch extraction, follow these three key principles:

• Maximize file capacity Each batch job file can hold up to 50,000 records. However, Amazon Bedrock enforces input file size limits, typically between 1-5 GB. Check the specific limits for your model in the Amazon Bedrock service quotas section (see the Batch inference job size quotas in the Amazon Bedrock service quotas section for the limits particular to the model you are using). Note that the toolkit doesn’t automatically verify file sizes, so jobs may fail if they exceed these quotas. You may need to use fewer records than the maximum limit to stay within file size boundaries. Configure the BatchConfig.max_batch_size to set the maximum number of records per batch job.
• Use larger, fewer files Focus on using a minimal number of large files rather than splitting the work across many smaller ones. For example, it’s more efficient to process 40,000 records in a single job than to divide them into four parallel jobs of 10,000 records each.
• Leverage parallel processing Take advantage of parallel job execution using GraphRAGConfig.extraction_num_workers and BatchConfig.max_num_concurrent_batches. The total number of jobs (number of workers × number of concurrent batches) must stay within Bedrock’s quota of 20 combined in-progress and submitted batch inference jobs per region. If you exceed this limit, additional jobs will wait in the queue until capacity becomes available.

Each batch inference job submitted to Amazon Bedrock must contain at least 100 records. If a batch job contains fewer than 100 records, batch extraction falls back to performing chunk-by-chunk extraction. (Batch extraction will also fall back to performing chunk-by-chunk extraction if there are a last few records left over after processing one or more batches containing max_batch_size records.) The batch extraction process issues a warning whenever it begins processing a set of records chunk-by-chunk.

Auto-tuning batch extraction

The performance guidance above requires you to reason about the relationship between extraction_batch_size (source documents pulled per iteration), the average number of chunks per document, extraction_num_workers, and max_batch_size in order to keep batch jobs well-utilized. Getting this wrong produces either many small, under-utilized jobs (slow, because of per-job overhead and the sub-100-record fallback) or jobs that need to be re-split.

Auto-tuning removes this manual calculation. Set auto_tune=True on your BatchConfig:

batch_config = BatchConfig(
    role_arn=batch_inference_role,
    region=aws_region_name,
    bucket_name=s3_results_bucket,
    max_batch_size=25000,
    max_num_concurrent_batches=3,
    auto_tune=True,
)

With auto-tuning enabled, the framework:

• Streams and chunks documents incrementally rather than materializing a fixed extraction_batch_size of documents up front.
• Accumulates chunks in document order and packs them contiguously into jobs of max_batch_size — the same count-based packing the fixed-batch_size path already uses — so every job except the last of a round is filled to max_batch_size.
• Submits a round of up to extraction_num_workers jobs once the accumulated chunks reach extraction_num_workers × max_batch_size (or before a document would overshoot that capacity), then resumes pulling documents for the next round.
• Consolidates on exhaustion: the remaining chunks are sliced into the minimum number of jobs (ceil(remaining / max_batch_size), at most extraction_num_workers). If the final remainder is below the 100-record Bedrock minimum, it is merged into the previous round’s last job (which may then slightly exceed max_batch_size, by at most 99 records, exactly as the fixed-batch_size path does) so it is still batched — unless that merge would push the job past Bedrock’s hard per-job record limit, in which case the tail is left as its own round. Only a corpus small enough that it never fills a single round falls back to chunk-by-chunk extraction.

Under auto-tuning you only need to specify extraction_num_workers and max_batch_size. extraction_batch_size becomes a derived value; if you set it explicitly it is treated as an optional upper bound on the number of documents consumed per round (useful for bounding memory). The extracted output is equivalent to the fixed-batch_size path for the same input — auto-tuning changes how jobs are sized and submitted, not the extraction results.

Concurrency under auto-tuning. Auto-tuning submits up to extraction_num_workers batch jobs per round (each up to max_batch_size records) and treats extraction_num_workers as the unit of job concurrency. BatchConfig.max_num_concurrent_batches — which governs per-worker job concurrency in the fixed-batch_size path — is not used when auto-tuning is enabled. Size extraction_num_workers with Bedrock’s per-region quota on in-progress and submitted batch jobs in mind.

Document integrity and large documents. A document’s chunks may be split across job boundaries (and, for very large documents, across rounds), exactly as in the fixed-batch_size path. Output integrity is preserved by source_id-keyed storage, which reassembles each document from its chunks regardless of how they were batched — so splitting does not fragment results. The framework still flushes a round before a document would overshoot capacity, keeping a document within one round when it fits. If a single document produces more chunks than a full round’s capacity (extraction_num_workers × max_batch_size), it is split across rounds and the framework logs a warning; increase max_batch_size or extraction_num_workers to keep such documents within a single round.