Skip to content

Commit 23908ca

Browse files
rahulgurnanirahulgurnani
committed
Update test and remove duplicate AttributeMap
1 parent a7354ca commit 23908ca

File tree

7 files changed

+91
-109
lines changed

7 files changed

+91
-109
lines changed

‎pkg/epp/backend/metrics/fake.go‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,9 @@ import (
3232

3333
// FakePodMetrics is an implementation of PodMetrics that doesn't run the async refresh loop.
3434
type FakePodMetrics struct {
35-
Pod *backend.Pod
36-
Metrics *MetricsState
35+
Pod *backend.Pod
36+
Metrics *MetricsState
37+
Attributes *datalayer.Attributes
3738
}
3839

3940
func (fpm *FakePodMetrics) String() string {
@@ -51,6 +52,9 @@ func (fpm *FakePodMetrics) GetMetrics() *MetricsState {
5152
func (fpm *FakePodMetrics) UpdatePod(pod *datalayer.PodInfo) {
5253
fpm.Pod = pod
5354
}
55+
func (fpm *FakePodMetrics) GetAttributes() *datalayer.Attributes {
56+
return fpm.Attributes
57+
}
5458

5559
func (*FakePodMetrics) Put(string, datalayer.Cloneable) {}
5660
func (*FakePodMetrics) Get(string) (datalayer.Cloneable, bool) { return nil, false }

‎pkg/epp/backend/metrics/pod_metrics.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,9 @@ func (pm *podMetrics) stopRefreshLoop() {
126126
func (*podMetrics) Put(string, datalayer.Cloneable) {}
127127
func (*podMetrics) Get(string) (datalayer.Cloneable, bool) { return nil, false }
128128
func (*podMetrics) Keys() []string { return nil }
129+
func (*podMetrics) GetAttributes() *datalayer.Attributes {
130+
return nil
131+
}
129132

130133
func (pm *podMetrics) UpdateMetrics(updated *MetricsState) {
131134
updated.UpdateTime = time.Now()

‎pkg/epp/backend/metrics/types.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,4 +75,8 @@ func (f *PodMetricsFactory) ReleaseEndpoint(ep PodMetrics) {
7575
}
7676
}
7777

78+
func (f *PodMetricsFactory) GetAttributes() *datalayer.Attributes {
79+
return datalayer.NewAttributes()
80+
}
81+
7882
type PodMetrics = datalayer.Endpoint

‎pkg/epp/datalayer/endpoint.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import (
2525
type EndpointPodState interface {
2626
GetPod() *PodInfo
2727
UpdatePod(*PodInfo)
28+
GetAttributes() *Attributes
2829
}
2930

3031
// EndpointMetricsState allows management of the Metrics related attributes.
@@ -89,6 +90,10 @@ func (srv *ModelServer) Keys() []string {
8990
return srv.attributes.Keys()
9091
}
9192

93+
func (srv *ModelServer) GetAttributes() *Attributes {
94+
return srv.attributes
95+
}
96+
9297
func (srv *ModelServer) Clone() *ModelServer {
9398
clone := &ModelServer{
9499
attributes: srv.attributes.Clone(),

‎pkg/epp/requestcontrol/director.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import (
3333
"sigs.k8s.io/gateway-api-inference-extension/apix/v1alpha2"
3434
"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend"
3535
backendmetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/metrics"
36+
"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datalayer"
3637
"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/handlers"
3738
"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metadata"
3839
"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metrics"
@@ -276,7 +277,11 @@ func (d *Director) prepareRequest(ctx context.Context, reqCtx *handlers.RequestC
276277
func (d *Director) toSchedulerPodMetrics(pods []backendmetrics.PodMetrics) []schedulingtypes.Pod {
277278
pm := make([]schedulingtypes.Pod, len(pods))
278279
for i, pod := range pods {
279-
pm[i] = &schedulingtypes.PodMetrics{Pod: pod.GetPod().Clone(), MetricsState: pod.GetMetrics().Clone()}
280+
if pod.GetAttributes() != nil {
281+
pm[i] = &schedulingtypes.PodMetrics{Pod: pod.GetPod().Clone(), MetricsState: pod.GetMetrics().Clone(), AttributeMap: pod.GetAttributes().Clone()}
282+
} else {
283+
pm[i] = &schedulingtypes.PodMetrics{Pod: pod.GetPod().Clone(), MetricsState: pod.GetMetrics().Clone(), AttributeMap: datalayer.NewAttributes()}
284+
}
280285
}
281286

282287
return pm

‎pkg/epp/requestcontrol/director_test.go‎

Lines changed: 63 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ import (
3737
"sigs.k8s.io/gateway-api-inference-extension/apix/v1alpha2"
3838
"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend"
3939
backendmetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/metrics"
40+
"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datalayer"
4041
"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datastore"
4142
"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/handlers"
4243
"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metadata"
@@ -48,6 +49,10 @@ import (
4849
testutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/testing"
4950
)
5051

52+
const (
53+
mockProducedDataKey = "producedDataKey"
54+
)
55+
5156
// --- Mocks ---
5257

5358
type mockAdmissionController struct {
@@ -66,9 +71,16 @@ func (m *mockAdmissionController) Admit(
6671
type mockScheduler struct {
6772
scheduleResults *schedulingtypes.SchedulingResult
6873
scheduleErr error
74+
dataProduced bool // denotes whether data production is expected.
6975
}
7076

71-
func (m *mockScheduler) Schedule(_ context.Context, _ *schedulingtypes.LLMRequest, _ []schedulingtypes.Pod) (*schedulingtypes.SchedulingResult, error) {
77+
func (m *mockScheduler) Schedule(_ context.Context, _ *schedulingtypes.LLMRequest, pods []schedulingtypes.Pod) (*schedulingtypes.SchedulingResult, error) {
78+
if pods != nil && m.dataProduced {
79+
data, ok := pods[0].Get(mockProducedDataKey)
80+
if !ok || data.(mockProducedDataType).value != 42 {
81+
return nil, fmt.Errorf("expected produced data not found in pod")
82+
}
83+
}
7284
return m.scheduleResults, m.scheduleErr
7385
}
7486

@@ -94,11 +106,10 @@ func (ds *mockDatastore) PodList(predicate func(backendmetrics.PodMetrics) bool)
94106
}
95107

96108
type mockDataProducerPlugin struct {
97-
tn plugins.TypedName
98-
prepareRequestDataCalled bool
109+
tn plugins.TypedName
99110
}
100111

101-
func newmockDataProducerPlugin(name string) *mockDataProducerPlugin {
112+
func newMockDataProducerPlugin(name string) *mockDataProducerPlugin {
102113
return &mockDataProducerPlugin{
103114
tn: plugins.TypedName{Type: "mock-prepare-request-data", Name: name},
104115
}
@@ -109,33 +120,44 @@ func (m *mockDataProducerPlugin) TypedName() plugins.TypedName {
109120
}
110121

111122
func (m *mockDataProducerPlugin) Produces() map[string]any {
112-
return map[string]any{}
123+
// Produces data of type int, 0 denotes it is int.
124+
return map[string]any{mockProducedDataKey: 0}
113125
}
114126

115127
func (m *mockDataProducerPlugin) PrepareRequestData(ctx context.Context, request *schedulingtypes.LLMRequest, pods []schedulingtypes.Pod) {
116-
m.prepareRequestDataCalled = true
128+
pods[0].Put(mockProducedDataKey, mockProducedDataType{value: 42})
117129
}
118130

119-
type mockAdmitRequestPlugins struct {
120-
tn plugins.TypedName
131+
type mockAdmitRequestPlugin struct {
132+
tn plugins.TypedName
133+
// TODO: Replace this will admission control.
121134
admitRequestCalled bool
122135
}
123136

124-
func newmockAdmitRequestPlugins(name string) *mockAdmitRequestPlugins {
125-
return &mockAdmitRequestPlugins{
137+
func newmockAdmitRequestPlugin(name string) *mockAdmitRequestPlugin {
138+
return &mockAdmitRequestPlugin{
126139
tn: plugins.TypedName{Type: "mock-admit-data", Name: name},
127140
}
128141
}
129142

130-
func (m *mockAdmitRequestPlugins) TypedName() plugins.TypedName {
143+
func (m *mockAdmitRequestPlugin) TypedName() plugins.TypedName {
131144
return m.tn
132145
}
133146

134-
func (m *mockAdmitRequestPlugins) AdmitRequest(ctx context.Context, request *schedulingtypes.LLMRequest, pods []schedulingtypes.Pod) error {
147+
func (m *mockAdmitRequestPlugin) AdmitRequest(ctx context.Context, request *schedulingtypes.LLMRequest, pods []schedulingtypes.Pod) error {
135148
m.admitRequestCalled = true
136149
return nil
137150
}
138151

152+
type mockProducedDataType struct {
153+
value int
154+
}
155+
156+
// Clone implements types.Cloneable.
157+
func (m mockProducedDataType) Clone() datalayer.Cloneable {
158+
return mockProducedDataType{value: m.value}
159+
}
160+
139161
func TestDirector_HandleRequest(t *testing.T) {
140162
ctx := logutil.NewTestLoggerIntoContext(context.Background())
141163

@@ -210,6 +232,7 @@ func TestDirector_HandleRequest(t *testing.T) {
210232
TargetPods: []schedulingtypes.Pod{
211233
&schedulingtypes.ScoredPod{
212234
Pod: &schedulingtypes.PodMetrics{
235+
AttributeMap: datalayer.NewAttributes(),
213236
Pod: &backend.Pod{
214237
Address: "192.168.1.100",
215238
Port: "8000",
@@ -220,6 +243,7 @@ func TestDirector_HandleRequest(t *testing.T) {
220243
},
221244
&schedulingtypes.ScoredPod{
222245
Pod: &schedulingtypes.PodMetrics{
246+
AttributeMap: datalayer.NewAttributes(),
223247
Pod: &backend.Pod{
224248
Address: "192.168.2.100",
225249
Port: "8000",
@@ -230,6 +254,7 @@ func TestDirector_HandleRequest(t *testing.T) {
230254
},
231255
&schedulingtypes.ScoredPod{
232256
Pod: &schedulingtypes.PodMetrics{
257+
AttributeMap: datalayer.NewAttributes(),
233258
Pod: &backend.Pod{
234259
Address: "192.168.4.100",
235260
Port: "8000",
@@ -245,19 +270,18 @@ func TestDirector_HandleRequest(t *testing.T) {
245270
}
246271

247272
tests := []struct {
248-
name string
249-
reqBodyMap map[string]any
250-
mockAdmissionController *mockAdmissionController
251-
inferenceObjectiveName string
252-
schedulerMockSetup func(m *mockScheduler)
253-
wantErrCode string // Expected errutil code string
254-
wantReqCtx *handlers.RequestContext // Fields to check in the returned RequestContext
255-
wantMutatedBodyModel string // Expected model in reqCtx.Request.Body after PostDispatch
256-
targetModelName string // Expected model name after target model resolution
257-
prepareRequestDataCalled bool
258-
admitRequestCalled bool
259-
dataProducerPlugins *mockDataProducerPlugin
260-
admitRequestPlugins *mockAdmitRequestPlugins
273+
name string
274+
reqBodyMap map[string]any
275+
mockAdmissionController *mockAdmissionController
276+
inferenceObjectiveName string
277+
schedulerMockSetup func(m *mockScheduler)
278+
wantErrCode string // Expected errutil code string
279+
wantReqCtx *handlers.RequestContext // Fields to check in the returned RequestContext
280+
wantMutatedBodyModel string // Expected model in reqCtx.Request.Body after PostDispatch
281+
targetModelName string // Expected model name after target model resolution
282+
admitRequestCalled bool
283+
dataProducerPlugin *mockDataProducerPlugin
284+
admitRequestPlugin *mockAdmitRequestPlugin
261285
}{
262286
{
263287
name: "successful completions request",
@@ -326,6 +350,7 @@ func TestDirector_HandleRequest(t *testing.T) {
326350
mockAdmissionController: &mockAdmissionController{admitErr: nil},
327351
schedulerMockSetup: func(m *mockScheduler) {
328352
m.scheduleResults = defaultSuccessfulScheduleResults
353+
m.dataProduced = true
329354
},
330355
wantReqCtx: &handlers.RequestContext{
331356
TargetModelName: model,
@@ -337,10 +362,9 @@ func TestDirector_HandleRequest(t *testing.T) {
337362
},
338363
TargetEndpoint: "192.168.1.100:8000,192.168.2.100:8000,192.168.4.100:8000",
339364
},
340-
wantMutatedBodyModel: model,
341-
targetModelName: model,
342-
prepareRequestDataCalled: true,
343-
dataProducerPlugins: newmockDataProducerPlugin("test-plugin"),
365+
wantMutatedBodyModel: model,
366+
targetModelName: model,
367+
dataProducerPlugin: newMockDataProducerPlugin("test-plugin"),
344368
},
345369
{
346370
name: "successful chat completions request with admit request plugins",
@@ -370,7 +394,7 @@ func TestDirector_HandleRequest(t *testing.T) {
370394
wantMutatedBodyModel: model,
371395
targetModelName: model,
372396
admitRequestCalled: true,
373-
admitRequestPlugins: newmockAdmitRequestPlugins("test-plugin"),
397+
admitRequestPlugin: newmockAdmitRequestPlugin("test-plugin"),
374398
},
375399
{
376400
name: "successful chat completions request with multiple messages",
@@ -522,11 +546,11 @@ func TestDirector_HandleRequest(t *testing.T) {
522546
test.schedulerMockSetup(mockSched)
523547
}
524548
config := NewConfig()
525-
if test.dataProducerPlugins != nil {
526-
config = config.WithDataProducers(test.dataProducerPlugins)
549+
if test.dataProducerPlugin != nil {
550+
config = config.WithDataProducers(test.dataProducerPlugin)
527551
}
528-
if test.admitRequestPlugins != nil {
529-
config = config.WithAdmitRequestPlugins(test.admitRequestPlugins)
552+
if test.admitRequestPlugin != nil {
553+
config = config.WithAdmitRequestPlugins(test.admitRequestPlugin)
530554
}
531555
director := NewDirectorWithConfig(ds, mockSched, test.mockAdmissionController, config)
532556

@@ -572,11 +596,12 @@ func TestDirector_HandleRequest(t *testing.T) {
572596
assert.Equal(t, test.wantMutatedBodyModel, returnedReqCtx.Request.Body["model"],
573597
"Mutated reqCtx.Request.Body model mismatch")
574598
}
575-
if test.admitRequestPlugins != nil {
576-
assert.True(t, test.admitRequestPlugins.admitRequestCalled, "AdmitRequest not called")
599+
if test.admitRequestPlugin != nil {
600+
assert.True(t, test.admitRequestPlugin.admitRequestCalled, "AdmitRequest not called")
577601
}
578-
if test.dataProducerPlugins != nil {
579-
assert.True(t, test.dataProducerPlugins.prepareRequestDataCalled, "PrepareRequestData not called")
602+
if test.dataProducerPlugin != nil && mockSched.dataProduced {
603+
// TODO: Test if the AttributeMap in pod got updated correctly.
604+
580605
}
581606
})
582607
}

0 commit comments

Comments
 (0)