unit test for exec validator

flyteorg · Aug 30, 2024 · 16474e5 · 16474e5
1 parent 532a9f6
commit 16474e5
Show file tree

Hide file tree

Showing 5 changed files with 58 additions and 5 deletions.
diff --git a/flyteadmin/pkg/manager/impl/testutils/mock_requests.go b/flyteadmin/pkg/manager/impl/testutils/mock_requests.go
@@ -246,6 +246,28 @@ func GetExecutionRequest() admin.ExecutionCreateRequest {
 	}
 }
 
+func GetExecutionRequestWithOffloadedInputs(inputParam string, literalValue *core.Literal) admin.ExecutionCreateRequest {
+	execReq := GetExecutionRequest()
+	execReq.Inputs = &core.LiteralMap{
+		Literals: map[string]*core.Literal{
+			"foo": {
+				Value: &core.Literal_OffloadedMetadata{
+					OffloadedMetadata: &core.LiteralOffloadedMetadata{
+						Uri:       "s3://bucket/key",
+						SizeBytes: 100,
+						InferredType: &core.LiteralType{
+							Type: &core.LiteralType_Simple{
+								Simple: core.SimpleType_STRING,
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+	return execReq
+}
+
 func GetSampleWorkflowSpecForTest() admin.WorkflowSpec {
 	return admin.WorkflowSpec{
 		Template: &core.WorkflowTemplate{

diff --git a/flyteadmin/pkg/manager/impl/validation/execution_validator_test.go b/flyteadmin/pkg/manager/impl/validation/execution_validator_test.go
@@ -105,6 +105,39 @@ func TestGetExecutionInputs(t *testing.T) {
 	assert.EqualValues(t, expectedMap, *actualInputs)
 }
 
+func TestGetExecutionWithOffloadedInputs(t *testing.T) {
+	execLiteral := &core.Literal{
+		Value: &core.Literal_OffloadedMetadata{
+			OffloadedMetadata: &core.LiteralOffloadedMetadata{
+				Uri:       "s3://bucket/key",
+				SizeBytes: 100,
+				InferredType: &core.LiteralType{
+					Type: &core.LiteralType_Simple{
+						Simple: core.SimpleType_STRING,
+					},
+				},
+			},
+		},
+	}
+	executionRequest := testutils.GetExecutionRequestWithOffloadedInputs("foo", execLiteral)
+	lpRequest := testutils.GetLaunchPlanRequest()
+
+	actualInputs, err := CheckAndFetchInputsForExecution(
+		executionRequest.Inputs,
+		lpRequest.Spec.FixedInputs,
+		lpRequest.Spec.DefaultInputs,
+	)
+	expectedMap := core.LiteralMap{
+		Literals: map[string]*core.Literal{
+			"foo": execLiteral,
+			"bar": coreutils.MustMakeLiteral("bar-value"),
+		},
+	}
+	assert.Nil(t, err)
+	assert.NotNil(t, actualInputs)
+	assert.EqualValues(t, expectedMap, *actualInputs)
+}
+
 func TestValidateExecInputsWrongType(t *testing.T) {
 	executionRequest := testutils.GetExecutionRequest()
 	lpRequest := testutils.GetLaunchPlanRequest()

diff --git a/flyteadmin/pkg/manager/impl/validation/validation.go b/flyteadmin/pkg/manager/impl/validation/validation.go
@@ -234,8 +234,7 @@ func validateLiteralMap(inputMap *core.LiteralMap, fieldName string) error {
 			if name == "" {
 				return errors.NewFlyteAdminErrorf(codes.InvalidArgument, "missing key in %s", fieldName)
 			}
-			fixedInput.GetOffloadedMetadata()
-			if fixedInput == nil || (fixedInput.GetValue() == nil && fixedInput.GetOffloadedMetadata() == nil) {
+			if fixedInput.GetValue() == nil && fixedInput.GetOffloadedMetadata() == nil {
 				return errors.NewFlyteAdminErrorf(codes.InvalidArgument, "missing valid literal in %s %s", fieldName, name)
 			}
 			if isDateTime(fixedInput) {

diff --git a/flytepropeller/pkg/controller/config/config.go b/flytepropeller/pkg/controller/config/config.go
@@ -192,13 +192,13 @@ func (l LiteralOffloadingConfig) IsSupportedSDKVersion(sdk string, versionString
 		c, err := semver.NewConstraint(fmt.Sprintf(">= %s", leastSupportedVersion))
 		if err != nil {
 			// This should never happen
-			logger.Errorf(context.TODO(), "Failed to parse version constraint %s", leastSupportedVersion)
+			logger.Warnf(context.TODO(), "Failed to parse version constraint %s", leastSupportedVersion)
 			return false
 		}
 		version, err := semver.NewVersion(versionString)
 		if err != nil {
 			// This should never happen
-			logger.Errorf(context.TODO(), "Failed to parse version %s", versionString)
+			logger.Warnf(context.TODO(), "Failed to parse version %s", versionString)
 			return false
 		}
 		return c.Check(version)

diff --git a/flytepropeller/pkg/controller/nodes/array/handler.go b/flytepropeller/pkg/controller/nodes/array/handler.go
@@ -657,7 +657,6 @@ func (a *arrayNodeHandler) Handle(ctx context.Context, nCtx interfaces.NodeExecu
 
 // OffloadLargeLiteral offloads the large literal if meets the threshold conditions
 func (a *arrayNodeHandler) OffloadLargeLiteral(ctx context.Context, datastore *storage.DataStore, dataReference storage.DataReference, toBeOffloaded *idlcore.Literal) error {
-	// TODO : also add sdk version checks to ensure that the literal is not offloaded if the sdk version does not support it
 	literalSizeBytes := uint64(proto.Size(toBeOffloaded))
 	literalSizeMB := literalSizeBytes / MB
 	// check if the literal is large