diff --git a/go/tasks/pluginmachinery/catalog/mocks/async_client.go b/go/tasks/pluginmachinery/catalog/mocks/async_client.go index 7ec246a001..46fc8c5eca 100644 --- a/go/tasks/pluginmachinery/catalog/mocks/async_client.go +++ b/go/tasks/pluginmachinery/catalog/mocks/async_client.go @@ -2,9 +2,13 @@ package mocks -import catalog "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/catalog" -import context "context" -import mock "github.com/stretchr/testify/mock" +import ( + context "context" + + catalog "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/catalog" + + mock "github.com/stretchr/testify/mock" +) // AsyncClient is an autogenerated mock type for the AsyncClient type type AsyncClient struct { diff --git a/go/tasks/pluginmachinery/catalog/mocks/client.go b/go/tasks/pluginmachinery/catalog/mocks/client.go index 04dcf7a3e5..d3bedd110d 100644 --- a/go/tasks/pluginmachinery/catalog/mocks/client.go +++ b/go/tasks/pluginmachinery/catalog/mocks/client.go @@ -2,10 +2,15 @@ package mocks -import catalog "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/catalog" -import context "context" -import io "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/io" -import mock "github.com/stretchr/testify/mock" +import ( + context "context" + + catalog "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/catalog" + + io "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/io" + + mock "github.com/stretchr/testify/mock" +) // Client is an autogenerated mock type for the Client type type Client struct { diff --git a/go/tasks/pluginmachinery/catalog/mocks/download_future.go b/go/tasks/pluginmachinery/catalog/mocks/download_future.go index 6c555a243d..ba08ced6f0 100644 --- a/go/tasks/pluginmachinery/catalog/mocks/download_future.go +++ b/go/tasks/pluginmachinery/catalog/mocks/download_future.go @@ -2,8 +2,10 @@ package mocks -import catalog "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/catalog" -import mock "github.com/stretchr/testify/mock" +import ( + catalog "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/catalog" + mock "github.com/stretchr/testify/mock" +) // DownloadFuture is an autogenerated mock type for the DownloadFuture type type DownloadFuture struct { diff --git a/go/tasks/pluginmachinery/catalog/mocks/download_response.go b/go/tasks/pluginmachinery/catalog/mocks/download_response.go index a58f93fa01..cff414ea80 100644 --- a/go/tasks/pluginmachinery/catalog/mocks/download_response.go +++ b/go/tasks/pluginmachinery/catalog/mocks/download_response.go @@ -2,9 +2,11 @@ package mocks -import bitarray "github.com/lyft/flytestdlib/bitarray" +import ( + bitarray "github.com/lyft/flytestdlib/bitarray" -import mock "github.com/stretchr/testify/mock" + mock "github.com/stretchr/testify/mock" +) // DownloadResponse is an autogenerated mock type for the DownloadResponse type type DownloadResponse struct { diff --git a/go/tasks/pluginmachinery/catalog/mocks/future.go b/go/tasks/pluginmachinery/catalog/mocks/future.go index c08d0b949a..710d46b753 100644 --- a/go/tasks/pluginmachinery/catalog/mocks/future.go +++ b/go/tasks/pluginmachinery/catalog/mocks/future.go @@ -2,8 +2,10 @@ package mocks -import catalog "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/catalog" -import mock "github.com/stretchr/testify/mock" +import ( + catalog "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/catalog" + mock "github.com/stretchr/testify/mock" +) // Future is an autogenerated mock type for the Future type type Future struct { diff --git a/go/tasks/pluginmachinery/catalog/mocks/upload_future.go b/go/tasks/pluginmachinery/catalog/mocks/upload_future.go index 42686bd74c..c6bd559186 100644 --- a/go/tasks/pluginmachinery/catalog/mocks/upload_future.go +++ b/go/tasks/pluginmachinery/catalog/mocks/upload_future.go @@ -2,8 +2,10 @@ package mocks -import catalog "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/catalog" -import mock "github.com/stretchr/testify/mock" +import ( + catalog "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/catalog" + mock "github.com/stretchr/testify/mock" +) // UploadFuture is an autogenerated mock type for the UploadFuture type type UploadFuture struct { diff --git a/go/tasks/pluginmachinery/core/mocks/events_recorder.go b/go/tasks/pluginmachinery/core/mocks/events_recorder.go index a2ba4a287d..ae493a1ec2 100644 --- a/go/tasks/pluginmachinery/core/mocks/events_recorder.go +++ b/go/tasks/pluginmachinery/core/mocks/events_recorder.go @@ -2,9 +2,12 @@ package mocks -import context "context" -import core "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/core" -import mock "github.com/stretchr/testify/mock" +import ( + context "context" + + core "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/core" + mock "github.com/stretchr/testify/mock" +) // EventsRecorder is an autogenerated mock type for the EventsRecorder type type EventsRecorder struct { diff --git a/go/tasks/pluginmachinery/core/mocks/kube_client.go b/go/tasks/pluginmachinery/core/mocks/kube_client.go index ca48835335..1404b317b5 100644 --- a/go/tasks/pluginmachinery/core/mocks/kube_client.go +++ b/go/tasks/pluginmachinery/core/mocks/kube_client.go @@ -2,10 +2,12 @@ package mocks -import cache "sigs.k8s.io/controller-runtime/pkg/cache" -import client "sigs.k8s.io/controller-runtime/pkg/client" +import ( + cache "sigs.k8s.io/controller-runtime/pkg/cache" + client "sigs.k8s.io/controller-runtime/pkg/client" -import mock "github.com/stretchr/testify/mock" + mock "github.com/stretchr/testify/mock" +) // KubeClient is an autogenerated mock type for the KubeClient type type KubeClient struct { diff --git a/go/tasks/pluginmachinery/core/mocks/plugin.go b/go/tasks/pluginmachinery/core/mocks/plugin.go index 74b6ae5b5a..1d925a798b 100644 --- a/go/tasks/pluginmachinery/core/mocks/plugin.go +++ b/go/tasks/pluginmachinery/core/mocks/plugin.go @@ -2,9 +2,12 @@ package mocks -import context "context" -import core "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/core" -import mock "github.com/stretchr/testify/mock" +import ( + context "context" + + core "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/core" + mock "github.com/stretchr/testify/mock" +) // Plugin is an autogenerated mock type for the Plugin type type Plugin struct { diff --git a/go/tasks/pluginmachinery/core/mocks/resource_manager.go b/go/tasks/pluginmachinery/core/mocks/resource_manager.go index 3657a69580..7478e13992 100644 --- a/go/tasks/pluginmachinery/core/mocks/resource_manager.go +++ b/go/tasks/pluginmachinery/core/mocks/resource_manager.go @@ -2,9 +2,12 @@ package mocks -import context "context" -import core "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/core" -import mock "github.com/stretchr/testify/mock" +import ( + context "context" + + core "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/core" + mock "github.com/stretchr/testify/mock" +) // ResourceManager is an autogenerated mock type for the ResourceManager type type ResourceManager struct { diff --git a/go/tasks/pluginmachinery/core/mocks/resource_registrar.go b/go/tasks/pluginmachinery/core/mocks/resource_registrar.go index 28e934801f..c2707c49d5 100644 --- a/go/tasks/pluginmachinery/core/mocks/resource_registrar.go +++ b/go/tasks/pluginmachinery/core/mocks/resource_registrar.go @@ -2,9 +2,12 @@ package mocks -import context "context" -import core "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/core" -import mock "github.com/stretchr/testify/mock" +import ( + context "context" + + core "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/core" + mock "github.com/stretchr/testify/mock" +) // ResourceRegistrar is an autogenerated mock type for the ResourceRegistrar type type ResourceRegistrar struct { diff --git a/go/tasks/pluginmachinery/core/mocks/secret_manager.go b/go/tasks/pluginmachinery/core/mocks/secret_manager.go index 3f2f123913..84953489b6 100644 --- a/go/tasks/pluginmachinery/core/mocks/secret_manager.go +++ b/go/tasks/pluginmachinery/core/mocks/secret_manager.go @@ -2,9 +2,11 @@ package mocks -import context "context" +import ( + context "context" -import mock "github.com/stretchr/testify/mock" + mock "github.com/stretchr/testify/mock" +) // SecretManager is an autogenerated mock type for the SecretManager type type SecretManager struct { diff --git a/go/tasks/pluginmachinery/core/mocks/setup_context.go b/go/tasks/pluginmachinery/core/mocks/setup_context.go index 49f6ad851c..46be413f34 100644 --- a/go/tasks/pluginmachinery/core/mocks/setup_context.go +++ b/go/tasks/pluginmachinery/core/mocks/setup_context.go @@ -2,9 +2,12 @@ package mocks -import core "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/core" -import mock "github.com/stretchr/testify/mock" -import promutils "github.com/lyft/flytestdlib/promutils" +import ( + core "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/core" + mock "github.com/stretchr/testify/mock" + + promutils "github.com/lyft/flytestdlib/promutils" +) // SetupContext is an autogenerated mock type for the SetupContext type type SetupContext struct { diff --git a/go/tasks/pluginmachinery/core/mocks/task_execution_context.go b/go/tasks/pluginmachinery/core/mocks/task_execution_context.go index 4a2bf2b477..c53bcd0c81 100644 --- a/go/tasks/pluginmachinery/core/mocks/task_execution_context.go +++ b/go/tasks/pluginmachinery/core/mocks/task_execution_context.go @@ -2,11 +2,16 @@ package mocks -import catalog "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/catalog" -import core "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/core" -import io "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/io" -import mock "github.com/stretchr/testify/mock" -import storage "github.com/lyft/flytestdlib/storage" +import ( + catalog "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/catalog" + core "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/core" + + io "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/io" + + mock "github.com/stretchr/testify/mock" + + storage "github.com/lyft/flytestdlib/storage" +) // TaskExecutionContext is an autogenerated mock type for the TaskExecutionContext type type TaskExecutionContext struct { diff --git a/go/tasks/pluginmachinery/core/mocks/task_execution_id.go b/go/tasks/pluginmachinery/core/mocks/task_execution_id.go index e07210e439..1b1332936e 100644 --- a/go/tasks/pluginmachinery/core/mocks/task_execution_id.go +++ b/go/tasks/pluginmachinery/core/mocks/task_execution_id.go @@ -2,8 +2,10 @@ package mocks -import flyteidlcore "github.com/lyft/flyteidl/gen/pb-go/flyteidl/core" -import mock "github.com/stretchr/testify/mock" +import ( + flyteidlcore "github.com/lyft/flyteidl/gen/pb-go/flyteidl/core" + mock "github.com/stretchr/testify/mock" +) // TaskExecutionID is an autogenerated mock type for the TaskExecutionID type type TaskExecutionID struct { diff --git a/go/tasks/pluginmachinery/core/mocks/task_execution_metadata.go b/go/tasks/pluginmachinery/core/mocks/task_execution_metadata.go index 3c844d9079..c7d7a0b886 100644 --- a/go/tasks/pluginmachinery/core/mocks/task_execution_metadata.go +++ b/go/tasks/pluginmachinery/core/mocks/task_execution_metadata.go @@ -2,10 +2,14 @@ package mocks -import core "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/core" -import mock "github.com/stretchr/testify/mock" -import types "k8s.io/apimachinery/pkg/types" -import v1 "k8s.io/apimachinery/pkg/apis/meta/v1" +import ( + core "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/core" + mock "github.com/stretchr/testify/mock" + + types "k8s.io/apimachinery/pkg/types" + + v1 "k8s.io/apimachinery/pkg/apis/meta/v1" +) // TaskExecutionMetadata is an autogenerated mock type for the TaskExecutionMetadata type type TaskExecutionMetadata struct { diff --git a/go/tasks/pluginmachinery/core/mocks/task_overrides.go b/go/tasks/pluginmachinery/core/mocks/task_overrides.go index 9f0167004c..dc35666c77 100644 --- a/go/tasks/pluginmachinery/core/mocks/task_overrides.go +++ b/go/tasks/pluginmachinery/core/mocks/task_overrides.go @@ -2,8 +2,10 @@ package mocks -import mock "github.com/stretchr/testify/mock" -import v1 "k8s.io/api/core/v1" +import ( + mock "github.com/stretchr/testify/mock" + v1 "k8s.io/api/core/v1" +) // TaskOverrides is an autogenerated mock type for the TaskOverrides type type TaskOverrides struct { diff --git a/go/tasks/pluginmachinery/core/mocks/task_reader.go b/go/tasks/pluginmachinery/core/mocks/task_reader.go index 62ea96ae91..c29cc26ab5 100644 --- a/go/tasks/pluginmachinery/core/mocks/task_reader.go +++ b/go/tasks/pluginmachinery/core/mocks/task_reader.go @@ -2,9 +2,12 @@ package mocks -import context "context" -import core "github.com/lyft/flyteidl/gen/pb-go/flyteidl/core" -import mock "github.com/stretchr/testify/mock" +import ( + context "context" + + core "github.com/lyft/flyteidl/gen/pb-go/flyteidl/core" + mock "github.com/stretchr/testify/mock" +) // TaskReader is an autogenerated mock type for the TaskReader type type TaskReader struct { diff --git a/go/tasks/pluginmachinery/io/mocks/input_file_paths.go b/go/tasks/pluginmachinery/io/mocks/input_file_paths.go index 23ace1dd90..91ea4938a4 100644 --- a/go/tasks/pluginmachinery/io/mocks/input_file_paths.go +++ b/go/tasks/pluginmachinery/io/mocks/input_file_paths.go @@ -2,8 +2,10 @@ package mocks -import mock "github.com/stretchr/testify/mock" -import storage "github.com/lyft/flytestdlib/storage" +import ( + storage "github.com/lyft/flytestdlib/storage" + mock "github.com/stretchr/testify/mock" +) // InputFilePaths is an autogenerated mock type for the InputFilePaths type type InputFilePaths struct { diff --git a/go/tasks/pluginmachinery/io/mocks/input_reader.go b/go/tasks/pluginmachinery/io/mocks/input_reader.go index 44b83e6a7b..83806da423 100644 --- a/go/tasks/pluginmachinery/io/mocks/input_reader.go +++ b/go/tasks/pluginmachinery/io/mocks/input_reader.go @@ -2,11 +2,15 @@ package mocks -import context "context" -import core "github.com/lyft/flyteidl/gen/pb-go/flyteidl/core" +import ( + context "context" -import mock "github.com/stretchr/testify/mock" -import storage "github.com/lyft/flytestdlib/storage" + core "github.com/lyft/flyteidl/gen/pb-go/flyteidl/core" + + mock "github.com/stretchr/testify/mock" + + storage "github.com/lyft/flytestdlib/storage" +) // InputReader is an autogenerated mock type for the InputReader type type InputReader struct { diff --git a/go/tasks/pluginmachinery/io/mocks/output_file_paths.go b/go/tasks/pluginmachinery/io/mocks/output_file_paths.go index e8438c59b1..bd65a3305f 100644 --- a/go/tasks/pluginmachinery/io/mocks/output_file_paths.go +++ b/go/tasks/pluginmachinery/io/mocks/output_file_paths.go @@ -2,8 +2,10 @@ package mocks -import mock "github.com/stretchr/testify/mock" -import storage "github.com/lyft/flytestdlib/storage" +import ( + storage "github.com/lyft/flytestdlib/storage" + mock "github.com/stretchr/testify/mock" +) // OutputFilePaths is an autogenerated mock type for the OutputFilePaths type type OutputFilePaths struct { diff --git a/go/tasks/pluginmachinery/io/mocks/output_reader.go b/go/tasks/pluginmachinery/io/mocks/output_reader.go index 4eab95124b..c2fbb8f2b9 100644 --- a/go/tasks/pluginmachinery/io/mocks/output_reader.go +++ b/go/tasks/pluginmachinery/io/mocks/output_reader.go @@ -2,10 +2,14 @@ package mocks -import context "context" -import core "github.com/lyft/flyteidl/gen/pb-go/flyteidl/core" -import io "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/io" -import mock "github.com/stretchr/testify/mock" +import ( + context "context" + + core "github.com/lyft/flyteidl/gen/pb-go/flyteidl/core" + io "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/io" + + mock "github.com/stretchr/testify/mock" +) // OutputReader is an autogenerated mock type for the OutputReader type type OutputReader struct { diff --git a/go/tasks/pluginmachinery/io/mocks/output_writer.go b/go/tasks/pluginmachinery/io/mocks/output_writer.go index 8f0770e4b0..0414271d01 100644 --- a/go/tasks/pluginmachinery/io/mocks/output_writer.go +++ b/go/tasks/pluginmachinery/io/mocks/output_writer.go @@ -2,10 +2,14 @@ package mocks -import context "context" -import io "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/io" -import mock "github.com/stretchr/testify/mock" -import storage "github.com/lyft/flytestdlib/storage" +import ( + context "context" + + io "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/io" + mock "github.com/stretchr/testify/mock" + + storage "github.com/lyft/flytestdlib/storage" +) // OutputWriter is an autogenerated mock type for the OutputWriter type type OutputWriter struct { diff --git a/go/tasks/pluginmachinery/k8s/mocks/plugin.go b/go/tasks/pluginmachinery/k8s/mocks/plugin.go index 2ad3481c39..5439f7f4fc 100644 --- a/go/tasks/pluginmachinery/k8s/mocks/plugin.go +++ b/go/tasks/pluginmachinery/k8s/mocks/plugin.go @@ -2,10 +2,14 @@ package mocks -import context "context" -import core "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/core" -import k8s "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/k8s" -import mock "github.com/stretchr/testify/mock" +import ( + context "context" + + core "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/core" + k8s "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/k8s" + + mock "github.com/stretchr/testify/mock" +) // Plugin is an autogenerated mock type for the Plugin type type Plugin struct { diff --git a/go/tasks/pluginmachinery/k8s/mocks/plugin_context.go b/go/tasks/pluginmachinery/k8s/mocks/plugin_context.go index 1009c8c430..a9011ddb0b 100644 --- a/go/tasks/pluginmachinery/k8s/mocks/plugin_context.go +++ b/go/tasks/pluginmachinery/k8s/mocks/plugin_context.go @@ -2,10 +2,12 @@ package mocks -import core "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/core" -import io "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/io" +import ( + core "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/core" + io "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/io" -import mock "github.com/stretchr/testify/mock" + mock "github.com/stretchr/testify/mock" +) // PluginContext is an autogenerated mock type for the PluginContext type type PluginContext struct { diff --git a/go/tasks/pluginmachinery/k8s/mocks/resource.go b/go/tasks/pluginmachinery/k8s/mocks/resource.go index 084e608889..85bfa9d21f 100644 --- a/go/tasks/pluginmachinery/k8s/mocks/resource.go +++ b/go/tasks/pluginmachinery/k8s/mocks/resource.go @@ -2,11 +2,16 @@ package mocks -import mock "github.com/stretchr/testify/mock" -import runtime "k8s.io/apimachinery/pkg/runtime" -import schema "k8s.io/apimachinery/pkg/runtime/schema" -import types "k8s.io/apimachinery/pkg/types" -import v1 "k8s.io/apimachinery/pkg/apis/meta/v1" +import ( + mock "github.com/stretchr/testify/mock" + runtime "k8s.io/apimachinery/pkg/runtime" + + schema "k8s.io/apimachinery/pkg/runtime/schema" + + types "k8s.io/apimachinery/pkg/types" + + v1 "k8s.io/apimachinery/pkg/apis/meta/v1" +) // Resource is an autogenerated mock type for the Resource type type Resource struct { diff --git a/go/tasks/pluginmachinery/workqueue/mocks/indexed_work_queue.go b/go/tasks/pluginmachinery/workqueue/mocks/indexed_work_queue.go index cec17ebf17..b6bd1f12ef 100644 --- a/go/tasks/pluginmachinery/workqueue/mocks/indexed_work_queue.go +++ b/go/tasks/pluginmachinery/workqueue/mocks/indexed_work_queue.go @@ -2,9 +2,12 @@ package mocks -import context "context" -import mock "github.com/stretchr/testify/mock" -import workqueue "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/workqueue" +import ( + context "context" + + workqueue "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/workqueue" + mock "github.com/stretchr/testify/mock" +) // IndexedWorkQueue is an autogenerated mock type for the IndexedWorkQueue type type IndexedWorkQueue struct { diff --git a/go/tasks/pluginmachinery/workqueue/mocks/processor.go b/go/tasks/pluginmachinery/workqueue/mocks/processor.go index 1cb9bce5c0..84b8bb9a47 100644 --- a/go/tasks/pluginmachinery/workqueue/mocks/processor.go +++ b/go/tasks/pluginmachinery/workqueue/mocks/processor.go @@ -2,9 +2,12 @@ package mocks -import context "context" -import mock "github.com/stretchr/testify/mock" -import workqueue "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/workqueue" +import ( + context "context" + + workqueue "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/workqueue" + mock "github.com/stretchr/testify/mock" +) // Processor is an autogenerated mock type for the Processor type type Processor struct { diff --git a/go/tasks/pluginmachinery/workqueue/mocks/work_item_info.go b/go/tasks/pluginmachinery/workqueue/mocks/work_item_info.go index 2aaad131a8..261ca9cc24 100644 --- a/go/tasks/pluginmachinery/workqueue/mocks/work_item_info.go +++ b/go/tasks/pluginmachinery/workqueue/mocks/work_item_info.go @@ -2,8 +2,10 @@ package mocks -import mock "github.com/stretchr/testify/mock" -import workqueue "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/workqueue" +import ( + workqueue "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/workqueue" + mock "github.com/stretchr/testify/mock" +) // WorkItemInfo is an autogenerated mock type for the WorkItemInfo type type WorkItemInfo struct { diff --git a/go/tasks/plugins/array/awsbatch/config/config.go b/go/tasks/plugins/array/awsbatch/config/config.go index c85872a980..9883c5c83b 100644 --- a/go/tasks/plugins/array/awsbatch/config/config.go +++ b/go/tasks/plugins/array/awsbatch/config/config.go @@ -18,6 +18,7 @@ type Config struct { MaxArrayJobSize int64 `json:"maxArrayJobSize" pflag:",Maximum size of array job."` MinRetries int32 `json:"minRetries" pflag:",Minimum number of retries"` MaxRetries int32 `json:"maxRetries" pflag:",Maximum number of retries"` + DefaultTimeOut config.Duration `json:"defaultTimeout" pflag:",Default timeout for the batch job."` // Provide additional environment variable pairs that plugin authors will provide to containers DefaultEnvVars map[string]string `json:"defaultEnvVars" pflag:"-,Additional environment variable that should be injected into every resource"` MaxErrorStringLength int `json:"maxErrLength" pflag:",Determines the maximum length of the error string returned for the array."` @@ -53,6 +54,7 @@ var ( }, MinRetries: 1, MaxRetries: 10, + DefaultTimeOut: config.Duration{Duration: 3 * 24 * time.Hour}, MaxErrorStringLength: 500, OutputAssembler: workqueue.Config{ IndexCacheMaxItems: 100000, diff --git a/go/tasks/plugins/array/awsbatch/config/config_flags.go b/go/tasks/plugins/array/awsbatch/config/config_flags.go index 63189f9fa9..80ffa552cb 100755 --- a/go/tasks/plugins/array/awsbatch/config/config_flags.go +++ b/go/tasks/plugins/array/awsbatch/config/config_flags.go @@ -53,6 +53,7 @@ func (cfg Config) GetPFlagSet(prefix string) *pflag.FlagSet { cmdFlags.Int64(fmt.Sprintf("%v%v", prefix, "maxArrayJobSize"), defaultConfig.MaxArrayJobSize, "Maximum size of array job.") cmdFlags.Int32(fmt.Sprintf("%v%v", prefix, "minRetries"), defaultConfig.MinRetries, "Minimum number of retries") cmdFlags.Int32(fmt.Sprintf("%v%v", prefix, "maxRetries"), defaultConfig.MaxRetries, "Maximum number of retries") + cmdFlags.String(fmt.Sprintf("%v%v", prefix, "defaultTimeout"), defaultConfig.DefaultTimeOut.String(), "Default timeout for the batch job.") cmdFlags.Int(fmt.Sprintf("%v%v", prefix, "maxErrLength"), defaultConfig.MaxErrorStringLength, "Determines the maximum length of the error string returned for the array.") cmdFlags.String(fmt.Sprintf("%v%v", prefix, "roleAnnotationKey"), defaultConfig.RoleAnnotationKey, "Map key to use to lookup role from task annotations.") cmdFlags.Int(fmt.Sprintf("%v%v", prefix, "outputAssembler.workers"), defaultConfig.OutputAssembler.Workers, "Number of concurrent workers to start processing the queue.") diff --git a/go/tasks/plugins/array/awsbatch/config/config_flags_test.go b/go/tasks/plugins/array/awsbatch/config/config_flags_test.go index 8ce8bfb28b..95fbf94aaf 100755 --- a/go/tasks/plugins/array/awsbatch/config/config_flags_test.go +++ b/go/tasks/plugins/array/awsbatch/config/config_flags_test.go @@ -363,6 +363,28 @@ func TestConfig_SetFlags(t *testing.T) { } }) }) + t.Run("Test_defaultTimeout", func(t *testing.T) { + t.Run("DefaultValue", func(t *testing.T) { + // Test that default value is set properly + if vString, err := cmdFlags.GetString("defaultTimeout"); err == nil { + assert.Equal(t, string(defaultConfig.DefaultTimeOut.String()), vString) + } else { + assert.FailNow(t, err.Error()) + } + }) + + t.Run("Override", func(t *testing.T) { + testValue := defaultConfig.DefaultTimeOut.String() + + cmdFlags.Set("defaultTimeout", testValue) + if vString, err := cmdFlags.GetString("defaultTimeout"); err == nil { + testDecodeJson_Config(t, fmt.Sprintf("%v", vString), &actual.DefaultTimeOut) + + } else { + assert.FailNow(t, err.Error()) + } + }) + }) t.Run("Test_maxErrLength", func(t *testing.T) { t.Run("DefaultValue", func(t *testing.T) { // Test that default value is set properly diff --git a/go/tasks/plugins/array/awsbatch/mocks/batch_service_client.go b/go/tasks/plugins/array/awsbatch/mocks/batch_service_client.go index 3f6a3d6c61..810ded0a61 100644 --- a/go/tasks/plugins/array/awsbatch/mocks/batch_service_client.go +++ b/go/tasks/plugins/array/awsbatch/mocks/batch_service_client.go @@ -2,10 +2,15 @@ package mocks -import batch "github.com/aws/aws-sdk-go/service/batch" -import context "context" -import mock "github.com/stretchr/testify/mock" -import request "github.com/aws/aws-sdk-go/aws/request" +import ( + context "context" + + batch "github.com/aws/aws-sdk-go/service/batch" + + mock "github.com/stretchr/testify/mock" + + request "github.com/aws/aws-sdk-go/aws/request" +) // BatchServiceClient is an autogenerated mock type for the BatchServiceClient type type BatchServiceClient struct { diff --git a/go/tasks/plugins/array/awsbatch/mocks/cache.go b/go/tasks/plugins/array/awsbatch/mocks/cache.go index 57e86864f3..c50fb1be35 100644 --- a/go/tasks/plugins/array/awsbatch/mocks/cache.go +++ b/go/tasks/plugins/array/awsbatch/mocks/cache.go @@ -2,8 +2,10 @@ package mocks -import definition "github.com/lyft/flyteplugins/go/tasks/plugins/array/awsbatch/definition" -import mock "github.com/stretchr/testify/mock" +import ( + definition "github.com/lyft/flyteplugins/go/tasks/plugins/array/awsbatch/definition" + mock "github.com/stretchr/testify/mock" +) // Cache is an autogenerated mock type for the Cache type type Cache struct { diff --git a/go/tasks/plugins/array/awsbatch/mocks/client.go b/go/tasks/plugins/array/awsbatch/mocks/client.go index 39f2983d37..70faae703c 100644 --- a/go/tasks/plugins/array/awsbatch/mocks/client.go +++ b/go/tasks/plugins/array/awsbatch/mocks/client.go @@ -2,9 +2,13 @@ package mocks -import batch "github.com/aws/aws-sdk-go/service/batch" -import context "context" -import mock "github.com/stretchr/testify/mock" +import ( + context "context" + + batch "github.com/aws/aws-sdk-go/service/batch" + + mock "github.com/stretchr/testify/mock" +) // Client is an autogenerated mock type for the Client type type Client struct { diff --git a/go/tasks/plugins/array/awsbatch/transformer.go b/go/tasks/plugins/array/awsbatch/transformer.go index 95f021b073..a82eaa75fe 100644 --- a/go/tasks/plugins/array/awsbatch/transformer.go +++ b/go/tasks/plugins/array/awsbatch/transformer.go @@ -3,6 +3,9 @@ package awsbatch import ( "context" "sort" + "time" + + "github.com/golang/protobuf/ptypes/duration" "k8s.io/apimachinery/pkg/api/resource" @@ -85,6 +88,7 @@ func FlyteTaskToBatchInput(ctx context.Context, tCtx pluginCore.TaskExecutionCon JobQueue: refStr(jobConfig.DynamicTaskQueue), RetryStrategy: toRetryStrategy(ctx, toBackoffLimit(taskTemplate.Metadata), cfg.MinRetries, cfg.MaxRetries), ContainerOverrides: toContainerOverrides(ctx, append(cmd, args...), res, envVars), + Timeout: toTimeout(taskTemplate.Metadata.GetTimeout(), cfg.DefaultTimeOut.Duration), }, nil } @@ -132,6 +136,18 @@ func getEnvVarsForTask(ctx context.Context, execID pluginCore.TaskExecutionID, c return finalEnvVars } +func toTimeout(templateTimeout *duration.Duration, defaultTimeout time.Duration) *batch.JobTimeout { + if templateTimeout != nil { + return (&batch.JobTimeout{}).SetAttemptDurationSeconds(templateTimeout.GetSeconds()) + } + + if defaultTimeout.Seconds() > 0 { + return (&batch.JobTimeout{}).SetAttemptDurationSeconds(int64(defaultTimeout.Seconds())) + } + + return nil +} + func toEnvironmentVariables(_ context.Context, envVars []v1.EnvVar) []*batch.KeyValuePair { res := make([]*batch.KeyValuePair, 0, len(envVars)) for _, pair := range envVars { @@ -192,7 +208,9 @@ func toRetryStrategy(_ context.Context, backoffLimit *int32, minRetryAttempts, m } return &batch.RetryStrategy{ - Attempts: refInt(int64(retries)), + // Attempts is the total number of attempts for a task, if retries is set to 1, attempts should be set to 2 to + // account for the first try. + Attempts: refInt(int64(retries) + 1), } } diff --git a/go/tasks/plugins/hive/client/mocks/qubole_client.go b/go/tasks/plugins/hive/client/mocks/qubole_client.go index c22c833071..9daa0ef226 100644 --- a/go/tasks/plugins/hive/client/mocks/qubole_client.go +++ b/go/tasks/plugins/hive/client/mocks/qubole_client.go @@ -2,9 +2,13 @@ package mocks -import client "github.com/lyft/flyteplugins/go/tasks/plugins/hive/client" -import context "context" -import mock "github.com/stretchr/testify/mock" +import ( + context "context" + + client "github.com/lyft/flyteplugins/go/tasks/plugins/hive/client" + + mock "github.com/stretchr/testify/mock" +) // QuboleClient is an autogenerated mock type for the QuboleClient type type QuboleClient struct {