Refactor scheduler

Author: liu-cong

pkg/epp/backend/metrics/metrics.go

@@ -109,6 +109,7 @@ func (p *PodMetricsClientImpl) promToPodMetrics(
 
 		if loraMetrics != nil {
 			updated.ActiveModels = make(map[string]int)
+			updated.WaitingModels = make(map[string]int)
 			for _, label := range loraMetrics.GetLabel() {
 				if label.GetName() == LoraInfoRunningAdaptersMetricName {
 					if label.GetValue() != "" {
@@ -122,7 +123,7 @@ func (p *PodMetricsClientImpl) promToPodMetrics(
 					if label.GetValue() != "" {
 						adapterList := strings.Split(label.GetValue(), ",")
 						for _, adapter := range adapterList {
-							updated.ActiveModels[adapter] = 0
+							updated.WaitingModels[adapter] = 0
 						}
 					}
 				}

pkg/epp/backend/metrics/metrics_test.go

@@ -404,7 +404,8 @@ func TestPromToPodMetrics(t *testing.T) {
 			expectedMetrics: &Metrics{
 				WaitingQueueSize:    7,
 				KVCacheUsagePercent: 0.8,
-				ActiveModels:        map[string]int{"lora1": 0, "lora2": 0, "lora3": 0},
+				ActiveModels:        map[string]int{"lora1": 0, "lora2": 0},
+				WaitingModels:       map[string]int{"lora3": 0},
 				MaxActiveModels:     3,
 			},
 		},
@@ -416,8 +417,8 @@ func TestPromToPodMetrics(t *testing.T) {
 				KVCacheUtilization:  &MetricSpec{MetricName: "vllm_usage"},
 				LoraRequestInfo:     &MetricSpec{MetricName: "vllm:lora_requests_info"},
 			},
-			existingMetrics: &Metrics{ActiveModels: map[string]int{}},
-			expectedMetrics: &Metrics{ActiveModels: map[string]int{}},
+			existingMetrics: &Metrics{ActiveModels: map[string]int{}, WaitingModels: map[string]int{}},
+			expectedMetrics: &Metrics{ActiveModels: map[string]int{}, WaitingModels: map[string]int{}},
 			expectedErr:     multierr.Combine(errors.New("metric family \"vllm_waiting\" not found"), errors.New("metric family \"vllm_usage\" not found"), errors.New("metric family \"vllm:lora_requests_info\" not found")),
 		},
 		{
@@ -439,7 +440,8 @@ func TestPromToPodMetrics(t *testing.T) {
 			expectedMetrics: &Metrics{
 				WaitingQueueSize:    0,
 				KVCacheUsagePercent: 0.8,
-				ActiveModels:        map[string]int{"lora1": 0, "lora2": 0, "lora3": 0},
+				ActiveModels:        map[string]int{"lora1": 0, "lora2": 0},
+				WaitingModels:       map[string]int{"lora3": 0},
 				MaxActiveModels:     3,
 			},
 			expectedErr: errors.New("metric family \"vllm_waiting\" not found"),
@@ -457,6 +459,7 @@ func TestPromToPodMetrics(t *testing.T) {
 			existingMetrics: &Metrics{},
 			expectedMetrics: &Metrics{
 				ActiveModels:    map[string]int{"lora1": 0},
+				WaitingModels:   map[string]int{},
 				MaxActiveModels: 0, // Should still default to 0.
 
 			},

pkg/epp/backend/metrics/pod_metrics_test.go

@@ -44,6 +44,7 @@ var (
 			"foo": 1,
 			"bar": 1,
 		},
+		WaitingModels: map[string]int{},
 	}
 	updated = &Metrics{
 		WaitingQueueSize:    9999,
@@ -53,6 +54,7 @@ var (
 			"foo": 1,
 			"bar": 1,
 		},
+		WaitingModels: map[string]int{},
 	}
 )
 

pkg/epp/backend/metrics/types.go

@@ -41,16 +41,17 @@ type PodMetricsFactory struct {
 }
 
 func (f *PodMetricsFactory) NewPodMetrics(parentCtx context.Context, in *corev1.Pod, ds Datastore) PodMetrics {
+	pod := toInternalPod(in)
 	pm := &podMetrics{
 		pmc:       f.pmc,
 		ds:        ds,
 		interval:  f.refreshMetricsInterval,
 		parentCtx: parentCtx,
 		once:      sync.Once{},
 		done:      make(chan struct{}),
-		logger:    log.FromContext(parentCtx),
+		logger:    log.FromContext(parentCtx).WithValues("pod", pod.NamespacedName),
 	}
-	pm.pod.Store(toInternalPod(in))
+	pm.pod.Store(pod)
 	pm.metrics.Store(newMetrics())
 
 	pm.startRefreshLoop()
@@ -77,9 +78,20 @@ func (p *Pod) String() string {
 	return fmt.Sprintf("%+v", *p)
 }
 
+func (p *Pod) Clone() *Pod {
+	return &Pod{
+		NamespacedName: types.NamespacedName{
+			Name:      p.NamespacedName.Name,
+			Namespace: p.NamespacedName.Namespace,
+		},
+		Address: p.Address,
+	}
+}
+
 type Metrics struct {
 	// ActiveModels is a set of models(including LoRA adapters) that are currently cached to GPU.
-	ActiveModels map[string]int
+	ActiveModels  map[string]int
+	WaitingModels map[string]int
 	// MaxActiveModels is the maximum number of models that can be loaded to GPU.
 	MaxActiveModels         int
 	RunningQueueSize        int
@@ -93,7 +105,8 @@ type Metrics struct {
 
 func newMetrics() *Metrics {
 	return &Metrics{
-		ActiveModels: make(map[string]int),
+		ActiveModels:  make(map[string]int),
+		WaitingModels: make(map[string]int),
 	}
 }
 
@@ -109,8 +122,13 @@ func (m *Metrics) Clone() *Metrics {
 	for k, v := range m.ActiveModels {
 		cm[k] = v
 	}
+	wm := make(map[string]int, len(m.WaitingModels))
+	for k, v := range m.WaitingModels {
+		wm[k] = v
+	}
 	clone := &Metrics{
 		ActiveModels:            cm,
+		WaitingModels:           wm,
 		MaxActiveModels:         m.MaxActiveModels,
 		RunningQueueSize:        m.RunningQueueSize,
 		WaitingQueueSize:        m.WaitingQueueSize,

pkg/epp/datastore/datastore_test.go

@@ -236,6 +236,7 @@ var (
 			"foo": 1,
 			"bar": 1,
 		},
+		WaitingModels: map[string]int{},
 	}
 	pod2 = &corev1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
@@ -250,6 +251,7 @@ var (
 			"foo1": 1,
 			"bar1": 1,
 		},
+		WaitingModels: map[string]int{},
 	}
 	pod1NamespacedName = types.NamespacedName{Name: pod1.Name, Namespace: pod1.Namespace}
 	pod2NamespacedName = types.NamespacedName{Name: pod2.Name, Namespace: pod2.Namespace}
@@ -305,6 +307,7 @@ func TestMetrics(t *testing.T) {
 				// Failed to fetch pod2 metrics so it remains the default values.
 				{
 					ActiveModels:        map[string]int{},
+					WaitingModels:       map[string]int{},
 					WaitingQueueSize:    0,
 					KVCacheUsagePercent: 0,
 					MaxActiveModels:     0,

pkg/epp/handlers/server.go

@@ -26,7 +26,6 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"sigs.k8s.io/controller-runtime/pkg/log"
-	backendmetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/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/scheduling"
@@ -57,7 +56,7 @@ type Server struct {
 }
 
 type Scheduler interface {
-	Schedule(ctx context.Context, b *scheduling.LLMRequest) (targetPod backendmetrics.PodMetrics, err error)
+	Schedule(ctx context.Context, b *scheduling.LLMRequest) (targetPod scheduling.Pod, err error)
 }
 
 func (s *Server) Process(srv extProcPb.ExternalProcessor_ProcessServer) error {