feat:add support to invalidate KV cache via AllBlocksCleared event

Author: yash9263

docs/architecture.md

@@ -96,7 +96,7 @@ sequenceDiagram
         else BlockRemoved
             Worker->>Index: Evict(key, podEntry)
         else AllBlocksCleared
-            Note over Worker: No-op
+            Worker->>Index: Clear()
         end
     end
 ```

examples/helper/events.go

@@ -109,6 +109,31 @@ func SimulateRemoveEvent(ctx context.Context, publisher *Publisher) error {
 	return nil
 }
 
+func SimulateClearAllBlocksEvent(ctx context.Context, publisher *Publisher) error {
+	logger := log.FromContext(ctx)
+	logger.Info("@@@ Simulating vLLM engine clear all blocks...")
+
+	clearAllBlocksEvent := []any{
+		"AllBlocksCleared",
+	}
+
+	clearAllBlocksPayload, _ := msgpack.Marshal(clearAllBlocksEvent)
+
+	clearAllBlockEventBatch := []any{
+		float64(time.Now().UnixNano()) / 1e9,
+		[]msgpack.RawMessage{clearAllBlocksPayload},
+		nil,
+	}
+
+	if err := publisher.PublishEvent(ctx, topic, clearAllBlockEventBatch); err != nil {
+		return fmt.Errorf("failed to publish AllBlocksCleared event: %w", err)
+	}
+	logger.Info("@@@ Published AllBlocksCleared event", "topic", topic)
+
+	time.Sleep(3 * time.Second)
+	return nil
+}
+
 func SetupEventsPool(ctx context.Context, kvBlockIndex kvblock.Index) (*kvevents.Pool, error) {
 	logger := log.FromContext(ctx)
 

examples/kv_events/offline/main.go

@@ -201,6 +201,32 @@ func RunEventsDemo(ctx context.Context, kvCacheIndexer *kvcache.Indexer, publish
 	}
 	logger.Info("@@@ Final pod scores after BlockRemoved events", "pods", pods)
 
+	// Simulate vLLM engine publishing BlockStored events
+	err = helper.SimulateProduceEvent(ctx, publisher)
+	if err != nil {
+		return err
+	}
+
+	// Query again to see the effect of the events
+	pods, err = kvCacheIndexer.GetPodScores(ctx, testdata.RenderReq, testdata.Prompt, testdata.ModelName, nil)
+	if err != nil {
+		return err
+	}
+	logger.Info("@@@ Pod scores after BlockStored events", "pods", pods)
+
+	// Simulate vLLM engine publishing AllBlocksCleared event
+	err = helper.SimulateClearAllBlocksEvent(ctx, publisher)
+	if err != nil {
+		return err
+	}
+
+	// Query again to see the effect of the events
+	pods, err = kvCacheIndexer.GetPodScores(ctx, testdata.RenderReq, testdata.Prompt, testdata.ModelName, nil)
+	if err != nil {
+		return err
+	}
+	logger.Info("@@@ Pod scores after AllBlocksCleared events", "pods", pods)
+
 	logger.Info("Events demo completed. Pool continues listening for more events...")
 	logger.Info("Press Ctrl+C to shutdown")
 

examples/valkey_example/main.go

@@ -201,6 +201,21 @@ func demonstrateValkeyOperations(ctx context.Context, indexer *kvcache.Indexer)
 
 	logger.Info("Cache lookup after eviction", "keysFound", len(lookupAfterEvict))
 
+	// Clear the cache
+	logger.Info("Clearing the cache")
+	err = indexer.KVBlockIndex().Clear(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to clear cache: %w", err)
+	}
+
+	// Lookup again after clear
+	lookupAfterClear, err := indexer.KVBlockIndex().Lookup(ctx, promptKeys, nil)
+	if err != nil {
+		return fmt.Errorf("failed to lookup after clear: %w", err)
+	}
+
+	logger.Info("Cache lookup after clear", "keysFound", len(lookupAfterClear))
+
 	// Final score check to see the difference
 	finalScores, err := indexer.GetPodScores(ctx, nil, prompt, modelName, []string{"demo-pod-1", "demo-pod-2"})
 	if err != nil {

pkg/kvcache/kvblock/cost_aware_memory.go

@@ -312,3 +312,17 @@ func (m *CostAwareMemoryIndex) GetRequestKey(ctx context.Context, engineKey Bloc
 	}
 	return requestKey, nil
 }
+
+// Clear removes all entries from the index backend.
+func (m *CostAwareMemoryIndex) Clear(ctx context.Context) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	traceLogger := log.FromContext(ctx).V(logging.TRACE).WithName("kvblock.CostAwareMemoryIndex.Clear")
+
+	m.requestKeys.Purge()
+	m.data.Clear()
+
+	m.data.Wait()
+	traceLogger.Info("Cleared CostAwareMemoryIndex")
+	return nil
+}

pkg/kvcache/kvblock/in_memory.go

@@ -277,3 +277,12 @@ func podsPerKeyPrintHelper(ks map[BlockHash][]PodEntry) string {
 	}
 	return b.String()
 }
+
+// Clear removes all entries from the index backend.
+func (m *InMemoryIndex) Clear(ctx context.Context) error {
+	traceLogger := log.FromContext(ctx).V(logging.TRACE).WithName("kvblock.InMemoryIndex.Clear")
+	m.engineToRequestKeys.Purge()
+	m.data.Purge()
+	traceLogger.Info("Cleared InMemoryIndex")
+	return nil
+}

pkg/kvcache/kvblock/index.go

@@ -133,6 +133,8 @@ type Index interface {
 	Evict(ctx context.Context, engineKey BlockHash, entries []PodEntry) error
 	// GetRequestKey returns the requestKey associated with the given engineKey.
 	GetRequestKey(ctx context.Context, engineKey BlockHash) (BlockHash, error)
+	// Clear removes all entries from the index backend.
+	Clear(ctx context.Context) error
 }
 
 // BlockHash struct represents a unique identifier for a KV-cache block.

pkg/kvcache/kvblock/index_test.go

@@ -63,6 +63,11 @@ func testCommonIndexBehavior(t *testing.T, indexFactory func(t *testing.T) Index
 		index := indexFactory(t)
 		testConcurrentOperations(t, ctx, index)
 	})
+
+	t.Run("Clear", func(t *testing.T) {
+		index := indexFactory(t)
+		testClear(t, ctx, index)
+	})
 }
 
 // testBasicAddAndLookup tests basic Add and Lookup functionality.
@@ -213,6 +218,30 @@ func testEvictBasic(t *testing.T, ctx context.Context, index Index) {
 	assert.ElementsMatch(t, expected, podsPerKey[requestKey])
 }
 
+func testClear(t *testing.T, ctx context.Context, index Index) {
+	t.Helper()
+	engineKey := BlockHash(17434655)
+	requestKey := BlockHash(59244875)
+	entries := []PodEntry{
+		{PodIdentifier: "pod1", DeviceTier: "gpu"},
+		{PodIdentifier: "pod2", DeviceTier: "gpu"},
+		{PodIdentifier: "pod3", DeviceTier: "cpu"},
+	}
+
+	// Add entries
+	err := index.Add(ctx, []BlockHash{engineKey}, []BlockHash{requestKey}, entries)
+	require.NoError(t, err)
+
+	// Clear the index
+	err = index.Clear(ctx)
+	require.NoError(t, err)
+
+	// Verify that the index is empty
+	podsPerKey, err := index.Lookup(ctx, []BlockHash{requestKey}, sets.Set[string]{})
+	require.NoError(t, err)
+	assert.Len(t, podsPerKey, 0)
+}
+
 // testConcurrentOperations tests thread safety with concurrent operations.
 func testConcurrentOperations(t *testing.T, ctx context.Context, index Index) {
 	t.Helper()
@@ -245,6 +274,10 @@ func testConcurrentOperations(t *testing.T, ctx context.Context, index Index) {
 					if err := index.Evict(ctx, engineKey, entries); err != nil {
 						errChan <- err
 					}
+				case 3: // Clear
+					if err := index.Clear(ctx); err != nil {
+						errChan <- err
+					}
 				}
 			}
 		}(goroutineID)

pkg/kvcache/kvblock/instrumented_index.go

@@ -27,7 +27,7 @@ type instrumentedIndex struct {
 }
 
 // NewInstrumentedIndex wraps an Index and emits metrics for Add, Evict, and
-// Lookup.
+// Lookup, Clear.
 func NewInstrumentedIndex(next Index) Index {
 	return &instrumentedIndex{next: next}
 }
@@ -90,3 +90,9 @@ func recordHitMetrics(keyToPods map[BlockHash][]PodEntry) {
 	metrics.MaxPodHitCount.Add(float64(maxHit))
 	metrics.LookupHits.Add(float64(maxHit))
 }
+
+func (m *instrumentedIndex) Clear(ctx context.Context) error {
+	err := m.next.Clear(ctx)
+	metrics.Clear.Add(1)
+	return err
+}

pkg/kvcache/kvblock/redis.go

@@ -292,3 +292,15 @@ func (r *RedisIndex) GetRequestKey(ctx context.Context, engineKey BlockHash) (Bl
 func redisEngineKey(engineKey BlockHash) string {
 	return "engine:" + engineKey.String()
 }
+
+// Clear removes all entries from the index backend.
+func (r *RedisIndex) Clear(ctx context.Context) error {
+	logger := log.FromContext(ctx).WithName("kvblock.RedisIndex.Clear")
+	err := r.RedisClient.FlushDB(ctx).Err()
+	if err != nil {
+		return fmt.Errorf("failed to flush %s db: %w", r.BackendType, err)
+	}
+	logger.Info("Cleared", "index", r.BackendType)
+
+	return nil
+}