implement kvblock.index metrics config (#57)

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

Author: vMaroon

docs/configuration.md

@@ -52,7 +52,8 @@ Here's a complete configuration example with all options:
       "size": 100000000,
       "podCacheSize": 10
     },
-    "enableMetrics": true
+    "enableMetrics": true,
+    "metricsLoggingInterval": "1m0s"
   },
   "tokenizersPoolConfig": {
     "workersCount": 8,
@@ -81,6 +82,7 @@ Configures the KV-block index backend. Multiple backends can be configured, but
 | `inMemoryConfig` | [InMemoryIndexConfig](#in-memory-index-configuration) | In-memory index configuration | See defaults |
 | `redisConfig` | [RedisIndexConfig](#redis-index-configuration)        | Redis index configuration | `null` |
 | `enableMetrics` | `boolean`                                             | Enable admissions/evictions/hits/misses recording | `false` |
+| `metricsLoggingInterval` | `string` (duration) | Interval at which metrics are logged (e.g., `"1m0s"`). If zero or omitted, metrics logging is disabled. Requires `enableMetrics` to be `true`. | `"0s"` |
 
 ### In-Memory Index Configuration (`InMemoryIndexConfig`)
 

examples/kv_cache_index/main.go

@@ -94,8 +94,7 @@ func setupKVCacheIndexer(ctx context.Context) (*kvcache.Indexer, error) {
 		return nil, err
 	}
 
-	//nolint:contextcheck // NewKVCacheIndexer does not accept context parameter
-	kvCacheIndexer, err := kvcache.NewKVCacheIndexer(config)
+	kvCacheIndexer, err := kvcache.NewKVCacheIndexer(ctx, config)
 	if err != nil {
 		return nil, err
 	}

examples/kv_events/offline/main.go

@@ -106,8 +106,7 @@ func main() {
 func setupKVCacheIndexer(ctx context.Context) (*kvcache.Indexer, error) {
 	logger := klog.FromContext(ctx)
 
-	//nolint:contextcheck // NewKVCacheIndexer does not accept context parameter
-	kvCacheIndexer, err := kvcache.NewKVCacheIndexer(getKVCacheIndexerConfig())
+	kvCacheIndexer, err := kvcache.NewKVCacheIndexer(ctx, getKVCacheIndexerConfig())
 	if err != nil {
 		return nil, err
 	}

examples/kv_events/online/main.go

@@ -27,7 +27,6 @@ import (
 	"syscall"
 	"time"
 
-	"github.com/llm-d/llm-d-kv-cache-manager/pkg/kvcache/metrics"
 	"k8s.io/klog/v2"
 
 	"github.com/llm-d/llm-d-kv-cache-manager/pkg/kvcache"
@@ -71,6 +70,9 @@ func getKVCacheIndexerConfig() *kvcache.Config {
 		config.TokenProcessorConfig.BlockSize = blockSize
 	}
 
+	config.KVBlockIndexConfig.EnableMetrics = true
+	config.KVBlockIndexConfig.MetricsLoggingInterval = 15 * time.Second
+
 	return config
 }
 
@@ -117,10 +119,6 @@ func main() {
 	eventsPool.Start(ctx)
 	logger.Info("Events pool started and listening for ZMQ messages")
 
-	metrics.Register()
-	metrics.StartMetricsLogging(ctx, time.Second*10)
-	logger.Info("Started metrics thread")
-
 	// Setup graceful shutdown
 	sigChan := make(chan os.Signal, 1)
 	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
@@ -171,8 +169,7 @@ func main() {
 func setupKVCacheIndexer(ctx context.Context) (*kvcache.Indexer, error) {
 	logger := klog.FromContext(ctx)
 
-	//nolint:contextcheck // NewKVCacheIndexer does not accept context parameter
-	kvCacheIndexer, err := kvcache.NewKVCacheIndexer(getKVCacheIndexerConfig())
+	kvCacheIndexer, err := kvcache.NewKVCacheIndexer(ctx, getKVCacheIndexerConfig())
 	if err != nil {
 		return nil, err
 	}

pkg/kvcache/indexer.go

@@ -53,6 +53,8 @@ func NewDefaultConfig() *Config {
 
 // Indexer is a concrete implementation of the KVCacheIndex interface.
 type Indexer struct {
+	config *Config
+
 	tokensIndexer   prefixstore.Indexer    // gets tokens for a prompt
 	tokensProcessor kvblock.TokenProcessor // turns tokens to kv block keys
 	kvBlockIndex    kvblock.Index          // looks up pods for block keys
@@ -62,15 +64,15 @@ type Indexer struct {
 }
 
 // NewKVCacheIndexer creates a KVCacheIndex given a Config.
-func NewKVCacheIndexer(config *Config) (*Indexer, error) {
+func NewKVCacheIndexer(ctx context.Context, config *Config) (*Indexer, error) {
 	tokensIndexer, err := prefixstore.NewLRUTokenStore(config.PrefixStoreConfig)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create prefixstore.Indexer: %w", err)
 	}
 
 	tokensProcessor := kvblock.NewChunkedTokenDatabase(config.TokenProcessorConfig)
 
-	kvBlockIndex, err := kvblock.NewIndex(config.KVBlockIndexConfig)
+	kvBlockIndex, err := kvblock.NewIndex(ctx, config.KVBlockIndexConfig)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create RedisKVBlockIndexer: %w", err)
 	}
@@ -86,6 +88,7 @@ func NewKVCacheIndexer(config *Config) (*Indexer, error) {
 	}
 
 	return &Indexer{
+		config:          config,
 		tokensIndexer:   tokensIndexer,
 		tokensProcessor: tokensProcessor,
 		kvBlockIndex:    kvBlockIndex,

pkg/kvcache/kvblock/index.go

@@ -19,7 +19,9 @@ package kvblock
 import (
 	"context"
 	"fmt"
+	"time"
 
+	"github.com/llm-d/llm-d-kv-cache-manager/pkg/kvcache/metrics"
 	"k8s.io/apimachinery/pkg/util/sets"
 )
 
@@ -34,6 +36,10 @@ type IndexConfig struct {
 	// EnableMetrics toggles whether admissions/evictions/hits/misses are
 	// recorded.
 	EnableMetrics bool `json:"enableMetrics"`
+	// MetricsLoggingInterval defines the interval at which metrics are logged.
+	// If zero, metrics logging is disabled.
+	// Requires `EnableMetrics` to be true.
+	MetricsLoggingInterval time.Duration `json:"metricsLoggingInterval"`
 }
 
 // DefaultIndexConfig returns a default configuration for the KV-block index.
@@ -45,7 +51,7 @@ func DefaultIndexConfig() *IndexConfig {
 }
 
 // NewIndex creates a new Index instance.
-func NewIndex(cfg *IndexConfig) (Index, error) {
+func NewIndex(ctx context.Context, cfg *IndexConfig) (Index, error) {
 	if cfg == nil {
 		cfg = DefaultIndexConfig()
 	}
@@ -60,6 +66,7 @@ func NewIndex(cfg *IndexConfig) (Index, error) {
 			return nil, fmt.Errorf("failed to create in-memory index: %w", err)
 		}
 	case cfg.RedisConfig != nil:
+		//nolint:contextcheck // NewKVCacheIndexer does not accept context parameter
 		idx, err = NewRedisIndex(cfg.RedisConfig)
 		if err != nil {
 			return nil, fmt.Errorf("failed to create Redis index: %w", err)
@@ -71,6 +78,11 @@ func NewIndex(cfg *IndexConfig) (Index, error) {
 	// wrap in metrics only if enabled
 	if cfg.EnableMetrics {
 		idx = NewInstrumentedIndex(idx)
+		metrics.Register()
+		if cfg.MetricsLoggingInterval > 0 {
+			// this is non-blocking
+			metrics.StartMetricsLogging(ctx, cfg.MetricsLoggingInterval)
+		}
 	}
 
 	return idx, nil

tests/e2e/redis_mock/e2e_suite_test.go

@@ -69,7 +69,7 @@ func (s *KVCacheSuite) SetupTest() {
 
 	s.Pod1IP = "10.0.0.1"
 
-	s.indexer, err = kvcache.NewKVCacheIndexer(s.config)
+	s.indexer, err = kvcache.NewKVCacheIndexer(s.ctx, s.config)
 	s.kvBlockIndex = s.indexer.KVBlockIndex()
 	s.Require().NoError(err)