Don't add finalizer if Workload doesn't have TopologyAssignment.

Author: mbobrovskyi

slice/config/rbac/role.yaml

@@ -30,14 +30,6 @@ rules:
   - list
   - update
   - watch
-- apiGroups:
-  - jobset.x-k8s.io
-  resources:
-  - jobsets
-  verbs:
-  - get
-  - list
-  - watch
 - apiGroups:
   - kueue.x-k8s.io
   resources:

slice/internal/controller/indexer.go

@@ -28,17 +28,17 @@ import (
 )
 
 const (
-	OwnerReferenceUID = "metadata.ownerReferences.uid"
+	OwnerReferenceName = "metadata.ownerReferences.name"
 )
 
-func IndexOwnerReferenceUID(obj client.Object) []string {
-	return slices.Map(obj.GetOwnerReferences(), func(o *metav1.OwnerReference) string { return string(o.UID) })
+func IndexOwnerReferenceName(obj client.Object) []string {
+	return slices.Map(obj.GetOwnerReferences(), func(o *metav1.OwnerReference) string { return o.Name })
 }
 
 // Setup sets the index with the given fields for core apis.
 func Setup(ctx context.Context, indexer client.FieldIndexer) error {
-	if err := indexer.IndexField(ctx, &kueue.Workload{}, OwnerReferenceUID, IndexOwnerReferenceUID); err != nil {
-		return fmt.Errorf("setting index on ownerReferences.uid for Workload: %w", err)
+	if err := indexer.IndexField(ctx, &kueue.Workload{}, OwnerReferenceName, IndexOwnerReferenceName); err != nil {
+		return fmt.Errorf("setting index on ownerReferences.name for Workload: %w", err)
 	}
 	return nil
 }

slice/internal/controller/workload_controller.go

@@ -35,8 +35,8 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/predicate"
 	"sigs.k8s.io/controller-runtime/pkg/reconcile"
 	jobset "sigs.k8s.io/jobset/api/jobset/v1alpha2"
+	kueuealpha "sigs.k8s.io/kueue/apis/kueue/v1alpha1"
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
-	"sigs.k8s.io/kueue/pkg/controller/core/indexer"
 	"sigs.k8s.io/kueue/pkg/workload"
 
 	"tpu-slice-controller/api/v1alpha1"
@@ -45,6 +45,8 @@ import (
 const (
 	CleanupSliceFinalizerName   = "accelerator.gke.io/slice"
 	TPUReservationSubblockLabel = "cloud.google.com/gke-tpu-reservation-subblock"
+
+	tasLabelValue = "true"
 )
 
 var (
@@ -69,7 +71,6 @@ func NewWorkloadReconciler(cl client.Client) *WorkloadReconciler {
 // +kubebuilder:rbac:groups=kueue.x-k8s.io,resources=workloads,verbs=get;list;watch;create;update;patch
 // +kubebuilder:rbac:groups=slice.accelerator.gke.io,resources=slices,verbs=get;list;watch;create;update;patch;delete
 // +kubebuilder:rbac:groups=slice.accelerator.gke.io,resources=slices/finalizers,verbs=update
-// +kubebuilder:rbac:groups=jobset.x-k8s.io,resources=jobsets,verbs=get;list;watch
 // +kubebuilder:rbac:groups="",resources=pods,verbs=get;list;watch
 
 func (r *WorkloadReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
@@ -137,6 +138,7 @@ func (r *WorkloadReconciler) finalize(ctx context.Context, wl *kueue.Workload) e
 
 	controllerutil.RemoveFinalizer(wl, CleanupSliceFinalizerName)
 	if err := r.client.Update(ctx, wl); err != nil {
+		log.Error(err, "Removing finalizer")
 		return err
 	}
 
@@ -151,25 +153,18 @@ func (r *WorkloadReconciler) podsGracefullyTerminated(ctx context.Context, wl *k
 	}
 
 	owner := metav1.GetControllerOf(wl)
-
-	jobSet := &jobset.JobSet{}
-	err := r.client.Get(ctx, types.NamespacedName{Name: owner.Name, Namespace: wl.Namespace}, jobSet)
-	if client.IgnoreNotFound(err) != nil {
-		return false, err
-	}
-	// That means the JobSet has already been deleted, along with all associated Jobs and Pods
-	// we should delete Slice and finalize Workload.
-	if err != nil {
-		//nolint:nilerr // Return nil to delete Slice and finalize Workload.
+	if owner == nil {
+		// That means the JobSet has already been deleted, along with all associated Jobs and Pods
+		// we should delete Slice and finalize Workload.
 		return true, nil
 	}
 
 	pods := &corev1.PodList{}
 	opts := []client.ListOption{
-		client.InNamespace(jobSet.Namespace),
-		client.MatchingLabels{jobset.JobSetNameKey: jobSet.Name},
+		client.InNamespace(wl.Namespace),
+		client.MatchingLabels{jobset.JobSetNameKey: owner.Name},
 	}
-	err = r.client.List(ctx, pods, opts...)
+	err := r.client.List(ctx, pods, opts...)
 	if err != nil {
 		return false, err
 	}
@@ -241,83 +236,100 @@ func (r *WorkloadReconciler) SetupWithManager(mgr ctrl.Manager) error {
 		For(&kueue.Workload{}).
 		Named("workload").
 		WithEventFilter(r).
-		Watches(&jobset.JobSet{}, &jobSetHandler{r: r}).
+		Watches(&corev1.Pod{}, &podHandler{r: r}).
 		Complete(r)
 }
 
 var _ predicate.Predicate = (*WorkloadReconciler)(nil)
 
 func (r *WorkloadReconciler) Create(e event.CreateEvent) bool {
-	return r.handle(e.Object)
+	return r.handleEvent(e.Object)
 }
 
-func (r *WorkloadReconciler) Delete(event.DeleteEvent) bool {
-	// Nothing handle for Delete event.
-	return false
+func (r *WorkloadReconciler) Delete(e event.DeleteEvent) bool {
+	return r.handleEvent(e.Object)
 }
 
 func (r *WorkloadReconciler) Update(e event.UpdateEvent) bool {
-	return r.handle(e.ObjectNew)
+	return r.handleEvent(e.ObjectNew)
 }
 
 func (r *WorkloadReconciler) Generic(event.GenericEvent) bool {
 	// Nothing handle for Generic event.
 	return false
 }
 
-func (r *WorkloadReconciler) handle(obj client.Object) bool {
+func (r *WorkloadReconciler) handleEvent(obj client.Object) bool {
 	wl, isWorkload := obj.(*kueue.Workload)
 	if !isWorkload {
 		return true
 	}
-	return r.handleWorkload(wl)
+	return r.shouldHandleWorkload(wl)
 }
 
-func (r *WorkloadReconciler) handleWorkload(wl *kueue.Workload) bool {
+func (r *WorkloadReconciler) shouldHandleWorkload(wl *kueue.Workload) bool {
 	// We should handle all Workloads that have the cleanup slice finalizer.
 	return controllerutil.ContainsFinalizer(wl, CleanupSliceFinalizerName) || !r.shouldFinalize(wl)
 }
 
-var _ handler.EventHandler = (*jobSetHandler)(nil)
+var _ handler.EventHandler = (*podHandler)(nil)
 
-type jobSetHandler struct {
+type podHandler struct {
 	r *WorkloadReconciler
 }
 
-func (h *jobSetHandler) Generic(context.Context, event.GenericEvent, workqueue.TypedRateLimitingInterface[reconcile.Request]) {
+func (h *podHandler) Generic(context.Context, event.GenericEvent, workqueue.TypedRateLimitingInterface[reconcile.Request]) {
 }
 
-func (h *jobSetHandler) Create(_ context.Context, e event.CreateEvent, q workqueue.TypedRateLimitingInterface[reconcile.Request]) {
-	h.handle(e.Object, q)
+func (h *podHandler) Create(context.Context, event.CreateEvent, workqueue.TypedRateLimitingInterface[reconcile.Request]) {
 }
 
-func (h *jobSetHandler) Update(_ context.Context, e event.UpdateEvent, q workqueue.TypedRateLimitingInterface[reconcile.Request]) {
-	h.handle(e.ObjectNew, q)
+func (h *podHandler) Delete(ctx context.Context, e event.DeleteEvent, q workqueue.TypedRateLimitingInterface[reconcile.Request]) {
+	h.handleEvent(ctx, e.Object, q)
 }
 
-func (h *jobSetHandler) Delete(context.Context, event.DeleteEvent, workqueue.TypedRateLimitingInterface[reconcile.Request]) {
+func (h *podHandler) Update(context.Context, event.UpdateEvent, workqueue.TypedRateLimitingInterface[reconcile.Request]) {
 }
 
-func (h *jobSetHandler) handle(obj client.Object, q workqueue.TypedRateLimitingInterface[reconcile.Request]) {
-	jobSet, isJobSet := obj.(*jobset.JobSet)
-	if !isJobSet {
+func (h *podHandler) handleEvent(ctx context.Context, obj client.Object, q workqueue.TypedRateLimitingInterface[reconcile.Request]) {
+	pod, isPod := obj.(*corev1.Pod)
+	// Only Pods with the TAS label should be handled.
+	if !isPod || pod.Labels[kueuealpha.TASLabel] != tasLabelValue {
 		return
 	}
 
-	ctx := context.Background()
+	jobSetName := pod.Labels[jobset.JobSetNameKey]
+	// Only pods owned by a JobSet should be handled.
+	if jobSetName == "" {
+		return
+	}
+
+	log := ctrl.LoggerFrom(ctx).WithValues(
+		"pod", klog.KObj(pod),
+		"jobSet", klog.ObjectRef{Name: jobSetName, Namespace: pod.Namespace},
+	)
+	ctrl.LoggerInto(ctx, log)
+
+	log.V(5).Info("Handle Pod event")
+
 	workloads := &kueue.WorkloadList{}
 	opts := []client.ListOption{
-		client.InNamespace(jobSet.Namespace),
-		client.MatchingFields{indexer.OwnerReferenceUID: string(jobSet.GetUID())},
+		client.InNamespace(pod.Namespace),
+		client.MatchingFields{OwnerReferenceName: jobSetName},
 	}
 	err := h.r.client.List(ctx, workloads, opts...)
 	if err != nil {
-		h.r.log.Error(err, "failed to list workloads", "jobset", klog.KObj(jobSet))
+		log.Error(err, "Failed to list workloads")
+		return
+	}
+
+	if len(workloads.Items) == 0 {
+		log.V(5).Info("No Workloads found for the JobSet – skipping handling")
 		return
 	}
 
 	for _, wl := range workloads.Items {
-		if h.r.handleWorkload(&wl) {
+		if h.r.shouldHandleWorkload(&wl) {
 			req := reconcile.Request{
 				NamespacedName: types.NamespacedName{
 					Name:      wl.Name,