scheduler: bugfix: Make pod_scenario_builder build scenarios for rest of elastic job (#132)

* add ginkgo test to pod_scenario_builder, add elastic scenario

* add fix for pod_scenario_builder.go for elastic jobs

* add new test cases for ensuring correct number of potential victims is still returned

* Fails - e2e test to reclaim elastic job for a distributed job

* Another elastic reclaim e2e test

* update changelog

---------

Co-authored-by: davidLif <[email protected]>

Author: ArmedGuy

CHANGELOG.md

@@ -12,3 +12,6 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 ### Changed
 - Queue order function now takes into account potential victims, resulting in better reclaim scenarios.
+
+### Fixes
+- Fixed preempt/reclaim of elastic workloads only taking one pod.

pkg/scheduler/actions/common/solvers/pod_scenario_builder.go

@@ -77,6 +77,19 @@ func (asb *PodAccumulatedScenarioBuilder) GetNextScenario() *solverscenario.ByNo
 	// Jump over recorded victims in potential victims generation
 	for _, potentialVictimTask := range potentialVictimTasks {
 		if _, ok := asb.recordedVictimsTasks[potentialVictimTask.UID]; ok {
+			// If any of the tasks of the victim job are recorded victims
+			// we still want to evaluate the job again if there are tasks
+			// that are not recorded victims yet, like elastic jobs
+			var remainingTasks []*pod_info.PodInfo
+			for _, task := range nextVictimJob.PodInfos {
+				if _, ok := asb.recordedVictimsTasks[task.UID]; !ok {
+					remainingTasks = append(remainingTasks, task)
+				}
+			}
+			if len(remainingTasks) != 0 {
+				jobToPush := nextVictimJob.CloneWithTasks(remainingTasks)
+				asb.victimsJobsQueue.PushJob(jobToPush)
+			}
 			return asb.GetNextScenario()
 		}
 	}

pkg/scheduler/actions/common/solvers/pod_scenario_builder_test.go

@@ -6,6 +6,7 @@ package solvers
 import (
 	"fmt"
 	"strconv"
+	"testing"
 
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
@@ -108,7 +109,7 @@ var _ = Describe("PodAccumulatedScenarioBuilder", func() {
 	})
 
 	Context("with recorded victims", func() {
-		It("All scenarios have the same recorded victims", func() {
+		It("returns scenarios that have the same recorded victims", func() {
 			ssn, _ = initializeSession(3, 2)
 			for _, podGroupInfo := range ssn.PodGroupInfos {
 				podGroupInfo.MinAvailable = int32(len(podGroupInfo.PodInfos))
@@ -141,6 +142,126 @@ var _ = Describe("PodAccumulatedScenarioBuilder", func() {
 
 			Expect(numberOfGeneratedScenarios).To(Equal(2))
 		})
+
+		It("returns scenarios that have correct number of potential victims", func() {
+			ssn, _ = initializeSession(3, 2)
+			for _, podGroupInfo := range ssn.PodGroupInfos {
+				podGroupInfo.MinAvailable = int32(len(podGroupInfo.PodInfos))
+				podGroupInfo.PodGroup.Spec.MinMember = int32(len(podGroupInfo.PodInfos))
+			}
+			submitQueue := createQueue("team-a")
+			ssn.Queues[submitQueue.UID] = submitQueue
+			reclaimerJob, _ = createJobWithTasks(1, 1, "team-a", v1.PodPending)
+
+			var recordedVictimsJobs []*podgroup_info.PodGroupInfo
+			recordedVictimIndexes := []int{0, 2}
+			podGroupIndex := 0
+
+			for _, podGroupInfo := range ssn.PodGroupInfos {
+				if slices.Contains(recordedVictimIndexes, podGroupIndex) {
+					recordedVictimsJobs = append(recordedVictimsJobs, podGroupInfo)
+				}
+				podGroupIndex += 1
+			}
+
+			victimsQueue := utils.GetVictimsQueue(ssn, nil)
+
+			scenarioBuilder = NewPodAccumulatedScenarioBuilder(ssn, reclaimerJob, recordedVictimsJobs, victimsQueue)
+
+			numberOfGeneratedScenarios := 0
+			potentialVictimsPerScenario := []int{0, 2}
+			for sn := scenarioBuilder.GetCurrentScenario(); sn != nil; sn = scenarioBuilder.GetNextScenario() {
+				Expect(numberOfGeneratedScenarios < len(potentialVictimsPerScenario)).To(BeTrue())
+				Expect(len(sn.PotentialVictimsTasks())).To(Equal(potentialVictimsPerScenario[numberOfGeneratedScenarios]))
+				numberOfGeneratedScenarios += 1
+			}
+			Expect(numberOfGeneratedScenarios).To(Equal(len(potentialVictimsPerScenario)))
+		})
+	})
+
+	Context("with recorded victims that are elastic", func() {
+		It("returns scenarios that have the same recorded victims", func() {
+			// run 1 job with 3 tasks, set minAvailable to 1 for elastic
+			ssn, _ = initializeSession(1, 3)
+			minAvailable := 1
+			for _, podGroupInfo := range ssn.PodGroupInfos {
+				podGroupInfo.MinAvailable = int32(minAvailable)
+				podGroupInfo.PodGroup.Spec.MinMember = int32(minAvailable)
+			}
+			submitQueue := createQueue("team-a")
+			ssn.Queues[submitQueue.UID] = submitQueue
+			reclaimerJob, _ = createJobWithTasks(1, 2, "team-a", v1.PodPending)
+
+			var recordedVictimsJobs []*podgroup_info.PodGroupInfo
+
+			// Only the first pod group with the last task is recordedVictimJobs
+			for _, podGroupInfo := range ssn.PodGroupInfos {
+				var partialTasks []*pod_info.PodInfo
+				for _, podInfo := range podGroupInfo.PodInfos {
+					// use last pod as recorded victim as sorting will be reversed
+					if podInfo.Name == "pod-2" {
+						partialTasks = append(partialTasks, podInfo)
+					}
+				}
+				recordedVictimsJobs = append(recordedVictimsJobs, podGroupInfo.CloneWithTasks(partialTasks))
+				// we only want to change the first pod group, break after this
+				break
+			}
+
+			victimsQueue := utils.GetVictimsQueue(ssn, nil)
+
+			scenarioBuilder = NewPodAccumulatedScenarioBuilder(ssn, reclaimerJob, recordedVictimsJobs, victimsQueue)
+
+			numberOfGeneratedScenarios := 0
+			for sn := scenarioBuilder.GetCurrentScenario(); sn != nil; sn = scenarioBuilder.GetNextScenario() {
+				Expect(len(sn.RecordedVictimsJobs())).To(Equal(len(recordedVictimsJobs)))
+				numberOfGeneratedScenarios += 1
+			}
+
+			Expect(numberOfGeneratedScenarios).To(Equal(3))
+		})
+
+		It("returns scenarios that have correct number of potential victims", func() {
+			// run 1 job with 4 tasks, set minAvailable to 2 for elastic
+			ssn, _ = initializeSession(1, 4)
+			minAvailable := 2
+			for _, podGroupInfo := range ssn.PodGroupInfos {
+				podGroupInfo.MinAvailable = int32(minAvailable)
+				podGroupInfo.PodGroup.Spec.MinMember = int32(minAvailable)
+			}
+			submitQueue := createQueue("team-a")
+			ssn.Queues[submitQueue.UID] = submitQueue
+			reclaimerJob, _ = createJobWithTasks(1, 2, "team-a", v1.PodPending)
+
+			var recordedVictimsJobs []*podgroup_info.PodGroupInfo
+
+			// Only the first pod group with the last task is recordedVictimJobs
+			for _, podGroupInfo := range ssn.PodGroupInfos {
+				var partialTasks []*pod_info.PodInfo
+				for _, podInfo := range podGroupInfo.PodInfos {
+					// use last pod as recorded victim as sorting will be reversed
+					if podInfo.Name == "pod-3" {
+						partialTasks = append(partialTasks, podInfo)
+					}
+				}
+				recordedVictimsJobs = append(recordedVictimsJobs, podGroupInfo.CloneWithTasks(partialTasks))
+				// we only want to change the first pod group, break after this
+				break
+			}
+
+			victimsQueue := utils.GetVictimsQueue(ssn, nil)
+
+			scenarioBuilder = NewPodAccumulatedScenarioBuilder(ssn, reclaimerJob, recordedVictimsJobs, victimsQueue)
+
+			numberOfGeneratedScenarios := 0
+			potentialVictimsPerScenario := []int{0, 1, 3}
+			for sn := scenarioBuilder.GetCurrentScenario(); sn != nil; sn = scenarioBuilder.GetNextScenario() {
+				Expect(numberOfGeneratedScenarios < len(potentialVictimsPerScenario)).To(BeTrue())
+				Expect(len(sn.PotentialVictimsTasks())).To(Equal(potentialVictimsPerScenario[numberOfGeneratedScenarios]))
+				numberOfGeneratedScenarios += 1
+			}
+			Expect(numberOfGeneratedScenarios).To(Equal(len(potentialVictimsPerScenario)))
+		})
 	})
 })
 
@@ -153,7 +274,6 @@ func initializeSession(jobsCount, tasksPerJob int) (*framework.Session, []*pod_i
 	queues := []*queue_info.QueueInfo{defaultQueue}
 
 	for jobID := 0; jobID < jobsCount; jobID++ {
-		jobTasks := []*pod_info.PodInfo{}
 		queueName := fmt.Sprintf("team-%d", jobID)
 		newJob, jobTasks := createJobWithTasks(tasksPerJob, jobID, queueName, v1.PodRunning)
 		jobs = append(jobs, newJob)
@@ -278,3 +398,8 @@ func requireOneGPU() v1.ResourceRequirements {
 		},
 	}
 }
+
+func TestScenarioSolvers(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "Scenario Solvers Suite")
+}

test/e2e/suites/reclaim/reclaim_elastic_test.go

@@ -145,4 +145,86 @@ var _ = Describe("Reclaim with Elastic Jobs", Ordered, func() {
 			return len(pods.Items) == 0
 		})
 	})
+
+	It("Reclaim elastic job for a distributed job", func(ctx context.Context) {
+		testCtx = testcontext.GetConnectivity(ctx, Default)
+		parentQueue, reclaimeeQueue, reclaimerQueue = createQueues(2, 0, 2)
+		reclaimeeQueue.Spec.Resources.GPU.OverQuotaWeight = 0
+		testCtx.InitQueues([]*v2.Queue{parentQueue, reclaimeeQueue, reclaimerQueue})
+		reclaimeeNamespace = queue.GetConnectedNamespaceToQueue(reclaimeeQueue)
+
+		// reclaimee job
+		reclaimeePodRequirements := v1.ResourceRequirements{
+			Limits: map[v1.ResourceName]resource.Quantity{
+				constants.GpuResource: resource.MustParse("1"),
+			},
+		}
+		reclaimeePodGroup, reclaimeePods := pod_group.CreateWithPods(ctx, testCtx.KubeClientset, testCtx.KubeAiSchedClientset,
+			"elastic-reclaimee-job", reclaimeeQueue, 2, nil,
+			reclaimeePodRequirements)
+		wait.ForPodsScheduled(ctx, testCtx.ControllerClient, reclaimeeNamespace, reclaimeePods)
+
+		// reclaimer job
+		reclaimerPodRequirements := v1.ResourceRequirements{
+			Limits: map[v1.ResourceName]resource.Quantity{
+				constants.GpuResource: resource.MustParse("1"),
+			},
+		}
+		_, reclaimerPods := pod_group.CreateDistributedJob(
+			ctx, testCtx.KubeClientset, testCtx.ControllerClient,
+			reclaimerQueue, 2, reclaimerPodRequirements, "",
+		)
+		reclaimerNamespace := queue.GetConnectedNamespaceToQueue(reclaimerQueue)
+		wait.ForPodsScheduled(ctx, testCtx.ControllerClient, reclaimerNamespace, reclaimerPods)
+
+		// verify results
+		wait.ForPodsWithCondition(ctx, testCtx.ControllerClient, func(watch.Event) bool {
+			pods, err := testCtx.KubeClientset.CoreV1().Pods(reclaimeeNamespace).List(ctx, metav1.ListOptions{
+				LabelSelector: fmt.Sprintf("%s=%s", podGroupLabelName, reclaimeePodGroup.Name),
+			})
+			Expect(err).To(Succeed())
+			return len(pods.Items) == 0
+		})
+	})
+
+	It("Reclaim elastic job partially for a distributed job", func(ctx context.Context) {
+		testCtx = testcontext.GetConnectivity(ctx, Default)
+		parentQueue, reclaimeeQueue, reclaimerQueue = createQueues(3, 1, 2)
+		reclaimeeQueue.Spec.Resources.GPU.OverQuotaWeight = 0
+		testCtx.InitQueues([]*v2.Queue{parentQueue, reclaimeeQueue, reclaimerQueue})
+		reclaimeeNamespace = queue.GetConnectedNamespaceToQueue(reclaimeeQueue)
+
+		// reclaimee job
+		reclaimeePodRequirements := v1.ResourceRequirements{
+			Limits: map[v1.ResourceName]resource.Quantity{
+				constants.GpuResource: resource.MustParse("1"),
+			},
+		}
+		reclaimeePodGroup, reclaimeePods := pod_group.CreateWithPods(ctx, testCtx.KubeClientset, testCtx.KubeAiSchedClientset,
+			"elastic-reclaimee-job", reclaimeeQueue, 3, nil,
+			reclaimeePodRequirements)
+		wait.ForPodsScheduled(ctx, testCtx.ControllerClient, reclaimeeNamespace, reclaimeePods)
+
+		// reclaimer job
+		reclaimerPodRequirements := v1.ResourceRequirements{
+			Limits: map[v1.ResourceName]resource.Quantity{
+				constants.GpuResource: resource.MustParse("1"),
+			},
+		}
+		_, reclaimerPods := pod_group.CreateDistributedJob(
+			ctx, testCtx.KubeClientset, testCtx.ControllerClient,
+			reclaimerQueue, 2, reclaimerPodRequirements, "",
+		)
+		reclaimerNamespace := queue.GetConnectedNamespaceToQueue(reclaimerQueue)
+		wait.ForPodsScheduled(ctx, testCtx.ControllerClient, reclaimerNamespace, reclaimerPods)
+
+		// verify results
+		wait.ForPodsWithCondition(ctx, testCtx.ControllerClient, func(watch.Event) bool {
+			pods, err := testCtx.KubeClientset.CoreV1().Pods(reclaimeeNamespace).List(ctx, metav1.ListOptions{
+				LabelSelector: fmt.Sprintf("%s=%s", podGroupLabelName, reclaimeePodGroup.Name),
+			})
+			Expect(err).To(Succeed())
+			return len(pods.Items) == 1
+		})
+	})
 })