awslabs / awslabs/graphrag-toolkit

[RFC] Local Execution Mode for Integration Tests

Open
#229 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

# RFC: Local Execution Mode for Integration Tests

## Summary

Enable developers to run integration tests locally against Docker-hosted services (Neo4j, PostgreSQL+pgvector) instead of requiring a full CloudFormation deployment to AWS. Amazon Bedrock remains the only required AWS service (accessed via local `~/.aws/credentials`).

_Aside: add more docker-hosted services (OpenSearch-KNN, Valkey-Search)_

## Motivation

The current integration test pipeline requires deploying a CloudFormation stack that provisions Neptune, AOSS, and a SageMaker notebook for every test run. This creates a slow feedback loop (stack creation alone takes ~15 minutes), incurs significant AWS costs, and prevents developers from iterating quickly on test changes.

The graphrag-toolkit library already natively supports Neo4j (`bolt://`) and PostgreSQL (`postgresql://`) connection strings. The test classes themselves have no hard dependency on SageMaker or CloudFormation — the AWS coupling is entirely in the runner layer (`test_suite.py` and `IntegrationTestHandler`). This means **24 of 35 tests can run locally with changes only to the orchestration code**.

## Proposal

### What Changes

1. **New local test runner** (`test_suite_local.py`) that bypasses CloudFormation polling, S3 result uploads, and SNS notifications.
2. **New local test handler** (`local_test_handler.py`) that writes results to the local filesystem only.
3. **Docker Compose file** for Neo4j and PostgreSQL+pgvector.
4. **Local test suite files** (`lexical.local.short`, `lexical.local.versioning`, etc.).
5. **Local environment template** (`env.local.template`).
6. **Entry-point script** (`run-local-tests.sh`).

### What Doesn't Change

- All existing test classes (`IntegrationTestBase` subclasses) remain untouched.
- Any test requiring cloud services will will excluded from test runs until an appropriate alternative is included.
- The existing CloudFormation-based pipeline continues to work as-is.
- No changes to the graphrag-toolkit library itself (Phase 1–2).

## Design

### AWS Dependency Analysis

| Dependency | Used By | Local Strategy |
|-----------|---------|----------------|
| CloudFormation (stack polling) | `test_suite.py` | **Remove** — skip in local runner |
| S3 (result uploads) | `test_suite.py`, `IntegrationTestHandler` | **Replace** — local filesystem only |
| SNS (notifications) | `test_suite.py` | **Remove** — not needed |
| Bedrock (LLM + embeddings) | All extraction/query tests | **Keep** — via local AWS credentials |
| Neptune DB / Analytics | `GRAPH_STORE` env var | **Replace** — Neo4j via Docker |
| OpenSearch Serverless | `VECTOR_STORE` env var | **Replace** — PostgreSQL+pgvector via Docker |
| S3 (document storage) | `S3BasedDocs`, `BatchConfig` | **Cannot replace** — skip affected tests |
| Bedrock Batch Inference | `BatchConfig` | **Cannot replace** — skip affected tests |
| Neptune Analytics (BYOKG) | `byokg_setup.py` | **Cannot replace** — skip |

### Connection Strings Already Supported by graphrag-toolkit

| Type | Connection String | Local? |
|------|------------------|--------|
| Neo4j graph store | `bolt://host:port` or `neo4j://host:port` | ✅ |
| PostgreSQL vector store | `postgresql://user:pass@host:port/db` | ✅ |
| Dummy graph store | `dummy://` | ✅ (no-op) |
| Dummy vector store | `dummy://` | ✅ (no-op) |
| OpenSearch Serverless | `aoss://endpoint` | ❌ (hardcoded AWS SigV4 auth) |

### OpenSearch Limitation

The OpenSearch client in `opensearch_vector_indexes.py` hardcodes `service = 'aoss'` with `Urllib3AWSV4SignerAuth`. This makes `aoss://` incompatible with local OpenSearch. **Recommendation:** use PostgreSQL+pgvector for Phase 1–2. A new `opensearch://` connection string with basic auth can be added in Phase 3.

### Local Configuration Profiles

**Profile 1 — Neo4j + PostgreSQL (recommended):**
```bash
export GRAPH_STORE="bolt://localhost:7687"
export VECTOR_STORE="postgresql://postgres:password@localhost:5432/graphrag"
export TEST_EXTRACTION_LLM="us.anthropic.claude-sonnet-4-20250514-v1:0"
export TEST_RESPONSE_LLM="us.anthropic.claude-sonnet-4-20250514-v1:0"
export AWS_REGION_NAME="us-east-1"
export FAIL_FAST="True"
```

**Profile 2 — Extraction only (no local services):**
```bash
export GRAPH_STORE="dummy://"
export VECTOR_STORE="dummy://"
export TEST_EXTRACTION_LLM="us.anthropic.claude-sonnet-4-20250514-v1:0"
export TEST_RESPONSE_LLM="us.anthropic.claude-sonnet-4-20250514-v1:0"
export AWS_REGION_NAME="us-east-1"
```

**Docker Compose:**
```yaml
services:
neo4j:
image: neo4j:5
ports:
- "7474:7474"
- "7687:7687"
environment:
NEO4J_AUTH: none
postgres:
image: pgvector/pgvector:pg16
ports:
- "5432:5432"
environment:
POSTGRES_PASSWORD: password
POSTGRES_DB: graphrag
```

## Local Test Suite Files

### `lexical.local.short` — 19 tests

Local equivalent of `lexical.short`. Excludes GPU, batch, and S3-dependent tests.

```
local_entities.BuildWithLocalEntities
local_entities.BuildWithoutLocalEntities
extract.ExtractToFileSystem
build.BuildFromFileSystem
build_facts.BuildFacts
query.TraversalBasedQuery
query.MetadataFilteringQuery
query.TraversalBasedQueryWithModelReranker
query.TraversalBasedQueryWithBedrockReranker
query.ChunkBasedTraversalQuery
query.SemanticGuidedQuery
query.SemanticGuidedQueryWithSubRetrievers
query.SemanticGuidedRerankingBeamSearchQuery
query.SemanticGuidedQueryWithPostProcessors
extract_and_build.ExtractAndBuild
checkpoint.ExtractWithCheckpoint
checkpoint.BuildWithCheckpoint
falkordb.TestFalkorDBContrib
delete_sources.DeleteSourceDocs
```

Requires: Neo4j + PostgreSQL + Bedrock credentials + internet access.

### `lexical.local.versioning` — 5 tests

Identical to `lexical.versioning` — all tests work locally.

```
versioning.CreateVersionedData
versioning.QueryVersionedData
versioning.DeleteVersionedData
versioning.AutoDeleteAllPrevVersionsData
versioning.DoNotAutoDeleteProtectedPrevVersionsData
```

Requires: Neo4j + PostgreSQL + Bedrock credentials + internet access.

### `lexical.local.extract` — 2 tests

Extraction-only. No graph or vector store needed (`dummy://`).

```
extract.ExtractToFileSystem
checkpoint.ExtractWithCheckpoint
```

Requires: Bedrock credentials + internet access only.

### `lexical.local.build` — 6 tests

Build from pre-extracted data. No Bedrock needed for the build step itself.

```
local_entities.BuildWithLocalEntities
local_entities.BuildWithoutLocalEntities
build.BuildFromFileSystem
build_facts.BuildFacts
checkpoint.BuildWithCheckpoint
delete_sources.DeleteSourceDocs
```

Requires: Neo4j + PostgreSQL. Note: `build.BuildFromFileSystem` depends on `extract.ExtractToFileSystem` having populated `params['collection_id']`. For standalone build testing, use `local_entities` and `build_facts` which read from `source-data/` directly.

## Tests Excluded from Local Execution

| Test | Cloud Suite | Reason |
|------|-------------|--------|
| `query_gpu.RerankingBeamGraphSearchGPU` | `lexical.short` | Requires GPU hardware |
| `batch_fallback.BatchExtractToS3Fallback` | `lexical.short` | Requires S3 + Bedrock batch inference |
| `batch_extract.BatchExtractToS3` | `lexical.long` | Requires S3 + Bedrock batch inference |
| `batch_extract.BatchExtractFromS3ToS3` | `lexical.long` | Requires S3 + Bedrock batch inference |
| `batch_build.BuildFromS3` | `lexical.long` | Requires `S3BasedDocs` |
| `query.MultiHopQuery` | `lexical.long` | Depends on `BuildFromS3` output |
| `query.ChunkBasedSemanticSearchMultiHopQuery` | `lexical.long` | Depends on `BuildFromS3` output |
| `byokg_setup.LoadBYOKGGraph` | `byokg.all` | Hardcoded to Neptune Analytics |
| `benchmark_build.CuadBenchmarkBuild` | `benchmark.cuad.prototype` | Could be adapted in Phase 4 |
| `benchmark_query.CuadBenchmarkQuery` | `benchmark.cuad.prototype` | Depends on benchmark build |
| `benchmark_evaluate.CuadBenchmarkEvaluate` | `benchmark.cuad.prototype` | Depends on benchmark query |

## Implementation Plan

### Phase 1: Minimal Local Runner (Extraction Only)

- [ ] Create `test-scripts/test_suite_local.py` — fork of `test_suite.py` that removes CloudFormation polling (`wait_till_stack_complete`), S3 result publishing (`publish_test_run_metadata`), SNS notifications (`notify`), and stack deletion (`delete_stack`). Make `S3_RESULTS_BUCKET`, `S3_RESULTS_PREFIX`, and `STACK_ID` optional.
- [ ] Create `test-scripts/local_test_handler.py` — subclass of `IntegrationTestHandler` that writes results to local `test-results/` only (skip `boto3.client('s3').upload_file()` in `__exit__`). Keep all assertion and timing logic.
- [ ] Create `env.local.template` with required variables: `GRAPH_STORE`, `VECTOR_STORE`, `TEST_EXTRACTION_LLM`, `TEST_RESPONSE_LLM`, `AWS_REGION_NAME`.
- [ ] Create `run-local-tests.sh` — sources `.env.local`, creates `test-results/` and `test-logs/` dirs, installs deps from local graphrag-toolkit clone, runs `test_suite_local.py`.
- [ ] Create `lexical.local.extract` suite file.
- [ ] Validate with `GRAPH_STORE=dummy://` and `VECTOR_STORE=dummy://` using only Bedrock credentials.

### Phase 2: Neo4j + PostgreSQL (Full Local Stack)

- [ ] Create `docker-compose.yml` with Neo4j 5 and PostgreSQL 16 + pgvector.
- [ ] Create `lexical.local.short`, `lexical.local.versioning`, and `lexical.local.build` suite files.
- [ ] Test full extract → build → query pipeline with Neo4j + PostgreSQL.
- [ ] Test the complete `lexical.local.short` suite (19 tests).
- [ ] Test `lexical.local.versioning` suite (5 tests).
- [ ] Document any Neo4j Cypher compatibility issues (double-underscore label prefixes like `__Source__`, `__Entity__`).

### Phase 3: Local OpenSearch (Optional)

- [ ] Add OpenSearch to `docker-compose.yml` (`opensearchproject/opensearch:2`).
- [ ] Add `LocalOpenSearchVectorIndexFactory` to graphrag-toolkit accepting `opensearch://host:port` with basic auth.
- [ ] Register new factory in `VectorStoreFactory`.
- [ ] Test with `VECTOR_STORE=opensearch://localhost:9200`.

### Phase 4: Benchmark Tests (Optional)

see: #228

- [ ] Adapt `CuadBenchmarkBuild` to work with local `FileBasedDocs`.
- [ ] Test benchmark build → query → evaluate pipeline locally.

### Phase 5: Documentation

- [ ] Add local testing instructions to `README.md`.
- [ ] Update `AGENTS.md` with local testing patterns.
- [ ] Consider GitHub Actions workflow for local test suite on PRs.

## Alternatives Considered

1. **LocalStack for S3/CloudFormation** — Would allow running batch tests and S3-based tests locally, but adds significant complexity and doesn't help with Neptune/OpenSearch Serverless. Not worth the overhead for the marginal test coverage gain.

2. **Modifying existing `test_suite.py` with a `--local` flag** — Simpler than a separate file, but risks breaking the production pipeline. A separate runner is safer and clearer.

3. **Starting with OpenSearch instead of PostgreSQL** — The OpenSearch client is hardcoded to AWS SigV4 auth, requiring changes to graphrag-toolkit itself. PostgreSQL+pgvector works out of the box.

Contributor guide

Open the contributing guide

Research direction

Start by reading test-scripts/test_suite.py and the IntegrationTestHandler flow, then compare the proposed responsibilities for test_suite_local.py and local_test_handler.py. Review lexical.local.extract and env.local.template before attempting the first phase; done means a scoped local extraction runner works with dummy stores, local result output, and the documented Bedrock credentials.

Written by the indexing model from the issue text.

Assessment

Tech stack
aws, docker-compose, neo4j, postgresql, python, shell
Domain
cloud, databases, devops, testing-qa
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.