feat(backend) implement retryStrategy for nested pipelines (kubeflow#11908)

Signed-off-by: agoins <[email protected]>

Author: alyssacgoins

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

@@ -101,6 +101,24 @@ func Test_argo_compiler(t *testing.T) {
 			argoYAMLPath:     "testdata/hello_world_cache_disabled.yaml",
 			compilerOptions:  argocompiler.Options{CacheDisabled: true},
 		},
+		// retry set at pipeline level only.
+		{
+			jobPath:          "../testdata/nested_pipeline_pipeline_retry.json",
+			platformSpecPath: "",
+			argoYAMLPath:     "testdata/nested_pipeline_pipeline_retry.yaml",
+		},
+		// retry set at component level only.
+		{
+			jobPath:          "../testdata/nested_pipeline_sub_component_retry.json",
+			platformSpecPath: "",
+			argoYAMLPath:     "testdata/nested_pipeline_sub_component_retry.yaml",
+		},
+		// retry set at both component and pipeline level.
+		{
+			jobPath:          "../testdata/nested_pipeline_all_level_retry.json",
+			platformSpecPath: "",
+			argoYAMLPath:     "testdata/nested_pipeline_all_level_retry.yaml",
+		},
 	}
 	for _, tt := range tests {
 		t.Run(fmt.Sprintf("%+v", tt), func(t *testing.T) {

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

@@ -260,41 +260,47 @@ type containerExecutorInputs struct {
 // name: argo workflows DAG task name
 // The other arguments are argo workflows task parameters, they can be either a
 // string or a placeholder.
-func (c *workflowCompiler) containerExecutorTask(name string, inputs containerExecutorInputs, refName string) *wfapi.DAGTask {
+func (c *workflowCompiler) containerExecutorTask(name string, inputs containerExecutorInputs, task *pipelinespec.PipelineTaskSpec) *wfapi.DAGTask {
 	when := ""
 	if inputs.condition != "" {
 		when = inputs.condition + " != false"
 	}
-	task := &wfapi.DAGTask{
+	dagTask := &wfapi.DAGTask{
 		Name:     name,
-		Template: c.addContainerExecutorTemplate(name, refName),
+		Template: c.addContainerExecutorTemplate(task),
 		When:     when,
 		Arguments: wfapi.Arguments{
 			Parameters: append([]wfapi.Parameter{
 				{Name: paramPodSpecPatch, Value: wfapi.AnyStringPtr(inputs.podSpecPatch)},
-				{Name: paramCachedDecision, Value: wfapi.AnyStringPtr(inputs.cachedDecision), Default: wfapi.AnyStringPtr("false")}},
-				c.getTaskRetryParametersWithValues(name)...,
+				{Name: paramCachedDecision, Value: wfapi.AnyStringPtr(inputs.cachedDecision), Default: wfapi.AnyStringPtr("false")},
+			},
+				c.getTaskRetryParametersWithValues(task)...,
 			),
 		},
 	}
 
-	addExitTask(task, inputs.exitTemplate, inputs.hookParentDagID)
+	addExitTask(dagTask, inputs.exitTemplate, inputs.hookParentDagID)
 
-	return task
+	return dagTask
 }
 
 // addContainerExecutorTemplate adds a generic container executor template for
 // any container component task.
 // During runtime, it's expected that pod-spec-patch will specify command, args
 // and resources etc, that are different for different tasks.
-func (c *workflowCompiler) addContainerExecutorTemplate(name string, refName string) string {
+func (c *workflowCompiler) addContainerExecutorTemplate(task *pipelinespec.PipelineTaskSpec) string {
 	// container template is parent of container implementation template
 	nameContainerExecutor := "system-container-executor"
 	nameContainerImpl := "system-container-impl"
-	taskRetrySpec := c.getTaskRetryPolicySpec(name)
+	taskRetrySpec := task.GetRetryPolicy()
+	refName := task.GetComponentRef().GetName()
 	if taskRetrySpec != nil {
-		nameContainerExecutor = "retry-" + nameContainerExecutor
-		nameContainerImpl = "retry-" + nameContainerImpl
+		task := c.spec.Components[refName]
+		// retry-system-container-executor/impl is used if retry is set at the component level (if task is not a DAG)
+		if task != nil && task.GetDag() == nil {
+			nameContainerExecutor = "retry-" + nameContainerExecutor
+			nameContainerImpl = "retry-" + nameContainerImpl
+		}
 	}
 	_, ok := c.templates[nameContainerExecutor]
 	if ok {
@@ -305,8 +311,9 @@ func (c *workflowCompiler) addContainerExecutorTemplate(name string, refName str
 		Inputs: wfapi.Inputs{
 			Parameters: append([]wfapi.Parameter{
 				{Name: paramPodSpecPatch},
-				{Name: paramCachedDecision, Default: wfapi.AnyStringPtr("false")}},
-				c.addParameterDefault(c.getTaskRetryParameters(name), "0")...,
+				{Name: paramCachedDecision, Default: wfapi.AnyStringPtr("false")},
+			},
+				c.addParameterDefault(c.getTaskRetryParameters(task), "0")...,
 			),
 		},
 		DAG: &wfapi.DAGTemplate{
@@ -316,8 +323,9 @@ func (c *workflowCompiler) addContainerExecutorTemplate(name string, refName str
 				Arguments: wfapi.Arguments{
 					Parameters: append([]wfapi.Parameter{{
 						Name:  paramPodSpecPatch,
-						Value: wfapi.AnyStringPtr(inputParameter(paramPodSpecPatch))}},
-						c.addParameterInputPath(c.getTaskRetryParameters(name))...,
+						Value: wfapi.AnyStringPtr(inputParameter(paramPodSpecPatch))},
+					},
+						c.addParameterInputPath(c.getTaskRetryParameters(task))...,
 					),
 				},
 				// When cached decision is true, the container
@@ -347,8 +355,9 @@ func (c *workflowCompiler) addContainerExecutorTemplate(name string, refName str
 		Name: nameContainerImpl,
 		Inputs: wfapi.Inputs{
 			Parameters: append([]wfapi.Parameter{
-				{Name: paramPodSpecPatch}},
-				c.getTaskRetryParameters(name)...),
+				{Name: paramPodSpecPatch},
+			},
+				c.getTaskRetryParameters(task)...),
 		},
 		// PodSpecPatch input param is where actual image, command and
 		// args come from. It is treated as a strategic merge patch on
@@ -530,34 +539,23 @@ func (c *workflowCompiler) addContainerExecutorTemplate(name string, refName str
 	return nameContainerExecutor
 }
 
-func (c *workflowCompiler) getTaskRetryPolicySpec(name string) *pipelinespec.PipelineTaskSpec_RetryPolicy {
-	if c.spec == nil || c.spec.Root == nil || c.spec.Root.GetDag() == nil {
-		return nil
-	}
-	rootDag := c.spec.Root.GetDag()
-	taskSpec := rootDag.Tasks[name]
-	if taskSpec == nil {
-		return nil
-	}
-	return taskSpec.RetryPolicy
-}
-
-func (c *workflowCompiler) getTaskRetryParameters(name string) []wfapi.Parameter {
-	retryPolicy := c.getTaskRetryPolicySpec(name)
-	if retryPolicy == nil {
+func (c *workflowCompiler) getTaskRetryParameters(task *pipelinespec.PipelineTaskSpec) []wfapi.Parameter {
+	if task != nil && task.RetryPolicy != nil {
+		// Create slice of parameters with "Name" field set.
+		parameters := []wfapi.Parameter{
+			{Name: paramRetryMaxCount},
+			{Name: paramRetryBackOffDuration},
+			{Name: paramRetryBackOffFactor},
+			{Name: paramRetryBackOffMaxDuration},
+		}
+		return parameters
+	} else {
 		return nil
 	}
-	// Create slice of parameters with "Name" field set.
-	parameters := []wfapi.Parameter{
-		{Name: paramRetryMaxCount},
-		{Name: paramRetryBackOffDuration},
-		{Name: paramRetryBackOffFactor},
-		{Name: paramRetryBackOffMaxDuration}}
-	return parameters
 }
 
-func (c *workflowCompiler) getTaskRetryParametersWithValues(name string) []wfapi.Parameter {
-	retryPolicy := c.getTaskRetryPolicySpec(name)
+func (c *workflowCompiler) getTaskRetryParametersWithValues(task *pipelinespec.PipelineTaskSpec) []wfapi.Parameter {
+	retryPolicy := task.GetRetryPolicy()
 	if retryPolicy == nil {
 		return []wfapi.Parameter{}
 	}
@@ -609,6 +607,26 @@ func (c *workflowCompiler) addParameterInputPath(parameters []wfapi.Parameter) [
 	return parameters
 }
 
+func (c *workflowCompiler) getTaskRetryStrategyFromInput(maxCount string, backOffDuration string, backOffFactor string,
+	backOffMaxDuration string) *wfapi.RetryStrategy {
+	backoff := &wfapi.Backoff{
+		Factor: &intstr.IntOrString{
+			Type:   intstr.String,
+			StrVal: backOffFactor,
+		},
+		MaxDuration: backOffMaxDuration,
+		Duration:    backOffDuration,
+	}
+
+	return &wfapi.RetryStrategy{
+		Limit: &intstr.IntOrString{
+			Type:   intstr.String,
+			StrVal: maxCount,
+		},
+		Backoff: backoff,
+	}
+}
+
 // Extends the PodMetadata to include Kubernetes-specific executor config.
 // Although the current podMetadata object is always empty, this function
 // doesn't overwrite the existing podMetadata because for security reasons
@@ -638,26 +656,6 @@ func extendPodMetadata(
 	}
 }
 
-func (c *workflowCompiler) getTaskRetryStrategyFromInput(maxCount string, backOffDuration string, backOffFactor string,
-	backOffMaxDuration string) *wfapi.RetryStrategy {
-	backoff := &wfapi.Backoff{
-		Factor: &intstr.IntOrString{
-			Type:   intstr.String,
-			StrVal: backOffFactor,
-		},
-		MaxDuration: backOffMaxDuration,
-		Duration:    backOffDuration,
-	}
-
-	return &wfapi.RetryStrategy{
-		Limit: &intstr.IntOrString{
-			Type:   intstr.String,
-			StrVal: maxCount,
-		},
-		Backoff: backoff,
-	}
-}
-
 // Extends metadata map values, highPriorityMap should overwrites lowPriorityMap values
 // The original Map inputs should have higher priority since its defined by admin
 // TODO: Use maps.Copy after moving to go 1.21+

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

@@ -18,6 +18,8 @@ import (
 	"os"
 	"testing"
 
+	"github.com/kubeflow/pipelines/api/v2alpha1/go/pipelinespec"
+
 	wfapi "github.com/argoproj/argo-workflows/v3/pkg/apis/workflow/v1alpha1"
 	"github.com/kubeflow/pipelines/kubernetes_platform/go/kubernetesplatform"
 	"github.com/stretchr/testify/assert"
@@ -59,7 +61,7 @@ func TestAddContainerExecutorTemplate(t *testing.T) {
 				},
 			}
 
-			c.addContainerExecutorTemplate("test-ref", "comp-test-ref")
+			c.addContainerExecutorTemplate(&pipelinespec.PipelineTaskSpec{ComponentRef: &pipelinespec.ComponentRef{Name: "comp-test-ref"}})
 			assert.NotEmpty(t, "system-container-impl", "Template name should not be empty")
 
 			executorTemplate, exists := c.templates["system-container-impl"]

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

@@ -311,7 +311,7 @@ func (c *workflowCompiler) task(name string, task *pipelinespec.PipelineTaskSpec
 				condition:       driverOutputs.condition,
 				exitTemplate:    inputs.exitTemplate,
 				hookParentDagID: inputs.parentDagID,
-			}, task.GetComponentRef().GetName())
+			}, task)
 			executor.Depends = depends([]string{driverTaskName})
 			return []wfapi.DAGTask{*driver, *executor}, nil
 		case *pipelinespec.PipelineDeploymentConfig_ExecutorSpec_Importer: