Add prefix cache aware scheduling

Author: liu-cong

cmd/epp/main.go

@@ -34,6 +34,7 @@ import (
 	"k8s.io/client-go/rest"
 	"k8s.io/component-base/metrics/legacyregistry"
 	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/controller-runtime/pkg/log/zap"
 	"sigs.k8s.io/controller-runtime/pkg/manager"
 	"sigs.k8s.io/controller-runtime/pkg/metrics/filters"
@@ -43,7 +44,9 @@ import (
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metrics"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metrics/collectors"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/plugins/prefix"
 	runserver "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/server"
+	envutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/env"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
 
@@ -107,8 +110,24 @@ var (
 		"Prometheus metric for the LoRA info metrics (must be in vLLM label format).")
 
 	setupLog = ctrl.Log.WithName("setup")
+
+	// Environment variables
+	schedulerV2       = envutil.GetEnvString("EXPERIMENTAL_USE_SCHEDULE_V2", "false", setupLog)
+	prefixCacheConfig = loadPrefixCacheConfig()
 )
 
+func loadPrefixCacheConfig() prefix.Config {
+	// logger := zap.New(zap.RawZapOpts(uberzap.AddCaller()))
+	// log.SetLogger(logger)
+	baseLogger := log.Log.WithName("env-config")
+
+	return prefix.Config{
+		HashBlockSize:          envutil.GetEnvInt("PREFIX_CACHE_HASH_BLOCK_SIZE", prefix.DefaultCacheBlockSize, baseLogger),
+		MaxPrefixBlocksToMatch: envutil.GetEnvInt("PREFIX_CACHE_MAX_PREFIX_BLOCKS", prefix.DefaultMaxPrefixBlocks, baseLogger),
+		LRUIndexerCapacity:     envutil.GetEnvInt("PREFIX_CACHE_MAX_CACHE_SIZE_MB", prefix.DefaultLRUIndexerCapacity, baseLogger),
+	}
+}
+
 func main() {
 	if err := run(); err != nil {
 		os.Exit(1)
@@ -172,6 +191,10 @@ func run() error {
 	datastore := datastore.NewDatastore(ctx, pmf)
 
 	scheduler := scheduling.NewScheduler(datastore)
+	if schedulerV2 == "true" {
+		setupLog.Info("Creating scheduler with prefixCache plugin", "prefix cache config", prefixCacheConfig)
+		scheduler = scheduling.NewSchedulerV2(datastore, prefixCacheConfig)
+	}
 	serverRunner := &runserver.ExtProcServerRunner{
 		GrpcPort:                                 *grpcPort,
 		DestinationEndpointHintMetadataNamespace: *destinationEndpointHintMetadataNamespace,

pkg/epp/metrics/metrics.go

@@ -18,6 +18,7 @@ package metrics
 
 import (
 	"context"
+	"runtime/debug"
 	"sync"
 	"time"
 
@@ -219,6 +220,40 @@ var (
 		},
 		[]string{"commit"},
 	)
+
+	// Prefix indexer Metrics
+	PrefixCacheSize = compbasemetrics.NewGaugeVec(
+		&compbasemetrics.GaugeOpts{
+			Subsystem:      InferenceExtension,
+			Name:           "prefix_indexer_size",
+			Help:           "Size of the prefix indexer.",
+			StabilityLevel: compbasemetrics.ALPHA,
+		},
+		[]string{},
+	)
+
+	PrefixCacheHitRatio = compbasemetrics.NewHistogramVec(
+		&compbasemetrics.HistogramOpts{
+			Subsystem: InferenceExtension,
+			Name:      "prefix_indexer_hit_ratio",
+			Help:      "Ratio of prefix length matched to total prefix length in the cache lookup.",
+			// Buckets from 0.0 to 1.0 in increments
+			Buckets:        []float64{0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0},
+			StabilityLevel: compbasemetrics.ALPHA,
+		},
+		[]string{},
+	)
+
+	PrefixCacheHitLength = compbasemetrics.NewHistogramVec(
+		&compbasemetrics.HistogramOpts{
+			Subsystem:      InferenceExtension,
+			Name:           "prefix_indexer_hit_bytes",
+			Help:           "Length of the prefix match in number of bytes in the cache lookup.",
+			Buckets:        []float64{0, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536},
+			StabilityLevel: compbasemetrics.ALPHA,
+		},
+		[]string{},
+	)
 )
 
 var registerMetrics sync.Once
@@ -244,6 +279,10 @@ func Register() {
 		legacyregistry.MustRegister(SchedulerE2ELatency)
 
 		legacyregistry.MustRegister(InferenceExtensionInfo)
+
+		legacyregistry.MustRegister(PrefixCacheSize)
+		legacyregistry.MustRegister(PrefixCacheHitRatio)
+		legacyregistry.MustRegister(PrefixCacheHitLength)
 	})
 }
 
@@ -352,8 +391,44 @@ func RecordSchedulerE2ELatency(duration time.Duration) {
 	SchedulerE2ELatency.WithLabelValues().Observe(duration.Seconds())
 }
 
+// RecordPrefixCacheSize records the size of the prefix indexer in megabytes.
+func RecordPrefixCacheSize(size int64) {
+	PrefixCacheSize.WithLabelValues().Set(float64(size))
+}
+
+// RecordPrefixCacheMatch records both the hit ratio and hit length for a prefix indexer match.
+// matchedLength is the number of characters that matched, and totalLength is the total prefix length.
+func RecordPrefixCacheMatch(matchedLength, totalLength int) {
+	// Record the hit length metric
+	PrefixCacheHitLength.WithLabelValues().Observe(float64(matchedLength))
+
+	// Record the hit ratio metric if totalLength is positive
+	if totalLength > 0 {
+		ratio := float64(matchedLength) / float64(totalLength)
+		PrefixCacheHitRatio.WithLabelValues().Observe(ratio)
+	}
+}
+
 func RecordInferenceExtensionInfo() {
 	if CommitSHA != "" {
 		InferenceExtensionInfo.WithLabelValues(CommitSHA).Set(1)
 	}
 }
+
+func init() {
+	info, ok := debug.ReadBuildInfo()
+	if !ok {
+		return
+	}
+
+	var Commit = func(i *debug.BuildInfo) string {
+		for _, setting := range i.Settings {
+			if setting.Key == "vcs.revision" {
+				return setting.Value
+			}
+		}
+		return ""
+	}(info)
+
+	CommitSHA = Commit
+}

pkg/epp/metrics/metrics_test.go

@@ -664,3 +664,106 @@ func TestSchedulerE2ELatency(t *testing.T) {
 		})
 	}
 }
+
+func TestPrefixCacheMetrics(t *testing.T) {
+	const (
+		PrefixCacheSizeMetric      = InferenceExtension + "_prefix_indexer_size"
+		PrefixCacheHitRatioMetric  = InferenceExtension + "_prefix_indexer_hit_ratio"
+		PrefixCacheHitLengthMetric = InferenceExtension + "_prefix_indexer_hit_bytes"
+	)
+
+	type cacheMatchRecord struct {
+		matchedLength int
+		totalLength   int
+	}
+
+	scenario := struct {
+		name         string
+		cacheSizes   []int64
+		cacheMatches []cacheMatchRecord
+	}{
+		name:       "multiple cache metrics",
+		cacheSizes: []int64{1024, 2048, 4096},
+		cacheMatches: []cacheMatchRecord{
+			{
+				matchedLength: 5,
+				totalLength:   10,
+			},
+			{
+				matchedLength: 0,
+				totalLength:   10,
+			},
+			{
+				matchedLength: 10,
+				totalLength:   10,
+			},
+			{
+				matchedLength: 7,
+				totalLength:   10,
+			},
+			{
+				matchedLength: 64,
+				totalLength:   128,
+			},
+			{
+				matchedLength: 0,
+				totalLength:   128,
+			},
+		},
+	}
+
+	Register()
+	t.Run(scenario.name, func(t *testing.T) {
+		// Record cache size metrics
+		for _, size := range scenario.cacheSizes {
+			RecordPrefixCacheSize(size)
+		}
+
+		// Record cache match metrics (both hit ratio and hit length)
+		for _, match := range scenario.cacheMatches {
+			RecordPrefixCacheMatch(match.matchedLength, match.totalLength)
+		}
+
+		// Verify cache size metrics
+		wantCacheSizeMetrics, err := os.Open("testdata/prefix_indexer_size_metric")
+		defer func() {
+			if err := wantCacheSizeMetrics.Close(); err != nil {
+				t.Error(err)
+			}
+		}()
+		if err != nil {
+			t.Fatal(err)
+		}
+		if err := testutil.GatherAndCompare(legacyregistry.DefaultGatherer, wantCacheSizeMetrics, PrefixCacheSizeMetric); err != nil {
+			t.Error(err)
+		}
+
+		// Verify hit ratio metrics
+		wantHitRatioMetrics, err := os.Open("testdata/prefix_indexer_hit_ratio_metric")
+		defer func() {
+			if err := wantHitRatioMetrics.Close(); err != nil {
+				t.Error(err)
+			}
+		}()
+		if err != nil {
+			t.Fatal(err)
+		}
+		if err := testutil.GatherAndCompare(legacyregistry.DefaultGatherer, wantHitRatioMetrics, PrefixCacheHitRatioMetric); err != nil {
+			t.Error(err)
+		}
+
+		// Verify hit length metrics
+		wantHitLengthMetrics, err := os.Open("testdata/prefix_indexer_hit_bytes_metric")
+		defer func() {
+			if err := wantHitLengthMetrics.Close(); err != nil {
+				t.Error(err)
+			}
+		}()
+		if err != nil {
+			t.Fatal(err)
+		}
+		if err := testutil.GatherAndCompare(legacyregistry.DefaultGatherer, wantHitLengthMetrics, PrefixCacheHitLengthMetric); err != nil {
+			t.Error(err)
+		}
+	})
+}

pkg/epp/metrics/testdata/prefix_indexer_hit_bytes_metric

@@ -0,0 +1,19 @@
+# HELP inference_extension_prefix_indexer_hit_bytes [ALPHA] Length of the prefix match in number of bytes in the cache lookup.
+# TYPE inference_extension_prefix_indexer_hit_bytes histogram
+inference_extension_prefix_indexer_hit_bytes_bucket{le="0"} 2
+inference_extension_prefix_indexer_hit_bytes_bucket{le="16"} 5
+inference_extension_prefix_indexer_hit_bytes_bucket{le="32"} 5
+inference_extension_prefix_indexer_hit_bytes_bucket{le="64"} 6
+inference_extension_prefix_indexer_hit_bytes_bucket{le="128"} 6
+inference_extension_prefix_indexer_hit_bytes_bucket{le="256"} 6
+inference_extension_prefix_indexer_hit_bytes_bucket{le="512"} 6
+inference_extension_prefix_indexer_hit_bytes_bucket{le="1024"} 6
+inference_extension_prefix_indexer_hit_bytes_bucket{le="2048"} 6
+inference_extension_prefix_indexer_hit_bytes_bucket{le="4096"} 6
+inference_extension_prefix_indexer_hit_bytes_bucket{le="8192"} 6
+inference_extension_prefix_indexer_hit_bytes_bucket{le="16384"} 6
+inference_extension_prefix_indexer_hit_bytes_bucket{le="32768"} 6
+inference_extension_prefix_indexer_hit_bytes_bucket{le="65536"} 6
+inference_extension_prefix_indexer_hit_bytes_bucket{le="+Inf"} 6
+inference_extension_prefix_indexer_hit_bytes_sum 86
+inference_extension_prefix_indexer_hit_bytes_count 6

pkg/epp/metrics/testdata/prefix_indexer_hit_ratio_metric

@@ -0,0 +1,16 @@
+# HELP inference_extension_prefix_indexer_hit_ratio [ALPHA] Ratio of prefix length matched to total prefix length in the cache lookup.
+# TYPE inference_extension_prefix_indexer_hit_ratio histogram
+inference_extension_prefix_indexer_hit_ratio_bucket{le="0"} 2
+inference_extension_prefix_indexer_hit_ratio_bucket{le="0.1"} 2
+inference_extension_prefix_indexer_hit_ratio_bucket{le="0.2"} 2
+inference_extension_prefix_indexer_hit_ratio_bucket{le="0.3"} 2
+inference_extension_prefix_indexer_hit_ratio_bucket{le="0.4"} 2
+inference_extension_prefix_indexer_hit_ratio_bucket{le="0.5"} 4
+inference_extension_prefix_indexer_hit_ratio_bucket{le="0.6"} 4
+inference_extension_prefix_indexer_hit_ratio_bucket{le="0.7"} 5
+inference_extension_prefix_indexer_hit_ratio_bucket{le="0.8"} 5
+inference_extension_prefix_indexer_hit_ratio_bucket{le="0.9"} 5
+inference_extension_prefix_indexer_hit_ratio_bucket{le="1"} 6
+inference_extension_prefix_indexer_hit_ratio_bucket{le="+Inf"} 6
+inference_extension_prefix_indexer_hit_ratio_sum 2.7
+inference_extension_prefix_indexer_hit_ratio_count 6

pkg/epp/metrics/testdata/prefix_indexer_size_metric

@@ -0,0 +1,3 @@
+# HELP inference_extension_prefix_indexer_size [ALPHA] Size of the prefix indexer.
+# TYPE inference_extension_prefix_indexer_size gauge
+inference_extension_prefix_indexer_size{} 4096

pkg/epp/scheduling/plugins/filter/filter.go

@@ -256,6 +256,14 @@ var HasCapacityFilter = &baseFilter{
 	filter: toFilterFunc(queueThresholdPredicate(config.Conf.QueueThresholdCritical).and(kvCacheThresholdPredicate(config.Conf.KVCacheThreshold))),
 }
 
+// NoopFilter is a filter that does not filter out any pods.
+var NoopFilter = &baseFilter{
+	name: "noop",
+	filter: toFilterFunc(func(req *types.LLMRequest, pod types.Pod) bool {
+		return true
+	}),
+}
+
 // podPredicate is a filter function to check whether a pod is desired.
 type podPredicate func(req *types.LLMRequest, pod types.Pod) bool