Propagate admission to child workload objects (#46)

If an AppWrapper contains a component that has its own Kueue-enabled controller, propagate the admission status
of the AppWrapper to the component's Workload object.

Author: dgrove-oss

internal/controller/appwrapper_config.go

@@ -22,6 +22,7 @@ import (
 	"os"
 
 	ctrl "sigs.k8s.io/controller-runtime"
+
 	"sigs.k8s.io/kueue/pkg/controller/jobframework"
 )
 

internal/controller/appwrapper_controller.go

@@ -28,21 +28,28 @@ import (
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
+
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
 	"sigs.k8s.io/controller-runtime/pkg/handler"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/controller-runtime/pkg/reconcile"
+
+	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
+	"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"
+	"sigs.k8s.io/kueue/pkg/workload"
 
 	workloadv1beta2 "github.com/project-codeflare/appwrapper/api/v1beta2"
 )
 
 const (
 	AppWrapperLabel     = "workload.codeflare.dev/appwrapper"
 	appWrapperFinalizer = "workload.codeflare.dev/finalizer"
+	childJobQueueName   = "workload.codeflare.dev.admitted"
 )
 
 // AppWrapperReconciler reconciles an appwrapper
@@ -212,6 +219,7 @@ func (r *AppWrapperReconciler) Reconcile(ctx context.Context, req ctrl.Request)
 			})
 			return ctrl.Result{RequeueAfter: time.Minute}, r.Status().Update(ctx, aw)
 		}
+		r.propagateAdmission(ctx, aw)
 		meta.SetStatusCondition(&aw.Status.Conditions, metav1.Condition{
 			Type:   string(workloadv1beta2.PodsReady),
 			Status: metav1.ConditionFalse,
@@ -266,14 +274,6 @@ func (r *AppWrapperReconciler) Reconcile(ctx context.Context, req ctrl.Request)
 	return ctrl.Result{}, nil
 }
 
-// SetupWithManager sets up the controller with the Manager.
-func (r *AppWrapperReconciler) SetupWithManager(mgr ctrl.Manager) error {
-	return ctrl.NewControllerManagedBy(mgr).
-		For(&workloadv1beta2.AppWrapper{}).
-		Watches(&v1.Pod{}, handler.EnqueueRequestsFromMapFunc(r.podMapFunc)).
-		Complete(r)
-}
-
 // podMapFunc maps pods to appwrappers
 func (r *AppWrapperReconciler) podMapFunc(ctx context.Context, obj client.Object) []reconcile.Request {
 	pod := obj.(*v1.Pod)
@@ -299,7 +299,8 @@ func parseComponent(aw *workloadv1beta2.AppWrapper, raw []byte) (*unstructured.U
 	return obj, nil
 }
 
-func materializeObject(aw *workloadv1beta2.AppWrapper, component *workloadv1beta2.AppWrapperComponent) (client.Object, error) {
+func (r *AppWrapperReconciler) createComponent(ctx context.Context, aw *workloadv1beta2.AppWrapper, componentIdx int) (*unstructured.Unstructured, error, bool) {
+	component := aw.Spec.Components[componentIdx]
 	toMap := func(x interface{}) map[string]string {
 		if x == nil {
 			return nil
@@ -322,19 +323,20 @@ func materializeObject(aw *workloadv1beta2.AppWrapper, component *workloadv1beta
 			}
 		}
 	}
-	awLabels := map[string]string{AppWrapperLabel: aw.Name}
 
 	obj, err := parseComponent(aw, component.Template.Raw)
 	if err != nil {
-		return nil, err
+		return nil, err, true
 	}
+	obj.SetLabels(utilmaps.MergeKeepFirst(obj.GetLabels(), map[string]string{AppWrapperLabel: aw.Name, constants.QueueLabel: childJobQueueName}))
 
+	awLabels := map[string]string{AppWrapperLabel: aw.Name}
 	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
+			return nil, err, true // Should not happen, path validity is enforced by validateAppWrapperInvariants
 		}
 		if md, ok := p["metadata"]; !ok || md == nil {
 			p["metadata"] = make(map[string]interface{})
@@ -346,7 +348,7 @@ func materializeObject(aw *workloadv1beta2.AppWrapper, component *workloadv1beta
 		if len(toInject.Annotations) > 0 {
 			existing := toMap(metadata["annotations"])
 			if err := utilmaps.HaveConflict(existing, toInject.Annotations); err != nil {
-				return nil, podset.BadPodSetsUpdateError("annotations", err)
+				return nil, podset.BadPodSetsUpdateError("annotations", err), true
 			}
 			metadata["annotations"] = utilmaps.MergeKeepFirst(existing, toInject.Annotations)
 		}
@@ -355,15 +357,15 @@ func materializeObject(aw *workloadv1beta2.AppWrapper, component *workloadv1beta
 		mergedLabels := utilmaps.MergeKeepFirst(toInject.Labels, awLabels)
 		existing := toMap(metadata["labels"])
 		if err := utilmaps.HaveConflict(existing, mergedLabels); err != nil {
-			return nil, podset.BadPodSetsUpdateError("labels", err)
+			return nil, podset.BadPodSetsUpdateError("labels", err), true
 		}
 		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)
+				return nil, podset.BadPodSetsUpdateError("nodeSelector", err), true
 			}
 			metadata["nodeSelector"] = utilmaps.MergeKeepFirst(existing, toInject.NodeSelector)
 		}
@@ -381,27 +383,57 @@ func materializeObject(aw *workloadv1beta2.AppWrapper, component *workloadv1beta
 		}
 	}
 
-	return obj, nil
+	if err := controllerutil.SetControllerReference(aw, obj, r.Scheme); err != nil {
+		return nil, err, true
+	}
+
+	if err := r.Create(ctx, obj); err != nil {
+		if !apierrors.IsAlreadyExists(err) {
+			return nil, err, meta.IsNoMatchError(err) || apierrors.IsInvalid(err) // fatal
+		}
+	}
+
+	return obj, nil, false
 }
 
 func (r *AppWrapperReconciler) createComponents(ctx context.Context, aw *workloadv1beta2.AppWrapper) (error, bool) {
-	for _, component := range aw.Spec.Components {
-		obj, err := materializeObject(aw, &component)
+	for componentIdx := range aw.Spec.Components {
+		_, err, fatal := r.createComponent(ctx, aw, componentIdx)
 		if err != nil {
-			return err, true
+			return err, fatal
 		}
+	}
+	return nil, false
+}
 
-		if err := controllerutil.SetControllerReference(aw, obj, r.Scheme); err != nil {
-			return err, true
-		}
-		if err := r.Create(ctx, obj); err != nil {
-			if apierrors.IsAlreadyExists(err) {
-				continue // ignore existing component
+func (r *AppWrapperReconciler) propagateAdmission(ctx context.Context, aw *workloadv1beta2.AppWrapper) {
+	for componentIdx, component := range aw.Spec.Components {
+		if len(component.PodSets) > 0 {
+			obj, err := parseComponent(aw, component.Template.Raw)
+			if err != nil {
+				return
+			}
+			wlName := jobframework.GetWorkloadNameForOwnerWithGVK(obj.GetName(), obj.GroupVersionKind())
+			wl := &kueue.Workload{}
+			if err := r.Client.Get(ctx, client.ObjectKey{Namespace: aw.Namespace, Name: wlName}, wl); err == nil {
+				if !workload.IsAdmitted(wl) {
+					admission := kueue.Admission{
+						ClusterQueue:      childJobQueueName,
+						PodSetAssignments: make([]kueue.PodSetAssignment, len(aw.Spec.Components[componentIdx].PodSets)),
+					}
+					for i := range admission.PodSetAssignments {
+						admission.PodSetAssignments[i].Name = wl.Spec.PodSets[i].Name
+					}
+					newWorkload := wl.DeepCopy()
+					workload.SetQuotaReservation(newWorkload, &admission)
+					_ = workload.SyncAdmittedCondition(newWorkload)
+					if err = workload.ApplyAdmissionStatus(ctx, r.Client, newWorkload, false); err != nil {
+						log.FromContext(ctx).Error(err, "syncing admission", "appwrapper", aw, "componentIdx", componentIdx, "workload", wl, "newworkload", newWorkload)
+					}
+				}
 			}
-			return err, meta.IsNoMatchError(err) || apierrors.IsInvalid(err) // fatal
 		}
 	}
-	return nil, false
 }
 
 func (r *AppWrapperReconciler) deleteComponents(ctx context.Context, aw *workloadv1beta2.AppWrapper) bool {
@@ -476,3 +508,11 @@ func ExpectedPodCount(aw *workloadv1beta2.AppWrapper) int32 {
 	}
 	return expected
 }
+
+// SetupWithManager sets up the controller with the Manager.
+func (r *AppWrapperReconciler) SetupWithManager(mgr ctrl.Manager) error {
+	return ctrl.NewControllerManagedBy(mgr).
+		For(&workloadv1beta2.AppWrapper{}).
+		Watches(&v1.Pod{}, handler.EnqueueRequestsFromMapFunc(r.podMapFunc)).
+		Complete(r)
+}

internal/controller/appwrapper_webhook.go

@@ -21,7 +21,6 @@ import (
 	"context"
 	"fmt"
 
-	workloadv1beta2 "github.com/project-codeflare/appwrapper/api/v1beta2"
 	authv1 "k8s.io/api/authorization/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
@@ -31,11 +30,15 @@ import (
 	discovery "k8s.io/client-go/discovery"
 	"k8s.io/client-go/kubernetes"
 	authClientv1 "k8s.io/client-go/kubernetes/typed/authorization/v1"
+
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/controller-runtime/pkg/webhook"
 	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
+
 	"sigs.k8s.io/kueue/pkg/controller/jobframework"
+
+	workloadv1beta2 "github.com/project-codeflare/appwrapper/api/v1beta2"
 )
 
 type AppWrapperWebhook struct {

internal/controller/suite_test.go

@@ -46,6 +46,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/webhook"
 
 	workloadv1beta2 "github.com/project-codeflare/appwrapper/api/v1beta2"
+	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 )
 
 // These tests use Ginkgo (BDD-style Go testing framework). Refer to
@@ -104,6 +105,8 @@ var _ = BeforeSuite(func() {
 	Expect(err).NotTo(HaveOccurred())
 	err = clientgoscheme.AddToScheme(scheme)
 	Expect(err).NotTo(HaveOccurred())
+	err = kueue.AddToScheme(scheme)
+	Expect(err).NotTo(HaveOccurred())
 
 	//+kubebuilder:scaffold:scheme
 

internal/controller/utils.go

@@ -23,6 +23,9 @@ import (
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime"
+
+	workloadv1beta2 "github.com/project-codeflare/appwrapper/api/v1beta2"
+	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 )
 
 // getPodTemplateSpec extracts a Kueue-compatible PodTemplateSpec at the given path within obj
@@ -56,6 +59,26 @@ func getPodTemplateSpec(obj *unstructured.Unstructured, path string) (*v1.PodTem
 	return template, nil
 }
 
+func getKueuePodSets(obj *unstructured.Unstructured, component *workloadv1beta2.AppWrapperComponent, awName string, componentIdx int) ([]kueue.PodSet, error) {
+	podSets := []kueue.PodSet{}
+	for psIdx, podSet := range component.PodSets {
+		replicas := int32(1)
+		if podSet.Replicas != nil {
+			replicas = *podSet.Replicas
+		}
+		template, err := getPodTemplateSpec(obj, podSet.Path)
+		if err != nil {
+			return nil, err
+		}
+		podSets = append(podSets, kueue.PodSet{
+			Name:     fmt.Sprintf("%s-%v-%v", awName, componentIdx, psIdx),
+			Template: *template,
+			Count:    replicas,
+		})
+	}
+	return podSets, nil
+}
+
 // return the subobject found at the given path, or nil if the path is invalid
 func getRawTemplate(obj map[string]interface{}, path string) (map[string]interface{}, error) {
 	parts := strings.Split(path, ".")

internal/controller/workload_controller.go

@@ -17,13 +17,13 @@ limitations under the License.
 package controller
 
 import (
-	"fmt"
-
 	"k8s.io/apimachinery/pkg/api/meta"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime/schema"
+
 	"sigs.k8s.io/controller-runtime/pkg/client"
+
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 	"sigs.k8s.io/kueue/pkg/controller/jobframework"
 	"sigs.k8s.io/kueue/pkg/podset"
@@ -68,26 +68,15 @@ func (aw *AppWrapper) GVK() schema.GroupVersionKind {
 
 func (aw *AppWrapper) PodSets() []kueue.PodSet {
 	podSets := []kueue.PodSet{}
-	i := 0
-	for _, component := range aw.Spec.Components {
-		obj := &unstructured.Unstructured{}
-		if _, _, err := unstructured.UnstructuredJSONScheme.Decode(component.Template.Raw, nil, obj); err != nil {
-			continue
-		}
-
-		for _, podSet := range component.PodSets {
-			replicas := int32(1)
-			if podSet.Replicas != nil {
-				replicas = *podSet.Replicas
+	for componentIdx, component := range aw.Spec.Components {
+		if len(component.PodSets) > 0 {
+			obj := &unstructured.Unstructured{}
+			if _, _, err := unstructured.UnstructuredJSONScheme.Decode(component.Template.Raw, nil, obj); err != nil {
+				continue // Should be unreachable; Template.Raw validated by our AdmissionController
 			}
-			template, err := getPodTemplateSpec(obj, podSet.Path)
+			toAdd, err := getKueuePodSets(obj, &component, aw.Name, componentIdx)
 			if err == nil {
-				podSets = append(podSets, kueue.PodSet{
-					Name:     aw.Name + "-" + fmt.Sprint(i),
-					Template: *template,
-					Count:    replicas,
-				})
-				i++
+				podSets = append(podSets, toAdd...)
 			}
 		}
 	}
@@ -99,13 +88,13 @@ func (aw *AppWrapper) RunWithPodSetsInfo(podSetsInfo []podset.PodSetInfo) error
 	podSetsInfoIndex := 0
 	for componentIdx := range aw.Spec.Components {
 		component := &aw.Spec.Components[componentIdx]
-		if len(component.PodSets) != len(component.PodSetInfos) {
+		if len(component.PodSetInfos) != len(component.PodSets) {
 			component.PodSetInfos = make([]workloadv1beta2.AppWrapperPodSetInfo, len(component.PodSets))
 		}
 		for podSetIdx := range component.PodSets {
 			podSetsInfoIndex += 1
 			if podSetsInfoIndex > len(podSetsInfo) {
-				continue // we're going to return an error below...just continuing to get an accurate count for the error message
+				continue // we will return an error below...continuing to get an accurate count for the error message
 			}
 			component.PodSetInfos[podSetIdx] = workloadv1beta2.AppWrapperPodSetInfo{
 				Annotations:  podSetsInfo[podSetIdx].Annotations,

samples/wrapped-job.yaml

@@ -22,7 +22,7 @@ spec:
             containers:
             - name: busybox
               image: quay.io/project-codeflare/busybox:1.36
-              command: ["sh", "-c", "sleep 600"]
+              command: ["sh", "-c", "sleep 30"]
               resources:
                 requests:
                   cpu: 1