Fix a bug on assigning AdmissionChecks for a Workload

pkg/cache/clusterqueue.go

@@ -418,8 +418,8 @@ func (c *ClusterQueue) updateLabelKeys(flavors map[kueue.ResourceFlavorReference
 // updateWithAdmissionChecks updates a ClusterQueue based on the passed AdmissionChecks set.
 func (c *ClusterQueue) updateWithAdmissionChecks(checks map[string]AdmissionCheck) {
 	hasMissing := false
-	singleInstanceControllers := sets.New[string]()
 	checksPerController := make(map[string]int, len(c.AdmissionChecks))
+	singleInstanceControllers := sets.New[string]()
 	for acName, _ := range c.AdmissionChecks {
 		if ac, found := checks[acName]; !found {
 			hasMissing = true

pkg/cache/snapshot.go

@@ -143,10 +143,8 @@ func (c *ClusterQueue) snapshot() *ClusterQueue {
 		Preemption:                    c.Preemption,
 		NamespaceSelector:             c.NamespaceSelector,
 		Status:                        c.Status,
+		AdmissionChecks:               utilmaps.DeepCopySets[string](c.AdmissionChecks),
 	}
-	admissionChecks := make(map[string]sets.Set[string], len(c.AdmissionChecks))
-	maps.Copy(admissionChecks, c.AdmissionChecks)
-	c.AdmissionChecks = admissionChecks
 
 	for fName, rUsage := range c.Usage {
 		cc.Usage[fName] = maps.Clone(rUsage)

pkg/controller/core/workload_controller.go

@@ -266,6 +266,7 @@ func (r *WorkloadReconciler) reconcileSyncAdmissionChecks(ctx context.Context, w
 	}
 
 	admissionChecks := workload.AdmissionChecksForWorkload(log, wl, utilac.NewAdmissionChecks(&cq))
+	log.V(2).Info("PATRYK created ac", "ac", admissionChecks)
 	newChecks, shouldUpdate := syncAdmissionCheckConditions(wl.Status.AdmissionChecks, admissionChecks)
 	if shouldUpdate {
 		log.V(3).Info("The workload needs admission checks updates", "clusterQueue", klog.KRef("", cqName), "admissionChecks", admissionChecks)

pkg/util/maps/maps.go

@@ -21,6 +21,8 @@ package maps
 import (
 	"fmt"
 	"maps"
+
+	"k8s.io/apimachinery/pkg/util/sets"
 )
 
 // Merge merges a and b while resolving the conflicts by calling commonKeyValue
@@ -108,3 +110,12 @@ func FilterKeys[K comparable, V any, M ~map[K]V](m M, k []K) M {
 	}
 	return ret
 }
+
+// DeepCopySets create a deep copy of map[string]Set which would otherwise be referenced
+func DeepCopySets[T comparable](src map[string]sets.Set[T]) map[string]sets.Set[T] {
+	copy := make(map[string]sets.Set[T], len(src))
+	for key, set := range src {
+		copy[key] = set.Clone()
+	}
+	return copy
+}

pkg/workload/workload.go

@@ -39,6 +39,7 @@ import (
 	"sigs.k8s.io/kueue/pkg/features"
 	"sigs.k8s.io/kueue/pkg/util/api"
 	"sigs.k8s.io/kueue/pkg/util/limitrange"
+	utilmaps "sigs.k8s.io/kueue/pkg/util/maps"
 )
 
 var (
@@ -550,6 +551,19 @@ func RemoveFinalizer(ctx context.Context, c client.Client, wl *kueue.Workload) e
 // AdmissionChecksForWorkload returns AdmissionChecks that should be assigned to a specific Workload based on
 // ClusterQueue configuration and ResourceFlavors
 func AdmissionChecksForWorkload(log logr.Logger, wl *kueue.Workload, admissionChecks map[string]sets.Set[string]) sets.Set[string] {
+	// If all admissionChecks should be run for all flavors we don't need to wait for Workload's Admission to be set.
+	// This is also the case if admissionChecks are specified with ClusterQueue.Spec.AdmissionChecks instead of
+	// ClusterQueue.Spec.AdmissionCheckStrategy
+	allFlavors := true
+	for _, flavors := range admissionChecks {
+		if len(flavors) != 0 {
+			allFlavors = false
+		}
+	}
+	if allFlavors {
+		return sets.New(utilmaps.Keys(admissionChecks)...)
+	}
+
 	// Kueue sets AdmissionChecks first based on ClusterQueue configuration and at this point Workload has no
 	// ResourceFlavors assigned, so we cannot match AdmissionChecks to ResourceFlavor.
 	// After Quota is reserved, another reconciliation happens and we can match AdmissionChecks to ResourceFlavors