[TCL-9479] feat(nodeset): reconcile compute class features

VERSION

@@ -1 +1 @@
-1.0.15
+1.0.16

internal/controller/nodeset/eventhandler/eventhandler_node.go

@@ -22,6 +22,8 @@ import (
 	"github.com/togethercomputer/slurm-operator/internal/utils/objectutils"
 )
 
+const computeClassLabel = "together.ai/compute-class"
+
 func NewNodeEventHandler(reader client.Reader) *NodeEventHandler {
 	return &NodeEventHandler{
 		Reader: reader,
@@ -40,7 +42,11 @@ func (h *NodeEventHandler) Create(
 	evt event.CreateEvent,
 	q workqueue.TypedRateLimitingInterface[reconcile.Request],
 ) {
-	// Intentionally blank
+	node, ok := evt.Object.(*corev1.Node)
+	if !ok {
+		return
+	}
+	h.enqueueNodeSetsForNode(ctx, node, q)
 }
 
 // Delete implements handler.EventHandler
@@ -85,6 +91,10 @@ func (h *NodeEventHandler) Update(
 	if !apiequality.Semantic.DeepEqual(oldNode.Annotations, newNode.Annotations) {
 		h.enqueueNodeSetsForNode(ctx, newNode, q)
 	}
+
+	if oldNode.Labels[computeClassLabel] != newNode.Labels[computeClassLabel] {
+		h.enqueueNodeSetsForNode(ctx, newNode, q)
+	}
 }
 
 func (h *NodeEventHandler) enqueueNodeSetsForNode(
@@ -114,13 +124,8 @@ func (h *NodeEventHandler) enqueueNodeSetsForNode(
 		if nodeset == nil {
 			continue
 		}
-		if node.Spec.Unschedulable {
-			logger.Info("Node was cordoned, reconcile NodeSet with Pod on Node",
-				"node", node.Name, "nodeset", klog.KObj(nodeset), "pod", klog.KObj(&pod))
-		} else {
-			logger.Info("Node was uncordoned, reconcile NodeSet with Pod on Node",
-				"node", node.Name, "nodeset", klog.KObj(nodeset), "pod", klog.KObj(&pod))
-		}
+		logger.Info("Node changed, reconcile NodeSet with Pod on Node",
+			"node", node.Name, "nodeset", klog.KObj(nodeset), "pod", klog.KObj(&pod))
 		objectutils.EnqueueRequest(q, nodeset)
 	}
 }

internal/controller/nodeset/eventhandler/eventhandler_node_test.go

@@ -47,6 +47,21 @@ func Test_NodeEventHandler_Create(t *testing.T) {
 			},
 			want: 0,
 		},
+		{
+			name: "Node created - should enqueue NodeSet",
+			fields: fields{
+				Reader: indexes.NewFakeClientBuilderWithIndexes(
+					newNodeSet("foo", "slurm", 0),
+					newNodeSetPod(newNodeSet("foo", "slurm", 0), 0, "test-node"),
+				).Build(),
+			},
+			args: args{
+				ctx: context.TODO(),
+				evt: event.CreateEvent{Object: newNode("test-node", false)},
+				q:   newQueue(),
+			},
+			want: 1,
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -205,6 +220,42 @@ func Test_NodeEventHandler_Update(t *testing.T) {
 			},
 			want: 0, // Should not enqueue anything
 		},
+		{
+			name: "Compute class changed - should enqueue NodeSet",
+			fields: fields{
+				Reader: indexes.NewFakeClientBuilderWithIndexes(
+					nodeset,
+					newNodeSetPod(nodeset, 0, "test-node"),
+				).Build(),
+			},
+			args: args{
+				ctx: context.TODO(),
+				evt: event.UpdateEvent{
+					ObjectOld: newNodeWithLabels("test-node", map[string]string{computeClassLabel: "preemptible"}),
+					ObjectNew: newNodeWithLabels("test-node", map[string]string{computeClassLabel: "non-preemptible"}),
+				},
+				q: newQueue(),
+			},
+			want: 1,
+		},
+		{
+			name: "Unrelated label changed - should not enqueue",
+			fields: fields{
+				Reader: indexes.NewFakeClientBuilderWithIndexes(
+					nodeset,
+					newNodeSetPod(nodeset, 0, "test-node"),
+				).Build(),
+			},
+			args: args{
+				ctx: context.TODO(),
+				evt: event.UpdateEvent{
+					ObjectOld: newNodeWithLabels("test-node", map[string]string{"example.com/label": "old"}),
+					ObjectNew: newNodeWithLabels("test-node", map[string]string{"example.com/label": "new"}),
+				},
+				q: newQueue(),
+			},
+			want: 0,
+		},
 		{
 			name: "No worker pods on node - should not enqueue",
 			fields: fields{
@@ -250,3 +301,9 @@ func newNode(name string, unschedulable bool) *corev1.Node {
 	}
 	return node
 }
+
+func newNodeWithLabels(name string, labels map[string]string) *corev1.Node {
+	node := newNode(name, false)
+	node.Labels = labels
+	return node
+}

internal/controller/nodeset/nodeset_sync_status.go

@@ -34,6 +34,8 @@ import (
 	slurmconditions "github.com/togethercomputer/slurm-operator/pkg/conditions"
 )
 
+const computeClassLabel = "together.ai/compute-class"
+
 // syncStatus handles synchronizing Slurm Nodes and NodeSet Status.
 func (r *NodeSetReconciler) syncStatus(
 	ctx context.Context,
@@ -74,7 +76,21 @@ func (r *NodeSetReconciler) syncSlurmStatus(
 		if !podutils.IsHealthy(pod) {
 			return nil
 		}
-		return r.slurmControl.UpdateNodeWithPodInfo(ctx, nodeset, pod)
+
+		var computeClass *string
+		if pod.Spec.NodeName != "" {
+			node := &corev1.Node{}
+			if err := r.Get(ctx, client.ObjectKey{Name: pod.Spec.NodeName}, node); err != nil {
+				if !apierrors.IsNotFound(err) {
+					return err
+				}
+			} else {
+				value := node.Labels[computeClassLabel]
+				computeClass = &value
+			}
+		}
+
+		return r.slurmControl.UpdateNodeWithPodInfo(ctx, nodeset, pod, computeClass)
 	}
 	if _, err := utils.SlowStartBatch(len(pods), utils.SlowStartInitialBatchSize, syncSlurmStatusFn); err != nil {
 		return err

internal/controller/nodeset/slurmcontrol/slurmcontrol.go

@@ -35,8 +35,8 @@ import (
 type SlurmControlInterface interface {
 	// RefreshNodeCache forces the Node cache to be refreshed
 	RefreshNodeCache(ctx context.Context, nodeset *slinkyv1beta1.NodeSet) error
-	// UpdateNodeWithPodInfo handles updating the Node with its pod info
-	UpdateNodeWithPodInfo(ctx context.Context, nodeset *slinkyv1beta1.NodeSet, pod *corev1.Pod) error
+	// UpdateNodeWithPodInfo handles updating the Node with its pod info and compute class.
+	UpdateNodeWithPodInfo(ctx context.Context, nodeset *slinkyv1beta1.NodeSet, pod *corev1.Pod, computeClass *string) error
 	// MakeNodeDrain handles adding the DRAIN state to the slurm node.
 	MakeNodeDrain(ctx context.Context, nodeset *slinkyv1beta1.NodeSet, pod *corev1.Pod, reason string) error
 	// MakeNodeUndrain handles removing the DRAIN state from the slurm node.
@@ -125,7 +125,12 @@ func (r *realSlurmControl) GetNodeNames(ctx context.Context, nodeset *slinkyv1be
 }
 
 // UpdateNodeWithPodInfo implements SlurmControlInterface.
-func (r *realSlurmControl) UpdateNodeWithPodInfo(ctx context.Context, nodeset *slinkyv1beta1.NodeSet, pod *corev1.Pod) error {
+func (r *realSlurmControl) UpdateNodeWithPodInfo(
+	ctx context.Context,
+	nodeset *slinkyv1beta1.NodeSet,
+	pod *corev1.Pod,
+	computeClass *string,
+) error {
 	logger := log.FromContext(ctx)
 
 	slurmClient := r.lookupClient(nodeset)
@@ -152,24 +157,37 @@ func (r *realSlurmControl) UpdateNodeWithPodInfo(ctx context.Context, nodeset *s
 	podInfoOld := &podinfo.PodInfo{}
 	_ = podinfo.ParseIntoPodInfo(slurmNode.Comment, podInfoOld)
 
-	if podInfoOld.Equal(podInfo) {
-		logger.V(3).Info("Node already contains podInfo, skipping update request",
-			"node", slurmNode.GetKey(), "podInfo", podInfo)
-		return nil
+	req := slurmapi.V0044UpdateNodeMsg{}
+	if !podInfoOld.Equal(podInfo) {
+		req.Comment = ptr.To(podInfo.ToString())
 	}
 
-	logger.Info("Update Slurm Node with Kubernetes Pod info",
-		"Node", slurmNode.Name, "podInfo", podInfo)
-	req := slurmapi.V0044UpdateNodeMsg{
-		Comment: ptr.To(podInfo.ToString()),
+	if computeClass != nil {
+		features := reconcileComputeClass(ptr.Deref(slurmNode.Features, nil), *computeClass)
+		if !featuresEqual(ptr.Deref(slurmNode.Features, nil), features) {
+			req.Features = ptr.To(features)
+		}
+
+		activeFeatures := reconcileComputeClass(ptr.Deref(slurmNode.ActiveFeatures, nil), *computeClass)
+		if !featuresEqual(ptr.Deref(slurmNode.ActiveFeatures, nil), activeFeatures) {
+			req.FeaturesAct = ptr.To(activeFeatures)
+		}
 	}
-	if err := slurmClient.Update(ctx, slurmNode, req); err != nil {
-		if !tolerateError(err) {
-			return err
+
+	if req.Comment != nil || req.Features != nil || req.FeaturesAct != nil {
+		logger.Info("Update Slurm Node with Kubernetes Pod info and features",
+			"Node", slurmNode.Name, "podInfo", podInfo, "computeClass", computeClass)
+		if err := slurmClient.Update(ctx, slurmNode, req); err != nil {
+			if !tolerateError(err) {
+				return err
+			}
 		}
+	} else {
+		logger.V(3).Info("Slurm Node pod info and features are already reconciled",
+			"node", slurmNode.GetKey(), "podInfo", podInfo)
 	}
 
-	if podInfoOld.Node != "" {
+	if !podInfoOld.Equal(podInfo) && podInfoOld.Node != "" {
 		logger.Info("Update Slurm Node state due to Kubernetes node migration", "Node", slurmNode.Name)
 		req := slurmapi.V0044UpdateNodeMsg{
 			State: ptr.To([]slurmapi.V0044UpdateNodeMsgState{slurmapi.V0044UpdateNodeMsgStateIDLE}),
@@ -184,6 +202,43 @@ func (r *realSlurmControl) UpdateNodeWithPodInfo(ctx context.Context, nodeset *s
 	return nil
 }
 
+const (
+	computeClassPreemptible    = "preemptible"
+	computeClassNonPreemptible = "non-preemptible"
+)
+
+func reconcileComputeClass(features []string, computeClass string) []string {
+	out := make([]string, 0, len(features)+1)
+	for _, feature := range features {
+		if feature != computeClassPreemptible && feature != computeClassNonPreemptible {
+			out = append(out, feature)
+		}
+	}
+	if computeClass == computeClassPreemptible || computeClass == computeClassNonPreemptible {
+		out = append(out, computeClass)
+	}
+	return out
+}
+
+func featuresEqual(left, right []string) bool {
+	if len(left) != len(right) {
+		return false
+	}
+
+	counts := make(map[string]int, len(left))
+	for _, feature := range left {
+		counts[feature]++
+	}
+	for _, feature := range right {
+		counts[feature]--
+		if counts[feature] < 0 {
+			return false
+		}
+	}
+
+	return true
+}
+
 const nodeReasonPrefix = "slurm-operator:"
 
 // MakeNodeDrain implements SlurmControlInterface.