Special-case type annotations to force cache invalidations for msgpac…

…k-binary literals (#6078) Signed-off-by: Eduardo Apolinario <[email protected]> Co-authored-by: Eduardo Apolinario <[email protected]>
flyteorg · Dec 5, 2024 · 4986a24 · 4986a24
1 parent 7fcbdf9
commit 4986a24
Show file tree

Hide file tree

Showing 2 changed files with 105 additions and 1 deletion.
diff --git 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
@@ -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 {