Merge branch 'main' into fix/per-prompt-mm-features

Author: albertoperdomo2

pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer.go

@@ -21,6 +21,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"sort"
 	"time"
 
 	"github.com/jellydator/ttlcache/v3"
@@ -64,7 +65,19 @@ type PluginConfig struct {
 	SpeculativeTTL string `json:"speculativeTTL"`
 }
 
-var _ requestcontrol.DataProducer = &Producer{}
+var (
+	_ requestcontrol.DataProducer = &Producer{}
+	_ plugin.StateDumper          = &Producer{}
+)
+
+// subscriberManager is the subset of kvevents.SubscriberManager the producer
+// relies on, narrowed so tests can substitute a fake.
+type subscriberManager interface {
+	EnsureSubscriber(ctx context.Context, podIdentifier, endpoint, topicFilter string, remoteSocket bool) error
+	RemoveSubscriber(ctx context.Context, podIdentifier string)
+	GetActiveSubscribers() ([]string, []string)
+	Shutdown(ctx context.Context)
+}
 
 // Producer is a DataProducer plugin that maintains a KV-block prefix-cache
 // index by subscribing to vLLM KV-events and writes per-endpoint
@@ -78,7 +91,7 @@ type Producer struct {
 	typedName      plugin.TypedName
 	kvCacheIndexer kvCacheIndexer
 
-	subscribersManager *kvevents.SubscriberManager
+	subscribersManager subscriberManager
 	kvEventsConfig     *kvevents.Config
 
 	kvBlockScorer kvcache.KVBlockScorer
@@ -202,6 +215,70 @@ func (p *Producer) TypedName() plugin.TypedName {
 	return p.typedName
 }
 
+// Debug-dump caps keep the payload bounded. A list is partial when its matching
+// TotalX exceeds MaxX.
+const (
+	maxDumpSubscribers        = 100
+	maxDumpSpeculativeEntries = 100
+)
+
+// precisePrefixState is the snapshot returned by DumpState. The KV-block index
+// is keyed by prompt-derived block hashes and is not enumerable, so it is not
+// reported; the active subscriber pod identities and the live speculative
+// request ids are enumerated (sorted and capped) for debugging.
+type precisePrefixState struct {
+	Subscribers             []string `json:"subscribers"`
+	TotalSubscribers        int      `json:"totalSubscribers"`
+	MaxSubscribers          int      `json:"maxSubscribers"`
+	SpeculativeIndexing     bool     `json:"speculativeIndexing"`
+	SpeculativeEntries      []string `json:"speculativeEntries"`
+	TotalSpeculativeEntries int      `json:"totalSpeculativeEntries"`
+	MaxSpeculativeEntries   int      `json:"maxSpeculativeEntries"`
+	BlockSizeTokens         int      `json:"blockSizeTokens"`
+}
+
+// DumpState reports the producer's bounded operational state: the active
+// KV-event subscriber pod identities, whether speculative indexing is on and
+// the live speculative request ids, and the block size in tokens. Both lists
+// are sorted and capped; the prompt-derived block index is not exposed.
+func (p *Producer) DumpState() (json.RawMessage, error) {
+	subscribers := []string{}
+	var totalSubscribers int
+	if p.subscribersManager != nil {
+		ids, _ := p.subscribersManager.GetActiveSubscribers()
+		totalSubscribers = len(ids)
+		subscribers = sortedCapped(ids, maxDumpSubscribers)
+	}
+	speculativeEntries := []string{}
+	var totalSpeculativeEntries int
+	if p.speculativeCache != nil {
+		keys := p.speculativeCache.Keys()
+		totalSpeculativeEntries = len(keys)
+		speculativeEntries = sortedCapped(keys, maxDumpSpeculativeEntries)
+	}
+	return json.Marshal(precisePrefixState{
+		Subscribers:             subscribers,
+		TotalSubscribers:        totalSubscribers,
+		MaxSubscribers:          maxDumpSubscribers,
+		SpeculativeIndexing:     p.speculativeEnabled,
+		SpeculativeEntries:      speculativeEntries,
+		TotalSpeculativeEntries: totalSpeculativeEntries,
+		MaxSpeculativeEntries:   maxDumpSpeculativeEntries,
+		BlockSizeTokens:         p.blockSizeTokens,
+	})
+}
+
+// sortedCapped returns a sorted copy of in (never nil, so empty lists serialize
+// as [] not null), truncated to limit entries.
+func sortedCapped(in []string, limit int) []string {
+	out := append([]string{}, in...)
+	sort.Strings(out)
+	if len(out) > limit {
+		out = out[:limit]
+	}
+	return out
+}
+
 // Produces declares the PrefixCacheMatchInfoDataKey published per endpoint,
 // name-bound to this producer instance.
 func (p *Producer) Produces() map[plugin.DataKey]any {

pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer_test.go

@@ -19,8 +19,10 @@ package preciseprefixcache
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"testing"
 
+	"github.com/jellydator/ttlcache/v3"
 	"github.com/llm-d/llm-d-router/pkg/kvcache"
 	"github.com/llm-d/llm-d-router/pkg/kvcache/kvblock"
 	"github.com/llm-d/llm-d-router/pkg/kvevents"
@@ -675,3 +677,99 @@ func TestNew_BlockSizeFlowsViaTokenProcessor(t *testing.T) {
 		})
 	}
 }
+
+type fakeSubscriberManager struct {
+	ids       []string
+	endpoints []string
+}
+
+func (f *fakeSubscriberManager) EnsureSubscriber(_ context.Context, _, _, _ string, _ bool) error {
+	return nil
+}
+func (f *fakeSubscriberManager) RemoveSubscriber(_ context.Context, _ string) {}
+func (f *fakeSubscriberManager) GetActiveSubscribers() ([]string, []string) {
+	return f.ids, f.endpoints
+}
+func (f *fakeSubscriberManager) Shutdown(_ context.Context) {}
+
+func TestDumpState(t *testing.T) {
+	// Start() is intentionally not called: NoTTL entries never expire, and the
+	// enumeration helpers work on an unstarted cache.
+	specCache := ttlcache.New[string, *speculativeEntries]()
+	specCache.Set("req-2", &speculativeEntries{}, ttlcache.NoTTL)
+	specCache.Set("req-1", &speculativeEntries{}, ttlcache.NoTTL)
+
+	p := &Producer{
+		typedName:          plugin.TypedName{Type: PluginType, Name: "x"},
+		subscribersManager: &fakeSubscriberManager{ids: []string{"ns/pod-b", "ns/pod-a"}},
+		speculativeCache:   specCache,
+		speculativeEnabled: true,
+		blockSizeTokens:    64,
+	}
+
+	payload, err := p.DumpState()
+	require.NoError(t, err)
+	// Subscriber pod identities and live request ids are enumerated for debugging.
+	assert.Contains(t, string(payload), "ns/pod-a")
+
+	var state precisePrefixState
+	require.NoError(t, json.Unmarshal(payload, &state))
+	assert.Equal(t, precisePrefixState{
+		Subscribers:             []string{"ns/pod-a", "ns/pod-b"},
+		TotalSubscribers:        2,
+		MaxSubscribers:          maxDumpSubscribers,
+		SpeculativeIndexing:     true,
+		SpeculativeEntries:      []string{"req-1", "req-2"},
+		TotalSpeculativeEntries: 2,
+		MaxSpeculativeEntries:   maxDumpSpeculativeEntries,
+		BlockSizeTokens:         64,
+	}, state)
+}
+
+func TestDumpStateEmpty(t *testing.T) {
+	p := &Producer{}
+
+	payload, err := p.DumpState()
+	require.NoError(t, err)
+	assert.True(t, json.Valid(payload))
+	// Empty lists serialize as [] not null, matching the documented response shape.
+	assert.Contains(t, string(payload), `"subscribers":[]`)
+	assert.Contains(t, string(payload), `"speculativeEntries":[]`)
+
+	var state precisePrefixState
+	require.NoError(t, json.Unmarshal(payload, &state))
+	assert.Equal(t, precisePrefixState{
+		Subscribers:           []string{},
+		SpeculativeEntries:    []string{},
+		MaxSubscribers:        maxDumpSubscribers,
+		MaxSpeculativeEntries: maxDumpSpeculativeEntries,
+	}, state)
+}
+
+func TestDumpStateCaps(t *testing.T) {
+	specCache := ttlcache.New[string, *speculativeEntries]()
+	ids := make([]string, 0, maxDumpSubscribers+5)
+	for i := 0; i < maxDumpSubscribers+5; i++ {
+		ids = append(ids, fmt.Sprintf("ns/pod-%03d", i))
+	}
+	for i := 0; i < maxDumpSpeculativeEntries+7; i++ {
+		specCache.Set(fmt.Sprintf("req-%04d", i), &speculativeEntries{}, ttlcache.NoTTL)
+	}
+
+	p := &Producer{
+		subscribersManager: &fakeSubscriberManager{ids: ids},
+		speculativeCache:   specCache,
+	}
+
+	payload, err := p.DumpState()
+	require.NoError(t, err)
+
+	var state precisePrefixState
+	require.NoError(t, json.Unmarshal(payload, &state))
+	// Lists are capped, but the totals report the full counts so a consumer can
+	// tell the dump is partial from totalX > maxX.
+	assert.Len(t, state.Subscribers, maxDumpSubscribers)
+	assert.Equal(t, maxDumpSubscribers+5, state.TotalSubscribers)
+	assert.Len(t, state.SpeculativeEntries, maxDumpSpeculativeEntries)
+	assert.Equal(t, maxDumpSpeculativeEntries+7, state.TotalSpeculativeEntries)
+}

pkg/epp/framework/plugins/scheduling/scorer/tokenload/token_load.go

@@ -23,6 +23,7 @@ import (
 
 	"sigs.k8s.io/controller-runtime/pkg/log"
 
+	logutil "github.com/llm-d/llm-d-router/pkg/common/observability/logging"
 	fwkplugin "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin"
 	fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
 	attrconcurrency "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/concurrency"
@@ -122,7 +123,7 @@ func (s *TokenLoadScorer) Score(ctx context.Context, _ *fwksched.InferenceReques
 			score = 1.0 - (tokenLoad / s.queueThresholdTokens)
 		}
 		scores[endpoint] = score
-		logger.V(1).Info("TokenLoadScorer scoring", "endpoint", endpointID, "tokenLoad", tokenLoad, "score", score)
+		logger.V(logutil.DEBUG).Info("TokenLoadScorer scoring", "endpoint", endpointID, "tokenLoad", tokenLoad, "score", score)
 	}
 
 	return scores