From 45be92aba2c0cf6d04b1c26c46b0a159cec64f77 Mon Sep 17 00:00:00 2001 From: Katrina Rogan Date: Wed, 21 Oct 2020 11:08:14 -0700 Subject: [PATCH] Introduce default plugin for task type (#121) --- flyteplugins/go/tasks/pluginmachinery/core/plugin.go | 4 ++++ flyteplugins/go/tasks/pluginmachinery/k8s/plugin.go | 4 ++++ flyteplugins/go/tasks/plugins/array/awsbatch/executor.go | 1 + flyteplugins/go/tasks/plugins/hive/executor.go | 1 + flyteplugins/go/tasks/plugins/k8s/container/container.go | 1 + .../go/tasks/plugins/k8s/kfoperators/pytorch/pytorch.go | 1 + flyteplugins/go/tasks/plugins/k8s/sagemaker/plugin.go | 3 +++ flyteplugins/go/tasks/plugins/k8s/sidecar/sidecar.go | 1 + flyteplugins/go/tasks/plugins/k8s/spark/spark.go | 1 + flyteplugins/go/tasks/plugins/presto/executor.go | 1 + 10 files changed, 18 insertions(+) diff --git a/flyteplugins/go/tasks/pluginmachinery/core/plugin.go b/flyteplugins/go/tasks/pluginmachinery/core/plugin.go index 5f7d05a269..a4314816d3 100644 --- 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 index 6a198ae0f4..81ce95ee32 100644 --- 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 index 6bbd93c1db..a7d5df10d8 100644 --- 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 index 556a301ab0..dffb596ec9 100644 --- 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 index 06646c50d9..c9c564815c 100755 --- 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 index 98f0b06f96..6b690cee70 100644 --- 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 index 4e87eb6657..af3e9cf884 100644 --- 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 index 462afe7c9a..83eddede16 100755 --- 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 index 37def6798a..c91b196694 100755 --- 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 index 3d34a6b99a..15ffe38a4b 100644 --- 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}, }) }