Repository: armadaproject/armada
Pull request #2179 (2023-02-24 13:29:36)

Set node selector when sending job to executor
The scheduler assigns pods to nodes

We now remember this value and set a node selector on jobs we send to the executor - to enforce the schedulers decision


Patch:
diff --git a/config/scheduler/config.yaml b/config/scheduler/config.yaml
index a5ba3a18..ed3ccbfc 100644
--- a/config/scheduler/config.yaml
+++ b/config/scheduler/config.yaml
@@ -37,6 +37,7 @@ grpc:
     permitWithoutStream: true
 scheduling:
   executorTimeout: 10m
+  nodeIdLabel: kubernetes.io/hostname
   preemption:
     enabled: true
     priorityClasses:
diff --git a/internal/scheduler/api.go b/internal/scheduler/api.go
index 84e26b81..20e50567 100644
--- a/internal/scheduler/api.go
+++ b/internal/scheduler/api.go
@@ -33,6 +33,7 @@ type ExecutorApi struct {
 	allowedPriorities        []int32 // allowed priority classes
 	maxJobsPerCall           uint    // maximum number of jobs that will be leased in a single call
 	maxPulsarMessageSize     uint    // maximum sizer of pulsar messages produced
+	nodeIdLabel              string
 	clock                    clock.Clock
 }
 
@@ -42,6 +43,7 @@ func NewExecutorApi(producer pulsar.Producer,
 	legacyExecutorRepository database.ExecutorRepository,
 	allowedPriorities []int32,
 	maxJobsPerCall uint,
+	nodeIdLabel string,
 ) (*ExecutorApi, error) {
 	if len(allowedPriorities) == 0 {
 		return nil, errors.New("allowedPriorities cannot be empty")
@@ -58,6 +60,7 @@ func NewExecutorApi(producer pulsar.Producer,
 		allowedPriorities:        allowedPriorities,
 		maxJobsPerCall:           maxJobsPerCall,
 		maxPulsarMessageSize:     1024 * 1024 * 2,
+		nodeIdLabel:              nodeIdLabel,
 		clock:                    clock.RealClock{},
 	}, nil
 }
@@ -131,6 +134,7 @@ func (srv *ExecutorApi) LeaseJobRuns(stream executorapi.ExecutorApi_LeaseJobRuns
 		if err != nil {
 			return err
 		}
+		srv.addNodeSelector(submitMsg, lease.Node)
 
 		var groups []string
 		if len(lease.Groups) > 0 {
@@ -168,6 +172,31 @@ func (srv *ExecutorApi) LeaseJobRuns(stream executorapi.ExecutorApi_LeaseJobRuns
 	return nil
 }
 
+func (srv *ExecutorApi) addNodeSelector(job *armadaevents.SubmitJob, nodeId string) {
+	if job == nil || nodeId == "" {
+		return
+	}
+
+	if job.MainObject != nil {
+		switch typed := job.MainObject.Object.(type) {
+		case *armadaevents.KubernetesMainObject_PodSpec:
+			addNodeSelector(typed.PodSpec, srv.nodeIdLabel, nodeId)
+		}
+	}
+}
+
+func addNodeSelector(podSpec *armadaevents.PodSpecWithAvoidList, key string, value string) {
+	if podSpec == nil || podSpec.PodSpec == nil || key == "" || value == "" {
+		return
+	}
+
+	if podSpec.PodSpec.NodeSelector == nil {
+		podSpec.PodSpec.NodeSelector = make(map[string]string, 1)
+	}
+
+	podSpec.PodSpec.NodeSelector[key] = value
+}
+
 // ReportEvents publishes all events to pulsar. The events are compacted for more efficient publishing
 func (srv *ExecutorApi) ReportEvents(ctx context.Context, list *executorapi.EventList) (*types.Empty, error) {
 	err := pulsarutils.CompactAndPublishSequences(ctx, list.Events, srv.producer, srv.maxPulsarMessageSize, schedulers.Pulsar)
diff --git a/internal/scheduler/publisher.go b/internal/scheduler/publisher.go
index 6d03a242..0ae05953 100644
--- a/internal/scheduler/publisher.go
+++ b/internal/scheduler/publisher.go
@@ -18,8 +18,8 @@ import (
 	"github.com/armadaproject/armada/pkg/armadaevents"
 )
 
-// This is half the default pulsar BatchingMaxSize
 const (
+	// This is half the default pulsar BatchingMaxSize
 	defaultMaxMessageBatchSize = 64 * 1024
 	explicitPartitionKey       = "armada_pulsar_partition"
 )
diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go
index 742fd5c4..c0f39aae 100644
--- a/internal/scheduler/scheduler.go
+++ b/internal/scheduler/scheduler.go
@@ -336,6 +336,7 @@ func (s *Scheduler) generateLeaseMessages(scheduledJobs []*jobdb.Job) ([]*armada
 							RunId:      armadaevents.ProtoUuidFromUuid(job.LatestRun().Id()),
 							JobId:      jobId,
 							ExecutorId: job.LatestRun().Executor(),
+							NodeId:     job.LatestRun().Node(),
 						},
 					},
 				},
diff --git a/internal/scheduler/schedulerapp.go b/internal/scheduler/schedulerapp.go
index 4a791be8..c98e27f0 100644
--- a/internal/scheduler/schedulerapp.go
+++ b/internal/scheduler/schedulerapp.go
@@ -117,7 +117,8 @@ func Run(config Configuration) error {
 		executorRepository,
 		legacyExecutorRepository,
 		allowedPcs,
-		config.Scheduling.MaximumJobsToSchedule)
+		config.Scheduling.MaximumJobsToSchedule,
+		config.Scheduling.Preemption.NodeIdLabel)
 	if err != nil {
 		return errors.WithMessage(err, "error creating executorApi")
 	}
diff --git a/internal/scheduleringester/instructions.go b/internal/scheduleringester/instructions.go
index 68eadacb..98b8fe8d 100644
--- a/internal/scheduleringester/instructions.go
+++ b/internal/scheduleringester/instructions.go
@@ -172,6 +172,7 @@ func (c *InstructionConverter) handleJobRunLeased(jobRunLeased *armadaevents.Job
 		JobID:    jobId,
 		JobSet:   meta.jobset,
 		Executor: jobRunLeased.GetExecutorId(),
+		Node:     jobRunLeased.GetNodeId(),
 	}}}, nil
 }
 


Test patch:
diff --git a/internal/common/ingest/testfixtures/event.go b/internal/common/ingest/testfixtures/event.go
index c660dd4f..aa314bb3 100644
--- a/internal/common/ingest/testfixtures/event.go
+++ b/internal/common/ingest/testfixtures/event.go
@@ -136,6 +136,7 @@ var Leased = &armadaevents.EventSequence_Event{
 			RunId:      RunIdProto,
 			JobId:      JobIdProto,
 			ExecutorId: ExecutorId,
+			NodeId:     NodeName,
 		},
 	},
 }
diff --git a/internal/scheduler/api_test.go b/internal/scheduler/api_test.go
index 1c8df295..33363707 100644
--- a/internal/scheduler/api_test.go
+++ b/internal/scheduler/api_test.go
@@ -11,6 +11,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/net/context"
+	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/util/clock"
 
 	"github.com/armadaproject/armada/internal/common/compress"
@@ -24,13 +25,14 @@ import (
 	"github.com/armadaproject/armada/pkg/executorapi"
 )
 
+const nodeIdName = "kubernetes.io/hostname"
+
 func TestExecutorApi_LeaseJobRuns(t *testing.T) {
 	const maxJobsPerCall = uint(100)
 	testClock := clock.NewFakeClock(time.Now())
 	runId1 := uuid.New()
 	runId2 := uuid.New()
 	runId3 := uuid.New()
-	submit, compressedSubmit := submitMsg(t)
 	groups, compressedGroups := groups(t)
 	defaultRequest := &executorapi.LeaseRequest{
 		ExecutorId: "test-executor",
@@ -72,15 +74,27 @@ func TestExecutorApi_LeaseJobRuns(t *testing.T) {
 		UnassignedJobRuns: []string{runId3.String()},
 	}
 
+	submit, compressedSubmit := submitMsg(t, "node-id")
 	defaultLease := &database.JobRunLease{
 		RunID:         uuid.New(),
 		Queue:         "test-queue",
 		JobSet:        "test-jobset",
 		UserID:        "test-user",
+		Node:          "node-id",
 		Groups:        compressedGroups,
 		SubmitMessage: compressedSubmit,
 	}
 
+	submitWithoutNodeSelector, compressedSubmitNoNodeSelector := submitMsg(t, "")
+	leaseWithoutNode := &database.JobRunLease{
+		RunID:         uuid.New(),
+		Queue:         "test-queue",
+		JobSet:        "test-jobset",
+		UserID:        "test-user",
+		Groups:        compressedGroups,
+		SubmitMessage: compressedSubmitNoNodeSelector,
+	}
+
 	tests := map[string]struct {
 		request          *executorapi.LeaseRequest
 		runsToCancel     []uuid.UUID
@@ -114,6 +128,26 @@ func TestExecutorApi_LeaseJobRuns(t *testing.T) {
 				},
 			},
 		},
+		"no node selector when missing in lease": {
+			request:          defaultRequest,
+			leases:           []*database.JobRunLease{leaseWithoutNode},
+			expectedExecutor: defaultExpectedExecutor,
+			expectedMsgs: []*executorapi.LeaseStreamMessage{
+				{
+					Event: &executorapi.LeaseStreamMessage_Lease{Lease: &executorapi.JobRunLease{
+						JobRunId: armadaevents.ProtoUuidFromUuid(leaseWithoutNode.RunID),
+						Queue:    leaseWithoutNode.Queue,
+						Jobset:   leaseWithoutNode.JobSet,
+						User:     leaseWithoutNode.UserID,
+						Groups:   groups,
+						Job:      submitWithoutNodeSelector,
+					}},
+				},
+				{
+					Event: &executorapi.LeaseStreamMessage_End{End: &executorapi.EndMarker{}},
+				},
+			},
+		},
 		"do nothing": {
 			request:          defaultRequest,
 			expectedExecutor: defaultExpectedExecutor,
@@ -165,6 +199,7 @@ func TestExecutorApi_LeaseJobRuns(t *testing.T) {
 				mockLegacyExecutorRepository,
 				[]int32{1000, 2000},
 				maxJobsPerCall,
+				"kubernetes.io/hostname",
 			)
 			require.NoError(t, err)
 			server.clock = testClock
@@ -177,6 +212,59 @@ func TestExecutorApi_LeaseJobRuns(t *testing.T) {
 	}
 }
 
+func TestAddNodeSelector(t *testing.T) {
+	withNodeSelector := &armadaevents.PodSpecWithAvoidList{
+		PodSpec: &v1.PodSpec{
+			NodeSelector: map[string]string{nodeIdName: "node-1"},
+		},
+	}
+
+	tests := map[string]struct {
+		input    *armadaevents.PodSpecWithAvoidList
+		expected *armadaevents.PodSpecWithAvoidList
+		key      string
+		value    string
+	}{
+		"Sets node selector": {
+			input:    &armadaevents.PodSpecWithAvoidList{PodSpec: &v1.PodSpec{}},
+			expected: withNodeSelector,
+			key:      nodeIdName,
+			value:    "node-1",
+		},
+		"input is nil": {
+			input:    nil,
+			expected: nil,
+			key:      nodeIdName,
+			value:    "node-1",
+		},
+		"podspec is nil": {
+			input:    &armadaevents.PodSpecWithAvoidList{},
+			expected: &armadaevents.PodSpecWithAvoidList{},
+			key:      nodeIdName,
+			value:    "node-1",
+		},
+		"key is not set": {
+			input:    &armadaevents.PodSpecWithAvoidList{PodSpec: &v1.PodSpec{}},
+			expected: &armadaevents.PodSpecWithAvoidList{PodSpec: &v1.PodSpec{}},
+			key:      "",
+			value:    "node-1",
+		},
+		"value is not set": {
+			input:    &armadaevents.PodSpecWithAvoidList{PodSpec: &v1.PodSpec{}},
+			expected: &armadaevents.PodSpecWithAvoidList{PodSpec: &v1.PodSpec{}},
+			key:      nodeIdName,
+			value:    "",
+		},
+	}
+
+	for name, tc := range tests {
+		t.Run(name, func(t *testing.T) {
+			addNodeSelector(tc.input, tc.key, tc.value)
+			assert.Equal(t, tc.expected, tc.input)
+		})
+	}
+}
+
 func TestExecutorApi_Publish(t *testing.T) {
 	tests := map[string]struct {
 		sequences []*armadaevents.EventSequence
@@ -237,6 +325,7 @@ func TestExecutorApi_Publish(t *testing.T) {
 				mockLegacyExecutorRepository,
 				[]int32{1000, 2000},
 				100,
+				"kubernetes.io/hostname",
 			)
 
 			require.NoError(t, err)
@@ -250,9 +339,21 @@ func TestExecutorApi_Publish(t *testing.T) {
 	}
 }
 
-func submitMsg(t *testing.T) (*armadaevents.SubmitJob, []byte) {
+func submitMsg(t *testing.T, nodeName string) (*armadaevents.SubmitJob, []byte) {
+	podSpec := &v1.PodSpec{}
+	if nodeName != "" {
+		podSpec.NodeSelector = map[string]string{}
+		podSpec.NodeSelector[nodeIdName] = nodeName
+	}
 	submitMsg := &armadaevents.SubmitJob{
 		JobId: armadaevents.ProtoUuidFromUuid(uuid.New()),
+		MainObject: &armadaevents.KubernetesMainObject{
+			Object: &armadaevents.KubernetesMainObject_PodSpec{
+				PodSpec: &armadaevents.PodSpecWithAvoidList{
+					PodSpec: podSpec,
+				},
+			},
+		},
 	}
 	bytes, err := proto.Marshal(submitMsg)
 	require.NoError(t, err)
diff --git a/internal/scheduleringester/instructions_test.go b/internal/scheduleringester/instructions_test.go
index 6e01b72d..dc4b9234 100644
--- a/internal/scheduleringester/instructions_test.go
+++ b/internal/scheduleringester/instructions_test.go
@@ -76,6 +76,7 @@ func TestConvertSequence(t *testing.T) {
 				JobID:    f.JobIdString,
 				JobSet:   f.JobSetName,
 				Executor: f.ExecutorId,
+				Node:     f.NodeName,
 			}}},
 		},
 		"job run running": {