Respect original task definition retry strategy for single task execu…

…tions Signed-off-by: Katrina Rogan <[email protected]>
flyteorg · Jul 22, 2024 · 947331a · 947331a
1 parent 337088e
commit 947331a
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/flyteadmin/pkg/manager/impl/util/single_task_execution.go b/flyteadmin/pkg/manager/impl/util/single_task_execution.go
@@ -85,6 +85,13 @@ func CreateOrGetWorkflowModel(
 		Version: workflowIdentifier.Version,
 	})
 
+	var retryStrategy *core.RetryStrategy
+	if task.GetClosure().GetCompiledTask().GetTemplate().GetMetadata().GetRetries().GetRetries() > 0 {
+		retryStrategy = task.GetClosure().GetCompiledTask().GetTemplate().GetMetadata().GetRetries()
+	} else {
+		retryStrategy = &defaultRetryStrategy
+	}
+
 	if err != nil {
 		if ferr, ok := err.(errors.FlyteAdminError); !ok || ferr.Code() != codes.NotFound {
 			return nil, err
@@ -99,7 +106,7 @@ func CreateOrGetWorkflowModel(
 						Id: generateNodeNameFromTask(taskIdentifier.Name),
 						Metadata: &core.NodeMetadata{
 							Name:    generateNodeNameFromTask(taskIdentifier.Name),
-							Retries: &defaultRetryStrategy,
+							Retries: retryStrategy,
 						},
 						Inputs: generateBindings(*task.Closure.CompiledTask.Template.Interface.Inputs, noInputNodeID),
 						Target: &core.Node_TaskNode{

diff --git a/flyteadmin/pkg/manager/impl/util/single_task_execution_test.go b/flyteadmin/pkg/manager/impl/util/single_task_execution_test.go
@@ -109,6 +109,9 @@ func TestCreateOrGetWorkflowModel(t *testing.T) {
 			Name:         ".flytegen.app.workflows.MyWorkflow.my_task",
 			Version:      "12345",
 		}), fmt.Sprintf("%+v", request.Id))
+		assert.Len(t, request.GetSpec().GetTemplate().GetNodes(), 1)
+		assert.Equal(t, request.GetSpec().GetTemplate().GetNodes()[0].GetMetadata().GetRetries().GetRetries(), uint32(2))
+
 		return &admin.WorkflowCreateResponse{}, nil
 	})
 	taskIdentifier := &core.Identifier{
@@ -148,6 +151,11 @@ func TestCreateOrGetWorkflowModel(t *testing.T) {
 							},
 						},
 					},
+					Metadata: &core.TaskMetadata{
+						Retries: &core.RetryStrategy{
+							Retries: 2,
+						},
+					},
 				},
 			},
 		},