Support SharedLister for NodeInfo

Makefile

@@ -9,7 +9,7 @@ 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 \
+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 guest/testdata/permit/main.wasm \
 			  internal/e2e/scheduler_perf/wasm/nodenumber/main.wasm 
 

examples/advanced/plugin/plugin.go

@@ -74,7 +74,7 @@ func (pl *NodeNumber) EventsToRegister() []api.ClusterEvent {
 }
 
 // PreScore implements api.PreScorePlugin
-func (pl *NodeNumber) PreScore(state api.CycleState, pod proto.Pod, _ proto.NodeList) *api.Status {
+func (pl *NodeNumber) PreScore(state api.CycleState, pod proto.Pod, _ api.NodeInfoList) *api.Status {
 	pl.klog.InfoS("execute PreScore on NodeNumber plugin", "pod", klog.KObj(pod))
 
 	podnum, ok := lastNumber(pod.Spec().GetNodeName())

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)
+}

examples/nodenumber/main.go

@@ -91,7 +91,7 @@ func (pl *NodeNumber) EventsToRegister() []api.ClusterEvent {
 }
 
 // PreScore implements api.PreScorePlugin
-func (pl *NodeNumber) PreScore(state api.CycleState, pod proto.Pod, _ proto.NodeList) *api.Status {
+func (pl *NodeNumber) PreScore(state api.CycleState, pod proto.Pod, _ api.NodeInfoList) *api.Status {
 	recorder := eventrecorder.Get()
 
 	klog.InfoS("execute PreScore on NodeNumber plugin", "pod", klog.KObj(pod))

guest/api/framework.go

@@ -0,0 +1,9 @@
+package api
+
+// ImageStateSummary provides summarized information about the state of an image.
+type ImageStateSummary struct {
+	// Size of the image
+	Size int64
+	// Used to track how many nodes have this image
+	NumNodes int
+}

guest/api/types.go

@@ -105,7 +105,7 @@ type EnqueueExtensions interface {
 type PreScorePlugin interface {
 	Plugin
 
-	PreScore(state CycleState, pod proto.Pod, nodeList proto.NodeList) *Status
+	PreScore(state CycleState, pod proto.Pod, nodeList NodeInfoList) *Status
 }
 
 // ScorePlugin is a WebAssembly implementation of framework.ScorePlugin.
@@ -161,11 +161,28 @@ type PostBindPlugin interface {
 	PostBind(state CycleState, pod proto.Pod, nodeName string)
 }
 
+type NodeInfoList interface {
+	// List lists all NodeInfo in this list.
+	List() []NodeInfo
+	// Get returns the NodeInfo with the given name.
+	// It returns nil if this list doesn't have the NodeInfo with the given name.
+	Get(name string) NodeInfo
+}
+
 type NodeInfo interface {
 	// Metadata is a convenience that triggers Get.
 	proto.Metadata
 
 	Node() proto.Node
+	ImageStates() map[string]*ImageStateSummary
+
+	// ... we'll support more fields of NodeInfo of the scheduling framework.
+}
+
+type PodInfoList interface {
+	// Get returns the PodInfo with the given name.
+	// It returns nil if this list doesn't have the PodInfo with the given name.
+	Get(name string) PodInfo
 }
 
 type PodInfo interface {

guest/bind/bind.go

@@ -60,7 +60,7 @@ func _bind() uint32 { //nolint
 		return 0
 	}
 
-	nodeName := imports.NodeName()
+	nodeName := imports.CurrentNodeName()
 	// The parameters passed are lazy with regard to host functions. This means
 	// a no-op plugin should not have any unmarshal penalty.
 	s := bind.Bind(cyclestate.Values, cyclestate.Pod, nodeName)

guest/filter/filter.go

@@ -20,12 +20,10 @@ 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"
-	internalproto "sigs.k8s.io/kube-scheduler-wasm-extension/guest/internal/proto"
-	protoapi "sigs.k8s.io/kube-scheduler-wasm-extension/kubernetes/proto/api"
 )
 
 // filter is the current plugin assigned with SetPlugin.
@@ -67,105 +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 {
-	node proto.Node
-}
-
-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 {
-	if node := n.node; node != nil {
-		return node
-	}
-
-	var msg protoapi.Node
-	if err := imports.Node(msg.UnmarshalVT); err != nil {
-		panic(err.Error())
+	nodename := imports.CurrentNodeName()
+	if nodename == "" {
+		return imports.StatusToCode(&api.Status{Code: api.StatusCodeError, Reason: "could not get current node name"})
 	}
-	n.node = &internalproto.Node{Msg: &msg}
-	return n.node
-}
-
-type PodInfo struct {
-	pod proto.Pod
-}
-
-func (p *PodInfo) GetApiVersion() string {
-	return p.lazyPod().GetApiVersion()
-}
-
-func (p *PodInfo) GetKind() string {
-	return p.lazyPod().GetKind()
-}
-
-func (p *PodInfo) GetName() string {
-	return p.lazyPod().GetName()
-}
-
-func (p *PodInfo) GetNamespace() string {
-	return p.lazyPod().GetNamespace()
-}
-
-func (p *PodInfo) GetResourceVersion() string {
-	return p.lazyPod().GetNamespace()
-}
 
-func (p *PodInfo) GetUid() string {
-	return p.lazyPod().GetUid()
-}
-
-func (p *PodInfo) Pod() proto.Pod {
-	return p.lazyPod()
-}
+	s := filter.Filter(cyclestate.Values, cyclestate.Pod, sharedlister.NodeInfos().Get(nodename))
 
-func (p *PodInfo) Spec() *protoapi.PodSpec {
-	return p.lazyPod().Spec()
-}
-
-func (p *PodInfo) Status() *protoapi.PodStatus {
-	return p.lazyPod().Status()
-}
-
-// lazyPod lazy initializes pod from imports.Pod.
-func (p *PodInfo) lazyPod() proto.Pod {
-	if pod := p.pod; pod != nil {
-		return pod
-	}
-
-	var msg protoapi.Pod
-	if err := imports.Pod(msg.UnmarshalVT); err != nil {
-		panic(err.Error())
-	}
-	p.pod = &internalproto.Pod{Msg: &msg}
-	return p.pod
+	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)
 }