[slice] set type and topology (#677)

Author: gabesaba

slice/internal/controller/workload_controller.go

@@ -516,6 +516,10 @@ func (r *WorkloadReconciler) createSlice(ctx context.Context, wl *kueue.Workload
 	}
 	parseTopologyAssignmentIntoNodeSelector(slice, psa.TopologyAssignment, nodes)
 
+	ps := podset.FindPodSetByName(wl.Spec.PodSets, psa.Name)
+	slice.Spec.AcceleratorType = core.GetTPUAccelerator(ps.Template)
+	slice.Spec.AcceleratorTopology = core.GetTPUTopology(ps.Template)
+
 	if err := r.client.Create(ctx, slice); err != nil {
 		msg := fmt.Sprintf("Error creating Slice %q: %v", client.ObjectKeyFromObject(slice), err)
 		log.Error(err, msg)

slice/internal/controller/workload_controller_test.go

@@ -134,9 +134,13 @@ func TestWorkloadReconciler(t *testing.T) {
 		UID(baseWorkloadName).
 		AdmissionCheck(buildAdmissionCheckState(kueue.CheckStatePending, ""))
 	baseSlice1Wrapper := utiltesting.MakeSliceWrapper(core.SliceName(baseWorkloadName, "ps1"), corev1.NamespaceDefault).
+		AcceleratorType("tpu-v7x").
+		AcceleratorTopology("4x4x12").
 		ControllerReference(workloadGVK, baseWorkloadName, baseWorkloadName).
 		NodeSelector(map[string][]string{"cloud.google.com/gke-tpu-slice-4x4x4-id": {"subblock1"}})
 	baseSlice2Wrapper := baseSlice1Wrapper.Clone().Name(core.SliceName(baseWorkloadName, "ps2")).
+		AcceleratorType("tpu-v7x").
+		AcceleratorTopology("4x4x12").
 		NodeSelector(map[string][]string{"cloud.google.com/gke-tpu-slice-4x4x4-id": {"subblock2"}})
 
 	worker1Node := utiltesting.MakeNode("worker1").Label("cloud.google.com/gke-tpu-slice-4x4x4-id", "subblock1")

slice/internal/core/constants.go

@@ -24,4 +24,6 @@ const (
 
 	TPUSliceHealthNodeSelectorKey   = "cloud.google.com/gke-tpu-slice-4x4x4-health"
 	TPUSliceHealthNodeSelectorValue = "true"
+
+	AcceleratorTpu7x = "tpu-v7x"
 )

slice/internal/core/core.go

@@ -34,10 +34,18 @@ func IsValidTPUTopology(tpuTopology string) bool {
 }
 
 func IsValidTPUAccelerator(tpuAccelerator string) bool {
-	return tpuAccelerator == "tpu-v7x"
+	return tpuAccelerator == AcceleratorTpu7x
 }
 
 func IsRelevantPodTemplateSpec(spec corev1.PodTemplateSpec) bool {
-	return IsValidTPUTopology(spec.Annotations[TPUTopologyAnnotation]) &&
-		IsValidTPUAccelerator(spec.Spec.NodeSelector[TPUAcceleratorLabel])
+	return IsValidTPUTopology(GetTPUTopology(spec)) &&
+		IsValidTPUAccelerator(GetTPUAccelerator(spec))
+}
+
+func GetTPUTopology(spec corev1.PodTemplateSpec) string {
+	return spec.Annotations[TPUTopologyAnnotation]
+}
+
+func GetTPUAccelerator(spec corev1.PodTemplateSpec) string {
+	return spec.Spec.NodeSelector[TPUAcceleratorLabel]
 }

slice/test/e2e/jobset_test.go

@@ -221,6 +221,8 @@ var _ = ginkgo.Describe("JobSet", func() {
 						g.Expect(k8sClient.Get(ctx, sliceKey, createdSlice)).To(gomega.Succeed())
 						g.Expect(createdSlice.Spec.NodeSelector).To(gomega.HaveLen(1))
 						g.Expect(createdSlice.Spec.NodeSelector).To(gomega.BeComparableTo(tc.wantSliceNodeSelector))
+						g.Expect(createdSlice.Spec.AcceleratorTopology).To(gomega.Equal(tc.tpuTopology))
+						g.Expect(createdSlice.Spec.AcceleratorType).To(gomega.Equal("tpu-v7x"))
 					}, utils.Timeout, utils.Interval).Should(gomega.Succeed())
 				})