fix(controller): delete all pods when removing workload and improve … (#159)

* feat(controller): delete all pods when removing workload and improve finalizer logic

* fix lint

* fix: workload may not exist, should check it in cleanup.

Author: 0x5457

internal/controller/tensorfusionworkload_controller.go

@@ -41,6 +41,7 @@ import (
 	"github.com/NexusGPU/tensor-fusion/internal/utils"
 	"github.com/NexusGPU/tensor-fusion/internal/worker"
 	"github.com/prometheus/client_golang/prometheus"
+	"github.com/samber/lo"
 )
 
 // TensorFusionWorkloadReconciler reconciles a TensorFusionWorkload object
@@ -80,6 +81,17 @@ func (r *TensorFusionWorkloadReconciler) Reconcile(ctx context.Context, req ctrl
 	}
 
 	shouldReturn, err := utils.HandleFinalizer(ctx, workload, r.Client, func(ctx context.Context, _ *tfv1.TensorFusionWorkload) (bool, error) {
+		// delete all pods
+		existsPods := lo.Filter(podList.Items, func(pod corev1.Pod, _ int) bool {
+			return pod.DeletionTimestamp == nil
+		})
+		if len(existsPods) > 0 {
+			if err := r.DeleteAllOf(ctx, &corev1.Pod{},
+				client.InNamespace(req.Namespace),
+				client.MatchingLabels{constants.WorkloadKey: workload.Name}); err != nil {
+				return false, fmt.Errorf("delete pods: %w", err)
+			}
+		}
 		// check if all pods are deleted
 		return len(podList.Items) == 0, nil
 	})

internal/controller/tensorfusionworkload_controller_test.go

@@ -24,6 +24,7 @@ import (
 	. "github.com/onsi/gomega"
 	"github.com/samber/lo"
 	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"sigs.k8s.io/controller-runtime/pkg/client"
@@ -184,6 +185,44 @@ var _ = Describe("TensorFusionWorkload Controller", func() {
 		})
 	})
 
+	Context("When deleting workload directly", func() {
+		It("Should delete all pods and the workload itself", func() {
+			pool := tfEnv.GetGPUPool(0)
+
+			workload := createTensorFusionWorkload(pool.Name, key, 2)
+			checkWorkerPodCount(workload)
+			checkWorkloadStatus(workload)
+
+			// wait for 2 pods to be created
+			Eventually(func(g Gomega) {
+				podList := &corev1.PodList{}
+				g.Expect(k8sClient.List(ctx, podList,
+					client.InNamespace(key.Namespace),
+					client.MatchingLabels{constants.WorkloadKey: key.Name})).To(Succeed())
+				g.Expect(podList.Items).To(HaveLen(2))
+			}, timeout, interval).Should(Succeed())
+
+			// delete workload
+			Expect(k8sClient.Delete(ctx, workload)).To(Succeed())
+
+			// wait for all pods to be deleted
+			Eventually(func(g Gomega) {
+				podList := &corev1.PodList{}
+				g.Expect(k8sClient.List(ctx, podList,
+					client.InNamespace(key.Namespace),
+					client.MatchingLabels{constants.WorkloadKey: key.Name})).To(Succeed())
+				g.Expect(podList.Items).Should(BeEmpty())
+			}, timeout, interval).Should(Succeed())
+
+			// wait for workload itself to be deleted
+			Eventually(func(g Gomega) {
+				w := &tfv1.TensorFusionWorkload{}
+				err := k8sClient.Get(ctx, key, w)
+				g.Expect(err).To(HaveOccurred())
+			}, timeout, interval).Should(Succeed())
+		})
+	})
+
 	Context("When GPUPool doesn't exist", func() {
 		It("Should not create worker pod when reconciling a workload with non-existent pool", func() {
 			workload := createTensorFusionWorkload("non-existent-pool", key, 1)
@@ -262,6 +301,13 @@ func cleanupWorkload(key client.ObjectKey) {
 	GinkgoHelper()
 	workload := &tfv1.TensorFusionWorkload{}
 
+	if err := k8sClient.Get(ctx, key, workload); err != nil {
+		if errors.IsNotFound(err) {
+			return
+		}
+		Expect(err).To(HaveOccurred())
+	}
+
 	// Set replicas to 0
 	Eventually(func(g Gomega) {
 		g.Expect(k8sClient.Get(ctx, key, workload)).Should(Succeed())

internal/utils/reconcile.go

@@ -60,11 +60,8 @@ func HandleFinalizer[T client.Object](
 				shouldReturn = true
 				return shouldReturn, err
 			}
-			// Cleanup not ready, wait for next reconcile
-			shouldReturn = true
-			return shouldReturn, err
 		}
-		// Finalizer already removed, continue with deletion
+		// continue with deletion logic
 		shouldReturn = false
 		return shouldReturn, err
 	}