Interacting with the Catalog Client

To interact with the Flyte Catalog for memoization, you use the catalog.Client or catalog.AsyncClient to manage task execution data. This allows Flyte to skip redundant computations by retrieving previously stored results based on a unique cache key.

Defining a Catalog Key

Every interaction with the catalog requires a catalog.Key. This structure uniquely identifies a task execution's inputs and configuration to determine if a cached result exists.

import (
    "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/catalog"
    corepb "github.com/flyteorg/flyteidl/gen/go/core"
)

// Constructing a key within a task execution context
key := catalog.Key{
    Identifier:           taskTemplate.GetId(),
    CacheVersion:         taskAction.Spec.CacheKey,
    CacheIgnoreInputVars: metadata.GetCacheIgnoreInputVars(),
    TypedInterface:       taskTemplate.GetInterface(),
    InputReader:          tCtx.InputReader(),
}

The InputReader is used by the catalog client to fetch and hash the actual input values. You can exclude specific variables from this hash by adding them to CacheIgnoreInputVars.

Synchronous Cache Operations

The catalog.Client provides synchronous methods to check for hits and store new results. This is typically used within task controllers to determine if a task needs to run.

Retrieving Cached Results

Use Get to check for an existing entry. If found, you can restore the outputs directly to the task's output writer.

// From executor/pkg/controller/taskaction_cache.go
entry, err := r.Catalog.Get(ctx, key)
if err == nil {
    // Cache Hit: Restore outputs
    if err := tCtx.OutputWriter().Put(ctx, entry.GetOutputs()); err != nil {
        return pluginsCore.UnknownTransition, false, fmt.Errorf("persisting cached outputs: %w", err)
    }
    return pluginsCore.DoTransition(pluginsCore.PhaseInfoSuccess(info)), true, nil
}

if catalog.IsNotFound(err) {
    // Cache Miss: Proceed with execution
}

Storing New Results

After a task completes successfully, use Put to store the outputs in the catalog.

// From executor/pkg/controller/taskaction_cache.go
func (r *TaskActionReconciler) writeTaskOutputsToCache(ctx context.Context, tCtx pluginsCore.TaskExecutionContext, key catalog.Key) error {
    outputPaths := ioutils.NewReadOnlyOutputFilePaths(ctx, r.DataStore, tCtx.OutputWriter().GetOutputPrefixPath())
    outputReader := ioutils.NewRemoteFileOutputReader(ctx, r.DataStore, outputPaths, 0)
    
    metadata := catalog.Metadata{
        WorkflowExecutionIdentifier: tCtx.TaskExecutionMetadata().GetTaskExecutionID().GetID().GetNodeExecutionId().GetExecutionId(),
        CreatedAt:                   timestamppb.Now(),
    }

    _, err := r.Catalog.Put(ctx, key, outputReader, metadata)
    return err
}

Managing Serializable Cache Reservations

When CacheSerializable is enabled, Flyte uses a reservation system to ensure that only one execution attempts to populate the cache for a specific key at a time. This prevents multiple concurrent executions from performing the same expensive computation.

// From executor/pkg/controller/taskaction_cache.go
reservation, err := r.Catalog.GetOrExtendReservation(ctx, key, ownerID, heartbeatInterval)
if err != nil {
    return pluginsCore.UnknownTransition, false, fmt.Errorf("acquiring cache reservation: %w", err)
}

if reservation.GetOwnerId() == ownerID {
    // Current execution owns the reservation; proceed to run the task and populate cache
    return pluginsCore.UnknownTransition, false, nil
}

// Another execution owns the reservation; wait for them to finish
info := cacheTaskInfo(corepb.CatalogCacheStatus_CACHE_MISS, "waiting for serialized cache owner")

Always call ReleaseReservation if the execution fails or completes without populating the cache to allow other executions to proceed.

Asynchronous Catalog Interactions

For high-throughput scenarios or plugins handling many sub-tasks (like array tasks), use AsyncClient. It offloads catalog requests to internal workqueues and returns futures.

// From flyteplugins/go/tasks/pluginmachinery/catalog/async_client_impl_test.go
asyncClient := catalog.AsyncClientImpl{
    Reader: readerWorkqueue,
    Writer: writerWorkqueue,
}

// Request an asynchronous download
request := catalog.DownloadRequest{
    Key:    key,
    Target: outputWriter,
}

future, err := asyncClient.Download(ctx, request)
if err != nil {
    return err
}

// Check status later in the execution loop
if future.GetResponseStatus() == catalog.ResponseStatusReady {
    resp, err := future.GetResponse()
    // Process response
}

Configuration

The behavior of the catalog client, especially the asynchronous workqueues, is controlled via the catalogCache configuration section.

Parameter	Default	Description
reader.workers	10	Number of parallel workers for catalog reads.
writer.workers	10	Number of parallel workers for catalog writes.
reader.maxRetries	3	Maximum retries for a failed catalog read.
reader.indexCacheMaxItems	10000	Size of the internal index for the reader workqueue.
maxCacheAge	0	Maximum age of a cache entry before it is considered a miss (0 = never expires).

Troubleshooting

Distinguishing Misses from Failures

Use catalog.IsNotFound(err) to identify a legitimate cache miss. Other errors (network timeouts, permission issues) should generally be logged but may allow the task to continue execution without memoization to avoid blocking the workflow.

Input Hashing Gotchas

If you experience unexpected cache misses, verify that the InputReader provided to the Key has access to the same data across executions. If certain inputs (like timestamps or random seeds) should not trigger a cache miss, ensure they are included in CacheIgnoreInputVars.