Managing Task Execution Reservations

Flyte provides a reservation system to coordinate task execution when "serializable" caching is enabled. This system ensures that if multiple workers attempt to execute the same task with the same inputs simultaneously, only one worker proceeds to run the task and populate the cache, while others wait for the result.

Coordinating Serializable Execution

When a task is marked as discoverable and serializable, Flyte must prevent redundant work. Without reservations, a "thundering herd" of workers might all see a cache miss and start executing the same expensive computation.

The TaskActionReconciler in executor/pkg/controller/taskaction_cache.go manages this by attempting to acquire a reservation before execution:

// From executor/pkg/controller/taskaction_cache.go

if cacheCfg.serializable {
    reservation, err := r.Catalog.GetOrExtendReservation(ctx, cacheCfg.key, cacheCfg.ownerID, cacheReservationHeartbeatInterval)
    if err != nil {
        return pluginsCore.UnknownTransition, false, fmt.Errorf("acquiring cache reservation: %w", err)
    }
    
    if reservation.GetOwnerId() == cacheCfg.ownerID {
        // We own the reservation, proceed to execute the task
        return pluginsCore.UnknownTransition, false, nil
    }

    // Someone else owns it, transition to a waiting state
    info := cacheTaskInfo(corepb.CatalogCacheStatus_CACHE_MISS, "waiting for serialized cache owner")
    phaseInfo := pluginsCore.PhaseInfoWaitingForCache(taskAction.Status.PluginPhaseVersion, info)
    phaseInfo.WithReason(fmt.Sprintf("waiting for cache to be populated by reservation owner %q", reservation.GetOwnerId()))
    return pluginsCore.DoTransition(phaseInfo), true, nil
}

The Reservation Lifecycle

The reservation system is implemented in the Manager class within cache_service/manager/manager.go. It follows a strict lifecycle of acquisition, heartbeating, and release.

1. Acquisition and Extension

The GetOrExtendReservation method is the primary entry point. It calculates an expiration time based on the requested heartbeat interval and a configured grace period.

Internally, the Manager uses the ReservationRepo to perform atomic updates. The UpdateIfExpiredOrOwned method ensures that a reservation can only be updated if the caller already owns it or if the current reservation has expired.

// 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)),
}

// ... logic to Create or UpdateIfExpiredOrOwned ...

2. Heartbeating

To maintain ownership during long-running tasks, the worker must periodically call GetOrExtendReservation. Each call pushes the ExpiresAt timestamp further into the future. If a worker stops heartbeating (e.g., due to a crash), the reservation eventually expires, allowing another worker to claim it.

3. Release

Once a task completes successfully and populates the cache, or if it fails terminally, the reservation must be released so other workers are not blocked indefinitely.

// From executor/pkg/controller/taskaction_cache.go

func (r *TaskActionReconciler) releaseCacheReservation(ctx context.Context, cacheCfg *taskCacheConfig) error {
    if r.Catalog == nil || cacheCfg == nil || !cacheCfg.serializable {
        return nil
    }

    return r.Catalog.ReleaseReservation(ctx, cacheCfg.key, cacheCfg.ownerID)
}

The Manager.ReleaseReservation method is idempotent. It uses DeleteByKeyAndOwner to ensure that an owner can only delete their own active reservation.

Client-Side State Tracking

The Flyte catalog client uses ReservationEntry and ReservationCache (defined in flyteplugins/go/tasks/pluginmachinery/catalog/client.go) to track the state of reservations in-memory.

• ReservationEntry: A data transfer object that encapsulates the expiresAt time, heartbeatInterval, and the current ownerID returned by the cache service.
• ReservationCache: Used by the client to store the last known ReservationStatus and a timestamp for a specific owner, helping the executor decide when to retry acquisition.

Configuration

The behavior of the reservation system is tuned via the Manager configuration:

Parameter	Default	Description
heartbeatGracePeriodMultiplier	3	The number of heartbeat intervals that must pass without an update before a reservation is considered expired.
maxReservationHeartbeatInterval	10s	The maximum interval allowed for heartbeats. If a client requests a longer interval, it is capped at this value.

These settings are applied during the initialization of the Manager in cache_service/manager/manager.go:

func New(cfg *cacheconfig.Config, outputs interfaces.CachedOutputRepo, reservations interfaces.ReservationRepo) *Manager {
    maxHeartbeat := cfg.MaxReservationHeartbeat.Duration
    if maxHeartbeat <= 0 {
        maxHeartbeat = 10 * time.Second
    }

    graceMultiplier := cfg.HeartbeatGracePeriodMultiplier
    if graceMultiplier <= 0 {
        graceMultiplier = 3
    }

    return &Manager{
        outputs:                         outputs,
        reservations:                    reservations,
        heartbeatGracePeriodMultiplier:  graceMultiplier,
        maxReservationHeartbeatInterval: maxHeartbeat,
    }
}

The actual expiration time is calculated as now + (heartbeatInterval * heartbeatGracePeriodMultiplier). This provides a buffer to account for network latency or temporary worker stalls.