Add support for native preemption retries

Signed-off-by: Jason Parraga <[email protected]>

Author: Jason Parraga

internal/common/preemption/utils.go

@@ -0,0 +1,39 @@
+package preemption
+
+import (
+	"github.com/armadaproject/armada/internal/server/configuration"
+	"strconv"
+)
+
+// AreRetriesEnabled determines whether preemption retries are enabled at the job level. Also returns whether the
+// annotation was set.
+func AreRetriesEnabled(annotations map[string]string) (bool, bool) {
+	preemptionRetryEnabledStr, exists := annotations[configuration.PreemptionRetryEnabledAnnotation]
+	if !exists {
+		return false, false
+	}
+
+	preemptionRetryEnabled, err := strconv.ParseBool(preemptionRetryEnabledStr)
+	if err != nil {
+		return false, true
+	}
+	return preemptionRetryEnabled, true
+}
+
+// GetMaxRetryCount gets the max preemption retry count at a job level. Also returns whether the annotation was set.
+func GetMaxRetryCount(annotations map[string]string) (uint, bool) {
+	var preemptionRetryCountMax uint = 0
+	preemptionRetryCountMaxStr, exists := annotations[configuration.PreemptionRetryCountMaxAnnotation]
+
+	if !exists {
+		return 0, false
+	}
+	maybePreemptionRetryCountMax, err := strconv.Atoi(preemptionRetryCountMaxStr)
+	if err != nil {
+		return 0, true
+	} else {
+		preemptionRetryCountMax = uint(maybePreemptionRetryCountMax)
+	}
+
+	return preemptionRetryCountMax, true
+}

internal/executor/service/job_state_reporter.go

@@ -93,6 +93,11 @@ func (stateReporter *JobStateReporter) reportCurrentStatus(pod *v1.Pod) {
 	}
 
 	if pod.Status.Phase == v1.PodFailed {
+
+		if util.IsPodPreempted(pod) {
+			return
+		}
+
 		hasIssue := stateReporter.podIssueHandler.HasIssue(util.ExtractJobRunId(pod))
 		if hasIssue {
 			// Pod already being handled by issue handler

internal/executor/service/job_state_reporter_test.go

@@ -166,6 +166,27 @@ func TestJobStateReporter_HandlesFailedPod_WithRetryableError(t *testing.T) {
 	assertExpectedEvents(t, before, eventReporter.ReceivedEvents, reflect.TypeOf(&armadaevents.EventSequence_Event_JobRunErrors{}))
 }
 
+func TestJobStateReporter_IgnoresPreemptedPod(t *testing.T) {
+	_, _, eventReporter, fakeClusterContext := setUpJobStateReporterTest(t)
+
+	before := makeTestPod(v1.PodStatus{Phase: v1.PodRunning})
+	after := before.DeepCopy()
+	after.Status = v1.PodStatus{
+		Phase: v1.PodFailed,
+		Conditions: []v1.PodCondition{
+			{
+				Type:   v1.DisruptionTarget,
+				Status: v1.ConditionTrue,
+				Reason: util.PreemptedReason,
+			},
+		},
+	}
+
+	fakeClusterContext.SimulateUpdateAddEvent(before, after)
+	time.Sleep(time.Millisecond * 100) // Give time for async routine to process message
+	assert.Len(t, eventReporter.ReceivedEvents, 0)
+}
+
 func setUpJobStateReporterTest(t *testing.T) (*JobStateReporter, *stubIssueHandler, *mocks.FakeEventReporter, *fakecontext.SyncFakeClusterContext) {
 	fakeClusterContext := fakecontext.NewSyncFakeClusterContext()
 	eventReporter := mocks.NewFakeEventReporter()

internal/executor/service/pod_issue_handler.go

@@ -521,7 +521,7 @@ func createStuckPodMessage(retryable bool, originalMessage string) string {
 func (p *PodIssueHandler) handleDeletedPod(pod *v1.Pod) {
 	jobId := util.ExtractJobId(pod)
 	if jobId != "" {
-		isUnexpectedDeletion := !util.IsMarkedForDeletion(pod) && !util.IsPodFinishedAndReported(pod)
+		isUnexpectedDeletion := !util.IsMarkedForDeletion(pod) && !util.IsPodFinishedAndReported(pod) && !util.IsPodPreempted(pod)
 		if isUnexpectedDeletion {
 			p.attemptToRegisterIssue(&runIssue{
 				JobId: jobId,

internal/executor/service/pod_issue_handler_test.go

@@ -329,6 +329,27 @@ func TestPodIssueService_ReportsFailed_IfDeletedExternally(t *testing.T) {
 	assert.Equal(t, jobId, failedEvent.JobRunErrors.JobId)
 }
 
+func TestPodIssueService_IgnoresFailed_IfDeletedExternallyDueToPreemption(t *testing.T) {
+	podIssueService, _, fakeClusterContext, eventsReporter, err := setupTestComponents([]*job.RunState{})
+	require.NoError(t, err)
+	preemptedPod := makeTestPod(v1.PodStatus{
+		Phase: v1.PodFailed,
+		Conditions: []v1.PodCondition{
+			{
+				Type:   v1.DisruptionTarget,
+				Status: v1.ConditionTrue,
+				Reason: util.PreemptedReason,
+			},
+		},
+	})
+	fakeClusterContext.SimulateDeletionEvent(preemptedPod)
+
+	podIssueService.HandlePodIssues()
+
+	assert.Len(t, eventsReporter.ReceivedEvents, 0)
+	assert.Len(t, podIssueService.knownPodIssues, 0)
+}
+
 func TestPodIssueService_ReportsFailed_IfPodOfActiveRunGoesMissing(t *testing.T) {
 	baseTime := time.Now()
 	fakeClock := clock.NewFakeClock(baseTime)

internal/executor/util/pod_status.go

@@ -15,6 +15,7 @@ const (
 	oomKilledReason  = "OOMKilled"
 	evictedReason    = "Evicted"
 	deadlineExceeded = "DeadlineExceeded"
+	PreemptedReason  = "PreemptionByScheduler"
 )
 
 // TODO: Need to detect pod preemption. So that job failed events can include a string indicating a pod was preempted.

internal/executor/util/pod_util.go

@@ -370,3 +370,15 @@ func GroupByQueue(pods []*v1.Pod) map[string][]*v1.Pod {
 	}
 	return podsByQueue
 }
+
+func IsPodPreempted(pod *v1.Pod) bool {
+	for _, containerCondition := range pod.Status.Conditions {
+		if containerCondition.Type == v1.DisruptionTarget && containerCondition.Status == v1.ConditionTrue {
+			if containerCondition.Reason == PreemptedReason {
+				return true
+			}
+		}
+	}
+
+	return false
+}

internal/executor/util/pod_util_test.go

@@ -13,6 +13,55 @@ import (
 	"github.com/armadaproject/armada/internal/server/configuration"
 )
 
+func TestIsPodPreempted(t *testing.T) {
+
+	t.Run("preempted pod is preempted", func(t *testing.T) {
+		pod := v1.Pod{
+			Status: v1.PodStatus{
+				Phase: v1.PodFailed,
+				Conditions: []v1.PodCondition{
+					{
+						Type:   v1.DisruptionTarget,
+						Status: v1.ConditionTrue,
+						Reason: PreemptedReason,
+					},
+				},
+			},
+		}
+
+		assert.True(t, IsPodPreempted(&pod))
+	})
+
+	t.Run("failed pod is not preempted", func(t *testing.T) {
+		pod := v1.Pod{
+			Status: v1.PodStatus{
+				Phase: v1.PodFailed,
+				ContainerStatuses: []v1.ContainerStatus{
+					{
+						State: v1.ContainerState{
+							Terminated: &v1.ContainerStateTerminated{
+								ExitCode: 1,
+							},
+						},
+					},
+				},
+			},
+		}
+
+		assert.False(t, IsPodPreempted(&pod))
+	})
+
+	t.Run("successful pod is not preempted", func(t *testing.T) {
+		pod := v1.Pod{
+			Status: v1.PodStatus{
+				Phase: v1.PodSucceeded,
+			},
+		}
+
+		assert.False(t, IsPodPreempted(&pod))
+	})
+}
+
 func TestIsInTerminalState_ShouldReturnTrueWhenPodInSucceededPhase(t *testing.T) {
 	pod := v1.Pod{
 		Status: v1.PodStatus{

internal/scheduler/configuration/configuration.go

@@ -248,6 +248,8 @@ type SchedulingConfig struct {
 	Pools                         []PoolConfig
 	ExperimentalIndicativePricing ExperimentalIndicativePricing
 	ExperimentalIndicativeShare   ExperimentalIndicativeShare
+	// Default preemption retries settings so you don't have to annotate all jobs with reties.
+	DefaultPreemptionRetry PreemptionRetryConfig
 }
 
 const (
@@ -354,3 +356,8 @@ type PriorityOverrideConfig struct {
 	ServiceUrl      string
 	ForceNoTls      bool
 }
+
+type PreemptionRetryConfig struct {
+	Enabled              bool
+	DefaultMaxRetryCount *uint
+}

internal/scheduler/database/job_repository.go

@@ -283,6 +283,7 @@ func (r *PostgresJobRepository) FindInactiveRuns(ctx *armadacontext.Context, run
 		WHERE runs.run_id IS NULL
 		OR runs.succeeded = true
  		OR runs.failed = true
+		OR runs.preempted = true
 		OR runs.cancelled = true;`
 
 		rows, err := tx.Query(ctx, fmt.Sprintf(query, tmpTable))
@@ -331,6 +332,7 @@ func (r *PostgresJobRepository) FetchJobRunLeases(ctx *armadacontext.Context, ex
 				AND jr.succeeded = false
 				AND jr.failed = false
 				AND jr.cancelled = false
+				AND jr.preempted = false
 				ORDER BY jr.serial
 				LIMIT %d;
 `