awslabs / awslabs/graphrag-toolkit
Feature: S3-backed Checkpoint for serverless/containerized environments
- Dominant language
- Python
- Stars
- 442
- Forks
- 106
- Avg merge
- 2d 3h
- Merged PRs (30d)
- 41
Description
### Package
lexical-graph
### Problem statement
The existing `Checkpoint` class (`graphrag_toolkit.lexical_graph.indexing.build.checkpoint`) uses local filesystem operations (`os.path.exists`, `os.makedirs`, `open`) to store checkpoint markers. This means checkpointing **does not work** in serverless or containerized environments where local disk is ephemeral:
- **ECS Fargate** — container restarts lose `/tmp` contents
- **AWS Lambda** — ephemeral `/tmp` cleared between invocations
- **EKS pods** — ephemeral storage unless PVC is configured
- **SQS-driven workers** — message retry may land on a different container
The toolkit already supports S3 for document storage (`S3BasedDocs`), vector storage (`S3VectorIndexes`), and batch inference I/O (`BatchConfig`). The checkpoint is the **only** component that cannot run without persistent local disk.
### Proposed solution
An `S3Checkpoint` class that provides the **same interface** as `Checkpoint` (`add_filter`, `add_writer`, `enabled`) but uses zero-byte S3 objects as markers instead of local files.
```python
from graphrag_toolkit.lexical_graph.indexing.build import S3Checkpoint
checkpoint = S3Checkpoint(
checkpoint_name="enrichment-run-001",
bucket_name="my-pipeline-bucket",
key_prefix="checkpoints/tenant-a"
)
# Same API as Checkpoint — drop-in replacement
graph_index.extract(docs, handler=extracted_docs, checkpoint=checkpoint)
graph_index.build(docs, checkpoint=checkpoint)
```
### Mapping (local → S3)
| Local `Checkpoint` | `S3Checkpoint` equivalent |
|---|---|
| `os.path.exists(path)` | `s3.head_object()` — catch 404 |
| `open(path, "a"); os.utime(path)` | `s3.put_object(Body=b"")` |
| `os.makedirs(dir)` | No-op (S3 has no directories) |
### Features
| Feature | Description |
|---|---|
| **Drop-in interface** | Same `add_filter` / `add_writer` / `enabled` as `Checkpoint` |
| **Zero-byte markers** | Minimal cost (HEAD + PUT, no data transfer) |
| **Survives restarts** | Markers persist in S3 across container restarts |
| **Safe defaults** | On S3 errors, includes node (re-processes rather than skips) |
| **Non-fatal writes** | Marker write failure logs warning, does not fail the job |
| **Multi-tenancy** | Uses `TenantId.rewrite_id()` same as local checkpoint |
| **No new dependencies** | Uses `boto3` already in the project via `GraphRAGConfig.s3` |
### Test coverage
15 unit tests covering:
- Initialization (3 tests)
- add_filter wrapping/skipping (3)
- add_writer wrapping/skipping (2)
- S3CheckpointFilter — HEAD 404 / exists / unexpected error / call filtering (5)
- S3CheckpointWriter — touch writes / failure non-fatal / accept yields + markers / skips INDEX_KEY nodes (4)
### Implementation
Branch: [`feat/s3-checkpoint`](https://github.com/evanerwee/graphrag-toolkit/tree/feat/s3-checkpoint)
Files:
- `lexical-graph/src/.../indexing/build/s3_checkpoint.py` (305 lines)
- `lexical-graph/src/.../indexing/build/__init__.py` (+1 line — export)
- `lexical-graph/tests/unit/indexing/build/test_s3_checkpoint.py` (327 lines)
### Alternatives considered
1. **EFS mount for Fargate** — adds infrastructure cost + complexity for a simple marker file
2. **DynamoDB-backed checkpoint** — heavier dependency, overkill for boolean markers
3. **S3 zero-byte objects (this proposal)** — minimal, same pattern as local, no new deps
Contributor guide
Research direction
Review the existing Checkpoint implementation and the proposed lexical-graph/src/.../indexing/build/s3_checkpoint.py, along with its export in __init__.py. Run lexical-graph/tests/unit/indexing/build/test_s3_checkpoint.py first to understand the 15 stated cases. Done means the S3 implementation preserves the Checkpoint interface and its marker, error, tenant, and INDEX_KEY behaviors are covered by passing tests.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- aws, python
- Domain
- backend, cloud
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100