Handle optional pipeline inputs in the driver (kubeflow#11788)

If the pipeline run is submitted without specifying an optional
parameter and there is no default, it was not handled by the driver.
The approach taken is explicitly set null for these values and let the
driver handle if the component parameter has a default that can be
used in the launcher.

Signed-off-by: mprahl <[email protected]>
(cherry picked from commit bb7a108)

Author: mprahl

backend/src/v2/compiler/argocompiler/argo.go

@@ -72,8 +72,12 @@ func Compile(jobArg *pipelinespec.PipelineJob, kubernetesSpecArg *pipelinespec.S
 	specParams := spec.GetRoot().GetInputDefinitions().GetParameters()
 	for name, param := range specParams {
 		_, ok := job.RuntimeConfig.ParameterValues[name]
-		if !ok && param.GetDefaultValue() != nil {
-			job.RuntimeConfig.ParameterValues[name] = param.GetDefaultValue()
+		if !ok {
+			if param.GetDefaultValue() != nil {
+				job.RuntimeConfig.ParameterValues[name] = param.GetDefaultValue()
+			} else if param.IsOptional {
+				job.RuntimeConfig.ParameterValues[name] = structpb.NewNullValue()
+			}
 		}
 	}
 

backend/src/v2/driver/driver.go

@@ -17,6 +17,7 @@ package driver
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"slices"
 	"strconv"
@@ -51,6 +52,8 @@ var dummyImages = map[string]string{
 	"argostub/deletepvc": "delete PVC",
 }
 
+var ErrResolvedParameterNull = errors.New("the resolved input parameter is null")
+
 // TODO(capri-xiyue): Move driver to component package
 // Driver options
 type Options struct {
@@ -191,6 +194,7 @@ func RootDAG(ctx context.Context, opts Options, mlmd *metadata.Client) (executio
 	if err != nil {
 		return nil, err
 	}
+
 	executorInput := &pipelinespec.ExecutorInput{
 		Inputs: &pipelinespec.ExecutorInput_Inputs{
 			ParameterValues: opts.RuntimeConfig.GetParameterValues(),
@@ -803,7 +807,7 @@ func extendPodSpecPatch(
 	for _, secretAsVolume := range kubernetesExecutorConfig.GetSecretAsVolume() {
 		var secretName string
 		if secretAsVolume.SecretNameParameter != nil {
-			resolvedSecretName, err := resolveInputParameter(ctx, dag, pipeline, opts, mlmd,
+			resolvedSecretName, err := resolveInputParameterStr(ctx, dag, pipeline, opts, mlmd,
 				secretAsVolume.SecretNameParameter, inputParams)
 			if err != nil {
 				return fmt.Errorf("failed to resolve secret name: %w", err)
@@ -847,7 +851,7 @@ func extendPodSpecPatch(
 
 			var secretName string
 			if secretAsEnv.SecretNameParameter != nil {
-				resolvedSecretName, err := resolveInputParameter(ctx, dag, pipeline, opts, mlmd,
+				resolvedSecretName, err := resolveInputParameterStr(ctx, dag, pipeline, opts, mlmd,
 					secretAsEnv.SecretNameParameter, inputParams)
 				if err != nil {
 					return fmt.Errorf("failed to resolve secret name: %w", err)
@@ -869,7 +873,7 @@ func extendPodSpecPatch(
 	for _, configMapAsVolume := range kubernetesExecutorConfig.GetConfigMapAsVolume() {
 		var configMapName string
 		if configMapAsVolume.ConfigMapNameParameter != nil {
-			resolvedSecretName, err := resolveInputParameter(ctx, dag, pipeline, opts, mlmd,
+			resolvedSecretName, err := resolveInputParameterStr(ctx, dag, pipeline, opts, mlmd,
 				configMapAsVolume.ConfigMapNameParameter, inputParams)
 			if err != nil {
 				return fmt.Errorf("failed to resolve configmap name: %w", err)
@@ -915,7 +919,7 @@ func extendPodSpecPatch(
 
 			var configMapName string
 			if configMapAsEnv.ConfigMapNameParameter != nil {
-				resolvedSecretName, err := resolveInputParameter(ctx, dag, pipeline, opts, mlmd,
+				resolvedSecretName, err := resolveInputParameterStr(ctx, dag, pipeline, opts, mlmd,
 					configMapAsEnv.ConfigMapNameParameter, inputParams)
 				if err != nil {
 					return fmt.Errorf("failed to resolve configmap name: %w", err)
@@ -937,7 +941,7 @@ func extendPodSpecPatch(
 	for _, imagePullSecret := range kubernetesExecutorConfig.GetImagePullSecret() {
 		var secretName string
 		if imagePullSecret.SecretNameParameter != nil {
-			resolvedSecretName, err := resolveInputParameter(ctx, dag, pipeline, opts, mlmd,
+			resolvedSecretName, err := resolveInputParameterStr(ctx, dag, pipeline, opts, mlmd,
 				imagePullSecret.SecretNameParameter, inputParams)
 			if err != nil {
 				return fmt.Errorf("failed to resolve image pull secret name: %w", err)
@@ -1497,8 +1501,20 @@ func resolveInputs(
 	for name, paramSpec := range task.GetInputs().GetParameters() {
 		v, err := resolveInputParameter(ctx, dag, pipeline, opts, mlmd, paramSpec, inputParams)
 		if err != nil {
+			if !errors.Is(err, ErrResolvedParameterNull) {
+				return nil, err
+			}
+
+			componentParam, ok := opts.Component.GetInputDefinitions().GetParameters()[name]
+			if ok && componentParam != nil && componentParam.IsOptional {
+				// If the resolved paramter was null and the component input parameter is optional, just skip setting
+				// it and the launcher will handle defaults.
+				continue
+			}
+
 			return nil, err
 		}
+
 		inputs.ParameterValues[name] = v
 	}
 
@@ -1515,7 +1531,9 @@ func resolveInputs(
 }
 
 // resolveInputParameter resolves an InputParameterSpec
-// using a given input context via InputParams.
+// using a given input context via InputParams. ErrResolvedParameterNull is returned if paramSpec
+// is a component input parameter and parameter resolves to a null value (i.e. an optional pipeline input with no
+// default). The caller can decide if this is allowed in that context.
 func resolveInputParameter(
 	ctx context.Context,
 	dag *metadata.DAG,
@@ -1539,6 +1557,13 @@ func resolveInputParameter(
 		if !ok {
 			return nil, paramError(fmt.Errorf("parent DAG does not have input parameter %s", componentInput))
 		}
+
+		if _, isNullValue := v.GetKind().(*structpb.Value_NullValue); isNullValue {
+			// Null values are only allowed for optional pipeline input parameters with no values. The caller has this
+			// context to know if this is allowed.
+			return nil, fmt.Errorf("%w: %s", ErrResolvedParameterNull, componentInput)
+		}
+
 		return v, nil
 
 	// This is the case where the input comes from the output of an upstream task.
@@ -1588,6 +1613,33 @@ func resolveInputParameter(
 	}
 }
 
+// resolveInputParameterStr is like resolveInputParameter but returns an error if the resolved value is not a non-empty
+// string.
+func resolveInputParameterStr(
+	ctx context.Context,
+	dag *metadata.DAG,
+	pipeline *metadata.Pipeline,
+	opts Options,
+	mlmd *metadata.Client,
+	paramSpec *pipelinespec.TaskInputsSpec_InputParameterSpec,
+	inputParams map[string]*structpb.Value,
+) (*structpb.Value, error) {
+	val, err := resolveInputParameter(ctx, dag, pipeline, opts, mlmd, paramSpec, inputParams)
+	if err != nil {
+		return nil, err
+	}
+
+	if typedVal, ok := val.GetKind().(*structpb.Value_StringValue); ok && typedVal != nil {
+		if typedVal.StringValue == "" {
+			return nil, fmt.Errorf("resolving input parameter with spec %s. Expected a non-empty string.", paramSpec)
+		}
+	} else {
+		return nil, fmt.Errorf("resolving input parameter with spec %s. Expected a string but got: %T", paramSpec, val.GetKind())
+	}
+
+	return val, nil
+}
+
 // resolveInputArtifact resolves an InputArtifactSpec
 // using a given input context via inputArtifacts.
 func resolveInputArtifact(
@@ -2366,7 +2418,7 @@ func makeVolumeMountPatch(
 			}
 		}
 
-		resolvedPvcName, err := resolveInputParameter(ctx, dag, pipeline, opts, mlmd,
+		resolvedPvcName, err := resolveInputParameterStr(ctx, dag, pipeline, opts, mlmd,
 			pvcNameParameter, inputParams)
 		if err != nil {
 			return nil, nil, fmt.Errorf("failed to resolve pvc name: %w", err)

samples/v2/component_with_optional_inputs.py

@@ -23,6 +23,9 @@ def component_op(
     input_str1: Optional[str] = 'string default value',
     input_str2: Optional[str] = None,
     input_str3: Optional[str] = None,
+    input_str4_from_pipeline: Optional[str] = "Some component default",
+    input_str5_from_pipeline: Optional[str] = "Some component default",
+    input_str6_from_pipeline: Optional[str] = None,
     input_bool1: Optional[bool] = True,
     input_bool2: Optional[bool] = None,
     input_dict: Optional[Dict[str, int]] = {"a": 1},
@@ -32,6 +35,9 @@ def component_op(
     print(f'input_str1: {input_str1}, type: {type(input_str1)}')
     print(f'input_str2: {input_str2}, type: {type(input_str2)}')
     print(f'input_str3: {input_str3}, type: {type(input_str3)}')
+    print(f'input_str4_from_pipeline: {input_str4_from_pipeline}, type: {type(input_str4_from_pipeline)}')
+    print(f'input_str5_from_pipeline: {input_str5_from_pipeline}, type: {type(input_str5_from_pipeline)}')
+    print(f'input_str6_from_pipeline: {input_str6_from_pipeline}, type: {type(input_str6_from_pipeline)}')
     print(f'input_bool1: {input_bool1}, type: {type(input_bool1)}')
     print(f'input_bool2: {input_bool2}, type: {type(input_bool2)}')
     print(f'input_bool: {input_dict}, type: {type(input_dict)}')
@@ -40,10 +46,13 @@ def component_op(
 
 
 @dsl.pipeline(name='v2-component-optional-input')
-def pipeline():
+def pipeline(input_str4: Optional[str] = None, input_str5: Optional[str] = "Some pipeline default", input_str6: Optional[str] = None):
     component_op(
         input_str1='Hello',
         input_str2='World',
+        input_str4_from_pipeline=input_str4,
+        input_str5_from_pipeline=input_str5,
+        input_str6_from_pipeline=input_str6,
     )
 
 

samples/v2/component_with_optional_inputs_test.py

@@ -38,6 +38,7 @@ def verify(t: unittest.TestCase, run: kfp_server_api.ApiRun,
                 'parameters': {
                     'input_str1': 'Hello',
                     'input_str2': 'World',
+                    'input_str5_from_pipeline': 'Some pipeline default',
                 },
             },
             'outputs': {},