resolve conflicts and fix minor issues with native preemption retries

Signed-off-by: Dejan Zele Pejchev <pejcev.dejan@gmail.com>

Author: dejanzele

config/scheduler/config.yaml

@@ -124,6 +124,8 @@ scheduling:
   executorTimeout: "10m"
   maxUnacknowledgedJobsPerExecutor: 2500
   executorUpdateFrequency: "60s"
+  preemptionRetry:
+    enabled: true
   experimentalIndicativePricing:
     basePrice: 100.0
     basePriority: 500.0

config/server/config.yaml

@@ -66,6 +66,9 @@ submission:
   defaultActiveDeadlineByResourceRequest:
     nvidia.com/gpu: "336h" # 14 days.
   assertInitContainersRequestFractionalCpu: true
+  preemptionRetry:
+    enabled: true
+    maxRetryCount: 10
 pulsar:
   URL: "pulsar://pulsar:6650"
   jobsetEventsTopic: "events"

internal/common/constants/constants.go

@@ -34,8 +34,12 @@ const (
 	PoolAnnotation      = "armadaproject.io/pool"
 	ReservationTaintKey = "armadaproject.io/reservation"
 
-	PreemptionRetryCountMaxAnnotation = "armadaproject.io/preemptionRetryCountMax"
-	PreemptionRetryEnabledAnnotation  = "armadaproject.io/preemptionRetryEnabled"
+	// PreemptionRetryEnabledAnnotation enables or disables preemption retries for a job.
+	// When set to "true", the job will be rescheduled if preempted, up to the max retry count.
+	PreemptionRetryEnabledAnnotation = "armadaproject.io/preemptionRetryEnabled"
+	// PreemptionMaxRetryCountAnnotation specifies the maximum number of times a job can be
+	// rescheduled after preemption. Must be a positive integer.
+	PreemptionMaxRetryCountAnnotation = "armadaproject.io/preemptionMaxRetryCount"
 )
 
 var schedulingAnnotations = map[string]bool{
@@ -44,6 +48,8 @@ var schedulingAnnotations = map[string]bool{
 	GangNodeUniformityLabelAnnotation: true,
 	FailFastAnnotation:                true,
 	JobPriceBand:                      true,
+	PreemptionRetryEnabledAnnotation:  true,
+	PreemptionMaxRetryCountAnnotation: true,
 }
 
 func IsSchedulingAnnotation(annotation string) bool {

internal/common/preemption/utils.go

@@ -6,6 +6,20 @@ import (
 	"github.com/armadaproject/armada/internal/common/constants"
 )
 
+// RetryConfig contains preemption retry settings.
+// Used by both server (submission validation) and scheduler (runtime behavior).
+type RetryConfig struct {
+	// Enabled toggles preemption retries globally.
+	// When false, jobs with preemption retry annotations are rejected at submission.
+	Enabled bool
+	// MaxRetryCount is the maximum retry count users can request via annotations.
+	// Used by server to validate job submissions. If nil, no upper bound is enforced.
+	MaxRetryCount *uint
+	// DefaultRetryCount is the retry count applied when jobs don't specify via annotation.
+	// Used by scheduler at runtime. If nil, defaults to 0 (no retries unless job specifies).
+	DefaultRetryCount *uint
+}
+
 // 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) {
@@ -23,7 +37,7 @@ func AreRetriesEnabled(annotations map[string]string) (enabled bool, annotationS
 
 // 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) {
-	maxRetryCountStr, exists := annotations[constants.PreemptionRetryCountMaxAnnotation]
+	maxRetryCountStr, exists := annotations[constants.PreemptionMaxRetryCountAnnotation]
 	if !exists {
 		return 0, false
 	}

internal/scheduler/configuration/configuration.go

@@ -9,6 +9,7 @@ import (
 	authconfig "github.com/armadaproject/armada/internal/common/auth/configuration"
 	commonconfig "github.com/armadaproject/armada/internal/common/config"
 	grpcconfig "github.com/armadaproject/armada/internal/common/grpc/configuration"
+	"github.com/armadaproject/armada/internal/common/preemption"
 	profilingconfig "github.com/armadaproject/armada/internal/common/profiling/configuration"
 	armadaresource "github.com/armadaproject/armada/internal/common/resource"
 	"github.com/armadaproject/armada/internal/common/types"
@@ -297,8 +298,8 @@ type SchedulingConfig struct {
 	DefaultPoolSchedulePriority int
 	Pools                       []PoolConfig
 	ExperimentalIndicativeShare ExperimentalIndicativeShare
-	// Default preemption retries settings so you don't have to annotate all jobs with retries.
-	DefaultPreemptionRetry PreemptionRetryConfig
+	// Preemption retry settings. Jobs can override via annotations.
+	PreemptionRetry preemption.RetryConfig
 }
 
 const (
@@ -482,8 +483,3 @@ type PricingApiConfig struct {
 	// It will stub the pricing api so it returns non-zero values but won't call and external service
 	DevModeEnabled bool
 }
-
-type PreemptionRetryConfig struct {
-	Enabled              bool
-	DefaultMaxRetryCount *uint
-}

internal/scheduler/database/migrations/022_add_job_run_index.sql

@@ -1 +1 @@
-ALTER TABLE runs ADD COLUMN IF NOT EXISTS run_index bigint;
+ALTER TABLE runs ADD COLUMN IF NOT EXISTS run_index bigint DEFAULT 0;

internal/scheduler/jobdb/job.go

@@ -4,9 +4,6 @@ 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"
@@ -15,6 +12,7 @@ import (
 
 	"github.com/armadaproject/armada/internal/common/constants"
 	armadamaps "github.com/armadaproject/armada/internal/common/maps"
+	"github.com/armadaproject/armada/internal/common/preemption"
 	"github.com/armadaproject/armada/internal/common/types"
 	"github.com/armadaproject/armada/internal/scheduler/adapters"
 	"github.com/armadaproject/armada/internal/scheduler/internaltypes"
@@ -846,9 +844,9 @@ func (job *Job) NumAttempts() uint {
 // 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 {
+func (job *Job) IsEligibleForPreemptionRetry(retryConfig preemption.RetryConfig) bool {
 	// Check if job explicitly enabled/disabled retries, falling back to platform default
-	enabled := defaultPreemptionRetryConfig.Enabled
+	enabled := retryConfig.Enabled
 	if jobRetryEnabled, exists := preemption.AreRetriesEnabled(job.Annotations()); exists {
 		enabled = jobRetryEnabled
 	}
@@ -857,7 +855,7 @@ func (job *Job) IsEligibleForPreemptionRetry(defaultPreemptionRetryConfig config
 		return false
 	}
 
-	return job.NumPreemptedRuns() <= job.MaxPreemptionRetryCount(defaultPreemptionRetryConfig)
+	return job.NumPreemptedRuns() <= job.MaxPreemptionRetryCount(retryConfig)
 }
 
 func (job *Job) NumPreemptedRuns() uint {
@@ -870,11 +868,11 @@ func (job *Job) NumPreemptedRuns() uint {
 	return preemptCount
 }
 
-func (job *Job) MaxPreemptionRetryCount(defaultPreemptionRetryConfig configuration.PreemptionRetryConfig) uint {
+func (job *Job) MaxPreemptionRetryCount(retryConfig preemption.RetryConfig) uint {
 	// Start with platform default
 	var maxRetryCount uint
-	if defaultPreemptionRetryConfig.DefaultMaxRetryCount != nil {
-		maxRetryCount = *defaultPreemptionRetryConfig.DefaultMaxRetryCount
+	if retryConfig.DefaultRetryCount != nil {
+		maxRetryCount = *retryConfig.DefaultRetryCount
 	}
 
 	// Allow jobs to override with a custom max retry count

internal/scheduler/jobdb/job_run_test.go

@@ -123,6 +123,51 @@ func TestJobRun_TestRunAttempted(t *testing.T) {
 	assert.True(t, attemptedRun.RunAttempted())
 }
 
+func TestJobRun_InTerminalState(t *testing.T) {
+	tests := map[string]struct {
+		run      *JobRun
+		expected bool
+	}{
+		"base run is not terminal": {
+			run:      baseJobRun,
+			expected: false,
+		},
+		"succeeded run is terminal": {
+			run:      baseJobRun.WithSucceeded(true),
+			expected: true,
+		},
+		"failed run is terminal": {
+			run:      baseJobRun.WithFailed(true),
+			expected: true,
+		},
+		"cancelled run is terminal": {
+			run:      baseJobRun.WithCancelled(true),
+			expected: true,
+		},
+		"returned run is terminal": {
+			run:      baseJobRun.WithReturned(true),
+			expected: true,
+		},
+		"preempted run is terminal": {
+			run:      baseJobRun.WithPreempted(true),
+			expected: true,
+		},
+		"running run is not terminal": {
+			run:      baseJobRun.WithRunning(true),
+			expected: false,
+		},
+		"pending run is not terminal": {
+			run:      baseJobRun.WithPending(true),
+			expected: false,
+		},
+	}
+	for name, tc := range tests {
+		t.Run(name, func(t *testing.T) {
+			assert.Equal(t, tc.expected, tc.run.InTerminalState())
+		})
+	}
+}
+
 func TestDeepCopy(t *testing.T) {
 	run := jobDb.CreateRun(
 		uuid.NewString(),