AI-Hypercomputer
diff --git a/‎slice/api/v1alpha1/slice_types.go‎
Lines changed: 29 additions & 25 deletions b/‎slice/api/v1alpha1/slice_types.go‎
Lines changed: 29 additions & 25 deletions
diff --git a/‎slice/api/v1alpha1/zz_generated.deepcopy.go‎
Lines changed: 4 additions & 4 deletions b/‎slice/api/v1alpha1/zz_generated.deepcopy.go‎
Lines changed: 4 additions & 4 deletions
diff --git a/‎slice/config/crd/bases/slice.accelerator.gke.io_slices.yaml‎
Lines changed: 16 additions & 8 deletions b/‎slice/config/crd/bases/slice.accelerator.gke.io_slices.yaml‎
Lines changed: 16 additions & 8 deletions
diff --git a/‎slice/internal/controller/workload_controller.go‎
Lines changed: 42 additions & 71 deletions b/‎slice/internal/controller/workload_controller.go‎
Lines changed: 42 additions & 71 deletions
@@ -21,21 +21,33 @@ import (
 )
 
 // NOTE: json tags are required.  Any new fields you add must have json tags for the fields to be serialized.
+type Type string
+
+const (
+	TypeV6e   Type = "v6e"
+	TypeTpu7x Type = "tpu-v7x"
+)
 
 // SliceSpec defines the desired state of Slice.
 type SliceSpec struct {
-	// AcceleratorType specifies the type of accelerator used in this slice.
+	// Type specifies the type of accelerator used in this slice, e.g., "v6e", "tpu-v7x".
 	// +kubebuilder:validation:Immutable
-	Type string `json:"type"`
+	// +kubebuilder:validation:Enum=v6e;tpu-v7x
+	Type Type `json:"type"`
 
 	// Topology represents the network topology of the slice.
+	// It defines the physical arrangement of TPU chips in a 2D or 3D mesg.
+	// The topology must be specified in `<X>x<Y>` or `<X>x<Y>x<Z>` format.
 	// +kubebuilder:validation:Immutable
+	// +kubebuilder:validation:Pattern=^\d+x\d+(x\d+)?$
 	Topology string `json:"topology"`
 
-	// Partition Ids denotes the set of partitions to use to form a slice
-	// For slices that span multiple partitions, it will be a list of 4x4x4 Ids
-	// For sub-partition topology, it will be a single entry corresponding to the ID with the partition Ids
-	PartitionIDs []string `json:"partitionIds"`
+	// PartitionIds denotes the set of partitions to use to form a slice
+	// For slices that span multiple partitions, it will be a list of 4x4x4 IDs
+	// For sub-partition topology, it will be a single entry corresponding to the ID of the partition
+	// +kubebuilder:validation:Immutable
+	// +kubebuilder:validation:MinItems=1
+	PartitionIds []string `json:"partitionIds"`
 }
 
 // SliceStatus defines the observed state of Slice.
@@ -45,17 +57,17 @@ type SliceStatus struct {
 	Conditions []metav1.Condition `json:"conditions,omitempty"`
 
 	// Populated to match the physical topology of block the Super-Slice is running on
-	BlockID string `json:"blockId,omitempty"`
+	BlockId string `json:"blockId,omitempty"`
 
 	// Populated to list of physical topology of sub-block the Super-Slice is running on
-	SubBlockIDs []string `json:"subBlockIds,omitempty"`
+	SubBlockIds []string `json:"subBlockIds,omitempty"`
 }
 
 // +kubebuilder:object:root=true
 // +kubebuilder:subresource:status
 // +kubebuilder:printcolumn:name="Type",type=string,JSONPath=`.spec.type`
 // +kubebuilder:printcolumn:name="Topology",type=string,JSONPath=`.spec.topology`
-// +kubebuilder:printcolumn:name="Status",type=string,JSONPath=`.status.conditions[0].type`
+// +kubebuilder:printcolumn:name="State",type=string,JSONPath=`.status.conditions[?(@.type=="Ready")].reason`
 // +kubebuilder:printcolumn:name="Age",type="date",JSONPath=".metadata.creationTimestamp"
 // Slice is the Schema for the slices API.
 type Slice struct {
@@ -75,22 +87,14 @@ type SliceList struct {
 	Items           []Slice `json:"items"`
 }
 
-// SliceConditionType defines the type of condition
-type SliceConditionType string
-
-const (
-	// Forming means the slice is being created and configured.
-	Forming SliceConditionType = "Forming"
-	// Ready means the slice is fully operational.
-	Ready SliceConditionType = "Ready"
-	// Degraded means the slice is operational but with reduced capacity or performance.
-	Degraded SliceConditionType = "Degraded"
-	// Deformed means the slice is being torn down.
-	Deformed SliceConditionType = "Deformed"
-	// Error means the slice has encountered an error and is not operational.
-	Error SliceConditionType = "Error"
-)
-
 func init() {
 	SchemeBuilder.Register(&Slice{}, &SliceList{})
 }
+
+const (
+	// Represent the underlying hardware readiness status
+	SliceStateConditionType = "Ready"
+	// Represent whether the user/scheduler should take action on the slice
+	// The slice is in an error state that can't not automatically recover
+	SliceCreationFailedConditionType = "SliceCreationFailed"
+)
@@ -21,8 +21,8 @@ spec:
     - jsonPath: .spec.topology
       name: Topology
       type: string
-    - jsonPath: .status.conditions[0].type
-      name: Status
+    - jsonPath: .status.conditions[?(@.type=="Ready")].reason
+      name: State
       type: string
     - jsonPath: .metadata.creationTimestamp
       name: Age
@@ -54,18 +54,26 @@ spec:
             properties:
               partitionIds:
                 description: |-
-                  Partition Ids denotes the set of partitions to use to form a slice
-                  For slices that span multiple partitions, it will be a list of 4x4x4 Ids
-                  For sub-partition topology, it will be a single entry corresponding to the ID with the partition Ids
+                  PartitionIds denotes the set of partitions to use to form a slice
+                  For slices that span multiple partitions, it will be a list of 4x4x4 IDs
+                  For sub-partition topology, it will be a single entry corresponding to the ID of the partition
                 items:
                   type: string
+                minItems: 1
                 type: array
               topology:
-                description: Topology represents the network topology of the slice.
+                description: |-
+                  Topology represents the network topology of the slice.
+                  It defines the physical arrangement of TPU chips in a 2D or 3D mesg.
+                  The topology must be specified in `<X>x<Y>` or `<X>x<Y>x<Z>` format.
+                pattern: ^\d+x\d+(x\d+)?$
                 type: string
               type:
-                description: AcceleratorType specifies the type of accelerator used
-                  in this slice.
+                description: Type specifies the type of accelerator used in this slice,
+                  e.g., "v6e", "tpu-v7x".
+                enum:
+                - v6e
+                - tpu-v7x
                 type: string
             required:
             - partitionIds
 
@@ -20,7 +20,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"slices"
 	"sort"
 	"strings"
 	"time"
@@ -162,7 +161,7 @@ func (r *WorkloadReconciler) Reconcile(ctx context.Context, req ctrl.Request) (c
 		return ctrl.Result{}, err
 	}
 
-	deleted, _, errored, _ := r.groupSlices(slices)
+	deleted, toDelete, _ := r.groupSlices(slices)
 	if len(deleted) > 0 {
 		log.V(3).Info(
 			"Waiting for deleted Slices to be cleaned up; skipping reconciliation for now",
@@ -181,12 +180,12 @@ func (r *WorkloadReconciler) Reconcile(ctx context.Context, req ctrl.Request) (c
 		return ctrl.Result{}, client.IgnoreNotFound(err)
 	}
 
-	if len(errored) > 0 {
+	if len(toDelete) > 0 {
 		log.V(3).Info(
-			"Deleting errored Slices",
-			"erroredSlices", klog.KObjSlice(errored),
+			"Deleting Slices",
+			"slices", klog.KObjSlice(toDelete),
 		)
-		err = r.deleteSlices(ctx, errored)
+		err = r.deleteSlices(ctx, toDelete)
 		if err != nil {
 			return ctrl.Result{}, err
 		}
@@ -256,29 +255,21 @@ func (r *WorkloadReconciler) cleanupSlices(ctx context.Context, wl *kueue.Worklo
 		return false, err
 	}
 
-	deleted, deformed, errored, other := r.groupSlices(slices)
+	deleted, toDelete, other := r.groupSlices(slices)
 
 	if len(deleted) == len(slices) {
 		log.V(3).Info("All slices already deleted; finishing cleanup")
 		return true, nil
 	}
 
-	if len(deformed) > 0 {
-		log.V(3).Info("Found Slices in deformed state; cleaning them up", "deformedSlices", klog.KObjSlice(deformed))
-		// We still need to delete deformed Slices because requeueing causes a conflict error during Slice creation.
-		err = r.deleteSlices(ctx, deformed)
-		if err != nil {
-			return false, err
-		}
-	}
-
-	if len(other)+len(errored) > 0 {
+	if len(other)+len(toDelete) > 0 {
 		terminated, err := r.ownerPodsFinished(ctx, wl)
 		if err != nil || !terminated {
 			return false, err
 		}
 	}
-	toDelete := append(errored, other...)
+	// after pods are terminated we should cleanup all the slices (including active ones)
+	toDelete = append(toDelete, other...)
 	log.V(3).Info("Deleting Slices", "slices", klog.KObjSlice(toDelete))
 	err = r.deleteSlices(ctx, toDelete)
 	if err != nil {
@@ -300,34 +291,31 @@ func (r *WorkloadReconciler) findWorkloadSlices(ctx context.Context, wl *kueue.W
 	return slices.Items, nil
 }
 
-// groupSlices categorizes a list of Slice objects into four groups based on their state.
-// It separates slices into deleted (marked for deletion), deformed (being torn down),
-// errored (has encountered an error) and other (active) slices.
+// groupSlices categorizes a list of Slice objects into three groups based on their state.
+// It separates slices into deleted (marked for deletion), ones that should be delete
+// (errored and stale) and other (active) slices.
 //
 // Parameters:
 //
 //	slices - A slice of v1alpha1.Slice objects to be categorized.
 //
 // Returns:
 //   - A slice containing deleted Slice objects (with non-zero DeletionTimestamp).
-//   - A slice containing deformed Slice objects (being torn down).
-//   - A slice containing errored Slice objects.
-//   - A slice containing other Slice objects (active/valid slices).
-func (r *WorkloadReconciler) groupSlices(slices []v1alpha1.Slice) ([]v1alpha1.Slice, []v1alpha1.Slice, []v1alpha1.Slice, []v1alpha1.Slice) {
-	var deleted, deformed, errored, other []v1alpha1.Slice
+//   - A slice containing Slice objects that should be deleted (errored and stale slices).
+//   - A slice containing other Slice objects (active/initializing slices).
+func (r *WorkloadReconciler) groupSlices(slices []v1alpha1.Slice) ([]v1alpha1.Slice, []v1alpha1.Slice, []v1alpha1.Slice) {
+	var deleted, toDelete, other []v1alpha1.Slice
 	for _, slice := range slices {
-		switch {
-		case !slice.DeletionTimestamp.IsZero():
+		switch core.GetSliceState(slice) {
+		case core.SliceStateDeleted:
 			deleted = append(deleted, slice)
-		case core.Deformed(&slice):
-			deformed = append(deformed, slice)
-		case core.IsError(&slice):
-			errored = append(errored, slice)
+		case core.SliceStateFailed, core.SliceStateStale:
+			toDelete = append(toDelete, slice)
 		default:
 			other = append(other, slice)
 		}
 	}
-	return deleted, deformed, errored, other
+	return deleted, toDelete, other
 }
 
 func (r *WorkloadReconciler) deleteSlices(ctx context.Context, slices []v1alpha1.Slice) error {
@@ -507,18 +495,16 @@ func shouldCreateSliceForPodSetAssignment(wl *kueue.Workload, psa kueue.PodSetAs
 	return false
 }
 
-func parseTopologyAssignmentIntoNodeSelector(slice *v1alpha1.Slice, topologyAssignment *kueue.TopologyAssignment, nodes map[string]corev1.Node) {
-	subBlockIDs := sets.New[string]()
+func parseTopologyAssignmentIntoPartitionIds(slice *v1alpha1.Slice, topologyAssignment *kueue.TopologyAssignment, nodes map[string]corev1.Node) {
+	subBlockIds := sets.New[string]()
 	// we already validated that all assignments have a valid level,
 	// in validateRelevantWorkload.
 	hostnameLevelIndex := topology.HostnameLevelIndex(topologyAssignment)
 	for _, domain := range topologyAssignment.Domains {
 		nodeName := domain.Values[hostnameLevelIndex]
-		subBlockIDs.Insert(topology.GetTPUSubBlockLabelValue(nodes, nodeName))
+		subBlockIds.Insert(topology.GetTPUSubBlockLabelValue(nodes, nodeName))
 	}
-	partitionIDs := subBlockIDs.UnsortedList()
-	slices.Sort(partitionIDs)
-	slice.Spec.PartitionIDs = partitionIDs
+	slice.Spec.PartitionIds = sets.List(subBlockIds)
 }
 
 func (r *WorkloadReconciler) createSlice(ctx context.Context, wl *kueue.Workload, ac *kueue.AdmissionCheckState, psa *kueue.PodSetAssignment, nodes map[string]corev1.Node) (*v1alpha1.Slice, error) {
@@ -529,10 +515,10 @@ func (r *WorkloadReconciler) createSlice(ctx context.Context, wl *kueue.Workload
 	if err := controllerutil.SetControllerReference(wl, slice, r.client.Scheme()); err != nil {
 		return nil, err
 	}
-	parseTopologyAssignmentIntoNodeSelector(slice, psa.TopologyAssignment, nodes)
+	parseTopologyAssignmentIntoPartitionIds(slice, psa.TopologyAssignment, nodes)
 
 	ps := podset.FindPodSetByName(wl.Spec.PodSets, psa.Name)
-	slice.Spec.Type = core.GetTPUAccelerator(ps.Template)
+	slice.Spec.Type = v1alpha1.Type(core.GetTPUAccelerator(ps.Template))
 	slice.Spec.Topology = core.GetTPUTopology(ps.Template)
 
 	if err := r.client.Create(ctx, slice); err != nil {
@@ -594,8 +580,8 @@ func (r *WorkloadReconciler) syncAdmissionCheckStatus(ctx context.Context, wl *k
 	if ac.Message != originalMessage {
 		// Logging error messages if exists
 		for _, slice := range slices {
-			cond := meta.FindStatusCondition(slice.Status.Conditions, string(v1alpha1.Error))
-			if cond != nil && cond.Status == metav1.ConditionTrue {
+			cond := meta.FindStatusCondition(slice.Status.Conditions, v1alpha1.SliceStateConditionType)
+			if cond != nil && cond.Status == metav1.ConditionFalse && cond.Reason == string(core.MMIGHealthStatusFailed) {
 				log.V(2).Info(
 					"WARNING: The Slice is not operational due to an error",
 					"slice", klog.KObj(&slice),
@@ -608,42 +594,27 @@ func (r *WorkloadReconciler) syncAdmissionCheckStatus(ctx context.Context, wl *k
 	return nil
 }
 
-func groupSlicesByState(slices []v1alpha1.Slice) (map[v1alpha1.SliceConditionType][]v1alpha1.Slice, []v1alpha1.Slice) {
-	slicesByState := make(map[v1alpha1.SliceConditionType][]v1alpha1.Slice)
-	var noState []v1alpha1.Slice
+func groupSlicesByState(slices []v1alpha1.Slice) map[core.SliceState][]v1alpha1.Slice {
+	slicesByState := make(map[core.SliceState][]v1alpha1.Slice)
 	for _, slice := range slices {
-		foundState := false
-		for _, status := range core.SliceStates {
-			if meta.IsStatusConditionTrue(slice.Status.Conditions, string(status)) {
-				slicesByState[status] = append(slicesByState[status], slice)
-				foundState = true
-				break
-			}
-		}
-		if !foundState {
-			noState = append(noState, slice)
-		}
+		slicesByState[core.GetSliceState(slice)] = append(slicesByState[core.GetSliceState(slice)], slice)
 	}
-	return slicesByState, noState
+	return slicesByState
 }
 
 func prepareAdmissionCheckStatus(ac *kueue.AdmissionCheckState, slices []v1alpha1.Slice) {
-	slicesByState, noState := groupSlicesByState(slices)
+	slicesByState := groupSlicesByState(slices)
 
 	switch {
-	case len(slicesByState[v1alpha1.Deformed]) > 0:
-		ac.State = kueue.CheckStateRejected
-	case len(slicesByState[v1alpha1.Error]) > 0:
-		ac.State = kueue.CheckStateRetry
-	case len(slices) == len(slicesByState[v1alpha1.Degraded])+len(slicesByState[v1alpha1.Ready]):
+	case len(slices) == len(slicesByState[core.SliceStateActive])+len(slicesByState[core.SliceStateActiveDegraded]):
 		ac.State = kueue.CheckStateReady
+	case len(slicesByState[core.SliceStateFailed]) > 0:
+		ac.State = kueue.CheckStateRetry
+	case len(slicesByState[core.SliceStateCreated])+len(slicesByState[core.SliceStateActivating]) > 0:
+		ac.State = kueue.CheckStatePending
 	}
 
 	var stateMessages []string
-	if len(noState) > 0 {
-		stateMessages = append(stateMessages, fmt.Sprintf("%d Created", len(noState)))
-	}
-
 	for _, state := range core.SliceStates {
 		if count := len(slicesByState[state]); count > 0 {
 			stateMessages = append(stateMessages, fmt.Sprintf("%d %s", count, state))
@@ -652,10 +623,10 @@ func prepareAdmissionCheckStatus(ac *kueue.AdmissionCheckState, slices []v1alpha
 
 	ac.Message = fmt.Sprintf("Slices are in states: %s", strings.Join(stateMessages, ", "))
 
-	if len(slicesByState[v1alpha1.Error]) > 0 {
+	if len(slicesByState[core.SliceStateFailed]) > 0 {
 		var errMessages []string
-		for _, slice := range slicesByState[v1alpha1.Error] {
-			cond := meta.FindStatusCondition(slice.Status.Conditions, string(v1alpha1.Error))
+		for _, slice := range slicesByState[core.SliceStateFailed] {
+			cond := meta.FindStatusCondition(slice.Status.Conditions, v1alpha1.SliceStateConditionType)
 			errMessages = append(errMessages, cond.Message)
 		}
 		ac.Message += ". Errors: " + strings.Join(errMessages, "; ")