fixed unit test except datastore

Signed-off-by: Nir Rozenbaum <[email protected]>

Author: nirrozenbaum

docs/proposals/003-model-server-protocol/README.md

@@ -4,7 +4,7 @@ This is the protocol between the EPP and the model servers.
 
 ### Inference API Protocol
 
-The model server MUST implement OpenAI’s [Completions](https://platform.openai.com/docs/api-reference/completions) 
+The model server MUST implement OpenAI’s [Completions](https://platform.openai.com/docs/api-reference/completions)
 and [Chat](https://platform.openai.com/docs/api-reference/chat) APIs.
 
 #### ** Experimental **   

pkg/epp/backend/metrics/fake_metrics_scraper.go

@@ -39,11 +39,11 @@ type FakeMetricsScraper struct {
 }
 
 func (s *FakeMetricsScraper) Name() string {
-	return "fake-metrics"
+	return MetricsDataKey
 }
 
 func (s *FakeMetricsScraper) InitData() podinfo.ScrapedData {
-	return s.existingMetrics
+	return NewMetrics()
 }
 
 func (s *FakeMetricsScraper) Scrape(ctx context.Context, pod *backend.Pod, port int) (podinfo.ScrapedData, error) {
@@ -63,7 +63,7 @@ func (s *FakeMetricsScraper) Scrape(ctx context.Context, pod *backend.Pod, port
 	return res.Clone(), nil
 }
 
-func (s *FakeMetricsScraper) ProcessResult(podinfo.ScrapedData) {} // noop
+func (s *FakeMetricsScraper) ProcessResult(ctx context.Context, podinfo podinfo.ScrapedData) {} // noop
 
 func (s *FakeMetricsScraper) SetRes(new map[types.NamespacedName]*Metrics) {
 	s.resMu.Lock()

pkg/epp/backend/metrics/metrics_scraper.go

@@ -22,12 +22,15 @@ import (
 	"net/http"
 	"strconv"
 	"strings"
+	"time"
 
 	dto "github.com/prometheus/client_model/go"
 	"github.com/prometheus/common/expfmt"
 	"go.uber.org/multierr"
+	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend"
 	podinfo "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/pod-info"
+	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
 
 const (
@@ -85,16 +88,18 @@ func (s *MetricsScraper) Scrape(ctx context.Context, pod *backend.Pod, port int)
 	return s.promToPodMetrics(metricFamilies)
 }
 
-func (s *MetricsScraper) ProcessResult(result podinfo.ScrapedData) {
-	if result == nil {
+func (s *MetricsScraper) ProcessResult(ctx context.Context, data podinfo.ScrapedData) {
+	if data == nil {
 		return
 	}
-	metrics := result.(*Metrics) //
+	metrics := data.(*Metrics)
 	if metrics == nil {
-		return // cannot convert to *Metrics
+		return
 	}
+	// if result cannot be converted to *Metrics, we cannot process result, otherwise
+	metrics.UpdateTime = time.Now()
 	s.existingMetrics = metrics
-	//TODO complete
+	log.FromContext(ctx).V(logutil.TRACE).Info("processed scraped metrics", "updated", metrics)
 }
 
 // promToPodMetrics updates internal pod metrics with scraped Prometheus metrics.

pkg/epp/backend/models/models_scraper.go

@@ -20,8 +20,10 @@ import (
 	"io"
 	"net/http"
 
+	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend"
 	podinfo "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/pod-info"
+	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
 
 const (
@@ -80,6 +82,6 @@ func (s *ModelsScraper) parseResponse(responseBody io.Reader) (podinfo.ScrapedDa
 	return &ModelsData{ModelIDs: modelIDs}, nil
 }
 
-func (s *ModelsScraper) ProcessResult(result podinfo.ScrapedData) {
-
+func (s *ModelsScraper) ProcessResult(ctx context.Context, data podinfo.ScrapedData) {
+	log.FromContext(ctx).V(logutil.TRACE).Info("processed scraped models", "models", data)
 }

pkg/epp/backend/pod-info/pod_info.go

@@ -49,7 +49,7 @@ type data struct {
 }
 
 func newPodInfo(pod *corev1.Pod, scrapers map[Scraper]*ScraperConfig, ds Datastore) *podInfo {
-	scraper := &podInfo{
+	info := &podInfo{
 		scrapers: scrapers,
 		ds:       ds,
 		once:     sync.Once{},
@@ -60,10 +60,10 @@ func newPodInfo(pod *corev1.Pod, scrapers map[Scraper]*ScraperConfig, ds Datasto
 		data.data[scraper.Name()] = scraper.InitData()
 	}
 
-	scraper.pod.Store(toInternalPod(pod))
-	scraper.data.Store(data)
+	info.pod.Store(toInternalPod(pod))
+	info.data.Store(data)
 
-	return scraper
+	return info
 }
 
 type podInfo struct {
@@ -97,6 +97,11 @@ func (pi *podInfo) String() string {
 	return fmt.Sprintf("Pod: %v; Data: %v", pi.GetPod(), pi.GetData())
 }
 
+func (pi *podInfo) Stop() {
+	pi.logger.V(logutil.DEFAULT).Info("Stopping scrape loop")
+	pi.stopFunc()
+}
+
 func (pi *podInfo) startScrapers(ctx context.Context) {
 	// The returned context'pi Done channel is closed when the returned cancel function is called or when the parent context'pi Done channel is closed, whichever happens first.
 	newCtx, cancel := context.WithCancel(ctx)
@@ -131,14 +136,19 @@ func (pi *podInfo) startScrapeLoop(ctx context.Context, scraper Scraper, interva
 			}
 			// allow processing partial results in case there was an error.
 			// the scraper can return nil in case of an error and then do nothing in the ProcessResult function in case no partial update is required
-			scraper.ProcessResult(scrapedData)
+			scraper.ProcessResult(ctx, scrapedData)
+			// store updated data if its valid
+			if scrapedData != nil {
+				pi.storeScrapedData(scraper.Name(), scrapedData)
+			}
 		}
 	}
 }
 
-func (pi *podInfo) Stop() {
-	pi.logger.V(logutil.DEFAULT).Info("Stopping scrape loop")
-	pi.stopFunc()
+func (pi *podInfo) storeScrapedData(key string, data ScrapedData) {
+	updated := pi.data.Load()
+	updated.data[key] = data
+	pi.data.Store(updated)
 }
 
 func toInternalPod(pod *corev1.Pod) *backend.Pod {

pkg/epp/backend/pod-info/types.go

@@ -48,8 +48,8 @@ type Scraper interface {
 	// Scrape scrapes infromation from a pod.
 	Scrape(ctx context.Context, pod *backend.Pod, port int) (ScrapedData, error)
 	// ProcessResult process the returned object from Scrape function.
-	// This function should update PodInfo data field with the new result.
-	ProcessResult(ScrapedData)
+	// This function is used to update internal fields in the scraper struct before storing the ScrapedData in the PodInfo data map.
+	ProcessResult(ctx context.Context, data ScrapedData)
 }
 
 // ScrapedData is a generic type for arbitrary scraperd data stored in PodInfo.

pkg/epp/datastore/datastore_test.go

@@ -353,8 +353,9 @@ func TestMetrics(t *testing.T) {
 				got := ds.PodGetAll()
 				metrics := []*backendmetrics.Metrics{}
 				for _, one := range got {
-					podMetrics := one.GetData()[backendmetrics.MetricsDataKey].(*backendmetrics.Metrics)
-					metrics = append(metrics, podMetrics)
+					if podMetrics, ok := one.GetData()[backendmetrics.MetricsDataKey].(*backendmetrics.Metrics); ok {
+						metrics = append(metrics, podMetrics)
+					}
 				}
 				diff := cmp.Diff(test.want, metrics, cmpopts.IgnoreFields(backendmetrics.Metrics{}, "UpdateTime"), cmpopts.SortSlices(func(a, b *backendmetrics.Metrics) bool {
 					return a.String() < b.String()

pkg/epp/scheduling/config.go

@@ -16,7 +16,9 @@ limitations under the License.
 
 package scheduling
 
-import "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/plugins"
+import (
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/plugins"
+)
 
 // NewSchedulerConfig creates a new SchedulerConfig object with the given plugins.
 func NewSchedulerConfig(preSchedulePlugins []plugins.PreSchedule, filters []plugins.Filter, scorers map[plugins.Scorer]int,

pkg/epp/scheduling/scheduler_test.go

@@ -394,7 +394,8 @@ func TestSchedulePlugins(t *testing.T) {
 
 			// Validate output
 			wantPod := &types.PodData{
-				Pod: &backend.Pod{NamespacedName: test.wantTargetPod},
+				Pod:  &backend.Pod{NamespacedName: test.wantTargetPod},
+				Data: map[string]podinfo.ScrapedData{},
 			}
 			wantRes := &types.Result{TargetPod: wantPod}
 			if diff := cmp.Diff(wantRes, got); diff != "" {