diff --git a/config/rbac/role.yaml b/config/rbac/role.yaml index 9587769ba..f963572e0 100644 --- a/config/rbac/role.yaml +++ b/config/rbac/role.yaml @@ -34,6 +34,15 @@ rules: - list - patch - watch +- apiGroups: + - "" + resources: + - nodes + verbs: + - get + - list + - patch + - watch - apiGroups: - apps resources: diff --git a/controllers/nodelabeler_controller.go b/controllers/nodelabeler_controller.go new file mode 100644 index 000000000..969e74ce8 --- /dev/null +++ b/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) +} diff --git a/controllers/nodelabeler_controller_test.go b/controllers/nodelabeler_controller_test.go new file mode 100644 index 000000000..7bc37a9ec --- /dev/null +++ b/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})) +} diff --git a/main.go b/main.go index 0f218fae2..9d71d1637 100644 --- a/main.go +++ b/main.go @@ -159,6 +159,7 @@ func main() { dcgmExporterImage string neuronMonitorImage string targetAllocatorImage string + enableNodeCapacityLabeler bool ) pflag.StringVar(&metricsAddr, "metrics-addr", ":8080", "The address the metric endpoint binds to.") @@ -175,6 +176,7 @@ func main() { stringFlagOrEnv(&dcgmExporterImage, "dcgm-exporter-image", "RELATED_IMAGE_DCGM_EXPORTER", fmt.Sprintf("%s:%s", dcgmExporterImageRepository, v.DcgmExporter), "The default DCGM Exporter image. This image is used when no image is specified in the CustomResource.") stringFlagOrEnv(&neuronMonitorImage, "neuron-monitor-image", "RELATED_IMAGE_NEURON_MONITOR", fmt.Sprintf("%s:%s", neuronMonitorImageRepository, v.NeuronMonitor), "The default Neuron monitor image. This image is used when no image is specified in the CustomResource.") stringFlagOrEnv(&targetAllocatorImage, "target-allocator-image", "RELATED_IMAGE_TARGET_ALLOCATOR", fmt.Sprintf("%s:%s", targetAllocatorImageRepository, v.TargetAllocator), "The default AmazonCloudWatchAgent target allocator image. This image is used when no image is specified in the CustomResource.") + pflag.BoolVar(&enableNodeCapacityLabeler, "enable-node-capacity-labeler", true, "Enable the controller that labels nodes with cloudwatch.aws.amazon.com/gpu.present and neuron.present based on advertised accelerator capacity.") pflag.Parse() // set instrumentation cpu and memory limits in environment variables to be used for default instrumentation; default values received from https://github.com/open-telemetry/opentelemetry-operator/blob/main/apis/v1alpha1/instrumentation_webhook.go @@ -264,6 +266,7 @@ func main() { TLSOpts: optionsTlSOptsFuncs, }), Cache: cache.Options{ + // DefaultNamespaces only restricts namespaced kinds; cluster-scoped kinds like Node are always watched cluster-wide, so no ByObject entry is needed. DefaultNamespaces: namespaces, }, } @@ -309,6 +312,17 @@ func main() { os.Exit(1) } + if enableNodeCapacityLabeler { + if err = (&controllers.NodeLabelerReconciler{ + Client: mgr.GetClient(), + Scheme: mgr.GetScheme(), + Log: ctrl.Log.WithName("controllers").WithName("NodeLabeler"), + }).SetupWithManager(mgr); err != nil { + setupLog.Error(err, "unable to create controller", "controller", "NodeLabeler") + os.Exit(1) + } + } + decoder := admission.NewDecoder(mgr.GetScheme()) instrumentationAnnotator := auto.CreateInstrumentationAnnotator(autoMonitorConfigStr, autoAnnotationConfigStr, ctx, mgr.GetClient(), mgr.GetAPIReader(), setupLog)