Add support for native preemption retries

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

Author: jparraga-stackav

internal/common/preemption/utils.go

@@ -0,0 +1,39 @@
+package preemption
+
+import (
+	"strconv"
+
+	"github.com/armadaproject/armada/internal/server/configuration"
+)
+
+// AreRetriesEnabled determines whether preemption retries are enabled at the job level. Also returns whether the
+// annotation was set.
+func AreRetriesEnabled(annotations map[string]string) (enabled bool, annotationSet 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) (maxRetryCount uint, annotationSet bool) {
+	var preemptionRetryCountMax uint = 0
+	preemptionRetryCountMaxStr, exists := annotations[configuration.PreemptionRetryCountMaxAnnotation]
+
+	if !exists {
+		return preemptionRetryCountMax, false
+	}
+	maybePreemptionRetryCountMax, err := strconv.Atoi(preemptionRetryCountMaxStr)
+	if err != nil {
+		return preemptionRetryCountMax, true
+	} else {
+		preemptionRetryCountMax = uint(maybePreemptionRetryCountMax)
+		return preemptionRetryCountMax, true
+	}
+}

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 retries.
+	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

@@ -54,7 +54,7 @@ type JobRepository interface {
 	CountReceivedPartitions(ctx *armadacontext.Context, groupId uuid.UUID) (uint32, error)
 
 	// FindInactiveRuns returns a slice containing all dbRuns that the scheduler does not currently consider active
-	// Runs are inactive if they don't exist or if they have succeeded, failed or been cancelled
+	// Runs are inactive if they don't exist or if they have succeeded, failed, preempted or been cancelled
 	FindInactiveRuns(ctx *armadacontext.Context, runIds []string) ([]string, error)
 
 	// FetchJobRunLeases fetches new job runs for a given executor.  A maximum of maxResults rows will be returned, while run
@@ -293,7 +293,7 @@ func (r *PostgresJobRepository) FetchJobUpdates(ctx *armadacontext.Context, jobS
 }
 
 // FindInactiveRuns returns a slice containing all dbRuns that the scheduler does not currently consider active
-// Runs are inactive if they don't exist or if they have succeeded, failed or been cancelled
+// Runs are inactive if they don't exist or if they have succeeded, failed, preempted or been cancelled
 func (r *PostgresJobRepository) FindInactiveRuns(ctx *armadacontext.Context, runIds []string) ([]string, error) {
 	var inactiveRuns []string
 	err := pgx.BeginTxFunc(ctx, r.db, pgx.TxOptions{
@@ -313,6 +313,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))
@@ -361,6 +362,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;
 `

internal/scheduler/database/job_repository_test.go

@@ -553,6 +553,15 @@ func TestFindInactiveRuns(t *testing.T) {
 			},
 			expectedInactive: []string{runIds[1]},
 		},
+		"run preempted": {
+			runsToCheck: runIds,
+			dbRuns: []Run{
+				{RunID: runIds[0]},
+				{RunID: runIds[1], Preempted: true},
+				{RunID: runIds[2]},
+			},
+			expectedInactive: []string{runIds[1]},
+		},
 		"run missing": {
 			runsToCheck: runIds,
 			dbRuns: []Run{
@@ -654,6 +663,14 @@ func TestFetchJobRunLeases(t *testing.T) {
 			Pool:      "test-pool",
 			Succeeded: true, // should be ignored as terminal
 		},
+		{
+			RunID:     uuid.NewString(),
+			JobID:     dbJobs[0].JobID,
+			JobSet:    "test-jobset",
+			Executor:  executorName,
+			Pool:      "test-pool",
+			Preempted: true, // should be ignored as terminal
+		},
 	}
 	expectedLeases := make([]*JobRunLease, 4)
 	for i := range expectedLeases {

internal/scheduler/jobdb/job.go

@@ -4,6 +4,9 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/armadaproject/armada/internal/common/preemption"
+	"github.com/armadaproject/armada/internal/scheduler/configuration"
+
 	"github.com/hashicorp/go-multierror"
 	"github.com/pkg/errors"
 	"golang.org/x/exp/maps"
@@ -734,6 +737,59 @@ func (job *Job) NumAttempts() uint {
 	return attempts
 }
 
+// IsEligibleForPreemptionRetry determines whether the job is eligible for preemption retries. It checks whether the
+// scheduler or the job has opted in for preemption retries. It then checks whether the job has exhausted the number
+// of retries.
+func (job *Job) IsEligibleForPreemptionRetry(defaultPreemptionRetryConfig configuration.PreemptionRetryConfig) bool {
+	enabled := false
+
+	// Check for platform default first
+	if defaultPreemptionRetryConfig.Enabled {
+		enabled = true
+	}
+
+	// Check if job explicitly enabled/disabled retries
+	jobRetryEnabled, exists := preemption.AreRetriesEnabled(job.Annotations())
+	if exists {
+		enabled = jobRetryEnabled
+	}
+
+	if !enabled {
+		return false
+	}
+
+	maxRetryCount := job.MaxPreemptionRetryCount(defaultPreemptionRetryConfig)
+
+	return job.NumPreemptedRuns() <= maxRetryCount
+}
+
+func (job *Job) NumPreemptedRuns() uint {
+	preemptCount := uint(0)
+	for _, run := range job.runsById {
+		if run.preempted {
+			preemptCount++
+		}
+	}
+	return preemptCount
+}
+
+func (job *Job) MaxPreemptionRetryCount(defaultPreemptionRetryConfig configuration.PreemptionRetryConfig) uint {
+	var maxRetryCount uint = 0
+
+	// Check for platform default first
+	if defaultPreemptionRetryConfig.DefaultMaxRetryCount != nil {
+		platformDefaultMaxRetryCount := *defaultPreemptionRetryConfig.DefaultMaxRetryCount
+		maxRetryCount = platformDefaultMaxRetryCount
+	}
+
+	// Allow jobs to set a custom max retry count
+	jobMaxRetryCount, exists := preemption.GetMaxRetryCount(job.Annotations())
+	if exists {
+		maxRetryCount = jobMaxRetryCount
+	}
+	return maxRetryCount
+}
+
 // AllRuns returns all runs associated with job.
 func (job *Job) AllRuns() []*JobRun {
 	return maps.Values(job.runsById)

internal/scheduler/jobdb/job_run.go

@@ -501,7 +501,7 @@ func (run *JobRun) WithoutTerminal() *JobRun {
 
 // InTerminalState returns true if the JobRun is in a terminal state
 func (run *JobRun) InTerminalState() bool {
-	return run.succeeded || run.failed || run.cancelled || run.returned
+	return run.succeeded || run.failed || run.cancelled || run.returned || run.preempted
 }
 
 func (run *JobRun) DeepCopy() *JobRun {

internal/scheduler/jobdb/job_test.go

@@ -10,6 +10,9 @@ import (
 
 	"github.com/armadaproject/armada/internal/common/types"
 	"github.com/armadaproject/armada/internal/scheduler/internaltypes"
+
+	configuration2 "github.com/armadaproject/armada/internal/scheduler/configuration"
+	"github.com/armadaproject/armada/internal/server/configuration"
 )
 
 var jobSchedulingInfo = &internaltypes.JobSchedulingInfo{
@@ -26,6 +29,35 @@ var jobSchedulingInfo = &internaltypes.JobSchedulingInfo{
 	},
 }
 
+var jobSchedulingInfoWithRetryEnabled = &internaltypes.JobSchedulingInfo{
+	PodRequirements: &internaltypes.PodRequirements{
+		ResourceRequirements: v1.ResourceRequirements{
+			Requests: v1.ResourceList{
+				"cpu":                 k8sResource.MustParse("1"),
+				"storage-connections": k8sResource.MustParse("1"),
+			},
+		},
+		Annotations: map[string]string{
+			configuration.PreemptionRetryEnabledAnnotation:  "true",
+			configuration.PreemptionRetryCountMaxAnnotation: "1",
+		},
+	},
+}
+
+var jobSchedulingInfoWithRetryDisabled = &internaltypes.JobSchedulingInfo{
+	PodRequirements: &internaltypes.PodRequirements{
+		ResourceRequirements: v1.ResourceRequirements{
+			Requests: v1.ResourceList{
+				"cpu":                 k8sResource.MustParse("1"),
+				"storage-connections": k8sResource.MustParse("1"),
+			},
+		},
+		Annotations: map[string]string{
+			configuration.PreemptionRetryEnabledAnnotation: "false",
+		},
+	},
+}
+
 var baseJob, _ = jobDb.NewJob(
 	"test-job",
 	"test-jobSet",
@@ -42,6 +74,38 @@ var baseJob, _ = jobDb.NewJob(
 	[]string{},
 )
 
+var baseJobWithRetryEnabled, _ = jobDb.NewJob(
+	"test-job",
+	"test-jobSet",
+	"test-queue",
+	2,
+	jobSchedulingInfoWithRetryEnabled,
+	true,
+	0,
+	false,
+	false,
+	false,
+	3,
+	false,
+	[]string{},
+)
+
+var baseJobWithRetryDisabled, _ = jobDb.NewJob(
+	"test-job",
+	"test-jobSet",
+	"test-queue",
+	2,
+	jobSchedulingInfoWithRetryDisabled,
+	true,
+	0,
+	false,
+	false,
+	false,
+	3,
+	false,
+	[]string{},
+)
+
 var baseRun = &JobRun{
 	id:        uuid.New().String(),
 	created:   3,
@@ -425,3 +489,85 @@ func TestJob_TestKubernetesResourceRequirements(t *testing.T) {
 	assert.Equal(t, int64(1000), baseJob.KubernetesResourceRequirements().GetByNameZeroIfMissing("cpu"))
 	assert.Equal(t, int64(0), baseJob.KubernetesResourceRequirements().GetByNameZeroIfMissing("storage-connections"))
 }
+
+func TestIsEligibleForPreemptionRetry(t *testing.T) {
+	premptedRun1 := &JobRun{
+		id:        uuid.New().String(),
+		created:   3,
+		executor:  "test-executor",
+		preempted: true,
+	}
+
+	premptedRun2 := &JobRun{
+		id:        uuid.New().String(),
+		created:   5,
+		executor:  "test-executor",
+		preempted: true,
+	}
+
+	defaultMaxRetryCountEnabled := uint(5)
+	platformDefaultEnabled := configuration2.PreemptionRetryConfig{
+		Enabled:              true,
+		DefaultMaxRetryCount: &defaultMaxRetryCountEnabled,
+	}
+
+	defaultMaxRetryCountDisabled := uint(0)
+	platformDefaultDisabled := configuration2.PreemptionRetryConfig{
+		Enabled:              false,
+		DefaultMaxRetryCount: &defaultMaxRetryCountDisabled,
+	}
+
+	// no runs
+	t.Run("job with retry enabled and platform disabled and no runs", func(t *testing.T) {
+		assert.True(t, baseJobWithRetryEnabled.IsEligibleForPreemptionRetry(platformDefaultDisabled))
+	})
+
+	t.Run("job with retry disabled and platform enabled and no runs", func(t *testing.T) {
+		assert.False(t, baseJobWithRetryDisabled.IsEligibleForPreemptionRetry(platformDefaultEnabled))
+	})
+
+	t.Run("job with platform retry enabled and no runs", func(t *testing.T) {
+		assert.True(t, baseJob.IsEligibleForPreemptionRetry(platformDefaultEnabled))
+	})
+
+	t.Run("job with platform retry disabled and no runs", func(t *testing.T) {
+		assert.False(t, baseJob.IsEligibleForPreemptionRetry(platformDefaultDisabled))
+	})
+
+	// runs but none are preempted
+	t.Run("job with retry enabled and platform disabled runs but no preempted runs", func(t *testing.T) {
+		updatedJob := baseJobWithRetryEnabled.WithUpdatedRun(baseRun).WithUpdatedRun(baseRun)
+		assert.True(t, updatedJob.IsEligibleForPreemptionRetry(platformDefaultDisabled))
+	})
+
+	t.Run("job with platform enabled runs but no preempted runs", func(t *testing.T) {
+		updatedJob := baseJob.WithUpdatedRun(baseRun).WithUpdatedRun(baseRun)
+		assert.True(t, updatedJob.IsEligibleForPreemptionRetry(platformDefaultEnabled))
+	})
+
+	t.Run("job with retry enabled and platform disabled and one run", func(t *testing.T) {
+		updatedJob := baseJobWithRetryEnabled.WithUpdatedRun(premptedRun1)
+		assert.True(t, updatedJob.IsEligibleForPreemptionRetry(platformDefaultDisabled))
+	})
+
+	t.Run("job with platform enabled and one run", func(t *testing.T) {
+		updatedJob := baseJob.WithUpdatedRun(premptedRun1)
+		assert.True(t, updatedJob.IsEligibleForPreemptionRetry(platformDefaultEnabled))
+	})
+
+	// runs that are preempted
+	t.Run("job with retry enabled and platform disabled and out of retries", func(t *testing.T) {
+		updatedJob := baseJobWithRetryEnabled.WithUpdatedRun(premptedRun1).WithUpdatedRun(premptedRun2)
+		assert.False(t, updatedJob.IsEligibleForPreemptionRetry(platformDefaultDisabled))
+	})
+
+	t.Run("job with retry enabled with platform enabled and out of retries", func(t *testing.T) {
+		updatedJob := baseJobWithRetryEnabled.WithUpdatedRun(premptedRun1).WithUpdatedRun(premptedRun2)
+		assert.False(t, updatedJob.IsEligibleForPreemptionRetry(platformDefaultEnabled))
+	})
+
+	t.Run("job with platform enabled and retries left", func(t *testing.T) {
+		updatedJob := baseJob.WithUpdatedRun(premptedRun1).WithUpdatedRun(premptedRun2)
+		assert.True(t, updatedJob.IsEligibleForPreemptionRetry(platformDefaultEnabled))
+	})
+}