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 <[email protected]>
Signed-off-by: Nelson Chen <[email protected]>
Signed-off-by: Nelson Chen <[email protected]>
Signed-off-by: Nelson Chen <[email protected]>
Signed-off-by: Nelson Chen <[email protected]>
|
#take |
Signed-off-by: Nelson Chen <[email protected]>
Signed-off-by: Nelson Chen <[email protected]>
Signed-off-by: Nelson Chen <[email protected]>
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 <[email protected]>
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) { |
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 and updates across the Flyte ecosystem (v1.14.0 to v1.14.1). Key features include primary_container_name field in K8sPod, pod_template in TaskNodeOverrides, and workflow node metadata support. The implementation spans multiple languages with protobuf definitions, enabling centralized pod template specification with downstream task override capability. Additional improvements include new resource management configurations, enhanced error handling, and a new DataMode enum for ArrayNodes.Unit tests added: True
Estimated effort to review (1-5, lower is better): 5