refactor PodSetInfo flow to enable immutability (#45)

Store the PodSetInfos provided by Kueue's RunWithPodSetInfo method in
a separate array in the AppWrapperComponent instead of immediately
injecting them into the raw templates.  This simplifies both the
Workload controller (which we hope to upstream to Kueue) and the
enforcement of RBAC by allowing us to require that Component.Template
is immutable.

Author: dgrove-oss

api/v1beta2/appwrapper_types.go

@@ -17,6 +17,7 @@ limitations under the License.
 package v1beta2
 
 import (
+	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	runtime "k8s.io/apimachinery/pkg/runtime"
 )
@@ -35,6 +36,9 @@ type AppWrapperComponent struct {
 	// PodSets contained in the component
 	PodSets []AppWrapperPodSet `json:"podSets"`
 
+	// PodSetInfos assigned to the Component by Kueue
+	PodSetInfos []AppWrapperPodSetInfo `json:"podSetInfos,omitempty"`
+
 	// +kubebuilder:pruning:PreserveUnknownFields
 	// +kubebuilder:validation:EmbeddedResource
 	// Template for the component
@@ -50,6 +54,13 @@ type AppWrapperPodSet struct {
 	Path string `json:"path"`
 }
 
+type AppWrapperPodSetInfo struct {
+	Annotations  map[string]string   `json:"annotations,omitempty"`
+	Labels       map[string]string   `json:"labels,omitempty"`
+	NodeSelector map[string]string   `json:"nodeSelector,omitempty"`
+	Tolerations  []corev1.Toleration `json:"tolerations,omitempty"`
+}
+
 // AppWrapperStatus defines the observed state of the appwrapper
 type AppWrapperStatus struct {
 	// Phase of the AppWrapper object

api/v1beta2/zz_generated.deepcopy.go

@@ -73,7 +73,6 @@ func main() {
 	flag.BoolVar(&enableHTTP2, "enable-http2", false,
 		"If set, HTTP/2 will be enabled for the metrics and webhook servers")
 	flag.BoolVar(&config.ManageJobsWithoutQueueName, "manage-no-queue", true, "Manage AppWrappers without queue names")
-	flag.StringVar(&config.ServiceAccountName, "service-account", "", "Service account name for controller")
 	opts := zap.Options{
 		Development: true,
 	}

cmd/main.go

@@ -48,6 +48,62 @@ spec:
                 items:
                   description: AppWrapperComponent describes a wrapped resource
                   properties:
+                    podSetInfos:
+                      description: PodSetInfos assigned to the Component by Kueue
+                      items:
+                        properties:
+                          annotations:
+                            additionalProperties:
+                              type: string
+                            type: object
+                          labels:
+                            additionalProperties:
+                              type: string
+                            type: object
+                          nodeSelector:
+                            additionalProperties:
+                              type: string
+                            type: object
+                          tolerations:
+                            items:
+                              description: |-
+                                The pod this Toleration is attached to tolerates any taint that matches
+                                the triple <key,value,effect> using the matching operator <operator>.
+                              properties:
+                                effect:
+                                  description: |-
+                                    Effect indicates the taint effect to match. Empty means match all taint effects.
+                                    When specified, allowed values are NoSchedule, PreferNoSchedule and NoExecute.
+                                  type: string
+                                key:
+                                  description: |-
+                                    Key is the taint key that the toleration applies to. Empty means match all taint keys.
+                                    If the key is empty, operator must be Exists; this combination means to match all values and all keys.
+                                  type: string
+                                operator:
+                                  description: |-
+                                    Operator represents a key's relationship to the value.
+                                    Valid operators are Exists and Equal. Defaults to Equal.
+                                    Exists is equivalent to wildcard for value, so that a pod can
+                                    tolerate all taints of a particular category.
+                                  type: string
+                                tolerationSeconds:
+                                  description: |-
+                                    TolerationSeconds represents the period of time the toleration (which must be
+                                    of effect NoExecute, otherwise this field is ignored) tolerates the taint. By default,
+                                    it is not set, which means tolerate the taint forever (do not evict). Zero and
+                                    negative values will be treated as 0 (evict immediately) by the system.
+                                  format: int64
+                                  type: integer
+                                value:
+                                  description: |-
+                                    Value is the taint value the toleration matches to.
+                                    If the operator is Exists, the value should be empty, otherwise just a regular string.
+                                  type: string
+                              type: object
+                            type: array
+                        type: object
+                      type: array
                     podSets:
                       description: PodSets contained in the component
                       items:

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

@@ -34,7 +34,6 @@ spec:
             memory: 64Mi
       - name: manager
         args:
-        - "--service-account=system:serviceaccount:$(MY_NAMESPACE):$(MY_SERVICE_ACCOUNT_NAME)"
         - "--health-probe-bind-address=:8081"
         - "--metrics-bind-address=127.0.0.1:8080"
         - "--leader-elect"

config/default/manager_auth_proxy_patch.yaml

@@ -66,7 +66,6 @@ spec:
         - /manager
         args:
         - --leader-elect
-        - "--service-account=system:serviceaccount:$(MY_NAMESPACE):$(MY_SERVICE_ACCOUNT_NAME)"
         image: controller:latest
         name: manager
         securityContext:
@@ -93,14 +92,5 @@ spec:
           requests:
             cpu: 10m
             memory: 64Mi
-        env:
-        - name: MY_SERVICE_ACCOUNT_NAME
-          valueFrom:
-            fieldRef:
-              fieldPath: spec.serviceAccountName
-        - name: MY_NAMESPACE
-          valueFrom:
-            fieldRef:
-              fieldPath: metadata.namespace
       serviceAccountName: controller-manager
       terminationGracePeriodSeconds: 10

config/manager/manager.yaml

@@ -26,8 +26,7 @@ import (
 )
 
 type AppWrapperConfig struct {
-	ManageJobsWithoutQueueName bool   `json:"manageJobsWithoutQueueName,omitempty"`
-	ServiceAccountName         string `json:"serviceAccountName,omitempty"`
+	ManageJobsWithoutQueueName bool `json:"manageJobsWithoutQueueName,omitempty"`
 }
 
 // SetupWithManager creates and configures all components of the AppWrapper controller

internal/controller/appwrapper_config.go

@@ -34,8 +34,8 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/handler"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/controller-runtime/pkg/reconcile"
-	"sigs.k8s.io/kueue/pkg/controller/constants"
-	"sigs.k8s.io/kueue/pkg/controller/jobframework"
+	"sigs.k8s.io/kueue/pkg/podset"
+	utilmaps "sigs.k8s.io/kueue/pkg/util/maps"
 
 	workloadv1beta2 "github.com/project-codeflare/appwrapper/api/v1beta2"
 )
@@ -299,38 +299,98 @@ func parseComponent(aw *workloadv1beta2.AppWrapper, raw []byte) (*unstructured.U
 	return obj, nil
 }
 
-func parseComponents(aw *workloadv1beta2.AppWrapper) ([]client.Object, error) {
-	components := aw.Spec.Components
-	objects := make([]client.Object, len(components))
-	for i, component := range components {
-		obj, err := parseComponent(aw, component.Template.Raw)
-		if err != nil {
-			return nil, err
+func materializeObject(aw *workloadv1beta2.AppWrapper, component *workloadv1beta2.AppWrapperComponent) (client.Object, error) {
+	toMap := func(x interface{}) map[string]string {
+		if x == nil {
+			return nil
+		} else {
+			if sm, ok := x.(map[string]string); ok {
+				return sm
+			} else if im, ok := x.(map[string]interface{}); ok {
+				sm := make(map[string]string, len(im))
+				for k, v := range im {
+					str, ok := v.(string)
+					if ok {
+						sm[k] = str
+					} else {
+						sm[k] = fmt.Sprint(v)
+					}
+				}
+				return sm
+			} else {
+				return nil
+			}
 		}
-		objects[i] = obj
 	}
-	return objects, nil
-}
+	awLabels := map[string]string{AppWrapperLabel: aw.Name}
 
-func (r *AppWrapperReconciler) createComponents(ctx context.Context, aw *workloadv1beta2.AppWrapper) (error, bool) {
-	objects, err := parseComponents(aw)
+	obj, err := parseComponent(aw, component.Template.Raw)
 	if err != nil {
-		return err, true // fatal
+		return nil, err
 	}
 
-	ref := &metav1.OwnerReference{APIVersion: GVK.GroupVersion().String(), Kind: GVK.Kind, Name: aw.Name, UID: aw.UID}
-	myWorkloadName, err := jobframework.GetWorkloadNameForOwnerRef(ref)
-	if err != nil {
-		return err, true
+	for podSetsIdx, podSet := range component.PodSets {
+		toInject := component.PodSetInfos[podSetsIdx]
+
+		p, err := getRawTemplate(obj.UnstructuredContent(), podSet.Path)
+		if err != nil {
+			return nil, err // Should not happen, path validity is enforced by validateAppWrapperInvariants
+		}
+		if md, ok := p["metadata"]; !ok || md == nil {
+			p["metadata"] = make(map[string]interface{})
+		}
+		metadata := p["metadata"].(map[string]interface{})
+		spec := p["spec"].(map[string]interface{}) // Must exist, enforced by validateAppWrapperInvariants
+
+		// Annotations
+		if len(toInject.Annotations) > 0 {
+			existing := toMap(metadata["annotations"])
+			if err := utilmaps.HaveConflict(existing, toInject.Annotations); err != nil {
+				return nil, podset.BadPodSetsUpdateError("annotations", err)
+			}
+			metadata["annotations"] = utilmaps.MergeKeepFirst(existing, toInject.Annotations)
+		}
+
+		// Labels
+		mergedLabels := utilmaps.MergeKeepFirst(toInject.Labels, awLabels)
+		existing := toMap(metadata["labels"])
+		if err := utilmaps.HaveConflict(existing, mergedLabels); err != nil {
+			return nil, podset.BadPodSetsUpdateError("labels", err)
+		}
+		metadata["labels"] = utilmaps.MergeKeepFirst(existing, mergedLabels)
+
+		// NodeSelectors
+		if len(toInject.NodeSelector) > 0 {
+			existing := toMap(metadata["nodeSelector"])
+			if err := utilmaps.HaveConflict(existing, toInject.NodeSelector); err != nil {
+				return nil, podset.BadPodSetsUpdateError("nodeSelector", err)
+			}
+			metadata["nodeSelector"] = utilmaps.MergeKeepFirst(existing, toInject.NodeSelector)
+		}
+
+		// Tolerations
+		if len(toInject.Tolerations) > 0 {
+			if _, ok := spec["tolerations"]; !ok {
+				spec["tolerations"] = []interface{}{}
+			}
+			tolerations := spec["tolerations"].([]interface{})
+			for _, addition := range toInject.Tolerations {
+				tolerations = append(tolerations, addition)
+			}
+			spec["tolerations"] = tolerations
+		}
 	}
 
-	for _, obj := range objects {
-		annotations := obj.GetAnnotations()
-		if annotations == nil {
-			annotations = make(map[string]string)
+	return obj, nil
+}
+
+func (r *AppWrapperReconciler) createComponents(ctx context.Context, aw *workloadv1beta2.AppWrapper) (error, bool) {
+	for _, component := range aw.Spec.Components {
+		obj, err := materializeObject(aw, &component)
+		if err != nil {
+			return err, true
 		}
-		annotations[constants.ParentWorkloadAnnotation] = myWorkloadName
-		obj.SetAnnotations(annotations)
+
 		if err := controllerutil.SetControllerReference(aw, obj, r.Scheme); err != nil {
 			return err, true
 		}

internal/controller/appwrapper_controller.go

@@ -84,7 +84,7 @@ func (w *AppWrapperWebhook) ValidateUpdate(ctx context.Context, oldObj, newObj r
 	oldAW := oldObj.(*workloadv1beta2.AppWrapper)
 	newAW := newObj.(*workloadv1beta2.AppWrapper)
 
-	allErrors := w.validateAppWrapperUpdate(ctx, oldAW, newAW)
+	allErrors := w.validateAppWrapperUpdate(oldAW, newAW)
 	if w.Config.ManageJobsWithoutQueueName || jobframework.QueueName((*AppWrapper)(newAW)) != "" {
 		allErrors = append(allErrors, jobframework.ValidateUpdateForQueueName((*AppWrapper)(oldAW), (*AppWrapper)(newAW))...)
 		allErrors = append(allErrors, jobframework.ValidateUpdateForWorkloadPriorityClassName((*AppWrapper)(oldAW), (*AppWrapper)(newAW))...)
@@ -201,16 +201,8 @@ func (w *AppWrapperWebhook) validateAppWrapperCreate(ctx context.Context, aw *wo
 	return allErrors
 }
 
-// validateAppWrapperUpdate enforces that AppWrapper.Spec.Components is deeply immutable
-func (w *AppWrapperWebhook) validateAppWrapperUpdate(ctx context.Context, old *workloadv1beta2.AppWrapper, new *workloadv1beta2.AppWrapper) field.ErrorList {
-	// The AppWrapper controller must be allowed to mutate Spec.Components
-	// to enable it to implement RunWithPodSetsInfo and RestorePodSetsInfo
-	if request, err := admission.RequestFromContext(ctx); err == nil {
-		if w.Config.ServiceAccountName != "" && request.UserInfo.Username == w.Config.ServiceAccountName {
-			return field.ErrorList{}
-		}
-	}
-
+// validateAppWrapperUpdate enforces deep immutablity of all fields that were validated by validateAppWrapperCreate
+func (w *AppWrapperWebhook) validateAppWrapperUpdate(old *workloadv1beta2.AppWrapper, new *workloadv1beta2.AppWrapper) field.ErrorList {
 	allErrors := field.ErrorList{}
 	msg := "attempt to change immutable field"
 	componentsPath := field.NewPath("spec").Child("components")

internal/controller/appwrapper_webhook.go

@@ -161,7 +161,7 @@ var _ = BeforeSuite(func() {
 	})
 	Expect(err).NotTo(HaveOccurred())
 
-	awConfig := AppWrapperConfig{ManageJobsWithoutQueueName: true, ServiceAccountName: ctrlUserName}
+	awConfig := AppWrapperConfig{ManageJobsWithoutQueueName: true}
 	err = (&AppWrapperWebhook{Config: &awConfig}).SetupWebhookWithManager(mgr)
 	Expect(err).NotTo(HaveOccurred())