add imagelocality plugin to exmaple

Makefile

@@ -9,8 +9,8 @@ examples/advanced/main.wasm: examples/advanced/main.go
 	@(cd $(@D); tinygo build -o main.wasm -scheduler=none --no-debug -target=wasi .)
 
 .PHONY: build-tinygo
-build-tinygo: examples/nodenumber/main.wasm examples/advanced/main.wasm guest/testdata/cyclestate/main.wasm guest/testdata/filter/main.wasm guest/testdata/score/main.wasm \
-			  guest/testdata/bind/main.wasm guest/testdata/reserve/main.wasm guest/testdata/handle/main.wasm
+build-tinygo: examples/nodenumber/main.wasm examples/advanced/main.wasm examples/imagelocality/main.wasm guest/testdata/cyclestate/main.wasm guest/testdata/filter/main.wasm guest/testdata/score/main.wasm \
+			  guest/testdata/bind/main.wasm guest/testdata/reserve/main.wasm guest/testdata/handle/main.wasm 
 
 %/main-debug.wasm: %/main.go
 	@(cd $(@D); tinygo build -o main-debug.wasm -gc=custom -tags=custommalloc -scheduler=none -target=wasi .)

examples/imagelocality/README.md

@@ -0,0 +1,3 @@
+# Wasm ImageLocality
+
+This is [ImageLocality plugin](https://github.com/kubernetes/kubernetes/tree/master/pkg/scheduler/framework/plugins/imagelocality) implemented with the wasm extension.

examples/imagelocality/go.mod

@@ -0,0 +1,18 @@
+module sigs.k8s.io/kube-scheduler-wasm-extension/examples/imagelocality
+
+go 1.20
+
+require (
+	github.com/wasilibs/nottinygc v0.7.1
+	sigs.k8s.io/kube-scheduler-wasm-extension/guest v0.0.0-00010101000000-000000000000
+)
+
+require (
+	github.com/magefile/mage v1.14.0 // indirect
+	google.golang.org/protobuf v1.30.0 // indirect
+	sigs.k8s.io/kube-scheduler-wasm-extension/kubernetes/proto v0.0.0-00010101000000-000000000000 // indirect
+)
+
+replace sigs.k8s.io/kube-scheduler-wasm-extension/guest => ./../../guest
+
+replace sigs.k8s.io/kube-scheduler-wasm-extension/kubernetes/proto => ./../../kubernetes/proto

examples/imagelocality/go.sum

@@ -0,0 +1,11 @@
+github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
+github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
+github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
+github.com/magefile/mage v1.14.0 h1:6QDX3g6z1YvJ4olPhT1wksUcSa/V0a1B+pJb73fBjyo=
+github.com/magefile/mage v1.14.0/go.mod h1:z5UZb/iS3GoOSn0JgWuiw7dxlurVYTu+/jHXqQg881A=
+github.com/wasilibs/nottinygc v0.7.1 h1:rKu19+SFniRNuSo5NX7/wxpSpXmMUmkcyt/YiWLJg8w=
+github.com/wasilibs/nottinygc v0.7.1/go.mod h1:oDcIotskuYNMpqMF23l7Z8uzD4TC0WXHK8jetlB3HIo=
+golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
+google.golang.org/protobuf v1.30.0 h1:kPPoIgf3TsEvrm0PFe15JQ+570QVxYzEvvHqChK+cng=
+google.golang.org/protobuf v1.30.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=

examples/imagelocality/image_locality.go

@@ -0,0 +1,99 @@
+package main
+
+import (
+	"fmt"
+	"strings"
+
+	guestapi "sigs.k8s.io/kube-scheduler-wasm-extension/guest/api"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api/proto"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/handle/sharedlister/api"
+)
+
+// The two thresholds are used as bounds for the image score range. They correspond to a reasonable size range for
+// container images compressed and stored in registries; 90%ile of images on dockerhub drops into this range.
+const (
+	mb                    int64 = 1024 * 1024
+	minThreshold          int64 = 23 * mb
+	maxContainerThreshold int64 = 1000 * mb
+)
+
+// imageLocality is a score plugin that favors nodes that already have requested pod container's images.
+type imageLocality struct {
+	sharedLister api.SharedLister
+}
+
+// Score invoked at the score extension point.
+func (pl *imageLocality) Score(state guestapi.CycleState, pod proto.Pod, nodeName string) (int32, *guestapi.Status) {
+	nodeInfo := pl.sharedLister.NodeInfos().Get(nodeName)
+	if nodeInfo == nil {
+		return 0, &guestapi.Status{Code: guestapi.StatusCodeError, Reason: fmt.Sprintf("failed to get node %q", nodeName)}
+	}
+
+	nodeInfos := pl.sharedLister.NodeInfos().List()
+	if nodeInfos == nil {
+		return 0, &guestapi.Status{Code: guestapi.StatusCodeError, Reason: "failed to list nodes"}
+	}
+	totalNumNodes := len(nodeInfos)
+
+	imageScores := sumImageScores(nodeInfo, pod, totalNumNodes)
+	score := calculatePriority(imageScores, len(pod.Spec().InitContainers)+len(pod.Spec().Containers))
+
+	return int32(score), nil
+}
+
+const (
+	// maxNodeScore is the maximum score a Score plugin is expected to return.
+	maxNodeScore int64 = 100
+)
+
+// calculatePriority returns the priority of a node. Given the sumScores of requested images on the node, the node's
+// priority is obtained by scaling the maximum priority value with a ratio proportional to the sumScores.
+func calculatePriority(sumScores int64, numContainers int) int64 {
+	maxThreshold := maxContainerThreshold * int64(numContainers)
+	if sumScores < minThreshold {
+		sumScores = minThreshold
+	} else if sumScores > maxThreshold {
+		sumScores = maxThreshold
+	}
+
+	return maxNodeScore * (sumScores - minThreshold) / (maxThreshold - minThreshold)
+}
+
+// sumImageScores returns the sum of image scores of all the containers 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 guestapi.NodeInfo, pod proto.Pod, totalNumNodes int) int64 {
+	var sum int64
+	for _, container := range pod.Spec().InitContainers {
+		if state, ok := nodeInfo.ImageStates()[normalizedImageName(*container.Image)]; ok {
+			sum += scaledImageScore(state, totalNumNodes)
+		}
+	}
+	for _, container := range pod.Spec().Containers {
+		if state, ok := nodeInfo.ImageStates()[normalizedImageName(*container.Image)]; ok {
+			sum += scaledImageScore(state, totalNumNodes)
+		}
+	}
+	return sum
+}
+
+// scaledImageScore returns an adaptively scaled score for the given state of an image.
+// The size of the image is used as the base score, scaled by a factor which considers how much nodes the image has "spread" to.
+// This heuristic aims to mitigate the undesirable "node heating problem", i.e., pods get assigned to the same or
+// a few nodes due to image locality.
+func scaledImageScore(imageState *guestapi.ImageStateSummary, totalNumNodes int) int64 {
+	spread := float64(imageState.NumNodes) / float64(totalNumNodes)
+	return int64(float64(imageState.Size) * spread)
+}
+
+// normalizedImageName returns the CRI compliant name for a given image.
+// TODO: cover the corner cases of missed matches, e.g,
+// 1. Using Docker as runtime and docker.io/library/test:tag in pod spec, but only test:tag will present in node status
+// 2. Using the implicit registry, i.e., test:tag or library/test:tag in pod spec but only docker.io/library/test:tag
+// in node status; note that if users consistently use one registry format, this should not happen.
+func normalizedImageName(name string) string {
+	if strings.LastIndex(name, ":") <= strings.LastIndex(name, "/") {
+		name = name + ":latest"
+	}
+	return name
+}

examples/imagelocality/main.go

@@ -0,0 +1,39 @@
+/*
+   Copyright 2023 The Kubernetes Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package main
+
+import (
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/handle/sharedlister"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/score"
+)
+
+// main is compiled to an exported Wasm function named "_start", called by the
+// Wasm scheduler plugin during initialization.
+func main() {
+	// The plugin package uses only normal Go code, which allows it to be
+	// unit testable via `tinygo test -target=wasi` as well normal `go test`.
+	//
+	// The real implementations, such as `config.Get()` use Wasm host functions
+	// (go:wasmimport), which cannot be tested with `tinygo test -target=wasi`.
+	plugin := &imageLocality{
+		sharedLister: sharedlister.Get(),
+	}
+	// Instead of using `plugin.Set`, this configures only the interfaces
+	// implemented by the plugin. The Wasm host only calls functions imported,
+	// so this prevents additional overhead.
+	score.SetPlugin(plugin)
+}

guest/api/types.go

@@ -156,6 +156,7 @@ type NodeInfo interface {
 	proto.Metadata
 
 	Node() proto.Node
+	ImageStates() map[string]*ImageStateSummary
 
 	// ... we'll support more fields of NodeInfo of the scheduling framework.
 }

guest/filter/filter.go

@@ -20,7 +20,7 @@ package filter
 
 import (
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api"
-	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api/proto"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/handle/sharedlister"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/internal/cyclestate"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/internal/imports"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/internal/plugin"
@@ -65,44 +65,12 @@ func _filter() uint32 { //nolint
 		return 0
 	}
 
-	s := filter.Filter(cyclestate.Values, cyclestate.Pod, &nodeInfo{})
-
-	return imports.StatusToCode(s)
-}
-
-var _ api.NodeInfo = (*nodeInfo)(nil)
-
-// nodeInfo is lazy so that a plugin which doesn't read fields avoids a
-// relatively expensive unmarshal penalty.
-//
-// Note: Unlike proto.Pod, this is not special cased for the scheduling cycle.
-type nodeInfo struct {
-	api.NodeInfo
-}
-
-func (n *nodeInfo) GetUid() string {
-	return n.lazyNode().GetUid()
-}
-
-func (n *nodeInfo) GetName() string {
-	return n.lazyNode().GetName()
-}
-
-func (n *nodeInfo) GetNamespace() string {
-	return n.lazyNode().GetNamespace()
-}
-
-func (n *nodeInfo) GetResourceVersion() string {
-	return n.lazyNode().GetResourceVersion()
-}
-
-func (n *nodeInfo) Node() proto.Node {
-	return n.lazyNode()
-}
-
-// lazyNode lazy initializes node from imports.Node.
-func (n *nodeInfo) lazyNode() proto.Node {
 	nodename := imports.CurrentNodeName()
+	if nodename == "" {
+		return imports.StatusToCode(&api.Status{Code: api.StatusCodeError, Reason: "could not get current node name"})
+	}
 
-	return cyclestate.Nodes.Get(nodename).Node()
+	s := filter.Filter(cyclestate.Values, cyclestate.Pod, sharedlister.NodeInfos().Get(nodename))
+
+	return imports.StatusToCode(s)
 }

guest/handle/eventrecorder/eventrecorder.go

@@ -46,8 +46,6 @@ func EventfFn(msg internal.EventMessage) {
 }
 
 // Eventf is a convenience that calls the same method documented on api.Eventf.
-//
-// Note: See Info for unit test and benchmarking impact.
 func Eventf(regarding proto.KObject, related proto.KObject, eventtype, reason, action, note string) {
 	eventRecorderInstance.Eventf(regarding, related, eventtype, reason, action, note)
 }

guest/handle/sharedlister/api/sharedlister.go

@@ -0,0 +1,15 @@
+package api
+
+import (
+	guestapi "sigs.k8s.io/kube-scheduler-wasm-extension/guest/api"
+)
+
+type SharedLister interface {
+	NodeInfos() guestapi.NodeInfoList
+}
+
+type UnimplementedSharedLister struct{}
+
+func (UnimplementedSharedLister) NodeInfos() guestapi.NodeInfoList {
+	return nil
+}

guest/handle/sharedlister/internal/sharedlister.go

@@ -0,0 +1,11 @@
+package internal
+
+import guestapi "sigs.k8s.io/kube-scheduler-wasm-extension/guest/api"
+
+type SharedLister struct {
+	NodeInfoList guestapi.NodeInfoList
+}
+
+func (s SharedLister) NodeInfos() guestapi.NodeInfoList {
+	return s.NodeInfoList
+}

guest/handle/sharedlister/sharedlister.go

@@ -0,0 +1,21 @@
+package sharedlister
+
+import (
+	guestapi "sigs.k8s.io/kube-scheduler-wasm-extension/guest/api"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/handle/sharedlister/api"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/handle/sharedlister/internal"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/internal/prefilter"
+)
+
+var sharedListerInstance api.SharedLister = &internal.SharedLister{
+	NodeInfoList: prefilter.Nodes,
+}
+
+func Get() api.SharedLister {
+	return sharedListerInstance
+}
+
+// NodeInfos is a convenience that calls the same method documented on api.NodeInfos.
+func NodeInfos() guestapi.NodeInfoList {
+	return sharedListerInstance.NodeInfos()
+}

guest/internal/cyclestate/cyclestate.go

@@ -19,5 +19,3 @@ import (
 var Pod proto.Pod = prefilter.Pod
 
 var Values api.CycleState = prefilter.CycleState
-
-var Nodes api.NodeInfoList = prefilter.Nodes

guest/internal/imports/host.go

@@ -54,6 +54,21 @@ func CurrentNodeName() string {
 	})
 }
 
+func NodeImageStates(nodeName string) map[string]*api.ImageStateSummary {
+	nodenamePtr, nodenameSize := mem.StringToPtr(nodeName)
+	// Wrap to avoid TinyGo 0.28: cannot use an exported function as value
+	jsonStr := mem.GetString(func(ptr uint32, limit mem.BufLimit) (len uint32) {
+		return k8sSchedulerNodeImageStates(nodenamePtr, nodenameSize, ptr, limit)
+	})
+	byte := []byte(jsonStr)
+	var nodeImageStates map[string]*api.ImageStateSummary
+	err := json.Unmarshal(byte, &nodeImageStates)
+	if err != nil {
+		panic(err)
+	}
+	return nodeImageStates
+}
+
 func Node(nodeName string, updater func([]byte) error) error {
 	nodenamePtr, nodenameSize := mem.StringToPtr(nodeName)
 	// Wrap to avoid TinyGo 0.28: cannot use an exported function as value

guest/internal/imports/imports.go

@@ -40,3 +40,6 @@ func k8sSchedulerCurrentNodeName(uint32, mem.BufLimit) (len uint32)
 
 //go:wasmimport k8s.io/scheduler currentPod
 func k8sSchedulerCurrentPod(ptr uint32, limit mem.BufLimit) (len uint32)
+
+//go:wasmimport k8s.io/scheduler nodeImageStates
+func k8sSchedulerNodeImageStates(uint32, uint32, uint32, mem.BufLimit) (len uint32)

guest/internal/imports/imports_stub.go

@@ -40,3 +40,6 @@ func k8sSchedulerCurrentNodeName(uint32, mem.BufLimit) (len uint32) { return }
 
 // k8sSchedulerCurrentPod is stubbed for compilation outside TinyGo.
 func k8sSchedulerCurrentPod(uint32, mem.BufLimit) (len uint32) { return }
+
+// k8sSchedulerNodeImageStates is stubbed for compilation outside TinyGo.
+func k8sSchedulerNodeImageStates(uint32, uint32, uint32, mem.BufLimit) (len uint32) { return }