feat: add preempt reason (#4640)

williamvega · JamesMurkin · web-flow · commit a4c8b5346e96 · 2026-02-17T14:06:41.000Z
#### What type of PR is this? Feature #### What this PR does / why we need it Currently the reason provided when preempting a job does not get stored, and the response to the UI is hardcoded for preempted jobs. This change adds the reason to the scheduler db runs table and appends the user to the run reason provided. <img width="1427" height="341" alt="image" src="https://github.com/user-attachments/assets/6f06effa-083e-471b-8285-34a375bcf53e" /> --------- Signed-off-by: williamvega <williamvega1006@gmail.com> Co-authored-by: JamesMurkin <jamesmurkin@hotmail.com>
diff --git a/internal/common/database/upsert.go b/internal/common/database/upsert.go
@@ -130,7 +130,7 @@ 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)
+	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)
diff --git a/internal/scheduler/database/migrations/036_add_preempt_reason.sql b/internal/scheduler/database/migrations/036_add_preempt_reason.sql
@@ -0,0 +1 @@
+ALTER TABLE runs ADD COLUMN preempt_reason TEXT;
diff --git a/internal/scheduler/database/models.go b/internal/scheduler/database/models.go
diff --git a/internal/scheduler/database/query.sql.go b/internal/scheduler/database/query.sql.go
diff --git a/internal/scheduler/database/query/query.sql b/internal/scheduler/database/query/query.sql
@@ -53,7 +53,7 @@ SELECT run_id FROM runs;
 SELECT * FROM runs WHERE serial > $1 AND job_id = ANY(sqlc.arg(job_ids)::text[]) ORDER BY serial;
 
 -- name: MarkJobRunsPreemptRequestedByJobId :exec
-UPDATE runs SET preempt_requested = true WHERE queue = sqlc.arg(queue) and job_set = sqlc.arg(job_set) and job_id = ANY(sqlc.arg(job_ids)::text[]) and terminated = false;
+UPDATE runs SET preempt_requested = true, preempt_reason = sqlc.arg(preempt_reason) WHERE queue = sqlc.arg(queue) and job_set = sqlc.arg(job_set) and job_id = ANY(sqlc.arg(job_ids)::text[]) and terminated = false;
 
 -- name: MarkJobRunsSucceededById :exec
 UPDATE runs SET succeeded = true WHERE run_id = ANY(sqlc.arg(run_ids)::text[]);
diff --git a/internal/scheduler/jobdb/job.go b/internal/scheduler/jobdb/job.go
@@ -784,6 +784,7 @@ func (job *Job) WithNewRun(executor, nodeId, nodeName, pool string, scheduledAtP
 		false,
 		false,
 		false,
+		nil,
 		false,
 		false,
 		false,
diff --git a/internal/scheduler/jobdb/job_run.go b/internal/scheduler/jobdb/job_run.go
@@ -46,6 +46,8 @@ type JobRun struct {
 	runningTime *time.Time
 	// True if a user has requested this run be preempted.
 	preemptRequested bool
+	// The reason provided when preemption was requested.
+	preemptReason *string
 	// True if the run has been reported as preempted by the executor.
 	preempted bool
 	// The time at which the run was reported as preempted by the executor.
@@ -224,6 +226,7 @@ func (jobDb *JobDb) CreateRun(
 	pending bool,
 	running bool,
 	preemptRequested bool,
+	preemptReason *string,
 	preempted bool,
 	succeeded bool,
 	failed bool,
@@ -249,6 +252,7 @@ func (jobDb *JobDb) CreateRun(
 		pending:             pending,
 		running:             running,
 		preemptRequested:    preemptRequested,
+		preemptReason:       preemptReason,
 		preempted:           preempted,
 		succeeded:           succeeded,
 		failed:              failed,
@@ -434,6 +438,18 @@ func (run *JobRun) WithPreemptRequested(preemptRequested bool) *JobRun {
 	return run
 }
 
+// PreemptReason returns the reason provided when preemption was requested, or nil if no reason was provided.
+func (run *JobRun) PreemptReason() *string {
+	return run.preemptReason
+}
+
+// WithPreemptReason returns a copy of the job run with the preemptReason updated.
+func (run *JobRun) WithPreemptReason(preemptReason *string) *JobRun {
+	run = run.DeepCopy()
+	run.preemptReason = preemptReason
+	return run
+}
+
 // Preempted Returns true if the executor has reported the job run as preempted
 func (run *JobRun) Preempted() bool {
 	return run.preempted
diff --git a/internal/scheduler/jobdb/job_run_test.go b/internal/scheduler/jobdb/job_run_test.go
@@ -40,6 +40,7 @@ var (
 		makeTestResourceListFactory(),
 	)
 	scheduledAtPriority = int32(5)
+	testPreemptReason   = "test preempt reason"
 )
 
 func init() {
@@ -59,6 +60,7 @@ var baseJobRun = jobDb.CreateRun(
 	false,
 	false,
 	false,
+	nil,
 	false,
 	false,
 	false,
@@ -136,6 +138,7 @@ func TestDeepCopy(t *testing.T) {
 		true,
 		true,
 		true,
+		&testPreemptReason,
 		true,
 		true,
 		true,
@@ -161,6 +164,7 @@ func TestDeepCopy(t *testing.T) {
 		true,
 		true,
 		true,
+		&testPreemptReason,
 		true,
 		true,
 		true,
diff --git a/internal/scheduler/jobdb/reconciliation.go b/internal/scheduler/jobdb/reconciliation.go
@@ -243,7 +243,7 @@ func (jobDb *JobDb) reconcileRunDifferences(jobRun *JobRun, jobRepoRun *database
 			rst.Running = true
 		}
 		if jobRepoRun.PreemptRequested && !jobRun.PreemptRequested() {
-			jobRun = jobRun.WithPreemptRequested(true)
+			jobRun = jobRun.WithPreemptRequested(true).WithPreemptReason(jobRepoRun.PreemptReason)
 			rst.PreemptionRequested = true
 		}
 		if jobRepoRun.Preempted && !jobRun.Preempted() {
@@ -359,6 +359,7 @@ func (jobDb *JobDb) schedulerRunFromDatabaseRun(dbRun *database.Run) *JobRun {
 		dbRun.Pending,
 		dbRun.Running,
 		dbRun.PreemptRequested,
+		dbRun.PreemptReason,
 		dbRun.Preempted,
 		dbRun.Succeeded,
 		dbRun.Failed,
diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go
@@ -1012,7 +1012,11 @@ func (s *Scheduler) generateUpdateMessagesFromJob(ctx *armadacontext.Context, jo
 			}
 		} else if lastRun.PreemptRequested() && job.PriorityClass().Preemptible {
 			job = job.WithQueued(false).WithFailed(true).WithUpdatedRun(lastRun.WithoutTerminal().WithFailed(true))
-			events = append(events, createEventsForPreemptedJob(job.Id(), lastRun.Id(), "Preempted - preemption requested via API", s.clock.Now())...)
+			reason := "Preempted - preemption requested via API"
+			if lastRun.PreemptReason() != nil && *lastRun.PreemptReason() != "" {
+				reason = *lastRun.PreemptReason()
+			}
+			events = append(events, createEventsForPreemptedJob(job.Id(), lastRun.Id(), reason, s.clock.Now())...)
 			s.metrics.ReportJobPreemptedWithType(job, schedulercontext.PreemptedViaApi)
 		}
 	}
diff --git a/internal/scheduler/scheduler_test.go b/internal/scheduler/scheduler_test.go
@@ -243,6 +243,7 @@ var returnedOnceLeasedJob = testfixtures.NewJob(
 	false,
 	false,
 	false,
+	nil,
 	false,
 	false,
 	true,
@@ -328,6 +329,7 @@ var (
 		false,
 		false,
 		false,
+		nil,
 		false,
 		false,
 		true,
@@ -1474,6 +1476,7 @@ func TestScheduler_TestSyncInitialState(t *testing.T) {
 						false,
 						false,
 						false,
+						nil,
 						false,
 						false,
 						false,
@@ -1634,6 +1637,7 @@ func TestScheduler_TestSyncState(t *testing.T) {
 						false,
 						false,
 						false,
+						nil,
 						false,
 						false,
 						false,
diff --git a/internal/scheduleringester/dbops.go b/internal/scheduleringester/dbops.go
@@ -203,7 +203,7 @@ type (
 	UpdateJobQueuedState           map[string]*JobQueuedStateUpdate
 	MarkRunsSucceeded              map[string]time.Time
 	MarkRunsFailed                 map[string]*JobRunFailed
-	MarkRunsForJobPreemptRequested map[JobSetKey][]string
+	MarkRunsForJobPreemptRequested map[JobSetKey]map[string]string
 	MarkRunsRunning                map[string]time.Time
 	MarkRunsPending                map[string]time.Time
 	MarkRunsPreempted              map[string]time.Time
@@ -289,7 +289,21 @@ func (a MarkJobsCancelRequested) Merge(b DbOperation) bool {
 }
 
 func (a MarkRunsForJobPreemptRequested) Merge(b DbOperation) bool {
-	return mergeListMaps(a, b)
+	switch op := b.(type) {
+	case MarkRunsForJobPreemptRequested:
+		for k, v := range op {
+			if _, present := a[k]; present {
+				// merge the inner maps
+				for jobId, reason := range v {
+					a[k][jobId] = reason
+				}
+			} else {
+				a[k] = v
+			}
+		}
+		return true
+	}
+	return false
 }
 
 func (a UpdateJobSchedulingInfo) Merge(b DbOperation) bool {
diff --git a/internal/scheduleringester/dbops_test.go b/internal/scheduleringester/dbops_test.go
@@ -36,10 +36,10 @@ func TestMerge_MarkRunsForJobPreemptRequested(t *testing.T) {
 	jobId2 := util.NewULID()
 	jobId3 := util.NewULID()
 	jobId4 := util.NewULID()
-	markPreemptRequested1 := MarkRunsForJobPreemptRequested{JobSetKey{queue: testQueueName, jobSet: "set1"}: []string{jobId1}}
-	markPreemptRequested2 := MarkRunsForJobPreemptRequested{JobSetKey{queue: testQueueName, jobSet: "set1"}: []string{jobId2}}
-	markPreemptRequested3 := MarkRunsForJobPreemptRequested{JobSetKey{queue: testQueueName, jobSet: "set2"}: []string{jobId3}}
-	markPreemptRequested4 := MarkRunsForJobPreemptRequested{JobSetKey{queue: "test-queue-2", jobSet: "set1"}: []string{jobId4}}
+	markPreemptRequested1 := MarkRunsForJobPreemptRequested{JobSetKey{queue: testQueueName, jobSet: "set1"}: map[string]string{jobId1: "test-preempt-reason"}}
+	markPreemptRequested2 := MarkRunsForJobPreemptRequested{JobSetKey{queue: testQueueName, jobSet: "set1"}: map[string]string{jobId2: "test-preempt-reason"}}
+	markPreemptRequested3 := MarkRunsForJobPreemptRequested{JobSetKey{queue: testQueueName, jobSet: "set2"}: map[string]string{jobId3: "test-preempt-reason"}}
+	markPreemptRequested4 := MarkRunsForJobPreemptRequested{JobSetKey{queue: "test-queue-2", jobSet: "set1"}: map[string]string{jobId4: "test-preempt-reason"}}
 
 	ok := markPreemptRequested1.Merge(markPreemptRequested2)
 	assert.True(t, ok)
@@ -51,11 +51,11 @@ func TestMerge_MarkRunsForJobPreemptRequested(t *testing.T) {
 	assert.Equal(t,
 		MarkRunsForJobPreemptRequested{
 			// When jobset key matches, values are merged
-			JobSetKey{queue: testQueueName, jobSet: "set1"}: []string{jobId1, jobId2},
+			JobSetKey{queue: testQueueName, jobSet: "set1"}: map[string]string{jobId1: "test-preempt-reason", jobId2: "test-preempt-reason"},
 			// Does not merge values when different jobset in key
-			JobSetKey{queue: testQueueName, jobSet: "set2"}: []string{jobId3},
+			JobSetKey{queue: testQueueName, jobSet: "set2"}: map[string]string{jobId3: "test-preempt-reason"},
 			// Does not merge values when different queue in key
-			JobSetKey{queue: "test-queue-2", jobSet: "set1"}: []string{jobId4},
+			JobSetKey{queue: "test-queue-2", jobSet: "set1"}: map[string]string{jobId4: "test-preempt-reason"},
 		},
 		markPreemptRequested1)
 }
@@ -225,12 +225,12 @@ func TestDbOperationOptimisation(t *testing.T) {
 			}, // 4
 		}},
 		"MarkRunsForJobPreemptRequested": {N: 2, Ops: []DbOperation{
-			InsertJobs{jobIds[0]: &schedulerdb.Job{JobID: jobIds[0], Queue: testQueueName, JobSet: "set1"}},      // 1
-			InsertJobs{jobIds[1]: &schedulerdb.Job{JobID: jobIds[1], Queue: testQueueName, JobSet: "set1"}},      // 1
-			MarkRunsForJobPreemptRequested{JobSetKey{queue: testQueueName, jobSet: "set1"}: []string{jobIds[0]}}, // 2                                                            // 2
-			MarkRunsForJobPreemptRequested{JobSetKey{queue: testQueueName, jobSet: "set1"}: []string{jobIds[1]}}, // 2                                                            // 2
-			InsertJobs{jobIds[2]: &schedulerdb.Job{JobID: jobIds[2], Queue: testQueueName, JobSet: "set1"}},      // 1
-			MarkRunsForJobPreemptRequested{JobSetKey{queue: testQueueName, jobSet: "set1"}: []string{jobIds[2]}}, // 2
+			InsertJobs{jobIds[0]: &schedulerdb.Job{JobID: jobIds[0], Queue: testQueueName, JobSet: "set1"}},                              // 1
+			InsertJobs{jobIds[1]: &schedulerdb.Job{JobID: jobIds[1], Queue: testQueueName, JobSet: "set1"}},                              // 1
+			MarkRunsForJobPreemptRequested{JobSetKey{queue: testQueueName, jobSet: "set1"}: map[string]string{jobIds[0]: "test-reason"}}, // 2                                                            // 2
+			MarkRunsForJobPreemptRequested{JobSetKey{queue: testQueueName, jobSet: "set1"}: map[string]string{jobIds[1]: "test-reason"}}, // 2                                                            // 2
+			InsertJobs{jobIds[2]: &schedulerdb.Job{JobID: jobIds[2], Queue: testQueueName, JobSet: "set1"}},                              // 1
+			MarkRunsForJobPreemptRequested{JobSetKey{queue: testQueueName, jobSet: "set1"}: map[string]string{jobIds[2]: "test-reason"}}, // 2
 		}},
 		"MarkJobsSucceeded": {N: 2, Ops: []DbOperation{
 			InsertJobs{jobIds[0]: &schedulerdb.Job{JobID: jobIds[0]}}, // 1
diff --git a/internal/scheduleringester/instructions.go b/internal/scheduleringester/instructions.go
@@ -316,11 +316,18 @@ func (c *JobSetEventsInstructionConverter) handleJobErrors(jobErrors *armadaeven
 }
 
 func (c *JobSetEventsInstructionConverter) handleJobPreemptionRequested(preemptionRequested *armadaevents.JobPreemptionRequested, meta eventSequenceCommon) ([]DbOperation, error) {
+	// combine the preemption reason and user
+	reasonWithUser := preemptionRequested.Reason
+	if meta.user != "" {
+		reasonWithUser += " | Preempted by " + meta.user
+	}
 	return []DbOperation{MarkRunsForJobPreemptRequested{
 		JobSetKey{
 			queue:  meta.queue,
 			jobSet: meta.jobset,
-		}: []string{preemptionRequested.JobId},
+		}: map[string]string{
+			preemptionRequested.JobId: reasonWithUser,
+		},
 	}}, nil
 }
 
diff --git a/internal/scheduleringester/instructions_test.go b/internal/scheduleringester/instructions_test.go
@@ -103,7 +103,7 @@ func TestConvertEventSequence(t *testing.T) {
 		},
 		"job preemption requested": {
 			events:   []*armadaevents.EventSequence_Event{f.JobPreemptionRequested},
-			expected: []DbOperation{MarkRunsForJobPreemptRequested{JobSetKey{queue: f.Queue, jobSet: f.JobsetName}: []string{f.JobId}}},
+			expected: []DbOperation{MarkRunsForJobPreemptRequested{JobSetKey{queue: f.Queue, jobSet: f.JobsetName}: map[string]string{f.JobId: " | Preempted by " + f.UserId}}},
 		},
 		"job run preempted": {
 			events:   []*armadaevents.EventSequence_Event{f.JobRunPreempted},
diff --git a/internal/scheduleringester/schedulerdb.go b/internal/scheduleringester/schedulerdb.go
@@ -228,13 +228,27 @@ func (s *SchedulerDb) WriteDbOp(ctx *armadacontext.Context, tx pgx.Tx, op DbOper
 			return errors.WithStack(err)
 		}
 	case MarkRunsForJobPreemptRequested:
-		for key, value := range o {
-			params := schedulerdb.MarkJobRunsPreemptRequestedByJobIdParams{
-				Queue:  key.queue,
-				JobSet: key.jobSet,
-				JobIds: value,
+		const batchSize = 5000
+		markRunsPreemptRequestedSqlStatement := "UPDATE runs SET preempt_requested = true, preempt_reason = $1 WHERE queue = $2 and job_set = $3 and job_id = ANY($4::text[]) and terminated = false"
+
+		for key, preemptReason := range o {
+			// group by reason
+			reasonToJobIds := make(map[string][]string)
+			for jobId, reason := range preemptReason {
+				reasonToJobIds[reason] = append(reasonToJobIds[reason], jobId)
+			}
+
+			// batch updates for jobs with the same reason
+			batch := &pgx.Batch{}
+			for reason, jobIds := range reasonToJobIds {
+				for i := 0; i < len(jobIds); i += batchSize {
+					end := min(i+batchSize, len(jobIds))
+					jobIdBatch := jobIds[i:end]
+					batch.Queue(markRunsPreemptRequestedSqlStatement, reason, key.queue, key.jobSet, jobIdBatch)
+				}
 			}
-			err := queries.MarkJobRunsPreemptRequestedByJobId(ctx, params)
+
+			err := execBatch(ctx, tx, batch)
 			if err != nil {
 				return errors.WithStack(err)
 			}
diff --git a/internal/scheduleringester/schedulerdb_test.go b/internal/scheduleringester/schedulerdb_test.go
diff --git a/testsuite/testcases/basic/gang.yaml b/testsuite/testcases/basic/gang.yaml
diff --git a/testsuite/testcases/basic/service_headless.yaml b/testsuite/testcases/basic/service_headless.yaml

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+ALTER TABLE runs ADD COLUMN preempt_reason TEXT;`
-Original file line number
+Diff line change
 		false,
 		false,
 		false,
 +		nil,
 		false,
 		false,
 		false,
Original file line number	Diff line number	Diff line change
`@@ -1012,7 +1012,11 @@ func (s Scheduler) generateUpdateMessagesFromJob(ctx armadacontext.Context, jo`
`1012`	`1012`	`}`
`1013`	`1013`	`} else if lastRun.PreemptRequested() && job.PriorityClass().Preemptible {`
`1014`	`1014`	`job = job.WithQueued(false).WithFailed(true).WithUpdatedRun(lastRun.WithoutTerminal().WithFailed(true))`
`1015`		`- events = append(events, createEventsForPreemptedJob(job.Id(), lastRun.Id(), "Preempted - preemption requested via API", s.clock.Now())...)`
	`1015`	`+ reason := "Preempted - preemption requested via API"`
	`1016`	`+ if lastRun.PreemptReason() != nil && *lastRun.PreemptReason() != "" {`
	`1017`	`+ reason = *lastRun.PreemptReason()`
	`1018`	`+ }`
	`1019`	`+ events = append(events, createEventsForPreemptedJob(job.Id(), lastRun.Id(), reason, s.clock.Now())...)`
`1016`	`1020`	`s.metrics.ReportJobPreemptedWithType(job, schedulercontext.PreemptedViaApi)`
`1017`	`1021`	`}`
`1018`	`1022`	`}`