Introduce default plugin for task type (flyteorg#121)

eapolinario · Oct 21, 2020 · 45be92a · 45be92a
1 parent e5362df
commit 45be92a
Show file tree

Hide file tree

Showing 10 changed files with 18 additions and 0 deletions.
diff --git a/flyteplugins/go/tasks/pluginmachinery/core/plugin.go b/flyteplugins/go/tasks/pluginmachinery/core/plugin.go
@@ -24,6 +24,10 @@ type PluginEntry struct {
 	// Boolean that indicates if this plugin can be used as the default for unknown task types. There can only be
 	// one default in the system
 	IsDefault bool
+	// A list of all task types for which this plugin should be default handler when multiple registered plugins
+	// support the same task type. This must be a subset of RegisteredTaskTypes and at most one default per task type
+	// is supported.
+	DefaultForTaskTypes []TaskType
 }
 
 // System level properties that this Plugin supports

diff --git a/flyteplugins/go/tasks/pluginmachinery/k8s/plugin.go b/flyteplugins/go/tasks/pluginmachinery/k8s/plugin.go
@@ -31,6 +31,10 @@ type PluginEntry struct {
 	// Boolean that indicates if this plugin can be used as the default for unknown task types. There can only be
 	// one default in the system
 	IsDefault bool
+	// A list of all task types for which this plugin should be default handler when multiple registered plugins
+	// support the same task type. This must be a subset of RegisteredTaskTypes and at most one default per task type
+	// is supported.
+	DefaultForTaskTypes []pluginsCore.TaskType
 }
 
 // A proxy object for k8s resource

diff --git a/flyteplugins/go/tasks/plugins/array/awsbatch/executor.go b/flyteplugins/go/tasks/plugins/array/awsbatch/executor.go
@@ -192,6 +192,7 @@ func init() {
 			RegisteredTaskTypes: []core.TaskType{arrayTaskType},
 			LoadPlugin:          createNewExecutorPlugin,
 			IsDefault:           false,
+			DefaultForTaskTypes: []core.TaskType{arrayTaskType},
 		})
 }
 

diff --git a/flyteplugins/go/tasks/plugins/hive/executor.go b/flyteplugins/go/tasks/plugins/hive/executor.go
@@ -163,5 +163,6 @@ func init() {
 			RegisteredTaskTypes: []core.TaskType{hiveTaskType},
 			LoadPlugin:          QuboleHiveExecutorLoader,
 			IsDefault:           false,
+			DefaultForTaskTypes: []core.TaskType{hiveTaskType},
 		})
 }
diff --git a/flyteplugins/go/tasks/plugins/k8s/container/container.go b/flyteplugins/go/tasks/plugins/k8s/container/container.go
@@ -80,5 +80,6 @@ func init() {
 			ResourceToWatch:     &v1.Pod{},
 			Plugin:              Plugin{},
 			IsDefault:           true,
+			DefaultForTaskTypes: []pluginsCore.TaskType{containerTaskType},
 		})
 }
diff --git a/flyteplugins/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch.go b/flyteplugins/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch.go
@@ -220,5 +220,6 @@ func init() {
 			ResourceToWatch:     &ptOp.PyTorchJob{},
 			Plugin:              pytorchOperatorResourceHandler{},
 			IsDefault:           false,
+			DefaultForTaskTypes: []pluginsCore.TaskType{pytorchTaskType},
 		})
 }
diff --git a/flyteplugins/go/tasks/plugins/k8s/sagemaker/plugin.go b/flyteplugins/go/tasks/plugins/k8s/sagemaker/plugin.go
@@ -75,6 +75,7 @@ func init() {
 			ResourceToWatch:     &hpojobv1.HyperparameterTuningJob{},
 			Plugin:              awsSagemakerPlugin{TaskType: hyperparameterTuningJobTaskType},
 			IsDefault:           false,
+			DefaultForTaskTypes: []pluginsCore.TaskType{hyperparameterTuningJobTaskType},
 		})
 
 	// Registering the plugin for standalone TrainingJob
@@ -85,6 +86,7 @@ func init() {
 			ResourceToWatch:     &trainingjobv1.TrainingJob{},
 			Plugin:              awsSagemakerPlugin{TaskType: trainingJobTaskType},
 			IsDefault:           false,
+			DefaultForTaskTypes: []pluginsCore.TaskType{trainingJobTaskType},
 		})
 
 	// Registering the plugin for custom TrainingJob
@@ -95,5 +97,6 @@ func init() {
 			ResourceToWatch:     &trainingjobv1.TrainingJob{},
 			Plugin:              awsSagemakerPlugin{TaskType: customTrainingJobTaskType},
 			IsDefault:           false,
+			DefaultForTaskTypes: []pluginsCore.TaskType{customTrainingJobTaskType},
 		})
 }
diff --git a/flyteplugins/go/tasks/plugins/k8s/sidecar/sidecar.go b/flyteplugins/go/tasks/plugins/k8s/sidecar/sidecar.go
@@ -175,5 +175,6 @@ func init() {
 			ResourceToWatch:     &k8sv1.Pod{},
 			Plugin:              sidecarResourceHandler{},
 			IsDefault:           false,
+			DefaultForTaskTypes: []pluginsCore.TaskType{sidecarTaskType},
 		})
 }
diff --git a/flyteplugins/go/tasks/plugins/k8s/spark/spark.go b/flyteplugins/go/tasks/plugins/k8s/spark/spark.go
@@ -389,5 +389,6 @@ func init() {
 			ResourceToWatch:     &sparkOp.SparkApplication{},
 			Plugin:              sparkResourceHandler{},
 			IsDefault:           false,
+			DefaultForTaskTypes: []pluginsCore.TaskType{sparkTaskType},
 		})
 }
diff --git a/flyteplugins/go/tasks/plugins/presto/executor.go b/flyteplugins/go/tasks/plugins/presto/executor.go
@@ -153,5 +153,6 @@ func init() {
 			RegisteredTaskTypes: []core.TaskType{prestoTaskType},
 			LoadPlugin:          ExecutorLoader,
 			IsDefault:           false,
+			DefaultForTaskTypes: []core.TaskType{prestoTaskType},
 		})
 }