Refactor controller into multiple packages (#52)

Cleaner separation between the portions that should be eventually upstreamed to Kueue and those that should be maintained outside of Kueue.

Author: dgrove-oss

Dockerfile

@@ -14,7 +14,7 @@ RUN go mod download
 # Copy the go source
 COPY cmd/main.go cmd/main.go
 COPY api/ api/
-COPY internal/controller/ internal/controller/
+COPY internal/ internal/
 
 # Build
 # the GOARCH has not a default value to allow the binary be built according to the host where the command

README.md

@@ -8,7 +8,9 @@ be jointly queued and admitted for execution using [Kueue](https://kueue.sigs.k8
 AppWrappers provide a flexible and workload-agnostic mechanism for enabling
 Kueue to manage a group of Kubernetes resources
 as a single logical unit without requiring any Kueue-specific support by
-the controllers of those resources.
+the controllers of those resources. If some of the contained resources
+are Kueue-enabled, the AppWrapper controller ensures that the admission
+by Kueue of the parent AppWrapper will be propagated appropriately.
 
 ## Description
 
@@ -26,7 +28,7 @@ set to true. We also leverage the Admission Controller to ensure that
 the user creating the AppWrapper is also entitled to create the contained resources
 and to validate AppWrapper-specific invariants.
 
-See [appwrapper_webhook.go](./internal/controller/appwrapper_webhook.go)
+See [appwrapper_webhook.go](./internal/webhook/appwrapper_webhook.go)
 for the implementation.
 
 #### Workload Controller
@@ -39,7 +41,15 @@ the two. This controller will make it possible for Kueue to suspend,
 resume, and constrain the placement of the AppWrapper. It will report
 the status of the AppWrapper to Kueue.
 
-See [workload_controller.go](./internal/controller/workload_controller.go)
+See [workload_controller.go](./internal/controller/workload/workload_controller.go)
+for the implementation.
+
+A small additional piece of logic is currently needed to generalize
+Kueue's ability to recognize parent/children relationships and enforce
+that admission by Kueue of the parent AppWrapper will be propagated to
+its immediate children.
+
+See [child_admission_controller.go](./internal/controller/workload/child_admission_controller.go)
 for the implementation.
 
 #### Framework Controller
@@ -52,7 +62,7 @@ status made by the Workload Controller described above.
 
 This [state transition diagram](docs/state-diagram.md) depicts the
 lifecycle of an AppWrapper; the implementation is found in
-[appwrapper_controller.go](./internal/controller/appwrapper_controller.go).
+[appwrapper_controller.go](./internal/controller/appwrapper/appwrapper_controller.go).
 
 ## Getting Started
 

cmd/main.go

@@ -36,6 +36,7 @@ import (
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 
 	workloadv1beta2 "github.com/project-codeflare/appwrapper/api/v1beta2"
+	"github.com/project-codeflare/appwrapper/internal/config"
 	"github.com/project-codeflare/appwrapper/internal/controller"
 	//+kubebuilder:scaffold:imports
 )
@@ -61,7 +62,7 @@ func main() {
 	var secureMetrics bool
 	var enableHTTP2 bool
 
-	config := controller.AppWrapperConfig{}
+	awConfig := config.AppWrapperConfig{}
 
 	flag.StringVar(&metricsAddr, "metrics-bind-address", ":8080", "The address the metric endpoint binds to.")
 	flag.StringVar(&probeAddr, "health-probe-bind-address", ":8081", "The address the probe endpoint binds to.")
@@ -72,7 +73,7 @@ func main() {
 		"If set the metrics endpoint is served securely")
 	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.BoolVar(&awConfig.ManageJobsWithoutQueueName, "manage-no-queue", true, "Manage AppWrappers without queue names")
 	opts := zap.Options{
 		Development: true,
 	}
@@ -81,7 +82,7 @@ func main() {
 
 	ctrl.SetLogger(zap.New(zap.UseFlagOptions(&opts)))
 	setupLog.Info("Build info", "version", BuildVersion, "date", BuildDate)
-	setupLog.Info("Configuration", "config", config)
+	setupLog.Info("Configuration", "config", awConfig)
 
 	// if the enable-http2 flag is false (the default), http/2 should be disabled
 	// due to its vulnerabilities. More specifically, disabling http/2 will
@@ -132,7 +133,7 @@ func main() {
 	}
 
 	ctx := ctrl.SetupSignalHandler()
-	err = controller.SetupWithManager(ctx, mgr, &config)
+	err = controller.SetupWithManager(ctx, mgr, &awConfig)
 	if err != nil {
 		setupLog.Error(err, "unable to start appwrapper controllers")
 		os.Exit(1)

internal/config/config.go

@@ -0,0 +1,21 @@
+/*
+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 config
+
+type AppWrapperConfig struct {
+	ManageJobsWithoutQueueName bool `json:"manageJobsWithoutQueueName,omitempty"`
+}

internal/controller/appwrapper_controller.go renamed to internal/controller/appwrapper/appwrapper_controller.go

@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package controller
+package appwrapper
 
 import (
 	"context"
@@ -36,27 +36,26 @@ import (
 	"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"
+	"github.com/project-codeflare/appwrapper/internal/config"
+	"github.com/project-codeflare/appwrapper/internal/utils"
 )
 
 const (
 	AppWrapperLabel     = "workload.codeflare.dev/appwrapper"
-	appWrapperFinalizer = "workload.codeflare.dev/finalizer"
+	AppWrapperFinalizer = "workload.codeflare.dev/finalizer"
 	childJobQueueName   = "workload.codeflare.dev.admitted"
 )
 
 // AppWrapperReconciler reconciles an appwrapper
 type AppWrapperReconciler struct {
 	client.Client
 	Scheme *runtime.Scheme
-	Config *AppWrapperConfig
+	Config *config.AppWrapperConfig
 }
 
 type podStatusSummary struct {
@@ -72,10 +71,6 @@ type podStatusSummary struct {
 //+kubebuilder:rbac:groups=workload.codeflare.dev,resources=appwrappers/status,verbs=get;update;patch
 //+kubebuilder:rbac:groups=workload.codeflare.dev,resources=appwrappers/finalizers,verbs=update
 
-// permission to manipulate workloads controlling appwrapper components to enable admitting them to our pseudo-clusterqueue
-// +kubebuilder:rbac:groups=kueue.x-k8s.io,resources=workloads,verbs=get
-// +kubebuilder:rbac:groups=kueue.x-k8s.io,resources=workloads/status,verbs=get;update;patch
-
 // permission to edit wrapped resources: pods, services, jobs, podgroups, pytorchjobs, rayclusters
 
 //+kubebuilder:rbac:groups="",resources=pods;services,verbs=get;list;watch;create;update;patch;delete
@@ -100,7 +95,7 @@ func (r *AppWrapperReconciler) Reconcile(ctx context.Context, req ctrl.Request)
 
 	// handle deletion first
 	if !aw.DeletionTimestamp.IsZero() {
-		if controllerutil.ContainsFinalizer(aw, appWrapperFinalizer) {
+		if controllerutil.ContainsFinalizer(aw, AppWrapperFinalizer) {
 			statusUpdated := false
 			if meta.IsStatusConditionTrue(aw.Status.Conditions, string(workloadv1beta2.ResourcesDeployed)) {
 				if !r.deleteComponents(ctx, aw) {
@@ -136,7 +131,7 @@ func (r *AppWrapperReconciler) Reconcile(ctx context.Context, req ctrl.Request)
 				}
 			}
 
-			if controllerutil.RemoveFinalizer(aw, appWrapperFinalizer) {
+			if controllerutil.RemoveFinalizer(aw, AppWrapperFinalizer) {
 				if err := r.Update(ctx, aw); err != nil {
 					return ctrl.Result{}, err
 				}
@@ -149,7 +144,7 @@ func (r *AppWrapperReconciler) Reconcile(ctx context.Context, req ctrl.Request)
 	switch aw.Status.Phase {
 
 	case workloadv1beta2.AppWrapperEmpty: // initial state, inject finalizer
-		if controllerutil.AddFinalizer(aw, appWrapperFinalizer) {
+		if controllerutil.AddFinalizer(aw, AppWrapperFinalizer) {
 			if err := r.Update(ctx, aw); err != nil {
 				return ctrl.Result{}, err
 			}
@@ -224,7 +219,6 @@ 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,
@@ -339,7 +333,7 @@ func (r *AppWrapperReconciler) createComponent(ctx context.Context, aw *workload
 	for podSetsIdx, podSet := range component.PodSets {
 		toInject := component.PodSetInfos[podSetsIdx]
 
-		p, err := getRawTemplate(obj.UnstructuredContent(), podSet.Path)
+		p, err := utils.GetRawTemplate(obj.UnstructuredContent(), podSet.Path)
 		if err != nil {
 			return nil, err, true // Should not happen, path validity is enforced by validateAppWrapperInvariants
 		}
@@ -411,36 +405,6 @@ func (r *AppWrapperReconciler) createComponents(ctx context.Context, aw *workloa
 	return nil, false
 }
 
-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)
-					}
-				}
-			}
-		}
-	}
-}
-
 func (r *AppWrapperReconciler) deleteComponents(ctx context.Context, aw *workloadv1beta2.AppWrapper) bool {
 	// TODO forceful deletion: See https://github.com/project-codeflare/appwrapper/issues/36
 	log := log.FromContext(ctx)
@@ -478,7 +442,7 @@ func (r *AppWrapperReconciler) workloadStatus(ctx context.Context, aw *workloadv
 		client.MatchingLabels{AppWrapperLabel: aw.Name}); err != nil {
 		return nil, err
 	}
-	summary := &podStatusSummary{expected: ExpectedPodCount(aw)}
+	summary := &podStatusSummary{expected: utils.ExpectedPodCount(aw)}
 
 	for _, pod := range pods.Items {
 		switch pod.Status.Phase {
@@ -496,24 +460,6 @@ func (r *AppWrapperReconciler) workloadStatus(ctx context.Context, aw *workloadv
 	return summary, nil
 }
 
-func replicas(ps workloadv1beta2.AppWrapperPodSet) int32 {
-	if ps.Replicas == nil {
-		return 1
-	} else {
-		return *ps.Replicas
-	}
-}
-
-func ExpectedPodCount(aw *workloadv1beta2.AppWrapper) int32 {
-	var expected int32
-	for _, c := range aw.Spec.Components {
-		for _, s := range c.PodSets {
-			expected += replicas(s)
-		}
-	}
-	return expected
-}
-
 // SetupWithManager sets up the controller with the Manager.
 func (r *AppWrapperReconciler) SetupWithManager(mgr ctrl.Manager) error {
 	return ctrl.NewControllerManagedBy(mgr).