Topology structs updates: (#352)

1- Remove resources fields from TopologyDomainInfo
2- Saves domains by level for TopologyInfo

Author: davidLif

pkg/scheduler/plugins/topology/topology_plugin.go

@@ -4,9 +4,10 @@
 package topology
 
 import (
+	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/common_info"
 	"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/framework"
+	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/k8s_internal"
 	kueuev1alpha1 "sigs.k8s.io/kueue/apis/kueue/v1alpha1"
 )
 
@@ -16,8 +17,11 @@ const (
 )
 
 type topologyPlugin struct {
-	enabled       bool
-	TopologyTrees map[string]*TopologyInfo
+	enabled            bool
+	taskOrderFunc      common_info.LessFn
+	sessionStateGetter k8s_internal.SessionStateProvider
+	nodesInfos         map[string]*node_info.NodeInfo
+	TopologyTrees      map[string]*TopologyInfo
 }
 
 func New(pluginArgs map[string]string) framework.Plugin {
@@ -33,19 +37,17 @@ func (t *topologyPlugin) Name() string {
 
 func (t *topologyPlugin) OnSessionOpen(ssn *framework.Session) {
 	topologies := ssn.Topologies
+	t.taskOrderFunc = ssn.TaskOrderFn
+	t.sessionStateGetter = ssn
+	t.nodesInfos = ssn.Nodes
 	t.initializeTopologyTree(topologies, ssn)
-
-	ssn.AddEventHandler(&framework.EventHandler{
-		AllocateFunc:   t.handleAllocate(ssn),
-		DeallocateFunc: t.handleDeallocate(ssn),
-	})
 }
 
 func (t *topologyPlugin) initializeTopologyTree(topologies []*kueuev1alpha1.Topology, ssn *framework.Session) {
 	for _, singleTopology := range topologies {
 		topologyTree := &TopologyInfo{
 			Name:             singleTopology.Name,
-			Domains:          map[TopologyDomainID]*TopologyDomainInfo{},
+			DomainsByLevel:   map[string]map[TopologyDomainID]*TopologyDomainInfo{},
 			Root:             NewTopologyDomainInfo(TopologyDomainID("root"), "datacenter", "cluster", 0),
 			TopologyResource: singleTopology,
 		}
@@ -69,10 +71,16 @@ func (*topologyPlugin) addNodeDataToTopology(topologyTree *TopologyInfo, singleT
 		}
 
 		domainId := calcDomainId(levelIndex, singleTopology.Spec.Levels, nodeInfo.Node.Labels)
-		domainInfo, foundLevelLabel := topologyTree.Domains[domainId]
+		domainLevel := level.NodeLabel
+		domainsForLevel, foundLevelLabel := topologyTree.DomainsByLevel[domainLevel]
 		if !foundLevelLabel {
+			topologyTree.DomainsByLevel[level.NodeLabel] = map[TopologyDomainID]*TopologyDomainInfo{}
+			domainsForLevel = topologyTree.DomainsByLevel[level.NodeLabel]
+		}
+		domainInfo, foundDomain := domainsForLevel[domainId]
+		if !foundDomain {
 			domainInfo = NewTopologyDomainInfo(domainId, domainName, level.NodeLabel, levelIndex+1)
-			topologyTree.Domains[domainId] = domainInfo
+			domainsForLevel[domainId] = domainInfo
 		}
 		domainInfo.AddNode(nodeInfo)
 
@@ -86,48 +94,4 @@ func (*topologyPlugin) addNodeDataToTopology(topologyTree *TopologyInfo, singleT
 	topologyTree.Root.AddNode(nodeInfo)
 }
 
-func (t *topologyPlugin) handleAllocate(ssn *framework.Session) func(event *framework.Event) {
-	return t.updateTopologyGivenPodEvent(ssn, func(domainInfo *TopologyDomainInfo, podInfo *pod_info.PodInfo) {
-		domainInfo.AllocatedResources.AddResourceRequirements(podInfo.AcceptedResource)
-		domainInfo.AllocatedResources.BaseResource.ScalarResources()["pods"] =
-			domainInfo.AllocatedResources.BaseResource.ScalarResources()["pods"] + 1
-	})
-}
-
-func (t *topologyPlugin) handleDeallocate(ssn *framework.Session) func(event *framework.Event) {
-	return t.updateTopologyGivenPodEvent(ssn, func(domainInfo *TopologyDomainInfo, podInfo *pod_info.PodInfo) {
-		domainInfo.AllocatedResources.SubResourceRequirements(podInfo.AcceptedResource)
-		domainInfo.AllocatedResources.BaseResource.ScalarResources()["pods"] =
-			domainInfo.AllocatedResources.BaseResource.ScalarResources()["pods"] - 1
-	})
-}
-
-func (t *topologyPlugin) updateTopologyGivenPodEvent(
-	ssn *framework.Session,
-	domainUpdater func(domainInfo *TopologyDomainInfo, podInfo *pod_info.PodInfo),
-) func(event *framework.Event) {
-	return func(event *framework.Event) {
-		pod := event.Task.Pod
-		nodeName := event.Task.NodeName
-		if nodeName == noNodeName {
-			return
-		}
-		node := ssn.Nodes[nodeName].Node
-		podInfo := ssn.Nodes[nodeName].PodInfos[pod_info.PodKey(pod)]
-
-		for _, topologyTree := range t.TopologyTrees {
-			leafDomainId := calcLeafDomainId(topologyTree.TopologyResource, node.Labels)
-			domainInfo := topologyTree.Domains[leafDomainId]
-			for domainInfo != nil {
-				domainUpdater(domainInfo, podInfo)
-
-				if domainInfo.Nodes[nodeName] != nil {
-					break
-				}
-				domainInfo = domainInfo.Parent
-			}
-		}
-	}
-}
-
 func (t *topologyPlugin) OnSessionClose(ssn *framework.Session) {}