Support overriding task annotations and labels via with_overrides #6171
Conversation
Signed-off-by: Nelson Chen <asd3431090@gmail.com>
Code Review Agent Run #d2a6a3Actionable Suggestions - 7
Additional Suggestions - 1
Review Details
AI Code Review powered by |
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #6171 +/- ##
==========================================
- Coverage 37.05% 34.39% -2.66%
==========================================
Files 1318 1097 -221
Lines 132583 114688 -17895
==========================================
- Hits 49122 39451 -9671
+ Misses 79211 71868 -7343
+ Partials 4250 3369 -881
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. |
Changelist by BitoThis pull request implements the following key changes.
|
| Annotations map[string]string `json:"annotation,omitempty" protobuf:"bytes,22,opt,name=annotation"` | ||
| // If specified, includes overrides for labels to allocate to the node | ||
| //+optional | ||
| Labels map[string]string `json:"label,omitempty" protobuf:"bytes,22,opt,name=label"` |
There was a problem hiding this comment.
Consider updating the protobuf field numbers for annotation and label fields as they both use field number 22 which is already used by tolerations. This could cause protobuf serialization issues.
Code suggestion
Check the AI-generated fix before applying
- Annotations map[string]string `json:"annotation,omitempty" protobuf:"bytes,22,opt,name=annotation"`
- Labels map[string]string `json:"label,omitempty" protobuf:"bytes,22,opt,name=label"`
+ Annotations map[string]string `json:"annotation,omitempty" protobuf:"bytes,23,opt,name=annotation"`
+ Labels map[string]string `json:"label,omitempty" protobuf:"bytes,24,opt,name=label"`
Code Review Run #d2a6a3
Is this a valid issue, or was it incorrectly flagged by the Agent?
- it was incorrectly flagged
| func (in *NodeSpec) GetAnnotations() map[string]string { | ||
| return in.Annotations | ||
| } | ||
|
|
||
| func (in *NodeSpec) GetLabels() map[string]string { | ||
| return in.Labels | ||
| } |
There was a problem hiding this comment.
Consider adding validation for nil maps in GetAnnotations() and GetLabels() to avoid potential nil pointer dereference. Similar issues were also found in:
- flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go (line 467)
- flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go (line 439)
- flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go (line 474)
Code suggestion
Check the AI-generated fix before applying
-func (in *NodeSpec) GetAnnotations() map[string]string {
- return in.Annotations
+func (in *NodeSpec) GetAnnotations() map[string]string {
+ if in.Annotations == nil {
+ return make(map[string]string)
+ }
+ return in.Annotations
}
-func (in *NodeSpec) GetLabels() map[string]string {
- return in.Labels
+func (in *NodeSpec) GetLabels() map[string]string {
+ if in.Labels == nil {
+ return make(map[string]string)
+ }
+ return in.Labels
}
Code Review Run #d2a6a3
Is this a valid issue, or was it incorrectly flagged by the Agent?
- it was incorrectly flagged
| GetAnnotations() map[string]string | ||
| GetLabels() map[string]string |
There was a problem hiding this comment.
Consider consolidating the GetAnnotations() and GetLabels() methods since they already exist in the Meta interface. This appears to be semantic duplication.
Code suggestion
Check the AI-generated fix before applying
- GetAnnotations() map[string]string
- GetLabels() map[string]string
+ Meta
}
Code Review Run #d2a6a3
Is this a valid issue, or was it incorrectly flagged by the Agent?
- it was incorrectly flagged
| to.On("GetAnnotations").Return(map[string]string{"annotation-1": "val1"}) | ||
| to.On("GetLabels").Return(map[string]string{"label-1": "val1"}) |
There was a problem hiding this comment.
Consider consolidating the mock setup for annotations and labels into the existing dummyTaskExecutionMetadata function since these are part of task metadata configuration.
Code suggestion
Check the AI-generated fix before applying
taskExecutionMetadata.On("GetNamespace").Return("test-namespace")
+ taskExecutionMetadata.On("GetAnnotations").Return(map[string]string{"annotation-1": "val1"})
+ taskExecutionMetadata.On("GetLabels").Return(map[string]string{"label-1": "val1"})
to := &pluginsCoreMock.TaskOverrides{}
to.On("GetResources").Return(resources)
to.On("GetExtendedResources").Return(extendedResources)
to.On("GetContainerImage").Return(containerImage)
- to.On("GetAnnotations").Return(map[string]string{"annotation-1": "val1"})
- to.On("GetLabels").Return(map[string]string{"label-1": "val1"})
Code Review Run #d2a6a3
Is this a valid issue, or was it incorrectly flagged by the Agent?
- it was incorrectly flagged
| var annotations map[string]string | ||
| var labels map[string]string |
There was a problem hiding this comment.
Consider initializing the maps annotations and labels with make() to avoid potential nil map issues when adding key-value pairs later.
Code suggestion
Check the AI-generated fix before applying
| var annotations map[string]string | |
| var labels map[string]string | |
| annotations := make(map[string]string) | |
| labels := make(map[string]string) |
Code Review Run #d2a6a3
Is this a valid issue, or was it incorrectly flagged by the Agent?
- it was incorrectly flagged
| if overrides.GetAnnotations() != nil { | ||
| annotations = overrides.GetAnnotations() | ||
| } | ||
|
|
||
| if overrides.GetLabels() != nil { | ||
| labels = overrides.GetLabels() | ||
| } |
There was a problem hiding this comment.
Consider combining the annotation and label override checks into a single function to improve code organization and reusability. The current implementation has similar logic repeated for both annotations and labels.
Code suggestion
Check the AI-generated fix before applying
| if overrides.GetAnnotations() != nil { | |
| annotations = overrides.GetAnnotations() | |
| } | |
| if overrides.GetLabels() != nil { | |
| labels = overrides.GetLabels() | |
| } | |
| func getOverrideMap(original, override map[string]string) map[string]string { | |
| if override != nil { | |
| return override | |
| } | |
| return original | |
| } | |
| annotations = getOverrideMap(annotations, overrides.GetAnnotations()) | |
| labels = getOverrideMap(labels, overrides.GetLabels()) |
Code Review Run #d2a6a3
Is this a valid issue, or was it incorrectly flagged by the Agent?
- it was incorrectly flagged
| // Optional labels to add to the pod definition. | ||
| Labels map[string]string `protobuf:"bytes,4,rep,name=labels,proto3" json:"labels,omitempty" protobuf_key:"bytes,1,opt,name=key,proto3" protobuf_val:"bytes,2,opt,name=value,proto3"` | ||
| // Optional annotations to add to the pod definition. | ||
| Annotations map[string]string `protobuf:"bytes,5,rep,name=annotations,proto3" json:"annotations,omitempty" protobuf_key:"bytes,1,opt,name=key,proto3" protobuf_val:"bytes,2,opt,name=value,proto3"` |
There was a problem hiding this comment.
Consider consolidating the pod metadata fields (labels and annotations) into a separate K8SObjectMetadata struct since these fields are already defined in that type. This would help maintain consistency and reduce duplication across the codebase.
Code suggestion
Check the AI-generated fix before applying
- // Optional labels to add to the pod definition.
- Labels map[string]string `protobuf:"bytes,4,rep,name=labels,proto3" json:"labels,omitempty" protobuf_key:"bytes,1,opt,name=key,proto3" protobuf_val:"bytes,2,opt,name=value,proto3"`
- // Optional annotations to add to the pod definition.
- Annotations map[string]string `protobuf:"bytes,5,rep,name=annotations,proto3" json:"annotations,omitempty" protobuf_key:"bytes,1,opt,name=key,proto3" protobuf_val:"bytes,2,opt,name=value,proto3"`
+ // Pod metadata including labels and annotations
+ PodMetadata *K8SObjectMetadata `protobuf:"bytes,4,opt,name=pod_metadata,json=podMetadata,proto3" json:"pod_metadata,omitempty"`
}
Code Review Run #d2a6a3
Is this a valid issue, or was it incorrectly flagged by the Agent?
- it was incorrectly flagged
arbaobao
force-pushed
the
override_label
branch
from
df815bc to
eeba0e5
Compare
Signed-off-by: Nelson Chen <asd3431090@gmail.com>
Code Review Agent Run #1ba60dActionable Suggestions - 2
Review Details
AI Code Review powered by |
| @@ -756,7 +756,7 @@ func updatePod(t *testing.T) { | |||
| } | |||
|
|
|||
| func TestUpdatePodWithDefaultAffinityAndInterruptibleNodeSelectorRequirement(t *testing.T) { | |||
| taskExecutionMetadata := dummyTaskExecutionMetadata(&v1.ResourceRequirements{}, nil, "") | |||
| taskExecutionMetadata := dummyTaskExecutionMetadata(&v1.ResourceRequirements{}, nil, "", nil, nil) | |||
| assert.NoError(t, config.SetK8sPluginConfig(&config.K8sPluginConfig{ | |||
| DefaultAffinity: &v1.Affinity{ | |||
There was a problem hiding this comment.
Consider updating all test cases to use the new function signature for dummyTaskExecutionMetadata with annotations and labels parameters. Some test cases are still using the old signature which may cause compilation errors.
Code suggestion
Check the AI-generated fix before applying
| DefaultAffinity: &v1.Affinity{ | |
| taskExecutionMetadata := dummyTaskExecutionMetadata(&v1.ResourceRequirements{}, nil, "", nil, nil) |
Code Review Run #1ba60d
Is this a valid issue, or was it incorrectly flagged by the Agent?
- it was incorrectly flagged
| func TestToK8sPodAnnotations(t *testing.T) { | ||
| t.Run("Override annotations", func(t *testing.T) { | ||
| taskContext := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{ | ||
| Requests: v1.ResourceList{ | ||
| v1.ResourceCPU: resource.MustParse("1024m"), | ||
| }}, nil, "foo:latest", map[string]string{"a": "b"}, nil) | ||
| _, m, _, err := ToK8sPodSpec(context.TODO(), taskContext) | ||
| assert.NoError(t, err) | ||
| assert.Equal(t, "b", m.Annotations["a"]) | ||
| }) | ||
| } | ||
|
|
||
| func TestToK8sPodLabels(t *testing.T) { | ||
| t.Run("Override labels", func(t *testing.T) { | ||
| taskContext := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{ | ||
| Requests: v1.ResourceList{ | ||
| v1.ResourceCPU: resource.MustParse("1024m"), | ||
| }}, nil, "foo:latest", nil, map[string]string{"a": "b"}) | ||
| _, m, _, err := ToK8sPodSpec(context.TODO(), taskContext) | ||
| assert.NoError(t, err) | ||
| assert.Equal(t, "b", m.Labels["a"]) | ||
| }) | ||
| } |
There was a problem hiding this comment.
Consider combining TestToK8sPodAnnotations and TestToK8sPodLabels into a single test function since they have very similar test logic and setup. This would reduce code duplication while still maintaining test coverage.
Code suggestion
Check the AI-generated fix before applying
| func TestToK8sPodAnnotations(t *testing.T) { | |
| t.Run("Override annotations", func(t *testing.T) { | |
| taskContext := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{ | |
| Requests: v1.ResourceList{ | |
| v1.ResourceCPU: resource.MustParse("1024m"), | |
| }}, nil, "foo:latest", map[string]string{"a": "b"}, nil) | |
| _, m, _, err := ToK8sPodSpec(context.TODO(), taskContext) | |
| assert.NoError(t, err) | |
| assert.Equal(t, "b", m.Annotations["a"]) | |
| }) | |
| } | |
| func TestToK8sPodLabels(t *testing.T) { | |
| t.Run("Override labels", func(t *testing.T) { | |
| taskContext := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{ | |
| Requests: v1.ResourceList{ | |
| v1.ResourceCPU: resource.MustParse("1024m"), | |
| }}, nil, "foo:latest", nil, map[string]string{"a": "b"}) | |
| _, m, _, err := ToK8sPodSpec(context.TODO(), taskContext) | |
| assert.NoError(t, err) | |
| assert.Equal(t, "b", m.Labels["a"]) | |
| }) | |
| } | |
| func TestToK8sPodMetadata(t *testing.T) { | |
| tests := []struct { | |
| name string | |
| annotations map[string]string | |
| labels map[string]string | |
| }{ | |
| {"Override annotations", map[string]string{"a": "b"}, nil}, | |
| {"Override labels", nil, map[string]string{"a": "b"}}, | |
| } | |
| for _, tt := range tests { | |
| t.Run(tt.name, func(t *testing.T) { | |
| taskContext := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{ | |
| Requests: v1.ResourceList{ | |
| v1.ResourceCPU: resource.MustParse("1024m"), | |
| }}, nil, "foo:latest", tt.annotations, tt.labels) | |
| _, m, _, err := ToK8sPodSpec(context.TODO(), taskContext) | |
| assert.NoError(t, err) | |
| if tt.annotations != nil { | |
| assert.Equal(t, "b", m.Annotations["a"]) | |
| } | |
| if tt.labels != nil { | |
| assert.Equal(t, "b", m.Labels["a"]) | |
| } | |
| }) | |
| } | |
| } |
Code Review Run #1ba60d
Is this a valid issue, or was it incorrectly flagged by the Agent?
- it was incorrectly flagged
Tracking issue
Why are the changes needed?
Improvement to allow setting labels and annotations dynamically at run time for things like cost allocation.
What changes were proposed in this pull request?
Adds labels and annotations to with_overrides().
How was this patch tested?
Excute a workflow and using with_override(labels={"lKeyA": "lValA"},annotations={"lKeyB": "lValB"}).
Setup process
It can be tested by using kubectl describe pods {pod_name} -n flytesnacks-development
Screenshots
Check all the applicable boxes
Related PRs
flytekit 3061
Docs link
Summary by Bito
This PR implements support for overriding task annotations and labels via TaskNodeOverrides and with_overrides(), enabling dynamic pod metadata configuration at runtime. The implementation includes protobuf definition updates across multiple language bindings (TypeScript, Go, JavaScript, Python), interface methods, and workflow compiler updates. The changes include updated test helper functions, new test cases, and mock interfaces for annotation and label getters. This enhancement particularly benefits scenarios like cost allocation through runtime pod metadata modification.Unit tests added: True
Estimated effort to review (1-5, lower is better): 5