diff --git a/pkg/defaults/timeouts.go b/pkg/defaults/timeouts.go index daeea3b9b..9dd457a28 100644 --- a/pkg/defaults/timeouts.go +++ b/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 diff --git a/validators/performance/consts.go b/validators/performance/consts.go index 097cdac24..8bdd93169 100644 --- a/validators/performance/consts.go +++ b/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 diff --git a/validators/performance/inference_perf_constraint.go b/validators/performance/inference_perf_constraint.go index d4fa5f883..9da3fab94 100644 --- a/validators/performance/inference_perf_constraint.go +++ b/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 diff --git a/validators/performance/nccl_all_reduce_bw_constraint.go b/validators/performance/nccl_all_reduce_bw_constraint.go index 8d2297ba8..92365f519 100644 --- a/validators/performance/nccl_all_reduce_bw_constraint.go +++ b/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 diff --git a/validators/performance/nccl_test.go b/validators/performance/nccl_test.go index 5db706ac9..17c1c1e3a 100644 --- a/validators/performance/nccl_test.go +++ b/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) } } } diff --git a/validators/performance/trainer_lifecycle.go b/validators/performance/trainer_lifecycle.go index d51ec0b5d..0f17f267b 100644 --- a/validators/performance/trainer_lifecycle.go +++ b/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 } diff --git a/validators/performance/trainer_lifecycle_test.go b/validators/performance/trainer_lifecycle_test.go index b03cfe7ef..6e85f7fd1 100644 --- a/validators/performance/trainer_lifecycle_test.go +++ b/validators/performance/trainer_lifecycle_test.go @@ -15,8 +15,13 @@ package main import ( + "fmt" "strings" "testing" + + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "sigs.k8s.io/kustomize/api/hasher" + "sigs.k8s.io/kustomize/api/resource" ) func TestRewriteJobSetStagingImage(t *testing.T) { @@ -76,3 +81,281 @@ 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 []any) *unstructured.Unstructured { + podSpec := map[string]any{ + "containers": []any{ + map[string]any{"name": "manager", "image": "example/manager:latest"}, + }, + } + if existingTolerations != nil { + podSpec["tolerations"] = existingTolerations + } + return &unstructured.Unstructured{Object: map[string]any{ + "apiVersion": "apps/v1", + "kind": "Deployment", + "metadata": map[string]any{"name": name}, + "spec": map[string]any{ + "template": map[string]any{ + "spec": podSpec, + }, + }, + }} +} + +// TestApplyControllerTolerations covers both controller names, the two +// mutation-failure paths, and that an unrelated Deployment is left untouched. +// TestApplyControllerTolerations_Isolation pins that the two controllers do +// not share a live toleration slice: mutating one Deployment's stamped +// tolerations in place must not affect the other, or the shared +// controllerTolerateAll package-level global. +func TestApplyControllerTolerations_Isolation(t *testing.T) { + trainerObj := deploymentFixture(trainerControllerDeployment, nil) + jobSetObj := deploymentFixture(jobSetControllerDeployment, nil) + + if err := applyControllerTolerations(trainerObj); err != nil { + t.Fatalf("applyControllerTolerations(trainer) error = %v", err) + } + if err := applyControllerTolerations(jobSetObj); err != nil { + t.Fatalf("applyControllerTolerations(jobset) error = %v", err) + } + + trainerTols, _, _ := unstructured.NestedSlice(trainerObj.Object, "spec", "template", "spec", "tolerations") + trainerTol, _ := trainerTols[0].(map[string]any) + trainerTol["key"] = "mutated-for-trainer-only" + + jobSetTols, _, _ := unstructured.NestedSlice(jobSetObj.Object, "spec", "template", "spec", "tolerations") + jobSetTol, _ := jobSetTols[0].(map[string]any) + if _, mutated := jobSetTol["key"]; mutated { + t.Errorf("mutating the Trainer Deployment's toleration leaked into the JobSet Deployment: %v", jobSetTol) + } + if _, mutated := controllerTolerateAll[0].(map[string]any)["key"]; mutated { + t.Errorf("mutating a stamped toleration leaked into the shared controllerTolerateAll global: %v", controllerTolerateAll[0]) + } +} + +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 []any + }{ + { + name: "Trainer controller Deployment with no tolerations gets tolerate-all", + obj: deploymentFixture(trainerControllerDeployment, nil), + wantTolerations: []any{ + map[string]any{"operator": "Exists"}, + }, + }, + { + name: "JobSet controller Deployment with no tolerations gets tolerate-all", + obj: deploymentFixture(jobSetControllerDeployment, nil), + wantTolerations: []any{ + map[string]any{"operator": "Exists"}, + }, + }, + { + name: "Deployment with existing tolerations is left untouched", + obj: deploymentFixture(trainerControllerDeployment, []any{ + map[string]any{"key": "dedicated", "operator": "Equal", "value": "trainer", "effect": "NoSchedule"}, + }), + wantTolerations: []any{ + map[string]any{"key": "dedicated", "operator": "Equal", "value": "trainer", "effect": "NoSchedule"}, + }, + }, + { + name: "non-controller Deployment is left untouched", + obj: deploymentFixture("some-other-deployment", nil), + wantTolerations: nil, + }, + { + // found && len(existing) > 0 is the guard: an explicit empty slice is + // "present but empty," which must fall through to being stamped, not + // be treated the same as "already tolerated." Pins this boundary + // against a future refactor that flips the guard to `found` alone. + name: "Deployment with present-but-empty tolerations gets tolerate-all", + obj: deploymentFixture(trainerControllerDeployment, []any{}), + wantTolerations: []any{ + map[string]any{"operator": "Exists"}, + }, + }, + { + name: "non-Deployment resource is left untouched", + obj: &unstructured.Unstructured{Object: map[string]any{ + "apiVersion": "v1", + "kind": "Service", + "metadata": map[string]any{"name": trainerControllerDeployment}, + "spec": map[string]any{}, + }}, + wantTolerations: nil, + }, + { + name: "Deployment-kind resource in a non-apps group is left untouched", + obj: &unstructured.Unstructured{Object: map[string]any{ + "apiVersion": "example.com/v1", + "kind": "Deployment", + "metadata": map[string]any{"name": trainerControllerDeployment}, + "spec": map[string]any{}, + }}, + wantTolerations: nil, + }, + { + name: "missing pod spec fails closed", + obj: &unstructured.Unstructured{Object: map[string]any{ + "apiVersion": "apps/v1", + "kind": "Deployment", + "metadata": map[string]any{"name": trainerControllerDeployment}, + "spec": map[string]any{}, + }}, + wantErr: true, + }, + { + name: "malformed tolerations field fails closed", + obj: &unstructured.Unstructured{Object: map[string]any{ + "apiVersion": "apps/v1", + "kind": "Deployment", + "metadata": map[string]any{"name": trainerControllerDeployment}, + "spec": map[string]any{ + "template": map[string]any{ + "spec": map[string]any{ + // 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]any) + wantTol, _ := tt.wantTolerations[i].(map[string]any) + for k, v := range wantTol { + if gotTol[k] != v { + t.Errorf("toleration[%d][%q] = %v, want %v", i, k, gotTol[k], v) + } + } + } + }) + } +} + +// TestDecodeTrainerObjects exercises decodeTrainerObjects end to end: the +// toleration lands only on a controller Deployment resource, not on an +// unrelated resource in the same manifest set, and the Kind=="" skip ordering +// above the applyControllerTolerations call does not interfere. +func TestDecodeTrainerObjects(t *testing.T) { + rf := resource.NewFactory(&hasher.Hasher{}) + + manifest := fmt.Sprintf(`apiVersion: apps/v1 +kind: Deployment +metadata: + name: %s +spec: + template: + spec: + containers: + - name: manager + image: example/trainer:latest +--- +apiVersion: v1 +kind: ConfigMap +metadata: + name: some-config +data: + foo: bar +`, trainerControllerDeployment) + + resources, err := rf.SliceFromBytes([]byte(manifest)) + if err != nil { + t.Fatalf("SliceFromBytes() error = %v", err) + } + + objs, err := decodeTrainerObjects(resources) + if err != nil { + t.Fatalf("decodeTrainerObjects() error = %v", err) + } + if len(objs) != 2 { + t.Fatalf("got %d decoded object(s), want 2", len(objs)) + } + + var trainerObj, unrelatedObj *unstructured.Unstructured + for _, o := range objs { + switch o.GetName() { + case trainerControllerDeployment: + trainerObj = o + case "some-config": + unrelatedObj = o + } + } + if trainerObj == nil { + t.Fatal("Trainer controller Deployment missing from decoded objects") + } + tols, found, _ := unstructured.NestedSlice(trainerObj.Object, "spec", "template", "spec", "tolerations") + if !found || len(tols) != 1 { + t.Errorf("Trainer Deployment tolerations = %v, want a single blanket tolerate-all entry", tols) + } + + if unrelatedObj == nil { + t.Fatal("unrelated ConfigMap missing from decoded objects") + } + if _, found, _ := unstructured.NestedSlice(unrelatedObj.Object, "spec", "template", "spec", "tolerations"); found { + t.Error("unrelated ConfigMap must not receive a toleration") + } +} + +// TestDecodeTrainerObjects_PropagatesTolerationError verifies a malformed +// resource that fails applyControllerTolerations aborts decoding rather than +// leaving a partial object list. +func TestDecodeTrainerObjects_PropagatesTolerationError(t *testing.T) { + rf := resource.NewFactory(&hasher.Hasher{}) + + manifest := fmt.Sprintf(`apiVersion: apps/v1 +kind: Deployment +metadata: + name: %s +spec: + template: + spec: + tolerations: not-a-slice +`, trainerControllerDeployment) + + resources, err := rf.SliceFromBytes([]byte(manifest)) + if err != nil { + t.Fatalf("SliceFromBytes() error = %v", err) + } + + if _, err := decodeTrainerObjects(resources); err == nil { + t.Fatal("decodeTrainerObjects() expected error for a malformed tolerations field, got nil") + } +}