diff --git a/cmd/controller/main.go b/cmd/controller/main.go index 1494491e0..89bee3e9d 100644 --- a/cmd/controller/main.go +++ b/cmd/controller/main.go @@ -9,6 +9,7 @@ import ( _ "github.com/flyteorg/flyteplugins/go/tasks/plugins/k8s/sagemaker" _ "github.com/flyteorg/flyteplugins/go/tasks/plugins/k8s/sidecar" _ "github.com/flyteorg/flyteplugins/go/tasks/plugins/k8s/spark" + _ "github.com/flyteorg/flytepropeller/pkg/controller/nodes/task" "github.com/flyteorg/flytestdlib/contextutils" "github.com/flyteorg/flytestdlib/promutils/labeled" diff --git a/config.yaml b/config.yaml index 01bb7501b..1eed2d781 100644 --- a/config.yaml +++ b/config.yaml @@ -33,6 +33,7 @@ tasks: - container - sidecar - K8S-ARRAY + - sleep # Uncomment to enable sagemaker plugin # - sagemaker_training # - sagemaker_hyperparameter_tuning diff --git a/pkg/controller/nodes/task/sleeper_plugin.go b/pkg/controller/nodes/task/sleeper_plugin.go new file mode 100644 index 000000000..95949b7f7 --- /dev/null +++ b/pkg/controller/nodes/task/sleeper_plugin.go @@ -0,0 +1,61 @@ +package task + +import ( + "context" + pluginMachinery "github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery" + pluginCore "github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/core" + "github.com/flyteorg/flytestdlib/logger" + "time" +) + +type SleeperPlugin struct { + +} + +func (s SleeperPlugin) GetID() string { + return "sleep" +} + +func (s SleeperPlugin) GetProperties() pluginCore.PluginProperties { + return pluginCore.PluginProperties{} +} + +func (s SleeperPlugin) Handle(ctx context.Context, tCtx pluginCore.TaskExecutionContext) (pluginCore.Transition, error) { + logger.Infof(ctx, "Sleeper plugin invoked!") + tk, err := tCtx.TaskReader().Read(ctx) + if err != nil { + return pluginCore.UnknownTransition, err + } + sleepTime := time.Millisecond * 1000 + if tk.GetConfig() != nil { + v, ok := tk.GetConfig()["sleep"] + if ok { + i, err := time.ParseDuration(v) + if err == nil { + sleepTime = i + } + } + } + logger.Infof(ctx, "Sleeping for %v", sleepTime) + time.Sleep(sleepTime) + return pluginCore.DoTransition(pluginCore.PhaseInfoSuccess(nil)), nil +} + +func (s SleeperPlugin) Abort(ctx context.Context, tCtx pluginCore.TaskExecutionContext) error { + return nil +} + +func (s SleeperPlugin) Finalize(ctx context.Context, tCtx pluginCore.TaskExecutionContext) error { + return nil +} + +func init() { + pluginMachinery.PluginRegistry().RegisterCorePlugin(pluginCore.PluginEntry{ + ID: "sleep", + RegisteredTaskTypes: []pluginCore.TaskType{"sleep"}, + LoadPlugin: func(ctx context.Context, iCtx pluginCore.SetupContext) (pluginCore.Plugin, error) { + return SleeperPlugin{}, nil + }, + DefaultForTaskTypes: []pluginCore.TaskType{"sleep"}, + }) +}