Repository: kubernetes-sigs/kueue

Issue:
TAS cleanup: move optional parameters for FindTopologyAssignmentsForFlavor to options
<!-- Please only use this template for submitting clean up requests -->

**What would you like to be cleaned**:

I would like to move optional parameters to FindTopologyAssignmentsForFlavor as [cache options](https://github.com/kubernetes-sigs/kueue/blob/9535437ade9479a691994dad69dbe093de50a3ef/pkg/cache/cache.go#L62-L99) (simulateEmpty and Workload). 

Similar pattern as for example 

**Why is this needed**:

Improve code readability by skipping optional parameters. 

These parameters are not essential for the function. For example when calling in [unit tests](https://github.com/kubernetes-sigs/kueue/blob/9535437ade9479a691994dad69dbe093de50a3ef/pkg/cache/tas_cache_test.go#L1928C26-L1928C58) we pass false and nil.

Discussion:
@mimowo:  @PBundyra @lchrzaszcz  ptal

cc @kaisoz @mbobrovskyi  @mimowo: I think we should also use the opts for FindTopologyAssignmentsForWorkload, then we can just propagate them when calling `FindTopologyAssignmentsForFlavor` from that function.@mbobrovskyi: /assign

Solution patch:
diff --git a/pkg/cache/clusterqueue_snapshot.go b/pkg/cache/clusterqueue_snapshot.go
index ffddf1a82..5cff48d73 100644
--- a/pkg/cache/clusterqueue_snapshot.go
+++ b/pkg/cache/clusterqueue_snapshot.go
@@ -208,14 +208,20 @@ type WorkloadTASRequests map[kueue.ResourceFlavorReference]FlavorTASRequests
 
 func (c *ClusterQueueSnapshot) FindTopologyAssignmentsForWorkload(
 	tasRequestsByFlavor WorkloadTASRequests,
-	simulateEmpty bool, wl *kueue.Workload) TASAssignmentsResult {
+	options ...FindTopologyAssignmentsOption,
+) TASAssignmentsResult {
+	opts := &findTopologyAssignmentsOption{}
+	for _, option := range options {
+		option(opts)
+	}
+
 	result := make(TASAssignmentsResult)
 	for tasFlavor, flavorTASRequests := range tasRequestsByFlavor {
 		// We assume the `tasFlavor` is already in the snapshot as this was
 		// already checked earlier during flavor assignment, and the set of
 		// flavors is immutable in snapshot.
 		tasFlavorCache := c.TASFlavors[tasFlavor]
-		flvResult := tasFlavorCache.FindTopologyAssignmentsForFlavor(flavorTASRequests, simulateEmpty, wl)
+		flvResult := tasFlavorCache.FindTopologyAssignmentsForFlavor(flavorTASRequests, options...)
 		maps.Copy(result, flvResult)
 	}
 	return result
diff --git a/pkg/cache/tas_flavor_snapshot.go b/pkg/cache/tas_flavor_snapshot.go
index 6f4450ce8..8ac4e5414 100644
--- a/pkg/cache/tas_flavor_snapshot.go
+++ b/pkg/cache/tas_flavor_snapshot.go
@@ -387,31 +387,55 @@ func (s *TASFlavorSnapshot) Fits(flavorUsage workload.TASFlavorUsage) bool {
 	return true
 }
 
+type findTopologyAssignmentsOption struct {
+	simulateEmpty bool
+	workload      *kueue.Workload
+}
+
+type FindTopologyAssignmentsOption func(*findTopologyAssignmentsOption)
+
+// WithSimulateEmpty sets parameter allows to look for the assignment under the
+// assumption that all TAS workloads are preempted.
+func WithSimulateEmpty(simulateEmpty bool) FindTopologyAssignmentsOption {
+	return func(o *findTopologyAssignmentsOption) {
+		o.simulateEmpty = simulateEmpty
+	}
+}
+
+func WithWorkload(wl *kueue.Workload) FindTopologyAssignmentsOption {
+	return func(o *findTopologyAssignmentsOption) {
+		o.workload = wl
+	}
+}
+
 // FindTopologyAssignmentsForFlavor returns TAS assignment, if possible, for all
 // the TAS requests in the flavor handled by the snapshot.
-// The simulateEmpty parameter allows to look for the assignment under the
-// assumption that all TAS workloads are preempted.
-func (s *TASFlavorSnapshot) FindTopologyAssignmentsForFlavor(flavorTASRequests FlavorTASRequests, simulateEmpty bool, wl *kueue.Workload) TASAssignmentsResult {
+func (s *TASFlavorSnapshot) FindTopologyAssignmentsForFlavor(flavorTASRequests FlavorTASRequests, options ...FindTopologyAssignmentsOption) TASAssignmentsResult {
+	opts := &findTopologyAssignmentsOption{}
+	for _, option := range options {
+		option(opts)
+	}
+
 	result := make(map[kueue.PodSetReference]tasPodSetAssignmentResult)
 	assumedUsage := make(map[utiltas.TopologyDomainID]resources.Requests)
 	for _, tr := range flavorTASRequests {
-		if workload.HasNodeToReplace(wl) {
+		if workload.HasNodeToReplace(opts.workload) {
 			// In case of looking for Node replacement, TopologyRequest has only
 			// PodSets with the Node to replace, so we match PodSetAssignment
-			psa := findPSA(wl, tr.PodSet.Name)
+			psa := findPSA(opts.workload, tr.PodSet.Name)
 			if psa == nil || psa.TopologyAssignment == nil {
 				continue
 			}
 			// We deepCopy the existing TopologyAssignment, so if we delete unwanted domain,
 			// And there is no fit, we have the original newAssignment to retry with
-			newAssignment, replacementAssignment, reason := s.findReplacementAssignment(&tr, psa.TopologyAssignment.DeepCopy(), wl, assumedUsage)
+			newAssignment, replacementAssignment, reason := s.findReplacementAssignment(&tr, psa.TopologyAssignment.DeepCopy(), opts.workload, assumedUsage)
 			result[tr.PodSet.Name] = tasPodSetAssignmentResult{TopologyAssignment: newAssignment, FailureReason: reason}
 			if reason != "" {
 				return result
 			}
 			addAssumedUsage(assumedUsage, replacementAssignment, &tr)
 		} else {
-			assignment, reason := s.findTopologyAssignment(tr, assumedUsage, simulateEmpty, "")
+			assignment, reason := s.findTopologyAssignment(tr, assumedUsage, opts.simulateEmpty, "")
 			result[tr.PodSet.Name] = tasPodSetAssignmentResult{TopologyAssignment: assignment, FailureReason: reason}
 			if reason != "" {
 				return result
diff --git a/pkg/scheduler/flavorassigner/flavorassigner.go b/pkg/scheduler/flavorassigner/flavorassigner.go
index 25a745710..e644e59da 100644
--- a/pkg/scheduler/flavorassigner/flavorassigner.go
+++ b/pkg/scheduler/flavorassigner/flavorassigner.go
@@ -467,7 +467,7 @@ func (a *FlavorAssigner) assignFlavors(log logr.Logger, counts []int32) Assignme
 	if features.Enabled(features.TopologyAwareScheduling) {
 		tasRequests := assignment.WorkloadsTopologyRequests(a.wl, a.cq)
 		if assignment.RepresentativeMode() == Fit {
-			result := a.cq.FindTopologyAssignmentsForWorkload(tasRequests, false, a.wl.Obj)
+			result := a.cq.FindTopologyAssignmentsForWorkload(tasRequests, cache.WithWorkload(a.wl.Obj))
 			if failure := result.Failure(); failure != nil {
 				// There is at least one PodSet which does not fit
 				psAssignment := assignment.podSetAssignmentByName(failure.PodSetName)
@@ -482,7 +482,7 @@ func (a *FlavorAssigner) assignFlavors(log logr.Logger, counts []int32) Assignme
 		}
 		if assignment.RepresentativeMode() == Preempt && !workload.HasNodeToReplace(a.wl.Obj) {
 			// Don't preempt other workloads if looking for a failed node replacement
-			result := a.cq.FindTopologyAssignmentsForWorkload(tasRequests, true, nil)
+			result := a.cq.FindTopologyAssignmentsForWorkload(tasRequests, cache.WithSimulateEmpty(true))
 			if failure := result.Failure(); failure != nil {
 				// There is at least one PodSet which does not fit even if
 				// all workloads are preempted.
diff --git a/pkg/scheduler/preemption/preemption.go b/pkg/scheduler/preemption/preemption.go
index bcaa0c87a..268ac9f9c 100644
--- a/pkg/scheduler/preemption/preemption.go
+++ b/pkg/scheduler/preemption/preemption.go
@@ -489,7 +489,7 @@ func workloadFits(preemptionCtx *preemptionCtx, allowBorrowing bool) bool {
 			return false
 		}
 	}
-	tasResult := preemptionCtx.preemptorCQ.FindTopologyAssignmentsForWorkload(preemptionCtx.tasRequests, false, nil)
+	tasResult := preemptionCtx.preemptorCQ.FindTopologyAssignmentsForWorkload(preemptionCtx.tasRequests)
 	return tasResult.Failure() == nil
 }
 
diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go
index 88399ea58..5347c6936 100644
--- a/pkg/scheduler/scheduler.go
+++ b/pkg/scheduler/scheduler.go
@@ -499,7 +499,7 @@ func updateAssignmentForTAS(cq *cache.ClusterQueueSnapshot, wl *workload.Info, a
 				targetWorkloads = append(targetWorkloads, target.WorkloadInfo)
 			}
 			revertUsage := cq.SimulateWorkloadRemoval(targetWorkloads)
-			tasResult = cq.FindTopologyAssignmentsForWorkload(tasRequests, false, nil)
+			tasResult = cq.FindTopologyAssignmentsForWorkload(tasRequests)
 			revertUsage()
 		} else {
 			// In this scenario we don't have any preemption candidates, yet we need
@@ -508,7 +508,7 @@ func updateAssignmentForTAS(cq *cache.ClusterQueueSnapshot, wl *workload.Info, a
 			// in the next scheduling cycle by the waiting workload. To obtain
 			// a TAS assignment for reserving the resources we run the algorithm
 			// assuming the cluster is empty.
-			tasResult = cq.FindTopologyAssignmentsForWorkload(tasRequests, true, nil)
+			tasResult = cq.FindTopologyAssignmentsForWorkload(tasRequests, cache.WithSimulateEmpty(true))
 		}
 		assignment.UpdateForTASResult(tasResult)
 	}



Test patch:
diff --git a/pkg/cache/tas_cache_test.go b/pkg/cache/tas_cache_test.go
index 0558323af..186d6e503 100644
--- a/pkg/cache/tas_cache_test.go
+++ b/pkg/cache/tas_cache_test.go
@@ -2863,7 +2863,7 @@ func TestFindTopologyAssignment(t *testing.T) {
 				wantMainPodSetResult.TopologyAssignment = tc.wantAssignment
 			}
 			wantResult[kueue.DefaultPodSetName] = wantMainPodSetResult
-			gotResult := snapshot.FindTopologyAssignmentsForFlavor(flavorTASRequests, false, nil)
+			gotResult := snapshot.FindTopologyAssignmentsForFlavor(flavorTASRequests)
 			if diff := cmp.Diff(wantResult, gotResult); diff != "" {
 				t.Errorf("unexpected topology assignment (-want,+got): %s", diff)
 			}
@@ -2957,7 +2957,7 @@ func TestFindTopologyAssignmentForTwoPodSets(t *testing.T) {
 		wantResult["podset1"] = buildWantedResult(wantAssignment1)
 		wantResult["podset2"] = buildWantedResult(wantAssignment2)
 
-		gotResult := snapshot.FindTopologyAssignmentsForFlavor(flavorTASRequests, false, nil)
+		gotResult := snapshot.FindTopologyAssignmentsForFlavor(flavorTASRequests)
 		if diff := cmp.Diff(wantResult, gotResult); diff != "" {
 			t.Errorf("unexpected topology assignment (-want,+got): %s", diff)
 		}