Enable pluggable datalayer as experimental feature (#1391)

* enable global metrics logging

Signed-off-by: Etai Lev Ran <[email protected]>

* enable v2 data layer

Signed-off-by: Etai Lev Ran <[email protected]>

* move feature flag to cmd

Signed-off-by: Etai Lev Ran <[email protected]>

* remove logger from collector

Signed-off-by: Etai Lev Ran <[email protected]>

* return a typed data source using generics

Signed-off-by: Etai Lev Ran <[email protected]>

* document client configuration constants

Signed-off-by: Etai Lev Ran <[email protected]>

* retrieve a typed data source

Signed-off-by: Etai Lev Ran <[email protected]>

* use structured logging

Signed-off-by: Etai Lev Ran <[email protected]>

* commandline specified port value has precedence

Signed-off-by: Etai Lev Ran <[email protected]>

* pass data sources in ctor

Signed-off-by: Etai Lev Ran <[email protected]>

* address remaining review comments

Signed-off-by: Etai Lev Ran <[email protected]>

---------

Signed-off-by: Etai Lev Ran <[email protected]>

Author: elevran

cmd/epp/runner/runner.go

@@ -46,6 +46,8 @@ import (
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/common"
 	backendmetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/metrics"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/config/loader"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datalayer"
+	dlmetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datalayer/metrics"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datastore"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metrics"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metrics/collectors"
@@ -59,10 +61,17 @@ import (
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/framework/plugins/scorer"
 	testfilter "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/framework/plugins/test/filter"
 	runserver "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/server"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/env"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 	"sigs.k8s.io/gateway-api-inference-extension/version"
 )
 
+const (
+	// enableExperimentalDatalayerV2 defines the environment variable
+	// used as feature flag for the pluggable data layer.
+	enableExperimentalDatalayerV2 = "ENABLE_EXPERIMENTAL_DATALAYER_V2"
+)
+
 var (
 	grpcPort = flag.Int(
 		"grpc-port",
@@ -245,40 +254,12 @@ func (r *Runner) Run(ctx context.Context) error {
 	}
 
 	// --- Setup Datastore ---
-	mapping, err := backendmetrics.NewMetricMapping(
-		*totalQueuedRequestsMetric,
-		*kvCacheUsagePercentageMetric,
-		*loraInfoMetric,
-	)
+	useDatalayerV2 := env.GetEnvBool(enableExperimentalDatalayerV2, false, setupLog)
+	epf, err := r.setupMetricsCollection(setupLog, useDatalayerV2)
 	if err != nil {
-		setupLog.Error(err, "Failed to create metric mapping from flags.")
 		return err
 	}
-	verifyMetricMapping(*mapping, setupLog)
-
-	var metricsHttpClient *http.Client
-	if *modelServerMetricsScheme == "https" {
-		metricsHttpClient = &http.Client{
-			Transport: &http.Transport{
-				TLSClientConfig: &tls.Config{
-					InsecureSkipVerify: *modelServerMetricsHttpsInsecureSkipVerify,
-				},
-			},
-		}
-	} else {
-		metricsHttpClient = http.DefaultClient
-	}
-
-	pmf := backendmetrics.NewPodMetricsFactory(&backendmetrics.PodMetricsClientImpl{
-		MetricMapping:            mapping,
-		ModelServerMetricsPort:   int32(*modelServerMetricsPort),
-		ModelServerMetricsPath:   *modelServerMetricsPath,
-		ModelServerMetricsScheme: *modelServerMetricsScheme,
-		Client:                   metricsHttpClient,
-	},
-		*refreshMetricsInterval)
-
-	datastore := datastore.NewDatastore(ctx, pmf)
+	datastore := datastore.NewDatastore(ctx, epf)
 
 	// --- Setup Metrics Server ---
 	customCollectors := []prometheus.Collector{collectors.NewInferencePoolMetricsCollector(datastore)}
@@ -371,6 +352,7 @@ func (r *Runner) Run(ctx context.Context) error {
 		MetricsStalenessThreshold:        *metricsStalenessThreshold,
 		Director:                         director,
 		SaturationDetector:               saturationDetector,
+		UseExperimentalDatalayerV2:       useDatalayerV2, // pluggable data layer feature flag
 	}
 	if err := serverRunner.SetupWithManager(ctx, mgr); err != nil {
 		setupLog.Error(err, "Failed to setup EPP controllers")
@@ -446,6 +428,81 @@ func (r *Runner) parsePluginsConfiguration(ctx context.Context) error {
 	return nil
 }
 
+func (r *Runner) setupMetricsCollection(setupLog logr.Logger, useExperimentalDatalayer bool) (datalayer.EndpointFactory, error) {
+	if useExperimentalDatalayer {
+		return setupDatalayer()
+	}
+
+	if len(datalayer.GetSources()) != 0 {
+		setupLog.Info("data sources registered but pluggable datalayer is disabled")
+	}
+	return setupMetricsV1(setupLog)
+}
+
+func setupMetricsV1(setupLog logr.Logger) (datalayer.EndpointFactory, error) {
+	mapping, err := backendmetrics.NewMetricMapping(
+		*totalQueuedRequestsMetric,
+		*kvCacheUsagePercentageMetric,
+		*loraInfoMetric,
+	)
+	if err != nil {
+		setupLog.Error(err, "Failed to create metric mapping from flags.")
+		return nil, err
+	}
+	verifyMetricMapping(*mapping, setupLog)
+
+	var metricsHttpClient *http.Client
+	if *modelServerMetricsScheme == "https" {
+		metricsHttpClient = &http.Client{
+			Transport: &http.Transport{
+				TLSClientConfig: &tls.Config{
+					InsecureSkipVerify: *modelServerMetricsHttpsInsecureSkipVerify,
+				},
+			},
+		}
+	} else {
+		metricsHttpClient = http.DefaultClient
+	}
+
+	pmf := backendmetrics.NewPodMetricsFactory(&backendmetrics.PodMetricsClientImpl{
+		MetricMapping:            mapping,
+		ModelServerMetricsPort:   int32(*modelServerMetricsPort),
+		ModelServerMetricsPath:   *modelServerMetricsPath,
+		ModelServerMetricsScheme: *modelServerMetricsScheme,
+		Client:                   metricsHttpClient,
+	},
+		*refreshMetricsInterval)
+	return pmf, nil
+}
+
+func setupDatalayer() (datalayer.EndpointFactory, error) {
+	// create and register a metrics data source and extractor. In the future,
+	// data sources and extractors might be configured via a file. Once done,
+	// this (and registering the sources with the endpoint factory) should
+	// be moved accordingly.
+	source := dlmetrics.NewDataSource(*modelServerMetricsScheme,
+		int32(*modelServerMetricsPort), // start with (optional) command line port value
+		*modelServerMetricsPath,
+		*modelServerMetricsHttpsInsecureSkipVerify,
+		nil)
+	extractor, err := dlmetrics.NewExtractor(*totalQueuedRequestsMetric,
+		*kvCacheUsagePercentageMetric,
+		*loraInfoMetric)
+
+	if err != nil {
+		return nil, err
+	}
+	if err := source.AddExtractor(extractor); err != nil {
+		return nil, err
+	}
+	if err := datalayer.RegisterSource(source); err != nil {
+		return nil, err
+	}
+
+	factory := datalayer.NewEndpointFactory(datalayer.GetSources(), *refreshMetricsInterval)
+	return factory, nil
+}
+
 func initLogging(opts *zap.Options) {
 	// Unless -zap-log-level is explicitly set, use -v
 	useV := true

pkg/epp/datalayer/collector.go

@@ -82,17 +82,18 @@ func NewCollector() *Collector {
 }
 
 // Start initiates data source collection for the endpoint.
+// TODO: pass PoolInfo for backward compatibility
 func (c *Collector) Start(ctx context.Context, ticker Ticker, ep Endpoint, sources []DataSource) error {
 	var ready chan struct{}
 	started := false
 
 	c.startOnce.Do(func() {
+		logger := log.FromContext(ctx).WithValues("endpoint", ep.GetPod().GetIPAddress())
 		c.ctx, c.cancel = context.WithCancel(ctx)
 		started = true
 		ready = make(chan struct{})
 
 		go func(endpoint Endpoint, sources []DataSource) {
-			logger := log.FromContext(ctx).WithValues("endpoint", ep.GetPod().GetIPAddress())
 			logger.V(logging.DEFAULT).Info("starting collection")
 
 			defer func() {
@@ -107,6 +108,7 @@ func (c *Collector) Start(ctx context.Context, ticker Ticker, ep Endpoint, sourc
 				case <-c.ctx.Done(): // per endpoint context cancelled
 					return
 				case <-ticker.Channel():
+					// TODO: do not collect if there's no pool specified?
 					for _, src := range sources {
 						ctx, cancel := context.WithTimeout(c.ctx, defaultCollectionTimeout)
 						_ = src.Collect(ctx, endpoint) // TODO: track errors per collector?

pkg/epp/datalayer/datasource.go

@@ -92,14 +92,27 @@ func (dsr *DataSourceRegistry) GetSources() []DataSource {
 
 // --- default registry accessors ---
 
+// RegisterSource adds a new data source to the default registry.
 func RegisterSource(src DataSource) error {
 	return defaultDataSources.Register(src)
 }
 
-func GetNamedSource(name string) (DataSource, bool) {
-	return defaultDataSources.GetNamedSource(name)
+// GetNamedSource returns a typed data source from the default registry.
+func GetNamedSource[T DataSource](name string) (T, bool) {
+	v, ok := defaultDataSources.GetNamedSource(name)
+	if !ok {
+		var zero T
+		return zero, false
+	}
+	src, ok := v.(T)
+	if !ok {
+		var zero T
+		return zero, false
+	}
+	return src, true
 }
 
+// GetSources returns the list of data sources registered in the default registry.
 func GetSources() []DataSource {
 	return defaultDataSources.GetSources()
 }

pkg/epp/datalayer/factory.go

@@ -18,8 +18,12 @@ package datalayer
 
 import (
 	"context"
+	"sync"
+	"time"
 
 	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/types"
+	"sigs.k8s.io/controller-runtime/pkg/log"
 
 	v1 "sigs.k8s.io/gateway-api-inference-extension/api/v1"
 )
@@ -44,3 +48,74 @@ type EndpointFactory interface {
 	NewEndpoint(parent context.Context, inpod *corev1.Pod, poolinfo PoolInfo) Endpoint
 	ReleaseEndpoint(ep Endpoint)
 }
+
+// EndpointLifecycle manages the life cycle (creation and termination) of
+// endpoints.
+type EndpointLifecycle struct {
+	sources         []DataSource  // data sources for collectors
+	collectors      sync.Map      // collectors map. key: Pod namespaced name, value: *Collector
+	refreshInterval time.Duration // metrics refresh interval
+}
+
+// NewEndpointFactory returns a new endpoint for factory, managing collectors for
+// its endpoints. This function assumes that sources are not modified afterwards.
+func NewEndpointFactory(sources []DataSource, refreshMetricsInterval time.Duration) *EndpointLifecycle {
+	return &EndpointLifecycle{
+		sources:         sources,
+		collectors:      sync.Map{},
+		refreshInterval: refreshMetricsInterval,
+	}
+}
+
+// NewEndpoint implements EndpointFactory.NewEndpoint.
+// Creates a new endpoint and starts its associated collector with its own ticker.
+// Guards against multiple concurrent calls for the same endpoint.
+func (lc *EndpointLifecycle) NewEndpoint(parent context.Context, inpod *corev1.Pod, _ PoolInfo) Endpoint {
+	key := types.NamespacedName{Namespace: inpod.Namespace, Name: inpod.Name}
+	logger := log.FromContext(parent).WithValues("pod", key)
+
+	if _, ok := lc.collectors.Load(key); ok {
+		logger.Info("collector already running for endpoint", "endpoint", key)
+		return nil
+	}
+
+	endpoint := NewEndpoint()
+	endpoint.UpdatePod(inpod)
+	collector := NewCollector() // for full backward compatibility, set the logger and poolinfo
+
+	if _, loaded := lc.collectors.LoadOrStore(key, collector); loaded {
+		// another goroutine already created and stored a collector for this endpoint.
+		// No need to start the new collector.
+		logger.Info("collector already running for endpoint", "endpoint", key)
+		return nil
+	}
+
+	ticker := NewTimeTicker(lc.refreshInterval)
+	if err := collector.Start(parent, ticker, endpoint, lc.sources); err != nil {
+		logger.Error(err, "failed to start collector for endpoint", "endpoint", key)
+		lc.collectors.Delete(key)
+	}
+
+	return endpoint
+}
+
+// ReleaseEndpoint implements EndpointFactory.ReleaseEndpoint
+// Stops the collector and cleans up resources for the endpoint
+func (lc *EndpointLifecycle) ReleaseEndpoint(ep Endpoint) {
+	key := ep.GetPod().GetNamespacedName()
+
+	if value, ok := lc.collectors.LoadAndDelete(key); ok {
+		collector := value.(*Collector)
+		_ = collector.Stop()
+	}
+}
+
+// Shutdown gracefully stops all collectors and cleans up all resources.
+func (lc *EndpointLifecycle) Shutdown() {
+	lc.collectors.Range(func(key, value any) bool {
+		collector := value.(*Collector)
+		_ = collector.Stop()
+		lc.collectors.Delete(key)
+		return true
+	})
+}

pkg/epp/datalayer/metrics/client.go

@@ -33,22 +33,31 @@ type Client interface {
 	Get(ctx context.Context, target *url.URL, ep datalayer.Addressable) (PrometheusMetricMap, error)
 }
 
-// -- package implementations --
 const (
+	// the maximum idle connection count is shared by all endpoints. The value is
+	// set high to ensure the number of idle connection is above the expected
+	// endpoint count. Setting it too low would cause thrashing of the idle connection
+	// pool and incur higher overheads for every GET (e.g., socket initiation, certificate
+	// exchange, connections in timed wait state, etc.).
 	maxIdleConnections = 5000
-	maxIdleTime        = 10 * time.Second
-	timeout            = 10 * time.Second
+	maxIdleTime        = 10 * time.Second // once a endpoint goes down, allow closing.
+	timeout            = 10 * time.Second // mostly guard against unresponsive endpoints.
+	// allow some grace when connections are not made idle immediately (e.g., parsing
+	// and updating might take some time). This allows maintaining up to two idle connections
+	// per endpoint (defined as scheme://host:port).
+	maxIdleConnsPerHost = 2
 )
 
 var (
+	baseTransport = &http.Transport{
+		MaxIdleConns:        maxIdleConnections,
+		MaxIdleConnsPerHost: maxIdleConnsPerHost,
+		// TODO: set additional timeouts, transport options, etc.
+	}
 	defaultClient = &client{
 		Client: http.Client{
-			Timeout: timeout,
-			Transport: &http.Transport{
-				MaxIdleConns:        maxIdleConnections,
-				MaxIdleConnsPerHost: 4, // host is defined as scheme://host:port
-			},
-			// TODO: set additional timeouts, transport options, etc.
+			Timeout:   timeout,
+			Transport: baseTransport,
 		},
 	}
 )