feat(operator): Add node capacity labeler for GPU/Neuron

config/rbac/role.yaml

@@ -34,6 +34,15 @@ rules:
   - list
   - patch
   - watch
+- apiGroups:
+  - ""
+  resources:
+  - nodes
+  verbs:
+  - get
+  - list
+  - patch
+  - watch
 - apiGroups:
   - apps
   resources:

controllers/nodelabeler_controller.go

@@ -0,0 +1,174 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+// Package controllers contains the main controller, where the reconciliation starts.
+package controllers
+
+import (
+	"context"
+	"encoding/json"
+	"reflect"
+
+	"github.com/go-logr/logr"
+	corev1 "k8s.io/api/core/v1"
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/types"
+	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/builder"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/event"
+	"sigs.k8s.io/controller-runtime/pkg/predicate"
+)
+
+const (
+	// GPUPresentLabel is stamped on a node that exposes an NVIDIA GPU capacity key.
+	GPUPresentLabel = "cloudwatch.aws.amazon.com/gpu.present"
+	// NeuronPresentLabel is stamped on a node that exposes an AWS Neuron capacity key.
+	NeuronPresentLabel = "cloudwatch.aws.amazon.com/neuron.present"
+
+	nvidiaGPUResource    corev1.ResourceName = "nvidia.com/gpu"
+	neuronResource       corev1.ResourceName = "aws.amazon.com/neuron"
+	neuronCoreResource   corev1.ResourceName = "aws.amazon.com/neuroncore"
+	neuronDeviceResource corev1.ResourceName = "aws.amazon.com/neurondevice"
+)
+
+// relevantCapacityResources are the node capacity keys the labeler projects into labels.
+var relevantCapacityResources = []corev1.ResourceName{
+	nvidiaGPUResource,
+	neuronResource,
+	neuronCoreResource,
+	neuronDeviceResource,
+}
+
+// NodeLabelerReconciler projects accelerator capacity present on a Node into
+// fixed-key labels that a DaemonSet affinity can select on. The scheduler cannot
+// match on status.capacity directly, so we mirror capacity presence into labels.
+type NodeLabelerReconciler struct {
+	client.Client
+	Scheme *runtime.Scheme
+	Log    logr.Logger
+}
+
+// +kubebuilder:rbac:groups="",resources=nodes,verbs=get;list;watch;patch
+
+// Reconcile stamps or removes the GPU/Neuron presence labels on a single Node so
+// that they match the accelerator capacity the kubelet currently advertises.
+func (r *NodeLabelerReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
+	log := r.Log.WithValues("node", req.Name)
+
+	var node corev1.Node
+	if err := r.Get(ctx, req.NamespacedName, &node); err != nil {
+		if !apierrors.IsNotFound(err) {
+			log.Error(err, "unable to fetch Node")
+		}
+		// we'll ignore not-found errors, since they can't be fixed by an immediate
+		// requeue (we'll need to wait for a new notification), and we can get them
+		// on deleted requests.
+		return ctrl.Result{}, client.IgnoreNotFound(err)
+	}
+
+	gpuDesired := hasCapacityKey(&node, nvidiaGPUResource)
+	neuronDesired := hasCapacityKey(&node, neuronResource) ||
+		hasCapacityKey(&node, neuronCoreResource) ||
+		hasCapacityKey(&node, neuronDeviceResource)
+
+	// Build a JSON merge patch touching only metadata.labels. We intentionally do
+	// NOT use MergeFrom on the fetched node object: that would send the full node
+	// (including our stale copy of status) and could race the kubelet's frequent
+	// status updates. A raw merge patch scoped to labels only touches what we own.
+	labels := map[string]interface{}{}
+	var changes []string
+	if value, changed, action := labelChange(&node, GPUPresentLabel, gpuDesired); changed {
+		labels[GPUPresentLabel] = value
+		changes = append(changes, action+" "+GPUPresentLabel)
+	}
+	if value, changed, action := labelChange(&node, NeuronPresentLabel, neuronDesired); changed {
+		labels[NeuronPresentLabel] = value
+		changes = append(changes, action+" "+NeuronPresentLabel)
+	}
+
+	if len(labels) == 0 {
+		log.V(1).Info("node capacity labels already up to date")
+		return ctrl.Result{}, nil
+	}
+
+	patch, err := json.Marshal(map[string]interface{}{"metadata": map[string]interface{}{"labels": labels}})
+	if err != nil {
+		return ctrl.Result{}, err
+	}
+	if err := r.Patch(ctx, &node, client.RawPatch(types.MergePatchType, patch)); err != nil {
+		log.Error(err, "unable to patch node labels")
+		return ctrl.Result{}, err
+	}
+
+	log.Info("updated node capacity labels", "changes", changes, "gpu.present", gpuDesired, "neuron.present", neuronDesired)
+	return ctrl.Result{}, nil
+}
+
+// hasCapacityKey reports whether the node advertises the given capacity key.
+//
+// We test for the *presence* of the key, not a non-zero quantity: when a device
+// plugin becomes unhealthy the kubelet keeps the capacity key but resets its
+// quantity to 0. We still want dcgm-exporter scheduled onto such a GPU node so it
+// can surface the degraded device, so presence of the key is the signal we use,
+// regardless of the advertised quantity.
+func hasCapacityKey(node *corev1.Node, name corev1.ResourceName) bool {
+	_, ok := node.Status.Capacity[name]
+	return ok
+}
+
+// labelChange decides how a single managed label must change to match desired.
+// It returns the merge-patch value ("true" to set, nil to delete), whether a
+// change is needed at all, and a short action string for logging.
+func labelChange(node *corev1.Node, key string, desired bool) (value interface{}, changed bool, action string) {
+	current, present := node.Labels[key]
+	switch {
+	case desired && current != "true":
+		return "true", true, "add"
+	case !desired && present:
+		return nil, true, "remove"
+	default:
+		return nil, false, ""
+	}
+}
+
+// nodeCapacityOrLabelPredicate limits reconciliation to events that can actually
+// change a managed label: every Create (so existing nodes get labeled on startup)
+// and only those Updates where a relevant capacity key or one of the two managed
+// labels differs. This filters out the frequent, irrelevant node status churn
+// (heartbeats, conditions, addresses). Delete and Generic are ignored.
+func nodeCapacityOrLabelPredicate() predicate.Funcs {
+	snapshot := func(node *corev1.Node) map[string]string {
+		snap := map[string]string{}
+		for _, name := range relevantCapacityResources {
+			if _, ok := node.Status.Capacity[name]; ok {
+				snap["cap/"+string(name)] = "present"
+			}
+		}
+		snap["lbl/"+GPUPresentLabel] = node.Labels[GPUPresentLabel]
+		snap["lbl/"+NeuronPresentLabel] = node.Labels[NeuronPresentLabel]
+		return snap
+	}
+	return predicate.Funcs{
+		CreateFunc: func(_ event.CreateEvent) bool { return true },
+		UpdateFunc: func(e event.UpdateEvent) bool {
+			oldNode, okOld := e.ObjectOld.(*corev1.Node)
+			newNode, okNew := e.ObjectNew.(*corev1.Node)
+			if !okOld || !okNew {
+				return false
+			}
+			return !reflect.DeepEqual(snapshot(oldNode), snapshot(newNode))
+		},
+		DeleteFunc:  func(_ event.DeleteEvent) bool { return false },
+		GenericFunc: func(_ event.GenericEvent) bool { return false },
+	}
+}
+
+// SetupWithManager tells the manager what our controller is interested in.
+func (r *NodeLabelerReconciler) SetupWithManager(mgr ctrl.Manager) error {
+	return ctrl.NewControllerManagedBy(mgr).
+		For(&corev1.Node{}, builder.WithPredicates(nodeCapacityOrLabelPredicate())).
+		Named("nodelabeler").
+		Complete(r)
+}

controllers/nodelabeler_controller_test.go

@@ -0,0 +1,169 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package controllers
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/types"
+	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/client/fake"
+	"sigs.k8s.io/controller-runtime/pkg/event"
+	logf "sigs.k8s.io/controller-runtime/pkg/log"
+)
+
+func newTestNodeLabelerReconciler(objs ...client.Object) *NodeLabelerReconciler {
+	scheme := runtime.NewScheme()
+	_ = corev1.AddToScheme(scheme)
+	c := fake.NewClientBuilder().WithScheme(scheme).WithObjects(objs...).Build()
+	return &NodeLabelerReconciler{
+		Client: c,
+		Scheme: scheme,
+		Log:    logf.Log.WithName("nodelabeler-test"),
+	}
+}
+
+func node(name string, labels map[string]string, capacity corev1.ResourceList) *corev1.Node {
+	return &corev1.Node{
+		ObjectMeta: metav1.ObjectMeta{Name: name, Labels: labels},
+		Status:     corev1.NodeStatus{Capacity: capacity},
+	}
+}
+
+func reconcileAndGet(t *testing.T, r *NodeLabelerReconciler, name string) *corev1.Node {
+	t.Helper()
+	_, err := r.Reconcile(context.Background(), ctrl.Request{NamespacedName: types.NamespacedName{Name: name}})
+	require.NoError(t, err)
+	var got corev1.Node
+	require.NoError(t, r.Get(context.Background(), types.NamespacedName{Name: name}, &got))
+	return &got
+}
+
+// (a) node with nvidia.com/gpu capacity -> gpu label added, neuron label absent.
+func TestNodeLabeler_GPUCapacityAddsLabel(t *testing.T) {
+	r := newTestNodeLabelerReconciler(node("gpu-node", nil, corev1.ResourceList{
+		nvidiaGPUResource:  resource.MustParse("1"),
+		corev1.ResourceCPU: resource.MustParse("8"),
+	}))
+	got := reconcileAndGet(t, r, "gpu-node")
+	assert.Equal(t, "true", got.Labels[GPUPresentLabel])
+	_, hasNeuron := got.Labels[NeuronPresentLabel]
+	assert.False(t, hasNeuron)
+}
+
+// Presence-of-key semantics: a GPU node whose device plugin died keeps the
+// capacity key at quantity 0 and must still be labeled.
+func TestNodeLabeler_GPUCapacityZeroQuantityStillLabels(t *testing.T) {
+	r := newTestNodeLabelerReconciler(node("gpu-unhealthy", nil, corev1.ResourceList{
+		nvidiaGPUResource: resource.MustParse("0"),
+	}))
+	got := reconcileAndGet(t, r, "gpu-unhealthy")
+	assert.Equal(t, "true", got.Labels[GPUPresentLabel])
+}
+
+// (b) node with no relevant capacity -> no managed labels.
+func TestNodeLabeler_NoRelevantCapacityNoLabels(t *testing.T) {
+	r := newTestNodeLabelerReconciler(node("plain-node", nil, corev1.ResourceList{
+		corev1.ResourceCPU:    resource.MustParse("4"),
+		corev1.ResourceMemory: resource.MustParse("16Gi"),
+	}))
+	got := reconcileAndGet(t, r, "plain-node")
+	_, hasGPU := got.Labels[GPUPresentLabel]
+	_, hasNeuron := got.Labels[NeuronPresentLabel]
+	assert.False(t, hasGPU)
+	assert.False(t, hasNeuron)
+}
+
+// (c) node carrying our gpu label but capacity key gone -> label removed.
+func TestNodeLabeler_StaleGPULabelRemoved(t *testing.T) {
+	r := newTestNodeLabelerReconciler(node("was-gpu", map[string]string{
+		GPUPresentLabel: "true",
+	}, corev1.ResourceList{
+		corev1.ResourceCPU: resource.MustParse("4"),
+	}))
+	got := reconcileAndGet(t, r, "was-gpu")
+	_, hasGPU := got.Labels[GPUPresentLabel]
+	assert.False(t, hasGPU)
+}
+
+// (d) node with a neuron capacity key -> neuron label added.
+func TestNodeLabeler_NeuronCapacityAddsLabel(t *testing.T) {
+	for _, res := range []corev1.ResourceName{neuronResource, neuronCoreResource, neuronDeviceResource} {
+		r := newTestNodeLabelerReconciler(node("neuron-node", nil, corev1.ResourceList{
+			res: resource.MustParse("1"),
+		}))
+		got := reconcileAndGet(t, r, "neuron-node")
+		assert.Equal(t, "true", got.Labels[NeuronPresentLabel], "resource %s should set neuron label", res)
+		_, hasGPU := got.Labels[GPUPresentLabel]
+		assert.False(t, hasGPU)
+	}
+}
+
+// (e) already-correct node -> reconcile is a no-op: no error, labels unchanged,
+// and no API write (resourceVersion unchanged).
+func TestNodeLabeler_AlreadyCorrectIsNoOp(t *testing.T) {
+	r := newTestNodeLabelerReconciler(node("gpu-labeled", map[string]string{
+		GPUPresentLabel: "true",
+	}, corev1.ResourceList{
+		nvidiaGPUResource: resource.MustParse("1"),
+	}))
+	var before corev1.Node
+	require.NoError(t, r.Get(context.Background(), types.NamespacedName{Name: "gpu-labeled"}, &before))
+
+	got := reconcileAndGet(t, r, "gpu-labeled")
+	assert.Equal(t, "true", got.Labels[GPUPresentLabel])
+	assert.Equal(t, before.ResourceVersion, got.ResourceVersion, "no-op reconcile must not write")
+}
+
+// Other existing labels must survive the metadata.labels merge patch.
+func TestNodeLabeler_PreservesOtherLabels(t *testing.T) {
+	r := newTestNodeLabelerReconciler(node("mixed", map[string]string{
+		"kubernetes.io/hostname":      "mixed",
+		"node.kubernetes.io/instance": "keep-me",
+	}, corev1.ResourceList{
+		nvidiaGPUResource: resource.MustParse("1"),
+	}))
+	got := reconcileAndGet(t, r, "mixed")
+	assert.Equal(t, "true", got.Labels[GPUPresentLabel])
+	assert.Equal(t, "mixed", got.Labels["kubernetes.io/hostname"])
+	assert.Equal(t, "keep-me", got.Labels["node.kubernetes.io/instance"])
+}
+
+// (f) predicate: an Update with unchanged capacity+labels is filtered out, while
+// an Update that adds a relevant capacity key passes.
+func TestNodeLabeler_Predicate(t *testing.T) {
+	pred := nodeCapacityOrLabelPredicate()
+
+	base := node("n", map[string]string{GPUPresentLabel: "true"}, corev1.ResourceList{
+		nvidiaGPUResource: resource.MustParse("1"),
+	})
+
+	// Unchanged relevant fields (only an irrelevant condition/heartbeat differs).
+	unchangedNew := base.DeepCopy()
+	unchangedNew.Status.Allocatable = corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("4")}
+	assert.False(t, pred.Update(event.UpdateEvent{ObjectOld: base, ObjectNew: unchangedNew}))
+
+	// New relevant capacity key appears.
+	changedNew := base.DeepCopy()
+	changedNew.Status.Capacity[neuronResource] = resource.MustParse("1")
+	assert.True(t, pred.Update(event.UpdateEvent{ObjectOld: base, ObjectNew: changedNew}))
+
+	// A managed label change also passes.
+	labelChangedNew := base.DeepCopy()
+	labelChangedNew.Labels = map[string]string{}
+	assert.True(t, pred.Update(event.UpdateEvent{ObjectOld: base, ObjectNew: labelChangedNew}))
+
+	// Create always reconciles; Delete/Generic never do.
+	assert.True(t, pred.Create(event.CreateEvent{Object: base}))
+	assert.False(t, pred.Delete(event.DeleteEvent{Object: base}))
+	assert.False(t, pred.Generic(event.GenericEvent{Object: base}))
+}