WIP

Author: pajakd

slice/internal/controller/workload_controller.go

@@ -50,6 +50,7 @@ import (
 
 	"tpu-slice-controller/api/v1beta1"
 	"tpu-slice-controller/internal/core"
+	"tpu-slice-controller/internal/features"
 	"tpu-slice-controller/internal/topology"
 	"tpu-slice-controller/internal/util/api"
 	"tpu-slice-controller/internal/util/node"
@@ -851,13 +852,8 @@ func buildPodSetUpdates(wl *kueue.Workload) []kueue.PodSetUpdate {
 	var podSetUpdates []kueue.PodSetUpdate
 	for _, ps := range wl.Spec.PodSets {
 		if topology := core.GetTPUTopology(ps.Template); topology != "" {
-			labels := make(map[string]string)
-			if features.Enabled(features.NodesInSlicesAntiAffinity) {
-				labels[core.PodWebhookLabelKey] = "true"
-			}
 			podSetUpdates = append(podSetUpdates, kueue.PodSetUpdate{
-				Name:   ps.Name,
-				Labels: labels,
+				Name: ps.Name,
 				NodeSelector: map[string]string{
 					core.TPUTopologyAnnotation: topology,
 				},

slice/internal/features/features.go

@@ -25,16 +25,11 @@ import (
 )
 
 const (
-	// Adds AntiAffinity that excludes nodes belonging to existing slices from scheduling.
-	NodesInSlicesAntiAffinity featuregate.Feature = "NodesInSlicesAntiAffinity"
 	// FailOnUntoleratedDegradedSlice treats degraded slices as failed if the workload requested only healthy slices.
 	FailOnUntoleratedDegradedSlice featuregate.Feature = "FailOnUntoleratedDegradedSlice"
 )
 
 var defaultVersionedFeatureGates = map[featuregate.Feature]featuregate.VersionedSpecs{
-	NodesInSlicesAntiAffinity: {
-		{Version: version.MustParse("0.1"), Default: true, PreRelease: featuregate.Alpha},
-	},
 	FailOnUntoleratedDegradedSlice: {
 		{Version: version.MustParse("0.1"), Default: true, PreRelease: featuregate.Alpha},
 	},

slice/test/e2e/jobset_test.go

@@ -1478,144 +1478,6 @@ var _ = ginkgo.Describe("JobSet", func() {
 			})
 		})
 
-		ginkgo.It("should only admit the workload if nodes are not used in some other slices", func() {
-			nodes := &corev1.NodeList{}
-			gomega.Expect(k8sClient.List(ctx, nodes, client.HasLabels{core.TPUSubBlockLabel})).To(gomega.Succeed())
-			var partitionIDs []string
-			for _, node := range nodes.Items {
-				partitionIDs = append(partitionIDs, node.Labels[core.TPUSubBlockLabel])
-			}
-			gomega.Expect(partitionIDs).ToNot(gomega.BeEmpty())
-			manualSlice := &slice.Slice{
-				ObjectMeta: metav1.ObjectMeta{
-					Name: "manual-slice-blocking",
-				},
-				Spec: slice.SliceSpec{
-					Type:         slice.TypeTpu7x,
-					Topology:     "4x4x4",
-					PartitionIds: partitionIDs,
-				},
-			}
-			utils.MustCreate(ctx, k8sClient, manualSlice)
-			utils.SetSliceReady(ctx, k8sClient, client.ObjectKeyFromObject(manualSlice), manualSlice.Spec.Topology)
-
-			jobSet := testingjobsjobset.MakeJobSet("jobset", ns.Name).
-				Queue(lq.Name).
-				ReplicatedJobs(
-					testingjobsjobset.ReplicatedJobRequirements{
-						Name:        "rj1",
-						Image:       utils.E2eTestAgnHostImage,
-						Args:        utils.BehaviorWaitForDeletion,
-						Replicas:    1,
-						Parallelism: 16,
-						Completions: 16,
-						PodAnnotations: map[string]string{
-							core.TPUSliceTopologyAnnotation: "4x4x4",
-						},
-						NodeSelector: map[string]string{
-							core.TPUAcceleratorLabel:           string(slice.TypeTpu7x),
-							core.TPUSliceHealthNodeSelectorKey: core.TPUSliceHealthNodeSelectorHealthy,
-						},
-					},
-				).
-				RequestAndLimit("rj1", core.TPUResourceName, "4").
-				Obj()
-
-			ginkgo.By("Creating a JobSet", func() {
-				utils.MustCreate(ctx, k8sClient, jobSet)
-			})
-
-			createdWorkload := &kueue.Workload{}
-			wlKey := types.NamespacedName{
-				Name:      jobsetcontroller.GetWorkloadNameForJobSet(jobSet.Name, jobSet.UID),
-				Namespace: ns.Name,
-			}
-
-			ginkgo.By("Check that the Workload is not admissible", func() {
-				gomega.Consistently(func(g gomega.Gomega) {
-					err := k8sClient.Get(ctx, wlKey, createdWorkload)
-					if err != nil {
-						g.Expect(client.IgnoreNotFound(err)).To(gomega.Succeed())
-					} else {
-						g.Expect(createdWorkload.Status.Admission).Should(gomega.BeNil())
-					}
-				}, utils.ConsistentDuration, utils.Interval).Should(gomega.Succeed())
-			})
-
-			ginkgo.By("Deleting the manual slice and unlabeling nodes", func() {
-				utils.ExpectObjectToBeDeleted(ctx, k8sClient, manualSlice, true)
-				for _, node := range nodes.Items {
-					gomega.Eventually(func(g gomega.Gomega) {
-						n := &corev1.Node{}
-						g.Expect(k8sClient.Get(ctx, client.ObjectKeyFromObject(&node), n)).To(gomega.Succeed())
-						delete(n.Labels, core.TPUSliceNodeLabel)
-						delete(n.Labels, core.TPUTopologyAnnotation)
-						g.Expect(k8sClient.Update(ctx, n)).To(gomega.Succeed())
-					}, utils.Timeout, utils.Interval).Should(gomega.Succeed())
-				}
-			})
-
-			ginkgo.By("Waiting for Admission of the Workload", func() {
-				gomega.Eventually(func(g gomega.Gomega) {
-					g.Expect(k8sClient.Get(ctx, wlKey, createdWorkload)).Should(gomega.Succeed())
-					g.Expect(createdWorkload.Status.Admission).ShouldNot(gomega.BeNil())
-				}, utils.Timeout, utils.Interval).Should(gomega.Succeed())
-			})
-
-			createdSlice := &slice.Slice{}
-			sliceKey := core.SliceKeyFromWorkload(createdWorkload, "rj1", 0)
-
-			ginkgo.By("Checking that Slice is created", func() {
-				gomega.Eventually(func(g gomega.Gomega) {
-					g.Expect(k8sClient.Get(ctx, sliceKey, createdSlice)).To(gomega.Succeed())
-					g.Expect(createdSlice.Spec.PartitionIds).To(gomega.HaveLen(1))
-					g.Expect(createdSlice.Spec.Topology).To(gomega.Equal("4x4x4"))
-					g.Expect(createdSlice.Spec.Type).To(gomega.Equal(slice.TypeTpu7x))
-				}, utils.Timeout, utils.Interval).Should(gomega.Succeed())
-			})
-
-			ginkgo.By("Adding Ready condition", func() {
-				utils.SetSliceReady(ctx, k8sClient, sliceKey, "4x4x4")
-			})
-
-			ginkgo.By("Checking that the Workload is admitted and admission check status is ready", func() {
-				gomega.Eventually(func(g gomega.Gomega) {
-					g.Expect(k8sClient.Get(ctx, wlKey, createdWorkload)).Should(gomega.Succeed())
-					g.Expect(workload.IsAdmitted(createdWorkload)).Should(gomega.BeTrue())
-					g.Expect(createdWorkload.Status.AdmissionChecks).Should(gomega.BeComparableTo([]kueue.AdmissionCheckState{{
-						Name:    kueue.AdmissionCheckReference(ac.Name),
-						State:   kueue.CheckStateReady,
-						Message: `Slices are in states: 1 ACTIVE`,
-					}}, cmpopts.IgnoreFields(kueue.AdmissionCheckState{}, "LastTransitionTime", "PodSetUpdates")))
-				}, utils.LongTimeout, utils.Timeout).Should(gomega.Succeed())
-			})
-
-			ginkgo.By("Checking that all pods are running with topology node selector and without anti-affinity", func() {
-				pods := &corev1.PodList{}
-				gomega.Eventually(func(g gomega.Gomega) {
-					g.Expect(k8sClient.List(ctx, pods, client.InNamespace(ns.Name))).To(gomega.Succeed())
-					g.Expect(pods.Items).Should(gomega.HaveLen(int(16)))
-					for _, pod := range pods.Items {
-						g.Expect(pod.Spec.NodeSelector).To(gomega.HaveKeyWithValue(core.TPUTopologyAnnotation, "4x4x4"))
-						g.Expect(pod.Spec.Affinity).To(gomega.BeNil())
-						g.Expect(pod.Status.Phase).To(gomega.Equal(corev1.PodRunning))
-					}
-				}, utils.LongTimeout, utils.Interval).Should(gomega.Succeed())
-			})
-
-			ginkgo.By("Deleting JobSet", func() {
-				utils.ExpectObjectToBeDeleted(ctx, k8sClient, jobSet, true)
-			})
-
-			ginkgo.By("Checking that Slice is deleted", func() {
-				utils.ExpectObjectToBeDeleted(ctx, k8sClient, createdSlice, false)
-			})
-
-			ginkgo.By("Checking that Workload is deleted", func() {
-				utils.ExpectObjectToBeDeleted(ctx, k8sClient, createdWorkload, false)
-			})
-		})
-
 		ginkgo.It("should handle mixed tolerance for degraded slices across multiple PodSets", func() {
 			jobSet := testingjobsjobset.MakeJobSet("jobset", ns.Name).
 				Queue(lq.Name).