diff --git a/.changes/unreleased/Fixed-20260728-181225.yaml b/.changes/unreleased/Fixed-20260728-181225.yaml new file mode 100644 index 000000000..f2e28c9dc --- /dev/null +++ b/.changes/unreleased/Fixed-20260728-181225.yaml @@ -0,0 +1,6 @@ +kind: Fixed +body: stalegangeviction no longer evicts remaining pods of a gang whose pods complete successfully +time: 2026-07-28T18:12:25.0211683+05:30 +custom: + Author: Thezone-1 + Issue: "1968" diff --git a/pkg/scheduler/actions/stalegangeviction/stalegangeviction_test.go b/pkg/scheduler/actions/stalegangeviction/stalegangeviction_test.go index 889b91d33..6ee66469d 100644 --- a/pkg/scheduler/actions/stalegangeviction/stalegangeviction_test.go +++ b/pkg/scheduler/actions/stalegangeviction/stalegangeviction_test.go @@ -709,6 +709,201 @@ func TestStaleGangEviction(t *testing.T) { }, }, }, + { + name: "Gang with a succeeded pod - no evict", + topology: test_utils.TestTopologyBasic{ + Jobs: []*jobs_fake.TestJobBasic{ + { + Name: "job-1", + QueueName: "q-1", + RootSubGroupSet: jobs_fake.DefaultSubGroup(3), + Tasks: []*tasks_fake.TestTaskBasic{ + { + Name: "job-1-0", + State: pod_status.Succeeded, + }, + { + Name: "job-1-1", + State: pod_status.Running, + NodeName: "node-1", + }, + { + Name: "job-1-2", + State: pod_status.Running, + NodeName: "node-1", + }, + }, + StaleDuration: pointer.Duration(61 * time.Second), + }, + }, + Nodes: map[string]nodes_fake.TestNodeBasic{ + "node-1": {}, + }, + Queues: []test_utils.TestQueueBasic{ + { + Name: "q-1", + ParentQueue: "d-1", + }, + }, + Departments: []test_utils.TestDepartmentBasic{ + { + Name: "d-1", + }, + }, + TaskExpectedResults: map[string]test_utils.TestExpectedResultBasic{ + "job-1-0": { + Status: pod_status.Succeeded, + }, + "job-1-1": { + NodeName: "node-1", + Status: pod_status.Running, + }, + "job-1-2": { + NodeName: "node-1", + Status: pod_status.Running, + }, + }, + Mocks: &test_utils.TestMock{ + CacheRequirements: &test_utils.CacheMocking{ + NumberOfCacheBinds: 0, + NumberOfCacheEvictions: 0, + NumberOfPipelineActions: 0, + }, + }, + }, + }, + { + name: "Gang with only one pod left running - no evict", + topology: test_utils.TestTopologyBasic{ + Jobs: []*jobs_fake.TestJobBasic{ + { + Name: "job-1", + QueueName: "q-1", + RootSubGroupSet: jobs_fake.DefaultSubGroup(3), + Tasks: []*tasks_fake.TestTaskBasic{ + { + Name: "job-1-0", + State: pod_status.Succeeded, + }, + { + Name: "job-1-1", + State: pod_status.Succeeded, + }, + { + Name: "job-1-2", + State: pod_status.Running, + NodeName: "node-1", + }, + }, + StaleDuration: pointer.Duration(61 * time.Second), + }, + }, + Nodes: map[string]nodes_fake.TestNodeBasic{ + "node-1": {}, + }, + Queues: []test_utils.TestQueueBasic{ + { + Name: "q-1", + ParentQueue: "d-1", + }, + }, + Departments: []test_utils.TestDepartmentBasic{ + { + Name: "d-1", + }, + }, + TaskExpectedResults: map[string]test_utils.TestExpectedResultBasic{ + "job-1-0": { + Status: pod_status.Succeeded, + }, + "job-1-1": { + Status: pod_status.Succeeded, + }, + "job-1-2": { + NodeName: "node-1", + Status: pod_status.Running, + }, + }, + Mocks: &test_utils.TestMock{ + CacheRequirements: &test_utils.CacheMocking{ + NumberOfCacheBinds: 0, + NumberOfCacheEvictions: 0, + NumberOfPipelineActions: 0, + }, + }, + }, + }, + { + name: "Gang with a succeeded pod in one sub group - no evict", + topology: test_utils.TestTopologyBasic{ + Jobs: []*jobs_fake.TestJobBasic{ + { + Name: "job-1", + QueueName: "q-1", + RootSubGroupSet: func() *subgroup_info.SubGroupSet { + root := subgroup_info.NewSubGroupSet(subgroup_info.RootSubGroupSetName, nil) + root.AddPodSet(subgroup_info.NewPodSet("sub-group-0", 2, nil)) + root.AddPodSet(subgroup_info.NewPodSet("sub-group-1", 1, nil)) + return root + }(), + Tasks: []*tasks_fake.TestTaskBasic{ + { + Name: "job-1-0", + SubGroupName: "sub-group-0", + State: pod_status.Succeeded, + }, + { + Name: "job-1-1", + SubGroupName: "sub-group-0", + State: pod_status.Running, + NodeName: "node-1", + }, + { + Name: "job-1-2", + SubGroupName: "sub-group-1", + State: pod_status.Running, + NodeName: "node-1", + }, + }, + StaleDuration: pointer.Duration(61 * time.Second), + }, + }, + Nodes: map[string]nodes_fake.TestNodeBasic{ + "node-1": {}, + }, + Queues: []test_utils.TestQueueBasic{ + { + Name: "q-1", + ParentQueue: "d-1", + }, + }, + Departments: []test_utils.TestDepartmentBasic{ + { + Name: "d-1", + }, + }, + TaskExpectedResults: map[string]test_utils.TestExpectedResultBasic{ + "job-1-0": { + Status: pod_status.Succeeded, + }, + "job-1-1": { + NodeName: "node-1", + Status: pod_status.Running, + }, + "job-1-2": { + NodeName: "node-1", + Status: pod_status.Running, + }, + }, + Mocks: &test_utils.TestMock{ + CacheRequirements: &test_utils.CacheMocking{ + NumberOfCacheBinds: 0, + NumberOfCacheEvictions: 0, + NumberOfPipelineActions: 0, + }, + }, + }, + }, } { t.Run(test.name, func(t *testing.T) { t.Logf("Running test number: %v, test name: %v,", i, test.name) diff --git a/pkg/scheduler/cache/cache.go b/pkg/scheduler/cache/cache.go index 633aeaaf8..082f71dc4 100644 --- a/pkg/scheduler/cache/cache.go +++ b/pkg/scheduler/cache/cache.go @@ -83,9 +83,13 @@ var terminalPodPhases = []v1.PodPhase{ v1.PodFailed, } -func filterTerminalPods(options *metav1.ListOptions) { - selectors := make([]string, 0, len(terminalPodPhases)) - for _, phase := range terminalPodPhases { +var watchFilteredPodPhases = []v1.PodPhase{ + v1.PodFailed, +} + +func filterFailedPods(options *metav1.ListOptions) { + selectors := make([]string, 0, len(watchFilteredPodPhases)) + for _, phase := range watchFilteredPodPhases { selectors = append(selectors, fmt.Sprintf("status.phase!=%s", phase)) } selector := strings.Join(selectors, ",") @@ -103,7 +107,7 @@ func registerSchedulerPodInformer(informerFactory informers.SharedInformerFactor metav1.NamespaceAll, resyncPeriod, k8scache.Indexers{k8scache.NamespaceIndex: k8scache.MetaNamespaceIndexFunc}, - filterTerminalPods, + filterFailedPods, ) }) } diff --git a/pkg/scheduler/cache/cache_test.go b/pkg/scheduler/cache/cache_test.go index 9c6ee8a19..5853ce511 100644 --- a/pkg/scheduler/cache/cache_test.go +++ b/pkg/scheduler/cache/cache_test.go @@ -58,7 +58,7 @@ func TestCache(t *testing.T) { var _ = Describe("Cache", func() { Describe("New", func() { Context("Pod informer filtering", func() { - It("should filter terminal pods without filtering pods by scheduler name", func() { + It("should filter failed pods while keeping succeeded pods, without filtering by scheduler name", func() { kubeClient := fake.NewSimpleClientset() cache := New(&SchedulerCacheParams{ KubeClient: kubeClient, @@ -95,8 +95,8 @@ var _ = Describe("Cache", func() { Expect(podSelectors).NotTo(BeEmpty()) for _, selector := range podSelectors { - Expect(selector).To(ContainSubstring("status.phase!=Succeeded")) Expect(selector).To(ContainSubstring("status.phase!=Failed")) + Expect(selector).NotTo(ContainSubstring("status.phase!=Succeeded")) Expect(selector).NotTo(ContainSubstring("spec.schedulerName")) } Expect(nonPodSelectors).To(BeEmpty()) diff --git a/pkg/scheduler/cache/pod_transform.go b/pkg/scheduler/cache/pod_transform.go index 361607698..f3fe64d58 100644 --- a/pkg/scheduler/cache/pod_transform.go +++ b/pkg/scheduler/cache/pod_transform.go @@ -5,7 +5,10 @@ package cache import ( v1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/client-go/tools/cache" + + commonconstants "github.com/kai-scheduler/KAI-scheduler/pkg/common/constants" ) func setSchedulerPodTransform(informer cache.SharedIndexInformer) error { @@ -18,6 +21,10 @@ func compactSchedulerPod(obj any) (any, error) { return obj, nil } + if pod.Status.Phase == v1.PodSucceeded { + return compactSucceededPod(pod), nil + } + compact := pod.DeepCopy() compact.ManagedFields = nil compact.Spec.Containers = compactContainers(compact.Spec.Containers) @@ -26,6 +33,28 @@ func compactSchedulerPod(obj any) (any, error) { return compact, nil } +func compactSucceededPod(pod *v1.Pod) *v1.Pod { + compact := &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: pod.Name, + Namespace: pod.Namespace, + UID: pod.UID, + }, + Status: v1.PodStatus{ + Phase: pod.Status.Phase, + }, + } + + if podGroup, found := pod.Annotations[commonconstants.PodGroupAnnotationForPod]; found { + compact.Annotations = map[string]string{commonconstants.PodGroupAnnotationForPod: podGroup} + } + if subGroup, found := pod.Labels[commonconstants.SubGroupLabelKey]; found { + compact.Labels = map[string]string{commonconstants.SubGroupLabelKey: subGroup} + } + + return compact +} + func compactContainers(containers []v1.Container) []v1.Container { compact := make([]v1.Container, 0, len(containers)) for _, container := range containers { diff --git a/pkg/scheduler/cache/pod_transform_test.go b/pkg/scheduler/cache/pod_transform_test.go new file mode 100644 index 000000000..6a9966044 --- /dev/null +++ b/pkg/scheduler/cache/pod_transform_test.go @@ -0,0 +1,152 @@ +// Copyright 2025 NVIDIA CORPORATION +// SPDX-License-Identifier: Apache-2.0 + +package cache + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + v1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/api/resource" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + + commonconstants "github.com/kai-scheduler/KAI-scheduler/pkg/common/constants" + "github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_info" + "github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_status" + "github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/podgroup_info" + "github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/resource_info" +) + +var _ = Describe("compactSchedulerPod", func() { + succeededPod := func() *v1.Pod { + return &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "worker-0", + Namespace: "team-a", + UID: "worker-0-uid", + Annotations: map[string]string{ + commonconstants.PodGroupAnnotationForPod: "pg-1", + "kubectl.kubernetes.io/last-applied-configuration": "{...}", + }, + Labels: map[string]string{ + commonconstants.SubGroupLabelKey: "workers", + "app": "trainer", + }, + ManagedFields: []metav1.ManagedFieldsEntry{{Manager: "kubelet"}}, + }, + Spec: v1.PodSpec{ + NodeName: "node-1", + NodeSelector: map[string]string{"gpu": "true"}, + Tolerations: []v1.Toleration{{Key: "nvidia.com/gpu"}}, + Volumes: []v1.Volume{{Name: "data"}}, + Containers: []v1.Container{{ + Name: "main", + Resources: v1.ResourceRequirements{ + Requests: v1.ResourceList{v1.ResourceCPU: resource.MustParse("1")}, + }, + }}, + }, + Status: v1.PodStatus{ + Phase: v1.PodSucceeded, + PodIP: "10.0.0.1", + Conditions: []v1.PodCondition{{Type: v1.PodReady}}, + ContainerStatuses: []v1.ContainerStatus{{Name: "main"}}, + }, + } + } + + Context("succeeded pods", func() { + It("retains identity, podgroup, subgroup and phase", func() { + transformed, err := compactSchedulerPod(succeededPod()) + Expect(err).NotTo(HaveOccurred()) + + compact, ok := transformed.(*v1.Pod) + Expect(ok).To(BeTrue()) + Expect(compact.Name).To(Equal("worker-0")) + Expect(compact.Namespace).To(Equal("team-a")) + Expect(compact.UID).To(BeEquivalentTo("worker-0-uid")) + Expect(compact.Annotations).To(HaveKeyWithValue(commonconstants.PodGroupAnnotationForPod, "pg-1")) + Expect(compact.Labels).To(HaveKeyWithValue(commonconstants.SubGroupLabelKey, "workers")) + Expect(compact.Status.Phase).To(Equal(v1.PodSucceeded)) + }) + + It("drops spec, status and metadata not read for succeeded pods", func() { + transformed, err := compactSchedulerPod(succeededPod()) + Expect(err).NotTo(HaveOccurred()) + + compact := transformed.(*v1.Pod) + Expect(compact.Annotations).To(HaveLen(1)) + Expect(compact.Labels).To(HaveLen(1)) + Expect(compact.ManagedFields).To(BeNil()) + Expect(compact.Spec).To(Equal(v1.PodSpec{})) + Expect(compact.Status.PodIP).To(BeEmpty()) + Expect(compact.Status.Conditions).To(BeNil()) + Expect(compact.Status.ContainerStatuses).To(BeNil()) + }) + + It("leaves the informer's copy untouched", func() { + pod := succeededPod() + _, err := compactSchedulerPod(pod) + Expect(err).NotTo(HaveOccurred()) + + Expect(pod.Spec.Containers).To(HaveLen(1)) + Expect(pod.Annotations).To(HaveLen(2)) + }) + + It("keeps enough for the podgroup of a completing gang to stay non-stale", func() { + vectorMap := resource_info.NewResourceVectorMap() + taskInfo := func(pod *v1.Pod) *pod_info.PodInfo { + delete(pod.Labels, commonconstants.SubGroupLabelKey) + transformed, err := compactSchedulerPod(pod) + Expect(err).NotTo(HaveOccurred()) + return pod_info.NewTaskInfo(transformed.(*v1.Pod), vectorMap) + } + + succeeded := taskInfo(succeededPod()) + Expect(succeeded.Status).To(Equal(pod_status.Succeeded)) + Expect(succeeded.Job).To(BeEquivalentTo("pg-1")) + + running := succeededPod() + running.Name = "worker-1" + running.UID = "worker-1-uid" + running.Status.Phase = v1.PodRunning + + podGroup := podgroup_info.NewPodGroupInfo("pg-1", succeeded, taskInfo(running)) + podGroup.GetAllPodSets()[podgroup_info.DefaultSubGroup].SetMinAvailable(2) + + Expect(podGroup.IsStale()).To(BeFalse()) + }) + + It("omits the podgroup annotation and subgroup label when the pod has neither", func() { + pod := succeededPod() + pod.Annotations = nil + pod.Labels = nil + + transformed, err := compactSchedulerPod(pod) + Expect(err).NotTo(HaveOccurred()) + + compact := transformed.(*v1.Pod) + Expect(compact.Annotations).To(BeNil()) + Expect(compact.Labels).To(BeNil()) + }) + }) + + Context("non-succeeded pods", func() { + It("keeps the fields scheduling depends on", func() { + pod := succeededPod() + pod.Status.Phase = v1.PodRunning + + transformed, err := compactSchedulerPod(pod) + Expect(err).NotTo(HaveOccurred()) + + compact := transformed.(*v1.Pod) + Expect(compact.Spec.NodeName).To(Equal("node-1")) + Expect(compact.Spec.NodeSelector).To(HaveKeyWithValue("gpu", "true")) + Expect(compact.Spec.Tolerations).To(HaveLen(1)) + Expect(compact.Spec.Containers).To(HaveLen(1)) + Expect(compact.Spec.Containers[0].Resources.Requests).To(HaveKey(v1.ResourceCPU)) + Expect(compact.Annotations).To(HaveLen(2)) + Expect(compact.ManagedFields).To(BeNil()) + }) + }) +})