diff --git a/.changes/unreleased/added-20261012-120909.yaml b/.changes/unreleased/added-20261012-120909.yaml new file mode 100644 index 0000000..04344e3 --- /dev/null +++ b/.changes/unreleased/added-20261012-120909.yaml @@ -0,0 +1,6 @@ +kind: Added +body: |- + PodGroup safeToDisrupt field and kai.scheduler/safe-to-disrupt annotation key to opt out of consolidation +custom: + Author: "sam-huang1223" + Issue: "22" diff --git a/config/crd/scheduling.run.ai_podgroups.yaml b/config/crd/scheduling.run.ai_podgroups.yaml index 638b8ba..9aaba25 100644 --- a/config/crd/scheduling.run.ai_podgroups.yaml +++ b/config/crd/scheduling.run.ai_podgroups.yaml @@ -98,6 +98,14 @@ spec: Queue defines the queue to allocate resource for PodGroup; if queue does not exist, the PodGroup will not be scheduled. type: string + safeToDisrupt: + description: |- + SafeToDisrupt is whether the scheduler may evict this PodGroup's pods for reasons other than + fairness or priority. Today that is consolidation, which moves pods to other nodes to make room + for a pending workload. A workload that does not survive a restart (for example replicas with + restartPolicy Never, or bare pods) can set it to false. It does not affect preempt or reclaim. + Defaults to true. + type: boolean schedulingBackoff: description: The number of scheduling cycles to try before marking the pod group as UnschedulableOnNodePool. Currently only supporting diff --git a/podgrouper/constants/constants.go b/podgrouper/constants/constants.go index 08199fb..8c8f147 100644 --- a/podgrouper/constants/constants.go +++ b/podgrouper/constants/constants.go @@ -12,6 +12,7 @@ const ( PreemptionDelayAnnotationKey = "kai.scheduler/preemption-delay" StalenessGracePeriodAnnotationKey = "kai.scheduler/staleness-grace-period" + SafeToDisruptAnnotationKey = "kai.scheduler/safe-to-disrupt" BuildPriorityClass = "build" TrainPriorityClass = "train" diff --git a/scheduling/v2alpha2/podgroup_types.go b/scheduling/v2alpha2/podgroup_types.go index ea81f85..3886308 100644 --- a/scheduling/v2alpha2/podgroup_types.go +++ b/scheduling/v2alpha2/podgroup_types.go @@ -21,6 +21,7 @@ package v2alpha2 import ( "fmt" + "strconv" "strings" "time" @@ -90,6 +91,14 @@ type PodGroupSpec struct { // eviction for this PodGroup. Defaults to the scheduler's global staleness grace period. // +optional StalenessGracePeriod *metav1.Duration `json:"stalenessGracePeriod,omitempty" protobuf:"bytes,10,opt,name=stalenessGracePeriod"` + + // SafeToDisrupt is whether the scheduler may evict this PodGroup's pods for reasons other than + // fairness or priority. Today that is consolidation, which moves pods to other nodes to make room + // for a pending workload. A workload that does not survive a restart (for example replicas with + // restartPolicy Never, or bare pods) can set it to false. It does not affect preempt or reclaim. + // Defaults to true. + // +optional + SafeToDisrupt *bool `json:"safeToDisrupt,omitempty" protobuf:"varint,11,opt,name=safeToDisrupt"` } // Preemptibility defines whether this PodGroup can be preempted @@ -151,6 +160,16 @@ func ParseStalenessGracePeriod(value string) (*metav1.Duration, error) { return &metav1.Duration{Duration: stale}, nil } +// ParseSafeToDisrupt parses a safe-to-disrupt value ("true" or "false"). +// Returns an error for anything strconv.ParseBool does not accept. +func ParseSafeToDisrupt(value string) (*bool, error) { + safe, err := strconv.ParseBool(value) + if err != nil { + return nil, err + } + return &safe, nil +} + type SubGroup struct { // Name uniquely identifies the SubGroup within the PodGroup. // Must be a valid DNS label (RFC 1123). diff --git a/scheduling/v2alpha2/podgroup_types_test.go b/scheduling/v2alpha2/podgroup_types_test.go new file mode 100644 index 0000000..e508885 --- /dev/null +++ b/scheduling/v2alpha2/podgroup_types_test.go @@ -0,0 +1,40 @@ +// Copyright 2025 NVIDIA CORPORATION +// SPDX-License-Identifier: Apache-2.0 + +package v2alpha2 + +import ( + "testing" +) + +func TestParseSafeToDisrupt(t *testing.T) { + tests := []struct { + value string + want bool + wantErr bool + }{ + {value: "true", want: true}, + {value: "false", want: false}, + {value: "False", want: false}, + {value: "0", want: false}, + {value: "", wantErr: true}, + {value: "no", wantErr: true}, + } + for _, tt := range tests { + t.Run(tt.value, func(t *testing.T) { + got, err := ParseSafeToDisrupt(tt.value) + if tt.wantErr { + if err == nil { + t.Fatalf("ParseSafeToDisrupt(%q) = %v, want an error", tt.value, *got) + } + return + } + if err != nil { + t.Fatalf("ParseSafeToDisrupt(%q) returned error %v", tt.value, err) + } + if got == nil || *got != tt.want { + t.Fatalf("ParseSafeToDisrupt(%q) = %v, want %v", tt.value, got, tt.want) + } + }) + } +} diff --git a/scheduling/v2alpha2/zz_generated.deepcopy.go b/scheduling/v2alpha2/zz_generated.deepcopy.go index b3aa181..2f94edd 100644 --- a/scheduling/v2alpha2/zz_generated.deepcopy.go +++ b/scheduling/v2alpha2/zz_generated.deepcopy.go @@ -187,6 +187,11 @@ func (in *PodGroupSpec) DeepCopyInto(out *PodGroupSpec) { *out = new(v1.Duration) **out = **in } + if in.SafeToDisrupt != nil { + in, out := &in.SafeToDisrupt, &out.SafeToDisrupt + *out = new(bool) + **out = **in + } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new PodGroupSpec.