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.
Introduce WebAPI Plugin Interface, an example & AWS athena plugin (fl…
…yteorg#146) * Proposal: Common plugin interface for all service plugins * Add remote plugin interface and implemention * make properties config-friendly * Rename remote.PluginContext * Simplify ResourceMeta * New task log interface and template plugin * Update spark and pytorch plugins * Proposal: Common plugin interface for all service plugins Co-authored-by: Ketan Umare <[email protected]>
- Loading branch information
Showing
59 changed files
with
5,714 additions
and
550 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
Large diffs are not rendered by default.
Oops, something went wrong.
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
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.
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
93 changes: 93 additions & 0 deletions
93
go/tasks/pluginmachinery/internal/webapi/allocation_token.go
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,93 @@ | ||
| package webapi | ||
|
|
||
| import ( | ||
| "context" | ||
| "fmt" | ||
|
|
||
| "k8s.io/utils/clock" | ||
|
|
||
| "github.com/lyft/flytestdlib/logger" | ||
|
|
||
| "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/webapi" | ||
|
|
||
| "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/core" | ||
| ) | ||
|
|
||
| type tokenAllocator struct { | ||
| clock clock.Clock | ||
| } | ||
|
|
||
| func newTokenAllocator(c clock.Clock) tokenAllocator { | ||
| return tokenAllocator{ | ||
| clock: c, | ||
| } | ||
| } | ||
|
|
||
| func (a tokenAllocator) allocateToken(ctx context.Context, p webapi.AsyncPlugin, tCtx core.TaskExecutionContext, state *State, metrics Metrics) ( | ||
| newState *State, phaseInfo core.PhaseInfo, err error) { | ||
| if len(p.GetConfig().ResourceQuotas) == 0 { | ||
| // No quota, return success | ||
| return &State{ | ||
| AllocationTokenRequestStartTime: a.clock.Now(), | ||
| Phase: PhaseAllocationTokenAcquired, | ||
| }, core.PhaseInfoQueued(a.clock.Now(), 0, "No allocation token required"), nil | ||
| } | ||
|
|
||
| ns, constraints, err := p.ResourceRequirements(ctx, tCtx) | ||
| if err != nil { | ||
| logger.Errorf(ctx, "Failed to calculate resource requirements for task. Error: %v", err) | ||
| return nil, core.PhaseInfo{}, err | ||
| } | ||
|
|
||
| token := tCtx.TaskExecutionMetadata().GetTaskExecutionID().GetGeneratedName() | ||
| allocationStatus, err := tCtx.ResourceManager().AllocateResource(ctx, ns, token, constraints) | ||
| if err != nil { | ||
| logger.Errorf(ctx, "Failed to allocate resources for task. Error: %v", err) | ||
| return nil, core.PhaseInfo{}, err | ||
| } | ||
|
|
||
| switch allocationStatus { | ||
| case core.AllocationStatusGranted: | ||
| metrics.AllocationGranted.Inc(ctx) | ||
| metrics.ResourceWaitTime.Observe(float64(a.clock.Since(state.AllocationTokenRequestStartTime).Milliseconds())) | ||
| return &State{ | ||
| AllocationTokenRequestStartTime: a.clock.Now(), | ||
| Phase: PhaseAllocationTokenAcquired, | ||
| }, core.PhaseInfoQueued(a.clock.Now(), 0, "Allocation token required"), nil | ||
| case core.AllocationStatusNamespaceQuotaExceeded: | ||
| case core.AllocationStatusExhausted: | ||
| metrics.AllocationNotGranted.Inc(ctx) | ||
| logger.Infof(ctx, "Couldn't allocate token because allocation status is [%v].", allocationStatus.String()) | ||
| startTime := state.AllocationTokenRequestStartTime | ||
| if startTime.IsZero() { | ||
| startTime = a.clock.Now() | ||
| } | ||
|
|
||
| return &State{ | ||
| AllocationTokenRequestStartTime: startTime, | ||
| Phase: PhaseNotStarted, | ||
| }, core.PhaseInfoQueued(a.clock.Now(), 0, | ||
| fmt.Sprintf("Quota for task has exceeded. The request is enqueued.")), nil | ||
| } | ||
|
|
||
| return nil, core.PhaseInfo{}, fmt.Errorf("allocation status undefined [%v]", allocationStatus) | ||
| } | ||
|
|
||
| func (a tokenAllocator) releaseToken(ctx context.Context, p webapi.AsyncPlugin, tCtx core.TaskExecutionContext, metrics Metrics) error { | ||
| ns, _, err := p.ResourceRequirements(ctx, tCtx) | ||
| if err != nil { | ||
| logger.Errorf(ctx, "Failed to calculate resource requirements for task. Error: %v", err) | ||
| return err | ||
| } | ||
|
|
||
| token := tCtx.TaskExecutionMetadata().GetTaskExecutionID().GetGeneratedName() | ||
| err = tCtx.ResourceManager().ReleaseResource(ctx, ns, token) | ||
| if err != nil { | ||
| metrics.ResourceReleaseFailed.Inc(ctx) | ||
| logger.Errorf(ctx, "Failed to release resources for task. Error: %v", err) | ||
| return err | ||
| } | ||
|
|
||
| metrics.ResourceReleased.Inc(ctx) | ||
| return nil | ||
| } |
Oops, something went wrong.