diff --git a/cmd/epp/runner/runner.go b/cmd/epp/runner/runner.go index 9f210bf66f..fc7f8c7e8c 100644 --- a/cmd/epp/runner/runner.go +++ b/cmd/epp/runner/runner.go @@ -499,6 +499,9 @@ func (r *Runner) registerInTreePlugins() { // data layer models source/extractor fwkplugin.Register(srcmodels.ModelsDataSourceType, srcmodels.ModelDataSourceFactory) fwkplugin.Register(attrmodels.ModelsExtractorType, extmodels.ModelServerExtractorFactory) + // Default producer of the /v1/models attribute: fetches once per endpoint on + // add rather than polling, and self-wires its endpoint-notification source. + fwkplugin.RegisterAsDefaultProducer(extmodels.ModelsEndpointExtractorType, extmodels.ModelEndpointExtractorFactory, attrmodels.ModelsAttributeKey) fwkplugin.Register(prefix.PrefixCacheScorerPluginType, prefix.PrefixCachePluginFactory) fwkplugin.Register(maxscore.MaxScorePickerType, maxscore.MaxScorePickerFactory) diff --git a/pkg/epp/framework/plugins/datalayer/attribute/models/data_types.go b/pkg/epp/framework/plugins/datalayer/attribute/models/data_types.go index 9667349757..5460cb49f4 100644 --- a/pkg/epp/framework/plugins/datalayer/attribute/models/data_types.go +++ b/pkg/epp/framework/plugins/datalayer/attribute/models/data_types.go @@ -37,6 +37,9 @@ type ModelDataCollection []ModelData type ModelData struct { ID string `json:"id"` Parent string `json:"parent,omitempty"` + // MaxModelLen is the model's context window in tokens as reported by the + // model server; zero when the server does not report it. + MaxModelLen int `json:"max_model_len,omitempty"` } // String returns a string representation of the model info diff --git a/pkg/epp/framework/plugins/datalayer/attribute/models/data_types_test.go b/pkg/epp/framework/plugins/datalayer/attribute/models/data_types_test.go new file mode 100644 index 0000000000..de98c24cd0 --- /dev/null +++ b/pkg/epp/framework/plugins/datalayer/attribute/models/data_types_test.go @@ -0,0 +1,38 @@ +/* +Copyright 2026 The Kubernetes 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 models + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +// TestModelDataCollectionClone verifies Clone copies entries (including +// MaxModelLen) and is independent of the original. +func TestModelDataCollectionClone(t *testing.T) { + assert.Nil(t, ModelDataCollection(nil).Clone()) + + orig := ModelDataCollection{{ID: "m1", MaxModelLen: 100}, {ID: "m2", MaxModelLen: 200}} + cloned, ok := orig.Clone().(ModelDataCollection) + assert.True(t, ok) + assert.Equal(t, orig, cloned) + + // Mutating the clone must not affect the original. + cloned[0].MaxModelLen = 999 + assert.Equal(t, 100, orig[0].MaxModelLen) +} diff --git a/pkg/epp/framework/plugins/datalayer/extractor/models/README.md b/pkg/epp/framework/plugins/datalayer/extractor/models/README.md index 69358a070e..a863569881 100644 --- a/pkg/epp/framework/plugins/datalayer/extractor/models/README.md +++ b/pkg/epp/framework/plugins/datalayer/extractor/models/README.md @@ -27,3 +27,26 @@ modelData, ok := attr.(models.ModelDataCollection) ## Configuration No configuration parameters. + +## Once-per-endpoint variant + +**Type:** `models-endpoint-extractor` + +`ModelEndpointExtractor` produces the same `ModelsAttributeKey` attribute but is +driven by endpoint lifecycle events instead of a poll loop. On endpoint add it +fetches `/v1/models` once and stores the parsed `ModelDataCollection`; the model +list is fixed for an endpoint's lifetime, so it is not re-fetched on a timer. A +failed fetch is logged and skipped, leaving the attribute unset so consumers fall +back to their default. + +It is registered as the default producer of `ModelsAttributeKey` and self-wires +to an `endpoint-notification-source` (auto-created when absent), so any consumer +that declares the attribute as a dependency gets it without extra configuration. + +### Configuration + +| Field | Default | Description | +|---|---|---| +| `scheme` | `http` | Scheme used to reach the model server. | +| `path` | `/v1/models` | Path fetched on the model server. | +| `insecureSkipVerify` | `true` | Skip TLS verification when `scheme` is `https`. | diff --git a/pkg/epp/framework/plugins/datalayer/extractor/models/endpoint_extractor.go b/pkg/epp/framework/plugins/datalayer/extractor/models/endpoint_extractor.go new file mode 100644 index 0000000000..d19882a692 --- /dev/null +++ b/pkg/epp/framework/plugins/datalayer/extractor/models/endpoint_extractor.go @@ -0,0 +1,172 @@ +/* +Copyright 2026 The Kubernetes 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 models + +import ( + "context" + "encoding/json" + "fmt" + "time" + + "sigs.k8s.io/controller-runtime/pkg/log" + + "github.com/llm-d/llm-d-router/pkg/common/observability/logging" + fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer" + fwkplugin "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin" + attrmodels "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/models" + srchttp "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/source/http" + srcnotifications "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/source/notifications" +) + +// ModelsEndpointExtractorType identifies the extractor that fetches /v1/models +// once per endpoint, as opposed to the polling models-data-extractor. +const ModelsEndpointExtractorType = "models-endpoint-extractor" + +// fetchTimeout bounds a single /v1/models fetch so a slow or unresponsive model +// server cannot keep the background fetch goroutine alive indefinitely. +const fetchTimeout = 10 * time.Second + +// Model-server connection defaults, matching the polling models source so both +// reach the same endpoint with the same TLS handling. +const ( + defaultModelsScheme = "http" + defaultModelsPath = "/v1/models" + defaultModelsInsecureSkipVerify = true +) + +// Assert the extractor produces an attribute, runs on endpoint lifecycle events, +// and self-wires its source dependency. +var ( + _ fwkplugin.ProducerPlugin = (*ModelEndpointExtractor)(nil) + _ fwkdl.EndpointExtractor = (*ModelEndpointExtractor)(nil) + _ fwkdl.Registrant = (*ModelEndpointExtractor)(nil) +) + +// modelsEndpointExtractorParams configures the model-server endpoint and TLS +// handling. It mirrors the polling models source so operators can point the +// extractor at an https model server. +type modelsEndpointExtractorParams struct { + Scheme string `json:"scheme"` + Path string `json:"path"` + InsecureSkipVerify bool `json:"insecureSkipVerify"` +} + +// ModelEndpointExtractor fetches /v1/models once when an endpoint is added and +// stores the parsed model list (including max_model_len) as an endpoint +// attribute. The model list is fixed for an endpoint's lifetime, so it does not +// re-fetch on a timer the way the polling extractor does. +type ModelEndpointExtractor struct { + typedName fwkplugin.TypedName + dk fwkplugin.DataKey + // fetcher performs the one-shot GET and JSON parse, reusing the HTTP source's + // scheme/TLS handling. Only its Poll method is used; the polling Dispatch loop + // is never driven. + fetcher *srchttp.HTTPDataSource[*ModelResponse] +} + +// NewModelEndpointExtractor builds an extractor that fetches from +// scheme:///path, verifying the server certificate unless insecure is set. +func NewModelEndpointExtractor(name, scheme, path string, insecure bool) (*ModelEndpointExtractor, error) { + fetcher, err := srchttp.NewHTTPDataSource(scheme, path, insecure, ModelsEndpointExtractorType, name, ParseModels) + if err != nil { + return nil, fmt.Errorf("failed to create models fetcher: %w", err) + } + return &ModelEndpointExtractor{ + typedName: fwkplugin.TypedName{Type: ModelsEndpointExtractorType, Name: name}, + dk: attrmodels.ModelsAttributeKey, + fetcher: fetcher, + }, nil +} + +// TypedName returns the plugin type and name. +func (me *ModelEndpointExtractor) TypedName() fwkplugin.TypedName { return me.typedName } + +// Produces declares the /v1/models attribute this extractor populates so the +// framework can wire it to consumers of that attribute. +func (me *ModelEndpointExtractor) Produces() map[fwkplugin.DataKey]any { + return map[fwkplugin.DataKey]any{me.dk: attrmodels.ModelDataCollection{}} +} + +// Extract records the model list for an added endpoint. The fetch runs in the +// background and returns immediately: endpoint add/update is processed serially, +// so a slow or unreachable model server must not stall endpoint provisioning. +// Until the fetch completes the attribute stays unset and the consumer falls +// back to its default. Delete events need no work; the attribute leaves with the +// endpoint. +func (me *ModelEndpointExtractor) Extract(ctx context.Context, event fwkdl.EndpointEvent) error { + // Only an add/update carries model data to extract; deletes carry none. + if event.Type != fwkdl.EventAddOrUpdate { + return nil + } + ep := event.Endpoint + if ep == nil || ep.GetMetadata() == nil { + return nil + } + // The model list is fixed for an endpoint's lifetime, so fetch only once even + // though add/update may fire repeatedly for the same endpoint. + if _, ok := ep.GetAttributes().Get(me.dk.String()); ok { + return nil + } + go me.fetchAndStore(ctx, ep) + return nil +} + +// fetchAndStore fetches /v1/models once, under a bounded timeout, and records the +// model list (including max_model_len) as an endpoint attribute. A failed fetch +// is logged and dropped, leaving the attribute unset so the consumer falls back +// to its default. +func (me *ModelEndpointExtractor) fetchAndStore(ctx context.Context, ep fwkdl.Endpoint) { + ctx, cancel := context.WithTimeout(ctx, fetchTimeout) + defer cancel() + resp, err := me.fetcher.Poll(ctx, ep) + if err != nil { + log.FromContext(ctx).V(logging.DEBUG).Info("failed to fetch /v1/models", + "endpoint", ep.GetMetadata().NamespacedName, "err", err) + return + } + ep.GetAttributes().Put(me.dk.String(), attrmodels.ModelDataCollection(resp.Data)) +} + +// RegisterDependencies binds this extractor to an endpoint-notification-source, +// auto-creating that source when the configuration does not already define one. +// This is what makes the extractor wire itself in as a default source. +func (me *ModelEndpointExtractor) RegisterDependencies(r fwkdl.Registrar) error { + return r.Register(fwkdl.PendingRegistration{ + Owner: me.TypedName(), + SourceType: srcnotifications.EndpointNotificationSourceType, + Extractor: me, + DefaultSource: srcnotifications.NewEndpointDataSource( + srcnotifications.EndpointNotificationSourceType, srcnotifications.EndpointNotificationSourceType), + }) +} + +// ModelEndpointExtractorFactory instantiates the extractor from optional JSON +// parameters, defaulting to an http model server with TLS verification skipped. +func ModelEndpointExtractorFactory(name string, parameters *json.Decoder, _ fwkplugin.Handle) (fwkplugin.Plugin, error) { + params := modelsEndpointExtractorParams{ + Scheme: defaultModelsScheme, + Path: defaultModelsPath, + InsecureSkipVerify: defaultModelsInsecureSkipVerify, + } + // Overlay the defaults with any configured values. + if parameters != nil { + if err := parameters.Decode(¶ms); err != nil { + return nil, err + } + } + return NewModelEndpointExtractor(name, params.Scheme, params.Path, params.InsecureSkipVerify) +} diff --git a/pkg/epp/framework/plugins/datalayer/extractor/models/endpoint_extractor_test.go b/pkg/epp/framework/plugins/datalayer/extractor/models/endpoint_extractor_test.go new file mode 100644 index 0000000000..285fcf95c5 --- /dev/null +++ b/pkg/epp/framework/plugins/datalayer/extractor/models/endpoint_extractor_test.go @@ -0,0 +1,200 @@ +/* +Copyright 2026 The Kubernetes 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 models + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + k8stypes "k8s.io/apimachinery/pkg/types" + + fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer" + fwkplugin "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin" + attrmodels "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/models" +) + +// modelsServer starts a stub /v1/models server returning body with status code, +// and returns an endpoint whose metrics host points at it. +func modelsServer(t *testing.T, status int, body string) (fwkdl.Endpoint, func()) { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(status) + _, _ = w.Write([]byte(body)) + })) + // MetricsHost is host:port without scheme; the fetcher prepends the scheme. + ep := fwkdl.NewEndpoint(&fwkdl.EndpointMetadata{ + NamespacedName: k8stypes.NamespacedName{Name: "pod1"}, + MetricsHost: strings.TrimPrefix(srv.URL, "http://"), + }, fwkdl.NewMetrics()) + return ep, srv.Close +} + +// readModels returns the /v1/models attribute on ep, or nil when unset. +func readModels(ep fwkdl.Endpoint) attrmodels.ModelDataCollection { + val, ok := ep.GetAttributes().Get(attrmodels.ModelsAttributeKey.String()) + if !ok { + return nil + } + models, _ := val.(attrmodels.ModelDataCollection) + return models +} + +// newTestExtractor builds an http extractor pointing at an insecure model server. +func newTestExtractor(t *testing.T) *ModelEndpointExtractor { + t.Helper() + me, err := NewModelEndpointExtractor("test", "http", "/v1/models", true) + require.NoError(t, err) + return me +} + +// TestFetchAndStoreRecordsModels verifies a successful fetch records every model +// (and its max_model_len) on the endpoint attribute. +func TestFetchAndStoreRecordsModels(t *testing.T) { + ep, closeSrv := modelsServer(t, http.StatusOK, + `{"object":"list","data":[{"id":"m1","max_model_len":131072},{"id":"m2","max_model_len":65536}]}`) + defer closeSrv() + + newTestExtractor(t).fetchAndStore(context.Background(), ep) + + models := readModels(ep) + require.Len(t, models, 2) + assert.Equal(t, "m1", models[0].ID) + assert.Equal(t, 131072, models[0].MaxModelLen) + assert.Equal(t, "m2", models[1].ID) + assert.Equal(t, 65536, models[1].MaxModelLen) +} + +// TestFetchAndStoreSwallowsError verifies a failing model server leaves the +// attribute unset rather than recording anything, so the consumer falls back. +func TestFetchAndStoreSwallowsError(t *testing.T) { + ep, closeSrv := modelsServer(t, http.StatusInternalServerError, "") + defer closeSrv() + + newTestExtractor(t).fetchAndStore(context.Background(), ep) + assert.Nil(t, readModels(ep)) +} + +// TestExtractFetchesOnAdd verifies an add event triggers a background fetch that +// eventually records the model list. +func TestExtractFetchesOnAdd(t *testing.T) { + ep, closeSrv := modelsServer(t, http.StatusOK, + `{"object":"list","data":[{"id":"m","max_model_len":131072}]}`) + defer closeSrv() + + require.NoError(t, newTestExtractor(t).Extract(context.Background(), + fwkdl.EndpointEvent{Type: fwkdl.EventAddOrUpdate, Endpoint: ep})) + + // The fetch runs in the background, so wait for the attribute to appear. + assert.Eventually(t, func() bool { + m := readModels(ep) + return len(m) == 1 && m[0].MaxModelLen == 131072 + }, time.Second, 5*time.Millisecond) +} + +// TestExtractSkipsWhenAlreadyPresent verifies a second add/update does not +// re-fetch once the attribute is set. The server would error, proving no fetch. +func TestExtractSkipsWhenAlreadyPresent(t *testing.T) { + ep, closeSrv := modelsServer(t, http.StatusInternalServerError, "") + defer closeSrv() + existing := attrmodels.ModelDataCollection{{ID: "m", MaxModelLen: 42}} + ep.GetAttributes().Put(attrmodels.ModelsAttributeKey.String(), existing) + + require.NoError(t, newTestExtractor(t).Extract(context.Background(), + fwkdl.EndpointEvent{Type: fwkdl.EventAddOrUpdate, Endpoint: ep})) + + // No goroutine is spawned, so the pre-existing value must remain untouched. + assert.Never(t, func() bool { + m := readModels(ep) + return len(m) != 1 || m[0].MaxModelLen != 42 + }, 100*time.Millisecond, 10*time.Millisecond) +} + +// TestExtractIgnoresDelete verifies a delete event neither fetches nor writes the +// attribute; the attribute leaves with the endpoint. +func TestExtractIgnoresDelete(t *testing.T) { + // A failing server proves no fetch is attempted on delete. + ep, closeSrv := modelsServer(t, http.StatusInternalServerError, "") + defer closeSrv() + + require.NoError(t, newTestExtractor(t).Extract(context.Background(), + fwkdl.EndpointEvent{Type: fwkdl.EventDelete, Endpoint: ep})) + assert.Nil(t, readModels(ep)) +} + +// TestExtractNilEndpoint verifies missing endpoint or metadata is a no-op and +// does not panic. +func TestExtractNilEndpoint(t *testing.T) { + me := newTestExtractor(t) + assert.NoError(t, me.Extract(context.Background(), + fwkdl.EndpointEvent{Type: fwkdl.EventAddOrUpdate, Endpoint: nil})) + assert.NoError(t, me.Extract(context.Background(), + fwkdl.EndpointEvent{Type: fwkdl.EventAddOrUpdate, Endpoint: fwkdl.NewEndpoint(nil, nil)})) +} + +// fakeRegistrar captures the registration a plugin submits. +type fakeRegistrar struct { + reg fwkdl.PendingRegistration +} + +func (f *fakeRegistrar) Register(reg fwkdl.PendingRegistration) error { + f.reg = reg + return nil +} + +// TestRegisterDependencies verifies the extractor self-wires to an +// endpoint-notification-source with a default source to auto-create. +func TestRegisterDependencies(t *testing.T) { + me := newTestExtractor(t) + r := &fakeRegistrar{} + require.NoError(t, me.RegisterDependencies(r)) + assert.Equal(t, "endpoint-notification-source", r.reg.SourceType) + assert.Equal(t, fwkplugin.Plugin(me), r.reg.Extractor) + assert.NotNil(t, r.reg.DefaultSource) +} + +// TestProduces verifies the extractor declares the /v1/models attribute so +// consumers of that key are wired to it. +func TestProduces(t *testing.T) { + produced := newTestExtractor(t).Produces() + _, ok := produced[attrmodels.ModelsAttributeKey] + assert.True(t, ok, "must produce the models attribute") +} + +// TestFactory verifies the factory builds with defaults, honours custom +// parameters, and rejects an unsupported scheme. +func TestFactory(t *testing.T) { + // Defaults. + p, err := ModelEndpointExtractorFactory("models", nil, nil) + require.NoError(t, err) + assert.Equal(t, ModelsEndpointExtractorType, p.TypedName().Type) + + // Custom https params are accepted. + custom := fwkplugin.StrictDecoder([]byte(`{"scheme":"https","path":"/models","insecureSkipVerify":false}`)) + _, err = ModelEndpointExtractorFactory("models", custom, nil) + require.NoError(t, err) + + // An unsupported scheme is rejected. + bad := fwkplugin.StrictDecoder([]byte(`{"scheme":"ftp"}`)) + _, err = ModelEndpointExtractorFactory("models", bad, nil) + assert.Error(t, err) +} diff --git a/pkg/epp/framework/plugins/datalayer/extractor/models/extractor.go b/pkg/epp/framework/plugins/datalayer/extractor/models/extractor.go index c27b0c0aab..ce5693e53d 100644 --- a/pkg/epp/framework/plugins/datalayer/extractor/models/extractor.go +++ b/pkg/epp/framework/plugins/datalayer/extractor/models/extractor.go @@ -3,6 +3,8 @@ package models import ( "context" "encoding/json" + "fmt" + "io" fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer" fwkplugin "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin" @@ -21,6 +23,20 @@ type ModelResponse struct { Data []attrmodels.ModelData `json:"data"` } +// ParseModels decodes a /v1/models response body into a ModelResponse. It is the +// shared parser for both the polling source and the once-per-endpoint extractor. +func ParseModels(data io.Reader) (*ModelResponse, error) { + body, err := io.ReadAll(data) + if err != nil { + return nil, fmt.Errorf("failed to read response body: %v", err) + } + var modelsResponse ModelResponse + if err := json.Unmarshal(body, &modelsResponse); err != nil { + return nil, err + } + return &modelsResponse, nil +} + // ModelExtractor implements the models extraction. type ModelExtractor struct { typedName fwkplugin.TypedName diff --git a/pkg/epp/framework/plugins/datalayer/extractor/models/extractor_test.go b/pkg/epp/framework/plugins/datalayer/extractor/models/extractor_test.go index 61b0e95f5a..204ba5c7f0 100644 --- a/pkg/epp/framework/plugins/datalayer/extractor/models/extractor_test.go +++ b/pkg/epp/framework/plugins/datalayer/extractor/models/extractor_test.go @@ -106,3 +106,32 @@ func TestExtractorExtract(t *testing.T) { }) } } + +// TestExtractorCapturesMaxModelLen verifies the extractor surfaces max_model_len +// from the /v1/models payload onto the endpoint attribute. +func TestExtractorCapturesMaxModelLen(t *testing.T) { + ctx := context.Background() + extPlugin, err := ModelServerExtractorFactory("test-extractor", nil, nil) + if err != nil { + t.Fatalf("failed to create extractor: %v", err) + } + extractor := extPlugin.(*ModelExtractor) + ep := fwkdl.NewEndpoint(nil, nil) + + resp := &ModelResponse{Data: []attrmodels.ModelData{{ID: "m", MaxModelLen: 131072}}} + if err := extractor.Extract(ctx, fwkdl.PollInput[*ModelResponse]{Payload: resp, Endpoint: ep}); err != nil { + t.Fatalf("unexpected error: %v", err) + } + + val, ok := ep.GetAttributes().Get(attrmodels.ModelsAttributeKey.String()) + if !ok { + t.Fatal("expected models attribute to be set") + } + models, ok := val.(attrmodels.ModelDataCollection) + if !ok { + t.Fatalf("unexpected attribute type %T", val) + } + if len(models) != 1 || models[0].MaxModelLen != 131072 { + t.Errorf("expected max_model_len 131072, got %+v", models) + } +} diff --git a/pkg/epp/framework/plugins/datalayer/source/models/datasource.go b/pkg/epp/framework/plugins/datalayer/source/models/datasource.go index 0109904838..5917e81ce9 100644 --- a/pkg/epp/framework/plugins/datalayer/source/models/datasource.go +++ b/pkg/epp/framework/plugins/datalayer/source/models/datasource.go @@ -3,7 +3,6 @@ package models import ( "encoding/json" "fmt" - "io" "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin" extmodels "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/extractor/models" @@ -35,7 +34,7 @@ type modelsDatasourceParams struct { // Use this function directly in tests to bypass JSON parameter marshaling. func NewHTTPModelsDataSource(scheme, path, name string) (*http.HTTPDataSource[*extmodels.ModelResponse], error) { return http.NewHTTPDataSource(scheme, path, defaultModelsInsecureSkipVerify, - ModelsDataSourceType, name, parseModels) + ModelsDataSourceType, name, extmodels.ParseModels) } // ModelDataSourceFactory is a factory function used to instantiate data layer's @@ -52,7 +51,7 @@ func ModelDataSourceFactory(name string, parameters *json.Decoder, _ plugin.Hand } ds, err := http.NewHTTPDataSource(cfg.Scheme, cfg.Path, cfg.InsecureSkipVerify, - ModelsDataSourceType, name, parseModels) + ModelsDataSourceType, name, extmodels.ParseModels) if err != nil { return nil, fmt.Errorf("failed to create HTTP data source: %w", err) } @@ -66,15 +65,3 @@ func defaultDataSourceConfigParams() *modelsDatasourceParams { InsecureSkipVerify: defaultModelsInsecureSkipVerify, } } - -func parseModels(data io.Reader) (*extmodels.ModelResponse, error) { - body, err := io.ReadAll(data) - if err != nil { - return nil, fmt.Errorf("failed to read response body: %v", err) - } - var modelsResponse extmodels.ModelResponse - if err := json.Unmarshal(body, &modelsResponse); err != nil { - return nil, err - } - return &modelsResponse, nil -} diff --git a/pkg/epp/framework/plugins/datalayer/source/models/parsercontract_test.go b/pkg/epp/framework/plugins/datalayer/source/models/parsercontract_test.go index bc89e91c8d..c4879c1618 100644 --- a/pkg/epp/framework/plugins/datalayer/source/models/parsercontract_test.go +++ b/pkg/epp/framework/plugins/datalayer/source/models/parsercontract_test.go @@ -3,11 +3,12 @@ package models import ( "testing" + extmodels "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/extractor/models" "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/source/http/httptest" ) func TestParseModels_Contract(t *testing.T) { - httptest.ParserContract(t, parseModels, + httptest.ParserContract(t, extmodels.ParseModels, []byte(`{"object":"list","data":[]}`), []byte(`{"object":"list","data":[{"id":"qwen2.5-7b"}]}`), []byte(`{"object":"list","data":[{"id":"a","parent":"a"},{"id":"b"}]}`), diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/autotune_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/autotune_test.go new file mode 100644 index 0000000000..c89f9e633b --- /dev/null +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/autotune_test.go @@ -0,0 +1,178 @@ +/* +Copyright 2026 The Kubernetes 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 approximateprefix + +import ( + "context" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + k8stypes "k8s.io/apimachinery/pkg/types" + + fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer" + "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin" + fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling" + attrmodels "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/models" + attrprefix "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/prefix" +) + +// modelLenEndpoint builds an endpoint whose /v1/models attribute reports the +// given max_model_len values, one per model entry. With no values, the attribute +// is left unset. +func modelLenEndpoint(maxModelLens ...int) fwksched.Endpoint { + attrs := fwkdl.NewAttributes() + if maxModelLens != nil { + models := make(attrmodels.ModelDataCollection, len(maxModelLens)) + for i, ml := range maxModelLens { + models[i] = attrmodels.ModelData{ID: "m", MaxModelLen: ml} + } + attrs.Put(attrmodels.ModelsAttributeKey.String(), models) + } + return fwksched.NewEndpoint(&fwkdl.EndpointMetadata{NamespacedName: k8stypes.NamespacedName{Name: "pod1"}}, fwkdl.NewMetrics(), attrs) +} + +// TestProduceAutoTunesFromModelLen verifies the prefix cap is derived from the +// endpoint's max_model_len attribute when auto-tuning, and falls back otherwise. +func TestProduceAutoTunesFromModelLen(t *testing.T) { + disableMinBlockSizeClamp(t) + + // A 20-block prompt at block size 1. + tokens := make([]uint32, 20) + for i := range tokens { + tokens[i] = uint32(i + 1) + } + + produceBlocks := func(t *testing.T, p *dataProducer, ep fwksched.Endpoint) int { + t.Helper() + req := &fwksched.InferenceRequest{RequestID: uuid.NewString(), TargetModel: "m", Body: tokenizedBody(tokens)} + assert.NoError(t, p.Produce(context.Background(), req, []fwksched.Endpoint{ep})) + state, err := plugin.ReadPluginStateKey[*SchedulingContextState]( + p.PluginState(), req.RequestID, plugin.StateKey(ApproxPrefixCachePluginType)) + assert.NoError(t, err) + return len(state.PerPromptHashes[0]) + } + + autoTuneConfig := config{ + AutoTune: true, + BlockSizeTokens: 1, + MaxPrefixBlocksToMatch: defaultMaxPrefixBlocks, + MaxPrefixTokensToMatch: defaultMaxPrefixTokens, // left at default -> auto-tune on + LRUCapacityPerServer: defaultLRUCapacityPerServer, + } + + t.Run("caps at max_model_len when known", func(t *testing.T) { + p, err := newDataProducer(context.Background(), ApproxPrefixCachePluginType, autoTuneConfig, testHandle()) + assert.NoError(t, err) + assert.True(t, p.config.autoTuneMaxPrefixTokens) + // max_model_len = 10 -> cap = 10/blockSize(1) = 10 blocks. + assert.Equal(t, 10, produceBlocks(t, p, modelLenEndpoint(10))) + }) + + t.Run("uses largest across models", func(t *testing.T) { + p, err := newDataProducer(context.Background(), ApproxPrefixCachePluginType, autoTuneConfig, testHandle()) + assert.NoError(t, err) + // Largest wins regardless of position (first entry here). + assert.Equal(t, 15, produceBlocks(t, p, modelLenEndpoint(15, 8, 12))) + }) + + t.Run("falls back on negative max_model_len", func(t *testing.T) { + p, err := newDataProducer(context.Background(), ApproxPrefixCachePluginType, autoTuneConfig, testHandle()) + assert.NoError(t, err) + assert.Equal(t, 20, produceBlocks(t, p, modelLenEndpoint(-5))) + }) + + t.Run("falls back on empty models collection", func(t *testing.T) { + p, err := newDataProducer(context.Background(), ApproxPrefixCachePluginType, autoTuneConfig, testHandle()) + assert.NoError(t, err) + // Attribute present but no model entries. + attrs := fwkdl.NewAttributes() + attrs.Put(attrmodels.ModelsAttributeKey.String(), attrmodels.ModelDataCollection{}) + ep := fwksched.NewEndpoint(&fwkdl.EndpointMetadata{NamespacedName: k8stypes.NamespacedName{Name: "pod1"}}, fwkdl.NewMetrics(), attrs) + assert.Equal(t, 20, produceBlocks(t, p, ep)) + }) + + t.Run("falls back to default when attribute absent", func(t *testing.T) { + p, err := newDataProducer(context.Background(), ApproxPrefixCachePluginType, autoTuneConfig, testHandle()) + assert.NoError(t, err) + // No models attribute -> default cap (131072) far exceeds 20 blocks. + assert.Equal(t, 20, produceBlocks(t, p, modelLenEndpoint())) + }) + + t.Run("falls back when max_model_len is zero", func(t *testing.T) { + p, err := newDataProducer(context.Background(), ApproxPrefixCachePluginType, autoTuneConfig, testHandle()) + assert.NoError(t, err) + // Attribute present but server did not report the field. + assert.Equal(t, 20, produceBlocks(t, p, modelLenEndpoint(0))) + }) + + t.Run("explicit cap is not auto-tuned", func(t *testing.T) { + explicit := autoTuneConfig + explicit.MaxPrefixTokensToMatch = 5 // operator-pinned, non-default + p, err := newDataProducer(context.Background(), ApproxPrefixCachePluginType, explicit, testHandle()) + assert.NoError(t, err) + assert.False(t, p.config.autoTuneMaxPrefixTokens) + // Even with a known model length, the explicit cap (5) wins. + assert.Equal(t, 5, produceBlocks(t, p, modelLenEndpoint(10))) + }) + + t.Run("empty pods uses default cap", func(t *testing.T) { + p, err := newDataProducer(context.Background(), ApproxPrefixCachePluginType, autoTuneConfig, testHandle()) + assert.NoError(t, err) + + req := &fwksched.InferenceRequest{RequestID: uuid.NewString(), TargetModel: "m", Body: tokenizedBody(tokens)} + assert.NoError(t, p.Produce(context.Background(), req, []fwksched.Endpoint{})) + state, err := plugin.ReadPluginStateKey[*SchedulingContextState]( + p.PluginState(), req.RequestID, plugin.StateKey(ApproxPrefixCachePluginType)) + assert.NoError(t, err) + assert.Equal(t, 20, len(state.PerPromptHashes[0])) + }) +} + +// TestConsumesDeclaresModelsDependency verifies the producer requires the +// /v1/models attribute when auto-tuning, so its default once-per-endpoint +// producer is wired in, and omits the dependency when not auto-tuning. +func TestConsumesDeclaresModelsDependency(t *testing.T) { + // Auto-tuning on (default config): the attribute is a required dependency. + p, err := newDataProducer(context.Background(), ApproxPrefixCachePluginType, defaultConfig, testHandle()) + assert.NoError(t, err) + assert.True(t, p.config.autoTuneMaxPrefixTokens) + _, ok := p.Consumes().Required[attrmodels.ModelsAttributeKey] + assert.True(t, ok, "models attribute must be a required dependency when auto-tuning") + + // An explicit cap disables auto-tuning, so the attribute is not requested. + explicit := defaultConfig + explicit.MaxPrefixTokensToMatch = 5 + p, err = newDataProducer(context.Background(), ApproxPrefixCachePluginType, explicit, testHandle()) + assert.NoError(t, err) + assert.False(t, p.config.autoTuneMaxPrefixTokens) + _, ok = p.Consumes().Required[attrmodels.ModelsAttributeKey] + assert.False(t, ok, "models attribute must not be requested without auto-tuning") +} + +// TestMaxModelLenHelper covers the attribute reader's non-collection and +// missing-attribute paths. +func TestMaxModelLenHelper(t *testing.T) { + // Missing attribute -> 0. + assert.Equal(t, 0, maxModelLen(modelLenEndpoint())) + + // Wrong type stored under the key -> 0. + attrs := fwkdl.NewAttributes() + attrs.Put(attrmodels.ModelsAttributeKey.String(), attrprefix.NewPrefixCacheMatchInfo(0, 0, 0)) + ep := fwksched.NewEndpoint(&fwkdl.EndpointMetadata{NamespacedName: k8stypes.NamespacedName{Name: "pod1"}}, fwkdl.NewMetrics(), attrs) + assert.Equal(t, 0, maxModelLen(ep)) +} diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin.go index 5b7f59694c..90d27e8262 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin.go @@ -30,6 +30,7 @@ import ( "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin" "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requestcontrol" fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling" + attrmodels "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/models" attrprefix "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/prefix" approxprefixconstants "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/constants" tokenproducer "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer" @@ -76,11 +77,17 @@ func (p *dataProducer) Produces() map[plugin.DataKey]any { // Consumes declares the TokenizedPrompt dependency so the data-layer DAG orders // the token-producer before this producer runs and auto-creates one when none -// is configured. +// is configured. When auto-tuning the cap, it also requires the /v1/models +// attribute so the default once-per-endpoint models producer is wired in; the +// attribute is not requested otherwise since Produce never reads it then. func (p *dataProducer) Consumes() plugin.DataDependencies { - return plugin.DataDependencies{ + deps := plugin.DataDependencies{ Required: map[plugin.DataKey]any{tokenproducer.TokenizedPromptDataKey: fwksched.TokenizedPrompt{}}, } + if p.config.autoTuneMaxPrefixTokens { + deps.Required[attrmodels.ModelsAttributeKey] = attrmodels.ModelDataCollection{} + } + return deps } // newDataProducer returns a new DataProducer plugin. @@ -118,6 +125,10 @@ func newDataProducer(ctx context.Context, name string, config config, handle plu indexer := newIndexer(ctx, config.LRUCapacityPerServer, name, ApproxPrefixCachePluginType) + // Auto-tune the cap from max_model_len only when AutoTune is on and the + // operator left MaxPrefixTokensToMatch at its default (no explicit cap). + config.autoTuneMaxPrefixTokens = config.AutoTune && config.MaxPrefixTokensToMatch == defaultMaxPrefixTokens + p := &dataProducer{ typedName: plugin.TypedName{ Type: ApproxPrefixCachePluginType, @@ -175,9 +186,18 @@ func (p *dataProducer) PluginState() *plugin.PluginState { // Produce is called by the director before scheduling requests. func (p *dataProducer) Produce(ctx context.Context, request *fwksched.InferenceRequest, pods []fwksched.Endpoint) error { blockSize := p.GetBlockSize(pods) + maxTokens := p.config.MaxPrefixTokensToMatch + // Raise the cap to the model's full context window when max_model_len is known. + // Candidate pods serve the same model, so the first pod's value is + // representative; fall back to the default when it isn't reported. + if p.config.autoTuneMaxPrefixTokens && len(pods) > 0 { + if modelLen := maxModelLen(pods[0]); modelLen > 0 { + maxTokens = modelLen + } + } maxBlocks := p.config.MaxPrefixBlocksToMatch - if p.config.MaxPrefixTokensToMatch > 0 && blockSize > 0 { - maxBlocks = p.config.MaxPrefixTokensToMatch / blockSize + if maxTokens > 0 && blockSize > 0 { + maxBlocks = maxTokens / blockSize } perPromptHashes := getBlockHashes(ctx, request, blockSize, maxBlocks) @@ -307,6 +327,26 @@ func (p *dataProducer) GetBlockSize(endpoints []fwksched.Endpoint) int { return blockSize } +// maxModelLen returns the largest max_model_len reported for ep via the +// /v1/models attribute, or 0 when absent or non-positive. +func maxModelLen(ep fwksched.Endpoint) int { + val, ok := ep.Get(attrmodels.ModelsAttributeKey.String()) + if !ok { + return 0 + } + models, ok := val.(attrmodels.ModelDataCollection) + if !ok { + return 0 + } + maxLen := 0 + for i := range models { + if models[i].MaxModelLen > maxLen { + maxLen = models[i].MaxModelLen + } + } + return maxLen +} + // ApproxPrefixCacheFactory is the factory function for the prefix cache data producer plugin. func ApproxPrefixCacheFactory(name string, rawParameters *json.Decoder, handle plugin.Handle) (plugin.Plugin, error) { parameters := defaultConfig diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/types.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/types.go index a9109bef5e..18afea1b3b 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/types.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/types.go @@ -140,6 +140,11 @@ type config struct { MaxPrefixTokensToMatch int `json:"maxPrefixTokensToMatch"` // Max capacity size of the LRU indexer in number of entries per server (pod). LRUCapacityPerServer int `json:"lruCapacityPerServer"` + + // autoTuneMaxPrefixTokens, set internally (never from JSON), derives the cap + // per endpoint from the model's max_model_len. Enabled when AutoTune is on and + // MaxPrefixTokensToMatch was left at its default (no operator-pinned cap). + autoTuneMaxPrefixTokens bool } // defaultConfig provides sensible defaults for the prefix cache plugins.