-
Notifications
You must be signed in to change notification settings - Fork 673
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge branch 'master' of github.com:flyteorg/flyte into failure_node
- Loading branch information
Showing
12 changed files
with
545 additions
and
376 deletions.
There are no files selected for viewing
35 changes: 0 additions & 35 deletions
35
flyteplugins/go/tasks/pluginmachinery/flytek8s/non_interruptible.go
This file was deleted.
Oops, something went wrong.
123 changes: 123 additions & 0 deletions
123
flyteplugins/go/tasks/pluginmachinery/flytek8s/plugin_exec_context.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,123 @@ | ||
package flytek8s | ||
|
||
import ( | ||
v1 "k8s.io/api/core/v1" | ||
|
||
"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core" | ||
pluginsCore "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/core" | ||
) | ||
|
||
type pluginTaskOverrides struct { | ||
pluginsCore.TaskOverrides | ||
resources *v1.ResourceRequirements | ||
extendedResources *core.ExtendedResources | ||
} | ||
|
||
func (to *pluginTaskOverrides) GetResources() *v1.ResourceRequirements { | ||
if to.resources != nil { | ||
return to.resources | ||
} | ||
return to.TaskOverrides.GetResources() | ||
} | ||
|
||
func (to *pluginTaskOverrides) GetExtendedResources() *core.ExtendedResources { | ||
if to.extendedResources != nil { | ||
return to.extendedResources | ||
} | ||
return to.TaskOverrides.GetExtendedResources() | ||
} | ||
|
||
type pluginTaskExecutionMetadata struct { | ||
pluginsCore.TaskExecutionMetadata | ||
interruptible *bool | ||
overrides *pluginTaskOverrides | ||
} | ||
|
||
func (tm *pluginTaskExecutionMetadata) IsInterruptible() bool { | ||
if tm.interruptible != nil { | ||
return *tm.interruptible | ||
} | ||
return tm.TaskExecutionMetadata.IsInterruptible() | ||
} | ||
|
||
func (tm *pluginTaskExecutionMetadata) GetOverrides() pluginsCore.TaskOverrides { | ||
if tm.overrides != nil { | ||
return tm.overrides | ||
} | ||
return tm.TaskExecutionMetadata.GetOverrides() | ||
} | ||
|
||
type pluginTaskExecutionContext struct { | ||
pluginsCore.TaskExecutionContext | ||
metadata *pluginTaskExecutionMetadata | ||
} | ||
|
||
func (tc *pluginTaskExecutionContext) TaskExecutionMetadata() pluginsCore.TaskExecutionMetadata { | ||
if tc.metadata != nil { | ||
return tc.metadata | ||
} | ||
return tc.TaskExecutionContext.TaskExecutionMetadata() | ||
} | ||
|
||
type PluginTaskExecutionContextOption func(*pluginTaskExecutionContext) | ||
|
||
func WithInterruptible(v bool) PluginTaskExecutionContextOption { | ||
return func(tc *pluginTaskExecutionContext) { | ||
if tc.metadata == nil { | ||
tc.metadata = &pluginTaskExecutionMetadata{ | ||
TaskExecutionMetadata: tc.TaskExecutionContext.TaskExecutionMetadata(), | ||
} | ||
} | ||
tc.metadata.interruptible = &v | ||
} | ||
} | ||
|
||
func WithResources(r *v1.ResourceRequirements) PluginTaskExecutionContextOption { | ||
return func(tc *pluginTaskExecutionContext) { | ||
if tc.metadata == nil { | ||
tc.metadata = &pluginTaskExecutionMetadata{ | ||
TaskExecutionMetadata: tc.TaskExecutionContext.TaskExecutionMetadata(), | ||
} | ||
} | ||
if tc.metadata.overrides == nil { | ||
tc.metadata.overrides = &pluginTaskOverrides{ | ||
TaskOverrides: tc.metadata.TaskExecutionMetadata.GetOverrides(), | ||
} | ||
} | ||
tc.metadata.overrides.resources = r | ||
} | ||
} | ||
|
||
func WithExtendedResources(er *core.ExtendedResources) PluginTaskExecutionContextOption { | ||
return func(tc *pluginTaskExecutionContext) { | ||
if tc.metadata == nil { | ||
tc.metadata = &pluginTaskExecutionMetadata{ | ||
TaskExecutionMetadata: tc.TaskExecutionContext.TaskExecutionMetadata(), | ||
} | ||
} | ||
if tc.metadata.overrides == nil { | ||
tc.metadata.overrides = &pluginTaskOverrides{ | ||
TaskOverrides: tc.metadata.TaskExecutionMetadata.GetOverrides(), | ||
} | ||
} | ||
tc.metadata.overrides.extendedResources = er | ||
} | ||
} | ||
|
||
func NewPluginTaskExecutionContext(tc pluginsCore.TaskExecutionContext, options ...PluginTaskExecutionContextOption) pluginsCore.TaskExecutionContext { | ||
tm := tc.TaskExecutionMetadata() | ||
to := tm.GetOverrides() | ||
ctx := &pluginTaskExecutionContext{ | ||
TaskExecutionContext: tc, | ||
metadata: &pluginTaskExecutionMetadata{ | ||
TaskExecutionMetadata: tm, | ||
overrides: &pluginTaskOverrides{ | ||
TaskOverrides: to, | ||
}, | ||
}, | ||
} | ||
for _, o := range options { | ||
o(ctx) | ||
} | ||
return ctx | ||
} |
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
Oops, something went wrong.