Update Kueue queues

config/kueue/kustomization.yaml

@@ -1,4 +1,4 @@
 apiVersion: kustomize.config.k8s.io/v1beta1
 kind: Kustomization
 resources:
-- "https://github.com/kubernetes-sigs/kueue/releases/download/v0.17.1/manifests.yaml"
+- "https://github.com/kubernetes-sigs/kueue/releases/download/v0.18.2/manifests.yaml"

internal/controller/deviceplugin_controller.go

@@ -71,6 +71,11 @@ func hexArgStr(s []string) string {
 }
 
 func addXpumdMounts(spec *core.PodSpec) {
+	if len(spec.Containers) == 0 {
+		klog.Warning("Cannot add xpumd mounts: no containers in pod spec")
+		return
+	}
+
 	for _, v := range spec.Volumes {
 		if v.Name == xpumdVolumeName {
 			return
@@ -98,6 +103,11 @@ func addXpumdMounts(spec *core.PodSpec) {
 }
 
 func removeXpumdMounts(spec *core.PodSpec) {
+	if len(spec.Containers) == 0 {
+		klog.Warning("Cannot remove xpumd mounts: no containers in pod spec")
+		return
+	}
+
 	for i, v := range spec.Volumes {
 		if v.Name == xpumdVolumeName {
 			spec.Volumes = slices.Delete(spec.Volumes, i, i+1)
@@ -145,6 +155,11 @@ func dpArgs(spec *v1alpha.ClusterPolicy) []string {
 }
 
 func (r *DevicePluginReconciler) updateDaemonSetObject(ds *apps.DaemonSet, spec *v1alpha.ClusterPolicy) {
+	if len(ds.Spec.Template.Spec.Containers) == 0 {
+		klog.Error(nil, "Cannot update DaemonSet: no containers in pod spec", "DaemonSet", ds.Name)
+		return
+	}
+
 	name := fmt.Sprintf("%s-device-plugin", spec.Name)
 
 	ds.Name = name

internal/controller/misc_controller.go

@@ -461,7 +461,7 @@ func (r *MiscReconciler) createResourceFlavor() *kueuev1beta2.ResourceFlavor {
 	}
 }
 
-func (r *MiscReconciler) modifyResourceFlavor(ctx context.Context) error {
+func (r *MiscReconciler) modifyResourceFlavor(ctx context.Context, cp *v1alpha.ClusterPolicy) error {
 	currentUnstructured := &unstructured.Unstructured{}
 	currentUnstructured.SetKind(kueueResourceFlavorKind)
 	currentUnstructured.SetAPIVersion(kueueAPIVersion)
@@ -476,7 +476,14 @@ func (r *MiscReconciler) modifyResourceFlavor(ctx context.Context) error {
 	}
 
 	if !found {
-		err := r.Create(ctx, r.createResourceFlavor())
+		resourceFlavor := r.createResourceFlavor()
+
+		if err := ctrl.SetControllerReference(cp, resourceFlavor, r.Scheme); err != nil {
+			klog.Error(err, "unable to set Kueue ResourceFlavor controller reference")
+			return err
+		}
+
+		err := r.Create(ctx, resourceFlavor)
 		if client.IgnoreAlreadyExists(err) != nil {
 			klog.Error(err, "unable to create Kueue ResourceFlavor '%s'", kueueFlavorName)
 			return err
@@ -590,7 +597,91 @@ func (r *MiscReconciler) createLocalQueue(clusterQueueName string, localQueueSpe
 	return &localQueue
 }
 
-func (r *MiscReconciler) modifyLocalQueues(ctx context.Context, clusterQueue *v1alpha.ClusterQueueSpec) error {
+func (r *MiscReconciler) keepLocalQueue(lq *kueuev1beta2.LocalQueue, keepLQs []v1alpha.LocalQueueSpec) bool {
+	for _, keep := range keepLQs {
+		if lq.Name == keep.Name && lq.Namespace == keep.Namespace {
+			return true
+		}
+	}
+	return false
+}
+
+func (r *MiscReconciler) pruneLocalQueues(ctx context.Context, lqList *kueuev1beta2.LocalQueueList, keepLQs []v1alpha.LocalQueueSpec, cqName string, cp *v1alpha.ClusterPolicy) {
+	for _, lq := range lqList.Items {
+		klog.V(3).Infof("Checking local queue '%s' attached to cluster queue '%s'", lq.Name, cqName)
+		if lq.Spec.ClusterQueue != kueuev1beta2.ClusterQueueReference(cqName) {
+			continue
+		}
+
+		if r.keepLocalQueue(&lq, keepLQs) {
+			klog.V(3).Infof("Keeping LocalQueue %s (in ClusterQueue %s)", lq.Name, cqName)
+			continue
+		}
+
+		// delete this local queue
+		if err := r.Delete(ctx, &lq); client.IgnoreNotFound(err) != nil {
+			klog.Warningf("Error when attempting to delete LocalQueue %s: %v", lq.Name, err)
+		} else {
+			klog.V(3).Infof("Deleted LocalQueue %s namespace %s (in ClusterQueue %s)", lq.Name, lq.Namespace, cqName)
+		}
+	}
+}
+
+func (r *MiscReconciler) clusterQueueExists(name string, cp *v1alpha.ClusterPolicy) []v1alpha.LocalQueueSpec {
+	kueue := cp.Spec.Kueue
+	if kueue == nil {
+		return nil
+	}
+
+	for _, eq := range kueue.EqualResources {
+		if name == eq.Name {
+			return eq.LocalQueues
+		}
+	}
+
+	return nil
+}
+
+func (r *MiscReconciler) pruneQueues(ctx context.Context, cp *v1alpha.ClusterPolicy) error {
+	matchLabels := map[string]string{
+		"app":   kueueAppLabel,
+		"owner": r.Opts.ReqName,
+	}
+
+	cqList := &kueuev1beta2.ClusterQueueList{}
+	if err := r.List(ctx, cqList, client.MatchingLabels(matchLabels)); client.IgnoreNotFound(err) != nil {
+		return fmt.Errorf("Could not fetch ClusterQueue list: %v", err)
+	}
+
+	lqList := &kueuev1beta2.LocalQueueList{}
+	if err := r.List(ctx, lqList, client.InNamespace(metav1.NamespaceAll), client.MatchingLabels(matchLabels)); err != nil {
+		if client.IgnoreNotFound(err) != nil {
+			return fmt.Errorf("Could not fetch LocalQueue list: %v", err)
+		} else {
+			klog.Warningf("No LocalQueues found")
+		}
+	}
+
+	for _, cq := range cqList.Items {
+		lq := r.clusterQueueExists(cq.Name, cp)
+		if lq == nil {
+			if err := r.Delete(ctx, &cq); client.IgnoreNotFound(err) != nil {
+				klog.Warningf("Error when attempting to delete ClusterQueue %s: %v", cq.Name, err)
+			} else {
+				klog.V(3).Infof("Deleted ClusterQueue %s", cq.Name)
+			}
+		} else {
+			klog.V(3).Infof("Keep ClusterQueue %s", cq.Name)
+		}
+		// delete lqs, lq = keep, lq = nil -> delete all
+		r.pruneLocalQueues(ctx, lqList, lq, cq.Name, cp)
+
+	}
+
+	return nil
+}
+
+func (r *MiscReconciler) modifyLocalQueues(ctx context.Context, clusterQueue *v1alpha.ClusterQueueSpec, cp *v1alpha.ClusterPolicy) error {
 	if len(clusterQueue.LocalQueues) == 0 {
 		klog.V(3).Infof("No LocalQueues defined for ClusterQueue '%s'", clusterQueue.Name)
 		return nil
@@ -613,6 +704,11 @@ func (r *MiscReconciler) modifyLocalQueues(ctx context.Context, clusterQueue *v1
 		newLocalQueue := r.createLocalQueue(clusterQueue.Name, &localQueue)
 
 		if !found {
+			if err := ctrl.SetControllerReference(cp, newLocalQueue, r.Scheme); err != nil {
+				klog.Error(err, "unable to set Kueue LocalQueue controller reference")
+				return err
+			}
+
 			err := r.Create(ctx, newLocalQueue)
 			if client.IgnoreAlreadyExists(err) != nil {
 				klog.Error(err, "unable to create Kueue LocalQueue '%s/%s'", localQueue.Namespace, localQueue.Name)
@@ -692,7 +788,9 @@ func (r *MiscReconciler) createClusterQueue(resources clusterResourceMap, cluste
 	return clusterQueue
 }
 
-func (r *MiscReconciler) modifyClusterQueue(ctx context.Context, resources clusterResourceMap, kueueSpec *v1alpha.KueueQueueSpec) error {
+func (r *MiscReconciler) modifyClusterQueue(ctx context.Context, resources clusterResourceMap, cp *v1alpha.ClusterPolicy) error {
+	kueueSpec := cp.Spec.Kueue
+
 	if len(kueueSpec.EqualResources) == 0 {
 		klog.V(3).Infof("At least one EqualResources queue should be configured")
 		return nil
@@ -703,6 +801,10 @@ func (r *MiscReconciler) modifyClusterQueue(ctx context.Context, resources clust
 		return err
 	}
 
+	if err := r.pruneQueues(ctx, cp); err != nil {
+		klog.Warningf("Pruning Queues returned error but continuing anyway: %v", err)
+	}
+
 	for n, clusterQueue := range kueueSpec.EqualResources {
 		currentUnstructured := &unstructured.Unstructured{}
 		currentUnstructured.SetKind(kueueClusterQueueKind)
@@ -720,6 +822,11 @@ func (r *MiscReconciler) modifyClusterQueue(ctx context.Context, resources clust
 		newClusterQueue := r.createClusterQueue(clusterResources[n], &clusterQueue)
 
 		if !found {
+			if err := ctrl.SetControllerReference(cp, newClusterQueue, r.Scheme); err != nil {
+				klog.Error(err, "unable to set Kueue ClusterQueue controller reference")
+				return err
+			}
+
 			err = r.Create(ctx, newClusterQueue)
 			if client.IgnoreAlreadyExists(err) != nil {
 				klog.Error(err, "unable to create Kueue ClusterQueue '%s'", clusterQueue.Name)
@@ -752,25 +859,25 @@ func (r *MiscReconciler) modifyClusterQueue(ctx context.Context, resources clust
 			}
 		}
 
-		if err := r.modifyLocalQueues(ctx, &clusterQueue); err != nil {
+		if err := r.modifyLocalQueues(ctx, &clusterQueue, cp); err != nil {
 			return err
 		}
 	}
 
 	return nil
 }
 
-func (r *MiscReconciler) createKueueQueues(ctx context.Context, resources clusterResourceMap, kueueSpec *v1alpha.KueueQueueSpec) error {
+func (r *MiscReconciler) createKueueQueues(ctx context.Context, resources clusterResourceMap, cp *v1alpha.ClusterPolicy) error {
 	if resources == nil {
 		klog.V(3).Infof("no resources defined when attempting to create Kueue queues")
 		return nil
 	}
 
-	if err := r.modifyResourceFlavor(ctx); err != nil {
+	if err := r.modifyResourceFlavor(ctx, cp); err != nil {
 		return err
 	}
 
-	if err := r.modifyClusterQueue(ctx, resources, kueueSpec); err != nil {
+	if err := r.modifyClusterQueue(ctx, resources, cp); err != nil {
 		return err
 	}
 
@@ -818,7 +925,11 @@ func (r *MiscReconciler) removeKueueObjects(ctx context.Context, crName string)
 
 	localQueues := &kueuev1beta2.LocalQueueList{}
 	if err := r.List(ctx, localQueues, client.InNamespace(metav1.NamespaceAll), client.MatchingLabels(matchLabels)); err != nil {
-		klog.Warningf("Error when deleting Kueue LocalQueues: %v", err)
+		if client.IgnoreNotFound(err) == nil {
+			klog.V(3).Infof("No Kueue LocalQueues to delete")
+		} else {
+			klog.Warningf("Error when deleting Kueue LocalQueues: %v", err)
+		}
 	} else {
 		for i := range localQueues.Items {
 
@@ -881,7 +992,7 @@ func (r *MiscReconciler) reconcileKueueObjects(ctx context.Context, cp *v1alpha.
 		return nil
 	}
 
-	return r.createKueueQueues(ctx, resources, cp.Spec.Kueue)
+	return r.createKueueQueues(ctx, resources, cp)
 }
 
 func (r *MiscReconciler) Reconcile(ctx context.Context, cp *v1alpha.ClusterPolicy) (ctrl.Result, error) {