stripping metadata from dataclass types (#626) (#6177)

* stripping metadata from dataclass types (#626) * stripping metadata from dataclass types Signed-off-by: Daniel Rammer <[email protected]> * actually setting on Structure Signed-off-by: Daniel Rammer <[email protected]> * added unit test Signed-off-by: Daniel Rammer <[email protected]> --------- Signed-off-by: Daniel Rammer <[email protected]> * fixed lint Signed-off-by: Daniel Rammer <[email protected]> --------- Signed-off-by: Daniel Rammer <[email protected]>
flyteorg · Jan 17, 2025 · ab463c3 · ab463c3
1 parent 729e71a
commit ab463c3
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 5 deletions.
diff --git a/flytepropeller/pkg/compiler/transformers/k8s/utils.go b/flytepropeller/pkg/compiler/transformers/k8s/utils.go
@@ -97,9 +97,15 @@ func StripTypeMetadata(t *core.LiteralType) *core.LiteralType {
 		}
 	}
 
-	// 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.
+	// strip metadata from dataclass types
+	if c.GetStructure() != nil && len(c.GetStructure().GetDataclassType()) > 0 {
+		dataclassTypes := make(map[string]*core.LiteralType, len(c.GetStructure().GetDataclassType()))
+		for k, v := range c.GetStructure().GetDataclassType() {
+			dataclassTypes[k] = StripTypeMetadata(v)
+		}
+
+		c.GetStructure().DataclassType = dataclassTypes
+	}
 
 	switch underlyingType := c.GetType().(type) {
 	case *core.LiteralType_UnionType:

diff --git a/flytepropeller/pkg/compiler/transformers/k8s/utils_test.go b/flytepropeller/pkg/compiler/transformers/k8s/utils_test.go
@@ -224,7 +224,25 @@ func TestStripTypeMetadata(t *testing.T) {
 											},
 										},
 									},
-									Structure: &core.TypeStructure{Tag: "str"},
+									Structure: &core.TypeStructure{
+										DataclassType: map[string]*core.LiteralType{
+											"foo": {
+												Type: &core.LiteralType_Simple{
+													Simple: core.SimpleType_STRING,
+												},
+												Metadata: &_struct.Struct{
+													Fields: map[string]*_struct.Value{
+														"foo": {
+															Kind: &_struct.Value_StringValue{
+																StringValue: "bar",
+															},
+														},
+													},
+												},
+											},
+										},
+										Tag: "str",
+									},
 								},
 							},
 						},
@@ -241,7 +259,16 @@ func TestStripTypeMetadata(t *testing.T) {
 									Type: &core.LiteralType_Simple{
 										Simple: core.SimpleType_STRING,
 									},
-									Structure: &core.TypeStructure{Tag: "str"},
+									Structure: &core.TypeStructure{
+										DataclassType: map[string]*core.LiteralType{
+											"foo": {
+												Type: &core.LiteralType_Simple{
+													Simple: core.SimpleType_STRING,
+												},
+											},
+										},
+										Tag: "str",
+									},
 								},
 							},
 						},