fix: misuse CreateOrUpdate (#85)

Author: knave

internal/controller/gpunode_controller.go

@@ -368,19 +368,59 @@ func (r *GPUNodeReconciler) reconcileNodeDiscoveryJob(
 }
 
 func (r *GPUNodeReconciler) reconcileHypervisorPod(ctx context.Context, node *tfv1.GPUNode, pool *tfv1.GPUPool) (string, error) {
-	if pool == nil {
-		return "", fmt.Errorf("failed to get tensor-fusion pool, can not create hypervisor pod")
+	log := log.FromContext(ctx)
+
+	log.Info("reconciling hypervisor pod")
+
+	if pool.Spec.ComponentConfig == nil || pool.Spec.ComponentConfig.Hypervisor == nil {
+		return "", fmt.Errorf("missing hypervisor config")
+	}
+
+	key := client.ObjectKey{
+		Namespace: utils.CurrentNamespace(),
+		Name:      fmt.Sprintf("hypervisor-%s", node.Name),
+	}
+
+	currentPod := &corev1.Pod{}
+	if err := r.Get(ctx, key, currentPod); err != nil {
+		if !errors.IsNotFound(err) {
+			return "", fmt.Errorf("failed to get current hypervisor pod: %w", err)
+		}
+	} else {
+		if !currentPod.DeletionTimestamp.IsZero() {
+			log.Info("hypervisor pod is being deleted", "name", key.Name)
+			return key.Name, nil
+		}
+
+		if currentPod.Labels[constants.LabelKeyPodTemplateHash] != utils.GetObjectHash(pool.Spec.ComponentConfig.Hypervisor) {
+			if err := r.Delete(ctx, currentPod); err != nil {
+				return "", fmt.Errorf("failed to delete old hypervisor pod: %w", err)
+			}
+			log.Info("old hypervisor pod deleted", "name", currentPod.Name)
+		} else {
+			return key.Name, nil
+		}
 	}
 
-	namespace := utils.CurrentNamespace()
+	// no existing pod or config changed, so create new one
+	if err := r.createHypervisorPod(ctx, key, node, pool); err != nil {
+		if errors.IsAlreadyExists(err) {
+			return "", nil
+		} else {
+			return "", fmt.Errorf("failed to create hypervisor pod: %w", err)
+		}
+	}
+
+	return key.Name, nil
+}
+
+func (r *GPUNodeReconciler) createHypervisorPod(ctx context.Context, key client.ObjectKey, node *tfv1.GPUNode, pool *tfv1.GPUPool) error {
 	log := log.FromContext(ctx)
-	hypervisorPodName := fmt.Sprintf("hypervisor-%s", node.Name)
 
-	hypervisorConfig := pool.Spec.ComponentConfig.Hypervisor
 	podTmpl := &corev1.PodTemplate{}
-	err := json.Unmarshal(hypervisorConfig.PodTemplate.Raw, podTmpl)
+	err := json.Unmarshal(pool.Spec.ComponentConfig.Hypervisor.PodTemplate.Raw, podTmpl)
 	if err != nil {
-		return "", fmt.Errorf("failed to unmarshal pod template: %w", err)
+		return fmt.Errorf("failed to unmarshal pod template: %w", err)
 	}
 	spec := podTmpl.Template.Spec.DeepCopy()
 	if spec.NodeSelector == nil {
@@ -402,10 +442,11 @@ func (r *GPUNodeReconciler) reconcileHypervisorPod(ctx context.Context, node *tf
 	})
 	newPod := &corev1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
-			Name:      hypervisorPodName,
-			Namespace: namespace,
+			Name:      key.Name,
+			Namespace: key.Namespace,
 			Labels: map[string]string{
 				fmt.Sprintf(constants.GPUNodePoolIdentifierLabelFormat, pool.Name): "true",
+				constants.LabelKeyPodTemplateHash:                                  utils.GetObjectHash(pool.Spec.ComponentConfig.Hypervisor),
 			},
 		},
 		Spec: *spec,
@@ -419,18 +460,18 @@ func (r *GPUNodeReconciler) reconcileHypervisorPod(ctx context.Context, node *tf
 		Operator: corev1.TolerationOpExists,
 	})
 
-	e := controllerutil.SetControllerReference(node, newPod, r.Scheme)
-	if e != nil {
-		return "", fmt.Errorf("failed to set controller reference: %w", e)
-	}
-	_, err = controllerutil.CreateOrUpdate(ctx, r.Client, newPod, func() error {
-		log.Info("Creating or Updating hypervisor pod", "name", hypervisorPodName)
-		return nil
-	})
+	err = controllerutil.SetControllerReference(node, newPod, r.Scheme)
 	if err != nil {
-		return "", fmt.Errorf("failed to create or update hypervisor pod: %w", err)
+		return fmt.Errorf("failed to set controller reference: %w", err)
+	}
+
+	if err = r.Create(ctx, newPod); err != nil {
+		return fmt.Errorf("failed to create hypervisor pod: %w", err)
 	}
-	return hypervisorPodName, nil
+
+	log.Info("hypervisor pod created", "name", key.Name)
+
+	return nil
 }
 
 func (r *GPUNodeReconciler) reconcileCloudVendorNode(ctx context.Context, node *tfv1.GPUNode, pool *tfv1.GPUPool) error {

internal/controller/gpunode_controller_test.go

@@ -18,13 +18,16 @@ package controller
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
+	"time"
 
 	tfv1 "github.com/NexusGPU/tensor-fusion/api/v1"
 	"github.com/NexusGPU/tensor-fusion/internal/constants"
 	"github.com/NexusGPU/tensor-fusion/internal/utils"
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
+	"github.com/samber/lo"
 	batchv1 "k8s.io/api/batch/v1"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/errors"
@@ -117,6 +120,31 @@ var _ = Describe("GPUNode Controller", func() {
 				Name:      fmt.Sprintf("hypervisor-%s", gpunode.Name),
 				Namespace: utils.CurrentNamespace(),
 			}, pod)).To(Succeed())
+
+			By("Verify the hypervior pod recreated after hypervisor config change")
+			pool := &tfv1.GPUPool{}
+			Expect(k8sClient.Get(ctx, types.NamespacedName{
+				Name:      "mock",
+				Namespace: "default",
+			}, pool)).To(Succeed())
+			podTmpl := &corev1.PodTemplate{}
+			err = json.Unmarshal(pool.Spec.ComponentConfig.Hypervisor.PodTemplate.Raw, podTmpl)
+			Expect(err).NotTo(HaveOccurred())
+			podTmpl.Template.Spec.Containers[0].Name = "foo"
+			pool.Spec.ComponentConfig.Hypervisor.PodTemplate.Raw = lo.Must(json.Marshal(podTmpl))
+			Expect(k8sClient.Update(ctx, pool)).To(Succeed())
+			Eventually(func() string {
+				_, err = controllerReconciler.Reconcile(ctx, reconcile.Request{
+					NamespacedName: typeNamespacedName,
+				})
+				if err = k8sClient.Get(ctx, types.NamespacedName{
+					Name:      fmt.Sprintf("hypervisor-%s", gpunode.Name),
+					Namespace: utils.CurrentNamespace(),
+				}, pod); err != nil {
+					return ""
+				}
+				return pod.Spec.Containers[0].Name
+			}, 5*time.Second, time.Second).Should(Equal("foo"))
 		})
 	})
 })