Merge branch 'master' into f/chrisma/remove-config-always-schedule

internal/lookoutingesterv2/lookoutdb/insertion.go

@@ -52,7 +52,18 @@ func (l *LookoutDb) Store(ctx *armadacontext.Context, instructions *model.Instru
 
  start := time.Now()
  // Jobs need to be ingested first as other updates may reference these
- l.CreateJobs(ctx, instructions.JobsToCreate)
+ wgJobIngestion := sync.WaitGroup{}
+ wgJobIngestion.Add(2)
+ go func() {
+ defer wgJobIngestion.Done()
+ l.CreateJobs(ctx, instructions.JobsToCreate)
+ }()
+ go func() {
+ defer wgJobIngestion.Done()
+ l.CreateJobSpecs(ctx, instructions.JobsToCreate)
+ }()
+
+ wgJobIngestion.Wait()
 
  // Now we can job updates, annotations and new job runs
  wg := sync.WaitGroup{}
@@ -98,6 +109,22 @@ func (l *LookoutDb) CreateJobs(ctx *armadacontext.Context, instructions []*model
  log.Infof("Inserted %d jobs in %s", len(instructions), taken)
 }
 
+func (l *LookoutDb) CreateJobSpecs(ctx *armadacontext.Context, instructions []*model.CreateJobInstruction) {
+ if len(instructions) == 0 {
+ return
+ }
+ start := time.Now()
+ err := l.CreateJobSpecsBatch(ctx, instructions)
+ if err != nil {
+ log.WithError(err).Warn("Creating job specs via batch failed, will attempt to insert serially (this might be slow).")
+ l.CreateJobSpecsScalar(ctx, instructions)
+ }
+ taken := time.Since(start)
+ l.metrics.RecordAvRowChangeTimeByOperation("job_spec", commonmetrics.DBOperationInsert, len(instructions), taken)
+ l.metrics.RecordRowsChange("job_spec", commonmetrics.DBOperationInsert, len(instructions))
+ log.Infof("Inserted %d job specs in %s", len(instructions), taken)
+}
+
 func (l *LookoutDb) UpdateJobs(ctx *armadacontext.Context, instructions []*model.UpdateJobInstruction) {
  if len(instructions) == 0 {
  return
@@ -185,7 +212,6 @@ func (l *LookoutDb) CreateJobsBatch(ctx *armadacontext.Context, instructions []*
  state smallint,
  last_transition_time timestamp,
  last_transition_time_seconds bigint,
- job_spec bytea,
  priority_class varchar(63),
  annotations jsonb
  ) ON COMMIT DROP;`, tmpTable))
@@ -213,7 +239,6 @@ func (l *LookoutDb) CreateJobsBatch(ctx *armadacontext.Context, instructions []*
  "state",
  "last_transition_time",
  "last_transition_time_seconds",
- "job_spec",
  "priority_class",
  "annotations",
  },
@@ -233,7 +258,6 @@ func (l *LookoutDb) CreateJobsBatch(ctx *armadacontext.Context, instructions []*
  instructions[i].State,
  instructions[i].LastTransitionTime,
  instructions[i].LastTransitionTimeSeconds,
- instructions[i].JobProto,
  instructions[i].PriorityClass,
  instructions[i].Annotations,
  }, nil
@@ -261,7 +285,6 @@ func (l *LookoutDb) CreateJobsBatch(ctx *armadacontext.Context, instructions []*
  state,
  last_transition_time,
  last_transition_time_seconds,
- job_spec,
  priority_class,
  annotations
  ) SELECT * from %s
@@ -294,11 +317,10 @@ func (l *LookoutDb) CreateJobsScalar(ctx *armadacontext.Context, instructions []
  state,
  last_transition_time,
  last_transition_time_seconds,
- job_spec,
  priority_class,
  annotations
  )
- VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17)
+ VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16)
  ON CONFLICT DO NOTHING`
  for _, i := range instructions {
  err := l.withDatabaseRetryInsert(func() error {
@@ -317,7 +339,6 @@ func (l *LookoutDb) CreateJobsScalar(ctx *armadacontext.Context, instructions []
  i.State,
  i.LastTransitionTime,
  i.LastTransitionTimeSeconds,
- i.JobProto,
  i.PriorityClass,
  i.Annotations,
  )
@@ -446,6 +467,83 @@ func (l *LookoutDb) UpdateJobsScalar(ctx *armadacontext.Context, instructions []
  }
 }
 
+func (l *LookoutDb) CreateJobSpecsBatch(ctx *armadacontext.Context, instructions []*model.CreateJobInstruction) error {
+ return l.withDatabaseRetryInsert(func() error {
+ tmpTable := "job_spec_create_tmp"
+
+ createTmp := func(tx pgx.Tx) error {
+ _, err := tx.Exec(ctx, fmt.Sprintf(`
+ CREATE TEMPORARY TABLE %s (
+ job_id varchar(32),
+ job_spec bytea
+ ) ON COMMIT DROP;`, tmpTable))
+ if err != nil {
+ l.metrics.RecordDBError(commonmetrics.DBOperationCreateTempTable)
+ }
+ return err
+ }
+
+ insertTmp := func(tx pgx.Tx) error {
+ _, err := tx.CopyFrom(ctx,
+ pgx.Identifier{tmpTable},
+ []string{
+ "job_id",
+ "job_spec",
+ },
+ pgx.CopyFromSlice(len(instructions), func(i int) ([]interface{}, error) {
+ return []interface{}{
+ instructions[i].JobId,
+ instructions[i].JobProto,
+ }, nil
+ }),
+ )
+ return err
+ }
+
+ copyToDest := func(tx pgx.Tx) error {
+ _, err := tx.Exec(
+ ctx,
+ fmt.Sprintf(`
+ INSERT INTO job_spec (
+ job_id,
+ job_spec
+ ) SELECT * from %s
+ ON CONFLICT DO NOTHING`, tmpTable),
+ )
+ if err != nil {
+ l.metrics.RecordDBError(commonmetrics.DBOperationInsert)
+ }
+ return err
+ }
+
+ return batchInsert(ctx, l.db, createTmp, insertTmp, copyToDest)
+ })
+}
+
+func (l *LookoutDb) CreateJobSpecsScalar(ctx *armadacontext.Context, instructions []*model.CreateJobInstruction) {
+ sqlStatement := `INSERT INTO job_spec (
+ job_id,
+ job_spec
+ )
+ VALUES ($1, $2)
+ ON CONFLICT DO NOTHING`
+ for _, i := range instructions {
+ err := l.withDatabaseRetryInsert(func() error {
+ _, err := l.db.Exec(ctx, sqlStatement,
+ i.JobId,
+ i.JobProto,
+ )
+ if err != nil {
+ l.metrics.RecordDBError(commonmetrics.DBOperationInsert)
+ }
+ return err
+ })
+ if err != nil {
+ log.WithError(err).Warnf("Create job spec for job %s, jobset %s failed", i.JobId, i.JobSet)
+ }
+ }
+}
+
 func (l *LookoutDb) CreateJobRunsBatch(ctx *armadacontext.Context, instructions []*model.CreateJobRunInstruction) error {
  return l.withDatabaseRetryInsert(func() error {
  tmpTable := "job_run_create_tmp"

internal/lookoutingesterv2/lookoutdb/insertion_test.go

@@ -79,6 +79,11 @@ type JobRow struct {
  Annotations map[string]string
 }
 
+type JobSpecRow struct {
+ JobId string
+ JobProto []byte
+}
+
 type JobRunRow struct {
  RunId string
  JobId string
@@ -147,7 +152,7 @@ var expectedJobAfterSubmit = JobRow{
  State: lookout.JobQueuedOrdinal,
  LastTransitionTime: baseTime,
  LastTransitionTimeSeconds: baseTime.Unix(),
- JobProto: []byte(jobProto),
+ JobProto: []byte(nil),
  Duplicate: false,
  PriorityClass: priorityClass,
  Annotations: annotations,
@@ -167,7 +172,7 @@ var expectedJobAfterUpdate = JobRow{
  State: lookout.JobFailedOrdinal,
  LastTransitionTime: updateTime,
  LastTransitionTimeSeconds: updateTime.Unix(),
- JobProto: []byte(jobProto),
+ JobProto: []byte(nil),
  Duplicate: false,
  PriorityClass: priorityClass,
  Annotations: annotations,
@@ -841,10 +846,10 @@ func TestStoreNullValue(t *testing.T) {
  err := ldb.Store(armadacontext.Background(), instructions)
  assert.NoError(t, err)
 
- job := getJob(t, ldb.db, jobIdString)
+ jobSpec := getJobSpec(t, ldb.db, jobIdString)
  jobRun := getJobRun(t, ldb.db, runIdString)
 
- assert.Equal(t, jobProto, job.JobProto)
+ assert.Equal(t, jobProto, jobSpec.JobProto)
  assert.Equal(t, errorMsg, jobRun.Error)
  assert.Equal(t, debugMsg, jobRun.Debug)
  return nil
@@ -988,6 +993,23 @@ func getJob(t *testing.T, db *pgxpool.Pool, jobId string) JobRow {
  return job
 }
 
+func getJobSpec(t *testing.T, db *pgxpool.Pool, jobId string) JobSpecRow {
+ jobSpec := JobSpecRow{}
+ r := db.QueryRow(
+ armadacontext.Background(),
+ `SELECT
+ job_id,
+ job_spec
+ FROM job_spec WHERE job_id = $1`,
+ jobId)
+ err := r.Scan(
+ &jobSpec.JobId,
+ &jobSpec.JobProto,
+ )
+ assert.Nil(t, err)
+ return jobSpec
+}
+
 func getJobRun(t *testing.T, db *pgxpool.Pool, runId string) JobRunRow {
  run := JobRunRow{}
  r := db.QueryRow(

internal/lookoutv2/pruner/pruner.go

@@ -132,6 +132,7 @@ func deleteBatch(ctx *armadacontext.Context, tx pgx.Tx, batchLimit int) (int, er
  }
  _, err = tx.Exec(ctx, `
  DELETE FROM job WHERE job_id in (SELECT job_id from batch);
+ DELETE FROM job_spec WHERE job_id in (SELECT job_id from batch);
  DELETE FROM job_run WHERE job_id in (SELECT job_id from batch);
  DELETE FROM job_ids_to_delete WHERE job_id in (SELECT job_id from batch);
  TRUNCATE TABLE batch;`)

internal/lookoutv2/pruner/pruner_test.go

@@ -145,6 +145,7 @@ func TestPruneDb(t *testing.T) {
 
  queriedJobIdsPerTable := []map[string]bool{
  selectStringSet(t, db, "SELECT job_id FROM job"),
+ selectStringSet(t, db, "SELECT job_id FROM job_spec"),
  selectStringSet(t, db, "SELECT DISTINCT job_id FROM job_run"),
  }
  for _, queriedJobs := range queriedJobIdsPerTable {

internal/lookoutv2/repository/getjobspec.go

@@ -30,10 +30,17 @@ func NewSqlGetJobSpecRepository(db *pgxpool.Pool, decompressor compress.Decompre
 
 func (r *SqlGetJobSpecRepository) GetJobSpec(ctx *armadacontext.Context, jobId string) (*api.Job, error) {
  var rawBytes []byte
- err := r.db.QueryRow(ctx, "SELECT job_spec FROM job WHERE job_id = $1", jobId).Scan(&rawBytes)
+
+ err := r.db.QueryRow(
+ ctx, `
+ SELECT
+ COALESCE(job_spec.job_spec, job.job_spec)
+ FROM job LEFT JOIN job_spec
+ ON job.job_id = job_spec.job_id
+ WHERE job.job_id = $1`, jobId).Scan(&rawBytes)
  if err != nil {
  if err == pgx.ErrNoRows {
- return nil, errors.Errorf("job with id %s not found", jobId)
+ return nil, errors.Errorf("job_spec with job id %s not found", jobId)
  }
  return nil, err
  }

internal/lookoutv2/repository/getjobspec_test.go

@@ -3,6 +3,7 @@ package repository
 import (
  "testing"
 
+ "github.com/gogo/protobuf/proto"
  "github.com/jackc/pgx/v5/pgxpool"
  "github.com/stretchr/testify/assert"
 
@@ -50,6 +51,84 @@ func TestGetJobSpec(t *testing.T) {
  assert.NoError(t, err)
 }
 
+func TestMIGRATEDGetJobSpec(t *testing.T) {
+ var migratedResult *api.Job
+ err := lookout.WithLookoutDb(func(db *pgxpool.Pool) error {
+ converter := instructions.NewInstructionConverter(metrics.Get().Metrics, userAnnotationPrefix, &compress.NoOpCompressor{})
+ store := lookoutdb.NewLookoutDb(db, nil, metrics.Get(), 10)
+
+ _ = NewJobSimulator(converter, store).
+ Submit(queue, jobSet, owner, namespace, baseTime, &JobOptions{
+ JobId: jobId,
+ Priority: priority,
+ PriorityClass: "other-default",
+ Cpu: cpu,
+ Memory: memory,
+ EphemeralStorage: ephemeralStorage,
+ Gpu: gpu,
+ Annotations: map[string]string{
+ "step_path": "/1/2/3",
+ "hello": "world",
+ },
+ }).
+ Pending(runId, cluster, baseTime).
+ Running(runId, node, baseTime).
+ RunSucceeded(runId, baseTime).
+ Succeeded(baseTime).
+ Build().
+ ApiJob()
+
+ repo := NewSqlGetJobSpecRepository(db, &compress.NoOpDecompressor{})
+ var err error
+ migratedResult, err = repo.GetJobSpec(armadacontext.TODO(), jobId)
+ assert.NoError(t, err)
+ return nil
+ })
+ assert.NoError(t, err)
+
+ var result *api.Job
+ err = lookout.WithLookoutDb(func(db *pgxpool.Pool) error {
+ bytes, err := proto.Marshal(migratedResult)
+ assert.NoError(t, err)
+
+ _, err = db.Exec(armadacontext.Background(),
+ `INSERT INTO job (
+ job_id, queue, owner, namespace, jobset,
+ cpu,
+ memory,
+ ephemeral_storage,
+ gpu,
+ priority,
+ submitted,
+ state,
+ last_transition_time,
+ last_transition_time_seconds,
+ job_spec,
+ priority_class,
+ annotations
+ )
+ VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17)
+ ON CONFLICT DO NOTHING`,
+ jobId, queue, owner, namespace, jobSet,
+ int64(15), int64(48*1024*1024*1024), int64(100*1024*1024*1024), 8,
+ priority, baseTime, 1, baseTime, baseTime.Unix(), bytes, "other-default",
+ map[string]string{
+ "step_path": "/1/2/3",
+ "hello": "world",
+ })
+ assert.NoError(t, err)
+
+ repo := NewSqlGetJobSpecRepository(db, &compress.NoOpDecompressor{})
+ result, err = repo.GetJobSpec(armadacontext.TODO(), jobId)
+ assert.NoError(t, err)
+
+ return nil
+ })
+ assert.NoError(t, err)
+
+ assertApiJobsEquivalent(t, migratedResult, result)
+}
+
 func TestGetJobSpecError(t *testing.T) {
  err := lookout.WithLookoutDb(func(db *pgxpool.Pool) error {
  repo := NewSqlGetJobSpecRepository(db, &compress.NoOpDecompressor{})

internal/lookoutv2/schema/migrations/013_add_job_spec_table.sql

@@ -0,0 +1,6 @@
+CREATE TABLE IF NOT EXISTS job_spec (
+ job_id varchar(32) NOT NULL PRIMARY KEY,
+ job_spec bytea NOT NULL
+ );
+ALTER TABLE job_spec ALTER COLUMN job_spec SET STORAGE EXTERNAL;
+ALTER TABLE job ALTER COLUMN job_spec DROP NOT NULL;