Avoid storing PodSet inference results in Spec (#147)

Store the results of PodSet inference in the Status, not the Spec. This means we end up doing the inference twice for each AppWrapper (once in the ValidatingWebHook for error checking and again later in the Controller when we have access to the Status), but it avoids mutating the user-provided Spec. Avoiding Spec mutations avoids spurious rejections when the user does multiple kubectl apply operations that do not change the Component itself.

Author: dgrove-oss

api/v1beta2/appwrapper_types.go

@@ -39,9 +39,9 @@ type AppWrapperComponent struct {
 	//+optional
 	Annotations map[string]string `json:"annotations,omitempty"`
 
-	// PodSets contained in the Component
+	// DeclaredPodSets for the Component (optional for known PodCreating GVKs)
 	//+optional
-	PodSets []AppWrapperPodSet `json:"podSets,omitempty"`
+	DeclaredPodSets []AppWrapperPodSet `json:"podSets,omitempty"`
 
 	// PodSetInfos assigned to the Component's PodSets by Kueue
 	//+optional
@@ -121,6 +121,9 @@ type AppWrapperComponentStatus struct {
 	// APIVersion is the APIVersion of the Component
 	APIVersion string `json:"apiVersion"`
 
+	// PodSets is the validated PodSets for the Component (either from AppWrapperComponent.DeclaredPodSets or inferred by the controller)
+	PodSets []AppWrapperPodSet `json:"podSets"`
+
 	// Conditions hold the latest available observations of the Component's current state.
 	//
 	// The type of the condition could be:

api/v1beta2/zz_generated.deepcopy.go

@@ -129,7 +129,8 @@ spec:
                         type: object
                       type: array
                     podSets:
-                      description: PodSets contained in the Component
+                      description: DeclaredPodSets for the Component (optional for
+                        known PodCreating GVKs)
                       items:
                         description: AppWrapperPodSet describes an homogeneous set
                           of pods
@@ -264,10 +265,31 @@ spec:
                     name:
                       description: Name is the name of the Component
                       type: string
+                    podSets:
+                      description: PodSets is the validated PodSets for the Component
+                        (either from AppWrapperComponent.DeclaredPodSets or inferred
+                        by the controller)
+                      items:
+                        description: AppWrapperPodSet describes an homogeneous set
+                          of pods
+                        properties:
+                          path:
+                            description: Path is the path Component.Template to the
+                              PodTemplateSpec for this PodSet
+                            type: string
+                          replicas:
+                            description: Replicas is the number of pods in this PodSet
+                            format: int32
+                            type: integer
+                        required:
+                        - path
+                        type: object
+                      type: array
                   required:
                   - apiVersion
                   - kind
                   - name
+                  - podSets
                   type: object
                 type: array
               conditions:

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

@@ -37,6 +37,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/reconcile"
 
 	workloadv1beta2 "github.com/project-codeflare/appwrapper/api/v1beta2"
+	"github.com/project-codeflare/appwrapper/internal/controller/awstatus"
 	"github.com/project-codeflare/appwrapper/pkg/config"
 	"github.com/project-codeflare/appwrapper/pkg/utils"
 )
@@ -150,13 +151,19 @@ func (r *AppWrapperReconciler) Reconcile(ctx context.Context, req ctrl.Request)
 				return ctrl.Result{}, err
 			}
 		}
-		aw.Status.ComponentStatus = make([]workloadv1beta2.AppWrapperComponentStatus, len(aw.Spec.Components))
+
 		return r.updateStatus(ctx, aw, workloadv1beta2.AppWrapperSuspended)
 
 	case workloadv1beta2.AppWrapperSuspended: // no components deployed
 		if aw.Spec.Suspend {
 			return ctrl.Result{}, nil // remain suspended
 		}
+
+		// Normally already done as a side-effect of Kueue calling PodSets(), but be absolutely certain before we start using it.
+		if err := awstatus.EnsureComponentStatusInitialized(ctx, aw); err != nil {
+			return ctrl.Result{}, err
+		}
+
 		// begin deployment
 		meta.SetStatusCondition(&aw.Status.Conditions, metav1.Condition{
 			Type:    string(workloadv1beta2.QuotaReserved),

internal/controller/appwrapper/appwrapper_controller.go

@@ -107,8 +107,8 @@ func pod(milliCPU int64) workloadv1beta2.AppWrapperComponent {
 	jsonBytes, err := yaml.YAMLToJSON([]byte(yamlString))
 	Expect(err).NotTo(HaveOccurred())
 	return workloadv1beta2.AppWrapperComponent{
-		PodSets:  []workloadv1beta2.AppWrapperPodSet{{Replicas: ptr.To(int32(1)), Path: "template"}},
-		Template: runtime.RawExtension{Raw: jsonBytes},
+		DeclaredPodSets: []workloadv1beta2.AppWrapperPodSet{{Replicas: ptr.To(int32(1)), Path: "template"}},
+		Template:        runtime.RawExtension{Raw: jsonBytes},
 	}
 }
 
@@ -134,7 +134,7 @@ func malformedPod(milliCPU int64) workloadv1beta2.AppWrapperComponent {
 	jsonBytes, err := yaml.YAMLToJSON([]byte(yamlString))
 	Expect(err).NotTo(HaveOccurred())
 	return workloadv1beta2.AppWrapperComponent{
-		PodSets:  []workloadv1beta2.AppWrapperPodSet{{Replicas: ptr.To(int32(1)), Path: "template"}},
-		Template: runtime.RawExtension{Raw: jsonBytes},
+		DeclaredPodSets: []workloadv1beta2.AppWrapperPodSet{{Replicas: ptr.To(int32(1)), Path: "template"}},
+		Template:        runtime.RawExtension{Raw: jsonBytes},
 	}
 }

internal/controller/appwrapper/fixtures_test.go

@@ -36,24 +36,25 @@ import (
 	utilmaps "sigs.k8s.io/kueue/pkg/util/maps"
 )
 
-func parseComponent(aw *workloadv1beta2.AppWrapper, raw []byte) (*unstructured.Unstructured, error) {
+func parseComponent(raw []byte, expectedNamespace string) (*unstructured.Unstructured, error) {
 	obj := &unstructured.Unstructured{}
 	if _, _, err := unstructured.UnstructuredJSONScheme.Decode(raw, nil, obj); err != nil {
 		return nil, err
 	}
 	namespace := obj.GetNamespace()
 	if namespace == "" {
-		obj.SetNamespace(aw.Namespace)
-	} else if namespace != aw.Namespace {
+		obj.SetNamespace(expectedNamespace)
+	} else if namespace != expectedNamespace {
 		// Should not happen, namespace equality checked by validateAppWrapperInvariants
-		return nil, fmt.Errorf("component namespace \"%s\" is different from appwrapper namespace \"%s\"", namespace, aw.Namespace)
+		return nil, fmt.Errorf("component namespace \"%s\" is different from appwrapper namespace \"%s\"", namespace, expectedNamespace)
 	}
 	return obj, nil
 }
 
 //gocyclo:ignore
 func (r *AppWrapperReconciler) createComponent(ctx context.Context, aw *workloadv1beta2.AppWrapper, componentIdx int) (*unstructured.Unstructured, error, bool) {
 	component := aw.Spec.Components[componentIdx]
+	componentStatus := aw.Status.ComponentStatus[componentIdx]
 	toMap := func(x interface{}) map[string]string {
 		if x == nil {
 			return nil
@@ -77,7 +78,7 @@ func (r *AppWrapperReconciler) createComponent(ctx context.Context, aw *workload
 		}
 	}
 
-	obj, err := parseComponent(aw, component.Template.Raw)
+	obj, err := parseComponent(component.Template.Raw, aw.Namespace)
 	if err != nil {
 		return nil, err, true
 	}
@@ -88,7 +89,7 @@ func (r *AppWrapperReconciler) createComponent(ctx context.Context, aw *workload
 	}
 
 	awLabels := map[string]string{AppWrapperLabel: aw.Name}
-	for podSetsIdx, podSet := range component.PodSets {
+	for podSetsIdx, podSet := range componentStatus.PodSets {
 		toInject := &workloadv1beta2.AppWrapperPodSetInfo{}
 		if r.Config.EnableKueueIntegrations {
 			if podSetsIdx < len(component.PodSetInfos) {

internal/controller/appwrapper/resource_management.go

@@ -39,6 +39,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/log/zap"
 
 	workloadv1beta2 "github.com/project-codeflare/appwrapper/api/v1beta2"
+	"github.com/project-codeflare/appwrapper/internal/controller/awstatus"
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 )
 
@@ -100,6 +101,7 @@ var _ = BeforeSuite(func() {
 	k8sClient, err = client.New(cfg, client.Options{Scheme: scheme})
 	Expect(err).NotTo(HaveOccurred())
 	Expect(k8sClient).NotTo(BeNil())
+	awstatus.CacheClient(k8sClient)
 })
 
 var _ = AfterSuite(func() {

internal/controller/appwrapper/suite_test.go

@@ -0,0 +1,88 @@
+/*
+Copyright 2024 IBM Corporation.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package awstatus
+
+import (
+	"context"
+
+	workloadv1beta2 "github.com/project-codeflare/appwrapper/api/v1beta2"
+	"github.com/project-codeflare/appwrapper/pkg/utils"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+)
+
+var (
+	cachedClient client.Client
+)
+
+const controllerName = "workload.codeflare.dev-appwrapper"
+
+// CacheClient initializes cachedClient; must be called during startup
+func CacheClient(k8sclient client.Client) {
+	cachedClient = k8sclient
+}
+
+// BaseSSAAppWrapper creates a new object based on the input AppWrapper that
+// only contains the fields necessary to identify the original object.
+// The object can be used as a base for Server-Side-Apply.
+func BaseSSAAppWrapper(aw *workloadv1beta2.AppWrapper) *workloadv1beta2.AppWrapper {
+	patch := &workloadv1beta2.AppWrapper{
+		ObjectMeta: metav1.ObjectMeta{
+			UID:       aw.UID,
+			Name:      aw.Name,
+			Namespace: aw.Namespace,
+		},
+		TypeMeta: metav1.TypeMeta{
+			APIVersion: workloadv1beta2.GroupVersion.String(),
+			Kind:       "AppWrapper",
+		},
+	}
+	return patch
+}
+
+// EnsureComponentStatusInitialized initializes aw.Status.ComponenetStatus, including performing PodSet inference for known GVKs
+func EnsureComponentStatusInitialized(ctx context.Context, aw *workloadv1beta2.AppWrapper) error {
+	if len(aw.Status.ComponentStatus) == len(aw.Spec.Components) {
+		return nil
+	}
+
+	// Construct definitive PodSets from the Spec + InferPodSets and cache in the Status (to avoid clashing with user updates to the Spec via apply)
+	compStatus := make([]workloadv1beta2.AppWrapperComponentStatus, len(aw.Spec.Components))
+	for idx := range aw.Spec.Components {
+		if len(aw.Spec.Components[idx].DeclaredPodSets) > 0 {
+			compStatus[idx].PodSets = aw.Spec.Components[idx].DeclaredPodSets
+		} else {
+			obj := &unstructured.Unstructured{}
+			if _, _, err := unstructured.UnstructuredJSONScheme.Decode(aw.Spec.Components[idx].Template.Raw, nil, obj); err != nil {
+				// Transient error; Template.Raw was validated by our AdmissionController
+				return err
+			}
+			podSets, err := utils.InferPodSets(obj)
+			if err != nil {
+				// Transient error; InferPodSets was validated by our AdmissionController
+				return err
+			}
+			compStatus[idx].PodSets = podSets
+		}
+	}
+	aw.Status.ComponentStatus = compStatus
+
+	patch := BaseSSAAppWrapper(aw)
+	patch.Status.ComponentStatus = compStatus
+	return cachedClient.Status().Patch(ctx, patch, client.Apply, client.FieldOwner(controllerName), client.ForceOwnership)
+}

internal/controller/awstatus/status_utils.go

@@ -21,8 +21,8 @@ import (
 	"time"
 
 	"k8s.io/apimachinery/pkg/api/meta"
-	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/runtime/schema"
 
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
@@ -70,20 +70,19 @@ func (r *ChildWorkloadReconciler) Reconcile(ctx context.Context, req ctrl.Reques
 		!meta.IsStatusConditionTrue(aw.Status.Conditions, string(workloadv1beta2.Unhealthy)) {
 		admittedChildren := 0
 		childrenWithPods := 0
-		for componentIdx, component := range aw.Spec.Components {
-			if len(component.PodSets) > 0 {
+		for componentIdx, componentStatus := range aw.Status.ComponentStatus {
+			if len(componentStatus.PodSets) > 0 {
 				childrenWithPods += 1
-				unstruct := &unstructured.Unstructured{}
-				if _, gvk, err := unstructured.UnstructuredJSONScheme.Decode(component.Template.Raw, nil, unstruct); err == nil {
-					wlName := jobframework.GetWorkloadNameForOwnerWithGVK(unstruct.GetName(), *gvk)
+				if gv, err := schema.ParseGroupVersion(componentStatus.APIVersion); err == nil {
+					wlName := jobframework.GetWorkloadNameForOwnerWithGVK(componentStatus.Name, gv.WithKind(componentStatus.Kind))
 					wl := &kueue.Workload{}
 					if err := r.Client.Get(ctx, client.ObjectKey{Namespace: aw.Namespace, Name: wlName}, wl); err == nil {
 						if workload.IsAdmitted(wl) {
 							admittedChildren += 1
 						} else {
 							admission := kueue.Admission{
 								ClusterQueue:      childJobQueueName,
-								PodSetAssignments: make([]kueue.PodSetAssignment, len(aw.Spec.Components[componentIdx].PodSets)),
+								PodSetAssignments: make([]kueue.PodSetAssignment, len(componentStatus.PodSets)),
 							}
 							for i := range admission.PodSetAssignments {
 								admission.PodSetAssignments[i].Name = wl.Spec.PodSets[i].Name

internal/controller/workload/child_admission_controller.go

@@ -17,6 +17,7 @@ limitations under the License.
 package workload
 
 import (
+	"context"
 	"fmt"
 
 	"k8s.io/apimachinery/pkg/api/meta"
@@ -32,6 +33,7 @@ import (
 	"sigs.k8s.io/kueue/pkg/podset"
 
 	workloadv1beta2 "github.com/project-codeflare/appwrapper/api/v1beta2"
+	"github.com/project-codeflare/appwrapper/internal/controller/awstatus"
 	"github.com/project-codeflare/appwrapper/pkg/utils"
 )
 
@@ -77,17 +79,21 @@ func (aw *AppWrapper) GVK() schema.GroupVersionKind {
 
 func (aw *AppWrapper) PodSets() []kueue.PodSet {
 	podSets := []kueue.PodSet{}
-	for componentIdx, component := range aw.Spec.Components {
-		if len(component.PodSets) > 0 {
+	if err := awstatus.EnsureComponentStatusInitialized(context.Background(), (*workloadv1beta2.AppWrapper)(aw)); err != nil {
+		// Kueue will raise an error on zero length PodSet.  Unfortunately, the Kueue API prevents propagating the actual error
+		return podSets
+	}
+	for idx := range aw.Status.ComponentStatus {
+		if len(aw.Status.ComponentStatus[idx].PodSets) > 0 {
 			obj := &unstructured.Unstructured{}
-			if _, _, err := unstructured.UnstructuredJSONScheme.Decode(component.Template.Raw, nil, obj); err != nil {
+			if _, _, err := unstructured.UnstructuredJSONScheme.Decode(aw.Spec.Components[idx].Template.Raw, nil, obj); err != nil {
 				continue // Should be unreachable; Template.Raw validated by our AdmissionController
 			}
-			for psIdx, podSet := range component.PodSets {
+			for psIdx, podSet := range aw.Status.ComponentStatus[idx].PodSets {
 				replicas := utils.Replicas(podSet)
 				if template, err := utils.GetPodTemplateSpec(obj, podSet.Path); err == nil {
 					podSets = append(podSets, kueue.PodSet{
-						Name:     fmt.Sprintf("%s-%v-%v", aw.Name, componentIdx, psIdx),
+						Name:     fmt.Sprintf("%s-%v-%v", aw.Name, idx, psIdx),
 						Template: *template,
 						Count:    replicas,
 					})
@@ -101,17 +107,16 @@ func (aw *AppWrapper) PodSets() []kueue.PodSet {
 // RunWithPodSetsInfo records the assigned PodSetInfos for each component and sets aw.spec.Suspend to false
 func (aw *AppWrapper) RunWithPodSetsInfo(podSetsInfo []podset.PodSetInfo) error {
 	podSetsInfoIndex := 0
-	for componentIdx := range aw.Spec.Components {
-		component := &aw.Spec.Components[componentIdx]
-		if len(component.PodSetInfos) != len(component.PodSets) {
-			component.PodSetInfos = make([]workloadv1beta2.AppWrapperPodSetInfo, len(component.PodSets))
+	for idx := range aw.Spec.Components {
+		if len(aw.Spec.Components[idx].PodSetInfos) != len(aw.Status.ComponentStatus[idx].PodSets) {
+			aw.Spec.Components[idx].PodSetInfos = make([]workloadv1beta2.AppWrapperPodSetInfo, len(aw.Status.ComponentStatus[idx].PodSets))
 		}
-		for podSetIdx := range component.PodSets {
+		for podSetIdx := range aw.Status.ComponentStatus[idx].PodSets {
 			podSetsInfoIndex += 1
 			if podSetsInfoIndex > len(podSetsInfo) {
 				continue // we will return an error below...continuing to get an accurate count for the error message
 			}
-			component.PodSetInfos[podSetIdx] = workloadv1beta2.AppWrapperPodSetInfo{
+			aw.Spec.Components[idx].PodSetInfos[podSetIdx] = workloadv1beta2.AppWrapperPodSetInfo{
 				Annotations:  podSetsInfo[podSetIdx].Annotations,
 				Labels:       podSetsInfo[podSetIdx].Labels,
 				NodeSelector: podSetsInfo[podSetIdx].NodeSelector,