Improve allocatable resource calculation - nodeDb (#4197)

* Improve allocatable resource calculation - nodeDb - We no longer add nodes to the nodeDb if they are unschedulable with no jobs assigned to them - These nodes should never get anything assigned to them, so no need to include them in the nodeDb - This means these nodes resource is not counted towards total allocatable resource, which is used for fairshare purposes - Don't allow Armada to schedule over unallocatableResources - node now has an allocatableResources resources field, which is calculated as allocatableResources = totalResources - sum(unallocatableResources) - This is how much resource is actually alloctable, and we no longer use totalResources for fairshare calculation The net result here is: - Resource used for fairshare calculate better represents the amount of resource available to be scheduled on - Unallocatable resource is properly respected - Nodes that cannot be scheduled on do not skew the fairshare calculation The fairshare calculation still is "wrong" in the case an unschedulable node has jobs running on it, but those jobs don't fill the node - The remaning resource is used in the fairshare calculation, but can never actually be allocated Signed-off-by: JamesMurkin <[email protected]> * Simplify setting evicted priority allocatable resource Signed-off-by: JamesMurkin <[email protected]> * Remove unused func Signed-off-by: JamesMurkin <[email protected]> --------- Signed-off-by: JamesMurkin <[email protected]>
armadaproject · Feb 10, 2025 · 6569e24 · 6569e24
1 parent eb1324c
commit 6569e24
Show file tree

Hide file tree

Showing 11 changed files with 83 additions and 27 deletions.
diff --git a/e2e/setup/kind.yaml b/e2e/setup/kind.yaml
@@ -10,6 +10,8 @@ nodes:
   image: kindest/node:v1.26.15
 - role: control-plane
   image: kindest/node:v1.26.15
+  labels:
+    armadaproject.io/pool: control-plane
   kubeadmConfigPatches:
   - |
     kind: InitConfiguration

diff --git a/internal/scheduler/internaltypes/node.go b/internal/scheduler/internaltypes/node.go
@@ -2,7 +2,6 @@ package internaltypes
 
 import (
 	"fmt"
-	"math"
 
 	"golang.org/x/exp/maps"
 	v1 "k8s.io/api/core/v1"
@@ -41,8 +40,13 @@ type Node struct {
 	taints []v1.Taint
 	labels map[string]string
 
-	// Total space allocatable on this node
+	unschedulable bool
+
+	// Total space on this node
 	totalResources ResourceList
+	// Total space allocatable by armada jobs on this node
+	// allocatableResources = totalResources - sum(unallocatableResources)
+	allocatableResources ResourceList
 
 	unallocatableResources map[int32]ResourceList
 
@@ -63,20 +67,19 @@ func FromSchedulerObjectsNode(node *schedulerobjects.Node,
 	resourceListFactory *ResourceListFactory,
 ) *Node {
 	totalResources := resourceListFactory.FromNodeProto(node.TotalResources.Resources)
-
-	allocatableByPriority := map[int32]ResourceList{}
-	for _, p := range allowedPriorities {
-		allocatableByPriority[p] = totalResources
-	}
+	allocatableResources := totalResources
+	unallocatableResources := make(map[int32]ResourceList, len(node.UnallocatableResources))
 	for p, rl := range node.UnallocatableResources {
-		MarkAllocated(allocatableByPriority, p, resourceListFactory.FromJobResourceListIgnoreUnknown(rl.Resources))
+		resource := resourceListFactory.FromJobResourceListIgnoreUnknown(rl.Resources)
+		allocatableResources = allocatableResources.Subtract(resource)
+		unallocatableResources[p] = resource
 	}
-	allocatableByPriority[EvictedPriority] = allocatableByPriority[minInt32(allowedPriorities)]
 
-	unallocatableResources := map[int32]ResourceList{}
-	for p, u := range node.UnallocatableResources {
-		unallocatableResources[p] = resourceListFactory.FromJobResourceListIgnoreUnknown(u.Resources)
+	allocatableByPriority := map[int32]ResourceList{}
+	for _, p := range allowedPriorities {
+		allocatableByPriority[p] = allocatableResources
 	}
+	allocatableByPriority[EvictedPriority] = allocatableResources
 
 	return CreateNodeAndType(
 		node.Id,
@@ -90,6 +93,7 @@ func FromSchedulerObjectsNode(node *schedulerobjects.Node,
 		indexedTaints,
 		indexedNodeLabels,
 		totalResources,
+		allocatableResources,
 		unallocatableResources,
 		allocatableByPriority,
 	)
@@ -107,6 +111,7 @@ func CreateNodeAndType(
 	indexedTaints map[string]bool,
 	indexedNodeLabels map[string]bool,
 	totalResources ResourceList,
+	allocatableResources ResourceList,
 	unallocatableResources map[int32]ResourceList,
 	allocatableByPriority map[int32]ResourceList,
 ) *Node {
@@ -137,7 +142,9 @@ func CreateNodeAndType(
 		pool,
 		taints,
 		labels,
+		unschedulable,
 		totalResources,
+		allocatableResources,
 		unallocatableResources,
 		allocatableByPriority,
 		map[string]ResourceList{},
@@ -155,7 +162,9 @@ func CreateNode(
 	pool string,
 	taints []v1.Taint,
 	labels map[string]string,
+	unschedulable bool,
 	totalResources ResourceList,
+	allocatableResources ResourceList,
 	unallocatableResources map[int32]ResourceList,
 	allocatableByPriority map[int32]ResourceList,
 	allocatedByQueue map[string]ResourceList,
@@ -172,7 +181,9 @@ func CreateNode(
 		pool:                   pool,
 		taints:                 koTaint.DeepCopyTaints(taints),
 		labels:                 deepCopyLabels(labels),
+		unschedulable:          unschedulable,
 		totalResources:         totalResources,
+		allocatableResources:   allocatableResources,
 		unallocatableResources: maps.Clone(unallocatableResources),
 		AllocatableByPriority:  maps.Clone(allocatableByPriority),
 		AllocatedByQueue:       maps.Clone(allocatedByQueue),
@@ -190,6 +201,10 @@ func (node *Node) GetName() string {
 	return node.name
 }
 
+func (node *Node) IsUnschedulable() bool {
+	return node.unschedulable
+}
+
 func (node *Node) GetPool() string {
 	return node.pool
 }
@@ -243,6 +258,10 @@ func (node *Node) GetTotalResources() ResourceList {
 	return node.totalResources
 }
 
+func (node *Node) GetAllocatableResources() ResourceList {
+	return node.allocatableResources
+}
+
 func (node *Node) GetUnallocatableResources() map[int32]ResourceList {
 	return maps.Clone(node.unallocatableResources)
 }
@@ -258,7 +277,9 @@ func (node *Node) DeepCopyNilKeys() *Node {
 		nodeType:               node.nodeType,
 		taints:                 node.taints,
 		labels:                 node.labels,
+		unschedulable:          node.unschedulable,
 		totalResources:         node.totalResources,
+		allocatableResources:   node.allocatableResources,
 		unallocatableResources: node.unallocatableResources,
 
 		// keys set to nil
@@ -282,7 +303,9 @@ func (node *Node) SummaryString() string {
 	result += fmt.Sprintf("Executor: %s\n", node.executor)
 	result += fmt.Sprintf("Name: %s\n", node.name)
 	result += fmt.Sprintf("Pool: %s\n", node.pool)
+	result += fmt.Sprintf("Unschedulable: %t\n", node.unschedulable)
 	result += fmt.Sprintf("TotalResources: %s\n", node.totalResources.String())
+	result += fmt.Sprintf("AllocatableResources: %s\n", node.allocatableResources.String())
 	result += fmt.Sprintf("Labels: %v\n", node.labels)
 	result += fmt.Sprintf("Taints: %v\n", node.taints)
 	for p, u := range node.unallocatableResources {
@@ -299,11 +322,3 @@ func deepCopyLabels(labels map[string]string) map[string]string {
 	}
 	return result
 }
-
-func minInt32(arr []int32) int32 {
-	result := int32(math.MaxInt32)
-	for _, val := range arr {
-		result = min(result, val)
-	}
-	return result
-}
diff --git a/internal/scheduler/internaltypes/node_factory.go b/internal/scheduler/internaltypes/node_factory.go
@@ -62,6 +62,7 @@ func (f *NodeFactory) CreateNodeAndType(
 	taints []v1.Taint,
 	labels map[string]string,
 	totalResources ResourceList,
+	allocatableResources ResourceList,
 	unallocatableResources map[int32]ResourceList,
 	allocatableByPriority map[int32]ResourceList,
 ) *Node {
@@ -77,6 +78,7 @@ func (f *NodeFactory) CreateNodeAndType(
 		f.indexedTaints,
 		f.indexedNodeLabels,
 		totalResources,
+		allocatableResources,
 		unallocatableResources,
 		allocatableByPriority,
 	)
@@ -125,6 +127,7 @@ func (f *NodeFactory) AddLabels(nodes []*Node, extraLabels map[string]string) []
 			f.indexedTaints,
 			f.indexedNodeLabels,
 			node.GetTotalResources(),
+			node.GetAllocatableResources(),
 			node.GetUnallocatableResources(),
 			node.AllocatableByPriority,
 		)
@@ -146,6 +149,7 @@ func (f *NodeFactory) AddTaints(nodes []*Node, extraTaints []v1.Taint) []*Node {
 			f.indexedTaints,
 			f.indexedNodeLabels,
 			node.GetTotalResources(),
+			node.GetAllocatableResources(),
 			node.GetUnallocatableResources(),
 			node.AllocatableByPriority,
 		)

diff --git a/internal/scheduler/internaltypes/node_test.go b/internal/scheduler/internaltypes/node_test.go
@@ -40,6 +40,12 @@ func TestNode(t *testing.T) {
 			"memory": resource.MustParse("32Gi"),
 		},
 	)
+	allocatableResources := resourceListFactory.FromNodeProto(
+		map[string]resource.Quantity{
+			"cpu":    resource.MustParse("8"),
+			"memory": resource.MustParse("16Gi"),
+		},
+	)
 	unallocatableResources := map[int32]ResourceList{
 		1: resourceListFactory.FromJobResourceListIgnoreUnknown(
 			map[string]resource.Quantity{
@@ -110,7 +116,9 @@ func TestNode(t *testing.T) {
 		pool,
 		taints,
 		labels,
+		false,
 		totalResources,
+		allocatableResources,
 		unallocatableResources,
 		allocatableByPriority,
 		allocatedByQueue,

diff --git a/internal/scheduler/nodedb/nodedb.go b/internal/scheduler/nodedb/nodedb.go
@@ -34,7 +34,7 @@ func (nodeDb *NodeDb) addNodeToStats(node *internaltypes.Node) {
 	}
 	nodeType := node.GetNodeType()
 	nodeDb.numNodesByNodeType[nodeType.GetId()]++
-	nodeDb.totalResources = nodeDb.totalResources.Add(node.GetTotalResources())
+	nodeDb.totalAllocatableResources = nodeDb.totalAllocatableResources.Add(node.GetAllocatableResources())
 	nodeDb.nodeTypes[node.GetNodeTypeId()] = nodeType
 }
 
@@ -123,8 +123,8 @@ type NodeDb struct {
 	numNodes uint64
 	// Number of nodes in the db by node type.
 	numNodesByNodeType map[uint64]int
-	// Total amount of resources, e.g., "cpu", "memory", "gpu", across all nodes in the db.
-	totalResources internaltypes.ResourceList
+	// Total amount of allocatable resources, e.g., "cpu", "memory", "gpu", across all nodes in the db.
+	totalAllocatableResources internaltypes.ResourceList
 	// Set of node types. Populated automatically as nodes are inserted.
 	// Node types are not cleaned up if all nodes of that type are removed from the NodeDb.
 	nodeTypes map[uint64]*internaltypes.NodeType
@@ -204,7 +204,7 @@ func NewNodeDb(
 		nodeTypes:                 make(map[uint64]*internaltypes.NodeType),
 		wellKnownNodeTypes:        make(map[string]*configuration.WellKnownNodeType),
 		numNodesByNodeType:        make(map[uint64]int),
-		totalResources:            resourceListFactory.MakeAllZero(),
+		totalAllocatableResources: resourceListFactory.MakeAllZero(),
 		db:                        db,
 		// Set the initial capacity (somewhat arbitrarily) to 128 reasons.
 		podRequirementsNotMetReasonStringCache: make(map[uint64]string, 128),
@@ -297,7 +297,7 @@ func (nodeDb *NodeDb) NumNodes() int {
 }
 
 func (nodeDb *NodeDb) TotalKubernetesResources() internaltypes.ResourceList {
-	return nodeDb.totalResources
+	return nodeDb.totalAllocatableResources
 }
 
 func (nodeDb *NodeDb) Txn(write bool) *memdb.Txn {
@@ -1023,7 +1023,7 @@ func (nodeDb *NodeDb) ClearAllocated() error {
 		node = node.DeepCopyNilKeys()
 		node.AllocatableByPriority = newAllocatableByPriorityAndResourceType(
 			nodeDb.nodeDbPriorities,
-			node.GetTotalResources(),
+			node.GetAllocatableResources(),
 		)
 		newNodes = append(newNodes, node)
 	}

diff --git a/internal/scheduler/nodedb/nodeidindex_test.go b/internal/scheduler/nodedb/nodeidindex_test.go
@@ -55,6 +55,8 @@ func makeTestNode(id string) *internaltypes.Node {
 		"pool",
 		[]v1.Taint{},
 		map[string]string{},
+		false,
+		internaltypes.ResourceList{},
 		internaltypes.ResourceList{},
 		nil,
 		map[int32]internaltypes.ResourceList{},

diff --git a/internal/scheduler/nodedb/nodematching.go b/internal/scheduler/nodedb/nodematching.go
@@ -179,6 +179,8 @@ func StaticJobRequirementsMet(node *internaltypes.Node, jctx *schedulercontext.J
 		return matches, reason, err
 	}
 
+	// We are using total resource here, as it is an optimistic initial check
+	// Ideally we'd use allocatable, however allocatable could be quite dynamic and have an adverse impact on the SubmitCheck
 	matches, reason = resourceRequirementsMet(node.GetTotalResources(), jctx.KubernetesResourceRequirements)
 	if !matches {
 		return matches, reason, nil

diff --git a/internal/scheduler/nodedb/nodematching_test.go b/internal/scheduler/nodedb/nodematching_test.go
@@ -661,6 +661,8 @@ func makeTestNodeTaintsLabels(taints []v1.Taint, labels map[string]string) *inte
 		"pool",
 		taints,
 		labels,
+		false,
+		internaltypes.ResourceList{},
 		internaltypes.ResourceList{},
 		nil,
 		map[int32]internaltypes.ResourceList{},
@@ -685,6 +687,8 @@ func makeTestNodeResources(t *testing.T, allocatableByPriority map[int32]interna
 		"pool",
 		[]v1.Taint{},
 		map[string]string{},
+		false,
+		totalResources,
 		totalResources,
 		nil,
 		allocatableByPriority,

diff --git a/internal/scheduler/scheduling/preempting_queue_scheduler_test.go b/internal/scheduler/scheduling/preempting_queue_scheduler_test.go
@@ -2498,8 +2498,10 @@ func testNodeWithTaints(node *internaltypes.Node, taints []v1.Taint) *internalty
 		node.GetPool(),
 		taints,
 		node.GetLabels(),
+		false,
 		node.GetTotalResources(),
-		nil,
+		node.GetAllocatableResources(),
+		node.GetUnallocatableResources(),
 		node.AllocatableByPriority,
 		node.AllocatedByQueue,
 		node.AllocatedByJobId,

diff --git a/internal/scheduler/scheduling/scheduling_algo.go b/internal/scheduler/scheduling/scheduling_algo.go
@@ -654,6 +654,14 @@ func (l *FairSchedulingAlgo) populateNodeDb(nodeDb *nodedb.NodeDb, currentPoolJo
 	}
 
 	for _, node := range nodes {
+		if node.IsUnschedulable() && len(jobsByNodeId[node.GetId()]) == 0 {
+			// Don't add nodes that cannot be scheduled on into the nodedb
+			// - For efficiency
+			// - So the resource of the node is not counted for fairshare
+			// NOTE - Unschedulable nodes with jobs already scheduled on to them still need to be added to the nodeDb,
+			//         so the jobs can be rescheduled onto them if evicted
+			continue
+		}
 		if err := nodeDb.CreateAndInsertWithJobDbJobsWithTxn(txn, jobsByNodeId[node.GetId()], node); err != nil {
 			return err
 		}

diff --git a/internal/scheduler/testfixtures/testfixtures.go b/internal/scheduler/testfixtures/testfixtures.go
@@ -307,7 +307,9 @@ func WithNodeTypeNodes(nodeType *internaltypes.NodeType, nodes []*internaltypes.
 			node.GetPool(),
 			node.GetTaints(),
 			node.GetLabels(),
+			false,
 			node.GetTotalResources(),
+			node.GetAllocatableResources(),
 			node.GetUnallocatableResources(),
 			node.AllocatableByPriority,
 			node.AllocatedByQueue,
@@ -329,7 +331,9 @@ func WithIdNodes(nodeId string, nodes []*internaltypes.Node) []*internaltypes.No
 			node.GetPool(),
 			node.GetTaints(),
 			node.GetLabels(),
+			false,
 			node.GetTotalResources(),
+			node.GetAllocatableResources(),
 			node.GetUnallocatableResources(),
 			node.AllocatableByPriority,
 			node.AllocatedByQueue,
@@ -350,7 +354,9 @@ func WithIndexNode(idx uint64, node *internaltypes.Node) *internaltypes.Node {
 		node.GetPool(),
 		node.GetTaints(),
 		node.GetLabels(),
+		false,
 		node.GetTotalResources(),
+		node.GetAllocatableResources(),
 		node.GetUnallocatableResources(),
 		node.AllocatableByPriority,
 		node.AllocatedByQueue,
@@ -743,6 +749,8 @@ func TestSimpleNode(id string) *internaltypes.Node {
 		"",
 		nil,
 		nil,
+		false,
+		internaltypes.ResourceList{},
 		internaltypes.ResourceList{},
 		nil,
 		nil,
@@ -766,6 +774,7 @@ func TestNode(priorities []int32, resources map[string]resource.Quantity) *inter
 			schedulerconfiguration.NodeIdLabel: id,
 		},
 		rl,
+		rl,
 		map[int32]internaltypes.ResourceList{},
 		internaltypes.NewAllocatableByPriorityAndResourceType(priorities, rl))
 }