From 9ee506c24c222e239d50b3c18c6e47935dc4997e Mon Sep 17 00:00:00 2001 From: pjdurden Date: Sun, 21 Jun 2026 23:38:28 +0530 Subject: [PATCH] fix(kvevent): purge pod prefixes on AllBlocksCleared (#2287) Signed-off-by: pjdurden --- pkg/kvevent/handler.go | 32 +++++++++-- pkg/kvevent/handler_test.go | 106 +++++++++++++++++++++++++++++++++--- 2 files changed, 125 insertions(+), 13 deletions(-) diff --git a/pkg/kvevent/handler.go b/pkg/kvevent/handler.go index ec6aab47f..7b00088c5 100644 --- a/pkg/kvevent/handler.go +++ b/pkg/kvevent/handler.go @@ -132,10 +132,32 @@ func (h *eventHandler) handleBlockRemoved(ctx context.Context, event *kvcache.Bl } func (h *eventHandler) handleAllBlocksCleared(ctx context.Context, event *kvcache.AllBlocksClearedEvent) error { - // Not implemented: AllBlocksCleared events are pod-local optimizations that - // clear the entire cache for a specific model. These events don't need cross-pod - // synchronization as they represent local memory management decisions. Each pod - // manages its own cache lifecycle independently based on its memory constraints. - klog.V(4).Infof("Received AllBlocksCleared event for pod %s (not implemented)", h.podKey) + // The engine wiped its entire prefix cache (/reset_prefix_cache, OOM-driven + // reset, or /sleep level>=1). The pod stays Running and /health still returns + // 200, so the pod-unsubscribe path that normally calls RemovePrefix never + // fires. If we don't purge the pod's entries here, the prefix router keeps + // preferring this pod for prefixes whose blocks are now gone, turning every + // such match into a cold miss. See issue #2287. + // + // Note: events alone are not a complete fix (vLLM only flushes queued KV + // events from inside a scheduler step, so an idle/sleeping engine may never + // emit this, and ZMQ delivery is lossy). A metric-driven reconcile backstop + // is tracked separately in #2287. + syncIndexer, err := h.manager.syncProvider.GetSyncIndexer(ctx) + if err != nil { + if IsTemporaryError(err) { + klog.V(4).Infof("Temporary error getting sync indexer: %v", err) + return nil // Don't fail on temporary errors + } + return fmt.Errorf("failed to get sync indexer: %w", err) + } + + if err := syncIndexer.RemovePrefix(ctx, h.modelName, h.loraID, h.podKey); err != nil { + klog.Errorf("Failed to process AllBlocksCleared event for pod %s: %v", h.podKey, err) + return err + } + + klog.V(4).Infof("Processed AllBlocksCleared event: purged prefix entries for pod %s", h.podKey) + return nil } diff --git a/pkg/kvevent/handler_test.go b/pkg/kvevent/handler_test.go index 331496482..364b55680 100644 --- a/pkg/kvevent/handler_test.go +++ b/pkg/kvevent/handler_test.go @@ -35,13 +35,21 @@ func int32SliceToBytes(tokens []int32) []byte { return result } +// removePrefixCall records the arguments of a RemovePrefix invocation. +type removePrefixCall struct { + modelName string + loraID int64 + podKey string +} + // mockSyncIndexerWithErrors allows simulating errors type mockSyncIndexerWithErrors struct { - blockStoredErr error - blockRemovedErr error - removePrefixErr error - storedEvents []BlockStoredEvent - removedEvents []BlockRemovedEvent + blockStoredErr error + blockRemovedErr error + removePrefixErr error + storedEvents []BlockStoredEvent + removedEvents []BlockRemovedEvent + removePrefixCalls []removePrefixCall } func (m *mockSyncIndexerWithErrors) ProcessBlockStored(ctx context.Context, event BlockStoredEvent) error { @@ -55,6 +63,7 @@ func (m *mockSyncIndexerWithErrors) ProcessBlockRemoved(ctx context.Context, eve } func (m *mockSyncIndexerWithErrors) RemovePrefix(ctx context.Context, modelName string, loraID int64, podKey string) error { + m.removePrefixCalls = append(m.removePrefixCalls, removePrefixCall{modelName, loraID, podKey}) return m.removePrefixErr } @@ -170,10 +179,19 @@ func TestHandleBlockRemovedEvent(t *testing.T) { } } -// Test HandleEvent with AllBlocksClearedEvent +// Test HandleEvent with AllBlocksClearedEvent. +// The engine wiped its whole prefix cache (e.g. /reset_prefix_cache, OOM reset, +// /sleep), so the pod must be purged from the cross-pod prefix index. Otherwise +// the prefix router keeps routing to a pod whose cache is now cold (#2287). func TestHandleAllBlocksClearedEvent(t *testing.T) { + syncIndexer := &mockSyncIndexerWithErrors{} + syncProvider := &mockSyncProvider{ + indexer: syncIndexer, + } + manager := &Manager{ - ctx: context.Background(), + syncProvider: syncProvider, + ctx: context.Background(), } handler := &eventHandler{ @@ -187,11 +205,83 @@ func TestHandleAllBlocksClearedEvent(t *testing.T) { ModelName: "test-model", } - // Should not return error (no-op implementation) err := handler.HandleEvent(event) if err != nil { t.Errorf("HandleEvent failed: %v", err) } + + // The pod's prefix entries must be removed exactly once, scoped to the + // handler's model/lora/pod. + if len(syncIndexer.removePrefixCalls) != 1 { + t.Fatalf("Expected 1 RemovePrefix call, got %d", len(syncIndexer.removePrefixCalls)) + } + call := syncIndexer.removePrefixCalls[0] + if call.modelName != "test-model" { + t.Errorf("Expected model name 'test-model', got %s", call.modelName) + } + if call.loraID != 789 { + t.Errorf("Expected LoraID 789, got %d", call.loraID) + } + if call.podKey != "default/test-pod" { + t.Errorf("Expected pod key 'default/test-pod', got %s", call.podKey) + } +} + +// A temporary error fetching the sync indexer must be swallowed (no error +// returned, no RemovePrefix call), matching the BlockStored/BlockRemoved paths. +func TestHandleAllBlocksClearedTemporaryError(t *testing.T) { + syncIndexer := &mockSyncIndexerWithErrors{} + syncProvider := &mockSyncProvider{ + indexer: syncIndexer, + err: ErrIndexerNotInitialized, + } + + manager := &Manager{ + syncProvider: syncProvider, + ctx: context.Background(), + } + + handler := &eventHandler{ + manager: manager, + podKey: "default/test-pod", + modelName: "test-model", + loraID: 789, + } + + err := handler.HandleEvent(&kvcache.AllBlocksClearedEvent{ModelName: "test-model"}) + if err != nil { + t.Errorf("Expected nil on temporary error, got %v", err) + } + if len(syncIndexer.removePrefixCalls) != 0 { + t.Errorf("Expected no RemovePrefix calls on temporary error, got %d", len(syncIndexer.removePrefixCalls)) + } +} + +// A RemovePrefix failure must propagate so the subscriber can react. +func TestHandleAllBlocksClearedProcessingError(t *testing.T) { + syncIndexer := &mockSyncIndexerWithErrors{ + removePrefixErr: errors.New("indexer down"), + } + syncProvider := &mockSyncProvider{ + indexer: syncIndexer, + } + + manager := &Manager{ + syncProvider: syncProvider, + ctx: context.Background(), + } + + handler := &eventHandler{ + manager: manager, + podKey: "default/test-pod", + modelName: "test-model", + loraID: 789, + } + + err := handler.HandleEvent(&kvcache.AllBlocksClearedEvent{ModelName: "test-model"}) + if err == nil { + t.Error("Expected error to propagate from RemovePrefix, got nil") + } } // Test HandleEvent with unknown event type