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 }