Read pod role from kubernetes pod info in reconciler (kubernetes-sigs#82)

* Read pod role from kubernetes pod info in reconciler

* Mark pod without role label as Both

* updated for review comments

Author: mayabar

pkg/epp/backend/metrics/pod_metrics.go

@@ -32,6 +32,10 @@ import (
 
 const (
 	fetchMetricsTimeout = 5 * time.Second
+	roleLabel           = "llmd.org/role"
+	rolePrefill         = "prefill"
+	roleDecode          = "decode"
+	roleBoth            = "both"
 )
 
 type podMetrics struct {
@@ -67,13 +71,34 @@ func (pm *podMetrics) UpdatePod(in *corev1.Pod) {
 	pm.pod.Store(toInternalPod(in))
 }
 
+func podLabelToRole(in *corev1.Pod) PodRole {
+	roleLabel, ok := in.ObjectMeta.Labels[roleLabel]
+
+	if ok {
+		switch roleLabel {
+		case rolePrefill:
+			return Prefill
+		case roleDecode:
+			return Decode
+		case roleBoth:
+			return Both
+		default:
+			return Unknown
+		}
+	}
+
+	// role label is missing
+	return Both
+}
+
 func toInternalPod(in *corev1.Pod) *Pod {
 	return &Pod{
 		NamespacedName: types.NamespacedName{
 			Name:      in.Name,
 			Namespace: in.Namespace,
 		},
 		Address: in.Status.PodIP,
+		Role:    podLabelToRole(in),
 	}
 }
 

pkg/epp/backend/metrics/types.go

@@ -65,9 +65,19 @@ type PodMetrics interface {
 	String() string
 }
 
+type PodRole int
+
+const (
+	Prefill PodRole = iota
+	Decode
+	Both
+	Unknown
+)
+
 type Pod struct {
 	NamespacedName types.NamespacedName
 	Address        string
+	Role           PodRole
 }
 
 func (p *Pod) String() string {
@@ -87,6 +97,7 @@ func (p *Pod) Clone() *Pod {
 			Namespace: p.NamespacedName.Namespace,
 		},
 		Address: p.Address,
+		Role:    p.Role,
 	}
 }