Elastic reclaim panic (#390)

* Fix panic in elastic reclaim scenarios

---------

Co-authored-by: Erez Freiberger <[email protected]>

Author: itsomri

CHANGELOG.md

@@ -6,6 +6,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 ## [Unreleased]
 
+### Fixed
+- Fixed scheduler panic in some elastic reclaim scenarios
+
 ## [v0.8.0] - 2025-08-18
 
 ### Added

pkg/scheduler/plugins/proportion/proportion.go

@@ -155,25 +155,46 @@ func (pp *proportionPlugin) reclaimableFn(
 
 func (pp *proportionPlugin) getVictimResources(victim *api.VictimInfo) []*resource_info.Resource {
 	var victimResources []*resource_info.Resource
-	if len(victim.Tasks) > int(victim.Job.GetDefaultMinAvailable()) {
-		elasticTasks := victim.Tasks[victim.Job.GetDefaultMinAvailable():]
-		for _, task := range elasticTasks {
-			resources := getResources(pp.allowConsolidatingReclaim, task)
-			if resources == nil {
-				continue
-			}
-			victimResources = append(victimResources, resources)
+
+	elasticTasks, coreTasks := splitVictimTasks(victim.Tasks, victim.Job.GetDefaultMinAvailable())
+
+	// Process elastic tasks individually
+	for _, task := range elasticTasks {
+		resources := getResources(pp.allowConsolidatingReclaim, task)
+		if resources == nil {
+			continue
 		}
+		victimResources = append(victimResources, resources)
 	}
 
-	resources := getResources(pp.allowConsolidatingReclaim, victim.Tasks[:victim.Job.GetDefaultMinAvailable()]...)
+	// Process core tasks as a group
+	resources := getResources(pp.allowConsolidatingReclaim, coreTasks...)
 	if resources != nil {
 		victimResources = append(victimResources, resources)
 	}
 
 	return victimResources
 }
 
+// splitVictimTasks safely splits victim tasks into elastic and core tasks
+// Returns (elasticTasks, coreTasks)
+func splitVictimTasks(tasks []*pod_info.PodInfo, minAvailable int32) ([]*pod_info.PodInfo, []*pod_info.PodInfo) {
+	totalTasks := len(tasks)
+	minAvailableInt := int(minAvailable)
+
+	// Handle case where minAvailable is greater than or equal to the number of tasks
+	if minAvailableInt >= totalTasks {
+		// All tasks are considered core tasks, no elastic tasks
+		return nil, tasks
+	}
+
+	// Normal case: split tasks into elastic and core
+	elasticTasks := tasks[minAvailableInt:]
+	coreTasks := tasks[:minAvailableInt]
+
+	return elasticTasks, coreTasks
+}
+
 func getResources(ignoreReallocatedTasks bool, pods ...*pod_info.PodInfo) *resource_info.Resource {
 	resources := make([]*resource_info.ResourceRequirements, 0, len(pods))
 	for _, task := range pods {

pkg/scheduler/plugins/proportion/proportion_test.go

@@ -22,6 +22,7 @@ import (
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/node_info"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/pod_info"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/pod_status"
+	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/podgroup_info"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/resource_info"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/cache"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/conf"
@@ -773,6 +774,128 @@ var _ = Describe("Set Fair Share in Proportion", func() {
 		}
 
 	})
+
+	Context("getVictimResources", func() {
+		It("should handle case where MinAvailable is greater than number of tasks (panic fix)", func() {
+			plugin := &proportionPlugin{
+				allowConsolidatingReclaim: true,
+			}
+
+			// Create a victim with only 1 task but MinAvailable = 2
+			// This should cause a slice bounds panic without the fix
+			victim := &api.VictimInfo{
+				Job: &podgroup_info.PodGroupInfo{},
+				Tasks: []*pod_info.PodInfo{
+					{
+						Status:           pod_status.Pending,
+						AcceptedResource: common_info.BuildResourceRequirements("1", "1Gi"),
+					},
+				},
+			}
+			victim.Job.SetDefaultMinAvailable(2)
+
+			// This should not panic
+			result := plugin.getVictimResources(victim)
+			// Should return resources for the single task that exists
+			Expect(len(result)).To(Equal(1))
+			Expect(result[0]).ToNot(BeNil())
+			Expect(result[0].Cpu()).To(Equal(1000.0))
+		})
+
+		It("should correctly split elastic and core tasks when MinAvailable is less than task count", func() {
+			plugin := &proportionPlugin{
+				allowConsolidatingReclaim: true,
+			}
+
+			// Create a victim with 3 tasks but MinAvailable = 1
+			victim := &api.VictimInfo{
+				Job: &podgroup_info.PodGroupInfo{},
+				Tasks: []*pod_info.PodInfo{
+					{
+						Status:           pod_status.Pending,
+						AcceptedResource: common_info.BuildResourceRequirements("1", "1Gi"),
+					},
+					{
+						Status:           pod_status.Pending,
+						AcceptedResource: common_info.BuildResourceRequirements("1", "1Gi"),
+					},
+					{
+						Status:           pod_status.Pending,
+						AcceptedResource: common_info.BuildResourceRequirements("1", "1Gi"),
+					},
+				},
+			}
+			victim.Job.SetDefaultMinAvailable(1)
+
+			result := plugin.getVictimResources(victim)
+
+			// Should return 3 resources: 2 elastic tasks + 1 core task group
+			Expect(len(result)).To(Equal(3))
+			for _, res := range result {
+				Expect(res).ToNot(BeNil())
+				Expect(res.Cpu()).To(Equal(1000.0))
+			}
+		})
+
+		It("should handle case where MinAvailable equals task count", func() {
+			plugin := &proportionPlugin{
+				allowConsolidatingReclaim: true,
+			}
+
+			// Create a victim with 2 tasks and MinAvailable = 2
+			victim := &api.VictimInfo{
+				Job: &podgroup_info.PodGroupInfo{},
+				Tasks: []*pod_info.PodInfo{
+					{
+						Status:           pod_status.Pending,
+						AcceptedResource: common_info.BuildResourceRequirements("1", "1Gi"),
+					},
+					{
+						Status:           pod_status.Pending,
+						AcceptedResource: common_info.BuildResourceRequirements("1", "1Gi"),
+					},
+				},
+			}
+			victim.Job.SetDefaultMinAvailable(2)
+
+			result := plugin.getVictimResources(victim)
+
+			// Should return 1 resource for all core tasks (no elastic tasks)
+			Expect(len(result)).To(Equal(1))
+			Expect(result[0]).ToNot(BeNil())
+			Expect(result[0].Cpu()).To(Equal(2000.0)) // Combined resources
+		})
+
+		It("should handle zero MinAvailable", func() {
+			plugin := &proportionPlugin{
+				allowConsolidatingReclaim: true,
+			}
+
+			victim := &api.VictimInfo{
+				Job: &podgroup_info.PodGroupInfo{},
+				Tasks: []*pod_info.PodInfo{
+					{
+						Status:           pod_status.Pending,
+						AcceptedResource: common_info.BuildResourceRequirements("1", "1Gi"),
+					},
+					{
+						Status:           pod_status.Pending,
+						AcceptedResource: common_info.BuildResourceRequirements("1", "1Gi"),
+					},
+				},
+			}
+			victim.Job.SetDefaultMinAvailable(0)
+
+			result := plugin.getVictimResources(victim)
+
+			// Should return 2 resources (each task individually as elastic)
+			Expect(len(result)).To(Equal(2))
+			for _, res := range result {
+				Expect(res).ToNot(BeNil())
+				Expect(res.Cpu()).To(Equal(1000.0))
+			}
+		})
+	})
 })
 
 var _ = Describe("New", func() {

test/e2e/suites/reclaim/reclaim_elastic_test.go

@@ -14,7 +14,9 @@ import (
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/apimachinery/pkg/watch"
+	"sigs.k8s.io/controller-runtime/pkg/client"
 
 	v2 "github.com/NVIDIA/KAI-scheduler/pkg/apis/scheduling/v2"
 	"github.com/NVIDIA/KAI-scheduler/pkg/common/constants"
@@ -190,7 +192,7 @@ var _ = Describe("Reclaim with Elastic Jobs", Ordered, func() {
 
 	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)
+		parentQueue, reclaimeeQueue, reclaimerQueue = createQueues(4, 2, 2)
 		reclaimeeQueue.Spec.Resources.GPU.OverQuotaWeight = 0
 		testCtx.InitQueues([]*v2.Queue{parentQueue, reclaimeeQueue, reclaimerQueue})
 		reclaimeeNamespace = queue.GetConnectedNamespaceToQueue(reclaimeeQueue)
@@ -204,6 +206,8 @@ var _ = Describe("Reclaim with Elastic Jobs", Ordered, func() {
 		reclaimeePodGroup, reclaimeePods := pod_group.CreateWithPods(ctx, testCtx.KubeClientset, testCtx.KubeAiSchedClientset,
 			"elastic-reclaimee-job", reclaimeeQueue, 3, nil,
 			reclaimeePodRequirements)
+		Expect(testCtx.ControllerClient.Patch(
+			ctx, reclaimeePodGroup, client.RawPatch(types.JSONPatchType, []byte(`[{"op": "replace", "path": "/spec/minMember", "value": 2}]`)))).To(Succeed())
 		wait.ForPodsScheduled(ctx, testCtx.ControllerClient, reclaimeeNamespace, reclaimeePods)
 
 		// reclaimer job
@@ -225,7 +229,7 @@ var _ = Describe("Reclaim with Elastic Jobs", Ordered, func() {
 				LabelSelector: fmt.Sprintf("%s=%s", podGroupLabelName, reclaimeePodGroup.Name),
 			})
 			Expect(err).To(Succeed())
-			return len(pods.Items) == 1
+			return len(pods.Items) == 2
 		})
 	})
 	It("Reclaim elastic job with min runtime protecting", func(ctx context.Context) {