Plumb run index to pod name

jparraga-stackav · jparraga-stackav · commit dbd0cef9cdaf · 2025-05-28T00:27:33.000Z
Signed-off-by: Jason Parraga &lt;jparraga+gh@stackav.com&gt;
diff --git a/internal/executor/domain/pod_metadata.go b/internal/executor/domain/pod_metadata.go
@@ -3,6 +3,7 @@ package domain
 const (
 	JobId                    = "armada_job_id"
 	JobRunId                 = "armada_job_run_id"
+	JobRunIndex              = "armada_job_run_index"
 	PodNumber                = "armada_pod_number"
 	PodCount                 = "armada_pod_count"
 	JobSetId                 = "armada_jobset_id"
diff --git a/internal/executor/util/kubernetes_object.go b/internal/executor/util/kubernetes_object.go
@@ -116,11 +116,12 @@ func CreatePodFromExecutorApiJob(job *executorapi.JobRunLease, defaults *configu
 	}
 
 	labels := util.MergeMaps(job.Job.ObjectMeta.Labels, map[string]string{
-		domain.JobId:     jobId,
-		domain.JobRunId:  runId,
-		domain.Queue:     job.Queue,
-		domain.PodNumber: strconv.Itoa(0),
-		domain.PodCount:  strconv.Itoa(1),
+		domain.JobId:       jobId,
+		domain.JobRunId:    runId,
+		domain.JobRunIndex: strconv.Itoa(int(job.JobRunIndex)),
+		domain.Queue:       job.Queue,
+		domain.PodNumber:   strconv.Itoa(0),
+		domain.PodCount:    strconv.Itoa(1),
 	})
 	annotation := util.MergeMaps(job.Job.ObjectMeta.Annotations, map[string]string{
 		domain.JobSetId: job.Jobset,
@@ -132,7 +133,7 @@ func CreatePodFromExecutorApiJob(job *executorapi.JobRunLease, defaults *configu
 
 	pod := &v1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
-			Name:        common.PodNamePrefix + job.Job.JobId + "-" + strconv.Itoa(0),
+			Name:        common.PodNamePrefix + job.Job.JobId + "-" + strconv.Itoa(0) + "-" + strconv.FormatUint(uint64(job.JobRunIndex), 10),
 			Labels:      labels,
 			Annotations: annotation,
 			Namespace:   job.Job.ObjectMeta.Namespace,
diff --git a/internal/executor/util/kubernetes_objects_test.go b/internal/executor/util/kubernetes_objects_test.go
@@ -2,6 +2,7 @@ package util
 
 import (
 	"fmt"
+	"strconv"
 	"testing"
 
 	"github.com/armadaproject/armada/internal/common/util"
@@ -142,12 +143,14 @@ func makePodSpec() *v1.PodSpec {
 func TestCreatePodFromExecutorApiJob(t *testing.T) {
 	runId := uuid.NewString()
 	jobId := util.NewULID()
+	runIndex := 0
 
 	validJobLease := &executorapi.JobRunLease{
-		JobRunId: runId,
-		Queue:    "queue",
-		Jobset:   "job-set",
-		User:     "user",
+		JobRunId:    runId,
+		JobRunIndex: uint32(runIndex),
+		Queue:       "queue",
+		Jobset:      "job-set",
+		User:        "user",
 		Job: &armadaevents.SubmitJob{
 			ObjectMeta: &armadaevents.ObjectMeta{
 				Labels:      map[string]string{},
@@ -167,14 +170,15 @@ func TestCreatePodFromExecutorApiJob(t *testing.T) {
 
 	expectedPod := &v1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
-			Name:      fmt.Sprintf("armada-%s-0", jobId),
+			Name:      fmt.Sprintf("armada-%s-0-%d", jobId, runIndex),
 			Namespace: "test-namespace",
 			Labels: map[string]string{
-				domain.JobId:     jobId,
-				domain.JobRunId:  runId,
-				domain.Queue:     "queue",
-				domain.PodNumber: "0",
-				domain.PodCount:  "1",
+				domain.JobId:       jobId,
+				domain.JobRunId:    runId,
+				domain.JobRunIndex: strconv.Itoa(runIndex),
+				domain.Queue:       "queue",
+				domain.PodNumber:   "0",
+				domain.PodCount:    "1",
 			},
 			Annotations: map[string]string{
 				domain.JobSetId:            "job-set",
diff --git a/internal/scheduler/api.go b/internal/scheduler/api.go
@@ -173,12 +173,13 @@ func (srv *ExecutorApi) LeaseJobRuns(stream executorapi.ExecutorApi_LeaseJobRuns
 		err := stream.Send(&executorapi.LeaseStreamMessage{
 			Event: &executorapi.LeaseStreamMessage_Lease{
 				Lease: &executorapi.JobRunLease{
-					JobRunId: lease.RunID,
-					Queue:    lease.Queue,
-					Jobset:   lease.JobSet,
-					User:     lease.UserID,
-					Groups:   groups,
-					Job:      submitMsg,
+					JobRunId:    lease.RunID,
+					JobRunIndex: lease.RunIndex,
+					Queue:       lease.Queue,
+					Jobset:      lease.JobSet,
+					User:        lease.UserID,
+					Groups:      groups,
+					Job:         submitMsg,
 				},
 			},
 		})
diff --git a/internal/scheduler/database/job_repository.go b/internal/scheduler/database/job_repository.go
@@ -24,6 +24,7 @@ type hasSerial interface {
 
 type JobRunLease struct {
 	RunID                  string
+	RunIndex               uint32
 	Queue                  string
 	Pool                   string
 	JobSet                 string
@@ -352,7 +353,7 @@ func (r *PostgresJobRepository) FetchJobRunLeases(ctx *armadacontext.Context, ex
 		}
 
 		query := `
-				SELECT jr.run_id, jr.node, j.queue, j.job_set,  jr.pool, j.user_id, j.groups, j.submit_message, jr.pod_requirements_overlay
+				SELECT jr.run_id, jr.run_index, jr.node, j.queue, j.job_set,  jr.pool, j.user_id, j.groups, j.submit_message, jr.pod_requirements_overlay
 				FROM runs jr
 				LEFT JOIN %s as tmp ON (tmp.run_id = jr.run_id)
 			    JOIN jobs j
@@ -374,7 +375,7 @@ func (r *PostgresJobRepository) FetchJobRunLeases(ctx *armadacontext.Context, ex
 		defer rows.Close()
 		for rows.Next() {
 			run := JobRunLease{}
-			err = rows.Scan(&run.RunID, &run.Node, &run.Queue, &run.JobSet, &run.Pool, &run.UserID, &run.Groups, &run.SubmitMessage, &run.PodRequirementsOverlay)
+			err = rows.Scan(&run.RunID, &run.RunIndex, &run.Node, &run.Queue, &run.JobSet, &run.Pool, &run.UserID, &run.Groups, &run.SubmitMessage, &run.PodRequirementsOverlay)
 			if err != nil {
 				return errors.WithStack(err)
 			}
diff --git a/internal/scheduler/database/migrations/022_add_job_run_index.sql b/internal/scheduler/database/migrations/022_add_job_run_index.sql
@@ -0,0 +1 @@
+ALTER TABLE runs ADD COLUMN IF NOT EXISTS run_index bigint NOT NULL DEFAULT 0;
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/jobdb/job.go b/internal/scheduler/jobdb/job.go
@@ -670,8 +670,10 @@ func (job *Job) ValidateResourceRequests() error {
 // WithNewRun creates a copy of the job with a new run on the given executor.
 func (job *Job) WithNewRun(executor, nodeId, nodeName, pool string, scheduledAtPriority int32) *Job {
 	now := job.jobDb.clock.Now()
+	nextRunIndex := len(job.runsById)
 	return job.WithUpdatedRun(job.jobDb.CreateRun(
 		job.jobDb.uuidProvider.New(),
+		uint32(nextRunIndex),
 		job.Id(),
 		now.UnixNano(),
 		executor,
diff --git a/internal/scheduler/jobdb/job_run.go b/internal/scheduler/jobdb/job_run.go
@@ -15,6 +15,8 @@ import (
 type JobRun struct {
 	// Unique identifier for the run.
 	id string
+
+	index uint32
 	// Id of the job this run is associated with.
 	jobId string
 	// Time at which the run was created.
@@ -157,6 +159,9 @@ func (run *JobRun) Equal(other *JobRun) bool {
 	if run.id != other.id {
 		return false
 	}
+	if run.index != other.index {
+		return false
+	}
 	if run.jobId != other.jobId {
 		return false
 	}
@@ -213,6 +218,7 @@ func MinimalRun(id string, creationTime int64) *JobRun {
 // CreateRun creates a new scheduler job run from a database job run
 func (jobDb *JobDb) CreateRun(
 	id string,
+	index uint32,
 	jobId string,
 	creationTime int64,
 	executor string,
@@ -238,6 +244,7 @@ func (jobDb *JobDb) CreateRun(
 ) *JobRun {
 	return &JobRun{
 		id:                  id,
+		index:               index,
 		jobId:               jobId,
 		created:             creationTime,
 		executor:            jobDb.stringInterner.Intern(executor),
@@ -268,6 +275,11 @@ func (run *JobRun) Id() string {
 	return run.id
 }
 
+// Index returns the index of the JobRun with respect to the Job.
+func (run *JobRun) Index() uint32 {
+	return run.index
+}
+
 // JobId returns the id of the job this run is associated with.
 func (run *JobRun) JobId() string {
 	return run.jobId
diff --git a/internal/scheduler/jobdb/job_run_test.go b/internal/scheduler/jobdb/job_run_test.go
@@ -48,6 +48,7 @@ func init() {
 
 var baseJobRun = jobDb.CreateRun(
 	uuid.New().String(),
+	1,
 	uuid.NewString(),
 	5,
 	"test-executor",
@@ -125,6 +126,7 @@ func TestJobRun_TestRunAttempted(t *testing.T) {
 func TestDeepCopy(t *testing.T) {
 	run := jobDb.CreateRun(
 		uuid.NewString(),
+		1,
 		"job id",
 		1,
 		"executor",
@@ -150,6 +152,7 @@ func TestDeepCopy(t *testing.T) {
 	)
 	expected := jobDb.CreateRun(
 		run.id,
+		1,
 		"job id",
 		1,
 		"executor",
diff --git a/internal/scheduler/jobdb/reconciliation.go b/internal/scheduler/jobdb/reconciliation.go
@@ -308,6 +308,7 @@ func (jobDb *JobDb) schedulerRunFromDatabaseRun(dbRun *database.Run) *JobRun {
 	nodeId := api.NodeIdFromExecutorAndNodeName(dbRun.Executor, dbRun.Node)
 	return jobDb.CreateRun(
 		dbRun.RunID,
+		uint32(dbRun.RunIndex),
 		dbRun.JobID,
 		dbRun.Created,
 		dbRun.Executor,
diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go
@@ -605,6 +605,7 @@ func AppendEventSequencesFromScheduledJobs(eventSequences []*armadaevents.EventS
 					Event: &armadaevents.EventSequence_Event_JobRunLeased{
 						JobRunLeased: &armadaevents.JobRunLeased{
 							RunId:      run.Id(),
+							RunIndex:   run.Index(),
 							JobId:      job.Id(),
 							ExecutorId: run.Executor(),
 							// NodeId here refers to the unique identifier of the node in an executor cluster,
diff --git a/internal/scheduler/scheduler_test.go b/internal/scheduler/scheduler_test.go
@@ -229,6 +229,7 @@ var returnedOnceLeasedJob = testfixtures.NewJob(
 	true,
 ).WithUpdatedRun(testfixtures.JobDb.CreateRun(
 	uuid.NewString(),
+	0,
 	"01h3w2wtdchtc80hgyp782shrv",
 	0,
 	"testExecutor",
@@ -309,6 +310,7 @@ var (
 		true,
 	).WithUpdatedRun(testfixtures.JobDb.CreateRun(
 		uuid.NewString(),
+		1,
 		requeuedJobId,
 		time.Now().Unix(),
 		"testExecutor",
@@ -1157,6 +1159,7 @@ func TestScheduler_TestSyncInitialState(t *testing.T) {
 				queuedJob.WithUpdatedRun(
 					testfixtures.JobDb.CreateRun(
 						leasedJob.LatestRun().Id(),
+						leasedJob.LatestRun().Index(),
 						queuedJob.Id(),
 						123,
 						"test-executor",
diff --git a/internal/scheduleringester/instructions.go b/internal/scheduleringester/instructions.go
@@ -195,6 +195,7 @@ func (c *JobSetEventsInstructionConverter) handleJobRunLeased(jobRunLeased *arma
 			Queue: meta.queue,
 			DbRun: &schedulerdb.Run{
 				RunID:                  runId,
+				RunIndex:               jobRunLeased.RunIndex,
 				JobID:                  jobRunLeased.JobId,
 				Created:                eventTime.UnixNano(),
 				JobSet:                 meta.jobset,
diff --git a/pkg/armadaevents/events.pb.go b/pkg/armadaevents/events.pb.go
diff --git a/pkg/armadaevents/events.proto b/pkg/armadaevents/events.proto
diff --git a/pkg/executorapi/executorapi.pb.go b/pkg/executorapi/executorapi.pb.go
diff --git a/pkg/executorapi/executorapi.proto b/pkg/executorapi/executorapi.proto

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+ALTER TABLE runs ADD COLUMN IF NOT EXISTS run_index bigint NOT NULL DEFAULT 0;`