remove setting max size bytes in node context

Signed-off-by: Paul Dittamo <[email protected]>

flyteplugins/go/tasks/pluginmachinery/core/exec_context.go

@@ -9,6 +9,8 @@ import (
 	"github.com/flyteorg/flyte/flytestdlib/storage"
 )
 
+//go:generate mockery -all -case=underscore
+
 // An interface to access a remote/sharable location that contains the serialized TaskTemplate
 type TaskTemplatePath interface {
 	// Returns the path
@@ -34,9 +36,6 @@ type TaskExecutionContext interface {
 	// Returns a method that allows a plugin to indicate that the task has a new update and can be invoked again to check for updates
 	TaskRefreshIndicator() SignalAsync
 
-	// Returns the max allowed dataset size that the outputwriter will accept
-	MaxDatasetSizeBytes() int64
-
 	// Returns a handle to the currently configured storage backend that can be used to communicate with the tasks or write metadata
 	DataStore() *storage.DataStore
 

flyteplugins/go/tasks/pluginmachinery/ioutils/remote_file_output_reader.go

@@ -123,9 +123,13 @@ func (r RemoteFileOutputReader) DeckExists(ctx context.Context) (bool, error) {
 }
 
 func NewRemoteFileOutputReader(_ context.Context, store storage.ComposedProtobufStore, outPaths io.OutputFilePaths, maxDatasetSize int64) RemoteFileOutputReader {
+	maxPayloadSize := maxDatasetSize
+	if maxPayloadSize == 0 {
+		maxPayloadSize = storage.GetConfig().Limits.GetLimitMegabytes * 1024 * 1024
+	}
 	return RemoteFileOutputReader{
 		outPath:        outPaths,
 		store:          store,
-		maxPayloadSize: maxDatasetSize,
+		maxPayloadSize: maxPayloadSize,
 	}
 }

flyteplugins/go/tasks/pluginmachinery/k8s/plugin.go

@@ -61,9 +61,6 @@ type PluginContext interface {
 	// Returns a handle to the currently configured storage backend that can be used to communicate with the tasks or write metadata
 	DataStore() *storage.DataStore
 
-	// Returns the max allowed dataset size that the outputwriter will accept
-	MaxDatasetSizeBytes() int64
-
 	// Returns a handle to the Task's execution metadata.
 	TaskExecutionMetadata() pluginsCore.TaskExecutionMetadata
 

flyteplugins/go/tasks/pluginmachinery/webapi/plugin.go

@@ -72,9 +72,6 @@ type TaskExecutionContext interface {
 
 	// Provides the raw datastore to enable persisting outputs.
 	DataStore() *storage.DataStore
-
-	// Returns the max allowed dataset size that the outputwriter will accept
-	MaxDatasetSizeBytes() int64
 }
 
 type GetContext interface {

flyteplugins/go/tasks/plugins/array/outputs.go

@@ -225,7 +225,7 @@ func AssembleFinalOutputs(ctx context.Context, assemblyQueue OutputAssembler, tC
 
 	switch w.Status() {
 	case workqueue.WorkStatusSucceeded:
-		or := ioutils.NewRemoteFileOutputReader(ctx, tCtx.DataStore(), tCtx.OutputWriter(), tCtx.MaxDatasetSizeBytes())
+		or := ioutils.NewRemoteFileOutputReader(ctx, tCtx.DataStore(), tCtx.OutputWriter(), 0)
 		if err = tCtx.OutputWriter().Put(ctx, or); err != nil {
 			return nil, err
 		}

flyteplugins/go/tasks/plugins/array/outputs_test.go

@@ -244,7 +244,6 @@ func TestAssembleFinalOutputs(t *testing.T) {
 		tCtx := &mocks3.TaskExecutionContext{}
 		tCtx.OnTaskExecutionMetadata().Return(tMeta)
 		tCtx.OnOutputWriter().Return(ow)
-		tCtx.OnMaxDatasetSizeBytes().Return(10000)
 		tCtx.OnDataStore().Return(d)
 
 		_, err = AssembleFinalOutputs(ctx, assemblyQueue, tCtx, arrayCore.PhaseSuccess, 1, s)
@@ -368,7 +367,6 @@ func TestAssembleFinalOutputs(t *testing.T) {
 		tCtx.OnTaskReader().Return(tReader)
 		tCtx.OnOutputWriter().Return(ow)
 		tCtx.OnDataStore().Return(ds)
-		tCtx.OnMaxDatasetSizeBytes().Return(10000)
 
 		_, err = AssembleFinalOutputs(ctx, assemblyQueue, tCtx, arrayCore.PhaseSuccess, 1, s)
 		assert.NoError(t, err)

flyteplugins/go/tasks/plugins/testing/echo.go

@@ -76,7 +76,7 @@ func (e *EchoPlugin) Handle(ctx context.Context, tCtx core.TaskExecutionContext)
 				return core.UnknownTransition, err
 			}
 
-			or := ioutils.NewRemoteFileOutputReader(ctx, tCtx.DataStore(), tCtx.OutputWriter(), tCtx.MaxDatasetSizeBytes())
+			or := ioutils.NewRemoteFileOutputReader(ctx, tCtx.DataStore(), tCtx.OutputWriter(), 0)
 			if err = tCtx.OutputWriter().Put(ctx, or); err != nil {
 				return core.UnknownTransition, err
 			}

flyteplugins/go/tasks/plugins/webapi/agent/plugin.go

@@ -332,7 +332,7 @@ func writeOutput(ctx context.Context, taskCtx webapi.StatusContext, outputs *fly
 		opReader = ioutils.NewInMemoryOutputReader(outputs, nil, nil)
 	} else {
 		logger.Debugf(ctx, "AgentDeployment didn't return any output, assuming file based outputs.")
-		opReader = ioutils.NewRemoteFileOutputReader(ctx, taskCtx.DataStore(), taskCtx.OutputWriter(), taskCtx.MaxDatasetSizeBytes())
+		opReader = ioutils.NewRemoteFileOutputReader(ctx, taskCtx.DataStore(), taskCtx.OutputWriter(), 0)
 	}
 	return taskCtx.OutputWriter().Put(ctx, opReader)
 }

flyteplugins/go/tasks/plugins/webapi/databricks/plugin.go

@@ -292,7 +292,7 @@ func writeOutput(ctx context.Context, taskCtx webapi.StatusContext) error {
 		return nil
 	}
 
-	outputReader := ioutils.NewRemoteFileOutputReader(ctx, taskCtx.DataStore(), taskCtx.OutputWriter(), taskCtx.MaxDatasetSizeBytes())
+	outputReader := ioutils.NewRemoteFileOutputReader(ctx, taskCtx.DataStore(), taskCtx.OutputWriter(), 0)
 	return taskCtx.OutputWriter().Put(ctx, outputReader)
 }
 

flyteplugins/tests/end_to_end.go

@@ -243,7 +243,6 @@ func RunPluginEndToEndTest(t *testing.T, executor pluginCore.Plugin, template *i
 	tCtx.OnCatalog().Return(cat)
 	tCtx.OnEventsRecorder().Return(eRecorder)
 	tCtx.OnResourceManager().Return(resourceManager)
-	tCtx.OnMaxDatasetSizeBytes().Return(1000000)
 	tCtx.OnSecretManager().Return(secretManager)
 
 	trns := pluginCore.DoTransition(pluginCore.PhaseInfoQueued(time.Now(), 0, ""))

flytepropeller/pkg/controller/controller.go

@@ -439,7 +439,7 @@ func New(ctx context.Context, cfg *config.Config, kubeClientset kubernetes.Inter
 	}
 
 	nodeExecutor, err := nodes.NewExecutor(ctx, cfg.NodeConfig, store, controller.enqueueWorkflowForNodeUpdates, eventSink,
-		launchPlanActor, launchPlanActor, sCfg.Limits.GetLimitMegabytes*1024*1024, storage.DataReference(cfg.DefaultRawOutputPrefix), kubeClient,
+		launchPlanActor, launchPlanActor, storage.DataReference(cfg.DefaultRawOutputPrefix), kubeClient,
 		catalogClient, recoveryClient, &cfg.EventConfig, cfg.ClusterID, signalClient, nodeHandlerFactory, scope)
 	if err != nil {
 		return nil, errors.Wrapf(err, "Failed to create Controller.")

flytepropeller/pkg/controller/nodes/array/handler.go

@@ -469,7 +469,7 @@ func (a *arrayNodeHandler) Handle(ctx context.Context, nCtx interfaces.NodeExecu
 				// checkpoint paths are not computed here because this function is only called when writing
 				// existing cached outputs. if this functionality changes this will need to be revisited.
 				outputPaths := ioutils.NewCheckpointRemoteFilePaths(ctx, nCtx.DataStore(), subOutputDir, ioutils.NewRawOutputPaths(ctx, subDataDir), "")
-				reader := ioutils.NewRemoteFileOutputReader(ctx, nCtx.DataStore(), outputPaths, nCtx.MaxDatasetSizeBytes())
+				reader := ioutils.NewRemoteFileOutputReader(ctx, nCtx.DataStore(), outputPaths, 0)
 
 				gatherOutputsRequest.reader = &reader
 				a.gatherOutputsRequestChannel <- gatherOutputsRequest
@@ -505,7 +505,7 @@ func (a *arrayNodeHandler) Handle(ctx context.Context, nCtx interfaces.NodeExecu
 				}
 			}
 		case v1alpha1.NodeKindWorkflow:
-			// TODO - to support launchplans we will need to process the output interface variables here 
+			// TODO - to support launchplans we will need to process the output interface variables here
 			fallthrough
 		default:
 			logger.Warnf(ctx, "ArrayNode does not support pre-populating outputLiteral collections for node kind '%s'", arrayNode.GetSubNodeSpec().GetKind())