refactor(datalayer): key AttributeMap by DataKey instead of string

Signed-off-by: satyamg1620 <Satyam.Gupta.3@ibm.com>

Author: satyamg1620

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

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)

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,38 +49,41 @@ 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")
 
 	dv, ok := got.(*dummy)
 	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,18 +137,19 @@ 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)
 
 	// Mutate the source value
 	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)
 }

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
 }
 

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)

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)
 }

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)
 }

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
 }