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.SetupIndexer(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/hack/kind-cluster.yaml

@@ -24,10 +24,15 @@ nodes:
   - role: worker
     labels:
       instance-type: on-demand
+      cloud.google.com/gke-tpu-topology: 2x2x2
+      cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
+      cloud.google.com/gke-tpu-reservation-subblock: tpu-subblock-1
   - role: worker
     labels:
-      instance-type: spot
-
+      instance-type: on-demand
+      cloud.google.com/gke-tpu-topology: 2x2x2
+      cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
+      cloud.google.com/gke-nodepool: tpu-v4-pool
 kubeadmConfigPatches:
   - |
     kind: JoinConfiguration

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) })
+}
+
+// SetupIndexer sets the index with the given fields for core apis.
+func SetupIndexer(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

@@ -18,9 +18,15 @@ package controller
 
 import (
 	"context"
+	"errors"
+	"fmt"
+	"maps"
 	"slices"
 
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/util/sets"
+	"k8s.io/klog/v2"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
@@ -33,7 +39,18 @@ import (
 
 const (
 	CleanupSliceFinalizerName   = "accelerator.gke.io/slice"
-	TPUReservationSubblockLabel = "cloud.google.com/gke-tpu-reservation-subblock"
+	TPUReservationSubBlockLabel = "cloud.google.com/gke-tpu-reservation-subblock"
+	NodePoolLabel               = "cloud.google.com/gke-nodepool"
+	TPUTopologyLabel            = "cloud.google.com/gke-tpu-topology"
+	TPUAcceleratorLabel         = "cloud.google.com/gke-tpu-accelerator"
+)
+
+var (
+	errPodSetNotFound              = errors.New("PodSet not found")
+	errPodSetAssignmentNotFound    = errors.New("PodSetAssignment not found")
+	errTPUTopologyLabelNotFound    = fmt.Errorf("%s label not found", TPUTopologyLabel)
+	errTPUAcceleratorLabelNotFound = fmt.Errorf("%s label not found", TPUAcceleratorLabel)
+	errTopologyAssignmentNotFound  = errors.New("TopologyAssignment not found")
 )
 
 // WorkloadReconciler reconciles a Workload object
@@ -64,18 +81,7 @@ func (r *WorkloadReconciler) Reconcile(ctx context.Context, req ctrl.Request) (c
 	log.V(2).Info("Reconcile Workload")
 
 	if r.shouldFinalize(wl) {
-		if controllerutil.ContainsFinalizer(wl, CleanupSliceFinalizerName) {
-			err = r.client.Delete(ctx, r.newEmptySlice(wl))
-			if client.IgnoreNotFound(err) != nil {
-				return ctrl.Result{}, err
-			}
-			controllerutil.RemoveFinalizer(wl, CleanupSliceFinalizerName)
-			if err := r.client.Update(ctx, wl); err != nil {
-				return ctrl.Result{}, err
-			}
-			log.V(5).Info("Removed finalizer")
-		}
-		return ctrl.Result{}, nil
+		return ctrl.Result{}, client.IgnoreNotFound(r.finalize(ctx, wl))
 	}
 
 	if controllerutil.AddFinalizer(wl, CleanupSliceFinalizerName) {
@@ -85,75 +91,212 @@ func (r *WorkloadReconciler) Reconcile(ctx context.Context, req ctrl.Request) (c
 		}
 	}
 
-	return ctrl.Result{}, r.createSliceIfNotExist(ctx, wl)
+	return ctrl.Result{}, r.createSlicesIfNotExist(ctx, wl)
 }
 
 func (r *WorkloadReconciler) shouldFinalize(wl *kueue.Workload) bool {
 	return !wl.DeletionTimestamp.IsZero() || workload.IsFinished(wl) || workload.IsEvicted(wl) || !workload.IsActive(wl)
 }
 
-func (r *WorkloadReconciler) newEmptySlice(wl *kueue.Workload) *v1alpha1.Slice {
-	return &v1alpha1.Slice{
-		ObjectMeta: metav1.ObjectMeta{
-			Name:      wl.Name,
-			Namespace: wl.Namespace,
-		},
+func (r *WorkloadReconciler) finalize(ctx context.Context, wl *kueue.Workload) error {
+	if !controllerutil.ContainsFinalizer(wl, CleanupSliceFinalizerName) {
+		return nil
+	}
+
+	log := ctrl.LoggerFrom(ctx)
+
+	slices, err := r.findWorkloadSlices(ctx, wl)
+	if err != nil {
+		log.Error(err, "Failed to find Slices")
+		return err
 	}
-}
 
-func (r *WorkloadReconciler) newSlice(wl *kueue.Workload) (*v1alpha1.Slice, error) {
-	slice := r.newEmptySlice(wl)
+	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
+		}
+	}
+
+	controllerutil.RemoveFinalizer(wl, CleanupSliceFinalizerName)
+	if err := r.client.Update(ctx, wl); err != nil {
+		if !apierrors.IsNotFound(err) {
+			log.Error(err, "Failed to remove finalizer")
+		}
+		return err
+	}
 
-	if err := controllerutil.SetControllerReference(wl, slice, r.client.Scheme()); err != nil {
+	log.V(5).Info("Removed finalizer")
+
+	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) createSlicesIfNotExist(ctx context.Context, wl *kueue.Workload) error {
+	log := ctrl.LoggerFrom(ctx)
+
+	createdSlices, err := r.findWorkloadSlices(ctx, wl)
+	if err != nil {
+		log.Error(err, "Failed to find Slices")
+		return err
+	}
 
-	if wl.Status.Admission != nil && wl.Status.Admission.PodSetAssignments != nil {
+	createdSlicesByName := make(map[string]*v1alpha1.Slice, len(createdSlices))
+	for _, slice := range createdSlices {
+		createdSlicesByName[slice.Name] = &slice
+	}
+
+	var slicesToCreate []*v1alpha1.Slice
+
+	if wl.Status.Admission != nil {
 		for _, psa := range wl.Status.Admission.PodSetAssignments {
-			for _, domain := range psa.TopologyAssignment.Domains {
-				if slice.Spec.NodeSelector == nil {
-					slice.Spec.NodeSelector = make(map[string][]string)
-				}
-				if slice.Spec.NodeSelector[TPUReservationSubblockLabel] == nil {
-					slice.Spec.NodeSelector[TPUReservationSubblockLabel] = []string{}
-				}
-				// make sure there are no duplicates in the nodeSelector
-				for _, v := range domain.Values {
-					exists := slices.Contains(slice.Spec.NodeSelector[TPUReservationSubblockLabel], v)
-					if !exists {
-						slice.Spec.NodeSelector[TPUReservationSubblockLabel] = append(slice.Spec.NodeSelector[TPUReservationSubblockLabel], v)
-					}
+			sliceName := GetSliceName(wl.Name, psa.Name)
+
+			if _, ok := createdSlicesByName[sliceName]; ok {
+				delete(createdSlicesByName, sliceName)
+				continue
+			}
+
+			slice, err := newSlice(wl, psa.Name)
+			if err != nil {
+				if !isUnsupportedPodSetError(err) {
+					log.Error(err, "Failed to create a Slice")
+					return err
 				}
+				log.V(8).Info("Failed to create Slice", "error", err)
+				continue
+			}
+
+			if err := controllerutil.SetControllerReference(wl, slice, r.client.Scheme()); err != nil {
+				return err
 			}
+
+			slicesToCreate = append(slicesToCreate, slice)
 		}
 	}
 
-	return slice, nil
+	for _, slice := range slicesToCreate {
+		err = r.client.Create(ctx, slice)
+		if err != nil {
+			log.Error(err, "Failed to create a Slice", "slice", klog.KObj(slice))
+			return err
+		}
+	}
+
+	for _, slice := range slices.Collect(maps.Values(createdSlicesByName)) {
+		err = r.client.Delete(ctx, slice)
+		if client.IgnoreNotFound(err) != nil {
+			log.Error(err, "Failed to delete the redundant Slice", "slice", klog.KObj(slice))
+			return err
+		}
+	}
+
+	return nil
 }
 
-func (r *WorkloadReconciler) createSliceIfNotExist(ctx context.Context, wl *kueue.Workload) error {
-	slice := r.newEmptySlice(wl)
+func GetSliceName(workloadName string, podSetName kueue.PodSetReference) string {
+	return fmt.Sprintf("%s-%s", workloadName, podSetName)
+}
 
-	err := r.client.Get(ctx, client.ObjectKeyFromObject(slice), slice)
-	if client.IgnoreNotFound(err) != nil {
-		return err
+func newSlice(wl *kueue.Workload, podSetName kueue.PodSetReference) (*v1alpha1.Slice, error) {
+	ps := findPodSet(wl, podSetName)
+	if findPodSet(wl, podSetName) == nil {
+		return nil, errPodSetNotFound
 	}
-	if err == nil {
-		return nil
+
+	if ps.Template.Spec.NodeSelector[TPUTopologyLabel] == "" {
+		return nil, errTPUTopologyLabelNotFound
 	}
 
-	slice, err = r.newSlice(wl)
-	if err != nil {
-		return err
+	if ps.Template.Spec.NodeSelector[TPUAcceleratorLabel] == "" {
+		return nil, errTPUAcceleratorLabelNotFound
 	}
 
-	return r.client.Create(ctx, slice)
+	psa := findPodSetAssignment(wl, podSetName)
+	if psa == nil {
+		return nil, errPodSetAssignmentNotFound
+	}
+
+	if psa.TopologyAssignment == nil {
+		return nil, errTopologyAssignmentNotFound
+	}
+
+	slice := &v1alpha1.Slice{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      GetSliceName(wl.Name, podSetName),
+			Namespace: wl.Namespace,
+		},
+		Spec: v1alpha1.SliceSpec{
+			AcceleratorTopology: ps.Template.Spec.NodeSelector[TPUTopologyLabel],
+			AcceleratorType:     ps.Template.Spec.NodeSelector[TPUAcceleratorLabel],
+			NodeSelector:        make(map[string][]string),
+		},
+	}
+
+	for _, domain := range psa.TopologyAssignment.Domains {
+		if ps.Template.Spec.NodeSelector[TPUReservationSubBlockLabel] != "" {
+			subBlockDomains := sets.New[string]()
+			for _, v := range domain.Values {
+				subBlockDomains.Insert(v)
+			}
+			if subBlockDomains.Len() > 0 {
+				slice.Spec.NodeSelector[TPUReservationSubBlockLabel] = sets.List(subBlockDomains)
+			}
+		}
+		if ps.Template.Spec.NodeSelector[NodePoolLabel] != "" {
+			nodePoolDomains := sets.New[string]()
+			for _, v := range domain.Values {
+				nodePoolDomains.Insert(v)
+			}
+			if nodePoolDomains.Len() > 0 {
+				slice.Spec.NodeSelector[NodePoolLabel] = sets.List(nodePoolDomains)
+			}
+		}
+	}
+
+	return slice, nil
+}
+
+func findPodSet(wl *kueue.Workload, podSetName kueue.PodSetReference) *kueue.PodSet {
+	for _, ps := range wl.Spec.PodSets {
+		if ps.Name == podSetName {
+			return &ps
+		}
+	}
+	return nil
+}
+
+func findPodSetAssignment(wl *kueue.Workload, podSetName kueue.PodSetReference) *kueue.PodSetAssignment {
+	for _, psa := range wl.Status.Admission.PodSetAssignments {
+		if psa.Name == podSetName {
+			return &psa
+		}
+	}
+	return nil
+}
+
+func isUnsupportedPodSetError(err error) bool {
+	return errors.Is(err, errTPUTopologyLabelNotFound) ||
+		errors.Is(err, errTPUAcceleratorLabelNotFound) ||
+		errors.Is(err, errTopologyAssignmentNotFound)
 }
 
 // SetupWithManager sets up the controller with the Manager.
 func (r *WorkloadReconciler) SetupWithManager(mgr ctrl.Manager) error {
 	return ctrl.NewControllerManagedBy(mgr).
 		For(&kueue.Workload{}).
-		Named("workload").
+		Named("workload_controller").
 		Complete(r)
 }