awslabs / awslabs/graphrag-toolkit
[FEATURE] Configurable write pacing for large-scale graph ingestion
- Dominant language
- Python
- Stars
- 442
- Forks
- 106
- Avg merge
- 2d 3h
- Merged PRs (30d)
- 41
Description
> **⚠️ DRAFT — This issue is under discussion and not yet ready for implementation.**
Related: #317 (ProgressMonitor interface covers progress reporting), #325 (DocumentStateStore covers resume for both phases)
## Problem Statement
The two-phase extraction-to-build architecture already exists:
```python
# Phase 1: Extract → S3 (already works via S3BasedDocs handler)
graph_index.extract(docs, handler=s3_docs, show_progress=True)
# Phase 2: Build from S3 (already works independently)
graph_index.build(s3_docs, show_progress=True)
```
**What's missing** is write pacing during Phase 2 (graph build):
`GraphBatchClient` retries on failure (`BATCH_MAX_ATTEMPTS=10`) but doesn't proactively throttle. At scale, Neptune DB and OpenSearch can return 429s faster than backoff recovers, causing cascading failures. (Neptune Analytics has provisioned throughput and is less affected; this primarily impacts Neptune DB + OpenSearch deployments.)
**Note:** Build-phase document-level resume is now handled by `DocumentStateStore` in #325 — this issue focuses solely on write pacing.
### What Already Exists (NOT to be reimplemented)
- Two-phase pipeline: `extract()` → S3 → `build()` ✅
- `S3BasedDocs` as intermediate store (upload + download) ✅
- Per-query retry with `execute_query_with_retry` (max_attempts=5, max_wait=7) ✅
- OpenSearch bulk retry (`max_retries=3, initial_backoff=5`) ✅
- `Checkpoint`/`CheckpointFilter` for per-node savepoints ✅
- Progress reporting: covered by #317 ✅
- Document-level resume (both phases): covered by #325 ✅
## Proposed Solution
**Write pacing** — configurable max writes/second to Neptune. Token bucket algorithm in `GraphBatchClient` to prevent overwhelming the graph store:
```python
GraphRAGConfig.build_max_writes_per_second = 500 # None = unlimited (current behavior)
```
In-memory token bucket — no external storage dependency. Pure concurrency control.
### Performance Impact
| Metric | Current | Proposed |
|--------|---------|----------|
| Neptune 429 cascade at scale | Retry storm → eventual failure | Proactive pacing → steady throughput |
| Sustained write throughput (1M docs) | Spiky (burst → throttle → backoff) | Smooth (paced at sustainable rate) |
**Why pacing improves throughput:** Reactive retry with `wait_random(0, 7s)` causes multiple workers to simultaneously retry after a throttle event, creating burst patterns that re-trigger throttling. Proactive pacing at a sustainable rate eliminates these bursts entirely, yielding higher sustained throughput over time.
### Acceptance Criteria
- [ ] Configurable write rate limit for graph store operations (`build_max_writes_per_second`)
- [ ] Token bucket implementation in `GraphBatchClient` (in-memory, no storage dep)
- [ ] Backward compatible — existing `build()` works unchanged (`None` = unlimited writes)
- [ ] Measured: paced writes show fewer 429s and higher sustained throughput than unlimited writes at scale
## Alternatives Considered
- **Current reactive retry** (`wait_random(0, 7s)`, 10 attempts) — works for low-volume writes but creates burst-throttle cycles at 1M+ document scale.
- **Reduce `num_workers`** — limits concurrency but doesn't solve the fundamental lack of admission control; individual workers can still burst.
Contributor guide
Assessment
This issue has not been assessed yet.