Propagate SchedulingGates in RunWithPodSetsInfo (#301)

Author: dgrove-oss

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

api/v1beta2/zz_generated.deepcopy.go

@@ -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:

config/crd/bases/workload.codeflare.dev_appwrappers.yaml

@@ -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 {

internal/controller/appwrapper/appwrapper_controller_test.go

@@ -163,6 +163,8 @@ spec:
             operator: NotIn
             values:
             - badHost1
+  schedulingGates:
+  - name: myComplexGate
   tolerations:
   - key: myComplexKey
     value: myComplexValue

internal/controller/appwrapper/fixtures_test.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 == "" {

internal/controller/appwrapper/resource_management.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,
 			}
 		}
 	}