Create Slice for each PodSet in Workload.

Author: mbobrovskyi

slice/cmd/main.go

@@ -232,12 +232,18 @@ func main() {
 		}
 	}
 
+	ctx := ctrl.SetupSignalHandler()
+	if err := controller.SetupSliceIndexer(ctx, mgr.GetFieldIndexer()); err != nil {
+		setupLog.Error(err, "unable to setup indexes")
+		os.Exit(1)
+	}
+
 	go setupControllers(mgr, certsReady)
 
 	setupProbeEndpoints(mgr, certsReady)
 
 	setupLog.Info("starting manager")
-	if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
+	if err := mgr.Start(ctx); err != nil {
 		setupLog.Error(err, "problem running manager")
 		os.Exit(1)
 	}

slice/internal/controller/indexer.go

@@ -0,0 +1,44 @@
+/*
+Copyright The Kubernetes Authors.
+
+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 controller
+
+import (
+	"context"
+	"fmt"
+
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+
+	"tpu-slice-controller/api/v1alpha1"
+	"tpu-slice-controller/internal/util/slices"
+)
+
+const (
+	OwnerReferenceUID = "metadata.ownerReferences.uid"
+)
+
+func indexOwnerReferenceUID(obj client.Object) []string {
+	return slices.Map(obj.GetOwnerReferences(), func(o *metav1.OwnerReference) string { return string(o.UID) })
+}
+
+// SetupSliceIndexer configures the indexer to index specific fields for v1alpha1.Slice resources.
+func SetupSliceIndexer(ctx context.Context, indexer client.FieldIndexer) error {
+	if err := indexer.IndexField(ctx, &v1alpha1.Slice{}, OwnerReferenceUID, indexOwnerReferenceUID); err != nil {
+		return fmt.Errorf("setting index on ownerReferences.uid for Slice: %w", err)
+	}
+	return nil
+}

slice/internal/controller/workload_controller.go

@@ -20,12 +20,12 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"k8s.io/apimachinery/pkg/api/meta"
+	"strings"
 	"time"
 
-	"github.com/go-logr/logr"
 	corev1 "k8s.io/api/core/v1"
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
-	"k8s.io/apimachinery/pkg/api/meta"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/apimachinery/pkg/util/sets"
@@ -100,22 +100,19 @@ func (r *WorkloadReconciler) Reconcile(ctx context.Context, req ctrl.Request) (c
 	log.V(3).Info("Reconcile Workload")
 
 	if finalize, reason := shouldFinalize(wl); finalize {
-		if controllerutil.ContainsFinalizer(wl, SliceControllerName) {
-			log.V(3).Info(fmt.Sprintf("Cleaning up the Slice and finalize the Workload because %s", reason))
-			err = r.client.Delete(ctx, core.SliceWithMetadata(wl))
-			if client.IgnoreNotFound(err) != nil {
-				return ctrl.Result{}, err
-			}
-			controllerutil.RemoveFinalizer(wl, SliceControllerName)
-			if err := r.client.Update(ctx, wl); err != nil {
-				if !apierrors.IsNotFound(err) {
-					log.Error(err, "Failed to remove finalizer")
-				}
-				return ctrl.Result{}, client.IgnoreNotFound(err)
-			}
-			log.V(3).Info("Removed finalizer")
+		if !controllerutil.ContainsFinalizer(wl, SliceControllerName) {
+			return ctrl.Result{}, nil
 		}
-		return ctrl.Result{}, nil
+
+		log.V(3).Info(fmt.Sprintf("Cleaning up the Slice and finalize the Workload because %s", reason))
+
+		err := r.cleanupSlices(ctx, wl)
+		if err != nil {
+			return ctrl.Result{}, err
+		}
+
+		err = r.finalizeWorkload(ctx, wl)
+		return ctrl.Result{}, client.IgnoreNotFound(err)
 	}
 
 	if err = validateRelevantWorkload(wl); err != nil {
@@ -146,19 +143,12 @@ func (r *WorkloadReconciler) Reconcile(ctx context.Context, req ctrl.Request) (c
 		return ctrl.Result{}, nil
 	}
 
-	slice := v1alpha1.Slice{}
-	err = r.client.Get(ctx, core.SliceKeyFromWorkload(wl), &slice)
-	if apierrors.IsNotFound(err) {
-		// slice not found, create it and exit.
-		err = r.createSlice(ctx, log, wl, ac)
-		return ctrl.Result{}, err
-	} else if err != nil {
-		// error fetching slice
-		log.Error(err, "Failed to fetch the Slice")
+	slices, changed, err := r.syncSlices(ctx, wl, ac)
+	if err != nil || changed {
 		return ctrl.Result{}, err
 	}
 
-	err = r.syncAdmissionCheckStatus(ctx, wl, ac, &slice)
+	err = r.syncAdmissionCheckStatus(ctx, wl, ac, slices)
 	return ctrl.Result{}, client.IgnoreNotFound(err)
 }
 
@@ -182,6 +172,68 @@ func shouldFinalize(wl *kueue.Workload) (bool, string) {
 	return false, ""
 }
 
+func (r *WorkloadReconciler) cleanupSlices(ctx context.Context, wl *kueue.Workload) error {
+	log := ctrl.LoggerFrom(ctx)
+
+	slices, err := r.findWorkloadSlices(ctx, wl)
+	if err != nil {
+		log.Error(err, "Failed to find Slices")
+		return err
+	}
+
+	for _, slice := range slices {
+		err = r.client.Delete(ctx, &slice)
+		if client.IgnoreNotFound(err) != nil {
+			log.Error(err, "Failed to delete the Slice", "slice", klog.KObj(&slice))
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (r *WorkloadReconciler) findWorkloadSlices(ctx context.Context, wl *kueue.Workload) ([]v1alpha1.Slice, error) {
+	slices := &v1alpha1.SliceList{}
+	opts := []client.ListOption{
+		client.InNamespace(wl.Namespace),
+		client.MatchingFields{OwnerReferenceUID: string(wl.UID)},
+	}
+	if err := r.client.List(ctx, slices, opts...); err != nil {
+		return nil, err
+	}
+	return slices.Items, nil
+}
+
+func (r *WorkloadReconciler) findWorkloadSlicesByName(ctx context.Context, wl *kueue.Workload) (map[string]*v1alpha1.Slice, error) {
+	existingSlices, err := r.findWorkloadSlices(ctx, wl)
+	if err != nil {
+		return nil, err
+	}
+
+	existingSlicesByName := make(map[string]*v1alpha1.Slice, len(existingSlices))
+	for _, slice := range existingSlices {
+		existingSlicesByName[slice.Name] = &slice
+	}
+
+	return existingSlicesByName, nil
+}
+
+func (r *WorkloadReconciler) finalizeWorkload(ctx context.Context, wl *kueue.Workload) error {
+	log := ctrl.LoggerFrom(ctx)
+
+	controllerutil.RemoveFinalizer(wl, SliceControllerName)
+	if err := r.client.Update(ctx, wl); err != nil {
+		if !apierrors.IsNotFound(err) {
+			log.Error(err, "Failed to remove the finalizer")
+		}
+		return err
+	}
+
+	log.V(5).Info("Removed finalizer")
+
+	return nil
+}
+
 func validateRelevantWorkload(wl *kueue.Workload) error {
 	if !hasRelevantPodSet(wl.Spec.PodSets) {
 		return errors.New("does not have a relevant podset")
@@ -228,46 +280,93 @@ func (r *WorkloadReconciler) sliceAC(ctx context.Context, wl *kueue.Workload) (*
 	return workload.FindAdmissionCheck(wl.Status.AdmissionChecks, relevantChecks[0]), nil
 }
 
-func parseTopologyAssignmentIntoNodeSelector(slice *v1alpha1.Slice, wl *kueue.Workload) {
-	nodeSelectors := sets.New[string]()
+func (r *WorkloadReconciler) syncSlices(ctx context.Context, wl *kueue.Workload, ac *kueue.AdmissionCheckState) ([]v1alpha1.Slice, bool, error) {
+	log := ctrl.LoggerFrom(ctx)
+
+	existingSlices, err := r.findWorkloadSlicesByName(ctx, wl)
+	if err != nil {
+		log.Error(err, "Failed to list Slices")
+		return nil, false, err
+	}
+
+	createdSlices := make([]v1alpha1.Slice, 0, len(existingSlices))
+	changed := false
+
 	for _, psa := range wl.Status.Admission.PodSetAssignments {
-		// we already validated that all assignments have a valid level,
-		// in validateRelevantWorkload.
-		subblockLevelIndex := topology.SubblockLevelIndex(&psa)
-		for _, domain := range psa.TopologyAssignment.Domains {
-			nodeSelectors.Insert(domain.Values[subblockLevelIndex])
+		sliceKey := core.SliceKeyFromWorkload(wl, psa.Name)
+
+		if createdSlice, ok := existingSlices[sliceKey.Name]; ok {
+			createdSlices = append(createdSlices, *createdSlice)
+			continue
+		}
+
+		createdSlice, err := r.createSlice(ctx, wl, ac, psa.Name)
+		if err != nil {
+			return nil, false, err
 		}
+
+		createdSlices = append(createdSlices, *createdSlice)
+		changed = true
+	}
+
+	if changed {
+		msg := fmt.Sprintf("The Slices %v have been created", joinSliceNames(createdSlices))
+		log.V(3).Info(msg)
+		r.record.Event(wl, corev1.EventTypeNormal, SliceCreatedEventType, msg)
+		ac.Message = msg
+		return createdSlices, true, r.updateWorkloadAdmissionCheckStatus(ctx, wl, ac)
+	}
+
+	return createdSlices, false, nil
+}
+
+func parseTopologyAssignmentIntoNodeSelector(slice *v1alpha1.Slice, wl *kueue.Workload, podSetName kueue.PodSetReference) {
+	psa := findPodSetAssignmentByName(wl.Status.Admission.PodSetAssignments, podSetName)
+	if psa == nil {
+		return
+	}
+	nodeSelectors := sets.New[string]()
+	// we already validated that all assignments have a valid level,
+	// in validateRelevantWorkload.
+	subblockLevelIndex := topology.SubblockLevelIndex(psa)
+	for _, domain := range psa.TopologyAssignment.Domains {
+		nodeSelectors.Insert(domain.Values[subblockLevelIndex])
 	}
 	slice.Spec.NodeSelector = map[string][]string{
 		TPUReservationSubblockLabel: sets.List(nodeSelectors),
 	}
 }
 
-func (r *WorkloadReconciler) createSlice(ctx context.Context, log logr.Logger, wl *kueue.Workload, ac *kueue.AdmissionCheckState) error {
-	slice := core.SliceWithMetadata(wl)
-	log = log.WithValues("slice", klog.KObj(slice))
+func findPodSetAssignmentByName(podSetAssignments []kueue.PodSetAssignment, name kueue.PodSetReference) *kueue.PodSetAssignment {
+	for _, psa := range podSetAssignments {
+		if psa.Name == name {
+			return &psa
+		}
+	}
+	return nil
+}
+
+func (r *WorkloadReconciler) createSlice(ctx context.Context, wl *kueue.Workload, ac *kueue.AdmissionCheckState, podSetName kueue.PodSetReference) (*v1alpha1.Slice, error) {
+	slice := core.SliceWithMetadata(wl, podSetName)
+	log := ctrl.LoggerFrom(ctx).WithValues("slice", klog.KObj(slice))
 	log.V(3).Info("Creating Slice")
 
 	if err := controllerutil.SetControllerReference(wl, slice, r.client.Scheme()); err != nil {
-		return err
+		return nil, err
 	}
-	parseTopologyAssignmentIntoNodeSelector(slice, wl)
+	parseTopologyAssignmentIntoNodeSelector(slice, wl, podSetName)
 
 	if err := r.client.Create(ctx, slice); err != nil {
-		msg := fmt.Sprintf("Error creating Slice %q: %v", slice.Name, err)
+		msg := fmt.Sprintf("Error creating Slice %q: %v", client.ObjectKeyFromObject(slice), err)
 		log.Error(err, msg)
 		r.record.Event(wl, corev1.EventTypeWarning, FailedCreateSliceEventType, api.TruncateEventMessage(msg))
+		ac.State = kueue.CheckStatePending
 		ac.Message = api.TruncateConditionMessage(msg)
 		patchErr := r.updateWorkloadAdmissionCheckStatus(ctx, wl, ac)
-		return errors.Join(err, patchErr)
+		return nil, errors.Join(err, patchErr)
 	}
 
-	msg := fmt.Sprintf("The Slice %s has been created", client.ObjectKeyFromObject(slice))
-	log.V(3).Info(msg)
-	r.record.Event(wl, corev1.EventTypeNormal, SliceCreatedEventType, msg)
-	ac.Message = msg
-
-	return r.updateWorkloadAdmissionCheckStatus(ctx, wl, ac)
+	return slice, nil
 }
 
 func (r *WorkloadReconciler) updateWorkloadAdmissionCheckStatus(ctx context.Context, wl *kueue.Workload, ac *kueue.AdmissionCheckState) error {
@@ -280,27 +379,49 @@ func (r *WorkloadReconciler) updateWorkloadAdmissionCheckStatus(ctx context.Cont
 	return err
 }
 
-// syncAdmissionCheckStatus syncs the admission check status with the state of the slice.
-func (r *WorkloadReconciler) syncAdmissionCheckStatus(ctx context.Context, wl *kueue.Workload, ac *kueue.AdmissionCheckState, slice *v1alpha1.Slice) error {
+func joinSliceNames(slices []v1alpha1.Slice) string {
+	sliceNames := make([]string, len(slices))
+	for index, slice := range slices {
+		sliceNames[index] = fmt.Sprintf("%q", client.ObjectKeyFromObject(&slice))
+	}
+	return strings.Join(sliceNames, ", ")
+}
+
+// syncAdmissionCheckStatus syncs the admission check status with the state of the Slices.
+func (r *WorkloadReconciler) syncAdmissionCheckStatus(ctx context.Context, wl *kueue.Workload, ac *kueue.AdmissionCheckState, slices []v1alpha1.Slice) error {
 	originalState := ac.State
+	originalMessage := ac.Message
 
-	errCond := meta.FindStatusCondition(slice.Status.Conditions, string(v1alpha1.Error))
+	slicesByStatus := make(map[v1alpha1.SliceConditionType][]v1alpha1.Slice)
+	for _, slice := range slices {
+		for _, status := range core.SliceStatuses {
+			if meta.IsStatusConditionTrue(slice.Status.Conditions, string(status)) {
+				slicesByStatus[status] = append(slicesByStatus[status], slice)
+			}
+		}
+	}
 
 	switch {
-	case meta.IsStatusConditionTrue(slice.Status.Conditions, string(v1alpha1.Forming)):
-		ac.Message = fmt.Sprintf("The Slice %q is being formed", slice.Name)
-	case meta.IsStatusConditionTrue(slice.Status.Conditions, string(v1alpha1.Ready)):
-		ac.State = kueue.CheckStateReady
-		ac.Message = fmt.Sprintf("The Slice %q is fully operational", slice.Name)
-	case meta.IsStatusConditionTrue(slice.Status.Conditions, string(v1alpha1.Degraded)):
-		ac.State = kueue.CheckStateReady
-		ac.Message = fmt.Sprintf("The Slice %q is running with reduced capacity or performance", slice.Name)
-	case meta.IsStatusConditionTrue(slice.Status.Conditions, string(v1alpha1.Deformed)):
+	case len(slicesByStatus[v1alpha1.Error]) > 0:
 		ac.State = kueue.CheckStateRejected
-		ac.Message = fmt.Sprintf("The Slice %q is being torn down", slice.Name)
-	case errCond != nil && errCond.Status == metav1.ConditionTrue:
+		ac.Message = fmt.Sprintf("The Slices %s are not operational due to an errors", joinSliceNames(slicesByStatus[v1alpha1.Error]))
+	case len(slicesByStatus[v1alpha1.Deformed]) > 0:
 		ac.State = kueue.CheckStateRejected
-		ac.Message = fmt.Sprintf("The Slice %q is not operational due to an error: %s", slice.Name, errCond.Message)
+		ac.Message = fmt.Sprintf("The Slices %s are being torn down", joinSliceNames(slicesByStatus[v1alpha1.Deformed]))
+	case len(slicesByStatus[v1alpha1.Forming]) > 0:
+		ac.State = kueue.CheckStatePending
+		ac.Message = fmt.Sprintf("The Slices %s are being formed", joinSliceNames(slicesByStatus[v1alpha1.Forming]))
+	case len(slicesByStatus[v1alpha1.Degraded]) > 0:
+		ac.State = kueue.CheckStateReady
+		ac.Message = fmt.Sprintf("The Slices %s are running with reduced capacity or performance", joinSliceNames(slicesByStatus[v1alpha1.Degraded]))
+	case len(slicesByStatus[v1alpha1.Ready]) > 0:
+		ac.State = kueue.CheckStateReady
+		ac.Message = fmt.Sprintf("The Slices %s are fully operational", joinSliceNames(slicesByStatus[v1alpha1.Ready]))
+	}
+
+	// No changes.
+	if originalState == ac.State && ac.Message == originalMessage {
+		return nil
 	}
 
 	err := r.updateWorkloadAdmissionCheckStatus(ctx, wl, ac)
@@ -330,8 +451,7 @@ type sliceHandler struct {
 func (h *sliceHandler) Generic(context.Context, event.GenericEvent, workqueue.TypedRateLimitingInterface[reconcile.Request]) {
 }
 
-func (h *sliceHandler) Create(ctx context.Context, e event.CreateEvent, q workqueue.TypedRateLimitingInterface[reconcile.Request]) {
-	h.handleEvent(ctx, e.Object, q)
+func (h *sliceHandler) Create(context.Context, event.CreateEvent, workqueue.TypedRateLimitingInterface[reconcile.Request]) {
 }
 
 func (h *sliceHandler) Delete(ctx context.Context, e event.DeleteEvent, q workqueue.TypedRateLimitingInterface[reconcile.Request]) {