feat: improve TensorFusion workload finalizer handling (#107)

Author: 0x5457

internal/controller/tensorfusionworkload_controller.go

@@ -73,19 +73,29 @@ func (r *TensorFusionWorkloadReconciler) Reconcile(ctx context.Context, req ctrl
 		return ctrl.Result{}, err
 	}
 
-	// First, handle pods with finalizers that need GPU resource cleanup
 	podList := &corev1.PodList{}
 	if err := r.List(ctx, podList,
 		client.InNamespace(req.Namespace),
 		client.MatchingLabels{constants.WorkloadKey: workload.Name}); err != nil {
 		return ctrl.Result{}, fmt.Errorf("list pods: %w", err)
 	}
 
-	hasdeletion := false
+	deleted, err := utils.HandleFinalizer(ctx, workload, r.Client, func(ctx context.Context, _ *tfv1.TensorFusionWorkload) (bool, error) {
+		// check if all pods are deleted
+		return len(podList.Items) == 0, nil
+	})
+	if err != nil {
+		return ctrl.Result{}, fmt.Errorf("handle finalizer: %w", err)
+	}
+	if deleted {
+		return ctrl.Result{}, nil
+	}
+
+	// Handle pods with finalizers that need GPU resource cleanup
+	hasDeletion := false
 	// Process pods with our finalizer
 	for i := range podList.Items {
 		pod := &podList.Items[i]
-
 		// Handle our GPU resource cleanup finalizer
 		deleted, err := utils.HandleFinalizer(ctx, pod, r.Client, func(ctx context.Context, obj *corev1.Pod) (bool, error) {
 			return r.handlePodGPUCleanup(ctx, pod, workload)
@@ -94,10 +104,10 @@ func (r *TensorFusionWorkloadReconciler) Reconcile(ctx context.Context, req ctrl
 		if err != nil {
 			return ctrl.Result{}, err
 		}
-		hasdeletion = hasdeletion || deleted
+		hasDeletion = hasDeletion || deleted
 	}
 
-	if hasdeletion {
+	if hasDeletion {
 		return ctrl.Result{Requeue: true, RequeueAfter: constants.PendingRequeueDuration}, nil
 	}
 

internal/controller/tensorfusionworkload_controller_test.go

@@ -112,9 +112,16 @@ var _ = Describe("TensorFusionWorkload Controller", func() {
 
 	AfterEach(func() {
 		// Clean up workload resources
+
 		resource := &tensorfusionaiv1.TensorFusionWorkload{}
 		err := k8sClient.Get(ctx, typeNamespacedName, resource)
 		if err == nil {
+			By("remove finalizers from workload")
+			if len(resource.Finalizers) > 0 {
+				resource.Finalizers = []string{}
+				Expect(k8sClient.Update(ctx, resource)).To(Succeed())
+			}
+
 			By("Cleaning up the test workload")
 			Expect(k8sClient.Delete(ctx, resource)).To(Succeed())
 		}