Skip to content

Commit 0450532

Browse files
romanbaronromanbaron
authored
Made node label key configurable (#217)
1 parent 85f046c commit 0450532

File tree

13 files changed

+48
-17
lines changed

13 files changed

+48
-17
lines changed

‎cmd/scheduler/app/options/options.go‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,9 @@ const (
2727
DefaultPyroscopeBlockProfilerRate = 5
2828
defaultNumOfStatusRecordingWorkers = 5
2929
defaultNodePoolLabelKey = ""
30+
defaultCPUWorkerNodeLabelKey = "node-role.kubernetes.io/cpu-worker"
31+
defaultGPUWorkerNodeLabelKey = "node-role.kubernetes.io/gpu-worker"
32+
defaultMIGWorkerNodeLabelKey = "node-role.kubernetes.io/mig-enabled"
3033
)
3134

3235
// ServerOption is the main context object for the controller manager.
@@ -59,6 +62,9 @@ type ServerOption struct {
5962
NumOfStatusRecordingWorkers int
6063
GlobalDefaultStalenessGracePeriod time.Duration
6164
PluginServerPort int
65+
CPUWorkerNodeLabelKey string
66+
GPUWorkerNodeLabelKey string
67+
MIGWorkerNodeLabelKey string
6268

6369
QPS int
6470
Burst int
@@ -105,6 +111,9 @@ func (s *ServerOption) AddFlags(fs *pflag.FlagSet) {
105111
fs.IntVar(&s.NumOfStatusRecordingWorkers, "num-of-status-recording-workers", defaultNumOfStatusRecordingWorkers, "specifies the max number of go routines spawned to update pod and podgroups conditions and events. Defaults to 5")
106112
fs.DurationVar(&s.GlobalDefaultStalenessGracePeriod, "default-staleness-grace-period", defaultStalenessGracePeriod, "Global default staleness grace period duration. Negative values means infinite. Defaults to 60s")
107113
fs.IntVar(&s.PluginServerPort, "plugin-server-port", 8081, "The port to bind for plugin server requests")
114+
fs.StringVar(&s.CPUWorkerNodeLabelKey, "cpu-worker-node-label-key", defaultCPUWorkerNodeLabelKey, "The label key for CPU worker nodes")
115+
fs.StringVar(&s.GPUWorkerNodeLabelKey, "gpu-worker-node-label-key", defaultGPUWorkerNodeLabelKey, "The label key for GPU worker nodes")
116+
fs.StringVar(&s.MIGWorkerNodeLabelKey, "mig-worker-node-label-key", defaultMIGWorkerNodeLabelKey, "The label key for MIG enabled worker nodes")
108117

109118
utilfeature.DefaultMutableFeatureGate.AddFlag(fs)
110119
}

‎cmd/scheduler/app/options/options_test.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,9 @@ func TestAddFlags(t *testing.T) {
4949
NumOfStatusRecordingWorkers: defaultNumOfStatusRecordingWorkers,
5050
NodePoolLabelKey: defaultNodePoolLabelKey,
5151
PluginServerPort: 8081,
52+
CPUWorkerNodeLabelKey: defaultCPUWorkerNodeLabelKey,
53+
GPUWorkerNodeLabelKey: defaultGPUWorkerNodeLabelKey,
54+
MIGWorkerNodeLabelKey: defaultMIGWorkerNodeLabelKey,
5255
}
5356

5457
if !reflect.DeepEqual(expected, s) {

‎cmd/scheduler/app/server.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,9 @@ func setupLogging(so *options.ServerOption) error {
132132
func setConfig(so *options.ServerOption) {
133133
config := conf.GetConfig()
134134
config.ResourceReservationAppLabelValue = so.ResourceReservationAppLabel
135+
config.CPUWorkerNodeLabelKey = so.CPUWorkerNodeLabelKey
136+
config.GPUWorkerNodeLabelKey = so.GPUWorkerNodeLabelKey
137+
config.MIGWorkerNodeLabelKey = so.MIGWorkerNodeLabelKey
135138
}
136139

137140
func Run(opt *options.ServerOption, config *restclient.Config, mux *http.ServeMux) error {

‎pkg/common/constants/constants.go‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@ const (
2929
// Labels
3030
GPUGroup = "runai-gpu-group"
3131
MultiGpuGroupLabelPrefix = GPUGroup + "/"
32-
MigEnabledLabel = "node-role.kubernetes.io/runai-mig-enabled"
3332
MigStrategyLabel = "nvidia.com/mig.strategy"
3433
GpuCountLabel = "nvidia.com/gpu.count"
3534
)

‎pkg/scheduler/api/node_info/node_info.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,15 +24,14 @@ import (
2424
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/resource_info"
2525
sc_info "github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/storagecapacity_info"
2626
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/storageclaim_info"
27+
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/conf"
2728
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/log"
2829
)
2930

3031
const (
3132
DefaultGpuMemory = 100 // The default value is 100 because it allows all the calculation of (memory = fractional * GpuMemory) to work, if it was 0 the result will always be zero too
3233
GpuMemoryLabel = "nvidia.com/gpu.memory"
3334
GpuCountLabel = "nvidia.com/gpu.count"
34-
CpuWorkerNode = "node-role.kubernetes.io/runai-cpu-worker"
35-
GpuWorkerNode = "node-role.kubernetes.io/runai-gpu-worker"
3635
MbToBRatio = 1000000
3736
BitToMib = 1024 * 1024
3837
TibInMib = 1024 * 1024
@@ -633,7 +632,8 @@ func (ni *NodeInfo) IsCPUOnlyNode() bool {
633632
}
634633

635634
func (ni *NodeInfo) IsMIGEnabled() bool {
636-
enabled, found := ni.Node.Labels[commonconstants.MigEnabledLabel]
635+
migWorkerLabelKey := conf.GetConfig().MIGWorkerNodeLabelKey
636+
enabled, found := ni.Node.Labels[migWorkerLabelKey]
637637
if found {
638638
isMig, err := strconv.ParseBool(enabled)
639639
return err == nil && isMig

‎pkg/scheduler/api/node_info/node_info_test.go‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,10 @@ import (
2727
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/storagecapacity_info"
2828
)
2929

30-
const MibToMbScale = 1.048576
30+
const (
31+
MibToMbScale = 1.048576
32+
migEnabledLabelKey = "node-role.kubernetes.io/mig-enabled"
33+
)
3134

3235
func nodeInfoEqual(l, r *NodeInfo) bool {
3336
l.PodAffinityInfo = nil
@@ -535,10 +538,10 @@ func TestIsTaskAllocatable(t *testing.T) {
535538

536539
func TestIsTaskAllocatableOnReleasingOrIdle(t *testing.T) {
537540
singleMigNode := common_info.BuildNode("single-mig", common_info.BuildResourceListWithGPU("2000m", "2G", "8"))
538-
singleMigNode.Labels[commonconstants.MigEnabledLabel] = "true"
541+
singleMigNode.Labels[migEnabledLabelKey] = "true"
539542

540543
mixedMigNode := common_info.BuildNode("mixed-mig", common_info.BuildResourceListWithGPU("2000m", "2G", "8"))
541-
mixedMigNode.Labels[commonconstants.MigEnabledLabel] = "true"
544+
mixedMigNode.Labels[migEnabledLabelKey] = "true"
542545
mixedMigNode.Labels[commonconstants.GpuCountLabel] = "8"
543546
mixedMigNode.Labels[GpuMemoryLabel] = "40"
544547

‎pkg/scheduler/cache/cluster_info/cluster_info.go‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -429,10 +429,12 @@ func getPodGroupPriority(
429429
}
430430

431431
func filterUnmarkedNodes(nodes []*v1.Node) []*v1.Node {
432+
cpuWorkerLabelKey := conf.GetConfig().CPUWorkerNodeLabelKey
433+
gpuWorkerLabelKey := conf.GetConfig().GPUWorkerNodeLabelKey
432434
markedNodes := []*v1.Node{}
433435
for _, node := range nodes {
434-
_, foundGpuNode := node.Labels[node_info.GpuWorkerNode]
435-
_, foundCpuNode := node.Labels[node_info.CpuWorkerNode]
436+
_, foundGpuNode := node.Labels[gpuWorkerLabelKey]
437+
_, foundCpuNode := node.Labels[cpuWorkerLabelKey]
436438
if foundGpuNode || foundCpuNode {
437439
markedNodes = append(markedNodes, node)
438440
log.InfraLogger.V(6).Infof("Node: <%v> is considered by cpu or gpu label", node.Name)

‎pkg/scheduler/conf/global_config.go‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,18 @@ var (
1515
type GlobalConfig struct {
1616
ResourceReservationAppLabelValue string
1717
ScalingPodAppLabelValue string
18+
CPUWorkerNodeLabelKey string
19+
GPUWorkerNodeLabelKey string
20+
MIGWorkerNodeLabelKey string
1821
}
1922

2023
func newGlobalConfig() *GlobalConfig {
2124
return &GlobalConfig{
22-
ResourceReservationAppLabelValue: "kai-resource-reservation",
25+
ResourceReservationAppLabelValue: "runai-reservation",
2326
ScalingPodAppLabelValue: "scaling-pod",
27+
CPUWorkerNodeLabelKey: "node-role.kubernetes.io/cpu-worker",
28+
GPUWorkerNodeLabelKey: "node-role.kubernetes.io/gpu-worker",
29+
MIGWorkerNodeLabelKey: "node-role.kubernetes.io/mig-enabled",
2430
}
2531
}
2632

‎pkg/scheduler/plugins/predicates/predicates.go‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import (
1616
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/pod_info"
1717
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/podgroup_info"
1818
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/cache/cluster_info"
19+
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/conf"
1920
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/framework"
2021
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/k8s_internal"
2122
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/k8s_internal/predicates"
@@ -229,14 +230,16 @@ func (pp *predicatesPlugin) evaluateTaskOnPredicates(
229230
}
230231

231232
if isRestrictNodeSchedulingEnabled() {
233+
gpuWorkerLabelKey := conf.GetConfig().GPUWorkerNodeLabelKey
234+
cpuWorkerLabelKey := conf.GetConfig().CPUWorkerNodeLabelKey
232235
if task.IsRequireAnyKindOfGPU() {
233-
if _, found := node.Node.Labels[node_info.GpuWorkerNode]; !found {
236+
if _, found := node.Node.Labels[gpuWorkerLabelKey]; !found {
234237
log.InfraLogger.V(6).Infof("Task <%s/%s> is a GPU job and will not be allocated to a non GPU <%s>",
235238
task.Namespace, task.Name, node.Name)
236239
return fmt.Errorf("gpu task: <%v/%v> can't run on non gpu nodes, node: <%v>", task.Namespace, task.Name, node.Name)
237240
}
238241
} else {
239-
if _, found := node.Node.Labels[node_info.CpuWorkerNode]; !found {
242+
if _, found := node.Node.Labels[cpuWorkerLabelKey]; !found {
240243
log.InfraLogger.V(6).Infof("Task <%s/%s> is a CPU job and will not be allocated to a GPU node <%s>",
241244
task.Namespace, task.Name, node.Name)
242245
return fmt.Errorf("cpu task: <%v/%v> can't run on non cpu nodes, node: <%v>", task.Namespace, task.Name, node.Name)

‎pkg/scheduler/plugins/proportion/proportion.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import (
1616
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/podgroup_info"
1717
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/queue_info"
1818
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/resource_info"
19+
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/conf"
1920
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/framework"
2021
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/log"
2122
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/metrics"
@@ -193,7 +194,8 @@ func getNodeResources(ssn *framework.Session, node *node_info.NodeInfo) rs.Resou
193194
return nodeResource
194195
}
195196

196-
_, found := node.Node.Labels[node_info.GpuWorkerNode]
197+
gpuWorkerLabelKey := conf.GetConfig().GPUWorkerNodeLabelKey
198+
_, found := node.Node.Labels[gpuWorkerLabelKey]
197199
shouldIgnoreGPUs := ssn.IsRestrictNodeSchedulingEnabled() && !found
198200
if shouldIgnoreGPUs {
199201
nodeResource.Add(rs.NewResourceQuantities(node.Allocatable.Cpu(), node.Allocatable.Memory(), 0))

0 commit comments

Comments
 (0)