Fix pool reconciliation (#3647)

The pool field is not correctly populated when the scheduler restarts, this PR just makes it so we properly populate the field Signed-off-by: JamesMurkin <jamesmurkin@hotmail.com>
armadaproject · Jun 4, 2024 · 05b211a · 05b211a
1 parent 5cc5061
commit 05b211a
Show file tree

Hide file tree

Showing 8 changed files with 14 additions and 4 deletions.
diff --git a/internal/scheduler/database/job_repository.go b/internal/scheduler/database/job_repository.go
@@ -171,6 +171,7 @@ func (r *PostgresJobRepository) FetchJobUpdates(ctx *armadacontext.Context, jobS
  SchedulingInfo: row.SchedulingInfo,
  SchedulingInfoVersion: row.SchedulingInfoVersion,
  Serial: row.Serial,
+ Pools: row.Pools,
  }
  }
 

diff --git a/internal/scheduler/jobdb/job_test.go b/internal/scheduler/jobdb/job_test.go
@@ -38,6 +38,7 @@ var baseJob, _ = jobDb.NewJob(
  false,
  3,
  false,
+ []string{},
 )
 
 var baseRun = &JobRun{
@@ -308,10 +309,10 @@ func TestJob_TestWithCreated(t *testing.T) {
 }
 
 func TestJob_DeepCopy(t *testing.T) {
- original, err := jobDb.NewJob("test-job", "test-jobSet", "test-queue", 2, jobSchedulingInfo, true, 0, false, false, false, 3, false)
+ original, err := jobDb.NewJob("test-job", "test-jobSet", "test-queue", 2, jobSchedulingInfo, true, 0, false, false, false, 3, false, []string{})
  assert.Nil(t, err)
  original = original.WithUpdatedRun(baseJobRun.DeepCopy())
- expected, err := jobDb.NewJob("test-job", "test-jobSet", "test-queue", 2, jobSchedulingInfo, true, 0, false, false, false, 3, false)
+ expected, err := jobDb.NewJob("test-job", "test-jobSet", "test-queue", 2, jobSchedulingInfo, true, 0, false, false, false, 3, false, []string{})
  assert.Nil(t, err)
  expected = expected.WithUpdatedRun(baseJobRun.DeepCopy())
 
@@ -364,7 +365,7 @@ func TestJobSchedulingInfoFieldsInitialised(t *testing.T) {
  assert.Nil(t, infoWithNilFields.GetPodRequirements().NodeSelector)
  assert.Nil(t, infoWithNilFields.GetPodRequirements().Annotations)
 
- job, err := jobDb.NewJob("test-job", "test-jobSet", "test-queue", 2, infoWithNilFieldsCopy, true, 0, false, false, false, 3, false)
+ job, err := jobDb.NewJob("test-job", "test-jobSet", "test-queue", 2, infoWithNilFieldsCopy, true, 0, false, false, false, 3, false, []string{})
  assert.Nil(t, err)
  assert.NotNil(t, job.NodeSelector())
  assert.NotNil(t, job.Annotations())

diff --git a/internal/scheduler/jobdb/jobdb.go b/internal/scheduler/jobdb/jobdb.go
@@ -138,6 +138,7 @@ func (jobDb *JobDb) NewJob(
  cancelled bool,
  created int64,
  validated bool,
+ pools []string,
 ) (*Job, error) {
  priorityClass, ok := jobDb.priorityClasses[schedulingInfo.PriorityClassName]
  if !ok {
@@ -164,6 +165,7 @@ func (jobDb *JobDb) NewJob(
  cancelled: cancelled,
  validated: validated,
  runsById: map[uuid.UUID]*JobRun{},
+ pools: pools,
  }
  job.ensureJobSchedulingInfoFieldsInitialised()
  job.schedulingKey = SchedulingKeyFromJob(jobDb.schedulingKeyGenerator, job)

diff --git a/internal/scheduler/jobdb/jobdb_test.go b/internal/scheduler/jobdb/jobdb_test.go
@@ -269,7 +269,7 @@ func TestJobDb_SchedulingKeyIsPopulated(t *testing.T) {
  },
  }
  jobDb := NewTestJobDb()
- job, err := jobDb.NewJob("jobId", "jobSet", "queue", 1, jobSchedulingInfo, false, 0, false, false, false, 2, false)
+ job, err := jobDb.NewJob("jobId", "jobSet", "queue", 1, jobSchedulingInfo, false, 0, false, false, false, 2, false, []string{})
  assert.Nil(t, err)
  assert.Equal(t, SchedulingKeyFromJob(jobDb.schedulingKeyGenerator, job), job.SchedulingKey())
 }

diff --git a/internal/scheduler/jobdb/reconciliation.go b/internal/scheduler/jobdb/reconciliation.go
@@ -275,6 +275,7 @@ func (jobDb *JobDb) schedulerJobFromDatabaseJob(dbJob *database.Job) (*Job, erro
  dbJob.Cancelled,
  dbJob.Submitted,
  dbJob.Validated,
+ dbJob.Pools,
  )
  if err != nil {
  return nil, err

diff --git a/internal/scheduler/scheduler_test.go b/internal/scheduler/scheduler_test.go
@@ -1751,6 +1751,7 @@ func jobDbJobFromDbJob(resourceListFactory *internaltypes.ResourceListFactory, j
  job.Cancelled,
  0,
  job.Validated,
+ job.Pools,
  )
  if err != nil {
  panic(err)

diff --git a/internal/scheduler/simulator/simulator.go b/internal/scheduler/simulator/simulator.go
@@ -675,6 +675,7 @@ func (s *Simulator) handleSubmitJob(txn *jobdb.Txn, e *armadaevents.SubmitJob, t
  false,
  s.logicalJobCreatedTimestamp.Add(1),
  false,
+ []string{},
  )
  if err != nil {
  return nil, false, err

diff --git a/internal/scheduler/testfixtures/testfixtures.go b/internal/scheduler/testfixtures/testfixtures.go
@@ -134,6 +134,7 @@ func NewJob(
  cancelled,
  created,
  validated,
+ []string{},
  )
  if err != nil {
  panic(err)
@@ -482,6 +483,7 @@ func TestJob(queue string, jobId ulid.ULID, priorityClassName string, req *sched
  false,
  created,
  false,
+ []string{},
  )
  return job
 }
@@ -838,6 +840,7 @@ func TestQueuedJobDbJob() *jobdb.Job {
  false,
  BaseTime.UnixNano(),
  false,
+ []string{},
  )
  return job
 }