diff --git a/pkg/epp/datalayer/data_graph_test.go b/pkg/epp/datalayer/data_graph_test.go index d1f437f9ff..cef9f80b23 100644 --- a/pkg/epp/datalayer/data_graph_test.go +++ b/pkg/epp/datalayer/data_graph_test.go @@ -35,7 +35,7 @@ import ( "github.com/llm-d/llm-d-router/pkg/epp/util" ) -const mockProducedDataKey = "mockProducedData" +var mockProducedDataKey = fwkplugin.NewDataKey("mockProducedData", "mock") type mockDataProducerP struct { name string diff --git a/pkg/epp/framework/interface/datalayer/attributemap.go b/pkg/epp/framework/interface/datalayer/attributemap.go index abb61d7191..1794149e3e 100644 --- a/pkg/epp/framework/interface/datalayer/attributemap.go +++ b/pkg/epp/framework/interface/datalayer/attributemap.go @@ -18,6 +18,8 @@ package datalayer import ( "sync" + + fwkplugin "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin" ) // Cloneable types support cloning of the value. @@ -41,16 +43,21 @@ func (d *DynamicAttribute) Clone() Cloneable { // AttributeMap is used to store flexible metadata or traits // across different aspects of an inference server. // Stored values must be Cloneable. +// +// Keys are DataKey values rather than strings so that a plugin can only reach +// an attribute through a key it holds -- the same value it names in Produces() +// or Consumes(). This removes the raw-string escape hatch by which a plugin +// could read or write an attribute unrelated to its declaration. type AttributeMap interface { - Put(string, Cloneable) - Get(string) (Cloneable, bool) - Keys() []string + Put(fwkplugin.DataKey, Cloneable) + Get(fwkplugin.DataKey) (Cloneable, bool) + Keys() []fwkplugin.DataKey Clone() AttributeMap } // Attributes provides a goroutine-safe implementation of AttributeMap. type Attributes struct { - data sync.Map // key: attribute name (string), value: attribute value (opaque, Cloneable) + data sync.Map // key: DataKey, value: attribute value (opaque, Cloneable) } // NewAttributes returns a new instance of Attributes. @@ -59,14 +66,14 @@ func NewAttributes() AttributeMap { } // Put adds or updates an attribute in the map. -func (a *Attributes) Put(key string, value Cloneable) { +func (a *Attributes) Put(key fwkplugin.DataKey, value Cloneable) { if value != nil { a.data.Store(key, value) // TODO: Clone into map to ensure isolation } } // Get retrieves an attribute by key, returning a cloned copy (or resolving it dynamically). -func (a *Attributes) Get(key string) (Cloneable, bool) { +func (a *Attributes) Get(key fwkplugin.DataKey) (Cloneable, bool) { val, ok := a.data.Load(key) if !ok { return nil, false @@ -87,11 +94,11 @@ func (a *Attributes) Get(key string) (Cloneable, bool) { } // Keys returns all keys in the attribute map. -func (a *Attributes) Keys() []string { - var keys []string +func (a *Attributes) Keys() []fwkplugin.DataKey { + var keys []fwkplugin.DataKey a.data.Range(func(key, _ any) bool { - if sk, ok := key.(string); ok { - keys = append(keys, sk) + if dk, ok := key.(fwkplugin.DataKey); ok { + keys = append(keys, dk) } return true }) @@ -102,9 +109,9 @@ func (a *Attributes) Keys() []string { func (a *Attributes) Clone() AttributeMap { clone := NewAttributes() a.data.Range(func(key, value any) bool { - if sk, ok := key.(string); ok { + if dk, ok := key.(fwkplugin.DataKey); ok { if v, ok := value.(Cloneable); ok { - clone.Put(sk, v) + clone.Put(dk, v) } } return true @@ -114,7 +121,7 @@ func (a *Attributes) Clone() AttributeMap { // ReadAttribute retrieves attribute with the given key from AttributeMap and asserts it to type T. // Second return value is 'false' if the key is not found or the type assertion fails. -func ReadAttribute[T Cloneable](attributeMap AttributeMap, key string) (T, bool) { +func ReadAttribute[T Cloneable](attributeMap AttributeMap, key fwkplugin.DataKey) (T, bool) { var zero T raw, ok := attributeMap.Get(key) diff --git a/pkg/epp/framework/interface/datalayer/attributemap_test.go b/pkg/epp/framework/interface/datalayer/attributemap_test.go index 6e96f2530d..d7f734391d 100644 --- a/pkg/epp/framework/interface/datalayer/attributemap_test.go +++ b/pkg/epp/framework/interface/datalayer/attributemap_test.go @@ -21,6 +21,13 @@ import ( "github.com/google/go-cmp/cmp" "github.com/stretchr/testify/assert" + + fwkplugin "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin" +) + +var ( + keyA = fwkplugin.NewDataKey("a", "test-producer") + keyB = fwkplugin.NewDataKey("b", "test-producer") ) type dummy struct { @@ -42,9 +49,9 @@ func (d *anotherDummy) Clone() Cloneable { func TestExpectPutThenGetToMatch(t *testing.T) { attrs := NewAttributes() original := &dummy{"foo"} - attrs.Put("a", original) + attrs.Put(keyA, original) - got, ok := attrs.Get("a") + got, ok := attrs.Get(keyA) assert.True(t, ok, "expected key to exist") assert.NotSame(t, original, got, "expected Get to return a clone, not original") @@ -52,28 +59,31 @@ func TestExpectPutThenGetToMatch(t *testing.T) { assert.True(t, ok, "expected value to be of type *dummy") assert.Equal(t, "foo", dv.Text) - _, ok = attrs.Get("b") + _, ok = attrs.Get(keyB) assert.False(t, ok, "expected key not to exist") } func TestExpectKeysToMatchAdded(t *testing.T) { attrs := NewAttributes() - attrs.Put("x", &dummy{"1"}) - attrs.Put("y", &dummy{"2"}) + keyX := fwkplugin.NewDataKey("x", "test-producer") + keyY := fwkplugin.NewDataKey("y", "test-producer") + attrs.Put(keyX, &dummy{"1"}) + attrs.Put(keyY, &dummy{"2"}) keys := attrs.Keys() assert.Len(t, keys, 2) - assert.ElementsMatch(t, keys, []string{"x", "y"}) + assert.ElementsMatch(t, keys, []fwkplugin.DataKey{keyX, keyY}) } func TestCloneReturnsCopy(t *testing.T) { original := NewAttributes() - original.Put("k", &dummy{"value"}) + keyK := fwkplugin.NewDataKey("k", "test-producer") + original.Put(keyK, &dummy{"value"}) cloned := original.Clone() - kOrig, _ := original.Get("k") - kClone, _ := cloned.Get("k") + kOrig, _ := original.Get(keyK) + kClone, _ := cloned.Get(keyK) assert.NotSame(t, kOrig, kClone, "expected cloned value to be a different instance") if diff := cmp.Diff(kOrig, kClone); diff != "" { @@ -85,24 +95,38 @@ func TestReadAttribute(t *testing.T) { // successful retrieval attrs := NewAttributes() original := &dummy{"foo"} - attrs.Put("a", original) + attrs.Put(keyA, original) - got, ok := ReadAttribute[*dummy](attrs, "a") + got, ok := ReadAttribute[*dummy](attrs, keyA) assert.True(t, ok, "expected key to exist and value to be of type *dummy") assert.NotSame(t, original, got, "expected Get to return a clone, not original") assert.Equal(t, "foo", got.Text) // missing key - _, ok = ReadAttribute[*dummy](attrs, "b") + _, ok = ReadAttribute[*dummy](attrs, keyB) assert.False(t, ok, "expected key not to exist") // type mismatch - other, ok := ReadAttribute[*anotherDummy](attrs, "a") + other, ok := ReadAttribute[*anotherDummy](attrs, keyA) assert.False(t, ok, "expected type mismatch") assert.Nil(t, other) // zero value of pointer is nil } +func TestKeysWithDifferentProducersAreDistinct(t *testing.T) { + attrs := NewAttributes() + attrs.Put(keyA, &dummy{"foo"}) + + other := keyA.WithNonEmptyProducerName("other-producer") + _, ok := attrs.Get(other) + assert.False(t, ok, "expected key with a different producer name to be a distinct entry") + + attrs.Put(other, &dummy{"bar"}) + got, ok := attrs.Get(keyA) + assert.True(t, ok) + assert.Equal(t, "foo", got.(*dummy).Text, "expected the original entry to be untouched") +} + func TestDynamicAttribute(t *testing.T) { attrs := NewAttributes() @@ -113,10 +137,11 @@ func TestDynamicAttribute(t *testing.T) { }, } - attrs.Put("dynamic_key", dynamic) + dynamicKey := fwkplugin.NewDataKey("dynamic", "test-producer") + attrs.Put(dynamicKey, dynamic) // First read - got, ok := attrs.Get("dynamic_key") + got, ok := attrs.Get(dynamicKey) assert.True(t, ok) assert.Equal(t, "live", got.(*dummy).Text) @@ -124,7 +149,7 @@ func TestDynamicAttribute(t *testing.T) { val.Text = "changed" // Second read should reflect change - got2, ok := attrs.Get("dynamic_key") + got2, ok := attrs.Get(dynamicKey) assert.True(t, ok) assert.Equal(t, "changed", got2.(*dummy).Text) } diff --git a/pkg/epp/framework/interface/scheduling/types.go b/pkg/epp/framework/interface/scheduling/types.go index 056c8e9029..d6379c269d 100644 --- a/pkg/epp/framework/interface/scheduling/types.go +++ b/pkg/epp/framework/interface/scheduling/types.go @@ -23,6 +23,7 @@ import ( "sync" 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" fwkrh "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requesthandling" ) @@ -79,9 +80,9 @@ type Endpoint interface { GetMetadata() *fwkdl.EndpointMetadata GetMetrics() *fwkdl.Metrics String() string - Get(string) (fwkdl.Cloneable, bool) - Put(string, fwkdl.Cloneable) - Keys() []string + Get(fwkplugin.DataKey) (fwkdl.Cloneable, bool) + Put(fwkplugin.DataKey, fwkdl.Cloneable) + Keys() []fwkplugin.DataKey Clone() fwkdl.AttributeMap } diff --git a/pkg/epp/framework/interface/scheduling/types_test.go b/pkg/epp/framework/interface/scheduling/types_test.go index 49d4e44e23..9b2e149269 100644 --- a/pkg/epp/framework/interface/scheduling/types_test.go +++ b/pkg/epp/framework/interface/scheduling/types_test.go @@ -23,9 +23,16 @@ import ( "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" fwkrh "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requesthandling" ) +var ( + testKeyK = fwkplugin.NewDataKey("k", "test-producer") + testKeyK1 = fwkplugin.NewDataKey("k1", "test-producer") + testKeyExtra = fwkplugin.NewDataKey("extra", "test-producer") +) + type cloneableString string func (s cloneableString) Clone() fwkdl.Cloneable { return s } @@ -71,7 +78,7 @@ func TestNewEndpoint_CopiesInputs(t *testing.T) { meta := newTestMetadata("pod-a") metrics := newTestMetrics() attr := fwkdl.NewAttributes() - attr.Put("key", cloneableString("value")) + attr.Put(testKeyK, cloneableString("value")) ep := NewEndpoint(meta, metrics, attr) assert.NotNil(t, ep) @@ -85,7 +92,7 @@ func TestNewEndpoint_CopiesInputs(t *testing.T) { assert.Equal(t, 3, ep.GetMetrics().RunningRequestsSize) // values from attribute map should be retrievable - v, ok := ep.Get("key") + v, ok := ep.Get(testKeyK) assert.True(t, ok) assert.Equal(t, cloneableString("value"), v) } @@ -96,8 +103,8 @@ func TestNewEndpoint_NilAttributeUsesDefault(t *testing.T) { assert.Empty(t, ep.Keys()) // Should still be safe to write into the default-allocated attribute map - ep.Put("k", cloneableString("v")) - v, ok := ep.Get("k") + ep.Put(testKeyK, cloneableString("v")) + v, ok := ep.Get(testKeyK) assert.True(t, ok) assert.Equal(t, cloneableString("v"), v) } @@ -115,20 +122,20 @@ func TestEndpoint_StringContainsFields(t *testing.T) { func TestEndpoint_Clone(t *testing.T) { attr := fwkdl.NewAttributes() - attr.Put("k", cloneableString("v")) + attr.Put(testKeyK, cloneableString("v")) ep := NewEndpoint(newTestMetadata("pod-d"), newTestMetrics(), attr) cloned := ep.Clone() - v, ok := cloned.Get("k") + v, ok := cloned.Get(testKeyK) assert.True(t, ok) assert.Equal(t, cloneableString("v"), v) } func TestEndpointComparer_Equal(t *testing.T) { attrA := fwkdl.NewAttributes() - attrA.Put("k", cloneableString("v")) + attrA.Put(testKeyK, cloneableString("v")) attrB := fwkdl.NewAttributes() - attrB.Put("k", cloneableString("v")) + attrB.Put(testKeyK, cloneableString("v")) a := NewEndpoint(newTestMetadata("pod"), newTestMetrics(), attrA) b := NewEndpoint(newTestMetadata("pod"), newTestMetrics(), attrB) @@ -153,10 +160,10 @@ func TestEndpointComparer_DifferByMetrics(t *testing.T) { func TestEndpointComparer_DifferByAttributeKeys(t *testing.T) { attrA := fwkdl.NewAttributes() - attrA.Put("k1", cloneableString("v")) + attrA.Put(testKeyK1, cloneableString("v")) attrB := fwkdl.NewAttributes() - attrB.Put("k1", cloneableString("v")) - attrB.Put("extra", cloneableString("x")) + attrB.Put(testKeyK1, cloneableString("v")) + attrB.Put(testKeyExtra, cloneableString("x")) a := NewEndpoint(newTestMetadata("pod"), newTestMetrics(), attrA) b := NewEndpoint(newTestMetadata("pod"), newTestMetrics(), attrB) @@ -166,9 +173,9 @@ func TestEndpointComparer_DifferByAttributeKeys(t *testing.T) { func TestEndpointComparer_DifferByAttributeValues(t *testing.T) { attrA := fwkdl.NewAttributes() - attrA.Put("k", cloneableString("v1")) + attrA.Put(testKeyK, cloneableString("v1")) attrB := fwkdl.NewAttributes() - attrB.Put("k", cloneableString("v2")) + attrB.Put(testKeyK, cloneableString("v2")) a := NewEndpoint(newTestMetadata("pod"), newTestMetrics(), attrA) b := NewEndpoint(newTestMetadata("pod"), newTestMetrics(), attrB) diff --git a/pkg/epp/framework/plugins/datalayer/attribute/gpu/data_types.go b/pkg/epp/framework/plugins/datalayer/attribute/gpu/data_types.go index e71af7b69a..9a980222df 100644 --- a/pkg/epp/framework/plugins/datalayer/attribute/gpu/data_types.go +++ b/pkg/epp/framework/plugins/datalayer/attribute/gpu/data_types.go @@ -40,6 +40,6 @@ func (v GPUUtilization) Clone() fwkdl.Cloneable { } // ReadGPUUtilization retrieves GPU utilization from an endpoint's AttributeMap. -func ReadGPUUtilization(attrs fwkdl.AttributeMap, key string) (GPUUtilization, bool) { +func ReadGPUUtilization(attrs fwkdl.AttributeMap, key plugin.DataKey) (GPUUtilization, bool) { return fwkdl.ReadAttribute[GPUUtilization](attrs, key) } diff --git a/pkg/epp/framework/plugins/datalayer/attribute/metrics/data_types.go b/pkg/epp/framework/plugins/datalayer/attribute/metrics/data_types.go index 6b6c24f0f8..a194c565e0 100644 --- a/pkg/epp/framework/plugins/datalayer/attribute/metrics/data_types.go +++ b/pkg/epp/framework/plugins/datalayer/attribute/metrics/data_types.go @@ -16,7 +16,16 @@ limitations under the License. package metrics -import fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer" +import ( + 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" +) + +const ( + // MetricsExtractorType is the plugin type for the core metrics extractor, + // which publishes the custom scalar metric attributes. + MetricsExtractorType = "core-metrics-extractor" +) // ScalarMetricValue is a numeric endpoint attribute extracted from a configured scalar metric. type ScalarMetricValue float64 @@ -25,6 +34,14 @@ func (v ScalarMetricValue) Clone() fwkdl.Cloneable { return v } -func ReadScalarMetricValue(attrs fwkdl.AttributeMap, key string) (ScalarMetricValue, bool) { +// ScalarMetricDataKey returns the key under which the core metrics extractor +// publishes the custom scalar metric configured as attributeKey. The data type +// of a custom metric is chosen in configuration rather than in code, so the key +// is derived at construction time instead of being declared as a package var. +func ScalarMetricDataKey(attributeKey string) plugin.DataKey { + return plugin.NewDataKey(attributeKey, MetricsExtractorType) +} + +func ReadScalarMetricValue(attrs fwkdl.AttributeMap, key plugin.DataKey) (ScalarMetricValue, bool) { return fwkdl.ReadAttribute[ScalarMetricValue](attrs, key) } diff --git a/pkg/epp/framework/plugins/datalayer/extractor/dcgm/extractor.go b/pkg/epp/framework/plugins/datalayer/extractor/dcgm/extractor.go index b9fd463526..940b0dead7 100644 --- a/pkg/epp/framework/plugins/datalayer/extractor/dcgm/extractor.go +++ b/pkg/epp/framework/plugins/datalayer/extractor/dcgm/extractor.go @@ -102,7 +102,7 @@ func (e *Extractor) Extract(_ context.Context, in fwkdl.PollInput[sourcemetrics. } normalized := attrgpu.GPUUtilization(maxUtil / 100.0) - in.Endpoint.GetAttributes().Put(e.dk.String(), normalized) + in.Endpoint.GetAttributes().Put(e.dk, normalized) return nil } diff --git a/pkg/epp/framework/plugins/datalayer/extractor/dcgm/extractor_test.go b/pkg/epp/framework/plugins/datalayer/extractor/dcgm/extractor_test.go index c61003e939..fd0645c633 100644 --- a/pkg/epp/framework/plugins/datalayer/extractor/dcgm/extractor_test.go +++ b/pkg/epp/framework/plugins/datalayer/extractor/dcgm/extractor_test.go @@ -87,7 +87,7 @@ func TestExtractorExtract(t *testing.T) { t.Error("empty extractor name") } - key := attrgpu.GPUUtilizationDataKey.WithNonEmptyProducerName(attrgpu.DCGMExtractorType).String() + key := attrgpu.GPUUtilizationDataKey.WithNonEmptyProducerName(attrgpu.DCGMExtractorType) gaugeType := dto.MetricType_GAUGE tests := []struct { diff --git a/pkg/epp/framework/plugins/datalayer/extractor/metrics/extractor.go b/pkg/epp/framework/plugins/datalayer/extractor/metrics/extractor.go index c9fbcc1f58..c4858b4783 100644 --- a/pkg/epp/framework/plugins/datalayer/extractor/metrics/extractor.go +++ b/pkg/epp/framework/plugins/datalayer/extractor/metrics/extractor.go @@ -178,7 +178,7 @@ func (ext *Extractor) Extract(ctx context.Context, in fwkdl.PollInput[sourcemetr errs = append(errs, fmt.Errorf("custom metric %q: %w", custom.AttributeKey, err)) continue } - ep.GetAttributes().Put(custom.AttributeKey, attrmetrics.ScalarMetricValue(extractValue(metric))) + ep.GetAttributes().Put(attrmetrics.ScalarMetricDataKey(custom.AttributeKey), attrmetrics.ScalarMetricValue(extractValue(metric))) updated = true } diff --git a/pkg/epp/framework/plugins/datalayer/extractor/metrics/factories.go b/pkg/epp/framework/plugins/datalayer/extractor/metrics/factories.go index 1f783b7cbf..806d2ea684 100644 --- a/pkg/epp/framework/plugins/datalayer/extractor/metrics/factories.go +++ b/pkg/epp/framework/plugins/datalayer/extractor/metrics/factories.go @@ -25,10 +25,11 @@ import ( "sigs.k8s.io/controller-runtime/pkg/log" fwkplugin "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin" + attrmetrics "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/metrics" ) const ( - MetricsExtractorType = "core-metrics-extractor" + MetricsExtractorType = attrmetrics.MetricsExtractorType ) // Configuration parameters for metrics data source and extractor. diff --git a/pkg/epp/framework/plugins/datalayer/extractor/metrics/metrics_extraction_from_config_test.go b/pkg/epp/framework/plugins/datalayer/extractor/metrics/metrics_extraction_from_config_test.go index cd4bb694f1..5f14ac61c5 100644 --- a/pkg/epp/framework/plugins/datalayer/extractor/metrics/metrics_extraction_from_config_test.go +++ b/pkg/epp/framework/plugins/datalayer/extractor/metrics/metrics_extraction_from_config_test.go @@ -378,7 +378,7 @@ func TestMetricsExtractionCustomScalarFromConfig(t *testing.T) { require.NoError(t, p.Poll(context.Background(), ep)) - got, ok := attrmetrics.ReadScalarMetricValue(ep.GetAttributes(), attributeKey) + got, ok := attrmetrics.ReadScalarMetricValue(ep.GetAttributes(), attrmetrics.ScalarMetricDataKey(attributeKey)) require.True(t, ok, "custom scalar metric should be stored as an endpoint attribute") assert.InDelta(t, 42.5, float64(got), 0.001) assert.Zero(t, ep.GetMetrics().WaitingQueueSize) @@ -419,7 +419,7 @@ func TestMetricsExtractionCustomCounterFromConfig(t *testing.T) { }) require.NoError(t, err) - got, ok := attrmetrics.ReadScalarMetricValue(ep.GetAttributes(), attributeKey) + got, ok := attrmetrics.ReadScalarMetricValue(ep.GetAttributes(), attrmetrics.ScalarMetricDataKey(attributeKey)) require.True(t, ok, "custom counter should be stored as an endpoint attribute") assert.InDelta(t, 12.0, float64(got), 0.001) } diff --git a/pkg/epp/framework/plugins/datalayer/extractor/models/extractor.go b/pkg/epp/framework/plugins/datalayer/extractor/models/extractor.go index c27b0c0aab..3bac54a7d3 100644 --- a/pkg/epp/framework/plugins/datalayer/extractor/models/extractor.go +++ b/pkg/epp/framework/plugins/datalayer/extractor/models/extractor.go @@ -53,7 +53,7 @@ func ModelServerExtractorFactory(name string, _ *json.Decoder, _ fwkplugin.Handl // Extract stores the model list as an endpoint attribute. func (me *ModelExtractor) Extract(_ context.Context, in fwkdl.PollInput[*ModelResponse]) error { - in.Endpoint.GetAttributes().Put(me.dk.String(), attrmodels.ModelDataCollection(in.Payload.Data)) + in.Endpoint.GetAttributes().Put(me.dk, attrmodels.ModelDataCollection(in.Payload.Data)) return nil } 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..46f02360f6 100644 --- a/pkg/epp/framework/plugins/datalayer/extractor/models/extractor_test.go +++ b/pkg/epp/framework/plugins/datalayer/extractor/models/extractor_test.go @@ -67,7 +67,7 @@ func TestExtractorExtract(t *testing.T) { }, } - key := attrmodels.ModelsAttributeKey.WithNonEmptyProducerName(attrmodels.ModelsExtractorType).String() + key := attrmodels.ModelsAttributeKey.WithNonEmptyProducerName(attrmodels.ModelsExtractorType) for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { defer func() { diff --git a/pkg/epp/framework/plugins/datalayer/extractor/topology/extractor.go b/pkg/epp/framework/plugins/datalayer/extractor/topology/extractor.go index be4ddaba03..54655656a8 100644 --- a/pkg/epp/framework/plugins/datalayer/extractor/topology/extractor.go +++ b/pkg/epp/framework/plugins/datalayer/extractor/topology/extractor.go @@ -218,7 +218,7 @@ func (h *endpointHandler) Extract(_ context.Context, event fwkdl.EndpointEvent) if hn == "" && rack == "" && zone == "" && region == "" { return nil } - event.Endpoint.GetAttributes().Put(h.ext.dk.String(), &attrtopology.Topology{ + event.Endpoint.GetAttributes().Put(h.ext.dk, &attrtopology.Topology{ Hostname: hn, Rack: rack, Zone: zone, @@ -288,14 +288,14 @@ func (h *podNotificationHandler) Extract(_ context.Context, event fwkdl.Notifica for _, ep := range eps { // Preserve rack/zone/region already set from the endpoint event. topo := &attrtopology.Topology{Hostname: hostname} - if raw, ok := ep.GetAttributes().Get(h.ext.dk.String()); ok { + if raw, ok := ep.GetAttributes().Get(h.ext.dk); ok { if existing, ok := raw.(*attrtopology.Topology); ok { topo.Rack = existing.Rack topo.Zone = existing.Zone topo.Region = existing.Region } } - ep.GetAttributes().Put(h.ext.dk.String(), topo) + ep.GetAttributes().Put(h.ext.dk, topo) } return nil } diff --git a/pkg/epp/framework/plugins/datalayer/extractor/topology/extractor_test.go b/pkg/epp/framework/plugins/datalayer/extractor/topology/extractor_test.go index fa8f5db985..ed2b6d7a99 100644 --- a/pkg/epp/framework/plugins/datalayer/extractor/topology/extractor_test.go +++ b/pkg/epp/framework/plugins/datalayer/extractor/topology/extractor_test.go @@ -70,7 +70,7 @@ const ( ) func readTopology(ep fwkdl.Endpoint) (*attrtopology.Topology, bool) { - dk := attrtopology.TopologyAttributeKey.WithNonEmptyProducerName(testPluginName).String() + dk := attrtopology.TopologyAttributeKey.WithNonEmptyProducerName(testPluginName) raw, ok := ep.GetAttributes().Get(dk) if !ok { return nil, false diff --git a/pkg/epp/framework/plugins/flowcontrol/saturationdetector/concurrency/detector.go b/pkg/epp/framework/plugins/flowcontrol/saturationdetector/concurrency/detector.go index 5dce0e6484..fec9bfce96 100644 --- a/pkg/epp/framework/plugins/flowcontrol/saturationdetector/concurrency/detector.go +++ b/pkg/epp/framework/plugins/flowcontrol/saturationdetector/concurrency/detector.go @@ -110,7 +110,7 @@ func (d *detector) Consumes() fwkplugin.DataDependencies { } func (d *detector) getLoad(m datalayer.AttributeMap) *attrconcurrency.InFlightLoad { - if val, ok := m.Get(d.inFlightLoadDataKey.String()); ok { + if val, ok := m.Get(d.inFlightLoadDataKey); ok { if load, ok := val.(*attrconcurrency.InFlightLoad); ok { return load } diff --git a/pkg/epp/framework/plugins/flowcontrol/saturationdetector/concurrency/detector_test.go b/pkg/epp/framework/plugins/flowcontrol/saturationdetector/concurrency/detector_test.go index 0126420a9c..b8b76db601 100644 --- a/pkg/epp/framework/plugins/flowcontrol/saturationdetector/concurrency/detector_test.go +++ b/pkg/epp/framework/plugins/flowcontrol/saturationdetector/concurrency/detector_test.go @@ -642,15 +642,17 @@ func (e *liveEndpoint) UpdateMetrics(*datalayer.Metrics) {} func (e *liveEndpoint) String() string { return e.id } // liveEndpoint also implements AttributeMap. -func (e *liveEndpoint) Get(key string) (datalayer.Cloneable, bool) { - if key == attrconcurrency.InFlightLoadDataKey.String() { +func (e *liveEndpoint) Get(key fwkplugin.DataKey) (datalayer.Cloneable, bool) { + if key == attrconcurrency.InFlightLoadDataKey { return e.reg.get(e.id), true } return nil, false } -func (e *liveEndpoint) Put(string, datalayer.Cloneable) {} -func (e *liveEndpoint) Keys() []string { return []string{attrconcurrency.InFlightLoadDataKey.String()} } -func (e *liveEndpoint) Clone() datalayer.AttributeMap { return e } +func (e *liveEndpoint) Put(fwkplugin.DataKey, datalayer.Cloneable) {} +func (e *liveEndpoint) Keys() []fwkplugin.DataKey { + return []fwkplugin.DataKey{attrconcurrency.InFlightLoadDataKey} +} +func (e *liveEndpoint) Clone() datalayer.AttributeMap { return e } func newFakeEndpoint(reg *localRegistry, name string) datalayer.Endpoint { id := fullEndpointName(name) @@ -678,15 +680,15 @@ func newStubSchedulingEndpoint(reg *localRegistry, name string) *liveSchedulingE } func (f *liveSchedulingEndpoint) GetMetadata() *datalayer.EndpointMetadata { return f.metadata } -func (f *liveSchedulingEndpoint) Get(key string) (datalayer.Cloneable, bool) { - if key == attrconcurrency.InFlightLoadDataKey.String() { +func (f *liveSchedulingEndpoint) Get(key fwkplugin.DataKey) (datalayer.Cloneable, bool) { + if key == attrconcurrency.InFlightLoadDataKey { return f.reg.get(f.id), true } return nil, false } -func (f *liveSchedulingEndpoint) Put(string, datalayer.Cloneable) {} -func (f *liveSchedulingEndpoint) Keys() []string { - return []string{attrconcurrency.InFlightLoadDataKey.String()} +func (f *liveSchedulingEndpoint) Put(fwkplugin.DataKey, datalayer.Cloneable) {} +func (f *liveSchedulingEndpoint) Keys() []fwkplugin.DataKey { + return []fwkplugin.DataKey{attrconcurrency.InFlightLoadDataKey} } func (f *liveSchedulingEndpoint) String() string { return f.id } func (f *liveSchedulingEndpoint) Clone() datalayer.AttributeMap { return f } @@ -706,16 +708,16 @@ func makeTokenRequest(requestID string, inputTokens int) *fwksched.InferenceRequ // has not been populated yet. type nilMetadataEndpoint struct{} -func (e *nilMetadataEndpoint) GetMetadata() *datalayer.EndpointMetadata { return nil } -func (e *nilMetadataEndpoint) UpdateMetadata(m *datalayer.EndpointMetadata) {} -func (e *nilMetadataEndpoint) GetAttributes() datalayer.AttributeMap { return e } -func (e *nilMetadataEndpoint) GetMetrics() *datalayer.Metrics { return nil } -func (e *nilMetadataEndpoint) UpdateMetrics(*datalayer.Metrics) {} -func (e *nilMetadataEndpoint) String() string { return "nil-metadata" } -func (e *nilMetadataEndpoint) Get(string) (datalayer.Cloneable, bool) { return nil, false } -func (e *nilMetadataEndpoint) Put(string, datalayer.Cloneable) {} -func (e *nilMetadataEndpoint) Keys() []string { return nil } -func (e *nilMetadataEndpoint) Clone() datalayer.AttributeMap { return e } +func (e *nilMetadataEndpoint) GetMetadata() *datalayer.EndpointMetadata { return nil } +func (e *nilMetadataEndpoint) UpdateMetadata(m *datalayer.EndpointMetadata) {} +func (e *nilMetadataEndpoint) GetAttributes() datalayer.AttributeMap { return e } +func (e *nilMetadataEndpoint) GetMetrics() *datalayer.Metrics { return nil } +func (e *nilMetadataEndpoint) UpdateMetrics(*datalayer.Metrics) {} +func (e *nilMetadataEndpoint) String() string { return "nil-metadata" } +func (e *nilMetadataEndpoint) Get(fwkplugin.DataKey) (datalayer.Cloneable, bool) { return nil, false } +func (e *nilMetadataEndpoint) Put(fwkplugin.DataKey, datalayer.Cloneable) {} +func (e *nilMetadataEndpoint) Keys() []fwkplugin.DataKey { return nil } +func (e *nilMetadataEndpoint) Clone() datalayer.AttributeMap { return e } func TestDetector_NilEndpointInList(t *testing.T) { t.Parallel() diff --git a/pkg/epp/framework/plugins/requestcontrol/admitter/latencyslo/plugin.go b/pkg/epp/framework/plugins/requestcontrol/admitter/latencyslo/plugin.go index a3786be082..3d851a5387 100644 --- a/pkg/epp/framework/plugins/requestcontrol/admitter/latencyslo/plugin.go +++ b/pkg/epp/framework/plugins/requestcontrol/admitter/latencyslo/plugin.go @@ -141,7 +141,7 @@ func (p *LatencyAdmission) Admit(ctx context.Context, request *fwksched.Inferenc } // Valid prediction: both TTFT and TPOT within SLO. - if latencyInfoRaw, ok := endpoint.Get(p.latencyPredictionInfoDataKey.String()); ok { + if latencyInfoRaw, ok := endpoint.Get(p.latencyPredictionInfoDataKey); ok { hasPredictions = true latencyInfo := latencyInfoRaw.(*attrlatency.LatencyPredictionInfo) if latencyInfo.IsValid() { diff --git a/pkg/epp/framework/plugins/requestcontrol/admitter/latencyslo/plugin_test.go b/pkg/epp/framework/plugins/requestcontrol/admitter/latencyslo/plugin_test.go index 9028eecad9..be96be0738 100644 --- a/pkg/epp/framework/plugins/requestcontrol/admitter/latencyslo/plugin_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/admitter/latencyslo/plugin_test.go @@ -82,7 +82,7 @@ func TestAdmit(t *testing.T) { }, setupFn: func(endpoints []fwksched.Endpoint) { // All invalid predictions - endpoints[0].Put(attrlatency.LatencyPredictionInfoDataKey.String(), + endpoints[0].Put(attrlatency.LatencyPredictionInfoDataKey, attrlatency.NewLatencyPredictionInfo(false, false, -50, -10, 150, 40, 0)) }, wantErr: false, @@ -103,9 +103,9 @@ func TestAdmit(t *testing.T) { makeLatencyAdmissionEndpoint("pod2", 0.4, 3), }, setupFn: func(endpoints []fwksched.Endpoint) { - endpoints[0].Put(attrlatency.LatencyPredictionInfoDataKey.String(), + endpoints[0].Put(attrlatency.LatencyPredictionInfoDataKey, attrlatency.NewLatencyPredictionInfo(false, false, -50, -10, 150, 40, 0)) - endpoints[1].Put(attrlatency.LatencyPredictionInfoDataKey.String(), + endpoints[1].Put(attrlatency.LatencyPredictionInfoDataKey, attrlatency.NewLatencyPredictionInfo(false, false, -30, -5, 130, 35, 0)) }, wantErr: true, @@ -123,7 +123,7 @@ func TestAdmit(t *testing.T) { makeLatencyAdmissionEndpoint("pod1", 0.5, 5), }, setupFn: func(endpoints []fwksched.Endpoint) { - endpoints[0].Put(attrlatency.LatencyPredictionInfoDataKey.String(), + endpoints[0].Put(attrlatency.LatencyPredictionInfoDataKey, attrlatency.NewLatencyPredictionInfo(false, false, -50, -10, 150, 40, 0)) }, wantErr: true, @@ -143,7 +143,7 @@ func TestAdmit(t *testing.T) { makeLatencyAdmissionEndpoint("pod1", 0.5, 5), }, setupFn: func(endpoints []fwksched.Endpoint) { - endpoints[0].Put(attrlatency.LatencyPredictionInfoDataKey.String(), + endpoints[0].Put(attrlatency.LatencyPredictionInfoDataKey, attrlatency.NewLatencyPredictionInfo(false, false, -50, -10, 150, 40, 0)) }, wantErr: false, @@ -156,9 +156,9 @@ func TestAdmit(t *testing.T) { makeLatencyAdmissionEndpoint("pod2", 0.4, 0), // idle }, setupFn: func(endpoints []fwksched.Endpoint) { - endpoints[0].Put(attrlatency.LatencyPredictionInfoDataKey.String(), + endpoints[0].Put(attrlatency.LatencyPredictionInfoDataKey, attrlatency.NewLatencyPredictionInfo(false, false, -50, -10, 150, 40, 0)) - endpoints[1].Put(attrlatency.LatencyPredictionInfoDataKey.String(), + endpoints[1].Put(attrlatency.LatencyPredictionInfoDataKey, attrlatency.NewLatencyPredictionInfo(false, false, -30, -5, 130, 35, 0)) }, wantErr: false, @@ -171,9 +171,9 @@ func TestAdmit(t *testing.T) { makeLatencyAdmissionEndpoint("pod2", 0.01, 3), // cold }, setupFn: func(endpoints []fwksched.Endpoint) { - endpoints[0].Put(attrlatency.LatencyPredictionInfoDataKey.String(), + endpoints[0].Put(attrlatency.LatencyPredictionInfoDataKey, attrlatency.NewLatencyPredictionInfo(false, false, -50, -10, 150, 40, 0)) - endpoints[1].Put(attrlatency.LatencyPredictionInfoDataKey.String(), + endpoints[1].Put(attrlatency.LatencyPredictionInfoDataKey, attrlatency.NewLatencyPredictionInfo(false, false, -30, -5, 130, 35, 0)) }, wantErr: false, @@ -186,9 +186,9 @@ func TestAdmit(t *testing.T) { makeLatencyAdmissionEndpoint("pod2", 0.4, 3), }, setupFn: func(endpoints []fwksched.Endpoint) { - endpoints[0].Put(attrlatency.LatencyPredictionInfoDataKey.String(), + endpoints[0].Put(attrlatency.LatencyPredictionInfoDataKey, attrlatency.NewLatencyPredictionInfo(false, false, -50, -10, 150, 40, 0)) - endpoints[1].Put(attrlatency.LatencyPredictionInfoDataKey.String(), + endpoints[1].Put(attrlatency.LatencyPredictionInfoDataKey, attrlatency.NewLatencyPredictionInfo(true, true, 20, 5, 80, 25, 0)) // valid }, wantErr: false, diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin.go index d8f0709d18..103f88b512 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin.go @@ -240,7 +240,7 @@ func (p *dataProducer) Produce(ctx context.Context, request *fwksched.InferenceR for _, pod := range pods { matchLen := prefixCacheServers[ServerID(pod.GetMetadata().NamespacedName)] - pod.Put(p.dk.String(), attrprefix.NewPrefixCacheMatchInfo(matchLen, totalBlocks, blockSize)) + pod.Put(p.dk, attrprefix.NewPrefixCacheMatchInfo(matchLen, totalBlocks, blockSize)) } state := &SchedulingContextState{ diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin_test.go index 6d22e99079..a4d433e4cc 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin_test.go @@ -92,7 +92,7 @@ func TestProduce(t *testing.T) { assert.Equal(t, 2, len(state.PerPromptHashes[0])) // 2 token IDs at blockSize 1 -> 2 blocks // Verify pod match info was set (should be 0 match since indexer is empty) - key := attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(ApproxPrefixCachePluginType).String() + key := attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(ApproxPrefixCachePluginType) for _, ep := range endpoints { info, ok := ep.Get(key) assert.True(t, ok) @@ -269,7 +269,7 @@ func TestPrefixPluginPartialPrefixMatch(t *testing.T) { } _ = p.Produce(context.Background(), req3, endpoints) - key := attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(ApproxPrefixCachePluginType).String() + key := attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(ApproxPrefixCachePluginType) // Verify pod1 has the correct prefix match info info1, _ := endpoint1.Get(key) prefixInfo1 := info1.(*attrprefix.PrefixCacheMatchInfo) @@ -332,7 +332,7 @@ func TestPrefixPluginPrefixGrowth(t *testing.T) { extendedHashCount := len(state2.PerPromptHashes[0]) assert.Greater(t, extendedHashCount, initialHashCount) - key := attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(ApproxPrefixCachePluginType).String() + key := attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(ApproxPrefixCachePluginType) info, _ := endpoint1.Get(key) prefixInfo := info.(*attrprefix.PrefixCacheMatchInfo) assert.Greater(t, prefixInfo.MatchBlocks(), 0, "should have prefix cache hit") @@ -678,7 +678,7 @@ func TestProduce_MultiPrompt(t *testing.T) { assert.Equal(t, 3, len(state.PerPromptHashes[0]), "first prompt: 3 tokens at blockSize 1") assert.Equal(t, 2, len(state.PerPromptHashes[1]), "second prompt: 2 tokens at blockSize 1") - key := attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(ApproxPrefixCachePluginType).String() + key := attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(ApproxPrefixCachePluginType) info, ok := endpoint.Get(key) assert.True(t, ok) prefixInfo := info.(*attrprefix.PrefixCacheMatchInfo) @@ -732,7 +732,7 @@ func TestMultiPromptMatchAggregation(t *testing.T) { } _ = p.Produce(context.Background(), req2, endpoints) - key := attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(ApproxPrefixCachePluginType).String() + key := attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(ApproxPrefixCachePluginType) info, _ := endpoint.Get(key) prefixInfo := info.(*attrprefix.PrefixCacheMatchInfo) assert.Equal(t, 5, prefixInfo.MatchBlocks(), "all 5 blocks (3+2) should match") @@ -785,7 +785,7 @@ func TestMultiPromptPartialMatch(t *testing.T) { } _ = p.Produce(context.Background(), req2, endpoints) - key := attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(ApproxPrefixCachePluginType).String() + key := attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(ApproxPrefixCachePluginType) info, _ := endpoint.Get(key) prefixInfo := info.(*attrprefix.PrefixCacheMatchInfo) assert.Equal(t, 2, prefixInfo.MatchBlocks(), "only first prompt's 2 blocks should match") @@ -824,7 +824,7 @@ func TestPrefixPluginTokenizedRequest(t *testing.T) { assert.Equal(t, 4, len(state.PerPromptHashes[0])) // Verify match info was set on the endpoint (0 match since indexer is empty). - key := attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(ApproxPrefixCachePluginType).String() + key := attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(ApproxPrefixCachePluginType) info, ok := endpoint.Get(key) assert.True(t, ok) prefixInfo := info.(*attrprefix.PrefixCacheMatchInfo) diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin.go index 7eb3619710..62011ffed8 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin.go @@ -171,7 +171,7 @@ func (p *dataProducer) Produce(ctx context.Context, request *fwksched.InferenceR matchLen = total matched = true } - pod.Put(p.dk.String(), attrprefix.NewPrefixCacheMatchInfo(matchLen, total, p.config.BlockSizeTokens)) + pod.Put(p.dk, attrprefix.NewPrefixCacheMatchInfo(matchLen, total, p.config.BlockSizeTokens)) } if e.assigned != nil && !matched { // The endpoint set changed between batching and release (rolling update or diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin_test.go index 21f641268e..7de71e9af2 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin_test.go @@ -42,7 +42,7 @@ func tokenizedRequest(tokens []uint32) *fwksched.InferenceRequest { func (p *dataProducer) assignedReplica(t *testing.T, endpoints []fwksched.Endpoint) string { t.Helper() for _, ep := range endpoints { - v, ok := ep.Get(p.dk.String()) + v, ok := ep.Get(p.dk) require.True(t, ok, "PrefixCacheMatchInfo must be attached to every endpoint") info, ok := v.(*attrprefix.PrefixCacheMatchInfo) require.True(t, ok) @@ -121,7 +121,7 @@ func TestProduce_CancelledContextBeforeSeal(t *testing.T) { require.Error(t, err, "a context cancelled before seal must return an error") for _, ep := range endpoints { - v, ok := ep.Get(p.dk.String()) + v, ok := ep.Get(p.dk) if !ok { continue // no affinity attached is expected when Produce returns early } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/producer.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/producer.go index 9b5b27dce5..a7a4fe74b3 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/producer.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/producer.go @@ -301,7 +301,7 @@ func (p *InFlightLoadProducer) Extract(ctx context.Context, event datalayer.Endp log.FromContext(ctx).V(logutil.DEFAULT).Info("Cleaned up in-flight load for deleted endpoint", "endpoint", id) case datalayer.EventAddOrUpdate: p.registeredEndpoints.Store(id, event.Endpoint) - event.Endpoint.GetAttributes().Put(p.dk.String(), &datalayer.DynamicAttribute{ + event.Endpoint.GetAttributes().Put(p.dk, &datalayer.DynamicAttribute{ Get: func() datalayer.Cloneable { return &attrconcurrency.InFlightLoad{ Tokens: p.GetTokens(id), @@ -326,7 +326,7 @@ func (p *InFlightLoadProducer) Produce(_ context.Context, request *fwksched.Infe } if request != nil { tokens := p.estimateRequestTokens(e, request, inputTokens) - e.Put(p.uncachedRequestTokensDk.String(), &attrconcurrency.UncachedRequestTokens{ + e.Put(p.uncachedRequestTokensDk, &attrconcurrency.UncachedRequestTokens{ Tokens: tokens, }) } @@ -391,7 +391,7 @@ func (p *InFlightLoadProducer) PreRequest(ctx context.Context, request *fwksched } func (p *InFlightLoadProducer) estimateRequestTokens(endpoint fwksched.Endpoint, request *fwksched.InferenceRequest, inputTokens int64) int64 { - adjustedInput := uncachedInputTokens(endpoint, inputTokens, p.prefixMatchInfoDK.String()) + adjustedInput := uncachedInputTokens(endpoint, inputTokens, p.prefixMatchInfoDK) tokens := adjustedInput if p.addEstimatedOutputTokens { var maxOutputTokens *int64 @@ -518,7 +518,7 @@ func addedTokensKey(endpointID, profileName string) string { // still reflected. // // When the attribute is missing, we fall back to the estimated inputTokens. -func uncachedInputTokens(endpoint fwksched.Endpoint, inputTokens int64, prefixMatchInfoKey string) int64 { +func uncachedInputTokens(endpoint fwksched.Endpoint, inputTokens int64, prefixMatchInfoKey fwkplugin.DataKey) int64 { if endpoint == nil { return nonNeg(inputTokens) } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/producer_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/producer_test.go index ff16c93837..6c4b9bda63 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/producer_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/producer_test.go @@ -93,12 +93,12 @@ func TestInFlightLoadProducer_PrefixMatchInfoProducerName(t *testing.T) { // The discount reads PrefixCacheMatchInfo from the configured producer's key // (indexed 2*4=8, matched 1*4=4 -> uncached 4). hit := newStubSchedulingEndpoint("ep-hit") - hit.Put(preciseKey.String(), attrprefix.NewPrefixCacheMatchInfo(1, 2, 4)) + hit.Put(preciseKey, attrprefix.NewPrefixCacheMatchInfo(1, 2, 4)) require.Equal(t, int64(4), producer.estimateRequestTokens(hit, nil, 5)) // Data under the approx (default) key is ignored, so it falls back to inputTokens. miss := newStubSchedulingEndpoint("ep-miss") - miss.Put(attrprefix.PrefixCacheMatchInfoDataKey.String(), attrprefix.NewPrefixCacheMatchInfo(1, 2, 4)) + miss.Put(attrprefix.PrefixCacheMatchInfoDataKey, attrprefix.NewPrefixCacheMatchInfo(1, 2, 4)) require.Equal(t, int64(5), producer.estimateRequestTokens(miss, nil, 5)) } @@ -114,7 +114,7 @@ func TestInFlightLoadProducer_Produce(t *testing.T) { // 1. Produce with nil request -> should not put anything err := producer.Produce(context.Background(), nil, endpoints) require.NoError(t, err) - _, ok := endpoint.Get(producer.uncachedRequestTokensDk.String()) + _, ok := endpoint.Get(producer.uncachedRequestTokensDk) require.False(t, ok) // 2. Produce with request -> should put UncachedRequestTokens @@ -123,13 +123,13 @@ func TestInFlightLoadProducer_Produce(t *testing.T) { require.NoError(t, err) - val, ok := endpoint.Get(producer.uncachedRequestTokensDk.String()) + val, ok := endpoint.Get(producer.uncachedRequestTokensDk) require.True(t, ok) uncached := val.(*attrconcurrency.UncachedRequestTokens) require.Equal(t, int64(10), uncached.Tokens) // Verify that InFlightLoad was NOT put/overwritten by Produce - _, ok = endpoint.Get(producer.dk.String()) + _, ok = endpoint.Get(producer.dk) require.False(t, ok, "InFlightLoad should not be populated by Produce") } @@ -151,7 +151,7 @@ func TestInFlightLoadProducer_Extract(t *testing.T) { require.NoError(t, err) // Verify dynamic attribute is registered - key := producer.dk.String() + key := producer.dk val, ok := endpoint.Get(key) require.True(t, ok) @@ -575,11 +575,13 @@ func (f *stubSchedulingEndpoint) GetMetrics() *datalayer.Metrics { r func (f *stubSchedulingEndpoint) UpdateMetrics(*datalayer.Metrics) {} func (f *stubSchedulingEndpoint) GetAttributes() datalayer.AttributeMap { return f.attr } func (f *stubSchedulingEndpoint) String() string { return "" } -func (f *stubSchedulingEndpoint) Put(key string, val datalayer.Cloneable) { f.attr.Put(key, val) } -func (f *stubSchedulingEndpoint) Get(key string) (datalayer.Cloneable, bool) { +func (f *stubSchedulingEndpoint) Put(key fwkplugin.DataKey, val datalayer.Cloneable) { + f.attr.Put(key, val) +} +func (f *stubSchedulingEndpoint) Get(key fwkplugin.DataKey) (datalayer.Cloneable, bool) { return f.attr.Get(key) } -func (f *stubSchedulingEndpoint) Keys() []string { return f.attr.Keys() } +func (f *stubSchedulingEndpoint) Keys() []fwkplugin.DataKey { return f.attr.Keys() } // makeTokenRequest builds a request whose tokenized prompt carries inputTokens token IDs, // which is what the estimator reads to derive the input token count. @@ -663,7 +665,7 @@ func TestInFlightLoadProducer_PrefixCacheDiscount(t *testing.T) { // uncached_input = (2-1)*4 + max(0, 8-2*4) = 4 // total tokens = 4 + 12 = 16 endpoint := newStubSchedulingEndpoint(endpointName) - endpoint.Put(attrprefix.PrefixCacheMatchInfoDataKey.String(), attrprefix.NewPrefixCacheMatchInfo(1, 2, 4)) + endpoint.Put(attrprefix.PrefixCacheMatchInfoDataKey, attrprefix.NewPrefixCacheMatchInfo(1, 2, 4)) req := makeTokenRequest("req-prefix", 8) res := &fwksched.SchedulingResult{ @@ -701,9 +703,9 @@ func TestInFlightLoadProducer_PrefixCacheDiscount_PerEndpoint(t *testing.T) { // 8 input tokens, output 12. epA := newStubSchedulingEndpoint(podA) - epA.Put(attrprefix.PrefixCacheMatchInfoDataKey.String(), attrprefix.NewPrefixCacheMatchInfo(2, 2, 4)) // fully cached + epA.Put(attrprefix.PrefixCacheMatchInfoDataKey, attrprefix.NewPrefixCacheMatchInfo(2, 2, 4)) // fully cached epB := newStubSchedulingEndpoint(podB) - epB.Put(attrprefix.PrefixCacheMatchInfoDataKey.String(), attrprefix.NewPrefixCacheMatchInfo(0, 2, 4)) // none cached + epB.Put(attrprefix.PrefixCacheMatchInfoDataKey, attrprefix.NewPrefixCacheMatchInfo(0, 2, 4)) // none cached req := makeTokenRequest("req-multi-cache", 8) res := &fwksched.SchedulingResult{ @@ -923,11 +925,11 @@ func TestUncachedInputTokens_Overestimate(t *testing.T) { // matchedTokens = 1 * 4 = 4 endpoint := newStubSchedulingEndpoint("test-ep") - endpoint.Put(attrprefix.PrefixCacheMatchInfoDataKey.String(), attrprefix.NewPrefixCacheMatchInfo(1, 2, 4)) + endpoint.Put(attrprefix.PrefixCacheMatchInfoDataKey, attrprefix.NewPrefixCacheMatchInfo(1, 2, 4)) inputTokens := int64(5) - uncached := uncachedInputTokens(endpoint, inputTokens, attrprefix.PrefixCacheMatchInfoDataKey.String()) + uncached := uncachedInputTokens(endpoint, inputTokens, attrprefix.PrefixCacheMatchInfoDataKey) // When the prefix cache says 4 tokens are definitely uncached in the indexed portion (8-4), // we trust that over the smaller (approximate) estimate of 5. diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer.go index b11dea3302..123d4e11ea 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer.go @@ -285,7 +285,7 @@ func (p *Producer) Produce(ctx context.Context, request *scheduling.InferenceReq } matchedItems := p.matchedItemsForPod(metadata.NamespacedName.String(), requestItems) p.recordHitRatio(len(matchedItems), len(requestItems)) - endpoint.Put(p.dk.String(), attrmm.NewEncoderCacheMatchInfo( + endpoint.Put(p.dk, attrmm.NewEncoderCacheMatchInfo( matchedItems, requestItems, )) diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer_test.go index 2d9e808c59..526ea765eb 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/multimodal/producer_test.go @@ -284,7 +284,7 @@ func schedulingResult(target scheduling.Endpoint) *scheduling.SchedulingResult { func assertMatchInfo(t *testing.T, p *Producer, endpoint scheduling.Endpoint, matchedItems, requestItems []attrmm.MatchItem) { t.Helper() - raw, ok := endpoint.Get(p.dk.String()) + raw, ok := endpoint.Get(p.dk) require.True(t, ok) info, ok := raw.(*attrmm.EncoderCacheMatchInfo) require.True(t, ok) diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/p2psource/producer.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/p2psource/producer.go index 00669d3560..f1dce71b92 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/p2psource/producer.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/p2psource/producer.go @@ -219,7 +219,7 @@ func (p *Producer) PreRequest(ctx context.Context, request *scheduling.Inference // cached-block count times the block size) from its PrefixCacheMatchInfo, // or 0 when absent. func (p *Producer) cachedTokenCount(ep scheduling.Endpoint) int { - raw, ok := ep.Get(p.prefixMatchDataKey.String()) + raw, ok := ep.Get(p.prefixMatchDataKey) if !ok { return 0 } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/p2psource/producer_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/p2psource/producer_test.go index 2ab11af60e..ec6d7c1019 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/p2psource/producer_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/p2psource/producer_test.go @@ -43,7 +43,7 @@ func endpoint(p *Producer, name, address string, cachedBlocks int) scheduling.En Address: address, Port: "8080", }, nil, nil) - e.Put(p.prefixMatchDataKey.String(), + e.Put(p.prefixMatchDataKey, attrprefix.NewPrefixCacheMatchInfo(cachedBlocks, 4, testBlockSize).WithCachedBlockCount(cachedBlocks)) return e } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer.go index 75fe087ac3..a399a97d80 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer.go @@ -381,10 +381,9 @@ func (p *Producer) produceFromBlockKeys(ctx context.Context, span trace.Span, cachedBlocksByTier[tier] += count } } - ep.Put(p.dk.String(), - attrprefix.NewPrefixCacheMatchInfo(matchLen, totalBlocks, p.blockSizeTokens). - WithCachedBlockCount(cachedBlocks). - WithCachedBlocksByTier(cachedBlocksByTier)) + ep.Put(p.dk, attrprefix.NewPrefixCacheMatchInfo(matchLen, totalBlocks, p.blockSizeTokens). + WithCachedBlockCount(cachedBlocks). + WithCachedBlocksByTier(cachedBlocksByTier)) } if p.speculativeEnabled { diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer_test.go index 757ca27d04..42a8a61f18 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer_test.go @@ -190,7 +190,7 @@ func TestProduce_UsesTokenizedPrompt(t *testing.T) { require.NoError(t, p.Produce(ctx, req, testEndpoints)) require.Equal(t, tokens, capturedTokens) - raw, ok := testEndpoints[0].Get(attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName("test").String()) + raw, ok := testEndpoints[0].Get(attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName("test")) require.True(t, ok) info, ok := raw.(*attrprefix.PrefixCacheMatchInfo) require.True(t, ok) @@ -198,7 +198,7 @@ func TestProduce_UsesTokenizedPrompt(t *testing.T) { assert.Equal(t, 1, info.TotalBlocks()) assert.Equal(t, 16, info.BlockSizeTokens()) - raw2, ok := testEndpoints[1].Get(attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName("test").String()) + raw2, ok := testEndpoints[1].Get(attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName("test")) require.True(t, ok) info2 := raw2.(*attrprefix.PrefixCacheMatchInfo) assert.Equal(t, 0, info2.MatchBlocks()) @@ -294,7 +294,7 @@ func TestProduce_MultiPromptEmptyBlockKeys_NoOp(t *testing.T) { require.NoError(t, p.Produce(ctx, req, endpoints)) require.Equal(t, [][]uint32{promptA, promptB}, computeCalls) - _, ok := endpoints[0].Get(attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName("test").String()) + _, ok := endpoints[0].Get(attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName("test")) assert.False(t, ok) } @@ -353,14 +353,14 @@ func TestProduce_MultiPromptSkipsEmptyPromptKeys(t *testing.T) { require.Equal(t, [][]uint32{shortPrompt, fullPrompt}, computeCalls) require.Equal(t, [][]kvblock.BlockHash{{wantKey}}, lookupCalls) - raw, ok := endpoints[0].Get(attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName("test").String()) + raw, ok := endpoints[0].Get(attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName("test")) require.True(t, ok) info, ok := raw.(*attrprefix.PrefixCacheMatchInfo) require.True(t, ok) assert.Equal(t, 1, info.MatchBlocks()) assert.Equal(t, 1, info.TotalBlocks()) - raw, ok = endpoints[1].Get(attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName("test").String()) + raw, ok = endpoints[1].Get(attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName("test")) require.True(t, ok) info, ok = raw.(*attrprefix.PrefixCacheMatchInfo) require.True(t, ok) @@ -424,7 +424,7 @@ func TestProduce_WritesCachedBlocksByTier(t *testing.T) { require.NoError(t, p.Produce(ctx, req, endpoints)) - raw, ok := endpoints[0].Get(attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName("test").String()) + raw, ok := endpoints[0].Get(attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName("test")) require.True(t, ok) info, ok := raw.(*attrprefix.PrefixCacheMatchInfo) require.True(t, ok) @@ -432,7 +432,7 @@ func TestProduce_WritesCachedBlocksByTier(t *testing.T) { // gpu: 2 (prompt A) + 1 (prompt B); cpu: 1 (prompt A) + 1 (prompt B). assert.Equal(t, map[string]int{"gpu": 3, "cpu": 2}, info.CachedBlocksByTier()) - raw, ok = endpoints[1].Get(attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName("test").String()) + raw, ok = endpoints[1].Get(attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName("test")) require.True(t, ok) info, ok = raw.(*attrprefix.PrefixCacheMatchInfo) require.True(t, ok) @@ -666,7 +666,7 @@ func TestNew_BlockSizeFlowsViaTokenProcessor(t *testing.T) { } require.NoError(t, p.Produce(ctx, req, []scheduling.Endpoint{endpoint})) - raw, ok := endpoint.Get(attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(name).String()) + raw, ok := endpoint.Get(attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(name)) require.True(t, ok) info, ok := raw.(*attrprefix.PrefixCacheMatchInfo) require.True(t, ok, "expected *PrefixCacheMatchInfo, got %T", raw) diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/dataproducer_hooks.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/dataproducer_hooks.go index f1c556de2f..1f708a0d66 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/dataproducer_hooks.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/dataproducer_hooks.go @@ -44,7 +44,7 @@ func (pl *PredictedLatency) Produce(ctx context.Context, request *fwksched.Infer var prefixCacheScore float64 for _, endpoint := range endpoints { - if prefixCacheInfoRaw, ok := endpoint.Get(pl.prefixMatchDataKey.String()); ok { + if prefixCacheInfoRaw, ok := endpoint.Get(pl.prefixMatchDataKey); ok { prefixCacheInfo := prefixCacheInfoRaw.(*attrprefix.PrefixCacheMatchInfo) prefixCacheScore = float64(prefixCacheInfo.MatchBlocks()) / float64(prefixCacheInfo.TotalBlocks()) if !math.IsNaN(prefixCacheScore) { @@ -90,7 +90,7 @@ func (pl *PredictedLatency) Produce(ctx context.Context, request *fwksched.Infer pred.TPOT, pl.getEndpointRunningRequestCount(pred.Endpoint), ) - pred.Endpoint.Put(pl.latencyPredictionInfoDataKey.String(), latencyInfo) + pred.Endpoint.Put(pl.latencyPredictionInfoDataKey, latencyInfo) logger.V(logutil.DEBUG).Info("Stored latency prediction in endpoint", "pod", pred.Endpoint.GetMetadata().NamespacedName.Name, "ttft", pred.TTFT, diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/plugin.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/plugin.go index e957209bdf..6d4d5beae6 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/plugin.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/plugin.go @@ -85,7 +85,7 @@ type PredictedLatency struct { // flight and Requests the active request count. Returns false when the // attribute is absent (e.g. an endpoint added before the producer injected it). func (pl *PredictedLatency) endpointInFlightLoad(endpoint fwksched.Endpoint) (*attrconcurrency.InFlightLoad, bool) { - if raw, ok := endpoint.Get(pl.inFlightLoadDataKey.String()); ok { + if raw, ok := endpoint.Get(pl.inFlightLoadDataKey); ok { if load, ok := raw.(*attrconcurrency.InFlightLoad); ok && load != nil { return load, true } diff --git a/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/plugin_test.go b/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/plugin_test.go index 993c093c78..a307f19475 100644 --- a/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/plugin_test.go +++ b/pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/plugin_test.go @@ -121,7 +121,7 @@ func TestReadInFlightLoad(t *testing.T) { // Attribute present: both fields come from InFlightLoad. epWithLoad := createTestEndpoint("pod2", 0.5, 5, 0) - epWithLoad.Put(attrconcurrency.InFlightLoadDataKey.String(), &attrconcurrency.InFlightLoad{Tokens: 42, Requests: 3}) + epWithLoad.Put(attrconcurrency.InFlightLoadDataKey, &attrconcurrency.InFlightLoad{Tokens: 42, Requests: 3}) got = pl.readInFlightLoad(epWithLoad) assert.Equal(t, int64(42), got.tokens) assert.Equal(t, 3, got.requests) diff --git a/pkg/epp/framework/plugins/scheduling/filter/endpointattribute/filter.go b/pkg/epp/framework/plugins/scheduling/filter/endpointattribute/filter.go index 09befeed53..60adab66aa 100644 --- a/pkg/epp/framework/plugins/scheduling/filter/endpointattribute/filter.go +++ b/pkg/epp/framework/plugins/scheduling/filter/endpointattribute/filter.go @@ -164,7 +164,7 @@ func (f *EndpointAttributeFilter) Consumes() map[string]any { func (f *EndpointAttributeFilter) Filter(_ context.Context, _ *scheduling.InferenceRequest, endpoints []scheduling.Endpoint) []scheduling.Endpoint { filtered := make([]scheduling.Endpoint, 0, len(endpoints)) for _, endpoint := range endpoints { - value, ok := attrmetrics.ReadScalarMetricValue(endpoint, f.attribute) + value, ok := attrmetrics.ReadScalarMetricValue(endpoint, attrmetrics.ScalarMetricDataKey(f.attribute)) if !ok { if f.passOnMissing { filtered = append(filtered, endpoint) diff --git a/pkg/epp/framework/plugins/scheduling/filter/endpointattribute/filter_test.go b/pkg/epp/framework/plugins/scheduling/filter/endpointattribute/filter_test.go index 4ed017f878..a63345de1c 100644 --- a/pkg/epp/framework/plugins/scheduling/filter/endpointattribute/filter_test.go +++ b/pkg/epp/framework/plugins/scheduling/filter/endpointattribute/filter_test.go @@ -108,7 +108,7 @@ func TestEndpointAttributeFilterFactory(t *testing.T) { func newEndpointWithValue(value float64) scheduling.Endpoint { attrs := fwkdl.NewAttributes() - attrs.Put(testAttribute, attrmetrics.ScalarMetricValue(value)) + attrs.Put(attrmetrics.ScalarMetricDataKey(testAttribute), attrmetrics.ScalarMetricValue(value)) return scheduling.NewEndpoint(&fwkdl.EndpointMetadata{}, &fwkdl.Metrics{}, attrs) } diff --git a/pkg/epp/framework/plugins/scheduling/filter/prefixcacheaffinity/plugin.go b/pkg/epp/framework/plugins/scheduling/filter/prefixcacheaffinity/plugin.go index f6db22de18..3af5c3a3bd 100644 --- a/pkg/epp/framework/plugins/scheduling/filter/prefixcacheaffinity/plugin.go +++ b/pkg/epp/framework/plugins/scheduling/filter/prefixcacheaffinity/plugin.go @@ -225,7 +225,7 @@ func (p *Plugin) Consumes() fwkplugin.DataDependencies { } func (p *Plugin) prefixCacheScore(ep fwksched.Endpoint) float64 { - if raw, ok := ep.Get(p.prefixMatchDataKey.String()); ok { + if raw, ok := ep.Get(p.prefixMatchDataKey); ok { info := raw.(*attrprefix.PrefixCacheMatchInfo) if info.TotalBlocks() > 0 { score := float64(info.MatchBlocks()) / float64(info.TotalBlocks()) @@ -255,7 +255,7 @@ func (p *Plugin) bestTTFT(endpoints []fwksched.Endpoint) float64 { // the throughput path (no observed load). func (p *Plugin) endpointTTFT(ep fwksched.Endpoint) float64 { if p.config.usesLatencyPredictor() { - if raw, ok := ep.Get(p.latencyPredictionInfoDataKey.String()); ok { + if raw, ok := ep.Get(p.latencyPredictionInfoDataKey); ok { info := raw.(*attrlatency.LatencyPredictionInfo) return info.TTFT() } @@ -267,7 +267,7 @@ func (p *Plugin) endpointTTFT(ep fwksched.Endpoint) float64 { // inFlightTokens returns an endpoint's in-flight token count, or 0 when the // attribute is absent (no observed load). func (p *Plugin) inFlightTokens(ep fwksched.Endpoint) int64 { - if raw, ok := ep.Get(p.inFlightLoadDataKey.String()); ok { + if raw, ok := ep.Get(p.inFlightLoadDataKey); ok { if load, ok := raw.(*attrconcurrency.InFlightLoad); ok && load != nil { return load.Tokens } diff --git a/pkg/epp/framework/plugins/scheduling/filter/prefixcacheaffinity/plugin_test.go b/pkg/epp/framework/plugins/scheduling/filter/prefixcacheaffinity/plugin_test.go index b7f45a7fa3..4e3c95a2ed 100644 --- a/pkg/epp/framework/plugins/scheduling/filter/prefixcacheaffinity/plugin_test.go +++ b/pkg/epp/framework/plugins/scheduling/filter/prefixcacheaffinity/plugin_test.go @@ -39,13 +39,13 @@ func makeEndpoint(name string, prefixMatch int, ttft float64, tokens int64) fwks } ep := fwksched.NewEndpoint(meta, &fwkdl.Metrics{}, fwkdl.NewAttributes()) if prefixMatch >= 0 { - ep.Put(attrprefix.PrefixCacheMatchInfoDataKey.String(), attrprefix.NewPrefixCacheMatchInfo(prefixMatch, 100, 16)) + ep.Put(attrprefix.PrefixCacheMatchInfoDataKey, attrprefix.NewPrefixCacheMatchInfo(prefixMatch, 100, 16)) } if ttft >= 0 { - ep.Put(attrlatency.LatencyPredictionInfoDataKey.String(), attrlatency.NewLatencyPredictionInfo(true, true, 0, 0, ttft, 0, 0)) + ep.Put(attrlatency.LatencyPredictionInfoDataKey, attrlatency.NewLatencyPredictionInfo(true, true, 0, 0, ttft, 0, 0)) } if tokens >= 0 { - ep.Put(attrconcurrency.InFlightLoadDataKey.String(), &attrconcurrency.InFlightLoad{Tokens: tokens}) + ep.Put(attrconcurrency.InFlightLoadDataKey, &attrconcurrency.InFlightLoad{Tokens: tokens}) } return ep } diff --git a/pkg/epp/framework/plugins/scheduling/filter/sloheadroomtier/plugin.go b/pkg/epp/framework/plugins/scheduling/filter/sloheadroomtier/plugin.go index 383f9c5093..113f44c192 100644 --- a/pkg/epp/framework/plugins/scheduling/filter/sloheadroomtier/plugin.go +++ b/pkg/epp/framework/plugins/scheduling/filter/sloheadroomtier/plugin.go @@ -94,7 +94,7 @@ func (p *Plugin) Filter(ctx context.Context, _ *fwksched.InferenceRequest, endpo var positive, negative, noPrediction []fwksched.Endpoint for _, ep := range endpoints { - raw, ok := ep.Get(p.latencyPredictionInfoDataKey.String()) + raw, ok := ep.Get(p.latencyPredictionInfoDataKey) if !ok { noPrediction = append(noPrediction, ep) continue diff --git a/pkg/epp/framework/plugins/scheduling/filter/sloheadroomtier/plugin_test.go b/pkg/epp/framework/plugins/scheduling/filter/sloheadroomtier/plugin_test.go index 8363758dbe..087178f60d 100644 --- a/pkg/epp/framework/plugins/scheduling/filter/sloheadroomtier/plugin_test.go +++ b/pkg/epp/framework/plugins/scheduling/filter/sloheadroomtier/plugin_test.go @@ -35,7 +35,7 @@ func makeEndpoint(name string, ttftHeadroom, tpotHeadroom float64, hasPrediction } ep := fwksched.NewEndpoint(meta, &fwkdl.Metrics{}, fwkdl.NewAttributes()) if hasPrediction { - ep.Put(attrlatency.LatencyPredictionInfoDataKey.String(), attrlatency.NewLatencyPredictionInfo( + ep.Put(attrlatency.LatencyPredictionInfoDataKey, attrlatency.NewLatencyPredictionInfo( ttftHeadroom >= 0, tpotHeadroom >= 0, ttftHeadroom, tpotHeadroom, 100, 10, 0)) } return ep diff --git a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/disagg_profile_handler_test.go b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/disagg_profile_handler_test.go index 16c182504d..26ba3c540c 100644 --- a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/disagg_profile_handler_test.go +++ b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/disagg_profile_handler_test.go @@ -127,7 +127,7 @@ func injectPrefixCache(profileResults map[string]*scheduling.ProfileRunResult, c return } for _, ep := range res.TargetEndpoints { - ep.Put(attrprefix.PrefixCacheMatchInfoDataKey.String(), + ep.Put(attrprefix.PrefixCacheMatchInfoDataKey, attrprefix.NewPrefixCacheMatchInfo(cachedTokens, inputTokens, 1)) } } diff --git a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/pd_profile_handler_test.go b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/pd_profile_handler_test.go index efde5d9860..83cf37c6b6 100644 --- a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/pd_profile_handler_test.go +++ b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/pd_profile_handler_test.go @@ -337,7 +337,7 @@ func TestPdProfileHandler_Pick(t *testing.T) { for profileName, profileRes := range tt.profileResults { if profileName == defaultDecodeProfile && profileRes != nil { for _, pod := range profileRes.TargetEndpoints { - pod.Put(attrprefix.PrefixCacheMatchInfoDataKey.String(), + pod.Put(attrprefix.PrefixCacheMatchInfoDataKey, attrprefix.NewPrefixCacheMatchInfo(tt.cachedTokens, inputTokens, 1)) } } @@ -439,7 +439,7 @@ func TestPdProfileHandler_PickSeries(t *testing.T) { for profileName, profileRes := range profileResults { if profileName == defaultDecodeProfile && profileRes != nil { for _, endpoint := range profileRes.TargetEndpoints { - endpoint.Put(attrprefix.PrefixCacheMatchInfoDataKey.String(), + endpoint.Put(attrprefix.PrefixCacheMatchInfoDataKey, attrprefix.NewPrefixCacheMatchInfo(innerTest.cachedTokens, inputTokens, 1)) } } diff --git a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider.go b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider.go index b2be45f5b9..072cb69f1f 100644 --- a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider.go +++ b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider.go @@ -132,7 +132,7 @@ func (d *PrefixBasedPDDecider) disaggregate(ctx context.Context, request *schedu } // inspect the decode endpoint to disaggregate if prefill should run or not. // if the non-cached part is short enough - no disaggregation. - prefixInfoRaw, ok := endpoint.Get(attrprefix.PrefixCacheMatchInfoDataKey.String()) + prefixInfoRaw, ok := endpoint.Get(attrprefix.PrefixCacheMatchInfoDataKey) if !ok || prefixInfoRaw == nil { logger.Error(nil, "unable to read prefix cache state") return false diff --git a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider_test.go b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider_test.go index fad76a2426..aeafd7e64e 100644 --- a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider_test.go +++ b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider_test.go @@ -50,7 +50,7 @@ func makeTestEndpointBase() scheduling.Endpoint { func makeTestEndpoint(cachedTokens int) scheduling.Endpoint { ep := makeTestEndpointBase() - ep.Put(attrprefix.PrefixCacheMatchInfoDataKey.String(), + ep.Put(attrprefix.PrefixCacheMatchInfoDataKey, attrprefix.NewPrefixCacheMatchInfo(cachedTokens, testTotalTokens, testBlockSize)) return ep } @@ -424,7 +424,7 @@ func TestDisaggregate_UsesUnweightedCachedBlockCount(t *testing.T) { newEndpoint := func(info *attrprefix.PrefixCacheMatchInfo) scheduling.Endpoint { ep := makeTestEndpointBase() - ep.Put(attrprefix.PrefixCacheMatchInfoDataKey.String(), info) + ep.Put(attrprefix.PrefixCacheMatchInfoDataKey, info) return ep } // Exact token count via the tokenized-prompt path the decider reads. @@ -463,7 +463,7 @@ func TestDisaggregateWrongPrefixInfoType(t *testing.T) { ctx := utils.NewTestContext(t) ep := makeTestEndpointBase() - ep.Put(attrprefix.PrefixCacheMatchInfoDataKey.String(), ¬PrefixCacheMatchInfo{}) + ep.Put(attrprefix.PrefixCacheMatchInfoDataKey, ¬PrefixCacheMatchInfo{}) decider, err := NewPrefixBasedPDDecider(PrefixBasedPDDeciderConfig{NonCachedTokens: 5}) require.NoError(t, err) diff --git a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/scheduler_test.go b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/scheduler_test.go index b4d78ad991..17f255a6db 100644 --- a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/scheduler_test.go +++ b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/scheduler_test.go @@ -252,7 +252,7 @@ func TestPDSchedule(t *testing.T) { inputTokens := len(test.req.Body.Completions.Prompt.Raw) / averageCharactersPerToken for _, pod := range test.input { - pod.Put(attrprefix.PrefixCacheMatchInfoDataKey.String(), attrprefix.NewPrefixCacheMatchInfo(0, inputTokens, 1)) + pod.Put(attrprefix.PrefixCacheMatchInfoDataKey, attrprefix.NewPrefixCacheMatchInfo(0, inputTokens, 1)) } got, err := scheduler.Schedule(ctx, test.req, test.input) @@ -266,7 +266,7 @@ func TestPDSchedule(t *testing.T) { if test.wantRes2 != nil { // Checking the prefix match in the decode pod. // update number of cached tokens for the following schedule call for _, pod := range test.input { - pod.Put(attrprefix.PrefixCacheMatchInfoDataKey.String(), attrprefix.NewPrefixCacheMatchInfo(inputTokens, inputTokens, 1)) + pod.Put(attrprefix.PrefixCacheMatchInfoDataKey, attrprefix.NewPrefixCacheMatchInfo(inputTokens, inputTokens, 1)) } got, err = scheduler.Schedule(ctx, test.req, test.input) diff --git a/pkg/epp/framework/plugins/scheduling/scorer/activerequest/active_request.go b/pkg/epp/framework/plugins/scheduling/scorer/activerequest/active_request.go index 17324efd24..3a903a3cc2 100644 --- a/pkg/epp/framework/plugins/scheduling/scorer/activerequest/active_request.go +++ b/pkg/epp/framework/plugins/scheduling/scorer/activerequest/active_request.go @@ -187,7 +187,7 @@ func (s *ActiveRequest) Score(ctx context.Context, _ *scheduling.InferenceReques } func (s *ActiveRequest) requestCount(ctx context.Context, endpoint scheduling.Endpoint) int64 { - val, ok := endpoint.Get(s.inFlightLoadDataKey.String()) + val, ok := endpoint.Get(s.inFlightLoadDataKey) if !ok { return 0 } diff --git a/pkg/epp/framework/plugins/scheduling/scorer/activerequest/active_request_test.go b/pkg/epp/framework/plugins/scheduling/scorer/activerequest/active_request_test.go index ecb8396f4c..631e1f5a34 100644 --- a/pkg/epp/framework/plugins/scheduling/scorer/activerequest/active_request_test.go +++ b/pkg/epp/framework/plugins/scheduling/scorer/activerequest/active_request_test.go @@ -8,6 +8,7 @@ import ( k8stypes "k8s.io/apimachinery/pkg/types" "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" "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requestcontrol" "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling" attrconcurrency "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/concurrency" @@ -35,17 +36,17 @@ func newStubEndpoint(name string, queueSize int) *stubEndpoint { } } -func (f *stubEndpoint) GetMetadata() *datalayer.EndpointMetadata { return f.metadata } -func (f *stubEndpoint) UpdateMetadata(*datalayer.EndpointMetadata) {} -func (f *stubEndpoint) GetMetrics() *datalayer.Metrics { return f.metrics } -func (f *stubEndpoint) UpdateMetrics(*datalayer.Metrics) {} -func (f *stubEndpoint) GetAttributes() datalayer.AttributeMap { return f.attr } -func (f *stubEndpoint) String() string { return f.metadata.NamespacedName.String() } -func (f *stubEndpoint) Put(key string, val datalayer.Cloneable) { f.attr.Put(key, val) } -func (f *stubEndpoint) Get(key string) (datalayer.Cloneable, bool) { +func (f *stubEndpoint) GetMetadata() *datalayer.EndpointMetadata { return f.metadata } +func (f *stubEndpoint) UpdateMetadata(*datalayer.EndpointMetadata) {} +func (f *stubEndpoint) GetMetrics() *datalayer.Metrics { return f.metrics } +func (f *stubEndpoint) UpdateMetrics(*datalayer.Metrics) {} +func (f *stubEndpoint) GetAttributes() datalayer.AttributeMap { return f.attr } +func (f *stubEndpoint) String() string { return f.metadata.NamespacedName.String() } +func (f *stubEndpoint) Put(key fwkplugin.DataKey, val datalayer.Cloneable) { f.attr.Put(key, val) } +func (f *stubEndpoint) Get(key fwkplugin.DataKey) (datalayer.Cloneable, bool) { return f.attr.Get(key) } -func (f *stubEndpoint) Keys() []string { return f.attr.Keys() } +func (f *stubEndpoint) Keys() []fwkplugin.DataKey { return f.attr.Keys() } func (f *stubEndpoint) Clone() datalayer.AttributeMap { return f.attr.Clone() } func newTestEndpoint(name string, queueSize int) scheduling.Endpoint { @@ -54,7 +55,7 @@ func newTestEndpoint(name string, queueSize int) scheduling.Endpoint { func newTestEndpointWithLoad(name string, requests int64) scheduling.Endpoint { ep := newStubEndpoint(name, 0) - ep.Put(attrconcurrency.InFlightLoadDataKey.String(), &attrconcurrency.InFlightLoad{Requests: requests}) + ep.Put(attrconcurrency.InFlightLoadDataKey, &attrconcurrency.InFlightLoad{Requests: requests}) return ep } @@ -218,7 +219,7 @@ func TestActiveRequest_IdleThresholdAndMaxBusyScore(t *testing.T) { assert.Equal(t, 1.0, scores[podA]) assert.Equal(t, 1.0, scores[podB]) - podA.Put(attrconcurrency.InFlightLoadDataKey.String(), &attrconcurrency.InFlightLoad{Requests: 1}) + podA.Put(attrconcurrency.InFlightLoadDataKey, &attrconcurrency.InFlightLoad{Requests: 1}) scores = scorer.Score(ctx, nil, []scheduling.Endpoint{podA, podB}) assert.Equal(t, 0.0, scores[podA], "Busy pod scores 0.0 in binary mode") @@ -263,7 +264,7 @@ func TestActiveRequest_DefaultParamsProduceContinuousScores(t *testing.T) { func inFlightRequests(t *testing.T, endpoint scheduling.Endpoint) int64 { t.Helper() - val, ok := endpoint.Get(attrconcurrency.InFlightLoadDataKey.String()) + val, ok := endpoint.Get(attrconcurrency.InFlightLoadDataKey) require.True(t, ok) load, ok := val.(*attrconcurrency.InFlightLoad) require.True(t, ok) diff --git a/pkg/epp/framework/plugins/scheduling/scorer/endpointattribute/endpointattribute.go b/pkg/epp/framework/plugins/scheduling/scorer/endpointattribute/endpointattribute.go index b18697a1f3..d0f3223afd 100644 --- a/pkg/epp/framework/plugins/scheduling/scorer/endpointattribute/endpointattribute.go +++ b/pkg/epp/framework/plugins/scheduling/scorer/endpointattribute/endpointattribute.go @@ -155,7 +155,7 @@ func (s *EndpointAttributeScorer) Score(_ context.Context, _ *fwksched.Inference func (s *EndpointAttributeScorer) scoreFixedRange(endpoints []fwksched.Endpoint) map[fwksched.Endpoint]float64 { scores := make(map[fwksched.Endpoint]float64, len(endpoints)) for _, endpoint := range endpoints { - value, ok := attrmetrics.ReadScalarMetricValue(endpoint, s.attributeKey) + value, ok := attrmetrics.ReadScalarMetricValue(endpoint, attrmetrics.ScalarMetricDataKey(s.attributeKey)) if !ok { scores[endpoint] = 0.0 continue @@ -180,7 +180,7 @@ func (s *EndpointAttributeScorer) scoreAdaptiveRange(endpoints []fwksched.Endpoi maxValue := math.Inf(-1) for _, endpoint := range endpoints { - value, ok := attrmetrics.ReadScalarMetricValue(endpoint, s.attributeKey) + value, ok := attrmetrics.ReadScalarMetricValue(endpoint, attrmetrics.ScalarMetricDataKey(s.attributeKey)) if !ok { continue } diff --git a/pkg/epp/framework/plugins/scheduling/scorer/endpointattribute/endpointattribute_test.go b/pkg/epp/framework/plugins/scheduling/scorer/endpointattribute/endpointattribute_test.go index 09d65b81a9..502d08ca5d 100644 --- a/pkg/epp/framework/plugins/scheduling/scorer/endpointattribute/endpointattribute_test.go +++ b/pkg/epp/framework/plugins/scheduling/scorer/endpointattribute/endpointattribute_test.go @@ -105,7 +105,7 @@ func TestEndpointAttributeScorerFactory(t *testing.T) { func newEndpointWithValue(value float64) fwksched.Endpoint { attrs := fwkdl.NewAttributes() - attrs.Put(testAttributeKey, attrmetrics.ScalarMetricValue(value)) + attrs.Put(attrmetrics.ScalarMetricDataKey(testAttributeKey), attrmetrics.ScalarMetricValue(value)) return fwksched.NewEndpoint(&fwkdl.EndpointMetadata{}, &fwkdl.Metrics{}, attrs) } diff --git a/pkg/epp/framework/plugins/scheduling/scorer/latency/plugin.go b/pkg/epp/framework/plugins/scheduling/scorer/latency/plugin.go index 8845392123..281d2b2501 100644 --- a/pkg/epp/framework/plugins/scheduling/scorer/latency/plugin.go +++ b/pkg/epp/framework/plugins/scheduling/scorer/latency/plugin.go @@ -162,7 +162,7 @@ func (s *Plugin) Score(ctx context.Context, _ *fwksched.InferenceRequest, endpoi hasPredictions := false for _, ep := range endpoints { d := epData{endpoint: ep} - if raw, ok := ep.Get(s.latencyPredictionInfoDataKey.String()); ok { + if raw, ok := ep.Get(s.latencyPredictionInfoDataKey); ok { info := raw.(*attrlatency.LatencyPredictionInfo) d.info = info d.ttftHeadroom = info.TTFTHeadroom() @@ -393,7 +393,7 @@ func normalizedWeights(a, b float64) (float64, float64) { } func (s *Plugin) prefixCacheScore(ep fwksched.Endpoint) float64 { - if raw, ok := ep.Get(s.prefixMatchDataKey.String()); ok { + if raw, ok := ep.Get(s.prefixMatchDataKey); ok { info := raw.(*attrprefix.PrefixCacheMatchInfo) if info.TotalBlocks() > 0 { score := float64(info.MatchBlocks()) / float64(info.TotalBlocks()) diff --git a/pkg/epp/framework/plugins/scheduling/scorer/latency/plugin_test.go b/pkg/epp/framework/plugins/scheduling/scorer/latency/plugin_test.go index d3b1b794d8..f5ea58b406 100644 --- a/pkg/epp/framework/plugins/scheduling/scorer/latency/plugin_test.go +++ b/pkg/epp/framework/plugins/scheduling/scorer/latency/plugin_test.go @@ -40,7 +40,7 @@ func makeLatencyScorerEndpoint(name string, kvCache float64, queueSize, runningR } func setLatencyPrediction(ep fwksched.Endpoint, ttftValid, tpotValid bool, ttftHeadroom, tpotHeadroom, ttft, tpot float64) { - ep.Put(attrlatency.LatencyPredictionInfoDataKey.String(), + ep.Put(attrlatency.LatencyPredictionInfoDataKey, attrlatency.NewLatencyPredictionInfo(ttftValid, tpotValid, ttftHeadroom, tpotHeadroom, ttft, tpot, 0)) } @@ -140,9 +140,9 @@ func TestScoreIdlePodPreference(t *testing.T) { // Same predictions — both negative, same deficit. // Use dispatch count for idle detection (matches EPP internal queue). - epBusy.Put(attrlatency.LatencyPredictionInfoDataKey.String(), + epBusy.Put(attrlatency.LatencyPredictionInfoDataKey, attrlatency.NewLatencyPredictionInfo(false, false, -50, -10, 150, 40, 5)) - epIdle.Put(attrlatency.LatencyPredictionInfoDataKey.String(), + epIdle.Put(attrlatency.LatencyPredictionInfoDataKey, attrlatency.NewLatencyPredictionInfo(false, false, -50, -10, 150, 40, 0)) endpoints := []fwksched.Endpoint{epBusy, epIdle} diff --git a/pkg/epp/framework/plugins/scheduling/scorer/mmcacheaffinity/scorer.go b/pkg/epp/framework/plugins/scheduling/scorer/mmcacheaffinity/scorer.go index 0c8a013558..b9414eaabc 100644 --- a/pkg/epp/framework/plugins/scheduling/scorer/mmcacheaffinity/scorer.go +++ b/pkg/epp/framework/plugins/scheduling/scorer/mmcacheaffinity/scorer.go @@ -106,7 +106,7 @@ func (s *Scorer) Score(ctx context.Context, req *scheduling.InferenceRequest, en if meta := endpoint.GetMetadata(); meta != nil { pod = meta.PodName } - info, ok := endpoint.Get(s.mmMatchDataKey.String()) + info, ok := endpoint.Get(s.mmMatchDataKey) if !ok { traceLogger.Info("mm-embeddings-cache: no match info, score 0", "requestID", requestID, "pod", pod, "scorer", s.typedName) continue diff --git a/pkg/epp/framework/plugins/scheduling/scorer/mmcacheaffinity/scorer_test.go b/pkg/epp/framework/plugins/scheduling/scorer/mmcacheaffinity/scorer_test.go index 2547f492fe..e8fd320d55 100644 --- a/pkg/epp/framework/plugins/scheduling/scorer/mmcacheaffinity/scorer_test.go +++ b/pkg/epp/framework/plugins/scheduling/scorer/mmcacheaffinity/scorer_test.go @@ -53,9 +53,9 @@ func TestScoreFromProducedMatchInfo(t *testing.T) { endpointB := newEndpoint("pod-b") endpointC := newEndpoint("pod-c") requestItems := []attrmm.MatchItem{{Hash: "image", Size: 80}, {Hash: "icon", Size: 20}} - endpointA.Put(attrmm.EncoderCacheMatchInfoKey.String(), attrmm.NewEncoderCacheMatchInfo([]attrmm.MatchItem{{Hash: "image", Size: 80}}, requestItems)) - endpointB.Put(attrmm.EncoderCacheMatchInfoKey.String(), attrmm.NewEncoderCacheMatchInfo([]attrmm.MatchItem{{Hash: "icon", Size: 20}}, requestItems)) - endpointC.Put(attrmm.EncoderCacheMatchInfoKey.String(), attrmm.NewEncoderCacheMatchInfo(nil, requestItems)) + endpointA.Put(attrmm.EncoderCacheMatchInfoKey, attrmm.NewEncoderCacheMatchInfo([]attrmm.MatchItem{{Hash: "image", Size: 80}}, requestItems)) + endpointB.Put(attrmm.EncoderCacheMatchInfoKey, attrmm.NewEncoderCacheMatchInfo([]attrmm.MatchItem{{Hash: "icon", Size: 20}}, requestItems)) + endpointC.Put(attrmm.EncoderCacheMatchInfoKey, attrmm.NewEncoderCacheMatchInfo(nil, requestItems)) scores := scorer.Score(context.Background(), nil, []scheduling.Endpoint{endpointA, endpointB, endpointC}) @@ -68,7 +68,7 @@ func TestScoreMissingOrInvalidMatchInfoReturnsZero(t *testing.T) { scorer := New(testName, "") endpointA := newEndpoint("pod-a") endpointB := newEndpoint("pod-b") - endpointB.Put(attrmm.EncoderCacheMatchInfoKey.String(), attrmm.NewEncoderCacheMatchInfo([]attrmm.MatchItem{{Hash: "image", Size: 1}}, nil)) + endpointB.Put(attrmm.EncoderCacheMatchInfoKey, attrmm.NewEncoderCacheMatchInfo([]attrmm.MatchItem{{Hash: "image", Size: 1}}, nil)) scores := scorer.Score(context.Background(), nil, []scheduling.Endpoint{endpointA, endpointB}) diff --git a/pkg/epp/framework/plugins/scheduling/scorer/nohitlru/no_hit_lru.go b/pkg/epp/framework/plugins/scheduling/scorer/nohitlru/no_hit_lru.go index 4655826de1..7f1af45e8e 100644 --- a/pkg/epp/framework/plugins/scheduling/scorer/nohitlru/no_hit_lru.go +++ b/pkg/epp/framework/plugins/scheduling/scorer/nohitlru/no_hit_lru.go @@ -169,7 +169,7 @@ func (s *NoHitLRU) isColdRequest(ctx context.Context, endpoints []scheduling.End logger := log.FromContext(ctx).V(logging.DEBUG) for _, ep := range endpoints { - attr, ok := ep.Get(s.dk.String()) + attr, ok := ep.Get(s.dk) if !ok { continue } diff --git a/pkg/epp/framework/plugins/scheduling/scorer/nohitlru/no_hit_lru_test.go b/pkg/epp/framework/plugins/scheduling/scorer/nohitlru/no_hit_lru_test.go index b41345b091..732cbd5f5c 100644 --- a/pkg/epp/framework/plugins/scheduling/scorer/nohitlru/no_hit_lru_test.go +++ b/pkg/epp/framework/plugins/scheduling/scorer/nohitlru/no_hit_lru_test.go @@ -92,7 +92,7 @@ func newColdNS(name string) scheduling.Endpoint { // newWarm returns an endpoint with prefix-cache match info indicating a cache hit. func newWarm(name string) scheduling.Endpoint { ep := newCold(name) - ep.Put(attrprefix.PrefixCacheMatchInfoDataKey.String(), attrprefix.NewPrefixCacheMatchInfo(5, 10, 1)) + ep.Put(attrprefix.PrefixCacheMatchInfoDataKey, attrprefix.NewPrefixCacheMatchInfo(5, 10, 1)) return ep } @@ -287,7 +287,7 @@ func TestNoHitLRUPreferLeastRecentlyUsedAfterColdRequests(t *testing.T) { &fwkdl.Metrics{}, nil, ) - w.Put(attrprefix.PrefixCacheMatchInfoDataKey.String(), attrprefix.NewPrefixCacheMatchInfo(5, 10, 1)) + w.Put(attrprefix.PrefixCacheMatchInfoDataKey, attrprefix.NewPrefixCacheMatchInfo(5, 10, 1)) return []scheduling.Endpoint{w, endpointB, endpointC} } diff --git a/pkg/epp/framework/plugins/scheduling/scorer/preciseprefixcache/precise_prefix_cache_test.go b/pkg/epp/framework/plugins/scheduling/scorer/preciseprefixcache/precise_prefix_cache_test.go index 5d81518b66..a937558eef 100644 --- a/pkg/epp/framework/plugins/scheduling/scorer/preciseprefixcache/precise_prefix_cache_test.go +++ b/pkg/epp/framework/plugins/scheduling/scorer/preciseprefixcache/precise_prefix_cache_test.go @@ -226,7 +226,7 @@ func TestLegacyProducer_TokensFlowToEndpointAttribute(t *testing.T) { require.NoError(t, lp.Produce(ctx, req, []scheduling.Endpoint{endpoint})) - key := attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName("inner").String() + key := attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName("inner") raw, ok := endpoint.Get(key) require.True(t, ok, "endpoint should have PrefixCacheMatchInfo set") info, ok := raw.(*attrprefix.PrefixCacheMatchInfo) diff --git a/pkg/epp/framework/plugins/scheduling/scorer/prefix/plugin.go b/pkg/epp/framework/plugins/scheduling/scorer/prefix/plugin.go index dbbca58faf..6a4b7d2453 100644 --- a/pkg/epp/framework/plugins/scheduling/scorer/prefix/plugin.go +++ b/pkg/epp/framework/plugins/scheduling/scorer/prefix/plugin.go @@ -135,7 +135,7 @@ func (p *Plugin) Score(ctx context.Context, _ *fwksched.InferenceRequest, endpoi for _, endpoint := range endpoints { // Default to score 0 if PrefixCacheMatchInfo is missing or invalid. scores[endpoint] = 0.0 - info, ok := endpoint.Get(p.prefixMatchDataKey.String()) + info, ok := endpoint.Get(p.prefixMatchDataKey) if !ok { logger.V(logutil.DEFAULT).Error(nil, "PrefixCacheMatchInfo not found for endpoint, assigning score 0", "endpoint", endpoint, "key", p.prefixMatchDataKey.String()) continue diff --git a/pkg/epp/framework/plugins/scheduling/scorer/prefix/plugin_test.go b/pkg/epp/framework/plugins/scheduling/scorer/prefix/plugin_test.go index c7918dacf2..df92ea1b17 100644 --- a/pkg/epp/framework/plugins/scheduling/scorer/prefix/plugin_test.go +++ b/pkg/epp/framework/plugins/scheduling/scorer/prefix/plugin_test.go @@ -35,7 +35,7 @@ func TestPrefixPluginScore(t *testing.T) { producerName := "approx-prefix-cache-producer" p, _ := New(context.Background(), PrefixCacheScorerPluginType, producerName) - key := attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(producerName).String() + key := attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(producerName) endpoint1 := fwksched.NewEndpoint(&fwkdl.EndpointMetadata{NamespacedName: k8stypes.NamespacedName{Name: "pod1"}}, fwkdl.NewMetrics(), nil) endpoint1.Put(key, attrprefix.NewPrefixCacheMatchInfo(5, 10, 1)) @@ -57,7 +57,7 @@ func TestPrefixPluginScoreWithWeights(t *testing.T) { p.matchLengthWeight = 0.5 p.matchLengthScaleTokens = 100 - key := attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(producerName).String() + key := attrprefix.PrefixCacheMatchInfoDataKey.WithNonEmptyProducerName(producerName) // Endpoint 1: match 5, total 10, block size 1 // matchRatio = 5/10 = 0.5 diff --git a/pkg/epp/framework/plugins/scheduling/scorer/tokenload/token_load.go b/pkg/epp/framework/plugins/scheduling/scorer/tokenload/token_load.go index 1b272f18e0..6c739cd390 100644 --- a/pkg/epp/framework/plugins/scheduling/scorer/tokenload/token_load.go +++ b/pkg/epp/framework/plugins/scheduling/scorer/tokenload/token_load.go @@ -101,12 +101,12 @@ func (s *TokenLoadScorer) Score(ctx context.Context, _ *fwksched.InferenceReques // Read both accumulated in-flight load and the projected impact of the // request being scored, which are now carried on separate attributes. var tokens int64 - if val, ok := endpoint.Get(s.inFlightLoadDataKey.String()); ok { + if val, ok := endpoint.Get(s.inFlightLoadDataKey); ok { if load, ok := val.(*attrconcurrency.InFlightLoad); ok && load != nil { tokens += load.Tokens } } - if val, ok := endpoint.Get(s.uncachedRequestTokensDataKey.String()); ok { + if val, ok := endpoint.Get(s.uncachedRequestTokensDataKey); ok { if uncached, ok := val.(*attrconcurrency.UncachedRequestTokens); ok && uncached != nil { tokens += uncached.Tokens } diff --git a/pkg/epp/framework/plugins/scheduling/scorer/tokenload/token_load_test.go b/pkg/epp/framework/plugins/scheduling/scorer/tokenload/token_load_test.go index 4c172bc3ea..c753d39f33 100644 --- a/pkg/epp/framework/plugins/scheduling/scorer/tokenload/token_load_test.go +++ b/pkg/epp/framework/plugins/scheduling/scorer/tokenload/token_load_test.go @@ -52,9 +52,9 @@ func TestTokenLoadScorer(t *testing.T) { // pod1: 0 in-flight, no current-request impact. Score = 1 - 0/1000 = 1.0 // pod2: 500 in-flight, no current-request impact. Score = 1 - 500/1000 = 0.5 - endpoints[1].Put(scorer.inFlightLoadDataKey.String(), &attrconcurrency.InFlightLoad{Tokens: 500}) + endpoints[1].Put(scorer.inFlightLoadDataKey, &attrconcurrency.InFlightLoad{Tokens: 500}) // pod3: 1000 in-flight, no current-request impact. Score = 1 - 1000/1000 = 0.0 - endpoints[2].Put(scorer.inFlightLoadDataKey.String(), &attrconcurrency.InFlightLoad{Tokens: 1000}) + endpoints[2].Put(scorer.inFlightLoadDataKey, &attrconcurrency.InFlightLoad{Tokens: 1000}) scores := scorer.Score(context.Background(), &fwksched.InferenceRequest{}, endpoints) @@ -75,16 +75,16 @@ func TestTokenLoadScorer(t *testing.T) { } // pod1: 0 in-flight + 250 current = 250. Score = 1 - 250/1000 = 0.75 - endpoints[0].Put(scorer.inFlightLoadDataKey.String(), &attrconcurrency.InFlightLoad{Tokens: 0}) - endpoints[0].Put(scorer.uncachedRequestTokensDataKey.String(), &attrconcurrency.UncachedRequestTokens{Tokens: 250}) + endpoints[0].Put(scorer.inFlightLoadDataKey, &attrconcurrency.InFlightLoad{Tokens: 0}) + endpoints[0].Put(scorer.uncachedRequestTokensDataKey, &attrconcurrency.UncachedRequestTokens{Tokens: 250}) // pod2: 250 in-flight + 250 current = 500. Score = 1 - 500/1000 = 0.5 - endpoints[1].Put(scorer.inFlightLoadDataKey.String(), &attrconcurrency.InFlightLoad{Tokens: 250}) - endpoints[1].Put(scorer.uncachedRequestTokensDataKey.String(), &attrconcurrency.UncachedRequestTokens{Tokens: 250}) + endpoints[1].Put(scorer.inFlightLoadDataKey, &attrconcurrency.InFlightLoad{Tokens: 250}) + endpoints[1].Put(scorer.uncachedRequestTokensDataKey, &attrconcurrency.UncachedRequestTokens{Tokens: 250}) // pod3: 750 in-flight + 250 current = 1000. Score = 1 - 1000/1000 = 0.0 - endpoints[2].Put(scorer.inFlightLoadDataKey.String(), &attrconcurrency.InFlightLoad{Tokens: 750}) - endpoints[2].Put(scorer.uncachedRequestTokensDataKey.String(), &attrconcurrency.UncachedRequestTokens{Tokens: 250}) + endpoints[2].Put(scorer.inFlightLoadDataKey, &attrconcurrency.InFlightLoad{Tokens: 750}) + endpoints[2].Put(scorer.uncachedRequestTokensDataKey, &attrconcurrency.UncachedRequestTokens{Tokens: 250}) scores := scorer.Score(context.Background(), &fwksched.InferenceRequest{}, endpoints) @@ -111,9 +111,9 @@ func TestTokenLoadScorer(t *testing.T) { } var nilLoad *attrconcurrency.InFlightLoad - endpoints[0].Put(scorer.inFlightLoadDataKey.String(), nilLoad) + endpoints[0].Put(scorer.inFlightLoadDataKey, nilLoad) var nilUncached *attrconcurrency.UncachedRequestTokens - endpoints[0].Put(scorer.uncachedRequestTokensDataKey.String(), nilUncached) + endpoints[0].Put(scorer.uncachedRequestTokensDataKey, nilUncached) // Typed nil; should score as if 0 token load (no panic) scores := scorer.Score(context.Background(), &fwksched.InferenceRequest{}, endpoints) diff --git a/pkg/epp/requestcontrol/director.go b/pkg/epp/requestcontrol/director.go index 4e33e980f0..8e5ad5dbc8 100644 --- a/pkg/epp/requestcontrol/director.go +++ b/pkg/epp/requestcontrol/director.go @@ -89,7 +89,7 @@ func primaryEndpointHasCachedPrefix(logger logr.Logger, result *fwksched.Schedul debug.Info("conditional-decode: primary endpoint is nil") return false } - raw, ok := endpoint.Get(attrprefix.PrefixCacheMatchInfoDataKey.String()) + raw, ok := endpoint.Get(attrprefix.PrefixCacheMatchInfoDataKey) if !ok || raw == nil { debug.Info("conditional-decode: endpoint has no prefix-cache match attribute (no scorer attached?)") return false diff --git a/pkg/epp/requestcontrol/director_test.go b/pkg/epp/requestcontrol/director_test.go index d1b2677510..280eaaab0a 100644 --- a/pkg/epp/requestcontrol/director_test.go +++ b/pkg/epp/requestcontrol/director_test.go @@ -86,7 +86,7 @@ type mockScheduler struct { func (m *mockScheduler) Schedule(_ context.Context, _ *fwksched.InferenceRequest, endpoints []fwksched.Endpoint) (*fwksched.SchedulingResult, error) { if endpoints != nil && m.dataProduced { - data, ok := endpoints[0].Get(mockProducedDataKey.String()) + data, ok := endpoints[0].Get(mockProducedDataKey) if !ok || data.(mockProducedDataType).value != 42 { return nil, errors.New("expected produced data not found in pod") } @@ -135,7 +135,7 @@ func (m *mockDataProducerPlugin) Consumes() fwkplugin.DataDependencies { } func (m *mockDataProducerPlugin) Produce(ctx context.Context, request *fwksched.InferenceRequest, endpoints []fwksched.Endpoint) error { - endpoints[0].Put(mockProducedDataKey.String(), mockProducedDataType{value: 42}) + endpoints[0].Put(mockProducedDataKey, mockProducedDataType{value: 42}) return nil } @@ -1804,7 +1804,7 @@ func (w wrongTypeAttr) Clone() fwkdl.Cloneable { return w } func TestPrimaryEndpointHasCachedPrefix(t *testing.T) { endpointWith := func(matched, total int) fwksched.Endpoint { attrs := fwkdl.NewAttributes() - attrs.Put(attrprefix.PrefixCacheMatchInfoDataKey.String(), + attrs.Put(attrprefix.PrefixCacheMatchInfoDataKey, attrprefix.NewPrefixCacheMatchInfo(matched, total, 1)) return fwksched.NewEndpoint( &fwkdl.EndpointMetadata{NamespacedName: types.NamespacedName{Namespace: "default", Name: "p"}}, @@ -1819,7 +1819,7 @@ func TestPrimaryEndpointHasCachedPrefix(t *testing.T) { } endpointWithWrongType := func() fwksched.Endpoint { attrs := fwkdl.NewAttributes() - attrs.Put(attrprefix.PrefixCacheMatchInfoDataKey.String(), wrongTypeAttr{}) + attrs.Put(attrprefix.PrefixCacheMatchInfoDataKey, wrongTypeAttr{}) return fwksched.NewEndpoint( &fwkdl.EndpointMetadata{NamespacedName: types.NamespacedName{Namespace: "default", Name: "p"}}, nil, attrs, @@ -1907,7 +1907,7 @@ func TestDirector_HandleRequest_ConditionalDecode(t *testing.T) { scheduleResultWith := func(matched, total int) *fwksched.SchedulingResult { attrs := fwkdl.NewAttributes() if matched >= 0 { - attrs.Put(attrprefix.PrefixCacheMatchInfoDataKey.String(), + attrs.Put(attrprefix.PrefixCacheMatchInfoDataKey, attrprefix.NewPrefixCacheMatchInfo(matched, total, 1)) } return &fwksched.SchedulingResult{