diff --git a/flyteidl/clients/go/assets/admin.swagger.json b/flyteidl/clients/go/assets/admin.swagger.json index 45f471248b..c3f23ea125 100644 --- a/flyteidl/clients/go/assets/admin.swagger.json +++ b/flyteidl/clients/go/assets/admin.swagger.json @@ -8334,6 +8334,20 @@ "container_image": { "type": "string", "description": "Override for the image used by task pods." + }, + "labels": { + "type": "object", + "additionalProperties": { + "type": "string" + }, + "description": "Optional labels to add to the pod definition." + }, + "annotations": { + "type": "object", + "additionalProperties": { + "type": "string" + }, + "description": "Optional annotations to add to the pod definition." } }, "description": "Optional task node overrides that will be applied at task execution time." diff --git a/flyteidl/gen/pb-es/flyteidl/core/workflow_pb.ts b/flyteidl/gen/pb-es/flyteidl/core/workflow_pb.ts index 04295d0a16..d207846572 100644 --- a/flyteidl/gen/pb-es/flyteidl/core/workflow_pb.ts +++ b/flyteidl/gen/pb-es/flyteidl/core/workflow_pb.ts @@ -1214,6 +1214,20 @@ export class TaskNodeOverrides extends Message { */ containerImage = ""; + /** + * Optional labels to add to the pod definition. + * + * @generated from field: map labels = 4; + */ + labels: { [key: string]: string } = {}; + + /** + * Optional annotations to add to the pod definition. + * + * @generated from field: map annotations = 5; + */ + annotations: { [key: string]: string } = {}; + constructor(data?: PartialMessage) { super(); proto3.util.initPartial(data, this); @@ -1225,6 +1239,8 @@ export class TaskNodeOverrides extends Message { { no: 1, name: "resources", kind: "message", T: Resources }, { no: 2, name: "extended_resources", kind: "message", T: ExtendedResources }, { no: 3, name: "container_image", kind: "scalar", T: 9 /* ScalarType.STRING */ }, + { no: 4, name: "labels", kind: "map", K: 9 /* ScalarType.STRING */, V: {kind: "scalar", T: 9 /* ScalarType.STRING */} }, + { no: 5, name: "annotations", kind: "map", K: 9 /* ScalarType.STRING */, V: {kind: "scalar", T: 9 /* ScalarType.STRING */} }, ]); static fromBinary(bytes: Uint8Array, options?: Partial): TaskNodeOverrides { diff --git a/flyteidl/gen/pb-go/flyteidl/core/workflow.pb.go b/flyteidl/gen/pb-go/flyteidl/core/workflow.pb.go index 26392969fe..89128ee2f0 100644 --- a/flyteidl/gen/pb-go/flyteidl/core/workflow.pb.go +++ b/flyteidl/gen/pb-go/flyteidl/core/workflow.pb.go @@ -1652,6 +1652,10 @@ type TaskNodeOverrides struct { ExtendedResources *ExtendedResources `protobuf:"bytes,2,opt,name=extended_resources,json=extendedResources,proto3" json:"extended_resources,omitempty"` // Override for the image used by task pods. ContainerImage string `protobuf:"bytes,3,opt,name=container_image,json=containerImage,proto3" json:"container_image,omitempty"` + // 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"` } func (x *TaskNodeOverrides) Reset() { @@ -1707,6 +1711,20 @@ func (x *TaskNodeOverrides) GetContainerImage() string { return "" } +func (x *TaskNodeOverrides) GetLabels() map[string]string { + if x != nil { + return x.Labels + } + return nil +} + +func (x *TaskNodeOverrides) GetAnnotations() map[string]string { + if x != nil { + return x.Annotations + } + return nil +} + // A structure that uniquely identifies a launch plan in the system. type LaunchPlanTemplate struct { state protoimpl.MessageState @@ -2026,7 +2044,7 @@ var file_flyteidl_core_workflow_proto_rawDesc = []byte{ 0x63, 0x6f, 0x72, 0x65, 0x2e, 0x57, 0x6f, 0x72, 0x6b, 0x66, 0x6c, 0x6f, 0x77, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x44, 0x65, 0x66, 0x61, 0x75, 0x6c, 0x74, 0x73, 0x52, 0x10, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x44, 0x65, 0x66, 0x61, 0x75, 0x6c, 0x74, 0x73, 0x22, - 0xc5, 0x01, 0x0a, 0x11, 0x54, 0x61, 0x73, 0x6b, 0x4e, 0x6f, 0x64, 0x65, 0x4f, 0x76, 0x65, 0x72, + 0xdb, 0x03, 0x0a, 0x11, 0x54, 0x61, 0x73, 0x6b, 0x4e, 0x6f, 0x64, 0x65, 0x4f, 0x76, 0x65, 0x72, 0x72, 0x69, 0x64, 0x65, 0x73, 0x12, 0x36, 0x0a, 0x09, 0x72, 0x65, 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x73, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x18, 0x2e, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2e, 0x63, 0x6f, 0x72, 0x65, 0x2e, 0x52, 0x65, 0x73, 0x6f, 0x75, 0x72, 0x63, @@ -2038,31 +2056,48 @@ var file_flyteidl_core_workflow_proto_rawDesc = []byte{ 0x65, 0x6e, 0x64, 0x65, 0x64, 0x52, 0x65, 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x73, 0x12, 0x27, 0x0a, 0x0f, 0x63, 0x6f, 0x6e, 0x74, 0x61, 0x69, 0x6e, 0x65, 0x72, 0x5f, 0x69, 0x6d, 0x61, 0x67, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0e, 0x63, 0x6f, 0x6e, 0x74, 0x61, 0x69, 0x6e, - 0x65, 0x72, 0x49, 0x6d, 0x61, 0x67, 0x65, 0x22, 0xba, 0x01, 0x0a, 0x12, 0x4c, 0x61, 0x75, 0x6e, - 0x63, 0x68, 0x50, 0x6c, 0x61, 0x6e, 0x54, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x12, 0x29, - 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x19, 0x2e, 0x66, 0x6c, 0x79, - 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2e, 0x63, 0x6f, 0x72, 0x65, 0x2e, 0x49, 0x64, 0x65, 0x6e, 0x74, - 0x69, 0x66, 0x69, 0x65, 0x72, 0x52, 0x02, 0x69, 0x64, 0x12, 0x3b, 0x0a, 0x09, 0x69, 0x6e, 0x74, - 0x65, 0x72, 0x66, 0x61, 0x63, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1d, 0x2e, 0x66, - 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2e, 0x63, 0x6f, 0x72, 0x65, 0x2e, 0x54, 0x79, 0x70, - 0x65, 0x64, 0x49, 0x6e, 0x74, 0x65, 0x72, 0x66, 0x61, 0x63, 0x65, 0x52, 0x09, 0x69, 0x6e, 0x74, - 0x65, 0x72, 0x66, 0x61, 0x63, 0x65, 0x12, 0x3c, 0x0a, 0x0c, 0x66, 0x69, 0x78, 0x65, 0x64, 0x5f, - 0x69, 0x6e, 0x70, 0x75, 0x74, 0x73, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x19, 0x2e, 0x66, - 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2e, 0x63, 0x6f, 0x72, 0x65, 0x2e, 0x4c, 0x69, 0x74, - 0x65, 0x72, 0x61, 0x6c, 0x4d, 0x61, 0x70, 0x52, 0x0b, 0x66, 0x69, 0x78, 0x65, 0x64, 0x49, 0x6e, - 0x70, 0x75, 0x74, 0x73, 0x42, 0xb3, 0x01, 0x0a, 0x11, 0x63, 0x6f, 0x6d, 0x2e, 0x66, 0x6c, 0x79, - 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2e, 0x63, 0x6f, 0x72, 0x65, 0x42, 0x0d, 0x57, 0x6f, 0x72, 0x6b, - 0x66, 0x6c, 0x6f, 0x77, 0x50, 0x72, 0x6f, 0x74, 0x6f, 0x50, 0x01, 0x5a, 0x3a, 0x67, 0x69, 0x74, - 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x6f, 0x72, 0x67, - 0x2f, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x2f, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2f, - 0x67, 0x65, 0x6e, 0x2f, 0x70, 0x62, 0x2d, 0x67, 0x6f, 0x2f, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69, - 0x64, 0x6c, 0x2f, 0x63, 0x6f, 0x72, 0x65, 0xa2, 0x02, 0x03, 0x46, 0x43, 0x58, 0xaa, 0x02, 0x0d, - 0x46, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2e, 0x43, 0x6f, 0x72, 0x65, 0xca, 0x02, 0x0d, - 0x46, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x5c, 0x43, 0x6f, 0x72, 0x65, 0xe2, 0x02, 0x19, - 0x46, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x5c, 0x43, 0x6f, 0x72, 0x65, 0x5c, 0x47, 0x50, - 0x42, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0xea, 0x02, 0x0e, 0x46, 0x6c, 0x79, 0x74, - 0x65, 0x69, 0x64, 0x6c, 0x3a, 0x3a, 0x43, 0x6f, 0x72, 0x65, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, - 0x6f, 0x33, + 0x65, 0x72, 0x49, 0x6d, 0x61, 0x67, 0x65, 0x12, 0x44, 0x0a, 0x06, 0x6c, 0x61, 0x62, 0x65, 0x6c, + 0x73, 0x18, 0x04, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x2c, 0x2e, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69, + 0x64, 0x6c, 0x2e, 0x63, 0x6f, 0x72, 0x65, 0x2e, 0x54, 0x61, 0x73, 0x6b, 0x4e, 0x6f, 0x64, 0x65, + 0x4f, 0x76, 0x65, 0x72, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2e, 0x4c, 0x61, 0x62, 0x65, 0x6c, 0x73, + 0x45, 0x6e, 0x74, 0x72, 0x79, 0x52, 0x06, 0x6c, 0x61, 0x62, 0x65, 0x6c, 0x73, 0x12, 0x53, 0x0a, + 0x0b, 0x61, 0x6e, 0x6e, 0x6f, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x73, 0x18, 0x05, 0x20, 0x03, + 0x28, 0x0b, 0x32, 0x31, 0x2e, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2e, 0x63, 0x6f, + 0x72, 0x65, 0x2e, 0x54, 0x61, 0x73, 0x6b, 0x4e, 0x6f, 0x64, 0x65, 0x4f, 0x76, 0x65, 0x72, 0x72, + 0x69, 0x64, 0x65, 0x73, 0x2e, 0x41, 0x6e, 0x6e, 0x6f, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x73, + 0x45, 0x6e, 0x74, 0x72, 0x79, 0x52, 0x0b, 0x61, 0x6e, 0x6e, 0x6f, 0x74, 0x61, 0x74, 0x69, 0x6f, + 0x6e, 0x73, 0x1a, 0x39, 0x0a, 0x0b, 0x4c, 0x61, 0x62, 0x65, 0x6c, 0x73, 0x45, 0x6e, 0x74, 0x72, + 0x79, 0x12, 0x10, 0x0a, 0x03, 0x6b, 0x65, 0x79, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, + 0x6b, 0x65, 0x79, 0x12, 0x14, 0x0a, 0x05, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x18, 0x02, 0x20, 0x01, + 0x28, 0x09, 0x52, 0x05, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x3a, 0x02, 0x38, 0x01, 0x1a, 0x3e, 0x0a, + 0x10, 0x41, 0x6e, 0x6e, 0x6f, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x73, 0x45, 0x6e, 0x74, 0x72, + 0x79, 0x12, 0x10, 0x0a, 0x03, 0x6b, 0x65, 0x79, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, + 0x6b, 0x65, 0x79, 0x12, 0x14, 0x0a, 0x05, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x18, 0x02, 0x20, 0x01, + 0x28, 0x09, 0x52, 0x05, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x3a, 0x02, 0x38, 0x01, 0x22, 0xba, 0x01, + 0x0a, 0x12, 0x4c, 0x61, 0x75, 0x6e, 0x63, 0x68, 0x50, 0x6c, 0x61, 0x6e, 0x54, 0x65, 0x6d, 0x70, + 0x6c, 0x61, 0x74, 0x65, 0x12, 0x29, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, + 0x32, 0x19, 0x2e, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2e, 0x63, 0x6f, 0x72, 0x65, + 0x2e, 0x49, 0x64, 0x65, 0x6e, 0x74, 0x69, 0x66, 0x69, 0x65, 0x72, 0x52, 0x02, 0x69, 0x64, 0x12, + 0x3b, 0x0a, 0x09, 0x69, 0x6e, 0x74, 0x65, 0x72, 0x66, 0x61, 0x63, 0x65, 0x18, 0x02, 0x20, 0x01, + 0x28, 0x0b, 0x32, 0x1d, 0x2e, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2e, 0x63, 0x6f, + 0x72, 0x65, 0x2e, 0x54, 0x79, 0x70, 0x65, 0x64, 0x49, 0x6e, 0x74, 0x65, 0x72, 0x66, 0x61, 0x63, + 0x65, 0x52, 0x09, 0x69, 0x6e, 0x74, 0x65, 0x72, 0x66, 0x61, 0x63, 0x65, 0x12, 0x3c, 0x0a, 0x0c, + 0x66, 0x69, 0x78, 0x65, 0x64, 0x5f, 0x69, 0x6e, 0x70, 0x75, 0x74, 0x73, 0x18, 0x03, 0x20, 0x01, + 0x28, 0x0b, 0x32, 0x19, 0x2e, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2e, 0x63, 0x6f, + 0x72, 0x65, 0x2e, 0x4c, 0x69, 0x74, 0x65, 0x72, 0x61, 0x6c, 0x4d, 0x61, 0x70, 0x52, 0x0b, 0x66, + 0x69, 0x78, 0x65, 0x64, 0x49, 0x6e, 0x70, 0x75, 0x74, 0x73, 0x42, 0xb3, 0x01, 0x0a, 0x11, 0x63, + 0x6f, 0x6d, 0x2e, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2e, 0x63, 0x6f, 0x72, 0x65, + 0x42, 0x0d, 0x57, 0x6f, 0x72, 0x6b, 0x66, 0x6c, 0x6f, 0x77, 0x50, 0x72, 0x6f, 0x74, 0x6f, 0x50, + 0x01, 0x5a, 0x3a, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x66, 0x6c, + 0x79, 0x74, 0x65, 0x6f, 0x72, 0x67, 0x2f, 0x66, 0x6c, 0x79, 0x74, 0x65, 0x2f, 0x66, 0x6c, 0x79, + 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2f, 0x67, 0x65, 0x6e, 0x2f, 0x70, 0x62, 0x2d, 0x67, 0x6f, 0x2f, + 0x66, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2f, 0x63, 0x6f, 0x72, 0x65, 0xa2, 0x02, 0x03, + 0x46, 0x43, 0x58, 0xaa, 0x02, 0x0d, 0x46, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x2e, 0x43, + 0x6f, 0x72, 0x65, 0xca, 0x02, 0x0d, 0x46, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x5c, 0x43, + 0x6f, 0x72, 0x65, 0xe2, 0x02, 0x19, 0x46, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x5c, 0x43, + 0x6f, 0x72, 0x65, 0x5c, 0x47, 0x50, 0x42, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0xea, + 0x02, 0x0e, 0x46, 0x6c, 0x79, 0x74, 0x65, 0x69, 0x64, 0x6c, 0x3a, 0x3a, 0x43, 0x6f, 0x72, 0x65, + 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, } var ( @@ -2078,7 +2113,7 @@ func file_flyteidl_core_workflow_proto_rawDescGZIP() []byte { } var file_flyteidl_core_workflow_proto_enumTypes = make([]protoimpl.EnumInfo, 3) -var file_flyteidl_core_workflow_proto_msgTypes = make([]protoimpl.MessageInfo, 19) +var file_flyteidl_core_workflow_proto_msgTypes = make([]protoimpl.MessageInfo, 21) var file_flyteidl_core_workflow_proto_goTypes = []interface{}{ (ArrayNode_ExecutionMode)(0), // 0: flyteidl.core.ArrayNode.ExecutionMode (ArrayNode_DataMode)(0), // 1: flyteidl.core.ArrayNode.DataMode @@ -2102,71 +2137,75 @@ var file_flyteidl_core_workflow_proto_goTypes = []interface{}{ (*TaskNodeOverrides)(nil), // 19: flyteidl.core.TaskNodeOverrides (*LaunchPlanTemplate)(nil), // 20: flyteidl.core.LaunchPlanTemplate nil, // 21: flyteidl.core.WorkflowMetadata.TagsEntry - (*BooleanExpression)(nil), // 22: flyteidl.core.BooleanExpression - (*Error)(nil), // 23: flyteidl.core.Error - (*Identifier)(nil), // 24: flyteidl.core.Identifier - (*LiteralType)(nil), // 25: flyteidl.core.LiteralType - (*durationpb.Duration)(nil), // 26: google.protobuf.Duration - (*wrapperspb.BoolValue)(nil), // 27: google.protobuf.BoolValue - (*RetryStrategy)(nil), // 28: flyteidl.core.RetryStrategy - (*Binding)(nil), // 29: flyteidl.core.Binding - (*QualityOfService)(nil), // 30: flyteidl.core.QualityOfService - (*TypedInterface)(nil), // 31: flyteidl.core.TypedInterface - (*Resources)(nil), // 32: flyteidl.core.Resources - (*ExtendedResources)(nil), // 33: flyteidl.core.ExtendedResources - (*LiteralMap)(nil), // 34: flyteidl.core.LiteralMap + nil, // 22: flyteidl.core.TaskNodeOverrides.LabelsEntry + nil, // 23: flyteidl.core.TaskNodeOverrides.AnnotationsEntry + (*BooleanExpression)(nil), // 24: flyteidl.core.BooleanExpression + (*Error)(nil), // 25: flyteidl.core.Error + (*Identifier)(nil), // 26: flyteidl.core.Identifier + (*LiteralType)(nil), // 27: flyteidl.core.LiteralType + (*durationpb.Duration)(nil), // 28: google.protobuf.Duration + (*wrapperspb.BoolValue)(nil), // 29: google.protobuf.BoolValue + (*RetryStrategy)(nil), // 30: flyteidl.core.RetryStrategy + (*Binding)(nil), // 31: flyteidl.core.Binding + (*QualityOfService)(nil), // 32: flyteidl.core.QualityOfService + (*TypedInterface)(nil), // 33: flyteidl.core.TypedInterface + (*Resources)(nil), // 34: flyteidl.core.Resources + (*ExtendedResources)(nil), // 35: flyteidl.core.ExtendedResources + (*LiteralMap)(nil), // 36: flyteidl.core.LiteralMap } var file_flyteidl_core_workflow_proto_depIdxs = []int32{ - 22, // 0: flyteidl.core.IfBlock.condition:type_name -> flyteidl.core.BooleanExpression + 24, // 0: flyteidl.core.IfBlock.condition:type_name -> flyteidl.core.BooleanExpression 15, // 1: flyteidl.core.IfBlock.then_node:type_name -> flyteidl.core.Node 3, // 2: flyteidl.core.IfElseBlock.case:type_name -> flyteidl.core.IfBlock 3, // 3: flyteidl.core.IfElseBlock.other:type_name -> flyteidl.core.IfBlock 15, // 4: flyteidl.core.IfElseBlock.else_node:type_name -> flyteidl.core.Node - 23, // 5: flyteidl.core.IfElseBlock.error:type_name -> flyteidl.core.Error + 25, // 5: flyteidl.core.IfElseBlock.error:type_name -> flyteidl.core.Error 4, // 6: flyteidl.core.BranchNode.if_else:type_name -> flyteidl.core.IfElseBlock - 24, // 7: flyteidl.core.TaskNode.reference_id:type_name -> flyteidl.core.Identifier + 26, // 7: flyteidl.core.TaskNode.reference_id:type_name -> flyteidl.core.Identifier 19, // 8: flyteidl.core.TaskNode.overrides:type_name -> flyteidl.core.TaskNodeOverrides - 24, // 9: flyteidl.core.WorkflowNode.launchplan_ref:type_name -> flyteidl.core.Identifier - 24, // 10: flyteidl.core.WorkflowNode.sub_workflow_ref:type_name -> flyteidl.core.Identifier - 25, // 11: flyteidl.core.SignalCondition.type:type_name -> flyteidl.core.LiteralType - 26, // 12: flyteidl.core.SleepCondition.duration:type_name -> google.protobuf.Duration + 26, // 9: flyteidl.core.WorkflowNode.launchplan_ref:type_name -> flyteidl.core.Identifier + 26, // 10: flyteidl.core.WorkflowNode.sub_workflow_ref:type_name -> flyteidl.core.Identifier + 27, // 11: flyteidl.core.SignalCondition.type:type_name -> flyteidl.core.LiteralType + 28, // 12: flyteidl.core.SleepCondition.duration:type_name -> google.protobuf.Duration 8, // 13: flyteidl.core.GateNode.approve:type_name -> flyteidl.core.ApproveCondition 9, // 14: flyteidl.core.GateNode.signal:type_name -> flyteidl.core.SignalCondition 10, // 15: flyteidl.core.GateNode.sleep:type_name -> flyteidl.core.SleepCondition 15, // 16: flyteidl.core.ArrayNode.node:type_name -> flyteidl.core.Node 0, // 17: flyteidl.core.ArrayNode.execution_mode:type_name -> flyteidl.core.ArrayNode.ExecutionMode - 27, // 18: flyteidl.core.ArrayNode.is_original_sub_node_interface:type_name -> google.protobuf.BoolValue + 29, // 18: flyteidl.core.ArrayNode.is_original_sub_node_interface:type_name -> google.protobuf.BoolValue 1, // 19: flyteidl.core.ArrayNode.data_mode:type_name -> flyteidl.core.ArrayNode.DataMode - 26, // 20: flyteidl.core.NodeMetadata.timeout:type_name -> google.protobuf.Duration - 28, // 21: flyteidl.core.NodeMetadata.retries:type_name -> flyteidl.core.RetryStrategy + 28, // 20: flyteidl.core.NodeMetadata.timeout:type_name -> google.protobuf.Duration + 30, // 21: flyteidl.core.NodeMetadata.retries:type_name -> flyteidl.core.RetryStrategy 13, // 22: flyteidl.core.Node.metadata:type_name -> flyteidl.core.NodeMetadata - 29, // 23: flyteidl.core.Node.inputs:type_name -> flyteidl.core.Binding + 31, // 23: flyteidl.core.Node.inputs:type_name -> flyteidl.core.Binding 14, // 24: flyteidl.core.Node.output_aliases:type_name -> flyteidl.core.Alias 6, // 25: flyteidl.core.Node.task_node:type_name -> flyteidl.core.TaskNode 7, // 26: flyteidl.core.Node.workflow_node:type_name -> flyteidl.core.WorkflowNode 5, // 27: flyteidl.core.Node.branch_node:type_name -> flyteidl.core.BranchNode 11, // 28: flyteidl.core.Node.gate_node:type_name -> flyteidl.core.GateNode 12, // 29: flyteidl.core.Node.array_node:type_name -> flyteidl.core.ArrayNode - 30, // 30: flyteidl.core.WorkflowMetadata.quality_of_service:type_name -> flyteidl.core.QualityOfService + 32, // 30: flyteidl.core.WorkflowMetadata.quality_of_service:type_name -> flyteidl.core.QualityOfService 2, // 31: flyteidl.core.WorkflowMetadata.on_failure:type_name -> flyteidl.core.WorkflowMetadata.OnFailurePolicy 21, // 32: flyteidl.core.WorkflowMetadata.tags:type_name -> flyteidl.core.WorkflowMetadata.TagsEntry - 24, // 33: flyteidl.core.WorkflowTemplate.id:type_name -> flyteidl.core.Identifier + 26, // 33: flyteidl.core.WorkflowTemplate.id:type_name -> flyteidl.core.Identifier 16, // 34: flyteidl.core.WorkflowTemplate.metadata:type_name -> flyteidl.core.WorkflowMetadata - 31, // 35: flyteidl.core.WorkflowTemplate.interface:type_name -> flyteidl.core.TypedInterface + 33, // 35: flyteidl.core.WorkflowTemplate.interface:type_name -> flyteidl.core.TypedInterface 15, // 36: flyteidl.core.WorkflowTemplate.nodes:type_name -> flyteidl.core.Node - 29, // 37: flyteidl.core.WorkflowTemplate.outputs:type_name -> flyteidl.core.Binding + 31, // 37: flyteidl.core.WorkflowTemplate.outputs:type_name -> flyteidl.core.Binding 15, // 38: flyteidl.core.WorkflowTemplate.failure_node:type_name -> flyteidl.core.Node 17, // 39: flyteidl.core.WorkflowTemplate.metadata_defaults:type_name -> flyteidl.core.WorkflowMetadataDefaults - 32, // 40: flyteidl.core.TaskNodeOverrides.resources:type_name -> flyteidl.core.Resources - 33, // 41: flyteidl.core.TaskNodeOverrides.extended_resources:type_name -> flyteidl.core.ExtendedResources - 24, // 42: flyteidl.core.LaunchPlanTemplate.id:type_name -> flyteidl.core.Identifier - 31, // 43: flyteidl.core.LaunchPlanTemplate.interface:type_name -> flyteidl.core.TypedInterface - 34, // 44: flyteidl.core.LaunchPlanTemplate.fixed_inputs:type_name -> flyteidl.core.LiteralMap - 45, // [45:45] is the sub-list for method output_type - 45, // [45:45] is the sub-list for method input_type - 45, // [45:45] is the sub-list for extension type_name - 45, // [45:45] is the sub-list for extension extendee - 0, // [0:45] is the sub-list for field type_name + 34, // 40: flyteidl.core.TaskNodeOverrides.resources:type_name -> flyteidl.core.Resources + 35, // 41: flyteidl.core.TaskNodeOverrides.extended_resources:type_name -> flyteidl.core.ExtendedResources + 22, // 42: flyteidl.core.TaskNodeOverrides.labels:type_name -> flyteidl.core.TaskNodeOverrides.LabelsEntry + 23, // 43: flyteidl.core.TaskNodeOverrides.annotations:type_name -> flyteidl.core.TaskNodeOverrides.AnnotationsEntry + 26, // 44: flyteidl.core.LaunchPlanTemplate.id:type_name -> flyteidl.core.Identifier + 33, // 45: flyteidl.core.LaunchPlanTemplate.interface:type_name -> flyteidl.core.TypedInterface + 36, // 46: flyteidl.core.LaunchPlanTemplate.fixed_inputs:type_name -> flyteidl.core.LiteralMap + 47, // [47:47] is the sub-list for method output_type + 47, // [47:47] is the sub-list for method input_type + 47, // [47:47] is the sub-list for extension type_name + 47, // [47:47] is the sub-list for extension extendee + 0, // [0:47] is the sub-list for field type_name } func init() { file_flyteidl_core_workflow_proto_init() } @@ -2440,7 +2479,7 @@ func file_flyteidl_core_workflow_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: file_flyteidl_core_workflow_proto_rawDesc, NumEnums: 3, - NumMessages: 19, + NumMessages: 21, NumExtensions: 0, NumServices: 0, }, diff --git a/flyteidl/gen/pb-go/gateway/flyteidl/service/admin.swagger.json b/flyteidl/gen/pb-go/gateway/flyteidl/service/admin.swagger.json index 45f471248b..c3f23ea125 100644 --- a/flyteidl/gen/pb-go/gateway/flyteidl/service/admin.swagger.json +++ b/flyteidl/gen/pb-go/gateway/flyteidl/service/admin.swagger.json @@ -8334,6 +8334,20 @@ "container_image": { "type": "string", "description": "Override for the image used by task pods." + }, + "labels": { + "type": "object", + "additionalProperties": { + "type": "string" + }, + "description": "Optional labels to add to the pod definition." + }, + "annotations": { + "type": "object", + "additionalProperties": { + "type": "string" + }, + "description": "Optional annotations to add to the pod definition." } }, "description": "Optional task node overrides that will be applied at task execution time." diff --git a/flyteidl/gen/pb-go/gateway/flyteidl/service/agent.swagger.json b/flyteidl/gen/pb-go/gateway/flyteidl/service/agent.swagger.json index 40fab85799..af6a87b29e 100644 --- a/flyteidl/gen/pb-go/gateway/flyteidl/service/agent.swagger.json +++ b/flyteidl/gen/pb-go/gateway/flyteidl/service/agent.swagger.json @@ -1881,6 +1881,20 @@ "container_image": { "type": "string", "description": "Override for the image used by task pods." + }, + "labels": { + "type": "object", + "additionalProperties": { + "type": "string" + }, + "description": "Optional labels to add to the pod definition." + }, + "annotations": { + "type": "object", + "additionalProperties": { + "type": "string" + }, + "description": "Optional annotations to add to the pod definition." } }, "description": "Optional task node overrides that will be applied at task execution time." diff --git a/flyteidl/gen/pb-js/flyteidl.d.ts b/flyteidl/gen/pb-js/flyteidl.d.ts index b51cfad923..20a4f4eb29 100644 --- a/flyteidl/gen/pb-js/flyteidl.d.ts +++ b/flyteidl/gen/pb-js/flyteidl.d.ts @@ -5175,6 +5175,12 @@ export namespace flyteidl { /** TaskNodeOverrides containerImage */ containerImage?: (string|null); + + /** TaskNodeOverrides labels */ + labels?: ({ [k: string]: string }|null); + + /** TaskNodeOverrides annotations */ + annotations?: ({ [k: string]: string }|null); } /** Represents a TaskNodeOverrides. */ @@ -5195,6 +5201,12 @@ export namespace flyteidl { /** TaskNodeOverrides containerImage. */ public containerImage: string; + /** TaskNodeOverrides labels. */ + public labels: { [k: string]: string }; + + /** TaskNodeOverrides annotations. */ + public annotations: { [k: string]: string }; + /** * Creates a new TaskNodeOverrides instance using the specified properties. * @param [properties] Properties to set diff --git a/flyteidl/gen/pb-js/flyteidl.js b/flyteidl/gen/pb-js/flyteidl.js index 9ff0d83c82..9bbf3309b1 100644 --- a/flyteidl/gen/pb-js/flyteidl.js +++ b/flyteidl/gen/pb-js/flyteidl.js @@ -12561,6 +12561,8 @@ * @property {flyteidl.core.IResources|null} [resources] TaskNodeOverrides resources * @property {flyteidl.core.IExtendedResources|null} [extendedResources] TaskNodeOverrides extendedResources * @property {string|null} [containerImage] TaskNodeOverrides containerImage + * @property {Object.|null} [labels] TaskNodeOverrides labels + * @property {Object.|null} [annotations] TaskNodeOverrides annotations */ /** @@ -12572,6 +12574,8 @@ * @param {flyteidl.core.ITaskNodeOverrides=} [properties] Properties to set */ function TaskNodeOverrides(properties) { + this.labels = {}; + this.annotations = {}; if (properties) for (var keys = Object.keys(properties), i = 0; i < keys.length; ++i) if (properties[keys[i]] != null) @@ -12602,6 +12606,22 @@ */ TaskNodeOverrides.prototype.containerImage = ""; + /** + * TaskNodeOverrides labels. + * @member {Object.} labels + * @memberof flyteidl.core.TaskNodeOverrides + * @instance + */ + TaskNodeOverrides.prototype.labels = $util.emptyObject; + + /** + * TaskNodeOverrides annotations. + * @member {Object.} annotations + * @memberof flyteidl.core.TaskNodeOverrides + * @instance + */ + TaskNodeOverrides.prototype.annotations = $util.emptyObject; + /** * Creates a new TaskNodeOverrides instance using the specified properties. * @function create @@ -12632,6 +12652,12 @@ $root.flyteidl.core.ExtendedResources.encode(message.extendedResources, writer.uint32(/* id 2, wireType 2 =*/18).fork()).ldelim(); if (message.containerImage != null && message.hasOwnProperty("containerImage")) writer.uint32(/* id 3, wireType 2 =*/26).string(message.containerImage); + if (message.labels != null && message.hasOwnProperty("labels")) + for (var keys = Object.keys(message.labels), i = 0; i < keys.length; ++i) + writer.uint32(/* id 4, wireType 2 =*/34).fork().uint32(/* id 1, wireType 2 =*/10).string(keys[i]).uint32(/* id 2, wireType 2 =*/18).string(message.labels[keys[i]]).ldelim(); + if (message.annotations != null && message.hasOwnProperty("annotations")) + for (var keys = Object.keys(message.annotations), i = 0; i < keys.length; ++i) + writer.uint32(/* id 5, wireType 2 =*/42).fork().uint32(/* id 1, wireType 2 =*/10).string(keys[i]).uint32(/* id 2, wireType 2 =*/18).string(message.annotations[keys[i]]).ldelim(); return writer; }; @@ -12649,7 +12675,7 @@ TaskNodeOverrides.decode = function decode(reader, length) { if (!(reader instanceof $Reader)) reader = $Reader.create(reader); - var end = length === undefined ? reader.len : reader.pos + length, message = new $root.flyteidl.core.TaskNodeOverrides(); + var end = length === undefined ? reader.len : reader.pos + length, message = new $root.flyteidl.core.TaskNodeOverrides(), key; while (reader.pos < end) { var tag = reader.uint32(); switch (tag >>> 3) { @@ -12662,6 +12688,22 @@ case 3: message.containerImage = reader.string(); break; + case 4: + reader.skip().pos++; + if (message.labels === $util.emptyObject) + message.labels = {}; + key = reader.string(); + reader.pos++; + message.labels[key] = reader.string(); + break; + case 5: + reader.skip().pos++; + if (message.annotations === $util.emptyObject) + message.annotations = {}; + key = reader.string(); + reader.pos++; + message.annotations[key] = reader.string(); + break; default: reader.skipType(tag & 7); break; @@ -12694,6 +12736,22 @@ if (message.containerImage != null && message.hasOwnProperty("containerImage")) if (!$util.isString(message.containerImage)) return "containerImage: string expected"; + if (message.labels != null && message.hasOwnProperty("labels")) { + if (!$util.isObject(message.labels)) + return "labels: object expected"; + var key = Object.keys(message.labels); + for (var i = 0; i < key.length; ++i) + if (!$util.isString(message.labels[key[i]])) + return "labels: string{k:string} expected"; + } + if (message.annotations != null && message.hasOwnProperty("annotations")) { + if (!$util.isObject(message.annotations)) + return "annotations: object expected"; + var key = Object.keys(message.annotations); + for (var i = 0; i < key.length; ++i) + if (!$util.isString(message.annotations[key[i]])) + return "annotations: string{k:string} expected"; + } return null; }; diff --git a/flyteidl/gen/pb_python/flyteidl/core/workflow_pb2.py b/flyteidl/gen/pb_python/flyteidl/core/workflow_pb2.py index 4f30f4d632..cb06b95a27 100644 --- a/flyteidl/gen/pb_python/flyteidl/core/workflow_pb2.py +++ b/flyteidl/gen/pb_python/flyteidl/core/workflow_pb2.py @@ -23,7 +23,7 @@ from google.protobuf import wrappers_pb2 as google_dot_protobuf_dot_wrappers__pb2 -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1c\x66lyteidl/core/workflow.proto\x12\rflyteidl.core\x1a\x1d\x66lyteidl/core/condition.proto\x1a\x1d\x66lyteidl/core/execution.proto\x1a\x1e\x66lyteidl/core/identifier.proto\x1a\x1d\x66lyteidl/core/interface.proto\x1a\x1c\x66lyteidl/core/literals.proto\x1a\x19\x66lyteidl/core/tasks.proto\x1a\x19\x66lyteidl/core/types.proto\x1a\x1c\x66lyteidl/core/security.proto\x1a\x1egoogle/protobuf/duration.proto\x1a\x1egoogle/protobuf/wrappers.proto\"{\n\x07IfBlock\x12>\n\tcondition\x18\x01 \x01(\x0b\x32 .flyteidl.core.BooleanExpressionR\tcondition\x12\x30\n\tthen_node\x18\x02 \x01(\x0b\x32\x13.flyteidl.core.NodeR\x08thenNode\"\xd4\x01\n\x0bIfElseBlock\x12*\n\x04\x63\x61se\x18\x01 \x01(\x0b\x32\x16.flyteidl.core.IfBlockR\x04\x63\x61se\x12,\n\x05other\x18\x02 \x03(\x0b\x32\x16.flyteidl.core.IfBlockR\x05other\x12\x32\n\telse_node\x18\x03 \x01(\x0b\x32\x13.flyteidl.core.NodeH\x00R\x08\x65lseNode\x12,\n\x05\x65rror\x18\x04 \x01(\x0b\x32\x14.flyteidl.core.ErrorH\x00R\x05\x65rrorB\t\n\x07\x64\x65\x66\x61ult\"A\n\nBranchNode\x12\x33\n\x07if_else\x18\x01 \x01(\x0b\x32\x1a.flyteidl.core.IfElseBlockR\x06ifElse\"\x97\x01\n\x08TaskNode\x12>\n\x0creference_id\x18\x01 \x01(\x0b\x32\x19.flyteidl.core.IdentifierH\x00R\x0breferenceId\x12>\n\toverrides\x18\x02 \x01(\x0b\x32 .flyteidl.core.TaskNodeOverridesR\toverridesB\x0b\n\treference\"\xa6\x01\n\x0cWorkflowNode\x12\x42\n\x0elaunchplan_ref\x18\x01 \x01(\x0b\x32\x19.flyteidl.core.IdentifierH\x00R\rlaunchplanRef\x12\x45\n\x10sub_workflow_ref\x18\x02 \x01(\x0b\x32\x19.flyteidl.core.IdentifierH\x00R\x0esubWorkflowRefB\x0b\n\treference\"/\n\x10\x41pproveCondition\x12\x1b\n\tsignal_id\x18\x01 \x01(\tR\x08signalId\"\x90\x01\n\x0fSignalCondition\x12\x1b\n\tsignal_id\x18\x01 \x01(\tR\x08signalId\x12.\n\x04type\x18\x02 \x01(\x0b\x32\x1a.flyteidl.core.LiteralTypeR\x04type\x12\x30\n\x14output_variable_name\x18\x03 \x01(\tR\x12outputVariableName\"G\n\x0eSleepCondition\x12\x35\n\x08\x64uration\x18\x01 \x01(\x0b\x32\x19.google.protobuf.DurationR\x08\x64uration\"\xc5\x01\n\x08GateNode\x12;\n\x07\x61pprove\x18\x01 \x01(\x0b\x32\x1f.flyteidl.core.ApproveConditionH\x00R\x07\x61pprove\x12\x38\n\x06signal\x18\x02 \x01(\x0b\x32\x1e.flyteidl.core.SignalConditionH\x00R\x06signal\x12\x35\n\x05sleep\x18\x03 \x01(\x0b\x32\x1d.flyteidl.core.SleepConditionH\x00R\x05sleepB\x0b\n\tcondition\"\xb9\x04\n\tArrayNode\x12\'\n\x04node\x18\x01 \x01(\x0b\x32\x13.flyteidl.core.NodeR\x04node\x12\"\n\x0bparallelism\x18\x02 \x01(\rH\x00R\x0bparallelism\x12%\n\rmin_successes\x18\x03 \x01(\rH\x01R\x0cminSuccesses\x12,\n\x11min_success_ratio\x18\x04 \x01(\x02H\x01R\x0fminSuccessRatio\x12M\n\x0e\x65xecution_mode\x18\x05 \x01(\x0e\x32&.flyteidl.core.ArrayNode.ExecutionModeR\rexecutionMode\x12^\n\x1eis_original_sub_node_interface\x18\x06 \x01(\x0b\x32\x1a.google.protobuf.BoolValueR\x1aisOriginalSubNodeInterface\x12>\n\tdata_mode\x18\x07 \x01(\x0e\x32!.flyteidl.core.ArrayNode.DataModeR\x08\x64\x61taMode\"2\n\rExecutionMode\x12\x11\n\rMINIMAL_STATE\x10\x00\x12\x0e\n\nFULL_STATE\x10\x01\"=\n\x08\x44\x61taMode\x12\x15\n\x11SINGLE_INPUT_FILE\x10\x00\x12\x1a\n\x16INDIVIDUAL_INPUT_FILES\x10\x01\x42\x14\n\x12parallelism_optionB\x12\n\x10success_criteria\"\x8c\x03\n\x0cNodeMetadata\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12\x33\n\x07timeout\x18\x04 \x01(\x0b\x32\x19.google.protobuf.DurationR\x07timeout\x12\x36\n\x07retries\x18\x05 \x01(\x0b\x32\x1c.flyteidl.core.RetryStrategyR\x07retries\x12&\n\rinterruptible\x18\x06 \x01(\x08H\x00R\rinterruptible\x12\x1e\n\tcacheable\x18\x07 \x01(\x08H\x01R\tcacheable\x12%\n\rcache_version\x18\x08 \x01(\tH\x02R\x0c\x63\x61\x63heVersion\x12/\n\x12\x63\x61\x63he_serializable\x18\t \x01(\x08H\x03R\x11\x63\x61\x63heSerializableB\x15\n\x13interruptible_valueB\x11\n\x0f\x63\x61\x63heable_valueB\x15\n\x13\x63\x61\x63he_version_valueB\x1a\n\x18\x63\x61\x63he_serializable_value\"/\n\x05\x41lias\x12\x10\n\x03var\x18\x01 \x01(\tR\x03var\x12\x14\n\x05\x61lias\x18\x02 \x01(\tR\x05\x61lias\"\x9f\x04\n\x04Node\x12\x0e\n\x02id\x18\x01 \x01(\tR\x02id\x12\x37\n\x08metadata\x18\x02 \x01(\x0b\x32\x1b.flyteidl.core.NodeMetadataR\x08metadata\x12.\n\x06inputs\x18\x03 \x03(\x0b\x32\x16.flyteidl.core.BindingR\x06inputs\x12*\n\x11upstream_node_ids\x18\x04 \x03(\tR\x0fupstreamNodeIds\x12;\n\x0eoutput_aliases\x18\x05 \x03(\x0b\x32\x14.flyteidl.core.AliasR\routputAliases\x12\x36\n\ttask_node\x18\x06 \x01(\x0b\x32\x17.flyteidl.core.TaskNodeH\x00R\x08taskNode\x12\x42\n\rworkflow_node\x18\x07 \x01(\x0b\x32\x1b.flyteidl.core.WorkflowNodeH\x00R\x0cworkflowNode\x12<\n\x0b\x62ranch_node\x18\x08 \x01(\x0b\x32\x19.flyteidl.core.BranchNodeH\x00R\nbranchNode\x12\x36\n\tgate_node\x18\t \x01(\x0b\x32\x17.flyteidl.core.GateNodeH\x00R\x08gateNode\x12\x39\n\narray_node\x18\n \x01(\x0b\x32\x18.flyteidl.core.ArrayNodeH\x00R\tarrayNodeB\x08\n\x06target\"\xfc\x02\n\x10WorkflowMetadata\x12M\n\x12quality_of_service\x18\x01 \x01(\x0b\x32\x1f.flyteidl.core.QualityOfServiceR\x10qualityOfService\x12N\n\non_failure\x18\x02 \x01(\x0e\x32/.flyteidl.core.WorkflowMetadata.OnFailurePolicyR\tonFailure\x12=\n\x04tags\x18\x03 \x03(\x0b\x32).flyteidl.core.WorkflowMetadata.TagsEntryR\x04tags\x1a\x37\n\tTagsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01\"Q\n\x0fOnFailurePolicy\x12\x14\n\x10\x46\x41IL_IMMEDIATELY\x10\x00\x12(\n$FAIL_AFTER_EXECUTABLE_NODES_COMPLETE\x10\x01\"@\n\x18WorkflowMetadataDefaults\x12$\n\rinterruptible\x18\x01 \x01(\x08R\rinterruptible\"\xa2\x03\n\x10WorkflowTemplate\x12)\n\x02id\x18\x01 \x01(\x0b\x32\x19.flyteidl.core.IdentifierR\x02id\x12;\n\x08metadata\x18\x02 \x01(\x0b\x32\x1f.flyteidl.core.WorkflowMetadataR\x08metadata\x12;\n\tinterface\x18\x03 \x01(\x0b\x32\x1d.flyteidl.core.TypedInterfaceR\tinterface\x12)\n\x05nodes\x18\x04 \x03(\x0b\x32\x13.flyteidl.core.NodeR\x05nodes\x12\x30\n\x07outputs\x18\x05 \x03(\x0b\x32\x16.flyteidl.core.BindingR\x07outputs\x12\x36\n\x0c\x66\x61ilure_node\x18\x06 \x01(\x0b\x32\x13.flyteidl.core.NodeR\x0b\x66\x61ilureNode\x12T\n\x11metadata_defaults\x18\x07 \x01(\x0b\x32\'.flyteidl.core.WorkflowMetadataDefaultsR\x10metadataDefaults\"\xc5\x01\n\x11TaskNodeOverrides\x12\x36\n\tresources\x18\x01 \x01(\x0b\x32\x18.flyteidl.core.ResourcesR\tresources\x12O\n\x12\x65xtended_resources\x18\x02 \x01(\x0b\x32 .flyteidl.core.ExtendedResourcesR\x11\x65xtendedResources\x12\'\n\x0f\x63ontainer_image\x18\x03 \x01(\tR\x0e\x63ontainerImage\"\xba\x01\n\x12LaunchPlanTemplate\x12)\n\x02id\x18\x01 \x01(\x0b\x32\x19.flyteidl.core.IdentifierR\x02id\x12;\n\tinterface\x18\x02 \x01(\x0b\x32\x1d.flyteidl.core.TypedInterfaceR\tinterface\x12<\n\x0c\x66ixed_inputs\x18\x03 \x01(\x0b\x32\x19.flyteidl.core.LiteralMapR\x0b\x66ixedInputsB\xb3\x01\n\x11\x63om.flyteidl.coreB\rWorkflowProtoP\x01Z:github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core\xa2\x02\x03\x46\x43X\xaa\x02\rFlyteidl.Core\xca\x02\rFlyteidl\\Core\xe2\x02\x19\x46lyteidl\\Core\\GPBMetadata\xea\x02\x0e\x46lyteidl::Coreb\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1c\x66lyteidl/core/workflow.proto\x12\rflyteidl.core\x1a\x1d\x66lyteidl/core/condition.proto\x1a\x1d\x66lyteidl/core/execution.proto\x1a\x1e\x66lyteidl/core/identifier.proto\x1a\x1d\x66lyteidl/core/interface.proto\x1a\x1c\x66lyteidl/core/literals.proto\x1a\x19\x66lyteidl/core/tasks.proto\x1a\x19\x66lyteidl/core/types.proto\x1a\x1c\x66lyteidl/core/security.proto\x1a\x1egoogle/protobuf/duration.proto\x1a\x1egoogle/protobuf/wrappers.proto\"{\n\x07IfBlock\x12>\n\tcondition\x18\x01 \x01(\x0b\x32 .flyteidl.core.BooleanExpressionR\tcondition\x12\x30\n\tthen_node\x18\x02 \x01(\x0b\x32\x13.flyteidl.core.NodeR\x08thenNode\"\xd4\x01\n\x0bIfElseBlock\x12*\n\x04\x63\x61se\x18\x01 \x01(\x0b\x32\x16.flyteidl.core.IfBlockR\x04\x63\x61se\x12,\n\x05other\x18\x02 \x03(\x0b\x32\x16.flyteidl.core.IfBlockR\x05other\x12\x32\n\telse_node\x18\x03 \x01(\x0b\x32\x13.flyteidl.core.NodeH\x00R\x08\x65lseNode\x12,\n\x05\x65rror\x18\x04 \x01(\x0b\x32\x14.flyteidl.core.ErrorH\x00R\x05\x65rrorB\t\n\x07\x64\x65\x66\x61ult\"A\n\nBranchNode\x12\x33\n\x07if_else\x18\x01 \x01(\x0b\x32\x1a.flyteidl.core.IfElseBlockR\x06ifElse\"\x97\x01\n\x08TaskNode\x12>\n\x0creference_id\x18\x01 \x01(\x0b\x32\x19.flyteidl.core.IdentifierH\x00R\x0breferenceId\x12>\n\toverrides\x18\x02 \x01(\x0b\x32 .flyteidl.core.TaskNodeOverridesR\toverridesB\x0b\n\treference\"\xa6\x01\n\x0cWorkflowNode\x12\x42\n\x0elaunchplan_ref\x18\x01 \x01(\x0b\x32\x19.flyteidl.core.IdentifierH\x00R\rlaunchplanRef\x12\x45\n\x10sub_workflow_ref\x18\x02 \x01(\x0b\x32\x19.flyteidl.core.IdentifierH\x00R\x0esubWorkflowRefB\x0b\n\treference\"/\n\x10\x41pproveCondition\x12\x1b\n\tsignal_id\x18\x01 \x01(\tR\x08signalId\"\x90\x01\n\x0fSignalCondition\x12\x1b\n\tsignal_id\x18\x01 \x01(\tR\x08signalId\x12.\n\x04type\x18\x02 \x01(\x0b\x32\x1a.flyteidl.core.LiteralTypeR\x04type\x12\x30\n\x14output_variable_name\x18\x03 \x01(\tR\x12outputVariableName\"G\n\x0eSleepCondition\x12\x35\n\x08\x64uration\x18\x01 \x01(\x0b\x32\x19.google.protobuf.DurationR\x08\x64uration\"\xc5\x01\n\x08GateNode\x12;\n\x07\x61pprove\x18\x01 \x01(\x0b\x32\x1f.flyteidl.core.ApproveConditionH\x00R\x07\x61pprove\x12\x38\n\x06signal\x18\x02 \x01(\x0b\x32\x1e.flyteidl.core.SignalConditionH\x00R\x06signal\x12\x35\n\x05sleep\x18\x03 \x01(\x0b\x32\x1d.flyteidl.core.SleepConditionH\x00R\x05sleepB\x0b\n\tcondition\"\xb9\x04\n\tArrayNode\x12\'\n\x04node\x18\x01 \x01(\x0b\x32\x13.flyteidl.core.NodeR\x04node\x12\"\n\x0bparallelism\x18\x02 \x01(\rH\x00R\x0bparallelism\x12%\n\rmin_successes\x18\x03 \x01(\rH\x01R\x0cminSuccesses\x12,\n\x11min_success_ratio\x18\x04 \x01(\x02H\x01R\x0fminSuccessRatio\x12M\n\x0e\x65xecution_mode\x18\x05 \x01(\x0e\x32&.flyteidl.core.ArrayNode.ExecutionModeR\rexecutionMode\x12^\n\x1eis_original_sub_node_interface\x18\x06 \x01(\x0b\x32\x1a.google.protobuf.BoolValueR\x1aisOriginalSubNodeInterface\x12>\n\tdata_mode\x18\x07 \x01(\x0e\x32!.flyteidl.core.ArrayNode.DataModeR\x08\x64\x61taMode\"2\n\rExecutionMode\x12\x11\n\rMINIMAL_STATE\x10\x00\x12\x0e\n\nFULL_STATE\x10\x01\"=\n\x08\x44\x61taMode\x12\x15\n\x11SINGLE_INPUT_FILE\x10\x00\x12\x1a\n\x16INDIVIDUAL_INPUT_FILES\x10\x01\x42\x14\n\x12parallelism_optionB\x12\n\x10success_criteria\"\x8c\x03\n\x0cNodeMetadata\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12\x33\n\x07timeout\x18\x04 \x01(\x0b\x32\x19.google.protobuf.DurationR\x07timeout\x12\x36\n\x07retries\x18\x05 \x01(\x0b\x32\x1c.flyteidl.core.RetryStrategyR\x07retries\x12&\n\rinterruptible\x18\x06 \x01(\x08H\x00R\rinterruptible\x12\x1e\n\tcacheable\x18\x07 \x01(\x08H\x01R\tcacheable\x12%\n\rcache_version\x18\x08 \x01(\tH\x02R\x0c\x63\x61\x63heVersion\x12/\n\x12\x63\x61\x63he_serializable\x18\t \x01(\x08H\x03R\x11\x63\x61\x63heSerializableB\x15\n\x13interruptible_valueB\x11\n\x0f\x63\x61\x63heable_valueB\x15\n\x13\x63\x61\x63he_version_valueB\x1a\n\x18\x63\x61\x63he_serializable_value\"/\n\x05\x41lias\x12\x10\n\x03var\x18\x01 \x01(\tR\x03var\x12\x14\n\x05\x61lias\x18\x02 \x01(\tR\x05\x61lias\"\x9f\x04\n\x04Node\x12\x0e\n\x02id\x18\x01 \x01(\tR\x02id\x12\x37\n\x08metadata\x18\x02 \x01(\x0b\x32\x1b.flyteidl.core.NodeMetadataR\x08metadata\x12.\n\x06inputs\x18\x03 \x03(\x0b\x32\x16.flyteidl.core.BindingR\x06inputs\x12*\n\x11upstream_node_ids\x18\x04 \x03(\tR\x0fupstreamNodeIds\x12;\n\x0eoutput_aliases\x18\x05 \x03(\x0b\x32\x14.flyteidl.core.AliasR\routputAliases\x12\x36\n\ttask_node\x18\x06 \x01(\x0b\x32\x17.flyteidl.core.TaskNodeH\x00R\x08taskNode\x12\x42\n\rworkflow_node\x18\x07 \x01(\x0b\x32\x1b.flyteidl.core.WorkflowNodeH\x00R\x0cworkflowNode\x12<\n\x0b\x62ranch_node\x18\x08 \x01(\x0b\x32\x19.flyteidl.core.BranchNodeH\x00R\nbranchNode\x12\x36\n\tgate_node\x18\t \x01(\x0b\x32\x17.flyteidl.core.GateNodeH\x00R\x08gateNode\x12\x39\n\narray_node\x18\n \x01(\x0b\x32\x18.flyteidl.core.ArrayNodeH\x00R\tarrayNodeB\x08\n\x06target\"\xfc\x02\n\x10WorkflowMetadata\x12M\n\x12quality_of_service\x18\x01 \x01(\x0b\x32\x1f.flyteidl.core.QualityOfServiceR\x10qualityOfService\x12N\n\non_failure\x18\x02 \x01(\x0e\x32/.flyteidl.core.WorkflowMetadata.OnFailurePolicyR\tonFailure\x12=\n\x04tags\x18\x03 \x03(\x0b\x32).flyteidl.core.WorkflowMetadata.TagsEntryR\x04tags\x1a\x37\n\tTagsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01\"Q\n\x0fOnFailurePolicy\x12\x14\n\x10\x46\x41IL_IMMEDIATELY\x10\x00\x12(\n$FAIL_AFTER_EXECUTABLE_NODES_COMPLETE\x10\x01\"@\n\x18WorkflowMetadataDefaults\x12$\n\rinterruptible\x18\x01 \x01(\x08R\rinterruptible\"\xa2\x03\n\x10WorkflowTemplate\x12)\n\x02id\x18\x01 \x01(\x0b\x32\x19.flyteidl.core.IdentifierR\x02id\x12;\n\x08metadata\x18\x02 \x01(\x0b\x32\x1f.flyteidl.core.WorkflowMetadataR\x08metadata\x12;\n\tinterface\x18\x03 \x01(\x0b\x32\x1d.flyteidl.core.TypedInterfaceR\tinterface\x12)\n\x05nodes\x18\x04 \x03(\x0b\x32\x13.flyteidl.core.NodeR\x05nodes\x12\x30\n\x07outputs\x18\x05 \x03(\x0b\x32\x16.flyteidl.core.BindingR\x07outputs\x12\x36\n\x0c\x66\x61ilure_node\x18\x06 \x01(\x0b\x32\x13.flyteidl.core.NodeR\x0b\x66\x61ilureNode\x12T\n\x11metadata_defaults\x18\x07 \x01(\x0b\x32\'.flyteidl.core.WorkflowMetadataDefaultsR\x10metadataDefaults\"\xdb\x03\n\x11TaskNodeOverrides\x12\x36\n\tresources\x18\x01 \x01(\x0b\x32\x18.flyteidl.core.ResourcesR\tresources\x12O\n\x12\x65xtended_resources\x18\x02 \x01(\x0b\x32 .flyteidl.core.ExtendedResourcesR\x11\x65xtendedResources\x12\'\n\x0f\x63ontainer_image\x18\x03 \x01(\tR\x0e\x63ontainerImage\x12\x44\n\x06labels\x18\x04 \x03(\x0b\x32,.flyteidl.core.TaskNodeOverrides.LabelsEntryR\x06labels\x12S\n\x0b\x61nnotations\x18\x05 \x03(\x0b\x32\x31.flyteidl.core.TaskNodeOverrides.AnnotationsEntryR\x0b\x61nnotations\x1a\x39\n\x0bLabelsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01\x1a>\n\x10\x41nnotationsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01\"\xba\x01\n\x12LaunchPlanTemplate\x12)\n\x02id\x18\x01 \x01(\x0b\x32\x19.flyteidl.core.IdentifierR\x02id\x12;\n\tinterface\x18\x02 \x01(\x0b\x32\x1d.flyteidl.core.TypedInterfaceR\tinterface\x12<\n\x0c\x66ixed_inputs\x18\x03 \x01(\x0b\x32\x19.flyteidl.core.LiteralMapR\x0b\x66ixedInputsB\xb3\x01\n\x11\x63om.flyteidl.coreB\rWorkflowProtoP\x01Z:github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core\xa2\x02\x03\x46\x43X\xaa\x02\rFlyteidl.Core\xca\x02\rFlyteidl\\Core\xe2\x02\x19\x46lyteidl\\Core\\GPBMetadata\xea\x02\x0e\x46lyteidl::Coreb\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) @@ -34,6 +34,10 @@ DESCRIPTOR._serialized_options = b'\n\021com.flyteidl.coreB\rWorkflowProtoP\001Z:github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core\242\002\003FCX\252\002\rFlyteidl.Core\312\002\rFlyteidl\\Core\342\002\031Flyteidl\\Core\\GPBMetadata\352\002\016Flyteidl::Core' _WORKFLOWMETADATA_TAGSENTRY._options = None _WORKFLOWMETADATA_TAGSENTRY._serialized_options = b'8\001' + _TASKNODEOVERRIDES_LABELSENTRY._options = None + _TASKNODEOVERRIDES_LABELSENTRY._serialized_options = b'8\001' + _TASKNODEOVERRIDES_ANNOTATIONSENTRY._options = None + _TASKNODEOVERRIDES_ANNOTATIONSENTRY._serialized_options = b'8\001' _globals['_IFBLOCK']._serialized_start=350 _globals['_IFBLOCK']._serialized_end=473 _globals['_IFELSEBLOCK']._serialized_start=476 @@ -75,7 +79,11 @@ _globals['_WORKFLOWTEMPLATE']._serialized_start=3565 _globals['_WORKFLOWTEMPLATE']._serialized_end=3983 _globals['_TASKNODEOVERRIDES']._serialized_start=3986 - _globals['_TASKNODEOVERRIDES']._serialized_end=4183 - _globals['_LAUNCHPLANTEMPLATE']._serialized_start=4186 - _globals['_LAUNCHPLANTEMPLATE']._serialized_end=4372 + _globals['_TASKNODEOVERRIDES']._serialized_end=4461 + _globals['_TASKNODEOVERRIDES_LABELSENTRY']._serialized_start=4340 + _globals['_TASKNODEOVERRIDES_LABELSENTRY']._serialized_end=4397 + _globals['_TASKNODEOVERRIDES_ANNOTATIONSENTRY']._serialized_start=4399 + _globals['_TASKNODEOVERRIDES_ANNOTATIONSENTRY']._serialized_end=4461 + _globals['_LAUNCHPLANTEMPLATE']._serialized_start=4464 + _globals['_LAUNCHPLANTEMPLATE']._serialized_end=4650 # @@protoc_insertion_point(module_scope) diff --git a/flyteidl/gen/pb_python/flyteidl/core/workflow_pb2.pyi b/flyteidl/gen/pb_python/flyteidl/core/workflow_pb2.pyi index bc93b230b3..e3e81a81ac 100644 --- a/flyteidl/gen/pb_python/flyteidl/core/workflow_pb2.pyi +++ b/flyteidl/gen/pb_python/flyteidl/core/workflow_pb2.pyi @@ -218,14 +218,32 @@ class WorkflowTemplate(_message.Message): def __init__(self, id: _Optional[_Union[_identifier_pb2.Identifier, _Mapping]] = ..., metadata: _Optional[_Union[WorkflowMetadata, _Mapping]] = ..., interface: _Optional[_Union[_interface_pb2.TypedInterface, _Mapping]] = ..., nodes: _Optional[_Iterable[_Union[Node, _Mapping]]] = ..., outputs: _Optional[_Iterable[_Union[_literals_pb2.Binding, _Mapping]]] = ..., failure_node: _Optional[_Union[Node, _Mapping]] = ..., metadata_defaults: _Optional[_Union[WorkflowMetadataDefaults, _Mapping]] = ...) -> None: ... class TaskNodeOverrides(_message.Message): - __slots__ = ["resources", "extended_resources", "container_image"] + __slots__ = ["resources", "extended_resources", "container_image", "labels", "annotations"] + class LabelsEntry(_message.Message): + __slots__ = ["key", "value"] + KEY_FIELD_NUMBER: _ClassVar[int] + VALUE_FIELD_NUMBER: _ClassVar[int] + key: str + value: str + def __init__(self, key: _Optional[str] = ..., value: _Optional[str] = ...) -> None: ... + class AnnotationsEntry(_message.Message): + __slots__ = ["key", "value"] + KEY_FIELD_NUMBER: _ClassVar[int] + VALUE_FIELD_NUMBER: _ClassVar[int] + key: str + value: str + def __init__(self, key: _Optional[str] = ..., value: _Optional[str] = ...) -> None: ... RESOURCES_FIELD_NUMBER: _ClassVar[int] EXTENDED_RESOURCES_FIELD_NUMBER: _ClassVar[int] CONTAINER_IMAGE_FIELD_NUMBER: _ClassVar[int] + LABELS_FIELD_NUMBER: _ClassVar[int] + ANNOTATIONS_FIELD_NUMBER: _ClassVar[int] resources: _tasks_pb2.Resources extended_resources: _tasks_pb2.ExtendedResources container_image: str - def __init__(self, resources: _Optional[_Union[_tasks_pb2.Resources, _Mapping]] = ..., extended_resources: _Optional[_Union[_tasks_pb2.ExtendedResources, _Mapping]] = ..., container_image: _Optional[str] = ...) -> None: ... + labels: _containers.ScalarMap[str, str] + annotations: _containers.ScalarMap[str, str] + def __init__(self, resources: _Optional[_Union[_tasks_pb2.Resources, _Mapping]] = ..., extended_resources: _Optional[_Union[_tasks_pb2.ExtendedResources, _Mapping]] = ..., container_image: _Optional[str] = ..., labels: _Optional[_Mapping[str, str]] = ..., annotations: _Optional[_Mapping[str, str]] = ...) -> None: ... class LaunchPlanTemplate(_message.Message): __slots__ = ["id", "interface", "fixed_inputs"] diff --git a/flyteidl/gen/pb_rust/flyteidl.core.rs b/flyteidl/gen/pb_rust/flyteidl.core.rs index 135c5d1f65..2e437a77b8 100644 --- a/flyteidl/gen/pb_rust/flyteidl.core.rs +++ b/flyteidl/gen/pb_rust/flyteidl.core.rs @@ -2781,6 +2781,12 @@ pub struct TaskNodeOverrides { /// Override for the image used by task pods. #[prost(string, tag="3")] pub container_image: ::prost::alloc::string::String, + /// Optional labels to add to the pod definition. + #[prost(map="string, string", tag="4")] + pub labels: ::std::collections::HashMap<::prost::alloc::string::String, ::prost::alloc::string::String>, + /// Optional annotations to add to the pod definition. + #[prost(map="string, string", tag="5")] + pub annotations: ::std::collections::HashMap<::prost::alloc::string::String, ::prost::alloc::string::String>, } /// A structure that uniquely identifies a launch plan in the system. #[allow(clippy::derive_partial_eq_without_eq)] diff --git a/flyteidl/protos/flyteidl/core/workflow.proto b/flyteidl/protos/flyteidl/core/workflow.proto index 5e95531976..1d06ab7df0 100644 --- a/flyteidl/protos/flyteidl/core/workflow.proto +++ b/flyteidl/protos/flyteidl/core/workflow.proto @@ -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 labels = 4; + + // Optional annotations to add to the pod definition. + map annotations = 5; } // A structure that uniquely identifies a launch plan in the system. diff --git a/flyteplugins/go/tasks/pluginmachinery/core/exec_metadata.go b/flyteplugins/go/tasks/pluginmachinery/core/exec_metadata.go index 9a020bd188..5b0c642479 100644 --- a/flyteplugins/go/tasks/pluginmachinery/core/exec_metadata.go +++ b/flyteplugins/go/tasks/pluginmachinery/core/exec_metadata.go @@ -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 diff --git a/flyteplugins/go/tasks/pluginmachinery/core/mocks/task_overrides.go b/flyteplugins/go/tasks/pluginmachinery/core/mocks/task_overrides.go index ab42b1189f..7cc46d35ef 100644 --- a/flyteplugins/go/tasks/pluginmachinery/core/mocks/task_overrides.go +++ b/flyteplugins/go/tasks/pluginmachinery/core/mocks/task_overrides.go @@ -14,6 +14,40 @@ type TaskOverrides struct { mock.Mock } +type TaskOverrides_GetAnnotations struct { + *mock.Call +} + +func (_m TaskOverrides_GetAnnotations) Return(_a0 map[string]string) *TaskOverrides_GetAnnotations { + return &TaskOverrides_GetAnnotations{Call: _m.Call.Return(_a0)} +} + +func (_m *TaskOverrides) OnGetAnnotations() *TaskOverrides_GetAnnotations { + c_call := _m.On("GetAnnotations") + return &TaskOverrides_GetAnnotations{Call: c_call} +} + +func (_m *TaskOverrides) OnGetAnnotationsMatch(matchers ...interface{}) *TaskOverrides_GetAnnotations { + c_call := _m.On("GetAnnotations", matchers...) + return &TaskOverrides_GetAnnotations{Call: c_call} +} + +// GetAnnotations provides a mock function with given fields: +func (_m *TaskOverrides) GetAnnotations() map[string]string { + ret := _m.Called() + + var r0 map[string]string + if rf, ok := ret.Get(0).(func() map[string]string); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(map[string]string) + } + } + + return r0 +} + type TaskOverrides_GetConfig struct { *mock.Call } @@ -114,6 +148,40 @@ func (_m *TaskOverrides) GetExtendedResources() *flyteidlcore.ExtendedResources return r0 } +type TaskOverrides_GetLabels struct { + *mock.Call +} + +func (_m TaskOverrides_GetLabels) Return(_a0 map[string]string) *TaskOverrides_GetLabels { + return &TaskOverrides_GetLabels{Call: _m.Call.Return(_a0)} +} + +func (_m *TaskOverrides) OnGetLabels() *TaskOverrides_GetLabels { + c_call := _m.On("GetLabels") + return &TaskOverrides_GetLabels{Call: c_call} +} + +func (_m *TaskOverrides) OnGetLabelsMatch(matchers ...interface{}) *TaskOverrides_GetLabels { + c_call := _m.On("GetLabels", matchers...) + return &TaskOverrides_GetLabels{Call: c_call} +} + +// GetLabels provides a mock function with given fields: +func (_m *TaskOverrides) GetLabels() map[string]string { + ret := _m.Called() + + var r0 map[string]string + if rf, ok := ret.Get(0).(func() map[string]string); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(map[string]string) + } + } + + return r0 +} + type TaskOverrides_GetResources struct { *mock.Call } diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/config/k8spluginconfig_flags.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/config/k8spluginconfig_flags.go index caa485ff39..7e94782749 100755 --- a/flyteplugins/go/tasks/pluginmachinery/flytek8s/config/k8spluginconfig_flags.go +++ b/flyteplugins/go/tasks/pluginmachinery/flytek8s/config/k8spluginconfig_flags.go @@ -70,5 +70,6 @@ func (cfg K8sPluginConfig) GetPFlagSet(prefix string) *pflag.FlagSet { cmdFlags.Int(fmt.Sprintf("%v%v", prefix, "update-base-backoff-duration"), defaultK8sConfig.UpdateBaseBackoffDuration, "Initial delay in exponential backoff when updating a resource in milliseconds.") cmdFlags.Int(fmt.Sprintf("%v%v", prefix, "update-backoff-retries"), defaultK8sConfig.UpdateBackoffRetries, "Number of retries for exponential backoff when updating a resource.") cmdFlags.StringSlice(fmt.Sprintf("%v%v", prefix, "add-tolerations-for-extended-resources"), defaultK8sConfig.AddTolerationsForExtendedResources, "Name of the extended resources for which tolerations should be added.") + cmdFlags.Bool(fmt.Sprintf("%v%v", prefix, "enable-distributed-error-aggregation"), defaultK8sConfig.EnableDistributedErrorAggregation, "If true, will aggregate errors of different worker pods for distributed tasks.") return cmdFlags } diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/config/k8spluginconfig_flags_test.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/config/k8spluginconfig_flags_test.go index cb50078620..3f68cfa1d7 100755 --- a/flyteplugins/go/tasks/pluginmachinery/flytek8s/config/k8spluginconfig_flags_test.go +++ b/flyteplugins/go/tasks/pluginmachinery/flytek8s/config/k8spluginconfig_flags_test.go @@ -379,4 +379,18 @@ func TestK8sPluginConfig_SetFlags(t *testing.T) { } }) }) + 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()) + } + }) + }) } diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/plugin_exec_context.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/plugin_exec_context.go index ff9ae2ba9f..7e5e04893f 100644 --- a/flyteplugins/go/tasks/pluginmachinery/flytek8s/plugin_exec_context.go +++ b/flyteplugins/go/tasks/pluginmachinery/flytek8s/plugin_exec_context.go @@ -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 diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper_test.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper_test.go index 0a70cdd895..151197b3c9 100644 --- a/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper_test.go +++ b/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper_test.go @@ -27,11 +27,9 @@ import ( "github.com/flyteorg/flyte/flytestdlib/storage" ) -func dummyTaskExecutionMetadata(resources *v1.ResourceRequirements, extendedResources *core.ExtendedResources, containerImage string) pluginsCore.TaskExecutionMetadata { +func dummyTaskExecutionMetadata(resources *v1.ResourceRequirements, extendedResources *core.ExtendedResources, containerImage string, annotations map[string]string, labels map[string]string) pluginsCore.TaskExecutionMetadata { taskExecutionMetadata := &pluginsCoreMock.TaskExecutionMetadata{} 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"}) taskExecutionMetadata.On("GetOwnerReference").Return(metav1.OwnerReference{ Kind: "node", Name: "blah", @@ -54,6 +52,8 @@ func dummyTaskExecutionMetadata(resources *v1.ResourceRequirements, extendedReso to.On("GetResources").Return(resources) to.On("GetExtendedResources").Return(extendedResources) to.On("GetContainerImage").Return(containerImage) + to.On("GetAnnotations").Return(annotations) + to.On("GetLabels").Return(labels) taskExecutionMetadata.On("GetOverrides").Return(to) taskExecutionMetadata.On("IsInterruptible").Return(true) taskExecutionMetadata.OnGetPlatformResources().Return(&v1.ResourceRequirements{}) @@ -82,7 +82,7 @@ func dummyInputReader() io.InputReader { return inputReader } -func dummyExecContext(taskTemplate *core.TaskTemplate, r *v1.ResourceRequirements, rm *core.ExtendedResources, containerImage string) pluginsCore.TaskExecutionContext { +func dummyExecContext(taskTemplate *core.TaskTemplate, r *v1.ResourceRequirements, rm *core.ExtendedResources, containerImage string, annotations map[string]string, labels map[string]string) pluginsCore.TaskExecutionContext { ow := &pluginsIOMock.OutputWriter{} ow.OnGetOutputPrefixPath().Return("") ow.OnGetRawOutputPrefix().Return("") @@ -90,7 +90,7 @@ func dummyExecContext(taskTemplate *core.TaskTemplate, r *v1.ResourceRequirement ow.OnGetPreviousCheckpointsPrefix().Return("/prev") tCtx := &pluginsCoreMock.TaskExecutionContext{} - tCtx.OnTaskExecutionMetadata().Return(dummyTaskExecutionMetadata(r, rm, containerImage)) + tCtx.OnTaskExecutionMetadata().Return(dummyTaskExecutionMetadata(r, rm, containerImage, annotations, labels)) tCtx.OnInputReader().Return(dummyInputReader()) tCtx.OnOutputWriter().Return(ow) @@ -703,7 +703,7 @@ func updatePod(t *testing.T) { v1.ResourceCPU: resource.MustParse("1024m"), v1.ResourceEphemeralStorage: resource.MustParse("100M"), }, - }, nil, "") + }, nil, "", nil, nil) pod := &v1.Pod{ Spec: v1.PodSpec{ @@ -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{ NodeAffinity: &v1.NodeAffinity{ @@ -820,7 +820,7 @@ func toK8sPodInterruptible(t *testing.T) { v1.ResourceCPU: resource.MustParse("1024m"), v1.ResourceEphemeralStorage: resource.MustParse("100M"), }, - }, nil, "") + }, nil, "", nil, nil) p, _, _, err := ToK8sPodSpec(ctx, x) assert.NoError(t, err) @@ -887,7 +887,7 @@ func TestToK8sPod(t *testing.T) { v1.ResourceCPU: resource.MustParse("1024m"), v1.ResourceEphemeralStorage: resource.MustParse("100M"), }, - }, nil, "") + }, nil, "", nil, nil) p, _, _, err := ToK8sPodSpec(ctx, x) assert.NoError(t, err) @@ -904,7 +904,7 @@ func TestToK8sPod(t *testing.T) { v1.ResourceCPU: resource.MustParse("1024m"), v1.ResourceEphemeralStorage: resource.MustParse("100M"), }, - }, nil, "") + }, nil, "", nil, nil) p, _, _, err := ToK8sPodSpec(ctx, x) assert.NoError(t, err) @@ -922,7 +922,7 @@ func TestToK8sPod(t *testing.T) { v1.ResourceCPU: resource.MustParse("1024m"), v1.ResourceEphemeralStorage: resource.MustParse("100M"), }, - }, nil, "") + }, nil, "", nil, nil) assert.NoError(t, config.SetK8sPluginConfig(&config.K8sPluginConfig{ DefaultNodeSelector: map[string]string{ @@ -949,7 +949,7 @@ func TestToK8sPod(t *testing.T) { }, })) - x := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{}, nil, "") + x := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{}, nil, "", nil, nil) p, _, _, err := ToK8sPodSpec(ctx, x) assert.NoError(t, err) assert.NotNil(t, p.SecurityContext) @@ -961,7 +961,7 @@ func TestToK8sPod(t *testing.T) { assert.NoError(t, config.SetK8sPluginConfig(&config.K8sPluginConfig{ EnableHostNetworkingPod: &enabled, })) - x := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{}, nil, "") + x := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{}, nil, "", nil, nil) p, _, _, err := ToK8sPodSpec(ctx, x) assert.NoError(t, err) assert.True(t, p.HostNetwork) @@ -972,7 +972,7 @@ func TestToK8sPod(t *testing.T) { assert.NoError(t, config.SetK8sPluginConfig(&config.K8sPluginConfig{ EnableHostNetworkingPod: &enabled, })) - x := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{}, nil, "") + x := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{}, nil, "", nil, nil) p, _, _, err := ToK8sPodSpec(ctx, x) assert.NoError(t, err) assert.False(t, p.HostNetwork) @@ -980,7 +980,7 @@ func TestToK8sPod(t *testing.T) { t.Run("skipSettingHostNetwork", func(t *testing.T) { assert.NoError(t, config.SetK8sPluginConfig(&config.K8sPluginConfig{})) - x := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{}, nil, "") + x := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{}, nil, "", nil, nil) p, _, _, err := ToK8sPodSpec(ctx, x) assert.NoError(t, err) assert.False(t, p.HostNetwork) @@ -1014,7 +1014,7 @@ func TestToK8sPod(t *testing.T) { }, })) - x := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{}, nil, "") + x := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{}, nil, "", nil, nil) p, _, _, err := ToK8sPodSpec(ctx, x) assert.NoError(t, err) assert.NotNil(t, p.DNSConfig) @@ -1035,7 +1035,7 @@ func TestToK8sPod(t *testing.T) { "foo": "bar", }, })) - x := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{}, nil, "") + x := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{}, nil, "", nil, nil) p, _, _, err := ToK8sPodSpec(ctx, x) assert.NoError(t, err) for _, c := range p.Containers { @@ -1055,7 +1055,7 @@ func TestToK8sPodContainerImage(t *testing.T) { taskContext := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{ Requests: v1.ResourceList{ v1.ResourceCPU: resource.MustParse("1024m"), - }}, nil, "foo:latest") + }}, nil, "foo:latest", nil, nil) p, _, _, err := ToK8sPodSpec(context.TODO(), taskContext) assert.NoError(t, err) assert.Equal(t, "foo:latest", p.Containers[0].Image) @@ -1167,7 +1167,7 @@ func TestToK8sPodExtendedResources(t *testing.T) { t.Run(f.name, func(t *testing.T) { taskTemplate := dummyTaskTemplate() taskTemplate.ExtendedResources = f.extendedResourcesBase - taskContext := dummyExecContext(taskTemplate, f.resources, f.extendedResourcesOverride, "") + taskContext := dummyExecContext(taskTemplate, f.resources, f.extendedResourcesOverride, "", nil, nil) p, _, _, err := ToK8sPodSpec(context.TODO(), taskContext) assert.NoError(t, err) @@ -1829,7 +1829,7 @@ func TestGetPodTemplate(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, nil)) tCtx.OnTaskReader().Return(taskReader) // initialize PodTemplateStore @@ -1855,7 +1855,7 @@ func TestGetPodTemplate(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, nil)) tCtx.OnTaskReader().Return(taskReader) // initialize PodTemplateStore @@ -1882,7 +1882,7 @@ func TestGetPodTemplate(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, nil)) tCtx.OnTaskReader().Return(taskReader) // initialize PodTemplateStore @@ -1910,7 +1910,7 @@ func TestGetPodTemplate(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, nil)) tCtx.OnTaskReader().Return(taskReader) // initialize PodTemplateStore @@ -1960,7 +1960,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, nil)) tCtx.OnTaskReader().Return(taskReader) resultPodSpec, resultObjectMeta, err := MergeWithBasePodTemplate(context.TODO(), tCtx, &podSpec, &objectMeta, "foo", "foo-init") @@ -2022,7 +2022,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, nil)) tCtx.OnTaskReader().Return(taskReader) resultPodSpec, resultObjectMeta, err := MergeWithBasePodTemplate(context.TODO(), tCtx, &podSpec, &objectMeta, "foo", "foo-init") diff --git a/flyteplugins/tests/end_to_end.go b/flyteplugins/tests/end_to_end.go index b045deae13..4d21196d1a 100644 --- a/flyteplugins/tests/end_to_end.go +++ b/flyteplugins/tests/end_to_end.go @@ -148,6 +148,8 @@ func RunPluginEndToEndTest(t *testing.T, executor pluginCore.Plugin, template *i }) overrides.OnGetExtendedResources().Return(&idlCore.ExtendedResources{}) overrides.OnGetContainerImage().Return("") + overrides.OnGetAnnotations().Return(map[string]string{}) + overrides.OnGetLabels().Return(map[string]string{}) tMeta := &coreMocks.TaskExecutionMetadata{} tMeta.OnGetTaskExecutionID().Return(tID) diff --git a/flytepropeller/pkg/apis/flyteworkflow/v1alpha1/iface.go b/flytepropeller/pkg/apis/flyteworkflow/v1alpha1/iface.go index c2022dea25..f60f7115ee 100644 --- a/flytepropeller/pkg/apis/flyteworkflow/v1alpha1/iface.go +++ b/flytepropeller/pkg/apis/flyteworkflow/v1alpha1/iface.go @@ -443,6 +443,8 @@ type ExecutableNode interface { IsInterruptible() *bool GetName() string GetContainerImage() string + GetAnnotations() map[string]string + GetLabels() map[string]string } // ExecutableWorkflowStatus is an interface for the Workflow p. This is the mutable portion for a Workflow diff --git a/flytepropeller/pkg/apis/flyteworkflow/v1alpha1/mocks/ExecutableNode.go b/flytepropeller/pkg/apis/flyteworkflow/v1alpha1/mocks/ExecutableNode.go index e6a03b7deb..386305481c 100644 --- a/flytepropeller/pkg/apis/flyteworkflow/v1alpha1/mocks/ExecutableNode.go +++ b/flytepropeller/pkg/apis/flyteworkflow/v1alpha1/mocks/ExecutableNode.go @@ -52,6 +52,40 @@ func (_m *ExecutableNode) GetActiveDeadline() *time.Duration { return r0 } +type ExecutableNode_GetAnnotations struct { + *mock.Call +} + +func (_m ExecutableNode_GetAnnotations) Return(_a0 map[string]string) *ExecutableNode_GetAnnotations { + return &ExecutableNode_GetAnnotations{Call: _m.Call.Return(_a0)} +} + +func (_m *ExecutableNode) OnGetAnnotations() *ExecutableNode_GetAnnotations { + c_call := _m.On("GetAnnotations") + return &ExecutableNode_GetAnnotations{Call: c_call} +} + +func (_m *ExecutableNode) OnGetAnnotationsMatch(matchers ...interface{}) *ExecutableNode_GetAnnotations { + c_call := _m.On("GetAnnotations", matchers...) + return &ExecutableNode_GetAnnotations{Call: c_call} +} + +// GetAnnotations provides a mock function with given fields: +func (_m *ExecutableNode) GetAnnotations() map[string]string { + ret := _m.Called() + + var r0 map[string]string + if rf, ok := ret.Get(0).(func() map[string]string); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(map[string]string) + } + } + + return r0 +} + type ExecutableNode_GetArrayNode struct { *mock.Call } @@ -386,6 +420,40 @@ func (_m *ExecutableNode) GetKind() v1alpha1.NodeKind { return r0 } +type ExecutableNode_GetLabels struct { + *mock.Call +} + +func (_m ExecutableNode_GetLabels) Return(_a0 map[string]string) *ExecutableNode_GetLabels { + return &ExecutableNode_GetLabels{Call: _m.Call.Return(_a0)} +} + +func (_m *ExecutableNode) OnGetLabels() *ExecutableNode_GetLabels { + c_call := _m.On("GetLabels") + return &ExecutableNode_GetLabels{Call: c_call} +} + +func (_m *ExecutableNode) OnGetLabelsMatch(matchers ...interface{}) *ExecutableNode_GetLabels { + c_call := _m.On("GetLabels", matchers...) + return &ExecutableNode_GetLabels{Call: c_call} +} + +// GetLabels provides a mock function with given fields: +func (_m *ExecutableNode) GetLabels() map[string]string { + ret := _m.Called() + + var r0 map[string]string + if rf, ok := ret.Get(0).(func() map[string]string); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(map[string]string) + } + } + + return r0 +} + type ExecutableNode_GetName struct { *mock.Call } diff --git a/flytepropeller/pkg/apis/flyteworkflow/v1alpha1/nodes.go b/flytepropeller/pkg/apis/flyteworkflow/v1alpha1/nodes.go index 743327e373..96ece9a4a3 100644 --- a/flytepropeller/pkg/apis/flyteworkflow/v1alpha1/nodes.go +++ b/flytepropeller/pkg/apis/flyteworkflow/v1alpha1/nodes.go @@ -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"` } 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 +} diff --git a/flytepropeller/pkg/compiler/transformers/k8s/node.go b/flytepropeller/pkg/compiler/transformers/k8s/node.go index 18ec1ba02f..a97c63fb11 100644 --- a/flytepropeller/pkg/compiler/transformers/k8s/node.go +++ b/flytepropeller/pkg/compiler/transformers/k8s/node.go @@ -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 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() + } } } @@ -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) { diff --git a/flytepropeller/pkg/controller/nodes/task/k8s/plugin_manager.go b/flytepropeller/pkg/controller/nodes/task/k8s/plugin_manager.go index 6d558f818e..b623905095 100644 --- a/flytepropeller/pkg/controller/nodes/task/k8s/plugin_manager.go +++ b/flytepropeller/pkg/controller/nodes/task/k8s/plugin_manager.go @@ -112,8 +112,8 @@ type PluginManager struct { func (e *PluginManager) addObjectMetadata(taskCtx pluginsCore.TaskExecutionMetadata, o client.Object, cfg *config.K8sPluginConfig) { o.SetNamespace(taskCtx.GetNamespace()) - o.SetAnnotations(pluginsUtils.UnionMaps(cfg.DefaultAnnotations, o.GetAnnotations(), pluginsUtils.CopyMap(taskCtx.GetAnnotations()))) - o.SetLabels(pluginsUtils.UnionMaps(cfg.DefaultLabels, o.GetLabels(), pluginsUtils.CopyMap(taskCtx.GetLabels()))) + o.SetAnnotations(pluginsUtils.UnionMaps(cfg.DefaultAnnotations, o.GetAnnotations(), pluginsUtils.CopyMap(taskCtx.GetAnnotations()), pluginsUtils.CopyMap(taskCtx.GetOverrides().GetAnnotations()))) + o.SetLabels(pluginsUtils.UnionMaps(cfg.DefaultLabels, o.GetLabels(), pluginsUtils.CopyMap(taskCtx.GetLabels()), pluginsUtils.CopyMap(taskCtx.GetOverrides().GetLabels()))) o.SetName(taskCtx.GetTaskExecutionID().GetGeneratedName()) if !e.plugin.GetProperties().DisableInjectOwnerReferences { diff --git a/flytepropeller/pkg/controller/nodes/task/k8s/plugin_manager_test.go b/flytepropeller/pkg/controller/nodes/task/k8s/plugin_manager_test.go index a02c32dca1..24c1025986 100644 --- a/flytepropeller/pkg/controller/nodes/task/k8s/plugin_manager_test.go +++ b/flytepropeller/pkg/controller/nodes/task/k8s/plugin_manager_test.go @@ -207,6 +207,12 @@ func getMockTaskExecutionMetadata() pluginsCore.TaskExecutionMetadata { taskExecutionMetadata.On("GetOwnerReference").Return(metav1.OwnerReference{Name: "x"}) taskExecutionMetadata.On("GetSecurityContext").Return(core.SecurityContext{RunAs: &core.Identity{}}) + to := &pluginsCoreMock.TaskOverrides{} + to.On("GetAnnotations").Return(map[string]string{}) + to.On("GetLabels").Return(map[string]string{}) + + taskExecutionMetadata.On("GetOverrides").Return(to) + id := &pluginsCoreMock.TaskExecutionID{} id.On("GetGeneratedName").Return("test") id.On("GetID").Return(core.TaskExecutionIdentifier{}) @@ -219,13 +225,21 @@ func getMockTaskExecutionMetadataCustom( ns string, annotations map[string]string, labels map[string]string, - ownerRef metav1.OwnerReference) pluginsCore.TaskExecutionMetadata { + ownerRef metav1.OwnerReference, + annotationsOverride map[string]string, + labelsOverride map[string]string) pluginsCore.TaskExecutionMetadata { taskExecutionMetadata := &pluginsCoreMock.TaskExecutionMetadata{} taskExecutionMetadata.On("GetNamespace").Return(ns) taskExecutionMetadata.On("GetAnnotations").Return(annotations) taskExecutionMetadata.On("GetLabels").Return(labels) taskExecutionMetadata.On("GetOwnerReference").Return(ownerRef) + to := &pluginsCoreMock.TaskOverrides{} + to.On("GetAnnotations").Return(annotationsOverride) + to.On("GetLabels").Return(labelsOverride) + + taskExecutionMetadata.On("GetOverrides").Return(to) + id := &pluginsCoreMock.TaskExecutionID{} id.On("GetGeneratedName").Return(tid) id.On("GetID").Return(core.TaskExecutionIdentifier{}) @@ -908,7 +922,9 @@ func TestPluginManager_AddObjectMetadata(t *testing.T) { or := metav1.OwnerReference{} l := map[string]string{"l1": "lv1"} a := map[string]string{"aKey": "aVal"} - tm := getMockTaskExecutionMetadataCustom(genName, ns, a, l, or) + lo := map[string]string{} + ao := map[string]string{} + tm := getMockTaskExecutionMetadataCustom(genName, ns, a, l, or, lo, ao) cfg := config.GetK8sPluginConfig() @@ -1008,6 +1024,24 @@ func TestPluginManager_AddObjectMetadata(t *testing.T) { assert.Contains(t, o.GetFinalizers(), finalizer) }) + t.Run("with override annotations and labels", func(t *testing.T) { + lo = map[string]string{"l1": "lv2"} + ao = map[string]string{"aKey": "bVal"} + tm = getMockTaskExecutionMetadataCustom(genName, ns, a, l, or, ao, lo) + + cfg = config.GetK8sPluginConfig() + + o := &v1.Pod{} + p := pluginsk8sMock.Plugin{} + p.OnGetProperties().Return(k8s.PluginProperties{}) + pluginManager := PluginManager{plugin: &p} + pluginManager.addObjectMetadata(tm, o, cfg) + assert.Equal(t, map[string]string{ + "cluster-autoscaler.kubernetes.io/safe-to-evict": "false", + "aKey": "bVal", + }, o.GetAnnotations()) + assert.Equal(t, lo, o.GetLabels()) + }) } func TestResourceManagerConstruction(t *testing.T) {