Update imageLocality plugin

Barakmor1 · Barakmor1 · commit a9ec92412545 · 2025-02-24T16:44:19.000+02:00
to account for ImageVolume images when scoring
and prioritizing nodes with required pod images

Signed-off-by: bmordeha &lt;bmordeha@redhat.com&gt;
diff --git a/pkg/scheduler/framework/plugins/imagelocality/image_locality.go b/pkg/scheduler/framework/plugins/imagelocality/image_locality.go
@@ -23,6 +23,8 @@ import (
 
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/runtime"
+	utilfeature "k8s.io/apiserver/pkg/util/feature"
+	"k8s.io/kubernetes/pkg/features"
 	"k8s.io/kubernetes/pkg/scheduler/framework"
 	"k8s.io/kubernetes/pkg/scheduler/framework/plugins/names"
 )
@@ -92,7 +94,8 @@ func calculatePriority(sumScores int64, numContainers int) int64 {
 	return framework.MaxNodeScore * (sumScores - minThreshold) / (maxThreshold - minThreshold)
 }
 
-// sumImageScores returns the sum of image scores of all the containers that are already on the node.
+// sumImageScores returns the sum of image scores of all the containers, including init containers and
+// imageVolumes specified in the pod's spec, that are already on the node.
 // Each image receives a raw score of its size, scaled by scaledImageScore. The raw scores are later used to calculate
 // the final score.
 func sumImageScores(nodeInfo *framework.NodeInfo, pod *v1.Pod, totalNumNodes int) int64 {
@@ -107,6 +110,16 @@ func sumImageScores(nodeInfo *framework.NodeInfo, pod *v1.Pod, totalNumNodes int
 			sum += scaledImageScore(state, totalNumNodes)
 		}
 	}
+	if utilfeature.DefaultFeatureGate.Enabled(features.ImageVolume) {
+		for _, volume := range pod.Spec.Volumes {
+			if volume.Image == nil {
+				continue
+			}
+			if state, ok := nodeInfo.ImageStates[normalizedImageName(volume.Image.Reference)]; ok {
+				sum += scaledImageScore(state, totalNumNodes)
+			}
+		}
+	}
 	return sum
 }
 
diff --git a/pkg/scheduler/framework/plugins/imagelocality/image_locality_test.go b/pkg/scheduler/framework/plugins/imagelocality/image_locality_test.go
@@ -92,6 +92,24 @@ func TestImageLocalityPriority(t *testing.T) {
 		},
 	}
 
+	testImageVolume := v1.PodSpec{
+		Containers: []v1.Container{
+			{
+				Image: "gcr.io/30",
+			},
+		},
+		Volumes: []v1.Volume{
+			{
+				Name: "imageVolume",
+				VolumeSource: v1.VolumeSource{
+					Image: &v1.ImageVolumeSource{
+						Reference: "gcr.io/300",
+					},
+				},
+			},
+		},
+	}
+
 	test30Init300 := v1.PodSpec{
 		Containers: []v1.Container{
 			{
@@ -340,6 +358,21 @@ func TestImageLocalityPriority(t *testing.T) {
 			expectedList: []framework.NodeScore{{Name: "node1", Score: 1}, {Name: "node2", Score: 0}},
 			name:         "pod with multiple small images",
 		},
+		{
+			// Pod: gcr.io/300 gcr.io/30
+
+			// Node1
+			// Image: gcr.io/300:latest 300MB
+			// Score: 100 * (300M * 1/2 - 23M) / (1000M - 23M) = 12
+
+			// Node2
+			// Image: gcr.io/30:latest 30MB
+			// Score:  100 * (30M - 23M) / (1000M - 23M) = 0
+			pod:          &v1.Pod{Spec: testImageVolume},
+			nodes:        []*v1.Node{makeImageNode("node1", node300600900), makeImageNode("node2", node400030)},
+			expectedList: []framework.NodeScore{{Name: "node1", Score: 12}, {Name: "node2", Score: 0}},
+			name:         "pod with ImageVolume",
+		},
 		{
 			// Pod: gcr.io/30  InitContainers: gcr.io/300
 
