fix(validator): tolerate all taints on Trainer/JobSet controller Deployments

The Kubeflow Trainer/JobSet controller-manager Deployments ship with no
tolerations. On a cluster where every node pool carries a taint (e.g. an
arch-tainted GPU pool plus a system pool GKE reserves for its own managed
components once no untainted pool remains), the controllers have nowhere
to schedule and installTrainer times out waiting for a Deployment that can
never become Ready.

applyControllerTolerations stamps a blanket tolerate-all onto the Trainer
and JobSet controller-manager Deployments specifically (by name) when
either has no existing tolerations; a Deployment that already declares
tolerations, or any other Deployment in the manifest set, is left
untouched. Scoping by name rather than by Kind alone matters here: this
is called for every Deployment decoded from the installer's manifest set,
and a future addition to that set must not silently inherit a blanket
{operator: Exists} it never asked for.

Extract the repeated "operator" toleration-key literal into keyOperator to
satisfy golangci-lint's goconst threshold across the package.

Signed-off-by: Mike Cook <micook@nvidia.com>

Author: mikecook

validators/performance/consts.go

@@ -21,6 +21,7 @@ const (
 	versionV1alpha1          = "v1alpha1"
 	versionV1beta1           = "v1beta1"
 	keyName                  = "name"
+	keyOperator              = "operator"
 	checkNameNCCLAllReduceBW = "nccl-all-reduce-bw"
 
 	// nodeJobName is the name of both the NCCL worker replicatedJob and its

validators/performance/inference_perf_constraint.go

@@ -1850,7 +1850,7 @@ func tolerationsToUnstructured(tolerations []v1.Toleration) []interface{} {
 	tolList := make([]interface{}, 0, len(tolerations))
 	for _, t := range tolerations {
 		tolMap := map[string]interface{}{
-			"operator": string(t.Operator),
+			keyOperator: string(t.Operator),
 		}
 		if t.Key != "" {
 			tolMap["key"] = t.Key

validators/performance/nccl_all_reduce_bw_constraint.go

@@ -1451,7 +1451,7 @@ func applyNCCLWorkerScheduling(obj *unstructured.Unstructured, nodeSelector map[
 			tolList := make([]interface{}, 0, len(tolerations))
 			for _, t := range tolerations {
 				tolMap := map[string]interface{}{
-					"operator": string(t.Operator),
+					keyOperator: string(t.Operator),
 				}
 				if t.Key != "" {
 					tolMap["key"] = t.Key

validators/performance/trainer_lifecycle.go

@@ -67,6 +67,11 @@ const (
 	// trainerControllerDeployment is the Deployment name for the Trainer controller-manager.
 	trainerControllerDeployment = "kubeflow-trainer-controller-manager"
 
+	// jobSetControllerDeployment is the JobSet controller-manager Deployment name
+	// emitted by this package's kustomize overlay (see jobSetNameLabel for why
+	// the Helm chart's release-derived name doesn't apply here).
+	jobSetControllerDeployment = "jobset-controller-manager"
+
 	// trainerControllerService is the Service fronting the controller-manager's
 	// webhook port. Without it the admission webhooks have no endpoints and every
 	// TrainJob create is rejected.
@@ -158,6 +163,43 @@ const (
 	jobSetPromotedImageRepo = "registry.k8s.io/jobset/jobset"
 )
 
+// controllerTolerateAll lets a Trainer/JobSet controller-manager Deployment
+// schedule on any node pool, regardless of taints.
+var controllerTolerateAll = []interface{}{
+	map[string]interface{}{keyOperator: "Exists"},
+}
+
+// applyControllerTolerations stamps controllerTolerateAll onto the Trainer and
+// JobSet controller-manager Deployments' pod template, unless one already
+// declares tolerations. Scoped to those two names so an unrelated Deployment
+// in the manifest set never gets a blanket toleration it didn't ask for.
+func applyControllerTolerations(obj *unstructured.Unstructured) error {
+	if obj.GroupVersionKind().Kind != "Deployment" {
+		return nil
+	}
+	switch obj.GetName() {
+	case trainerControllerDeployment, jobSetControllerDeployment:
+	default:
+		return nil
+	}
+
+	if existing, found, err := unstructured.NestedSlice(obj.Object, "spec", "template", "spec", "tolerations"); err != nil {
+		return aicrErrors.Wrap(aicrErrors.ErrCodeInternal,
+			fmt.Sprintf("failed to read tolerations from Deployment %q", obj.GetName()), err)
+	} else if found && len(existing) > 0 {
+		return nil
+	}
+
+	podSpec, found := nestedMap(obj.Object, "spec", "template", "spec")
+	if !found {
+		return aicrErrors.New(aicrErrors.ErrCodeInternal,
+			fmt.Sprintf("pod spec not found in Deployment %q", obj.GetName()))
+	}
+	podSpec["tolerations"] = controllerTolerateAll
+	slog.Info("Applying blanket toleration to controller Deployment", "name", obj.GetName())
+	return nil
+}
+
 // GVRs for the objects the Trainer lifecycle probes and waits on.
 var (
 	trainerCRDGVR = schema.GroupVersionResource{
@@ -806,6 +848,9 @@ func decodeTrainerObjects(resources []*resource.Resource) ([]*unstructured.Unstr
 		if obj.GroupVersionKind().Kind == "" {
 			continue
 		}
+		if tolErr := applyControllerTolerations(obj); tolErr != nil {
+			return nil, tolErr
+		}
 		objs = append(objs, obj)
 	}
 	return objs, nil

validators/performance/trainer_lifecycle_test.go

@@ -17,6 +17,8 @@ package main
 import (
 	"strings"
 	"testing"
+
+	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 )
 
 func TestRewriteJobSetStagingImage(t *testing.T) {
@@ -76,3 +78,141 @@ func TestRewriteJobSetStagingImage_PreservesTag(t *testing.T) {
 		t.Errorf("got %q, want %q", got, want)
 	}
 }
+
+// deploymentFixture returns a minimal unstructured Deployment, optionally with an
+// existing tolerations list, for exercising applyControllerTolerations.
+func deploymentFixture(name string, existingTolerations []interface{}) *unstructured.Unstructured {
+	podSpec := map[string]interface{}{
+		"containers": []interface{}{
+			map[string]interface{}{"name": "manager", "image": "example/manager:latest"},
+		},
+	}
+	if existingTolerations != nil {
+		podSpec["tolerations"] = existingTolerations
+	}
+	return &unstructured.Unstructured{Object: map[string]interface{}{
+		"apiVersion": "apps/v1",
+		"kind":       "Deployment",
+		"metadata":   map[string]interface{}{"name": name},
+		"spec": map[string]interface{}{
+			"template": map[string]interface{}{
+				"spec": podSpec,
+			},
+		},
+	}}
+}
+
+// TestApplyControllerTolerations covers both controller names, the two
+// mutation-failure paths, and that an unrelated Deployment is left untouched.
+func TestApplyControllerTolerations(t *testing.T) {
+	tests := []struct {
+		name    string
+		obj     *unstructured.Unstructured
+		wantErr bool
+		// wantTolerations is checked only when wantErr is false. nil means "the
+		// tolerations field must not be present at all" (untouched, not merely
+		// empty).
+		wantTolerations []interface{}
+	}{
+		{
+			name: "Trainer controller Deployment with no tolerations gets tolerate-all",
+			obj:  deploymentFixture(trainerControllerDeployment, nil),
+			wantTolerations: []interface{}{
+				map[string]interface{}{"operator": "Exists"},
+			},
+		},
+		{
+			name: "JobSet controller Deployment with no tolerations gets tolerate-all",
+			obj:  deploymentFixture(jobSetControllerDeployment, nil),
+			wantTolerations: []interface{}{
+				map[string]interface{}{"operator": "Exists"},
+			},
+		},
+		{
+			name: "Deployment with existing tolerations is left untouched",
+			obj: deploymentFixture(trainerControllerDeployment, []interface{}{
+				map[string]interface{}{"key": "dedicated", "operator": "Equal", "value": "trainer", "effect": "NoSchedule"},
+			}),
+			wantTolerations: []interface{}{
+				map[string]interface{}{"key": "dedicated", "operator": "Equal", "value": "trainer", "effect": "NoSchedule"},
+			},
+		},
+		{
+			name:            "non-controller Deployment is left untouched",
+			obj:             deploymentFixture("some-other-deployment", nil),
+			wantTolerations: nil,
+		},
+		{
+			name: "non-Deployment resource is left untouched",
+			obj: &unstructured.Unstructured{Object: map[string]interface{}{
+				"apiVersion": "v1",
+				"kind":       "Service",
+				"metadata":   map[string]interface{}{"name": trainerControllerDeployment},
+				"spec":       map[string]interface{}{},
+			}},
+			wantTolerations: nil,
+		},
+		{
+			name: "missing pod spec fails closed",
+			obj: &unstructured.Unstructured{Object: map[string]interface{}{
+				"apiVersion": "apps/v1",
+				"kind":       "Deployment",
+				"metadata":   map[string]interface{}{"name": trainerControllerDeployment},
+				"spec":       map[string]interface{}{},
+			}},
+			wantErr: true,
+		},
+		{
+			name: "malformed tolerations field fails closed",
+			obj: &unstructured.Unstructured{Object: map[string]interface{}{
+				"apiVersion": "apps/v1",
+				"kind":       "Deployment",
+				"metadata":   map[string]interface{}{"name": trainerControllerDeployment},
+				"spec": map[string]interface{}{
+					"template": map[string]interface{}{
+						"spec": map[string]interface{}{
+							// A string, not a slice: NestedSlice's type assertion fails.
+							"tolerations": "not-a-slice",
+						},
+					},
+				},
+			}},
+			wantErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := applyControllerTolerations(tt.obj)
+			if (err != nil) != tt.wantErr {
+				t.Fatalf("error = %v, wantErr %v", err, tt.wantErr)
+			}
+			if tt.wantErr {
+				return
+			}
+
+			got, found, _ := unstructured.NestedSlice(tt.obj.Object, "spec", "template", "spec", "tolerations")
+			if tt.wantTolerations == nil {
+				if found {
+					t.Errorf("expected no tolerations field, got %v", got)
+				}
+				return
+			}
+			if !found {
+				t.Fatalf("expected tolerations %v, found none", tt.wantTolerations)
+			}
+			if len(got) != len(tt.wantTolerations) {
+				t.Fatalf("got %d toleration(s) %v, want %d %v", len(got), got, len(tt.wantTolerations), tt.wantTolerations)
+			}
+			for i := range got {
+				gotTol, _ := got[i].(map[string]interface{})
+				wantTol, _ := tt.wantTolerations[i].(map[string]interface{})
+				for k, v := range wantTol {
+					if gotTol[k] != v {
+						t.Errorf("toleration[%d][%q] = %v, want %v", i, k, gotTol[k], v)
+					}
+				}
+			}
+		})
+	}
+}