adjust appwrapper/kueue api to push kueue types into workload controller

Author: dgrove-oss

internal/controller/appwrapper/appwrapper_controller_test.go

@@ -29,7 +29,6 @@ import (
 	"k8s.io/client-go/tools/record"
 	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
 	"sigs.k8s.io/controller-runtime/pkg/reconcile"
-	"sigs.k8s.io/kueue/pkg/podset"
 
 	awv1beta2 "github.com/project-codeflare/appwrapper/api/v1beta2"
 	"github.com/project-codeflare/appwrapper/pkg/config"
@@ -39,7 +38,7 @@ import (
 var _ = Describe("AppWrapper Controller", func() {
 	var awReconciler *AppWrapperReconciler
 	var awName types.NamespacedName
-	markerPodSet := podset.PodSetInfo{
+	markerPodSet := awv1beta2.AppWrapperPodSetInfo{
 		Labels:          map[string]string{"testkey1": "value1"},
 		Annotations:     map[string]string{"test2": "test2"},
 		NodeSelector:    map[string]string{"nodeName": "myNode"},
@@ -79,7 +78,7 @@ var _ = Describe("AppWrapper Controller", func() {
 		Expect(aw.Status.Phase).Should(Equal(awv1beta2.AppWrapperSuspended))
 
 		By("Updating aw.Spec by invoking utils.SetPodSetInfos and setting suspend to false")
-		Expect(utils.SetPodSetInfos(aw, []podset.PodSetInfo{markerPodSet, markerPodSet})).To(Succeed())
+		Expect(utils.SetPodSetInfos(aw, []awv1beta2.AppWrapperPodSetInfo{markerPodSet, markerPodSet})).To(Succeed())
 		aw.Spec.Suspend = false
 		Expect(k8sClient.Update(ctx, aw)).To(Succeed())
 

internal/controller/workload/workload_controller.go

@@ -17,6 +17,8 @@ limitations under the License.
 package workload
 
 import (
+	"fmt"
+
 	"k8s.io/apimachinery/pkg/api/meta"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 
@@ -72,21 +74,35 @@ func (aw *AppWrapper) GVK() schema.GroupVersionKind {
 }
 
 func (aw *AppWrapper) PodSets() []kueue.PodSet {
-	podSets, err := utils.GetPodSets((*awv1beta2.AppWrapper)(aw))
+	podSpecTemplates, awPodSets, err := utils.GetComponentPodSpecs((*awv1beta2.AppWrapper)(aw))
 	if err != nil {
 		// Kueue will raise an error on zero length PodSet; the Kueue GenericJob API prevents propagating the actual error.
 		return []kueue.PodSet{}
 	}
-	for psIndex := range podSets {
-		podSets[psIndex].TopologyRequest = jobframework.PodSetTopologyRequest(&podSets[psIndex].Template.ObjectMeta, nil, nil, nil)
+	podSets := []kueue.PodSet{}
+	for psIndex := range podSpecTemplates {
+		podSets = append(podSets, kueue.PodSet{
+			Name:            fmt.Sprintf("%s-%v", aw.Name, psIndex),
+			Template:        *podSpecTemplates[psIndex],
+			Count:           utils.Replicas(awPodSets[psIndex]),
+			TopologyRequest: jobframework.PodSetTopologyRequest(&(podSpecTemplates[psIndex].ObjectMeta), nil, nil, nil),
+		})
 	}
-
 	return podSets
 }
 
 func (aw *AppWrapper) RunWithPodSetsInfo(podSetsInfo []podset.PodSetInfo) error {
-	if err := utils.SetPodSetInfos((*awv1beta2.AppWrapper)(aw), podSetsInfo); err != nil {
-		return err
+	awPodSetsInfo := make([]awv1beta2.AppWrapperPodSetInfo, len(podSetsInfo))
+	for idx := range podSetsInfo {
+		awPodSetsInfo[idx].Annotations = podSetsInfo[idx].Annotations
+		awPodSetsInfo[idx].Labels = podSetsInfo[idx].Labels
+		awPodSetsInfo[idx].NodeSelector = podSetsInfo[idx].NodeSelector
+		awPodSetsInfo[idx].Tolerations = podSetsInfo[idx].Tolerations
+		awPodSetsInfo[idx].SchedulingGates = podSetsInfo[idx].SchedulingGates
+	}
+
+	if err := utils.SetPodSetInfos((*awv1beta2.AppWrapper)(aw), awPodSetsInfo); err != nil {
+		return fmt.Errorf("%w: %v", podset.ErrInvalidPodsetInfo, err)
 	}
 	aw.Spec.Suspend = false
 	return nil

pkg/utils/utils.go

@@ -38,9 +38,6 @@ import (
 	"k8s.io/utils/ptr"
 	jobsetapi "sigs.k8s.io/jobset/api/jobset/v1alpha2"
 
-	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
-	"sigs.k8s.io/kueue/pkg/podset"
-
 	awv1beta2 "github.com/project-codeflare/appwrapper/api/v1beta2"
 )
 
@@ -349,36 +346,32 @@ func EnsureComponentStatusInitialized(aw *awv1beta2.AppWrapper) error {
 	return nil
 }
 
-// GetPodSets constructs the kueue.PodSets for an AppWrapper
-func GetPodSets(aw *awv1beta2.AppWrapper) ([]kueue.PodSet, error) {
-	podSets := []kueue.PodSet{}
+func GetComponentPodSpecs(aw *awv1beta2.AppWrapper) ([]*v1.PodTemplateSpec, []awv1beta2.AppWrapperPodSet, error) {
+	templates := []*v1.PodTemplateSpec{}
+	podSets := []awv1beta2.AppWrapperPodSet{}
 	if err := EnsureComponentStatusInitialized(aw); err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	for idx := range aw.Status.ComponentStatus {
 		if len(aw.Status.ComponentStatus[idx].PodSets) > 0 {
 			obj := &unstructured.Unstructured{}
 			if _, _, err := unstructured.UnstructuredJSONScheme.Decode(aw.Spec.Components[idx].Template.Raw, nil, obj); err != nil {
 				// Should be unreachable; Template.Raw validated by AppWrapper AdmissionController
-				return nil, err
+				return nil, nil, err
 			}
-			for psIdx, podSet := range aw.Status.ComponentStatus[idx].PodSets {
-				replicas := Replicas(podSet)
+			for _, podSet := range aw.Status.ComponentStatus[idx].PodSets {
 				if template, err := GetPodTemplateSpec(obj, podSet.Path); err == nil {
-					podSets = append(podSets, kueue.PodSet{
-						Name:     fmt.Sprintf("%s-%v-%v", aw.Name, idx, psIdx),
-						Template: *template,
-						Count:    replicas,
-					})
+					templates = append(templates, template)
+					podSets = append(podSets, podSet)
 				}
 			}
 		}
 	}
-	return podSets, nil
+	return templates, podSets, nil
 }
 
 // SetPodSetInfos propagates podSetsInfo into the PodSetInfos of aw.Spec.Components
-func SetPodSetInfos(aw *awv1beta2.AppWrapper, podSetsInfo []podset.PodSetInfo) error {
+func SetPodSetInfos(aw *awv1beta2.AppWrapper, podSetsInfo []awv1beta2.AppWrapperPodSetInfo) error {
 	if err := EnsureComponentStatusInitialized(aw); err != nil {
 		return err
 	}
@@ -392,18 +385,12 @@ func SetPodSetInfos(aw *awv1beta2.AppWrapper, podSetsInfo []podset.PodSetInfo) e
 			if podSetsInfoIndex > len(podSetsInfo) {
 				continue // we will return an error below...continuing to get an accurate count for the error message
 			}
-			aw.Spec.Components[idx].PodSetInfos[podSetIdx] = awv1beta2.AppWrapperPodSetInfo{
-				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,
-			}
+			aw.Spec.Components[idx].PodSetInfos[podSetIdx] = podSetsInfo[podSetIdx]
 		}
 	}
 
 	if podSetsInfoIndex != len(podSetsInfo) {
-		return podset.BadPodSetsInfoLenError(podSetsInfoIndex, len(podSetsInfo))
+		return fmt.Errorf("expecting %d podsets, got %d", podSetsInfoIndex, len(podSetsInfo))
 	}
 	return nil
 }