Implementing Core Plugins

Implementing a core plugin in Flyte allows you to define custom task execution logic that runs directly within the FlytePropeller process. Unlike container-based tasks, core plugins are compiled into the Flyte binary and are ideal for lightweight operations or orchestrating external services.

This guide walks you through building a core plugin by implementing the Plugin interface, using the Sleep plugin as a reference.

Define the Plugin Structure

First, define a struct that will hold any state or configuration needed by your plugin. For a simple plugin like Sleep, you might only need to track internal state like task start times.

package sleep

import (
	"time"
	"sync"
	"github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/core"
)

type Plugin struct {
	// Example state: tracking start times for idempotency
	taskStartTimes map[string]time.Time
	mu             sync.Mutex
}

func (p *Plugin) GetID() string {
	return "sleep-plugin"
}

func (p *Plugin) GetProperties() core.PluginProperties {
	return core.PluginProperties{}
}

Implement the Execution Logic

The Handle method is the heart of your plugin. It is invoked repeatedly by FlytePropeller until the task reaches a terminal state (Success or Failure).

Critical Requirements:

1. Non-blocking: Handle must return quickly. Do not perform long-running synchronous work here.
2. Idempotency: Handle may be called multiple times for the same task execution.

You use the TaskExecutionContext to access task metadata, inputs, and state.

func (p *Plugin) Handle(ctx context.Context, tCtx core.TaskExecutionContext) (core.Transition, error) {
	// 1. Read task inputs using the InputReader
	inputs, err := tCtx.InputReader().Get(ctx)
	if err != nil {
		return core.UnknownTransition, err
	}

	// 2. Determine the current state of the task
	// In this example, we check if we've waited long enough
	startTime := p.getOrAddTaskStartTime(tCtx)
	duration := time.Second * 10 // Simplified logic

	if time.Since(startTime) >= duration {
		// 3. Transition to Success
		return core.DoTransition(core.PhaseInfoSuccess(nil)), nil
	}

	// 4. Transition to Running to indicate work is still in progress
	return core.DoTransition(core.PhaseInfoRunning(core.DefaultPhaseVersion, nil)), nil
}

Manage Task Transitions

Flyte uses the Transition and PhaseInfo types to manage the task lifecycle. When you return a transition, FlytePropeller updates the task status in the Admin service.

Common transitions include:

• core.PhaseInfoQueued(version, reason): Task is waiting for resources.
• core.PhaseInfoRunning(version, info): Task is actively executing.
• core.PhaseInfoSuccess(info): Task completed successfully.
• core.PhaseInfoFailure(errorCode, message, info): Task failed with a specific error.

The PhaseInfo struct (defined in flyteplugins/go/tasks/pluginmachinery/core/phase.go) allows you to attach TaskInfo which can include external logs or custom metadata for the Flyte Console.

Handle Abort and Finalize

You must implement Abort and Finalize to ensure resources are cleaned up.

• Abort: Called when a user cancels the workflow or a downstream failure occurs.
• Finalize: Always called after Handle or Abort completes.

func (p *Plugin) Abort(ctx context.Context, tCtx core.TaskExecutionContext) error {
	// Clean up external resources if necessary
	return nil
}

func (p *Plugin) Finalize(ctx context.Context, tCtx core.TaskExecutionContext) error {
	// Final cleanup, e.g., removing internal state
	return nil
}

Register the Plugin

To make Flyte aware of your plugin, register it using the PluginRegistry. This is typically done in an init() function. You provide a LoadPlugin function that receives a SetupContext, which grants access to Kubernetes clients and metrics scopes.

import (
	"github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery"
)

func init() {
	pluginmachinery.PluginRegistry().RegisterCorePlugin(
		core.PluginEntry{
			ID:                  "sleep",
			RegisteredTaskTypes: []core.TaskType{"sleep"},
			LoadPlugin: func(ctx context.Context, iCtx core.SetupContext) (core.Plugin, error) {
				// Initialize your plugin using SetupContext
				return &Plugin{
					taskStartTimes: make(map[string]time.Time),
				}, nil
			},
			IsDefault: false,
		},
	)
}

Complete Example Summary

By implementing the Plugin interface, you have created a component that:

1. Registers itself for a specific TaskType via PluginEntry.
2. Initializes its environment using SetupContext.
3. Executes non-blocking logic in Handle using TaskExecutionContext.
4. Communicates state changes back to Flyte using Transition and PhaseInfo.
5. Ensures reliability through idempotent Abort and Finalize methods.