Build image size map upon node info updates

pkg/scheduler/algorithm/priorities/image_locality.go

@@ -42,7 +42,7 @@ func ImageLocalityPriorityMap(pod *v1.Pod, meta interface{}, nodeInfo *scheduler
 		return schedulerapi.HostPriority{}, fmt.Errorf("node not found")
 	}
 
-	sumSize := totalImageSize(node, pod.Spec.Containers)
+	sumSize := totalImageSize(nodeInfo, pod.Spec.Containers)
 
 	return schedulerapi.HostPriority{
 		Host:  node.Name,
@@ -69,15 +69,10 @@ func calculateScoreFromSize(sumSize int64) int {
 }
 
 // totalImageSize returns the total image size of all the containers that are already on the node.
-func totalImageSize(node *v1.Node, containers []v1.Container) int64 {
+func totalImageSize(nodeInfo *schedulercache.NodeInfo, containers []v1.Container) int64 {
-	imageSizes := make(map[string]int64)
-	for _, image := range node.Status.Images {
-		for _, name := range image.Names {
-			imageSizes[name] = image.SizeBytes
-		}
-	}
-
 	var total int64
+
+	imageSizes := nodeInfo.Images()
 	for _, container := range containers {
 		if size, ok := imageSizes[container.Image]; ok {
 			total += size

pkg/scheduler/schedulercache/cache_test.go

@@ -108,6 +108,7 @@ func TestAssumePodScheduled(t *testing.T) {
 			allocatableResource: &Resource{},
 			pods:                []*v1.Pod{testPods[0]},
 			usedPorts:           newHostPortInfoBuilder().add("TCP", "127.0.0.1", 80).build(),
+			imageSizes:          map[string]int64{},
 		},
 	}, {
 		pods: []*v1.Pod{testPods[1], testPods[2]},
@@ -124,6 +125,7 @@ func TestAssumePodScheduled(t *testing.T) {
 			allocatableResource: &Resource{},
 			pods:                []*v1.Pod{testPods[1], testPods[2]},
 			usedPorts:           newHostPortInfoBuilder().add("TCP", "127.0.0.1", 80).add("TCP", "127.0.0.1", 8080).build(),
+			imageSizes:          map[string]int64{},
 		},
 	}, { // test non-zero request
 		pods: []*v1.Pod{testPods[3]},
@@ -140,6 +142,7 @@ func TestAssumePodScheduled(t *testing.T) {
 			allocatableResource: &Resource{},
 			pods:                []*v1.Pod{testPods[3]},
 			usedPorts:           newHostPortInfoBuilder().add("TCP", "127.0.0.1", 80).build(),
+			imageSizes:          map[string]int64{},
 		},
 	}, {
 		pods: []*v1.Pod{testPods[4]},
@@ -157,6 +160,7 @@ func TestAssumePodScheduled(t *testing.T) {
 			allocatableResource: &Resource{},
 			pods:                []*v1.Pod{testPods[4]},
 			usedPorts:           newHostPortInfoBuilder().add("TCP", "127.0.0.1", 80).build(),
+			imageSizes:          map[string]int64{},
 		},
 	}, {
 		pods: []*v1.Pod{testPods[4], testPods[5]},
@@ -174,6 +178,7 @@ func TestAssumePodScheduled(t *testing.T) {
 			allocatableResource: &Resource{},
 			pods:                []*v1.Pod{testPods[4], testPods[5]},
 			usedPorts:           newHostPortInfoBuilder().add("TCP", "127.0.0.1", 80).add("TCP", "127.0.0.1", 8080).build(),
+			imageSizes:          map[string]int64{},
 		},
 	}, {
 		pods: []*v1.Pod{testPods[6]},
@@ -190,6 +195,7 @@ func TestAssumePodScheduled(t *testing.T) {
 			allocatableResource: &Resource{},
 			pods:                []*v1.Pod{testPods[6]},
 			usedPorts:           newHostPortInfoBuilder().build(),
+			imageSizes:          map[string]int64{},
 		},
 	},
 	}
@@ -269,6 +275,7 @@ func TestExpirePod(t *testing.T) {
 			allocatableResource: &Resource{},
 			pods:                []*v1.Pod{testPods[1]},
 			usedPorts:           newHostPortInfoBuilder().add("TCP", "127.0.0.1", 8080).build(),
+			imageSizes:          map[string]int64{},
 		},
 	}}
 
@@ -321,6 +328,7 @@ func TestAddPodWillConfirm(t *testing.T) {
 			allocatableResource: &Resource{},
 			pods:                []*v1.Pod{testPods[0]},
 			usedPorts:           newHostPortInfoBuilder().add("TCP", "127.0.0.1", 80).build(),
+			imageSizes:          map[string]int64{},
 		},
 	}}
 
@@ -417,6 +425,7 @@ func TestAddPodWillReplaceAssumed(t *testing.T) {
 				allocatableResource: &Resource{},
 				pods:                []*v1.Pod{updatedPod.DeepCopy()},
 				usedPorts:           newHostPortInfoBuilder().add("TCP", "0.0.0.0", 90).build(),
+				imageSizes:          map[string]int64{},
 			},
 		},
 	}}
@@ -472,6 +481,7 @@ func TestAddPodAfterExpiration(t *testing.T) {
 			allocatableResource: &Resource{},
 			pods:                []*v1.Pod{basePod},
 			usedPorts:           newHostPortInfoBuilder().add("TCP", "127.0.0.1", 80).build(),
+			imageSizes:          map[string]int64{},
 		},
 	}}
 
@@ -528,6 +538,7 @@ func TestUpdatePod(t *testing.T) {
 			allocatableResource: &Resource{},
 			pods:                []*v1.Pod{testPods[1]},
 			usedPorts:           newHostPortInfoBuilder().add("TCP", "127.0.0.1", 8080).build(),
+			imageSizes:          map[string]int64{},
 		}, {
 			requestedResource: &Resource{
 				MilliCPU: 100,
@@ -541,6 +552,7 @@ func TestUpdatePod(t *testing.T) {
 			allocatableResource: &Resource{},
 			pods:                []*v1.Pod{testPods[0]},
 			usedPorts:           newHostPortInfoBuilder().add("TCP", "127.0.0.1", 80).build(),
+			imageSizes:          map[string]int64{},
 		}},
 	}}
 
@@ -599,6 +611,7 @@ func TestExpireAddUpdatePod(t *testing.T) {
 			allocatableResource: &Resource{},
 			pods:                []*v1.Pod{testPods[1]},
 			usedPorts:           newHostPortInfoBuilder().add("TCP", "127.0.0.1", 8080).build(),
+			imageSizes:          map[string]int64{},
 		}, {
 			requestedResource: &Resource{
 				MilliCPU: 100,
@@ -612,6 +625,7 @@ func TestExpireAddUpdatePod(t *testing.T) {
 			allocatableResource: &Resource{},
 			pods:                []*v1.Pod{testPods[0]},
 			usedPorts:           newHostPortInfoBuilder().add("TCP", "127.0.0.1", 80).build(),
+			imageSizes:          map[string]int64{},
 		}},
 	}}
 
@@ -689,6 +703,7 @@ func TestEphemeralStorageResource(t *testing.T) {
 				allocatableResource: &Resource{},
 				pods:                []*v1.Pod{podE},
 				usedPorts:           schedutil.HostPortInfo{},
+				imageSizes:          map[string]int64{},
 			},
 		},
 	}
@@ -735,6 +750,7 @@ func TestRemovePod(t *testing.T) {
 			allocatableResource: &Resource{},
 			pods:                []*v1.Pod{basePod},
 			usedPorts:           newHostPortInfoBuilder().add("TCP", "127.0.0.1", 80).build(),
+			imageSizes:          map[string]int64{},
 		},
 	}}
 

pkg/scheduler/schedulercache/node_info.go

@@ -58,6 +58,10 @@ type NodeInfo struct {
 	taints    []v1.Taint
 	taintsErr error
 
+	// This is a map from image name to image size, also for checking image existence on the node
+	// Cache it here to avoid rebuilding the map during scheduling, e.g., in image_locality.go
+	imageSizes map[string]int64
+
 	// TransientInfo holds the information pertaining to a scheduling cycle. This will be destructed at the end of
 	// scheduling cycle.
 	// TODO: @ravig. Remove this once we have a clear approach for message passing across predicates and priorities.
@@ -226,6 +230,7 @@ func NewNodeInfo(pods ...*v1.Pod) *NodeInfo {
 		TransientInfo:       newTransientSchedulerInfo(),
 		generation:          nextGeneration(),
 		usedPorts:           make(util.HostPortInfo),
+		imageSizes:          make(map[string]int64),
 	}
 	for _, pod := range pods {
 		ni.AddPod(pod)
@@ -257,6 +262,14 @@ func (n *NodeInfo) UsedPorts() util.HostPortInfo {
 	return n.usedPorts
 }
 
+// Images returns the image size information on this node.
+func (n *NodeInfo) Images() map[string]int64 {
+	if n == nil {
+		return nil
+	}
+	return n.imageSizes
+}
+
 // PodsWithAffinity return all pods with (anti)affinity constraints on this node.
 func (n *NodeInfo) PodsWithAffinity() []*v1.Pod {
 	if n == nil {
@@ -348,6 +361,7 @@ func (n *NodeInfo) Clone() *NodeInfo {
 		diskPressureCondition:   n.diskPressureCondition,
 		pidPressureCondition:    n.pidPressureCondition,
 		usedPorts:               make(util.HostPortInfo),
+		imageSizes:              n.imageSizes,
 		generation:              n.generation,
 	}
 	if len(n.pods) > 0 {
@@ -491,6 +505,17 @@ func (n *NodeInfo) updateUsedPorts(pod *v1.Pod, add bool) {
 	}
 }
 
+func (n *NodeInfo) updateImageSizes() {
+	node := n.Node()
+	imageSizes := make(map[string]int64)
+	for _, image := range node.Status.Images {
+		for _, name := range image.Names {
+			imageSizes[name] = image.SizeBytes
+		}
+	}
+	n.imageSizes = imageSizes
+}
+
 // SetNode sets the overall node information.
 func (n *NodeInfo) SetNode(node *v1.Node) error {
 	n.node = node
@@ -512,6 +537,7 @@ func (n *NodeInfo) SetNode(node *v1.Node) error {
 		}
 	}
 	n.TransientInfo = newTransientSchedulerInfo()
+	n.updateImageSizes()
 	n.generation = nextGeneration()
 	return nil
 }

pkg/scheduler/schedulercache/node_info_test.go

@@ -216,6 +216,7 @@ func TestNewNodeInfo(t *testing.T) {
 				{Protocol: "TCP", Port: 8080}: {},
 			},
 		},
+		imageSizes: map[string]int64{},
 		pods: []*v1.Pod{
 			{
 				ObjectMeta: metav1.ObjectMeta{
@@ -304,6 +305,9 @@ func TestNodeInfoClone(t *testing.T) {
 						{Protocol: "TCP", Port: 8080}: {},
 					},
 				},
+				imageSizes: map[string]int64{
+					"gcr.io/10": 10 * 1024 * 1024,
+				},
 				pods: []*v1.Pod{
 					{
 						ObjectMeta: metav1.ObjectMeta{
@@ -373,6 +377,9 @@ func TestNodeInfoClone(t *testing.T) {
 						{Protocol: "TCP", Port: 8080}: {},
 					},
 				},
+				imageSizes: map[string]int64{
+					"gcr.io/10": 10 * 1024 * 1024,
+				},
 				pods: []*v1.Pod{
 					{
 						ObjectMeta: metav1.ObjectMeta{
@@ -530,6 +537,7 @@ func TestNodeInfoAddPod(t *testing.T) {
 				{Protocol: "TCP", Port: 8080}: {},
 			},
 		},
+		imageSizes: map[string]int64{},
 		pods: []*v1.Pod{
 			{
 				ObjectMeta: metav1.ObjectMeta{
@@ -648,6 +656,7 @@ func TestNodeInfoRemovePod(t *testing.T) {
 						{Protocol: "TCP", Port: 8080}: {},
 					},
 				},
+				imageSizes: map[string]int64{},
 				pods: []*v1.Pod{
 					{
 						ObjectMeta: metav1.ObjectMeta{
@@ -763,6 +772,7 @@ func TestNodeInfoRemovePod(t *testing.T) {
 						{Protocol: "TCP", Port: 8080}: {},
 					},
 				},
+				imageSizes: map[string]int64{},
 				pods: []*v1.Pod{
 					{
 						ObjectMeta: metav1.ObjectMeta{