build: image for PR #43 (compute class features) — do not merge

.github/workflows/container-images-1.0.yaml

@@ -0,0 +1,115 @@
+# Build and push slurm-operator container images to Docker Hub.
+#
+# Tagging:
+#   - push to main: togethercomputer/slurm-operator:<VERSION> (VERSION from ./VERSION, no suffix)
+#   - pull_request (same-repo only): ...:<VERSION>-dev-<short_sha> ; re-runs on every PR commit
+#
+# Required repository secrets (see together-web dockerhub-release-tag, substrate-metadata deploy):
+#   - DOCKERHUB_USERNAME
+#   - DOCKERHUB_TOKEN
+# Optional (private Go modules / togethercomputer forks):
+#   - ROBOT_GITHUB_TOKEN — git insteadOf for github.com/togethercomputer/ (same pattern as substrate-metadata)
+#
+# Fork PRs: build runs; push and Docker Hub login are skipped (no secrets on fork workflows).
+# CI avoids double image builds: no-push uses `make build` once; push uses `make build-chart` then a single `make push-images` (one bake --push).
+
+name: Container images v1.0
+
+on:
+  pull_request:
+    branches:
+      - slurm-1.0-together-changes
+  push:
+    branches:
+      - slurm-1.0-together-changes
+  workflow_dispatch:
+
+concurrency:
+  group: container-images-${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: true
+
+env:
+  REGISTRY: togethercomputer
+
+jobs:
+  build-and-push:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+        with:
+          # PRs: use the tip of the PR branch so the image tag matches the commit under review.
+          ref: ${{ github.event.pull_request.head.sha || github.sha }}
+          fetch-depth: 0
+
+      - name: Set image VERSION for docker-bake
+        id: image-version
+        run: |
+          set -euo pipefail
+          BASE="$(tr -d ' \n\r\t' < VERSION)"
+          SHORT="$(git rev-parse --short HEAD)"
+          if [ "${{ github.event_name }}" = "pull_request" ]; then
+            IMG="${BASE}-dev-${SHORT}"
+          else
+            IMG="${BASE}"
+          fi
+          echo "version=${IMG}" >> "$GITHUB_OUTPUT"
+          echo "Image VERSION (make/docker-bake): ${IMG}"
+
+      - name: Whether to push to Docker Hub
+        id: should-push
+        run: |
+          if [ "${{ github.event_name }}" = "push" ] && [ "${{ github.ref }}" = "refs/heads/main" ]; then
+            echo "push=true" >> "$GITHUB_OUTPUT"
+          elif [ "${{ github.event_name }}" = "pull_request" ] && \
+               [ "${{ github.event.pull_request.head.repo.full_name }}" = "${{ github.repository }}" ]; then
+            echo "push=true" >> "$GITHUB_OUTPUT"
+          elif [ "${{ github.event_name }}" = "workflow_dispatch" ] && [ "${{ github.ref }}" = "refs/heads/main" ]; then
+            echo "push=true" >> "$GITHUB_OUTPUT"
+          else
+            echo "push=false" >> "$GITHUB_OUTPUT"
+          fi
+
+      - name: Set up Go
+        uses: actions/setup-go@v5
+        with:
+          go-version-file: go.mod
+
+      - name: Configure git for private Go modules
+        run: |
+          git config --global url."https://together-robot:${ROBOT_GITHUB_TOKEN}@github.com/togethercomputer/".insteadOf "https://github.com/togethercomputer/"
+        env:
+          ROBOT_GITHUB_TOKEN: ${{ secrets.ROBOT_GITHUB_TOKEN }}
+
+      - name: Set up Helm
+        uses: azure/setup-helm@v4
+        with:
+          version: v3.16.3
+
+      - name: Set up QEMU
+        uses: docker/setup-qemu-action@v3
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Log in to Docker Hub
+        if: steps.should-push.outputs.push == 'true'
+        uses: docker/login-action@v3
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
+
+      # Push path: one bake (--push) only. Non-push: make build = build-images + build-chart (single bake).
+      - name: Build (images + charts, no registry push)
+        if: steps.should-push.outputs.push != 'true'
+        run: make build VERSION="${{ steps.image-version.outputs.version }}"
+
+      - name: Package charts (push path)
+        if: steps.should-push.outputs.push == 'true'
+        run: make build-chart
+
+      - name: Build and push images
+        if: steps.should-push.outputs.push == 'true'
+        run: make push-images VERSION="${{ steps.image-version.outputs.version }}"

go.mod

@@ -1,4 +1,4 @@
-module github.com/SlinkyProject/slurm-operator
+module github.com/togethercomputer/slurm-operator
 
 go 1.25.0
 

internal/controller/nodeset/eventhandler/eventhandler_node.go

@@ -22,6 +22,8 @@ import (
 	"github.com/SlinkyProject/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/SlinkyProject/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

@@ -34,8 +34,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.
@@ -112,7 +112,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)
@@ -139,24 +144,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}),
@@ -171,6 +189,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.