awslabs / awslabs/graphrag-toolkit

Feature: S3-backed Checkpoint for serverless/containerized environments

Open
#401 1 comment 0 reactions 0 assignees 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 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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.