Add default cluster pool assignments to config (#600)

flyteadmin/pkg/manager/impl/execution_manager.go

@@ -444,8 +444,11 @@ func (m *ExecutionManager) getClusterAssignment(ctx context.Context, request *ad
 	if resource != nil && resource.Attributes.GetClusterAssignment() != nil {
 		return resource.Attributes.GetClusterAssignment(), nil
 	}
-	// Defaults to empty assignment with no selectors
-	return &admin.ClusterAssignment{}, nil
+	clusterPoolAssignment := m.config.ClusterPoolAssignmentConfiguration().GetClusterPoolAssignments()[request.GetDomain()]
+
+	return &admin.ClusterAssignment{
+		ClusterPoolName: clusterPoolAssignment.Pool,
+	}, nil
 }
 
 func (m *ExecutionManager) launchSingleTaskExecution(

flyteadmin/pkg/manager/impl/execution_manager_test.go

@@ -5440,6 +5440,30 @@ func TestGetClusterAssignment(t *testing.T) {
 		assert.NoError(t, err)
 		assert.True(t, proto.Equal(ca, &reqClusterAssignment))
 	})
+	t.Run("value from config", func(t *testing.T) {
+		customCP := "my_cp"
+		clusterPoolAsstProvider := &runtimeIFaceMocks.ClusterPoolAssignmentConfiguration{}
+		clusterPoolAsstProvider.OnGetClusterPoolAssignments().Return(runtimeInterfaces.ClusterPoolAssignments{
+			workflowIdentifier.GetDomain(): runtimeInterfaces.ClusterPoolAssignment{
+				Pool: customCP,
+			},
+		})
+		mockConfig := getMockExecutionsConfigProvider()
+		mockConfig.(*runtimeMocks.MockConfigurationProvider).AddClusterPoolAssignmentConfiguration(clusterPoolAsstProvider)
+
+		executionManager := ExecutionManager{
+			resourceManager: &managerMocks.MockResourceManager{},
+			config:          mockConfig,
+		}
+
+		ca, err := executionManager.getClusterAssignment(context.TODO(), &admin.ExecutionCreateRequest{
+			Project: workflowIdentifier.Project,
+			Domain:  workflowIdentifier.Domain,
+			Spec:    &admin.ExecutionSpec{},
+		})
+		assert.NoError(t, err)
+		assert.Equal(t, customCP, ca.GetClusterPoolName())
+	})
 }
 
 func TestResolvePermissions(t *testing.T) {

flyteadmin/pkg/runtime/cluster_pool_assignment_provider.go

@@ -0,0 +1,24 @@
+package runtime
+
+import (
+	"github.com/flyteorg/flyteadmin/pkg/runtime/interfaces"
+
+	"github.com/flyteorg/flytestdlib/config"
+)
+
+const clusterPoolsKey = "clusterPools"
+
+var clusterPoolsConfig = config.MustRegisterSection(clusterPoolsKey, &interfaces.ClusterPoolAssignmentConfig{
+	ClusterPoolAssignments: make(interfaces.ClusterPoolAssignments),
+})
+
+// Implementation of an interfaces.ClusterPoolAssignmentConfiguration
+type ClusterPoolAssignmentConfigurationProvider struct{}
+
+func (p *ClusterPoolAssignmentConfigurationProvider) GetClusterPoolAssignments() interfaces.ClusterPoolAssignments {
+	return clusterPoolsConfig.GetConfig().(*interfaces.ClusterPoolAssignmentConfig).ClusterPoolAssignments
+}
+
+func NewClusterPoolAssignmentConfigurationProvider() interfaces.ClusterPoolAssignmentConfiguration {
+	return &ClusterPoolAssignmentConfigurationProvider{}
+}

flyteadmin/pkg/runtime/configuration_provider.go

@@ -15,6 +15,7 @@ type ConfigurationProvider struct {
 	clusterResourceConfiguration        interfaces.ClusterResourceConfiguration
 	namespaceMappingConfiguration       interfaces.NamespaceMappingConfiguration
 	qualityOfServiceConfiguration       interfaces.QualityOfServiceConfiguration
+	clusterPoolAssignmentConfiguration  interfaces.ClusterPoolAssignmentConfiguration
 }
 
 func (p *ConfigurationProvider) ApplicationConfiguration() interfaces.ApplicationConfiguration {
@@ -53,6 +54,10 @@ func (p *ConfigurationProvider) QualityOfServiceConfiguration() interfaces.Quali
 	return p.qualityOfServiceConfiguration
 }
 
+func (p *ConfigurationProvider) ClusterPoolAssignmentConfiguration() interfaces.ClusterPoolAssignmentConfiguration {
+	return p.clusterPoolAssignmentConfiguration
+}
+
 func NewConfigurationProvider() interfaces.Configuration {
 	return &ConfigurationProvider{
 		applicationConfiguration:            NewApplicationConfigurationProvider(),
@@ -64,5 +69,6 @@ func NewConfigurationProvider() interfaces.Configuration {
 		clusterResourceConfiguration:        NewClusterResourceConfigurationProvider(),
 		namespaceMappingConfiguration:       NewNamespaceMappingConfigurationProvider(),
 		qualityOfServiceConfiguration:       NewQualityOfServiceConfigProvider(),
+		clusterPoolAssignmentConfiguration:  NewClusterPoolAssignmentConfigurationProvider(),
 	}
 }

flyteadmin/pkg/runtime/interfaces/cluster_pools.go

@@ -0,0 +1,17 @@
+package interfaces
+
+//go:generate mockery -name ClusterPoolAssignmentConfiguration -output=mocks -case=underscore
+
+type ClusterPoolAssignment struct {
+	Pool string `json:"pool"`
+}
+
+type ClusterPoolAssignments = map[DomainName]ClusterPoolAssignment
+
+type ClusterPoolAssignmentConfig struct {
+	ClusterPoolAssignments ClusterPoolAssignments `json:"clusterPoolAssignments"`
+}
+
+type ClusterPoolAssignmentConfiguration interface {
+	GetClusterPoolAssignments() ClusterPoolAssignments
+}

flyteadmin/pkg/runtime/interfaces/configuration.go

@@ -11,4 +11,5 @@ type Configuration interface {
 	ClusterResourceConfiguration() ClusterResourceConfiguration
 	NamespaceMappingConfiguration() NamespaceMappingConfiguration
 	QualityOfServiceConfiguration() QualityOfServiceConfiguration
+	ClusterPoolAssignmentConfiguration() ClusterPoolAssignmentConfiguration
 }

flyteadmin/pkg/runtime/mocks/mock_configuration_provider.go

@@ -16,6 +16,7 @@ type MockConfigurationProvider struct {
 	clusterResourceConfiguration        interfaces.ClusterResourceConfiguration
 	namespaceMappingConfiguration       interfaces.NamespaceMappingConfiguration
 	qualityOfServiceConfiguration       interfaces.QualityOfServiceConfiguration
+	clusterPoolAssignmentConfiguration  interfaces.ClusterPoolAssignmentConfiguration
 }
 
 func (p *MockConfigurationProvider) ApplicationConfiguration() interfaces.ApplicationConfiguration {
@@ -70,6 +71,14 @@ func (p *MockConfigurationProvider) AddQualityOfServiceConfiguration(config inte
 	p.qualityOfServiceConfiguration = config
 }
 
+func (p *MockConfigurationProvider) ClusterPoolAssignmentConfiguration() interfaces.ClusterPoolAssignmentConfiguration {
+	return p.clusterPoolAssignmentConfiguration
+}
+
+func (p *MockConfigurationProvider) AddClusterPoolAssignmentConfiguration(cfg interfaces.ClusterPoolAssignmentConfiguration) {
+	p.clusterPoolAssignmentConfiguration = cfg
+}
+
 func NewMockConfigurationProvider(
 	applicationConfiguration interfaces.ApplicationConfiguration,
 	queueConfiguration interfaces.QueueConfiguration,
@@ -82,13 +91,17 @@ func NewMockConfigurationProvider(
 	mockQualityOfServiceConfiguration.OnGetDefaultTiers().Return(make(map[string]core.QualityOfService_Tier))
 	mockQualityOfServiceConfiguration.OnGetTierExecutionValues().Return(make(map[core.QualityOfService_Tier]core.QualityOfServiceSpec))
 
+	mockClusterPoolAssignmentConfiguration := &ifaceMocks.ClusterPoolAssignmentConfiguration{}
+	mockClusterPoolAssignmentConfiguration.OnGetClusterPoolAssignments().Return(make(map[string]interfaces.ClusterPoolAssignment))
+
 	return &MockConfigurationProvider{
-		applicationConfiguration:      applicationConfiguration,
-		queueConfiguration:            queueConfiguration,
-		clusterConfiguration:          clusterConfiguration,
-		taskResourceConfiguration:     taskResourceConfiguration,
-		whitelistConfiguration:        whitelistConfiguration,
-		namespaceMappingConfiguration: namespaceMappingConfiguration,
-		qualityOfServiceConfiguration: mockQualityOfServiceConfiguration,
+		applicationConfiguration:           applicationConfiguration,
+		queueConfiguration:                 queueConfiguration,
+		clusterConfiguration:               clusterConfiguration,
+		taskResourceConfiguration:          taskResourceConfiguration,
+		whitelistConfiguration:             whitelistConfiguration,
+		namespaceMappingConfiguration:      namespaceMappingConfiguration,
+		qualityOfServiceConfiguration:      mockQualityOfServiceConfiguration,
+		clusterPoolAssignmentConfiguration: mockClusterPoolAssignmentConfiguration,
 	}
 }