From 735f527c08f80cfb183e250a4bd0fb33285f2ae8 Mon Sep 17 00:00:00 2001 From: Mayur Das Date: Tue, 23 Jun 2026 12:34:41 +0530 Subject: [PATCH] feat(prefix): auto-tune prefix match cap from model max_model_len The approximate-prefix-cache producer caps prefix matching at a fixed maxPrefixTokensToMatch (default 128k). On models with a larger context window that under-matches long shared prefixes, and it's one more knob to set per model. When autoTune is on and maxPrefixTokensToMatch is left at its default, derive the cap from each endpoint's max_model_len instead, falling back to the existing default when it isn't known. An explicitly configured cap is always respected. max_model_len rides the existing /v1/models attribute. Rather than poll that endpoint on a timer for information that's fixed for an endpoint's lifetime, a new models-endpoint-extractor fetches /v1/models once when an endpoint is added and stores the model list (now carrying max_model_len) as an attribute. It registers as the default producer of that attribute and self-wires an endpoint-notification source, so the producer only declares the dependency and reads the value. A failed or slow fetch leaves the attribute unset and the cap falls back to the default. Signed-off-by: Mayur Das --- cmd/epp/runner/runner.go | 3 + .../datalayer/attribute/models/data_types.go | 3 + .../attribute/models/data_types_test.go | 38 ++++ .../datalayer/extractor/models/README.md | 23 ++ .../extractor/models/endpoint_extractor.go | 172 +++++++++++++++ .../models/endpoint_extractor_test.go | 200 ++++++++++++++++++ .../datalayer/extractor/models/extractor.go | 16 ++ .../extractor/models/extractor_test.go | 29 +++ .../datalayer/source/models/datasource.go | 17 +- .../source/models/parsercontract_test.go | 3 +- .../approximateprefix/autotune_test.go | 178 ++++++++++++++++ .../dataproducer/approximateprefix/plugin.go | 48 ++++- .../dataproducer/approximateprefix/types.go | 5 + 13 files changed, 715 insertions(+), 20 deletions(-) create mode 100644 pkg/epp/framework/plugins/datalayer/attribute/models/data_types_test.go create mode 100644 pkg/epp/framework/plugins/datalayer/extractor/models/endpoint_extractor.go create mode 100644 pkg/epp/framework/plugins/datalayer/extractor/models/endpoint_extractor_test.go create mode 100644 pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/autotune_test.go 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.