Support overriding task pod_template via with_overrides #6118
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Signed-off-by: Nelson Chen <asd3431090@gmail.com>
Signed-off-by: Nelson Chen <asd3431090@gmail.com>
|
#take |
Code Review Agent Run #9af1d5Actionable Suggestions - 7
Review Details
AI Code Review powered by |
Changelist by BitoThis pull request implements the following key changes.
|
flyte-bot
reviewed
Dec 29, 2024
Comment on lines
+382
to
+395
| t.Run("Test_enable-distributed-error-aggregation", func(t *testing.T) { | ||
|
|
||
| t.Run("Override", func(t *testing.T) { | ||
| testValue := "1" | ||
|
|
||
| cmdFlags.Set("enable-distributed-error-aggregation", testValue) | ||
| if vBool, err := cmdFlags.GetBool("enable-distributed-error-aggregation"); err == nil { | ||
| testDecodeJson_K8sPluginConfig(t, fmt.Sprintf("%v", vBool), &actual.EnableDistributedErrorAggregation) | ||
|
|
||
| } else { | ||
| assert.FailNow(t, err.Error()) | ||
| } | ||
| }) | ||
| }) |
There was a problem hiding this comment.
Consider adding more test cases
Consider adding test cases for both true and false values for the enable-distributed-error-aggregation flag to ensure proper boolean parsing behavior.
Code suggestion
Check the AI-generated fix before applying
@@ -382,14 +382,40 @@ func TestK8sPluginConfig_SetFlags(t *testing.T) {
t.Run("Test_enable-distributed-error-aggregation", func(t *testing.T) {
t.Run("Override with 1", func(t *testing.T) {
testValue := "1"
cmdFlags.Set("enable-distributed-error-aggregation", testValue)
if vBool, err := cmdFlags.GetBool("enable-distributed-error-aggregation"); err == nil {
testDecodeJson_K8sPluginConfig(t, fmt.Sprintf("%v", vBool), &actual.EnableDistributedErrorAggregation)
} else {
assert.FailNow(t, err.Error())
}
})
t.Run("Override with true", func(t *testing.T) {
testValue := "true"
cmdFlags.Set("enable-distributed-error-aggregation", testValue)
if vBool, err := cmdFlags.GetBool("enable-distributed-error-aggregation"); err == nil {
testDecodeJson_K8sPluginConfig(t, fmt.Sprintf("%v", vBool), &actual.EnableDistributedErrorAggregation)
} else {
assert.FailNow(t, err.Error())
}
})
t.Run("Override with false", func(t *testing.T) {
testValue := "false"
cmdFlags.Set("enable-distributed-error-aggregation", testValue)
if vBool, err := cmdFlags.GetBool("enable-distributed-error-aggregation"); err == nil {
testDecodeJson_K8sPluginConfig(t, fmt.Sprintf("%v", vBool), &actual.EnableDistributedErrorAggregation)
} else {
assert.FailNow(t, err.Error())
}
})
})
Code Review Run #9af1d5
Is this a valid issue, or was it incorrectly flagged by the Agent?
- it was incorrectly flagged
flyte-bot
reviewed
Dec 29, 2024
flyteidl/gen/pb-js/flyteidl.d.ts
Outdated
| @@ -6953,6 +6959,9 @@ export namespace flyteidl { | |||
|
|
|||
| /** K8sPod dataConfig */ | |||
| dataConfig?: (flyteidl.core.IDataLoadingConfig|null); | |||
|
|
|||
| /** K8sPod primaryContainerName */ | |||
| primaryContainerName?: (string|null); | |||
There was a problem hiding this comment.
Consider making primaryContainerName field required
Consider making primaryContainerName required instead of optional since it appears to be a critical field for K8sPod configuration. The interface shows it as optional with ? but the class implementation has it as required. A similar issue was also found in flyteidl/gen/pb-js/flyteidl.d.ts (line 6985-6986).
Code suggestion
Check the AI-generated fix before applying
Suggested change
| primaryContainerName?: (string|null); | |
| primaryContainerName: string; |
Code Review Run #9af1d5
Is this a valid issue, or was it incorrectly flagged by the Agent?
- it was incorrectly flagged
flyte-bot
reviewed
Dec 29, 2024
| @@ -31,6 +31,7 @@ func buildNodeSpec(n *core.Node, tasks []*core.CompiledTask, errs errors.Compile | |||
| var resources *core.Resources | |||
| var extendedResources *v1alpha1.ExtendedResources | |||
| var containerImage string | |||
| var podtemplate *core.K8SPod | |||
There was a problem hiding this comment.
Consider initializing podtemplate variable
Consider initializing the podtemplate variable with a default value to avoid potential nil pointer dereference issues.
Code suggestion
Check the AI-generated fix before applying
Suggested change
| var podtemplate *core.K8SPod | |
| var podtemplate *core.K8SPod = &core.K8SPod{} |
Code Review Run #9af1d5
Is this a valid issue, or was it incorrectly flagged by the Agent?
- it was incorrectly flagged
flyte-bot
reviewed
Dec 29, 2024
Comment on lines
+65
to
+67
| if overrides.GetPodTemplate() != nil { | ||
| podtemplate = overrides.GetPodTemplate() | ||
| } |
There was a problem hiding this comment.
Consider validating pod template before assignment
Consider adding validation for podtemplate before assignment. The current code directly assigns the pod template without any validation which could lead to issues if an invalid template is provided.
Code suggestion
Check the AI-generated fix before applying
Suggested change
| if overrides.GetPodTemplate() != nil { | |
| podtemplate = overrides.GetPodTemplate() | |
| } | |
| if template := overrides.GetPodTemplate(); template != nil { | |
| if err := validatePodTemplate(template); err != nil { | |
| return nil, err | |
| } | |
| podtemplate = template | |
| } |
Code Review Run #9af1d5
Is this a valid issue, or was it incorrectly flagged by the Agent?
- it was incorrectly flagged
flyte-bot
reviewed
Dec 29, 2024
flyte-bot
reviewed
Dec 29, 2024
| @@ -148,6 +148,7 @@ func RunPluginEndToEndTest(t *testing.T, executor pluginCore.Plugin, template *i | |||
| }) | |||
| overrides.OnGetExtendedResources().Return(&idlCore.ExtendedResources{}) | |||
| overrides.OnGetContainerImage().Return("") | |||
| overrides.OnGetPodTemplate().Return(&idlCore.K8SPod{}) | |||
There was a problem hiding this comment.
Consider adding error handling for mock
Consider adding error handling for the mock return value setup. The OnGetPodTemplate() call may need error handling to properly simulate failure scenarios.
Code suggestion
Check the AI-generated fix before applying
Suggested change
| overrides.OnGetPodTemplate().Return(&idlCore.K8SPod{}) | |
| overrides.OnGetPodTemplate().Return(&idlCore.K8SPod{}, nil) |
Code Review Run #9af1d5
Is this a valid issue, or was it incorrectly flagged by the Agent?
- it was incorrectly flagged
Signed-off-by: Nelson Chen <asd3431090@gmail.com>
Code Review Agent Run #81015aActionable Suggestions - 0Additional Suggestions - 10
Review Details
AI Code Review powered by |
pingsutw
reviewed
Jan 11, 2025
Comment on lines
340
to
341
| // Override for the pod template used by task pods | ||
| K8sPod pod_template = 4; |
There was a problem hiding this comment.
Suggested change
| // Override for the pod template used by task pods | |
| K8sPod pod_template = 4; | |
| // Override for the pod template used by task pods. | |
| //+optional | |
| K8sPod pod_template = 4; |
| @@ -286,6 +286,13 @@ func BuildRawPod(ctx context.Context, tCtx pluginsCore.TaskExecutionContext) (*v | |||
| *c, | |||
| }, | |||
| } | |||
| if tCtx.TaskExecutionMetadata().GetOverrides().GetPodTemplate() != nil { | |||
There was a problem hiding this comment.
nit: could we create a variable for podTemplate here? like
podTemplate := tCtx.TaskExecutionMetadata().GetOverrides().GetPodTemplate
if podTemplate != nil {
...
}| if podtemplate.Metadata.Labels != nil { | ||
| mergeMapInto(podtemplate.Metadata.Labels, objectMeta.Labels) | ||
| } | ||
| var podspec_override *v1.PodSpec |
There was a problem hiding this comment.
Suggested change
| var podspec_override *v1.PodSpec | |
| var podpec_override *v1.PodSpec |
| @@ -2022,7 +2023,7 @@ func TestMergeWithBasePodTemplate(t *testing.T) { | |||
| taskReader.On("Read", mock.Anything).Return(task, nil) | |||
|
|
|||
| tCtx := &pluginsCoreMock.TaskExecutionContext{} | |||
| tCtx.OnTaskExecutionMetadata().Return(dummyTaskExecutionMetadata(&v1.ResourceRequirements{}, nil, "")) | |||
| tCtx.OnTaskExecutionMetadata().Return(dummyTaskExecutionMetadata(&v1.ResourceRequirements{}, nil, "", nil)) | |||
There was a problem hiding this comment.
Could you add a test that has a podTemplate in the dummyTaskExecution?
| @@ -446,6 +453,25 @@ func ApplyContainerImageOverride(podSpec *v1.PodSpec, containerImage string, pri | |||
| } | |||
| } | |||
|
|
|||
| func ApplyPodTemplateOverride(podSpec *v1.PodSpec, objectMeta metav1.ObjectMeta, podtemplate *core.K8SPod) (*v1.PodSpec, metav1.ObjectMeta, error) { | |||
There was a problem hiding this comment.
Suggested change
| func ApplyPodTemplateOverride(podSpec *v1.PodSpec, objectMeta metav1.ObjectMeta, podtemplate *core.K8SPod) (*v1.PodSpec, metav1.ObjectMeta, error) { | |
| func ApplyPodTemplateOverride(podSpec *v1.PodSpec, objectMeta metav1.ObjectMeta, podTemplate *core.K8SPod) (*v1.PodSpec, metav1.ObjectMeta, error) { |
Signed-off-by: Nelson Chen <asd3431090@gmail.com>
|
@pingsutw Hi Kevin, I have added a unit test in pod_helper_test.go and fixed some typos that you mentioned. |
Code Review Agent Run #7e1782Actionable Suggestions - 1
Review Details
AI Code Review powered by |
flyte-bot
reviewed
Jan 14, 2025
Comment on lines
+1090
to
+1105
| t.Run("Override pod template", func(t *testing.T) { | ||
| taskContext := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{ | ||
| Requests: v1.ResourceList{ | ||
| v1.ResourceCPU: resource.MustParse("1024m"), | ||
| }}, nil, "", &core.K8SPod{ | ||
| PrimaryContainerName: "foo", | ||
| PodSpec: podSpecStruct, | ||
| Metadata: metadata, | ||
| }) | ||
| p, m, _, err := ToK8sPodSpec(context.TODO(), taskContext) | ||
| assert.NoError(t, err) | ||
| assert.Equal(t, "a", m.Labels["l"]) | ||
| assert.Equal(t, "b", m.Annotations["a"]) | ||
| assert.Equal(t, "foo:latest", p.Containers[0].Image) | ||
| assert.Equal(t, "foo", p.Containers[0].Name) | ||
| }) |
There was a problem hiding this comment.
Consider adding error test cases
Consider adding test cases to verify error scenarios and edge cases in TestPodTemplateOverride. For example, testing with invalid pod specs, missing container names, or empty metadata.
Code suggestion
Check the AI-generated fix before applying
Suggested change
| t.Run("Override pod template", func(t *testing.T) { | |
| taskContext := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{ | |
| Requests: v1.ResourceList{ | |
| v1.ResourceCPU: resource.MustParse("1024m"), | |
| }}, nil, "", &core.K8SPod{ | |
| PrimaryContainerName: "foo", | |
| PodSpec: podSpecStruct, | |
| Metadata: metadata, | |
| }) | |
| p, m, _, err := ToK8sPodSpec(context.TODO(), taskContext) | |
| assert.NoError(t, err) | |
| assert.Equal(t, "a", m.Labels["l"]) | |
| assert.Equal(t, "b", m.Annotations["a"]) | |
| assert.Equal(t, "foo:latest", p.Containers[0].Image) | |
| assert.Equal(t, "foo", p.Containers[0].Name) | |
| }) | |
| t.Run("Override pod template", func(t *testing.T) { | |
| taskContext := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{ | |
| Requests: v1.ResourceList{ | |
| v1.ResourceCPU: resource.MustParse("1024m"), | |
| }}, nil, "", &core.K8SPod{ | |
| PrimaryContainerName: "foo", | |
| PodSpec: podSpecStruct, | |
| Metadata: metadata, | |
| }) | |
| p, m, _, err := ToK8sPodSpec(context.TODO(), taskContext) | |
| assert.NoError(t, err) | |
| assert.Equal(t, "a", m.Labels["l"]) | |
| assert.Equal(t, "b", m.Annotations["a"]) | |
| assert.Equal(t, "foo:latest", p.Containers[0].Image) | |
| assert.Equal(t, "foo", p.Containers[0].Name) | |
| }) | |
| t.Run("Invalid pod spec", func(t *testing.T) { | |
| taskContext := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{ | |
| Requests: v1.ResourceList{ | |
| v1.ResourceCPU: resource.MustParse("1024m"), | |
| }}, nil, "", &core.K8SPod{ | |
| PrimaryContainerName: "foo", | |
| PodSpec: nil, | |
| Metadata: metadata, | |
| }) | |
| _, _, _, err := ToK8sPodSpec(context.TODO(), taskContext) | |
| assert.Error(t, err) | |
| }) |
Code Review Run #7e1782
Is this a valid issue, or was it incorrectly flagged by the Agent?
- it was incorrectly flagged
Signed-off-by: Nelson Chen <asd3431090@gmail.com>
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #6118 +/- ##
==========================================
- Coverage 37.05% 31.43% -5.62%
==========================================
Files 1318 781 -537
Lines 132583 84280 -48303
==========================================
- Hits 49122 26494 -22628
+ Misses 79211 55673 -23538
+ Partials 4250 2113 -2137
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. |
Code Review Agent Run #1518a4Actionable Suggestions - 0Additional Suggestions - 1
Review Details
AI Code Review powered by |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Tracking issue
Related to #5683
Why are the changes needed?
If we can support pod_template in with_overrides, this would reduce a lot of toil since we can supply pod templates in a central location and override downstream tasks, similar to how we can do so for resources.
What changes were proposed in this pull request?
We can use with_override() to override podtemplate, just like resources.
How was this patch tested?
Excute a workflow and using with_override(pod_template=PodTemplate(xxx)) to override the default podtemplate
Setup process
I ran flyte on my local machine and tested my code with this workflow and task:
Screenshots
Check all the applicable boxes
Related PRs
Docs link
Summary by Bito
This PR implements pod template override functionality with with_overrides() method across multiple languages, including security enhancements through configurable security contexts in deployment templates. The implementation includes protobuf definitions, pod template metadata handling, labels, annotations, and container specifications. The changes also improve the email notification system with retry mechanisms, add container port naming support, and enhance array node event handling in Flytepropeller. Updates include stow dependency changes and comprehensive security and configuration improvements.Unit tests added: True
Estimated effort to review (1-5, lower is better): 5