State Management and Transitions

Flyte plugins operate on a round-based execution model where the Handle method is invoked periodically by the Flyte engine (Propeller or the TaskAction controller). Because these invocations are discrete and the plugin instances themselves are stateless between rounds, Flyte provides a robust state management system to persist internal plugin data and drive the task's state machine through transitions.

State Persistence Interfaces

Plugins interact with their persisted state through two primary interfaces defined in flyteplugins/go/tasks/pluginmachinery/core/state.go: PluginStateReader and PluginStateWriter.

• PluginStateReader: Used at the beginning of a reconciliation round to retrieve the state persisted in the previous round. It provides Get(t interface{}) to deserialize the state into a provided struct.
• PluginStateWriter: Used during the round to record a new state. The Put(stateVersion uint8, v interface{}) method captures the state to be persisted.

A critical design constraint is that state written via Put is not accessible until the next reconciliation round. Furthermore, only the last call to Put within a single Handle execution is recorded; all previous calls in that same round are overwritten.

The Plugin State Manager

The PluginStateManager in executor/pkg/plugin/state_manager.go is the concrete implementation used by the Flyte executor to bridge the gap between the plugin's in-memory structures and the persisted storage (typically a Kubernetes CRD like TaskAction).

It uses Go's encoding/gob package to serialize and deserialize state into byte buffers. When the executor prepares to call a plugin's Handle method, it initializes the manager with the bytes stored in the task's status:

// From executor/pkg/controller/taskaction_controller.go
stateMgr := plugin.NewPluginStateManager(
    taskAction.Status.PluginState,
    taskAction.Status.PluginStateVersion,
)

After the plugin completes its work, the executor checks if any new state was written and persists it back:

// From executor/pkg/plugin/state_manager.go
func (m *PluginStateManager) GetNewState() (stateBytes []byte, version uint8, written bool) {
    return m.newStateBytes, m.newStateVersion, m.stateWritten
}

Transitions and the State Machine

The result of every Handle call is a Transition object. This object tells the Flyte engine what to do next: whether to requeue the task for another round, mark it as successful, or handle a failure.

A Transition consists of two parts:

1. PhaseInfo: Contains the current phase (e.g., PhaseRunning, PhaseSuccess), a version, and a reason string.
2. TransitionType: Defines the consistency model for the transition.

Transition Types and Consistency

Flyte defines two transition types in flyteplugins/go/tasks/pluginmachinery/core/transition.go:

• TransitionTypeEphemeral: The default and recommended type. It assumes the plugin logic is idempotent. It is eventually consistent, meaning the state written might not be immediately visible in the very next call in high-concurrency scenarios, but it is the most performant.
• TransitionTypeBarrier: (Deprecated) Attempted to provide stronger consistency guarantees but is no longer supported in favor of idempotent plugin design.

Standard State Structures

While plugins can define any serializable struct for their state, Flyte provides standard structures for common plugin types to ensure consistency across the ecosystem.

Kubernetes Plugin State

For plugins interacting with Kubernetes resources, flyteplugins/go/tasks/pluginmachinery/k8s/plugin.go defines a PluginState that tracks the resource's phase and version:

type PluginState struct {
    Phase        pluginsCore.Phase
    PhaseVersion uint32
    Reason       string
}

Web API Plugin State

For plugins interacting with external REST or gRPC APIs, flyteplugins/go/tasks/pluginmachinery/webapi/state.go provides a more complex State structure:

type State struct {
    Phase            Phase        `json:"phase,omitempty"`
    PhaseVersion     uint32       `json:"phaseVersion,omitempty"`
    ResourceMeta     ResourceMeta `json:"resourceMeta,omitempty"`
    SyncFailureCount int          `json:"syncFailureCount,omitempty"`
    ErrorMessage     string       `json:"error_message,omitempty"`
}

This structure includes ResourceMeta for tracking external identifiers and a SyncFailureCount to track transient API errors separately from Flyte's system-level retries.

Implementation Example: K8s Plugin Manager

The following example from executor/pkg/plugin/k8s/plugin_manager.go demonstrates the typical flow: reading state, performing logic, updating state, and returning a transition.

func (pm *PluginManager) Handle(ctx context.Context, tCtx pluginsCore.TaskExecutionContext) (pluginsCore.Transition, error) {
    // 1. Read previous state
    pluginState := PluginState{}
    if _, err := tCtx.PluginStateReader().Get(&pluginState); err != nil {
        return pluginsCore.UnknownTransition, err
    }

    // ... logic to check K8s resource status ...

    // 2. Prepare new state
    newPluginState := PluginState{
        Phase: pluginPhase,
        K8sPluginState: k8s.PluginState{
            Phase:        phaseInfo.Phase(),
            PhaseVersion: phaseInfo.Version(),
            Reason:       phaseInfo.Reason(),
        },
    }

    // 3. Persist state if changed
    if pluginState != newPluginState {
        if err := tCtx.PluginStateWriter().Put(pluginStateVersion, &newPluginState); err != nil {
            return pluginsCore.UnknownTransition, err
        }
    }

    // 4. Return transition to drive the engine
    return pluginsCore.DoTransition(phaseInfo), nil
}

This pattern ensures that Flyte can safely interrupt and resume task execution across different worker nodes or engine restarts without losing progress.