diff --git a/flytepropeller/pkg/compiler/transformers/k8s/utils.go b/flytepropeller/pkg/compiler/transformers/k8s/utils.go index bd08be3a9a..b8e15e176b 100644 --- a/flytepropeller/pkg/compiler/transformers/k8s/utils.go +++ b/flytepropeller/pkg/compiler/transformers/k8s/utils.go @@ -86,7 +86,17 @@ func StripTypeMetadata(t *core.LiteralType) *core.LiteralType { c := *t c.Metadata = nil - c.Annotation = nil + + // Special-case the presence of cache-key-metadata. This is a special field that is used to store metadata about + // used in the cache key generation. This does not affect compatibility and it is purely used for cache key calculations. + if c.GetAnnotation() != nil { + annotations := c.GetAnnotation().GetAnnotations().GetFields() + // The presence of the key `cache-key-metadata` indicates that we should leave the metadata intact. + if _, ok := annotations["cache-key-metadata"]; !ok { + c.Annotation = nil + } + } + // Note that we cannot strip `Structure` from the type because the dynamic node output type is used to validate the // interface of the dynamically compiled workflow. `Structure` is used to extend type checking information on // different Flyte types and is therefore required to ensure correct type validation. diff --git a/flytepropeller/pkg/compiler/transformers/k8s/utils_test.go b/flytepropeller/pkg/compiler/transformers/k8s/utils_test.go index 0a7e991399..0a8866ff8d 100644 --- a/flytepropeller/pkg/compiler/transformers/k8s/utils_test.go +++ b/flytepropeller/pkg/compiler/transformers/k8s/utils_test.go @@ -6,6 +6,7 @@ import ( "github.com/go-test/deep" _struct "github.com/golang/protobuf/ptypes/struct" "github.com/stretchr/testify/assert" + "google.golang.org/protobuf/types/known/structpb" "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core" ) @@ -248,6 +249,99 @@ func TestStripTypeMetadata(t *testing.T) { }, }, }, + { + name: "cache-key-metadata set", + args: &core.LiteralType{ + Type: &core.LiteralType_Simple{ + Simple: core.SimpleType_INTEGER, + }, + Annotation: &core.TypeAnnotation{ + Annotations: &structpb.Struct{ + Fields: map[string]*structpb.Value{ + "cache-key-metadata": { + Kind: &_struct.Value_StructValue{ + StructValue: &structpb.Struct{ + Fields: map[string]*structpb.Value{ + "foo": { + Kind: &_struct.Value_StringValue{ + StringValue: "bar", + }, + }, + }, + }, + }, + }, + }, + }, + }, + Metadata: &_struct.Struct{ + Fields: map[string]*_struct.Value{ + "foo": { + Kind: &_struct.Value_StringValue{ + StringValue: "bar", + }, + }, + }, + }, + }, + want: &core.LiteralType{ + Type: &core.LiteralType_Simple{ + Simple: core.SimpleType_INTEGER, + }, + Annotation: &core.TypeAnnotation{ + Annotations: &structpb.Struct{ + Fields: map[string]*structpb.Value{ + "cache-key-metadata": { + Kind: &_struct.Value_StructValue{ + StructValue: &structpb.Struct{ + Fields: map[string]*structpb.Value{ + "foo": { + Kind: &_struct.Value_StringValue{ + StringValue: "bar", + }, + }, + }, + }, + }, + }, + }, + }, + }, + }, + }, + { + name: "cache-key-metadata not present in annotation", + args: &core.LiteralType{ + Type: &core.LiteralType_Simple{ + Simple: core.SimpleType_INTEGER, + }, + Annotation: &core.TypeAnnotation{ + Annotations: &structpb.Struct{ + Fields: map[string]*structpb.Value{ + "some-key": { + Kind: &_struct.Value_StringValue{ + StringValue: "some-value", + }, + }, + }, + }, + }, + Metadata: &_struct.Struct{ + Fields: map[string]*_struct.Value{ + "foo": { + Kind: &_struct.Value_StringValue{ + StringValue: "bar", + }, + }, + }, + }, + }, + want: &core.LiteralType{ + Type: &core.LiteralType_Simple{ + Simple: core.SimpleType_INTEGER, + }, + }, + }, } for _, tt := range tests {