diff --git a/pkg/cache/clusterqueue_snapshot.go b/pkg/cache/clusterqueue_snapshot.go index b081cee98b..4f1e131399 100644 --- a/pkg/cache/clusterqueue_snapshot.go +++ b/pkg/cache/clusterqueue_snapshot.go @@ -23,6 +23,7 @@ import ( "k8s.io/apimachinery/pkg/util/sets" kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1" + "sigs.k8s.io/kueue/pkg/hierarchy" "sigs.k8s.io/kueue/pkg/metrics" "sigs.k8s.io/kueue/pkg/resources" "sigs.k8s.io/kueue/pkg/workload" @@ -30,7 +31,6 @@ import ( type ClusterQueueSnapshot struct { Name string - Cohort *CohortSnapshot ResourceGroups []ResourceGroup Workloads map[string]*workload.Info WorkloadsNotReady sets.Set[string] @@ -48,6 +48,7 @@ type ClusterQueueSnapshot struct { AllocatableResourceGeneration int64 ResourceNode ResourceNode + hierarchy.ClusterQueue[*CohortSnapshot] } // RGByResource returns the ResourceGroup which contains capacity @@ -108,17 +109,9 @@ func (c *ClusterQueueSnapshot) PotentialAvailable(fr resources.FlavorResource) i return potentialAvailable(c, fr) } -func (c *ClusterQueueSnapshot) Parent() *CohortSnapshot { - return c.Cohort -} - // The methods below implement several interfaces. See // dominantResourceShareNode, resourceGroupNode, and netQuotaNode. -func (c *ClusterQueueSnapshot) HasParent() bool { - return c.Cohort != nil -} - func (c *ClusterQueueSnapshot) fairWeight() *resource.Quantity { return &c.FairWeight } @@ -135,6 +128,10 @@ func (c *ClusterQueueSnapshot) parentResources() ResourceNode { return c.Parent().ResourceNode } +func (c *ClusterQueueSnapshot) GetName() string { + return c.Name +} + // The methods below implement hierarchicalResourceNode interface. func (c *ClusterQueueSnapshot) getResourceNode() ResourceNode { diff --git a/pkg/cache/cohort_snapshot.go b/pkg/cache/cohort_snapshot.go index 132006c416..43cdfadbd2 100644 --- a/pkg/cache/cohort_snapshot.go +++ b/pkg/cache/cohort_snapshot.go @@ -16,27 +16,25 @@ limitations under the License. package cache -import ( - "k8s.io/apimachinery/pkg/util/sets" -) +import "sigs.k8s.io/kueue/pkg/hierarchy" type CohortSnapshot struct { - Name string - Members sets.Set[*ClusterQueueSnapshot] + Name string ResourceNode ResourceNode + hierarchy.Cohort[*ClusterQueueSnapshot, *CohortSnapshot] } -// The methods below implement hierarchicalResourceNode interface. - -func (c *CohortSnapshot) HasParent() bool { - return false +func (c *CohortSnapshot) GetName() string { + return c.Name } +// The methods below implement hierarchicalResourceNode interface. + func (c *CohortSnapshot) getResourceNode() ResourceNode { return c.ResourceNode } func (c *CohortSnapshot) parentHRN() hierarchicalResourceNode { - return nil + return c.Parent() } diff --git a/pkg/cache/snapshot.go b/pkg/cache/snapshot.go index c620d5064d..43cf225546 100644 --- a/pkg/cache/snapshot.go +++ b/pkg/cache/snapshot.go @@ -24,12 +24,13 @@ import ( "k8s.io/klog/v2" kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1" + "sigs.k8s.io/kueue/pkg/hierarchy" utilmaps "sigs.k8s.io/kueue/pkg/util/maps" "sigs.k8s.io/kueue/pkg/workload" ) type Snapshot struct { - ClusterQueues map[string]*ClusterQueueSnapshot + hierarchy.Manager[*ClusterQueueSnapshot, *CohortSnapshot] ResourceFlavors map[kueue.ResourceFlavorReference]*kueue.ResourceFlavor InactiveClusterQueueSets sets.Set[string] } @@ -51,12 +52,10 @@ func (s *Snapshot) AddWorkload(wl *workload.Info) { } func (s *Snapshot) Log(log logr.Logger) { - cohorts := make(map[string]*CohortSnapshot) for name, cq := range s.ClusterQueues { cohortName := "" - if cq.Cohort != nil { - cohortName = cq.Cohort.Name - cohorts[cohortName] = cq.Cohort + if cq.HasParent() { + cohortName = cq.Parent().Name } log.Info("Found ClusterQueue", @@ -67,7 +66,7 @@ func (s *Snapshot) Log(log logr.Logger) { "workloads", utilmaps.Keys(cq.Workloads), ) } - for name, cohort := range cohorts { + for name, cohort := range s.Cohorts { log.Info("Found cohort", "cohort", name, "resources", cohort.ResourceNode.SubtreeQuota, @@ -81,30 +80,33 @@ func (c *Cache) Snapshot() Snapshot { defer c.RUnlock() snap := Snapshot{ - ClusterQueues: make(map[string]*ClusterQueueSnapshot, len(c.hm.ClusterQueues)), + Manager: hierarchy.NewManager(newCohortSnapshot), ResourceFlavors: make(map[kueue.ResourceFlavorReference]*kueue.ResourceFlavor, len(c.resourceFlavors)), InactiveClusterQueueSets: sets.New[string](), } + for _, cohort := range c.hm.Cohorts { + snap.AddCohort(snapshotCohort(cohort)) + } for _, cq := range c.hm.ClusterQueues { if !cq.Active() { snap.InactiveClusterQueueSets.Insert(cq.Name) continue } - snap.ClusterQueues[cq.Name] = cq.snapshot() + snap.AddClusterQueue(snapshotClusterQueue(cq)) + if cq.HasParent() { + snap.UpdateClusterQueueEdge(cq.Name, cq.Parent().Name) + } } for name, rf := range c.resourceFlavors { // Shallow copy is enough snap.ResourceFlavors[name] = rf } - for _, cohort := range c.hm.Cohorts { - cohort.snapshotInto(snap.ClusterQueues) - } return snap } -// snapshot creates a copy of ClusterQueue that includes references to immutable -// objects and deep copies of changing ones. A reference to the cohort is not included. -func (c *clusterQueue) snapshot() *ClusterQueueSnapshot { +// snapshotClusterQueue creates a copy of ClusterQueue that includes +// references to immutable objects and deep copies of changing ones. +func snapshotClusterQueue(c *clusterQueue) *ClusterQueueSnapshot { cc := &ClusterQueueSnapshot{ Name: c.Name, ResourceGroups: make([]ResourceGroup, len(c.ResourceGroups)), @@ -124,17 +126,15 @@ func (c *clusterQueue) snapshot() *ClusterQueueSnapshot { return cc } -func (c *cohort) snapshotInto(cqs map[string]*ClusterQueueSnapshot) { - cohortSnap := &CohortSnapshot{ - Name: c.Name, - Members: make(sets.Set[*ClusterQueueSnapshot], len(c.ChildCQs())), - ResourceNode: c.resourceNode.Clone(), - } - for _, cq := range c.ChildCQs() { - if cq.Active() { - cqSnap := cqs[cq.Name] - cqSnap.Cohort = cohortSnap - cohortSnap.Members.Insert(cqSnap) - } +func snapshotCohort(c *cohort) *CohortSnapshot { + snapshot := newCohortSnapshot(c.Name) + snapshot.ResourceNode = c.resourceNode.Clone() + return snapshot +} + +func newCohortSnapshot(name string) *CohortSnapshot { + return &CohortSnapshot{ + Name: name, + Cohort: hierarchy.NewCohort[*ClusterQueueSnapshot, *CohortSnapshot](), } } diff --git a/pkg/cache/snapshot_test.go b/pkg/cache/snapshot_test.go index b4f7f1a3c7..9c2d5f2cfa 100644 --- a/pkg/cache/snapshot_test.go +++ b/pkg/cache/snapshot_test.go @@ -32,6 +32,7 @@ import ( kueuealpha "sigs.k8s.io/kueue/apis/kueue/v1alpha1" kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1" "sigs.k8s.io/kueue/pkg/features" + "sigs.k8s.io/kueue/pkg/hierarchy" "sigs.k8s.io/kueue/pkg/resources" utiltesting "sigs.k8s.io/kueue/pkg/util/testing" "sigs.k8s.io/kueue/pkg/workload" @@ -39,7 +40,9 @@ import ( var snapCmpOpts = []cmp.Option{ cmpopts.EquateEmpty(), - cmpopts.IgnoreFields(CohortSnapshot{}, "Members"), // avoid recursion. + cmpopts.IgnoreUnexported(hierarchy.Cohort[*ClusterQueueSnapshot, *CohortSnapshot]{}), + cmpopts.IgnoreUnexported(hierarchy.ClusterQueue[*CohortSnapshot]{}), + cmpopts.IgnoreUnexported(hierarchy.Manager[*ClusterQueueSnapshot, *CohortSnapshot]{}), cmpopts.IgnoreFields(metav1.Condition{}, "LastTransitionTime"), } @@ -65,34 +68,36 @@ func TestSnapshot(t *testing.T) { ReserveQuota(&kueue.Admission{ClusterQueue: "b"}).Obj(), }, wantSnapshot: Snapshot{ - ClusterQueues: map[string]*ClusterQueueSnapshot{ - "a": { - Name: "a", - NamespaceSelector: labels.Everything(), - Status: active, - FlavorFungibility: defaultFlavorFungibility, - AllocatableResourceGeneration: 1, - Workloads: map[string]*workload.Info{ - "/alpha": workload.NewInfo( - utiltesting.MakeWorkload("alpha", ""). - ReserveQuota(&kueue.Admission{ClusterQueue: "a"}).Obj()), + Manager: hierarchy.Manager[*ClusterQueueSnapshot, *CohortSnapshot]{ + ClusterQueues: map[string]*ClusterQueueSnapshot{ + "a": { + Name: "a", + NamespaceSelector: labels.Everything(), + Status: active, + FlavorFungibility: defaultFlavorFungibility, + AllocatableResourceGeneration: 1, + Workloads: map[string]*workload.Info{ + "/alpha": workload.NewInfo( + utiltesting.MakeWorkload("alpha", ""). + ReserveQuota(&kueue.Admission{ClusterQueue: "a"}).Obj()), + }, + Preemption: defaultPreemption, + FairWeight: oneQuantity, }, - Preemption: defaultPreemption, - FairWeight: oneQuantity, - }, - "b": { - Name: "b", - NamespaceSelector: labels.Everything(), - Status: active, - FlavorFungibility: defaultFlavorFungibility, - AllocatableResourceGeneration: 1, - Workloads: map[string]*workload.Info{ - "/beta": workload.NewInfo( - utiltesting.MakeWorkload("beta", ""). - ReserveQuota(&kueue.Admission{ClusterQueue: "b"}).Obj()), + "b": { + Name: "b", + NamespaceSelector: labels.Everything(), + Status: active, + FlavorFungibility: defaultFlavorFungibility, + AllocatableResourceGeneration: 1, + Workloads: map[string]*workload.Info{ + "/beta": workload.NewInfo( + utiltesting.MakeWorkload("beta", ""). + ReserveQuota(&kueue.Admission{ClusterQueue: "b"}).Obj()), + }, + Preemption: defaultPreemption, + FairWeight: oneQuantity, }, - Preemption: defaultPreemption, - FairWeight: oneQuantity, }, }, }, @@ -121,7 +126,9 @@ func TestSnapshot(t *testing.T) { utiltesting.MakeResourceFlavor("default").Obj(), }, wantSnapshot: Snapshot{ - ClusterQueues: map[string]*ClusterQueueSnapshot{}, + Manager: hierarchy.Manager[*ClusterQueueSnapshot, *CohortSnapshot]{ + ClusterQueues: map[string]*ClusterQueueSnapshot{}, + }, ResourceFlavors: map[kueue.ResourceFlavorReference]*kueue.ResourceFlavor{ "demand": utiltesting.MakeResourceFlavor("demand"). NodeLabel("a", "b"). @@ -218,130 +225,133 @@ func TestSnapshot(t *testing.T) { }, } return Snapshot{ - ClusterQueues: map[string]*ClusterQueueSnapshot{ - "a": { - Name: "a", - Cohort: cohort, - AllocatableResourceGeneration: 2, - ResourceGroups: []ResourceGroup{ - { - CoveredResources: sets.New(corev1.ResourceCPU), - Flavors: []kueue.ResourceFlavorReference{"demand", "spot"}, - LabelKeys: sets.New("instance"), - }, - }, - ResourceNode: ResourceNode{ - Quotas: map[resources.FlavorResource]ResourceQuota{ - {Flavor: "demand", Resource: corev1.ResourceCPU}: {Nominal: 100_000}, - {Flavor: "spot", Resource: corev1.ResourceCPU}: {Nominal: 200_000}, - }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "demand", Resource: "cpu"}: 100_000, - {Flavor: "spot", Resource: "cpu"}: 200_000, - }, - Usage: resources.FlavorResourceQuantities{ - {Flavor: "demand", Resource: corev1.ResourceCPU}: 10_000, - }, - }, - FlavorFungibility: defaultFlavorFungibility, - Workloads: map[string]*workload.Info{ - "/alpha": workload.NewInfo(utiltesting.MakeWorkload("alpha", ""). - PodSets(*utiltesting.MakePodSet("main", 5). - Request(corev1.ResourceCPU, "2").Obj()). - ReserveQuota(utiltesting.MakeAdmission("a", "main"). - Assignment(corev1.ResourceCPU, "demand", "10000m"). - AssignmentPodCount(5). - Obj()). - Obj()), - }, - Preemption: defaultPreemption, - FairWeight: oneQuantity, - NamespaceSelector: labels.Everything(), - Status: active, + Manager: hierarchy.Manager[*ClusterQueueSnapshot, *CohortSnapshot]{ + Cohorts: map[string]*CohortSnapshot{ + "borrowing": cohort, }, - "b": { - Name: "b", - Cohort: cohort, - AllocatableResourceGeneration: 1, - ResourceGroups: []ResourceGroup{ - { - CoveredResources: sets.New(corev1.ResourceCPU), - Flavors: []kueue.ResourceFlavorReference{"spot"}, - LabelKeys: sets.New("instance"), - }, - { - CoveredResources: sets.New[corev1.ResourceName]("example.com/gpu"), - Flavors: []kueue.ResourceFlavorReference{"default"}, - }, + ClusterQueues: map[string]*ClusterQueueSnapshot{ + "a": { + Name: "a", + AllocatableResourceGeneration: 2, + ResourceGroups: []ResourceGroup{ + { + CoveredResources: sets.New(corev1.ResourceCPU), + Flavors: []kueue.ResourceFlavorReference{"demand", "spot"}, + LabelKeys: sets.New("instance"), + }, + }, + ResourceNode: ResourceNode{ + Quotas: map[resources.FlavorResource]ResourceQuota{ + {Flavor: "demand", Resource: corev1.ResourceCPU}: {Nominal: 100_000}, + {Flavor: "spot", Resource: corev1.ResourceCPU}: {Nominal: 200_000}, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "demand", Resource: "cpu"}: 100_000, + {Flavor: "spot", Resource: "cpu"}: 200_000, + }, + Usage: resources.FlavorResourceQuantities{ + {Flavor: "demand", Resource: corev1.ResourceCPU}: 10_000, + }, + }, + FlavorFungibility: defaultFlavorFungibility, + Workloads: map[string]*workload.Info{ + "/alpha": workload.NewInfo(utiltesting.MakeWorkload("alpha", ""). + PodSets(*utiltesting.MakePodSet("main", 5). + Request(corev1.ResourceCPU, "2").Obj()). + ReserveQuota(utiltesting.MakeAdmission("a", "main"). + Assignment(corev1.ResourceCPU, "demand", "10000m"). + AssignmentPodCount(5). + Obj()). + Obj()), + }, + Preemption: defaultPreemption, + FairWeight: oneQuantity, + NamespaceSelector: labels.Everything(), + Status: active, }, - ResourceNode: ResourceNode{ - Quotas: map[resources.FlavorResource]ResourceQuota{ - {Flavor: "spot", Resource: corev1.ResourceCPU}: {Nominal: 100_000}, - {Flavor: "default", Resource: "example.com/gpu"}: {Nominal: 50}, - }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "spot", Resource: "cpu"}: 100_000, - {Flavor: "default", Resource: "example.com/gpu"}: 50, - }, - Usage: resources.FlavorResourceQuantities{ - {Flavor: "spot", Resource: corev1.ResourceCPU}: 10_000, - {Flavor: "default", Resource: "example.com/gpu"}: 15, - }, + "b": { + Name: "b", + AllocatableResourceGeneration: 1, + ResourceGroups: []ResourceGroup{ + { + CoveredResources: sets.New(corev1.ResourceCPU), + Flavors: []kueue.ResourceFlavorReference{"spot"}, + LabelKeys: sets.New("instance"), + }, + { + CoveredResources: sets.New[corev1.ResourceName]("example.com/gpu"), + Flavors: []kueue.ResourceFlavorReference{"default"}, + }, + }, + ResourceNode: ResourceNode{ + Quotas: map[resources.FlavorResource]ResourceQuota{ + {Flavor: "spot", Resource: corev1.ResourceCPU}: {Nominal: 100_000}, + {Flavor: "default", Resource: "example.com/gpu"}: {Nominal: 50}, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "spot", Resource: "cpu"}: 100_000, + {Flavor: "default", Resource: "example.com/gpu"}: 50, + }, + Usage: resources.FlavorResourceQuantities{ + {Flavor: "spot", Resource: corev1.ResourceCPU}: 10_000, + {Flavor: "default", Resource: "example.com/gpu"}: 15, + }, + }, + FlavorFungibility: defaultFlavorFungibility, + Workloads: map[string]*workload.Info{ + "/beta": workload.NewInfo(utiltesting.MakeWorkload("beta", ""). + PodSets(*utiltesting.MakePodSet("main", 5). + Request(corev1.ResourceCPU, "1"). + Request("example.com/gpu", "2"). + Obj()). + ReserveQuota(utiltesting.MakeAdmission("b", "main"). + Assignment(corev1.ResourceCPU, "spot", "5000m"). + Assignment("example.com/gpu", "default", "10"). + AssignmentPodCount(5). + Obj()). + Obj()), + "/gamma": workload.NewInfo(utiltesting.MakeWorkload("gamma", ""). + PodSets(*utiltesting.MakePodSet("main", 5). + Request(corev1.ResourceCPU, "1"). + Request("example.com/gpu", "1"). + Obj(), + ). + ReserveQuota(utiltesting.MakeAdmission("b", "main"). + Assignment(corev1.ResourceCPU, "spot", "5000m"). + Assignment("example.com/gpu", "default", "5"). + AssignmentPodCount(5). + Obj()). + Obj()), + }, + Preemption: defaultPreemption, + FairWeight: oneQuantity, + NamespaceSelector: labels.Everything(), + Status: active, }, - FlavorFungibility: defaultFlavorFungibility, - Workloads: map[string]*workload.Info{ - "/beta": workload.NewInfo(utiltesting.MakeWorkload("beta", ""). - PodSets(*utiltesting.MakePodSet("main", 5). - Request(corev1.ResourceCPU, "1"). - Request("example.com/gpu", "2"). - Obj()). - ReserveQuota(utiltesting.MakeAdmission("b", "main"). - Assignment(corev1.ResourceCPU, "spot", "5000m"). - Assignment("example.com/gpu", "default", "10"). - AssignmentPodCount(5). - Obj()). - Obj()), - "/gamma": workload.NewInfo(utiltesting.MakeWorkload("gamma", ""). - PodSets(*utiltesting.MakePodSet("main", 5). - Request(corev1.ResourceCPU, "1"). - Request("example.com/gpu", "1"). - Obj(), - ). - ReserveQuota(utiltesting.MakeAdmission("b", "main"). - Assignment(corev1.ResourceCPU, "spot", "5000m"). - Assignment("example.com/gpu", "default", "5"). - AssignmentPodCount(5). - Obj()). - Obj()), + "c": { + Name: "c", + AllocatableResourceGeneration: 1, + ResourceGroups: []ResourceGroup{ + { + CoveredResources: sets.New(corev1.ResourceCPU), + Flavors: []kueue.ResourceFlavorReference{"default"}, + }, + }, + ResourceNode: ResourceNode{ + Quotas: map[resources.FlavorResource]ResourceQuota{ + {Flavor: "default", Resource: corev1.ResourceCPU}: {Nominal: 100_000}, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: "cpu"}: 100_000, + }, + Usage: resources.FlavorResourceQuantities{}, + }, + FlavorFungibility: defaultFlavorFungibility, + Preemption: defaultPreemption, + FairWeight: oneQuantity, + NamespaceSelector: labels.Everything(), + Status: active, }, - Preemption: defaultPreemption, - FairWeight: oneQuantity, - NamespaceSelector: labels.Everything(), - Status: active, - }, - "c": { - Name: "c", - AllocatableResourceGeneration: 1, - ResourceGroups: []ResourceGroup{ - { - CoveredResources: sets.New(corev1.ResourceCPU), - Flavors: []kueue.ResourceFlavorReference{"default"}, - }, - }, - ResourceNode: ResourceNode{ - Quotas: map[resources.FlavorResource]ResourceQuota{ - {Flavor: "default", Resource: corev1.ResourceCPU}: {Nominal: 100_000}, - }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: "cpu"}: 100_000, - }, - Usage: resources.FlavorResourceQuantities{}, - }, - FlavorFungibility: defaultFlavorFungibility, - Preemption: defaultPreemption, - FairWeight: oneQuantity, - NamespaceSelector: labels.Everything(), - Status: active, }, }, ResourceFlavors: map[kueue.ResourceFlavorReference]*kueue.ResourceFlavor{ @@ -361,19 +371,21 @@ func TestSnapshot(t *testing.T) { }).Obj(), }, wantSnapshot: Snapshot{ - ClusterQueues: map[string]*ClusterQueueSnapshot{ - "with-preemption": { - Name: "with-preemption", - NamespaceSelector: labels.Everything(), - AllocatableResourceGeneration: 1, - Status: active, - Workloads: map[string]*workload.Info{}, - FlavorFungibility: defaultFlavorFungibility, - Preemption: kueue.ClusterQueuePreemption{ - ReclaimWithinCohort: kueue.PreemptionPolicyAny, - WithinClusterQueue: kueue.PreemptionPolicyLowerPriority, + Manager: hierarchy.Manager[*ClusterQueueSnapshot, *CohortSnapshot]{ + ClusterQueues: map[string]*ClusterQueueSnapshot{ + "with-preemption": { + Name: "with-preemption", + NamespaceSelector: labels.Everything(), + AllocatableResourceGeneration: 1, + Status: active, + Workloads: map[string]*workload.Info{}, + FlavorFungibility: defaultFlavorFungibility, + Preemption: kueue.ClusterQueuePreemption{ + ReclaimWithinCohort: kueue.PreemptionPolicyAny, + WithinClusterQueue: kueue.PreemptionPolicyLowerPriority, + }, + FairWeight: oneQuantity, }, - FairWeight: oneQuantity, }, }, }, @@ -383,16 +395,18 @@ func TestSnapshot(t *testing.T) { utiltesting.MakeClusterQueue("with-preemption").FairWeight(resource.MustParse("3")).Obj(), }, wantSnapshot: Snapshot{ - ClusterQueues: map[string]*ClusterQueueSnapshot{ - "with-preemption": { - Name: "with-preemption", - NamespaceSelector: labels.Everything(), - AllocatableResourceGeneration: 1, - Status: active, - Workloads: map[string]*workload.Info{}, - FlavorFungibility: defaultFlavorFungibility, - Preemption: defaultPreemption, - FairWeight: resource.MustParse("3"), + Manager: hierarchy.Manager[*ClusterQueueSnapshot, *CohortSnapshot]{ + ClusterQueues: map[string]*ClusterQueueSnapshot{ + "with-preemption": { + Name: "with-preemption", + NamespaceSelector: labels.Everything(), + AllocatableResourceGeneration: 1, + Status: active, + Workloads: map[string]*workload.Info{}, + FlavorFungibility: defaultFlavorFungibility, + Preemption: defaultPreemption, + FairWeight: resource.MustParse("3"), + }, }, }, }, @@ -456,88 +470,91 @@ func TestSnapshot(t *testing.T) { }, } return Snapshot{ - ClusterQueues: map[string]*ClusterQueueSnapshot{ - "a": { - Name: "a", - Cohort: cohort, - AllocatableResourceGeneration: 2, - ResourceGroups: []ResourceGroup{ - { - CoveredResources: sets.New(corev1.ResourceCPU), - Flavors: []kueue.ResourceFlavorReference{"arm", "x86"}, - LabelKeys: sets.New("arch"), - }, - }, - ResourceNode: ResourceNode{ - Quotas: map[resources.FlavorResource]ResourceQuota{ - {Flavor: "arm", Resource: corev1.ResourceCPU}: {Nominal: 10_000, BorrowingLimit: nil, LendingLimit: ptr.To[int64](5_000)}, - {Flavor: "x86", Resource: corev1.ResourceCPU}: {Nominal: 20_000, BorrowingLimit: nil, LendingLimit: ptr.To[int64](10_000)}, - }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "arm", Resource: corev1.ResourceCPU}: 10_000, - {Flavor: "x86", Resource: corev1.ResourceCPU}: 20_000, - }, - Usage: resources.FlavorResourceQuantities{ - {Flavor: "arm", Resource: corev1.ResourceCPU}: 15_000, - {Flavor: "x86", Resource: corev1.ResourceCPU}: 10_000, - }, - }, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - Workloads: map[string]*workload.Info{ - "/alpha": workload.NewInfo(utiltesting.MakeWorkload("alpha", ""). - PodSets(*utiltesting.MakePodSet("main", 5). - Request(corev1.ResourceCPU, "2").Obj()). - ReserveQuota(utiltesting.MakeAdmission("a", "main"). - Assignment(corev1.ResourceCPU, "arm", "10000m"). - AssignmentPodCount(5).Obj()). - Obj()), - "/beta": workload.NewInfo(utiltesting.MakeWorkload("beta", ""). - PodSets(*utiltesting.MakePodSet("main", 5). - Request(corev1.ResourceCPU, "1").Obj()). - ReserveQuota(utiltesting.MakeAdmission("a", "main"). - Assignment(corev1.ResourceCPU, "arm", "5000m"). - AssignmentPodCount(5).Obj()). - Obj()), - "/gamma": workload.NewInfo(utiltesting.MakeWorkload("gamma", ""). - PodSets(*utiltesting.MakePodSet("main", 5). - Request(corev1.ResourceCPU, "2").Obj()). - ReserveQuota(utiltesting.MakeAdmission("a", "main"). - Assignment(corev1.ResourceCPU, "x86", "10000m"). - AssignmentPodCount(5).Obj()). - Obj()), - }, - Preemption: defaultPreemption, - NamespaceSelector: labels.Everything(), - Status: active, + Manager: hierarchy.Manager[*ClusterQueueSnapshot, *CohortSnapshot]{ + Cohorts: map[string]*CohortSnapshot{ + "lending": cohort, }, - "b": { - Name: "b", - Cohort: cohort, - AllocatableResourceGeneration: 1, - ResourceGroups: []ResourceGroup{ - { - CoveredResources: sets.New(corev1.ResourceCPU), - Flavors: []kueue.ResourceFlavorReference{"arm", "x86"}, - LabelKeys: sets.New("arch"), - }, + ClusterQueues: map[string]*ClusterQueueSnapshot{ + "a": { + Name: "a", + AllocatableResourceGeneration: 2, + ResourceGroups: []ResourceGroup{ + { + CoveredResources: sets.New(corev1.ResourceCPU), + Flavors: []kueue.ResourceFlavorReference{"arm", "x86"}, + LabelKeys: sets.New("arch"), + }, + }, + ResourceNode: ResourceNode{ + Quotas: map[resources.FlavorResource]ResourceQuota{ + {Flavor: "arm", Resource: corev1.ResourceCPU}: {Nominal: 10_000, BorrowingLimit: nil, LendingLimit: ptr.To[int64](5_000)}, + {Flavor: "x86", Resource: corev1.ResourceCPU}: {Nominal: 20_000, BorrowingLimit: nil, LendingLimit: ptr.To[int64](10_000)}, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "arm", Resource: corev1.ResourceCPU}: 10_000, + {Flavor: "x86", Resource: corev1.ResourceCPU}: 20_000, + }, + Usage: resources.FlavorResourceQuantities{ + {Flavor: "arm", Resource: corev1.ResourceCPU}: 15_000, + {Flavor: "x86", Resource: corev1.ResourceCPU}: 10_000, + }, + }, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + Workloads: map[string]*workload.Info{ + "/alpha": workload.NewInfo(utiltesting.MakeWorkload("alpha", ""). + PodSets(*utiltesting.MakePodSet("main", 5). + Request(corev1.ResourceCPU, "2").Obj()). + ReserveQuota(utiltesting.MakeAdmission("a", "main"). + Assignment(corev1.ResourceCPU, "arm", "10000m"). + AssignmentPodCount(5).Obj()). + Obj()), + "/beta": workload.NewInfo(utiltesting.MakeWorkload("beta", ""). + PodSets(*utiltesting.MakePodSet("main", 5). + Request(corev1.ResourceCPU, "1").Obj()). + ReserveQuota(utiltesting.MakeAdmission("a", "main"). + Assignment(corev1.ResourceCPU, "arm", "5000m"). + AssignmentPodCount(5).Obj()). + Obj()), + "/gamma": workload.NewInfo(utiltesting.MakeWorkload("gamma", ""). + PodSets(*utiltesting.MakePodSet("main", 5). + Request(corev1.ResourceCPU, "2").Obj()). + ReserveQuota(utiltesting.MakeAdmission("a", "main"). + Assignment(corev1.ResourceCPU, "x86", "10000m"). + AssignmentPodCount(5).Obj()). + Obj()), + }, + Preemption: defaultPreemption, + NamespaceSelector: labels.Everything(), + Status: active, }, - ResourceNode: ResourceNode{ - Quotas: map[resources.FlavorResource]ResourceQuota{ - {Flavor: "arm", Resource: corev1.ResourceCPU}: {Nominal: 10_000, BorrowingLimit: nil, LendingLimit: ptr.To[int64](5_000)}, - {Flavor: "x86", Resource: corev1.ResourceCPU}: {Nominal: 20_000, BorrowingLimit: nil, LendingLimit: ptr.To[int64](10_000)}, - }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "arm", Resource: corev1.ResourceCPU}: 10_000, - {Flavor: "x86", Resource: corev1.ResourceCPU}: 20_000, - }, - Usage: resources.FlavorResourceQuantities{}, + "b": { + Name: "b", + AllocatableResourceGeneration: 1, + ResourceGroups: []ResourceGroup{ + { + CoveredResources: sets.New(corev1.ResourceCPU), + Flavors: []kueue.ResourceFlavorReference{"arm", "x86"}, + LabelKeys: sets.New("arch"), + }, + }, + ResourceNode: ResourceNode{ + Quotas: map[resources.FlavorResource]ResourceQuota{ + {Flavor: "arm", Resource: corev1.ResourceCPU}: {Nominal: 10_000, BorrowingLimit: nil, LendingLimit: ptr.To[int64](5_000)}, + {Flavor: "x86", Resource: corev1.ResourceCPU}: {Nominal: 20_000, BorrowingLimit: nil, LendingLimit: ptr.To[int64](10_000)}, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "arm", Resource: corev1.ResourceCPU}: 10_000, + {Flavor: "x86", Resource: corev1.ResourceCPU}: 20_000, + }, + Usage: resources.FlavorResourceQuantities{}, + }, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + Preemption: defaultPreemption, + NamespaceSelector: labels.Everything(), + Status: active, }, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - Preemption: defaultPreemption, - NamespaceSelector: labels.Everything(), - Status: active, }, }, ResourceFlavors: map[kueue.ResourceFlavorReference]*kueue.ResourceFlavor{ @@ -608,87 +625,90 @@ func TestSnapshot(t *testing.T) { }, } return Snapshot{ - ClusterQueues: map[string]*ClusterQueueSnapshot{ - "a": { - Name: "a", - Cohort: cohort, - AllocatableResourceGeneration: 2, - ResourceGroups: []ResourceGroup{ - { - CoveredResources: sets.New(corev1.ResourceCPU), - Flavors: []kueue.ResourceFlavorReference{"arm", "x86"}, - LabelKeys: sets.New("arch"), - }, - }, - ResourceNode: ResourceNode{ - Quotas: map[resources.FlavorResource]ResourceQuota{ - {Flavor: "arm", Resource: corev1.ResourceCPU}: {Nominal: 10_000, BorrowingLimit: nil, LendingLimit: nil}, - {Flavor: "x86", Resource: corev1.ResourceCPU}: {Nominal: 20_000, BorrowingLimit: nil, LendingLimit: nil}, - }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "arm", Resource: corev1.ResourceCPU}: 10_000, - {Flavor: "x86", Resource: corev1.ResourceCPU}: 20_000, - }, - Usage: resources.FlavorResourceQuantities{ - {Flavor: "arm", Resource: corev1.ResourceCPU}: 15_000, - {Flavor: "x86", Resource: corev1.ResourceCPU}: 10_000, - }, - }, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - Workloads: map[string]*workload.Info{ - "/alpha": workload.NewInfo(utiltesting.MakeWorkload("alpha", ""). - PodSets(*utiltesting.MakePodSet("main", 5). - Request(corev1.ResourceCPU, "2").Obj()). - ReserveQuota(utiltesting.MakeAdmission("a", "main"). - Assignment(corev1.ResourceCPU, "arm", "10000m"). - AssignmentPodCount(5).Obj()). - Obj()), - "/beta": workload.NewInfo(utiltesting.MakeWorkload("beta", ""). - PodSets(*utiltesting.MakePodSet("main", 5). - Request(corev1.ResourceCPU, "1").Obj()). - ReserveQuota(utiltesting.MakeAdmission("a", "main"). - Assignment(corev1.ResourceCPU, "arm", "5000m"). - AssignmentPodCount(5).Obj()). - Obj()), - "/gamma": workload.NewInfo(utiltesting.MakeWorkload("gamma", ""). - PodSets(*utiltesting.MakePodSet("main", 5). - Request(corev1.ResourceCPU, "2").Obj()). - ReserveQuota(utiltesting.MakeAdmission("a", "main"). - Assignment(corev1.ResourceCPU, "x86", "10000m"). - AssignmentPodCount(5).Obj()). - Obj()), - }, - Preemption: defaultPreemption, - NamespaceSelector: labels.Everything(), - Status: active, + Manager: hierarchy.Manager[*ClusterQueueSnapshot, *CohortSnapshot]{ + Cohorts: map[string]*CohortSnapshot{ + "lending": cohort, }, - "b": { - Name: "b", - Cohort: cohort, - AllocatableResourceGeneration: 1, - ResourceGroups: []ResourceGroup{ - { - CoveredResources: sets.New(corev1.ResourceCPU), - Flavors: []kueue.ResourceFlavorReference{"arm", "x86"}, - LabelKeys: sets.New("arch"), - }, + ClusterQueues: map[string]*ClusterQueueSnapshot{ + "a": { + Name: "a", + AllocatableResourceGeneration: 2, + ResourceGroups: []ResourceGroup{ + { + CoveredResources: sets.New(corev1.ResourceCPU), + Flavors: []kueue.ResourceFlavorReference{"arm", "x86"}, + LabelKeys: sets.New("arch"), + }, + }, + ResourceNode: ResourceNode{ + Quotas: map[resources.FlavorResource]ResourceQuota{ + {Flavor: "arm", Resource: corev1.ResourceCPU}: {Nominal: 10_000, BorrowingLimit: nil, LendingLimit: nil}, + {Flavor: "x86", Resource: corev1.ResourceCPU}: {Nominal: 20_000, BorrowingLimit: nil, LendingLimit: nil}, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "arm", Resource: corev1.ResourceCPU}: 10_000, + {Flavor: "x86", Resource: corev1.ResourceCPU}: 20_000, + }, + Usage: resources.FlavorResourceQuantities{ + {Flavor: "arm", Resource: corev1.ResourceCPU}: 15_000, + {Flavor: "x86", Resource: corev1.ResourceCPU}: 10_000, + }, + }, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + Workloads: map[string]*workload.Info{ + "/alpha": workload.NewInfo(utiltesting.MakeWorkload("alpha", ""). + PodSets(*utiltesting.MakePodSet("main", 5). + Request(corev1.ResourceCPU, "2").Obj()). + ReserveQuota(utiltesting.MakeAdmission("a", "main"). + Assignment(corev1.ResourceCPU, "arm", "10000m"). + AssignmentPodCount(5).Obj()). + Obj()), + "/beta": workload.NewInfo(utiltesting.MakeWorkload("beta", ""). + PodSets(*utiltesting.MakePodSet("main", 5). + Request(corev1.ResourceCPU, "1").Obj()). + ReserveQuota(utiltesting.MakeAdmission("a", "main"). + Assignment(corev1.ResourceCPU, "arm", "5000m"). + AssignmentPodCount(5).Obj()). + Obj()), + "/gamma": workload.NewInfo(utiltesting.MakeWorkload("gamma", ""). + PodSets(*utiltesting.MakePodSet("main", 5). + Request(corev1.ResourceCPU, "2").Obj()). + ReserveQuota(utiltesting.MakeAdmission("a", "main"). + Assignment(corev1.ResourceCPU, "x86", "10000m"). + AssignmentPodCount(5).Obj()). + Obj()), + }, + Preemption: defaultPreemption, + NamespaceSelector: labels.Everything(), + Status: active, }, - ResourceNode: ResourceNode{ - Quotas: map[resources.FlavorResource]ResourceQuota{ - {Flavor: "arm", Resource: corev1.ResourceCPU}: {Nominal: 10_000, BorrowingLimit: nil, LendingLimit: nil}, - {Flavor: "x86", Resource: corev1.ResourceCPU}: {Nominal: 20_000, BorrowingLimit: nil, LendingLimit: nil}, - }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "arm", Resource: corev1.ResourceCPU}: 10_000, - {Flavor: "x86", Resource: corev1.ResourceCPU}: 20_000, - }, + "b": { + Name: "b", + AllocatableResourceGeneration: 1, + ResourceGroups: []ResourceGroup{ + { + CoveredResources: sets.New(corev1.ResourceCPU), + Flavors: []kueue.ResourceFlavorReference{"arm", "x86"}, + LabelKeys: sets.New("arch"), + }, + }, + ResourceNode: ResourceNode{ + Quotas: map[resources.FlavorResource]ResourceQuota{ + {Flavor: "arm", Resource: corev1.ResourceCPU}: {Nominal: 10_000, BorrowingLimit: nil, LendingLimit: nil}, + {Flavor: "x86", Resource: corev1.ResourceCPU}: {Nominal: 20_000, BorrowingLimit: nil, LendingLimit: nil}, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "arm", Resource: corev1.ResourceCPU}: 10_000, + {Flavor: "x86", Resource: corev1.ResourceCPU}: 20_000, + }, + }, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + Preemption: defaultPreemption, + NamespaceSelector: labels.Everything(), + Status: active, }, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - Preemption: defaultPreemption, - NamespaceSelector: labels.Everything(), - Status: active, }, }, ResourceFlavors: map[kueue.ResourceFlavorReference]*kueue.ResourceFlavor{ @@ -724,32 +744,36 @@ func TestSnapshot(t *testing.T) { ).Obj(), }, wantSnapshot: Snapshot{ - ClusterQueues: map[string]*ClusterQueueSnapshot{ - "cq": { - Name: "cq", - AllocatableResourceGeneration: 2, - ResourceGroups: []ResourceGroup{ - { - CoveredResources: sets.New(corev1.ResourceCPU), - Flavors: []kueue.ResourceFlavorReference{"arm", "x86"}, - }, - }, - ResourceNode: ResourceNode{ - Quotas: map[resources.FlavorResource]ResourceQuota{ - {Flavor: "arm", Resource: corev1.ResourceCPU}: {Nominal: 7_000, BorrowingLimit: nil, LendingLimit: ptr.To[int64](3_000)}, - {Flavor: "x86", Resource: corev1.ResourceCPU}: {Nominal: 5_000, BorrowingLimit: nil, LendingLimit: nil}, + Manager: hierarchy.Manager[*ClusterQueueSnapshot, *CohortSnapshot]{ + ClusterQueues: map[string]*ClusterQueueSnapshot{ + "cq": { + Name: "cq", + AllocatableResourceGeneration: 2, + ResourceGroups: []ResourceGroup{ + { + CoveredResources: sets.New(corev1.ResourceCPU), + Flavors: []kueue.ResourceFlavorReference{"arm", "x86"}, + }, }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "arm", Resource: corev1.ResourceCPU}: 7_000, - {Flavor: "x86", Resource: corev1.ResourceCPU}: 5_000, + ResourceNode: ResourceNode{ + Quotas: map[resources.FlavorResource]ResourceQuota{ + {Flavor: "arm", Resource: corev1.ResourceCPU}: {Nominal: 7_000, BorrowingLimit: nil, LendingLimit: ptr.To[int64](3_000)}, + {Flavor: "x86", Resource: corev1.ResourceCPU}: {Nominal: 5_000, BorrowingLimit: nil, LendingLimit: nil}, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "arm", Resource: corev1.ResourceCPU}: 7_000, + {Flavor: "x86", Resource: corev1.ResourceCPU}: 5_000, + }, }, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + Preemption: defaultPreemption, + NamespaceSelector: labels.Everything(), + Status: active, }, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - Preemption: defaultPreemption, - NamespaceSelector: labels.Everything(), - Status: active, - Cohort: &CohortSnapshot{ + }, + Cohorts: map[string]*CohortSnapshot{ + "cohort": { Name: "cohort", ResourceNode: ResourceNode{ Quotas: map[resources.FlavorResource]ResourceQuota{ @@ -881,7 +905,7 @@ func TestSnapshotAddRemoveWorkload(t *testing.T) { } } initialSnapshot := cqCache.Snapshot() - initialCohortResources := initialSnapshot.ClusterQueues["c1"].Cohort.ResourceNode.SubtreeQuota + initialCohortResources := initialSnapshot.ClusterQueues["c1"].Parent().ResourceNode.SubtreeQuota cases := map[string]struct { remove []string add []string @@ -907,41 +931,44 @@ func TestSnapshotAddRemoveWorkload(t *testing.T) { }, } return Snapshot{ - ClusterQueues: map[string]*ClusterQueueSnapshot{ - "c1": { - Name: "c1", - Cohort: cohort, - Workloads: make(map[string]*workload.Info), - ResourceGroups: cqCache.hm.ClusterQueues["c1"].ResourceGroups, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - ResourceNode: ResourceNode{ - Usage: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 0, - {Flavor: "alpha", Resource: corev1.ResourceMemory}: 0, - {Flavor: "beta", Resource: corev1.ResourceMemory}: 0, - }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 6_000, - {Flavor: "alpha", Resource: corev1.ResourceMemory}: utiltesting.Gi * 6, - {Flavor: "beta", Resource: corev1.ResourceMemory}: utiltesting.Gi * 6, - }, - }, + Manager: hierarchy.Manager[*ClusterQueueSnapshot, *CohortSnapshot]{ + Cohorts: map[string]*CohortSnapshot{ + "cohort": cohort, }, - "c2": { - Name: "c2", - Cohort: cohort, - Workloads: make(map[string]*workload.Info), - ResourceGroups: cqCache.hm.ClusterQueues["c2"].ResourceGroups, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - AllocatableResourceGeneration: 1, - ResourceNode: ResourceNode{ - Usage: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 0, + ClusterQueues: map[string]*ClusterQueueSnapshot{ + "c1": { + Name: "c1", + Workloads: make(map[string]*workload.Info), + ResourceGroups: cqCache.hm.ClusterQueues["c1"].ResourceGroups, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + ResourceNode: ResourceNode{ + Usage: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 0, + {Flavor: "alpha", Resource: corev1.ResourceMemory}: 0, + {Flavor: "beta", Resource: corev1.ResourceMemory}: 0, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 6_000, + {Flavor: "alpha", Resource: corev1.ResourceMemory}: utiltesting.Gi * 6, + {Flavor: "beta", Resource: corev1.ResourceMemory}: utiltesting.Gi * 6, + }, }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 6_000, + }, + "c2": { + Name: "c2", + Workloads: make(map[string]*workload.Info), + ResourceGroups: cqCache.hm.ClusterQueues["c2"].ResourceGroups, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + AllocatableResourceGeneration: 1, + ResourceNode: ResourceNode{ + Usage: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 0, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 6_000, + }, }, }, }, @@ -964,47 +991,50 @@ func TestSnapshotAddRemoveWorkload(t *testing.T) { }, } return Snapshot{ - ClusterQueues: map[string]*ClusterQueueSnapshot{ - "c1": { - Name: "c1", - Cohort: cohort, - Workloads: map[string]*workload.Info{ - "/c1-memory-alpha": nil, - "/c1-memory-beta": nil, - }, - ResourceGroups: cqCache.hm.ClusterQueues["c1"].ResourceGroups, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - ResourceNode: ResourceNode{ - Usage: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 0, - {Flavor: "alpha", Resource: corev1.ResourceMemory}: utiltesting.Gi, - {Flavor: "beta", Resource: corev1.ResourceMemory}: utiltesting.Gi, - }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 6_000, - {Flavor: "alpha", Resource: corev1.ResourceMemory}: utiltesting.Gi * 6, - {Flavor: "beta", Resource: corev1.ResourceMemory}: utiltesting.Gi * 6, - }, - }, + Manager: hierarchy.Manager[*ClusterQueueSnapshot, *CohortSnapshot]{ + Cohorts: map[string]*CohortSnapshot{ + "cohort": cohort, }, - "c2": { - Name: "c2", - Cohort: cohort, - Workloads: map[string]*workload.Info{ - "/c2-cpu-1": nil, - "/c2-cpu-2": nil, - }, - ResourceGroups: cqCache.hm.ClusterQueues["c2"].ResourceGroups, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - AllocatableResourceGeneration: 1, - ResourceNode: ResourceNode{ - Usage: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 2_000, + ClusterQueues: map[string]*ClusterQueueSnapshot{ + "c1": { + Name: "c1", + Workloads: map[string]*workload.Info{ + "/c1-memory-alpha": nil, + "/c1-memory-beta": nil, + }, + ResourceGroups: cqCache.hm.ClusterQueues["c1"].ResourceGroups, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + ResourceNode: ResourceNode{ + Usage: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 0, + {Flavor: "alpha", Resource: corev1.ResourceMemory}: utiltesting.Gi, + {Flavor: "beta", Resource: corev1.ResourceMemory}: utiltesting.Gi, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 6_000, + {Flavor: "alpha", Resource: corev1.ResourceMemory}: utiltesting.Gi * 6, + {Flavor: "beta", Resource: corev1.ResourceMemory}: utiltesting.Gi * 6, + }, }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 6_000, + }, + "c2": { + Name: "c2", + Workloads: map[string]*workload.Info{ + "/c2-cpu-1": nil, + "/c2-cpu-2": nil, + }, + ResourceGroups: cqCache.hm.ClusterQueues["c2"].ResourceGroups, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + AllocatableResourceGeneration: 1, + ResourceNode: ResourceNode{ + Usage: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 2_000, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 6_000, + }, }, }, }, @@ -1027,46 +1057,49 @@ func TestSnapshotAddRemoveWorkload(t *testing.T) { }, } return Snapshot{ - ClusterQueues: map[string]*ClusterQueueSnapshot{ - "c1": { - Name: "c1", - Cohort: cohort, - Workloads: map[string]*workload.Info{ - "/c1-memory-alpha": nil, - "/c1-memory-beta": nil, - }, - ResourceGroups: cqCache.hm.ClusterQueues["c1"].ResourceGroups, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - ResourceNode: ResourceNode{ - Usage: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 1_000, - {Flavor: "alpha", Resource: corev1.ResourceMemory}: 0, - {Flavor: "beta", Resource: corev1.ResourceMemory}: utiltesting.Gi, - }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 6_000, - {Flavor: "alpha", Resource: corev1.ResourceMemory}: utiltesting.Gi * 6, - {Flavor: "beta", Resource: corev1.ResourceMemory}: utiltesting.Gi * 6, - }, - }, + Manager: hierarchy.Manager[*ClusterQueueSnapshot, *CohortSnapshot]{ + Cohorts: map[string]*CohortSnapshot{ + "cohort": cohort, }, - "c2": { - Name: "c2", - Cohort: cohort, - Workloads: map[string]*workload.Info{ - "/c2-cpu-1": nil, - "/c2-cpu-2": nil, - }, - ResourceGroups: cqCache.hm.ClusterQueues["c2"].ResourceGroups, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - ResourceNode: ResourceNode{ - Usage: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 2_000, + ClusterQueues: map[string]*ClusterQueueSnapshot{ + "c1": { + Name: "c1", + Workloads: map[string]*workload.Info{ + "/c1-memory-alpha": nil, + "/c1-memory-beta": nil, + }, + ResourceGroups: cqCache.hm.ClusterQueues["c1"].ResourceGroups, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + ResourceNode: ResourceNode{ + Usage: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 1_000, + {Flavor: "alpha", Resource: corev1.ResourceMemory}: 0, + {Flavor: "beta", Resource: corev1.ResourceMemory}: utiltesting.Gi, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 6_000, + {Flavor: "alpha", Resource: corev1.ResourceMemory}: utiltesting.Gi * 6, + {Flavor: "beta", Resource: corev1.ResourceMemory}: utiltesting.Gi * 6, + }, }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 6_000, + }, + "c2": { + Name: "c2", + Workloads: map[string]*workload.Info{ + "/c2-cpu-1": nil, + "/c2-cpu-2": nil, + }, + ResourceGroups: cqCache.hm.ClusterQueues["c2"].ResourceGroups, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + ResourceNode: ResourceNode{ + Usage: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 2_000, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 6_000, + }, }, }, }, @@ -1160,7 +1193,7 @@ func TestSnapshotAddRemoveWorkloadWithLendingLimit(t *testing.T) { } } initialSnapshot := cqCache.Snapshot() - initialCohortResources := initialSnapshot.ClusterQueues["lend-a"].Cohort.ResourceNode.SubtreeQuota + initialCohortResources := initialSnapshot.ClusterQueues["lend-a"].Parent().ResourceNode.SubtreeQuota cases := map[string]struct { remove []string add []string @@ -1184,36 +1217,39 @@ func TestSnapshotAddRemoveWorkloadWithLendingLimit(t *testing.T) { }, } return Snapshot{ - ClusterQueues: map[string]*ClusterQueueSnapshot{ - "lend-a": { - Name: "lend-a", - Cohort: cohort, - Workloads: make(map[string]*workload.Info), - ResourceGroups: cqCache.hm.ClusterQueues["lend-a"].ResourceGroups, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - ResourceNode: ResourceNode{ - Usage: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 0, - }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, - }, - }, + Manager: hierarchy.Manager[*ClusterQueueSnapshot, *CohortSnapshot]{ + Cohorts: map[string]*CohortSnapshot{ + "lend": cohort, }, - "lend-b": { - Name: "lend-b", - Cohort: cohort, - Workloads: make(map[string]*workload.Info), - ResourceGroups: cqCache.hm.ClusterQueues["lend-b"].ResourceGroups, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - ResourceNode: ResourceNode{ - Usage: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 0, + ClusterQueues: map[string]*ClusterQueueSnapshot{ + "lend-a": { + Name: "lend-a", + Workloads: make(map[string]*workload.Info), + ResourceGroups: cqCache.hm.ClusterQueues["lend-a"].ResourceGroups, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + ResourceNode: ResourceNode{ + Usage: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 0, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, + }, }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, + }, + "lend-b": { + Name: "lend-b", + Workloads: make(map[string]*workload.Info), + ResourceGroups: cqCache.hm.ClusterQueues["lend-b"].ResourceGroups, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + ResourceNode: ResourceNode{ + Usage: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 0, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, + }, }, }, }, @@ -1234,37 +1270,40 @@ func TestSnapshotAddRemoveWorkloadWithLendingLimit(t *testing.T) { }, } return Snapshot{ - ClusterQueues: map[string]*ClusterQueueSnapshot{ - "lend-a": { - Name: "lend-a", - Cohort: cohort, - Workloads: make(map[string]*workload.Info), - ResourceGroups: cqCache.hm.ClusterQueues["lend-a"].ResourceGroups, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - ResourceNode: ResourceNode{ - Usage: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 7_000, - }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, - }, - }, + Manager: hierarchy.Manager[*ClusterQueueSnapshot, *CohortSnapshot]{ + Cohorts: map[string]*CohortSnapshot{ + "lend": cohort, }, - "lend-b": { - Name: "lend-b", - Cohort: cohort, - Workloads: make(map[string]*workload.Info), - ResourceGroups: cqCache.hm.ClusterQueues["lend-b"].ResourceGroups, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - AllocatableResourceGeneration: 1, - ResourceNode: ResourceNode{ - Usage: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 4_000, + ClusterQueues: map[string]*ClusterQueueSnapshot{ + "lend-a": { + Name: "lend-a", + Workloads: make(map[string]*workload.Info), + ResourceGroups: cqCache.hm.ClusterQueues["lend-a"].ResourceGroups, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + ResourceNode: ResourceNode{ + Usage: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 7_000, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, + }, }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, + }, + "lend-b": { + Name: "lend-b", + Workloads: make(map[string]*workload.Info), + ResourceGroups: cqCache.hm.ClusterQueues["lend-b"].ResourceGroups, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + AllocatableResourceGeneration: 1, + ResourceNode: ResourceNode{ + Usage: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 4_000, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, + }, }, }, }, @@ -1285,37 +1324,40 @@ func TestSnapshotAddRemoveWorkloadWithLendingLimit(t *testing.T) { }, } return Snapshot{ - ClusterQueues: map[string]*ClusterQueueSnapshot{ - "lend-a": { - Name: "lend-a", - Cohort: cohort, - Workloads: make(map[string]*workload.Info), - ResourceGroups: cqCache.hm.ClusterQueues["lend-a"].ResourceGroups, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - ResourceNode: ResourceNode{ - Usage: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 6_000, - }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, - }, - }, + Manager: hierarchy.Manager[*ClusterQueueSnapshot, *CohortSnapshot]{ + Cohorts: map[string]*CohortSnapshot{ + "lend": cohort, }, - "lend-b": { - Name: "lend-b", - Cohort: cohort, - Workloads: make(map[string]*workload.Info), - ResourceGroups: cqCache.hm.ClusterQueues["lend-b"].ResourceGroups, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - AllocatableResourceGeneration: 1, - ResourceNode: ResourceNode{ - Usage: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 4_000, + ClusterQueues: map[string]*ClusterQueueSnapshot{ + "lend-a": { + Name: "lend-a", + Workloads: make(map[string]*workload.Info), + ResourceGroups: cqCache.hm.ClusterQueues["lend-a"].ResourceGroups, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + ResourceNode: ResourceNode{ + Usage: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 6_000, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, + }, }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, + }, + "lend-b": { + Name: "lend-b", + Workloads: make(map[string]*workload.Info), + ResourceGroups: cqCache.hm.ClusterQueues["lend-b"].ResourceGroups, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + AllocatableResourceGeneration: 1, + ResourceNode: ResourceNode{ + Usage: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 4_000, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, + }, }, }, }, @@ -1336,37 +1378,40 @@ func TestSnapshotAddRemoveWorkloadWithLendingLimit(t *testing.T) { }, } return Snapshot{ - ClusterQueues: map[string]*ClusterQueueSnapshot{ - "lend-a": { - Name: "lend-a", - Cohort: cohort, - Workloads: make(map[string]*workload.Info), - ResourceGroups: cqCache.hm.ClusterQueues["lend-a"].ResourceGroups, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - ResourceNode: ResourceNode{ - Usage: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 1_000, - }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, - }, - }, + Manager: hierarchy.Manager[*ClusterQueueSnapshot, *CohortSnapshot]{ + Cohorts: map[string]*CohortSnapshot{ + "lend": cohort, }, - "lend-b": { - Name: "lend-b", - Cohort: cohort, - Workloads: make(map[string]*workload.Info), - ResourceGroups: cqCache.hm.ClusterQueues["lend-b"].ResourceGroups, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - AllocatableResourceGeneration: 1, - ResourceNode: ResourceNode{ - Usage: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 4_000, + ClusterQueues: map[string]*ClusterQueueSnapshot{ + "lend-a": { + Name: "lend-a", + Workloads: make(map[string]*workload.Info), + ResourceGroups: cqCache.hm.ClusterQueues["lend-a"].ResourceGroups, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + ResourceNode: ResourceNode{ + Usage: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 1_000, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, + }, }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, + }, + "lend-b": { + Name: "lend-b", + Workloads: make(map[string]*workload.Info), + ResourceGroups: cqCache.hm.ClusterQueues["lend-b"].ResourceGroups, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + AllocatableResourceGeneration: 1, + ResourceNode: ResourceNode{ + Usage: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 4_000, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, + }, }, }, }, @@ -1388,37 +1433,40 @@ func TestSnapshotAddRemoveWorkloadWithLendingLimit(t *testing.T) { }, } return Snapshot{ - ClusterQueues: map[string]*ClusterQueueSnapshot{ - "lend-a": { - Name: "lend-a", - Cohort: cohort, - Workloads: make(map[string]*workload.Info), - ResourceGroups: cqCache.hm.ClusterQueues["lend-a"].ResourceGroups, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - ResourceNode: ResourceNode{ - Usage: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 1_000, - }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, - }, - }, + Manager: hierarchy.Manager[*ClusterQueueSnapshot, *CohortSnapshot]{ + Cohorts: map[string]*CohortSnapshot{ + "lend": cohort, }, - "lend-b": { - Name: "lend-b", - Cohort: cohort, - Workloads: make(map[string]*workload.Info), - ResourceGroups: cqCache.hm.ClusterQueues["lend-b"].ResourceGroups, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - AllocatableResourceGeneration: 1, - ResourceNode: ResourceNode{ - Usage: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 0, + ClusterQueues: map[string]*ClusterQueueSnapshot{ + "lend-a": { + Name: "lend-a", + Workloads: make(map[string]*workload.Info), + ResourceGroups: cqCache.hm.ClusterQueues["lend-a"].ResourceGroups, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + ResourceNode: ResourceNode{ + Usage: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 1_000, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, + }, }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, + }, + "lend-b": { + Name: "lend-b", + Workloads: make(map[string]*workload.Info), + ResourceGroups: cqCache.hm.ClusterQueues["lend-b"].ResourceGroups, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + AllocatableResourceGeneration: 1, + ResourceNode: ResourceNode{ + Usage: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 0, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, + }, }, }, }, @@ -1440,36 +1488,39 @@ func TestSnapshotAddRemoveWorkloadWithLendingLimit(t *testing.T) { }, } return Snapshot{ - ClusterQueues: map[string]*ClusterQueueSnapshot{ - "lend-a": { - Name: "lend-a", - Cohort: cohort, - Workloads: make(map[string]*workload.Info), - ResourceGroups: cqCache.hm.ClusterQueues["lend-a"].ResourceGroups, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - ResourceNode: ResourceNode{ - Usage: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 6_000, - }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, - }, - }, + Manager: hierarchy.Manager[*ClusterQueueSnapshot, *CohortSnapshot]{ + Cohorts: map[string]*CohortSnapshot{ + "lend": cohort, }, - "lend-b": { - Name: "lend-b", - Cohort: cohort, - Workloads: make(map[string]*workload.Info), - ResourceGroups: cqCache.hm.ClusterQueues["lend-b"].ResourceGroups, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - ResourceNode: ResourceNode{ - Usage: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 0, + ClusterQueues: map[string]*ClusterQueueSnapshot{ + "lend-a": { + Name: "lend-a", + Workloads: make(map[string]*workload.Info), + ResourceGroups: cqCache.hm.ClusterQueues["lend-a"].ResourceGroups, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + ResourceNode: ResourceNode{ + Usage: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 6_000, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, + }, }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, + }, + "lend-b": { + Name: "lend-b", + Workloads: make(map[string]*workload.Info), + ResourceGroups: cqCache.hm.ClusterQueues["lend-b"].ResourceGroups, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + ResourceNode: ResourceNode{ + Usage: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 0, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, + }, }, }, }, @@ -1491,37 +1542,40 @@ func TestSnapshotAddRemoveWorkloadWithLendingLimit(t *testing.T) { }, } return Snapshot{ - ClusterQueues: map[string]*ClusterQueueSnapshot{ - "lend-a": { - Name: "lend-a", - Cohort: cohort, - Workloads: make(map[string]*workload.Info), - ResourceGroups: cqCache.hm.ClusterQueues["lend-a"].ResourceGroups, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - ResourceNode: ResourceNode{ - Usage: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 9_000, - }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, - }, - }, + Manager: hierarchy.Manager[*ClusterQueueSnapshot, *CohortSnapshot]{ + Cohorts: map[string]*CohortSnapshot{ + "lend": cohort, }, - "lend-b": { - Name: "lend-b", - Cohort: cohort, - Workloads: make(map[string]*workload.Info), - ResourceGroups: cqCache.hm.ClusterQueues["lend-b"].ResourceGroups, - FlavorFungibility: defaultFlavorFungibility, - FairWeight: oneQuantity, - AllocatableResourceGeneration: 1, - ResourceNode: ResourceNode{ - Usage: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 0, + ClusterQueues: map[string]*ClusterQueueSnapshot{ + "lend-a": { + Name: "lend-a", + Workloads: make(map[string]*workload.Info), + ResourceGroups: cqCache.hm.ClusterQueues["lend-a"].ResourceGroups, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + ResourceNode: ResourceNode{ + Usage: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 9_000, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, + }, }, - SubtreeQuota: resources.FlavorResourceQuantities{ - {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, + }, + "lend-b": { + Name: "lend-b", + Workloads: make(map[string]*workload.Info), + ResourceGroups: cqCache.hm.ClusterQueues["lend-b"].ResourceGroups, + FlavorFungibility: defaultFlavorFungibility, + FairWeight: oneQuantity, + AllocatableResourceGeneration: 1, + ResourceNode: ResourceNode{ + Usage: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 0, + }, + SubtreeQuota: resources.FlavorResourceQuantities{ + {Flavor: "default", Resource: corev1.ResourceCPU}: 10_000, + }, }, }, }, diff --git a/pkg/scheduler/flavorassigner/flavorassigner.go b/pkg/scheduler/flavorassigner/flavorassigner.go index bac68990b5..d38b77ce02 100644 --- a/pkg/scheduler/flavorassigner/flavorassigner.go +++ b/pkg/scheduler/flavorassigner/flavorassigner.go @@ -613,7 +613,7 @@ func (a *FlavorAssigner) fitsResourceQuota(log logr.Logger, fr resources.FlavorR lackQuantity := resources.ResourceQuantity(fr.Resource, lack) msg := fmt.Sprintf("insufficient unused quota in cohort for %s in flavor %s, %s more needed", fr.Resource, fr.Flavor, &lackQuantity) - if a.cq.Cohort == nil { + if !a.cq.HasParent() { if mode == noFit { msg = fmt.Sprintf("insufficient quota for %s in flavor %s in ClusterQueue", fr.Resource, fr.Flavor) } else { diff --git a/pkg/scheduler/flavorassigner/flavorassigner_test.go b/pkg/scheduler/flavorassigner/flavorassigner_test.go index b23031b8fd..c2493f2353 100644 --- a/pkg/scheduler/flavorassigner/flavorassigner_test.go +++ b/pkg/scheduler/flavorassigner/flavorassigner_test.go @@ -1954,11 +1954,11 @@ func TestAssignFlavors(t *testing.T) { } if tc.cohortResources != nil { - if clusterQueue.Cohort == nil { + if !clusterQueue.HasParent() { t.Fatalf("Test case has cohort resources, but cluster queue doesn't have cohort") } - clusterQueue.Cohort.ResourceNode.Usage = tc.cohortResources.usage - clusterQueue.Cohort.ResourceNode.SubtreeQuota = tc.cohortResources.requestableResources + clusterQueue.Parent().ResourceNode.Usage = tc.cohortResources.usage + clusterQueue.Parent().ResourceNode.SubtreeQuota = tc.cohortResources.requestableResources } clusterQueue.ResourceNode.Usage = tc.clusterQueueUsage @@ -2269,7 +2269,6 @@ func TestLastAssignmentOutdated(t *testing.T) { }, }, cq: &cache.ClusterQueueSnapshot{ - Cohort: nil, AllocatableResourceGeneration: 1, }, }, diff --git a/pkg/scheduler/preemption/preemption.go b/pkg/scheduler/preemption/preemption.go index f370afa577..3b061c7cdd 100644 --- a/pkg/scheduler/preemption/preemption.go +++ b/pkg/scheduler/preemption/preemption.go @@ -510,9 +510,9 @@ func (p *Preemptor) findCandidates(wl *kueue.Workload, cq *cache.ClusterQueueSna } } - if cq.Cohort != nil && cq.Preemption.ReclaimWithinCohort != kueue.PreemptionPolicyNever { + if cq.HasParent() && cq.Preemption.ReclaimWithinCohort != kueue.PreemptionPolicyNever { onlyLowerPriority := cq.Preemption.ReclaimWithinCohort != kueue.PreemptionPolicyAny - for cohortCQ := range cq.Cohort.Members { + for _, cohortCQ := range cq.Parent().ChildCQs() { if cq == cohortCQ || !cqIsBorrowing(cohortCQ, frsNeedPreemption) { // Can't reclaim quota from itself or ClusterQueues that are not borrowing. continue @@ -532,7 +532,7 @@ func (p *Preemptor) findCandidates(wl *kueue.Workload, cq *cache.ClusterQueueSna } func cqIsBorrowing(cq *cache.ClusterQueueSnapshot, frsNeedPreemption sets.Set[resources.FlavorResource]) bool { - if cq.Cohort == nil { + if !cq.HasParent() { return false } for fr := range frsNeedPreemption { diff --git a/pkg/scheduler/preemption/preemption_test.go b/pkg/scheduler/preemption/preemption_test.go index ee5020205d..4d6a533a8b 100644 --- a/pkg/scheduler/preemption/preemption_test.go +++ b/pkg/scheduler/preemption/preemption_test.go @@ -41,6 +41,7 @@ import ( "sigs.k8s.io/kueue/pkg/cache" "sigs.k8s.io/kueue/pkg/constants" "sigs.k8s.io/kueue/pkg/features" + "sigs.k8s.io/kueue/pkg/hierarchy" "sigs.k8s.io/kueue/pkg/scheduler/flavorassigner" "sigs.k8s.io/kueue/pkg/util/slices" utiltesting "sigs.k8s.io/kueue/pkg/util/testing" @@ -49,7 +50,9 @@ import ( var snapCmpOpts = []cmp.Option{ cmpopts.EquateEmpty(), - cmpopts.IgnoreUnexported(cache.ClusterQueueSnapshot{}), + cmpopts.IgnoreUnexported(hierarchy.Cohort[*cache.ClusterQueueSnapshot, *cache.CohortSnapshot]{}), + cmpopts.IgnoreUnexported(hierarchy.ClusterQueue[*cache.CohortSnapshot]{}), + cmpopts.IgnoreUnexported(hierarchy.Manager[*cache.ClusterQueueSnapshot, *cache.CohortSnapshot]{}), cmpopts.IgnoreFields(cache.ClusterQueueSnapshot{}, "AllocatableResourceGeneration"), cmp.Transformer("Cohort.Members", func(s sets.Set[*cache.ClusterQueueSnapshot]) sets.Set[string] { result := make(sets.Set[string], len(s)) diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go index ff34f79565..f225603f6e 100644 --- a/pkg/scheduler/scheduler.go +++ b/pkg/scheduler/scheduler.go @@ -273,17 +273,17 @@ func (s *Scheduler) schedule(ctx context.Context) wait.SpeedSignal { } preemptedWorkloads.Insert(pendingPreemptions...) cq.AddUsage(usage) - } else if cq.Cohort != nil { - sum := cycleCohortsUsage.totalUsageForCommonFlavorResources(cq.Cohort.Name, e.assignment.Usage) + } else if cq.HasParent() { + sum := cycleCohortsUsage.totalUsageForCommonFlavorResources(cq.Parent().Name, e.assignment.Usage) // Check whether there was an assignment in this cycle that could render the next assignments invalid: // - If the workload no longer fits in the cohort. // - If there was another assignment in the cohort, then the preemption calculation is no longer valid. - if cycleCohortsUsage.hasCommonFlavorResources(cq.Cohort.Name, e.assignment.Usage) { + if cycleCohortsUsage.hasCommonFlavorResources(cq.Parent().Name, e.assignment.Usage) { if mode == flavorassigner.Fit && !cq.FitInCohort(sum) { setSkipped(e, "Workload no longer fits after processing another workload") continue } - if mode == flavorassigner.Preempt && cycleCohortsSkipPreemption.Has(cq.Cohort.Name) { + if mode == flavorassigner.Preempt && cycleCohortsSkipPreemption.Has(cq.Parent().Name) { setSkipped(e, "Workload skipped because its preemption calculations were invalidated by another workload") skippedPreemptions[cq.Name]++ continue @@ -291,7 +291,7 @@ func (s *Scheduler) schedule(ctx context.Context) wait.SpeedSignal { } // Even if the workload will not be admitted after this point, due to preemption pending or other failures, // we should still account for its usage. - cycleCohortsUsage.add(cq.Cohort.Name, resourcesToReserve(e, cq)) + cycleCohortsUsage.add(cq.Parent().Name, resourcesToReserve(e, cq)) } if e.assignment.RepresentativeMode() != flavorassigner.Fit { if len(e.preemptionTargets) != 0 { @@ -305,8 +305,8 @@ func (s *Scheduler) schedule(ctx context.Context) wait.SpeedSignal { e.inadmissibleMsg += fmt.Sprintf(". Pending the preemption of %d workload(s)", preempted) e.requeueReason = queue.RequeueReasonPendingPreemption } - if cq.Cohort != nil { - cycleCohortsSkipPreemption.Insert(cq.Cohort.Name) + if cq.HasParent() { + cycleCohortsSkipPreemption.Insert(cq.Parent().Name) } } else { log.V(2).Info("Workload requires preemption, but there are no candidate workloads allowed for preemption", "preemption", cq.Preemption) @@ -329,8 +329,8 @@ func (s *Scheduler) schedule(ctx context.Context) wait.SpeedSignal { if err := s.admit(ctx, e, cq); err != nil { e.inadmissibleMsg = fmt.Sprintf("Failed to admit workload: %v", err) } - if cq.Cohort != nil { - cycleCohortsSkipPreemption.Insert(cq.Cohort.Name) + if cq.HasParent() { + cycleCohortsSkipPreemption.Insert(cq.Parent().Name) } }