Skip to content

Commit 6cad145

Browse files
LukeAVanDrieLukeAVanDrie
committed
split InFlightLoad into live and projected keys
Split InFlightLoad attribute into two separate keys: - InFlightLoadDataKey (dynamic, live endpoint load tokens/requests). - UncachedRequestTokensDataKey (static, request-specific projected load). This avoids overwriting dynamic attributes with static ones during scheduling. Signed-off-by: Luke Van Drie <lukevandrie@google.com>
1 parent 5713345 commit 6cad145

6 files changed

Lines changed: 163 additions & 78 deletions

File tree

‎pkg/epp/framework/plugins/datalayer/attribute/concurrency/data_types.go‎

Lines changed: 30 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -26,17 +26,16 @@ import (
2626
// load-aware scorers. Populated by InFlightLoadProducer.Produce on each
2727
// candidate endpoint at the start of every scheduling cycle.
2828
//
29-
// The snapshot is per-cycle (stored on the per-cycle endpoint clone produced
30-
// by fwksched.NewEndpoint) and combines two semantically distinct sources:
31-
// - Tokens / Requests: accumulated load from already in-flight requests,
32-
// read from the producer's persistent trackers.
33-
// - UncachedRequestTokens: the projected work this request would add to
34-
// the endpoint if scheduled here, computed fresh from the request being
35-
// scored and the endpoint's prefix-cache state. Not persisted.
29+
// Represents the live, real-time load on the endpoint (tokens and requests),
30+
// injected dynamically and never overwritten during a scheduling cycle.
3631
var InFlightLoadDataKey = plugin.NewDataKey("InFlightLoadDataKey", inflightloadconstants.InFlightLoadProducerType)
3732

33+
// UncachedRequestTokensDataKey carries the projected impact of the current request
34+
// on a specific endpoint. Populated statically during scheduling.
35+
var UncachedRequestTokensDataKey = plugin.NewDataKey("UncachedRequestTokensDataKey", inflightloadconstants.InFlightLoadProducerType)
36+
3837
// InFlightLoad captures the current real-time load of an endpoint as tracked
39-
// by the EPP, plus the projected impact of the request currently being scheduled.
38+
// by the EPP.
4039
type InFlightLoad struct {
4140
// Tokens is the in-flight token count this endpoint has committed to,
4241
// accumulated from past scheduling decisions. Updated by PreRequest (when
@@ -47,20 +46,6 @@ type InFlightLoad struct {
4746
// Requests is the in-flight request count this endpoint has committed to,
4847
// maintained with the same lifecycle as Tokens.
4948
Requests int64
50-
51-
// UncachedRequestTokens is a speculative projection: the work the request
52-
// currently being scheduled would add to this endpoint if it landed here.
53-
// Includes the uncached input portion (accounting for prefix-cache hits)
54-
// plus the estimated output when the producer is configured with
55-
// AddEstimatedOutputTokens=true. Computed fresh by Produce on every cycle
56-
// from the request being scored and this endpoint's prefix-cache state;
57-
// never committed to endpoint state and not decremented on stream end.
58-
// Zero when no request is in scope (e.g., background snapshots).
59-
//
60-
// Could equivalently be computed inside any scorer that needs it, but
61-
// it's produced here so every consumer reads the same value — single
62-
// source of truth for the current request's per-endpoint impact.
63-
UncachedRequestTokens int64
6449
}
6550

6651
// Clone returns an independent copy of the InFlightLoad. The value-copy
@@ -74,3 +59,26 @@ func (l *InFlightLoad) Clone() fwkdl.Cloneable {
7459
cp := *l
7560
return &cp
7661
}
62+
63+
// UncachedRequestTokens represents the projected impact of the current request
64+
// on a specific endpoint, populated statically during scheduling.
65+
type UncachedRequestTokens struct {
66+
// Tokens is a speculative projection: the work the request currently being
67+
// scheduled would add to this endpoint if it landed here. Includes the
68+
// uncached input portion (accounting for prefix-cache hits) plus the
69+
// estimated output when the producer is configured with
70+
// AddEstimatedOutputTokens=true. Computed fresh by Produce on every cycle
71+
// from the request being scored and this endpoint's prefix-cache state;
72+
// never committed to endpoint state and not decremented on stream end.
73+
// Zero when no request is in scope (e.g., background snapshots).
74+
Tokens int64
75+
}
76+
77+
// Clone returns an independent copy of the UncachedRequestTokens.
78+
func (u *UncachedRequestTokens) Clone() fwkdl.Cloneable {
79+
if u == nil {
80+
return nil
81+
}
82+
cp := *u
83+
return &cp
84+
}

‎pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/producer.go‎

Lines changed: 8 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,7 @@ func InFlightLoadProducerFactory(name string, decoder *json.Decoder, handle fwkp
8383
addEstimatedOutputTokens: cfg.AddEstimatedOutputTokens,
8484
dk: attrconcurrency.InFlightLoadDataKey.WithNonEmptyProducerName(name),
8585
prefixMatchInfoDK: attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(cfg.PrefixMatchInfoProducerName),
86+
uncachedRequestTokensDk: attrconcurrency.UncachedRequestTokensDataKey.WithNonEmptyProducerName(name),
8687
PluginState: fwkplugin.NewPluginState(ctx),
8788
}, nil
8889
}
@@ -105,6 +106,7 @@ type InFlightLoadProducer struct {
105106
PluginState *fwkplugin.PluginState
106107
dk fwkplugin.DataKey
107108
prefixMatchInfoDK fwkplugin.DataKey
109+
uncachedRequestTokensDk fwkplugin.DataKey
108110
}
109111

110112
// addedTokensEntry tracks a request's contribution to the global token and
@@ -220,20 +222,12 @@ func (p *InFlightLoadProducer) Produce(_ context.Context, request *fwksched.Infe
220222
if e == nil || e.GetMetadata() == nil {
221223
continue
222224
}
223-
endpointID := e.GetMetadata().NamespacedName.String()
224-
225-
load := &attrconcurrency.InFlightLoad{
226-
Tokens: p.tokenTracker.get(endpointID),
227-
Requests: p.requestTracker.get(endpointID),
228-
}
229225
if request != nil {
230-
// Project this request's additional work onto the endpoint: uncached
231-
// input tokens (per its prefix-cache state) plus estimated output
232-
// when the producer is configured to include it. Per-cycle, not
233-
// persisted in the tracker — that happens only on PreRequest.
234-
load.UncachedRequestTokens = p.estimateRequestTokens(e, inputTokens)
226+
tokens := p.estimateRequestTokens(e, inputTokens)
227+
e.Put(p.uncachedRequestTokensDk.String(), &attrconcurrency.UncachedRequestTokens{
228+
Tokens: tokens,
229+
})
235230
}
236-
e.Put(p.dk.String(), load)
237231
}
238232
return nil
239233
}
@@ -459,7 +453,8 @@ func nonNeg(v int64) int64 {
459453

460454
func (p *InFlightLoadProducer) Produces() map[fwkplugin.DataKey]any {
461455
return map[fwkplugin.DataKey]any{
462-
p.dk: attrconcurrency.InFlightLoad{},
456+
p.dk: attrconcurrency.InFlightLoad{},
457+
p.uncachedRequestTokensDk: attrconcurrency.UncachedRequestTokens{},
463458
}
464459
}
465460

‎pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/producer_test.go‎

Lines changed: 53 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -107,25 +107,68 @@ func TestInFlightLoadProducer_Produce(t *testing.T) {
107107
producer := newTestProducer(t)
108108

109109
endpointName := "test-endpoint"
110-
endpointID := fullEndpointName(endpointName)
110+
endpoint := newStubSchedulingEndpoint(endpointName)
111+
endpoints := []fwksched.Endpoint{endpoint}
111112

112-
// Mock some initial load
113-
producer.requestTracker.add(endpointID, 5)
114-
producer.tokenTracker.add(endpointID, 500)
113+
// 1. Produce with nil request -> should not put anything
114+
err := producer.Produce(context.Background(), nil, endpoints)
115+
require.NoError(t, err)
116+
_, ok := endpoint.Get(producer.uncachedRequestTokensDk.String())
117+
require.False(t, ok)
118+
119+
// 2. Produce with request -> should put UncachedRequestTokens
120+
req := makeTokenRequest("req1", 4) // 4 input tokens -> 10 total (with output)
121+
err = producer.Produce(context.Background(), req, endpoints)
122+
123+
require.NoError(t, err)
124+
125+
val, ok := endpoint.Get(producer.uncachedRequestTokensDk.String())
126+
require.True(t, ok)
127+
uncached := val.(*attrconcurrency.UncachedRequestTokens)
128+
require.Equal(t, int64(10), uncached.Tokens)
129+
130+
// Verify that InFlightLoad was NOT put/overwritten by Produce
131+
_, ok = endpoint.Get(producer.dk.String())
132+
require.False(t, ok, "InFlightLoad should not be populated by Produce")
133+
}
134+
135+
func TestInFlightLoadProducer_Extract(t *testing.T) {
136+
t.Parallel()
115137

138+
producer := newTestProducer(t)
139+
endpointName := "test-endpoint"
140+
endpointID := fullEndpointName(endpointName)
141+
142+
endpoint := newStubSchedulingEndpoint(endpointName)
116143
ctx := context.Background()
117-
endpoints := []fwksched.Endpoint{newStubSchedulingEndpoint(endpointName)}
118144

119-
err := producer.Produce(ctx, nil, endpoints)
145+
// Simulate Add event
146+
err := producer.Extract(ctx, datalayer.EndpointEvent{
147+
Type: datalayer.EventAddOrUpdate,
148+
Endpoint: endpoint,
149+
})
120150
require.NoError(t, err)
121151

122-
// Verify AttributeMap population
152+
// Verify dynamic attribute is registered
123153
key := producer.dk.String()
124-
val, ok := endpoints[0].Get(key)
154+
val, ok := endpoint.Get(key)
125155
require.True(t, ok)
156+
157+
// Verify initial values (should be 0)
126158
load := val.(*attrconcurrency.InFlightLoad)
127-
require.Equal(t, int64(5), load.Requests)
128-
require.Equal(t, int64(500), load.Tokens)
159+
require.Equal(t, int64(0), load.Requests)
160+
require.Equal(t, int64(0), load.Tokens)
161+
162+
// Update trackers
163+
producer.requestTracker.add(endpointID, 5)
164+
producer.tokenTracker.add(endpointID, 500)
165+
166+
// Verify values are updated dynamically without calling Produce
167+
val2, ok := endpoint.Get(key)
168+
require.True(t, ok)
169+
load2 := val2.(*attrconcurrency.InFlightLoad)
170+
require.Equal(t, int64(5), load2.Requests)
171+
require.Equal(t, int64(500), load2.Tokens)
129172
}
130173

131174
func TestInFlightLoadProducer_Lifecycle(t *testing.T) {

‎pkg/epp/framework/plugins/scheduling/scorer/activerequest/active_request_test.go‎

Lines changed: 36 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -20,18 +20,41 @@ import (
2020

2121
func float64Ptr(v float64) *float64 { return &v }
2222

23-
func newTestEndpoint(name string, queueSize int) scheduling.Endpoint {
24-
return scheduling.NewEndpoint(
25-
&datalayer.EndpointMetadata{NamespacedName: k8stypes.NamespacedName{Name: name, Namespace: "default"}},
26-
&datalayer.Metrics{
23+
type stubEndpoint struct {
24+
metadata *datalayer.EndpointMetadata
25+
metrics *datalayer.Metrics
26+
attr datalayer.AttributeMap
27+
}
28+
29+
func newStubEndpoint(name string, queueSize int) *stubEndpoint {
30+
return &stubEndpoint{
31+
metadata: &datalayer.EndpointMetadata{NamespacedName: k8stypes.NamespacedName{Name: name, Namespace: "default"}},
32+
metrics: &datalayer.Metrics{
2733
WaitingQueueSize: queueSize,
2834
},
29-
nil,
30-
)
35+
attr: datalayer.NewAttributes(),
36+
}
37+
}
38+
39+
func (f *stubEndpoint) GetMetadata() *datalayer.EndpointMetadata { return f.metadata }
40+
func (f *stubEndpoint) UpdateMetadata(*datalayer.EndpointMetadata) {}
41+
func (f *stubEndpoint) GetMetrics() *datalayer.Metrics { return f.metrics }
42+
func (f *stubEndpoint) UpdateMetrics(*datalayer.Metrics) {}
43+
func (f *stubEndpoint) GetAttributes() datalayer.AttributeMap { return f.attr }
44+
func (f *stubEndpoint) String() string { return f.metadata.NamespacedName.String() }
45+
func (f *stubEndpoint) Put(key string, val datalayer.Cloneable) { f.attr.Put(key, val) }
46+
func (f *stubEndpoint) Get(key string) (datalayer.Cloneable, bool) {
47+
return f.attr.Get(key)
48+
}
49+
func (f *stubEndpoint) Keys() []string { return f.attr.Keys() }
50+
func (f *stubEndpoint) Clone() datalayer.AttributeMap { return f.attr.Clone() }
51+
52+
func newTestEndpoint(name string, queueSize int) scheduling.Endpoint {
53+
return newStubEndpoint(name, queueSize)
3154
}
3255

3356
func newTestEndpointWithLoad(name string, requests int64) scheduling.Endpoint {
34-
ep := newTestEndpoint(name, 0)
57+
ep := newStubEndpoint(name, 0)
3558
ep.Put(attrconcurrency.InFlightLoadDataKey.String(), &attrconcurrency.InFlightLoad{Requests: requests})
3659
return ep
3760
}
@@ -106,6 +129,12 @@ func TestActiveRequestScorer_UsesInFlightLoadProducerLifecycle(t *testing.T) {
106129
podB := newTestEndpoint("pod-b", 0)
107130
endpoints := []scheduling.Endpoint{podA, podB}
108131

132+
// Simulate Extract to inject the dynamic attribute
133+
err = producer.Extract(ctx, datalayer.EndpointEvent{Type: datalayer.EventAddOrUpdate, Endpoint: podA.(datalayer.Endpoint)})
134+
require.NoError(t, err)
135+
err = producer.Extract(ctx, datalayer.EndpointEvent{Type: datalayer.EventAddOrUpdate, Endpoint: podB.(datalayer.Endpoint)})
136+
require.NoError(t, err)
137+
109138
req := &scheduling.InferenceRequest{RequestID: "req-1", RequestSizeBytes: 4}
110139
result := &scheduling.SchedulingResult{
111140
PrimaryProfileName: "default",

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

Lines changed: 22 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -45,9 +45,10 @@ type Config struct {
4545
var _ fwksched.Scorer = &TokenLoadScorer{}
4646

4747
type TokenLoadScorer struct {
48-
typedName fwkplugin.TypedName
49-
queueThresholdTokens float64
50-
inFlightLoadDataKey fwkplugin.DataKey
48+
typedName fwkplugin.TypedName
49+
queueThresholdTokens float64
50+
inFlightLoadDataKey fwkplugin.DataKey
51+
uncachedRequestTokensDataKey fwkplugin.DataKey
5152
}
5253

5354
func TokenLoadScorerFactory(name string, params *json.Decoder, _ fwkplugin.Handle) (fwkplugin.Plugin, error) {
@@ -64,9 +65,10 @@ func TokenLoadScorerFactory(name string, params *json.Decoder, _ fwkplugin.Handl
6465
}
6566

6667
return &TokenLoadScorer{
67-
typedName: fwkplugin.TypedName{Type: TokenLoadScorerType, Name: name},
68-
queueThresholdTokens: float64(cfg.QueueThresholdTokens),
69-
inFlightLoadDataKey: attrconcurrency.InFlightLoadDataKey.WithNonEmptyProducerName(cfg.InFlightLoadProducerName),
68+
typedName: fwkplugin.TypedName{Type: TokenLoadScorerType, Name: name},
69+
queueThresholdTokens: float64(cfg.QueueThresholdTokens),
70+
inFlightLoadDataKey: attrconcurrency.InFlightLoadDataKey.WithNonEmptyProducerName(cfg.InFlightLoadProducerName),
71+
uncachedRequestTokensDataKey: attrconcurrency.UncachedRequestTokensDataKey.WithNonEmptyProducerName(cfg.InFlightLoadProducerName),
7072
}, nil
7173
}
7274

@@ -80,7 +82,10 @@ func (s *TokenLoadScorer) Category() fwksched.ScorerCategory {
8082

8183
func (s *TokenLoadScorer) Consumes() fwkplugin.DataDependencies {
8284
return fwkplugin.DataDependencies{
83-
Required: map[fwkplugin.DataKey]any{s.inFlightLoadDataKey: attrconcurrency.InFlightLoad{}},
85+
Required: map[fwkplugin.DataKey]any{
86+
s.inFlightLoadDataKey: attrconcurrency.InFlightLoad{},
87+
s.uncachedRequestTokensDataKey: attrconcurrency.UncachedRequestTokens{},
88+
},
8489
}
8590
}
8691

@@ -92,14 +97,20 @@ func (s *TokenLoadScorer) Score(ctx context.Context, _ *fwksched.InferenceReques
9297
endpointID := endpoint.GetMetadata().NamespacedName.String()
9398
tokenLoad := 0.0
9499

95-
// Single read: accumulated in-flight load plus the projected impact
96-
// of the request being scored, both carried on the same InFlightLoad
97-
// struct populated by InFlightLoadProducer.Produce.
100+
// Read both accumulated in-flight load and the projected impact of the
101+
// request being scored, which are now carried on separate attributes.
102+
var tokens int64
98103
if val, ok := endpoint.Get(s.inFlightLoadDataKey.String()); ok {
99104
if load, ok := val.(*attrconcurrency.InFlightLoad); ok && load != nil {
100-
tokenLoad = float64(load.Tokens + load.UncachedRequestTokens)
105+
tokens += load.Tokens
101106
}
102107
}
108+
if val, ok := endpoint.Get(s.uncachedRequestTokensDataKey.String()); ok {
109+
if uncached, ok := val.(*attrconcurrency.UncachedRequestTokens); ok && uncached != nil {
110+
tokens += uncached.Tokens
111+
}
112+
}
113+
tokenLoad = float64(tokens)
103114

104115
score := 0.0
105116
if tokenLoad <= 0 {

‎pkg/epp/framework/plugins/scheduling/scorer/tokenload/token_load_test.go‎

Lines changed: 14 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -33,9 +33,10 @@ func TestTokenLoadScorer(t *testing.T) {
3333
threshold := 1000.0
3434

3535
scorer := &TokenLoadScorer{
36-
typedName: fwkplugin.TypedName{Type: TokenLoadScorerType, Name: TokenLoadScorerType},
37-
queueThresholdTokens: threshold,
38-
inFlightLoadDataKey: attrconcurrency.InFlightLoadDataKey.WithNonEmptyProducerName(""),
36+
typedName: fwkplugin.TypedName{Type: TokenLoadScorerType, Name: TokenLoadScorerType},
37+
queueThresholdTokens: threshold,
38+
inFlightLoadDataKey: attrconcurrency.InFlightLoadDataKey.WithNonEmptyProducerName(""),
39+
uncachedRequestTokensDataKey: attrconcurrency.UncachedRequestTokensDataKey.WithNonEmptyProducerName(""),
3940
}
4041

4142
t.Run("in-flight load only", func(t *testing.T) {
@@ -74,20 +75,16 @@ func TestTokenLoadScorer(t *testing.T) {
7475
}
7576

7677
// pod1: 0 in-flight + 250 current = 250. Score = 1 - 250/1000 = 0.75
77-
endpoints[0].Put(scorer.inFlightLoadDataKey.String(), &attrconcurrency.InFlightLoad{
78-
Tokens: 0,
79-
UncachedRequestTokens: 250,
80-
})
78+
endpoints[0].Put(scorer.inFlightLoadDataKey.String(), &attrconcurrency.InFlightLoad{Tokens: 0})
79+
endpoints[0].Put(scorer.uncachedRequestTokensDataKey.String(), &attrconcurrency.UncachedRequestTokens{Tokens: 250})
80+
8181
// pod2: 250 in-flight + 250 current = 500. Score = 1 - 500/1000 = 0.5
82-
endpoints[1].Put(scorer.inFlightLoadDataKey.String(), &attrconcurrency.InFlightLoad{
83-
Tokens: 250,
84-
UncachedRequestTokens: 250,
85-
})
82+
endpoints[1].Put(scorer.inFlightLoadDataKey.String(), &attrconcurrency.InFlightLoad{Tokens: 250})
83+
endpoints[1].Put(scorer.uncachedRequestTokensDataKey.String(), &attrconcurrency.UncachedRequestTokens{Tokens: 250})
84+
8685
// pod3: 750 in-flight + 250 current = 1000. Score = 1 - 1000/1000 = 0.0
87-
endpoints[2].Put(scorer.inFlightLoadDataKey.String(), &attrconcurrency.InFlightLoad{
88-
Tokens: 750,
89-
UncachedRequestTokens: 250,
90-
})
86+
endpoints[2].Put(scorer.inFlightLoadDataKey.String(), &attrconcurrency.InFlightLoad{Tokens: 750})
87+
endpoints[2].Put(scorer.uncachedRequestTokensDataKey.String(), &attrconcurrency.UncachedRequestTokens{Tokens: 250})
9188

9289
scores := scorer.Score(context.Background(), &fwksched.InferenceRequest{}, endpoints)
9390

@@ -115,6 +112,8 @@ func TestTokenLoadScorer(t *testing.T) {
115112

116113
var nilLoad *attrconcurrency.InFlightLoad
117114
endpoints[0].Put(scorer.inFlightLoadDataKey.String(), nilLoad)
115+
var nilUncached *attrconcurrency.UncachedRequestTokens
116+
endpoints[0].Put(scorer.uncachedRequestTokensDataKey.String(), nilUncached)
118117

119118
// Typed nil; should score as if 0 token load (no panic)
120119
scores := scorer.Score(context.Background(), &fwksched.InferenceRequest{}, endpoints)

0 commit comments

Comments
 (0)