Address review comments:

- DataSource is no longer tracking Endpoints.
- Removed unused interfaces.
- Changes to exported scope.
- Move Addressable to podinfo.go.

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

Author: elevran

pkg/epp/datalayer/datasource.go

@@ -17,59 +17,56 @@ limitations under the License.
 package datalayer
 
 import (
-	"context"
 	"errors"
 	"fmt"
 	"reflect"
 	"sync"
 )
 
-// DataSource is an interface required from all datalayer data collection
+// DataSource is an interface required from all data layer data collection
 // sources.
 type DataSource interface {
 	// Name returns the name of this datasource.
 	Name() string
 
-	// Start begins the collection process.
-	Start(ctx context.Context) error
-
-	// Stop stops the collection process.
-	Stop()
-
 	// AddExtractor adds an extractor to the data source.
-	// The extractor will be called whenever the data source might
+	// The extractor will be called whenever the Collector might
 	// have some new raw information regarding an endpoint.
 	// The Extractor's expected input type should be validated against
-	// the data source output type upon registration.
+	// the data source's output type upon registration.
 	AddExtractor(extractor Extractor) error
 
-	// AddEndpoint adds an endpoint to collect from.
-	AddEndpoint(ep Endpoint) error
-
-	// RemoveEndpoint removes an endpoint from collection.
-	RemoveEndpoint(ep Endpoint) error
+	// Collect is triggered by the data layer framework to fetch potentially new
+	// data for an endpoint. It passes retrieved data to registered Extractors.
+	Collect(ep Endpoint)
 }
 
+// Extractor is used to convert raw data into relevant data layer information
+// for an endpoint. They are called by data sources whenever new data might be
+// be available. Multiple Extractors can be registered with a source. Extractors
+// are expected to save their output with an endpoint so it becomes accessible
+// to consumers in other subsystem of the inference gateway (e.g., when making
+// scheduling decisions).
 type Extractor interface {
 	// Name returns the name of the extractor.
 	Name() string
 
 	// ExpectedType defines the type expected by the extractor. It must match
-	// the DataSource.OutputType() the extractor registers for.
-	ExpectedType() reflect.Type
+	// the output type of the data source where the extractor is registered.
+	ExpectedInputType() reflect.Type
 
 	// Extract transforms the data source output into a concrete attribute that
 	// is stored on the given endpoint.
 	Extract(data any, ep Endpoint)
 }
 
 var (
-	// DefaultDataSources is the system default data source registry.
-	DefaultDataSources = DataSourceRegistry{}
+	// defaultDataSources is the system default data source registry.
+	defaultDataSources = DataSourceRegistry{}
 )
 
 // DataSourceRegistry stores named data sources and makes them
-// accessible to GIE subsystems.
+// accessible to other subsystems in the inference gateway.
 type DataSourceRegistry struct {
 	sources sync.Map
 }
@@ -101,70 +98,32 @@ func (dsr *DataSourceRegistry) GetNamedSource(name string) (DataSource, bool) {
 	return nil, false
 }
 
-// AddEndpoint adds a new endpoint to all registered sources.
-// Endpoints are not tracked and DataSources are only notified of
-// endpoints added after the data source has been registered.
-//
-// TODO: track endpoints and update on later source registrations? It seems safe
-// to assume that all sources are registered before endpoints are
-// discovered and added to the system.
-func (dsr *DataSourceRegistry) AddEndpoint(ep Endpoint) error {
-	if ep == nil {
-		return nil
-	}
-
-	errs := []error{}
-	dsr.sources.Range(func(_, val interface{}) bool {
-		if ds, ok := val.(DataSource); ok {
-			if err := ds.AddEndpoint(ep); err != nil {
-				errs = append(errs, err)
-			}
-		}
-		return true
-	})
-	return errors.Join(errs...)
-}
-
-// RemoveEndpoint removes an endpoint from all registered sources.
-// A source may be called to remove an endpoint it has not added - this
-// is should not result in an error.
-func (dsr *DataSourceRegistry) RemoveEndpoint(ep Endpoint) error {
-	if ep == nil {
-		return nil
-	}
-
-	errs := []error{}
-	dsr.sources.Range(func(_, val interface{}) bool {
+// GetSources returns all sources registered.
+func (dsr *DataSourceRegistry) GetSources() []DataSource {
+	sources := []DataSource{}
+	dsr.sources.Range(func(_, val any) bool {
 		if ds, ok := val.(DataSource); ok {
-			if err := ds.RemoveEndpoint(ep); err != nil {
-				errs = append(errs, err)
-			}
+			sources = append(sources, ds)
 		}
-		return true
+		return true // continue iteration
 	})
-	return errors.Join(errs...)
+	return sources
 }
 
 // RegisterSource adds the data source to the default registry.
 func RegisterSource(src DataSource) error {
-	return DefaultDataSources.Register(src)
+	return defaultDataSources.Register(src)
 }
 
 // GetNamedSource returns the named source from the default registry,
 // if found.
 func GetNamedSource(name string) (DataSource, bool) {
-	return DefaultDataSources.GetNamedSource(name)
-}
-
-// AddEndpoint adds an endpoint to all sources in the default source registry.
-func AddEndpoint(ep Endpoint) error {
-	return DefaultDataSources.AddEndpoint(ep)
+	return defaultDataSources.GetNamedSource(name)
 }
 
-// RemoveEndpoint removes an endpoint from all sources in the default source
-// registry.
-func RemoveEndpoint(ep Endpoint) error {
-	return DefaultDataSources.RemoveEndpoint(ep)
+// GetSources returns all sources in the default registry.
+func GetSources() []DataSource {
+	return defaultDataSources.GetSources()
 }
 
 // ValidateExtractorType checks if an extractor can handle

pkg/epp/datalayer/endpoint.go

@@ -20,13 +20,6 @@ import (
 	corev1 "k8s.io/api/core/v1"
 )
 
-// MetricsLoop is a temporary bridge interface to match backend.metrics.PodMetrics
-// collection of metrics. It will be phased out as we move towards decoupled collection
-// via a DataSource.
-type MetricsLoop interface {
-	StopRefreshLoop()
-}
-
 // EndpointPodState allows management of the Pod related attributes.
 type EndpointPodState interface {
 	GetPod() *PodInfo
@@ -41,7 +34,6 @@ type EndpointMetricsState interface {
 
 // Endpoint represents an inference serving endpoint and its related attributes.
 type Endpoint interface {
-	MetricsLoop // TODO: remove once transition over to independent data source
 	EndpointPodState
 	EndpointMetricsState
 	AttributeMap

pkg/epp/datalayer/metrics.go

@@ -17,26 +17,10 @@ limitations under the License.
 package datalayer
 
 import (
-	"context"
 	"fmt"
 	"time"
-
-	"k8s.io/apimachinery/pkg/types"
 )
 
-// Addressable supports getting an IP address and a namespaced name.
-type Addressable interface {
-	GetIPAddress() string
-	GetNamespacedName() types.NamespacedName
-}
-
-// MetricsClient supports fetching the metrics of the endpoint. It is provided
-// with a snapshot of the last read metrics value (if any), under the assumption
-// that stale metrics are better than none.
-type MetricsClient interface {
-	FetchMetrics(c context.Context, ep Addressable, last *Metrics, port int32) (*Metrics, error)
-}
-
 // Metrics holds the latest metrics snapshot scraped from a pod.
 type Metrics struct {
 	// ActiveModels is a set of models(including LoRA adapters) that are currently cached to GPU.

pkg/epp/datalayer/podinfo.go

@@ -23,6 +23,12 @@ import (
 	"k8s.io/apimachinery/pkg/types"
 )
 
+// Addressable supports getting an IP address and a namespaced name.
+type Addressable interface {
+	GetIPAddress() string
+	GetNamespacedName() types.NamespacedName
+}
+
 // PodInfo represents the relevant Kubernetes Pod state of an inference server.
 type PodInfo struct {
 	NamespacedName types.NamespacedName