predicate and nodeOrder for the topology plugin + topology result caching (#356)

Author: davidLif

pkg/scheduler/actions/common/allocate.go

@@ -73,6 +73,8 @@ func allocateTask(ssn *framework.Session, stmt *framework.Statement, nodes []*no
 			task.Namespace, task.Name, node.Name)
 	}
 
+	ssn.CleanAllocationAttemptCache(job)
+
 	if success {
 		log.InfraLogger.V(6).Infof("Allocation succeeded for task: <%v/%v>", task.Namespace, task.Name)
 	} else {

pkg/scheduler/api/types.go

@@ -68,6 +68,9 @@ type OnJobSolutionStartFn func()
 // BindRequestMutateFn allows plugins to add annotations before BindRequest creation.
 type BindRequestMutateFn func(pod *pod_info.PodInfo, nodeName string) map[string]string
 
+// CleanAllocationAttemptCacheFn is used to clean the cycle cache after an allocation attempt for a job.
+type CleanAllocationAttemptCacheFn func(job *podgroup_info.PodGroupInfo) error
+
 type SchedulableResult struct {
 	IsSchedulable bool
 	Reason        v2alpha2.UnschedulableReason

pkg/scheduler/framework/session.go

@@ -81,30 +81,24 @@ type Session struct {
 	PrePredicateFns                       []api.PrePredicateFn
 	PredicateFns                          []api.PredicateFn
 	BindRequestMutateFns                  []api.BindRequestMutateFn
+	CleanAllocationAttemptCacheFns        []api.CleanAllocationAttemptCacheFn
 
 	Config          *conf.SchedulerConfiguration
 	plugins         map[string]Plugin
 	eventHandlers   []*EventHandler
 	SchedulerParams conf.SchedulerParams
 	mux             *http.ServeMux
 
-	k8sPodState map[types.UID]k8s_internal.SessionState
+	k8sResourceStateCache sync.Map
 }
 
 func (ssn *Session) Statement() *Statement {
 	return &Statement{ssn: ssn, sessionUID: ssn.UID}
 }
 
-func (ssn *Session) GetK8sStateForPod(uid types.UID) k8s_internal.SessionState {
-	if ssn.k8sPodState == nil {
-		ssn.k8sPodState = make(map[types.UID]k8s_internal.SessionState)
-	}
-	state, found := ssn.k8sPodState[uid]
-	if found {
-		return state
-	}
-	ssn.k8sPodState[uid] = k8s_internal.NewSessionState()
-	return ssn.k8sPodState[uid]
+func (ssn *Session) GetSessionStateForResource(uid types.UID) k8s_internal.SessionState {
+	state, _ := ssn.k8sResourceStateCache.LoadOrStore(uid, k8s_internal.NewSessionState())
+	return state.(k8s_internal.SessionState)
 }
 
 func (ssn *Session) BindPod(pod *pod_info.PodInfo) error {
@@ -346,10 +340,10 @@ func openSession(cache cache.Cache, sessionId types.UID, schedulerParams conf.Sc
 		Queues:        map[common_info.QueueID]*queue_info.QueueInfo{},
 		Topologies:    []*kueuev1alpha1.Topology{},
 
-		plugins:         map[string]Plugin{},
-		SchedulerParams: schedulerParams,
-		mux:             mux,
-		k8sPodState:     map[types.UID]k8s_internal.SessionState{},
+		plugins:               map[string]Plugin{},
+		SchedulerParams:       schedulerParams,
+		mux:                   mux,
+		k8sResourceStateCache: sync.Map{},
 	}
 
 	log.InfraLogger.V(2).Infof("Taking cluster snapshot ...")

pkg/scheduler/framework/session_plugins.go

@@ -375,6 +375,20 @@ func (ssn *Session) NodeOrderFn(task *pod_info.PodInfo, node *node_info.NodeInfo
 	return priorityScore, nil
 }
 
+func (ssn *Session) AddCleanAllocationAttemptCacheFn(fn api.CleanAllocationAttemptCacheFn) {
+	ssn.CleanAllocationAttemptCacheFns = append(ssn.CleanAllocationAttemptCacheFns, fn)
+}
+
+func (ssn *Session) CleanAllocationAttemptCache(job *podgroup_info.PodGroupInfo) {
+	for _, cleaner := range ssn.CleanAllocationAttemptCacheFns {
+		err := cleaner(job)
+		if err != nil {
+			log.InfraLogger.V(6).Infof(
+				"Failed to run CleanAllocationAttemptCache on job %s", job.Name)
+		}
+	}
+}
+
 func (ssn *Session) IsRestrictNodeSchedulingEnabled() bool {
 	return ssn.SchedulerParams.RestrictSchedulingNodes
 }

pkg/scheduler/k8s_internal/k8s_internal.go

@@ -52,7 +52,7 @@ func FitPrePredicateConverter(
 	nodePreFilter NodePreFilter,
 ) FitPredicatePreFilter {
 	return func(pod *v1.Pod) (sets.Set[string], *k8sframework.Status) {
-		state := stateProvider.GetK8sStateForPod(pod.UID)
+		state := stateProvider.GetSessionStateForResource(pod.UID)
 		result, status := nodePreFilter.PreFilter(context.Background(), state, pod)
 		if status != nil {
 			return nil, status
@@ -69,7 +69,7 @@ func FitPredicateConverter(
 	nodeFilter NodeFilter,
 ) FitPredicateFilter {
 	return func(pod *v1.Pod, nodeInfo *k8sframework.NodeInfo) (bool, []string, error) {
-		state := stateProvider.GetK8sStateForPod(pod.UID)
+		state := stateProvider.GetSessionStateForResource(pod.UID)
 		result := nodeFilter.Filter(context.Background(), state, pod, nodeInfo)
 		if result == nil {
 			return true, nil, nil
@@ -83,7 +83,7 @@ func PreScorePluginConverter(
 	nodeScorer ExtendedNodeScorer,
 ) PreScoreFn {
 	return func(pod *v1.Pod, fittingNodes []*k8sframework.NodeInfo) *k8sframework.Status {
-		state := stateProvider.GetK8sStateForPod(pod.UID)
+		state := stateProvider.GetSessionStateForResource(pod.UID)
 		status := nodeScorer.PreScore(context.Background(), state, pod, fittingNodes)
 		return status
 	}
@@ -94,7 +94,7 @@ func ScorePluginConverter(
 	nodeScorer NodeScorer,
 ) ScorePredicate {
 	return func(pod *v1.Pod, nodeInfo *k8sframework.NodeInfo) (int64, []string, error) {
-		state := stateProvider.GetK8sStateForPod(pod.UID)
+		state := stateProvider.GetSessionStateForResource(pod.UID)
 
 		score, result := nodeScorer.Score(context.Background(), state, pod, nodeInfo.Node().Name)
 		if result == nil {

pkg/scheduler/k8s_internal/types.go

@@ -58,5 +58,5 @@ type SessionScoreFns struct {
 type SessionState *k8sframework.CycleState
 
 type SessionStateProvider interface {
-	GetK8sStateForPod(podUID types.UID) SessionState
+	GetSessionStateForResource(podUID types.UID) SessionState
 }

pkg/scheduler/plugins/scores/scores.go

@@ -8,6 +8,7 @@ const (
 	ResourceType   = 10
 	Availability   = 100
 	GpuSharing     = 1000
-	K8sPlugins     = 10000
-	NominatedNode  = 100000
+	Topology       = 10000
+	K8sPlugins     = 100000
+	NominatedNode  = 1000000
 )

pkg/scheduler/plugins/topology/topology_plugin.go

@@ -46,6 +46,12 @@ func (t *topologyPlugin) OnSessionOpen(ssn *framework.Session) {
 
 	//pre-predicate to generate the whole topology tree and store per workload
 	ssn.AddPrePredicateFn(t.prePredicateFn)
+	//predicate to filter nodes that are related to parts of the tree that cannot accommodate the workload - this is for "required" use only
+	ssn.AddPredicateFn(t.predicateFn)
+	//node order to sort the nodes according to topology nodes score - this is for "prefer" use only
+	ssn.AddNodeOrderFn(t.nodeOrderFn)
+	//clean cycle cache after an allocation attempt for a job
+	ssn.AddCleanAllocationAttemptCacheFn(t.cleanAllocationAttemptCache)
 }
 
 func (t *topologyPlugin) initializeTopologyTree(topologies []*kueuev1alpha1.Topology, ssn *framework.Session) {

pkg/scheduler/plugins/topology/topology_plugin_job_filtering.go

@@ -4,17 +4,32 @@
 package topology
 
 import (
+	"errors"
 	"fmt"
 	"slices"
 	"sort"
+	"strings"
 
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/node_info"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/pod_info"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/podgroup_info"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/resource_info"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/log"
+	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/plugins/scores"
+	"k8s.io/apimachinery/pkg/types"
+	k8sframework "k8s.io/kubernetes/pkg/scheduler/framework"
 )
 
+type topologyStateData struct {
+	relevantDomains []*TopologyDomainInfo
+}
+
+func (t *topologyStateData) Clone() k8sframework.StateData {
+	return &topologyStateData{
+		relevantDomains: t.relevantDomains,
+	}
+}
+
 type jobAllocationMetaData struct {
 	maxPodResources    *resource_info.ResourceRequirements
 	allocationTestPods []*pod_info.PodInfo
@@ -30,6 +45,16 @@ func (t *topologyPlugin) prePredicateFn(_ *pod_info.PodInfo, job *podgroup_info.
 		return nil
 	}
 
+	// Check in cache if the job has already been allocated to a domain
+	jobAllocatableDomains, err := t.loadAllocatableDomainsFromCache(types.UID(job.PodGroupUID))
+	if err != nil {
+		return err
+	}
+	if len(jobAllocatableDomains) > 0 {
+		// Cache is already populated, no need to calculate anything
+		return nil
+	}
+
 	defer t.treeAllocatableCleanup(topologyTree)
 	maxAllocatablePods, err := t.calcTreeAllocatable(job, topologyTree)
 	if err != nil {
@@ -42,11 +67,18 @@ func (t *topologyPlugin) prePredicateFn(_ *pod_info.PodInfo, job *podgroup_info.
 		return nil
 	}
 
-	_, err = t.getBestJobAllocatableDomains(job, topologyTree)
+	jobAllocatableDomain, err := t.getBestJobAllocatableDomains(job, topologyTree)
 	if err != nil {
 		return err
 	}
 
+	//Save results to cycle cache
+	cycleJobState := (*k8sframework.CycleState)(t.sessionStateGetter.GetSessionStateForResource(job.PodGroupUID))
+	cycleJobState.Write(
+		k8sframework.StateKey(topologyPluginName),
+		&topologyStateData{relevantDomains: jobAllocatableDomain},
+	)
+
 	return nil
 }
 
@@ -278,3 +310,73 @@ func (*topologyPlugin) treeAllocatableCleanup(topologyTree *TopologyInfo) {
 		}
 	}
 }
+
+func (t *topologyPlugin) predicateFn(pod *pod_info.PodInfo, job *podgroup_info.PodGroupInfo, node *node_info.NodeInfo) error {
+	jobAllocatableDomains, err := t.loadAllocatableDomainsFromCache(job.PodGroupUID)
+	if err != nil {
+		return err
+	}
+
+	if len(jobAllocatableDomains) > 0 {
+		jobDomainsNames := []string{}
+		for _, domain := range jobAllocatableDomains {
+			if domain.Nodes[node.Node.Name] != nil {
+				return nil
+			}
+			jobDomainsNames = append(jobDomainsNames, domain.Name)
+		}
+		return fmt.Errorf("the node %s is not part of the chosen topology domain for the job %s. The chosen domains are %s",
+			node.Node.Name, job.PodGroup.Name, strings.Join(jobDomainsNames, ", "))
+	}
+
+	return nil
+}
+
+func (t *topologyPlugin) nodeOrderFn(pod *pod_info.PodInfo, node *node_info.NodeInfo) (float64, error) {
+	score := 0.0
+
+	jobAllocatableDomains, err := t.loadAllocatableDomainsFromCache(types.UID(pod.Job))
+	if err != nil {
+		return score, err
+	}
+
+	if len(jobAllocatableDomains) > 0 {
+		for _, domain := range jobAllocatableDomains {
+			if domain.Nodes[node.Node.Name] != nil {
+				score = scores.Topology
+				break
+			}
+		}
+	}
+
+	return score, nil
+}
+
+func (t *topologyPlugin) loadAllocatableDomainsFromCache(podGroupUID types.UID) ([]*TopologyDomainInfo, error) {
+	cycleJobState := (*k8sframework.CycleState)(t.sessionStateGetter.GetSessionStateForResource(podGroupUID))
+	if cycleJobState == nil {
+		return nil, nil
+	}
+	jobTopologyStateData, err := cycleJobState.Read(k8sframework.StateKey(topologyPluginName))
+	if err != nil {
+		if errors.Is(err, k8sframework.ErrNotFound) {
+			return nil, nil
+		}
+		return nil, err
+	}
+	jobAllocatableDomains := jobTopologyStateData.(*topologyStateData).relevantDomains
+	return jobAllocatableDomains, nil
+}
+
+func (t *topologyPlugin) cleanAllocationAttemptCache(job *podgroup_info.PodGroupInfo) error {
+	if job.PodGroup.Spec.TopologyConstraint.Topology == "" {
+		return nil
+	}
+
+	cycleJobState := (*k8sframework.CycleState)(t.sessionStateGetter.GetSessionStateForResource(job.PodGroupUID))
+	if cycleJobState == nil {
+		return nil
+	}
+	cycleJobState.Delete(k8sframework.StateKey(topologyPluginName))
+	return nil
+}