diff --git a/VERSION b/VERSION index a9707166b..b668c3b2c 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -1.0.15 +1.0.16 diff --git a/internal/controller/nodeset/eventhandler/eventhandler_node.go b/internal/controller/nodeset/eventhandler/eventhandler_node.go index 65018ffbe..f316161fe 100644 --- a/internal/controller/nodeset/eventhandler/eventhandler_node.go +++ b/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) } } diff --git a/internal/controller/nodeset/eventhandler/eventhandler_node_test.go b/internal/controller/nodeset/eventhandler/eventhandler_node_test.go index 8f2d6d13e..8410fa8db 100644 --- a/internal/controller/nodeset/eventhandler/eventhandler_node_test.go +++ b/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 +} diff --git a/internal/controller/nodeset/nodeset_sync_status.go b/internal/controller/nodeset/nodeset_sync_status.go index 25bf71d66..79a7d1bb7 100644 --- a/internal/controller/nodeset/nodeset_sync_status.go +++ b/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 diff --git a/internal/controller/nodeset/slurmcontrol/slurmcontrol.go b/internal/controller/nodeset/slurmcontrol/slurmcontrol.go index 0e01846c5..37d6bda01 100644 --- a/internal/controller/nodeset/slurmcontrol/slurmcontrol.go +++ b/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. diff --git a/internal/controller/nodeset/slurmcontrol/slurmcontrol_test.go b/internal/controller/nodeset/slurmcontrol/slurmcontrol_test.go index 66097bff2..85d03bc74 100644 --- a/internal/controller/nodeset/slurmcontrol/slurmcontrol_test.go +++ b/internal/controller/nodeset/slurmcontrol/slurmcontrol_test.go @@ -8,6 +8,7 @@ import ( "errors" "net/http" "reflect" + "slices" "testing" "time" @@ -89,8 +90,18 @@ var _ = Describe("SlurmControlInterface", func() { } } o.State = ptr.To(stateSet.UnsortedList()) - o.Comment = r.Comment - o.Reason = r.Reason + if r.Comment != nil { + o.Comment = r.Comment + } + if r.Reason != nil { + o.Reason = r.Reason + } + if r.Features != nil { + o.Features = r.Features + } + if r.FeaturesAct != nil { + o.ActiveFeatures = r.FeaturesAct + } default: return errors.New("failed to cast slurm object") } @@ -116,7 +127,7 @@ var _ = Describe("SlurmControlInterface", func() { slurmcontrol = NewSlurmControl(controllers) By("Update Slurm pod info") - err := slurmcontrol.UpdateNodeWithPodInfo(ctx, nodeset, pod) + err := slurmcontrol.UpdateNodeWithPodInfo(ctx, nodeset, pod, nil) Expect(err).ToNot(HaveOccurred()) By("Check Slurm Node podInfo") @@ -187,7 +198,7 @@ var _ = Describe("SlurmControlInterface", func() { slurmcontrol = NewSlurmControl(controllers) By("Update Slurm pod info") - err := slurmcontrol.UpdateNodeWithPodInfo(ctx, nodeset, pod) + err := slurmcontrol.UpdateNodeWithPodInfo(ctx, nodeset, pod, nil) Expect(err).ToNot(HaveOccurred()) By("Check Slurm Node podInfo") @@ -213,7 +224,7 @@ var _ = Describe("SlurmControlInterface", func() { pod.Spec.NodeName = "bar" By("Update Slurm pod info") - err = slurmcontrol.UpdateNodeWithPodInfo(ctx, nodeset, pod) + err = slurmcontrol.UpdateNodeWithPodInfo(ctx, nodeset, pod, nil) Expect(err).ToNot(HaveOccurred()) By("Check Slurm Node podInfo") @@ -633,6 +644,142 @@ var _ = Describe("SlurmControlInterface", func() { }) }) +func TestUpdateNodeWithPodInfo_ReconcilesComputeClassFeatures(t *testing.T) { + ctx := context.Background() + controller := &slinkyv1beta1.Controller{ObjectMeta: metav1.ObjectMeta{Name: "slurm"}} + nodeset := newNodeSet("foo", controller.Name, 1) + pod := nodesetutils.NewNodeSetPod(nodeset, controller, 0, "") + pod.Spec.NodeName = "kube-node" + slurmNodeName := nodesetutils.GetNodeName(pod) + info := podinfo.PodInfo{ + Namespace: pod.Namespace, + PodName: pod.Name, + Node: pod.Spec.NodeName, + } + node := &types.V0044Node{ + V0044Node: api.V0044Node{ + Name: ptr.To(slurmNodeName), + Comment: ptr.To(info.ToString()), + Features: ptr.To(api.V0044CsvString{"nodeset-static", "configured"}), + ActiveFeatures: ptr.To(api.V0044CsvString{"nodeset-static", "active-configured"}), + }, + } + + updateCount := 0 + updateFn := func(_ context.Context, obj object.Object, req any, _ ...client.UpdateOption) error { + updateCount++ + slurmNode := obj.(*types.V0044Node) + update := req.(api.V0044UpdateNodeMsg) + if update.Features != nil { + slurmNode.Features = update.Features + } + if update.FeaturesAct != nil { + slurmNode.ActiveFeatures = update.FeaturesAct + } + return nil + } + sclient := fake.NewClientBuilder().WithUpdateFn(updateFn).WithObjects(node).Build() + slurmControl := NewSlurmControl(newSlurmClientMap(controller.Name, sclient)) + + assertFeatures := func(want, wantActive []string) { + t.Helper() + got := &types.V0044Node{} + if err := sclient.Get(ctx, object.ObjectKey(slurmNodeName), got); err != nil { + t.Fatal(err) + } + if !slices.Equal(ptr.Deref(got.Features, nil), want) { + t.Errorf("features = %v, want %v", ptr.Deref(got.Features, nil), want) + } + if !slices.Equal(ptr.Deref(got.ActiveFeatures, nil), wantActive) { + t.Errorf("active features = %v, want %v", ptr.Deref(got.ActiveFeatures, nil), wantActive) + } + } + + preemptible := computeClassPreemptible + if err := slurmControl.UpdateNodeWithPodInfo(ctx, nodeset, pod, &preemptible); err != nil { + t.Fatal(err) + } + assertFeatures( + []string{"nodeset-static", "configured", "preemptible"}, + []string{"nodeset-static", "active-configured", "preemptible"}, + ) + if updateCount != 1 { + t.Fatalf("update count = %d, want 1", updateCount) + } + + if err := slurmControl.UpdateNodeWithPodInfo(ctx, nodeset, pod, &preemptible); err != nil { + t.Fatal(err) + } + if updateCount != 1 { + t.Fatalf("idempotent update count = %d, want 1", updateCount) + } + + nonPreemptible := computeClassNonPreemptible + if err := slurmControl.UpdateNodeWithPodInfo(ctx, nodeset, pod, &nonPreemptible); err != nil { + t.Fatal(err) + } + assertFeatures( + []string{"nodeset-static", "configured", "non-preemptible"}, + []string{"nodeset-static", "active-configured", "non-preemptible"}, + ) + if updateCount != 2 { + t.Fatalf("label-change update count = %d, want 2", updateCount) + } + + unsupportedValue := "arbitrary-feature" + if err := slurmControl.UpdateNodeWithPodInfo(ctx, nodeset, pod, &unsupportedValue); err != nil { + t.Fatal(err) + } + assertFeatures( + []string{"nodeset-static", "configured"}, + []string{"nodeset-static", "active-configured"}, + ) + if updateCount != 3 { + t.Fatalf("unsupported-label update count = %d, want 3", updateCount) + } + + missingLabel := "" + if err := slurmControl.UpdateNodeWithPodInfo(ctx, nodeset, pod, &missingLabel); err != nil { + t.Fatal(err) + } + if updateCount != 3 { + t.Fatalf("missing-label update count = %d, want 3", updateCount) + } + + if err := slurmControl.UpdateNodeWithPodInfo(ctx, nodeset, pod, nil); err != nil { + t.Fatal(err) + } + if updateCount != 3 { + t.Fatalf("missing-node update count = %d, want 3", updateCount) + } +} + +func TestFeaturesEqual(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + left []string + right []string + want bool + }{ + {name: "same order", left: []string{"a", "b"}, right: []string{"a", "b"}, want: true}, + {name: "different order", left: []string{"a", "b"}, right: []string{"b", "a"}, want: true}, + {name: "different features", left: []string{"a", "b"}, right: []string{"a", "c"}, want: false}, + {name: "different duplicate counts", left: []string{"a", "a"}, right: []string{"a", "b"}, want: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + if got := featuresEqual(tt.left, tt.right); got != tt.want { + t.Fatalf("featuresEqual(%v, %v) = %t, want %t", tt.left, tt.right, got, tt.want) + } + }) + } +} + func Test_realSlurmControl_IsNodeDrain(t *testing.T) { ctx := context.Background() controller := &slinkyv1beta1.Controller{