Repository: armadaproject/armada
Pull request #2116 (2023-02-14 15:58:31)

[Pulsar Scheduler] Better Cancel By Jobset
Support cancel by jobset and send cancelling messages when this is the case


Patch:
diff --git a/internal/scheduler/database/job_repository.go b/internal/scheduler/database/job_repository.go
index 567e28dc..94a4a232 100644
--- a/internal/scheduler/database/job_repository.go
+++ b/internal/scheduler/database/job_repository.go
@@ -111,17 +111,18 @@ func (r *PostgresJobRepository) FetchJobUpdates(ctx context.Context, jobSerial i
 		updatedJobs = make([]Job, len(updatedJobRows))
 		for i, row := range updatedJobRows {
 			updatedJobs[i] = Job{
-				JobID:           row.JobID,
-				JobSet:          row.JobSet,
-				Queue:           row.Queue,
-				Priority:        row.Priority,
-				Submitted:       row.Submitted,
-				CancelRequested: row.CancelRequested,
-				Cancelled:       row.Cancelled,
-				Succeeded:       row.Succeeded,
-				Failed:          row.Failed,
-				SchedulingInfo:  row.SchedulingInfo,
-				Serial:          row.Serial,
+				JobID:                   row.JobID,
+				JobSet:                  row.JobSet,
+				Queue:                   row.Queue,
+				Priority:                row.Priority,
+				Submitted:               row.Submitted,
+				CancelRequested:         row.CancelRequested,
+				Cancelled:               row.Cancelled,
+				CancelByJobsetRequested: row.CancelByJobsetRequested,
+				Succeeded:               row.Succeeded,
+				Failed:                  row.Failed,
+				SchedulingInfo:          row.SchedulingInfo,
+				Serial:                  row.Serial,
 			}
 		}
 
diff --git a/internal/scheduler/database/migrations/001_initialize_schema.up.sql b/internal/scheduler/database/migrations/001_initialize_schema.up.sql
index d80e10b7..150f8d2b 100644
--- a/internal/scheduler/database/migrations/001_initialize_schema.up.sql
+++ b/internal/scheduler/database/migrations/001_initialize_schema.up.sql
@@ -16,6 +16,8 @@ CREATE TABLE jobs (
     cancel_requested boolean NOT NULL DEFAULT false,
     -- Indicates if this job has been cancelled
     cancelled boolean NOT NULL DEFAULT false,
+    -- Indicates if this job has has been cancelled as part of a cancel jobset op
+    cancel_by_jobset_requested boolean NOT NULL DEFAULT false,
     -- Set to true when a JobSucceeded event has been received for this job by the ingester.
     succeeded boolean NOT NULL DEFAULT false,
     -- Set to true when a terminal JobErrors event has been received for this job by the ingester.
diff --git a/internal/scheduler/database/models.go b/internal/scheduler/database/models.go
index 0a7c376f..e8c7b722 100644
--- a/internal/scheduler/database/models.go
+++ b/internal/scheduler/database/models.go
@@ -17,21 +17,22 @@ type Executor struct {
 }
 
 type Job struct {
-	JobID           string    `db:"job_id"`
-	JobSet          string    `db:"job_set"`
-	Queue           string    `db:"queue"`
-	UserID          string    `db:"user_id"`
-	Submitted       int64     `db:"submitted"`
-	Groups          []byte    `db:"groups"`
-	Priority        int64     `db:"priority"`
-	CancelRequested bool      `db:"cancel_requested"`
-	Cancelled       bool      `db:"cancelled"`
-	Succeeded       bool      `db:"succeeded"`
-	Failed          bool      `db:"failed"`
-	SubmitMessage   []byte    `db:"submit_message"`
-	SchedulingInfo  []byte    `db:"scheduling_info"`
-	Serial          int64     `db:"serial"`
-	LastModified    time.Time `db:"last_modified"`
+	JobID                   string    `db:"job_id"`
+	JobSet                  string    `db:"job_set"`
+	Queue                   string    `db:"queue"`
+	UserID                  string    `db:"user_id"`
+	Submitted               int64     `db:"submitted"`
+	Groups                  []byte    `db:"groups"`
+	Priority                int64     `db:"priority"`
+	CancelRequested         bool      `db:"cancel_requested"`
+	Cancelled               bool      `db:"cancelled"`
+	CancelByJobsetRequested bool      `db:"cancel_by_jobset_requested"`
+	Succeeded               bool      `db:"succeeded"`
+	Failed                  bool      `db:"failed"`
+	SubmitMessage           []byte    `db:"submit_message"`
+	SchedulingInfo          []byte    `db:"scheduling_info"`
+	Serial                  int64     `db:"serial"`
+	LastModified            time.Time `db:"last_modified"`
 }
 
 type JobRunError struct {
diff --git a/internal/scheduler/database/query.sql.go b/internal/scheduler/database/query.sql.go
index 2ffa6f6d..d13076c1 100644
--- a/internal/scheduler/database/query.sql.go
+++ b/internal/scheduler/database/query.sql.go
@@ -118,7 +118,7 @@ func (q *Queries) MarkJobsCancelRequestedById(ctx context.Context, jobIds []stri
 }
 
 const markJobsCancelRequestedBySets = `-- name: MarkJobsCancelRequestedBySets :exec
-UPDATE jobs SET cancel_requested = true WHERE job_set = ANY($1::text[])
+UPDATE jobs SET cancel_by_jobset_requested = true WHERE job_set = ANY($1::text[])
 `
 
 func (q *Queries) MarkJobsCancelRequestedBySets(ctx context.Context, jobSets []string) error {
@@ -354,7 +354,7 @@ func (q *Queries) SelectJobsForExecutor(ctx context.Context, arg SelectJobsForEx
 }
 
 const selectNewJobs = `-- name: SelectNewJobs :many
-SELECT job_id, job_set, queue, user_id, submitted, groups, priority, cancel_requested, cancelled, succeeded, failed, submit_message, scheduling_info, serial, last_modified FROM jobs WHERE serial > $1 ORDER BY serial LIMIT $2
+SELECT job_id, job_set, queue, user_id, submitted, groups, priority, cancel_requested, cancelled, cancel_by_jobset_requested, succeeded, failed, submit_message, scheduling_info, serial, last_modified FROM jobs WHERE serial > $1 ORDER BY serial LIMIT $2
 `
 
 type SelectNewJobsParams struct {
@@ -381,6 +381,7 @@ func (q *Queries) SelectNewJobs(ctx context.Context, arg SelectNewJobsParams) ([
 			&i.Priority,
 			&i.CancelRequested,
 			&i.Cancelled,
+			&i.CancelByJobsetRequested,
 			&i.Succeeded,
 			&i.Failed,
 			&i.SubmitMessage,
@@ -508,7 +509,7 @@ func (q *Queries) SelectRunErrorsById(ctx context.Context, runIds []uuid.UUID) (
 }
 
 const selectUpdatedJobs = `-- name: SelectUpdatedJobs :many
-SELECT job_id, job_set, queue, priority, submitted, cancel_requested, cancelled, succeeded, failed, scheduling_info, serial FROM jobs WHERE serial > $1 ORDER BY serial LIMIT $2
+SELECT job_id, job_set, queue, priority, submitted, cancel_requested, cancel_by_jobset_requested, cancelled, succeeded, failed, scheduling_info, serial FROM jobs WHERE serial > $1 ORDER BY serial LIMIT $2
 `
 
 type SelectUpdatedJobsParams struct {
@@ -517,17 +518,18 @@ type SelectUpdatedJobsParams struct {
 }
 
 type SelectUpdatedJobsRow struct {
-	JobID           string `db:"job_id"`
-	JobSet          string `db:"job_set"`
-	Queue           string `db:"queue"`
-	Priority        int64  `db:"priority"`
-	Submitted       int64  `db:"submitted"`
-	CancelRequested bool   `db:"cancel_requested"`
-	Cancelled       bool   `db:"cancelled"`
-	Succeeded       bool   `db:"succeeded"`
-	Failed          bool   `db:"failed"`
-	SchedulingInfo  []byte `db:"scheduling_info"`
-	Serial          int64  `db:"serial"`
+	JobID                   string `db:"job_id"`
+	JobSet                  string `db:"job_set"`
+	Queue                   string `db:"queue"`
+	Priority                int64  `db:"priority"`
+	Submitted               int64  `db:"submitted"`
+	CancelRequested         bool   `db:"cancel_requested"`
+	CancelByJobsetRequested bool   `db:"cancel_by_jobset_requested"`
+	Cancelled               bool   `db:"cancelled"`
+	Succeeded               bool   `db:"succeeded"`
+	Failed                  bool   `db:"failed"`
+	SchedulingInfo          []byte `db:"scheduling_info"`
+	Serial                  int64  `db:"serial"`
 }
 
 func (q *Queries) SelectUpdatedJobs(ctx context.Context, arg SelectUpdatedJobsParams) ([]SelectUpdatedJobsRow, error) {
@@ -546,6 +548,7 @@ func (q *Queries) SelectUpdatedJobs(ctx context.Context, arg SelectUpdatedJobsPa
 			&i.Priority,
 			&i.Submitted,
 			&i.CancelRequested,
+			&i.CancelByJobsetRequested,
 			&i.Cancelled,
 			&i.Succeeded,
 			&i.Failed,
diff --git a/internal/scheduler/database/query/query.sql b/internal/scheduler/database/query/query.sql
index e7239a70..9281d431 100644
--- a/internal/scheduler/database/query/query.sql
+++ b/internal/scheduler/database/query/query.sql
@@ -5,13 +5,13 @@ SELECT * FROM jobs WHERE serial > $1 ORDER BY serial LIMIT $2;
 SELECT job_id FROM jobs;
 
 -- name: SelectUpdatedJobs :many
-SELECT job_id, job_set, queue, priority, submitted, cancel_requested, cancelled, succeeded, failed, scheduling_info, serial FROM jobs WHERE serial > $1 ORDER BY serial LIMIT $2;
+SELECT job_id, job_set, queue, priority, submitted, cancel_requested, cancel_by_jobset_requested, cancelled, succeeded, failed, scheduling_info, serial FROM jobs WHERE serial > $1 ORDER BY serial LIMIT $2;
 
 -- name: UpdateJobPriorityByJobSet :exec
 UPDATE jobs SET priority = $1 WHERE job_set = $2;
 
 -- name: MarkJobsCancelRequestedBySets :exec
-UPDATE jobs SET cancel_requested = true WHERE job_set = ANY(sqlc.arg(job_sets)::text[]);
+UPDATE jobs SET cancel_by_jobset_requested = true WHERE job_set = ANY(sqlc.arg(job_sets)::text[]);
 
 -- name: MarkJobsSucceededById :exec
 UPDATE jobs SET succeeded = true WHERE job_id = ANY(sqlc.arg(job_ids)::text[]);
diff --git a/internal/scheduler/jobdb/job.go b/internal/scheduler/jobdb/job.go
index a578487c..a6eae9cd 100644
--- a/internal/scheduler/jobdb/job.go
+++ b/internal/scheduler/jobdb/job.go
@@ -34,6 +34,8 @@ type Job struct {
 	jobSchedulingInfo *schedulerobjects.JobSchedulingInfo
 	// True if the user has requested this job be cancelled
 	cancelRequested bool
+	// True if the user has requested this job's jobset be cancelled
+	cancelByJobsetRequested bool
 	// True if the scheduler has cancelled the job
 	cancelled bool
 	// True if the scheduler has failed the job
@@ -50,27 +52,29 @@ type Job struct {
 
 // NewJob creates a new scheduler job
 func NewJob(
-	id string,
+	jobId string,
 	jobset string,
 	queue string,
 	priority uint32,
 	schedulingInfo *schedulerobjects.JobSchedulingInfo,
 	cancelRequested bool,
+	cancelByJobsetRequested bool,
 	cancelled bool,
 	created int64,
 ) *Job {
 	return &Job{
-		id:                id,
-		jobset:            jobset,
-		queue:             queue,
-		queued:            true,
-		priority:          priority,
-		requestedPriority: priority,
-		jobSchedulingInfo: schedulingInfo,
-		cancelRequested:   cancelRequested,
-		cancelled:         cancelled,
-		created:           created,
-		runsById:          map[uuid.UUID]*JobRun{},
+		id:                      jobId,
+		jobset:                  jobset,
+		queue:                   queue,
+		queued:                  true,
+		priority:                priority,
+		requestedPriority:       priority,
+		jobSchedulingInfo:       schedulingInfo,
+		cancelRequested:         cancelRequested,
+		cancelByJobsetRequested: cancelByJobsetRequested,
+		cancelled:               cancelled,
+		created:                 created,
+		runsById:                map[uuid.UUID]*JobRun{},
 	}
 }
 
@@ -158,6 +162,11 @@ func (job *Job) CancelRequested() bool {
 	return job.cancelRequested
 }
 
+// CancelByJobsetRequested returns true if the user has requested this job's jobset be cancelled.
+func (job *Job) CancelByJobsetRequested() bool {
+	return job.cancelByJobsetRequested
+}
+
 // WithCancelRequested returns a copy of the job with the cancelRequested status updated.
 func (job *Job) WithCancelRequested(cancelRequested bool) *Job {
 	j := copyJob(*job)
@@ -165,6 +174,13 @@ func (job *Job) WithCancelRequested(cancelRequested bool) *Job {
 	return j
 }
 
+// WithCancelByJobsetRequested returns a copy of the job with the cancelByJobsetRequested status updated.
+func (job *Job) WithCancelByJobsetRequested(cancelByJobsetRequested bool) *Job {
+	j := copyJob(*job)
+	j.cancelByJobsetRequested = cancelByJobsetRequested
+	return j
+}
+
 // Cancelled Returns true if the scheduler has cancelled the job
 func (job *Job) Cancelled() bool {
 	return job.cancelled
diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go
index e6bd1583..5b551b2d 100644
--- a/internal/scheduler/scheduler.go
+++ b/internal/scheduler/scheduler.go
@@ -400,6 +400,21 @@ func (s *Scheduler) generateUpdateMessagesFromJob(job *jobdb.Job, jobRunErrors m
 			},
 		}
 		events = append(events, cancel)
+	} else if job.CancelByJobsetRequested() {
+		job = job.WithCancelled(true).WithQueued(false)
+		cancelRequest := &armadaevents.EventSequence_Event{
+			Created: s.now(),
+			Event: &armadaevents.EventSequence_Event_CancelJob{
+				CancelJob: &armadaevents.CancelJob{JobId: jobId},
+			},
+		}
+		cancel := &armadaevents.EventSequence_Event{
+			Created: s.now(),
+			Event: &armadaevents.EventSequence_Event_CancelledJob{
+				CancelledJob: &armadaevents.CancelledJob{JobId: jobId},
+			},
+		}
+		events = append(events, cancelRequest, cancel)
 	} else if job.HasRuns() {
 		lastRun := job.LatestRun()
 		// InTerminalState states. Can only have one of these
@@ -637,6 +652,7 @@ func (s *Scheduler) createSchedulerJob(dbJob *database.Job) (*jobdb.Job, error)
 		uint32(dbJob.Priority),
 		schedulingInfo,
 		dbJob.CancelRequested,
+		dbJob.CancelByJobsetRequested,
 		dbJob.Cancelled,
 		dbJob.Submitted,
 	), nil
@@ -692,6 +708,9 @@ func updateSchedulerJob(job *jobdb.Job, dbJob *database.Job) *jobdb.Job {
 	if dbJob.CancelRequested && !job.CancelRequested() {
 		job = job.WithCancelRequested(true)
 	}
+	if dbJob.CancelByJobsetRequested && !job.CancelByJobsetRequested() {
+		job = job.WithCancelByJobsetRequested(true)
+	}
 	if dbJob.Cancelled && !job.Cancelled() {
 		job = job.WithCancelled(true)
 	}


Test patch:
diff --git a/internal/scheduler/database/job_repository_test.go b/internal/scheduler/database/job_repository_test.go
index 87202720..5b28991e 100644
--- a/internal/scheduler/database/job_repository_test.go
+++ b/internal/scheduler/database/job_repository_test.go
@@ -256,33 +256,35 @@ func createTestJobs(numJobs int) ([]Job, []Job) {
 
 	for i := 0; i < numJobs; i++ {
 		dbJobs[i] = Job{
-			JobID:           util.NewULID(),
-			JobSet:          "test-jobset",
-			Queue:           "test-queue",
-			Submitted:       int64(i),
-			Priority:        int64(i),
-			CancelRequested: true,
-			Cancelled:       true,
-			Succeeded:       true,
-			Failed:          true,
-			SchedulingInfo:  []byte{byte(i)},
-			SubmitMessage:   []byte{},
+			JobID:                   util.NewULID(),
+			JobSet:                  "test-jobset",
+			Queue:                   "test-queue",
+			Submitted:               int64(i),
+			Priority:                int64(i),
+			CancelRequested:         true,
+			Cancelled:               true,
+			CancelByJobsetRequested: true,
+			Succeeded:               true,
+			Failed:                  true,
+			SchedulingInfo:          []byte{byte(i)},
+			SubmitMessage:           []byte{},
 		}
 	}
 
 	for i, job := range dbJobs {
 		expectedJobs[i] = Job{
-			JobID:           job.JobID,
-			JobSet:          job.JobSet,
-			Queue:           job.Queue,
-			Submitted:       job.Submitted,
-			Priority:        job.Priority,
-			CancelRequested: job.CancelRequested,
-			Cancelled:       job.Cancelled,
-			Succeeded:       job.Succeeded,
-			Failed:          job.Failed,
-			SchedulingInfo:  job.SchedulingInfo,
-			Serial:          int64(i + 1),
+			JobID:                   job.JobID,
+			JobSet:                  job.JobSet,
+			Queue:                   job.Queue,
+			Submitted:               job.Submitted,
+			Priority:                job.Priority,
+			CancelRequested:         job.CancelRequested,
+			CancelByJobsetRequested: true,
+			Cancelled:               job.Cancelled,
+			Succeeded:               job.Succeeded,
+			Failed:                  job.Failed,
+			SchedulingInfo:          job.SchedulingInfo,
+			Serial:                  int64(i + 1),
 		}
 	}
 	return dbJobs, expectedJobs
diff --git a/internal/scheduler/jobdb/job_test.go b/internal/scheduler/jobdb/job_test.go
index de129ee7..0b5eb783 100644
--- a/internal/scheduler/jobdb/job_test.go
+++ b/internal/scheduler/jobdb/job_test.go
@@ -31,6 +31,7 @@ var baseJob = NewJob(
 	schedulingInfo,
 	false,
 	false,
+	false,
 	3)
 
 var baseRun = &JobRun{
@@ -82,6 +83,12 @@ func TestJob_TestCancelRequested(t *testing.T) {
 	assert.Equal(t, true, newJob.CancelRequested())
 }
 
+func TestJob_TestCancelByJobsetRequested(t *testing.T) {
+	newJob := baseJob.WithCancelByJobsetRequested(true)
+	assert.Equal(t, false, baseJob.CancelByJobsetRequested())
+	assert.Equal(t, true, newJob.CancelByJobsetRequested())
+}
+
 func TestJob_TestCancelled(t *testing.T) {
 	newJob := baseJob.WithCancelled(true)
 	assert.Equal(t, false, baseJob.Cancelled())
diff --git a/internal/scheduler/scheduler_test.go b/internal/scheduler/scheduler_test.go
index cbd7fa36..b753763d 100644
--- a/internal/scheduler/scheduler_test.go
+++ b/internal/scheduler/scheduler_test.go
@@ -38,6 +38,7 @@ var queuedJob = jobdb.NewJob(
 	schedulingInfo,
 	false,
 	false,
+	false,
 	1)
 
 var leasedJob = jobdb.NewJob(
@@ -48,6 +49,7 @@ var leasedJob = jobdb.NewJob(
 	schedulingInfo,
 	false,
 	false,
+	false,
 	1).WithQueued(false).WithNewRun("testExecutor")
 
 // Test a single scheduler cycle
diff --git a/internal/scheduler/scheduling_algo_test.go b/internal/scheduler/scheduling_algo_test.go
index 0c951911..3a7263dc 100644
--- a/internal/scheduler/scheduling_algo_test.go
+++ b/internal/scheduler/scheduling_algo_test.go
@@ -262,6 +262,7 @@ func OneCpuJob(creationTime int64) *jobdb.Job {
 		schedulingInfo,
 		false,
 		false,
+		false,
 		creationTime).WithQueued(true)
 }
 
diff --git a/internal/scheduleringester/schedulerdb_test.go b/internal/scheduleringester/schedulerdb_test.go
index d52a5cfb..79553b6d 100644
--- a/internal/scheduleringester/schedulerdb_test.go
+++ b/internal/scheduleringester/schedulerdb_test.go
@@ -344,7 +344,7 @@ func assertOpSuccess(t *testing.T, schedulerDb *SchedulerDb, serials map[string]
 		jobIds := make([]string, 0)
 		for _, job := range jobs {
 			if _, ok := expected[job.JobSet]; ok {
-				assert.True(t, job.CancelRequested)
+				assert.True(t, job.CancelByJobsetRequested)
 				numChanged++
 				jobIds = append(jobIds, job.JobID)
 			}