Skip to content
This repository was archived by the owner on Oct 9, 2023. It is now read-only.

Commit d5295d2

Browse files
pingsutwpingsutw
authored
The status of the AWS batch job should become failed once the retry limit exceeded (#291)
* Turn PhaseRetryableFailure into PhaseRetryLimitExceededFailure Signed-off-by: Kevin Su <pingsutw@apache.org> * nit Signed-off-by: Kevin Su <pingsutw@apache.org> * update Signed-off-by: Kevin Su <pingsutw@apache.org> * update test Signed-off-by: Kevin Su <pingsutw@apache.org> * lint Signed-off-by: Kevin Su <pingsutw@apache.org> * update Signed-off-by: Kevin Su <pingsutw@apache.org> * update tests Signed-off-by: Kevin Su <pingsutw@apache.org> * lint Signed-off-by: Kevin Su <pingsutw@apache.org> * wip Signed-off-by: Kevin Su <pingsutw@apache.org> * udpate Signed-off-by: Kevin Su <pingsutw@apache.org> * address comment Signed-off-by: Kevin Su <pingsutw@apache.org> * nit Signed-off-by: Kevin Su <pingsutw@apache.org> * fix tests Signed-off-by: Kevin Su <pingsutw@apache.org> * nit Signed-off-by: Kevin Su <pingsutw@apache.org> Signed-off-by: Kevin Su <pingsutw@apache.org>
1 parent b0f20e8 commit d5295d2

4 files changed

Lines changed: 115 additions & 19 deletions

File tree

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

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -80,9 +80,7 @@ func (e Executor) Handle(ctx context.Context, tCtx core.TaskExecutionContext) (c
8080
pluginState, err = LaunchSubTasks(ctx, tCtx, e.jobStore, pluginConfig, pluginState, e.metrics)
8181

8282
case arrayCore.PhaseCheckingSubTaskExecutions:
83-
pluginState, err = CheckSubTasksState(ctx, tCtx.TaskExecutionMetadata(),
84-
tCtx.OutputWriter().GetOutputPrefixPath(), tCtx.OutputWriter().GetRawOutputPrefix(),
85-
e.jobStore, tCtx.DataStore(), pluginConfig, pluginState, e.metrics)
83+
pluginState, err = CheckSubTasksState(ctx, tCtx, e.jobStore, pluginConfig, pluginState, e.metrics)
8684

8785
case arrayCore.PhaseAssembleFinalOutput:
8886
pluginState.State, err = array.AssembleFinalOutputs(ctx, e.outputAssembler, tCtx, arrayCore.PhaseSuccess, version, pluginState.State)

‎go/tasks/plugins/array/awsbatch/monitor.go‎

Lines changed: 22 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,8 @@ import (
44
"context"
55

66
core2 "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/core"
7-
"github.com/flyteorg/flytestdlib/storage"
87

8+
"github.com/flyteorg/flyteplugins/go/tasks/errors"
99
"github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/io"
1010
"github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/ioutils"
1111
"github.com/flyteorg/flyteplugins/go/tasks/plugins/array"
@@ -34,19 +34,32 @@ func createSubJobList(count int) []*Job {
3434
return res
3535
}
3636

37-
func CheckSubTasksState(ctx context.Context, taskMeta core.TaskExecutionMetadata, outputPrefix, baseOutputSandbox storage.DataReference, jobStore *JobStore,
38-
dataStore *storage.DataStore, cfg *config.Config, currentState *State, metrics ExecutorMetrics) (newState *State, err error) {
37+
func CheckSubTasksState(ctx context.Context, tCtx core.TaskExecutionContext, jobStore *JobStore,
38+
cfg *config.Config, currentState *State, metrics ExecutorMetrics) (newState *State, err error) {
3939
newState = currentState
4040
parentState := currentState.State
41-
jobName := taskMeta.GetTaskExecutionID().GetGeneratedName()
41+
jobName := tCtx.TaskExecutionMetadata().GetTaskExecutionID().GetGeneratedName()
4242
job := jobStore.Get(jobName)
43+
outputPrefix := tCtx.OutputWriter().GetOutputPrefixPath()
44+
baseOutputSandbox := tCtx.OutputWriter().GetRawOutputPrefix()
45+
dataStore := tCtx.DataStore()
46+
// Check that the taskTemplate is valid
47+
var taskTemplate *core2.TaskTemplate
48+
taskTemplate, err = tCtx.TaskReader().Read(ctx)
49+
if err != nil {
50+
return nil, errors.Wrapf(errors.CorruptedPluginState, err, "Failed to read task template")
51+
} else if taskTemplate == nil {
52+
return nil, errors.Errorf(errors.BadTaskSpecification, "Required value not set, taskTemplate is nil")
53+
}
54+
retry := toRetryStrategy(ctx, toBackoffLimit(taskTemplate.Metadata), cfg.MinRetries, cfg.MaxRetries)
55+
4356
// If job isn't currently being monitored (recovering from a restart?), add it to the sync-cache and return
4457
if job == nil {
4558
logger.Info(ctx, "Job not found in cache, adding it. [%v]", jobName)
4659

4760
_, err = jobStore.GetOrCreate(jobName, &Job{
4861
ID: *currentState.ExternalJobID,
49-
OwnerReference: taskMeta.GetOwnerID(),
62+
OwnerReference: tCtx.TaskExecutionMetadata().GetOwnerID(),
5063
SubJobs: createSubJobList(currentState.GetExecutionArraySize()),
5164
})
5265

@@ -108,6 +121,10 @@ func CheckSubTasksState(ctx context.Context, taskMeta core.TaskExecutionMetadata
108121
} else {
109122
msg.Collect(childIdx, "Job failed")
110123
}
124+
125+
if subJob.Status.Phase == core.PhaseRetryableFailure && *retry.Attempts == int64(len(subJob.Attempts)) {
126+
actualPhase = core.PhasePermanentFailure
127+
}
111128
} else if subJob.Status.Phase.IsSuccess() {
112129
actualPhase, err = array.CheckTaskOutput(ctx, dataStore, outputPrefix, baseOutputSandbox, childIdx, originalIdx)
113130
if err != nil {

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

Lines changed: 76 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@ package awsbatch
33
import (
44
"testing"
55

6+
"github.com/stretchr/testify/mock"
7+
68
"github.com/flyteorg/flytestdlib/contextutils"
79
"github.com/flyteorg/flytestdlib/promutils/labeled"
810

@@ -11,6 +13,7 @@ import (
1113

1214
"github.com/flyteorg/flyteplugins/go/tasks/plugins/array/arraystatus"
1315

16+
flyteIdl "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/core"
1417
"github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/core"
1518

1619
"github.com/aws/aws-sdk-go/aws/request"
@@ -19,6 +22,7 @@ import (
1922
arrayCore "github.com/flyteorg/flyteplugins/go/tasks/plugins/array/core"
2023

2124
"github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/core/mocks"
25+
ioMocks "github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/io/mocks"
2226
"github.com/flyteorg/flyteplugins/go/tasks/plugins/array/awsbatch/config"
2327
batchMocks "github.com/flyteorg/flyteplugins/go/tasks/plugins/array/awsbatch/mocks"
2428
"github.com/flyteorg/flytestdlib/utils"
@@ -35,15 +39,39 @@ func init() {
3539

3640
func TestCheckSubTasksState(t *testing.T) {
3741
ctx := context.Background()
42+
tCtx := &mocks.TaskExecutionContext{}
3843
tID := &mocks.TaskExecutionID{}
3944
tID.OnGetGeneratedName().Return("generated-name")
40-
4145
tMeta := &mocks.TaskExecutionMetadata{}
4246
tMeta.OnGetOwnerID().Return(types.NamespacedName{
4347
Namespace: "domain",
4448
Name: "name",
4549
})
4650
tMeta.OnGetTaskExecutionID().Return(tID)
51+
inMemDatastore, err := storage.NewDataStore(&storage.Config{Type: storage.TypeMemory}, promutils.NewTestScope())
52+
assert.NoError(t, err)
53+
54+
outputWriter := &ioMocks.OutputWriter{}
55+
outputWriter.OnGetOutputPrefixPath().Return("")
56+
outputWriter.OnGetRawOutputPrefix().Return("")
57+
58+
taskReader := &mocks.TaskReader{}
59+
task := &flyteIdl.TaskTemplate{
60+
Type: "test",
61+
Target: &flyteIdl.TaskTemplate_Container{
62+
Container: &flyteIdl.Container{
63+
Command: []string{"command"},
64+
Args: []string{"{{.Input}}"},
65+
},
66+
},
67+
Metadata: &flyteIdl.TaskMetadata{Retries: &flyteIdl.RetryStrategy{Retries: 3}},
68+
}
69+
taskReader.On("Read", mock.Anything).Return(task, nil)
70+
71+
tCtx.OnOutputWriter().Return(outputWriter)
72+
tCtx.OnTaskReader().Return(taskReader)
73+
tCtx.OnDataStore().Return(inMemDatastore)
74+
tCtx.OnTaskExecutionMetadata().Return(tMeta)
4775

4876
t.Run("Not in cache", func(t *testing.T) {
4977
mBatchClient := batchMocks.NewMockAwsBatchClient()
@@ -52,7 +80,7 @@ func TestCheckSubTasksState(t *testing.T) {
5280
utils.NewRateLimiter("", 10, 20))
5381

5482
jobStore := newJobsStore(t, batchClient)
55-
newState, err := CheckSubTasksState(ctx, tMeta, "", "", jobStore, nil, &config.Config{}, &State{
83+
newState, err := CheckSubTasksState(ctx, tCtx, jobStore, &config.Config{}, &State{
5684
State: &arrayCore.State{
5785
CurrentPhase: arrayCore.PhaseCheckingSubTaskExecutions,
5886
ExecutionArraySize: 5,
@@ -98,7 +126,7 @@ func TestCheckSubTasksState(t *testing.T) {
98126

99127
assert.NoError(t, err)
100128

101-
newState, err := CheckSubTasksState(ctx, tMeta, "", "", jobStore, nil, &config.Config{}, &State{
129+
newState, err := CheckSubTasksState(ctx, tCtx, jobStore, &config.Config{}, &State{
102130
State: &arrayCore.State{
103131
CurrentPhase: arrayCore.PhaseCheckingSubTaskExecutions,
104132
ExecutionArraySize: 5,
@@ -133,13 +161,10 @@ func TestCheckSubTasksState(t *testing.T) {
133161

134162
assert.NoError(t, err)
135163

136-
inMemDatastore, err := storage.NewDataStore(&storage.Config{Type: storage.TypeMemory}, promutils.NewTestScope())
137-
assert.NoError(t, err)
138-
139164
retryAttemptsArray, err := bitarray.NewCompactArray(1, bitarray.Item(1))
140165
assert.NoError(t, err)
141166

142-
newState, err := CheckSubTasksState(ctx, tMeta, "", "", jobStore, inMemDatastore, &config.Config{}, &State{
167+
newState, err := CheckSubTasksState(ctx, tCtx, jobStore, &config.Config{}, &State{
143168
State: &arrayCore.State{
144169
CurrentPhase: arrayCore.PhaseCheckingSubTaskExecutions,
145170
ExecutionArraySize: 1,
@@ -181,13 +206,10 @@ func TestCheckSubTasksState(t *testing.T) {
181206

182207
assert.NoError(t, err)
183208

184-
inMemDatastore, err := storage.NewDataStore(&storage.Config{Type: storage.TypeMemory}, promutils.NewTestScope())
185-
assert.NoError(t, err)
186-
187209
retryAttemptsArray, err := bitarray.NewCompactArray(2, bitarray.Item(1))
188210
assert.NoError(t, err)
189211

190-
newState, err := CheckSubTasksState(ctx, tMeta, "", "", jobStore, inMemDatastore, &config.Config{}, &State{
212+
newState, err := CheckSubTasksState(ctx, tCtx, jobStore, &config.Config{}, &State{
191213
State: &arrayCore.State{
192214
CurrentPhase: arrayCore.PhaseCheckingSubTaskExecutions,
193215
ExecutionArraySize: 2,
@@ -206,6 +228,49 @@ func TestCheckSubTasksState(t *testing.T) {
206228
assert.NoError(t, err)
207229
p, _ := newState.GetPhase()
208230
assert.Equal(t, arrayCore.PhaseCheckingSubTaskExecutions.String(), p.String())
231+
})
232+
233+
t.Run("retry limit exceeded", func(t *testing.T) {
234+
mBatchClient := batchMocks.NewMockAwsBatchClient()
235+
batchClient := NewCustomBatchClient(mBatchClient, "", "",
236+
utils.NewRateLimiter("", 10, 20),
237+
utils.NewRateLimiter("", 10, 20))
238+
239+
jobStore := newJobsStore(t, batchClient)
240+
_, err := jobStore.GetOrCreate(tID.GetGeneratedName(), &Job{
241+
ID: "job-id",
242+
Status: JobStatus{
243+
Phase: core.PhaseRunning,
244+
},
245+
SubJobs: []*Job{
246+
{Status: JobStatus{Phase: core.PhaseRetryableFailure}, Attempts: []Attempt{{LogStream: "failed"}}},
247+
{Status: JobStatus{Phase: core.PhaseSuccess}},
248+
},
249+
})
250+
251+
assert.NoError(t, err)
252+
253+
retryAttemptsArray, err := bitarray.NewCompactArray(2, bitarray.Item(1))
254+
assert.NoError(t, err)
209255

256+
newState, err := CheckSubTasksState(ctx, tCtx, jobStore, &config.Config{}, &State{
257+
State: &arrayCore.State{
258+
CurrentPhase: arrayCore.PhaseWriteToDiscoveryThenFail,
259+
ExecutionArraySize: 2,
260+
OriginalArraySize: 2,
261+
OriginalMinSuccesses: 2,
262+
ArrayStatus: arraystatus.ArrayStatus{
263+
Detailed: arrayCore.NewPhasesCompactArray(2),
264+
},
265+
IndexesToCache: bitarray.NewBitSet(2),
266+
RetryAttempts: retryAttemptsArray,
267+
},
268+
ExternalJobID: refStr("job-id"),
269+
JobDefinitionArn: "",
270+
}, getAwsBatchExecutorMetrics(promutils.NewTestScope()))
271+
272+
assert.NoError(t, err)
273+
p, _ := newState.GetPhase()
274+
assert.Equal(t, arrayCore.PhaseWriteToDiscoveryThenFail, p)
210275
})
211276
}

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -334,6 +334,22 @@ func TestSummaryToPhase(t *testing.T) {
334334
core.PhaseSuccess: 10,
335335
},
336336
},
337+
{
338+
"FailedToRetry",
339+
PhaseWriteToDiscoveryThenFail,
340+
map[core.Phase]int64{
341+
core.PhaseSuccess: 5,
342+
core.PhasePermanentFailure: 5,
343+
},
344+
},
345+
{
346+
"Retrying",
347+
PhaseCheckingSubTaskExecutions,
348+
map[core.Phase]int64{
349+
core.PhaseSuccess: 5,
350+
core.PhaseRetryableFailure: 5,
351+
},
352+
},
337353
}
338354

339355
for _, tt := range tests {

0 commit comments

Comments
 (0)