Bug/abort map task subtasks

flyteplugins/go/tasks/pluginmachinery/core/phase.go

@@ -36,6 +36,8 @@
 	PhasePermanentFailure
 	// Indicates the task is waiting for the cache to be populated so it can reuse results
 	PhaseWaitingForCache
+	// Indicate the task has been aborted
+	PhaseAborted
 )
 
 var Phases = []Phase{
@@ -49,11 +51,12 @@
 	PhaseRetryableFailure,
 	PhasePermanentFailure,
 	PhaseWaitingForCache,
+	PhaseAborted,
 }
 
 // Returns true if the given phase is failure, retryable failure or success
 func (p Phase) IsTerminal() bool {
-	return p.IsFailure() || p.IsSuccess()
+	return p.IsFailure() || p.IsSuccess() || p.IsAborted()
 }
 
 func (p Phase) IsFailure() bool {
@@ -64,6 +67,10 @@
 	return p == PhaseSuccess
 }
 
+func (p Phase) IsAborted() bool {
+	return p == PhaseAborted
+}
+
 func (p Phase) IsWaitingForResources() bool {
 	return p == PhaseWaitingForResources
 }
@@ -257,10 +264,18 @@
 	return PhaseInfoFailed(PhasePermanentFailure, &core.ExecutionError{Code: code, Message: reason, Kind: core.ExecutionError_SYSTEM}, info)
 }
 
+func PhaseInfoSystemFailureWithCleanup(code, reason string, info *TaskInfo) PhaseInfo {
+	return phaseInfoFailed(PhasePermanentFailure, &core.ExecutionError{Code: code, Message: reason, Kind: core.ExecutionError_USER}, info, true)
+}
+
 func PhaseInfoFailure(code, reason string, info *TaskInfo) PhaseInfo {
 	return PhaseInfoFailed(PhasePermanentFailure, &core.ExecutionError{Code: code, Message: reason, Kind: core.ExecutionError_USER}, info)
 }
 
+func PhaseInfoFailureWithCleanup(code, reason string, info *TaskInfo) PhaseInfo {
+	return phaseInfoFailed(PhasePermanentFailure, &core.ExecutionError{Code: code, Message: reason, Kind: core.ExecutionError_USER}, info, true)
+}
+
 func PhaseInfoRetryableFailure(code, reason string, info *TaskInfo) PhaseInfo {
 	return PhaseInfoFailed(PhaseRetryableFailure, &core.ExecutionError{Code: code, Message: reason, Kind: core.ExecutionError_USER}, info)
 }

flyteplugins/go/tasks/plugins/array/awsbatch/executor.go

@@ -78,6 +78,9 @@
 
 	case arrayCore.PhaseAssembleFinalOutput:
 		pluginState.State, err = array.AssembleFinalOutputs(ctx, e.outputAssembler, tCtx, arrayCore.PhaseSuccess, version+1, pluginState.State)
+
+	case arrayCore.PhaseAbortSubTasks:
+		fallthrough
 
 	case arrayCore.PhaseWriteToDiscoveryThenFail:
 		pluginState.State, externalResources, err = array.WriteToDiscovery(ctx, tCtx, pluginState.State, arrayCore.PhaseAssembleFinalError, version+1)

flyteplugins/go/tasks/plugins/array/awsbatch/monitor_test.go

@@ -248,7 +248,7 @@ func TestCheckSubTasksState(t *testing.T) {
 
 		newState, err := CheckSubTasksState(ctx, tCtx, jobStore, &config.Config{}, &State{
 			State: &arrayCore.State{
-				CurrentPhase:         arrayCore.PhaseWriteToDiscoveryThenFail,
+				CurrentPhase:         arrayCore.PhaseCheckingSubTaskExecutions,
 				ExecutionArraySize:   2,
 				OriginalArraySize:    2,
 				OriginalMinSuccesses: 2,
@@ -264,6 +264,6 @@ func TestCheckSubTasksState(t *testing.T) {
 
 		assert.NoError(t, err)
 		p, _ := newState.GetPhase()
-		assert.Equal(t, arrayCore.PhaseWriteToDiscoveryThenFail, p)
+		assert.Equal(t, arrayCore.PhaseAbortSubTasks, p)
 	})
 }

flyteplugins/go/tasks/plugins/array/core/state.go

@@ -35,6 +35,7 @@ const (
 	PhaseAssembleFinalError
 	PhaseRetryableFailure
 	PhasePermanentFailure
+	PhaseAbortSubTasks
 )
 
 type State struct {
@@ -204,6 +205,9 @@ func MapArrayStateToPluginPhase(_ context.Context, state *State, logLinks []*idl
 	case PhaseAssembleFinalError:
 		fallthrough
 
+	case PhaseAbortSubTasks:
+		fallthrough
+
 	case PhaseWriteToDiscoveryThenFail:
 		fallthrough
 
@@ -259,14 +263,14 @@ func SummaryToPhase(ctx context.Context, minSuccesses int64, summary arraystatus
 
 	if totalCount < minSuccesses {
 		logger.Infof(ctx, "Array failed because totalCount[%v] < minSuccesses[%v]", totalCount, minSuccesses)
-		return PhaseWriteToDiscoveryThenFail
+		return PhaseAbortSubTasks
 	}
 
 	// No chance to reach the required success numbers.
 	if totalRunning+totalSuccesses+totalWaitingForResources+totalRetryableFailures < minSuccesses {
 		logger.Infof(ctx, "Array failed early because total failures > minSuccesses[%v]. Snapshot totalRunning[%v] + totalSuccesses[%v] + totalWaitingForResource[%v] + totalRetryableFailures[%v]",
 			minSuccesses, totalRunning, totalSuccesses, totalWaitingForResources, totalRetryableFailures)
-		return PhaseWriteToDiscoveryThenFail
+		return PhaseAbortSubTasks
 	}
 
 	if totalWaitingForResources > 0 {

flyteplugins/go/tasks/plugins/array/core/state_test.go

@@ -291,7 +291,7 @@ func TestSummaryToPhase(t *testing.T) {
 	}{
 		{
 			"FailOnTooFewTasks",
-			PhaseWriteToDiscoveryThenFail,
+			PhaseAbortSubTasks,
 			map[core.Phase]int64{},
 		},
 		{
@@ -304,7 +304,7 @@ func TestSummaryToPhase(t *testing.T) {
 		},
 		{
 			"FailOnTooManyPermanentFailures",
-			PhaseWriteToDiscoveryThenFail,
+			PhaseAbortSubTasks,
 			map[core.Phase]int64{
 				core.PhasePermanentFailure: 1,
 				core.PhaseSuccess:          9,
@@ -335,7 +335,7 @@ func TestSummaryToPhase(t *testing.T) {
 		},
 		{
 			"FailedToRetry",
-			PhaseWriteToDiscoveryThenFail,
+			PhaseAbortSubTasks,
 			map[core.Phase]int64{
 				core.PhaseSuccess:          5,
 				core.PhasePermanentFailure: 5,

flyteplugins/go/tasks/plugins/array/k8s/executor.go

@@ -117,6 +117,9 @@
 		nextState, externalResources, err = LaunchAndCheckSubTasksState(ctx, tCtx, e.kubeClient, pluginConfig,
 			tCtx.DataStore(), tCtx.OutputWriter().GetOutputPrefixPath(), tCtx.OutputWriter().GetRawOutputPrefix(), pluginState)
 
+	case arrayCore.PhaseAbortSubTasks:
+		nextState, externalResources, err = TerminateSubTasks(ctx, tCtx, e.kubeClient, GetConfig(), abortSubtask, pluginState)
+
 	case arrayCore.PhaseAssembleFinalOutput:
 		nextState, err = array.AssembleFinalOutputs(ctx, e.outputsAssembler, tCtx, arrayCore.PhaseSuccess, version+1, pluginState)
 
@@ -156,7 +159,7 @@
 		return errors.Wrapf(errors.CorruptedPluginState, err, "Failed to read unmarshal custom state")
 	}
 
-	return TerminateSubTasks(ctx, tCtx, e.kubeClient, GetConfig(), abortSubtask, pluginState)
+	return TerminateSubTasksOnAbort(ctx, tCtx, e.kubeClient, GetConfig(), abortSubtask, pluginState)
 }
 
 func (e Executor) Finalize(ctx context.Context, tCtx core.TaskExecutionContext) error {
@@ -165,7 +168,8 @@
 		return errors.Wrapf(errors.CorruptedPluginState, err, "Failed to read unmarshal custom state")
 	}
 
-	return TerminateSubTasks(ctx, tCtx, e.kubeClient, GetConfig(), finalizeSubtask, pluginState)
+	_, _, err := TerminateSubTasks(ctx, tCtx, e.kubeClient, GetConfig(), finalizeSubtask, pluginState)
+	return err
 }
 
 func (e Executor) Start(ctx context.Context) error {

flyteplugins/go/tasks/plugins/array/k8s/management.go

@@ -324,16 +324,37 @@
 	return newState, externalResources, nil
 }
 
+func TerminateSubTasksOnAbort(ctx context.Context, tCtx core.TaskExecutionContext, kubeClient core.KubeClient, config *Config,
+	terminateFunction func(context.Context, SubTaskExecutionContext, *Config, core.KubeClient) error, currentState *arrayCore.State) error {
+
+	_, externalResources, err := TerminateSubTasks(ctx, tCtx, kubeClient, GetConfig(), terminateFunction, currentState)
+	if err != nil {
+		return err
+	}
+
+	taskInfo := &core.TaskInfo{
+		ExternalResources: externalResources,
+	}
+	executionErr := &idlCore.ExecutionError{
+		Code:    "ArraySubtasksAborted",
+		Message: "Array subtasks were aborted",
+	}
+	phaseInfo := core.PhaseInfoFailed(core.PhaseAborted, executionErr, taskInfo)
+
+	return tCtx.EventsRecorder().RecordRaw(ctx, phaseInfo)
+}
+
 // TerminateSubTasks performs operations to gracefully terminate all subtasks. This may include
 // aborting and finalizing active k8s resources.
 func TerminateSubTasks(ctx context.Context, tCtx core.TaskExecutionContext, kubeClient core.KubeClient, config *Config,
-	terminateFunction func(context.Context, SubTaskExecutionContext, *Config, core.KubeClient) error, currentState *arrayCore.State) error {
+	terminateFunction func(context.Context, SubTaskExecutionContext, *Config, core.KubeClient) error, currentState *arrayCore.State) (*arrayCore.State, []*core.ExternalResource, error) {
 
 	taskTemplate, err := tCtx.TaskReader().Read(ctx)
+	externalResources := make([]*core.ExternalResource, 0, len(currentState.GetArrayStatus().Detailed.GetItems()))
 	if err != nil {
-		return err
+		return currentState, externalResources, err
 	} else if taskTemplate == nil {
-		return errors.Errorf(errors.BadTaskSpecification, "Required value not set, taskTemplate is nil")
+		return currentState, externalResources, errors.Errorf(errors.BadTaskSpecification, "Required value not set, taskTemplate is nil")
 	}
 
 	messageCollector := errorcollector.NewErrorMessageCollector()
@@ -353,18 +374,25 @@
 		originalIdx := arrayCore.CalculateOriginalIndex(childIdx, currentState.GetIndexesToCache())
 		stCtx, err := NewSubTaskExecutionContext(ctx, tCtx, taskTemplate, childIdx, originalIdx, retryAttempt, 0)
 		if err != nil {
-			return err
+			return currentState, externalResources, err
 		}
 
 		err = terminateFunction(ctx, stCtx, config, kubeClient)
 		if err != nil {
 			messageCollector.Collect(childIdx, err.Error())
+		} else {
+			externalResources = append(externalResources, &core.ExternalResource{
+				ExternalID:       stCtx.TaskExecutionMetadata().GetTaskExecutionID().GetGeneratedName(),
+				Index:            uint32(originalIdx),
+				RetryAttempt:     uint32(retryAttempt),
+				Phase:            core.PhaseAborted,
+			})
 		}
 	}
 
 	if messageCollector.Length() > 0 {
-		return fmt.Errorf(messageCollector.Summary(config.MaxErrorStringLength))
+		return currentState, externalResources, fmt.Errorf(messageCollector.Summary(config.MaxErrorStringLength))
 	}
 
-	return nil
+	return currentState.SetPhase(arrayCore.PhaseWriteToDiscoveryThenFail, currentState.PhaseVersion+1), externalResources, nil
 }