Wire the retry policy engine into the scheduler behind a feature flag

client/rust/src/gen/api.rs

@@ -1791,6 +1791,13 @@ pub struct JobFailedEvent {
     pub failure_category: ::prost::alloc::string::String,
     #[prost(string, tag = "17")]
     pub failure_subcategory: ::prost::alloc::string::String,
+    /// retryable indicates the scheduler emitted this failure for an
+    /// intermediate (non-terminal) run that will be retried. When true,
+    /// a subsequent leased/succeeded/failed event for the same job is
+    /// expected. Default false preserves the prior behavior where every
+    /// emitted JobFailedEvent was terminal.
+    #[prost(bool, tag = "18")]
+    pub retryable: bool,
 }
 #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)]
 pub struct JobPreemptingEvent {

config/scheduler/config.yaml

@@ -108,6 +108,9 @@ scheduling:
   maximumPerQueueSchedulingBurst: 1000
   maxJobSchedulingContextsPerExecutor: 10000
   maxRetries: 3
+  retryPolicy:
+    enabled: false
+    globalMaxRetries: 5
   dominantResourceFairnessResourcesToConsider:
     - "cpu"
     - "memory"

internal/scheduler/api.go

@@ -3,6 +3,7 @@ package scheduler
 import (
 	"context"
 	"strconv"
+	"sync"
 
 	"github.com/gogo/protobuf/proto"
 	"github.com/gogo/protobuf/types"
@@ -48,6 +49,8 @@ type ExecutorApi struct {
 	// This is needed to ensure floating resources are not passed to k8s.
 	allowedResources map[string]bool
 	nodeIdLabel      string
+	// Executors already warned about nodes missing the node id label.
+	nodeIdLabelWarnings sync.Map
 	// See scheduling schedulingConfig.
 	priorityClassNameOverride *string
 	clock                     clock.Clock
@@ -432,15 +435,30 @@ func (srv *ExecutorApi) authorize(ctx *armadacontext.Context) error {
 func (srv *ExecutorApi) executorFromLeaseRequest(ctx *armadacontext.Context, req *executorapi.LeaseRequest) *schedulerobjects.Executor {
 	nodes := make([]*schedulerobjects.Node, 0, len(req.Nodes))
 	now := srv.clock.Now().UTC()
+	nodeIdLabelMissing := false
 	for _, nodeInfo := range req.Nodes {
 		if node, err := executorapi.NewNodeFromNodeInfo(nodeInfo, req.ExecutorId, srv.allowedPriorities, now); err != nil {
 			ctx.Logger().WithStacktrace(err).Warnf(
 				"skipping node %s from executor %s", nodeInfo.GetName(), req.GetExecutorId(),
 			)
 		} else {
+			if _, ok := node.Labels[srv.nodeIdLabel]; !ok {
+				nodeIdLabelMissing = true
+			}
 			nodes = append(nodes, node)
 		}
 	}
+	// Node anti-affinity on retries matches against the node id label. A node
+	// without it satisfies every avoidance vacuously, so warn loudly. Once per
+	// executor: the executor must add the label to trackedNodeLabels.
+	if nodeIdLabelMissing {
+		if _, warned := srv.nodeIdLabelWarnings.LoadOrStore(req.ExecutorId, true); !warned {
+			ctx.Warnf(
+				"executor %s reports nodes without the node id label %q: retry node anti-affinity cannot work; add the label to the executor's trackedNodeLabels",
+				req.ExecutorId, srv.nodeIdLabel,
+			)
+		}
+	}
 	return &schedulerobjects.Executor{
 		Id:                req.ExecutorId,
 		Pool:              req.Pool,

internal/scheduler/configuration/configuration.go

@@ -293,6 +293,8 @@ type SchedulingConfig struct {
 	MaximumPerQueueSchedulingBurst int `validate:"gt=0"`
 	// Maximum number of times a job is retried before considered failed.
 	MaxRetries uint
+	// RetryPolicy controls the policy-based retry engine (disabled by default).
+	RetryPolicy RetryPolicyConfig
 	// List of resource names, e.g., []string{"cpu", "memory"}, to consider when computing DominantResourceFairness costs.
 	// Dominant resource fairness is the algorithm used to assign a cost to jobs and queues.
 	DominantResourceFairnessResourcesToConsider []string
@@ -347,13 +349,33 @@ type SchedulingConfig struct {
 	ExperimentalIndicativeShare ExperimentalIndicativeShare
 }
 
+// RetryPolicyConfig controls the scheduler's retry policy behavior.
+type RetryPolicyConfig struct {
+	// Enabled controls whether the retry policy engine is active.
+	Enabled bool
+	// GlobalMaxRetries is the scheduler-wide cap on genuine-failure retries per
+	// job, on top of every policy. It counts retries, not runs: the initial
+	// failure consumes no budget, only subsequent retries do. Preempted and
+	// lease-returned runs are never charged. A value of 0 is the kill switch:
+	// no job is ever retried by the policy engine. There is no unlimited
+	// setting for the global cap.
+	GlobalMaxRetries uint
+	// DefaultPolicyName is the retry policy applied to jobs whose queue has no
+	// policy of its own. It lets an operator turn on retry policies fleet-wide
+	// with a single named policy before per-queue attachment is configured.
+	// Optional: when empty, only queues with an attached policy get engine
+	// decisions and every other queue keeps the existing behaviour.
+	DefaultPolicyName string
+}
+
 const (
 	DuplicateWellKnownNodeTypeErrorMessage              = "duplicate well-known node type name"
 	AwayNodeTypesWithoutPreemptionErrorMessage          = "priority class has away node types but is not preemptible"
 	UnknownWellKnownNodeTypeErrorMessage                = "priority class refers to unknown well-known node type"
 	WildCardWellKnownNodeTypeValue                      = "*"
 	InvalidAwayNodeTypeConditionOperatorErrorMessage    = "away node type condition has invalid operator; must be one of >, <, =="
 	PreemptionRateLimitWithMarketSchedulingErrorMessage = "preemption rate limit is not supported with market scheduling enabled on the same pool"
+	NodeIdLabelNotIndexedErrorMessage                   = "nodeIdLabel must be in indexedNodeLabels when the retry policy engine is enabled, so avoidSameNode retries can match nodes efficiently"
 )
 
 // ResourceType represents a resource the scheduler indexes for efficient lookup.

internal/scheduler/configuration/validation.go

@@ -2,6 +2,7 @@ package configuration
 
 import (
 	"fmt"
+	"slices"
 
 	"github.com/go-playground/validator/v10"
 
@@ -22,6 +23,10 @@ func (c *Configuration) Mutate() (config.Config, error) {
 		c.Scheduling.MaxNewJobSchedulingDuration = c.NewJobsSchedulingTimeout
 	}
 
+	if c.Scheduling.RetryPolicy.Enabled && c.Scheduling.RetryPolicy.GlobalMaxRetries == 0 {
+		log.Warnf("scheduling.retryPolicy.enabled is true but globalMaxRetries is 0: the retry engine is active but will never grant a retry (kill-switch semantics). Set globalMaxRetries above 0 to allow policy retries.")
+	}
+
 	return c, nil
 }
 
@@ -34,6 +39,13 @@ func (c *Configuration) Validate() error {
 func SchedulingConfigValidation(sl validator.StructLevel) {
 	c := sl.Current().Interface().(SchedulingConfig)
 
+	// avoidSameNode retries express node avoidance through nodeIdLabel, and an
+	// unindexed label forces a per-node scan for every job that carries the
+	// anti-affinity. Reject the config instead of running slow.
+	if c.RetryPolicy.Enabled && !slices.Contains(c.IndexedNodeLabels, c.NodeIdLabel) {
+		sl.ReportError(c.IndexedNodeLabels, "IndexedNodeLabels", "", NodeIdLabelNotIndexedErrorMessage, "")
+	}
+
 	for i, pool := range c.Pools {
 		// The preemption rate limit relies on rescheduling evicted jobs before new jobs, which the
 		// market-driven scheduler does not support. Reject the combination rather than silently no-op.

internal/scheduler/configuration/validation_test.go

@@ -273,6 +273,52 @@ func TestValidate_PreemptionRateLimitWithMarketScheduling(t *testing.T) {
 	}
 }
 
+func TestValidate_RetryPolicyRequiresIndexedNodeIdLabel(t *testing.T) {
+	tests := map[string]struct {
+		retryPolicyEnabled bool
+		indexedNodeLabels  []string
+		expectErr          bool
+	}{
+		"retry policy disabled without indexed nodeIdLabel is allowed": {
+			retryPolicyEnabled: false,
+			indexedNodeLabels:  nil,
+			expectErr:          false,
+		},
+		"retry policy enabled with indexed nodeIdLabel is allowed": {
+			retryPolicyEnabled: true,
+			indexedNodeLabels:  []string{"nodeid"},
+			expectErr:          false,
+		},
+		"retry policy enabled without indexed nodeIdLabel is rejected": {
+			retryPolicyEnabled: true,
+			indexedNodeLabels:  []string{"zone"},
+			expectErr:          true,
+		},
+		"retry policy enabled with empty indexedNodeLabels is rejected": {
+			retryPolicyEnabled: true,
+			indexedNodeLabels:  nil,
+			expectErr:          true,
+		},
+	}
+
+	for name, tc := range tests {
+		t.Run(name, func(t *testing.T) {
+			c := createValidMinimalConfig()
+			c.Scheduling.RetryPolicy.Enabled = tc.retryPolicyEnabled
+			c.Scheduling.IndexedNodeLabels = tc.indexedNodeLabels
+
+			err := c.Validate()
+
+			if tc.expectErr {
+				assert.Error(t, err)
+				assert.Contains(t, err.Error(), NodeIdLabelNotIndexedErrorMessage)
+			} else {
+				assert.NoError(t, err)
+			}
+		})
+	}
+}
+
 func TestSchedulingConfigValidate(t *testing.T) {
 	c := Configuration{
 		Scheduling: SchedulingConfig{

internal/scheduler/jobdb/job.go

@@ -786,6 +786,7 @@ func (job *Job) ValidateResourceRequests() error {
 }
 
 // WithNewRun creates a copy of the job with a new run on the given executor.
+// Each run gets a fresh id. Runs are ordered by their creation time.
 func (job *Job) WithNewRun(executor, nodeId, nodeName, pool string, scheduledAtPriority int32) *Job {
 	now := job.jobDb.clock.Now()
 	return job.WithUpdatedRun(job.jobDb.CreateRun(
@@ -861,6 +862,22 @@ func (job *Job) AllRuns() []*JobRun {
 	return maps.Values(job.runsById)
 }
 
+// FailureCount returns the number of runs of this job that genuinely failed.
+// The retry engine charges this count against a policy's retry budgets.
+// Preempted and lease-returned runs are marked failed but never ran to a
+// genuine failure, so they do not count: neither is something the job did.
+// The count derives from run history, which keeps it correct across
+// scheduler restarts.
+func (job *Job) FailureCount() uint32 {
+	count := uint32(0)
+	for _, run := range job.runsById {
+		if run.failed && !run.everPreempted && !run.returned {
+			count++
+		}
+	}
+	return count
+}
+
 // LatestRun returns the currently active job run or nil if there are no runs yet.
 // Callers should either guard against nil values explicitly or via HasRuns.
 func (job *Job) LatestRun() *JobRun {

internal/scheduler/jobdb/job_run.go

@@ -50,6 +50,12 @@ type JobRun struct {
 	preemptReason *string
 	// True if the run has been reported as preempted by the executor.
 	preempted bool
+	// True if the run was ever preempted, even after another terminal state
+	// takes precedence over preempted during reconciliation. Job.FailureCount
+	// reads this to exclude a run that is both failed and preempted (the
+	// API-preemption path marks preempted runs failed to terminate them on the
+	// executor) from the job's genuine-failure count.
+	everPreempted bool
 	// The time at which the run was reported as preempted by the executor.
 	preemptedTime *time.Time
 	// True if the job has been reported as succeeded by the executor.
@@ -159,6 +165,9 @@ func (run *JobRun) Equal(other *JobRun) bool {
 	if run.id != other.id {
 		return false
 	}
+	if run.everPreempted != other.everPreempted {
+		return false
+	}
 	if run.jobId != other.jobId {
 		return false
 	}
@@ -254,6 +263,7 @@ func (jobDb *JobDb) CreateRun(
 		preemptRequested:    preemptRequested,
 		preemptReason:       preemptReason,
 		preempted:           preempted,
+		everPreempted:       preempted,
 		succeeded:           succeeded,
 		failed:              failed,
 		cancelled:           cancelled,
@@ -463,9 +473,19 @@ func (run *JobRun) PreemptedTime() *time.Time {
 func (run *JobRun) WithPreempted(preempted bool) *JobRun {
 	run = run.DeepCopy()
 	run.preempted = preempted
+	if preempted {
+		run.everPreempted = true
+	}
 	return run
 }
 
+// EverPreempted returns true if the run was preempted at any point, even when
+// a later terminal state (typically failed) replaced preempted as the run's
+// single terminal state. Unlike Preempted, this survives WithoutTerminal.
+func (run *JobRun) EverPreempted() bool {
+	return run.everPreempted
+}
+
 func (run *JobRun) WithPreemptedTime(preemptedTime *time.Time) *JobRun {
 	run = run.DeepCopy()
 	run.preemptedTime = preemptedTime

internal/scheduler/jobdb/job_test.go

@@ -422,6 +422,35 @@ func TestJob_TestNumAttempts(t *testing.T) {
 	assert.Equal(t, uint(2), returned3.NumAttempts())
 }
 
+func TestJob_FailureCount(t *testing.T) {
+	newRun := func(mutate func(*JobRun)) *JobRun {
+		run := &JobRun{id: uuid.New().String(), created: baseRun.created}
+		mutate(run)
+		return run
+	}
+	failed := func() *JobRun { return newRun(func(r *JobRun) { r.failed = true }) }
+	// Preempted and lease-returned runs are marked failed but never ran to a
+	// genuine failure. everPreempted is set alongside preempted, mirroring
+	// WithPreempted and CreateRun.
+	preempted := func() *JobRun {
+		return newRun(func(r *JobRun) { r.failed = true; r.preempted = true; r.everPreempted = true })
+	}
+	leaseReturned := func() *JobRun { return newRun(func(r *JobRun) { r.failed = true; r.returned = true }) }
+	succeeded := func() *JobRun { return newRun(func(r *JobRun) { r.succeeded = true }) }
+
+	assert.Equal(t, uint32(0), baseJob.FailureCount())
+
+	job := baseJob.WithUpdatedRun(failed())
+	assert.Equal(t, uint32(1), job.FailureCount())
+
+	// Preempted, lease-returned, and succeeded runs are not genuine failures.
+	job = job.WithUpdatedRun(preempted()).WithUpdatedRun(leaseReturned()).WithUpdatedRun(succeeded())
+	assert.Equal(t, uint32(1), job.FailureCount())
+
+	job = job.WithUpdatedRun(failed())
+	assert.Equal(t, uint32(2), job.FailureCount())
+}
+
 func TestJob_TestRunsById(t *testing.T) {
 	runs := make([]*JobRun, 10)
 	job := baseJob

internal/scheduler/metrics/constants.go

@@ -28,6 +28,8 @@ const (
 	unschedulableReasonLabel = "unschedulable_reason"
 	outcomeLabel             = "outcome"
 	terminationReasonLabel   = "termination_reason"
+	retryPolicyLabel         = "policy"
+	retryDecisionLabel       = "decision"
 
 	SchedulingOutcomeSuccess = "success"
 	SchedulingOutcomeFailure = "failure"

internal/scheduler/metrics/state_metrics.go

@@ -29,6 +29,7 @@ type jobStateMetrics struct {
 	jobErrorsByQueue                          *prometheus.CounterVec
 	jobErrorsByNode                           *prometheus.CounterVec
 	jobResourceSecondsLostToPreemptionByQueue *prometheus.CounterVec
+	retryPolicyDecisionsByQueue               *prometheus.CounterVec
 	allMetrics                                []resettableMetric
 }
 
@@ -108,6 +109,13 @@ func newJobStateMetrics(
 		},
 		[]string{queueLabel, poolLabel, checkpointLabel, resourceLabel},
 	)
+	retryPolicyDecisionsByQueue := prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Name: ArmadaSchedulerMetricsPrefix + "retry_policy_decisions_total",
+			Help: "Retry policy engine decisions at the queue level",
+		},
+		[]string{queueLabel, poolLabel, retryPolicyLabel, retryDecisionLabel},
+	)
 	return &jobStateMetrics{
 		trackedResourceNames:                      trackedResourceNames,
 		jobCheckpointIntervals:                    jobCheckpointIntervals,
@@ -124,6 +132,7 @@ func newJobStateMetrics(
 		jobErrorsByQueue:                          jobErrorsByQueue,
 		jobErrorsByNode:                           jobErrorsByNode,
 		jobResourceSecondsLostToPreemptionByQueue: jobResourceSecondsLostToPreemptionByQueue,
+		retryPolicyDecisionsByQueue:               retryPolicyDecisionsByQueue,
 		allMetrics: []resettableMetric{
 			completedRunDurations,
 			jobStateCounterByQueue,
@@ -135,6 +144,7 @@ func newJobStateMetrics(
 			jobErrorsByQueue,
 			jobErrorsByNode,
 			jobResourceSecondsLostToPreemptionByQueue,
+			retryPolicyDecisionsByQueue,
 		},
 	}
 }
@@ -180,6 +190,14 @@ func (m *jobStateMetrics) ReportJobPreempted(job *jobdb.Job) {
 	}
 }
 
+// ReportRetryPolicyDecision records the final outcome of a retry engine
+// decision. The scheduler reports after the mutation gates, so a granted retry
+// it abandoned counts as what actually happened, not as a retry.
+func (m *jobStateMetrics) ReportRetryPolicyDecision(job *jobdb.Job, policyName string, decision string) {
+	run := job.LatestRun()
+	m.retryPolicyDecisionsByQueue.WithLabelValues(job.Queue(), run.Pool(), policyName, decision).Inc()
+}
+
 func (m *jobStateMetrics) recordPreemptedSecondsLost(job *jobdb.Job, duration float64, checkpointLabel string) {
 	run := job.LatestRun()
 	requests := job.AllResourceRequirements()