WIP on appwrapper/kueue api adjustment

Author: dgrove-oss

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,15 +74,20 @@ 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
 }
 

pkg/utils/utils.go

@@ -38,7 +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,32 +348,28 @@ 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