WIP: Repeat cohort tests fixed; experiment after fix

pkg/cache/scheduler/cache.go

@@ -540,6 +540,8 @@ func (c *Cache) AddOrUpdateCohort(apiCohort *kueue.Cohort) error {
 	return nil
 }
 
+// DeleteCohort removes the cohort from the cache and updates the SubtreeQuota
+// of ancestor cohorts to reflect the removal.
 func (c *Cache) DeleteCohort(cohortName kueue.CohortReference) {
 	c.Lock()
 	defer c.Unlock()
@@ -560,6 +562,14 @@ func (c *Cache) DeleteCohort(cohortName kueue.CohortReference) {
 	if cohort := c.hm.Cohort(cohortName); cohort != nil {
 		updateCohortResourceNode(cohort)
 	}
+
+	// Update the SubtreeQuota of all ancestors to reflect the removal of
+	// the deleted cohort.
+	if parent != nil {
+		if updatedParent := c.hm.Cohort(parent.GetName()); updatedParent != nil {
+			updateCohortTreeResourcesIfNoCycle(updatedParent)
+		}
+	}
 }
 
 func (c *Cache) handleParentUpdate(cachedParent *cohort) {

pkg/cache/scheduler/cache_test.go

@@ -3570,6 +3570,150 @@ func TestCohortCycles(t *testing.T) {
 	})
 }
 
+func TestDeleteCohortUpdatesAncestorSubtreeQuota(t *testing.T) {
+	t.Run("deleting child cohort updates parent SubtreeQuota", func(t *testing.T) {
+		cache := New(utiltesting.NewFakeClient())
+
+		root := utiltestingapi.MakeCohort("root").
+			ResourceGroup(*utiltestingapi.MakeFlavorQuotas("default").Resource(corev1.ResourceCPU, "10").Obj()).
+			Obj()
+		if err := cache.AddOrUpdateCohort(root); err != nil {
+			t.Fatal(err)
+		}
+
+		child := utiltestingapi.MakeCohort("child").
+			Parent("root").
+			ResourceGroup(*utiltestingapi.MakeFlavorQuotas("default").Resource(corev1.ResourceCPU, "5").Obj()).
+			Obj()
+		if err := cache.AddOrUpdateCohort(child); err != nil {
+			t.Fatal(err)
+		}
+
+		// before deletion: root SubtreeQuota = 10 (own) + 5 (child) = 15
+		wantBefore := resources.FlavorResourceQuantities{
+			{Flavor: "default", Resource: corev1.ResourceCPU}: 15_000,
+		}
+		if diff := cmp.Diff(wantBefore, cache.hm.Cohort("root").getResourceNode().SubtreeQuota); diff != "" {
+			t.Errorf("before deletion, unexpected SubtreeQuota (-want,+got):\n%s", diff)
+		}
+
+		cache.DeleteCohort("child")
+
+		// after deletion: root SubtreeQuota = 10 (own only)
+		wantAfter := resources.FlavorResourceQuantities{
+			{Flavor: "default", Resource: corev1.ResourceCPU}: 10_000,
+		}
+		if diff := cmp.Diff(wantAfter, cache.hm.Cohort("root").getResourceNode().SubtreeQuota); diff != "" {
+			t.Errorf("after deletion, unexpected SubtreeQuota (-want,+got):\n%s", diff)
+		}
+	})
+
+	t.Run("deleting grandchild cohort updates all ancestors SubtreeQuota", func(t *testing.T) {
+		cache := New(utiltesting.NewFakeClient())
+
+		root := utiltestingapi.MakeCohort("root").
+			ResourceGroup(*utiltestingapi.MakeFlavorQuotas("default").Resource(corev1.ResourceCPU, "10").Obj()).
+			Obj()
+		if err := cache.AddOrUpdateCohort(root); err != nil {
+			t.Fatal(err)
+		}
+
+		mid := utiltestingapi.MakeCohort("mid").
+			Parent("root").
+			ResourceGroup(*utiltestingapi.MakeFlavorQuotas("default").Resource(corev1.ResourceCPU, "5").Obj()).
+			Obj()
+		if err := cache.AddOrUpdateCohort(mid); err != nil {
+			t.Fatal(err)
+		}
+
+		leaf := utiltestingapi.MakeCohort("leaf").
+			Parent("mid").
+			ResourceGroup(*utiltestingapi.MakeFlavorQuotas("default").Resource(corev1.ResourceCPU, "3").Obj()).
+			Obj()
+		if err := cache.AddOrUpdateCohort(leaf); err != nil {
+			t.Fatal(err)
+		}
+
+		// before: mid = 5+3 = 8, root = 10+8 = 18
+		wantMidBefore := resources.FlavorResourceQuantities{
+			{Flavor: "default", Resource: corev1.ResourceCPU}: 8_000,
+		}
+		wantRootBefore := resources.FlavorResourceQuantities{
+			{Flavor: "default", Resource: corev1.ResourceCPU}: 18_000,
+		}
+		if diff := cmp.Diff(wantMidBefore, cache.hm.Cohort("mid").getResourceNode().SubtreeQuota); diff != "" {
+			t.Errorf("before deletion, mid unexpected SubtreeQuota (-want,+got):\n%s", diff)
+		}
+		if diff := cmp.Diff(wantRootBefore, cache.hm.Cohort("root").getResourceNode().SubtreeQuota); diff != "" {
+			t.Errorf("before deletion, root unexpected SubtreeQuota (-want,+got):\n%s", diff)
+		}
+
+		cache.DeleteCohort("leaf")
+
+		// after: mid = 5, root = 10+5 = 15
+		wantMidAfter := resources.FlavorResourceQuantities{
+			{Flavor: "default", Resource: corev1.ResourceCPU}: 5_000,
+		}
+		wantRootAfter := resources.FlavorResourceQuantities{
+			{Flavor: "default", Resource: corev1.ResourceCPU}: 15_000,
+		}
+		if diff := cmp.Diff(wantMidAfter, cache.hm.Cohort("mid").getResourceNode().SubtreeQuota); diff != "" {
+			t.Errorf("after deletion, mid unexpected SubtreeQuota (-want,+got):\n%s", diff)
+		}
+		if diff := cmp.Diff(wantRootAfter, cache.hm.Cohort("root").getResourceNode().SubtreeQuota); diff != "" {
+			t.Errorf("after deletion, root unexpected SubtreeQuota (-want,+got):\n%s", diff)
+		}
+	})
+
+	t.Run("deleting cohort with children updates parent SubtreeQuota", func(t *testing.T) {
+		cache := New(utiltesting.NewFakeClient())
+		ctx, _ := utiltesting.ContextWithLog(t)
+
+		root := utiltestingapi.MakeCohort("root").
+			ResourceGroup(*utiltestingapi.MakeFlavorQuotas("default").Resource(corev1.ResourceCPU, "10").Obj()).
+			Obj()
+		if err := cache.AddOrUpdateCohort(root); err != nil {
+			t.Fatal(err)
+		}
+
+		mid := utiltestingapi.MakeCohort("mid").
+			Parent("root").
+			ResourceGroup(*utiltestingapi.MakeFlavorQuotas("default").Resource(corev1.ResourceCPU, "5").Obj()).
+			Obj()
+		if err := cache.AddOrUpdateCohort(mid); err != nil {
+			t.Fatal(err)
+		}
+
+		// cq is a child of mid; after mid is deleted, mid becomes implicit (no parent)
+		// and is detached from root. So root loses mid's entire subtree contribution.
+		if err := cache.AddClusterQueue(ctx, utiltestingapi.MakeClusterQueue("cq").
+			Cohort("mid").
+			ResourceGroup(*utiltestingapi.MakeFlavorQuotas("default").Resource(corev1.ResourceCPU, "7").Obj()).
+			Obj()); err != nil {
+			t.Fatal(err)
+		}
+
+		// before: mid = 5+7 = 12, root = 10+12 = 22
+		wantRootBefore := resources.FlavorResourceQuantities{
+			{Flavor: "default", Resource: corev1.ResourceCPU}: 22_000,
+		}
+		if diff := cmp.Diff(wantRootBefore, cache.hm.Cohort("root").getResourceNode().SubtreeQuota); diff != "" {
+			t.Errorf("before deletion, root unexpected SubtreeQuota (-want,+got):\n%s", diff)
+		}
+
+		// Delete mid (explicit cohort). It becomes implicit because cq still references it.
+		cache.DeleteCohort("mid")
+
+		// after: mid is detached from root entirely (implicit mid has no parent),
+		// so root = 10 (own quota only)
+		wantRootAfter := resources.FlavorResourceQuantities{
+			{Flavor: "default", Resource: corev1.ResourceCPU}: 10_000,
+		}
+		if diff := cmp.Diff(wantRootAfter, cache.hm.Cohort("root").getResourceNode().SubtreeQuota); diff != "" {
+			t.Errorf("after deletion, root unexpected SubtreeQuota (-want,+got):\n%s", diff)
+		}
+	})
+}
 func TestClusterQueueAncestors(t *testing.T) {
 	testCases := map[string]struct {
 		cohorts       []*kueue.Cohort