From 2b351519652973c159fe4daad9027534e2c70686 Mon Sep 17 00:00:00 2001 From: szedan Date: Wed, 12 Aug 2026 11:37:44 +0300 Subject: [PATCH 1/2] Replace CycleState with request attributes (#288) Store plugin-shared data directly on InferenceRequest via a sync.Map attribute store, matching EPP's pattern. Pass request to response plugin signatures so they can read attributes set during request phase. Delete cycle_state.go entirely. Signed-off-by: szedan --- pkg/config/loader/configloader_test.go | 12 +-- .../interface/datalayer/datasource/types.go | 12 ++- .../interface/modelselector/plugins.go | 6 +- .../interface/modelselector/types.go | 3 +- pkg/framework/interface/plugin/cycle_state.go | 87 ------------------- .../interface/requesthandling/plugins.go | 12 +-- .../interface/requesthandling/types.go | 28 ++++++ .../modelselector/filter/modelgroup/filter.go | 2 +- .../filter/modelgroup/filter_test.go | 12 +-- .../modelselector/picker/maxscore/picker.go | 2 +- .../picker/maxscore/picker_test.go | 3 +- .../modelselector/picker/random/picker.go | 2 +- .../picker/random/picker_test.go | 5 +- .../picker/weightedrandom/picker.go | 4 +- .../picker/weightedrandom/picker_test.go | 9 +- .../modelselector/scorer/costaware/plugin.go | 2 +- .../scorer/costaware/plugin_test.go | 4 +- .../modelselector/scorer/costguard/plugin.go | 2 +- .../scorer/costguard/plugin_test.go | 19 ++-- .../scorer/inflightrequests/plugin.go | 2 +- .../scorer/sessionaffinity/plugin.go | 11 ++- .../scorer/sessionaffinity/plugin_test.go | 80 ++++++++--------- .../base_model_to_header.go | 2 +- .../base_model_to_header_test.go | 4 +- .../bodyfieldtoheader/body_field_to_header.go | 2 +- .../body_field_to_header_test.go | 4 +- .../requesthandling/modelselector/plugin.go | 12 +-- .../modelselector/plugin_test.go | 16 ++-- .../single/single_profile_picker.go | 2 +- .../single/single_profile_picker_test.go | 2 +- .../modelnametoheader/plugin.go | 12 +-- .../modelnametoheader/plugin_test.go | 36 ++++---- pkg/handlers/request.go | 13 ++- pkg/handlers/request_test.go | 22 ++--- pkg/handlers/response.go | 21 +++-- pkg/handlers/response_test.go | 27 +++--- pkg/handlers/server.go | 18 ++-- pkg/handlers/server_test.go | 6 +- pkg/modelselector/model_selector_pipeline.go | 30 +++---- .../model_selector_pipeline_bench_test.go | 7 +- .../model_selector_pipeline_test.go | 10 +-- pkg/modelselector/modelselector.go | 5 +- pkg/modelselector/modelselector_test.go | 4 +- test/integration/body_mutation_test.go | 2 +- 44 files changed, 233 insertions(+), 343 deletions(-) delete mode 100644 pkg/framework/interface/plugin/cycle_state.go diff --git a/pkg/config/loader/configloader_test.go b/pkg/config/loader/configloader_test.go index c8041201..5db55779 100644 --- a/pkg/config/loader/configloader_test.go +++ b/pkg/config/loader/configloader_test.go @@ -475,7 +475,7 @@ type mockProfilePicker struct{ mockPlugin } // compile-time type assertion var _ requesthandling.ProfilePicker = &mockProfilePicker{} -func (m *mockProfilePicker) Pick(ctx context.Context, cycleState *plugin.CycleState, request *requesthandling.InferenceRequest, +func (m *mockProfilePicker) Pick(ctx context.Context, request *requesthandling.InferenceRequest, profiles map[string]*requesthandling.Profile) (*requesthandling.Profile, error) { return nil, nil } @@ -486,7 +486,7 @@ type mockRequestProcessor struct{ mockPlugin } // compile-time type assertion var _ requesthandling.RequestProcessor = &mockRequestProcessor{} -func (m *mockRequestProcessor) ProcessRequest(ctx context.Context, cycleState *plugin.CycleState, request *requesthandling.InferenceRequest) error { +func (m *mockRequestProcessor) ProcessRequest(ctx context.Context, request *requesthandling.InferenceRequest) error { return nil } @@ -496,7 +496,7 @@ type mockResponseProcessor struct{ mockPlugin } // compile-time type assertion var _ requesthandling.ResponseProcessor = &mockResponseProcessor{} -func (m *mockResponseProcessor) ProcessResponse(ctx context.Context, cycleState *plugin.CycleState, request *requesthandling.InferenceResponse) error { +func (m *mockResponseProcessor) ProcessResponse(ctx context.Context, request *requesthandling.InferenceRequest, response *requesthandling.InferenceResponse) error { return nil } @@ -506,7 +506,7 @@ type mockFilter struct{ mockPlugin } // compile-time type assertion var _ modelselector.Filter = &mockFilter{} -func (m *mockFilter) Filter(_ context.Context, _ *plugin.CycleState, _ *requesthandling.InferenceRequest, models []datalayer.Model) []datalayer.Model { +func (m *mockFilter) Filter(_ context.Context, _ *requesthandling.InferenceRequest, models []datalayer.Model) []datalayer.Model { return models } @@ -516,7 +516,7 @@ type mockScorer struct{ mockPlugin } // compile-time type assertion var _ modelselector.Scorer = &mockScorer{} -func (m *mockScorer) Score(ctx context.Context, cycleState *plugin.CycleState, request *requesthandling.InferenceRequest, models []datalayer.Model) map[datalayer.Model]float64 { +func (m *mockScorer) Score(ctx context.Context, request *requesthandling.InferenceRequest, models []datalayer.Model) map[datalayer.Model]float64 { return nil } @@ -549,7 +549,7 @@ type mockPicker struct{ mockPlugin } // compile-time type assertion var _ modelselector.Picker = &mockPicker{} -func (m *mockPicker) Pick(ctx context.Context, cycleState *plugin.CycleState, scoredModels []*modelselector.ScoredModel) *modelselector.PipelineRunResult { +func (m *mockPicker) Pick(ctx context.Context, scoredModels []*modelselector.ScoredModel) *modelselector.PipelineRunResult { return nil } diff --git a/pkg/framework/interface/datalayer/datasource/types.go b/pkg/framework/interface/datalayer/datasource/types.go index b22f36f3..016605ce 100644 --- a/pkg/framework/interface/datalayer/datasource/types.go +++ b/pkg/framework/interface/datalayer/datasource/types.go @@ -46,17 +46,15 @@ const ( // RequestPayload is the Payload for RequestEventType. type RequestPayload struct { - Request *requesthandling.InferenceRequest - CycleState *plugin.CycleState + Request *requesthandling.InferenceRequest } // ResponsePayload is the Payload for ResponseEventType. type ResponsePayload struct { - Request *requesthandling.InferenceRequest - Response *requesthandling.InferenceResponse - CycleState *plugin.CycleState - Duration time.Duration - TTFT time.Duration + Request *requesthandling.InferenceRequest + Response *requesthandling.InferenceResponse + Duration time.Duration + TTFT time.Duration } type DatalayerProcessor interface { diff --git a/pkg/framework/interface/modelselector/plugins.go b/pkg/framework/interface/modelselector/plugins.go index 7fb12444..e70e4557 100644 --- a/pkg/framework/interface/modelselector/plugins.go +++ b/pkg/framework/interface/modelselector/plugins.go @@ -27,7 +27,7 @@ import ( // Filter defines the interface for filtering a list of candidate models based on context. type Filter interface { plugin.Plugin - Filter(ctx context.Context, cycleState *plugin.CycleState, request *requesthandling.InferenceRequest, models []datalayer.Model) []datalayer.Model + Filter(ctx context.Context, request *requesthandling.InferenceRequest, models []datalayer.Model) []datalayer.Model } // Scorer defines the interface for scoring a list of models based on context. @@ -36,11 +36,11 @@ type Filter interface { // If a scorer returns value lower than 0, it will be treated as score 0. type Scorer interface { plugin.Plugin - Score(ctx context.Context, cycleState *plugin.CycleState, request *requesthandling.InferenceRequest, models []datalayer.Model) map[datalayer.Model]float64 + Score(ctx context.Context, request *requesthandling.InferenceRequest, models []datalayer.Model) map[datalayer.Model]float64 } // Picker picks the final model(s) to send the request to. type Picker interface { plugin.Plugin - Pick(ctx context.Context, cycleState *plugin.CycleState, scoredModels []*ScoredModel) *PipelineRunResult + Pick(ctx context.Context, scoredModels []*ScoredModel) *PipelineRunResult } diff --git a/pkg/framework/interface/modelselector/types.go b/pkg/framework/interface/modelselector/types.go index 7840573a..a50d3716 100644 --- a/pkg/framework/interface/modelselector/types.go +++ b/pkg/framework/interface/modelselector/types.go @@ -20,7 +20,6 @@ import ( "context" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/datalayer" - "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/requesthandling" ) @@ -35,5 +34,5 @@ type PipelineRunResult struct { } type ModelSelectorPipeline interface { - Run(ctx context.Context, request *requesthandling.InferenceRequest, cycleState *plugin.CycleState, candidateModels []datalayer.Model) (*PipelineRunResult, error) + Run(ctx context.Context, request *requesthandling.InferenceRequest, candidateModels []datalayer.Model) (*PipelineRunResult, error) } diff --git a/pkg/framework/interface/plugin/cycle_state.go b/pkg/framework/interface/plugin/cycle_state.go deleted file mode 100644 index 91d151b5..00000000 --- a/pkg/framework/interface/plugin/cycle_state.go +++ /dev/null @@ -1,87 +0,0 @@ -/* -Copyright 2026 The llm-d Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -*/ - -package plugin - -import ( - "errors" - "fmt" - "sync" -) - -var ( - // ErrNotFound is the not found error message. - ErrNotFound = errors.New("not found") -) - -// NewCycleState initializes a new CycleState and returns its pointer. -func NewCycleState() *CycleState { - return &CycleState{} -} - -// CycleState provides a mechanism for plugins to store and retrieve arbitrary data. -// Data stored by one plugin can be read, altered, or deleted by another plugin. -// CycleState does not provide any data protection, as all plugins are assumed to be -// trusted. -// Note: CycleState uses a sync.Map to back the storage, because it is thread safe. -// It's aimed to optimize for the "write once and read many times" scenarios. -type CycleState struct { - // key: string, value: any - storage sync.Map -} - -// Read retrieves data with the given "key" from CycleState. If the key is not -// present, ErrNotFound is returned. -// -// See CycleState for notes on concurrency. -func (c *CycleState) Read(key string) (any, error) { - if v, ok := c.storage.Load(key); ok { - return v, nil - } - return nil, ErrNotFound -} - -// Write stores the given "val" in CycleState with the given "key". -// -// See CycleState for notes on concurrency. -func (c *CycleState) Write(key string, val any) { - c.storage.Store(key, val) -} - -// Delete deletes data with the given key from CycleState. -// -// See CycleState for notes on concurrency. -func (c *CycleState) Delete(key string) { - 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 any](c *CycleState, key string) (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 -} diff --git a/pkg/framework/interface/requesthandling/plugins.go b/pkg/framework/interface/requesthandling/plugins.go index 2275c27f..41d3e7cc 100644 --- a/pkg/framework/interface/requesthandling/plugins.go +++ b/pkg/framework/interface/requesthandling/plugins.go @@ -26,14 +26,14 @@ type ProfilePicker interface { plugin.Plugin // Pick selects the Profile to run from a list of candidate profiles, while taking into consideration the request properties. - Pick(ctx context.Context, cycleState *plugin.CycleState, request *InferenceRequest, profiles map[string]*Profile) (*Profile, error) + Pick(ctx context.Context, request *InferenceRequest, profiles map[string]*Profile) (*Profile, error) } type RequestProcessor interface { plugin.Plugin // ProcessRequest runs the RequestProcessor plugin. // RequestProcessor can mutate the headers and/or the body of the request. - ProcessRequest(ctx context.Context, cycleState *plugin.CycleState, request *InferenceRequest) error + ProcessRequest(ctx context.Context, request *InferenceRequest) error } // ResponseProcessor processes the complete buffered response body. @@ -41,16 +41,16 @@ type RequestProcessor interface { // the entire response before calling ProcessResponse on each such plugin. type ResponseProcessor interface { plugin.Plugin - ProcessResponse(ctx context.Context, cycleState *plugin.CycleState, response *InferenceResponse) error + ProcessResponse(ctx context.Context, request *InferenceRequest, response *InferenceResponse) error } // ResponseHeadersProcessor processes response headers before the body arrives. // Plugins implementing this interface run during HandleResponseHeaders, so they // work for both streaming and non-streaming responses. Use this when a plugin -// only needs CycleState and header access (not the response body). +// only needs request attributes and header access (not the response body). type ResponseHeadersProcessor interface { plugin.Plugin - ProcessResponseHeaders(ctx context.Context, cycleState *plugin.CycleState, response *InferenceResponse) error + ProcessResponseHeaders(ctx context.Context, request *InferenceRequest, response *InferenceResponse) error } // ResponseChunkProcessor processes individual response body chunks as they @@ -59,5 +59,5 @@ type ResponseHeadersProcessor interface { // and mutate it via response.SetChunk(). type ResponseChunkProcessor interface { plugin.Plugin - ProcessResponseChunk(ctx context.Context, cycleState *plugin.CycleState, response *InferenceResponse, isFinal bool) error + ProcessResponseChunk(ctx context.Context, request *InferenceRequest, response *InferenceResponse, isFinal bool) error } diff --git a/pkg/framework/interface/requesthandling/types.go b/pkg/framework/interface/requesthandling/types.go index 15377a43..67261076 100644 --- a/pkg/framework/interface/requesthandling/types.go +++ b/pkg/framework/interface/requesthandling/types.go @@ -17,6 +17,9 @@ limitations under the License. package requesthandling import ( + "fmt" + "sync" + "k8s.io/apimachinery/pkg/util/sets" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin" @@ -89,6 +92,31 @@ func (r *InferenceMessage) BodyMutated() bool { type InferenceRequest struct { InferenceMessage + attributes sync.Map +} + +func (r *InferenceRequest) SetAttribute(key string, val any) { + r.attributes.Store(key, val) +} + +func (r *InferenceRequest) GetAttribute(key string) (any, bool) { + return r.attributes.Load(key) +} + +func ReadRequestAttribute[T any](r *InferenceRequest, key string) (T, error) { + var zero T + + raw, ok := r.attributes.Load(key) + if !ok { + return zero, fmt.Errorf("attribute %q: not found", key) + } + + val, ok := raw.(T) + if !ok { + return zero, fmt.Errorf("unexpected type for key %q: got %T, want %T", key, raw, zero) + } + + return val, nil } type InferenceResponse struct { diff --git a/pkg/framework/plugins/modelselector/filter/modelgroup/filter.go b/pkg/framework/plugins/modelselector/filter/modelgroup/filter.go index 7b7622d1..092df86d 100644 --- a/pkg/framework/plugins/modelselector/filter/modelgroup/filter.go +++ b/pkg/framework/plugins/modelselector/filter/modelgroup/filter.go @@ -95,7 +95,7 @@ func (f *ModelGroupFilter) WithName(name string) *ModelGroupFilter { // - a plain non-"auto"-prefixed string: the single candidate matching that name. // - "auto/" (empty group name), unknown group, unmatched name, or non-string // type: no candidates (pipeline rejects with 429). -func (f *ModelGroupFilter) Filter(ctx context.Context, _ *plugin.CycleState, request *requesthandling.InferenceRequest, models []datalayer.Model) []datalayer.Model { +func (f *ModelGroupFilter) Filter(ctx context.Context, request *requesthandling.InferenceRequest, models []datalayer.Model) []datalayer.Model { logger := log.FromContext(ctx) raw := request.Body[requestModelField] diff --git a/pkg/framework/plugins/modelselector/filter/modelgroup/filter_test.go b/pkg/framework/plugins/modelselector/filter/modelgroup/filter_test.go index 0b13cc21..653eb30f 100644 --- a/pkg/framework/plugins/modelselector/filter/modelgroup/filter_test.go +++ b/pkg/framework/plugins/modelselector/filter/modelgroup/filter_test.go @@ -90,7 +90,7 @@ func TestModelGroupFilter_NoGroupsConfigured(t *testing.T) { t.Run("auto/somegroup fails with no groups defined", func(t *testing.T) { req := requestWithModel("auto/somegroup") - got := modelNames(f.Filter(context.Background(), nil, req, candidates)) + got := modelNames(f.Filter(context.Background(), req, candidates)) if len(got) != 0 { t.Errorf("Filter() = %v, want empty", got) } @@ -98,7 +98,7 @@ func TestModelGroupFilter_NoGroupsConfigured(t *testing.T) { t.Run("explicit valid model name succeeds with no groups defined", func(t *testing.T) { req := requestWithModel("qwen3-8b") - got := modelNames(f.Filter(context.Background(), nil, req, candidates)) + got := modelNames(f.Filter(context.Background(), req, candidates)) want := []string{"qwen3-8b"} if len(got) != len(want) || got[0] != want[0] { t.Errorf("Filter() = %v, want %v", got, want) @@ -200,7 +200,7 @@ func TestModelGroupFilter_Filter(t *testing.T) { f := NewModelGroupFilter() req := requestWithModel(tt.modelBody) - got := modelNames(f.Filter(context.Background(), nil, req, candidateModels(membership, all...))) + got := modelNames(f.Filter(context.Background(), req, candidateModels(membership, all...))) want := append([]string{}, tt.want...) sort.Strings(want) @@ -227,7 +227,7 @@ func TestModelGroupFilter_GroupModelsNotInCandidates(t *testing.T) { f := NewModelGroupFilter() req := requestWithModel("auto/qwen3models") - got := modelNames(f.Filter(context.Background(), nil, req, candidates)) + got := modelNames(f.Filter(context.Background(), req, candidates)) if len(got) != 0 { t.Errorf("Filter() = %v, want empty", got) } @@ -247,7 +247,7 @@ func TestModelGroupFilter_PartialGroupInCandidates(t *testing.T) { f := NewModelGroupFilter() req := requestWithModel("auto/qwen3models") - got := modelNames(f.Filter(context.Background(), nil, req, candidates)) + got := modelNames(f.Filter(context.Background(), req, candidates)) want := []string{"qwen3-72b", "qwen3-8b"} sort.Strings(want) @@ -273,7 +273,7 @@ func TestModelGroupFilter_ModelInMultipleGroups(t *testing.T) { for _, group := range []string{"qwen3models", "large-models"} { req := requestWithModel("auto/" + group) - got := modelNames(f.Filter(context.Background(), nil, req, candidates)) + got := modelNames(f.Filter(context.Background(), req, candidates)) if len(got) != 1 || got[0] != "qwen3-32b" { t.Errorf("Filter() for group %q = %v, want [qwen3-32b]", group, got) } diff --git a/pkg/framework/plugins/modelselector/picker/maxscore/picker.go b/pkg/framework/plugins/modelselector/picker/maxscore/picker.go index 152dd025..1f081c50 100644 --- a/pkg/framework/plugins/modelselector/picker/maxscore/picker.go +++ b/pkg/framework/plugins/modelselector/picker/maxscore/picker.go @@ -66,7 +66,7 @@ func (p *MaxScorePicker) TypedName() plugin.TypedName { } // Pick selects the model with the highest score. -func (p *MaxScorePicker) Pick(ctx context.Context, _ *plugin.CycleState, scoredModels []*modelselector.ScoredModel) *modelselector.PipelineRunResult { +func (p *MaxScorePicker) Pick(ctx context.Context, scoredModels []*modelselector.ScoredModel) *modelselector.PipelineRunResult { if debugLogger := log.FromContext(ctx).V(logutil.DEBUG); debugLogger.Enabled() { debugLogger.Info("selecting model from candidates by max score", "numCandidates", len(scoredModels), "scoredModels", scoredModels) diff --git a/pkg/framework/plugins/modelselector/picker/maxscore/picker_test.go b/pkg/framework/plugins/modelselector/picker/maxscore/picker_test.go index 090a1806..c773e9e1 100644 --- a/pkg/framework/plugins/modelselector/picker/maxscore/picker_test.go +++ b/pkg/framework/plugins/modelselector/picker/maxscore/picker_test.go @@ -22,7 +22,6 @@ import ( "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/datalayer" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/modelselector" - "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin" ) func TestMaxScorePicker(t *testing.T) { @@ -71,7 +70,7 @@ func TestMaxScorePicker(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { p := NewMaxScorePicker() - result := p.Pick(context.Background(), plugin.NewCycleState(), tt.input) + result := p.Pick(context.Background(), tt.input) if result == nil { t.Fatal("expected result, got nil") diff --git a/pkg/framework/plugins/modelselector/picker/random/picker.go b/pkg/framework/plugins/modelselector/picker/random/picker.go index 1da1d546..0c5e0ca9 100644 --- a/pkg/framework/plugins/modelselector/picker/random/picker.go +++ b/pkg/framework/plugins/modelselector/picker/random/picker.go @@ -69,7 +69,7 @@ func (p *RandomPicker) TypedName() plugin.TypedName { } // Pick selects random model from the list of candidates. -func (p *RandomPicker) Pick(ctx context.Context, _ *plugin.CycleState, scoredModels []*modelselector.ScoredModel) *modelselector.PipelineRunResult { +func (p *RandomPicker) Pick(ctx context.Context, scoredModels []*modelselector.ScoredModel) *modelselector.PipelineRunResult { if debugLogger := log.FromContext(ctx).V(logutil.DEBUG); debugLogger.Enabled() { debugLogger.Info("Selecting model from candidates randomly", "numOfCandidates", len(scoredModels), "scoredModels", scoredModels) diff --git a/pkg/framework/plugins/modelselector/picker/random/picker_test.go b/pkg/framework/plugins/modelselector/picker/random/picker_test.go index 1044a2e8..b6f45e4a 100644 --- a/pkg/framework/plugins/modelselector/picker/random/picker_test.go +++ b/pkg/framework/plugins/modelselector/picker/random/picker_test.go @@ -22,7 +22,6 @@ import ( "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/datalayer" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/modelselector" - "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin" ) func TestRandomPicker_Pick(t *testing.T) { @@ -60,7 +59,7 @@ func TestRandomPicker_Pick(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { p := NewRandomPicker() - result := p.Pick(context.Background(), plugin.NewCycleState(), tt.input) + result := p.Pick(context.Background(), tt.input) if result == nil { t.Fatal("expected result, got nil") @@ -103,7 +102,7 @@ func TestRandomPicker_Pick_IgnoresScores(t *testing.T) { counts := map[string]int{} for range iterations { - result := p.Pick(context.Background(), plugin.NewCycleState(), input) + result := p.Pick(context.Background(), input) counts[result.TargetModel.GetName()]++ } diff --git a/pkg/framework/plugins/modelselector/picker/weightedrandom/picker.go b/pkg/framework/plugins/modelselector/picker/weightedrandom/picker.go index 0ce21027..dddc035b 100644 --- a/pkg/framework/plugins/modelselector/picker/weightedrandom/picker.go +++ b/pkg/framework/plugins/modelselector/picker/weightedrandom/picker.go @@ -92,7 +92,7 @@ func (p *WeightedRandomPicker) TypedName() plugin.TypedName { // Pick selects the model randomly from the list of candidates, where the probability of the model to get picked is derived // from its weighted score. -func (p *WeightedRandomPicker) Pick(ctx context.Context, cycleState *plugin.CycleState, scoredModels []*modelselector.ScoredModel) *modelselector.PipelineRunResult { +func (p *WeightedRandomPicker) Pick(ctx context.Context, scoredModels []*modelselector.ScoredModel) *modelselector.PipelineRunResult { debugLogger := log.FromContext(ctx).V(logutil.DEBUG) debugEnabled := debugLogger.Enabled() @@ -101,7 +101,7 @@ func (p *WeightedRandomPicker) Pick(ctx context.Context, cycleState *plugin.Cycl if debugEnabled { debugLogger.Info("All scores are zero, delegating to RandomPicker for uniform selection") } - return p.randomPicker.Pick(ctx, cycleState, scoredModels) + return p.randomPicker.Pick(ctx, scoredModels) } if debugEnabled { diff --git a/pkg/framework/plugins/modelselector/picker/weightedrandom/picker_test.go b/pkg/framework/plugins/modelselector/picker/weightedrandom/picker_test.go index 47fae320..900473de 100644 --- a/pkg/framework/plugins/modelselector/picker/weightedrandom/picker_test.go +++ b/pkg/framework/plugins/modelselector/picker/weightedrandom/picker_test.go @@ -22,7 +22,6 @@ import ( "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/datalayer" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/modelselector" - "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin" ) func TestWeightedRandomPicker(t *testing.T) { @@ -36,7 +35,7 @@ func TestWeightedRandomPicker(t *testing.T) { {Model: modelB, Score: 0.1}, } - result := p.Pick(context.Background(), plugin.NewCycleState(), input) + result := p.Pick(context.Background(), input) if result == nil { t.Fatal("expected result, got nil") } @@ -51,7 +50,7 @@ func TestWeightedRandomPicker(t *testing.T) { {Model: modelA, Score: 1.0}, } - result := p.Pick(context.Background(), plugin.NewCycleState(), input) + result := p.Pick(context.Background(), input) if result.TargetModel.GetName() != "model-a" { t.Errorf("expected model-a, got %q", result.TargetModel.GetName()) } @@ -64,7 +63,7 @@ func TestWeightedRandomPicker(t *testing.T) { {Model: modelB, Score: 0}, } - result := p.Pick(context.Background(), plugin.NewCycleState(), input) + result := p.Pick(context.Background(), input) if result == nil || result.TargetModel == nil { t.Fatal("expected a result even with zero scores") } @@ -80,7 +79,7 @@ func TestWeightedRandomPicker(t *testing.T) { {Model: modelA, Score: 0.99}, {Model: modelB, Score: 0.01}, } - result := p.Pick(context.Background(), plugin.NewCycleState(), input) + result := p.Pick(context.Background(), input) counts[result.TargetModel.GetName()]++ } diff --git a/pkg/framework/plugins/modelselector/scorer/costaware/plugin.go b/pkg/framework/plugins/modelselector/scorer/costaware/plugin.go index a7d1e5d5..a3f33a6d 100644 --- a/pkg/framework/plugins/modelselector/scorer/costaware/plugin.go +++ b/pkg/framework/plugins/modelselector/scorer/costaware/plugin.go @@ -78,7 +78,7 @@ func (s *CostScorer) WithName(name string) *CostScorer { // - Higher score indicates better (cheaper) model // - If only one model, it receives neutral score 0.5 // - If all models have zero price, each receives score 1.0 -func (s *CostScorer) Score(_ context.Context, _ *plugin.CycleState, _ *requesthandling.InferenceRequest, models []datalayer.Model) map[datalayer.Model]float64 { +func (s *CostScorer) Score(_ context.Context, _ *requesthandling.InferenceRequest, models []datalayer.Model) map[datalayer.Model]float64 { // Create a map to hold the score of each model candidate scores := make(map[datalayer.Model]float64, len(models)) diff --git a/pkg/framework/plugins/modelselector/scorer/costaware/plugin_test.go b/pkg/framework/plugins/modelselector/scorer/costaware/plugin_test.go index 1ebd1656..992658fe 100644 --- a/pkg/framework/plugins/modelselector/scorer/costaware/plugin_test.go +++ b/pkg/framework/plugins/modelselector/scorer/costaware/plugin_test.go @@ -23,7 +23,6 @@ import ( "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/datalayer" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/datalayer/pricing" - "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/requesthandling" ) @@ -84,7 +83,6 @@ func TestWithName(t *testing.T) { // TestScore tests the Score method with various scenarios func TestScore(t *testing.T) { ctx := context.Background() - cycleState := plugin.NewCycleState() request := requesthandling.NewInferenceRequest() tests := []struct { @@ -200,7 +198,7 @@ func TestScore(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { scorer := NewCostScorer() - scores := scorer.Score(ctx, cycleState, request, tt.models) + scores := scorer.Score(ctx, request, tt.models) // Check that we got the expected number of scores if len(scores) != len(tt.expectedScores) { diff --git a/pkg/framework/plugins/modelselector/scorer/costguard/plugin.go b/pkg/framework/plugins/modelselector/scorer/costguard/plugin.go index bfb627fc..a8257ccb 100644 --- a/pkg/framework/plugins/modelselector/scorer/costguard/plugin.go +++ b/pkg/framework/plugins/modelselector/scorer/costguard/plugin.go @@ -171,7 +171,7 @@ func (s *CostGuardScorer) WithName(name string) *CostGuardScorer { // Conditional Tail Expectation (tail mean above alpha). Under-explored, // missing, or malformed digests yield neutralScore. Ranks map to scores via // a sigmoid centred at the median. -func (s *CostGuardScorer) Score(_ context.Context, _ *plugin.CycleState, _ *requesthandling.InferenceRequest, models []datalayer.Model) map[datalayer.Model]float64 { +func (s *CostGuardScorer) Score(_ context.Context, _ *requesthandling.InferenceRequest, models []datalayer.Model) map[datalayer.Model]float64 { if len(models) == 0 { return map[datalayer.Model]float64{} } diff --git a/pkg/framework/plugins/modelselector/scorer/costguard/plugin_test.go b/pkg/framework/plugins/modelselector/scorer/costguard/plugin_test.go index 1837f96e..042c998d 100644 --- a/pkg/framework/plugins/modelselector/scorer/costguard/plugin_test.go +++ b/pkg/framework/plugins/modelselector/scorer/costguard/plugin_test.go @@ -27,7 +27,6 @@ import ( "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/datalayer" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/datalayer/accumulator" - "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/requesthandling" ) @@ -121,7 +120,7 @@ func TestScore_AllNeutral(t *testing.T) { for i, spec := range tt.models { models[i] = modelWithDigest(t, spec.name, newCostDigestN(t, spec.cost, spec.count)) } - scores := s.Score(context.Background(), plugin.NewCycleState(), requesthandling.NewInferenceRequest(), models) + scores := s.Score(context.Background(), requesthandling.NewInferenceRequest(), models) require.Len(t, scores, len(models)) for _, m := range models { assert.Equal(t, neutralScore, scores[m]) @@ -139,7 +138,7 @@ func TestScore_TwoDistinctRanks(t *testing.T) { cheap := modelWithDigest(t, "cheap", newCostDigestN(t, 1.0, overCount)) expensive := modelWithDigest(t, "expensive", newCostDigestN(t, 3.0, overCount)) models := []datalayer.Model{cheap, expensive} - scores := s.Score(context.Background(), plugin.NewCycleState(), requesthandling.NewInferenceRequest(), models) + scores := s.Score(context.Background(), requesthandling.NewInferenceRequest(), models) require.Len(t, scores, len(models)) assert.Greater(t, scores[cheap], 0.5, "cheaper model should score above neutral") assert.Less(t, scores[expensive], 0.5, "more expensive model should score below neutral") @@ -161,7 +160,7 @@ func TestScore_ThreeDistinctRanks(t *testing.T) { mid := modelWithDigest(t, "mid", newCostDigestN(t, 2.0, overCount)) expensive := modelWithDigest(t, "expensive", newCostDigestN(t, 3.0, overCount)) models := []datalayer.Model{cheap, mid, expensive} - scores := s.Score(context.Background(), plugin.NewCycleState(), requesthandling.NewInferenceRequest(), models) + scores := s.Score(context.Background(), requesthandling.NewInferenceRequest(), models) require.Len(t, scores, len(models)) assert.InDelta(t, 0.5, scores[mid], 1e-9, "median-rank model should score exactly neutral") assert.Greater(t, scores[cheap], scores[mid]) @@ -183,14 +182,14 @@ func TestScore_SelfCalibratesToScale(t *testing.T) { smallCheap := modelWithDigest(t, "small-cheap", newCostDigestN(t, 1.0, overCount)) smallMid := modelWithDigest(t, "small-mid", newCostDigestN(t, 2.0, overCount)) smallExpensive := modelWithDigest(t, "small-expensive", newCostDigestN(t, 3.0, overCount)) - smallScores := s.Score(context.Background(), plugin.NewCycleState(), requesthandling.NewInferenceRequest(), + smallScores := s.Score(context.Background(), requesthandling.NewInferenceRequest(), []datalayer.Model{smallCheap, smallMid, smallExpensive}) // Large-scale fixture — same relative geometry, costs scaled by 1000. largeCheap := modelWithDigest(t, "large-cheap", newCostDigestN(t, 1000.0, overCount)) largeMid := modelWithDigest(t, "large-mid", newCostDigestN(t, 2000.0, overCount)) largeExpensive := modelWithDigest(t, "large-expensive", newCostDigestN(t, 3000.0, overCount)) - largeScores := s.Score(context.Background(), plugin.NewCycleState(), requesthandling.NewInferenceRequest(), + largeScores := s.Score(context.Background(), requesthandling.NewInferenceRequest(), []datalayer.Model{largeCheap, largeMid, largeExpensive}) assert.InDelta(t, smallScores[smallCheap], largeScores[largeCheap], 1e-9) @@ -213,7 +212,7 @@ func TestScore_UnevenSpread(t *testing.T) { cheap := modelWithDigest(t, "cheap", newCostDigestN(t, 1.0, overCount)) mid := modelWithDigest(t, "mid", newCostDigestN(t, 2.0, overCount)) expensive := modelWithDigest(t, "expensive", newCostDigestN(t, 100.0, overCount)) - scores := s.Score(context.Background(), plugin.NewCycleState(), requesthandling.NewInferenceRequest(), + scores := s.Score(context.Background(), requesthandling.NewInferenceRequest(), []datalayer.Model{cheap, mid, expensive}) require.Len(t, scores, 3) assert.InDelta(t, 0.5, scores[mid], 1e-9, "median-rank model still scores exactly neutral") @@ -270,7 +269,7 @@ func TestScore_UnderExplored(t *testing.T) { expensive = modelWithDigest(t, "expensive", newCostDigestN(t, 3.0, over)) models = []datalayer.Model{cheap, expensive, u} } - scores := s.Score(context.Background(), plugin.NewCycleState(), requesthandling.NewInferenceRequest(), models) + scores := s.Score(context.Background(), requesthandling.NewInferenceRequest(), models) require.Len(t, scores, len(models)) assert.Equal(t, neutralScore, scores[u]) if ctx.withExplored { @@ -494,7 +493,7 @@ func TestScore_Explore(t *testing.T) { } observed := make(map[datalayer.Model]struct{}, len(pool)) for i := 0; i < trials; i++ { - scores := s.Score(context.Background(), plugin.NewCycleState(), requesthandling.NewInferenceRequest(), models) + scores := s.Score(context.Background(), requesthandling.NewInferenceRequest(), models) require.Len(t, scores, len(models)) var pick datalayer.Model pickCount := 0 @@ -529,7 +528,7 @@ func TestScore_ExploitPathUnchanged(t *testing.T) { models := []datalayer.Model{cheap, expensive, u} for i := 0; i < trials; i++ { - scores := s.Score(context.Background(), plugin.NewCycleState(), requesthandling.NewInferenceRequest(), models) + scores := s.Score(context.Background(), requesthandling.NewInferenceRequest(), models) require.Len(t, scores, len(models)) assert.Equal(t, neutralScore, scores[u], "under-explored model must remain neutral on the exploit path") assert.Greater(t, scores[cheap], 0.5, "cheap model scores above neutral") diff --git a/pkg/framework/plugins/modelselector/scorer/inflightrequests/plugin.go b/pkg/framework/plugins/modelselector/scorer/inflightrequests/plugin.go index 5329323b..bb2e0519 100644 --- a/pkg/framework/plugins/modelselector/scorer/inflightrequests/plugin.go +++ b/pkg/framework/plugins/modelselector/scorer/inflightrequests/plugin.go @@ -64,7 +64,7 @@ func (s *InflightRequestsScorer) WithName(name string) *InflightRequestsScorer { // Score returns a score in [0,1] for each model based on its in-flight request count. // Formula: score = (max - count) / (max - min) -func (s *InflightRequestsScorer) Score(_ context.Context, _ *plugin.CycleState, _ *requesthandling.InferenceRequest, models []datalayer.Model) map[datalayer.Model]float64 { +func (s *InflightRequestsScorer) Score(_ context.Context, _ *requesthandling.InferenceRequest, models []datalayer.Model) map[datalayer.Model]float64 { var minCount int64 = math.MaxInt64 var maxCount int64 = math.MinInt64 diff --git a/pkg/framework/plugins/modelselector/scorer/sessionaffinity/plugin.go b/pkg/framework/plugins/modelselector/scorer/sessionaffinity/plugin.go index c3cbaf8e..51a11788 100644 --- a/pkg/framework/plugins/modelselector/scorer/sessionaffinity/plugin.go +++ b/pkg/framework/plugins/modelselector/scorer/sessionaffinity/plugin.go @@ -43,7 +43,7 @@ const ( defaultMaxSessions = 10000 defaultTTL = time.Hour - cycleStateSessionIDKey = "session-affinity/session-id" + sessionIDAttributeKey = "session-affinity/session-id" responseModelField = "model" @@ -172,7 +172,7 @@ func (s *SessionAffinityScorer) TypedName() plugin.TypedName { } // Score returns a score in [0,1] for each model based on session affinity. -func (s *SessionAffinityScorer) Score(ctx context.Context, cycleState *plugin.CycleState, request *requesthandling.InferenceRequest, models []datalayer.Model) map[datalayer.Model]float64 { +func (s *SessionAffinityScorer) Score(ctx context.Context, request *requesthandling.InferenceRequest, models []datalayer.Model) map[datalayer.Model]float64 { logger := log.FromContext(ctx) scores := make(map[datalayer.Model]float64, len(models)) @@ -185,7 +185,7 @@ func (s *SessionAffinityScorer) Score(ctx context.Context, cycleState *plugin.Cy return scores } - cycleState.Write(cycleStateSessionIDKey, sessionID) + request.SetAttribute(sessionIDAttributeKey, sessionID) previousModel, known := s.cache.Get(sessionID) @@ -228,7 +228,7 @@ func (s *SessionAffinityScorer) Score(ctx context.Context, cycleState *plugin.Cy // When no session ID was present in the request, a new UUID is generated // optimistically and echoed back — if the client sends it on the next request, // session affinity will be established automatically. -func (s *SessionAffinityScorer) ProcessResponse(ctx context.Context, cycleState *plugin.CycleState, response *requesthandling.InferenceResponse) error { +func (s *SessionAffinityScorer) ProcessResponse(ctx context.Context, request *requesthandling.InferenceRequest, response *requesthandling.InferenceResponse) error { logger := log.FromContext(ctx) modelName, ok := response.Body[responseModelField].(string) @@ -237,9 +237,8 @@ func (s *SessionAffinityScorer) ProcessResponse(ctx context.Context, cycleState return nil } - sessionID, err := plugin.ReadCycleStateKey[string](cycleState, cycleStateSessionIDKey) + sessionID, err := requesthandling.ReadRequestAttribute[string](request, sessionIDAttributeKey) if err != nil { - // No session ID in request — generate one optimistically sessionID = uuid.New().String() logger.V(logutil.VERBOSE).Info("generated optimistic session ID", "sessionId", sessionID, "model", modelName) diff --git a/pkg/framework/plugins/modelselector/scorer/sessionaffinity/plugin_test.go b/pkg/framework/plugins/modelselector/scorer/sessionaffinity/plugin_test.go index b9b0538c..f8c1791f 100644 --- a/pkg/framework/plugins/modelselector/scorer/sessionaffinity/plugin_test.go +++ b/pkg/framework/plugins/modelselector/scorer/sessionaffinity/plugin_test.go @@ -27,7 +27,6 @@ import ( "github.com/stretchr/testify/require" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/datalayer" - "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/requesthandling" ) @@ -82,10 +81,9 @@ func TestFactory_InvalidJSON(t *testing.T) { func TestScore_NoSessionID_NoOpinion(t *testing.T) { s := newTestScorer(t) req := requesthandling.NewInferenceRequest() - cs := plugin.NewCycleState() models := testModels("model-a", "model-b", "model-c") - scores := s.Score(context.Background(), cs, req, models) + scores := s.Score(context.Background(), req, models) require.Len(t, scores, 3) for _, m := range models { @@ -99,10 +97,9 @@ func TestScore_FirstTurn_NoOpinion(t *testing.T) { s := newTestScorer(t) req := requesthandling.NewInferenceRequest() req.Headers["x-session-id"] = "sess-new" - cs := plugin.NewCycleState() models := testModels("model-a", "model-b") - scores := s.Score(context.Background(), cs, req, models) + scores := s.Score(context.Background(), req, models) require.Len(t, scores, 2) for _, m := range models { @@ -118,10 +115,9 @@ func TestScore_FollowUpTurn_PrefersKnownModel(t *testing.T) { req := requesthandling.NewInferenceRequest() req.Headers["x-session-id"] = "sess-123" - cs := plugin.NewCycleState() models := testModels("model-a", "model-b", "model-c") - scores := s.Score(context.Background(), cs, req, models) + scores := s.Score(context.Background(), req, models) require.Len(t, scores, 3) for _, m := range models { @@ -140,10 +136,9 @@ func TestScore_KnownModelNoLongerInCandidates_NoOpinion(t *testing.T) { req := requesthandling.NewInferenceRequest() req.Headers["x-session-id"] = "sess-123" - cs := plugin.NewCycleState() models := testModels("model-a", "model-b") - scores := s.Score(context.Background(), cs, req, models) + scores := s.Score(context.Background(), req, models) require.Len(t, scores, 2) for _, m := range models { @@ -157,24 +152,22 @@ func TestScore_EmptyModels(t *testing.T) { s := newTestScorer(t) req := requesthandling.NewInferenceRequest() req.Headers["x-session-id"] = "sess-1" - cs := plugin.NewCycleState() - scores := s.Score(context.Background(), cs, req, nil) + scores := s.Score(context.Background(), req, nil) assert.Empty(t, scores) } -// Session ID is stored in CycleState for the ResponseProcessor. -func TestScore_StoresSessionIDInCycleState(t *testing.T) { +// Session ID is stored in request attributes for the ResponseProcessor. +func TestScore_StoresSessionIDInRequestAttributes(t *testing.T) { s := newTestScorer(t) req := requesthandling.NewInferenceRequest() req.Headers["x-session-id"] = "my-session" - cs := plugin.NewCycleState() models := testModels("model-a") - s.Score(context.Background(), cs, req, models) + s.Score(context.Background(), req, models) - val, err := plugin.ReadCycleStateKey[string](cs, cycleStateSessionIDKey) + val, err := requesthandling.ReadRequestAttribute[string](req, sessionIDAttributeKey) require.NoError(t, err) assert.Equal(t, "my-session", val) } @@ -189,10 +182,9 @@ func TestScore_UsesCustomSessionIDKey(t *testing.T) { req := requesthandling.NewInferenceRequest() req.Headers["x-conv-id"] = "conv-123" - cs := plugin.NewCycleState() models := testModels("model-a", "model-b") - scores := s.Score(context.Background(), cs, req, models) + scores := s.Score(context.Background(), req, models) assert.Equal(t, preferredScore, scores[models[0]]) assert.Equal(t, noOpinionScore, scores[models[1]]) @@ -203,13 +195,13 @@ func TestScore_UsesCustomSessionIDKey(t *testing.T) { // Records a new session-to-model mapping on first turn. func TestProcessResponse_RecordsMapping(t *testing.T) { s := newTestScorer(t) - cs := plugin.NewCycleState() - cs.Write(cycleStateSessionIDKey, "sess-456") + req := requesthandling.NewInferenceRequest() + req.SetAttribute(sessionIDAttributeKey, "sess-456") resp := requesthandling.NewInferenceResponse() resp.Body[responseModelField] = "model-a" - err := s.ProcessResponse(context.Background(), cs, resp) + err := s.ProcessResponse(context.Background(), req, resp) require.NoError(t, err) model, ok := s.cache.Get("sess-456") @@ -217,15 +209,15 @@ func TestProcessResponse_RecordsMapping(t *testing.T) { assert.Equal(t, "model-a", model) } -// No session ID in CycleState → generates optimistic UUID and stores mapping. +// No session ID in request attributes → generates optimistic UUID and stores mapping. func TestProcessResponse_NoSessionID_GeneratesOptimistic(t *testing.T) { s := newTestScorer(t) - cs := plugin.NewCycleState() + req := requesthandling.NewInferenceRequest() resp := requesthandling.NewInferenceResponse() resp.Body[responseModelField] = "model-a" - err := s.ProcessResponse(context.Background(), cs, resp) + err := s.ProcessResponse(context.Background(), req, resp) require.NoError(t, err) assert.Equal(t, 1, s.cache.Len(), "should have stored optimistic session") @@ -244,12 +236,12 @@ func TestProcessResponse_NoSessionID_GeneratesOptimistic(t *testing.T) { // No model in response body → no-op (no optimistic ID generated either). func TestProcessResponse_NoModelInBody_Skips(t *testing.T) { s := newTestScorer(t) - cs := plugin.NewCycleState() - cs.Write(cycleStateSessionIDKey, "sess-789") + req := requesthandling.NewInferenceRequest() + req.SetAttribute(sessionIDAttributeKey, "sess-789") resp := requesthandling.NewInferenceResponse() - err := s.ProcessResponse(context.Background(), cs, resp) + err := s.ProcessResponse(context.Background(), req, resp) require.NoError(t, err) assert.Equal(t, 0, s.cache.Len()) @@ -260,13 +252,13 @@ func TestProcessResponse_UpdatesExistingMapping(t *testing.T) { s := newTestScorer(t) s.cache.Add("sess-1", "model-old") - cs := plugin.NewCycleState() - cs.Write(cycleStateSessionIDKey, "sess-1") + req := requesthandling.NewInferenceRequest() + req.SetAttribute(sessionIDAttributeKey, "sess-1") resp := requesthandling.NewInferenceResponse() resp.Body[responseModelField] = "model-new" - err := s.ProcessResponse(context.Background(), cs, resp) + err := s.ProcessResponse(context.Background(), req, resp) require.NoError(t, err) model, ok := s.cache.Get("sess-1") @@ -279,13 +271,13 @@ func TestProcessResponse_SameModel_NoUpdate(t *testing.T) { s := newTestScorer(t) s.cache.Add("sess-1", "model-a") - cs := plugin.NewCycleState() - cs.Write(cycleStateSessionIDKey, "sess-1") + req := requesthandling.NewInferenceRequest() + req.SetAttribute(sessionIDAttributeKey, "sess-1") resp := requesthandling.NewInferenceResponse() resp.Body[responseModelField] = "model-a" - err := s.ProcessResponse(context.Background(), cs, resp) + err := s.ProcessResponse(context.Background(), req, resp) require.NoError(t, err) model, ok := s.cache.Get("sess-1") @@ -296,13 +288,13 @@ func TestProcessResponse_SameModel_NoUpdate(t *testing.T) { // Session ID is echoed back as a response header using the configured key. func TestProcessResponse_EchoesSessionID(t *testing.T) { s := newTestScorer(t) - cs := plugin.NewCycleState() - cs.Write(cycleStateSessionIDKey, "sess-echo") + req := requesthandling.NewInferenceRequest() + req.SetAttribute(sessionIDAttributeKey, "sess-echo") resp := requesthandling.NewInferenceResponse() resp.Body[responseModelField] = "model-a" - err := s.ProcessResponse(context.Background(), cs, resp) + err := s.ProcessResponse(context.Background(), req, resp) require.NoError(t, err) assert.Equal(t, "sess-echo", resp.Headers["x-session-id"]) @@ -319,9 +311,8 @@ func TestEndToEnd_FirstTurnThenFollowUp(t *testing.T) { // Turn 1: session ID present but no prior model → no opinion req1 := requesthandling.NewInferenceRequest() req1.Headers["x-session-id"] = "conv-abc" - cs1 := plugin.NewCycleState() - scores1 := s.Score(context.Background(), cs1, req1, models) + scores1 := s.Score(context.Background(), req1, models) for _, m := range models { assert.Equal(t, noOpinionScore, scores1[m]) } @@ -329,15 +320,14 @@ func TestEndToEnd_FirstTurnThenFollowUp(t *testing.T) { // Simulate model selection + response: model-b was picked resp1 := requesthandling.NewInferenceResponse() resp1.Body[responseModelField] = "model-b" - err := s.ProcessResponse(context.Background(), cs1, resp1) + err := s.ProcessResponse(context.Background(), req1, resp1) require.NoError(t, err) // Turn 2: same session → prefers model-b req2 := requesthandling.NewInferenceRequest() req2.Headers["x-session-id"] = "conv-abc" - cs2 := plugin.NewCycleState() - scores2 := s.Score(context.Background(), cs2, req2, models) + scores2 := s.Score(context.Background(), req2, models) assert.Equal(t, noOpinionScore, scores2[models[0]]) // model-a assert.Equal(t, preferredScore, scores2[models[1]]) // model-b } @@ -349,9 +339,8 @@ func TestEndToEnd_OptimisticSessionID(t *testing.T) { // Turn 1: no session ID at all req1 := requesthandling.NewInferenceRequest() - cs1 := plugin.NewCycleState() - scores1 := s.Score(context.Background(), cs1, req1, models) + scores1 := s.Score(context.Background(), req1, models) for _, m := range models { assert.Equal(t, noOpinionScore, scores1[m]) } @@ -359,7 +348,7 @@ func TestEndToEnd_OptimisticSessionID(t *testing.T) { // Response generates optimistic session ID resp1 := requesthandling.NewInferenceResponse() resp1.Body[responseModelField] = "model-b" - err := s.ProcessResponse(context.Background(), cs1, resp1) + err := s.ProcessResponse(context.Background(), req1, resp1) require.NoError(t, err) generatedID := resp1.Headers["x-session-id"] @@ -368,9 +357,8 @@ func TestEndToEnd_OptimisticSessionID(t *testing.T) { // Turn 2: client echoes back the generated session ID → affinity kicks in req2 := requesthandling.NewInferenceRequest() req2.Headers["x-session-id"] = generatedID - cs2 := plugin.NewCycleState() - scores2 := s.Score(context.Background(), cs2, req2, models) + scores2 := s.Score(context.Background(), req2, models) assert.Equal(t, noOpinionScore, scores2[models[0]]) // model-a assert.Equal(t, preferredScore, scores2[models[1]]) // model-b } diff --git a/pkg/framework/plugins/requesthandling/basemodelextractor/base_model_to_header.go b/pkg/framework/plugins/requesthandling/basemodelextractor/base_model_to_header.go index 573c84f9..1736aeb1 100644 --- a/pkg/framework/plugins/requesthandling/basemodelextractor/base_model_to_header.go +++ b/pkg/framework/plugins/requesthandling/basemodelextractor/base_model_to_header.go @@ -85,7 +85,7 @@ func (p *BaseModelToHeaderPlugin) WithName(name string) *BaseModelToHeaderPlugin } // ProcessRequest sets base model name on the header -func (p *BaseModelToHeaderPlugin) ProcessRequest(ctx context.Context, _ *plugin.CycleState, request *requesthandling.InferenceRequest) error { +func (p *BaseModelToHeaderPlugin) ProcessRequest(ctx context.Context, request *requesthandling.InferenceRequest) error { // extract raw field value from body rawFieldValue, exists := request.Body[modelField] if !exists { diff --git a/pkg/framework/plugins/requesthandling/basemodelextractor/base_model_to_header_test.go b/pkg/framework/plugins/requesthandling/basemodelextractor/base_model_to_header_test.go index b142ddf7..b4c20dfb 100644 --- a/pkg/framework/plugins/requesthandling/basemodelextractor/base_model_to_header_test.go +++ b/pkg/framework/plugins/requesthandling/basemodelextractor/base_model_to_header_test.go @@ -242,7 +242,7 @@ func TestBaseModelToHeaderPlugin_ProcessRequest(t *testing.T) { } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - err := p.ProcessRequest(context.Background(), nil, tt.request) + err := p.ProcessRequest(context.Background(), tt.request) if tt.wantErr { if err == nil { t.Fatal("expected error, got nil") @@ -288,7 +288,7 @@ func TestBaseModelToHeaderPlugin_ProcessRequest_MutatedHeaders(t *testing.T) { request := requesthandling.NewInferenceRequest() request.Body["model"] = testAdapter - if err := p.ProcessRequest(context.Background(), nil, request); err != nil { + if err := p.ProcessRequest(context.Background(), request); err != nil { t.Fatalf("unexpected error: %v", err) } diff --git a/pkg/framework/plugins/requesthandling/bodyfieldtoheader/body_field_to_header.go b/pkg/framework/plugins/requesthandling/bodyfieldtoheader/body_field_to_header.go index e669e566..f3a6ecc9 100644 --- a/pkg/framework/plugins/requesthandling/bodyfieldtoheader/body_field_to_header.go +++ b/pkg/framework/plugins/requesthandling/bodyfieldtoheader/body_field_to_header.go @@ -102,7 +102,7 @@ func (p *BodyFieldToHeaderPlugin) WithName(name string) *BodyFieldToHeaderPlugin } // ProcessRequest extracts value from a given body field and sets it as HTTP header. -func (p *BodyFieldToHeaderPlugin) ProcessRequest(ctx context.Context, _ *plugin.CycleState, request *requesthandling.InferenceRequest) error { +func (p *BodyFieldToHeaderPlugin) ProcessRequest(ctx context.Context, request *requesthandling.InferenceRequest) error { // extract raw field value from body rawFieldValue, exists := request.Body[p.fieldName] if !exists { diff --git a/pkg/framework/plugins/requesthandling/bodyfieldtoheader/body_field_to_header_test.go b/pkg/framework/plugins/requesthandling/bodyfieldtoheader/body_field_to_header_test.go index 281a5aea..54b29918 100644 --- a/pkg/framework/plugins/requesthandling/bodyfieldtoheader/body_field_to_header_test.go +++ b/pkg/framework/plugins/requesthandling/bodyfieldtoheader/body_field_to_header_test.go @@ -257,7 +257,7 @@ func TestBodyFieldToHeaderPlugin_ProcessRequest(t *testing.T) { t.Fatalf("failed to create plugin: %v", err) } - err = p.ProcessRequest(context.Background(), nil, tt.request) + err = p.ProcessRequest(context.Background(), tt.request) if tt.wantErr { if err == nil { t.Fatal("expected error, got nil") @@ -285,7 +285,7 @@ func TestBodyFieldToHeaderPlugin_ProcessRequest_MutatedHeaders(t *testing.T) { request := requesthandling.NewInferenceRequest() request.Body["model"] = testModelValue - if err := p.ProcessRequest(context.Background(), nil, request); err != nil { + if err := p.ProcessRequest(context.Background(), request); err != nil { t.Fatalf("unexpected error: %v", err) } diff --git a/pkg/framework/plugins/requesthandling/modelselector/plugin.go b/pkg/framework/plugins/requesthandling/modelselector/plugin.go index 3cbb0500..4503e460 100644 --- a/pkg/framework/plugins/requesthandling/modelselector/plugin.go +++ b/pkg/framework/plugins/requesthandling/modelselector/plugin.go @@ -33,8 +33,8 @@ import ( const ( ModelSelectorPluginType = "model-selector" - // SelectedModelCycleStateKey is the CycleState key where the selected model name is stored. - SelectedModelCycleStateKey = "model-selector/selected-model" + // SelectedModelAttributeKey is the request attribute key where the selected model name is stored. + SelectedModelAttributeKey = "model-selector/selected-model" ) var _ requesthandling.RequestProcessor = &ModelSelectorPlugin{} @@ -76,8 +76,8 @@ func (p *ModelSelectorPlugin) WithName(name string) *ModelSelectorPlugin { } // ProcessRequest reads candidate models from the Datastore, runs model -// selection, and writes the selected model into the request body and CycleState. -func (p *ModelSelectorPlugin) ProcessRequest(ctx context.Context, cycleState *plugin.CycleState, request *requesthandling.InferenceRequest) error { +// selection, and writes the selected model into the request body and attributes. +func (p *ModelSelectorPlugin) ProcessRequest(ctx context.Context, request *requesthandling.InferenceRequest) error { logger := log.FromContext(ctx) candidateModels := p.datastore.GetModels(datalayer.AllModelsPredicate) @@ -85,7 +85,7 @@ func (p *ModelSelectorPlugin) ProcessRequest(ctx context.Context, cycleState *pl return errors.New("no candidate models available in datastore") } - result, err := p.selector.Select(ctx, request, cycleState, candidateModels) + result, err := p.selector.Select(ctx, request, candidateModels) if err != nil { return err } @@ -94,7 +94,7 @@ func (p *ModelSelectorPlugin) ProcessRequest(ctx context.Context, cycleState *pl logger.V(logutil.VERBOSE).Info("Model selected", "model", selectedName) request.SetBodyField("model", selectedName) - cycleState.Write(SelectedModelCycleStateKey, selectedName) + request.SetAttribute(SelectedModelAttributeKey, selectedName) return nil } diff --git a/pkg/framework/plugins/requesthandling/modelselector/plugin_test.go b/pkg/framework/plugins/requesthandling/modelselector/plugin_test.go index f27395ca..4b1f82e2 100644 --- a/pkg/framework/plugins/requesthandling/modelselector/plugin_test.go +++ b/pkg/framework/plugins/requesthandling/modelselector/plugin_test.go @@ -103,9 +103,8 @@ func TestProcessRequestSelectsFromDatastoreModels(t *testing.T) { request := requesthandling.NewInferenceRequest() request.Body["model"] = "auto" - cycleState := fwkplugin.NewCycleState() - if err := p.ProcessRequest(context.Background(), cycleState, request); err != nil { + if err := p.ProcessRequest(context.Background(), request); err != nil { t.Fatalf("ProcessRequest failed: %v", err) } @@ -114,12 +113,12 @@ func TestProcessRequestSelectsFromDatastoreModels(t *testing.T) { t.Errorf("selected model %q is not in datastore models %v", selectedModel, candidates) } - storedModel, err := fwkplugin.ReadCycleStateKey[string](cycleState, SelectedModelCycleStateKey) + storedModel, err := requesthandling.ReadRequestAttribute[string](request, SelectedModelAttributeKey) if err != nil { - t.Fatalf("expected selected model in CycleState: %v", err) + t.Fatalf("expected selected model in request attributes: %v", err) } if storedModel != selectedModel { - t.Errorf("CycleState model %q != body model %q", storedModel, selectedModel) + t.Errorf("request attribute model %q != body model %q", storedModel, selectedModel) } } @@ -130,9 +129,8 @@ func TestProcessRequestFailsWithEmptyDatastore(t *testing.T) { request := requesthandling.NewInferenceRequest() request.Body["model"] = "auto" - cycleState := fwkplugin.NewCycleState() - if err := p.ProcessRequest(context.Background(), cycleState, request); err == nil { + if err := p.ProcessRequest(context.Background(), request); err == nil { t.Fatal("expected error with empty datastore") } } @@ -197,14 +195,14 @@ func TestAddPluginsRejectsScorerWithoutWeight(t *testing.T) { type fakeScorerFilter struct{ typedName fwkplugin.TypedName } func (f *fakeScorerFilter) TypedName() fwkplugin.TypedName { return f.typedName } -func (f *fakeScorerFilter) Score(_ context.Context, _ *fwkplugin.CycleState, _ *requesthandling.InferenceRequest, models []datalayer.Model) map[datalayer.Model]float64 { +func (f *fakeScorerFilter) Score(_ context.Context, _ *requesthandling.InferenceRequest, models []datalayer.Model) map[datalayer.Model]float64 { out := make(map[datalayer.Model]float64, len(models)) for _, m := range models { out[m] = 1.0 } return out } -func (f *fakeScorerFilter) Filter(_ context.Context, _ *fwkplugin.CycleState, _ *requesthandling.InferenceRequest, models []datalayer.Model) []datalayer.Model { +func (f *fakeScorerFilter) Filter(_ context.Context, _ *requesthandling.InferenceRequest, models []datalayer.Model) []datalayer.Model { return models } diff --git a/pkg/framework/plugins/requesthandling/profilepicker/single/single_profile_picker.go b/pkg/framework/plugins/requesthandling/profilepicker/single/single_profile_picker.go index 217478ba..a80c8bb6 100644 --- a/pkg/framework/plugins/requesthandling/profilepicker/single/single_profile_picker.go +++ b/pkg/framework/plugins/requesthandling/profilepicker/single/single_profile_picker.go @@ -62,7 +62,7 @@ func (p *SingleProfilePicker) WithName(name string) *SingleProfilePicker { // Pick selects the Profile to run from the list of candidate profiles, while taking into consideration the request properties and the // previously executed cycles along with their results. -func (p *SingleProfilePicker) Pick(ctx context.Context, cycleState *plugin.CycleState, request *requesthandling.InferenceRequest, profiles map[string]*requesthandling.Profile) (*requesthandling.Profile, error) { +func (p *SingleProfilePicker) Pick(ctx context.Context, request *requesthandling.InferenceRequest, profiles map[string]*requesthandling.Profile) (*requesthandling.Profile, error) { if len(profiles) != 1 { return nil, fmt.Errorf("failed to select a single profile from %d profiles", len(profiles)) } diff --git a/pkg/framework/plugins/requesthandling/profilepicker/single/single_profile_picker_test.go b/pkg/framework/plugins/requesthandling/profilepicker/single/single_profile_picker_test.go index f961db1b..61848a4d 100644 --- a/pkg/framework/plugins/requesthandling/profilepicker/single/single_profile_picker_test.go +++ b/pkg/framework/plugins/requesthandling/profilepicker/single/single_profile_picker_test.go @@ -103,7 +103,7 @@ func TestPick(t *testing.T) { picker := NewSingleProfilePicker() for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - got, err := picker.Pick(context.Background(), plugin.NewCycleState(), requesthandling.NewInferenceRequest(), tt.profiles) + got, err := picker.Pick(context.Background(), requesthandling.NewInferenceRequest(), tt.profiles) if tt.wantErr { if err == nil { diff --git a/pkg/framework/plugins/responsehandling/modelnametoheader/plugin.go b/pkg/framework/plugins/responsehandling/modelnametoheader/plugin.go index 68630f5b..d8c957f2 100644 --- a/pkg/framework/plugins/responsehandling/modelnametoheader/plugin.go +++ b/pkg/framework/plugins/responsehandling/modelnametoheader/plugin.go @@ -19,7 +19,6 @@ package modelnametoheader import ( "context" "encoding/json" - "errors" "fmt" "sigs.k8s.io/controller-runtime/pkg/log" @@ -74,14 +73,11 @@ func (p *ModelNameToHeaderPlugin) TypedName() plugin.TypedName { return p.typedName } -func (p *ModelNameToHeaderPlugin) ProcessResponseHeaders(ctx context.Context, cycleState *plugin.CycleState, response *requesthandling.InferenceResponse) error { - selectedModel, err := plugin.ReadCycleStateKey[string](cycleState, modelselector.SelectedModelCycleStateKey) +func (p *ModelNameToHeaderPlugin) ProcessResponseHeaders(ctx context.Context, request *requesthandling.InferenceRequest, response *requesthandling.InferenceResponse) error { + selectedModel, err := requesthandling.ReadRequestAttribute[string](request, modelselector.SelectedModelAttributeKey) if err != nil { - if errors.Is(err, plugin.ErrNotFound) { - log.FromContext(ctx).V(logutil.VERBOSE).Info("no selected model in CycleState, skipping") - return nil - } - return fmt.Errorf("failed to read selected model from CycleState: %w", err) + log.FromContext(ctx).V(logutil.VERBOSE).Info("no selected model in request attributes, skipping") + return nil } response.SetHeader(p.headerName, selectedModel) diff --git a/pkg/framework/plugins/responsehandling/modelnametoheader/plugin_test.go b/pkg/framework/plugins/responsehandling/modelnametoheader/plugin_test.go index 819cb6da..f61266a5 100644 --- a/pkg/framework/plugins/responsehandling/modelnametoheader/plugin_test.go +++ b/pkg/framework/plugins/responsehandling/modelnametoheader/plugin_test.go @@ -21,24 +21,23 @@ import ( "encoding/json" "testing" - "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/requesthandling" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/plugins/requesthandling/modelselector" ) -func TestProcessResponseHeaders_SetsHeaderFromCycleState(t *testing.T) { +func TestProcessResponseHeaders_SetsHeaderFromRequestAttribute(t *testing.T) { p, err := PluginFactory("test", nil, nil) if err != nil { t.Fatalf("PluginFactory failed: %v", err) } - cycleState := plugin.NewCycleState() - cycleState.Write(modelselector.SelectedModelCycleStateKey, "llama-70b") + request := requesthandling.NewInferenceRequest() + request.SetAttribute(modelselector.SelectedModelAttributeKey, "llama-70b") response := requesthandling.NewInferenceResponse() rhp := p.(requesthandling.ResponseHeadersProcessor) - if err := rhp.ProcessResponseHeaders(context.Background(), cycleState, response); err != nil { + if err := rhp.ProcessResponseHeaders(context.Background(), request, response); err != nil { t.Fatalf("ProcessResponseHeaders failed: %v", err) } @@ -48,17 +47,17 @@ func TestProcessResponseHeaders_SetsHeaderFromCycleState(t *testing.T) { } } -func TestProcessResponseHeaders_NoOpWithoutCycleStateEntry(t *testing.T) { +func TestProcessResponseHeaders_NoOpWithoutRequestAttribute(t *testing.T) { p, err := PluginFactory("test", nil, nil) if err != nil { t.Fatalf("PluginFactory failed: %v", err) } - cycleState := plugin.NewCycleState() + request := requesthandling.NewInferenceRequest() response := requesthandling.NewInferenceResponse() rhp := p.(requesthandling.ResponseHeadersProcessor) - if err := rhp.ProcessResponseHeaders(context.Background(), cycleState, response); err != nil { + if err := rhp.ProcessResponseHeaders(context.Background(), request, response); err != nil { t.Fatalf("ProcessResponseHeaders failed: %v", err) } @@ -67,21 +66,24 @@ func TestProcessResponseHeaders_NoOpWithoutCycleStateEntry(t *testing.T) { } } -func TestProcessResponseHeaders_ReturnsErrorOnUnexpectedType(t *testing.T) { +func TestProcessResponseHeaders_NoOpOnWrongType(t *testing.T) { p, err := PluginFactory("test", nil, nil) if err != nil { t.Fatalf("PluginFactory failed: %v", err) } - cycleState := plugin.NewCycleState() - cycleState.Write(modelselector.SelectedModelCycleStateKey, 12345) // wrong type + request := requesthandling.NewInferenceRequest() + request.SetAttribute(modelselector.SelectedModelAttributeKey, 12345) // wrong type response := requesthandling.NewInferenceResponse() rhp := p.(requesthandling.ResponseHeadersProcessor) - err = rhp.ProcessResponseHeaders(context.Background(), cycleState, response) - if err == nil { - t.Fatal("expected error for wrong-typed CycleState value, got nil") + if err := rhp.ProcessResponseHeaders(context.Background(), request, response); err != nil { + t.Fatalf("ProcessResponseHeaders should not error on wrong type: %v", err) + } + + if len(response.MutatedHeaders()) != 0 { + t.Errorf("expected no mutated headers for wrong type, got %v", response.MutatedHeaders()) } } @@ -115,13 +117,13 @@ func TestPluginFactory_CustomHeaderName(t *testing.T) { t.Errorf("expected header %q, got %q", "X-Custom-Model", mnp.headerName) } - cycleState := plugin.NewCycleState() - cycleState.Write(modelselector.SelectedModelCycleStateKey, "gpt-4") + request := requesthandling.NewInferenceRequest() + request.SetAttribute(modelselector.SelectedModelAttributeKey, "gpt-4") response := requesthandling.NewInferenceResponse() rhp := p.(requesthandling.ResponseHeadersProcessor) - if err := rhp.ProcessResponseHeaders(context.Background(), cycleState, response); err != nil { + if err := rhp.ProcessResponseHeaders(context.Background(), request, response); err != nil { t.Fatalf("ProcessResponseHeaders failed: %v", err) } diff --git a/pkg/handlers/request.go b/pkg/handlers/request.go index f8613d44..329c8b79 100644 --- a/pkg/handlers/request.go +++ b/pkg/handlers/request.go @@ -33,7 +33,6 @@ import ( errcommon "github.com/llm-d/llm-d-inference-payload-processor/pkg/common/error" logutil "github.com/llm-d/llm-d-inference-payload-processor/pkg/common/observability/logging" datasource "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/datalayer/datasource" - "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/requesthandling" "github.com/llm-d/llm-d-inference-payload-processor/pkg/metrics" ) @@ -93,16 +92,16 @@ func (s *Server) HandleRequestBody(ctx context.Context, reqCtx *RequestContext, return nil, errcommon.Error{Code: errcommon.BadRequest, Msg: fmt.Sprintf("failed to parse request body: %v", err)} } - if err := s.runRequestPlugins(ctx, reqCtx.CycleState, reqCtx.Request, s.preProcessors); err != nil { + if err := s.runRequestPlugins(ctx, reqCtx.Request, s.preProcessors); err != nil { return nil, err } var err error - reqCtx.Profile, err = s.profilePicker.Pick(ctx, reqCtx.CycleState, reqCtx.Request, s.profiles) + reqCtx.Profile, err = s.profilePicker.Pick(ctx, reqCtx.Request, s.profiles) if err != nil { return nil, errcommon.Error{Code: errcommon.Internal, Msg: fmt.Sprintf("failed to pick a profile: %v", err)} } - if err := s.runRequestPlugins(ctx, reqCtx.CycleState, reqCtx.Request, reqCtx.Profile.RequestPlugins); err != nil { + if err := s.runRequestPlugins(ctx, reqCtx.Request, reqCtx.Profile.RequestPlugins); err != nil { return nil, err } @@ -123,7 +122,7 @@ func (s *Server) HandleRequestBody(ctx context.Context, reqCtx *RequestContext, // Notify the data layer of the incoming request after headers are fully formed. s.eventNotifier.Notify(datasource.Event{ Type: datasource.RequestEventType, - Payload: datasource.RequestPayload{Request: reqCtx.Request, CycleState: reqCtx.CycleState}, + Payload: datasource.RequestPayload{Request: reqCtx.Request}, }) metrics.RecordSuccessCounter() @@ -151,7 +150,7 @@ func (s *Server) HandleRequestBody(ctx context.Context, reqCtx *RequestContext, } // runRequestPlugins executes request plugins in the order they were registered. -func (s *Server) runRequestPlugins(ctx context.Context, cycleState *plugin.CycleState, request *requesthandling.InferenceRequest, +func (s *Server) runRequestPlugins(ctx context.Context, request *requesthandling.InferenceRequest, reqPlugins []requesthandling.RequestProcessor) error { logger := log.FromContext(ctx).V(logutil.DEFAULT) @@ -165,7 +164,7 @@ func (s *Server) runRequestPlugins(ctx context.Context, cycleState *plugin.Cycle verboseLogger.Info("Executing request plugin", "plugin", reqPlugin.TypedName()) } before := time.Now() - err := reqPlugin.ProcessRequest(ctx, cycleState, request) + err := reqPlugin.ProcessRequest(ctx, request) metrics.RecordPluginProcessingLatency(requestPluginExtensionPoint, reqPlugin.TypedName().Type, reqPlugin.TypedName().Name, time.Since(before)) if err != nil { logger.Error(err, "Failed to execute request plugin", "plugin", reqPlugin.TypedName()) diff --git a/pkg/handlers/request_test.go b/pkg/handlers/request_test.go index 5436e73a..2672b473 100644 --- a/pkg/handlers/request_test.go +++ b/pkg/handlers/request_test.go @@ -574,8 +574,7 @@ func TestHandleRequestBody_BuiltInPlugins(t *testing.T) { addRequestPlugins(profiles, modelToHeaderPlugin, baseModelToHeaderPlugin) server := newServerForTest(profiles) reqCtx := &RequestContext{ - CycleState: plugin.NewCycleState(), - Request: requesthandling.NewInferenceRequest(), + Request: requesthandling.NewInferenceRequest(), } bodyBytes, _ := json.Marshal(test.body) resp, err := server.HandleRequestBody(ctx, reqCtx, bodyBytes) @@ -624,8 +623,7 @@ func TestHandleRequestBodyWithPluginMetrics(t *testing.T) { addRequestPlugins(profiles, modelToHeaderPlugin, baseModelToHeaderPlugin) server := newServerForTest(profiles) reqCtx := &RequestContext{ - CycleState: plugin.NewCycleState(), - Request: requesthandling.NewInferenceRequest(), + Request: requesthandling.NewInferenceRequest(), } bodyBytes, _ := json.Marshal(map[string]any{ @@ -675,7 +673,7 @@ func (p *fakeRequestPlugin) TypedName() plugin.TypedName { return plugin.TypedName{Type: "fake", Name: p.name} } -func (p *fakeRequestPlugin) ProcessRequest(ctx context.Context, _ *plugin.CycleState, request *requesthandling.InferenceRequest) error { +func (p *fakeRequestPlugin) ProcessRequest(ctx context.Context, request *requesthandling.InferenceRequest) error { return p.mutateFn(ctx, request) } @@ -897,8 +895,7 @@ func TestHandleRequestBody_MultiPluginHeaderMutations(t *testing.T) { addRequestPlugins(profiles, tc.plugins...) server := newServerForTest(profiles) reqCtx := &RequestContext{ - Request: requesthandling.NewInferenceRequest(), - CycleState: plugin.NewCycleState(), + Request: requesthandling.NewInferenceRequest(), } for k, v := range tc.initialHeaders { reqCtx.Request.Headers[k] = v @@ -984,15 +981,15 @@ func buildStreamingResponse(bodyBytes []byte, setHeaders map[string]string, remo type bodyMutatingPlugin struct { name string - mutateFn func(ctx context.Context, cycleState *plugin.CycleState, request *requesthandling.InferenceRequest) error + mutateFn func(ctx context.Context, request *requesthandling.InferenceRequest) error } func (p *bodyMutatingPlugin) TypedName() plugin.TypedName { return plugin.TypedName{Type: "fake", Name: p.name} } -func (p *bodyMutatingPlugin) ProcessRequest(ctx context.Context, cycleState *plugin.CycleState, request *requesthandling.InferenceRequest) error { - return p.mutateFn(ctx, cycleState, request) +func (p *bodyMutatingPlugin) ProcessRequest(ctx context.Context, request *requesthandling.InferenceRequest) error { + return p.mutateFn(ctx, request) } var _ requesthandling.RequestProcessor = &bodyMutatingPlugin{} @@ -1003,7 +1000,7 @@ func TestHandleRequestBody_BodyMutation(t *testing.T) { bodyPlugin := &bodyMutatingPlugin{ name: "body-mutator", - mutateFn: func(_ context.Context, _ *plugin.CycleState, request *requesthandling.InferenceRequest) error { + mutateFn: func(_ context.Context, request *requesthandling.InferenceRequest) error { request.SetBodyField("injected", "value") return nil }, @@ -1055,8 +1052,7 @@ func TestHandleRequestBody_BodyMutation(t *testing.T) { addRequestPlugins(profiles, bodyPlugin, baseModelPlugin) server := newServerForTest(profiles) reqCtx := &RequestContext{ - CycleState: plugin.NewCycleState(), - Request: requesthandling.NewInferenceRequest(), + Request: requesthandling.NewInferenceRequest(), } bodyBytes, _ := json.Marshal(body) resp, err := server.HandleRequestBody(ctx, reqCtx, bodyBytes) diff --git a/pkg/handlers/response.go b/pkg/handlers/response.go index 678cb30a..ebc27c54 100644 --- a/pkg/handlers/response.go +++ b/pkg/handlers/response.go @@ -29,7 +29,6 @@ import ( envoy "github.com/llm-d/llm-d-inference-payload-processor/pkg/common/envoy" logutil "github.com/llm-d/llm-d-inference-payload-processor/pkg/common/observability/logging" - "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/requesthandling" "github.com/llm-d/llm-d-inference-payload-processor/pkg/metrics" ) @@ -43,7 +42,7 @@ func (s *Server) HandleResponseHeaders(ctx context.Context, reqCtx *RequestConte } } - if err := s.runResponseHeadersProcessors(ctx, reqCtx.CycleState, reqCtx.Response); err != nil { + if err := s.runResponseHeadersProcessors(ctx, reqCtx.Request, reqCtx.Response); err != nil { return nil, err } @@ -83,12 +82,12 @@ func (s *Server) HandleResponseBody(ctx context.Context, reqCtx *RequestContext, } if hasProfilePlugins { - if err := s.runResponsePlugins(ctx, reqCtx.CycleState, reqCtx.Response, reqCtx.Profile.ResponsePlugins); err != nil { + if err := s.runResponsePlugins(ctx, reqCtx.Request, reqCtx.Response, reqCtx.Profile.ResponsePlugins); err != nil { return nil, err } } - if err := s.runResponsePlugins(ctx, reqCtx.CycleState, reqCtx.Response, s.postProcessors); err != nil { + if err := s.runResponsePlugins(ctx, reqCtx.Request, reqCtx.Response, s.postProcessors); err != nil { return nil, err } @@ -174,7 +173,7 @@ func (s *Server) HandleResponseChunk(ctx context.Context, reqCtx *RequestContext chunk := string(chunkBytes) reqCtx.Response.ResetChunkState(chunk) - if err := s.runResponseChunkProcessors(ctx, reqCtx.CycleState, reqCtx.Response, endOfStream, reqCtx.Profile.ResponseChunkProcessors); err != nil { + if err := s.runResponseChunkProcessors(ctx, reqCtx.Request, reqCtx.Response, endOfStream, reqCtx.Profile.ResponseChunkProcessors); err != nil { logger.Error(err, "Failed to run response chunk processors") return nil, err } @@ -190,7 +189,7 @@ func (s *Server) HandleResponseChunk(ctx context.Context, reqCtx *RequestContext // runResponseChunkProcessors executes chunk processors in the order they were registered. // Each plugin receives response.CurrentChunk so mutations from earlier plugins are visible // to later ones in the chain. -func (s *Server) runResponseChunkProcessors(ctx context.Context, cycleState *plugin.CycleState, response *requesthandling.InferenceResponse, isFinal bool, processors []requesthandling.ResponseChunkProcessor) error { +func (s *Server) runResponseChunkProcessors(ctx context.Context, request *requesthandling.InferenceRequest, response *requesthandling.InferenceResponse, isFinal bool, processors []requesthandling.ResponseChunkProcessor) error { logger := log.FromContext(ctx).V(logutil.DEFAULT) verboseLogger := logger.V(logutil.VERBOSE) @@ -199,7 +198,7 @@ func (s *Server) runResponseChunkProcessors(ctx context.Context, cycleState *plu verboseLogger.Info("Executing response chunk plugin", "plugin", cp.TypedName()) } before := time.Now() - err := cp.ProcessResponseChunk(ctx, cycleState, response, isFinal) + err := cp.ProcessResponseChunk(ctx, request, response, isFinal) metrics.RecordPluginProcessingLatency(responsePluginExtensionPoint, cp.TypedName().Type, cp.TypedName().Name, time.Since(before)) if err != nil { return err @@ -254,7 +253,7 @@ func (s *Server) HandleResponseTrailers(trailers *eppb.HttpTrailers) ([]*eppb.Pr } // runResponseHeadersProcessors executes response-headers post-processors in order. -func (s *Server) runResponseHeadersProcessors(ctx context.Context, cycleState *plugin.CycleState, response *requesthandling.InferenceResponse) error { +func (s *Server) runResponseHeadersProcessors(ctx context.Context, request *requesthandling.InferenceRequest, response *requesthandling.InferenceResponse) error { if len(s.responseHeadersPostProcessors) == 0 { return nil } @@ -267,7 +266,7 @@ func (s *Server) runResponseHeadersProcessors(ctx context.Context, cycleState *p verboseLogger.Info("Executing response headers plugin", "plugin", hp.TypedName()) } before := time.Now() - if err := hp.ProcessResponseHeaders(ctx, cycleState, response); err != nil { + if err := hp.ProcessResponseHeaders(ctx, request, response); err != nil { logger.Error(err, "Failed to execute response headers plugin", "plugin", hp.TypedName()) return err } @@ -278,7 +277,7 @@ func (s *Server) runResponseHeadersProcessors(ctx context.Context, cycleState *p } // runResponsePlugins executes response plugins in the order they were registered. -func (s *Server) runResponsePlugins(ctx context.Context, cycleState *plugin.CycleState, response *requesthandling.InferenceResponse, respPlugins []requesthandling.ResponseProcessor) error { +func (s *Server) runResponsePlugins(ctx context.Context, request *requesthandling.InferenceRequest, response *requesthandling.InferenceResponse, respPlugins []requesthandling.ResponseProcessor) error { logger := log.FromContext(ctx).V(logutil.DEFAULT) // Cache verbose logger and check Enabled() once to avoid per-iteration @@ -292,7 +291,7 @@ func (s *Server) runResponsePlugins(ctx context.Context, cycleState *plugin.Cycl verboseLogger.Info("Executing response plugin", "plugin", respPlugin.TypedName()) } before := time.Now() - err = respPlugin.ProcessResponse(ctx, cycleState, response) + err = respPlugin.ProcessResponse(ctx, request, response) metrics.RecordPluginProcessingLatency(responsePluginExtensionPoint, respPlugin.TypedName().Type, respPlugin.TypedName().Name, time.Since(before)) if err != nil { logger.Error(err, "Failed to execute response plugin", "plugin", respPlugin.TypedName()) diff --git a/pkg/handlers/response_test.go b/pkg/handlers/response_test.go index 305f16e7..76aa45b8 100644 --- a/pkg/handlers/response_test.go +++ b/pkg/handlers/response_test.go @@ -36,28 +36,27 @@ import ( const testPluginValue = "done" -// fakeResponsePlugin implements requesthandling.PayloadProcessor for testing response plugin execution. +// fakeResponsePlugin implements requesthandling.ResponseProcessor for testing. type fakeResponsePlugin struct { name string - mutateFn func(ctx context.Context, cycleState *plugin.CycleState, response *requesthandling.InferenceResponse) error + mutateFn func(ctx context.Context, request *requesthandling.InferenceRequest, response *requesthandling.InferenceResponse) error } func (p *fakeResponsePlugin) TypedName() plugin.TypedName { return plugin.TypedName{Type: "fake", Name: p.name} } -func (p *fakeResponsePlugin) ProcessResponse(ctx context.Context, cycleState *plugin.CycleState, response *requesthandling.InferenceResponse) error { - return p.mutateFn(ctx, cycleState, response) +func (p *fakeResponsePlugin) ProcessResponse(ctx context.Context, request *requesthandling.InferenceRequest, response *requesthandling.InferenceResponse) error { + return p.mutateFn(ctx, request, response) } var _ requesthandling.ResponseProcessor = &fakeResponsePlugin{} func newTestRequestContext(profiles map[string]*requesthandling.Profile) *RequestContext { return &RequestContext{ - Profile: profiles[testProfileName], - CycleState: plugin.NewCycleState(), - Request: requesthandling.NewInferenceRequest(), - Response: requesthandling.NewInferenceResponse(), + Profile: profiles[testProfileName], + Request: requesthandling.NewInferenceRequest(), + Response: requesthandling.NewInferenceResponse(), } } @@ -106,7 +105,7 @@ func TestHandleResponseBody_SinglePlugin(t *testing.T) { mutatePlugin := &fakeResponsePlugin{ name: "mutator", - mutateFn: func(_ context.Context, _ *plugin.CycleState, response *requesthandling.InferenceResponse) error { + mutateFn: func(_ context.Context, _ *requesthandling.InferenceRequest, response *requesthandling.InferenceResponse) error { response.SetBodyField("mutated", true) return nil }, @@ -139,14 +138,14 @@ func TestHandleResponseBody_MultiplePlugins(t *testing.T) { plugin1 := &fakeResponsePlugin{ name: "plugin1", - mutateFn: func(_ context.Context, _ *plugin.CycleState, response *requesthandling.InferenceResponse) error { + mutateFn: func(_ context.Context, _ *requesthandling.InferenceRequest, response *requesthandling.InferenceResponse) error { response.SetBodyField("p1", testPluginValue) return nil }, } plugin2 := &fakeResponsePlugin{ name: "plugin2", - mutateFn: func(_ context.Context, _ *plugin.CycleState, response *requesthandling.InferenceResponse) error { + mutateFn: func(_ context.Context, _ *requesthandling.InferenceRequest, response *requesthandling.InferenceResponse) error { response.SetBodyField("p2", testPluginValue) return nil }, @@ -180,7 +179,7 @@ func TestHandleResponseBody_PluginError(t *testing.T) { failingPlugin := &fakeResponsePlugin{ name: "failing", - mutateFn: func(_ context.Context, _ *plugin.CycleState, _ *requesthandling.InferenceResponse) error { + mutateFn: func(_ context.Context, _ *requesthandling.InferenceRequest, _ *requesthandling.InferenceResponse) error { return errors.New("failed to execute plugin") }, } @@ -204,7 +203,7 @@ func TestHandleResponseBody_StreamingWithPlugin(t *testing.T) { mutatePlugin := &fakeResponsePlugin{ name: "mutator", - mutateFn: func(_ context.Context, _ *plugin.CycleState, response *requesthandling.InferenceResponse) error { + mutateFn: func(_ context.Context, _ *requesthandling.InferenceRequest, response *requesthandling.InferenceResponse) error { response.SetBodyField("mutated", true) return nil }, @@ -237,7 +236,7 @@ func TestHandleResponseBody_PluginNoBodyMutation(t *testing.T) { headerOnlyPlugin := &fakeResponsePlugin{ name: "header-only", - mutateFn: func(_ context.Context, _ *plugin.CycleState, response *requesthandling.InferenceResponse) error { + mutateFn: func(_ context.Context, _ *requesthandling.InferenceRequest, response *requesthandling.InferenceResponse) error { response.SetHeader("X-Custom-Response", "added") return nil }, diff --git a/pkg/handlers/server.go b/pkg/handlers/server.go index a12e0389..d8e16413 100644 --- a/pkg/handlers/server.go +++ b/pkg/handlers/server.go @@ -38,7 +38,6 @@ import ( errcommon "github.com/llm-d/llm-d-inference-payload-processor/pkg/common/error" logutil "github.com/llm-d/llm-d-inference-payload-processor/pkg/common/observability/logging" datasource "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/datalayer/datasource" - "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/requesthandling" "github.com/llm-d/llm-d-inference-payload-processor/pkg/metrics" "github.com/llm-d/llm-d-inference-payload-processor/version" @@ -92,7 +91,6 @@ type RequestContext struct { ResponseCompleteTimestamp time.Time ResponseHeadersSent bool Profile *requesthandling.Profile - CycleState *plugin.CycleState Request *requesthandling.InferenceRequest Response *requesthandling.InferenceResponse } @@ -147,10 +145,9 @@ func (s *Server) Process(srv extProcPb.ExternalProcessor_ProcessServer) error { loggerVerbose.Info("Processing") reqCtx := &RequestContext{ - Request: requesthandling.NewInferenceRequest(), - Response: requesthandling.NewInferenceResponse(), - Profile: s.emptyProfile, // request is always initialized with an empty profile to avoid nil pointer - CycleState: plugin.NewCycleState(), + Request: requesthandling.NewInferenceRequest(), + Response: requesthandling.NewInferenceResponse(), + Profile: s.emptyProfile, // request is always initialized with an empty profile to avoid nil pointer } // TODO set a max cap on these. // both requestBody and responseBody accumulate without an upper bound. @@ -265,11 +262,10 @@ func (s *Server) Process(srv extProcPb.ExternalProcessor_ProcessServer) error { s.eventNotifier.Notify(datasource.Event{ Type: datasource.ResponseEventType, Payload: datasource.ResponsePayload{ - Request: reqCtx.Request, - Response: reqCtx.Response, - CycleState: reqCtx.CycleState, - Duration: reqCtx.ResponseCompleteTimestamp.Sub(reqCtx.RequestReceivedTimestamp), - TTFT: reqCtx.ResponseFirstChunkTimestamp.Sub(reqCtx.RequestSentTimestamp), + Request: reqCtx.Request, + Response: reqCtx.Response, + Duration: reqCtx.ResponseCompleteTimestamp.Sub(reqCtx.RequestReceivedTimestamp), + TTFT: reqCtx.ResponseFirstChunkTimestamp.Sub(reqCtx.RequestSentTimestamp), }, }) } diff --git a/pkg/handlers/server_test.go b/pkg/handlers/server_test.go index 53079d8c..7bc1274e 100644 --- a/pkg/handlers/server_test.go +++ b/pkg/handlers/server_test.go @@ -35,7 +35,6 @@ import ( envoytest "github.com/llm-d/llm-d-inference-payload-processor/pkg/common/envoy/test" logutil "github.com/llm-d/llm-d-inference-payload-processor/pkg/common/observability/logging" datasource "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/datalayer/datasource" - "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/requesthandling" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/plugins/requesthandling/basemodelextractor" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/plugins/requesthandling/bodyfieldtoheader" @@ -184,8 +183,7 @@ func TestHandleRequestBody(t *testing.T) { addRequestPlugins(profiles, modelToHeaderPlugin, baseModelToHeaderPlugin) srv := newServerForTest(profiles) reqCtx := &RequestContext{ - CycleState: plugin.NewCycleState(), - Request: requesthandling.NewInferenceRequest(), + Request: requesthandling.NewInferenceRequest(), } got, err := srv.HandleRequestBody(ctx, reqCtx, b) if err != nil { @@ -484,7 +482,7 @@ func TestProcess_ResponseBodyPopulatedAtEndOfStream(t *testing.T) { profiles[testProfileName].NeedsResponseBuffering = true withResponsePlugins(profiles, &fakeResponsePlugin{ name: "usage-mutator", - mutateFn: func(_ context.Context, _ *plugin.CycleState, response *requesthandling.InferenceResponse) error { + mutateFn: func(_ context.Context, _ *requesthandling.InferenceRequest, response *requesthandling.InferenceResponse) error { usage, _ := response.Body["usage"].(map[string]any) usage["mutated"] = true return nil diff --git a/pkg/modelselector/model_selector_pipeline.go b/pkg/modelselector/model_selector_pipeline.go index 4afefff6..da04966c 100644 --- a/pkg/modelselector/model_selector_pipeline.go +++ b/pkg/modelselector/model_selector_pipeline.go @@ -143,26 +143,22 @@ func (p *ModelSelectorPipeline) String() string { } // Run runs the ModelSelectorPipeline: Filter → Score → Pick. -func (p *ModelSelectorPipeline) Run(ctx context.Context, request *requesthandling.InferenceRequest, cycleState *plugin.CycleState, candidateModels []datalayer.Model) (*modelselector.PipelineRunResult, error) { - models := p.runFilterPlugins(ctx, request, cycleState, candidateModels) +func (p *ModelSelectorPipeline) Run(ctx context.Context, request *requesthandling.InferenceRequest, candidateModels []datalayer.Model) (*modelselector.PipelineRunResult, error) { + models := p.runFilterPlugins(ctx, request, candidateModels) if len(models) == 0 { - // Typed so the handler maps it to an HTTP ImmediateResponse instead of - // failing the ext_proc stream. return nil, errcommon.Error{Code: errcommon.ResourceExhausted, Msg: "no models available after filtering"} } - weightedScorePerModel := p.runScorerPlugins(ctx, request, cycleState, models) + weightedScorePerModel := p.runScorerPlugins(ctx, request, models) - result := p.runPickerPlugin(ctx, cycleState, weightedScorePerModel) + result := p.runPickerPlugin(ctx, weightedScorePerModel) return result, nil } -func (p *ModelSelectorPipeline) runFilterPlugins(ctx context.Context, request *requesthandling.InferenceRequest, cycleState *plugin.CycleState, models []datalayer.Model) []datalayer.Model { +func (p *ModelSelectorPipeline) runFilterPlugins(ctx context.Context, request *requesthandling.InferenceRequest, models []datalayer.Model) []datalayer.Model { logger := log.FromContext(ctx) - // Cache loggers and check Enabled() once to avoid per-iteration allocations - // from argument boxing when logging at that level is disabled. verboseLogger := logger.V(logutil.VERBOSE) verboseEnabled := verboseLogger.Enabled() debugLogger := logger.V(logutil.DEBUG) @@ -179,7 +175,7 @@ func (p *ModelSelectorPipeline) runFilterPlugins(ctx context.Context, request *r verboseLogger.Info("Running filter plugin", "plugin", filter.TypedName()) } before := time.Now() - filteredModels = filter.Filter(ctx, cycleState, request, filteredModels) + filteredModels = filter.Filter(ctx, request, filteredModels) metrics.RecordPluginProcessingLatency(filterExtensionPoint, filter.TypedName().Type, filter.TypedName().Name, time.Since(before)) if debugEnabled { debugLogger.Info("Completed running filter plugin", "plugin", filter.TypedName(), "remainingModels", len(filteredModels)) @@ -196,18 +192,14 @@ func (p *ModelSelectorPipeline) runFilterPlugins(ctx context.Context, request *r return filteredModels } -func (p *ModelSelectorPipeline) runScorerPlugins(ctx context.Context, request *requesthandling.InferenceRequest, cycleState *plugin.CycleState, models []datalayer.Model) map[string]*modelselector.ScoredModel { +func (p *ModelSelectorPipeline) runScorerPlugins(ctx context.Context, request *requesthandling.InferenceRequest, models []datalayer.Model) map[string]*modelselector.ScoredModel { logger := log.FromContext(ctx) - // Cache loggers and check Enabled() once to avoid per-iteration allocations - // from argument boxing when logging at that level is disabled. verboseLogger := logger.V(logutil.VERBOSE) verboseEnabled := verboseLogger.Enabled() debugLogger := logger.V(logutil.DEBUG) debugEnabled := debugLogger.Enabled() - // Create one big array for all ScoredModels instead of allocating each one - // separately. This reduces memory allocations from N to 1. n := len(models) storage := make([]modelselector.ScoredModel, n) scoredModels := make(map[string]*modelselector.ScoredModel, n) @@ -221,7 +213,7 @@ func (p *ModelSelectorPipeline) runScorerPlugins(ctx context.Context, request *r verboseLogger.Info("Running scorer plugin", "plugin", scorer.TypedName()) } before := time.Now() - scores := scorer.Score(ctx, cycleState, request, models) + scores := scorer.Score(ctx, request, models) metrics.RecordPluginProcessingLatency(scorerExtensionPoint, scorer.TypedName().Type, scorer.TypedName().Name, time.Since(before)) for model, score := range scores { if sm, exists := scoredModels[model.GetName()]; exists { @@ -237,11 +229,9 @@ func (p *ModelSelectorPipeline) runScorerPlugins(ctx context.Context, request *r return scoredModels } -func (p *ModelSelectorPipeline) runPickerPlugin(ctx context.Context, cycleState *plugin.CycleState, scoredModelMap map[string]*modelselector.ScoredModel) *modelselector.PipelineRunResult { +func (p *ModelSelectorPipeline) runPickerPlugin(ctx context.Context, scoredModelMap map[string]*modelselector.ScoredModel) *modelselector.PipelineRunResult { logger := log.FromContext(ctx) - // Cache loggers and check Enabled() once to avoid allocations from argument - // boxing when logging at that level is disabled. verboseLogger := logger.V(logutil.VERBOSE) verboseEnabled := verboseLogger.Enabled() debugLogger := logger.V(logutil.DEBUG) @@ -258,7 +248,7 @@ func (p *ModelSelectorPipeline) runPickerPlugin(ctx context.Context, cycleState verboseLogger.Info("Running picker plugin", "plugin", p.picker.TypedName()) } before := time.Now() - result := p.picker.Pick(ctx, cycleState, scoredModels) + result := p.picker.Pick(ctx, scoredModels) metrics.RecordPluginProcessingLatency(pickerExtensionPoint, p.picker.TypedName().Type, p.picker.TypedName().Name, time.Since(before)) if debugEnabled { debugLogger.Info("Completed running picker plugin", "plugin", p.picker.TypedName(), "result", result) diff --git a/pkg/modelselector/model_selector_pipeline_bench_test.go b/pkg/modelselector/model_selector_pipeline_bench_test.go index e1359133..d04feda6 100644 --- a/pkg/modelselector/model_selector_pipeline_bench_test.go +++ b/pkg/modelselector/model_selector_pipeline_bench_test.go @@ -81,8 +81,7 @@ func BenchmarkModelSelectorPipelineRun(b *testing.B) { b.ReportAllocs() b.ResetTimer() for i := 0; i < b.N; i++ { - cycleState := plugin.NewCycleState() - result, err := pipeline.Run(ctx, request, cycleState, models) + result, err := pipeline.Run(ctx, request, models) if err != nil { b.Fatalf("Run failed: %v", err) } @@ -111,7 +110,7 @@ type benchScorer struct { func (s *benchScorer) TypedName() plugin.TypedName { return s.typedName } -func (s *benchScorer) Score(_ context.Context, _ *plugin.CycleState, _ *requesthandling.InferenceRequest, models []datalayer.Model) map[datalayer.Model]float64 { +func (s *benchScorer) Score(_ context.Context, _ *requesthandling.InferenceRequest, models []datalayer.Model) map[datalayer.Model]float64 { scores := make(map[datalayer.Model]float64, len(models)) for i, m := range models { // Produce varied but deterministic scores @@ -127,7 +126,7 @@ type benchPicker struct { func (p *benchPicker) TypedName() plugin.TypedName { return p.typedName } -func (p *benchPicker) Pick(_ context.Context, _ *plugin.CycleState, scoredModels []*modelselector.ScoredModel) *modelselector.PipelineRunResult { +func (p *benchPicker) Pick(_ context.Context, scoredModels []*modelselector.ScoredModel) *modelselector.PipelineRunResult { if len(scoredModels) == 0 { return nil } diff --git a/pkg/modelselector/model_selector_pipeline_test.go b/pkg/modelselector/model_selector_pipeline_test.go index 7b57dfe2..810a0fe5 100644 --- a/pkg/modelselector/model_selector_pipeline_test.go +++ b/pkg/modelselector/model_selector_pipeline_test.go @@ -33,7 +33,7 @@ type testFilter struct { } func (f *testFilter) TypedName() plugin.TypedName { return f.typedName } -func (f *testFilter) Filter(_ context.Context, _ *plugin.CycleState, _ *requesthandling.InferenceRequest, models []datalayer.Model) []datalayer.Model { +func (f *testFilter) Filter(_ context.Context, _ *requesthandling.InferenceRequest, models []datalayer.Model) []datalayer.Model { f.callCount++ if f.filterFn != nil { return f.filterFn(models) @@ -48,7 +48,7 @@ type testScorer struct { } func (s *testScorer) TypedName() plugin.TypedName { return s.typedName } -func (s *testScorer) Score(_ context.Context, _ *plugin.CycleState, _ *requesthandling.InferenceRequest, models []datalayer.Model) map[datalayer.Model]float64 { +func (s *testScorer) Score(_ context.Context, _ *requesthandling.InferenceRequest, models []datalayer.Model) map[datalayer.Model]float64 { s.callCount++ if s.scoreFn != nil { return s.scoreFn(models) @@ -66,7 +66,7 @@ type testPicker struct { } func (p *testPicker) TypedName() plugin.TypedName { return p.typedName } -func (p *testPicker) Pick(_ context.Context, _ *plugin.CycleState, scoredModels []*modelselector.ScoredModel) *modelselector.PipelineRunResult { +func (p *testPicker) Pick(_ context.Context, scoredModels []*modelselector.ScoredModel) *modelselector.PipelineRunResult { p.callCount++ if len(scoredModels) == 0 { return nil @@ -181,7 +181,7 @@ func TestPipelineRun(t *testing.T) { } } - result, err := pipeline.Run(context.Background(), requesthandling.NewInferenceRequest(), plugin.NewCycleState(), tt.models) + result, err := pipeline.Run(context.Background(), requesthandling.NewInferenceRequest(), tt.models) if tt.wantErr { if err == nil { @@ -242,7 +242,7 @@ func TestScoreWeightAccumulation(t *testing.T) { t.Fatalf("AddPlugins failed: %v", err) } - result, err := pipeline.Run(context.Background(), requesthandling.NewInferenceRequest(), plugin.NewCycleState(), []datalayer.Model{modelA, modelB}) + result, err := pipeline.Run(context.Background(), requesthandling.NewInferenceRequest(), []datalayer.Model{modelA, modelB}) if err != nil { t.Fatalf("unexpected error: %v", err) } diff --git a/pkg/modelselector/modelselector.go b/pkg/modelselector/modelselector.go index 5361803a..3b5a9eab 100644 --- a/pkg/modelselector/modelselector.go +++ b/pkg/modelselector/modelselector.go @@ -26,7 +26,6 @@ import ( logutil "github.com/llm-d/llm-d-inference-payload-processor/pkg/common/observability/logging" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/datalayer" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/modelselector" - "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/requesthandling" "github.com/llm-d/llm-d-inference-payload-processor/pkg/metrics" ) @@ -49,7 +48,7 @@ func (s *ModelSelector) Pipeline() *ModelSelectorPipeline { } // Select runs the model selection pipeline (Filter → Score → Pick) and returns the selected model. -func (s *ModelSelector) Select(ctx context.Context, request *requesthandling.InferenceRequest, cycleState *plugin.CycleState, candidateModels []datalayer.Model) (result *modelselector.PipelineRunResult, err error) { +func (s *ModelSelector) Select(ctx context.Context, request *requesthandling.InferenceRequest, candidateModels []datalayer.Model) (result *modelselector.PipelineRunResult, err error) { logger := log.FromContext(ctx) logger.V(logutil.VERBOSE).Info("Starting model selection", "candidateModels", len(candidateModels)) @@ -64,7 +63,7 @@ func (s *ModelSelector) Select(ctx context.Context, request *requesthandling.Inf return nil, err } - result, err = s.pipeline.Run(ctx, request, cycleState, candidateModels) + result, err = s.pipeline.Run(ctx, request, candidateModels) if err != nil { logger.V(logutil.VERBOSE).Info("Model selection failed", "error", err.Error()) return nil, err diff --git a/pkg/modelselector/modelselector_test.go b/pkg/modelselector/modelselector_test.go index 35eac1f3..1f018c8c 100644 --- a/pkg/modelselector/modelselector_test.go +++ b/pkg/modelselector/modelselector_test.go @@ -77,7 +77,7 @@ func TestSelect(t *testing.T) { selector := NewModelSelector(pipeline) - result, err := selector.Select(context.Background(), requesthandling.NewInferenceRequest(), plugin.NewCycleState(), tt.models) + result, err := selector.Select(context.Background(), requesthandling.NewInferenceRequest(), tt.models) if tt.wantErr { if err == nil { @@ -141,7 +141,7 @@ func TestSelectWithFilterAndScorer(t *testing.T) { selector := NewModelSelector(pipeline) - result, err := selector.Select(context.Background(), requesthandling.NewInferenceRequest(), plugin.NewCycleState(), []datalayer.Model{modelA, modelB, modelC}) + result, err := selector.Select(context.Background(), requesthandling.NewInferenceRequest(), []datalayer.Model{modelA, modelB, modelC}) if err != nil { t.Fatalf("unexpected error: %v", err) } diff --git a/test/integration/body_mutation_test.go b/test/integration/body_mutation_test.go index 8d088aef..e115a8ac 100644 --- a/test/integration/body_mutation_test.go +++ b/test/integration/body_mutation_test.go @@ -45,7 +45,7 @@ func (p *bodyMutatingPlugin) TypedName() plugin.TypedName { return plugin.TypedName{Type: "test-body-mutator", Name: "test-body-mutator"} } -func (p *bodyMutatingPlugin) ProcessRequest(_ context.Context, _ *plugin.CycleState, request *requesthandling.InferenceRequest) error { +func (p *bodyMutatingPlugin) ProcessRequest(_ context.Context, request *requesthandling.InferenceRequest) error { request.SetBodyField(p.fieldName, p.fieldValue) return nil } From 9c83ca5327c06079195b208bafa9d620f52687b6 Mon Sep 17 00:00:00 2001 From: szedan Date: Wed, 12 Aug 2026 13:24:08 +0300 Subject: [PATCH 2/2] =?UTF-8?q?Align=20FactoryFunc=20signature=20with=20EP?= =?UTF-8?q?P:=20json.RawMessage=20=E2=86=92=20*json.Decoder?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace json.RawMessage with *json.Decoder in FactoryFunc and all 16 factory implementations. Config loader now wraps parameters via StrictDecoder (imported from EPP) to enable DisallowUnknownFields validation on plugin config parsing. Closes #293 Signed-off-by: szedan --- pkg/config/loader/configloader.go | 2 +- pkg/config/loader/configloader_test.go | 27 +++++++++---------- pkg/framework/interface/plugin/epp.go | 4 +++ pkg/framework/interface/plugin/registry.go | 7 +++-- .../datalayer/modelconfigcollector/plugin.go | 7 +++-- .../modelconfigcollector/plugin_test.go | 14 +++++----- .../datalayer/requestcostmetadata/plugin.go | 6 ++--- .../requestcostmetadata/plugin_test.go | 12 ++++----- .../datalayer/requestmetadata/plugin.go | 2 +- .../datalayer/requestmetadata/plugin_test.go | 3 +-- .../modelselector/filter/modelgroup/filter.go | 2 +- .../modelselector/picker/maxscore/picker.go | 2 +- .../modelselector/picker/random/picker.go | 2 +- .../picker/random/picker_test.go | 4 ++- .../picker/weightedrandom/picker.go | 2 +- .../modelselector/scorer/costaware/plugin.go | 2 +- .../scorer/costaware/plugin_test.go | 5 ++-- .../modelselector/scorer/costguard/plugin.go | 6 ++--- .../scorer/costguard/plugin_test.go | 17 ++++++------ .../scorer/inflightrequests/plugin.go | 2 +- .../scorer/sessionaffinity/plugin.go | 6 ++--- .../scorer/sessionaffinity/plugin_test.go | 25 ++++++++--------- .../base_model_to_header.go | 2 +- .../base_model_to_header_test.go | 8 +++--- .../bodyfieldtoheader/body_field_to_header.go | 6 ++--- .../body_field_to_header_test.go | 17 ++++++------ .../requesthandling/modelselector/plugin.go | 2 +- .../modelselector/plugin_test.go | 3 +-- .../single/single_profile_picker.go | 2 +- .../modelnametoheader/plugin.go | 6 ++--- .../modelnametoheader/plugin_test.go | 5 ++-- 31 files changed, 112 insertions(+), 98 deletions(-) diff --git a/pkg/config/loader/configloader.go b/pkg/config/loader/configloader.go index 28d8c62a..4e16d44e 100644 --- a/pkg/config/loader/configloader.go +++ b/pkg/config/loader/configloader.go @@ -199,7 +199,7 @@ func instantiatePlugins(configuredPlugins []configapi.PluginSpec, handle plugin. if !ok { return fmt.Errorf("plugin type '%s' is not registered", spec.Type) } - plugin, err := factory(spec.Name, spec.Parameters, handle) + plugin, err := factory(spec.Name, plugin.StrictDecoder(spec.Parameters), handle) if err != nil { return fmt.Errorf("failed to create plugin '%s' (type: %s): %w", spec.Name, spec.Type, err) } diff --git a/pkg/config/loader/configloader_test.go b/pkg/config/loader/configloader_test.go index 5db55779..ac6e3279 100644 --- a/pkg/config/loader/configloader_test.go +++ b/pkg/config/loader/configloader_test.go @@ -558,37 +558,36 @@ func registerTestPlugins(t *testing.T) { // Register standard test mocks. plugin.Register(testPluginType, - func(name string, params json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) { + func(name string, _ *json.Decoder, _ plugin.Handle) (plugin.Plugin, error) { return &mockPlugin{t: plugin.TypedName{Name: name, Type: testPluginType}}, nil }) plugin.Register(testProfilePicker, - func(name string, params json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) { + func(name string, _ *json.Decoder, _ plugin.Handle) (plugin.Plugin, error) { return &mockProfilePicker{mockPlugin{t: plugin.TypedName{Name: name, Type: testProfilePicker}}}, nil }) plugin.Register(testRequestProcType, - func(name string, params json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) { + func(name string, _ *json.Decoder, _ plugin.Handle) (plugin.Plugin, error) { return &mockRequestProcessor{mockPlugin{t: plugin.TypedName{Name: name, Type: testRequestProcType}}}, nil }) plugin.Register(testResponseProcType, - func(name string, params json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) { + func(name string, _ *json.Decoder, _ plugin.Handle) (plugin.Plugin, error) { return &mockResponseProcessor{mockPlugin{t: plugin.TypedName{Name: name, Type: testResponseProcType}}}, nil }) plugin.Register(testPickerType, - func(name string, params json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) { + func(name string, _ *json.Decoder, _ plugin.Handle) (plugin.Plugin, error) { return &mockPicker{mockPlugin{t: plugin.TypedName{Name: name, Type: testPickerType}}}, nil }) - plugin.Register(testScorerType, func(name string, params json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) { - // Attempt to unmarshal to trigger errors for invalid JSON in tests. - if len(params) > 0 { + plugin.Register(testScorerType, func(name string, params *json.Decoder, _ plugin.Handle) (plugin.Plugin, error) { + if params != nil { var p struct { Cost float32 `json:"cost"` } - if err := json.Unmarshal(params, &p); err != nil { + if err := params.Decode(&p); err != nil { return nil, err } } @@ -596,22 +595,22 @@ func registerTestPlugins(t *testing.T) { }) plugin.Register(testFilterType, - func(name string, _ json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) { + func(name string, _ *json.Decoder, _ plugin.Handle) (plugin.Plugin, error) { return &mockFilter{mockPlugin{t: plugin.TypedName{Name: name, Type: testFilterType}}}, nil }) plugin.Register(testExtractorType, - func(name string, _ json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) { + func(name string, _ *json.Decoder, _ plugin.Handle) (plugin.Plugin, error) { return &mockExtractor{mockPlugin{t: plugin.TypedName{Name: name, Type: testExtractorType}}}, nil }) plugin.Register(testCollectorType, - func(name string, _ json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) { + func(name string, _ *json.Decoder, _ plugin.Handle) (plugin.Plugin, error) { return &mockCollector{mockPlugin{t: plugin.TypedName{Name: name, Type: testCollectorType}}}, nil }) plugin.Register(testDataSourceType, - func(name string, _ json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) { + func(name string, _ *json.Decoder, _ plugin.Handle) (plugin.Plugin, error) { return &mockDataSource{mockPlugin{t: plugin.TypedName{Name: name, Type: testDataSourceType}}}, nil }) } @@ -624,7 +623,7 @@ func registerModelSelectorPlugins(t *testing.T) { plugin.Register(costaware.CostScorerType, costaware.CostScorerFactory) plugin.Register(maxscore.MaxScorePickerType, maxscore.MaxScorePickerFactory) plugin.Register(testFilterType, - func(name string, _ json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) { + func(name string, _ *json.Decoder, _ plugin.Handle) (plugin.Plugin, error) { return &mockFilter{mockPlugin{t: plugin.TypedName{Name: name, Type: testFilterType}}}, nil }) } diff --git a/pkg/framework/interface/plugin/epp.go b/pkg/framework/interface/plugin/epp.go index 8dbc0678..19a260d2 100644 --- a/pkg/framework/interface/plugin/epp.go +++ b/pkg/framework/interface/plugin/epp.go @@ -26,3 +26,7 @@ type Plugin = eppplugin.Plugin type TypedName = eppplugin.TypedName type HandlePlugins = eppplugin.HandlePlugins + +// StrictDecoder converts raw JSON plugin parameters into a strict *json.Decoder +// (DisallowUnknownFields), or returns nil when raw is empty. +var StrictDecoder = eppplugin.StrictDecoder diff --git a/pkg/framework/interface/plugin/registry.go b/pkg/framework/interface/plugin/registry.go index e56d01a0..0984d8c8 100644 --- a/pkg/framework/interface/plugin/registry.go +++ b/pkg/framework/interface/plugin/registry.go @@ -21,8 +21,11 @@ import ( ) // Factory is the definition of the factory functions that are used to instantiate plugins -// specified in a configuration. -type FactoryFunc func(name string, parameters json.RawMessage, handle Handle) (Plugin, error) +// specified in a configuration. The framework provides a strict decoder +// (DisallowUnknownFields) over the plugin's raw parameters, or nil when the plugin was +// instantiated without parameters. Factories that ignore parameters can take the decoder +// as `_ *json.Decoder`. +type FactoryFunc func(name string, parameters *json.Decoder, handle Handle) (Plugin, error) // Register is a static function that can be called to register plugin factory functions. func Register(pluginType string, factory FactoryFunc) { diff --git a/pkg/framework/plugins/datalayer/modelconfigcollector/plugin.go b/pkg/framework/plugins/datalayer/modelconfigcollector/plugin.go index 244b7529..ba4b4e66 100644 --- a/pkg/framework/plugins/datalayer/modelconfigcollector/plugin.go +++ b/pkg/framework/plugins/datalayer/modelconfigcollector/plugin.go @@ -80,9 +80,12 @@ type ModelConfigDataSource struct { // DatasourceFactory creates a ModelConfigDataSource from the plugin handle and raw JSON config. // It validates that modelsPath is set and that the file exists; content parsing happens in Start. -func DatasourceFactory(name string, rawCfg json.RawMessage, h plugin.Handle) (plugin.Plugin, error) { +func DatasourceFactory(name string, parameters *json.Decoder, h plugin.Handle) (plugin.Plugin, error) { + if parameters == nil { + return nil, errors.New("modelsPath is required") + } var cfg PluginConfig - if err := json.Unmarshal(rawCfg, &cfg); err != nil { + if err := parameters.Decode(&cfg); err != nil { return nil, err } if cfg.ModelsPath == "" { diff --git a/pkg/framework/plugins/datalayer/modelconfigcollector/plugin_test.go b/pkg/framework/plugins/datalayer/modelconfigcollector/plugin_test.go index 6a8f4e8d..2abe0876 100644 --- a/pkg/framework/plugins/datalayer/modelconfigcollector/plugin_test.go +++ b/pkg/framework/plugins/datalayer/modelconfigcollector/plugin_test.go @@ -51,7 +51,7 @@ func (f *fakeHandle) GetAllPluginsWithNames() map[string]plugin.Plugin { return func useFactory(t *testing.T, path string, ds datalayer.Datastore) *ModelConfigDataSource { t.Helper() rawCfg, _ := json.Marshal(PluginConfig{ModelsPath: path}) - p, err := DatasourceFactory("test", rawCfg, &fakeHandle{ds: ds}) + p, err := DatasourceFactory("test", plugin.StrictDecoder(rawCfg), &fakeHandle{ds: ds}) if err != nil { t.Fatalf("DatasourceFactory: %v", err) } @@ -134,7 +134,7 @@ func waitForUpdatedConfig(t *testing.T, ds datalayer.Datastore, wantCount int, t // that is not valid JSON at all. func TestDatasourceFactory_InvalidJSON(t *testing.T) { ds := datastore.NewFakeDataStore() - _, err := DatasourceFactory("x", json.RawMessage(`not-json`), &fakeHandle{ds: ds}) + _, err := DatasourceFactory("x", plugin.StrictDecoder(json.RawMessage(`not-json`)), &fakeHandle{ds: ds}) if err == nil { t.Error("expected error for invalid JSON plugin config, got nil") } @@ -143,7 +143,7 @@ func TestDatasourceFactory_InvalidJSON(t *testing.T) { // TestDatasourceFactory_EmptyInput ensures the factory rejects an empty config payload. func TestDatasourceFactory_EmptyInput(t *testing.T) { ds := datastore.NewFakeDataStore() - _, err := DatasourceFactory("x", json.RawMessage(``), &fakeHandle{ds: ds}) + _, err := DatasourceFactory("x", nil, &fakeHandle{ds: ds}) if err == nil { t.Error("expected error for empty plugin config input, got nil") } @@ -154,7 +154,7 @@ func TestDatasourceFactory_EmptyInput(t *testing.T) { func TestDatasourceFactory_MissingModelsPath(t *testing.T) { ds := datastore.NewFakeDataStore() rawCfg, _ := json.Marshal(PluginConfig{}) // modelsPath omitted → empty string - _, err := DatasourceFactory("x", rawCfg, &fakeHandle{ds: ds}) + _, err := DatasourceFactory("x", plugin.StrictDecoder(rawCfg), &fakeHandle{ds: ds}) if err == nil { t.Error("expected error for missing modelsPath, got nil") } @@ -165,7 +165,7 @@ func TestDatasourceFactory_MissingModelsPath(t *testing.T) { func TestDatasourceFactory_FileNotExist(t *testing.T) { ds := datastore.NewFakeDataStore() rawCfg, _ := json.Marshal(PluginConfig{ModelsPath: "/no/such/file.json"}) - _, err := DatasourceFactory("x", rawCfg, &fakeHandle{ds: ds}) + _, err := DatasourceFactory("x", plugin.StrictDecoder(rawCfg), &fakeHandle{ds: ds}) if err == nil { t.Error("expected error for non-existent file, got nil") } @@ -177,7 +177,7 @@ func TestDatasourceFactory_DirectoryNotFile(t *testing.T) { ds := datastore.NewFakeDataStore() dir := t.TempDir() rawCfg, _ := json.Marshal(PluginConfig{ModelsPath: dir}) - _, err := DatasourceFactory("x", rawCfg, &fakeHandle{ds: ds}) + _, err := DatasourceFactory("x", plugin.StrictDecoder(rawCfg), &fakeHandle{ds: ds}) if err == nil { t.Error("expected error for directory path, got nil") } @@ -206,7 +206,7 @@ func TestStart_InvalidFileContent(t *testing.T) { ds := datastore.NewFakeDataStore() path := writeTempRaw(t, `this is not valid json {{{`) rawCfg, _ := json.Marshal(PluginConfig{ModelsPath: path}) - p, err := DatasourceFactory("x", rawCfg, &fakeHandle{ds: ds}) + p, err := DatasourceFactory("x", plugin.StrictDecoder(rawCfg), &fakeHandle{ds: ds}) if err != nil { t.Fatalf("DatasourceFactory: %v", err) } diff --git a/pkg/framework/plugins/datalayer/requestcostmetadata/plugin.go b/pkg/framework/plugins/datalayer/requestcostmetadata/plugin.go index 1d4e0ec9..5637064b 100644 --- a/pkg/framework/plugins/datalayer/requestcostmetadata/plugin.go +++ b/pkg/framework/plugins/datalayer/requestcostmetadata/plugin.go @@ -87,14 +87,14 @@ type RequestCostMetadataExtractorConfig struct { } // ExtractorFactory creates a RequestCostMetadataExtractor wired to the shared Datastore. -func ExtractorFactory(name string, parameters json.RawMessage, h plugin.Handle) (plugin.Plugin, error) { +func ExtractorFactory(name string, parameters *json.Decoder, h plugin.Handle) (plugin.Plugin, error) { config := RequestCostMetadataExtractorConfig{ Compression: defaultCompression, FlushIntervalDuration: defaultFlushIntervalDuration.String(), WindowDuration: defaultWindowDuration.String(), } - if len(parameters) > 0 { - if err := json.Unmarshal(parameters, &config); err != nil { + if parameters != nil { + if err := parameters.Decode(&config); err != nil { return nil, fmt.Errorf("failed to parse parameters for plugin %q: %w", name, err) } } diff --git a/pkg/framework/plugins/datalayer/requestcostmetadata/plugin_test.go b/pkg/framework/plugins/datalayer/requestcostmetadata/plugin_test.go index eaddf1fa..c7b1b300 100644 --- a/pkg/framework/plugins/datalayer/requestcostmetadata/plugin_test.go +++ b/pkg/framework/plugins/datalayer/requestcostmetadata/plugin_test.go @@ -151,7 +151,7 @@ func newTestExtractor(t *testing.T) (*RequestCostMetadataExtractor, datalayer.Da func TestExtractorFactory_HonorsConfig(t *testing.T) { ds := datastore.NewFakeDataStore() raw := json.RawMessage(`{"compression":50,"flushIntervalDuration":"1m","windowDuration":"30m"}`) - p, err := ExtractorFactory("x", raw, &fakeHandle{ds: ds}) + p, err := ExtractorFactory("x", plugin.StrictDecoder(raw), &fakeHandle{ds: ds}) if err != nil { t.Fatalf("ExtractorFactory: %v", err) } @@ -170,12 +170,12 @@ func TestExtractorFactory_HonorsConfig(t *testing.T) { func TestExtractorFactory_RejectsInvalidDurations(t *testing.T) { tests := []struct { name string - raw json.RawMessage + raw *json.Decoder }{ - {"flushIntervalDuration malformed", json.RawMessage(`{"compression":200,"flushIntervalDuration":"not-a-duration"}`)}, - {"flushIntervalDuration negative", json.RawMessage(`{"compression":200,"flushIntervalDuration":"-1s"}`)}, - {"windowDuration malformed", json.RawMessage(`{"compression":200,"windowDuration":"not-a-duration"}`)}, - {"windowDuration negative", json.RawMessage(`{"compression":200,"windowDuration":"-1s"}`)}, + {"flushIntervalDuration malformed", plugin.StrictDecoder(json.RawMessage(`{"compression":200,"flushIntervalDuration":"not-a-duration"}`))}, + {"flushIntervalDuration negative", plugin.StrictDecoder(json.RawMessage(`{"compression":200,"flushIntervalDuration":"-1s"}`))}, + {"windowDuration malformed", plugin.StrictDecoder(json.RawMessage(`{"compression":200,"windowDuration":"not-a-duration"}`))}, + {"windowDuration negative", plugin.StrictDecoder(json.RawMessage(`{"compression":200,"windowDuration":"-1s"}`))}, } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { diff --git a/pkg/framework/plugins/datalayer/requestmetadata/plugin.go b/pkg/framework/plugins/datalayer/requestmetadata/plugin.go index 830e8dc7..0a6b6e72 100644 --- a/pkg/framework/plugins/datalayer/requestmetadata/plugin.go +++ b/pkg/framework/plugins/datalayer/requestmetadata/plugin.go @@ -37,7 +37,7 @@ const ( var _ dlsrc.Extractor = &RequestMetadataExtractor{} // ExtractorFactory creates a RequestMetadataExtractor wired to the shared DataStore. -func ExtractorFactory(name string, _ json.RawMessage, h plugin.Handle) (plugin.Plugin, error) { +func ExtractorFactory(name string, _ *json.Decoder, h plugin.Handle) (plugin.Plugin, error) { return NewRequestMetadataExtractor(h.Datastore()).WithName(name), nil } diff --git a/pkg/framework/plugins/datalayer/requestmetadata/plugin_test.go b/pkg/framework/plugins/datalayer/requestmetadata/plugin_test.go index afcc246e..f939cdfb 100644 --- a/pkg/framework/plugins/datalayer/requestmetadata/plugin_test.go +++ b/pkg/framework/plugins/datalayer/requestmetadata/plugin_test.go @@ -18,7 +18,6 @@ package requestmetadata import ( "context" - "encoding/json" "testing" "time" @@ -206,7 +205,7 @@ func TestExtractorFactoryWiresDatastore(t *testing.T) { ds := datastore.NewFakeDataStore() h := &fakeHandle{ds: ds} - p, err := ExtractorFactory("my-extractor", json.RawMessage(`{}`), h) + p, err := ExtractorFactory("my-extractor", nil, h) if err != nil { t.Fatalf("ExtractorFactory returned error: %v", err) } diff --git a/pkg/framework/plugins/modelselector/filter/modelgroup/filter.go b/pkg/framework/plugins/modelselector/filter/modelgroup/filter.go index 092df86d..67280cb2 100644 --- a/pkg/framework/plugins/modelselector/filter/modelgroup/filter.go +++ b/pkg/framework/plugins/modelselector/filter/modelgroup/filter.go @@ -59,7 +59,7 @@ var _ modelselector.Filter = &ModelGroupFilter{} // no parameters: group membership is resolved at filter time from each candidate // model's modelgroups.GroupsAttributeKey attribute, populated by the // model-config-datasource plugin from the shared config file's "groups" list. -func ModelGroupFilterFactory(name string, _ json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) { +func ModelGroupFilterFactory(name string, _ *json.Decoder, _ plugin.Handle) (plugin.Plugin, error) { return NewModelGroupFilter().WithName(name), nil } diff --git a/pkg/framework/plugins/modelselector/picker/maxscore/picker.go b/pkg/framework/plugins/modelselector/picker/maxscore/picker.go index 1f081c50..3806b17b 100644 --- a/pkg/framework/plugins/modelselector/picker/maxscore/picker.go +++ b/pkg/framework/plugins/modelselector/picker/maxscore/picker.go @@ -38,7 +38,7 @@ const ( var _ modelselector.Picker = &MaxScorePicker{} // MaxScorePickerFactory defines the factory function for MaxScorePicker. -func MaxScorePickerFactory(name string, _ json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) { +func MaxScorePickerFactory(name string, _ *json.Decoder, _ plugin.Handle) (plugin.Plugin, error) { return NewMaxScorePicker().WithName(name), nil } diff --git a/pkg/framework/plugins/modelselector/picker/random/picker.go b/pkg/framework/plugins/modelselector/picker/random/picker.go index 0c5e0ca9..10a8d495 100644 --- a/pkg/framework/plugins/modelselector/picker/random/picker.go +++ b/pkg/framework/plugins/modelselector/picker/random/picker.go @@ -41,7 +41,7 @@ const ( var _ modelselector.Picker = &RandomPicker{} // RandomPickerFactory defines the factory function for RandomPicker. -func RandomPickerFactory(name string, _ json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) { +func RandomPickerFactory(name string, _ *json.Decoder, _ plugin.Handle) (plugin.Plugin, error) { return NewRandomPicker().WithName(name), nil } diff --git a/pkg/framework/plugins/modelselector/picker/random/picker_test.go b/pkg/framework/plugins/modelselector/picker/random/picker_test.go index b6f45e4a..b889f944 100644 --- a/pkg/framework/plugins/modelselector/picker/random/picker_test.go +++ b/pkg/framework/plugins/modelselector/picker/random/picker_test.go @@ -18,10 +18,12 @@ package random import ( "context" + "encoding/json" "testing" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/datalayer" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/modelselector" + "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin" ) func TestRandomPicker_Pick(t *testing.T) { @@ -183,7 +185,7 @@ func TestRandomPickerFactory(t *testing.T) { }) t.Run("ignores config parameter", func(t *testing.T) { - _, err := RandomPickerFactory("test", []byte(`{"some": "config"}`), nil) + _, err := RandomPickerFactory("test", plugin.StrictDecoder(json.RawMessage(`{"some": "config"}`)), nil) if err != nil { t.Fatalf("config should be ignored, got error: %v", err) } diff --git a/pkg/framework/plugins/modelselector/picker/weightedrandom/picker.go b/pkg/framework/plugins/modelselector/picker/weightedrandom/picker.go index dddc035b..82ec07d0 100644 --- a/pkg/framework/plugins/modelselector/picker/weightedrandom/picker.go +++ b/pkg/framework/plugins/modelselector/picker/weightedrandom/picker.go @@ -47,7 +47,7 @@ type weightedScoredModel struct { } // WeightedRandomPickerFactory defines the factory function for WeightedRandomPicker. -func WeightedRandomPickerFactory(name string, _ json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) { +func WeightedRandomPickerFactory(name string, _ *json.Decoder, _ plugin.Handle) (plugin.Plugin, error) { return NewWeightedRandomPicker().WithName(name), nil } diff --git a/pkg/framework/plugins/modelselector/scorer/costaware/plugin.go b/pkg/framework/plugins/modelselector/scorer/costaware/plugin.go index a3f33a6d..3f33a6fc 100644 --- a/pkg/framework/plugins/modelselector/scorer/costaware/plugin.go +++ b/pkg/framework/plugins/modelselector/scorer/costaware/plugin.go @@ -43,7 +43,7 @@ const CostScorerType = "cost-scorer" var _ modelselector.Scorer = &CostScorer{} // CostScorerFactory defines the factory function for the CostScorer scorer -func CostScorerFactory(name string, _ json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) { +func CostScorerFactory(name string, _ *json.Decoder, _ plugin.Handle) (plugin.Plugin, error) { return NewCostScorer().WithName(name), nil } diff --git a/pkg/framework/plugins/modelselector/scorer/costaware/plugin_test.go b/pkg/framework/plugins/modelselector/scorer/costaware/plugin_test.go index 992658fe..b1c70df3 100644 --- a/pkg/framework/plugins/modelselector/scorer/costaware/plugin_test.go +++ b/pkg/framework/plugins/modelselector/scorer/costaware/plugin_test.go @@ -23,6 +23,7 @@ import ( "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/datalayer" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/datalayer/pricing" + "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/requesthandling" ) @@ -31,7 +32,7 @@ func TestFactory(t *testing.T) { tests := []struct { name string pluginName string - rawParameters json.RawMessage + rawParameters *json.Decoder expectError bool }{ { @@ -43,7 +44,7 @@ func TestFactory(t *testing.T) { { name: "factory with empty parameters", pluginName: "my-scorer", - rawParameters: json.RawMessage(`{}`), + rawParameters: plugin.StrictDecoder(json.RawMessage(`{}`)), expectError: false, }, } diff --git a/pkg/framework/plugins/modelselector/scorer/costguard/plugin.go b/pkg/framework/plugins/modelselector/scorer/costguard/plugin.go index a8257ccb..bc4a5606 100644 --- a/pkg/framework/plugins/modelselector/scorer/costguard/plugin.go +++ b/pkg/framework/plugins/modelselector/scorer/costguard/plugin.go @@ -104,15 +104,15 @@ type CostGuardScorer struct { } // ScorerFactory validates rawParameters into a Config and constructs a scorer. -func ScorerFactory(name string, rawParameters json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) { +func ScorerFactory(name string, parameters *json.Decoder, _ plugin.Handle) (plugin.Plugin, error) { config := Config{ Epsilon: defaultEpsilon, Alpha: defaultAlpha, Lambda: defaultLambda, PercentileMarginError: defaultPercentileMarginError, } - if len(rawParameters) > 0 { - if err := json.Unmarshal(rawParameters, &config); err != nil { + if parameters != nil { + if err := parameters.Decode(&config); err != nil { return nil, fmt.Errorf("costguard %q: failed to parse parameters: %w", name, err) } } diff --git a/pkg/framework/plugins/modelselector/scorer/costguard/plugin_test.go b/pkg/framework/plugins/modelselector/scorer/costguard/plugin_test.go index 042c998d..8bcfb4ef 100644 --- a/pkg/framework/plugins/modelselector/scorer/costguard/plugin_test.go +++ b/pkg/framework/plugins/modelselector/scorer/costguard/plugin_test.go @@ -27,6 +27,7 @@ import ( "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/datalayer" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/datalayer/accumulator" + "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/requesthandling" ) @@ -286,15 +287,15 @@ func TestScore_UnderExplored(t *testing.T) { // with an empty JSON object; both must produce the same defaulted scorer. func TestFactory_DefaultConfig(t *testing.T) { tests := []struct { - name string - raw json.RawMessage + name string + params *json.Decoder }{ {"nil parameters", nil}, - {"empty object", json.RawMessage(`{}`)}, + {"empty object", plugin.StrictDecoder(json.RawMessage(`{}`))}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - p, err := ScorerFactory("test-cg", tt.raw, nil) + p, err := ScorerFactory("test-cg", tt.params, nil) require.NoError(t, err) s, ok := p.(*CostGuardScorer) require.True(t, ok) @@ -311,8 +312,8 @@ func TestFactory_DefaultConfig(t *testing.T) { // TestFactory_CustomConfig verifies that custom parameters override defaults. func TestFactory_CustomConfig(t *testing.T) { - raw := json.RawMessage(`{"epsilon":0.2,"alpha":0.9,"lambda":2.0,"percentileMarginError":0.05}`) - p, err := ScorerFactory("custom", raw, nil) + params := plugin.StrictDecoder(json.RawMessage(`{"epsilon":0.2,"alpha":0.9,"lambda":2.0,"percentileMarginError":0.05}`)) + p, err := ScorerFactory("custom", params, nil) require.NoError(t, err) s := p.(*CostGuardScorer) assert.Equal(t, 0.2, s.epsilon) @@ -341,7 +342,7 @@ func TestFactory_ValidationErrors(t *testing.T) { } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - _, err := ScorerFactory("bad", json.RawMessage(tt.raw), nil) + _, err := ScorerFactory("bad", plugin.StrictDecoder(json.RawMessage(tt.raw)), nil) require.Error(t, err) }) } @@ -387,7 +388,7 @@ func TestFactory_AcceptedBoundaries(t *testing.T) { } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - p, err := ScorerFactory("boundary", json.RawMessage(tt.raw), nil) + p, err := ScorerFactory("boundary", plugin.StrictDecoder(json.RawMessage(tt.raw)), nil) require.NoError(t, err) tt.check(t, p.(*CostGuardScorer)) }) diff --git a/pkg/framework/plugins/modelselector/scorer/inflightrequests/plugin.go b/pkg/framework/plugins/modelselector/scorer/inflightrequests/plugin.go index bb2e0519..d26a5317 100644 --- a/pkg/framework/plugins/modelselector/scorer/inflightrequests/plugin.go +++ b/pkg/framework/plugins/modelselector/scorer/inflightrequests/plugin.go @@ -42,7 +42,7 @@ type InflightRequestsScorer struct { } // ScorerFactory is the factory function for InflightRequestsScorer. -func ScorerFactory(name string, _ json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) { +func ScorerFactory(name string, _ *json.Decoder, _ plugin.Handle) (plugin.Plugin, error) { return NewInflightRequestsScorer().WithName(name), nil } diff --git a/pkg/framework/plugins/modelselector/scorer/sessionaffinity/plugin.go b/pkg/framework/plugins/modelselector/scorer/sessionaffinity/plugin.go index 51a11788..de27e4bb 100644 --- a/pkg/framework/plugins/modelselector/scorer/sessionaffinity/plugin.go +++ b/pkg/framework/plugins/modelselector/scorer/sessionaffinity/plugin.go @@ -113,11 +113,11 @@ type SessionAffinityScorerConfig struct { } // ScorerFactory creates a new SessionAffinityScorer from config. -func ScorerFactory(name string, rawParameters json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) { +func ScorerFactory(name string, parameters *json.Decoder, _ plugin.Handle) (plugin.Plugin, error) { var config SessionAffinityScorerConfig - if len(rawParameters) > 0 { - if err := json.Unmarshal(rawParameters, &config); err != nil { + if parameters != nil { + if err := parameters.Decode(&config); err != nil { return nil, fmt.Errorf("failed to parse parameters for %q plugin: %w", PluginType, err) } } diff --git a/pkg/framework/plugins/modelselector/scorer/sessionaffinity/plugin_test.go b/pkg/framework/plugins/modelselector/scorer/sessionaffinity/plugin_test.go index f8c1791f..38f18700 100644 --- a/pkg/framework/plugins/modelselector/scorer/sessionaffinity/plugin_test.go +++ b/pkg/framework/plugins/modelselector/scorer/sessionaffinity/plugin_test.go @@ -27,6 +27,7 @@ import ( "github.com/stretchr/testify/require" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/datalayer" + "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/requesthandling" ) @@ -60,8 +61,8 @@ func TestFactory_DefaultConfig(t *testing.T) { // Verify custom parameters override defaults. func TestFactory_CustomConfig(t *testing.T) { - raw := json.RawMessage(`{"sessionIdKey":"x-custom-id","maxSessions":500,"ttlSeconds":300}`) - p, err := ScorerFactory("custom", raw, nil) + params := plugin.StrictDecoder(json.RawMessage(`{"sessionIdKey":"x-custom-id","maxSessions":500,"ttlSeconds":300}`)) + p, err := ScorerFactory("custom", params, nil) require.NoError(t, err) s := p.(*SessionAffinityScorer) assert.Equal(t, "x-custom-id", s.sessionIDKey) @@ -70,8 +71,8 @@ func TestFactory_CustomConfig(t *testing.T) { // Verify factory rejects malformed JSON. func TestFactory_InvalidJSON(t *testing.T) { - raw := json.RawMessage(`{invalid}`) - _, err := ScorerFactory("bad", raw, nil) + params := plugin.StrictDecoder(json.RawMessage(`{invalid}`)) + _, err := ScorerFactory("bad", params, nil) assert.Error(t, err) } @@ -174,8 +175,8 @@ func TestScore_StoresSessionIDInRequestAttributes(t *testing.T) { // Session ID lookup uses the configured custom key. func TestScore_UsesCustomSessionIDKey(t *testing.T) { - raw := json.RawMessage(`{"sessionIdKey":"x-conv-id"}`) - p, err := ScorerFactory("custom-key", raw, nil) + params := plugin.StrictDecoder(json.RawMessage(`{"sessionIdKey":"x-conv-id"}`)) + p, err := ScorerFactory("custom-key", params, nil) require.NoError(t, err) s := p.(*SessionAffinityScorer) s.cache.Add("conv-123", "model-a") @@ -367,8 +368,8 @@ func TestEndToEnd_OptimisticSessionID(t *testing.T) { // LRU eviction removes the least recently used entry when at capacity. func TestEviction_LRU_RemovesLeastRecent(t *testing.T) { - raw := json.RawMessage(`{"maxSessions":3}`) - p, err := ScorerFactory("evict-test", raw, nil) + params := plugin.StrictDecoder(json.RawMessage(`{"maxSessions":3}`)) + p, err := ScorerFactory("evict-test", params, nil) require.NoError(t, err) s := p.(*SessionAffinityScorer) @@ -394,8 +395,8 @@ func TestEviction_LRU_RemovesLeastRecent(t *testing.T) { // Get does not return entries after TTL has expired. func TestTTL_GetDoesNotReturnExpiredEntry(t *testing.T) { - raw := json.RawMessage(`{"maxSessions":100,"ttlSeconds":1}`) - p, err := ScorerFactory("ttl-test", raw, nil) + params := plugin.StrictDecoder(json.RawMessage(`{"maxSessions":100,"ttlSeconds":1}`)) + p, err := ScorerFactory("ttl-test", params, nil) require.NoError(t, err) s := p.(*SessionAffinityScorer) @@ -415,8 +416,8 @@ func TestTTL_GetDoesNotReturnExpiredEntry(t *testing.T) { // Cache respects maxSessions limit. func TestCapacity_DoesNotExceedMax(t *testing.T) { - raw := json.RawMessage(`{"maxSessions":10}`) - p, err := ScorerFactory("cap-test", raw, nil) + params := plugin.StrictDecoder(json.RawMessage(`{"maxSessions":10}`)) + p, err := ScorerFactory("cap-test", params, nil) require.NoError(t, err) s := p.(*SessionAffinityScorer) diff --git a/pkg/framework/plugins/requesthandling/basemodelextractor/base_model_to_header.go b/pkg/framework/plugins/requesthandling/basemodelextractor/base_model_to_header.go index 1736aeb1..c761a6a5 100644 --- a/pkg/framework/plugins/requesthandling/basemodelextractor/base_model_to_header.go +++ b/pkg/framework/plugins/requesthandling/basemodelextractor/base_model_to_header.go @@ -46,7 +46,7 @@ type BaseModelToHeaderPlugin struct { } // BaseModelToHeaderPluginFactory defines the factory function for BaseModelToHeaderPlugin -func BaseModelToHeaderPluginFactory(name string, _ json.RawMessage, handle plugin.Handle) (plugin.Plugin, error) { +func BaseModelToHeaderPluginFactory(name string, _ *json.Decoder, handle plugin.Handle) (plugin.Plugin, error) { plugin, err := NewBaseModelToHeaderPlugin(handle.ReconcilerBuilder, handle.Client()) if err != nil { return nil, fmt.Errorf("failed to create plugin '%s' - %w", BaseModelToHeaderPluginType, err) diff --git a/pkg/framework/plugins/requesthandling/basemodelextractor/base_model_to_header_test.go b/pkg/framework/plugins/requesthandling/basemodelextractor/base_model_to_header_test.go index b4c20dfb..74e41c04 100644 --- a/pkg/framework/plugins/requesthandling/basemodelextractor/base_model_to_header_test.go +++ b/pkg/framework/plugins/requesthandling/basemodelextractor/base_model_to_header_test.go @@ -77,13 +77,13 @@ func TestBaseModelToHeaderPluginFactory(t *testing.T) { tests := []struct { name string pluginName string - rawParams json.RawMessage + rawParams *json.Decoder wantName string }{ { name: "valid empty config", pluginName: "my-base-model-plugin", - rawParams: json.RawMessage(`{}`), + rawParams: plugin.StrictDecoder(json.RawMessage(`{}`)), wantName: "my-base-model-plugin", }, { @@ -95,13 +95,13 @@ func TestBaseModelToHeaderPluginFactory(t *testing.T) { { name: "JSON null", pluginName: "my-plugin", - rawParams: json.RawMessage(`null`), + rawParams: plugin.StrictDecoder(json.RawMessage(`null`)), wantName: "my-plugin", }, { name: "empty parameters", pluginName: "my-plugin", - rawParams: json.RawMessage(``), + rawParams: nil, wantName: "my-plugin", }, } diff --git a/pkg/framework/plugins/requesthandling/bodyfieldtoheader/body_field_to_header.go b/pkg/framework/plugins/requesthandling/bodyfieldtoheader/body_field_to_header.go index f3a6ecc9..46b137fe 100644 --- a/pkg/framework/plugins/requesthandling/bodyfieldtoheader/body_field_to_header.go +++ b/pkg/framework/plugins/requesthandling/bodyfieldtoheader/body_field_to_header.go @@ -46,11 +46,11 @@ type BodyFieldToHeaderConfig struct { } // BodyFieldToHeaderPluginFactory defines the factory function for NewBodyFieldToHeaderPlugin. -func BodyFieldToHeaderPluginFactory(name string, rawParameters json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) { +func BodyFieldToHeaderPluginFactory(name string, parameters *json.Decoder, _ plugin.Handle) (plugin.Plugin, error) { var config BodyFieldToHeaderConfig - if len(rawParameters) > 0 { - if err := json.Unmarshal(rawParameters, &config); err != nil { + if parameters != nil { + if err := parameters.Decode(&config); err != nil { return nil, fmt.Errorf("failed to parse the parameters of the '%s' plugin - %w", BodyFieldToHeaderPluginType, err) } } diff --git a/pkg/framework/plugins/requesthandling/bodyfieldtoheader/body_field_to_header_test.go b/pkg/framework/plugins/requesthandling/bodyfieldtoheader/body_field_to_header_test.go index 54b29918..aefd8982 100644 --- a/pkg/framework/plugins/requesthandling/bodyfieldtoheader/body_field_to_header_test.go +++ b/pkg/framework/plugins/requesthandling/bodyfieldtoheader/body_field_to_header_test.go @@ -21,6 +21,7 @@ import ( "encoding/json" "testing" + "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/requesthandling" ) @@ -101,38 +102,38 @@ func TestBodyFieldToHeaderPluginFactory(t *testing.T) { tests := []struct { name string pluginName string - rawParams json.RawMessage + rawParams *json.Decoder wantErr bool wantName string }{ { name: "valid config", pluginName: "my-plugin", - rawParams: json.RawMessage(`{"fieldName":"model","headerName":"X-Gateway-Model"}`), + rawParams: plugin.StrictDecoder(json.RawMessage(`{"fieldName":"model","headerName":"X-Gateway-Model"}`)), wantName: "my-plugin", }, { name: "invalid JSON", pluginName: "my-plugin", - rawParams: json.RawMessage(`{invalid`), + rawParams: plugin.StrictDecoder(json.RawMessage(`{invalid`)), wantErr: true, }, { name: "missing fieldName", pluginName: "my-plugin", - rawParams: json.RawMessage(`{"header_name":"X-Gateway-Model"}`), + rawParams: plugin.StrictDecoder(json.RawMessage(`{"header_name":"X-Gateway-Model"}`)), wantErr: true, }, { name: "missing headerName", pluginName: "my-plugin", - rawParams: json.RawMessage(`{"field_name":"model"}`), + rawParams: plugin.StrictDecoder(json.RawMessage(`{"field_name":"model"}`)), wantErr: true, }, { name: "empty parameters", pluginName: "my-plugin", - rawParams: json.RawMessage(``), + rawParams: nil, wantErr: true, }, { @@ -144,13 +145,13 @@ func TestBodyFieldToHeaderPluginFactory(t *testing.T) { { name: "JSON null", pluginName: "my-plugin", - rawParams: json.RawMessage(`null`), + rawParams: plugin.StrictDecoder(json.RawMessage(`null`)), wantErr: true, }, { name: "empty JSON object", pluginName: "my-plugin", - rawParams: json.RawMessage(`{}`), + rawParams: plugin.StrictDecoder(json.RawMessage(`{}`)), wantErr: true, }, } diff --git a/pkg/framework/plugins/requesthandling/modelselector/plugin.go b/pkg/framework/plugins/requesthandling/modelselector/plugin.go index 4503e460..0bf461cc 100644 --- a/pkg/framework/plugins/requesthandling/modelselector/plugin.go +++ b/pkg/framework/plugins/requesthandling/modelselector/plugin.go @@ -41,7 +41,7 @@ var _ requesthandling.RequestProcessor = &ModelSelectorPlugin{} // ModelSelectorPluginFactory is the factory function for the ModelSelector RequestProcessor plugin. // It creates a plugin with an empty pipeline; plugins are wired in by the configuration loader. -func ModelSelectorPluginFactory(name string, _ json.RawMessage, handle plugin.Handle) (plugin.Plugin, error) { +func ModelSelectorPluginFactory(name string, _ *json.Decoder, handle plugin.Handle) (plugin.Plugin, error) { return NewModelSelectorPlugin(ms.NewModelSelectorPipeline(), handle.Datastore()).WithName(name), nil } diff --git a/pkg/framework/plugins/requesthandling/modelselector/plugin_test.go b/pkg/framework/plugins/requesthandling/modelselector/plugin_test.go index 4b1f82e2..46b6c038 100644 --- a/pkg/framework/plugins/requesthandling/modelselector/plugin_test.go +++ b/pkg/framework/plugins/requesthandling/modelselector/plugin_test.go @@ -18,7 +18,6 @@ package modelselector import ( "context" - "encoding/json" "slices" "testing" @@ -80,7 +79,7 @@ func newFakeHandle(modelNames ...string) *fakeHandle { // mustFactory calls ModelSelectorPluginFactory and fails the test on error. func mustFactory(t *testing.T, handle *fakeHandle) *ModelSelectorPlugin { t.Helper() - plug, err := ModelSelectorPluginFactory(ModelSelectorPluginType, json.RawMessage(`{}`), handle) + plug, err := ModelSelectorPluginFactory(ModelSelectorPluginType, nil, handle) if err != nil { t.Fatalf("ModelSelectorPluginFactory failed: %v", err) } diff --git a/pkg/framework/plugins/requesthandling/profilepicker/single/single_profile_picker.go b/pkg/framework/plugins/requesthandling/profilepicker/single/single_profile_picker.go index a80c8bb6..2571ce4e 100644 --- a/pkg/framework/plugins/requesthandling/profilepicker/single/single_profile_picker.go +++ b/pkg/framework/plugins/requesthandling/profilepicker/single/single_profile_picker.go @@ -33,7 +33,7 @@ const ( var _ requesthandling.ProfilePicker = &SingleProfilePicker{} // SingleProfilePickerFactory defines the factory function for SingleProfilePicker. -func SingleProfilePickerFactory(name string, _ json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) { +func SingleProfilePickerFactory(name string, _ *json.Decoder, _ plugin.Handle) (plugin.Plugin, error) { return NewSingleProfilePicker().WithName(name), nil } diff --git a/pkg/framework/plugins/responsehandling/modelnametoheader/plugin.go b/pkg/framework/plugins/responsehandling/modelnametoheader/plugin.go index d8c957f2..8f9ebe62 100644 --- a/pkg/framework/plugins/responsehandling/modelnametoheader/plugin.go +++ b/pkg/framework/plugins/responsehandling/modelnametoheader/plugin.go @@ -41,10 +41,10 @@ type modelNameToHeaderConfig struct { } // PluginFactory creates a new ModelNameToHeaderPlugin. -func PluginFactory(name string, rawParameters json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) { +func PluginFactory(name string, parameters *json.Decoder, _ plugin.Handle) (plugin.Plugin, error) { var cfg modelNameToHeaderConfig - if len(rawParameters) > 0 { - if err := json.Unmarshal(rawParameters, &cfg); err != nil { + if parameters != nil { + if err := parameters.Decode(&cfg); err != nil { return nil, fmt.Errorf("failed to parse the parameters of the '%s' plugin - %w", PluginType, err) } } diff --git a/pkg/framework/plugins/responsehandling/modelnametoheader/plugin_test.go b/pkg/framework/plugins/responsehandling/modelnametoheader/plugin_test.go index f61266a5..ab585a7e 100644 --- a/pkg/framework/plugins/responsehandling/modelnametoheader/plugin_test.go +++ b/pkg/framework/plugins/responsehandling/modelnametoheader/plugin_test.go @@ -21,6 +21,7 @@ import ( "encoding/json" "testing" + "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/requesthandling" "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/plugins/requesthandling/modelselector" ) @@ -106,7 +107,7 @@ func TestPluginFactory_DefaultHeaderName(t *testing.T) { } func TestPluginFactory_CustomHeaderName(t *testing.T) { - params := json.RawMessage(`{"headerName": "X-Custom-Model"}`) + params := plugin.StrictDecoder(json.RawMessage(`{"headerName": "X-Custom-Model"}`)) p, err := PluginFactory("custom", params, nil) if err != nil { t.Fatalf("PluginFactory failed: %v", err) @@ -134,7 +135,7 @@ func TestPluginFactory_CustomHeaderName(t *testing.T) { } func TestPluginFactory_InvalidConfig(t *testing.T) { - params := json.RawMessage(`{invalid`) + params := plugin.StrictDecoder(json.RawMessage(`{invalid`)) _, err := PluginFactory("bad", params, nil) if err == nil { t.Fatal("expected error for invalid JSON config, got nil")