fix: add node hash for gpu k8s node, owner ref for hypervisor, isolate shm (#352)

Author: Code2Life

internal/constants/constants.go

@@ -30,6 +30,7 @@ const (
 	LabelKeyClusterOwner    = Domain + "/cluster"
 	LabelKeyNodeClass       = Domain + "/node-class"
 	LabelKeyPodTemplateHash = Domain + "/pod-template-hash"
+	LabelNodeSelectorHash   = Domain + "/node-selector-hash"
 	LabelComponent          = Domain + "/component"
 	// used by TF connection, for matching the related connections when worker Pod state changed
 	LabelWorkerName = Domain + "/worker-name"

internal/controller/gpunode_controller.go

@@ -140,7 +140,7 @@ func (r *GPUNodeReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ct
 		return ctrl.Result{}, nil
 	}
 
-	hypervisorName, err := r.reconcileHypervisorPod(ctx, node, poolObj)
+	hypervisorName, err := r.reconcileHypervisorPod(ctx, node, poolObj, coreNode)
 	if err != nil {
 		return ctrl.Result{}, err
 	}
@@ -319,7 +319,12 @@ func (r *GPUNodeReconciler) reconcileNodeDiscoveryJob(
 	return nil
 }
 
-func (r *GPUNodeReconciler) reconcileHypervisorPod(ctx context.Context, node *tfv1.GPUNode, pool *tfv1.GPUPool) (string, error) {
+func (r *GPUNodeReconciler) reconcileHypervisorPod(
+	ctx context.Context,
+	node *tfv1.GPUNode,
+	pool *tfv1.GPUPool,
+	k8sNode *corev1.Node,
+) (string, error) {
 	log := log.FromContext(ctx)
 
 	if pool.Spec.ComponentConfig == nil || pool.Spec.ComponentConfig.Hypervisor == nil {
@@ -361,7 +366,7 @@ func (r *GPUNodeReconciler) reconcileHypervisorPod(ctx context.Context, node *tf
 	}
 
 	log.Info("hypervisor pod not found, creating new one", "node", node.Name)
-	if err := r.createHypervisorPod(ctx, key, node, pool); err != nil {
+	if err := r.createHypervisorPod(ctx, key, node, pool, k8sNode); err != nil {
 		if errors.IsAlreadyExists(err) {
 			log.Info("hypervisor pod already exists, skip creation", "node", node.Name)
 			return "", nil
@@ -372,7 +377,13 @@ func (r *GPUNodeReconciler) reconcileHypervisorPod(ctx context.Context, node *tf
 	return key.Name, nil
 }
 
-func (r *GPUNodeReconciler) createHypervisorPod(ctx context.Context, key client.ObjectKey, node *tfv1.GPUNode, pool *tfv1.GPUPool) error {
+func (r *GPUNodeReconciler) createHypervisorPod(
+	ctx context.Context,
+	key client.ObjectKey,
+	node *tfv1.GPUNode,
+	pool *tfv1.GPUPool,
+	k8sNode *corev1.Node,
+) error {
 	log := log.FromContext(ctx)
 
 	podTmpl := &corev1.PodTemplate{}
@@ -447,7 +458,11 @@ func (r *GPUNodeReconciler) createHypervisorPod(ctx context.Context, key client.
 	})
 	err = controllerutil.SetControllerReference(node, newPod, r.Scheme)
 	if err != nil {
-		return fmt.Errorf("failed to set controller reference: %w", err)
+		return fmt.Errorf("failed to set controller reference for hypervisor: %w", err)
+	}
+	// also set node owned by k8s node to allow Karpenter to delete the node while hypervisor exists
+	if err := controllerutil.SetOwnerReference(k8sNode, newPod, r.Scheme); err != nil {
+		return fmt.Errorf("failed to set owner reference for hypervisor: %w", err)
 	}
 
 	// create hypervisor pod

internal/controller/gpupool_controller.go

@@ -30,13 +30,16 @@ import (
 	"github.com/NexusGPU/tensor-fusion/internal/metrics"
 	utils "github.com/NexusGPU/tensor-fusion/internal/utils"
 	"golang.org/x/time/rate"
+	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/equality"
 	"k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/api/resource"
 	"k8s.io/apimachinery/pkg/runtime"
 	utilerrors "k8s.io/apimachinery/pkg/util/errors"
 	"k8s.io/client-go/tools/record"
+	"k8s.io/client-go/util/retry"
 	"k8s.io/client-go/util/workqueue"
+	schedulingcorev1 "k8s.io/component-helpers/scheduling/corev1"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/controller"
@@ -83,6 +86,9 @@ type GPUPoolReconciler struct {
 // and requeue until current time after that, start provisioning loop
 var provisioningInitializationMinTime = map[string]time.Time{}
 
+// When GPU nodeSelector changed, trigger all node update
+var poolSelectorChangeMap = map[string]string{}
+
 // +kubebuilder:rbac:groups=tensor-fusion.ai,resources=gpupools,verbs=get;list;watch;create;update;patch;delete
 // +kubebuilder:rbac:groups=tensor-fusion.ai,resources=gpupools/status,verbs=get;update;patch
 // +kubebuilder:rbac:groups=tensor-fusion.ai,resources=gpupools/finalizers,verbs=update
@@ -116,6 +122,10 @@ func (r *GPUPoolReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ct
 		return ctrl.Result{}, nil
 	}
 
+	if err := r.reconcilePoolSelectorChange(ctx, pool); err != nil {
+		return ctrl.Result{}, err
+	}
+
 	if err := r.reconcilePoolCurrentCapacityAndReadiness(ctx, pool); err != nil {
 		return ctrl.Result{}, err
 	}
@@ -404,6 +414,59 @@ func (r *GPUPoolReconciler) reconcilePoolComponents(ctx context.Context, pool *t
 	return ctrlResult, utilerrors.NewAggregate(errs)
 }
 
+func (r *GPUPoolReconciler) reconcilePoolSelectorChange(ctx context.Context, pool *tfv1.GPUPool) error {
+	if pool.Spec.NodeManagerConfig != nil && pool.Spec.NodeManagerConfig.NodeSelector != nil {
+		hash := utils.GetObjectHash(pool.Spec.NodeManagerConfig.NodeSelector)
+		if poolSelectorChangeMap[pool.Name] == hash {
+			return nil
+		}
+
+		// hash has changed, or first reconcile, should check all k8s nodes
+		nodes := &corev1.NodeList{}
+		selectors := utils.GetInitialGPUNodeSelector()
+		if err := r.List(ctx, nodes, client.MatchingLabels{selectors[0]: selectors[1]}); err != nil {
+			return err
+		}
+		for _, node := range nodes.Items {
+			// skip no label or deleting nodes
+			if node.Labels == nil || !node.DeletionTimestamp.IsZero() {
+				continue
+			}
+			matches, err := schedulingcorev1.MatchNodeSelectorTerms(&node, pool.Spec.NodeManagerConfig.NodeSelector)
+			if err != nil {
+				return err
+			}
+			if matches {
+				if err := UpdateK8SNodeSelectorHash(ctx, r.Client, &node, hash); err != nil {
+					return err
+				}
+			}
+		}
+		poolSelectorChangeMap[pool.Name] = hash
+		return nil
+	}
+	return nil
+}
+
+func UpdateK8SNodeSelectorHash(ctx context.Context, k8sClient client.Client, node *corev1.Node, hash string) error {
+	// skip nodes that already injected the hash
+	if node.Labels[constants.LabelNodeSelectorHash] == hash {
+		return nil
+	}
+	// update label to trigger the GPUNode reconcile
+	if err := retry.RetryOnConflict(retry.DefaultBackoff, func() error {
+		latest := &corev1.Node{}
+		if err := k8sClient.Get(ctx, client.ObjectKey{Name: node.Name}, latest); err != nil {
+			return err
+		}
+		latest.Labels[constants.LabelNodeSelectorHash] = hash
+		return k8sClient.Update(ctx, latest)
+	}); err != nil {
+		return err
+	}
+	return nil
+}
+
 func (r *GPUPoolReconciler) cleanUpPool(ctx context.Context, pool *tfv1.GPUPool) (bool, error) {
 	log := log.FromContext(ctx)
 	log.Info("TensorFusionGPUPool is being deleted", "name", pool.Name)

internal/controller/gpupool_controller_test.go

@@ -42,6 +42,14 @@ var _ = Describe("GPUPool Controller", func() {
 				pool := tfEnv.GetGPUPool(0)
 				g.Expect(pool.Status.Phase).Should(Equal(tfv1.TensorFusionPoolPhaseRunning))
 			}).Should(Succeed())
+			Eventually(func(g Gomega) {
+				nodeList := tfEnv.GetGPUNodeList(0)
+				for _, gpuNode := range nodeList.Items {
+					node := &corev1.Node{}
+					g.Expect(k8sClient.Get(ctx, client.ObjectKey{Name: gpuNode.Name}, node)).Should(Succeed())
+					g.Expect(node.Labels).To(HaveKey(constants.LabelNodeSelectorHash))
+				}
+			}).Should(Succeed())
 			tfEnv.Cleanup()
 		})
 	})

internal/controller/node_controller.go

@@ -19,8 +19,6 @@ package controller
 import (
 	"context"
 	"fmt"
-	"os"
-	"strings"
 
 	tfv1 "github.com/NexusGPU/tensor-fusion/api/v1"
 	"github.com/NexusGPU/tensor-fusion/internal/constants"
@@ -86,12 +84,15 @@ func (r *NodeReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.
 		return ctrl.Result{}, err
 	}
 	if !matched {
-		// delete gpunode if no matched pool
-		if err := r.Delete(ctx, &tfv1.GPUNode{
-			ObjectMeta: metav1.ObjectMeta{
-				Name: node.Name,
-			},
-		}); err != nil {
+		existingGPUNode := &tfv1.GPUNode{}
+		if err := r.Get(ctx, client.ObjectKey{Name: node.Name}, existingGPUNode); err != nil {
+			if errors.IsNotFound(err) {
+				return ctrl.Result{}, nil
+			}
+			return ctrl.Result{}, fmt.Errorf("can not get gpuNode(%s) : %w", node.Name, err)
+		}
+		// delete existing gpunode if no matched pool
+		if err := r.Delete(ctx, existingGPUNode); err != nil {
 			// requeue if the gpunode is not generated
 			if errors.IsNotFound(err) {
 				return ctrl.Result{}, nil
@@ -121,6 +122,14 @@ func (r *NodeReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.
 		return ctrl.Result{}, nil
 	}
 
+	// update k8s node hash
+	hash := utils.GetObjectHash(pool.Spec.NodeManagerConfig.NodeSelector)
+	if node.Labels[constants.LabelNodeSelectorHash] != hash {
+		if err := UpdateK8SNodeSelectorHash(ctx, r.Client, node, hash); err != nil {
+			return ctrl.Result{}, fmt.Errorf("failed to update k8s node hash: %w", err)
+		}
+	}
+
 	provisioningMode := pool.Spec.NodeManagerConfig.ProvisioningMode
 	isDirectManagedMode := provisioningMode == tfv1.ProvisioningModeProvisioned
 	isManagedNode := isDirectManagedMode || provisioningMode == tfv1.ProvisioningModeKarpenter
@@ -199,11 +208,7 @@ func (r *NodeReconciler) generateGPUNode(node *corev1.Node, pool *tfv1.GPUPool,
 // SetupWithManager sets up the controller with the Manager.
 func (r *NodeReconciler) SetupWithManager(mgr ctrl.Manager) error {
 	// must choose an initial label selector to avoid performance impact in large Kubernetes clusters
-	selector := os.Getenv("INITIAL_GPU_NODE_LABEL_SELECTOR")
-	if selector == "" {
-		selector = constants.InitialGPUNodeSelector
-	}
-	selectors := strings.Split(selector, "=")
+	selectors := utils.GetInitialGPUNodeSelector()
 	p, err := predicate.LabelSelectorPredicate(metav1.LabelSelector{
 		MatchLabels: map[string]string{
 			selectors[0]: selectors[1],

internal/utils/compose.go

@@ -229,10 +229,9 @@ func AddTFDefaultClientConfBeforePatch(
 			pod.Spec.Containers[injectContainerIndex].VolumeMounts = append(
 				pod.Spec.Containers[injectContainerIndex].VolumeMounts,
 				v1.VolumeMount{
-					Name:      constants.DataVolumeName,
-					MountPath: constants.SharedMemDeviceName,
-					SubPath:   constants.SharedMemMountSubPath,
-					//  + constants.TFLibsVolumeMountPath, SubPathExpr:      constants.TFDataPathWorkerExpr,
+					Name:             constants.DataVolumeName,
+					MountPath:        constants.TFLibsVolumeMountPath,
+					SubPathExpr:      constants.TFDataPathWorkerExpr,
 					MountPropagation: ptr.To(v1.MountPropagationHostToContainer),
 				})
 
@@ -682,12 +681,9 @@ func AddWorkerConfAfterTemplate(ctx context.Context, spec *v1.PodSpec, workerCon
 	spec.Containers[0].VolumeMounts = append(
 		spec.Containers[0].VolumeMounts,
 		v1.VolumeMount{
-			Name:      constants.DataVolumeName,
-			MountPath: constants.SharedMemDeviceName,
-			// TODO not working.
-			// + constants.TFLibsVolumeMountPath
-			// SubPathExpr: constants.TFDataPathWorkerExpr,
-			SubPath:          constants.SharedMemMountSubPath,
+			Name:             constants.DataVolumeName,
+			MountPath:        constants.TFLibsVolumeMountPath,
+			SubPathExpr:      constants.TFDataPathWorkerExpr,
 			MountPropagation: ptr.To(v1.MountPropagationHostToContainer),
 		})
 	spec.Containers[0].Env = append(spec.Containers[0].Env, v1.EnvVar{

internal/utils/reconcile.go

@@ -214,6 +214,15 @@ func IsTensorFusionWorker(pod *corev1.Pod) bool {
 	return pod.Labels[constants.LabelComponent] == constants.ComponentWorker
 }
 
+func GetInitialGPUNodeSelector() []string {
+	selector := os.Getenv("INITIAL_GPU_NODE_LABEL_SELECTOR")
+	if selector == "" {
+		selector = constants.InitialGPUNodeSelector
+	}
+	selectors := strings.Split(selector, "=")
+	return selectors
+}
+
 var GPUResourceNames = []corev1.ResourceName{
 	"nvidia.com/gpu",
 	"amd.com/gpu",