expand preemption retries to a general retry mechanism

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

Author: dejanzele

_local/scheduler/config.yaml

@@ -22,5 +22,37 @@ postgres:
 scheduling:
   indexedNodeLabels:
     - "kubernetes.io/hostname"
+  retryPolicy:
+    enabled: true
+    globalMaxRetries: 5
+    # Default policy: no rules = everything fails immediately
+    default:
+      retryLimit: 3
+      rules: []
+    policies:
+      # ML training: retry OOMKilled only
+      ml-training:
+        retryLimit: 5
+        rules:
+          - action: Retry
+            onExitCodes:
+              operator: In
+              values: [137]
+      # Retry-all: traditional batch behavior (retry everything except success)
+      retry-all:
+        retryLimit: 5
+        rules:
+          - action: Retry
+            onExitCodes:
+              operator: NotIn
+              values: [0]
+      # High-retry: for testing global limit caps
+      high-retry:
+        retryLimit: 10
+        rules:
+          - action: Retry
+            onExitCodes:
+              operator: NotIn
+              values: [0]
 metrics:
   port: 9001

cmd/armadactl/cmd/queue.go

@@ -63,13 +63,19 @@ Job priority is evaluated inside queue, queue has its own priority.  Any labels
 				return fmt.Errorf("error converting queue labels to map: %s", err)
 			}
 
+			retryPolicy, err := cmd.Flags().GetString("retry-policy")
+			if err != nil {
+				return fmt.Errorf("error reading retry-policy: %s", err)
+			}
+
 			newQueue, err := queue.NewQueue(&api.Queue{
 				Name:           name,
 				PriorityFactor: priorityFactor,
 				UserOwners:     owners,
 				GroupOwners:    groups,
 				Cordoned:       cordoned,
 				Labels:         labelsAsMap,
+				RetryPolicy:    retryPolicy,
 			})
 			if err != nil {
 				return fmt.Errorf("invalid queue data: %s", err)
@@ -83,6 +89,7 @@ Job priority is evaluated inside queue, queue has its own priority.  Any labels
 	cmd.Flags().StringSlice("group-owners", []string{}, "Comma separated list of queue group owners, defaults to empty list.")
 	cmd.Flags().Bool("cordon", false, "Used to pause scheduling on specified queue. Defaults to false.")
 	cmd.Flags().StringSliceP("labels", "l", []string{}, "Comma separated list of key-value queue labels, for example: armadaproject.io/submitter=airflow. Defaults to empty list.")
+	cmd.Flags().String("retry-policy", "", "Name of the retry policy from scheduler config to use for this queue. Defaults to empty (uses default policy).")
 	return cmd
 }
 
@@ -230,13 +237,19 @@ func queueUpdateCmdWithApp(a *armadactl.App) *cobra.Command {
 				return fmt.Errorf("error converting queue labels to map: %s", err)
 			}
 
+			retryPolicy, err := cmd.Flags().GetString("retry-policy")
+			if err != nil {
+				return fmt.Errorf("error reading retry-policy: %s", err)
+			}
+
 			newQueue, err := queue.NewQueue(&api.Queue{
 				Name:           name,
 				PriorityFactor: priorityFactor,
 				UserOwners:     owners,
 				GroupOwners:    groups,
 				Cordoned:       cordoned,
 				Labels:         labelsAsMap,
+				RetryPolicy:    retryPolicy,
 			})
 			if err != nil {
 				return fmt.Errorf("invalid queue data: %s", err)
@@ -251,5 +264,6 @@ func queueUpdateCmdWithApp(a *armadactl.App) *cobra.Command {
 	cmd.Flags().StringSlice("group-owners", []string{}, "Comma separated list of queue group owners, defaults to empty list.")
 	cmd.Flags().Bool("cordon", false, "Used to pause scheduling on specified queue. Defaults to false.")
 	cmd.Flags().StringSliceP("labels", "l", []string{}, "Comma separated list of key-value queue labels, for example: armadaproject.io/submitter=airflow. Defaults to empty list.")
+	cmd.Flags().String("retry-policy", "", "Name of the retry policy from scheduler config to use for this queue. Defaults to empty (uses default policy).")
 	return cmd
 }

config/scheduler/config.yaml

@@ -107,6 +107,25 @@ scheduling:
   maximumPerQueueSchedulingBurst: 1000
   maxJobSchedulingContextsPerExecutor: 10000
   maxRetries: 3
+  # Retry policy provides fine-grained control over job retry behavior.
+  # When enabled, it takes precedence over maxRetries.
+  # Default behavior: all failures fail immediately (no retry).
+  # Add Retry rules to opt-in to retry behavior for specific conditions.
+  retryPolicy:
+    enabled: false  # Set to true to enable retry policy
+    globalMaxRetries: 20  # Hard cap on total retries across all policies
+    default:
+      retryLimit: 3
+      rules: []  # No rules = all failures fail immediately
+    # Named policies can be referenced by queue's retry_policy field
+    # policies:
+    #   ml-training:
+    #     retryLimit: 10
+    #     rules:
+    #       - action: Retry
+    #         onExitCodes:
+    #           operator: In
+    #           values: [137, 143]
   dominantResourceFairnessResourcesToConsider:
     - "cpu"
     - "memory"

go.mod

@@ -72,6 +72,7 @@ require (
 	github.com/redis/go-redis/v9 v9.17.3
 	github.com/rs/zerolog v1.34.0
 	github.com/segmentio/fasthash v1.0.3
+	github.com/sirupsen/logrus v1.9.4
 	github.com/xitongsys/parquet-go v1.6.2
 	github.com/zalando/go-keyring v0.2.6
 	go.uber.org/atomic v1.11.0
@@ -218,7 +219,6 @@ require (
 	github.com/sagikazarmark/locafero v0.12.0 // indirect
 	github.com/sergi/go-diff v1.4.0 // indirect
 	github.com/shopspring/decimal v1.4.0 // indirect
-	github.com/sirupsen/logrus v1.9.4 // indirect
 	github.com/skeema/knownhosts v1.3.2 // indirect
 	github.com/spaolacci/murmur3 v1.1.0 // indirect
 	github.com/spf13/afero v1.15.0 // indirect

internal/common/database/upsert.go

@@ -130,7 +130,10 @@ func Upsert[T any](ctx *armadacontext.Context, tx pgx.Tx, tableName string, reco
 	// Move those rows into the main table, using ON CONFLICT rules to over-write existing rows.
 	var b strings.Builder
 
-	fmt.Fprintf(&b, "INSERT INTO %s SELECT %s from %s ", tableName, strings.Join(writableNames, ","), tempTableName)
+	// Explicitly specify target columns to handle GENERATED ALWAYS columns correctly.
+	// Without explicit columns, INSERT matches values by position, which fails when
+	// columns are excluded (e.g., GENERATED columns) and remaining values shift positions.
+	fmt.Fprintf(&b, "INSERT INTO %s (%s) SELECT %s from %s ", tableName, strings.Join(writableNames, ","), strings.Join(writableNames, ","), tempTableName)
 	fmt.Fprintf(&b, "ON CONFLICT (%s) DO UPDATE SET ", writableNames[0])
 	for i, name := range writableNames {
 		fmt.Fprintf(&b, "%s = EXCLUDED.%s", name, name)

internal/executor/job/submit.go

@@ -4,9 +4,11 @@ import (
 	"fmt"
 	"regexp"
 	"sync"
+	"time"
 
 	"github.com/pkg/errors"
 	v1 "k8s.io/api/core/v1"
+	networking "k8s.io/api/networking/v1"
 	k8s_errors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
@@ -19,6 +21,15 @@ import (
 	util2 "github.com/armadaproject/armada/internal/executor/util"
 )
 
+// When a job retries, the previous pod's deletion triggers async garbage collection
+// of owned resources (services/ingresses via OwnerReference). Creating new resources
+// may fail with AlreadyExists if GC hasn't completed. We use delete-then-create
+// with retries to handle this race condition.
+const (
+	serviceIngressMaxRetries = 3
+	serviceIngressRetryDelay = 100 * time.Millisecond
+)
+
 type Submitter interface {
 	SubmitJobs(jobsToSubmit []*SubmitJob) []*FailedSubmissionDetails
 }
@@ -130,15 +141,15 @@ func (submitService *SubmitService) submitPod(job *SubmitJob) (*v1.Pod, error) {
 
 	for _, service := range job.Services {
 		service.ObjectMeta.OwnerReferences = []metav1.OwnerReference{util2.CreateOwnerReference(submittedPod)}
-		_, err = submitService.clusterContext.SubmitService(service)
+		err = submitService.createServiceWithRetry(service)
 		if err != nil {
 			return pod, err
 		}
 	}
 
 	for _, ingress := range job.Ingresses {
 		ingress.ObjectMeta.OwnerReferences = []metav1.OwnerReference{util2.CreateOwnerReference(submittedPod)}
-		_, err = submitService.clusterContext.SubmitIngress(ingress)
+		err = submitService.createIngressWithRetry(ingress)
 		if err != nil {
 			return pod, err
 		}
@@ -200,3 +211,59 @@ func (submitService *SubmitService) isRecoverable(err error) bool {
 
 	return false
 }
+
+// createServiceWithRetry creates a service with retry handling for GC race conditions.
+func (submitService *SubmitService) createServiceWithRetry(service *v1.Service) error {
+	return submitService.createWithRetry(
+		"service",
+		service.Namespace,
+		service.Name,
+		func() error { return submitService.clusterContext.DeleteService(service) },
+		func() error { _, err := submitService.clusterContext.SubmitService(service); return err },
+	)
+}
+
+// createIngressWithRetry creates an ingress with retry handling for GC race conditions.
+func (submitService *SubmitService) createIngressWithRetry(ingress *networking.Ingress) error {
+	return submitService.createWithRetry(
+		"ingress",
+		ingress.Namespace,
+		ingress.Name,
+		func() error { return submitService.clusterContext.DeleteIngress(ingress) },
+		func() error { _, err := submitService.clusterContext.SubmitIngress(ingress); return err },
+	)
+}
+
+// createWithRetry implements a delete-then-create pattern with retries to handle
+// race conditions when garbage collection is still running from a previous run.
+func (submitService *SubmitService) createWithRetry(
+	resourceType string,
+	namespace string,
+	name string,
+	deleteFn func() error,
+	createFn func() error,
+) error {
+	for attempt := 0; attempt < serviceIngressMaxRetries; attempt++ {
+		if err := deleteFn(); err != nil {
+			return errors.Wrapf(err, "failed to delete existing %s %s/%s", resourceType, namespace, name)
+		}
+
+		err := createFn()
+		if err == nil {
+			return nil
+		}
+
+		if !k8s_errors.IsAlreadyExists(err) {
+			return err
+		}
+
+		if attempt < serviceIngressMaxRetries-1 {
+			log.Warnf("%s %s/%s already exists after delete (attempt %d/%d), retrying after %v",
+				resourceType, namespace, name, attempt+1, serviceIngressMaxRetries, serviceIngressRetryDelay)
+			time.Sleep(serviceIngressRetryDelay)
+		}
+	}
+
+	return errors.Errorf("failed to create %s %s/%s: still exists after %d delete attempts",
+		resourceType, namespace, name, serviceIngressMaxRetries)
+}