Inherit execution cluster label from source execution (#5431)

Signed-off-by: Vinayak Agarwal <[email protected]>
flyteorg · Jun 20, 2024 · 16e7780 · 16e7780
1 parent 297cdf8
commit 16e7780
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/flyteadmin/pkg/manager/impl/execution_manager.go b/flyteadmin/pkg/manager/impl/execution_manager.go
@@ -305,6 +305,12 @@ func (m *ExecutionManager) getInheritedExecMetadata(ctx context.Context, request
 	} else {
 		requestSpec.Metadata.Nesting = 1
 	}
+
+	// If the source execution has a cluster label, inherit it.
+	if sourceExecution.Spec.ExecutionClusterLabel != nil {
+		logger.Infof(ctx, "Inherited execution label from source execution [%+v]", sourceExecution.Spec.ExecutionClusterLabel.Value)
+		requestSpec.ExecutionClusterLabel = sourceExecution.Spec.ExecutionClusterLabel
+	}
 	return parentNodeExecutionID, sourceExecutionID, nil
 }
 

diff --git a/flyteadmin/pkg/manager/impl/execution_manager_test.go b/flyteadmin/pkg/manager/impl/execution_manager_test.go
@@ -428,6 +428,7 @@ func TestCreateExecutionFromWorkflowNode(t *testing.T) {
 	)
 
 	getExecutionCalled := false
+	var clusterLabel = &admin.ExecutionClusterLabel{Value: executionClusterLabel}
 	repository.ExecutionRepo().(*repositoryMocks.MockExecutionRepo).SetGetCallback(
 		func(ctx context.Context, input interfaces.Identifier) (models.Execution, error) {
 			assert.EqualValues(t, input.Project, parentNodeExecutionID.ExecutionId.Project)
@@ -437,6 +438,7 @@ func TestCreateExecutionFromWorkflowNode(t *testing.T) {
 				Metadata: &admin.ExecutionMetadata{
 					Nesting: 1,
 				},
+				ExecutionClusterLabel: clusterLabel,
 			}
 			specBytes, _ := proto.Marshal(spec)
 			getExecutionCalled = true
@@ -462,6 +464,7 @@ func TestCreateExecutionFromWorkflowNode(t *testing.T) {
 			assert.EqualValues(t, input.SourceExecutionID, 2)
 			assert.Equal(t, 2, int(spec.Metadata.Nesting))
 			assert.Equal(t, principal, spec.Metadata.Principal)
+			assert.Equal(t, executionClusterLabel, spec.ExecutionClusterLabel.Value)
 			assert.Equal(t, principal, input.User)
 			return nil
 		},