fix(validator): tolerate taints, widen Trainer readiness timeout

pkg/defaults/timeouts.go

@@ -666,8 +666,15 @@ const (
 	TrainerCRDEstablishedTimeout = 2 * time.Minute
 
 	// TrainerControllerReadyTimeout is the time to wait for the Kubeflow Trainer
-	// controller-manager Deployment to have at least one ready replica after installation.
-	TrainerControllerReadyTimeout = 2 * time.Minute
+	// controller-manager Deployment to have at least one ready replica after
+	// installation. Widened from 2m to 3m: on cold start the cert-controller
+	// sidecar's webhook-cert get-or-create can race a not-yet-synced informer
+	// cache, producing a resourceVersion conflict that the sidecar's own
+	// reconcile loop retries and self-heals from unassisted. This is expected
+	// behavior under cert-controller's optimistic-concurrency retry, not a
+	// defect in Trainer or in this validator — but each retry adds latency
+	// that could otherwise push first-ready past a tighter budget.
+	TrainerControllerReadyTimeout = 3 * time.Minute
 
 	// TrainerInstallPollInterval is the sleep between checks that a
 	// recipe-declared Kubeflow Trainer installation has become complete. The

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

@@ -1823,7 +1823,7 @@ func tolerationsToUnstructured(tolerations []v1.Toleration) []any {
 	tolList := make([]any, 0, len(tolerations))
 	for _, t := range tolerations {
 		tolMap := map[string]any{
-			"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([]any, 0, len(tolerations))
 			for _, t := range tolerations {
 				tolMap := map[string]any{
-					"operator": string(t.Operator),
+					keyOperator: string(t.Operator),
 				}
 				if t.Key != "" {
 					tolMap["key"] = t.Key
@@ -1479,7 +1479,7 @@ func applyNCCLWorkerScheduling(obj *unstructured.Unstructured, nodeSelector map[
 	return unstructured.SetNestedSlice(obj.Object, replicatedJobs, "spec", "template", "spec", "replicatedJobs")
 }
 
-// nestedMap navigates a chain of string keys through nested map[string]interface{} values.
+// nestedMap navigates a chain of string keys through nested map[string]any values.
 // Returns the target map and true if found, nil and false otherwise.
 func nestedMap(m map[string]any, keys ...string) (map[string]any, bool) {
 	current := m

validators/performance/nccl_test.go

@@ -142,8 +142,8 @@ func TestApplyNCCLWorkerScheduling_Tolerations(t *testing.T) {
 			t.Fatalf("tolerations count = %d, want 1", len(tolsRaw))
 		}
 		tol, _ := tolsRaw[0].(map[string]any)
-		if tol["key"] != "gpu-type" || tol["value"] != "h100" || tol["effect"] != "NoSchedule" {
-			t.Errorf("toleration = %v, want gpu-type=h100:NoSchedule", tol)
+		if tol["key"] != "gpu-type" || tol["value"] != "h100" || tol["effect"] != "NoSchedule" || tol["operator"] != "Equal" {
+			t.Errorf("toleration = %v, want gpu-type=h100:NoSchedule operator=Equal", tol)
 		}
 	}
 }
@@ -209,8 +209,8 @@ func TestApplyNCCLWorkerScheduling_Both(t *testing.T) {
 			t.Fatalf("tolerations count = %d, want 1", len(tolsRaw))
 		}
 		tol, _ := tolsRaw[0].(map[string]any)
-		if tol["key"] != "custom-taint" || tol["value"] != "true" || tol["effect"] != "NoSchedule" {
-			t.Errorf("toleration = %v, want custom-taint=true:NoSchedule", tol)
+		if tol["key"] != "custom-taint" || tol["value"] != "true" || tol["effect"] != "NoSchedule" || tol["operator"] != "Equal" {
+			t.Errorf("toleration = %v, want custom-taint=true:NoSchedule operator=Equal", tol)
 		}
 	}
 }

validators/performance/trainer_lifecycle.go

@@ -32,8 +32,10 @@ import (
 	"strings"
 	"time"
 
+	"github.com/NVIDIA/aicr/pkg/component"
 	"github.com/NVIDIA/aicr/pkg/defaults"
 	aicrErrors "github.com/NVIDIA/aicr/pkg/errors"
+	corev1 "k8s.io/api/core/v1"
 	k8serrors "k8s.io/apimachinery/pkg/api/errors"
 	apimeta "k8s.io/apimachinery/pkg/api/meta"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -67,6 +69,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.
@@ -159,6 +166,76 @@ const (
 	jobSetPromotedImageRepo = "registry.k8s.io/jobset/jobset"
 )
 
+// controllerTolerateAll lets a Trainer/JobSet controller-manager Deployment
+// schedule on any node pool, regardless of taints. Built through
+// component.TolerationsToPodSpec, the same converter used for Helm-values
+// toleration overrides, so there is one canonical place that knows the
+// toleration-to-map shape.
+var controllerTolerateAll = tolerationsToAnySlice(
+	component.TolerationsToPodSpec([]corev1.Toleration{{Operator: corev1.TolerationOpExists}}),
+)
+
+// tolerationsToAnySlice widens []map[string]any to []any: unstructured pod
+// specs (podSpec["tolerations"]) must hold []any, not []map[string]any, to
+// match how NestedSlice reads and how JSON round-tripping serializes it.
+func tolerationsToAnySlice(tolerations []map[string]any) []any {
+	result := make([]any, len(tolerations))
+	for i, t := range tolerations {
+		result[i] = t
+	}
+	return result
+}
+
+// cloneControllerTolerateAll returns a fresh copy of controllerTolerateAll.
+// It is stamped onto both the Trainer and JobSet Deployments, so a per-call
+// copy keeps a future in-place edit of one controller's tolerations from
+// silently corrupting the other Deployment's live object, or the shared
+// process-wide global itself.
+func cloneControllerTolerateAll() []any {
+	cloned := make([]any, len(controllerTolerateAll))
+	for i, t := range controllerTolerateAll {
+		m, _ := t.(map[string]any)
+		clonedMap := make(map[string]any, len(m))
+		for k, v := range m {
+			clonedMap[k] = v
+		}
+		cloned[i] = clonedMap
+	}
+	return cloned
+}
+
+// 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 gvk := obj.GroupVersionKind(); gvk.Kind != "Deployment" || gvk.Group != "apps" {
+		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 {
+		slog.Debug("Controller Deployment already declares tolerations; leaving untouched", "name", obj.GetName())
+		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"] = cloneControllerTolerateAll()
+	slog.Info("Applying blanket toleration to controller Deployment", "name", obj.GetName(), "namespace", obj.GetNamespace())
+	return nil
+}
+
 // GVRs for the objects the Trainer lifecycle probes and waits on.
 var (
 	trainerCRDGVR = schema.GroupVersionResource{
@@ -791,6 +868,7 @@ func installTrainer(ctx context.Context, dynamicClient dynamic.Interface, discov
 // before the first apply so a malformed manifest cannot leave a partial install.
 func decodeTrainerObjects(resources []*resource.Resource) ([]*unstructured.Unstructured, error) {
 	objs := make([]*unstructured.Unstructured, 0, len(resources))
+	seenControllers := make(map[string]bool, 2)
 	for _, res := range resources {
 		// Convert to unstructured via YAML round-trip (guarantees plain Go types).
 		yamlBytes, err := res.AsYAML()
@@ -807,8 +885,27 @@ func decodeTrainerObjects(resources []*resource.Resource) ([]*unstructured.Unstr
 		if obj.GroupVersionKind().Kind == "" {
 			continue
 		}
+		if tolErr := applyControllerTolerations(obj); tolErr != nil {
+			return nil, tolErr
+		}
+		if obj.GroupVersionKind().Kind == "Deployment" {
+			seenControllers[obj.GetName()] = true
+		}
 		objs = append(objs, obj)
 	}
+
+	// A future Trainer archive bump (or a JobSet overlay variant) that renames
+	// either controller Deployment makes applyControllerTolerations's name
+	// switch miss it silently, reverting that controller to unschedulable on
+	// an all-tainted cluster. Surface the mismatch here instead of letting it
+	// resurface only as a bare readiness timeout downstream.
+	for _, name := range []string{trainerControllerDeployment, jobSetControllerDeployment} {
+		if !seenControllers[name] {
+			slog.Warn("Expected controller Deployment not found in Trainer manifest set; "+
+				"it will not receive the blanket toleration and may be unschedulable on all-tainted clusters",
+				"deployment", name)
+		}
+	}
 	return objs, nil
 }