From 16e77801f7bd32377e9f55cf8f2f6f006ddc3acf Mon Sep 17 00:00:00 2001 From: Vinayak Agarwal Date: Thu, 20 Jun 2024 03:27:43 -0700 Subject: [PATCH] Inherit execution cluster label from source execution (#5431) Signed-off-by: Vinayak Agarwal --- flyteadmin/pkg/manager/impl/execution_manager.go | 6 ++++++ flyteadmin/pkg/manager/impl/execution_manager_test.go | 3 +++ 2 files changed, 9 insertions(+) diff --git a/flyteadmin/pkg/manager/impl/execution_manager.go b/flyteadmin/pkg/manager/impl/execution_manager.go index 977e1f8179..da7489258d 100644 --- 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 index 52ff607725..58e50c0444 100644 --- 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 },