router: make the kvcache GC configurable and stoppable

pkg/kthena-router/scheduler/plugins/kvcache_aware.go

@@ -33,6 +33,7 @@ import (
 	"encoding/binary"
 	"fmt"
 	"strconv"
+	"sync"
 	"time"
 
 	"github.com/redis/go-redis/v9"
@@ -85,6 +86,12 @@ type KVCacheAwareArgs struct {
 	VLLMTokenizerPort int `yaml:"vllmTokenizerPort,omitempty"`
 	// SGLangTokenizerPort overrides the default SGLang tokenizer port (30000).
 	SGLangTokenizerPort int `yaml:"sglangTokenizerPort,omitempty"`
+	// GCInterval overrides how often the ownership GC runs (default 1h).
+	GCInterval string `yaml:"gcInterval,omitempty"`
+	// GCFieldFreshDuration overrides how long an unrefreshed ownership field survives (default 24h).
+	GCFieldFreshDuration string `yaml:"gcFieldFreshDuration,omitempty"`
+	// GCScanSize overrides the SCAN COUNT hint per round (default 100).
+	GCScanSize int64 `yaml:"gcScanSize,omitempty"`
 }
 
 type KVCacheAware struct {
@@ -95,6 +102,11 @@ type KVCacheAware struct {
 	processor        *TokenBlockProcessor
 	tokenizerManager *tokenization.TokenizerManager
 	gcCursor         uint64
+	gcInterval       time.Duration
+	gcFreshDuration  time.Duration
+	gcScanSize       int64
+	gcStopCh         chan struct{}
+	gcStopOnce       sync.Once
 }
 
 var _ framework.ScorePlugin = &KVCacheAware{}
@@ -146,8 +158,15 @@ func NewKVCacheAware(pluginArg runtime.RawExtension) *KVCacheAware {
 	vllmPort := normalizeTokenizerPort(tokenization.EngineVLLM, args.VLLMTokenizerPort, defaultVLLMTokenizerPort)
 	sglangPort := normalizeTokenizerPort(tokenization.EngineSGLang, args.SGLangTokenizerPort, defaultSGLangTokenizerPort)
 
-	klog.Infof("KVCacheAware: config blockSizeToHash=%d, maxBlocksToMatch=%d, vllmTokenizerPort=%d, sglangTokenizerPort=%d",
-		blockSizeToHash, maxBlocksToMatch, vllmPort, sglangPort)
+	gcInterval := parseGCDurationArg("gcInterval", args.GCInterval, kvCacheGCInterval)
+	gcFreshDuration := parseGCDurationArg("gcFieldFreshDuration", args.GCFieldFreshDuration, kvCacheFieldFreshDuration)
+	gcScanSize := args.GCScanSize
+	if gcScanSize <= 0 {
+		gcScanSize = kvCacheGCScanSize
+	}
+
+	klog.Infof("KVCacheAware: config blockSizeToHash=%d, maxBlocksToMatch=%d, vllmTokenizerPort=%d, sglangTokenizerPort=%d, gcInterval=%v, gcFieldFreshDuration=%v, gcScanSize=%d",
+		blockSizeToHash, maxBlocksToMatch, vllmPort, sglangPort, gcInterval, gcFreshDuration, gcScanSize)
 
 	managerConfig := tokenization.TokenizerManagerConfig{
 		EndpointPorts: map[string]int{
@@ -171,6 +190,10 @@ func NewKVCacheAware(pluginArg runtime.RawExtension) *KVCacheAware {
 		redisClient:      redisClient,
 		processor:        &TokenBlockProcessor{blockSize: blockSizeToHash},
 		tokenizerManager: manager,
+		gcInterval:       gcInterval,
+		gcFreshDuration:  gcFreshDuration,
+		gcScanSize:       gcScanSize,
+		gcStopCh:         make(chan struct{}),
 	}
 	plugin.startGC()
 	return plugin
@@ -186,6 +209,55 @@ func normalizeTokenizerPort(engine string, configuredPort, defaultPort int) int
 	return defaultPort
 }
 
+// Stop halts the GC goroutine; safe to call more than once.
+func (t *KVCacheAware) Stop() {
+	t.gcStopOnce.Do(func() {
+		if t.gcStopCh != nil {
+			close(t.gcStopCh)
+		}
+	})
+}
+
+// parseGCDurationArg parses a duration arg, falling back to the default when unset, unparsable, or not positive.
+func parseGCDurationArg(field, raw string, fallback time.Duration) time.Duration {
+	if raw == "" {
+		return fallback
+	}
+	value, err := time.ParseDuration(raw)
+	if err != nil {
+		klog.Warningf("KVCacheAware: ignoring unparsable %s %q, using %v: %v", field, raw, fallback, err)
+		return fallback
+	}
+	if value <= 0 {
+		klog.Warningf("KVCacheAware: ignoring non-positive %s %q, using %v", field, raw, fallback)
+		return fallback
+	}
+	return value
+}
+
+// The effective* helpers keep the documented defaults for a zero-valued KVCacheAware.
+
+func (t *KVCacheAware) effectiveGCInterval() time.Duration {
+	if t.gcInterval <= 0 {
+		return kvCacheGCInterval
+	}
+	return t.gcInterval
+}
+
+func (t *KVCacheAware) effectiveGCFreshDuration() time.Duration {
+	if t.gcFreshDuration <= 0 {
+		return kvCacheFieldFreshDuration
+	}
+	return t.gcFreshDuration
+}
+
+func (t *KVCacheAware) effectiveGCScanSize() int64 {
+	if t.gcScanSize <= 0 {
+		return kvCacheGCScanSize
+	}
+	return t.gcScanSize
+}
+
 func (t *KVCacheAware) Name() string {
 	return t.name
 }
@@ -371,10 +443,15 @@ func (t *KVCacheAware) startGC() {
 }
 
 func (t *KVCacheAware) runGC() {
-	ticker := time.NewTicker(kvCacheGCInterval)
+	ticker := time.NewTicker(t.effectiveGCInterval())
 	defer ticker.Stop()
-	for range ticker.C {
-		t.gcStaleFields()
+	for {
+		select {
+		case <-t.gcStopCh:
+			return
+		case <-ticker.C:
+			t.gcStaleFields()
+		}
 	}
 }
 
@@ -386,9 +463,9 @@ func (t *KVCacheAware) gcStaleFields() {
 	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 	defer cancel()
 
-	keys, nextCursor, err := t.redisClient.Scan(ctx, t.gcCursor, t.keyPrefix+"*", kvCacheGCScanSize).Result()
+	keys, nextCursor, err := t.redisClient.Scan(ctx, t.gcCursor, t.keyPrefix+"*", t.effectiveGCScanSize()).Result()
 	if err != nil {
-		klog.V(4).Infof("KVCacheAware.gcStaleFields: scan failed: %v", err)
+		klog.Warningf("KVCacheAware.gcStaleFields: scan failed: %v", err)
 		return
 	}
 	t.gcCursor = nextCursor
@@ -398,12 +475,12 @@ func (t *KVCacheAware) gcStaleFields() {
 	for _, key := range keys {
 		podTimes, err := t.redisClient.HGetAll(ctx, key).Result()
 		if err != nil {
-			klog.V(4).Infof("KVCacheAware.gcStaleFields: failed to read %s: %v", key, err)
+			klog.Warningf("KVCacheAware.gcStaleFields: failed to read %s: %v", key, err)
 			continue
 		}
 		for pod, ts := range podTimes {
 			updatedAt, err := strconv.ParseInt(ts, 10, 64)
-			if err == nil && now.Sub(time.Unix(updatedAt, 0)) > kvCacheFieldFreshDuration {
+			if err == nil && now.Sub(time.Unix(updatedAt, 0)) > t.effectiveGCFreshDuration() {
 				staleFields[key] = append(staleFields[key], pod)
 			}
 		}
@@ -422,7 +499,7 @@ func (t *KVCacheAware) deleteStaleFields(staleFields map[string][]string) {
 		pipe.HDel(ctx, key, pods...)
 	}
 	if _, err := pipe.Exec(ctx); err != nil {
-		klog.V(4).Infof("KVCacheAware.gcStaleFields: failed to delete stale fields: %v", err)
+		klog.Warningf("KVCacheAware.gcStaleFields: failed to delete stale fields: %v", err)
 	}
 }
 

pkg/kthena-router/scheduler/plugins/kvcache_aware_test.go

@@ -503,6 +503,132 @@ func TestKVCacheAware_GCStaleFields(t *testing.T) {
 	}
 }
 
+func TestParseGCDurationArg(t *testing.T) {
+	fallback := time.Hour
+	tests := []struct {
+		name string
+		raw  string
+		want time.Duration
+	}{
+		{name: "unset uses the fallback", raw: "", want: fallback},
+		{name: "a valid duration is honoured", raw: "5s", want: 5 * time.Second},
+		{name: "zero falls back", raw: "0s", want: fallback},
+		{name: "negative falls back", raw: "-1m", want: fallback},
+		{name: "a bare number falls back", raw: "24", want: fallback},
+		{name: "nonsense falls back", raw: "soon", want: fallback},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			if got := parseGCDurationArg("gcInterval", tc.raw, fallback); got != tc.want {
+				t.Errorf("parseGCDurationArg(%q) = %v, want %v", tc.raw, got, tc.want)
+			}
+		})
+	}
+}
+
+func TestKVCacheAware_GCKnobsFromArgs(t *testing.T) {
+	// One end to end case, because the argument names have to survive the YAML to
+	// JSON bridge the plugin args go through. parseGCDurationArg is unit tested above.
+	plugin := NewKVCacheAware(runtime.RawExtension{
+		Raw: []byte("gcInterval: 5s\ngcFieldFreshDuration: 2m\ngcScanSize: 7\n"),
+	})
+	defer plugin.Stop()
+
+	if got := plugin.effectiveGCInterval(); got != 5*time.Second {
+		t.Errorf("gc interval = %v, want %v", got, 5*time.Second)
+	}
+	if got := plugin.effectiveGCFreshDuration(); got != 2*time.Minute {
+		t.Errorf("gc fresh duration = %v, want %v", got, 2*time.Minute)
+	}
+	if got := plugin.effectiveGCScanSize(); got != 7 {
+		t.Errorf("gc scan size = %d, want 7", got)
+	}
+}
+
+func TestKVCacheAware_ZeroValuePluginKeepsGCDefaults(t *testing.T) {
+	// Several tests build the plugin as a struct literal. A zero value must keep
+	// behaving the way it did before these knobs existed.
+	plugin := &KVCacheAware{}
+
+	if got := plugin.effectiveGCInterval(); got != kvCacheGCInterval {
+		t.Errorf("gc interval = %v, want %v", got, kvCacheGCInterval)
+	}
+	if got := plugin.effectiveGCFreshDuration(); got != kvCacheFieldFreshDuration {
+		t.Errorf("gc fresh duration = %v, want %v", got, kvCacheFieldFreshDuration)
+	}
+	if got := plugin.effectiveGCScanSize(); got != kvCacheGCScanSize {
+		t.Errorf("gc scan size = %d, want %d", got, kvCacheGCScanSize)
+	}
+}
+
+func TestKVCacheAware_GCStaleFieldsHonoursConfiguredFreshness(t *testing.T) {
+	mr, err := miniredis.Run()
+	if err != nil {
+		t.Fatalf("failed to start miniredis: %v", err)
+	}
+	defer mr.Close()
+
+	client := redis.NewClient(&redis.Options{Addr: mr.Addr()})
+	defer client.Close()
+
+	// A one minute window: the field below is 5 minutes old, so it is stale here
+	// while the 24h default would keep it.
+	plugin := &KVCacheAware{
+		keyPrefix:       kvCacheKeyPrefix,
+		redisClient:     client,
+		gcFreshDuration: time.Minute,
+	}
+
+	ctx := context.Background()
+	key := KVCacheAwareBlock{ModelName: "qwen", ChunkHash: 987}.String(kvCacheKeyPrefix)
+	if err := client.HSet(ctx, key,
+		"recent-pod.default", fmt.Sprintf("%d", time.Now().Unix()),
+		"older-pod.default", fmt.Sprintf("%d", time.Now().Add(-5*time.Minute).Unix()),
+	).Err(); err != nil {
+		t.Fatalf("failed to seed redis: %v", err)
+	}
+
+	plugin.gcStaleFields()
+
+	stale, err := client.HExists(ctx, key, "older-pod.default").Result()
+	if err != nil {
+		t.Fatalf("failed to check older-pod.default: %v", err)
+	}
+	if stale {
+		t.Error("expected the field older than the configured window to be removed")
+	}
+	fresh, err := client.HExists(ctx, key, "recent-pod.default").Result()
+	if err != nil {
+		t.Fatalf("failed to check recent-pod.default: %v", err)
+	}
+	if !fresh {
+		t.Error("expected the field inside the configured window to remain")
+	}
+}
+
+func TestKVCacheAware_StopIsIdempotentAndEndsTheGCLoop(t *testing.T) {
+	plugin := &KVCacheAware{
+		gcInterval: time.Millisecond,
+		gcStopCh:   make(chan struct{}),
+	}
+
+	done := make(chan struct{})
+	go func() {
+		plugin.runGC()
+		close(done)
+	}()
+
+	plugin.Stop()
+	plugin.Stop() // second call must not panic on a closed channel
+
+	select {
+	case <-done:
+	case <-time.After(5 * time.Second):
+		t.Fatal("runGC did not return after Stop")
+	}
+}
+
 // Helper function to create test pods
 func createTestPods(names ...string) []*datastore.PodInfo {
 	pods := make([]*datastore.PodInfo, len(names))