scheduler: Add LastStartTimestamp to PodGroup (#153)

* add podgroup starttime annotation

* call it LastStartTimeStamp instead

* set LastStartTimestamp after successful commit on allocate

* unit test to ensure start time is set

* add proper unit test for lastStartTime validation for elastic

Author: ArmedGuy

CHANGELOG.md

@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 ### Added
 - Added support for [k8s pod scheduling gates](https://kubernetes.io/docs/concepts/scheduling-eviction/pod-scheduling-readiness/)
 - nodeSelector, affinity and tolerations configurable with global value definitions
+- Scheduler now adds a "LastStartTimestamp" to podgroup on allocation
 
 ### Changed
 - Queue order function now takes into account potential victims, resulting in better reclaim scenarios.

pkg/common/constants/constants.go

@@ -23,6 +23,7 @@ const (
 	RunaiGpuLimit            = "runai-gpu-limit"
 	MpsAnnotation            = "mps"
 	StalePodgroupTimeStamp   = "kai.scheduler/stale-podgroup-timestamp"
+	LastStartTimeStamp       = "kai.scheduler/last-start-timestamp"
 
 	// Labels
 	NodePoolNameLabel        = "runai/node-pool"

pkg/scheduler/actions/allocate/allocate.go

@@ -4,6 +4,8 @@
 package allocate
 
 import (
+	"time"
+
 	"golang.org/x/exp/maps"
 
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/actions/common"
@@ -41,19 +43,24 @@ func (alloc *allocateAction) Execute(ssn *framework.Session) {
 	for !jobsOrderByQueues.IsEmpty() {
 		job := jobsOrderByQueues.PopNextJob()
 		stmt := ssn.Statement()
-		if attemptToAllocateJob(ssn, stmt, job) {
+		alreadyAllocated := job.GetNumAllocatedTasks() > 0
+		if ok, pipelined := attemptToAllocateJob(ssn, stmt, job); ok {
 			metrics.IncPodgroupScheduledByAction()
 			err := stmt.Commit()
+			if err == nil && !pipelined && !alreadyAllocated {
+				setLastStartTimestamp(job)
+			}
 			if err == nil && podgroup_info.HasTasksToAllocate(job, true) {
 				jobsOrderByQueues.PushJob(job)
+				continue
 			}
 		} else {
 			stmt.Discard()
 		}
 	}
 }
 
-func attemptToAllocateJob(ssn *framework.Session, stmt *framework.Statement, job *podgroup_info.PodGroupInfo) bool {
+func attemptToAllocateJob(ssn *framework.Session, stmt *framework.Statement, job *podgroup_info.PodGroupInfo) (allocated, pipelined bool) {
 	queue := ssn.Queues[job.Queue]
 
 	resReq := podgroup_info.GetTasksToAllocateInitResource(job, ssn.TaskOrderFn, true)
@@ -64,9 +71,9 @@ func attemptToAllocateJob(ssn *framework.Session, stmt *framework.Statement, job
 	if !common.AllocateJob(ssn, stmt, nodes, job, false) {
 		log.InfraLogger.V(3).Infof("Could not allocate resources for job: <%v/%v> of queue <%v>",
 			job.Namespace, job.Name, job.Queue)
-		return false
+		return false, false
 	}
-
+	pipelined = false
 	if job.ShouldPipelineJob() {
 		log.InfraLogger.V(3).Infof(
 			"Some tasks were pipelined, setting all job to be pipelined for job: <%v/%v>",
@@ -76,12 +83,18 @@ func attemptToAllocateJob(ssn *framework.Session, stmt *framework.Statement, job
 			log.InfraLogger.Errorf(
 				"Failed to covert tasks from allocated to pipelined for job: <%v/%v>, error: <%v>",
 				job.Namespace, job.Name, err)
-			return false
+			return false, false
 		}
+		pipelined = true
 	} else {
 		log.InfraLogger.V(3).Infof("Succesfully allocated resources for job: <%v/%v>",
 			job.Namespace, job.Name)
 	}
 
-	return true
+	return true, pipelined
+}
+
+func setLastStartTimestamp(job *podgroup_info.PodGroupInfo) {
+	timeNow := time.Now()
+	job.LastStartTimestamp = &timeNow
 }

pkg/scheduler/actions/allocate/allocateElastic_test.go

@@ -5,6 +5,7 @@ package allocate_test
 
 import (
 	"testing"
+	"time"
 
 	. "go.uber.org/mock/gomock"
 	"k8s.io/utils/pointer"
@@ -280,9 +281,10 @@ func getElasticTestsMetadata() []integration_tests_utils.TestTopologyMetadata {
 				},
 				TaskExpectedResults: map[string]test_utils.TestExpectedResultBasic{
 					"pending_job0-0": {
-						NodeName:     "node0",
-						GPUsRequired: 1,
-						Status:       pod_status.Running,
+						NodeName:                    "node0",
+						GPUsRequired:                1,
+						Status:                      pod_status.Running,
+						LastStartTimestampOlderThan: pointer.Duration(time.Second * 30),
 					},
 					"pending_job0-1": {
 						NodeName:     "node0",

pkg/scheduler/api/podgroup_info/job_info.go

@@ -66,9 +66,10 @@ type PodGroupInfo struct {
 
 	Allocated *resource_info.Resource
 
-	CreationTimestamp metav1.Time
-	PodGroup          *enginev2alpha2.PodGroup
-	PodGroupUID       types.UID
+	CreationTimestamp  metav1.Time
+	LastStartTimestamp *time.Time
+	PodGroup           *enginev2alpha2.PodGroup
+	PodGroupUID        types.UID
 
 	// TODO(k82cn): keep backward compatibility, removed it when v1alpha1 finalized.
 	PDB *policyv1.PodDisruptionBudget
@@ -100,6 +101,8 @@ func NewPodGroupInfo(uid common_info.PodGroupID, tasks ...*pod_info.PodInfo) *Po
 			TimeStamp: nil,
 			Stale:     false,
 		},
+
+		LastStartTimestamp:   nil,
 		activeAllocatedCount: ptr.To(0),
 	}
 
@@ -140,6 +143,16 @@ func (podGroupInfo *PodGroupInfo) SetPodGroup(pg *enginev2alpha2.PodGroup) {
 		}
 	}
 
+	if pg.Annotations[commonconstants.LastStartTimeStamp] != "" {
+		startTime, err := time.Parse(time.RFC3339, pg.Annotations[commonconstants.LastStartTimeStamp])
+		if err != nil {
+			log.InfraLogger.V(7).Warnf("Failed to parse start timestamp for podgroup <%s/%s> err: %v",
+				podGroupInfo.Namespace, podGroupInfo.Name, err)
+		} else {
+			podGroupInfo.LastStartTimestamp = &startTime
+		}
+	}
+
 	log.InfraLogger.V(7).Infof(
 		"SetPodGroup. podGroupName=<%s>, PodGroupUID=<%s> podGroupInfo.PodGroupIndex=<%d>",
 		podGroupInfo.Name, podGroupInfo.PodGroupUID)
@@ -273,6 +286,16 @@ func (podGroupInfo *PodGroupInfo) GetNumActiveUsedTasks() int {
 	return numTasks
 }
 
+func (podGroupInfo *PodGroupInfo) GetNumAllocatedTasks() int {
+	numTasks := 0
+	for _, task := range podGroupInfo.PodInfos {
+		if pod_status.AllocatedStatus(task.Status) {
+			numTasks++
+		}
+	}
+	return numTasks
+}
+
 func (podGroupInfo *PodGroupInfo) GetPendingTasks() []*pod_info.PodInfo {
 	var pendingTasks []*pod_info.PodInfo
 	for _, task := range podGroupInfo.PodInfos {

pkg/scheduler/cache/status_updater/concurrency.go

@@ -60,6 +60,15 @@ func (su *defaultStatusUpdater) syncPodGroup(inFlightPodGroup, snapshotPodGroup
 		}
 		snapshotPodGroup.Annotations[commonconstants.StalePodgroupTimeStamp] = inFlightPodGroup.Annotations[commonconstants.StalePodgroupTimeStamp]
 	}
+	lastStartTimestampUpdated := false
+	if snapshotPodGroup.Annotations[commonconstants.LastStartTimeStamp] == inFlightPodGroup.Annotations[commonconstants.LastStartTimeStamp] {
+		lastStartTimestampUpdated = true
+	} else {
+		if snapshotPodGroup.Annotations == nil {
+			snapshotPodGroup.Annotations = make(map[string]string)
+		}
+		snapshotPodGroup.Annotations[commonconstants.LastStartTimeStamp] = inFlightPodGroup.Annotations[commonconstants.LastStartTimeStamp]
+	}
 
 	updatedSchedulingCondition := false
 	lastSchedulingCondition := utils.GetLastSchedulingCondition(inFlightPodGroup)
@@ -74,7 +83,7 @@ func (su *defaultStatusUpdater) syncPodGroup(inFlightPodGroup, snapshotPodGroup
 	if !updatedSchedulingCondition {
 		snapshotPodGroup.Status.SchedulingConditions = inFlightPodGroup.Status.SchedulingConditions
 	}
-	return staleTimeStampUpdated && updatedSchedulingCondition
+	return lastStartTimestampUpdated && staleTimeStampUpdated && updatedSchedulingCondition
 }
 
 func (su *defaultStatusUpdater) keyForPayload(name, namespace string, uid types.UID) updatePayloadKey {

pkg/scheduler/cache/status_updater/concurrency_test.go

@@ -105,6 +105,7 @@ var _ = Describe("Status Updater Concurrency", func() {
 				if jobIndex, _ := strconv.Atoi(strings.Split(job.Name, "-")[0]); jobIndex%2 == 0 {
 					job.StalenessInfo.TimeStamp = ptr.To(time.Now())
 					job.StalenessInfo.Stale = true
+					job.LastStartTimestamp = ptr.To(time.Now())
 				}
 				Expect(statusUpdater.RecordJobStatusEvent(job)).To(Succeed())
 			}

pkg/scheduler/cache/status_updater/default_status_updater.go

@@ -176,8 +176,8 @@ func (su *defaultStatusUpdater) PatchPodLabels(pod *v1.Pod, labels map[string]an
 func (su *defaultStatusUpdater) RecordJobStatusEvent(job *podgroup_info.PodGroupInfo) error {
 	var err error
 	var patchData []byte
-	if err, patchData = su.updatePodGroupStaleTimeStamp(job.PodGroup, job.StalenessInfo.TimeStamp); err != nil {
-		log.InfraLogger.V(7).Warnf("Failed to update podgroup stale time. error:: %s", err)
+	if patchData, err = su.updatePodGroupAnnotations(job); err != nil {
+		log.InfraLogger.V(7).Warnf("Failed to update podgroup annotations, error: %s", err)
 	}
 	if job.StalenessInfo.Stale {
 		su.recordStaleJobEvent(job)
@@ -317,23 +317,23 @@ func (su *defaultStatusUpdater) recordUnschedulablePodsEvents(job *podgroup_info
 	return errors.Join(errs...)
 }
 
-func (su *defaultStatusUpdater) updatePodGroupStaleTimeStamp(podGroup *enginev2alpha2.PodGroup, staleTimeStamp *time.Time) (error, []byte) {
-	old := podGroup.DeepCopy()
-	updated := setPodGroupStaleTimeStamp(podGroup, staleTimeStamp)
-
-	if !updated {
+func (su *defaultStatusUpdater) updatePodGroupAnnotations(job *podgroup_info.PodGroupInfo) ([]byte, error) {
+	old := job.PodGroup.DeepCopy()
+	updatedStaleTime := setPodGroupStaleTimeStamp(job.PodGroup, job.StalenessInfo.TimeStamp)
+	updatedStartTime := setPodGroupLastStartTimeStamp(job.PodGroup, job.LastStartTimestamp)
+	if !updatedStaleTime && !updatedStartTime {
 		return nil, nil
 	}
 
-	patchData, err := getPodGroupPatch(old, podGroup)
+	patchData, err := getPodGroupPatch(old, job.PodGroup)
 	if err != nil {
-		return err, nil
+		return nil, err
 	}
 
 	if patchData == nil {
 		return nil, nil
 	}
-	return nil, patchData
+	return patchData, nil
 }
 
 func (su *defaultStatusUpdater) recordUnschedulablePodGroup(job *podgroup_info.PodGroupInfo) bool {
@@ -393,6 +393,34 @@ func setPodGroupStaleTimeStamp(podGroup *enginev2alpha2.PodGroup, staleTimeStamp
 	return true
 }
 
+func setPodGroupLastStartTimeStamp(podGroup *enginev2alpha2.PodGroup, startTimeStamp *time.Time) bool {
+	if podGroup.Annotations == nil {
+		podGroup.Annotations = make(map[string]string)
+	}
+
+	if startTimeStamp == nil {
+		if _, found := podGroup.Annotations[commonconstants.LastStartTimeStamp]; !found {
+			return false
+		}
+
+		delete(podGroup.Annotations, commonconstants.LastStartTimeStamp)
+		return true
+	}
+
+	currTimeStamp, found := podGroup.Annotations[commonconstants.LastStartTimeStamp]
+	if !found {
+		podGroup.Annotations[commonconstants.LastStartTimeStamp] = startTimeStamp.UTC().Format(time.RFC3339)
+		return true
+	}
+
+	if currTimeStamp == startTimeStamp.Format(time.RFC3339) {
+		return false
+	}
+
+	podGroup.Annotations[commonconstants.LastStartTimeStamp] = startTimeStamp.Format(time.RFC3339)
+	return true
+}
+
 func setPodGroupSchedulingCondition(podGroup *enginev2alpha2.PodGroup, schedulingCondition *enginev2alpha2.SchedulingCondition) bool {
 	currentSchedulingConditionIndex := utils.GetSchedulingConditionIndex(podGroup, schedulingCondition.NodePool)
 	lastSchedulingCondition := utils.GetLastSchedulingCondition(podGroup)

pkg/scheduler/test_utils/jobs_fake/jobs.go

@@ -135,6 +135,10 @@ func BuildJobInfo(
 		staleTime := time.Now().Add(-1 * *staleDuration)
 		result.StalenessInfo.TimeStamp = &staleTime
 	}
+	if result.LastStartTimestamp == nil && result.GetNumAllocatedTasks() > 0 {
+		startTime := time.Now().Add(-1 * time.Minute * 1)
+		result.LastStartTimestamp = &startTime
+	}
 	return result
 }
 

pkg/scheduler/test_utils/test_utils.go

@@ -9,6 +9,7 @@ import (
 	"fmt"
 	"strconv"
 	"testing"
+	"time"
 
 	// lint:ignore ST1001 we want to use gomock here
 	. "go.uber.org/mock/gomock"
@@ -96,14 +97,15 @@ type TestSessionConfig struct {
 }
 
 type TestExpectedResultBasic struct {
-	NodeName             string
-	GPUsRequired         float64
-	GPUsAccepted         float64
-	MilliCpuRequired     float64
-	MemoryRequired       float64
-	Status               pod_status.PodStatus
-	GPUGroups            []string
-	DontValidateGPUGroup bool
+	NodeName                    string
+	GPUsRequired                float64
+	GPUsAccepted                float64
+	MilliCpuRequired            float64
+	MemoryRequired              float64
+	Status                      pod_status.PodStatus
+	GPUGroups                   []string
+	DontValidateGPUGroup        bool
+	LastStartTimestampOlderThan *time.Duration
 }
 
 type TestExpectedNodesResources struct {
@@ -163,6 +165,19 @@ func MatchExpectedAndRealTasks(t *testing.T, testNumber int, testMetadata TestTo
 				}
 			}
 		}
+
+		if pod_status.AllocatedStatus(jobExpectedResult.Status) {
+			if job.LastStartTimestamp == nil {
+				t.Errorf("Test number: %d, name: %v, has failed. Task name: %v, actual last start timestamp is not set expecting pod_status.%v", testNumber, testMetadata.Name, jobName, jobExpectedResult.Status.String())
+			} else if jobExpectedResult.LastStartTimestampOlderThan != nil {
+				now := time.Now()
+				if now.Sub(*job.LastStartTimestamp) < *jobExpectedResult.LastStartTimestampOlderThan {
+					t.Errorf("Test number: %d, name: %v, has failed. Task name: %v, actual last start timestamp is not older than %v", testNumber, testMetadata.Name, jobName,
+						*jobExpectedResult.LastStartTimestampOlderThan)
+				}
+			}
+		}
+
 		if sumOfAcceptedGpus != jobExpectedResult.GPUsAccepted && jobExpectedResult.GPUsAccepted != 0 {
 			t.Errorf("Test number: %d, name: %v, has failed. Task name: %v, actual accept GPUs: %v, was expecting GPUs: %v", testNumber, testMetadata.Name, jobName, sumOfAcceptedGpus, jobExpectedResult.GPUsAccepted)
 		}
@@ -178,7 +193,7 @@ func MatchExpectedAndRealTasks(t *testing.T, testNumber int, testMetadata TestTo
 	}
 
 	if len(testMetadata.TaskExpectedResults) > 0 {
-		for jobId := range ssn.PodGroupInfos {
+		for jobId, job := range ssn.PodGroupInfos {
 			for taskId, task := range ssn.PodGroupInfos[jobId].PodInfos {
 				taskExpectedResult, found := testMetadata.TaskExpectedResults[string(taskId)]
 				if !found {
@@ -225,6 +240,18 @@ func MatchExpectedAndRealTasks(t *testing.T, testNumber int, testMetadata TestTo
 						taskExpectedResult.MemoryRequired)
 				}
 
+				if pod_status.AllocatedStatus(taskExpectedResult.Status) {
+					if job.LastStartTimestamp == nil {
+						t.Errorf("Test number: %d, name: %v, has failed. Task name: %v, actual last start timestamp is not set expecting pod_status.%v", testNumber, testMetadata.Name, taskId, taskExpectedResult.Status.String())
+					} else if taskExpectedResult.LastStartTimestampOlderThan != nil {
+						now := time.Now()
+						if now.Sub(*job.LastStartTimestamp) < *taskExpectedResult.LastStartTimestampOlderThan {
+							t.Errorf("Test number: %d, name: %v, has failed. Task name: %v, actual last start timestamp is not older than %v", testNumber, testMetadata.Name, taskId,
+								*taskExpectedResult.LastStartTimestampOlderThan)
+						}
+					}
+				}
+
 				// verify fractional GPUs index
 				if pod_status.IsActiveUsedStatus(task.Status) &&
 					!taskExpectedResult.DontValidateGPUGroup &&