Implement dynamic attributes for EPP AttributeMap

Introduce dynamic attributes (getter providers) to the AttributeMap
implementation. This allows retrieving live in-flight load data
directly from InFlightLoadProducer's trackers, resolving the issue
where concurrency-detector always saw 0 load during admission control.

Signed-off-by: Luke Van Drie <lukevandrie@google.com>

Author: LukeAVanDrie

cmd/epp/runner/runner.go

@@ -664,6 +664,24 @@ func (r *Runner) parseConfigurationPhaseTwo(ctx context.Context, rawConfig *conf
 	r.parsers = handlers.NewParsers(cfg.ParserConfig)
 	logger.Info("loaded configuration from file/text successfully")
 
+	var loadProducer *inflightload.InFlightLoadProducer
+	for _, p := range handle.GetAllPlugins() {
+		if prod, ok := p.(*inflightload.InFlightLoadProducer); ok {
+			loadProducer = prod
+			break
+		}
+	}
+
+	if loadProducer != nil {
+		key := attrconcurrency.InFlightLoadDataKey.WithNonEmptyProducerName(loadProducer.TypedName().Name)
+		r.dlRuntime.RegisterAttributeProvider(key.String(), func(eid string) fwkdl.Cloneable {
+			return &attrconcurrency.InFlightLoad{
+				Tokens:   loadProducer.GetTokens(eid),
+				Requests: loadProducer.GetRequests(eid),
+			}
+		})
+	}
+
 	return cfg, nil
 }
 

pkg/epp/datalayer/runtime.go

@@ -59,6 +59,9 @@ type Runtime struct {
 
 	collectors *collectorManager // per-endpoint poller, keyed by namespaced name
 	logger     logr.Logger       // Set in Configure; used where no context is available (e.g. ReleaseEndpoint).
+
+	provMu    sync.RWMutex
+	providers map[string]func(endpointID string) fwkdl.Cloneable
 }
 
 const (
@@ -80,6 +83,7 @@ func NewRuntime(pollingInterval time.Duration) *Runtime {
 		extractors:      newExtractorMap(),
 		collectors:      newCollectorManager(),
 		logger:          logr.Discard(),
+		providers:       make(map[string]func(endpointID string) fwkdl.Cloneable),
 	}
 }
 
@@ -222,6 +226,16 @@ func (r *Runtime) Register(reg fwkdl.PendingRegistration) error {
 	return nil
 }
 
+// RegisterAttributeProvider registers a provider function for a dynamic attribute.
+func (r *Runtime) RegisterAttributeProvider(key string, provider func(endpointID string) fwkdl.Cloneable) {
+	r.provMu.Lock()
+	defer r.provMu.Unlock()
+	if r.providers == nil {
+		r.providers = make(map[string]func(endpointID string) fwkdl.Cloneable)
+	}
+	r.providers[key] = provider
+}
+
 // registerSource dispatches src to the matching variant manager. g enforces
 // per-Configure-call GVK uniqueness for NotificationSources. src may be a
 // PollingDispatcher (not a DataSource), so the parameter is plugin.Plugin.
@@ -371,18 +385,29 @@ func (r *Runtime) NewEndpoint(ctx context.Context, endpointMetadata *fwkdl.Endpo
 	logger, _ := logr.FromContext(ctx)
 	logger = logger.WithValues("endpoint", endpointMetadata.GetNamespacedName())
 
+	endpoint := fwkdl.NewEndpoint(endpointMetadata, nil)
+	eid := endpointMetadata.GetNamespacedName().String()
+
+	r.provMu.RLock()
+	for key, provider := range r.providers {
+		endpoint.GetAttributes().Put(key, &fwkdl.DynamicAttribute{
+			Get: func() fwkdl.Cloneable {
+				return provider(eid)
+			},
+		})
+	}
+	r.provMu.RUnlock()
+
 	dispatchers := make([]fwkdl.PollingDispatcher, 0, r.dispatchers.Count())
 	for _, d := range r.dispatchers.Dispatchers() {
 		dispatchers = append(dispatchers, d)
 	}
 	if len(dispatchers) == 0 {
 		logger.Info("No polling sources configured, creating endpoint without collector")
-		endpoint := fwkdl.NewEndpoint(endpointMetadata, nil)
 		r.dispatchEndpointEvent(ctx, logger, fwkdl.EndpointEvent{Type: fwkdl.EventAddOrUpdate, Endpoint: endpoint})
 		return endpoint
 	}
 
-	endpoint := fwkdl.NewEndpoint(endpointMetadata, nil)
 	collector := NewCollector()
 
 	key := endpointMetadata.GetNamespacedName()

pkg/epp/datalayer/runtime_endpoint_dispatch_test.go

@@ -30,6 +30,17 @@ import (
 	"github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/source/notifications"
 )
 
+type dummy struct {
+	Text string
+}
+
+func (d *dummy) Clone() fwkdl.Cloneable {
+	if d == nil {
+		return nil
+	}
+	return &dummy{Text: d.Text}
+}
+
 // TestNewEndpointDispatchesEventWithNoPollers verifies that endpoint lifecycle
 // events are dispatched to EndpointSource even when no PollingDataSource is configured.
 // Regression test for: endpoint-notification-source silently drops events when
@@ -101,3 +112,26 @@ func TestUpdateEndpointDispatchesEvent(t *testing.T) {
 	assert.Equal(t, fwkdl.EventAddOrUpdate, events[0].Type)
 	assert.Equal(t, "5.6.7.8", events[0].Endpoint.GetMetadata().Address)
 }
+
+func TestRuntimeRegisterAttributeProvider(t *testing.T) {
+	r := NewRuntime(1)
+
+	r.RegisterAttributeProvider("test-key", func(endpointID string) fwkdl.Cloneable {
+		return &dummy{Text: endpointID + "-value"}
+	})
+
+	pod := &fwkdl.EndpointMetadata{
+		NamespacedName: types.NamespacedName{Name: "pod1", Namespace: "default"},
+		Address:        "1.2.3.4",
+	}
+
+	endpoint := r.NewEndpoint(context.Background(), pod)
+	assert.NotNil(t, endpoint)
+
+	val, ok := endpoint.GetAttributes().Get("test-key")
+	assert.True(t, ok)
+
+	dVal, ok := val.(*dummy)
+	assert.True(t, ok)
+	assert.Equal(t, "default/pod1-value", dVal.Text)
+}

pkg/epp/framework/interface/datalayer/attributemap.go

@@ -25,6 +25,19 @@ type Cloneable interface {
 	Clone() Cloneable
 }
 
+// DynamicAttribute wraps a getter function to allow on-demand resolution of attributes.
+type DynamicAttribute struct {
+	Get func() Cloneable
+}
+
+// Clone implements Cloneable. It copies the function pointer.
+func (d *DynamicAttribute) Clone() Cloneable {
+	if d == nil {
+		return nil
+	}
+	return &DynamicAttribute{Get: d.Get}
+}
+
 // AttributeMap is used to store flexible metadata or traits
 // across different aspects of an inference server.
 // Stored values must be Cloneable.
@@ -52,12 +65,21 @@ func (a *Attributes) Put(key string, value Cloneable) {
 	}
 }
 
-// Get retrieves an attribute by key, returning a cloned copy.
+// Get retrieves an attribute by key, returning a cloned copy (or resolving it dynamically).
 func (a *Attributes) Get(key string) (Cloneable, bool) {
 	val, ok := a.data.Load(key)
 	if !ok {
 		return nil, false
 	}
+
+	if dynamic, ok := val.(*DynamicAttribute); ok {
+		realVal := dynamic.Get()
+		if realVal == nil {
+			return nil, false
+		}
+		return realVal.Clone(), true
+	}
+
 	if cloneable, ok := val.(Cloneable); ok {
 		return cloneable.Clone(), true
 	}

pkg/epp/framework/interface/datalayer/attributemap_test.go

@@ -102,3 +102,29 @@ func TestReadAttribute(t *testing.T) {
 	assert.Nil(t, other) // zero value of pointer is nil
 
 }
+
+func TestDynamicAttribute(t *testing.T) {
+	attrs := NewAttributes()
+
+	val := &dummy{"live"}
+	dynamic := &DynamicAttribute{
+		Get: func() Cloneable {
+			return val
+		},
+	}
+
+	attrs.Put("dynamic_key", dynamic)
+
+	// First read
+	got, ok := attrs.Get("dynamic_key")
+	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")
+	assert.True(t, ok)
+	assert.Equal(t, "changed", got2.(*dummy).Text)
+}

pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/producer.go

@@ -454,6 +454,14 @@ func (p *InFlightLoadProducer) DeleteEndpoint(endpointID string) {
 	p.tokenTracker.delete(endpointID)
 }
 
+func (p *InFlightLoadProducer) GetTokens(eid string) int64 {
+	return p.tokenTracker.get(eid)
+}
+
+func (p *InFlightLoadProducer) GetRequests(eid string) int64 {
+	return p.requestTracker.get(eid)
+}
+
 // concurrencyTracker manages thread-safe counters for inflight requests.
 type concurrencyTracker struct {
 	mu     sync.RWMutex