fix: Add retry logic with delay for fetching DAG tasks during TaskFinalStatus resolution

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

Author: khushiiagrawal

backend/src/v2/driver/resolve.go

@@ -20,6 +20,7 @@ import (
 	"errors"
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/golang/glog"
 	"github.com/kubeflow/pipelines/api/v2alpha1/go/pipelinespec"
@@ -440,17 +441,35 @@ func resolveInputParameter(
 			return nil, paramError(fmt.Errorf("param runtime value spec of type %T not implemented", t))
 		}
 	case *pipelinespec.TaskInputsSpec_InputParameterSpec_TaskFinalStatus_:
-		tasks, err := getDAGTasks(ctx, dag, pipeline, mlmd, nil)
-		if err != nil {
-			return nil, err
+		// Retry logic for obtaining DAG tasks to handle potential race conditions or delays
+		// in MLMD propagation. We try up to 5 times with a 2-second sleep between attempts.
+		var tasks map[string]*metadata.Execution
+		for i := 0; i < 5; i++ {
+			var err error
+			tasks, err = getDAGTasks(ctx, dag, pipeline, mlmd, nil)
+			if err != nil {
+				return nil, err
+			}
+
+			if len(opts.Task.DependentTasks) < 1 {
+				return nil, fmt.Errorf("task %v has no dependent tasks", opts.Task.TaskInfo.GetName())
+			}
+
+			producerName := metadata.GetTaskNameWithDagID(opts.Task.DependentTasks[0], dag.Execution.GetID())
+			if _, ok := tasks[producerName]; ok {
+				break
+			}
+			glog.Warningf("Unable to obtain the node for %s, taskName %s. Retrying...", producerName, opts.TaskName)
+			time.Sleep(2 * time.Second)
 		}
 
 		if len(opts.Task.DependentTasks) < 1 {
 			return nil, fmt.Errorf("task %v has no dependent tasks", opts.Task.TaskInfo.GetName())
 		}
-		producer, ok := tasks[metadata.GetTaskNameWithDagID(opts.Task.DependentTasks[0], dag.Execution.GetID())]
+		producerName := metadata.GetTaskNameWithDagID(opts.Task.DependentTasks[0], dag.Execution.GetID())
+		producer, ok := tasks[producerName]
 		if !ok {
-			return nil, fmt.Errorf("producer task, %v, not in tasks", producer.TaskName())
+			return nil, fmt.Errorf("producer task, %s, not in tasks", producerName) // Fixed potential panic
 		}
 		finalStatus := pipelinespec.PipelineTaskFinalStatus{
 			State:                   producer.GetExecution().GetLastKnownState().String(),