Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support overriding task annotations and labels via with_overrides #6171

Open
wants to merge 7 commits into
base: master
Choose a base branch
from
Open
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
add override for annotation and labels
Signed-off-by: Nelson Chen <asd3431090@gmail.com>
arbaobao committed Jan 15, 2025

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature.
commit 5f9eb9d461357c1ec5b6edd3463728f40c3d6f5b
17 changes: 14 additions & 3 deletions flyteidl/clients/go/assets/admin.swagger.json

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

16 changes: 16 additions & 0 deletions flyteidl/gen/pb-es/flyteidl/core/workflow_pb.ts

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

171 changes: 105 additions & 66 deletions flyteidl/gen/pb-go/flyteidl/core/workflow.pb.go

Large diffs are not rendered by default.

14 changes: 14 additions & 0 deletions flyteidl/gen/pb-go/gateway/flyteidl/service/admin.swagger.json

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

14 changes: 14 additions & 0 deletions flyteidl/gen/pb-go/gateway/flyteidl/service/agent.swagger.json

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

12 changes: 12 additions & 0 deletions flyteidl/gen/pb-js/flyteidl.d.ts

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

60 changes: 59 additions & 1 deletion flyteidl/gen/pb-js/flyteidl.js

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

16 changes: 12 additions & 4 deletions flyteidl/gen/pb_python/flyteidl/core/workflow_pb2.py

Large diffs are not rendered by default.

22 changes: 20 additions & 2 deletions flyteidl/gen/pb_python/flyteidl/core/workflow_pb2.pyi

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

6 changes: 6 additions & 0 deletions flyteidl/gen/pb_rust/flyteidl.core.rs

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

6 changes: 6 additions & 0 deletions flyteidl/protos/flyteidl/core/workflow.proto
Original file line number Diff line number Diff line change
@@ -336,6 +336,12 @@ message TaskNodeOverrides {

// Override for the image used by task pods.
string container_image = 3;

// Optional labels to add to the pod definition.
map<string, string> labels = 4;

// Optional annotations to add to the pod definition.
map<string, string> annotations = 5;
}

// A structure that uniquely identifies a launch plan in the system.
2 changes: 2 additions & 0 deletions flyteplugins/go/tasks/pluginmachinery/core/exec_metadata.go
Original file line number Diff line number Diff line change
@@ -14,6 +14,8 @@ type TaskOverrides interface {
GetExtendedResources() *core.ExtendedResources
GetContainerImage() string
GetConfig() *v1.ConfigMap
GetAnnotations() map[string]string
GetLabels() map[string]string
}

// TaskExecutionID is a simple Interface to expose the ExecutionID of the running Task
Original file line number Diff line number Diff line change
@@ -31,6 +31,14 @@ func (to *pluginTaskOverrides) GetContainerImage() string {
return to.TaskOverrides.GetContainerImage()
}

func (to *pluginTaskOverrides) GetAnnotations() map[string]string {
return to.TaskOverrides.GetAnnotations()
}

func (to *pluginTaskOverrides) GetLabels() map[string]string {
return to.TaskOverrides.GetLabels()
}

type pluginTaskExecutionMetadata struct {
pluginsCore.TaskExecutionMetadata
interruptible *bool
30 changes: 30 additions & 0 deletions flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go
Original file line number Diff line number Diff line change
@@ -434,6 +434,22 @@ func ApplyFlytePodConfiguration(ctx context.Context, tCtx pluginsCore.TaskExecut
ApplyContainerImageOverride(podSpec, tCtx.TaskExecutionMetadata().GetOverrides().GetContainerImage(), primaryContainerName)
}

// Override annotation if necessary
if tCtx.TaskExecutionMetadata().GetOverrides().GetAnnotations() != nil {
objectMeta, err = ApplyAnnotationOverride(objectMeta, tCtx.TaskExecutionMetadata().GetOverrides().GetAnnotations())
if err != nil {
return nil, nil, err
}
}

// Override label if necessary
if tCtx.TaskExecutionMetadata().GetOverrides().GetLabels() != nil {
objectMeta, err = ApplyLabelOverride(objectMeta, tCtx.TaskExecutionMetadata().GetOverrides().GetLabels())
if err != nil {
return nil, nil, err
}
}

return podSpec, objectMeta, nil
}

@@ -446,6 +462,20 @@ func ApplyContainerImageOverride(podSpec *v1.PodSpec, containerImage string, pri
}
}

func ApplyAnnotationOverride(objectMeta *metav1.ObjectMeta, annotations map[string]string) (*metav1.ObjectMeta, error) {
if annotations != nil {
mergeMapInto(annotations, objectMeta.Annotations)
}
return objectMeta, nil
}

func ApplyLabelOverride(objectMeta *metav1.ObjectMeta, labels map[string]string) (*metav1.ObjectMeta, error) {
if labels != nil {
mergeMapInto(labels, objectMeta.Labels)
}
return objectMeta, nil
}

func addTolerationInPodSpec(podSpec *v1.PodSpec, toleration *v1.Toleration) *v1.PodSpec {
podTolerations := podSpec.Tolerations

4 changes: 3 additions & 1 deletion flytepropeller/pkg/apis/flyteworkflow/v1alpha1/iface.go
Original file line number Diff line number Diff line change
@@ -442,7 +442,9 @@ type ExecutableNode interface {
GetActiveDeadline() *time.Duration
IsInterruptible() *bool
GetName() string
GetContainerImage() string
GetContainerImage()
GetAnnotations() map[string]string
GetLabels() map[string]string
Comment on lines +446 to +447
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider consolidating duplicate interface methods

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

Sorry, something went wrong.

}

// ExecutableWorkflowStatus is an interface for the Workflow p. This is the mutable portion for a Workflow
14 changes: 14 additions & 0 deletions flytepropeller/pkg/apis/flyteworkflow/v1alpha1/nodes.go
Original file line number Diff line number Diff line change
@@ -159,6 +159,12 @@ type NodeSpec struct {
Interruptible *bool `json:"interruptible,omitempty"`

ContainerImage string `json:"containerImage,omitempty"`
// If specified, includes overrides for annotations to allocate to the node
//+optional
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"`
Comment on lines +164 to +167
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Duplicate protobuf field numbers need updating

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

Sorry, something went wrong.

}

func (in *NodeSpec) GetName() string {
@@ -268,3 +274,11 @@ func (in *NodeSpec) GetInputBindings() []*Binding {
func (in *NodeSpec) GetContainerImage() string {
return in.ContainerImage
}

func (in *NodeSpec) GetAnnotations() map[string]string {
return in.Annotations
}

func (in *NodeSpec) GetLabels() map[string]string {
return in.Labels
}
Comment on lines +278 to +284
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider adding nil check for maps

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

Sorry, something went wrong.

12 changes: 12 additions & 0 deletions flytepropeller/pkg/compiler/transformers/k8s/node.go
Original file line number Diff line number Diff line change
@@ -31,6 +31,8 @@ func buildNodeSpec(n *core.Node, tasks []*core.CompiledTask, errs errors.Compile
var resources *core.Resources
var extendedResources *v1alpha1.ExtendedResources
var containerImage string
var annotations map[string]string
var labels map[string]string
Comment on lines +34 to +35
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider initializing maps with make()

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
Suggested change
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

Sorry, something went wrong.

if n.GetTaskNode() != nil {
taskID := n.GetTaskNode().GetReferenceId().String()
// TODO: Use task index for quick lookup
@@ -60,6 +62,14 @@ func buildNodeSpec(n *core.Node, tasks []*core.CompiledTask, errs errors.Compile
if len(overrides.GetContainerImage()) > 0 {
containerImage = overrides.GetContainerImage()
}

if overrides.GetAnnotations() != nil {
annotations = overrides.GetAnnotations()
}

if overrides.GetLabels() != nil {
labels = overrides.GetLabels()
}
Comment on lines +66 to +72
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider consolidating override handling logic

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
Suggested change
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

Sorry, something went wrong.

}
}

@@ -102,6 +112,8 @@ func buildNodeSpec(n *core.Node, tasks []*core.CompiledTask, errs errors.Compile
ActiveDeadline: activeDeadline,
Interruptible: interruptible,
ContainerImage: containerImage,
Annotations: annotations,
Labels: labels,
}

switch v := n.GetTarget().(type) {