Add tolerations to away jobs (#3362)

internal/armada/repository/apimessages/conversions_test.go

@@ -11,6 +11,7 @@ import (
  v11 "k8s.io/api/networking/v1"
  "k8s.io/apimachinery/pkg/api/resource"
 
+ "github.com/armadaproject/armada/internal/scheduler/schedulerobjects"
  "github.com/armadaproject/armada/pkg/api"
  "github.com/armadaproject/armada/pkg/armadaevents"
 )
@@ -278,6 +279,16 @@ func TestConvertLeased(t *testing.T) {
  JobRunLeased: &armadaevents.JobRunLeased{
  JobId: jobIdProto,
  ExecutorId: executorId,
+ PodRequirementsOverlay: &schedulerobjects.PodRequirements{
+ Tolerations: []v1.Toleration{
+ {
+ Key: "whale",
+ Value: "true",
+ Effect: v1.TaintEffectNoSchedule,
+ },
+ },
+ Priority: 1000,
+ },
  },
  },
  }

internal/common/ingest/testfixtures/event.go

@@ -208,6 +208,16 @@ var Leased = &armadaevents.EventSequence_Event{
  HasScheduledAtPriority: true,
  ScheduledAtPriority: 15,
  UpdateSequenceNumber: 1,
+ PodRequirementsOverlay: &schedulerobjects.PodRequirements{
+ Tolerations: []v1.Toleration{
+ {
+ Key: "whale",
+ Value: "true",
+ Effect: v1.TaintEffectNoSchedule,
+ },
+ },
+ Priority: 15,
+ },
  },
  },
 }

internal/scheduler/api.go

@@ -9,6 +9,7 @@ import (
  "github.com/gogo/protobuf/types"
  "github.com/google/uuid"
  "github.com/pkg/errors"
+ v1 "k8s.io/api/core/v1"
  "k8s.io/apimachinery/pkg/util/clock"
 
  "github.com/armadaproject/armada/internal/common/armadacontext"
@@ -130,18 +131,29 @@ func (srv *ExecutorApi) LeaseJobRuns(stream executorapi.ExecutorApi_LeaseJobRuns
  if err := unmarshalFromCompressedBytes(lease.SubmitMessage, decompressor, submitMsg); err != nil {
  return err
  }
+
  if srv.priorityClassNameOverride != nil {
  srv.setPriorityClassName(submitMsg, *srv.priorityClassNameOverride)
  }
+
  srv.addNodeIdSelector(submitMsg, lease.Node)
 
+ if len(lease.PodRequirementsOverlay) > 0 {
+ PodRequirementsOverlay := schedulerobjects.PodRequirements{}
+ if err := proto.Unmarshal(lease.PodRequirementsOverlay, &PodRequirementsOverlay); err != nil {
+ return err
+ }
+ addTolerations(submitMsg, PodRequirementsOverlay.Tolerations)
+ }
+
  var groups []string
  if len(lease.Groups) > 0 {
  groups, err = compress.DecompressStringArray(lease.Groups, decompressor)
  if err != nil {
  return err
  }
  }
+
  err := stream.Send(&executorapi.LeaseStreamMessage{
  Event: &executorapi.LeaseStreamMessage_Lease{
  Lease: &executorapi.JobRunLease{
@@ -183,6 +195,13 @@ func (srv *ExecutorApi) setPriorityClassName(job *armadaevents.SubmitJob, priori
  }
 }
 
+func setPriorityClassName(podSpec *armadaevents.PodSpecWithAvoidList, priorityClassName string) {
+ if podSpec == nil || podSpec.PodSpec == nil {
+ return
+ }
+ podSpec.PodSpec.PriorityClassName = priorityClassName
+}
+
 func (srv *ExecutorApi) addNodeIdSelector(job *armadaevents.SubmitJob, nodeId string) {
  if job == nil || nodeId == "" {
  return
@@ -206,11 +225,18 @@ func addNodeSelector(podSpec *armadaevents.PodSpecWithAvoidList, key string, val
  }
 }
 
-func setPriorityClassName(podSpec *armadaevents.PodSpecWithAvoidList, priorityClassName string) {
- if podSpec == nil || podSpec.PodSpec == nil {
+func addTolerations(job *armadaevents.SubmitJob, tolerations []v1.Toleration) {
+ if job == nil || len(tolerations) == 0 {
  return
  }
- podSpec.PodSpec.PriorityClassName = priorityClassName
+ if job.MainObject != nil {
+ switch typed := job.MainObject.Object.(type) {
+ case *armadaevents.KubernetesMainObject_PodSpec:
+ if typed.PodSpec != nil && typed.PodSpec.PodSpec != nil {
+ typed.PodSpec.PodSpec.Tolerations = append(typed.PodSpec.PodSpec.Tolerations, tolerations...)
+ }
+ }
+ }
 }
 
 // ReportEvents publishes all eventSequences to Pulsar. The eventSequences are compacted for more efficient publishing.

internal/scheduler/api_test.go

@@ -17,6 +17,7 @@ import (
  "github.com/armadaproject/armada/internal/common/armadacontext"
  "github.com/armadaproject/armada/internal/common/compress"
  "github.com/armadaproject/armada/internal/common/mocks"
+ protoutil "github.com/armadaproject/armada/internal/common/proto"
  "github.com/armadaproject/armada/internal/common/pulsarutils"
  "github.com/armadaproject/armada/internal/scheduler/database"
  schedulermocks "github.com/armadaproject/armada/internal/scheduler/mocks"
@@ -80,7 +81,12 @@ func TestExecutorApi_LeaseJobRuns(t *testing.T) {
  UnassignedJobRuns: []string{runId3.String()},
  }
 
- submit, compressedSubmit := submitMsg(t, "node-id")
+ submit, compressedSubmit := submitMsg(
+ t,
+ &v1.PodSpec{
+ NodeSelector: map[string]string{nodeIdName: "node-id"},
+ },
+ )
  defaultLease := &database.JobRunLease{
  RunID: uuid.New(),
  Queue: "test-queue",
@@ -91,7 +97,7 @@ func TestExecutorApi_LeaseJobRuns(t *testing.T) {
  SubmitMessage: compressedSubmit,
  }
 
- submitWithoutNodeSelector, compressedSubmitNoNodeSelector := submitMsg(t, "")
+ submitWithoutNodeSelector, compressedSubmitNoNodeSelector := submitMsg(t, &v1.PodSpec{})
  leaseWithoutNode := &database.JobRunLease{
  RunID: uuid.New(),
  Queue: "test-queue",
@@ -101,6 +107,39 @@ func TestExecutorApi_LeaseJobRuns(t *testing.T) {
  SubmitMessage: compressedSubmitNoNodeSelector,
  }
 
+ tolerations := []v1.Toleration{
+ {
+ Key: "whale",
+ Value: "true",
+ Effect: v1.TaintEffectNoSchedule,
+ },
+ }
+ leaseWithTolerations := &database.JobRunLease{
+ RunID: uuid.New(),
+ Queue: "test-queue",
+ JobSet: "test-jobset",
+ UserID: "test-user",
+ Node: "node-id",
+ Groups: compressedGroups,
+ // We use the submit message without tolerations here, because the
+ // run-level tolerations are stored in PodRequirementsOverlay.
+ SubmitMessage: compressedSubmit,
+ PodRequirementsOverlay: protoutil.MustMarshall(
+ &schedulerobjects.PodRequirements{
+ Tolerations: tolerations,
+ Priority: 1000,
+ },
+ ),
+ }
+ submitWithTolerations, _ := submitMsg(
+ t,
+ &v1.PodSpec{
+ NodeSelector: map[string]string{nodeIdName: "node-id"},
+ Tolerations: tolerations,
+ },
+ )
+ submitWithTolerations.JobId = submit.JobId
+
  tests := map[string]struct {
  request *executorapi.LeaseRequest
  runsToCancel []uuid.UUID
@@ -154,6 +193,26 @@ func TestExecutorApi_LeaseJobRuns(t *testing.T) {
  },
  },
  },
+ "run with PodRequirementsOverlay": {
+ request: defaultRequest,
+ leases: []*database.JobRunLease{leaseWithTolerations},
+ expectedExecutor: defaultExpectedExecutor,
+ expectedMsgs: []*executorapi.LeaseStreamMessage{
+ {
+ Event: &executorapi.LeaseStreamMessage_Lease{Lease: &executorapi.JobRunLease{
+ JobRunId: armadaevents.ProtoUuidFromUuid(leaseWithTolerations.RunID),
+ Queue: leaseWithTolerations.Queue,
+ Jobset: leaseWithTolerations.JobSet,
+ User: leaseWithTolerations.UserID,
+ Groups: groups,
+ Job: submitWithTolerations,
+ }},
+ },
+ {
+ Event: &executorapi.LeaseStreamMessage_End{End: &executorapi.EndMarker{}},
+ },
+ },
+ },
  "do nothing": {
  request: defaultRequest,
  expectedExecutor: defaultExpectedExecutor,
@@ -347,12 +406,7 @@ func TestExecutorApi_Publish(t *testing.T) {
  }
 }
 
-func submitMsg(t *testing.T, nodeName string) (*armadaevents.SubmitJob, []byte) {
- podSpec := &v1.PodSpec{}
- if nodeName != "" {
- podSpec.NodeSelector = map[string]string{}
- podSpec.NodeSelector[nodeIdName] = nodeName
- }
+func submitMsg(t *testing.T, podSpec *v1.PodSpec) (*armadaevents.SubmitJob, []byte) {
  submitMsg := &armadaevents.SubmitJob{
  JobId: armadaevents.ProtoUuidFromUuid(uuid.New()),
  MainObject: &armadaevents.KubernetesMainObject{

internal/scheduler/database/job_repository.go

@@ -23,13 +23,14 @@ type hasSerial interface {
 }
 
 type JobRunLease struct {
- RunID uuid.UUID
- Queue string
- JobSet string
- UserID string
- Node string
- Groups []byte
- SubmitMessage []byte
+ RunID uuid.UUID
+ Queue string
+ JobSet string
+ UserID string
+ Node string
+ Groups []byte
+ SubmitMessage []byte
+ PodRequirementsOverlay []byte
 }
 
 // JobRepository is an interface to be implemented by structs which provide job and run information
@@ -246,7 +247,7 @@ func (r *PostgresJobRepository) FetchJobRunLeases(ctx *armadacontext.Context, ex
  }
 
  query := `
- SELECT jr.run_id, jr.node, j.queue, j.job_set, j.user_id, j.groups, j.submit_message
+ SELECT jr.run_id, jr.node, j.queue, j.job_set, 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
@@ -267,7 +268,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.UserID, &run.Groups, &run.SubmitMessage)
+ err = rows.Scan(&run.RunID, &run.Node, &run.Queue, &run.JobSet, &run.UserID, &run.Groups, &run.SubmitMessage, &run.PodRequirementsOverlay)
  if err != nil {
  return errors.WithStack(err)
  }

internal/scheduler/database/job_repository_test.go

@@ -11,12 +11,14 @@ import (
  "github.com/stretchr/testify/assert"
  "github.com/stretchr/testify/require"
  "golang.org/x/exp/slices"
+ v1 "k8s.io/api/core/v1"
 
  "github.com/armadaproject/armada/internal/common/armadacontext"
  "github.com/armadaproject/armada/internal/common/compress"
  "github.com/armadaproject/armada/internal/common/database"
  protoutil "github.com/armadaproject/armada/internal/common/proto"
  "github.com/armadaproject/armada/internal/common/util"
+ "github.com/armadaproject/armada/internal/scheduler/schedulerobjects"
  "github.com/armadaproject/armada/pkg/armadaevents"
 )
 
@@ -402,6 +404,24 @@ func TestFetchJobRunLeases(t *testing.T) {
  JobSet: "test-jobset",
  Executor: executorName,
  },
+ {
+ RunID: uuid.New(),
+ JobID: dbJobs[2].JobID,
+ JobSet: "test-jobset",
+ Executor: executorName,
+ PodRequirementsOverlay: protoutil.MustMarshall(
+ &schedulerobjects.PodRequirements{
+ Tolerations: []v1.Toleration{
+ {
+ Key: "whale",
+ Value: "true",
+ Effect: v1.TaintEffectNoSchedule,
+ },
+ },
+ Priority: 1000,
+ },
+ ),
+ },
  {
  RunID: uuid.New(),
  JobID: dbJobs[0].JobID,
@@ -424,15 +444,16 @@ func TestFetchJobRunLeases(t *testing.T) {
  Succeeded: true, // should be ignored as terminal
  },
  }
- expectedLeases := make([]*JobRunLease, 3)
+ expectedLeases := make([]*JobRunLease, 4)
  for i := range expectedLeases {
  expectedLeases[i] = &JobRunLease{
- RunID: dbRuns[i].RunID,
- Queue: dbJobs[i].Queue,
- JobSet: dbJobs[i].JobSet,
- UserID: dbJobs[i].UserID,
- Groups: dbJobs[i].Groups,
- SubmitMessage: dbJobs[i].SubmitMessage,
+ RunID: dbRuns[i].RunID,
+ Queue: dbJobs[i].Queue,
+ JobSet: dbJobs[i].JobSet,
+ UserID: dbJobs[i].UserID,
+ Groups: dbJobs[i].Groups,
+ SubmitMessage: dbJobs[i].SubmitMessage,
+ PodRequirementsOverlay: dbRuns[i].PodRequirementsOverlay,
  }
  }
  tests := map[string]struct {
@@ -465,12 +486,12 @@ func TestFetchJobRunLeases(t *testing.T) {
  excludedRuns: []uuid.UUID{dbRuns[1].RunID},
  maxRowsToFetch: 100,
  executor: executorName,
- expectedLeases: []*JobRunLease{expectedLeases[0], expectedLeases[2]},
+ expectedLeases: []*JobRunLease{expectedLeases[0], expectedLeases[2], expectedLeases[3]},
  },
  "exclude everything": {
  dbJobs: dbJobs,
  dbRuns: dbRuns,
- excludedRuns: []uuid.UUID{dbRuns[0].RunID, dbRuns[1].RunID, dbRuns[2].RunID},
+ excludedRuns: []uuid.UUID{dbRuns[0].RunID, dbRuns[1].RunID, dbRuns[2].RunID, dbRuns[3].RunID},
  maxRowsToFetch: 100,
  executor: executorName,
  expectedLeases: nil,

internal/scheduler/database/migrations/006_add_run_pod_requirements.up.sql

@@ -0,0 +1,5 @@
+-- The pod_requirements_overlay column holds a serialized PodRequirements
+-- object that can be used to modify the pod requirements of a job; for
+-- example, it is used to add additional tolerations to runs that are scheduled
+-- as away jobs.
+ALTER TABLE runs ADD COLUMN pod_requirements_overlay bytea;