Rainbond/gpushare-scheduler-extender/pkg/cache/cache.go

package cache

import (
	"github.com/AliyunContainerService/gpushare-scheduler-extender/pkg/log"
	"sync"

	"github.com/AliyunContainerService/gpushare-scheduler-extender/pkg/utils"
	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/types"
	corelisters "k8s.io/client-go/listers/core/v1"
)

type SchedulerCache struct {

	// a map from pod key to podState.
	nodes map[string]*NodeInfo

	// nodeLister can list/get nodes from the shared informer's store.
	nodeLister corelisters.NodeLister

	//
	podLister corelisters.PodLister

	// record the knownPod, it will be added when annotation ALIYUN_GPU_ID is added, and will be removed when complete and deleted
	knownPods map[types.UID]*v1.Pod
	nLock     *sync.RWMutex
}

func NewSchedulerCache(nLister corelisters.NodeLister, pLister corelisters.PodLister) *SchedulerCache {
	return &SchedulerCache{
		nodes:      make(map[string]*NodeInfo),
		nodeLister: nLister,
		podLister:  pLister,
		knownPods:  make(map[types.UID]*v1.Pod),
		nLock:      new(sync.RWMutex),
	}
}

func (cache *SchedulerCache) GetNodeinfos() []*NodeInfo {
	nodes := []*NodeInfo{}
	for _, n := range cache.nodes {
		nodes = append(nodes, n)
	}
	return nodes
}

// build cache when initializing
func (cache *SchedulerCache) BuildCache() error {
	log.V(5).Info("debug: begin to build scheduler cache")
	pods, err := cache.podLister.List(labels.Everything())

	if err != nil {
		return err
	} else {
		for _, pod := range pods {
			if utils.GetGPUMemoryFromPodAnnotation(pod) <= uint(0) {
				continue
			}

			if len(pod.Spec.NodeName) == 0 {
				continue
			}

			err = cache.AddOrUpdatePod(pod)
			if err != nil {
				return err
			}
		}

	}

	return nil
}

func (cache *SchedulerCache) GetPod(name, namespace string) (*v1.Pod, error) {
	return cache.podLister.Pods(namespace).Get(name)
}

// Get known pod from the pod UID
func (cache *SchedulerCache) KnownPod(podUID types.UID) bool {
	cache.nLock.RLock()
	defer cache.nLock.RUnlock()

	_, found := cache.knownPods[podUID]
	return found
}

func (cache *SchedulerCache) AddOrUpdatePod(pod *v1.Pod) error {
	log.V(100).Info("debug: Add or update pod info: %v", pod)
	log.V(100).Info("debug: Node %v", cache.nodes)
	if len(pod.Spec.NodeName) == 0 {
		log.V(100).Info("debug: pod %s in ns %s is not assigned to any node, skip", pod.Name, pod.Namespace)
		return nil
	}

	n, err := cache.GetNodeInfo(pod.Spec.NodeName)
	if err != nil {
		return err
	}
	podCopy := pod.DeepCopy()
	if n.addOrUpdatePod(podCopy) {
		// put it into known pod
		cache.rememberPod(pod.UID, podCopy)
	} else {
		log.V(100).Info("debug: pod %s in ns %s's gpu id is %d, it's illegal, skip",
			pod.Name,
			pod.Namespace,
			utils.GetGPUIDFromAnnotation(pod))
	}

	return nil
}

// The lock is in cacheNode
func (cache *SchedulerCache) RemovePod(pod *v1.Pod) {
	log.V(100).Info("debug: Remove pod info: %v", pod)
	log.V(100).Info("debug: Node %v", cache.nodes)
	n, err := cache.GetNodeInfo(pod.Spec.NodeName)
	if err == nil {
		n.removePod(pod)
	} else {
		log.V(10).Info("debug: Failed to get node %s due to %v", pod.Spec.NodeName, err)
	}

	cache.forgetPod(pod.UID)
}

// Get or build nodeInfo if it doesn't exist
func (cache *SchedulerCache) GetNodeInfo(name string) (*NodeInfo, error) {
	node, err := cache.nodeLister.Get(name)
	if err != nil {
		return nil, err
	}

	cache.nLock.Lock()
	defer cache.nLock.Unlock()
	n, ok := cache.nodes[name]

	if !ok {
		n = NewNodeInfo(node)
		cache.nodes[name] = n
	} else {
		// if the existing node turn from non gpushare to gpushare
		// if (utils.GetTotalGPUMemory(n.node) <= 0 && utils.GetTotalGPUMemory(node) > 0) ||
		// 	(utils.GetGPUCountInNode(n.node) <= 0 && utils.GetGPUCountInNode(node) > 0) ||
		// 	// if the existing node turn from gpushare to non gpushare
		// 	(utils.GetTotalGPUMemory(n.node) > 0 && utils.GetTotalGPUMemory(node) <= 0) ||
		// 	(utils.GetGPUCountInNode(n.node) > 0 && utils.GetGPUCountInNode(node) <= 0) {
		if len(cache.nodes[name].devs) == 0 ||
			utils.GetTotalGPUMemory(n.node) <= 0 ||
			utils.GetGPUCountInNode(n.node) <= 0 {
			log.V(10).Info("info: GetNodeInfo() need update node %s",
				name)

			// fix the scenario that the number of devices changes from 0 to an positive number
			cache.nodes[name].Reset(node)
			log.V(10).Info("info: node: %s, labels from cache after been updated: %v", n.node.Name, n.node.Labels)
		} else {
			log.V(10).Info("info: GetNodeInfo() uses the existing nodeInfo for %s", name)
		}
		log.V(100).Info("debug: node %s with devices %v", name, n.devs)
	}
	return n, nil
}

func (cache *SchedulerCache) forgetPod(uid types.UID) {
	cache.nLock.Lock()
	defer cache.nLock.Unlock()
	delete(cache.knownPods, uid)
}

func (cache *SchedulerCache) rememberPod(uid types.UID, pod *v1.Pod) {
	cache.nLock.Lock()
	defer cache.nLock.Unlock()
	cache.knownPods[pod.UID] = pod
}