[partialAdmission] Add feature gate

apis/kueue/v1beta1/workload_types.go

@@ -114,6 +114,9 @@ type PodSet struct {
 	//
 	// If not provided, partial admission for the current PodSet is not
 	// enabled.
+	//
+	// This is an alpha field and requires enabling PartialAdmission feature gate.
+	//
 	// +optional
 	MinCount *int32 `json:"minCount,omitempty"`
 }

charts/kueue/templates/crd/kueue.x-k8s.io_workloads.yaml

@@ -78,7 +78,8 @@ spec:
                       description: "minCount is the minimum number of pods for the
                         spec acceptable if the workload supports partial admission.
                         \n If not provided, partial admission for the current PodSet
-                        is not enabled."
+                        is not enabled. \n This is an alpha field and requires enabling
+                        PartialAdmission feature gate."
                       format: int32
                       type: integer
                     name:

config/components/crd/bases/kueue.x-k8s.io_workloads.yaml

@@ -65,7 +65,8 @@ spec:
                       description: "minCount is the minimum number of pods for the
                         spec acceptable if the workload supports partial admission.
                         \n If not provided, partial admission for the current PodSet
-                        is not enabled."
+                        is not enabled. \n This is an alpha field and requires enabling
+                        PartialAdmission feature gate."
                       format: int32
                       type: integer
                     name:

keps/420-partial-admission/kep.yaml

@@ -13,7 +13,7 @@ approvers:
 
 
 # The target maturity stage in the current dev cycle for this KEP.
-stage: stable
+stage: alpha
 
 # The most recent milestone for which work toward delivery of this KEP has been
 # done. This can be the current (upcoming) milestone, if it is being actively
@@ -22,5 +22,9 @@ latest-milestone: "v0.4"
 
 # The milestone at which this feature was, or is targeted to be, at each stage.
 milestone:
-  stable: "v0.5"
+  alpha: "v0.4"
+  beta: "v0.5"
+  stable: "v0.6"
 
+feature-gates:
+  - name: PartialAdmission

pkg/features/kube_features.go

@@ -17,18 +17,22 @@ limitations under the License.
 package features
 
 import (
+	"fmt"
+	"testing"
+
 	"k8s.io/apimachinery/pkg/util/runtime"
 	utilfeature "k8s.io/apiserver/pkg/util/feature"
 	"k8s.io/component-base/featuregate"
+	featuregatetesting "k8s.io/component-base/featuregate/testing"
 )
 
 const (
-	// owner: @kerthcet
-	// kep: <link to kep>
-	// alpha: v0.3.1
+	// owner: @trasc
+	// kep: https://github.com/kubernetes-sigs/kueue/tree/main/keps/420-partial-admission
+	// alpha: v0.4
 	//
-	// TODO: This is for reference. Remove it once we have real feature gates.
-	FeatureGateForTest featuregate.Feature = "FeatureGateForTest"
+	// Enables partial admission.
+	PartialAdmission featuregate.Feature = "PartialAdmission"
 )
 
 func init() {
@@ -42,5 +46,18 @@ func init() {
 // Entries are separated from each other with blank lines to avoid sweeping gofmt changes
 // when adding or removing one entry.
 var defaultFeatureGates = map[featuregate.Feature]featuregate.FeatureSpec{
-	FeatureGateForTest: {Default: false, PreRelease: featuregate.Alpha},
+	PartialAdmission: {Default: false, PreRelease: featuregate.Alpha},
+}
+
+func SetFeatureGateDuringTest(tb testing.TB, f featuregate.Feature, value bool) func() {
+	return featuregatetesting.SetFeatureGateDuringTest(tb, utilfeature.DefaultFeatureGate, f, value)
+}
+
+// Helper for `utilfeature.DefaultFeatureGate.Enabled()`
+func Enabled(f featuregate.Feature) bool {
+	return utilfeature.DefaultFeatureGate.Enabled(f)
+}
+
+func SetEnable(f featuregate.Feature, v bool) error {
+	return utilfeature.DefaultMutableFeatureGate.Set(fmt.Sprintf("%s=%v", f, v))
 }

pkg/features/kube_features_test.go

@@ -24,9 +24,9 @@ import (
 )
 
 func TestFeatureGate(t *testing.T) {
-	defer featuregatetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, FeatureGateForTest, true)()
+	defer featuregatetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, PartialAdmission, true)()
 
-	if !utilfeature.DefaultFeatureGate.Enabled(FeatureGateForTest) {
+	if !utilfeature.DefaultFeatureGate.Enabled(PartialAdmission) {
 		t.Error("feature gate should be enabled")
 	}
 }

pkg/scheduler/scheduler.go

@@ -38,6 +38,7 @@ import (
 
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 	"sigs.k8s.io/kueue/pkg/cache"
+	"sigs.k8s.io/kueue/pkg/features"
 	"sigs.k8s.io/kueue/pkg/metrics"
 	"sigs.k8s.io/kueue/pkg/queue"
 	"sigs.k8s.io/kueue/pkg/scheduler/flavorassigner"
@@ -287,8 +288,8 @@ func (s *Scheduler) getAssignments(log logr.Logger, wl *workload.Info, snap *cac
 		fullAssignmentTargets = s.preemptor.GetTargets(*wl, fullAssignment, snap)
 	}
 
-	// if we can preempt
-	if len(fullAssignmentTargets) > 0 {
+	// if the feature gate is not enabled or we can preempt
+	if !features.Enabled(features.PartialAdmission) || len(fullAssignmentTargets) > 0 {
 		return fullAssignment, fullAssignmentTargets
 	}
 

pkg/scheduler/scheduler_test.go

@@ -39,6 +39,7 @@ import (
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 	"sigs.k8s.io/kueue/pkg/cache"
 	"sigs.k8s.io/kueue/pkg/constants"
+	"sigs.k8s.io/kueue/pkg/features"
 	"sigs.k8s.io/kueue/pkg/queue"
 	"sigs.k8s.io/kueue/pkg/scheduler/flavorassigner"
 	"sigs.k8s.io/kueue/pkg/util/routine"
@@ -191,6 +192,9 @@ func TestSchedule(t *testing.T) {
 		// additional*Queues can hold any extra queues needed by the tc
 		additionalClusterQueues []kueue.ClusterQueue
 		additionalLocalQueues   []kueue.LocalQueue
+
+		// enable partial admission
+		enablePartialAdmission bool
 	}{
 		"workload fits in single clusterQueue": {
 			workloads: []kueue.Workload{
@@ -742,7 +746,8 @@ func TestSchedule(t *testing.T) {
 					},
 				},
 			},
-			wantScheduled: []string{"sales/new"},
+			wantScheduled:          []string{"sales/new"},
+			enablePartialAdmission: true,
 		},
 		"partial admission single variable pod set, preempt first": {
 			workloads: []kueue.Workload{
@@ -783,6 +788,7 @@ func TestSchedule(t *testing.T) {
 			wantLeft: map[string]sets.Set[string]{
 				"eng-beta": sets.New("eng-beta/new"),
 			},
+			enablePartialAdmission: true,
 		},
 		"partial admission multiple variable pod sets": {
 			workloads: []kueue.Workload{
@@ -840,11 +846,16 @@ func TestSchedule(t *testing.T) {
 					},
 				},
 			},
-			wantScheduled: []string{"sales/new"},
+			wantScheduled:          []string{"sales/new"},
+			enablePartialAdmission: true,
 		},
 	}
+
 	for name, tc := range cases {
 		t.Run(name, func(t *testing.T) {
+			if tc.enablePartialAdmission {
+				defer features.SetFeatureGateDuringTest(t, features.PartialAdmission, true)()
+			}
 			log := testr.NewWithOptions(t, testr.Options{
 				Verbosity: 2,
 			})

pkg/webhooks/workload_webhook.go

@@ -32,6 +32,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/webhook"
 
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
+	"sigs.k8s.io/kueue/pkg/features"
 	"sigs.k8s.io/kueue/pkg/workload"
 )
 
@@ -63,6 +64,13 @@ func (w *WorkloadWebhook) Default(ctx context.Context, obj runtime.Object) error
 			podSet.Name = kueue.DefaultPodSetName
 		}
 	}
+
+	// drop minCounts if PartialAdmission is not enabled
+	if !features.Enabled(features.PartialAdmission) {
+		for i := range wl.Spec.PodSets {
+			wl.Spec.PodSets[i].MinCount = nil
+		}
+	}
 	return nil
 }
 

test/integration/controller/job/job_controller_test.go

@@ -35,6 +35,7 @@ import (
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 	"sigs.k8s.io/kueue/pkg/controller/jobframework"
 	workloadjob "sigs.k8s.io/kueue/pkg/controller/jobs/job"
+	"sigs.k8s.io/kueue/pkg/features"
 	"sigs.k8s.io/kueue/pkg/util/pointer"
 	"sigs.k8s.io/kueue/pkg/util/testing"
 	testingjob "sigs.k8s.io/kueue/pkg/util/testingjobs/job"
@@ -1028,6 +1029,11 @@ var _ = ginkgo.Describe("Job controller interacting with scheduler", func() {
 	})
 
 	ginkgo.It("Should schedule jobs when partial admission is enabled", func() {
+		origPartialAdmission := features.Enabled(features.PartialAdmission)
+		ginkgo.By("enable partial admission", func() {
+			gomega.Expect(features.SetEnable(features.PartialAdmission, true)).To(gomega.Succeed())
+		})
+
 		prodLocalQ = testing.MakeLocalQueue("prod-queue", ns.Name).ClusterQueue(prodClusterQ.Name).Obj()
 		job1 := testingjob.MakeJob("job1", ns.Name).
 			Queue(prodLocalQ.Name).
@@ -1089,5 +1095,9 @@ var _ = ginkgo.Describe("Job controller interacting with scheduler", func() {
 			}, util.Timeout, util.Interval).Should(gomega.Equal(pointer.Bool(true)))
 			gomega.Expect(*createdJob.Spec.Parallelism).To(gomega.BeEquivalentTo(5))
 		})
+
+		ginkgo.By("restore partial admission", func() {
+			gomega.Expect(features.SetEnable(features.PartialAdmission, origPartialAdmission)).To(gomega.Succeed())
+		})
 	})
 })