Update IsManagedPod to check for job run id and job id

internal/executor/util/pod_util.go

@@ -61,8 +61,9 @@ func IsInTerminalState(pod *v1.Pod) bool {
 }
 
 func IsManagedPod(pod *v1.Pod) bool {
-	_, ok := pod.Labels[domain.JobId]
-	return ok
+	_, jobRunIdOk := pod.Labels[domain.JobRunId]
+	_, jobIdOk := pod.Labels[domain.JobId]
+	return jobRunIdOk && jobIdOk
 }
 
 func GetManagedPodSelector() labels.Selector {

internal/executor/util/pod_util_test.go

@@ -71,10 +71,10 @@ func TestHasIngress_WhenAnnotationNotSetToTrue(t *testing.T) {
 	assert.False(t, HasIngress(pod))
 }
 
-func TestIsManagedPod_ReturnsTrueIfJobIdLabelPresent(t *testing.T) {
+func TestIsManagedPod_ReturnsTrueIfJobRunIdAndJobIdLabelPresent(t *testing.T) {
 	pod := v1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
-			Labels: map[string]string{domain.JobId: "label"},
+			Labels: map[string]string{domain.JobRunId: "label", domain.JobId: "label"},
 		},
 	}
 
@@ -86,7 +86,19 @@ func TestIsManagedPod_ReturnsTrueIfJobIdLabelPresent(t *testing.T) {
 func TestIsManagedPod_ReturnsFalseIfJobIdLabelNotPresent(t *testing.T) {
 	pod := v1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
-			Labels: map[string]string{},
+			Labels: map[string]string{domain.JobRunId: "label"},
+		},
+	}
+
+	result := IsManagedPod(&pod)
+
+	assert.False(t, result)
+}
+
+func TestIsManagedPod_ReturnsFalseIfJobRunIdLabelNotPresent(t *testing.T) {
+	pod := v1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Labels: map[string]string{domain.JobId: "label"},
 		},
 	}
 

internal/executor/utilisation/cluster_utilisation_test.go

@@ -84,7 +84,7 @@ func TestGetAllPodsUsingResourceOnProcessingNodes_ShouldIncludeManagedPodsOnNode
 func TestGetAllPodsUsingResourceOnProcessingNodes_ShouldIncludeManagedPodNotAssignedToAnyNode(t *testing.T) {
 	podOnNode := v1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
-			Labels: map[string]string{domain.JobId: "label"},
+			Labels: map[string]string{domain.JobRunId: "label", domain.JobId: "label"},
 		},
 		Spec: v1.PodSpec{
 			NodeName: "",

internal/executor/utilisation/pod_utilisation_test.go

@@ -118,7 +118,8 @@ func createPod(phase v1.PodPhase, nodeName string, isManaged bool) *v1.Pod {
 	}
 	if isManaged {
 		pod.Labels = map[string]string{
-			domain.JobId: "jobid" + util2.NewULID(),
+			domain.JobRunId: "jobrunid" + util2.NewULID(),
+			domain.JobId:    "jobid" + util2.NewULID(),
 		}
 	}
 	return pod

internal/scheduler/api.go

@@ -87,6 +87,10 @@ func (srv *ExecutorApi) LeaseJobRuns(stream executorapi.ExecutorApi_LeaseJobRuns
 
 	ctx := armadacontext.WithLogField(armadacontext.FromGrpcCtx(stream.Context()), "executor", req.ExecutorId)
 
+	if err := validateLeaseRequest(ctx, req); err != nil {
+		return err
+	}
+
 	executor := srv.executorFromLeaseRequest(ctx, req)
 	if err := srv.executorRepository.StoreExecutor(ctx, executor); err != nil {
 		return err

internal/scheduler/api_test.go

@@ -5,6 +5,9 @@ import (
 	"testing"
 	"time"
 
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
 	"github.com/apache/pulsar-client-go/pulsar"
 	"github.com/gogo/protobuf/proto"
 	"github.com/golang/mock/gomock"
@@ -337,6 +340,70 @@ func TestExecutorApi_LeaseJobRuns(t *testing.T) {
 	}
 }
 
+func TestExecutorApi_LeaseJobRunsValidation(t *testing.T) {
+	const maxJobsPerCall = uint(100)
+	testClock := clock.NewFakeClock(time.Now())
+	runId2 := uuid.New()
+	runId3 := uuid.New()
+
+	tests := map[string]struct {
+		request     *executorapi.LeaseRequest
+		expectedErr error
+	}{
+		"invalid run id in run ids by state for node": {
+			request: &executorapi.LeaseRequest{
+				ExecutorId: "test-executor",
+				Pool:       "test-pool",
+				Nodes: []*executorapi.NodeInfo{
+					{
+						Name: "test-node",
+						RunIdsByState: map[string]api.JobState{
+							"":              api.JobState_RUNNING, // invalid job run id
+							runId2.String(): api.JobState_RUNNING,
+						},
+						NodeType: "node-type-1",
+					},
+				},
+				UnassignedJobRunIds: []armadaevents.Uuid{*armadaevents.ProtoUuidFromUuid(runId3)},
+				MaxJobsToLease:      uint32(maxJobsPerCall),
+			},
+			expectedErr: status.Error(codes.InvalidArgument, "runIdsByState for node test-node includes invalid run id ''"),
+		},
+	}
+	for name, tc := range tests {
+		t.Run(name, func(t *testing.T) {
+			ctx, cancel := armadacontext.WithTimeout(armadacontext.Background(), 5*time.Second)
+			ctrl := gomock.NewController(t)
+			mockPulsarProducer := mocks.NewMockProducer(ctrl)
+			mockJobRepository := schedulermocks.NewMockJobRepository(ctrl)
+			mockExecutorRepository := schedulermocks.NewMockExecutorRepository(ctrl)
+			mockStream := schedulermocks.NewMockExecutorApi_LeaseJobRunsServer(ctrl)
+
+			// set up mocks
+			mockStream.EXPECT().Context().Return(ctx).AnyTimes()
+			mockStream.EXPECT().Recv().Return(tc.request, nil).Times(1)
+
+			server, err := NewExecutorApi(
+				mockPulsarProducer,
+				mockJobRepository,
+				mockExecutorRepository,
+				[]int32{1000, 2000},
+				"kubernetes.io/hostname",
+				nil,
+				priorityClasses,
+				4*1024*1024,
+			)
+			require.NoError(t, err)
+			server.clock = testClock
+
+			err = server.LeaseJobRuns(mockStream)
+			require.Error(t, err)
+			require.Equal(t, tc.expectedErr, err)
+			cancel()
+		})
+	}
+}
+
 func TestAddNodeSelector(t *testing.T) {
 	withNodeSelector := &armadaevents.PodSpecWithAvoidList{
 		PodSpec: &v1.PodSpec{

internal/scheduler/api_validation.go

@@ -0,0 +1,25 @@
+package scheduler
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/google/uuid"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
+	"github.com/armadaproject/armada/pkg/executorapi"
+)
+
+func validateLeaseRequest(ctx context.Context, req *executorapi.LeaseRequest) error {
+	for _, node := range req.Nodes {
+		for runId := range node.RunIdsByState {
+			_, err := uuid.Parse(runId)
+			if err != nil {
+				return status.Error(codes.InvalidArgument, fmt.Sprintf("runIdsByState for node %s includes invalid run id '%s'", node.Name, runId))
+			}
+		}
+	}
+
+	return nil
+}