Skip to content
This repository has been archived by the owner on Oct 9, 2023. It is now read-only.

Commit

Permalink
Add default cluster pool assignments to config
Browse files Browse the repository at this point in the history
Signed-off-by: Katrina Rogan <katroganGH@gmail.com>
  • Loading branch information
katrogan committed Aug 10, 2023
1 parent 478b64c commit 14c1cfe
Show file tree
Hide file tree
Showing 8 changed files with 144 additions and 9 deletions.
7 changes: 5 additions & 2 deletions pkg/manager/impl/execution_manager.go
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
24 changes: 24 additions & 0 deletions pkg/manager/impl/execution_manager_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
24 changes: 24 additions & 0 deletions pkg/runtime/cluster_pool_assignment_provider.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
package runtime

import (
"github.com/flyteorg/flyteadmin/pkg/runtime/interfaces"

"github.com/flyteorg/flytestdlib/config"
)

const clusterPoolsKey = "cluster_pools"

var clusterPoolsConfig = config.MustRegisterSection(clusterPoolsKey, &interfaces.ClusterPoolAssignmentConfig{
ClusterPoolAssignments: make(interfaces.ClusterPoolAssignments),
})

// Implementation of an interfaces.QueueConfiguration
type ClusterPoolAssignmentConfigurationProvider struct{}

func (p *ClusterPoolAssignmentConfigurationProvider) GetClusterPoolAssignments() interfaces.ClusterPoolAssignments {
return clusterPoolsConfig.GetConfig().(*interfaces.ClusterPoolAssignmentConfig).ClusterPoolAssignments

Check warning on line 19 in pkg/runtime/cluster_pool_assignment_provider.go

View check run for this annotation

Codecov / codecov/patch

pkg/runtime/cluster_pool_assignment_provider.go#L18-L19

Added lines #L18 - L19 were not covered by tests
}

func NewClusterPoolAssignmentConfigurationProvider() interfaces.ClusterPoolAssignmentConfiguration {
return &ClusterPoolAssignmentConfigurationProvider{}
}
6 changes: 6 additions & 0 deletions pkg/runtime/configuration_provider.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -53,6 +54,10 @@ func (p *ConfigurationProvider) QualityOfServiceConfiguration() interfaces.Quali
return p.qualityOfServiceConfiguration
}

func (p *ConfigurationProvider) ClusterPoolAssignmentConfiguration() interfaces.ClusterPoolAssignmentConfiguration {
return p.clusterPoolAssignmentConfiguration

Check warning on line 58 in pkg/runtime/configuration_provider.go

View check run for this annotation

Codecov / codecov/patch

pkg/runtime/configuration_provider.go#L57-L58

Added lines #L57 - L58 were not covered by tests
}

func NewConfigurationProvider() interfaces.Configuration {
return &ConfigurationProvider{
applicationConfiguration: NewApplicationConfigurationProvider(),
Expand All @@ -64,5 +69,6 @@ func NewConfigurationProvider() interfaces.Configuration {
clusterResourceConfiguration: NewClusterResourceConfigurationProvider(),
namespaceMappingConfiguration: NewNamespaceMappingConfigurationProvider(),
qualityOfServiceConfiguration: NewQualityOfServiceConfigProvider(),
clusterPoolAssignmentConfiguration: NewClusterPoolAssignmentConfigurationProvider(),
}
}
17 changes: 17 additions & 0 deletions pkg/runtime/interfaces/cluster_pools.go
Original file line number Diff line number Diff line change
@@ -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
}
1 change: 1 addition & 0 deletions pkg/runtime/interfaces/configuration.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,4 +11,5 @@ type Configuration interface {
ClusterResourceConfiguration() ClusterResourceConfiguration
NamespaceMappingConfiguration() NamespaceMappingConfiguration
QualityOfServiceConfiguration() QualityOfServiceConfiguration
ClusterPoolAssignmentConfiguration() ClusterPoolAssignmentConfiguration
}

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

27 changes: 20 additions & 7 deletions pkg/runtime/mocks/mock_configuration_provider.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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,
Expand All @@ -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,
}
}

0 comments on commit 14c1cfe

Please sign in to comment.