From cabb5696304c9f2754d49b912da6561fb2c8b144 Mon Sep 17 00:00:00 2001 From: Alay Patel Date: Thu, 9 Jul 2026 16:51:08 -0400 Subject: [PATCH 1/2] feat: add Kueue prebuilt workload backend Add a Kueue scheduler backend path that creates prebuilt Workloads from Grove PodGangs so PodGroup minReplicas can map to Kueue PodSet minCount. Require PodCliqueSet label kueue.x-k8s.io/queue-name for queue selection and propagate it to PodCliques/Pods. Include focused demo manifests for standalone minCount and PCSG all-or-nothing behavior. Signed-off-by: Alay Patel --- .../demo/grove-kind-topology.yaml | 11 + .../demo/grove-kueue-pcsg.yaml | 57 +++ .../demo/grove-kueue-simple-mincount.yaml | 34 ++ .../demo/kueue-kind-queue.yaml | 44 ++ .../demo/validate-grove-kueue-mincount.sh | 82 ++++ .../demo/validate-grove-kueue-pcsg.sh | 104 +++++ operator/api/config/v1alpha1/types.go | 17 +- operator/charts/templates/clusterrole.yaml | 15 + operator/charts/values.yaml | 3 +- .../components/pod/kueue_poc_test.go | 143 ++++++ .../podclique/components/pod/pod.go | 62 ++- .../podclique/components/pod/syncflow.go | 1 - .../components/podclique/podclique.go | 6 + .../components/podclique/podclique_test.go | 34 ++ operator/internal/scheduler/kueue/backend.go | 328 ++++++++++++++ .../internal/scheduler/kueue/backend_test.go | 417 ++++++++++++++++++ operator/internal/scheduler/kueue/topology.go | 142 ++++++ .../internal/scheduler/registry/registry.go | 3 + .../scheduler/registry/registry_test.go | 8 + 19 files changed, 1505 insertions(+), 6 deletions(-) create mode 100644 docs/community/kueue-tas-backend-slides/demo/grove-kind-topology.yaml create mode 100644 docs/community/kueue-tas-backend-slides/demo/grove-kueue-pcsg.yaml create mode 100644 docs/community/kueue-tas-backend-slides/demo/grove-kueue-simple-mincount.yaml create mode 100644 docs/community/kueue-tas-backend-slides/demo/kueue-kind-queue.yaml create mode 100755 docs/community/kueue-tas-backend-slides/demo/validate-grove-kueue-mincount.sh create mode 100755 docs/community/kueue-tas-backend-slides/demo/validate-grove-kueue-pcsg.sh create mode 100644 operator/internal/controller/podclique/components/pod/kueue_poc_test.go create mode 100644 operator/internal/scheduler/kueue/backend.go create mode 100644 operator/internal/scheduler/kueue/backend_test.go create mode 100644 operator/internal/scheduler/kueue/topology.go diff --git a/docs/community/kueue-tas-backend-slides/demo/grove-kind-topology.yaml b/docs/community/kueue-tas-backend-slides/demo/grove-kind-topology.yaml new file mode 100644 index 000000000..10c25eab1 --- /dev/null +++ b/docs/community/kueue-tas-backend-slides/demo/grove-kind-topology.yaml @@ -0,0 +1,11 @@ +# Grove source-of-truth topology for the CPU-only kind TAS smoke test. +apiVersion: grove.io/v1alpha1 +kind: ClusterTopologyBinding +metadata: + name: grove-kind-topology +spec: + levels: + - domain: rack + key: topology.ai-dynamo.io/rack + - domain: host + key: kubernetes.io/hostname diff --git a/docs/community/kueue-tas-backend-slides/demo/grove-kueue-pcsg.yaml b/docs/community/kueue-tas-backend-slides/demo/grove-kueue-pcsg.yaml new file mode 100644 index 000000000..4a6b47278 --- /dev/null +++ b/docs/community/kueue-tas-backend-slides/demo/grove-kueue-pcsg.yaml @@ -0,0 +1,57 @@ +# Grove PodCliqueSet with a PodCliqueScalingGroup (PCSG). +# POC assumption: PCSG.minAvailable == PCSG.replicas, so the whole gang is +# all-or-nothing. Grove builds one prebuilt Kueue Workload per PodGang and +# omits minCount for the PCSG podSets. +# Apply after kueue-kind-queue.yaml and grove-kind-topology.yaml. +apiVersion: grove.io/v1alpha1 +kind: PodCliqueSet +metadata: + name: grove-kueue-pcsg + labels: + app: grove-kueue-pcsg + kueue.x-k8s.io/queue-name: grove-poc +spec: + replicas: 1 + template: + topologyConstraint: + topologyName: grove-kind-topology + pack: + required: rack + podCliqueScalingGroups: + - name: decode + replicas: 2 + minAvailable: 2 + cliqueNames: + - decode-leader + - decode-worker + cliques: + - name: decode-leader + spec: + roleName: decode-leader + replicas: 1 + minAvailable: 1 + podSpec: + schedulerName: kueue + terminationGracePeriodSeconds: 5 + containers: + - name: pause + image: registry.k8s.io/pause:3.9 + resources: + requests: + cpu: 50m + memory: 32Mi + - name: decode-worker + spec: + roleName: decode-worker + replicas: 2 + minAvailable: 2 + podSpec: + schedulerName: kueue + terminationGracePeriodSeconds: 5 + containers: + - name: pause + image: registry.k8s.io/pause:3.9 + resources: + requests: + cpu: 50m + memory: 32Mi diff --git a/docs/community/kueue-tas-backend-slides/demo/grove-kueue-simple-mincount.yaml b/docs/community/kueue-tas-backend-slides/demo/grove-kueue-simple-mincount.yaml new file mode 100644 index 000000000..8f7ddf32b --- /dev/null +++ b/docs/community/kueue-tas-backend-slides/demo/grove-kueue-simple-mincount.yaml @@ -0,0 +1,34 @@ +# Grove PodCliqueSet with minCount semantics and no PodCliqueScalingGroup. +# Grove's Kueue backend builds one prebuilt Kueue Workload per PodGang, with a +# podSet whose count == clique replicas and minCount == clique minAvailable. +# Apply after kueue-kind-queue.yaml and grove-kind-topology.yaml. +apiVersion: grove.io/v1alpha1 +kind: PodCliqueSet +metadata: + name: grove-kueue-mincount + labels: + app: grove-kueue-mincount + kueue.x-k8s.io/queue-name: grove-poc +spec: + replicas: 1 + template: + topologyConstraint: + topologyName: grove-kind-topology + pack: + required: rack + cliques: + - name: frontend + spec: + roleName: frontend + replicas: 4 + minAvailable: 2 + podSpec: + schedulerName: kueue + terminationGracePeriodSeconds: 5 + containers: + - name: pause + image: registry.k8s.io/pause:3.9 + resources: + requests: + cpu: 100m + memory: 64Mi diff --git a/docs/community/kueue-tas-backend-slides/demo/kueue-kind-queue.yaml b/docs/community/kueue-tas-backend-slides/demo/kueue-kind-queue.yaml new file mode 100644 index 000000000..99b21d02a --- /dev/null +++ b/docs/community/kueue-tas-backend-slides/demo/kueue-kind-queue.yaml @@ -0,0 +1,44 @@ +# CPU-only Kueue queue setup for kind smoke tests. +# This is intentionally not the slide-deck GPU/TAS platform example. +apiVersion: kueue.x-k8s.io/v1beta2 +kind: ResourceFlavor +metadata: + name: kind-cpu +spec: + nodeLabels: + topology.ai-dynamo.io/tas-node: "true" + topologyName: grove-kind-topology +--- +apiVersion: kueue.x-k8s.io/v1beta2 +kind: ClusterQueue +metadata: + name: kind-cpu +spec: + namespaceSelector: {} + resourceGroups: + - coveredResources: + - cpu + - memory + flavors: + - name: kind-cpu + resources: + - name: cpu + nominalQuota: "4" + - name: memory + nominalQuota: 8Gi +--- +apiVersion: kueue.x-k8s.io/v1beta2 +kind: LocalQueue +metadata: + name: production-gpu + namespace: default +spec: + clusterQueue: kind-cpu +--- +apiVersion: kueue.x-k8s.io/v1beta2 +kind: LocalQueue +metadata: + name: grove-poc + namespace: default +spec: + clusterQueue: kind-cpu diff --git a/docs/community/kueue-tas-backend-slides/demo/validate-grove-kueue-mincount.sh b/docs/community/kueue-tas-backend-slides/demo/validate-grove-kueue-mincount.sh new file mode 100755 index 000000000..4aceddb2d --- /dev/null +++ b/docs/community/kueue-tas-backend-slides/demo/validate-grove-kueue-mincount.sh @@ -0,0 +1,82 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Validates the Grove Kueue "simple PCS with minCount" path: Grove builds one +# prebuilt Kueue Workload per PodGang with a podSet whose count == clique +# replicas and minCount == clique minAvailable, and stamps the Pods with +# kueue.x-k8s.io/prebuilt-workload-name. + +KUBECTL="${KUBECTL:-kubectl}" +NAMESPACE="${NAMESPACE:-default}" +TIMEOUT_SECONDS="${TIMEOUT_SECONDS:-120}" +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +MANIFEST="${MANIFEST:-${SCRIPT_DIR}/grove-kueue-simple-mincount.yaml}" +PCS_NAME="grove-kueue-mincount" +PODGANG_NAME="${PCS_NAME}-0" + +cleanup() { + "${KUBECTL}" delete -f "${MANIFEST}" --ignore-not-found >/dev/null 2>&1 || true + "${KUBECTL}" -n "${NAMESPACE}" delete workload "${PODGANG_NAME}" --ignore-not-found >/dev/null 2>&1 || true +} + +wait_for_jsonpath() { + local name="$1" + local jsonpath="$2" + local expected="$3" + local timeout_seconds="${4:-120}" + local deadline=$((SECONDS + timeout_seconds)) + + while (( SECONDS < deadline )); do + local actual + actual="$("${KUBECTL}" -n "${NAMESPACE}" get "${name}" -o "jsonpath=${jsonpath}" 2>/dev/null || true)" + if [[ "${actual}" == "${expected}" ]]; then + return 0 + fi + sleep 2 + done + + echo "Timed out waiting for ${name} jsonpath ${jsonpath} to equal ${expected}" >&2 + return 1 +} + +cleanup + +"${KUBECTL}" apply -f "${MANIFEST}" >/dev/null + +# Grove builds one prebuilt Workload named after the PodGang. +wait_for_jsonpath \ + "workload/${PODGANG_NAME}" \ + "{.metadata.name}" \ + "${PODGANG_NAME}" \ + "${TIMEOUT_SECONDS}" + +wait_for_jsonpath "workload/${PODGANG_NAME}" "{.spec.podSets[0].count}" "4" "${TIMEOUT_SECONDS}" +wait_for_jsonpath "workload/${PODGANG_NAME}" "{.spec.podSets[0].minCount}" "2" "${TIMEOUT_SECONDS}" + +# Pods reference the prebuilt Workload. +pod="$("${KUBECTL}" -n "${NAMESPACE}" get pods \ + -l "app.kubernetes.io/part-of=${PCS_NAME}" \ + -o "jsonpath={.items[0].metadata.name}")" +if [[ -z "${pod}" ]]; then + echo "Expected at least one Pod for PodCliqueSet ${PCS_NAME}" >&2 + exit 1 +fi + +prebuilt="$("${KUBECTL}" -n "${NAMESPACE}" get pod "${pod}" \ + -o "jsonpath={.metadata.labels.kueue\\.x-k8s\\.io/prebuilt-workload-name}")" +if [[ "${prebuilt}" != "${PODGANG_NAME}" ]]; then + echo "Expected Pod ${pod} to carry kueue.x-k8s.io/prebuilt-workload-name=${PODGANG_NAME}, got '${prebuilt}'" >&2 + exit 1 +fi + +echo "PASS: Grove created one prebuilt Workload named ${PODGANG_NAME}." +echo "PASS: the Workload podSet has count=4 and minCount=2." +echo "PASS: Pods reference the prebuilt Workload via kueue.x-k8s.io/prebuilt-workload-name." +echo +echo "Inspect:" +echo " ${KUBECTL} -n ${NAMESPACE} get workload ${PODGANG_NAME} -o yaml" +echo " ${KUBECTL} -n ${NAMESPACE} get pods -l app.kubernetes.io/part-of=${PCS_NAME} -o wide" +echo +echo "Cleanup:" +echo " ${KUBECTL} delete -f ${MANIFEST} --ignore-not-found" +echo " ${KUBECTL} -n ${NAMESPACE} delete workload ${PODGANG_NAME} --ignore-not-found" diff --git a/docs/community/kueue-tas-backend-slides/demo/validate-grove-kueue-pcsg.sh b/docs/community/kueue-tas-backend-slides/demo/validate-grove-kueue-pcsg.sh new file mode 100755 index 000000000..f7873949e --- /dev/null +++ b/docs/community/kueue-tas-backend-slides/demo/validate-grove-kueue-pcsg.sh @@ -0,0 +1,104 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Validates the Grove Kueue "PCS with PCSG" path. Grove pre-builds one Kueue +# Workload per PodGang for every PodCliqueSet. PodCliqueScalingGroup cliques are +# all-or-nothing: their podSets use minCount == count (POC assumption +# PCSG.minAvailable == PCSG.replicas). Pods reference the prebuilt Workload. + +KUBECTL="${KUBECTL:-kubectl}" +NAMESPACE="${NAMESPACE:-default}" +TIMEOUT_SECONDS="${TIMEOUT_SECONDS:-120}" +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +MANIFEST="${MANIFEST:-${SCRIPT_DIR}/grove-kueue-pcsg.yaml}" +PCS_NAME="grove-kueue-pcsg" +PODGANG_NAME="${PCS_NAME}-0" +# decode-leader(1) + decode-worker(2) per PCSG replica, times 2 replicas = 6. +EXPECTED_TOTAL="6" + +cleanup() { + "${KUBECTL}" delete -f "${MANIFEST}" --ignore-not-found >/dev/null 2>&1 || true + "${KUBECTL}" -n "${NAMESPACE}" delete workload "${PODGANG_NAME}" --ignore-not-found >/dev/null 2>&1 || true +} + +wait_for_jsonpath() { + local name="$1" + local jsonpath="$2" + local expected="$3" + local timeout_seconds="${4:-120}" + local deadline=$((SECONDS + timeout_seconds)) + + while (( SECONDS < deadline )); do + local actual + actual="$("${KUBECTL}" -n "${NAMESPACE}" get "${name}" -o "jsonpath=${jsonpath}" 2>/dev/null || true)" + if [[ "${actual}" == "${expected}" ]]; then + return 0 + fi + sleep 2 + done + + echo "Timed out waiting for ${name} jsonpath ${jsonpath} to equal ${expected}" >&2 + return 1 +} + +cleanup + +"${KUBECTL}" apply -f "${MANIFEST}" >/dev/null + +# Grove builds one prebuilt Workload named after the PodGang, even for PCSG. +wait_for_jsonpath \ + "workload/${PODGANG_NAME}" \ + "{.metadata.name}" \ + "${PODGANG_NAME}" \ + "${TIMEOUT_SECONDS}" + +# Every PCSG podSet is all-or-nothing: minCount is omitted (Kueue defaults it to count, and Kueue rejects +# Workloads where more than one podSet sets minCount). +min_counts="$("${KUBECTL}" -n "${NAMESPACE}" get "workload/${PODGANG_NAME}" \ + -o "jsonpath={range .spec.podSets[*]}{.minCount}{' '}{end}" | tr -d ' ')" +if [[ -n "${min_counts}" ]]; then + echo "Expected no PCSG podSet to set minCount (all-or-nothing), got minCounts='${min_counts}'" >&2 + exit 1 +fi + +# Pods reference the prebuilt Workload. +pod="$("${KUBECTL}" -n "${NAMESPACE}" get pods \ + -l "app.kubernetes.io/part-of=${PCS_NAME}" \ + -o "jsonpath={.items[0].metadata.name}")" +if [[ -z "${pod}" ]]; then + echo "Expected at least one Pod for PodCliqueSet ${PCS_NAME}" >&2 + exit 1 +fi + +group="$("${KUBECTL}" -n "${NAMESPACE}" get pod "${pod}" \ + -o "jsonpath={.metadata.labels.kueue\\.x-k8s\\.io/pod-group-name}")" +if [[ "${group}" != "${PODGANG_NAME}" ]]; then + echo "Expected Pod ${pod} to use kueue.x-k8s.io/pod-group-name=${PODGANG_NAME}, got '${group}'" >&2 + exit 1 +fi + +prebuilt="$("${KUBECTL}" -n "${NAMESPACE}" get pod "${pod}" \ + -o "jsonpath={.metadata.labels.kueue\\.x-k8s\\.io/prebuilt-workload-name}")" +if [[ "${prebuilt}" != "${PODGANG_NAME}" ]]; then + echo "Expected Pod ${pod} to carry kueue.x-k8s.io/prebuilt-workload-name=${PODGANG_NAME}, got '${prebuilt}'" >&2 + exit 1 +fi + +total="$("${KUBECTL}" -n "${NAMESPACE}" get pod "${pod}" \ + -o "jsonpath={.metadata.annotations.kueue\\.x-k8s\\.io/pod-group-total-count}")" +if [[ "${total}" != "${EXPECTED_TOTAL}" ]]; then + echo "Expected pod-group-total-count=${EXPECTED_TOTAL}, got '${total}'" >&2 + exit 1 +fi + +echo "PASS: Grove created one prebuilt Workload named ${PODGANG_NAME} for the PCSG PodCliqueSet." +echo "PASS: every PCSG podSet is all-or-nothing (minCount omitted; Kueue allows minCount on at most one podSet)." +echo "PASS: Pods reference the prebuilt Workload and share pod-group-name=${PODGANG_NAME} (total ${EXPECTED_TOTAL})." +echo +echo "Inspect:" +echo " ${KUBECTL} -n ${NAMESPACE} get workload ${PODGANG_NAME} -o yaml" +echo " ${KUBECTL} -n ${NAMESPACE} get pods -l app.kubernetes.io/part-of=${PCS_NAME} -o wide" +echo +echo "Cleanup:" +echo " ${KUBECTL} delete -f ${MANIFEST} --ignore-not-found" +echo " ${KUBECTL} -n ${NAMESPACE} delete workload ${PODGANG_NAME} --ignore-not-found" diff --git a/operator/api/config/v1alpha1/types.go b/operator/api/config/v1alpha1/types.go index 4857f3864..c0c1171bf 100644 --- a/operator/api/config/v1alpha1/types.go +++ b/operator/api/config/v1alpha1/types.go @@ -56,6 +56,8 @@ type SchedulerName string const ( // SchedulerNameKai is the KAI scheduler backend. SchedulerNameKai SchedulerName = "kai-scheduler" + // SchedulerNameKueue is the Kueue plain Pod integration backend. + SchedulerNameKueue SchedulerName = "kueue" // SchedulerNameKube is the profile name for the Kubernetes default scheduler in OperatorConfiguration. SchedulerNameKube SchedulerName = "default-scheduler" // SchedulerNameVolcano is the Volcano scheduler backend. It supports gang scheduling via Volcano PodGroup. @@ -68,6 +70,7 @@ var ( // SupportedSchedulerNames is the list of profile names allowed in scheduler.profiles[].name. SupportedSchedulerNames = []SchedulerName{ SchedulerNameKai, + SchedulerNameKueue, SchedulerNameKube, SchedulerNameVolcano, SchedulerNameLPX, @@ -79,7 +82,7 @@ type SchedulerConfiguration struct { // Profiles is the list of scheduler profiles. Each profile has a backend name and an optional config. // The default-scheduler backend is always enabled to ensure that the kubernetes default scheduler is always enabled and supported. // Use profile name "default-scheduler" to configure or set it as default. - // Valid profile names: "default-scheduler", "kai-scheduler", "volcano", "lpx-scheduler". + // Valid profile names: "default-scheduler", "kai-scheduler", "kueue", "volcano", "lpx-scheduler". // Use defaultProfileName to designate the default backend. // +optional Profiles []SchedulerProfile `json:"profiles,omitempty"` @@ -95,7 +98,7 @@ type SchedulerProfile struct { // For the Kubernetes default scheduler use the standard "default-scheduler". // Ensure that the name chosen is a valid scheduler name. The name will also be directly set in `Pod.Spec.SchedulerName`. // +kubebuilder:validation:Required - // +kubebuilder:validation:Enum=kai-scheduler;default-scheduler;volcano;lpx-scheduler + // +kubebuilder:validation:Enum=kai-scheduler;kueue;default-scheduler;volcano;lpx-scheduler Name SchedulerName `json:"name"` // Config holds backend-specific options. The operator unmarshals it into the config type for this backend (see backend config types). @@ -108,6 +111,16 @@ type KaiSchedulerConfiguration struct { // Reserved for future kai-scheduler-specific options. } +// KueueSchedulerConfiguration defines the configuration for the kueue backend. +type KueueSchedulerConfiguration struct { + // RequiredTopologyKey is stamped as kueue.x-k8s.io/podset-required-topology when set. + // +optional + RequiredTopologyKey string `json:"requiredTopologyKey,omitempty"` + // UnderlyingSchedulerName is the schedulerName used after Kueue admits the Pods. + // +optional + UnderlyingSchedulerName string `json:"underlyingSchedulerName,omitempty"` +} + // KubeSchedulerConfig holds the configuration for the default scheduler. // Used when unmarshalling SchedulerProfile.Config for default-scheduler. type KubeSchedulerConfig struct { diff --git a/operator/charts/templates/clusterrole.yaml b/operator/charts/templates/clusterrole.yaml index 4d9b3c278..a33b11b41 100644 --- a/operator/charts/templates/clusterrole.yaml +++ b/operator/charts/templates/clusterrole.yaml @@ -145,6 +145,21 @@ rules: - update - delete {{- end }} +{{- if include "grove.scheduler.hasProfile" (list . "kueue") }} +- apiGroups: + - kueue.x-k8s.io + resources: + - topologies + - workloads + verbs: + - create + - get + - list + - watch + - patch + - update + - delete +{{- end }} # DRA (Dynamic Resource Allocation) integration for resource sharing - apiGroups: - resource.k8s.io diff --git a/operator/charts/values.yaml b/operator/charts/values.yaml index 1a99cfdcf..bc750e07c 100644 --- a/operator/charts/values.yaml +++ b/operator/charts/values.yaml @@ -99,13 +99,14 @@ config: profiles: - name: default-scheduler - name: kai-scheduler + - name: kueue # Uncomment optional scheduler profiles only when their dependencies are installed. # - name: volcano # - name: lpx-scheduler logLevel: info logFormat: json topologyAwareScheduling: - enabled: false + enabled: true authorizer: enabled: false exemptServiceAccountUserNames: diff --git a/operator/internal/controller/podclique/components/pod/kueue_poc_test.go b/operator/internal/controller/podclique/components/pod/kueue_poc_test.go new file mode 100644 index 000000000..9f9592e51 --- /dev/null +++ b/operator/internal/controller/podclique/components/pod/kueue_poc_test.go @@ -0,0 +1,143 @@ +// /* +// 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 pod + +import ( + "testing" + + apicommon "github.com/ai-dynamo/grove/operator/api/common" + 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" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "sigs.k8s.io/controller-runtime/pkg/client" +) + +func newKueuePOCResource(cl client.Client) _resource { + return _resource{ + client: cl, + scheme: cl.Scheme(), + schedRegistry: &testutils.FakeSchedulerRegistry{ + Backends: map[string]scheduler.Backend{ + string(configv1alpha1.SchedulerNameKueue): testutils.NewFakeSchedulerBackend(string(configv1alpha1.SchedulerNameKueue)), + }, + DefaultBackend: string(configv1alpha1.SchedulerNameKueue), + }, + } +} + +func TestBuildResource_StampsPrebuiltWorkloadMetadataForSimplePCS(t *testing.T) { + cl := testutils.CreateDefaultFakeClient(nil) + r := newKueuePOCResource(cl) + pcs := &grovecorev1alpha1.PodCliqueSet{ + ObjectMeta: metav1.ObjectMeta{ + Name: "demo", + Namespace: "default", + Labels: map[string]string{"kueue.x-k8s.io/queue-name": "grove-poc"}, + }, + Spec: grovecorev1alpha1.PodCliqueSetSpec{ + Template: grovecorev1alpha1.PodCliqueSetTemplateSpec{ + Cliques: []*grovecorev1alpha1.PodCliqueTemplateSpec{ + {Name: "prefill", Spec: grovecorev1alpha1.PodCliqueSpec{Replicas: 2}}, + {Name: "worker", Spec: grovecorev1alpha1.PodCliqueSpec{Replicas: 3}}, + }, + }, + }, + } + pclq := &grovecorev1alpha1.PodClique{ + ObjectMeta: metav1.ObjectMeta{ + Name: "demo-0-worker", + Namespace: "default", + Labels: map[string]string{ + apicommon.LabelPartOfKey: "demo", + apicommon.LabelPodCliqueSetReplicaIndex: "0", + apicommon.LabelPodGang: "demo-0", + "kueue.x-k8s.io/queue-name": "grove-poc", + }, + }, + Spec: grovecorev1alpha1.PodCliqueSpec{ + Replicas: 3, + PodSpec: corev1.PodSpec{SchedulerName: string(configv1alpha1.SchedulerNameKueue)}, + }, + } + pod := &corev1.Pod{} + + require.NoError(t, r.buildResource(pcs, pclq, "demo-0", pod, 0)) + + assert.Equal(t, "5", pod.Annotations[kueuePodGroupTotalCountAnnotation]) + assert.Equal(t, "demo-0", pod.Labels[kueuePodGroupNameLabel]) + assert.Equal(t, "demo-0", pod.Labels[kueuePrebuiltWorkloadNameLabel]) + assert.Equal(t, "grove-poc", pod.Labels["kueue.x-k8s.io/queue-name"]) + assert.Equal(t, "demo-0-worker", pod.Annotations[kueueRoleHashAnnotation]) +} + +func TestBuildResource_StampsPrebuiltWorkloadMetadataForPCSG(t *testing.T) { + cl := testutils.CreateDefaultFakeClient(nil) + r := newKueuePOCResource(cl) + pcsgReplicas := int32(2) + pcs := &grovecorev1alpha1.PodCliqueSet{ + ObjectMeta: metav1.ObjectMeta{ + Name: "demo", + Namespace: "default", + Labels: map[string]string{"kueue.x-k8s.io/queue-name": "grove-poc"}, + }, + Spec: grovecorev1alpha1.PodCliqueSetSpec{ + Template: grovecorev1alpha1.PodCliqueSetTemplateSpec{ + Cliques: []*grovecorev1alpha1.PodCliqueTemplateSpec{ + {Name: "leader", Spec: grovecorev1alpha1.PodCliqueSpec{Replicas: 1}}, + {Name: "worker", Spec: grovecorev1alpha1.PodCliqueSpec{Replicas: 2}}, + }, + PodCliqueScalingGroupConfigs: []grovecorev1alpha1.PodCliqueScalingGroupConfig{ + {Name: "decode", CliqueNames: []string{"leader", "worker"}, Replicas: &pcsgReplicas}, + }, + }, + }, + } + pclq := &grovecorev1alpha1.PodClique{ + ObjectMeta: metav1.ObjectMeta{ + Name: "demo-0-decode-0-worker", + Namespace: "default", + Labels: map[string]string{ + apicommon.LabelPartOfKey: "demo", + apicommon.LabelPodCliqueSetReplicaIndex: "0", + apicommon.LabelPodGang: "demo-0", + apicommon.LabelPodCliqueScalingGroup: "demo-0-decode", + apicommon.LabelPodCliqueScalingGroupReplicaIndex: "0", + "kueue.x-k8s.io/queue-name": "grove-poc", + }, + }, + Spec: grovecorev1alpha1.PodCliqueSpec{ + Replicas: 2, + PodSpec: corev1.PodSpec{SchedulerName: string(configv1alpha1.SchedulerNameKueue)}, + }, + } + pod := &corev1.Pod{} + + require.NoError(t, r.buildResource(pcs, pclq, "demo-0", pod, 0)) + + // (leader 1 + worker 2) * pcsg replicas 2 = 6 pods in the PodGang. + assert.Equal(t, "6", pod.Annotations[kueuePodGroupTotalCountAnnotation]) + assert.Equal(t, "demo-0", pod.Labels[kueuePodGroupNameLabel]) + // Grove pre-builds the Workload for PCSG too; Pods reference it and the role-hash is the PodClique FQN. + assert.Equal(t, "demo-0", pod.Labels[kueuePrebuiltWorkloadNameLabel]) + assert.Equal(t, "grove-poc", pod.Labels["kueue.x-k8s.io/queue-name"]) + assert.Equal(t, "demo-0-decode-0-worker", pod.Annotations[kueueRoleHashAnnotation]) +} diff --git a/operator/internal/controller/podclique/components/pod/pod.go b/operator/internal/controller/podclique/components/pod/pod.go index 9452827a6..610736238 100644 --- a/operator/internal/controller/podclique/components/pod/pod.go +++ b/operator/internal/controller/podclique/components/pod/pod.go @@ -23,6 +23,7 @@ import ( apicommon "github.com/ai-dynamo/grove/operator/api/common" "github.com/ai-dynamo/grove/operator/api/common/constants" + 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/controller/common/component" componentutils "github.com/ai-dynamo/grove/operator/internal/controller/common/component/utils" @@ -67,7 +68,11 @@ const ( ) const ( - podGangSchedulingGate = "grove.io/podgang-pending-creation" + podGangSchedulingGate = "grove.io/podgang-pending-creation" + kueuePodGroupTotalCountAnnotation = "kueue.x-k8s.io/pod-group-total-count" + kueuePodGroupNameLabel = "kueue.x-k8s.io/pod-group-name" + kueuePrebuiltWorkloadNameLabel = "kueue.x-k8s.io/prebuilt-workload-name" + kueueRoleHashAnnotation = "kueue.x-k8s.io/role-hash" ) type _resource struct { @@ -153,7 +158,7 @@ func (r _resource) buildResource(pcs *grovecorev1alpha1.PodCliqueSet, pclq *grov GenerateName: fmt.Sprintf("%s-", pclq.Name), Namespace: pclq.Namespace, Labels: labels, - Annotations: pclq.Annotations, + Annotations: lo.Assign(map[string]string{}, pclq.Annotations), } if err = controllerutil.SetControllerReference(pclq, pod, r.scheme); err != nil { return groveerr.WrapError(err, @@ -177,6 +182,9 @@ func (r _resource) buildResource(pcs *grovecorev1alpha1.PodCliqueSet, pclq *grov "failed to prepare pod spec with scheduler backend", ) } + if backend.Name() == string(configv1alpha1.SchedulerNameKueue) { + stampKueuePodGroupMetadata(pod, pcs, pclq, podGangName) + } if err = backend.PreparePod(pod); err != nil { return groveerr.WrapError( err, @@ -269,6 +277,56 @@ func injectPCLQResourceClaimRefs(podSpec *corev1.PodSpec, pclqName string, resou resourceclaim.InjectResourceClaimRefs(podSpec, pclqName, resourceSharers, &podIndex) } +// stampKueuePodGroupMetadata stamps PodGang-scoped Kueue metadata on a Pod. Grove always builds a prebuilt Kueue +// Workload per PodGang, so every Kueue-managed Pod is grouped by the PodGang name and references the prebuilt +// Workload via the prebuilt-workload-name label. The role-hash annotation is set to the PodClique FQN, which +// matches the corresponding Workload podSet name (the PodGroup name) so Kueue can map Pods to podSets uniquely, +// including across PodCliqueScalingGroup replicas of the same clique. Kueue reads role-hash from annotations (not +// labels); stamping it here prevents Kueue from computing a pod-spec hash that would never match the podSet name. +func stampKueuePodGroupMetadata(pod *corev1.Pod, pcs *grovecorev1alpha1.PodCliqueSet, pclq *grovecorev1alpha1.PodClique, podGangName string) { + if pod.Labels == nil { + pod.Labels = make(map[string]string) + } + if pod.Annotations == nil { + pod.Annotations = make(map[string]string) + } + pod.Labels[kueuePodGroupNameLabel] = podGangName + pod.Labels[kueuePrebuiltWorkloadNameLabel] = podGangName + pod.Annotations[kueuePodGroupTotalCountAnnotation] = strconv.Itoa(kueuePodGangTotalPodCount(pcs)) + pod.Annotations[kueueRoleHashAnnotation] = pclq.Name +} + +// kueuePodGangTotalPodCount returns the total number of Pods in a single PodCliqueSet-replica PodGang. Cliques that +// belong to a scaling group contribute clique.Replicas * scalingGroup.Replicas Pods; standalone cliques contribute +// clique.Replicas Pods. This assumes each PodGang corresponds to one PodCliqueSet replica. +func kueuePodGangTotalPodCount(pcs *grovecorev1alpha1.PodCliqueSet) int { + scalingGroupCliques := make(map[string]struct{}) + for _, config := range pcs.Spec.Template.PodCliqueScalingGroupConfigs { + for _, cliqueName := range config.CliqueNames { + scalingGroupCliques[cliqueName] = struct{}{} + } + } + total := 0 + for _, cliqueTemplate := range pcs.Spec.Template.Cliques { + if _, ok := scalingGroupCliques[cliqueTemplate.Name]; ok { + continue + } + total += int(cliqueTemplate.Spec.Replicas) + } + for _, config := range pcs.Spec.Template.PodCliqueScalingGroupConfigs { + groupReplicas := 1 + if config.Replicas != nil { + groupReplicas = int(*config.Replicas) + } + for _, cliqueName := range config.CliqueNames { + if cliqueTemplate := componentutils.FindPodCliqueTemplateSpecByName(pcs, cliqueName); cliqueTemplate != nil { + total += int(cliqueTemplate.Spec.Replicas) * groupReplicas + } + } + } + return total +} + // Delete removes all Pods associated with the specified PodClique func (r _resource) Delete(ctx context.Context, logger logr.Logger, pclqObjectMeta metav1.ObjectMeta) error { logger.Info("Triggering delete of all pods for the PodClique") diff --git a/operator/internal/controller/podclique/components/pod/syncflow.go b/operator/internal/controller/podclique/components/pod/syncflow.go index 1e3d63f6d..38cdfc16d 100644 --- a/operator/internal/controller/podclique/components/pod/syncflow.go +++ b/operator/internal/controller/podclique/components/pod/syncflow.go @@ -86,7 +86,6 @@ func (r _resource) prepareSyncFlow(ctx context.Context, logger logr.Logger, pclq if err = lo.Ternary(apierrors.IsNotFound(err), nil, err); err != nil { return nil, err } - // initialize the Pod names that are updated in the PodGang resource for this PCLQ. sc.podNamesUpdatedInPCLQPodGangs = r.getPodNamesUpdatedInAssociatedPodGang(existingPodGang, pclq.Name) sc.podNamesUpdatedInPCLQPodGangSet = componentutils.NewSet(sc.podNamesUpdatedInPCLQPodGangs) diff --git a/operator/internal/controller/podcliqueset/components/podclique/podclique.go b/operator/internal/controller/podcliqueset/components/podclique/podclique.go index a7348958d..34b56a404 100644 --- a/operator/internal/controller/podcliqueset/components/podclique/podclique.go +++ b/operator/internal/controller/podcliqueset/components/podclique/podclique.go @@ -50,6 +50,7 @@ const ( errDeletePodClique grovecorev1alpha1.ErrorCode = "ERR_DELETE_PODCLIQUE" errCodeListPodCliques grovecorev1alpha1.ErrorCode = "ERR_LIST_PODCLIQUES" errCodeCreateOrUpdatePodClique grovecorev1alpha1.ErrorCode = "ERR_CREATE_OR_UPDATE_PODCLIQUE" + kueueQueueNameLabel string = "kueue.x-k8s.io/queue-name" ) type _resource struct { @@ -394,8 +395,13 @@ func getLabels(pcs *grovecorev1alpha1.PodCliqueSet, pcsReplica int, pclqObjectKe apicommon.LabelPodGang: podGangName, apicommon.LabelPodTemplateHash: componentutils.ComputePCLQPodTemplateHash(pclqTemplateSpec, pcs.Spec.Template.PriorityClassName), } + pcsQueueLabels := map[string]string{} + if queueName := pcs.Labels[kueueQueueNameLabel]; queueName != "" { + pcsQueueLabels[kueueQueueNameLabel] = queueName + } return lo.Assign( pclqTemplateSpec.Labels, + pcsQueueLabels, apicommon.GetDefaultLabelsForPodCliqueSetManagedResources(pcs.Name), pclqComponentLabels, ) diff --git a/operator/internal/controller/podcliqueset/components/podclique/podclique_test.go b/operator/internal/controller/podcliqueset/components/podclique/podclique_test.go index a706cf5d7..42ed8f855 100644 --- a/operator/internal/controller/podcliqueset/components/podclique/podclique_test.go +++ b/operator/internal/controller/podcliqueset/components/podclique/podclique_test.go @@ -124,6 +124,40 @@ func TestGetExistingResourceNames(t *testing.T) { } } +func TestGetLabels_PropagatesKueueQueueNameFromPodCliqueSet(t *testing.T) { + pcs := &grovecorev1alpha1.PodCliqueSet{ + ObjectMeta: metav1.ObjectMeta{ + Name: testPCSName, + Namespace: testPCSNamespace, + Labels: map[string]string{kueueQueueNameLabel: "grove-poc"}, + }, + Spec: grovecorev1alpha1.PodCliqueSetSpec{ + Template: grovecorev1alpha1.PodCliqueSetTemplateSpec{ + Cliques: []*grovecorev1alpha1.PodCliqueTemplateSpec{ + { + Name: "howl", + Spec: grovecorev1alpha1.PodCliqueSpec{ + PodSpec: corev1.PodSpec{ + Containers: []corev1.Container{{Name: "main", Image: "registry.k8s.io/pause:3.9"}}, + }, + }, + }, + }, + }, + }, + } + + labels := getLabels( + pcs, + 0, + client.ObjectKey{Name: "coyote-0-howl", Namespace: testPCSNamespace}, + pcs.Spec.Template.Cliques[0], + "coyote-0", + ) + + assert.Equal(t, "grove-poc", labels[kueueQueueNameLabel]) +} + func TestDelete(t *testing.T) { testCases := []struct { description string diff --git a/operator/internal/scheduler/kueue/backend.go b/operator/internal/scheduler/kueue/backend.go new file mode 100644 index 000000000..9d19dc4ad --- /dev/null +++ b/operator/internal/scheduler/kueue/backend.go @@ -0,0 +1,328 @@ +// /* +// 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 kueue + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + apicommon "github.com/ai-dynamo/grove/operator/api/common" + 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" + groveschedulerv1alpha1 "github.com/ai-dynamo/grove/scheduler/api/core/v1alpha1" + corev1 "k8s.io/api/core/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/runtime/schema" + "k8s.io/client-go/tools/record" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil" + "sigs.k8s.io/controller-runtime/pkg/log" +) + +const ( + queueNameLabel = "kueue.x-k8s.io/queue-name" + podGroupNameLabel = "kueue.x-k8s.io/pod-group-name" + podGroupServingAnnotation = "kueue.x-k8s.io/pod-group-serving" + retriableInGroupAnnotation = "kueue.x-k8s.io/retriable-in-group" + podSetRequiredTopologyAnnotation = "kueue.x-k8s.io/podset-required-topology" +) + +var workloadGVK = schema.GroupVersionKind{ + Group: "kueue.x-k8s.io", + Version: "v1beta2", + Kind: "Workload", +} + +type schedulerBackend struct { + client client.Client + scheme *runtime.Scheme + name string + eventRecorder record.EventRecorder + profile configv1alpha1.SchedulerProfile + config configv1alpha1.KueueSchedulerConfiguration +} + +var _ scheduler.Backend = (*schedulerBackend)(nil) + +func New(cl client.Client, scheme *runtime.Scheme, eventRecorder record.EventRecorder, profile configv1alpha1.SchedulerProfile) scheduler.Backend { + return &schedulerBackend{ + client: cl, + scheme: scheme, + name: string(profile.Name), + eventRecorder: eventRecorder, + profile: profile, + config: defaultConfig(), + } +} + +func (b *schedulerBackend) Name() string { + return b.name +} + +func (b *schedulerBackend) Init(_ client.Client) error { + if b.profile.Config == nil || len(b.profile.Config.Raw) == 0 { + return nil + } + if err := json.Unmarshal(b.profile.Config.Raw, &b.config); err != nil { + return fmt.Errorf("failed to unmarshal kueue scheduler config: %w", err) + } + b.config = withDefaults(b.config) + return nil +} + +// SyncPodGang builds a prebuilt Kueue Workload for every PodGang, mapping each PodGroup to a Kueue podSet. +// Standalone PodCliques express partial-gang semantics via minCount == minAvailable, while PodCliques that +// belong to a PodCliqueScalingGroup are treated as all-or-nothing (minCount == count) under the POC +// assumption that PCSG.minAvailable == PCSG.replicas. +func (b *schedulerBackend) SyncPodGang(ctx context.Context, podGang *groveschedulerv1alpha1.PodGang) error { + pcsName := podGang.Labels[apicommon.LabelPartOfKey] + if pcsName == "" { + return nil + } + pcs := &grovecorev1alpha1.PodCliqueSet{} + if err := b.client.Get(ctx, client.ObjectKey{Namespace: podGang.Namespace, Name: pcsName}, pcs); err != nil { + return fmt.Errorf("failed to get PodCliqueSet %s/%s for PodGang %s: %w", podGang.Namespace, pcsName, podGang.Name, err) + } + + desired, err := b.buildPrebuiltWorkload(pcs, podGang) + if err != nil { + return err + } + existing := newKueueWorkload(podGang.Namespace, podGang.Name) + if err = b.client.Get(ctx, client.ObjectKeyFromObject(existing), existing); err != nil { + if apierrors.IsNotFound(err) { + if err = b.client.Create(ctx, desired); err != nil { + return fmt.Errorf("failed to create prebuilt kueue Workload %s/%s: %w", podGang.Namespace, podGang.Name, err) + } + log.FromContext(ctx).Info("Created prebuilt Kueue Workload", "workload", client.ObjectKeyFromObject(desired)) + return nil + } + return fmt.Errorf("failed to get prebuilt kueue Workload %s/%s: %w", podGang.Namespace, podGang.Name, err) + } + // Kueue Workload podSets are immutable once admitted, so the Workload is only created and not updated in this POC. + return nil +} + +func newKueueWorkload(namespace, name string) *unstructured.Unstructured { + workload := &unstructured.Unstructured{} + workload.SetGroupVersionKind(workloadGVK) + workload.SetNamespace(namespace) + workload.SetName(name) + return workload +} + +// buildPrebuiltWorkload constructs a prebuilt Kueue Workload from a PodGang, mapping each PodGroup to a Kueue +// podSet. The podSet name is the PodGroup FQN so scaling-group replicas of the same clique stay unique. count +// is the clique replicas; minCount is the PodGroup minReplicas for standalone cliques and is forced to count +// for scaling-group cliques (POC assumption PCSG.minAvailable == PCSG.replicas). +func (b *schedulerBackend) buildPrebuiltWorkload(pcs *grovecorev1alpha1.PodCliqueSet, podGang *groveschedulerv1alpha1.PodGang) (*unstructured.Unstructured, error) { + queueName := pcs.Labels[queueNameLabel] + if queueName == "" { + return nil, fmt.Errorf("PodCliqueSet %s/%s must set label %q for Kueue queue selection", pcs.Namespace, pcs.Name, queueNameLabel) + } + + // Kueue permits at most one podSet per Workload to set minCount (partial admission is limited to a single + // podSet). Grove therefore only sets minCount on the first podSet whose minCount is strictly less than its + // count; all other podSets (including every PodCliqueScalingGroup clique, which is all-or-nothing in this POC) + // omit minCount, which Kueue defaults to count. + minCountUsed := false + podSets := make([]any, 0, len(podGang.Spec.PodGroups)) + for _, podGroup := range podGang.Spec.PodGroups { + cliqueTemplate := findCliqueTemplateForPodGroup(pcs, podGroup.Name) + if cliqueTemplate == nil { + return nil, fmt.Errorf("no clique template found for PodGroup %q in PodCliqueSet %s/%s", podGroup.Name, pcs.Namespace, pcs.Name) + } + podSpec, err := runtime.DefaultUnstructuredConverter.ToUnstructured(&cliqueTemplate.Spec.PodSpec) + if err != nil { + return nil, fmt.Errorf("failed to convert pod spec for clique %q: %w", cliqueTemplate.Name, err) + } + topologyKey := podGangTopologyKey(podGang, podGroup, b.config.RequiredTopologyKey) + + template := map[string]any{"spec": podSpec} + if topologyKey != "" { + template["metadata"] = map[string]any{ + "annotations": map[string]any{podSetRequiredTopologyAnnotation: topologyKey}, + } + } + + count := int64(cliqueTemplate.Spec.Replicas) + minCount := int64(podGroup.MinReplicas) + + podSet := map[string]any{ + "name": podGroup.Name, + "count": count, + "template": template, + } + if !isCliqueInScalingGroup(pcs, cliqueTemplate.Name) && minCount > 0 && minCount < count && !minCountUsed { + podSet["minCount"] = minCount + minCountUsed = true + } + if topologyKey != "" { + podSet["topologyRequest"] = map[string]any{"required": topologyKey} + } + podSets = append(podSets, podSet) + } + + workload := newKueueWorkload(podGang.Namespace, podGang.Name) + if err := unstructured.SetNestedField(workload.Object, queueName, "spec", "queueName"); err != nil { + return nil, err + } + if err := unstructured.SetNestedSlice(workload.Object, podSets, "spec", "podSets"); err != nil { + return nil, err + } + // Own the Workload by the PodGang so Kubernetes garbage collection removes it once the PodGang and all + // member Pods (which Kueue adds as owners on admission) are gone. + if err := controllerutil.SetOwnerReference(podGang, workload, b.scheme); err != nil { + return nil, fmt.Errorf("failed to set owner reference on prebuilt kueue Workload %s/%s: %w", podGang.Namespace, podGang.Name, err) + } + return workload, nil +} + +// isCliqueInScalingGroup reports whether the named clique is a member of any PodCliqueScalingGroup. +func isCliqueInScalingGroup(pcs *grovecorev1alpha1.PodCliqueSet, cliqueName string) bool { + for _, config := range pcs.Spec.Template.PodCliqueScalingGroupConfigs { + for _, name := range config.CliqueNames { + if name == cliqueName { + return true + } + } + } + return false +} + +// findCliqueTemplateForPodGroup resolves the PodClique template whose role name matches the PodGroup FQN. +// PodGroup names are fully qualified (--), so the longest suffix match wins. +func findCliqueTemplateForPodGroup(pcs *grovecorev1alpha1.PodCliqueSet, podGroupName string) *grovecorev1alpha1.PodCliqueTemplateSpec { + var best *grovecorev1alpha1.PodCliqueTemplateSpec + for _, cliqueTemplate := range pcs.Spec.Template.Cliques { + if podGroupName == cliqueTemplate.Name || strings.HasSuffix(podGroupName, "-"+cliqueTemplate.Name) { + if best == nil || len(cliqueTemplate.Name) > len(best.Name) { + best = cliqueTemplate + } + } + } + return best +} + +// podGangTopologyKey returns the translated required topology node-label key for a PodGroup, preferring the +// PodGroup constraint, then the PodGang-wide constraint, then the backend configured default. +func podGangTopologyKey(podGang *groveschedulerv1alpha1.PodGang, podGroup groveschedulerv1alpha1.PodGroup, fallback string) string { + if key := packRequired(podGroup.TopologyConstraint); key != "" { + return key + } + if key := packRequired(podGang.Spec.TopologyConstraint); key != "" { + return key + } + return fallback +} + +func packRequired(tc *groveschedulerv1alpha1.TopologyConstraint) string { + if tc == nil || tc.PackConstraint == nil || tc.PackConstraint.Required == nil { + return "" + } + return *tc.PackConstraint.Required +} + +func (b *schedulerBackend) PreparePod(pod *corev1.Pod) error { + pod.Spec.SchedulerName = b.config.UnderlyingSchedulerName + if pod.Labels == nil { + pod.Labels = make(map[string]string) + } + if pod.Annotations == nil { + pod.Annotations = make(map[string]string) + } + + if pod.Labels[podGroupNameLabel] == "" { + pod.Labels[podGroupNameLabel] = pod.Labels[apicommon.LabelPodClique] + } + // Grove-managed pods are deliberately NOT marked as a Kueue "serving" pod group. A serving group is never + // considered finished, so Kueue would never remove the kueue.x-k8s.io/managed finalizer on teardown (the + // prebuilt Workload has no Kueue finalizer, so Workload deletion cannot trigger finalization either). As a + // non-serving unretriable group, Kueue finalizes the Pods once they all terminate, letting them drain. + if pod.Annotations[retriableInGroupAnnotation] == "" { + pod.Annotations[retriableInGroupAnnotation] = "false" + } + if b.config.RequiredTopologyKey != "" && pod.Annotations[podSetRequiredTopologyAnnotation] == "" { + pod.Annotations[podSetRequiredTopologyAnnotation] = b.config.RequiredTopologyKey + } + return nil +} + +// ValidatePodCliqueSet rejects PodCliqueSets that cannot be represented as a valid Kueue Workload. +// Grove builds one prebuilt Kueue Workload per PodGang and maps each standalone PodClique with +// minAvailable < replicas to a Kueue podSet that sets minCount. Kueue permits at most one podSet per +// Workload to set minCount, so at most one standalone PodClique may declare partial-gang +// (minAvailable < replicas) semantics. PodCliques that belong to a PodCliqueScalingGroup are +// all-or-nothing and never set minCount, so they are excluded from this check. +func (b *schedulerBackend) ValidatePodCliqueSet(_ context.Context, pcs *grovecorev1alpha1.PodCliqueSet) error { + var ( + partialGangCliques []string + partialGangPCSGCliques []string + ) + for _, cliqueTemplate := range pcs.Spec.Template.Cliques { + if cliqueTemplate == nil { + continue + } + minAvailable := cliqueTemplate.Spec.MinAvailable + // A nil minAvailable defaults to replicas, so it is a full gang. + if minAvailable == nil { + continue + } + isPartialGang := *minAvailable > 0 && *minAvailable < cliqueTemplate.Spec.Replicas + if !isPartialGang { + continue + } + if isCliqueInScalingGroup(pcs, cliqueTemplate.Name) { + partialGangPCSGCliques = append(partialGangPCSGCliques, cliqueTemplate.Name) + continue + } + partialGangCliques = append(partialGangCliques, cliqueTemplate.Name) + } + + // PodCliques that belong to a PodCliqueScalingGroup are treated as all-or-nothing by the Kueue backend: + // their podSets never set minCount, so a member clique's minAvailable < replicas would be silently + // ignored. Reject it so an accepted PodCliqueSet behaves as specified. + if len(partialGangPCSGCliques) > 0 { + return fmt.Errorf("kueue backend requires PodCliques that are members of a PodCliqueScalingGroup to set minAvailable == replicas (all-or-nothing), but the following set minAvailable < replicas: %s", strings.Join(partialGangPCSGCliques, ", ")) + } + + // Grove maps each standalone PodClique with minAvailable < replicas to a Kueue podSet that sets minCount, + // and Kueue permits minCount on at most one podSet per Workload. + if len(partialGangCliques) > 1 { + return fmt.Errorf("kueue backend allows at most one standalone PodClique with minAvailable < replicas because Kueue permits minCount on at most one podSet per Workload, but found %d: %s", len(partialGangCliques), strings.Join(partialGangCliques, ", ")) + } + return nil +} + +func defaultConfig() configv1alpha1.KueueSchedulerConfiguration { + return configv1alpha1.KueueSchedulerConfiguration{ + UnderlyingSchedulerName: string(configv1alpha1.SchedulerNameKube), + } +} + +func withDefaults(cfg configv1alpha1.KueueSchedulerConfiguration) configv1alpha1.KueueSchedulerConfiguration { + defaults := defaultConfig() + if cfg.UnderlyingSchedulerName == "" { + cfg.UnderlyingSchedulerName = defaults.UnderlyingSchedulerName + } + return cfg +} diff --git a/operator/internal/scheduler/kueue/backend_test.go b/operator/internal/scheduler/kueue/backend_test.go new file mode 100644 index 000000000..f8a750059 --- /dev/null +++ b/operator/internal/scheduler/kueue/backend_test.go @@ -0,0 +1,417 @@ +// /* +// 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 kueue + +import ( + "context" + "testing" + + apicommon "github.com/ai-dynamo/grove/operator/api/common" + 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" + groveschedulerv1alpha1 "github.com/ai-dynamo/grove/scheduler/api/core/v1alpha1" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/runtime/schema" + "k8s.io/apimachinery/pkg/util/uuid" + "k8s.io/client-go/tools/record" + "k8s.io/utils/ptr" + "sigs.k8s.io/controller-runtime/pkg/client" +) + +const podGroupTotalCountAnnotation = "kueue.x-k8s.io/pod-group-total-count" + +func TestBackend_PreparePod_Defaults(t *testing.T) { + cl := testutils.CreateDefaultFakeClient(nil) + recorder := record.NewFakeRecorder(10) + profile := configv1alpha1.SchedulerProfile{Name: configv1alpha1.SchedulerNameKueue} + b := New(cl, cl.Scheme(), recorder, profile) + assert.NoError(t, b.Init(nil)) + + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Labels: map[string]string{ + apicommon.LabelPodClique: "test-podclique", + apicommon.LabelPodGang: "test-podgang", + }, + Annotations: map[string]string{ + podGroupTotalCountAnnotation: "2", + }, + }, + Spec: corev1.PodSpec{SchedulerName: "kueue"}, + } + + require.NoError(t, b.PreparePod(pod)) + + assert.Equal(t, string(configv1alpha1.SchedulerNameKube), pod.Spec.SchedulerName) + assert.Empty(t, pod.Labels[queueNameLabel]) + assert.Equal(t, "test-podclique", pod.Labels[podGroupNameLabel]) + assert.Equal(t, "2", pod.Annotations[podGroupTotalCountAnnotation]) + // Grove pods are not marked as a Kueue serving group, so Kueue can finalize them on teardown. + assert.Empty(t, pod.Annotations[podGroupServingAnnotation]) + assert.Equal(t, "false", pod.Annotations[retriableInGroupAnnotation]) +} + +func TestBackend_PreparePod_ConfigAndExistingMetadata(t *testing.T) { + cl := testutils.CreateDefaultFakeClient(nil) + recorder := record.NewFakeRecorder(10) + profile := configv1alpha1.SchedulerProfile{ + Name: configv1alpha1.SchedulerNameKueue, + Config: &runtime.RawExtension{ + Raw: []byte(`{"requiredTopologyKey":"topology.ai-dynamo.io/rack","underlyingSchedulerName":"custom-scheduler"}`), + }, + } + b := New(cl, cl.Scheme(), recorder, profile) + assert.NoError(t, b.Init(nil)) + + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Labels: map[string]string{ + apicommon.LabelPodClique: "test-podclique", + apicommon.LabelPodGang: "test-podgang", + queueNameLabel: "existing-queue", + }, + Annotations: map[string]string{ + podGroupTotalCountAnnotation: "7", + }, + }, + } + + require.NoError(t, b.PreparePod(pod)) + + assert.Equal(t, "custom-scheduler", pod.Spec.SchedulerName) + assert.Equal(t, "existing-queue", pod.Labels[queueNameLabel]) + assert.Equal(t, "test-podclique", pod.Labels[podGroupNameLabel]) + assert.Equal(t, "7", pod.Annotations[podGroupTotalCountAnnotation]) + assert.Equal(t, "topology.ai-dynamo.io/rack", pod.Annotations[podSetRequiredTopologyAnnotation]) +} + +func TestBackend_SyncPodGang_RequiresPodCliqueSetQueueLabel(t *testing.T) { + pcs := &grovecorev1alpha1.PodCliqueSet{ + ObjectMeta: metav1.ObjectMeta{Name: "demo", Namespace: "default"}, + Spec: grovecorev1alpha1.PodCliqueSetSpec{ + Template: grovecorev1alpha1.PodCliqueSetTemplateSpec{ + Cliques: []*grovecorev1alpha1.PodCliqueTemplateSpec{ + {Name: "worker", Spec: grovecorev1alpha1.PodCliqueSpec{Replicas: 1}}, + }, + }, + }, + } + podGang := &groveschedulerv1alpha1.PodGang{ + ObjectMeta: metav1.ObjectMeta{ + Name: "demo-0", + Namespace: "default", + Labels: map[string]string{apicommon.LabelPartOfKey: "demo"}, + }, + Spec: groveschedulerv1alpha1.PodGangSpec{ + PodGroups: []groveschedulerv1alpha1.PodGroup{{Name: "demo-0-worker", MinReplicas: 1}}, + }, + } + cl := testutils.NewTestClientBuilder().WithObjects(pcs, podGang).Build() + recorder := record.NewFakeRecorder(10) + profile := configv1alpha1.SchedulerProfile{Name: configv1alpha1.SchedulerNameKueue} + b := New(cl, cl.Scheme(), recorder, profile) + require.NoError(t, b.Init(nil)) + + require.ErrorContains(t, b.SyncPodGang(context.Background(), podGang), `must set label "kueue.x-k8s.io/queue-name"`) +} + +func TestBackend_SyncPodGang_CreatesPrebuiltWorkloadForSimplePCS(t *testing.T) { + pcs := &grovecorev1alpha1.PodCliqueSet{ + ObjectMeta: metav1.ObjectMeta{ + Name: "demo", + Namespace: "default", + Labels: map[string]string{queueNameLabel: "grove-poc"}, + }, + Spec: grovecorev1alpha1.PodCliqueSetSpec{ + Template: grovecorev1alpha1.PodCliqueSetTemplateSpec{ + Cliques: []*grovecorev1alpha1.PodCliqueTemplateSpec{ + {Name: "worker", Spec: grovecorev1alpha1.PodCliqueSpec{Replicas: 4}}, + }, + }, + }, + } + cl := testutils.NewTestClientBuilder().WithObjects(pcs).Build() + recorder := record.NewFakeRecorder(10) + profile := configv1alpha1.SchedulerProfile{Name: configv1alpha1.SchedulerNameKueue} + b := New(cl, cl.Scheme(), recorder, profile) + require.NoError(t, b.Init(nil)) + + rackKey := "topology.ai-dynamo.io/rack" + podGang := &groveschedulerv1alpha1.PodGang{ + ObjectMeta: metav1.ObjectMeta{ + Name: "demo-0", + Namespace: "default", + UID: "demo-0-uid", + Labels: map[string]string{apicommon.LabelPartOfKey: "demo"}, + }, + Spec: groveschedulerv1alpha1.PodGangSpec{ + PodGroups: []groveschedulerv1alpha1.PodGroup{ + { + Name: "demo-0-worker", + MinReplicas: 2, + TopologyConstraint: &groveschedulerv1alpha1.TopologyConstraint{ + PackConstraint: &groveschedulerv1alpha1.TopologyPackConstraint{Required: &rackKey}, + }, + }, + }, + }, + } + + require.NoError(t, b.SyncPodGang(context.Background(), podGang)) + + got := &unstructured.Unstructured{} + got.SetGroupVersionKind(workloadGVK) + require.NoError(t, cl.Get(context.Background(), client.ObjectKey{Namespace: "default", Name: "demo-0"}, got)) + + queueName, _, _ := unstructured.NestedString(got.Object, "spec", "queueName") + assert.Equal(t, "grove-poc", queueName) + + ownerRefs := got.GetOwnerReferences() + require.Len(t, ownerRefs, 1) + assert.Equal(t, "PodGang", ownerRefs[0].Kind) + assert.Equal(t, "demo-0", ownerRefs[0].Name) + + podSets, _, _ := unstructured.NestedSlice(got.Object, "spec", "podSets") + require.Len(t, podSets, 1) + podSet := podSets[0].(map[string]any) + assert.Equal(t, "demo-0-worker", podSet["name"]) + assert.Equal(t, int64(4), podSet["count"]) + assert.Equal(t, int64(2), podSet["minCount"]) + topologyRequest := podSet["topologyRequest"].(map[string]any) + assert.Equal(t, rackKey, topologyRequest["required"]) +} + +func TestBackend_SyncPodGang_CreatesPrebuiltWorkloadForPCSGWithMinCountEqualsCount(t *testing.T) { + pcsgReplicas := int32(2) + pcs := &grovecorev1alpha1.PodCliqueSet{ + ObjectMeta: metav1.ObjectMeta{ + Name: "demo", + Namespace: "default", + Labels: map[string]string{queueNameLabel: "grove-poc"}, + }, + Spec: grovecorev1alpha1.PodCliqueSetSpec{ + Template: grovecorev1alpha1.PodCliqueSetTemplateSpec{ + Cliques: []*grovecorev1alpha1.PodCliqueTemplateSpec{ + {Name: "leader", Spec: grovecorev1alpha1.PodCliqueSpec{Replicas: 1}}, + {Name: "worker", Spec: grovecorev1alpha1.PodCliqueSpec{Replicas: 2}}, + }, + PodCliqueScalingGroupConfigs: []grovecorev1alpha1.PodCliqueScalingGroupConfig{ + {Name: "decode", CliqueNames: []string{"leader", "worker"}, Replicas: &pcsgReplicas}, + }, + }, + }, + } + cl := testutils.NewTestClientBuilder().WithObjects(pcs).Build() + recorder := record.NewFakeRecorder(10) + profile := configv1alpha1.SchedulerProfile{Name: configv1alpha1.SchedulerNameKueue} + b := New(cl, cl.Scheme(), recorder, profile) + require.NoError(t, b.Init(nil)) + + podGang := &groveschedulerv1alpha1.PodGang{ + ObjectMeta: metav1.ObjectMeta{ + Name: "demo-0", + Namespace: "default", + Labels: map[string]string{apicommon.LabelPartOfKey: "demo"}, + }, + Spec: groveschedulerv1alpha1.PodGangSpec{ + PodGroups: []groveschedulerv1alpha1.PodGroup{ + // MinReplicas is deliberately lower than the clique replicas to prove the PCSG + // all-or-nothing override forces minCount == count. + {Name: "demo-0-decode-0-worker", MinReplicas: 1}, + }, + }, + } + + require.NoError(t, b.SyncPodGang(context.Background(), podGang)) + + got := &unstructured.Unstructured{} + got.SetGroupVersionKind(workloadGVK) + require.NoError(t, cl.Get(context.Background(), client.ObjectKey{Namespace: "default", Name: "demo-0"}, got)) + + podSets, _, _ := unstructured.NestedSlice(got.Object, "spec", "podSets") + require.Len(t, podSets, 1) + podSet := podSets[0].(map[string]any) + assert.Equal(t, "demo-0-decode-0-worker", podSet["name"]) + assert.Equal(t, int64(2), podSet["count"]) + // PodCliqueScalingGroup cliques are all-or-nothing: minCount is omitted (Kueue defaults it to count). + // Kueue also rejects Workloads where more than one podSet sets minCount. + _, hasMinCount := podSet["minCount"] + assert.False(t, hasMinCount) +} + +func TestBackend_ValidatePodCliqueSet_MinCount(t *testing.T) { + testCases := []struct { + description string + cliques []*grovecorev1alpha1.PodCliqueTemplateSpec + pcsgConfigs []grovecorev1alpha1.PodCliqueScalingGroupConfig + wantErr bool + wantErrMsg string + }{ + { + description: "no partial-gang standalone clique is valid", + cliques: []*grovecorev1alpha1.PodCliqueTemplateSpec{ + {Name: "worker", Spec: grovecorev1alpha1.PodCliqueSpec{Replicas: 4, MinAvailable: ptr.To[int32](4)}}, + }, + }, + { + description: "single partial-gang standalone clique is valid", + cliques: []*grovecorev1alpha1.PodCliqueTemplateSpec{ + {Name: "worker", Spec: grovecorev1alpha1.PodCliqueSpec{Replicas: 4, MinAvailable: ptr.To[int32](2)}}, + {Name: "frontend", Spec: grovecorev1alpha1.PodCliqueSpec{Replicas: 2, MinAvailable: ptr.To[int32](2)}}, + }, + }, + { + description: "two partial-gang standalone cliques are rejected", + cliques: []*grovecorev1alpha1.PodCliqueTemplateSpec{ + {Name: "worker", Spec: grovecorev1alpha1.PodCliqueSpec{Replicas: 4, MinAvailable: ptr.To[int32](2)}}, + {Name: "frontend", Spec: grovecorev1alpha1.PodCliqueSpec{Replicas: 3, MinAvailable: ptr.To[int32](1)}}, + }, + wantErr: true, + wantErrMsg: "at most one standalone PodClique with minAvailable < replicas", + }, + { + description: "scaling-group cliques with minAvailable == replicas are valid", + cliques: []*grovecorev1alpha1.PodCliqueTemplateSpec{ + {Name: "leader", Spec: grovecorev1alpha1.PodCliqueSpec{Replicas: 1, MinAvailable: ptr.To[int32](1)}}, + {Name: "worker", Spec: grovecorev1alpha1.PodCliqueSpec{Replicas: 4, MinAvailable: ptr.To[int32](4)}}, + }, + pcsgConfigs: []grovecorev1alpha1.PodCliqueScalingGroupConfig{ + {Name: "decode", CliqueNames: []string{"leader", "worker"}}, + }, + }, + { + description: "scaling-group clique with minAvailable < replicas is rejected", + cliques: []*grovecorev1alpha1.PodCliqueTemplateSpec{ + {Name: "leader", Spec: grovecorev1alpha1.PodCliqueSpec{Replicas: 1, MinAvailable: ptr.To[int32](1)}}, + {Name: "worker", Spec: grovecorev1alpha1.PodCliqueSpec{Replicas: 4, MinAvailable: ptr.To[int32](3)}}, + }, + pcsgConfigs: []grovecorev1alpha1.PodCliqueScalingGroupConfig{ + {Name: "decode", CliqueNames: []string{"leader", "worker"}}, + }, + wantErr: true, + wantErrMsg: "members of a PodCliqueScalingGroup to set minAvailable == replicas", + }, + { + description: "nil minAvailable is treated as full gang", + cliques: []*grovecorev1alpha1.PodCliqueTemplateSpec{ + {Name: "worker", Spec: grovecorev1alpha1.PodCliqueSpec{Replicas: 4}}, + {Name: "frontend", Spec: grovecorev1alpha1.PodCliqueSpec{Replicas: 3, MinAvailable: ptr.To[int32](1)}}, + }, + }, + } + + for _, tc := range testCases { + t.Run(tc.description, func(t *testing.T) { + cl := testutils.CreateDefaultFakeClient(nil) + recorder := record.NewFakeRecorder(10) + profile := configv1alpha1.SchedulerProfile{Name: configv1alpha1.SchedulerNameKueue} + b := New(cl, cl.Scheme(), recorder, profile) + require.NoError(t, b.Init(nil)) + + pcs := &grovecorev1alpha1.PodCliqueSet{ + ObjectMeta: metav1.ObjectMeta{Name: "demo", Namespace: "default"}, + Spec: grovecorev1alpha1.PodCliqueSetSpec{ + Template: grovecorev1alpha1.PodCliqueSetTemplateSpec{ + Cliques: tc.cliques, + PodCliqueScalingGroupConfigs: tc.pcsgConfigs, + }, + }, + } + + err := b.ValidatePodCliqueSet(context.Background(), pcs) + if tc.wantErr { + require.ErrorContains(t, err, tc.wantErrMsg) + } else { + require.NoError(t, err) + } + }) + } +} + +func TestBackend_TopologyGVR(t *testing.T) { + b := newKueueBackend(testutils.CreateDefaultFakeClient(nil)) + + assert.Equal(t, schema.GroupVersionResource{ + Group: "kueue.x-k8s.io", + Version: "v1beta2", + Resource: "topologies", + }, b.TopologyGVR()) +} + +func TestBackend_SyncTopologyCreatesKueueTopology(t *testing.T) { + ctx := context.Background() + cl := testutils.CreateDefaultFakeClient(nil) + b := newKueueBackend(cl) + ct := testClusterTopology() + + require.NoError(t, b.SyncTopology(ctx, cl, ct)) + + topology := newKueueTopology(ct.Name) + require.NoError(t, cl.Get(ctx, client.ObjectKey{Name: ct.Name}, topology)) + assert.True(t, metav1.IsControlledBy(topology, ct)) + assert.Equal(t, []any{ + map[string]any{"nodeLabel": "topology.ai-dynamo.io/rack"}, + map[string]any{"nodeLabel": "kubernetes.io/hostname"}, + }, kueueTopologyLevels(topology)) +} + +func TestBackend_CheckTopologyDrift(t *testing.T) { + ctx := context.Background() + ct := testClusterTopology() + topology, err := buildKueueTopology(ct.Name, ct, testutils.CreateDefaultFakeClient(nil).Scheme()) + require.NoError(t, err) + cl := testutils.NewTestClientBuilder().WithObjects(ct, topology).Build() + b := newKueueBackend(cl) + + inSync, message, _, err := b.CheckTopologyDrift(ctx, ct, grovecorev1alpha1.SchedulerTopologyBinding{ + SchedulerName: string(configv1alpha1.SchedulerNameKueue), + TopologyReference: ct.Name, + }) + + require.NoError(t, err) + assert.True(t, inSync) + assert.Empty(t, message) +} + +func newKueueBackend(cl client.Client) scheduler.TopologyAwareBackend { + recorder := record.NewFakeRecorder(10) + profile := configv1alpha1.SchedulerProfile{Name: configv1alpha1.SchedulerNameKueue} + b := New(cl, cl.Scheme(), recorder, profile) + return b.(scheduler.TopologyAwareBackend) +} + +func testClusterTopology() *grovecorev1alpha1.ClusterTopologyBinding { + return &grovecorev1alpha1.ClusterTopologyBinding{ + ObjectMeta: metav1.ObjectMeta{ + Name: "grove-kind-topology", + UID: uuid.NewUUID(), + }, + Spec: grovecorev1alpha1.ClusterTopologyBindingSpec{ + Levels: []grovecorev1alpha1.TopologyLevel{ + {Domain: grovecorev1alpha1.TopologyDomainRack, Key: "topology.ai-dynamo.io/rack"}, + {Domain: grovecorev1alpha1.TopologyDomainHost, Key: "kubernetes.io/hostname"}, + }, + }, + } +} diff --git a/operator/internal/scheduler/kueue/topology.go b/operator/internal/scheduler/kueue/topology.go new file mode 100644 index 000000000..0af470f79 --- /dev/null +++ b/operator/internal/scheduler/kueue/topology.go @@ -0,0 +1,142 @@ +// /* +// 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 kueue + +import ( + "context" + "fmt" + "reflect" + + grovecorev1alpha1 "github.com/ai-dynamo/grove/operator/api/core/v1alpha1" + "github.com/ai-dynamo/grove/operator/internal/scheduler" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/runtime/schema" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil" + "sigs.k8s.io/controller-runtime/pkg/log" +) + +var _ scheduler.TopologyAwareBackend = (*schedulerBackend)(nil) + +var topologyGVK = schema.GroupVersionKind{ + Group: "kueue.x-k8s.io", + Version: "v1beta2", + Kind: "Topology", +} + +func (b *schedulerBackend) TopologyGVR() schema.GroupVersionResource { + return schema.GroupVersionResource{ + Group: "kueue.x-k8s.io", + Version: "v1beta2", + Resource: "topologies", + } +} + +func (b *schedulerBackend) TopologyResourceName(ct *grovecorev1alpha1.ClusterTopologyBinding) string { + return ct.Name +} + +func (b *schedulerBackend) SyncTopology(ctx context.Context, k8sClient client.Client, ct *grovecorev1alpha1.ClusterTopologyBinding) error { + if k8sClient == nil { + k8sClient = b.client + } + logger := log.FromContext(ctx) + + desiredTopology, err := buildKueueTopology(ct.Name, ct, b.scheme) + if err != nil { + return fmt.Errorf("failed to build Kueue Topology: %w", err) + } + + existingTopology := newKueueTopology(ct.Name) + if err = k8sClient.Get(ctx, client.ObjectKey{Name: ct.Name}, existingTopology); err != nil { + if apierrors.IsNotFound(err) { + if err = k8sClient.Create(ctx, desiredTopology); err != nil { + return fmt.Errorf("failed to create Kueue Topology %s: %w", ct.Name, err) + } + logger.Info("Created Kueue Topology", "name", ct.Name) + return nil + } + return fmt.Errorf("failed to get Kueue Topology %s: %w", ct.Name, err) + } + + if !metav1.IsControlledBy(existingTopology, ct) { + return fmt.Errorf("Kueue Topology %s is not owned by ClusterTopologyBinding %s", ct.Name, ct.Name) + } + if !reflect.DeepEqual(kueueTopologyLevels(existingTopology), desiredKueueTopologyLevels(ct)) { + if err = k8sClient.Delete(ctx, existingTopology); err != nil { + return fmt.Errorf("failed to recreate (action: delete) existing Kueue Topology %s: %w", ct.Name, err) + } + if err = k8sClient.Create(ctx, desiredTopology); err != nil { + return fmt.Errorf("failed to recreate (action: create) Kueue Topology %s: %w", ct.Name, err) + } + logger.Info("Recreated Kueue Topology with updated levels", "name", ct.Name) + } + return nil +} + +func (b *schedulerBackend) OnTopologyDelete(_ context.Context, _ client.Client, _ *grovecorev1alpha1.ClusterTopologyBinding) error { + return nil +} + +func (b *schedulerBackend) CheckTopologyDrift(ctx context.Context, ct *grovecorev1alpha1.ClusterTopologyBinding, ref grovecorev1alpha1.SchedulerTopologyBinding) (bool, string, int64, error) { + existingTopology := newKueueTopology(ref.TopologyReference) + if err := b.client.Get(ctx, client.ObjectKey{Name: ref.TopologyReference}, existingTopology); err != nil { + if apierrors.IsNotFound(err) { + return false, fmt.Sprintf("Kueue Topology %q not found", ref.TopologyReference), 0, nil + } + return false, "", 0, fmt.Errorf("failed to get Kueue Topology %s: %w", ref.TopologyReference, err) + } + if !reflect.DeepEqual(kueueTopologyLevels(existingTopology), desiredKueueTopologyLevels(ct)) { + return false, "Kueue Topology levels differ from ClusterTopologyBinding levels", existingTopology.GetGeneration(), nil + } + return true, "", existingTopology.GetGeneration(), nil +} + +func buildKueueTopology(name string, ct *grovecorev1alpha1.ClusterTopologyBinding, scheme *runtime.Scheme) (*unstructured.Unstructured, error) { + topology := newKueueTopology(name) + topology.SetOwnerReferences(nil) + if err := unstructured.SetNestedSlice(topology.Object, desiredKueueTopologyLevels(ct), "spec", "levels"); err != nil { + return nil, err + } + if err := controllerutil.SetControllerReference(ct, topology, scheme); err != nil { + return nil, fmt.Errorf("failed to set owner reference for Kueue Topology: %w", err) + } + return topology, nil +} + +func newKueueTopology(name string) *unstructured.Unstructured { + topology := &unstructured.Unstructured{} + topology.SetGroupVersionKind(topologyGVK) + topology.SetName(name) + return topology +} + +func desiredKueueTopologyLevels(ct *grovecorev1alpha1.ClusterTopologyBinding) []any { + levels := make([]any, 0, len(ct.Spec.Levels)) + for _, level := range ct.Spec.Levels { + levels = append(levels, map[string]any{"nodeLabel": level.Key}) + } + return levels +} + +func kueueTopologyLevels(topology *unstructured.Unstructured) []any { + levels, _, _ := unstructured.NestedSlice(topology.Object, "spec", "levels") + return levels +} diff --git a/operator/internal/scheduler/registry/registry.go b/operator/internal/scheduler/registry/registry.go index f59e5b13e..e8ea6744c 100644 --- a/operator/internal/scheduler/registry/registry.go +++ b/operator/internal/scheduler/registry/registry.go @@ -24,6 +24,7 @@ import ( "github.com/ai-dynamo/grove/operator/internal/scheduler" "github.com/ai-dynamo/grove/operator/internal/scheduler/kai" "github.com/ai-dynamo/grove/operator/internal/scheduler/kube" + "github.com/ai-dynamo/grove/operator/internal/scheduler/kueue" "github.com/ai-dynamo/grove/operator/internal/scheduler/lpx" "github.com/ai-dynamo/grove/operator/internal/scheduler/volcano" @@ -103,6 +104,8 @@ func newSchedulerBackend(cl, directClient client.Client, scheme *runtime.Scheme, b = kube.New(cl, scheme, rec, p) case configv1alpha1.SchedulerNameKai: b = kai.New(cl, scheme, rec, p) + case configv1alpha1.SchedulerNameKueue: + b = kueue.New(cl, scheme, rec, p) case configv1alpha1.SchedulerNameVolcano: b = volcano.New(cl, scheme, rec, p) case configv1alpha1.SchedulerNameLPX: diff --git a/operator/internal/scheduler/registry/registry_test.go b/operator/internal/scheduler/registry/registry_test.go index b0f5bd379..b3432bad1 100644 --- a/operator/internal/scheduler/registry/registry_test.go +++ b/operator/internal/scheduler/registry/registry_test.go @@ -56,6 +56,12 @@ func TestNewRegistry(t *testing.T) { wantErr: false, expectedName: "lpx-scheduler", }, + { + name: "kueue scheduler initialization", + schedulerName: configv1alpha1.SchedulerNameKueue, + wantErr: false, + expectedName: "kueue", + }, { name: "unsupported scheduler", schedulerName: "unknown-scheduler", @@ -102,6 +108,7 @@ func TestNewRegistry(t *testing.T) { {Name: configv1alpha1.SchedulerNameKai}, {Name: configv1alpha1.SchedulerNameVolcano}, {Name: configv1alpha1.SchedulerNameLPX}, + {Name: configv1alpha1.SchedulerNameKueue}, }, DefaultProfileName: string(configv1alpha1.SchedulerNameKai), } @@ -111,6 +118,7 @@ func TestNewRegistry(t *testing.T) { require.NotNil(t, reg.Get(string(configv1alpha1.SchedulerNameKube))) require.NotNil(t, reg.Get(string(configv1alpha1.SchedulerNameVolcano))) require.NotNil(t, reg.Get(string(configv1alpha1.SchedulerNameLPX))) + require.NotNil(t, reg.Get(string(configv1alpha1.SchedulerNameKueue))) assert.Equal(t, reg.GetDefault(), reg.Get(string(configv1alpha1.SchedulerNameKai))) assert.NotContains(t, reg.AllTopologyAware(), string(configv1alpha1.SchedulerNameLPX)) }) From 4f03065698606cb0700666ce9a29e98c1130758b Mon Sep 17 00:00:00 2001 From: Alay Patel Date: Thu, 16 Jul 2026 15:39:03 -0400 Subject: [PATCH 2/2] add 2 node demo with minAvailable of 2 Signed-off-by: Alay Patel --- .../kueue-tas-backend-slides/demo/README.md | 53 +++++++++ .../demo/grove-kind-topology.yaml | 5 +- .../demo/grove-kueue-pcsg.yaml | 4 +- .../demo/grove-kueue-simple-mincount.yaml | 15 +-- .../demo/kueue-kind-queue.yaml | 4 +- .../demo/setup-kueue-kind-demo.sh | 50 +++++++++ .../demo/validate-grove-kueue-mincount.sh | 101 +++++++++++++++--- operator/Makefile | 5 +- operator/hack/kind-up.sh | 13 +++ 9 files changed, 218 insertions(+), 32 deletions(-) create mode 100644 docs/community/kueue-tas-backend-slides/demo/README.md create mode 100755 docs/community/kueue-tas-backend-slides/demo/setup-kueue-kind-demo.sh diff --git a/docs/community/kueue-tas-backend-slides/demo/README.md b/docs/community/kueue-tas-backend-slides/demo/README.md new file mode 100644 index 000000000..385cda92e --- /dev/null +++ b/docs/community/kueue-tas-backend-slides/demo/README.md @@ -0,0 +1,53 @@ +# Grove + Kueue kind demo (2-node, hostname TAS) + +## Cluster + +Create a 2-node kind cluster (control-plane + 1 worker): + +```sh +cd operator +make kind-down +make kind-up WORKERS=1 +``` + +Install Kueue and deploy Grove as usual, then prepare the demo platform: + +```sh +docs/community/kueue-tas-backend-slides/demo/setup-kueue-kind-demo.sh +kubectl get topology grove-kind-topology +``` + +`setup-kueue-kind-demo.sh` labels every node with `topology.ai-dynamo.io/tas-node=true`, +applies `grove-kind-topology.yaml` (host -> `kubernetes.io/hostname`), and applies +`kueue-kind-queue.yaml`. Grove syncs the `ClusterTopologyBinding` to a Kueue +`Topology` with the same name. + +## Demos + +### Partial admission via minCount (primary) + +`grove-kueue-simple-mincount.yaml` requests `replicas: 4`, `minAvailable: 2`, and +`pack.required: host`, with each Pod requesting 1 CPU. The demo ClusterQueue has +only 2 CPU of quota, so Kueue admits the Workload at `minCount=2`: 2 of the 4 Pods +run (one per node via hostname TAS) and the other 2 stay scheduling-gated. + +```sh +docs/community/kueue-tas-backend-slides/demo/validate-grove-kueue-mincount.sh +``` + +### PCSG all-or-nothing + +`grove-kueue-pcsg.yaml` uses the same hostname topology. The base PodGang requests +six Pods across two PCSG replicas; it exercises prebuilt Workload shape more than +2-node placement. + +```sh +docs/community/kueue-tas-backend-slides/demo/validate-grove-kueue-pcsg.sh +``` + +## Cleanup + +```sh +kubectl delete -f docs/community/kueue-tas-backend-slides/demo/grove-kueue-simple-mincount.yaml --ignore-not-found +kubectl delete -f docs/community/kueue-tas-backend-slides/demo/grove-kueue-pcsg.yaml --ignore-not-found +``` diff --git a/docs/community/kueue-tas-backend-slides/demo/grove-kind-topology.yaml b/docs/community/kueue-tas-backend-slides/demo/grove-kind-topology.yaml index 10c25eab1..b3b6be63b 100644 --- a/docs/community/kueue-tas-backend-slides/demo/grove-kind-topology.yaml +++ b/docs/community/kueue-tas-backend-slides/demo/grove-kind-topology.yaml @@ -1,11 +1,10 @@ -# Grove source-of-truth topology for the CPU-only kind TAS smoke test. +# Grove source-of-truth topology for the 2-node kind TAS smoke test. +# Use required: host on PodCliqueSet so Kueue TAS packs at kubernetes.io/hostname. apiVersion: grove.io/v1alpha1 kind: ClusterTopologyBinding metadata: name: grove-kind-topology spec: levels: - - domain: rack - key: topology.ai-dynamo.io/rack - domain: host key: kubernetes.io/hostname diff --git a/docs/community/kueue-tas-backend-slides/demo/grove-kueue-pcsg.yaml b/docs/community/kueue-tas-backend-slides/demo/grove-kueue-pcsg.yaml index 4a6b47278..2c532a116 100644 --- a/docs/community/kueue-tas-backend-slides/demo/grove-kueue-pcsg.yaml +++ b/docs/community/kueue-tas-backend-slides/demo/grove-kueue-pcsg.yaml @@ -2,7 +2,7 @@ # POC assumption: PCSG.minAvailable == PCSG.replicas, so the whole gang is # all-or-nothing. Grove builds one prebuilt Kueue Workload per PodGang and # omits minCount for the PCSG podSets. -# Apply after kueue-kind-queue.yaml and grove-kind-topology.yaml. +# Apply after setup-kueue-kind-demo.sh. apiVersion: grove.io/v1alpha1 kind: PodCliqueSet metadata: @@ -16,7 +16,7 @@ spec: topologyConstraint: topologyName: grove-kind-topology pack: - required: rack + required: host podCliqueScalingGroups: - name: decode replicas: 2 diff --git a/docs/community/kueue-tas-backend-slides/demo/grove-kueue-simple-mincount.yaml b/docs/community/kueue-tas-backend-slides/demo/grove-kueue-simple-mincount.yaml index 8f7ddf32b..5c61dab03 100644 --- a/docs/community/kueue-tas-backend-slides/demo/grove-kueue-simple-mincount.yaml +++ b/docs/community/kueue-tas-backend-slides/demo/grove-kueue-simple-mincount.yaml @@ -1,7 +1,10 @@ -# Grove PodCliqueSet with minCount semantics and no PodCliqueScalingGroup. -# Grove's Kueue backend builds one prebuilt Kueue Workload per PodGang, with a -# podSet whose count == clique replicas and minCount == clique minAvailable. -# Apply after kueue-kind-queue.yaml and grove-kind-topology.yaml. +# Grove PodCliqueSet demonstrating Kueue partial admission via minCount. +# frontend requests replicas=4 with minAvailable=2. Each Pod requests 1 CPU and +# the demo ClusterQueue only has 2 CPU of quota, so Kueue admits the Workload at +# minCount=2 and leaves the other 2 Pods gated: 2 of 4 Pods run. +# With WORKERS=1 (control-plane + worker) and required: host, the 2 admitted Pods +# land on 2 distinct nodes. +# Apply after setup-kueue-kind-demo.sh. apiVersion: grove.io/v1alpha1 kind: PodCliqueSet metadata: @@ -15,7 +18,7 @@ spec: topologyConstraint: topologyName: grove-kind-topology pack: - required: rack + required: host cliques: - name: frontend spec: @@ -30,5 +33,5 @@ spec: image: registry.k8s.io/pause:3.9 resources: requests: - cpu: 100m + cpu: "1" memory: 64Mi diff --git a/docs/community/kueue-tas-backend-slides/demo/kueue-kind-queue.yaml b/docs/community/kueue-tas-backend-slides/demo/kueue-kind-queue.yaml index 99b21d02a..234a0cbf0 100644 --- a/docs/community/kueue-tas-backend-slides/demo/kueue-kind-queue.yaml +++ b/docs/community/kueue-tas-backend-slides/demo/kueue-kind-queue.yaml @@ -23,9 +23,9 @@ spec: - name: kind-cpu resources: - name: cpu - nominalQuota: "4" + nominalQuota: "2" - name: memory - nominalQuota: 8Gi + nominalQuota: 4Gi --- apiVersion: kueue.x-k8s.io/v1beta2 kind: LocalQueue diff --git a/docs/community/kueue-tas-backend-slides/demo/setup-kueue-kind-demo.sh b/docs/community/kueue-tas-backend-slides/demo/setup-kueue-kind-demo.sh new file mode 100755 index 000000000..c91476bb0 --- /dev/null +++ b/docs/community/kueue-tas-backend-slides/demo/setup-kueue-kind-demo.sh @@ -0,0 +1,50 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Prepares the 2-node kind cluster for Grove + Kueue TAS demos: +# - labels every node for the kind-cpu ResourceFlavor +# - applies ClusterTopologyBinding (host -> kubernetes.io/hostname) +# - applies Kueue queues/flavors +# +# Prereq: kind cluster with WORKERS=1 (control-plane + worker), Grove operator, +# and Kueue installed. Grove syncs ClusterTopologyBinding -> Kueue Topology. + +KUBECTL="${KUBECTL:-kubectl}" +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +require_two_nodes() { + local count + count="$("${KUBECTL}" get nodes --no-headers 2>/dev/null | wc -l | tr -d ' ')" + if [[ "${count}" -lt 2 ]]; then + echo "Expected at least 2 kind nodes for hostname TAS demos; found ${count}." >&2 + echo "Recreate the cluster with one worker:" >&2 + echo " cd operator && make kind-down && make kind-up WORKERS=1" >&2 + exit 1 + fi +} + +label_tas_nodes() { + local node + for node in $("${KUBECTL}" get nodes -o jsonpath='{.items[*].metadata.name}'); do + "${KUBECTL}" label node "${node}" topology.ai-dynamo.io/tas-node=true --overwrite >/dev/null + done +} + +require_two_nodes +label_tas_nodes + +"${KUBECTL}" apply -f "${SCRIPT_DIR}/grove-kind-topology.yaml" +"${KUBECTL}" apply -f "${SCRIPT_DIR}/kueue-kind-queue.yaml" + +echo "Labeled nodes:" +"${KUBECTL}" get nodes -L topology.ai-dynamo.io/tas-node +echo +echo "Grove topology binding:" +"${KUBECTL}" get clustertopologybinding grove-kind-topology +echo +echo "Wait for Grove to sync Kueue Topology, then verify:" +echo " kubectl get topology grove-kind-topology" +echo +echo "Run demos:" +echo " ${SCRIPT_DIR}/validate-grove-kueue-mincount.sh" +echo " ${SCRIPT_DIR}/validate-grove-kueue-pcsg.sh" diff --git a/docs/community/kueue-tas-backend-slides/demo/validate-grove-kueue-mincount.sh b/docs/community/kueue-tas-backend-slides/demo/validate-grove-kueue-mincount.sh index 4aceddb2d..bacd38204 100755 --- a/docs/community/kueue-tas-backend-slides/demo/validate-grove-kueue-mincount.sh +++ b/docs/community/kueue-tas-backend-slides/demo/validate-grove-kueue-mincount.sh @@ -1,24 +1,54 @@ #!/usr/bin/env bash set -euo pipefail -# Validates the Grove Kueue "simple PCS with minCount" path: Grove builds one -# prebuilt Kueue Workload per PodGang with a podSet whose count == clique -# replicas and minCount == clique minAvailable, and stamps the Pods with -# kueue.x-k8s.io/prebuilt-workload-name. +# Validates Kueue partial admission via minCount. +# grove-kueue-simple-mincount.yaml has one clique: replicas=4, minAvailable=2, each +# Pod requesting 1 CPU. The demo ClusterQueue (kueue-kind-queue.yaml) grants only +# 2 CPU, so Kueue admits the prebuilt Workload at minCount=2: Grove creates all 4 +# Pods, Kueue ungates 2 (Running) and leaves 2 scheduling-gated. Result: 2 of 4 run. +# +# With pack.required: host, the admitted Pods pack onto a single host, so this demo +# does NOT assert Pods land on distinct nodes. KUBECTL="${KUBECTL:-kubectl}" NAMESPACE="${NAMESPACE:-default}" -TIMEOUT_SECONDS="${TIMEOUT_SECONDS:-120}" +TIMEOUT_SECONDS="${TIMEOUT_SECONDS:-180}" SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" MANIFEST="${MANIFEST:-${SCRIPT_DIR}/grove-kueue-simple-mincount.yaml}" +QUEUE_MANIFEST="${QUEUE_MANIFEST:-${SCRIPT_DIR}/kueue-kind-queue.yaml}" PCS_NAME="grove-kueue-mincount" PODGANG_NAME="${PCS_NAME}-0" +CLUSTER_QUEUE="kind-cpu" cleanup() { "${KUBECTL}" delete -f "${MANIFEST}" --ignore-not-found >/dev/null 2>&1 || true "${KUBECTL}" -n "${NAMESPACE}" delete workload "${PODGANG_NAME}" --ignore-not-found >/dev/null 2>&1 || true } +diagnose() { + echo "---- diagnostics ----" >&2 + echo "Pods:" >&2 + "${KUBECTL}" -n "${NAMESPACE}" get pods -l "app.kubernetes.io/part-of=${PCS_NAME}" -o wide >&2 || true + echo "Workload podSets (count / minCount / request):" >&2 + "${KUBECTL}" -n "${NAMESPACE}" get "workload/${PODGANG_NAME}" \ + -o jsonpath='{range .spec.podSets[*]}{.name}{" count="}{.count}{" minCount="}{.minCount}{" cpu="}{.template.spec.containers[0].resources.requests.cpu}{"\n"}{end}' >&2 || true + echo >&2 + echo "Workload admission (admitted count):" >&2 + "${KUBECTL}" -n "${NAMESPACE}" get "workload/${PODGANG_NAME}" \ + -o jsonpath='{range .status.admission.podSetAssignments[*]}{.name}{" count="}{.count}{"\n"}{end}' >&2 || true + echo >&2 + echo "ClusterQueue ${CLUSTER_QUEUE} nominalQuota vs usage:" >&2 + "${KUBECTL}" get clusterqueue "${CLUSTER_QUEUE}" \ + -o jsonpath='nominal={range .spec.resourceGroups[*].flavors[*].resources[*]}{.name}={.nominalQuota}{" "}{end}{"\n"}' >&2 || true + "${KUBECTL}" get clusterqueue "${CLUSTER_QUEUE}" \ + -o jsonpath='usage={range .status.flavorsUsage[*].resources[*]}{.name}={.total}{" "}{end}{"\n"}' >&2 || true + echo >&2 + echo "If all 4 Pods are Running, the ClusterQueue quota is not capping admission." >&2 + echo "Re-apply the demo quota (2 CPU) and retry:" >&2 + echo " ${KUBECTL} apply -f ${QUEUE_MANIFEST}" >&2 + echo "---------------------" >&2 +} + wait_for_jsonpath() { local name="$1" local jsonpath="$2" @@ -39,11 +69,23 @@ wait_for_jsonpath() { return 1 } +running_pods() { + "${KUBECTL}" -n "${NAMESPACE}" get pods \ + -l "app.kubernetes.io/part-of=${PCS_NAME}" \ + --field-selector=status.phase=Running \ + -o name 2>/dev/null | wc -l | tr -d ' ' +} + +total_pods() { + "${KUBECTL}" -n "${NAMESPACE}" get pods \ + -l "app.kubernetes.io/part-of=${PCS_NAME}" \ + -o name 2>/dev/null | wc -l | tr -d ' ' +} + cleanup "${KUBECTL}" apply -f "${MANIFEST}" >/dev/null -# Grove builds one prebuilt Workload named after the PodGang. wait_for_jsonpath \ "workload/${PODGANG_NAME}" \ "{.metadata.name}" \ @@ -53,25 +95,50 @@ wait_for_jsonpath \ wait_for_jsonpath "workload/${PODGANG_NAME}" "{.spec.podSets[0].count}" "4" "${TIMEOUT_SECONDS}" wait_for_jsonpath "workload/${PODGANG_NAME}" "{.spec.podSets[0].minCount}" "2" "${TIMEOUT_SECONDS}" -# Pods reference the prebuilt Workload. -pod="$("${KUBECTL}" -n "${NAMESPACE}" get pods \ - -l "app.kubernetes.io/part-of=${PCS_NAME}" \ - -o "jsonpath={.items[0].metadata.name}")" -if [[ -z "${pod}" ]]; then - echo "Expected at least one Pod for PodCliqueSet ${PCS_NAME}" >&2 +wait_for_jsonpath \ + "workload/${PODGANG_NAME}" \ + '{.status.conditions[?(@.type=="Admitted")].status}' \ + "True" \ + "${TIMEOUT_SECONDS}" + +# Grove creates all 4 Pods; Kueue partial-admits at minCount=2, so exactly 2 become +# Running and the other 2 remain scheduling-gated. Wait for at least 2 Running, then +# assert the count settles at exactly 2 (quota only fits 2). +running_count=0 +deadline=$((SECONDS + TIMEOUT_SECONDS)) +while (( SECONDS < deadline )); do + running_count="$(running_pods)" + if [[ "${running_count}" -ge 2 ]]; then + break + fi + sleep 2 +done +if [[ "${running_count}" -lt 2 ]]; then + echo "Expected at least 2 Running Pods for ${PCS_NAME} (minCount=2), got ${running_count}" >&2 + diagnose + exit 1 +fi + +# Give Kueue a moment to (not) admit more, then confirm the Running count stays at 2. +sleep 5 +running_count="$(running_pods)" +if [[ "${running_count}" -ne 2 ]]; then + echo "Expected exactly 2 Running Pods (quota fits only 2), got ${running_count}" >&2 + diagnose exit 1 fi -prebuilt="$("${KUBECTL}" -n "${NAMESPACE}" get pod "${pod}" \ - -o "jsonpath={.metadata.labels.kueue\\.x-k8s\\.io/prebuilt-workload-name}")" -if [[ "${prebuilt}" != "${PODGANG_NAME}" ]]; then - echo "Expected Pod ${pod} to carry kueue.x-k8s.io/prebuilt-workload-name=${PODGANG_NAME}, got '${prebuilt}'" >&2 +total_count="$(total_pods)" +if [[ "${total_count}" -ne 4 ]]; then + echo "Expected 4 Pods total (2 Running, 2 gated), got ${total_count}" >&2 + diagnose exit 1 fi echo "PASS: Grove created one prebuilt Workload named ${PODGANG_NAME}." echo "PASS: the Workload podSet has count=4 and minCount=2." -echo "PASS: Pods reference the prebuilt Workload via kueue.x-k8s.io/prebuilt-workload-name." +echo "PASS: Kueue partial-admitted the Workload at minCount (quota fits only 2 CPU)." +echo "PASS: 2 of 4 Pods are Running; the other 2 remain scheduling-gated." echo echo "Inspect:" echo " ${KUBECTL} -n ${NAMESPACE} get workload ${PODGANG_NAME} -o yaml" diff --git a/operator/Makefile b/operator/Makefile index 6d8d73d46..cc25305f1 100644 --- a/operator/Makefile +++ b/operator/Makefile @@ -249,11 +249,12 @@ run-soak-test: # Make targets for local development and testing # ------------------------------------------------------------- # Starts a local k8s cluster using kind. -# Usage: make kind-up [FAKE_NODES=] +# Usage: make kind-up [FAKE_NODES=] [WORKERS=] # Example: make kind-up FAKE_NODES=20 +# Example: make kind-up WORKERS=1 # 2-node cluster (control-plane + 1 worker) .PHONY: kind-up kind-up: $(KIND) $(YQ) - @$(MODULE_HACK_DIR)/kind-up.sh $(if $(FAKE_NODES),--fake-nodes $(FAKE_NODES)) + @$(MODULE_HACK_DIR)/kind-up.sh $(if $(FAKE_NODES),--fake-nodes $(FAKE_NODES)) $(if $(WORKERS),--workers $(WORKERS)) # Stops the local k8s cluster. .PHONY: kind-down diff --git a/operator/hack/kind-up.sh b/operator/hack/kind-up.sh index 6bd86ac0b..55913254a 100755 --- a/operator/hack/kind-up.sh +++ b/operator/hack/kind-up.sh @@ -29,6 +29,7 @@ DEPLOY_REGISTRY=true RECREATE_CLUSTER=false FEATURE_GATES=() FAKE_NODES=0 +WORKER_NODES=0 USAGE="" function kind::create_usage() { @@ -40,6 +41,7 @@ function kind::create_usage() { -r | --recreate If this flag is specified then it will recreate the cluster if it already exists. -g | --feature-gates Comma separated list of feature gates to enable on the cluster. -f | --fake-nodes Number of fake nodes to create using KWOK. Default value is 0. + -w | --workers Number of kind worker nodes to add. Default value is 0. ") echo "${usage}" } @@ -83,6 +85,10 @@ function kind::parse_flags() { shift FAKE_NODES=$1 ;; + --workers | -w) + shift + WORKER_NODES=$1 + ;; -h | --help) shift echo "${USAGE}" @@ -114,6 +120,13 @@ nodes: - role: control-plane image: kindest/node:v1.35.1 EOF + local worker_index + for ((worker_index = 0; worker_index < WORKER_NODES; worker_index++)); do + cat >>"${KIND_CONFIG_DIR}/cluster-config.yaml" <