diff --git a/operator/e2e/tests/scheduler_test.go b/operator/e2e/tests/scheduler_test.go new file mode 100644 index 000000000..298255f31 --- /dev/null +++ b/operator/e2e/tests/scheduler_test.go @@ -0,0 +1,81 @@ +//go:build e2e + +// /* +// Copyright 2026 The Grove Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// */ + +package tests + +import ( + "context" + "testing" + + apicommon "github.com/ai-dynamo/grove/operator/api/common" + configv1alpha1 "github.com/ai-dynamo/grove/operator/api/config/v1alpha1" + "github.com/ai-dynamo/grove/operator/e2e/testctx" +) + +// Test_SR1_DisjointPodCliqueSchedulers verifies that each PodClique in a single +// PodCliqueSet can select a different enabled scheduler backend. Besides +// admission, this covers per-clique Pod preparation and successful scheduling. +func Test_SR1_DisjointPodCliqueSchedulers(t *testing.T) { + const ( + workloadName = "disjoint-schedulers" + expectedPods = 2 + ) + + ctx := context.Background() + tc, cleanup := testctx.PrepareTest(ctx, t, expectedPods, + testctx.WithWorkload(&testctx.WorkloadConfig{ + Name: workloadName, + YAMLPath: "../yaml/disjoint-schedulers.yaml", + Namespace: "default", + ExpectedPods: expectedPods, + }), + ) + defer cleanup() + + pods, err := tc.DeployAndVerifyWorkload() + if err != nil { + t.Fatalf("PodCliqueSet with disjoint scheduler names was not admitted and reconciled: %v", err) + } + + wantSchedulerByClique := map[string]string{ + workloadName + "-0-kube-worker": string(configv1alpha1.SchedulerNameKube), + workloadName + "-0-kai-worker": string(configv1alpha1.SchedulerNameKai), + } + for _, pod := range pods.Items { + cliqueName := pod.Labels[apicommon.LabelPodClique] + wantScheduler, ok := wantSchedulerByClique[cliqueName] + if !ok { + t.Errorf("pod %s has unexpected PodClique label %q", pod.Name, cliqueName) + continue + } + if pod.Spec.SchedulerName != wantScheduler { + t.Errorf("pod %s from PodClique %s uses scheduler %q, want %q", pod.Name, cliqueName, pod.Spec.SchedulerName, wantScheduler) + } + delete(wantSchedulerByClique, cliqueName) + } + for cliqueName := range wantSchedulerByClique { + t.Errorf("no pod was created for PodClique %s", cliqueName) + } + if t.Failed() { + return + } + + if err := tc.WaitForPods(expectedPods); err != nil { + t.Fatalf("pods using disjoint scheduler names did not become ready: %v", err) + } +} diff --git a/operator/e2e/yaml/disjoint-schedulers.yaml b/operator/e2e/yaml/disjoint-schedulers.yaml new file mode 100644 index 000000000..fc660f189 --- /dev/null +++ b/operator/e2e/yaml/disjoint-schedulers.yaml @@ -0,0 +1,72 @@ +# PodCliqueSet whose PodCliques select different enabled scheduler backends. +--- +apiVersion: grove.io/v1alpha1 +kind: PodCliqueSet +metadata: + name: disjoint-schedulers + labels: + app: disjoint-schedulers +spec: + replicas: 1 + template: + cliques: + - name: kube-worker + spec: + roleName: kube-worker + replicas: 1 + minAvailable: 1 + podSpec: + terminationGracePeriodSeconds: 5 + schedulerName: default-scheduler + affinity: + nodeAffinity: + requiredDuringSchedulingIgnoredDuringExecution: + nodeSelectorTerms: + - matchExpressions: + - key: node_role.e2e.grove.nvidia.com + operator: In + values: + - agent + tolerations: + - key: node_role.e2e.grove.nvidia.com + operator: Equal + value: agent + effect: NoSchedule + containers: + - name: kube-worker + image: registry:5001/busybox:latest + command: ["sleep", "infinity"] + resources: + requests: + memory: 80Mi + - name: kai-worker + labels: + kai.scheduler/queue: test + spec: + roleName: kai-worker + replicas: 1 + minAvailable: 1 + podSpec: + terminationGracePeriodSeconds: 5 + schedulerName: kai-scheduler + affinity: + nodeAffinity: + requiredDuringSchedulingIgnoredDuringExecution: + nodeSelectorTerms: + - matchExpressions: + - key: node_role.e2e.grove.nvidia.com + operator: In + values: + - agent + tolerations: + - key: node_role.e2e.grove.nvidia.com + operator: Equal + value: agent + effect: NoSchedule + containers: + - name: kai-worker + image: registry:5001/busybox:latest + command: ["sleep", "infinity"] + resources: + requests: + memory: 80Mi diff --git a/operator/internal/controller/common/component/utils/scheduler.go b/operator/internal/controller/common/component/utils/scheduler.go new file mode 100644 index 000000000..8fb30118c --- /dev/null +++ b/operator/internal/controller/common/component/utils/scheduler.go @@ -0,0 +1,76 @@ +// /* +// Copyright 2026 The Grove Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// */ + +package utils + +import ( + "crypto/sha256" + "fmt" + "maps" + "slices" + "strings" + + grovecorev1alpha1 "github.com/ai-dynamo/grove/operator/api/core/v1alpha1" + "github.com/ai-dynamo/grove/operator/internal/scheduler" +) + +const maxLabelValueLength = 63 + +// ResolveSchedulerName resolves an empty scheduler name to the configured +// default and otherwise returns the registered backend's canonical name. +func ResolveSchedulerName(schedRegistry scheduler.Registry, schedulerName string) string { + if schedRegistry == nil { + return schedulerName + } + if backend := schedRegistry.GetOrDefault(schedulerName); backend != nil { + return backend.Name() + } + return schedulerName +} + +// SchedulerNamesForPodCliqueSet returns the distinct resolved scheduler names +// selected by the PodClique templates in stable order. +func SchedulerNamesForPodCliqueSet(pcs *grovecorev1alpha1.PodCliqueSet, schedRegistry scheduler.Registry) []string { + names := make(map[string]struct{}, len(pcs.Spec.Template.Cliques)) + for _, clique := range pcs.Spec.Template.Cliques { + if clique == nil { + continue + } + name := ResolveSchedulerName(schedRegistry, clique.Spec.PodSpec.SchedulerName) + names[name] = struct{}{} + } + return slices.Sorted(maps.Keys(names)) +} + +// GenerateSchedulerScopedPodGangName returns the legacy PodGang name for a +// single-scheduler PCS and a scheduler-qualified name for a disjoint PCS. +func GenerateSchedulerScopedPodGangName(baseName string, pcs *grovecorev1alpha1.PodCliqueSet, schedulerName string, schedRegistry scheduler.Registry) string { + if len(SchedulerNamesForPodCliqueSet(pcs, schedRegistry)) <= 1 { + return baseName + } + + resolvedSchedulerName := ResolveSchedulerName(schedRegistry, schedulerName) + candidate := fmt.Sprintf("%s-%s", baseName, resolvedSchedulerName) + if len(candidate) <= maxLabelValueLength { + return candidate + } + + nameHash := sha256.Sum256([]byte(candidate)) + hashSuffix := fmt.Sprintf("-%x", nameHash[:4]) + maxBaseNameLength := maxLabelValueLength - len(hashSuffix) + trimmedBaseName := strings.TrimRight(baseName[:min(len(baseName), maxBaseNameLength)], "-.") + return trimmedBaseName + hashSuffix +} diff --git a/operator/internal/controller/common/component/utils/scheduler_test.go b/operator/internal/controller/common/component/utils/scheduler_test.go new file mode 100644 index 000000000..6ab87e0ae --- /dev/null +++ b/operator/internal/controller/common/component/utils/scheduler_test.go @@ -0,0 +1,81 @@ +// /* +// Copyright 2026 The Grove Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// */ + +package utils + +import ( + "strings" + "testing" + + configv1alpha1 "github.com/ai-dynamo/grove/operator/api/config/v1alpha1" + grovecorev1alpha1 "github.com/ai-dynamo/grove/operator/api/core/v1alpha1" + "github.com/ai-dynamo/grove/operator/internal/scheduler" + testutils "github.com/ai-dynamo/grove/operator/test/utils" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + corev1 "k8s.io/api/core/v1" +) + +func TestGenerateSchedulerScopedPodGangName(t *testing.T) { + registry := &testutils.FakeSchedulerRegistry{ + Backends: map[string]scheduler.Backend{ + string(configv1alpha1.SchedulerNameKube): testutils.NewFakeSchedulerBackend(string(configv1alpha1.SchedulerNameKube)), + string(configv1alpha1.SchedulerNameKai): testutils.NewFakeSchedulerBackend(string(configv1alpha1.SchedulerNameKai)), + }, + DefaultBackend: string(configv1alpha1.SchedulerNameKube), + } + + t.Run("preserves legacy name for a single scheduler", func(t *testing.T) { + pcs := podCliqueSetWithSchedulers("", string(configv1alpha1.SchedulerNameKube)) + assert.Equal(t, "workload-0", GenerateSchedulerScopedPodGangName("workload-0", pcs, "", registry)) + }) + + t.Run("qualifies each scheduler in a disjoint PCS", func(t *testing.T) { + pcs := podCliqueSetWithSchedulers("", string(configv1alpha1.SchedulerNameKai)) + assert.Equal(t, "workload-0-default-scheduler", GenerateSchedulerScopedPodGangName("workload-0", pcs, "", registry)) + assert.Equal(t, "workload-0-kai-scheduler", GenerateSchedulerScopedPodGangName("workload-0", pcs, string(configv1alpha1.SchedulerNameKai), registry)) + }) + + t.Run("keeps label values within the Kubernetes limit", func(t *testing.T) { + pcs := podCliqueSetWithSchedulers("", string(configv1alpha1.SchedulerNameKai)) + name := GenerateSchedulerScopedPodGangName(strings.Repeat("a", 60), pcs, string(configv1alpha1.SchedulerNameKai), registry) + require.Len(t, name, maxLabelValueLength) + assert.Regexp(t, `^[a-z0-9]([-a-z0-9]*[a-z0-9])?$`, name) + }) + + t.Run("hashes a short base name when the scheduler suffix exceeds the limit", func(t *testing.T) { + pcs := podCliqueSetWithSchedulers("", string(configv1alpha1.SchedulerNameKai)) + baseName := strings.Repeat("a", 44) + "-0" + + name := GenerateSchedulerScopedPodGangName(baseName, pcs, "", registry) + + require.Len(t, name, len(baseName)+9) + assert.Regexp(t, `^`+baseName+`-[a-f0-9]{8}$`, name) + }) +} + +func podCliqueSetWithSchedulers(schedulerNames ...string) *grovecorev1alpha1.PodCliqueSet { + pcs := &grovecorev1alpha1.PodCliqueSet{} + for _, schedulerName := range schedulerNames { + pcs.Spec.Template.Cliques = append(pcs.Spec.Template.Cliques, &grovecorev1alpha1.PodCliqueTemplateSpec{ + Spec: grovecorev1alpha1.PodCliqueSpec{ + PodSpec: corev1.PodSpec{SchedulerName: schedulerName}, + }, + }) + } + return pcs +} diff --git a/operator/internal/controller/podcliquescalinggroup/components/podclique/podclique.go b/operator/internal/controller/podcliquescalinggroup/components/podclique/podclique.go index 0b3eec388..136955964 100644 --- a/operator/internal/controller/podcliquescalinggroup/components/podclique/podclique.go +++ b/operator/internal/controller/podcliquescalinggroup/components/podclique/podclique.go @@ -33,6 +33,7 @@ import ( componentutils "github.com/ai-dynamo/grove/operator/internal/controller/common/component/utils" groveerr "github.com/ai-dynamo/grove/operator/internal/errors" "github.com/ai-dynamo/grove/operator/internal/mnnvl" + "github.com/ai-dynamo/grove/operator/internal/scheduler" "github.com/ai-dynamo/grove/operator/internal/utils" k8sutils "github.com/ai-dynamo/grove/operator/internal/utils/kubernetes" @@ -74,14 +75,16 @@ type _resource struct { client client.Client scheme *runtime.Scheme eventRecorder record.EventRecorder + schedRegistry scheduler.Registry } // New creates a new PodClique operator for managing PodClique resources within PodCliqueScalingGroups -func New(client client.Client, scheme *runtime.Scheme, eventRecorder record.EventRecorder) component.Operator[grovecorev1alpha1.PodCliqueScalingGroup] { +func New(client client.Client, scheme *runtime.Scheme, eventRecorder record.EventRecorder, schedRegistry scheduler.Registry) component.Operator[grovecorev1alpha1.PodCliqueScalingGroup] { return &_resource{ client: client, scheme: scheme, eventRecorder: eventRecorder, + schedRegistry: schedRegistry, } } @@ -313,9 +316,16 @@ func (r _resource) buildResource(logger logr.Logger, pcs *grovecorev1alpha1.PodC return err } - podGangName := apicommon.GeneratePodGangNameForPodCliqueOwnedByPCSG(pcs, pcsReplicaIndex, pcsg, pcsgReplicaIndex) + basePodGangName := apicommon.GeneratePodGangNameForPodCliqueOwnedByPCSG(pcs, pcsReplicaIndex, pcsg, pcsgReplicaIndex) + podGangName := componentutils.GenerateSchedulerScopedPodGangName(basePodGangName, pcs, pclqTemplateSpec.Spec.PodSpec.SchedulerName, r.schedRegistry) + scopedBasePodGangName := componentutils.GenerateSchedulerScopedPodGangName( + apicommon.GenerateBasePodGangName(apicommon.ResourceNameReplica{Name: pcs.Name, Replica: pcsReplicaIndex}), + pcs, + pclqTemplateSpec.Spec.PodSpec.SchedulerName, + r.schedRegistry, + ) - pclq.Labels = getLabels(pcs, pcsReplicaIndex, pcsg, pcsgReplicaIndex, pclqObjectKey, pclqTemplateSpec, podGangName) + pclq.Labels = getLabels(pcs, pcsReplicaIndex, pcsg, pcsgReplicaIndex, pclqObjectKey, pclqTemplateSpec, podGangName, scopedBasePodGangName) pclq.Annotations = maps.Clone(pclqTemplateSpec.Annotations) // PodGang owns topology selection; do not propagate a template topology annotation to PodClique pods. delete(pclq.Annotations, apiconstants.AnnotationTopologyName) @@ -469,7 +479,7 @@ func getPodCliqueSelectorLabels(pcsgObjectMeta metav1.ObjectMeta) map[string]str } // getLabels constructs the complete set of labels for a PodClique including Grove-specific, component, and template labels -func getLabels(pcs *grovecorev1alpha1.PodCliqueSet, pcsReplicaIndex int, pcsg *grovecorev1alpha1.PodCliqueScalingGroup, pcsgReplicaIndex int, pclqObjectKey client.ObjectKey, pclqTemplateSpec *grovecorev1alpha1.PodCliqueTemplateSpec, podGangName string) map[string]string { +func getLabels(pcs *grovecorev1alpha1.PodCliqueSet, pcsReplicaIndex int, pcsg *grovecorev1alpha1.PodCliqueScalingGroup, pcsgReplicaIndex int, pclqObjectKey client.ObjectKey, pclqTemplateSpec *grovecorev1alpha1.PodCliqueTemplateSpec, podGangName, basePodGangName string) map[string]string { pclqComponentLabels := map[string]string{ apicommon.LabelAppNameKey: pclqObjectKey.Name, apicommon.LabelComponentKey: apicommon.LabelComponentNamePodCliqueScalingGroupPodClique, @@ -481,9 +491,6 @@ func getLabels(pcs *grovecorev1alpha1.PodCliqueSet, pcsReplicaIndex int, pcsg *g } // Add base-podgang label for scaled PodGang pods (beyond minAvailable) - basePodGangName := apicommon.GenerateBasePodGangName( - apicommon.ResourceNameReplica{Name: pcs.Name, Replica: pcsReplicaIndex}, - ) if podGangName != basePodGangName { // This pod belongs to a scaled PodGang - add the base PodGang label pclqComponentLabels[apicommon.LabelBasePodGang] = basePodGangName diff --git a/operator/internal/controller/podcliquescalinggroup/components/podclique/podclique_test.go b/operator/internal/controller/podcliquescalinggroup/components/podclique/podclique_test.go index 2a01a738b..bdb46bb4d 100644 --- a/operator/internal/controller/podcliquescalinggroup/components/podclique/podclique_test.go +++ b/operator/internal/controller/podcliquescalinggroup/components/podclique/podclique_test.go @@ -54,7 +54,7 @@ func TestNew(t *testing.T) { client := fake.NewClientBuilder().WithScheme(scheme).Build() eventRecorder := &record.FakeRecorder{} - operator := New(client, scheme, eventRecorder) + operator := New(client, scheme, eventRecorder, testutils.NewDefaultFakeRegistry()) assert.NotNil(t, operator) resource, ok := operator.(*_resource) diff --git a/operator/internal/controller/podcliquescalinggroup/components/registry.go b/operator/internal/controller/podcliquescalinggroup/components/registry.go index d3b920db1..befdb2cf1 100644 --- a/operator/internal/controller/podcliquescalinggroup/components/registry.go +++ b/operator/internal/controller/podcliquescalinggroup/components/registry.go @@ -20,14 +20,15 @@ import ( "github.com/ai-dynamo/grove/operator/api/core/v1alpha1" "github.com/ai-dynamo/grove/operator/internal/controller/common/component" "github.com/ai-dynamo/grove/operator/internal/controller/podcliquescalinggroup/components/podclique" + "github.com/ai-dynamo/grove/operator/internal/scheduler" "k8s.io/client-go/tools/record" "sigs.k8s.io/controller-runtime/pkg/manager" ) // CreateOperatorRegistry initializes the operator registry for the PodCliqueScalingGroup reconciler. -func CreateOperatorRegistry(mgr manager.Manager, eventRecorder record.EventRecorder) component.OperatorRegistry[v1alpha1.PodCliqueScalingGroup] { +func CreateOperatorRegistry(mgr manager.Manager, eventRecorder record.EventRecorder, schedRegistry scheduler.Registry) component.OperatorRegistry[v1alpha1.PodCliqueScalingGroup] { reg := component.NewOperatorRegistry[v1alpha1.PodCliqueScalingGroup]() - reg.Register(component.KindPodClique, podclique.New(mgr.GetClient(), mgr.GetScheme(), eventRecorder)) + reg.Register(component.KindPodClique, podclique.New(mgr.GetClient(), mgr.GetScheme(), eventRecorder, schedRegistry)) return reg } diff --git a/operator/internal/controller/podcliquescalinggroup/components/registry_test.go b/operator/internal/controller/podcliquescalinggroup/components/registry_test.go index 497e44900..f1cc1b2f8 100644 --- a/operator/internal/controller/podcliquescalinggroup/components/registry_test.go +++ b/operator/internal/controller/podcliquescalinggroup/components/registry_test.go @@ -21,6 +21,7 @@ import ( grovecorev1alpha1 "github.com/ai-dynamo/grove/operator/api/core/v1alpha1" "github.com/ai-dynamo/grove/operator/internal/controller/common/component" + testutils "github.com/ai-dynamo/grove/operator/test/utils" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -44,7 +45,7 @@ func TestCreateOperatorRegistry(t *testing.T) { mgr := &mockManager{client: cl, scheme: scheme} eventRecorder := record.NewFakeRecorder(10) - registry := CreateOperatorRegistry(mgr, eventRecorder) + registry := CreateOperatorRegistry(mgr, eventRecorder, testutils.NewDefaultFakeRegistry()) require.NotNil(t, registry) diff --git a/operator/internal/controller/podcliquescalinggroup/reconciler.go b/operator/internal/controller/podcliquescalinggroup/reconciler.go index e0b1c1a6a..b0fa6d59a 100644 --- a/operator/internal/controller/podcliquescalinggroup/reconciler.go +++ b/operator/internal/controller/podcliquescalinggroup/reconciler.go @@ -27,6 +27,7 @@ import ( componentutils "github.com/ai-dynamo/grove/operator/internal/controller/common/component/utils" pcsgcomponent "github.com/ai-dynamo/grove/operator/internal/controller/podcliquescalinggroup/components" ctrlutils "github.com/ai-dynamo/grove/operator/internal/controller/utils" + "github.com/ai-dynamo/grove/operator/internal/scheduler" "k8s.io/client-go/tools/record" ctrl "sigs.k8s.io/controller-runtime" @@ -45,7 +46,7 @@ type Reconciler struct { } // NewReconciler creates a new instance of the PodClique Reconciler. -func NewReconciler(mgr ctrl.Manager, controllerCfg groveconfigv1alpha1.PodCliqueScalingGroupControllerConfiguration) *Reconciler { +func NewReconciler(mgr ctrl.Manager, controllerCfg groveconfigv1alpha1.PodCliqueScalingGroupControllerConfiguration, schedRegistry scheduler.Registry) *Reconciler { eventRecorder := mgr.GetEventRecorderFor(controllerName) client := mgr.GetClient() return &Reconciler{ @@ -53,7 +54,7 @@ func NewReconciler(mgr ctrl.Manager, controllerCfg groveconfigv1alpha1.PodClique client: client, eventRecorder: eventRecorder, reconcileStatusRecorder: ctrlcommon.NewReconcileErrorRecorder(client), - operatorRegistry: pcsgcomponent.CreateOperatorRegistry(mgr, eventRecorder), + operatorRegistry: pcsgcomponent.CreateOperatorRegistry(mgr, eventRecorder, schedRegistry), } } diff --git a/operator/internal/controller/podcliqueset/components/podclique/podclique.go b/operator/internal/controller/podcliqueset/components/podclique/podclique.go index a7348958d..499964caf 100644 --- a/operator/internal/controller/podcliqueset/components/podclique/podclique.go +++ b/operator/internal/controller/podcliqueset/components/podclique/podclique.go @@ -31,6 +31,7 @@ import ( componentutils "github.com/ai-dynamo/grove/operator/internal/controller/common/component/utils" groveerr "github.com/ai-dynamo/grove/operator/internal/errors" "github.com/ai-dynamo/grove/operator/internal/mnnvl" + "github.com/ai-dynamo/grove/operator/internal/scheduler" "github.com/ai-dynamo/grove/operator/internal/utils" k8sutils "github.com/ai-dynamo/grove/operator/internal/utils/kubernetes" @@ -56,14 +57,16 @@ type _resource struct { client client.Client scheme *runtime.Scheme eventRecorder record.EventRecorder + schedRegistry scheduler.Registry } // New creates an instance of PodClique components operator. -func New(client client.Client, scheme *runtime.Scheme, eventRecorder record.EventRecorder) component.Operator[grovecorev1alpha1.PodCliqueSet] { +func New(client client.Client, scheme *runtime.Scheme, eventRecorder record.EventRecorder, schedRegistry scheduler.Registry) component.Operator[grovecorev1alpha1.PodCliqueSet] { return &_resource{ client: client, scheme: scheme, eventRecorder: eventRecorder, + schedRegistry: schedRegistry, } } @@ -307,7 +310,9 @@ func (r _resource) buildResource(logger logr.Logger, pclq *grovecorev1alpha1.Pod } // Add finalizer at creation so PCLQ controller does not need a separate PATCH on first reconcile. controllerutil.AddFinalizer(pclq, apiconstants.FinalizerPodClique) - pclq.Labels = getLabels(pcs, pcsReplica, pclqObjectKey, pclqTemplateSpec, apicommon.GeneratePodGangNameForPodCliqueOwnedByPodCliqueSet(pcs, pcsReplica)) + basePodGangName := apicommon.GeneratePodGangNameForPodCliqueOwnedByPodCliqueSet(pcs, pcsReplica) + podGangName := componentutils.GenerateSchedulerScopedPodGangName(basePodGangName, pcs, pclqTemplateSpec.Spec.PodSpec.SchedulerName, r.schedRegistry) + pclq.Labels = getLabels(pcs, pcsReplica, pclqObjectKey, pclqTemplateSpec, podGangName) pclq.Annotations = maps.Clone(pclqTemplateSpec.Annotations) // PodGang owns topology selection; do not propagate a template topology annotation to PodClique pods. delete(pclq.Annotations, apiconstants.AnnotationTopologyName) diff --git a/operator/internal/controller/podcliqueset/components/podclique/podclique_test.go b/operator/internal/controller/podcliqueset/components/podclique/podclique_test.go index a706cf5d7..c95192d5f 100644 --- a/operator/internal/controller/podcliqueset/components/podclique/podclique_test.go +++ b/operator/internal/controller/podcliqueset/components/podclique/podclique_test.go @@ -112,7 +112,7 @@ func TestGetExistingResourceNames(t *testing.T) { existingObjects := createExistingPodCliquesFromPCS(pcs, tc.podCliqueNamesNotOwnedByPCS) // Create a fake client with PodCliques cl := testutils.CreateFakeClientForObjectsMatchingLabels(nil, tc.listErr, pcs.Namespace, grovecorev1alpha1.SchemeGroupVersion.WithKind("PodClique"), getPodCliqueSelectorLabels(pcs.ObjectMeta), existingObjects...) - operator := New(cl, groveclientscheme.Scheme, record.NewFakeRecorder(10)) + operator := New(cl, groveclientscheme.Scheme, record.NewFakeRecorder(10), testutils.NewDefaultFakeRegistry()) actualPCLQNames, err := operator.GetExistingResourceNames(context.Background(), logr.Discard(), pcs.ObjectMeta) if tc.expectedErr == nil { assert.NoError(t, err) @@ -165,7 +165,7 @@ func TestDelete(t *testing.T) { existingPodCliques := createDefaultPodCliques(pcsObjMeta, "howl", tc.numExistingPodCliques) // Create a fake client with PodCliques cl := testutils.CreateFakeClientForObjectsMatchingLabels(tc.deleteError, nil, testPCSNamespace, grovecorev1alpha1.SchemeGroupVersion.WithKind("PodClique"), getPodCliqueSelectorLabels(pcsObjMeta), existingPodCliques...) - operator := New(cl, groveclientscheme.Scheme, record.NewFakeRecorder(10)) + operator := New(cl, groveclientscheme.Scheme, record.NewFakeRecorder(10), testutils.NewDefaultFakeRegistry()) err := operator.Delete(context.Background(), logr.Discard(), pcsObjMeta) if tc.expectedError != nil { testutils.CheckGroveError(t, tc.expectedError, err) diff --git a/operator/internal/controller/podcliqueset/components/podgang/podgang.go b/operator/internal/controller/podcliqueset/components/podgang/podgang.go index c602d3567..b5e087b87 100644 --- a/operator/internal/controller/podcliqueset/components/podgang/podgang.go +++ b/operator/internal/controller/podcliqueset/components/podgang/podgang.go @@ -137,7 +137,13 @@ func (r _resource) buildResource(pcs *grovecorev1alpha1.PodCliqueSet, pgi *podGa pg.Labels = mirrorPCSMetadata(pg.Labels, pcs.Labels, getLabels(pcs.Name)) // Set scheduler name so the podgang controller can resolve the correct backend. // When no scheduler can be resolved, drop any stale label from a previous reconcile. - if schedName := r.getSchedulerNameForPCS(pcs); schedName != "" { + schedName := pgi.schedulerName + if schedName == "" { + // Preserve compatibility for directly constructed podGangInfo values in + // callers and tests. Reconciled PodGangs always carry schedulerName. + schedName = r.getSchedulerNameForPCS(pcs) + } + if schedName != "" { pg.Labels[apicommon.LabelSchedulerName] = schedName } else { delete(pg.Labels, apicommon.LabelSchedulerName) @@ -253,17 +259,19 @@ func podGangHasTranslatedTopologyConstraints(pgi *podGangInfo) bool { return false } -// getSchedulerNameForPCS returns the scheduler backend name for the PodCliqueSet: -// First check the PodClique templates to find any schedulerName configured. Validating webhook ensures -// that there cannot be more than one scheduler name configured for a PCS. -// the template's schedulerName if set (same across all cliques per validation), else the default backend. +// getSchedulerNameForPCS returns the first explicitly configured scheduler, or the default backend. +// Reconciled PodGangs carry their resolved scheduler directly; this is a compatibility fallback for +// callers that construct podGangInfo without a scheduler name. func (r _resource) getSchedulerNameForPCS(pcs *grovecorev1alpha1.PodCliqueSet) string { for _, c := range pcs.Spec.Template.Cliques { if c != nil && c.Spec.PodSpec.SchedulerName != "" { return c.Spec.PodSpec.SchedulerName } } - return r.schedRegistry.GetDefault().Name() + if defaultBackend := r.schedRegistry.GetDefault(); defaultBackend != nil { + return defaultBackend.Name() + } + return "" } // setOrUpdateInitializedCondition sets or updates the PodGangInitialized condition on the PodGang status. diff --git a/operator/internal/controller/podcliqueset/components/podgang/syncflow.go b/operator/internal/controller/podcliqueset/components/podgang/syncflow.go index 44bfdf94a..5a3ba5bbe 100644 --- a/operator/internal/controller/podcliqueset/components/podgang/syncflow.go +++ b/operator/internal/controller/podcliqueset/components/podgang/syncflow.go @@ -20,6 +20,8 @@ import ( "context" "errors" "fmt" + "maps" + "slices" apicommon "github.com/ai-dynamo/grove/operator/api/common" grovecorev1alpha1 "github.com/ai-dynamo/grove/operator/api/core/v1alpha1" @@ -28,6 +30,7 @@ import ( "github.com/ai-dynamo/grove/operator/internal/controller/common/component" componentutils "github.com/ai-dynamo/grove/operator/internal/controller/common/component/utils" groveerr "github.com/ai-dynamo/grove/operator/internal/errors" + "github.com/ai-dynamo/grove/operator/internal/scheduler" k8sutils "github.com/ai-dynamo/grove/operator/internal/utils/kubernetes" groveschedulerv1alpha1 "github.com/ai-dynamo/grove/scheduler/api/core/v1alpha1" @@ -50,6 +53,7 @@ func (r _resource) prepareSyncFlow(ctx context.Context, logger logr.Logger, pcs logger: logger, existingPCLQPods: make(map[string][]corev1.Pod), unassignedPodsByPCLQ: make(map[string][]corev1.Pod), + schedRegistry: r.schedRegistry, } sc.existingPCLQs, err = r.getExistingPCLQsForPCS(ctx, pcs) @@ -180,20 +184,19 @@ func (r _resource) computeExpectedPodGangs(sc *syncContext) error { func buildExpectedBasePodGangForPCSReplicas(sc *syncContext) ([]*podGangInfo, error) { expectedPodGangs := make([]*podGangInfo, 0, int(sc.pcs.Spec.Replicas)) for pcsReplica := range int(sc.pcs.Spec.Replicas) { - basePodGang, err := buildExpectedBasePodGangForPCSReplica(sc, pcsReplica) + basePodGangs, err := buildExpectedBasePodGangForPCSReplica(sc, pcsReplica) if err != nil { return nil, err } - expectedPodGangs = append(expectedPodGangs, basePodGang) + expectedPodGangs = append(expectedPodGangs, basePodGangs...) } return expectedPodGangs, nil } // buildExpectedBasePodGangForPCSReplica builds the base PodGang info for a given PodCliqueSet replica. -func buildExpectedBasePodGangForPCSReplica(sc *syncContext, pcsReplica int) (*podGangInfo, error) { +func buildExpectedBasePodGangForPCSReplica(sc *syncContext, pcsReplica int) ([]*podGangInfo, error) { podGangFQN := apicommon.GenerateBasePodGangName(apicommon.ResourceNameReplica{Name: sc.pcs.Name, Replica: pcsReplica}) pg := &podGangInfo{ - fqn: podGangFQN, // TopologyConstraint for the base PodGang comes from the topology constraint defined at the PCS level. topologyConstraint: createTopologyPackConstraint(sc, client.ObjectKeyFromObject(sc.pcs), sc.pcs.Spec.Template.TopologyConstraint), } @@ -210,7 +213,7 @@ func buildExpectedBasePodGangForPCSReplica(sc *syncContext, pcsReplica int) (*po pg.pcsgTopologyConstraints = pcsgPackConstraints pg.pclqs = pclqInfos - return pg, nil + return partitionPodGangInfoByScheduler(sc, podGangFQN, pg), nil } func buildStandalonePCLQInfosForBasePodGang(sc *syncContext, pcsReplica int) []pclqInfo { @@ -286,17 +289,18 @@ func (r _resource) buildExpectedScaledPodGangsForPCSG(sc *syncContext, pcsReplic minAvailable := int(*pcsgConfig.MinAvailable) scaledReplicas := replicas - minAvailable for podGangIndex, pcsgReplica := 0, minAvailable; podGangIndex < scaledReplicas; podGangIndex, pcsgReplica = podGangIndex+1, pcsgReplica+1 { - pg, err := doBuildExpectedScaledPodGangForPCSG(sc, pcsgFQN, pcsgConfig, pcsgReplica, podGangIndex) + pg, err := doBuildExpectedScaledPodGangForPCSG(sc, pcsgFQN, pcsgConfig, pcsgReplica) if err != nil { return nil, fmt.Errorf("failed to build expected scaled PodGang for PCSG %q replica %d: %w", pcsgFQN, pcsgReplica, err) } - expectedPodGangs = append(expectedPodGangs, pg) + basePodGangName := apicommon.CreatePodGangNameFromPCSGFQN(pcsgFQN, podGangIndex) + expectedPodGangs = append(expectedPodGangs, partitionPodGangInfoByScheduler(sc, basePodGangName, pg)...) } } return expectedPodGangs, nil } -func doBuildExpectedScaledPodGangForPCSG(sc *syncContext, pcsgFQN string, pcsgConfig grovecorev1alpha1.PodCliqueScalingGroupConfig, pcsgReplica int, podGangIndex int) (*podGangInfo, error) { +func doBuildExpectedScaledPodGangForPCSG(sc *syncContext, pcsgFQN string, pcsgConfig grovecorev1alpha1.PodCliqueScalingGroupConfig, pcsgReplica int) (*podGangInfo, error) { var ( pclqInfos = make([]pclqInfo, 0, len(pcsgConfig.CliqueNames)) topologyConstraint *groveschedulerv1alpha1.TopologyConstraint @@ -328,7 +332,6 @@ func doBuildExpectedScaledPodGangForPCSG(sc *syncContext, pcsgFQN string, pcsgCo } pg := &podGangInfo{ - fqn: apicommon.CreatePodGangNameFromPCSGFQN(pcsgFQN, podGangIndex), topologyConstraint: topologyConstraint, pclqs: pclqInfos, } @@ -340,14 +343,68 @@ func doBuildExpectedScaledPodGangForPCSG(sc *syncContext, pcsgFQN string, pcsgCo func buildPodCliqueInfo(sc *syncContext, pclqTemplateSpec *grovecorev1alpha1.PodCliqueTemplateSpec, pclqFQN string, belongsToPCSG bool) pclqInfo { replicas := determinePodCliqueReplicas(sc, pclqTemplateSpec, pclqFQN, belongsToPCSG) expectedPCLQ := pclqInfo{ - fqn: pclqFQN, - replicas: replicas, - minAvailable: *pclqTemplateSpec.Spec.MinAvailable, + fqn: pclqFQN, + replicas: replicas, + minAvailable: *pclqTemplateSpec.Spec.MinAvailable, + schedulerName: componentutils.ResolveSchedulerName(sc.schedRegistry, pclqTemplateSpec.Spec.PodSpec.SchedulerName), } expectedPCLQ.topologyConstraint = createTopologyPackConstraint(sc, types.NamespacedName{Namespace: sc.pcs.Namespace, Name: pclqFQN}, pclqTemplateSpec.TopologyConstraint) return expectedPCLQ } +// partitionPodGangInfoByScheduler splits a logical PodGang into one PodGang per +// selected scheduler. This prevents a scheduler backend from waiting on Pods +// that are intentionally assigned to a different scheduler. +func partitionPodGangInfoByScheduler(sc *syncContext, basePodGangName string, pg *podGangInfo) []*podGangInfo { + pclqsByScheduler := make(map[string][]pclqInfo) + for _, pclq := range pg.pclqs { + pclqsByScheduler[pclq.schedulerName] = append(pclqsByScheduler[pclq.schedulerName], pclq) + } + schedulerNames := slices.Sorted(maps.Keys(pclqsByScheduler)) + + result := make([]*podGangInfo, 0, len(schedulerNames)) + for _, schedulerName := range schedulerNames { + pclqs := pclqsByScheduler[schedulerName] + result = append(result, &podGangInfo{ + fqn: componentutils.GenerateSchedulerScopedPodGangName(basePodGangName, sc.pcs, schedulerName, sc.schedRegistry), + schedulerName: schedulerName, + pclqs: pclqs, + topologyConstraint: pg.topologyConstraint, + pcsgTopologyConstraints: filterPCSGTopologyConstraints(pg.pcsgTopologyConstraints, pclqs), + }) + } + return result +} + +func filterPCSGTopologyConstraints(configs []groveschedulerv1alpha1.TopologyConstraintGroupConfig, pclqs []pclqInfo) []groveschedulerv1alpha1.TopologyConstraintGroupConfig { + if len(configs) == 0 { + return nil + } + allowedPodGroups := make(map[string]struct{}, len(pclqs)) + for _, pclq := range pclqs { + allowedPodGroups[pclq.fqn] = struct{}{} + } + + result := make([]groveschedulerv1alpha1.TopologyConstraintGroupConfig, 0, len(configs)) + for _, config := range configs { + filteredPodGroupNames := make([]string, 0, len(config.PodGroupNames)) + for _, podGroupName := range config.PodGroupNames { + if _, ok := allowedPodGroups[podGroupName]; ok { + filteredPodGroupNames = append(filteredPodGroupNames, podGroupName) + } + } + if len(filteredPodGroupNames) == 0 { + continue + } + config.PodGroupNames = filteredPodGroupNames + result = append(result, config) + } + if len(result) == 0 { + return nil + } + return result +} + // createTopologyPackConstraint creates a TopologyPackConstraint based on the sync context and provided parameters for a resource. // PackConstraints are defined at multiple levels (PodCliqueSet, PodCliqueScalingGroup, PodClique). This function helps create a TopologyPackConstraint for any of these levels. func createTopologyPackConstraint(sc *syncContext, nsName types.NamespacedName, topologyConstraint *grovecorev1alpha1.TopologyConstraint) *groveschedulerv1alpha1.TopologyConstraint { @@ -659,6 +716,7 @@ type syncContext struct { unassignedPodsByPCLQ map[string][]corev1.Pod tasEnabled bool topologyLevels []grovecorev1alpha1.TopologyLevel + schedRegistry scheduler.Registry } // getPodGangNamesPendingCreation identifies PodGangs not yet created. @@ -782,6 +840,8 @@ func (sfr *syncFlowResult) getAggregatedError() error { type podGangInfo struct { // fqn is a fully qualified name of a PodGang. fqn string + // schedulerName is the resolved scheduler backend for this PodGang. + schedulerName string // pclqs holds the relevant information for all constituent PodCliques for this PodGang. pclqs []pclqInfo // topologyConstraint holds the topology pack constraint applicable at the PodGang level. @@ -810,6 +870,8 @@ type pclqInfo struct { replicas int32 // minAvailable is the minimum number of pods that are required for gang scheduling from this PodClique minAvailable int32 + // schedulerName is the resolved scheduler backend for this PodClique. + schedulerName string // associatedPodNames are Pod names (having this PodClique as an owner) that have already been associated to this PodGang. // This will be updated as and when pods are either deleted or new pods are associated. associatedPodNames []string diff --git a/operator/internal/controller/podcliqueset/components/podgang/syncflow_test.go b/operator/internal/controller/podcliqueset/components/podgang/syncflow_test.go index b68e3ad72..3744f09f7 100644 --- a/operator/internal/controller/podcliqueset/components/podgang/syncflow_test.go +++ b/operator/internal/controller/podcliqueset/components/podgang/syncflow_test.go @@ -46,6 +46,7 @@ import ( var defaultFakeSchedulerRegistry = &testutils.FakeSchedulerRegistry{ Backends: map[string]scheduler.Backend{ "default-scheduler": testutils.NewFakeSchedulerBackend("default-scheduler"), + "kai-scheduler": testutils.NewFakeSchedulerBackend("kai-scheduler"), }, DefaultBackend: "default-scheduler", } @@ -961,6 +962,62 @@ func TestComputeExpectedPodGangs(t *testing.T) { } } +func TestComputeExpectedPodGangsWithDisjointSchedulers(t *testing.T) { + pcs := &grovecorev1alpha1.PodCliqueSet{ + ObjectMeta: metav1.ObjectMeta{Name: "test-pcs", Namespace: "default", UID: "test-uid-123"}, + Spec: grovecorev1alpha1.PodCliqueSetSpec{ + Replicas: 1, + Template: grovecorev1alpha1.PodCliqueSetTemplateSpec{ + Cliques: []*grovecorev1alpha1.PodCliqueTemplateSpec{ + { + Name: "kube-worker", + Spec: grovecorev1alpha1.PodCliqueSpec{ + Replicas: 1, + MinAvailable: ptr.To(int32(1)), + PodSpec: v1.PodSpec{SchedulerName: "default-scheduler"}, + }, + }, + { + Name: "kai-worker", + Spec: grovecorev1alpha1.PodCliqueSpec{ + Replicas: 1, + MinAvailable: ptr.To(int32(1)), + PodSpec: v1.PodSpec{SchedulerName: "kai-scheduler"}, + }, + }, + }, + PodCliqueScalingGroupConfigs: []grovecorev1alpha1.PodCliqueScalingGroupConfig{ + { + Name: "sg", + Replicas: ptr.To(int32(2)), + MinAvailable: ptr.To(int32(1)), + CliqueNames: []string{"kube-worker", "kai-worker"}, + }, + }, + }, + }, + } + r := &_resource{schedRegistry: defaultFakeSchedulerRegistry} + sc := &syncContext{pcs: pcs, existingPCSGByName: map[string]grovecorev1alpha1.PodCliqueScalingGroup{}} + + require.NoError(t, r.computeExpectedPodGangs(sc)) + require.Len(t, sc.expectedPodGangs, 4) + + expectedSchedulersByPodGang := map[string]string{ + "test-pcs-0-default-scheduler": "default-scheduler", + "test-pcs-0-kai-scheduler": "kai-scheduler", + "test-pcs-0-sg-0-default-scheduler": "default-scheduler", + "test-pcs-0-sg-0-kai-scheduler": "kai-scheduler", + } + for _, podGang := range sc.expectedPodGangs { + expectedScheduler, ok := expectedSchedulersByPodGang[podGang.fqn] + require.True(t, ok, "unexpected PodGang %q", podGang.fqn) + assert.Equal(t, expectedScheduler, podGang.schedulerName) + require.Len(t, podGang.pclqs, 1) + assert.Equal(t, expectedScheduler, podGang.pclqs[0].schedulerName) + } +} + type expectedPodGangTopologyConstraints struct { fqn string topologyPackConstraint *expectedTopologyPackConstraint diff --git a/operator/internal/controller/podcliqueset/components/registry.go b/operator/internal/controller/podcliqueset/components/registry.go index ce75b12d6..247cd64be 100644 --- a/operator/internal/controller/podcliqueset/components/registry.go +++ b/operator/internal/controller/podcliqueset/components/registry.go @@ -42,7 +42,7 @@ import ( func CreateOperatorRegistry(mgr manager.Manager, eventRecorder record.EventRecorder, topologyAwareSchedulingConfig configv1alpha1.TopologyAwareSchedulingConfiguration, networkConfig configv1alpha1.NetworkAcceleration, schedRegistry scheduler.Registry) component.OperatorRegistry[v1alpha1.PodCliqueSet] { cl := mgr.GetClient() reg := component.NewOperatorRegistry[v1alpha1.PodCliqueSet]() - reg.Register(component.KindPodClique, podclique.New(cl, mgr.GetScheme(), eventRecorder)) + reg.Register(component.KindPodClique, podclique.New(cl, mgr.GetScheme(), eventRecorder, schedRegistry)) reg.Register(component.KindHeadlessService, service.New(cl, mgr.GetScheme())) reg.Register(component.KindRole, role.New(cl, mgr.GetScheme())) reg.Register(component.KindRoleBinding, rolebinding.New(cl, mgr.GetScheme())) diff --git a/operator/internal/controller/register.go b/operator/internal/controller/register.go index 1f6f54fcb..8377b577d 100644 --- a/operator/internal/controller/register.go +++ b/operator/internal/controller/register.go @@ -50,7 +50,7 @@ func RegisterControllers(mgr ctrl.Manager, config *configv1alpha1.OperatorConfig if err := pcReconciler.RegisterWithManager(mgr); err != nil { return err } - pcsgReconciler := podcliquescalinggroup.NewReconciler(mgr, config.Controllers.PodCliqueScalingGroup) + pcsgReconciler := podcliquescalinggroup.NewReconciler(mgr, config.Controllers.PodCliqueScalingGroup, schedRegistry) if err := pcsgReconciler.RegisterWithManager(mgr); err != nil { return err } diff --git a/operator/internal/webhook/admission/pcs/validation/handler.go b/operator/internal/webhook/admission/pcs/validation/handler.go index 137c48f2b..89f22472e 100644 --- a/operator/internal/webhook/admission/pcs/validation/handler.go +++ b/operator/internal/webhook/admission/pcs/validation/handler.go @@ -29,6 +29,7 @@ import ( "github.com/go-logr/logr" admissionv1 "k8s.io/api/admission/v1" "k8s.io/apimachinery/pkg/runtime" + utilerrors "k8s.io/apimachinery/pkg/util/errors" "k8s.io/apimachinery/pkg/util/validation/field" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/manager" @@ -129,22 +130,80 @@ func (h *Handler) ValidateDelete(_ context.Context, _ runtime.Object) (admission return nil, nil } -// validatePodCliqueSetWithBackend resolves the scheduler backend for the PCS and runs backend-specific validation. -// All cliques share the same (resolved) schedulerName after validateSchedulerNames, so we use the first clique. +// validatePodCliqueSetWithBackend resolves every scheduler backend selected by +// the PCS and runs backend-specific validation once per distinct backend. func (h *Handler) validatePodCliqueSetWithBackend(ctx context.Context, pcs *v1alpha1.PodCliqueSet) error { - schedulerName := "" - if len(pcs.Spec.Template.Cliques) > 0 && pcs.Spec.Template.Cliques[0] != nil { - schedulerName = pcs.Spec.Template.Cliques[0].Spec.PodSpec.SchedulerName + schedulerNames := make([]string, 0, len(pcs.Spec.Template.Cliques)) + for _, clique := range pcs.Spec.Template.Cliques { + if clique != nil { + schedulerNames = append(schedulerNames, clique.Spec.PodSpec.SchedulerName) + } + } + if len(schedulerNames) == 0 { + schedulerNames = append(schedulerNames, "") + } + + seenBackends := make(map[string]struct{}, len(schedulerNames)) + var validationErrs []error + for _, schedulerName := range schedulerNames { + backend := h.schedRegistry.GetOrDefault(schedulerName) + if backend == nil { + if schedulerName == "" { + validationErrs = append(validationErrs, fmt.Errorf("default scheduler backend is not configured")) + } else { + validationErrs = append(validationErrs, fmt.Errorf("schedulerName %q is not enabled in OperatorConfiguration", schedulerName)) + } + continue + } + if _, seen := seenBackends[backend.Name()]; seen { + continue + } + seenBackends[backend.Name()] = struct{}{} + if err := backend.ValidatePodCliqueSet(ctx, h.podCliqueSetForBackend(pcs, backend.Name())); err != nil { + validationErrs = append(validationErrs, fmt.Errorf("scheduler backend %q: %w", backend.Name(), err)) + } } - backend := h.schedRegistry.GetOrDefault(schedulerName) - if backend == nil { - if schedulerName == "" { - return fmt.Errorf("default scheduler backend is not configured") + return utilerrors.NewAggregate(validationErrs) +} + +// podCliqueSetForBackend returns a copy containing only the cliques and scaling +// groups assigned to backendName. PCS-level settings remain because they apply +// to every scheduler-specific PodGang produced from the PCS. +func (h *Handler) podCliqueSetForBackend(pcs *v1alpha1.PodCliqueSet, backendName string) *v1alpha1.PodCliqueSet { + scopedPCS := pcs.DeepCopy() + // Filter the already-deep-copied cliques in place instead of re-deep-copying the subset. + selectedCliques := scopedPCS.Spec.Template.Cliques[:0] + cliqueNames := make(map[string]struct{}) + for _, clique := range scopedPCS.Spec.Template.Cliques { + if clique == nil { + continue + } + backend := h.schedRegistry.GetOrDefault(clique.Spec.PodSpec.SchedulerName) + if backend == nil || backend.Name() != backendName { + continue + } + selectedCliques = append(selectedCliques, clique) + cliqueNames[clique.Name] = struct{}{} + } + scopedPCS.Spec.Template.Cliques = selectedCliques + + scopedScalingGroups := scopedPCS.Spec.Template.PodCliqueScalingGroupConfigs[:0] + for _, scalingGroup := range scopedPCS.Spec.Template.PodCliqueScalingGroupConfigs { + cliqueNamesForBackend := make([]string, 0, len(scalingGroup.CliqueNames)) + for _, cliqueName := range scalingGroup.CliqueNames { + if _, ok := cliqueNames[cliqueName]; ok { + cliqueNamesForBackend = append(cliqueNamesForBackend, cliqueName) + } + } + if len(cliqueNamesForBackend) == 0 { + continue } - return fmt.Errorf("schedulerName %q is not enabled in OperatorConfiguration", schedulerName) + scalingGroup.CliqueNames = cliqueNamesForBackend + scopedScalingGroups = append(scopedScalingGroups, scalingGroup) } - return backend.ValidatePodCliqueSet(ctx, pcs) + scopedPCS.Spec.Template.PodCliqueScalingGroupConfigs = scopedScalingGroups + return scopedPCS } // castToPodCliqueSet attempts to cast a runtime.Object to a PodCliqueSet. diff --git a/operator/internal/webhook/admission/pcs/validation/handler_test.go b/operator/internal/webhook/admission/pcs/validation/handler_test.go index a1d5d1b26..6f0fdc053 100644 --- a/operator/internal/webhook/admission/pcs/validation/handler_test.go +++ b/operator/internal/webhook/admission/pcs/validation/handler_test.go @@ -247,6 +247,103 @@ func TestValidatePodCliqueSetWithLPXBackend(t *testing.T) { assert.Contains(t, err.Error(), "does not support Grove topology constraints") } +// TestValidatePodCliqueSetWithDisjointBackends documents that backend-specific +// validation must run for every scheduler selected by a PodClique, not just the +// scheduler on the first PodClique in the PodCliqueSet. +func TestValidatePodCliqueSetWithDisjointBackends(t *testing.T) { + profile := groveconfigv1alpha1.SchedulerProfile{Name: groveconfigv1alpha1.SchedulerNameLPX} + registry := &testutils.FakeSchedulerRegistry{ + Backends: map[string]scheduler.Backend{ + string(groveconfigv1alpha1.SchedulerNameKube): testutils.NewFakeSchedulerBackend( + string(groveconfigv1alpha1.SchedulerNameKube), + ), + string(groveconfigv1alpha1.SchedulerNameLPX): lpx.New(profile), + }, + DefaultBackend: string(groveconfigv1alpha1.SchedulerNameKube), + } + handler := &Handler{schedRegistry: registry} + pcs := testutils.NewPodCliqueSetBuilder("test-pcs", "default", uuid.NewUUID()). + WithPodCliqueTemplateSpec( + testutils.NewPodCliqueTemplateSpecBuilder("kube-worker"). + WithRoleName("kube-worker"). + WithReplicas(1). + WithPodSpec(corev1.PodSpec{ + SchedulerName: string(groveconfigv1alpha1.SchedulerNameKube), + Containers: []corev1.Container{{ + Name: "worker", + Image: "worker", + }}, + }). + Build(), + ). + WithPodCliqueTemplateSpec( + testutils.NewPodCliqueTemplateSpecBuilder("lpx-worker"). + WithRoleName("lpx-worker"). + WithReplicas(1). + WithPodSpec(corev1.PodSpec{ + SchedulerName: string(groveconfigv1alpha1.SchedulerNameLPX), + Containers: []corev1.Container{{ + Name: "worker", + Image: "worker", + }}, + }). + Build(), + ). + Build() + pcs.Spec.Template.TopologyConstraint = &grovecorev1alpha1.TopologyConstraint{} + + err := handler.validatePodCliqueSetWithBackend(context.Background(), pcs) + require.Error(t, err) + assert.Contains(t, err.Error(), "does not support Grove topology constraints") +} + +func TestValidatePodCliqueSetScopesCliquesToTheirBackend(t *testing.T) { + registry := &testutils.FakeSchedulerRegistry{ + Backends: map[string]scheduler.Backend{ + string(groveconfigv1alpha1.SchedulerNameKai): testutils.NewFakeSchedulerBackend( + string(groveconfigv1alpha1.SchedulerNameKai), + ), + string(groveconfigv1alpha1.SchedulerNameLPX): lpx.New( + groveconfigv1alpha1.SchedulerProfile{Name: groveconfigv1alpha1.SchedulerNameLPX}, + ), + }, + DefaultBackend: string(groveconfigv1alpha1.SchedulerNameKai), + } + handler := &Handler{schedRegistry: registry} + pcs := testutils.NewPodCliqueSetBuilder("test-pcs", "default", uuid.NewUUID()). + WithPodCliqueTemplateSpec( + testutils.NewPodCliqueTemplateSpecBuilder("kai-worker"). + WithRoleName("kai-worker"). + WithReplicas(1). + WithTopologyConstraint(&grovecorev1alpha1.TopologyConstraint{}). + WithPodSpec(corev1.PodSpec{ + SchedulerName: string(groveconfigv1alpha1.SchedulerNameKai), + Containers: []corev1.Container{{Name: "worker", Image: "worker"}}, + }). + Build(), + ). + WithPodCliqueTemplateSpec( + testutils.NewPodCliqueTemplateSpecBuilder("lpx-worker"). + WithRoleName("lpx-worker"). + WithReplicas(1). + WithPodSpec(corev1.PodSpec{ + SchedulerName: string(groveconfigv1alpha1.SchedulerNameLPX), + Containers: []corev1.Container{{Name: "worker", Image: "worker"}}, + }). + Build(), + ). + Build() + pcs.Spec.Template.PodCliqueScalingGroupConfigs = []grovecorev1alpha1.PodCliqueScalingGroupConfig{ + { + Name: "kai-workers", + CliqueNames: []string{"kai-worker"}, + TopologyConstraint: &grovecorev1alpha1.TopologyConstraint{}, + }, + } + + require.NoError(t, handler.validatePodCliqueSetWithBackend(context.Background(), pcs)) +} + // TestValidateUpdate tests validation of PodCliqueSet update requests. func TestValidateUpdate(t *testing.T) { startupType := ptr.To(grovecorev1alpha1.CliqueStartupTypeAnyOrder) diff --git a/operator/internal/webhook/admission/pcs/validation/podcliqueset.go b/operator/internal/webhook/admission/pcs/validation/podcliqueset.go index 5d6758b78..1b1507c23 100644 --- a/operator/internal/webhook/admission/pcs/validation/podcliqueset.go +++ b/operator/internal/webhook/admission/pcs/validation/podcliqueset.go @@ -275,35 +275,34 @@ func (v *pcsValidator) validatePodCliqueTemplates(fldPath *field.Path) ([]string return warnings, allErrs } -// validateSchedulerNames ensures all pod scheduler names resolve to the same scheduler and that scheduler is enabled. +// validateSchedulerNames ensures every scheduler selected by a PodClique is enabled. // Empty schedulerName is resolved to the default backend name from the injected registry. func (v *pcsValidator) validateSchedulerNames(schedulerNames []string, fldPath *field.Path) field.ErrorList { allErrs := field.ErrorList{} specPath := fldPath.Child("spec").Child("podSpec").Child("schedulerName") - defaultSchedulerName := v.schedRegistry.GetDefault().Name() + defaultBackend := v.schedRegistry.GetDefault() + defaultSchedulerName := "" + if defaultBackend != nil { + defaultSchedulerName = defaultBackend.Name() + } - // Check-1: Check if the scheduler names are unique - // Resolve empty to default backend name; then require all resolved names to be the same. + // Resolve empty scheduler names to the default backend, then validate each + // distinct scheduler once. uniqueSchedulerNames := lo.Uniq(lo.Map(schedulerNames, func(item string, _ int) string { if item == "" { return defaultSchedulerName } return item })) - if len(uniqueSchedulerNames) > 1 { - allErrs = append(allErrs, field.Invalid(specPath, strings.Join(uniqueSchedulerNames, ", "), "the schedulerName for all pods have to be the same")) - } - - // Check-2: Validate that the resolved scheduler is enabled. - pcsSchedulerName := uniqueSchedulerNames[0] - // default-scheduler is always valid; any other name must appear in the registry of enabled OperatorConfiguration backends. - if v.schedRegistry.Get(pcsSchedulerName) == nil { - allErrs = append(allErrs, field.Invalid( - specPath, - pcsSchedulerName, - "schedulerName must be an enabled scheduler backend; this scheduler is not enabled in OperatorConfiguration", - )) + for _, schedulerName := range uniqueSchedulerNames { + if schedulerName == "" || v.schedRegistry.Get(schedulerName) == nil { + allErrs = append(allErrs, field.Invalid( + specPath, + schedulerName, + "schedulerName must be an enabled scheduler backend; this scheduler is not enabled in OperatorConfiguration", + )) + } } return allErrs } diff --git a/operator/internal/webhook/admission/pcs/validation/podcliqueset_test.go b/operator/internal/webhook/admission/pcs/validation/podcliqueset_test.go index 16a457b13..5eea0186d 100644 --- a/operator/internal/webhook/admission/pcs/validation/podcliqueset_test.go +++ b/operator/internal/webhook/admission/pcs/validation/podcliqueset_test.go @@ -167,7 +167,6 @@ func TestValidateSchedulerNames(t *testing.T) { schedulerConfig groveconfigv1alpha1.SchedulerConfiguration schedulerNames []string expectErrors int - expectInvalidSame bool expectInvalidEnabled bool }{ { @@ -219,7 +218,7 @@ func TestValidateSchedulerNames(t *testing.T) { expectErrors: 0, }, { - name: "mixed default-scheduler and kai-scheduler", + name: "different enabled scheduler names", schedulerConfig: groveconfigv1alpha1.SchedulerConfiguration{ Profiles: []groveconfigv1alpha1.SchedulerProfile{ {Name: groveconfigv1alpha1.SchedulerNameKube}, @@ -227,10 +226,8 @@ func TestValidateSchedulerNames(t *testing.T) { }, DefaultProfileName: string(groveconfigv1alpha1.SchedulerNameKube), }, - schedulerNames: []string{"default-scheduler", "kai-scheduler"}, - expectErrors: 1, - expectInvalidSame: true, - expectInvalidEnabled: false, + schedulerNames: []string{"default-scheduler", "kai-scheduler"}, + expectErrors: 0, }, { name: "single kai-scheduler when enabled (kube+kai)", @@ -264,7 +261,6 @@ func TestValidateSchedulerNames(t *testing.T) { }, schedulerNames: []string{"kai-scheduler"}, expectErrors: 1, - expectInvalidSame: false, expectInvalidEnabled: true, }, { @@ -278,11 +274,22 @@ func TestValidateSchedulerNames(t *testing.T) { }, schedulerNames: []string{"volcano"}, expectErrors: 1, - expectInvalidSame: false, expectInvalidEnabled: true, }, { - name: "mixed empty and kai when default is default-scheduler", + name: "empty resolves to default alongside a different enabled scheduler", + schedulerConfig: groveconfigv1alpha1.SchedulerConfiguration{ + Profiles: []groveconfigv1alpha1.SchedulerProfile{ + {Name: groveconfigv1alpha1.SchedulerNameKube}, + {Name: groveconfigv1alpha1.SchedulerNameKai}, + }, + DefaultProfileName: string(groveconfigv1alpha1.SchedulerNameKube), + }, + schedulerNames: []string{"", "kai-scheduler"}, + expectErrors: 0, + }, + { + name: "each distinct scheduler must be enabled", schedulerConfig: groveconfigv1alpha1.SchedulerConfiguration{ Profiles: []groveconfigv1alpha1.SchedulerProfile{ {Name: groveconfigv1alpha1.SchedulerNameKube}, @@ -290,10 +297,9 @@ func TestValidateSchedulerNames(t *testing.T) { }, DefaultProfileName: string(groveconfigv1alpha1.SchedulerNameKube), }, - schedulerNames: []string{"", "kai-scheduler"}, + schedulerNames: []string{"default-scheduler", "volcano"}, expectErrors: 1, - expectInvalidSame: true, - expectInvalidEnabled: false, + expectInvalidEnabled: true, }, } for _, tt := range tests { @@ -319,9 +325,6 @@ func TestValidateSchedulerNames(t *testing.T) { assert.Len(t, errs, tt.expectErrors, "validation errors: %v", errs) if tt.expectErrors > 0 { msgs := lo.Map(errs, func(e *field.Error, _ int) string { return e.ErrorBody() }) - if tt.expectInvalidSame { - assert.Contains(t, strings.Join(msgs, " "), "have to be the same") - } if tt.expectInvalidEnabled { assert.Contains(t, strings.Join(msgs, " "), "not enabled") }