-
Notifications
You must be signed in to change notification settings - Fork 132
/
node_group.go
122 lines (100 loc) · 3.36 KB
/
node_group.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
package node
import (
"sort"
"strings"
v1 "k8s.io/api/core/v1"
"github.com/armadaproject/armada/internal/common/util"
"github.com/armadaproject/armada/internal/executor/context"
util2 "github.com/armadaproject/armada/internal/executor/util"
"github.com/armadaproject/armada/pkg/api"
)
type NodeInfoService interface {
IsAvailableProcessingNode(*v1.Node) bool
GetAllAvailableProcessingNodes() ([]*v1.Node, error)
GetAllNodes() ([]*v1.Node, error)
GroupNodesByType(nodes []*v1.Node) []*NodeGroup
GetType(node *v1.Node) *api.NodeTypeIdentifier
}
type KubernetesNodeInfoService struct {
clusterContext context.ClusterContext
toleratedTaints map[string]bool
}
func NewKubernetesNodeInfoService(clusterContext context.ClusterContext, toleratedTaints []string) *KubernetesNodeInfoService {
return &KubernetesNodeInfoService{
clusterContext: clusterContext,
toleratedTaints: util.StringListToSet(toleratedTaints),
}
}
type NodeGroup struct {
NodeType *api.NodeTypeIdentifier
Nodes []*v1.Node
}
func (kubernetesNodeInfoService *KubernetesNodeInfoService) GroupNodesByType(nodes []*v1.Node) []*NodeGroup {
nodeGroupMap := map[string]*NodeGroup{}
for _, node := range nodes {
nodeType := kubernetesNodeInfoService.GetType(node)
if _, present := nodeGroupMap[nodeType.Id]; !present {
nodeGroupMap[nodeType.Id] = &NodeGroup{
NodeType: nodeType,
Nodes: []*v1.Node{},
}
}
nodeGroupMap[nodeType.Id].Nodes = append(nodeGroupMap[nodeType.Id].Nodes, node)
}
nodeGroups := make([]*NodeGroup, 0, len(nodeGroupMap))
for _, group := range nodeGroupMap {
nodeGroups = append(nodeGroups, group)
}
return nodeGroups
}
func (kubernetesNodeInfoService *KubernetesNodeInfoService) GetType(node *v1.Node) *api.NodeTypeIdentifier {
groupId := kubernetesNodeInfoService.clusterContext.GetClusterPool()
relevantTaints := kubernetesNodeInfoService.filterToleratedTaints(node.Spec.Taints)
if len(relevantTaints) > 0 {
groupId = nodeGroupId(relevantTaints)
}
return &api.NodeTypeIdentifier{
Id: groupId,
Taints: relevantTaints,
}
}
func (kubernetesNodeInfoService *KubernetesNodeInfoService) filterToleratedTaints(taints []v1.Taint) []v1.Taint {
result := []v1.Taint{}
for _, taint := range taints {
_, ok := kubernetesNodeInfoService.toleratedTaints[taint.Key]
if ok {
result = append(result, taint)
}
}
return result
}
func nodeGroupId(taints []v1.Taint) string {
idStrings := []string{}
for _, taint := range taints {
idStrings = append(idStrings, taint.Key)
}
sort.Strings(idStrings)
return strings.Join(idStrings, ",")
}
func (kubernetesNodeInfoService *KubernetesNodeInfoService) GetAllAvailableProcessingNodes() ([]*v1.Node, error) {
allNodes, err := kubernetesNodeInfoService.clusterContext.GetNodes()
if err != nil {
return []*v1.Node{}, err
}
return util2.FilterNodes(allNodes, kubernetesNodeInfoService.IsAvailableProcessingNode), nil
}
func (kubernetesNodeInfoService *KubernetesNodeInfoService) GetAllNodes() ([]*v1.Node, error) {
return kubernetesNodeInfoService.clusterContext.GetNodes()
}
func (kubernetesNodeInfoService *KubernetesNodeInfoService) IsAvailableProcessingNode(node *v1.Node) bool {
if node.Spec.Unschedulable {
return false
}
for _, taint := range node.Spec.Taints {
if taint.Effect == v1.TaintEffectNoSchedule &&
!kubernetesNodeInfoService.toleratedTaints[taint.Key] {
return false
}
}
return true
}