addressed review comments and fixed imports ordering across the codebase

Signed-off-by: Maroon Ayoub <[email protected]>

Author: vMaroon

examples/kv-cache-aware-scorer/kvcache-aware-scorer.go

@@ -23,7 +23,7 @@ import (
 	"fmt"
 	"os"
 
-	kvcache "github.com/llm-d/llm-d-kv-cache-manager/pkg/kv-cache"
+	"github.com/llm-d/llm-d-kv-cache-manager/pkg/kvcache"
 
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/plugins"

examples/kv-cache-index/main.go

@@ -21,7 +21,7 @@ import (
 	"os"
 	"time"
 
-	kvcache "github.com/llm-d/llm-d-kv-cache-manager/pkg/kv-cache"
+	"github.com/llm-d/llm-d-kv-cache-manager/pkg/kvcache"
 
 	"k8s.io/klog/v2"
 )

pkg/kv-cache/indexer.go renamed to pkg/kvcache/indexer.go

@@ -20,12 +20,12 @@ import (
 	"context"
 	"fmt"
 
-	kvblock "github.com/llm-d/llm-d-kv-cache-manager/pkg/kv-cache/kv-block"
-	"github.com/llm-d/llm-d-kv-cache-manager/pkg/tokenization"
-	"github.com/llm-d/llm-d-kv-cache-manager/pkg/tokenization/prefixstore"
 	"k8s.io/apimachinery/pkg/util/sets"
-
 	"k8s.io/klog/v2"
+
+	"github.com/llm-d/llm-d-kv-cache-manager/pkg/kvcache/kvblock"
+	"github.com/llm-d/llm-d-kv-cache-manager/pkg/tokenization"
+	"github.com/llm-d/llm-d-kv-cache-manager/pkg/tokenization/prefixstore"
 )
 
 // Config holds the configuration for the Indexer module.

pkg/kv-cache/kvblock-scorer.go renamed to pkg/kvcache/kvblock-scorer.go

@@ -19,9 +19,9 @@ package kvcache
 import (
 	"fmt"
 
-	kvblock "github.com/llm-d/llm-d-kv-cache-manager/pkg/kv-cache/kv-block"
-
 	"k8s.io/apimachinery/pkg/util/sets"
+
+	"github.com/llm-d/llm-d-kv-cache-manager/pkg/kvcache/kvblock"
 )
 
 // KVScoringStrategy defines the strategy used to score pods for KV cache block reuse.

pkg/kv-cache/kvblock-scorer_test.go renamed to pkg/kvcache/kvblock-scorer_test.go

@@ -19,25 +19,30 @@ package kvcache_test
 import (
 	"testing"
 
-	kvblock "github.com/llm-d/llm-d-kv-cache-manager/pkg/kv-cache/kv-block"
-
-	kvcache "github.com/llm-d/llm-d-kv-cache-manager/pkg/kv-cache"
+	"github.com/llm-d/llm-d-kv-cache-manager/pkg/kvcache"
+	"github.com/llm-d/llm-d-kv-cache-manager/pkg/kvcache/kvblock"
 
 	"github.com/stretchr/testify/assert"
 )
 
+const (
+	testModelName = "test-model"
+	podA          = "pod-a"
+	podB          = "pod-b"
+)
+
 // TestLongestPrefixScorer verifies scoring based on consecutive block hits from the start.
 func TestLongestPrefixScorer(t *testing.T) {
 	scorer := &kvcache.LongestPrefixScorer{}
 	blockKeys := stringKeysToKVBlockKeys([]string{"b1", "b2", "b3", "b4", "b5", "b6"})
 
 	hitmap := map[kvblock.Key][]string{
-		{ModelName: "test-model", ChunkHash: "b1"}: {"pod-a"},
-		{ModelName: "test-model", ChunkHash: "b2"}: {"pod-a"},
-		{ModelName: "test-model", ChunkHash: "b3"}: {"pod-a"},
-		{ModelName: "test-model", ChunkHash: "b4"}: {"pod-b"},
-		{ModelName: "test-model", ChunkHash: "b5"}: {"pod-b"},
-		{ModelName: "test-model", ChunkHash: "b6"}: {"pod-a"},
+		{ModelName: testModelName, ChunkHash: "b1"}: {podA},
+		{ModelName: testModelName, ChunkHash: "b2"}: {podA},
+		{ModelName: testModelName, ChunkHash: "b3"}: {podA},
+		{ModelName: testModelName, ChunkHash: "b4"}: {podB},
+		{ModelName: testModelName, ChunkHash: "b5"}: {podB},
+		{ModelName: testModelName, ChunkHash: "b6"}: {podA},
 	}
 
 	expected := map[string]int{
@@ -56,7 +61,7 @@ func stringKeysToKVBlockKeys(keys []string) []kvblock.Key {
 	kvKeys := make([]kvblock.Key, len(keys))
 	for i, key := range keys {
 		kvKeys[i] = kvblock.Key{
-			ModelName: "test-model",
+			ModelName: testModelName,
 			ChunkHash: key,
 		}
 	}

pkg/kv-cache/kv-block/in-memory.go renamed to pkg/kvcache/kvblock/in-memory.go

@@ -30,7 +30,7 @@ import (
 
 const (
 	defaultInMemoryIndexSize = 1e8 // TODO: change to memory-size based configuration
-	defaultPodCacheSize      = 10  // number of pods per key
+	defaultPodsPerKey        = 10  // number of pods per key
 )
 
 // InMemoryIndexConfig holds the configuration for the InMemoryIndex.
@@ -45,7 +45,7 @@ type InMemoryIndexConfig struct {
 func DefaultInMemoryIndexConfig() *InMemoryIndexConfig {
 	return &InMemoryIndexConfig{
 		Size:         defaultInMemoryIndexSize,
-		PodCacheSize: defaultPodCacheSize,
+		PodCacheSize: defaultPodsPerKey,
 	}
 }
 
@@ -107,7 +107,7 @@ func (m *InMemoryIndex) Lookup(ctx context.Context, keys []Key,
 	for idx, key := range keys {
 		if pods, found := m.data.Get(key); found { //nolint:nestif // TODO: can this be optimized?
 			if pods == nil || pods.cache.Len() == 0 {
-				traceLogger.Info("no pods found for key, cutting search", "key", key.String())
+				traceLogger.Info("no pods found for key, cutting search", "key", key)
 				return keys[:idx], podsPerKey, nil // early stop since prefix-chain breaks here
 			}
 
@@ -128,7 +128,7 @@ func (m *InMemoryIndex) Lookup(ctx context.Context, keys []Key,
 				}
 			}
 		} else {
-			traceLogger.Info("key not found in index", "key", key.String())
+			traceLogger.Info("key not found in index", "key", key)
 		}
 	}
 
@@ -166,7 +166,7 @@ func (m *InMemoryIndex) Add(ctx context.Context, keys []Key, entries []PodEntry)
 			m.data.Add(key, podCache) // TODO: conflicts here can lead to lost updates, fix locking
 		}
 
-		traceLogger.Info("added pods to key", "key", key.String(), "pods", entries)
+		traceLogger.Info("added pods to key", "key", key, "pods", entries)
 	}
 
 	return nil
@@ -182,19 +182,19 @@ func (m *InMemoryIndex) Evict(ctx context.Context, key Key, entries []PodEntry)
 
 	podCache, found := m.data.Get(key)
 	if !found || podCache == nil {
-		traceLogger.Info("key not found in index, nothing to evict", "key", key.String())
+		traceLogger.Info("key not found in index, nothing to evict", "key", key)
 		return nil
 	}
 
 	for _, entry := range entries {
 		podCache.cache.Remove(entry) // TODO: can this be batched to avoid multiple locks?
 	}
 
-	traceLogger.Info("evicted pods from key", "key", key.String(), "pods", entries)
+	traceLogger.Info("evicted pods from key", "key", key, "pods", entries)
 
 	if podCache.cache.Len() == 0 {
 		m.data.Remove(key)
-		traceLogger.Info("evicted key from index as no pods remain", "key", key.String())
+		traceLogger.Info("evicted key from index as no pods remain", "key", key)
 	}
 
 	return nil

pkg/kv-cache/kv-block/index.go renamed to pkg/kvcache/kvblock/index.go

@@ -127,7 +127,7 @@ func (r *RedisIndex) Lookup(ctx context.Context, keys []Key,
 		}
 
 		if len(filteredPods) == 0 {
-			logger.Info("no pods found for key, cutting search", "key", key.String())
+			logger.Info("no pods found for key, cutting search", "key", key)
 			return keys[:idx], podsPerKey, nil // early stop since prefix-chain breaks here
 		}
 

pkg/kv-cache/kv-block/redis.go renamed to pkg/kvcache/kvblock/redis.go

@@ -23,11 +23,10 @@ import (
 	"encoding/binary"
 	"encoding/hex"
 
-	kvblock "github.com/llm-d/llm-d-kv-cache-manager/pkg/kv-cache/kv-block"
+	"github.com/llm-d/llm-d-kv-cache-manager/pkg/kvcache/kvblock"
+	"github.com/llm-d/llm-d-kv-cache-manager/pkg/utils"
 
 	"k8s.io/klog/v2"
-
-	"github.com/llm-d/llm-d-kv-cache-manager/pkg/utils"
 )
 
 const defaultChunkSize = 256

pkg/kv-cache/token-processor.go renamed to pkg/kvcache/token-processor.go

@@ -21,11 +21,10 @@ import (
 	"fmt"
 	"sync"
 
-	"github.com/llm-d/llm-d-kv-cache-manager/pkg/tokenization/prefixstore"
-
+	"k8s.io/client-go/util/workqueue"
 	"k8s.io/klog/v2"
 
-	"k8s.io/client-go/util/workqueue"
+	"github.com/llm-d/llm-d-kv-cache-manager/pkg/tokenization/prefixstore"
 )
 
 const defaultWorkers = 5