Skip to content

Commit

Permalink
Fix pool validation (#369)
Browse files Browse the repository at this point in the history
* Apparently, its possible to store clusterassignment in Resource API with empty cluster pool string

*TODO: Summarize tests added, integration tests run, or other steps you took to validate this change. Include (or link to) relevant test output or screenshots.*

* no concerns

Should this change be upstreamed to OSS (flyteorg/flyte)? If not, please uncheck this box, which is used for auditing. Note, it is the responsibility of each developer to actually upstream their changes. See [this guide](https://unionai.atlassian.net/wiki/spaces/ENG/pages/447610883/Flyte+-+Union+Cloud+Development+Runbook/#When-are-versions-updated%3F).
- [x] To be upstreamed to OSS

*TODO: Link Linear issue(s) using [magic words](https://linear.app/docs/github#magic-words). `fixes` will move to merged status, while `ref` will only link the PR.*

* [ ] Added tests
* [ ] Ran a deploy dry run and shared the terraform plan
* [ ] Added logging and metrics
* [ ] Updated [dashboards](https://unionai.grafana.net/dashboards) and [alerts](https://unionai.grafana.net/alerting/list)
* [ ] Updated documentation
  • Loading branch information
iaroslav-ciupin committed Sep 26, 2024
1 parent 194288c commit e3fc6a3
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 5 deletions.
11 changes: 6 additions & 5 deletions flyteadmin/pkg/manager/impl/execution_manager.go
Original file line number Diff line number Diff line change
Expand Up @@ -409,16 +409,17 @@ func (m *ExecutionManager) getClusterAssignment(ctx context.Context, req *admin.
return nil, err

Check warning on line 409 in flyteadmin/pkg/manager/impl/execution_manager.go

View check run for this annotation

Codecov / codecov/patch

flyteadmin/pkg/manager/impl/execution_manager.go#L409

Added line #L409 was not covered by tests
}

if req.GetSpec().GetClusterAssignment() == nil {
reqAssignment := req.GetSpec().GetClusterAssignment()
reqPool := reqAssignment.GetClusterPoolName()
storedPool := storedAssignment.GetClusterPoolName()
if reqPool == "" {
return storedAssignment, nil
}

if storedAssignment == nil {
return req.GetSpec().GetClusterAssignment(), nil
if storedPool == "" {
return reqAssignment, nil
}

reqPool := req.Spec.ClusterAssignment.GetClusterPoolName()
storedPool := storedAssignment.GetClusterPoolName()
if reqPool != storedPool {
return nil, errors.NewFlyteAdminErrorf(codes.InvalidArgument, "execution with project %q and domain %q cannot run on cluster pool %q, because its configured to run on pool %q", req.Project, req.Domain, reqPool, storedPool)
}
Expand Down
26 changes: 26 additions & 0 deletions flyteadmin/pkg/manager/impl/execution_manager_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -5519,6 +5519,32 @@ func TestGetClusterAssignment(t *testing.T) {
assert.NoError(t, err)
assert.True(t, proto.Equal(ca, &reqClusterAssignment))
})
t.Run("empty value in DB, takes value from request", func(t *testing.T) {
clusterPoolAsstProvider := &runtimeIFaceMocks.ClusterPoolAssignmentConfiguration{}
clusterPoolAsstProvider.OnGetClusterPoolAssignments().Return(runtimeInterfaces.ClusterPoolAssignments{
workflowIdentifier.GetDomain(): runtimeInterfaces.ClusterPoolAssignment{
Pool: "",
},
})
mockConfig := getMockExecutionsConfigProvider()
mockConfig.(*runtimeMocks.MockConfigurationProvider).AddClusterPoolAssignmentConfiguration(clusterPoolAsstProvider)

executionManager := ExecutionManager{
resourceManager: &managerMocks.MockResourceManager{},
config: mockConfig,
}

reqClusterAssignment := admin.ClusterAssignment{ClusterPoolName: "gpu"}
ca, err := executionManager.getClusterAssignment(context.TODO(), &admin.ExecutionCreateRequest{
Project: workflowIdentifier.Project,
Domain: workflowIdentifier.Domain,
Spec: &admin.ExecutionSpec{
ClusterAssignment: &reqClusterAssignment,
},
})
assert.NoError(t, err)
assert.True(t, proto.Equal(ca, &reqClusterAssignment))
})
t.Run("value from request doesn't match value from config", func(t *testing.T) {
reqClusterAssignment := admin.ClusterAssignment{ClusterPoolName: "swimming-pool"}
_, err := executionManager.getClusterAssignment(context.TODO(), &admin.ExecutionCreateRequest{
Expand Down

0 comments on commit e3fc6a3

Please sign in to comment.