forked from flyteorg/flyte
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
sending all subtask info on map task start (flyteorg#252)
* added SubTaskMetadata type Signed-off-by: Daniel Rammer <[email protected]> * implemented SubTaskMetadata and integrated into k8s-array plugin Signed-off-by: Daniel Rammer <[email protected]> * using SubTaskMetadata in aws_batch plugin Signed-off-by: Daniel Rammer <[email protected]> * populating CacheStatus and Logs on ExternalResource Signed-off-by: Daniel Rammer <[email protected]> * removed SubTaskMetadata and instead directly using ExternalResource Signed-off-by: Daniel Rammer <[email protected]> * updated flyteidl version Signed-off-by: Daniel Rammer <[email protected]> * updated flyteidl Signed-off-by: Daniel Rammer <[email protected]> * updated flyteidl Signed-off-by: Daniel Rammer <[email protected]> * documentation Signed-off-by: Daniel Rammer <[email protected]> * documentation Signed-off-by: Daniel Rammer <[email protected]> * fixed unit tests Signed-off-by: Daniel Rammer <[email protected]> * fixed lint issue Signed-off-by: Daniel Rammer <[email protected]> * added unit test for InitializeExternalResources Signed-off-by: Daniel Rammer <[email protected]> * maintaing phase version on WaitingForResources phase change Signed-off-by: Daniel Rammer <[email protected]> * removing phase from log name Signed-off-by: Daniel Rammer <[email protected]> * fixed tests Signed-off-by: Daniel Rammer <[email protected]>
- Loading branch information
Showing
15 changed files
with
309 additions
and
179 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,66 @@ | ||
| package core | ||
|
|
||
| import ( | ||
| "context" | ||
|
|
||
| idlCore "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/core" | ||
|
|
||
| "github.com/flyteorg/flyteplugins/go/tasks/errors" | ||
| "github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/core" | ||
| ) | ||
|
|
||
| // InitializeExternalResources constructs an ExternalResource array where each element describes the | ||
| // initial state of the subtask. This involves labeling all cached subtasks as successful with a | ||
| // cache hit and initializing others to undefined state. | ||
| func InitializeExternalResources(ctx context.Context, tCtx core.TaskExecutionContext, state *State, | ||
| generateSubTaskID func(core.TaskExecutionContext, int) string) ([]*core.ExternalResource, error) { | ||
| externalResources := make([]*core.ExternalResource, state.GetOriginalArraySize()) | ||
|
|
||
| taskTemplate, err := tCtx.TaskReader().Read(ctx) | ||
| if err != nil { | ||
| return externalResources, err | ||
| } else if taskTemplate == nil { | ||
| return externalResources, errors.Errorf(errors.BadTaskSpecification, "Required value not set, taskTemplate is nil") | ||
| } | ||
|
|
||
| executeSubTaskCount := 0 | ||
| cachedSubTaskCount := 0 | ||
| for i := 0; i < int(state.GetOriginalArraySize()); i++ { | ||
| var cacheStatus idlCore.CatalogCacheStatus | ||
| var childIndex int | ||
| var phase core.Phase | ||
|
|
||
| if state.IndexesToCache.IsSet(uint(i)) { | ||
| // if not cached set to PhaseUndefined and set cacheStatus according to Discoverable | ||
| phase = core.PhaseUndefined | ||
| if taskTemplate.Metadata == nil || !taskTemplate.Metadata.Discoverable { | ||
| cacheStatus = idlCore.CatalogCacheStatus_CACHE_DISABLED | ||
| } else { | ||
| cacheStatus = idlCore.CatalogCacheStatus_CACHE_MISS | ||
| } | ||
|
|
||
| childIndex = executeSubTaskCount | ||
| executeSubTaskCount++ | ||
| } else { | ||
| // if cached set to PhaseSuccess and mark as CACHE_HIT | ||
| phase = core.PhaseSuccess | ||
| cacheStatus = idlCore.CatalogCacheStatus_CACHE_HIT | ||
|
|
||
| // child index is computed as a pseudo-value to ensure non-overlapping subTaskIDs | ||
| childIndex = state.GetExecutionArraySize() + cachedSubTaskCount | ||
| cachedSubTaskCount++ | ||
| } | ||
|
|
||
| subTaskID := generateSubTaskID(tCtx, childIndex) | ||
| externalResources[i] = &core.ExternalResource{ | ||
| ExternalID: subTaskID, | ||
| CacheStatus: cacheStatus, | ||
| Index: uint32(i), | ||
| Logs: nil, | ||
| RetryAttempt: 0, | ||
| Phase: phase, | ||
| } | ||
| } | ||
|
|
||
| return externalResources, nil | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,70 @@ | ||
| package core | ||
|
|
||
| import ( | ||
| "context" | ||
| "testing" | ||
|
|
||
| idlCore "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/core" | ||
|
|
||
| "github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/core" | ||
| "github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/core/mocks" | ||
|
|
||
| "github.com/flyteorg/flytestdlib/bitarray" | ||
|
|
||
| "github.com/stretchr/testify/assert" | ||
| ) | ||
|
|
||
| func TestInitializeExternalResources(t *testing.T) { | ||
| ctx := context.TODO() | ||
| subTaskCount := 10 | ||
| cachedCount := 4 | ||
|
|
||
| indexesToCache := InvertBitSet(bitarray.NewBitSet(uint(subTaskCount)), uint(subTaskCount)) | ||
| for i := 0; i < cachedCount; i++ { | ||
| indexesToCache.Clear(uint(i)) | ||
| } | ||
|
|
||
| tr := &mocks.TaskReader{} | ||
| tr.OnRead(ctx).Return(&idlCore.TaskTemplate{ | ||
| Metadata: &idlCore.TaskMetadata{ | ||
| Discoverable: true, | ||
| }, | ||
| }, nil) | ||
|
|
||
| tID := &mocks.TaskExecutionID{} | ||
| tID.OnGetGeneratedName().Return("notfound") | ||
|
|
||
| tMeta := &mocks.TaskExecutionMetadata{} | ||
| tMeta.OnGetTaskExecutionID().Return(tID) | ||
|
|
||
| tCtx := &mocks.TaskExecutionContext{} | ||
| tCtx.OnTaskReader().Return(tr) | ||
| tCtx.OnTaskExecutionMetadata().Return(tMeta) | ||
|
|
||
| state := State{ | ||
| OriginalArraySize: int64(subTaskCount), | ||
| ExecutionArraySize: subTaskCount - cachedCount, | ||
| IndexesToCache: indexesToCache, | ||
| } | ||
|
|
||
| externalResources, err := InitializeExternalResources(ctx, tCtx, &state, | ||
| func(_ core.TaskExecutionContext, i int) string { | ||
| return "" | ||
| }, | ||
| ) | ||
|
|
||
| assert.Nil(t, err) | ||
| assert.Equal(t, subTaskCount, len(externalResources)) | ||
| for i, externalResource := range externalResources { | ||
| assert.Equal(t, uint32(i), externalResource.Index) | ||
| assert.Equal(t, 0, len(externalResource.Logs)) | ||
| assert.Equal(t, uint32(0), externalResource.RetryAttempt) | ||
| if i < cachedCount { | ||
| assert.Equal(t, core.PhaseSuccess, externalResource.Phase) | ||
| assert.Equal(t, idlCore.CatalogCacheStatus_CACHE_HIT, externalResource.CacheStatus) | ||
| } else { | ||
| assert.Equal(t, core.PhaseUndefined, externalResource.Phase) | ||
| assert.Equal(t, idlCore.CatalogCacheStatus_CACHE_MISS, externalResource.CacheStatus) | ||
| } | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.