Data IO: Readers and Writers

When you develop a Flyte plugin, your code needs to read inputs and write outputs without being coupled to specific storage backends like S3 or GCS. The InputReader and OutputWriter interfaces provide a uniform way to handle data transfer while hiding the complexity of remote storage protocols and path construction.

Reading Task Inputs

To access the data passed to your task, you use the InputReader interface. In most plugin implementations, this is provided via the TaskExecutionContext. Instead of manually downloading files, you call Get to retrieve the inputs as a LiteralMap.

// Example usage in a plugin's Handle method
func (m MyPlugin) Handle(ctx context.Context, tCtx core.TaskExecutionContext) (core.Transition, error) {
    // Get the inputs for the task
    inputs, err := tCtx.InputReader().Get(ctx)
    if err != nil {
        return core.UnknownTransition, err
    }
    
    // Access specific input literals
    if inputs != nil && inputs.Literals != nil {
        // Process inputs...
    }
}

Internally, the InputReader (defined in flyteplugins/go/tasks/pluginmachinery/io/iface.go) also provides access to the raw input paths via InputFilePaths. This is useful if your task needs to know the exact URI of the inputs.pb file in remote storage.

Writing Task Outputs

Persisting results in Flyte follows a specific pattern: you don't write raw data directly to the OutputWriter. Instead, you provide it with an OutputReader. This design decouples the source of the output (which might be in memory, on disk, or in a remote buffer) from the final destination.

The OutputWriter interface defines the Put method for this purpose:

type OutputWriter interface {
	OutputFilePaths
	// Put indicates the output accessor to the framework once the task completes
	Put(ctx context.Context, reader OutputReader) error
}

Remote Persistence with RemoteFileOutputWriter

The standard implementation for remote storage is RemoteFileOutputWriter (found in flyteplugins/go/tasks/pluginmachinery/ioutils/remote_file_output_writer.go). When you call Put, it reads from your provided OutputReader and determines whether to write a success result (a LiteralMap) or an error document.

func (w RemoteFileOutputWriter) Put(ctx context.Context, reader io.OutputReader) error {
	literals, executionErr, err := reader.Read(ctx)
	if err != nil {
		return err
	}

	if executionErr != nil {
		// If the reader contains an execution error, write an ErrorDocument
		errDoc := &core.ErrorDocument{
			Error: &core.ContainerError{
				Code:    executionErr.Code,
				Message: executionErr.Message,
				Kind:    w.mapErrorKind(executionErr.IsRecoverable),
			},
		}
		return w.store.WriteProtobuf(ctx, w.GetErrorPath(), storage.Options{}, errDoc)
	}

	if literals != nil {
		// If the reader contains literals, write the outputs.pb
		return w.store.WriteProtobuf(ctx, w.GetOutputPath(), storage.Options{}, literals)
	}

	return fmt.Errorf("no data found to write")
}

Managing Remote Paths

Flyte organizes task metadata (inputs, outputs, errors, and decks) using a structured path convention. The RemoteFileOutputPaths struct in flyteplugins/go/tasks/pluginmachinery/ioutils/remote_file_output_writer.go manages these locations.

It uses a storage.ReferenceConstructor to build full URIs based on a base outputPrefix. This ensures that for any given task execution, the outputs.pb and errors.pb are always stored in predictable locations:

• Output Path: Usually [prefix]/outputs.pb
• Error Path: Usually [prefix]/errors.pb
• Deck Path: Usually [prefix]/deck.html

Handling Task Failures

When a task fails, Flyte looks for an error document in the remote storage to understand why. The RemoteFileOutputReader is responsible for reading these results back.

If a task fails, RemoteFileOutputReader.ReadError (in flyteplugins/go/tasks/pluginmachinery/ioutils/remote_file_output_reader.go) attempts to parse the ErrorDocument. If the file is missing, it treats the error as a system-level ErrorFileNotFound, which is marked as recoverable by default to allow for retries in case of transient storage issues.

func (r RemoteFileOutputReader) ReadError(ctx context.Context) (io.ExecutionError, error) {
	errorDoc := &core.ErrorDocument{}
	err := r.store.ReadProtobuf(ctx, r.OutPath.GetErrorPath(), errorDoc)
	if err != nil {
		if storage.IsNotFound(err) {
			return io.ExecutionError{
				IsRecoverable: true,
				ExecutionError: &core.ExecutionError{
					Code:    "ErrorFileNotFound",
					Message: err.Error(),
					Kind:    core.ExecutionError_SYSTEM,
				},
			}, nil
		}
		return io.ExecutionError{}, errors.Wrapf(err, "failed to read error data")
	}
    // ... returns the mapped error ...
}

In-Memory Buffering

In some scenarios, such as K8s-based plugins, you may not want to write to remote storage immediately. The BufferedOutputWriter allows you to "capture" the output reader in memory.

// Found in flyteplugins/go/tasks/pluginmachinery/ioutils/buffered_output_writer.go
type BufferedOutputWriter struct {
	io.OutputFilePaths
	outReader io.OutputReader
}

func (o *BufferedOutputWriter) Put(ctx context.Context, reader io.OutputReader) error {
	o.outReader = reader
	return nil
}

By using NewBufferedOutputWriter, a plugin can satisfy the OutputWriter interface requirements during the task's execution phase and then retrieve the stored OutputReader later using GetReader() to perform the actual persistence when the task transition is finalized.