Protobuf Message Persistence

When you need to persist structured data like task inputs, outputs, or error messages to remote storage, manually handling byte buffers and serialization logic is repetitive and error-prone. Flyte provides the ProtobufStore interface to automate the translation between Go Protobuf messages and the underlying blob storage (such as S3, GCS, or local disk).

The DataStore Entry Point

In Flyte, you typically interact with Protobuf persistence through the DataStore struct defined in flytestdlib/storage/storage.go. This struct implements the ComposedProtobufStore interface, which combines raw byte access (RawStore) with structured Protobuf access (ProtobufStore).

type DataStore struct {
	ComposedProtobufStore
	ReferenceConstructor
	metrics *dataStoreMetrics
}

By using DataStore, you gain a unified API for constructing paths, writing raw bytes, and serializing Protobuf messages.

Reading and Writing Messages

To store a Protobuf message, you pass a DataReference (a URI string like s3://my-bucket/data.pb) and the message itself to WriteProtobuf. To retrieve it, you pass the reference and a pointer to a message struct where the data will be unmarshaled.

The following example demonstrates a basic round-trip using a memory-backed store:

import (
    "context"
    "github.com/flyteorg/flyte/v2/flytestdlib/storage"
    "github.com/flyteorg/flyte/v2/flyteidl/gen/pb-go/flyteidl/core"
    "github.com/flyteorg/flyte/v2/flytestdlib/promutils"
)

func ExampleProtobufPersistence() {
    ctx := context.Background()
    testScope := promutils.NewTestScope()
    
    // Initialize a DataStore (e.g., using In-Memory storage for testing)
    s, _ := storage.NewDataStore(&storage.Config{Type: storage.TypeMemory}, testScope)

    // Define a message to store
    msg := &core.LiteralMap{
        Literals: map[string]*core.Literal{
            "key": {Value: &core.Literal_Scalar{Scalar: &core.Scalar{Value: &core.Scalar_Primitive{Primitive: &core.Primitive{Value: &core.Primitive_Integer{Integer: 5}}}}}},
        },
    }

    // Write the message to storage
    ref := storage.DataReference("mem://bucket/inputs.pb")
    err := s.WriteProtobuf(ctx, ref, storage.Options{}, msg)
    if err != nil {
        // Handle error
    }

    // Read the message back
    readMsg := &core.LiteralMap{}
    err = s.ReadProtobuf(ctx, ref, readMsg)
    if err != nil {
        // Handle error
    }
}

Internal Implementation: DefaultProtobufStore

The standard implementation of this logic is DefaultProtobufStore in flytestdlib/storage/protobuf_store.go. It acts as a decorator for any RawStore, handling the serialization lifecycle:

1. Writing: It uses proto.Marshal to convert the message into a byte slice and then calls WriteRaw on the underlying store.
2. Reading: It calls ReadRaw to get an io.ReadCloser, reads the entire stream into memory using ioutils.ReadAll, and finally populates the provided message via proto.Unmarshal.

Handling Cache Failures

A notable behavior in DefaultProtobufStore is how it handles caching errors. If the underlying RawStore is configured with a cache (e.g., via CachingRawStore), writes or reads might fail specifically at the cache layer.

In both ReadProtobuf and WriteProtobuf, Flyte checks if an error is a cache-specific failure using IsFailedWriteToCache(err) (defined in flytestdlib/storage/utils.go). If it is, the error is logged but suppressed, allowing the primary storage operation to be treated as a success.

// From flytestdlib/storage/protobuf_store.go
err = s.WriteRaw(ctx, reference, int64(len(raw)), opts, bytes.NewReader(raw))
if err != nil && !IsFailedWriteToCache(err) {
    logger.Errorf(ctx, "Failed to write to the raw store [%s] Error: %v", reference, err)
    s.metrics.WriteFailureUnrelatedToCache.Inc()
    return err
}

Observability and Metrics

Flyte automatically tracks the performance and reliability of Protobuf operations through the protoMetrics struct. These metrics are exported to Prometheus and include:

Metric Name	Type	Description
proto_fetch	StopWatch	Time spent reading raw bytes from storage before unmarshalling.
marshal	StopWatch	Time spent serializing the message to bytes.
unmarshal	StopWatch	Time spent deserializing bytes into the message struct.
marshal_failure	Counter	Number of times proto.Marshal failed.
unmarshal_failure	Counter	Number of times proto.Unmarshal failed.
write_failure_unrelated_to_cache	Counter	Failures during WriteRaw that were not caused by cache issues.

These metrics allow you to monitor the overhead of serialization and identify if specific storage backends are causing latency spikes during metadata retrieval.

Implementation Details

• Serialization Library: Flyte currently uses the legacy github.com/golang/protobuf/proto package for these operations.
• Memory Usage: ReadProtobuf reads the entire object into memory before unmarshalling. For extremely large Protobuf messages, ensure the environment has sufficient memory, or consider using ReadRaw for streaming if the structure allows.
• Data Integrity: When writing, you can pass storage.Options to include metadata like ContentMD5 if the underlying RawStore (like S3) supports it for integrity checks.