diff --git a/api/v1beta2/appwrapper_types.go b/api/v1beta2/appwrapper_types.go index 1a63ee8..f6ca43b 100644 --- a/api/v1beta2/appwrapper_types.go +++ b/api/v1beta2/appwrapper_types.go @@ -80,6 +80,9 @@ type AppWrapperPodSetInfo struct { // Tolerations to be added to the PodSpecTemplate //+optional Tolerations []corev1.Toleration `json:"tolerations,omitempty"` + // SchedulingGates to be added to the PodSpecTemplate + //+optional + SchedulingGates []corev1.PodSchedulingGate `json:"schedulingGates,omitempty"` } // AppWrapperStatus defines the observed state of the appwrapper diff --git a/api/v1beta2/zz_generated.deepcopy.go b/api/v1beta2/zz_generated.deepcopy.go index 92a1d59..593b521 100644 --- a/api/v1beta2/zz_generated.deepcopy.go +++ b/api/v1beta2/zz_generated.deepcopy.go @@ -202,6 +202,11 @@ func (in *AppWrapperPodSetInfo) DeepCopyInto(out *AppWrapperPodSetInfo) { (*in)[i].DeepCopyInto(&(*out)[i]) } } + if in.SchedulingGates != nil { + in, out := &in.SchedulingGates, &out.SchedulingGates + *out = make([]v1.PodSchedulingGate, len(*in)) + copy(*out, *in) + } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new AppWrapperPodSetInfo. diff --git a/config/crd/bases/workload.codeflare.dev_appwrappers.yaml b/config/crd/bases/workload.codeflare.dev_appwrappers.yaml index 4e7a065..58071fe 100644 --- a/config/crd/bases/workload.codeflare.dev_appwrappers.yaml +++ b/config/crd/bases/workload.codeflare.dev_appwrappers.yaml @@ -92,6 +92,21 @@ spec: type: string description: NodeSelectors to be added to the PodSpecTemplate type: object + schedulingGates: + description: SchedulingGates to be added to the PodSpecTemplate + items: + description: PodSchedulingGate is associated to a Pod + to guard its scheduling. + properties: + name: + description: |- + Name of the scheduling gate. + Each scheduling gate must have a unique name field. + type: string + required: + - name + type: object + type: array tolerations: description: Tolerations to be added to the PodSpecTemplate items: diff --git a/internal/controller/appwrapper/appwrapper_controller_test.go b/internal/controller/appwrapper/appwrapper_controller_test.go index b1df72f..77bafad 100644 --- a/internal/controller/appwrapper/appwrapper_controller_test.go +++ b/internal/controller/appwrapper/appwrapper_controller_test.go @@ -43,10 +43,11 @@ var _ = Describe("AppWrapper Controller", func() { var awReconciler *AppWrapperReconciler var awName types.NamespacedName markerPodSet := podset.PodSetInfo{ - Labels: map[string]string{"testkey1": "value1"}, - Annotations: map[string]string{"test2": "test2"}, - NodeSelector: map[string]string{"nodeName": "myNode"}, - Tolerations: []v1.Toleration{{Key: "aKey", Operator: "Exists", Effect: "NoSchedule"}}, + Labels: map[string]string{"testkey1": "value1"}, + Annotations: map[string]string{"test2": "test2"}, + NodeSelector: map[string]string{"nodeName": "myNode"}, + Tolerations: []v1.Toleration{{Key: "aKey", Operator: "Exists", Effect: "NoSchedule"}}, + SchedulingGates: []v1.PodSchedulingGate{{Name: "aGate"}}, } var kueuePodSets []kueue.PodSet @@ -172,6 +173,9 @@ var _ = Describe("AppWrapper Controller", func() { for k, v := range markerPodSet.NodeSelector { Expect(p.Spec.NodeSelector).Should(HaveKeyWithValue(k, v)) } + for _, v := range markerPodSet.SchedulingGates { + Expect(p.Spec.SchedulingGates).Should(ContainElement(v)) + } } validateAutopilot := func(p *v1.Pod) { @@ -393,6 +397,7 @@ var _ = Describe("AppWrapper Controller", func() { Expect(p.Annotations).Should(HaveKeyWithValue("myComplexAnnotation", "myComplexValue")) Expect(p.Spec.NodeSelector).Should(HaveKeyWithValue("myComplexSelector", "myComplexValue")) Expect(p.Spec.Tolerations).Should(ContainElement(v1.Toleration{Key: "myComplexKey", Value: "myComplexValue", Operator: v1.TolerationOpEqual, Effect: v1.TaintEffectNoSchedule})) + Expect(p.Spec.SchedulingGates).Should(ContainElement(v1.PodSchedulingGate{Name: "myComplexGate"})) mes := p.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms[0].MatchExpressions found := false for _, me := range mes { diff --git a/internal/controller/appwrapper/fixtures_test.go b/internal/controller/appwrapper/fixtures_test.go index 2c33dd0..aa90c32 100644 --- a/internal/controller/appwrapper/fixtures_test.go +++ b/internal/controller/appwrapper/fixtures_test.go @@ -163,6 +163,8 @@ spec: operator: NotIn values: - badHost1 + schedulingGates: + - name: myComplexGate tolerations: - key: myComplexKey value: myComplexValue diff --git a/internal/controller/appwrapper/resource_management.go b/internal/controller/appwrapper/resource_management.go index db3b1b3..0e30fb9 100644 --- a/internal/controller/appwrapper/resource_management.go +++ b/internal/controller/appwrapper/resource_management.go @@ -258,6 +258,31 @@ func (r *AppWrapperReconciler) createComponent(ctx context.Context, aw *workload spec["tolerations"] = tolerations } + // SchedulingGates + if len(toInject.SchedulingGates) > 0 { + if _, ok := spec["schedulingGates"]; !ok { + spec["schedulingGates"] = []interface{}{} + } + schedulingGates := spec["schedulingGates"].([]interface{}) + for _, addition := range toInject.SchedulingGates { + duplicate := false + for _, existing := range schedulingGates { + if imap, ok := existing.(map[string]interface{}); ok { + if iName, ok := imap["name"]; ok { + if sName, ok := iName.(string); ok && sName == addition.Name { + duplicate = true + break + } + } + } + } + if !duplicate { + schedulingGates = append(schedulingGates, map[string]interface{}{"name": addition.Name}) + } + } + spec["schedulingGates"] = schedulingGates + } + // Scheduler Name if r.Config.SchedulerName != "" { if existing, _ := spec["schedulerName"].(string); existing == "" { diff --git a/pkg/utils/utils.go b/pkg/utils/utils.go index b138826..62ba6cc 100644 --- a/pkg/utils/utils.go +++ b/pkg/utils/utils.go @@ -374,10 +374,11 @@ func SetPodSetInfos(aw *workloadv1beta2.AppWrapper, podSetsInfo []podset.PodSetI continue // we will return an error below...continuing to get an accurate count for the error message } aw.Spec.Components[idx].PodSetInfos[podSetIdx] = workloadv1beta2.AppWrapperPodSetInfo{ - Annotations: podSetsInfo[podSetsInfoIndex-1].Annotations, - Labels: podSetsInfo[podSetsInfoIndex-1].Labels, - NodeSelector: podSetsInfo[podSetsInfoIndex-1].NodeSelector, - Tolerations: podSetsInfo[podSetsInfoIndex-1].Tolerations, + Annotations: podSetsInfo[podSetsInfoIndex-1].Annotations, + Labels: podSetsInfo[podSetsInfoIndex-1].Labels, + NodeSelector: podSetsInfo[podSetsInfoIndex-1].NodeSelector, + Tolerations: podSetsInfo[podSetsInfoIndex-1].Tolerations, + SchedulingGates: podSetsInfo[podSetsInfoIndex-1].SchedulingGates, } } }