Caching & Memoization

Flyte implements task output caching (memoization) to avoid redundant computations. When a task is executed with a specific set of inputs, Flyte can store the resulting outputs and reuse them in subsequent executions with the same inputs. This mechanism is split between a Cache Service that manages metadata and a Catalog Client that coordinates with the execution engine.

The Memoization Flow

When the Flyte Propeller (or the TaskAction controller) prepares to run a task, it first checks if a valid cache entry exists for the given inputs. This logic is primarily implemented in the TaskActionReconciler within executor/pkg/controller/taskaction_cache.go.

The process follows these steps:

1. Build Cache Key: Flyte generates a unique key based on the task's identifier, cache version, and input values.
2. Evaluate Cache: The reconciler calls the Catalog client to check for a hit.
3. Short-circuit Execution: If a hit is found, the reconciler retrieves the output URIs from the cache and populates the task's output writer, skipping the actual task execution.

// From executor/pkg/controller/taskaction_cache.go
func (r *TaskActionReconciler) evaluateCacheBeforeExecution(
	ctx context.Context,
	taskAction *flyteorgv1.TaskAction,
	tCtx pluginsCore.TaskExecutionContext,
) (pluginsCore.Transition, bool, error) {
	cacheCfg, ok, err := buildTaskCacheConfig(ctx, taskAction, tCtx)
	if err != nil || !ok || r.Catalog == nil {
		return pluginsCore.UnknownTransition, false, err
	}

	// Check the catalog for an existing entry
	entry, err := r.Catalog.Get(ctx, cacheCfg.key)
	if err == nil {
		// Cache Hit: Persist cached outputs and transition to Success
		if err := tCtx.OutputWriter().Put(ctx, entry.GetOutputs()); err != nil {
			return pluginsCore.UnknownTransition, false, fmt.Errorf("persisting cached outputs: %w", err)
		}

		info := cacheTaskInfo(corepb.CatalogCacheStatus_CACHE_HIT, "cache hit")
		return pluginsCore.DoTransition(pluginsCore.PhaseInfoSuccess(info)), true, nil
	}
    // ... handles misses
}

Cache Service Architecture

Flyte's caching architecture distinguishes between metadata and data.

• Cache Service: A gRPC/HTTP service that stores the mapping of cache keys to output URIs and execution metadata. It does not store the actual data blobs.
• Object Storage: The actual task outputs (e.g., S3, GCS, or Minio blobs) remain in the configured data store. The Cache Service only stores the URI pointing to these blobs.

The Manager in cache_service/manager/manager.go handles the core logic for persisting these records. When you call Put, the service records the OutputUri and updates the LastUpdated timestamp.

// From cache_service/manager/manager.go
func (m *Manager) Put(ctx context.Context, request *cacheservicepb.PutCacheRequest) error {
	// ... validation logic ...
	model := &models.CachedOutput{
		Key:         request.GetKey(),
		OutputURI:   request.GetOutput().GetOutputUri(),
		Metadata:    metadataBytes,
		LastUpdated: metadata.GetLastUpdatedAt().AsTime(),
	}
	return m.outputs.Put(ctx, model)
}

Serialized Cache Population

When multiple concurrent executions of the same task occur (e.g., a large map task where many sub-tasks have identical inputs), Flyte prevents "thundering herd" problems using a Reservation System. This ensures that only one worker executes the task and populates the cache, while others wait for the result.

This behavior is triggered when metadata.GetCacheSerializable() is true in the task template.

How Reservations Work

1. Acquisition: On a cache miss, the worker attempts to acquire a reservation via GetOrExtendReservation.
2. Ownership: If the worker becomes the owner, it proceeds with execution.
3. Waiting: If another worker already holds a valid reservation, the current worker transitions to a WaitingForCache phase.
4. Heartbeats: The owner must periodically refresh the reservation. If the owner fails (e.g., the pod is deleted), the reservation expires after a grace period, allowing another worker to take over.

The Manager calculates expiration based on a heartbeat interval and a multiplier:

// From cache_service/manager/manager.go
reservation := &models.Reservation{
    Key:              reservationKey,
    OwnerID:          request.GetOwnerId(),
    HeartbeatSeconds: int64(heartbeat.Seconds()),
    ExpiresAt:        now.Add(heartbeat * time.Duration(m.heartbeatGracePeriodMultiplier)),
}

Asynchronous Caching for Plugins

For high-throughput scenarios, Flyte provides an AsyncClientImpl (found in flyteplugins/go/tasks/pluginmachinery/catalog/async_client_impl.go). This client uses internal workqueues to handle cache reads and writes without blocking the main execution loop of a plugin.

The AsyncClient manages two separate queues:

• Reader Queue: For Download requests (checking for cache hits).
• Writer Queue: For Upload requests (populating the cache after success).

This allows plugins to submit a cache request and check back later for a Future to be ready, improving the overall responsiveness of the Propeller.

Configuration

The Cache Service behavior is tuned via the cache_service configuration block. Key parameters include:

• server.port: The port for the Cache Service (default: 8094).
• heartbeatGracePeriodMultiplier: How many heartbeat intervals must pass before a reservation is considered abandoned (default: 3).
• maxReservationHeartbeat: The maximum allowed interval between heartbeats (default: 10s).

In the TaskActionReconciler, the heartbeat interval is tied to the controller's requeue duration:

const cacheReservationHeartbeatInterval = TaskActionDefaultRequeueDuration

If a worker dies without calling ReleaseReservation, other workers will be blocked until heartbeat * heartbeatGracePeriodMultiplier has elapsed. Once expired, the Manager allows a new owner to claim the key via UpdateIfExpiredOrOwned.