[wip] Perf test inmemory plugin - sleeper

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"
 

config.yaml

@@ -33,6 +33,7 @@ tasks:
       - container
       - sidecar
       - K8S-ARRAY
+      - sleep
     # Uncomment to enable sagemaker plugin
     #      - sagemaker_training
     #      - sagemaker_hyperparameter_tuning

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"},
+	})
+}