tnsm-plugin.go

package poweraware

import (
	"context"
	"math"
	"strconv"

	"gonum.org/v1/gonum/mat"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/klog/v2"
	"k8s.io/kubernetes/pkg/scheduler/framework"
	frameworkruntime "k8s.io/kubernetes/pkg/scheduler/framework/runtime"
)

// ---- Plugin identity ----
const (
	Name             = "PowerAware"
	preScoreStateKey = "PreScore" + Name
)

type CriteriaType string
type OptimizerType string

const (
	Positive CriteriaType = "Positive" // (kept for parity)
	Negative CriteriaType = "Negative"

	Topsis OptimizerType = "Topsis"
	Nsga2  OptimizerType = "Nsga2"
	DNsga2 OptimizerType = "DNsga2"
)

type Criteria struct {
	metav1.TypeMeta `json:",inline,omitempty"`
	Weight          float64      `json:"weight"`
	Type            CriteriaType `json:"type"` // Positive/Negative (TOPSIS sign)
}

type PowerAwareArgs struct {
	metav1.TypeMeta `json:",inline,omitempty"`

	// Objectives from the paper:
	PowerCriteria   Criteria `json:"powerCriteria"`            // P (avg power)           -> Cost
	CPULBCriteria   Criteria `json:"cpuLBCriteria"`            // L_cpu (CoV)             -> Cost
	MemLBCriteria   Criteria `json:"memLBCriteria"`            // L_mem (CoV)             -> Cost
	PodLBCriteria   Criteria `json:"podLBCriteria"`            // L_pod (CoV)             -> Cost
	CPUFragCriteria Criteria `json:"cpuFragmentationCriteria"` // F_cpu (avg) -> Cost
	MemFragCriteria Criteria `json:"memFragmentationCriteria"` // F_mem (avg) -> Cost

	Optimizer      OptimizerType `json:"optimizer"`
	PopulationSize int           `json:"populationSize,omitempty"`
	NumGenerations int           `json:"numGenerations,omitempty"`
}

// ---- Plugin ----
type PowerAware struct {
	logger klog.Logger
	handle framework.Handle

	PowerCriteria   Criteria
	CPULBCriteria   Criteria
	MemLBCriteria   Criteria
	PodLBCriteria   Criteria
	CPUFragCriteria Criteria
	MemFragCriteria Criteria

	Optimizer      OptimizerType
	PopulationSize int
	NumGenerations int
}

var (
	_ framework.PreScorePlugin = &PowerAware{}
	_ framework.ScorePlugin    = &PowerAware{}
)

func (pl *PowerAware) Name() string { return Name }

// ---- PreScore scratch ----
type preScoreState struct {
	// TOPSIS: relative closeness per node
	closeness map[string]float64

	// If NSGA-II/DNSGA-II picked a single best node, stash it here
	bestNode string
}

func (s *preScoreState) Clone() framework.StateData { return s }

// ---- Helpers: quantities, stats, power, decision matrix ----
type nodeSnapshot struct {
	name      string
	allocCPU  int64 // milli-CPU
	allocMem  int64 // bytes
	reqCPU    int64 // milli-CPU (sum of scheduled pods' requests)
	reqMem    int64 // bytes
	podCount  int
	idlePower float64 // from annotation "phd.uqtr.ca/power_idle"
	maxPower  float64 // from annotation "phd.uqtr.ca/power_max"
	alpha     float64 // from annotation "phd.uqtr.ca/alpha"
}

func mcpu(q v1.ResourceList, res v1.ResourceName) int64 {
	if qty, ok := q[res]; ok {
		return qty.MilliValue()
	}
	return 0
}

func bytesOf(q v1.ResourceList, res v1.ResourceName) int64 {
	if qty, ok := q[res]; ok {
		return qty.Value()
	}
	return 0
}

func podRequests(pod *v1.Pod) (mcpuReq int64, memReq int64) {
	for i := range pod.Spec.Containers {
		c := &pod.Spec.Containers[i]
		if cpu, ok := c.Resources.Limits[v1.ResourceCPU]; ok {
			mcpuReq += cpu.MilliValue()
		} else if cpu, ok := c.Resources.Requests[v1.ResourceCPU]; ok {
			mcpuReq += cpu.MilliValue()
		}
		if mem, ok := c.Resources.Limits[v1.ResourceMemory]; ok {
			memReq += mem.Value()
		} else if mem, ok := c.Resources.Requests[v1.ResourceMemory]; ok {
			memReq += mem.Value()
		}
	}
	return
}

// coefficient of variation = std/mean (returns 0 if mean==0)
func cov(vals []float64) float64 {
	if len(vals) == 0 {
		return 0
	}
	var sum float64
	for _, v := range vals {
		sum += v
	}
	mean := sum / float64(len(vals))
	if mean == 0 {
		return 0
	}
	var sq float64
	for _, v := range vals {
		d := v - mean
		sq += d * d
	}
	std := math.Sqrt(sq / float64(len(vals)))
	return std / mean
}

// ---- PreScore: build the decision matrix with projected placement ----
func (pl *PowerAware) PreScore(ctx context.Context, state *framework.CycleState, pod *v1.Pod, nodes []*v1.Node) *framework.Status {
	klog.InfoS("PowerAware.PreScore", "pod", klog.KObj(pod), "candidates", len(nodes))

	scoreState := &preScoreState{closeness: make(map[string]float64)}

	// Snapshot lister gives us per-node requested/allocatable and pods
	snapshot := pl.handle.SnapshotSharedLister()
	if snapshot == nil {
		klog.Error("SnapshotSharedLister is nil")
		state.Write(preScoreStateKey, scoreState)
		return nil
	}

	// Build snapshots for all candidate nodes
	nodeInfos, err := snapshot.NodeInfos().List()
	if err != nil {
		klog.ErrorS(err, "listing NodeInfos failed")
		state.Write(preScoreStateKey, scoreState)
		return nil
	}

	// index NodeInfo by name for quick lookup
	ninfo := map[string]framework.NodeInfo{}
	for _, ni := range nodeInfos {
		ninfo[ni.Node().Name] = *ni
	}

	// Gather base metrics per node
	ns := make([]nodeSnapshot, 0, len(nodes))
	for _, node := range nodes {
		ni, ok := ninfo[node.Name]
		if !ok {
			// If not in snapshot (rare), skip this candidate
			continue
		}

		allocCPU := ni.Allocatable.MilliCPU
		allocMem := ni.Allocatable.Memory

		reqCPU := ni.Requested.MilliCPU
		reqMem := ni.Requested.Memory

		podCount := len(ni.Pods)

		// power annotations
		idleStr := node.Annotations["phd.uqtr.ca/power_idle"]
		maxStr := node.Annotations["phd.uqtr.ca/power_max"]
		alphaStr := node.Annotations["phd.uqtr.ca/alpha"]
		idle, _ := strconv.ParseFloat(idleStr, 64)
		maxp, _ := strconv.ParseFloat(maxStr, 64)
		alpha, _ := strconv.ParseFloat(alphaStr, 64)
		if idle <= 0 {
			idle = 1
		}
		if maxp <= idle {
			maxp = idle + 1
		}
		if alpha <= 0 {
			alpha = 1
		}

		ns = append(ns, nodeSnapshot{
			name:      node.Name,
			allocCPU:  int64(allocCPU),
			allocMem:  int64(allocMem),
			reqCPU:    int64(reqCPU),
			reqMem:    int64(reqMem),
			podCount:  podCount,
			idlePower: idle,
			maxPower:  maxp,
			alpha:     alpha,
		})
	}

	if len(ns) == 0 {
		state.Write(preScoreStateKey, scoreState)
		return nil
	}

	// Pod requests
	podCPU, podMem := podRequests(pod)

	// Decision matrix columns in order:
	// 0: P (avg power) | 1: L_cpu | 2: L_mem | 3: L_pod | 4: F_cpu | 5: F_mem
	const numCriterias = 6
	dm := mat.NewDense(len(ns), numCriterias, nil)
	nodeIndexToName := make([]string, len(ns))
	for i := range ns {
		nodeIndexToName[i] = ns[i].name
	}

	// Precompute base arrays (without the new pod)
	baseCPUAlloc := make([]float64, len(ns)) // allocated CPU (mCPU)
	baseMemAlloc := make([]float64, len(ns)) // allocated MEM (bytes)
	basePods := make([]float64, len(ns))
	baseCPUCap := make([]float64, len(ns))
	baseMemCap := make([]float64, len(ns))
	for i, s := range ns {
		baseCPUAlloc[i] = float64(s.reqCPU)
		baseMemAlloc[i] = float64(s.reqMem)
		basePods[i] = float64(s.podCount)
		baseCPUCap[i] = float64(s.allocCPU)
		baseMemCap[i] = float64(s.allocMem)
	}

	// For each candidate node k, project the pod onto k and compute objectives
	for k := range ns {
		RCpu := make([]float64, len(ns))
		RMem := make([]float64, len(ns))
		RPod := make([]float64, len(ns))
		copy(RCpu, baseCPUAlloc)
		copy(RMem, baseMemAlloc)
		copy(RPod, basePods)

		// place pod on node k
		RCpu[k] += float64(podCPU)
		RMem[k] += float64(podMem)
		RPod[k] += 1

		// --- Power objective (cluster average) ---
		var pAvg float64
		for i, s := range ns {
			u := 0.0
			if s.allocCPU > 0 {
				u = RCpu[i] / float64(s.allocCPU) // utilization in [0,1]+
			}
			pAvg += PowerModel(s.idlePower, s.maxPower, s.alpha, u)
		}
		pAvg /= float64(len(ns))

		// --- Load balancing CoVs ---
		Lcpu := cov(RCpu)
		Lmem := cov(RMem)
		Lpod := cov(RPod)

		// --- Fragmentation averages ---
		var Fcpu, Fmem float64
		for i := range ns {
			fc := 0.0
			fm := 0.0
			if baseCPUCap[i] > 0 {
				fc = (baseCPUCap[i] - RCpu[i]) / baseCPUCap[i]
				if fc < 0 {
					fc = 0
				} // clamp if over-requested
			}
			if baseMemCap[i] > 0 {
				fm = (baseMemCap[i] - RMem[i]) / baseMemCap[i]
				if fm < 0 {
					fm = 0
				}
			}
			Fcpu += fc
			Fmem += fm
		}
		Fcpu /= float64(len(ns))
		Fmem /= float64(len(ns))

		// Fill row k
		dm.Set(k, 0, pAvg)
		dm.Set(k, 1, Lcpu)
		dm.Set(k, 2, Lmem)
		dm.Set(k, 3, Lpod)
		dm.Set(k, 4, Fcpu)
		dm.Set(k, 5, Fmem)
	}

	klog.Infof("Optimize using %s", pl.Optimizer)

	switch pl.Optimizer {
	case Topsis:
		// Weights in same order as columns
		weights := []float64{
			pl.PowerCriteria.Weight,
			pl.CPULBCriteria.Weight,
			pl.MemLBCriteria.Weight,
			pl.PodLBCriteria.Weight,
			pl.CPUFragCriteria.Weight,
			pl.MemFragCriteria.Weight,
		}

		// Normalize -> Weight -> Ideal/Nadir -> Distances -> Closeness
		n := normalizeMatrix(dm)
		w := weightMatrix(n, weights)
		ideal, nadir := idealSolutions(w)
		dI, dN := separationMeasure(w, ideal, nadir)
		closeness := relativeCloseness(dI, dN)

		for i, v := range closeness {
			scoreState.closeness[nodeIndexToName[i]] = v
		}

	case Nsga2, DNsga2:
		// Minimal wrapper to your existing optimize() interface.
		// It should interpret dm rows as alternatives and columns as costs.
		mode := "nsga2"
		if pl.Optimizer == DNsga2 {
			mode = "dnsga2"
		}
		opt, err := Nsga2Optimize(dm, len(ns), numCriterias, mode)
		if err != nil {
			klog.ErrorS(err, "optimize failed")
			state.Write(preScoreStateKey, scoreState)
			return nil
		}
		best := nodeIndexToName[opt.BestNode]
		scoreState.bestNode = best
		// give a high closeness to the chosen node; others 0
		for _, name := range nodeIndexToName {
			if name == best {
				scoreState.closeness[name] = 0.99
			} else {
				scoreState.closeness[name] = 0
			}
		}
	default:
		// Fallback: TOPSIS
		n := normalizeMatrix(dm)
		weights := []float64{
			pl.PowerCriteria.Weight,
			pl.CPULBCriteria.Weight,
			pl.MemLBCriteria.Weight,
			pl.PodLBCriteria.Weight,
			pl.CPUFragCriteria.Weight,
			pl.MemFragCriteria.Weight,
		}
		w := weightMatrix(n, weights)
		ideal, nadir := idealSolutions(w)
		dI, dN := separationMeasure(w, ideal, nadir)
		closeness := relativeCloseness(dI, dN)
		for i, v := range closeness {
			scoreState.closeness[nodeIndexToName[i]] = v
		}
	}

	state.Write(preScoreStateKey, scoreState)
	return nil
}

func (pl *PowerAware) Score(ctx context.Context, state *framework.CycleState, pod *v1.Pod, nodeName string) (int64, *framework.Status) {
	data, err := state.Read(preScoreStateKey)
	if err != nil {
		return 0, nil
	}
	s, ok := data.(*preScoreState)
	if !ok {
		return 0, framework.NewStatus(framework.Error, "invalid preScore state")
	}
	return int64(s.closeness[nodeName] * 100), nil
}

func (pl *PowerAware) ScoreExtensions() framework.ScoreExtensions { return nil }

func (pl *PowerAware) EventsToRegister() []framework.ClusterEvent {
	return []framework.ClusterEvent{
		{Resource: framework.Node, ActionType: framework.Add | framework.Update | framework.Delete},
		{Resource: framework.Pod, ActionType: framework.Add | framework.Update | framework.Delete},
	}
}

// ---- New: plugin initializer with defaults matching paper ----
func New(ctx context.Context, arg runtime.Object, h framework.Handle) (framework.Plugin, error) {
	typed := PowerAwareArgs{
		PowerCriteria:   Criteria{Weight: 0.5, Type: Negative},
		CPULBCriteria:   Criteria{Weight: 0.1, Type: Negative},
		MemLBCriteria:   Criteria{Weight: 0.1, Type: Negative},
		PodLBCriteria:   Criteria{Weight: 0.1, Type: Negative},
		CPUFragCriteria: Criteria{Weight: 0.1, Type: Negative},
		MemFragCriteria: Criteria{Weight: 0.1, Type: Negative},
		Optimizer:       Topsis,
		PopulationSize:  100,
		NumGenerations:  20,
	}
	if arg != nil {
		if err := frameworkruntime.DecodeInto(arg, &typed); err != nil {
			return nil, err
		}
	}

	return &PowerAware{
		handle:          h,
		PowerCriteria:   typed.PowerCriteria,
		CPULBCriteria:   typed.CPULBCriteria,
		MemLBCriteria:   typed.MemLBCriteria,
		PodLBCriteria:   typed.PodLBCriteria,
		CPUFragCriteria: typed.CPUFragCriteria,
		MemFragCriteria: typed.MemFragCriteria,
		Optimizer:       typed.Optimizer,
		PopulationSize:  typed.PopulationSize,
		NumGenerations:  typed.NumGenerations,
	}, nil
}