Use RequeueAfter instead of JobSet and Pods event handler.

Author: mbobrovskyi

slice/internal/controller/workload_controller.go

@@ -41,7 +41,6 @@ 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/util/admissioncheck"
 	"sigs.k8s.io/kueue/pkg/workload"
@@ -110,9 +109,9 @@ func (r *WorkloadReconciler) Reconcile(ctx context.Context, req ctrl.Request) (c
 		log.V(3).Info(fmt.Sprintf("Cleaning up the Slice and finalize the Workload because %s", reason))
 		cleanedUp, err := r.cleanupSlice(ctx, wl)
 		if err != nil || !cleanedUp {
-			return ctrl.Result{}, err
+			// Recheck again in 5 seconds.
+			return ctrl.Result{RequeueAfter: 5 * time.Second}, err
 		}
-
 		err = r.finalizeWorkload(ctx, wl)
 		return ctrl.Result{}, client.IgnoreNotFound(err)
 	}
@@ -449,8 +448,6 @@ func (r *WorkloadReconciler) SetupWithManager(mgr ctrl.Manager) error {
 		Named("workload_controller").
 		WithEventFilter(r).
 		Watches(&v1alpha1.Slice{}, &sliceHandler{client: r.client}).
-		Watches(&jobset.JobSet{}, &jobSetHandler{client: r.client}).
-		Watches(&corev1.Pod{}, &podHandler{client: r.client}).
 		Complete(r)
 }
 
@@ -500,6 +497,7 @@ func (h *sliceHandler) Generic(context.Context, event.GenericEvent, workqueue.Ty
 }
 
 func (h *sliceHandler) Create(context.Context, event.CreateEvent, workqueue.TypedRateLimitingInterface[reconcile.Request]) {
+	// No need to handle create event. We should wait for at least Forming state.
 }
 
 func (h *sliceHandler) Delete(ctx context.Context, e event.DeleteEvent, q workqueue.TypedRateLimitingInterface[reconcile.Request]) {
@@ -536,124 +534,3 @@ func (h *sliceHandler) handleEvent(ctx context.Context, obj client.Object, q wor
 
 	q.AddAfter(req, updatesBatchPeriod)
 }
-
-var _ handler.EventHandler = (*jobSetHandler)(nil)
-
-type jobSetHandler struct {
-	client client.Client
-}
-
-func (h *jobSetHandler) Generic(context.Context, event.GenericEvent, workqueue.TypedRateLimitingInterface[reconcile.Request]) {
-}
-
-func (h *jobSetHandler) Create(context.Context, event.CreateEvent, workqueue.TypedRateLimitingInterface[reconcile.Request]) {
-}
-
-func (h *jobSetHandler) Delete(ctx context.Context, e event.DeleteEvent, q workqueue.TypedRateLimitingInterface[reconcile.Request]) {
-	h.handleEvent(ctx, e.Object, q)
-}
-
-func (h *jobSetHandler) Update(context.Context, event.UpdateEvent, workqueue.TypedRateLimitingInterface[reconcile.Request]) {
-}
-
-func (h *jobSetHandler) handleEvent(ctx context.Context, obj client.Object, q workqueue.TypedRateLimitingInterface[reconcile.Request]) {
-	jobSet, isJobSet := obj.(*jobset.JobSet)
-	// Only JobSet should be handled.
-	if !isJobSet {
-		return
-	}
-
-	log := ctrl.LoggerFrom(ctx).WithValues("jobSet", klog.KRef(jobSet.Namespace, jobSet.Name))
-	ctrl.LoggerInto(ctx, log)
-
-	log.V(3).Info("Handle JobSet event")
-
-	handleEventForJobSet(ctx, h.client, jobSet, q)
-}
-
-func handleEventForJobSet(ctx context.Context, c client.Client, jobSet *jobset.JobSet, q workqueue.TypedRateLimitingInterface[reconcile.Request]) {
-	log := ctrl.LoggerFrom(ctx)
-
-	workloads := &kueue.WorkloadList{}
-	opts := []client.ListOption{
-		client.InNamespace(jobSet.Namespace),
-		client.MatchingFields{OwnerReferenceUID: string(jobSet.UID)},
-	}
-	err := c.List(ctx, workloads, opts...)
-	if err != nil {
-		log.Error(err, "Failed to list workloads")
-		return
-	}
-
-	if len(workloads.Items) == 0 {
-		log.V(3).Info("No Workloads found for the JobSet – skipping handling")
-		return
-	}
-
-	for _, wl := range workloads.Items {
-		if shouldHandleWorkload(&wl) {
-			req := reconcile.Request{
-				NamespacedName: types.NamespacedName{
-					Name:      wl.Name,
-					Namespace: wl.Namespace,
-				},
-			}
-			q.AddAfter(req, updatesBatchPeriod)
-		}
-	}
-}
-
-var _ handler.EventHandler = (*podHandler)(nil)
-
-type podHandler struct {
-	client client.Client
-}
-
-func (h *podHandler) Generic(context.Context, event.GenericEvent, workqueue.TypedRateLimitingInterface[reconcile.Request]) {
-}
-
-func (h *podHandler) Create(context.Context, event.CreateEvent, workqueue.TypedRateLimitingInterface[reconcile.Request]) {
-}
-
-func (h *podHandler) Delete(ctx context.Context, e event.DeleteEvent, q workqueue.TypedRateLimitingInterface[reconcile.Request]) {
-	h.handleEvent(ctx, e.Object, q)
-}
-
-func (h *podHandler) Update(context.Context, event.UpdateEvent, workqueue.TypedRateLimitingInterface[reconcile.Request]) {
-}
-
-func (h *podHandler) handleEvent(ctx context.Context, obj client.Object, q workqueue.TypedRateLimitingInterface[reconcile.Request]) {
-	pod, isPod := obj.(*corev1.Pod)
-	// Only Pods and Pods with the TAS label should be handled.
-	if !isPod || pod.Labels[kueuealpha.TASLabel] != tasLabelValue {
-		return
-	}
-
-	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.KRef(pod.Namespace, jobSetName),
-	)
-	ctrl.LoggerInto(ctx, log)
-
-	log.V(3).Info("Handle Pod event")
-
-	jobSet := &jobset.JobSet{}
-	jobSetKey := types.NamespacedName{Name: jobSetName, Namespace: pod.Namespace}
-	err := h.client.Get(ctx, jobSetKey, jobSet)
-	if err != nil {
-		if apierrors.IsNotFound(err) {
-			log.V(3).Info("JobSet not found")
-		} else {
-			log.Error(err, "Failed to get JobSet")
-		}
-		return
-	}
-
-	handleEventForJobSet(ctx, h.client, jobSet, q)
-}