diff --git a/docs/tech-specs/graphrag-performance-optimization.md b/docs/tech-specs/graphrag-performance-optimization.md new file mode 100644 index 00000000..da6e1667 --- /dev/null +++ b/docs/tech-specs/graphrag-performance-optimization.md @@ -0,0 +1,629 @@ +# GraphRAG Performance Optimisation Technical Specification + +## Overview + +This specification describes comprehensive performance optimisations for the GraphRAG (Graph Retrieval-Augmented Generation) algorithm in TrustGraph. The current implementation suffers from significant performance bottlenecks that limit scalability and response times. This specification addresses four primary optimisation areas: + +1. **Graph Traversal Optimisation**: Eliminate inefficient recursive database queries and implement batched graph exploration +2. **Label Resolution Optimisation**: Replace sequential label fetching with parallel/batched operations +3. **Caching Strategy Enhancement**: Implement intelligent caching with LRU eviction and prefetching +4. **Query Optimisation**: Add result memoisation and embedding caching for improved response times + +## Goals + +- **Reduce Database Query Volume**: Achieve 50-80% reduction in total database queries through batching and caching +- **Improve Response Times**: Target 3-5x faster subgraph construction and 2-3x faster label resolution +- **Enhance Scalability**: Support larger knowledge graphs with better memory management +- **Maintain Accuracy**: Preserve existing GraphRAG functionality and result quality +- **Enable Concurrency**: Improve parallel processing capabilities for multiple concurrent requests +- **Reduce Memory Footprint**: Implement efficient data structures and memory management +- **Add Observability**: Include performance metrics and monitoring capabilities +- **Ensure Reliability**: Add proper error handling and timeout mechanisms + +## Background + +The current GraphRAG implementation in `trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py` exhibits several critical performance issues that severely impact system scalability: + +### Current Performance Problems + +**1. Inefficient Graph Traversal (`follow_edges` function, lines 79-127)** +- Makes 3 separate database queries per entity per depth level +- Query pattern: subject-based, predicate-based, and object-based queries for each entity +- No batching: Each query processes only one entity at a time +- No cycle detection: Can revisit the same nodes multiple times +- Recursive implementation without memoisation leads to exponential complexity +- Time complexity: O(entities × max_path_length × triple_limit³) + +**2. Sequential Label Resolution (`get_labelgraph` function, lines 144-171)** +- Processes each triple component (subject, predicate, object) sequentially +- Each `maybe_label` call potentially triggers a database query +- No parallel execution or batching of label queries +- Results in up to 3 × subgraph_size individual database calls + +**3. Primitive Caching Strategy (`maybe_label` function, lines 62-77)** +- Simple dictionary cache without size limits or TTL +- No cache eviction policy leads to unbounded memory growth +- Cache misses trigger individual database queries +- No prefetching or intelligent cache warming + +**4. Suboptimal Query Patterns** +- Entity vector similarity queries not cached between similar requests +- No result memoisation for repeated query patterns +- Missing query optimisation for common access patterns + +**5. Critical Object Lifetime Issues (`rag.py:96-102`)** +- **GraphRag object recreated per request**: Fresh instance created for every query, losing all cache benefits +- **Query object extremely short-lived**: Created and destroyed within single query execution (lines 201-207) +- **Label cache reset per request**: Cache warming and accumulated knowledge lost between requests +- **Client recreation overhead**: Database clients potentially re-established for each request +- **No cross-request optimisation**: Cannot benefit from query patterns or result sharing + +### Performance Impact Analysis + +Current worst-case scenario for a typical query: +- **Entity Retrieval**: 1 vector similarity query +- **Graph Traversal**: entities × max_path_length × 3 × triple_limit queries +- **Label Resolution**: subgraph_size × 3 individual label queries + +For default parameters (50 entities, path length 2, 30 triple limit, 150 subgraph size): +- **Minimum queries**: 1 + (50 × 2 × 3 × 30) + (150 × 3) = **9,451 database queries** +- **Response time**: 15-30 seconds for moderate-sized graphs +- **Memory usage**: Unbounded cache growth over time +- **Cache effectiveness**: 0% - caches reset on every request +- **Object creation overhead**: GraphRag + Query objects created/destroyed per request + +This specification addresses these gaps by implementing batched queries, intelligent caching, and parallel processing. By optimizing query patterns and data access, TrustGraph can: +- Support enterprise-scale knowledge graphs with millions of entities +- Provide sub-second response times for typical queries +- Handle hundreds of concurrent GraphRAG requests +- Scale efficiently with graph size and complexity + +## Technical Design + +### Architecture + +The GraphRAG performance optimisation requires the following technical components: + +#### 1. **Object Lifetime Architectural Refactor** + - **Make GraphRag long-lived**: Move GraphRag instance to Processor level for persistence across requests + - **Preserve caches**: Maintain label cache, embedding cache, and query result cache between requests + - **Optimize Query object**: Refactor Query as lightweight execution context, not data container + - **Connection persistence**: Maintain database client connections across requests + + Module: `trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py` (modified) + +#### 2. **Optimized Graph Traversal Engine** + - Replace recursive `follow_edges` with iterative breadth-first search + - Implement batched entity processing at each traversal level + - Add cycle detection using visited node tracking + - Include early termination when limits are reached + + Module: `trustgraph-flow/trustgraph/retrieval/graph_rag/optimized_traversal.py` + +#### 3. **Parallel Label Resolution System** + - Batch label queries for multiple entities simultaneously + - Implement async/await patterns for concurrent database access + - Add intelligent prefetching for common label patterns + - Include label cache warming strategies + + Module: `trustgraph-flow/trustgraph/retrieval/graph_rag/label_resolver.py` + +#### 4. **Conservative Label Caching Layer** + - LRU cache with short TTL for labels only (5min) to balance performance vs consistency + - Cache metrics and hit ratio monitoring + - **No embedding caching**: Already cached per-query, no cross-query benefit + - **No query result caching**: Due to graph mutation consistency concerns + + Module: `trustgraph-flow/trustgraph/retrieval/graph_rag/cache_manager.py` + +#### 5. **Query Optimisation Framework** + - Query pattern analysis and optimisation suggestions + - Batch query coordinator for database access + - Connection pooling and query timeout management + - Performance monitoring and metrics collection + + Module: `trustgraph-flow/trustgraph/retrieval/graph_rag/query_optimizer.py` + +### Data Models + +#### Optimized Graph Traversal State + +The traversal engine maintains state to avoid redundant operations: + +```python +@dataclass +class TraversalState: + visited_entities: Set[str] + current_level_entities: Set[str] + next_level_entities: Set[str] + subgraph: Set[Tuple[str, str, str]] + depth: int + query_batch: List[TripleQuery] +``` + +This approach allows: +- Efficient cycle detection through visited entity tracking +- Batched query preparation at each traversal level +- Memory-efficient state management +- Early termination when size limits are reached + +#### Enhanced Cache Structure + +```python +@dataclass +class CacheEntry: + value: Any + timestamp: float + access_count: int + ttl: Optional[float] + +class CacheManager: + label_cache: LRUCache[str, CacheEntry] + embedding_cache: LRUCache[str, CacheEntry] + query_result_cache: LRUCache[str, CacheEntry] + cache_stats: CacheStatistics +``` + +#### Batch Query Structures + +```python +@dataclass +class BatchTripleQuery: + entities: List[str] + query_type: QueryType # SUBJECT, PREDICATE, OBJECT + limit_per_entity: int + +@dataclass +class BatchLabelQuery: + entities: List[str] + predicate: str = LABEL +``` + +### APIs + +#### New APIs: + +**GraphTraversal API** +```python +async def optimized_follow_edges_batch( + entities: List[str], + max_depth: int, + triple_limit: int, + max_subgraph_size: int +) -> Set[Tuple[str, str, str]] +``` + +**Batch Label Resolution API** +```python +async def resolve_labels_batch( + entities: List[str], + cache_manager: CacheManager +) -> Dict[str, str] +``` + +**Cache Management API** +```python +class CacheManager: + async def get_or_fetch_label(self, entity: str) -> str + async def get_or_fetch_embeddings(self, query: str) -> List[float] + async def cache_query_result(self, query_hash: str, result: Any, ttl: int) + def get_cache_statistics(self) -> CacheStatistics +``` + +#### Modified APIs: + +**GraphRag.query()** - Enhanced with performance optimisations: +- Add cache_manager parameter for cache control +- Include performance_metrics return value +- Add query_timeout parameter for reliability + +**Query class** - Refactored for batch processing: +- Replace individual entity processing with batch operations +- Add async context managers for resource cleanup +- Include progress callbacks for long-running operations + +### Implementation Details + +#### Phase 0: Critical Architectural Lifetime Refactor + +**Current Problematic Implementation:** +```python +# INEFFICIENT: GraphRag recreated every request +class Processor(FlowProcessor): + async def on_request(self, msg, consumer, flow): + # PROBLEM: New GraphRag instance per request! + self.rag = GraphRag( + embeddings_client = flow("embeddings-request"), + graph_embeddings_client = flow("graph-embeddings-request"), + triples_client = flow("triples-request"), + prompt_client = flow("prompt-request"), + verbose=True, + ) + # Cache starts empty every time - no benefit from previous requests + response = await self.rag.query(...) + +# VERY SHORT-LIVED: Query object created/destroyed per request +class GraphRag: + async def query(self, query, user="trustgraph", collection="default", ...): + q = Query(rag=self, user=user, collection=collection, ...) # Created + kg = await q.get_labelgraph(query) # Used briefly + # q automatically destroyed when function exits +``` + +**Optimized Long-Lived Architecture:** +```python +class Processor(FlowProcessor): + def __init__(self, **params): + super().__init__(**params) + self.rag_instance = None # Will be initialized once + self.client_connections = {} + + async def initialize_rag(self, flow): + """Initialize GraphRag once, reuse for all requests""" + if self.rag_instance is None: + self.rag_instance = LongLivedGraphRag( + embeddings_client=flow("embeddings-request"), + graph_embeddings_client=flow("graph-embeddings-request"), + triples_client=flow("triples-request"), + prompt_client=flow("prompt-request"), + verbose=True, + ) + return self.rag_instance + + async def on_request(self, msg, consumer, flow): + # REUSE the same GraphRag instance - caches persist! + rag = await self.initialize_rag(flow) + + # Query object becomes lightweight execution context + response = await rag.query_with_context( + query=v.query, + execution_context=QueryContext( + user=v.user, + collection=v.collection, + entity_limit=entity_limit, + # ... other params + ) + ) + +class LongLivedGraphRag: + def __init__(self, ...): + # CONSERVATIVE caches - balance performance vs consistency + self.label_cache = LRUCacheWithTTL(max_size=5000, ttl=300) # 5min TTL for freshness + # Note: No embedding cache - already cached per-query, no cross-query benefit + # Note: No query result cache due to consistency concerns + self.performance_metrics = PerformanceTracker() + + async def query_with_context(self, query: str, context: QueryContext): + # Use lightweight QueryExecutor instead of heavyweight Query object + executor = QueryExecutor(self, context) # Minimal object + return await executor.execute(query) + +@dataclass +class QueryContext: + """Lightweight execution context - no heavy operations""" + user: str + collection: str + entity_limit: int + triple_limit: int + max_subgraph_size: int + max_path_length: int + +class QueryExecutor: + """Lightweight execution context - replaces old Query class""" + def __init__(self, rag: LongLivedGraphRag, context: QueryContext): + self.rag = rag + self.context = context + # No heavy initialization - just references + + async def execute(self, query: str): + # All heavy lifting uses persistent rag caches + return await self.rag.execute_optimized_query(query, self.context) +``` + +This architectural change provides: +- **10-20% database query reduction** for graphs with common relationships (vs 0% currently) +- **Eliminated object creation overhead** for every request +- **Persistent connection pooling** and client reuse +- **Cross-request optimization** within cache TTL windows + +**Important Cache Consistency Limitation:** +Long-term caching introduces staleness risk when entities/labels are deleted or modified in the underlying graph. The LRU cache with TTL provides a balance between performance gains and data freshness, but cannot detect real-time graph changes. + +#### Phase 1: Graph Traversal Optimisation + +**Current Implementation Problems:** +```python +# INEFFICIENT: 3 queries per entity per level +async def follow_edges(self, ent, subgraph, path_length): + # Query 1: s=ent, p=None, o=None + res = await self.rag.triples_client.query(s=ent, p=None, o=None, limit=self.triple_limit) + # Query 2: s=None, p=ent, o=None + res = await self.rag.triples_client.query(s=None, p=ent, o=None, limit=self.triple_limit) + # Query 3: s=None, p=None, o=ent + res = await self.rag.triples_client.query(s=None, p=None, o=ent, limit=self.triple_limit) +``` + +**Optimized Implementation:** +```python +async def optimized_traversal(self, entities: List[str], max_depth: int) -> Set[Triple]: + visited = set() + current_level = set(entities) + subgraph = set() + + for depth in range(max_depth): + if not current_level or len(subgraph) >= self.max_subgraph_size: + break + + # Batch all queries for current level + batch_queries = [] + for entity in current_level: + if entity not in visited: + batch_queries.extend([ + TripleQuery(s=entity, p=None, o=None), + TripleQuery(s=None, p=entity, o=None), + TripleQuery(s=None, p=None, o=entity) + ]) + + # Execute all queries concurrently + results = await self.execute_batch_queries(batch_queries) + + # Process results and prepare next level + next_level = set() + for result in results: + subgraph.update(result.triples) + next_level.update(result.new_entities) + + visited.update(current_level) + current_level = next_level - visited + + return subgraph +``` + +#### Phase 2: Parallel Label Resolution + +**Current Sequential Implementation:** +```python +# INEFFICIENT: Sequential processing +for edge in subgraph: + s = await self.maybe_label(edge[0]) # Individual query + p = await self.maybe_label(edge[1]) # Individual query + o = await self.maybe_label(edge[2]) # Individual query +``` + +**Optimized Parallel Implementation:** +```python +async def resolve_labels_parallel(self, subgraph: List[Triple]) -> List[Triple]: + # Collect all unique entities needing labels + entities_to_resolve = set() + for s, p, o in subgraph: + entities_to_resolve.update([s, p, o]) + + # Remove already cached entities + uncached_entities = [e for e in entities_to_resolve if e not in self.label_cache] + + # Batch query for all uncached labels + if uncached_entities: + label_results = await self.batch_label_query(uncached_entities) + self.label_cache.update(label_results) + + # Apply labels to subgraph + return [ + (self.label_cache.get(s, s), self.label_cache.get(p, p), self.label_cache.get(o, o)) + for s, p, o in subgraph + ] +``` + +#### Phase 3: Advanced Caching Strategy + +**LRU Cache with TTL:** +```python +class LRUCacheWithTTL: + def __init__(self, max_size: int, default_ttl: int = 3600): + self.cache = OrderedDict() + self.max_size = max_size + self.default_ttl = default_ttl + self.access_times = {} + + async def get(self, key: str) -> Optional[Any]: + if key in self.cache: + # Check TTL expiration + if time.time() - self.access_times[key] > self.default_ttl: + del self.cache[key] + del self.access_times[key] + return None + + # Move to end (most recently used) + self.cache.move_to_end(key) + return self.cache[key] + return None + + async def put(self, key: str, value: Any): + if key in self.cache: + self.cache.move_to_end(key) + else: + if len(self.cache) >= self.max_size: + # Remove least recently used + oldest_key = next(iter(self.cache)) + del self.cache[oldest_key] + del self.access_times[oldest_key] + + self.cache[key] = value + self.access_times[key] = time.time() +``` + +#### Phase 4: Query Optimisation and Monitoring + +**Performance Metrics Collection:** +```python +@dataclass +class PerformanceMetrics: + total_queries: int + cache_hits: int + cache_misses: int + avg_response_time: float + subgraph_construction_time: float + label_resolution_time: float + total_entities_processed: int + memory_usage_mb: float +``` + +**Query Timeout and Circuit Breaker:** +```python +async def execute_with_timeout(self, query_func, timeout: int = 30): + try: + return await asyncio.wait_for(query_func(), timeout=timeout) + except asyncio.TimeoutError: + logger.error(f"Query timeout after {timeout}s") + raise GraphRagTimeoutError(f"Query exceeded timeout of {timeout}s") +``` + +## Cache Consistency Considerations + +**Data Staleness Trade-offs:** +- **Label cache (5min TTL)**: Risk of serving deleted/renamed entity labels +- **No embedding caching**: Not needed - embeddings already cached per-query +- **No result caching**: Prevents stale subgraph results from deleted entities/relationships + +**Mitigation Strategies:** +- **Conservative TTL values**: Balance performance gains (10-20%) with data freshness +- **Cache invalidation hooks**: Optional integration with graph mutation events +- **Monitoring dashboards**: Track cache hit rates vs staleness incidents +- **Configurable cache policies**: Allow per-deployment tuning based on mutation frequency + +**Recommended Cache Configuration by Graph Mutation Rate:** +- **High mutation (>100 changes/hour)**: TTL=60s, smaller cache sizes +- **Medium mutation (10-100 changes/hour)**: TTL=300s (default) +- **Low mutation (<10 changes/hour)**: TTL=600s, larger cache sizes + +## Security Considerations + +**Query Injection Prevention:** +- Validate all entity identifiers and query parameters +- Use parameterized queries for all database interactions +- Implement query complexity limits to prevent DoS attacks + +**Resource Protection:** +- Enforce maximum subgraph size limits +- Implement query timeouts to prevent resource exhaustion +- Add memory usage monitoring and limits + +**Access Control:** +- Maintain existing user and collection isolation +- Add audit logging for performance-impacting operations +- Implement rate limiting for expensive operations + +## Performance Considerations + +### Expected Performance Improvements + +**Query Reduction:** +- Current: ~9,000+ queries for typical request +- Optimized: ~50-100 batched queries (98% reduction) + +**Response Time Improvements:** +- Graph traversal: 15-20s → 3-5s (4-5x faster) +- Label resolution: 8-12s → 2-4s (3x faster) +- Overall query: 25-35s → 6-10s (3-4x improvement) + +**Memory Efficiency:** +- Bounded cache sizes prevent memory leaks +- Efficient data structures reduce memory footprint by ~40% +- Better garbage collection through proper resource cleanup + +**Realistic Performance Expectations:** +- **Label cache**: 10-20% query reduction for graphs with common relationships +- **Batching optimization**: 50-80% query reduction (primary optimization) +- **Object lifetime optimization**: Eliminate per-request creation overhead +- **Overall improvement**: 3-4x response time improvement primarily from batching + +**Scalability Improvements:** +- Support for 3-5x larger knowledge graphs (limited by cache consistency needs) +- 3-5x higher concurrent request capacity +- Better resource utilization through connection reuse + +### Performance Monitoring + +**Real-time Metrics:** +- Query execution times by operation type +- Cache hit ratios and effectiveness +- Database connection pool utilisation +- Memory usage and garbage collection impact + +**Performance Benchmarking:** +- Automated performance regression testing +- Load testing with realistic data volumes +- Comparison benchmarks against current implementation + +## Testing Strategy + +### Unit Testing +- Individual component testing for traversal, caching, and label resolution +- Mock database interactions for performance testing +- Cache eviction and TTL expiration testing +- Error handling and timeout scenarios + +### Integration Testing +- End-to-end GraphRAG query testing with optimisations +- Database interaction testing with real data +- Concurrent request handling and resource management +- Memory leak detection and resource cleanup verification + +### Performance Testing +- Benchmark testing against current implementation +- Load testing with varying graph sizes and complexities +- Stress testing for memory and connection limits +- Regression testing for performance improvements + +### Compatibility Testing +- Verify existing GraphRAG API compatibility +- Test with various graph database backends +- Validate result accuracy compared to current implementation + +## Implementation Plan + +### Direct Implementation Approach +Since APIs are allowed to change, implement optimizations directly without migration complexity: + +1. **Replace `follow_edges` method**: Rewrite with iterative batched traversal +2. **Optimize `get_labelgraph`**: Implement parallel label resolution +3. **Add long-lived GraphRag**: Modify Processor to maintain persistent instance +4. **Implement label caching**: Add LRU cache with TTL to GraphRag class + +### Scope of Changes +- **Query class**: Replace ~50 lines in `follow_edges`, add ~30 lines batch handling +- **GraphRag class**: Add caching layer (~40 lines) +- **Processor class**: Modify to use persistent GraphRag instance (~20 lines) +- **Total**: ~140 lines of focused changes, mostly within existing classes + +## Timeline + +**Week 1: Core Implementation** +- Replace `follow_edges` with batched iterative traversal +- Implement parallel label resolution in `get_labelgraph` +- Add long-lived GraphRag instance to Processor +- Implement label caching layer + +**Week 2: Testing and Integration** +- Unit tests for new traversal and caching logic +- Performance benchmarking against current implementation +- Integration testing with real graph data +- Code review and optimization + +**Week 3: Deployment** +- Deploy optimized implementation +- Monitor performance improvements +- Fine-tune cache TTL and batch sizes based on real usage + +## Open Questions + +- **Database Connection Pooling**: Should we implement custom connection pooling or rely on existing database client pooling? +- **Cache Persistence**: Should label and embedding caches persist across service restarts? +- **Distributed Caching**: For multi-instance deployments, should we implement distributed caching with Redis/Memcached? +- **Query Result Format**: Should we optimize the internal triple representation for better memory efficiency? +- **Monitoring Integration**: Which metrics should be exposed to existing monitoring systems (Prometheus, etc.)? + +## References + +- [GraphRAG Original Implementation](trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py) +- [TrustGraph Architecture Principles](architecture-principles.md) +- [Collection Management Specification](collection-management.md)