refactor: Allow export prefix SchedulingContextState for use across plugins (#1063)

* refactor: Allow export prefix SchedulingContextState for use across plugins

Signed-off-by: Kfir Toledo <[email protected]>

* refactor: replace getPrefixState with generic ReadCycleStateAs

Signed-off-by: Kfir Toledo <[email protected]>

---------

Signed-off-by: Kfir Toledo <[email protected]>

Author: kfirtoledo

pkg/epp/scheduling/framework/plugins/multi/prefix/plugin.go

@@ -90,25 +90,25 @@ func (s ServerID) String() string {
 }
 
 // compile-time type validation
-var _ types.StateData = &schedulingContextState{}
+var _ types.StateData = &SchedulingContextState{}
 
-// This is the state of this plugin to be used during a scheduling cycle.
-type schedulingContextState struct {
+// SchedulingContextState is the state of this plugin to be used during a scheduling cycle.
+type SchedulingContextState struct {
 	// PrefixHashes is a list of prefix hashes of the request prompt broken into blocks.
 	PrefixHashes []BlockHash
 	// A map of server to its longest prefix cache match length.
 	PrefixCacheServers map[ServerID]int
 }
 
-func (s *schedulingContextState) Clone() types.StateData {
+func (s *SchedulingContextState) Clone() types.StateData {
 	prefixHashes := make([]BlockHash, len(s.PrefixHashes))
 	copy(prefixHashes, s.PrefixHashes)
 	prefixCacheServers := make(map[ServerID]int, len(s.PrefixCacheServers))
 	for key, value := range s.PrefixCacheServers {
 		prefixCacheServers[key] = value
 	}
 
-	return &schedulingContextState{
+	return &SchedulingContextState{
 		PrefixHashes:       prefixHashes,
 		PrefixCacheServers: prefixCacheServers,
 	}
@@ -171,7 +171,7 @@ func (m *Plugin) Score(ctx context.Context, cycleState *types.CycleState, reques
 	loggerTrace := log.FromContext(ctx).V(logutil.TRACE)
 	// pre score step, hashing prompt and find longest prefix match.
 	hashes := hashPrompt(ctx, request, m.HashBlockSize, m.MaxPrefixBlocksToMatch)
-	state := &schedulingContextState{
+	state := &SchedulingContextState{
 		PrefixHashes:       hashes,
 		PrefixCacheServers: m.matchLongestPrefix(ctx, hashes),
 	}
@@ -199,7 +199,7 @@ func (m *Plugin) Score(ctx context.Context, cycleState *types.CycleState, reques
 // PostCycle records in the plugin cache the result of the scheduling selection.
 func (m *Plugin) PostCycle(ctx context.Context, cycleState *types.CycleState, res *types.ProfileRunResult) {
 	targetPod := res.TargetPod.GetPod()
-	state, err := m.getPrefixState(cycleState)
+	state, err := types.ReadCycleStateKey[*SchedulingContextState](cycleState, PrefixCachePluginType)
 	if err != nil {
 		log.FromContext(ctx).Error(err, "failed to read prefix plugin cycle state")
 		return
@@ -235,22 +235,6 @@ func (m *Plugin) matchLongestPrefix(ctx context.Context, hashes []BlockHash) map
 	return res
 }
 
-// getPrefixState returns the cycle state as a schedulingContextState.
-func (m *Plugin) getPrefixState(cycleState *types.CycleState) (*schedulingContextState, error) {
-	prefixStateKey := types.StateKey(m.Type())
-	state, err := cycleState.Read(prefixStateKey)
-	if err != nil {
-		return nil, fmt.Errorf("failed reading %q from CycleState: %w", prefixStateKey, err)
-	}
-
-	prefixSchedulingState, ok := state.(*schedulingContextState)
-	if !ok {
-		return nil, fmt.Errorf("invalid Prefix state, got type %T", state)
-	}
-
-	return prefixSchedulingState, nil
-}
-
 // hashPrompt divides the prompt into blocks and calculate the prefix cache for each block.
 // hash(0) is the hash of the model name, since different models generally don't share prefix cache.
 // For block i, hash(i) = hash(block i content, hash(i-1)).

pkg/epp/scheduling/framework/plugins/multi/prefix/plugin_test.go

@@ -50,7 +50,7 @@ func TestPrefixPlugin(t *testing.T) {
 	}
 	cycleState1 := types.NewCycleState()
 	scores := plugin.Score(context.Background(), cycleState1, req1, pods)
-	state, err := plugin.getPrefixState(cycleState1)
+	state, err := types.ReadCycleStateKey[*SchedulingContextState](cycleState1, PrefixCachePluginType)
 	assert.NoError(t, err)
 	t.Logf("Hashes %+v, cached servers: %+v", state.PrefixHashes, state.PrefixCacheServers)
 	// Input size is 6, hash block size is 4, the last 2 characters are ignored.
@@ -71,7 +71,7 @@ func TestPrefixPlugin(t *testing.T) {
 	}
 	cycleState2 := types.NewCycleState()
 	scores = plugin.Score(context.Background(), cycleState2, req2, pods)
-	state, err = plugin.getPrefixState(cycleState2)
+	state, err = types.ReadCycleStateKey[*SchedulingContextState](cycleState2, PrefixCachePluginType)
 	assert.NoError(t, err)
 	t.Logf("Hashes %+v, cached servers: %+v", state.PrefixHashes, state.PrefixCacheServers)
 	// Input size is 6, hash block size is 4, the last 2 characters are ignored.
@@ -91,7 +91,7 @@ func TestPrefixPlugin(t *testing.T) {
 	}
 	cycleState3 := types.NewCycleState()
 	scores = plugin.Score(context.Background(), cycleState3, req3, pods)
-	state, err = plugin.getPrefixState(cycleState3)
+	state, err = types.ReadCycleStateKey[*SchedulingContextState](cycleState3, PrefixCachePluginType)
 	assert.NoError(t, err)
 	t.Logf("Hashes %+v, cached servers: %+v", state.PrefixHashes, state.PrefixCacheServers)
 	// Input size is 8, hash block size is 4, so 2 hashes will be calculated.
@@ -110,7 +110,7 @@ func TestPrefixPlugin(t *testing.T) {
 	}
 	cycleState4 := types.NewCycleState()
 	scores = plugin.Score(context.Background(), cycleState4, req4, pods)
-	state, err = plugin.getPrefixState(cycleState4)
+	state, err = types.ReadCycleStateKey[*SchedulingContextState](cycleState4, PrefixCachePluginType)
 	assert.NoError(t, err)
 	t.Logf("Hashes %+v, cached servers: %+v", state.PrefixHashes, state.PrefixCacheServers)
 	// Input size is 8, hash block size is 4, so 2 hashes will be calculated.
@@ -129,7 +129,7 @@ func TestPrefixPlugin(t *testing.T) {
 	}
 	cycleState5 := types.NewCycleState()
 	scores = plugin.Score(context.Background(), cycleState5, req5, pods)
-	state, err = plugin.getPrefixState(cycleState5)
+	state, err = types.ReadCycleStateKey[*SchedulingContextState](cycleState5, PrefixCachePluginType)
 	assert.NoError(t, err)
 	t.Logf("Hashes %+v, cached servers: %+v", state.PrefixHashes, state.PrefixCacheServers)
 	// Input size is 12, hash block size is 4, so 3 hashes will be calculated.
@@ -183,7 +183,7 @@ func BenchmarkPrefixPluginStress(b *testing.B) {
 		plugin.PostCycle(context.Background(), cycleState, &types.ProfileRunResult{TargetPod: pod})
 
 		// Second cycle: validate internal state
-		state, err := plugin.getPrefixState(cycleState)
+		state, err := types.ReadCycleStateKey[*SchedulingContextState](cycleState, PrefixCachePluginType)
 		assert.NoError(b, err)
 		expectedHashes := int(math.Min(float64(maxPrefixBlocks+1), float64(len(req.Prompt)/blockSize+1))) // the extra one is for the model.
 		assert.Equal(b, expectedHashes, len(state.PrefixHashes), "number of hashes is incorrect")

pkg/epp/scheduling/types/cycle_state.go

@@ -18,6 +18,7 @@ package types
 
 import (
 	"errors"
+	"fmt"
 	"sync"
 )
 
@@ -90,3 +91,21 @@ func (c *CycleState) Write(key StateKey, val StateData) {
 func (c *CycleState) Delete(key StateKey) {
 	c.storage.Delete(key)
 }
+
+// ReadCycleStateKey  retrieves data with the given key from CycleState and asserts it to type T.
+// Returns an error if the key is not found or the type assertion fails.
+func ReadCycleStateKey[T StateData](c *CycleState, key StateKey) (T, error) {
+	var zero T
+
+	raw, err := c.Read(key)
+	if err != nil {
+		return zero, err
+	}
+
+	val, ok := raw.(T)
+	if !ok {
+		return zero, fmt.Errorf("unexpected type for key %q: got %T", key, raw)
+	}
+
+	return val, nil
+}