Skip to main content

Task Plugin System

The Flyte Task Plugin system allows developers to extend Flyte's execution capabilities by implementing specialized task types. While Flyte provides built-in support for containerized tasks, plugins enable deep integration with external compute systems like Spark, Ray, Dask, and various Kubernetes operators.

Plugin Types​

Flyte supports three primary categories of plugins:

  • Core Plugins: The base interface (core.Plugin) for any task execution logic. These require manual management of the execution loop, including idempotency and state transitions.
  • K8s Plugins: A specialized interface (k8s.Plugin) for managing Kubernetes Custom Resources (CRDs). This is the most common way to integrate with operators like the Spark or Ray operators.
  • WebAPI Plugins: Interfaces (webapi.AsyncPlugin and webapi.SyncPlugin) for interacting with external web services (REST/gRPC), managing the lifecycle of remote resources.

Implementing a K8s Plugin​

To implement a plugin for a Kubernetes operator, you must define a struct that implements the k8s.Plugin interface found in flyteplugins/go/tasks/pluginmachinery/k8s/plugin.go.

1. Define the Plugin Identity​

The BuildIdentityResource method provides a skeleton of the Kubernetes resource used for querying. It typically only includes the TypeMeta (Kind and APIVersion).

// From flyteplugins/go/tasks/plugins/k8s/spark/spark.go
func (sparkResourceHandler) BuildIdentityResource(ctx context.Context, taskCtx pluginsCore.TaskExecutionMetadata) (client.Object, error) {
	return &sparkOp.SparkApplication{
		TypeMeta: metav1.TypeMeta{
			Kind:       KindSparkApplication,
			APIVersion: sparkOp.SchemeGroupVersion.String(),
		},
	}, nil
}

2. Build the Resource Specification​

The BuildResource method is responsible for creating the actual Kubernetes object (e.g., a SparkApplication or RayJob) that will be submitted to the cluster. You use the TaskExecutionContext to read the task template and custom configuration.

// From flyteplugins/go/tasks/plugins/k8s/spark/spark.go
func (sparkResourceHandler) BuildResource(ctx context.Context, taskCtx pluginsCore.TaskExecutionContext) (client.Object, error) {
	taskTemplate, err := taskCtx.TaskReader().Read(ctx)
	if err != nil {
		return nil, errors.Errorf(errors.BadTaskSpecification, "unable to fetch task specification [%v]", err.Error())
	}

	// Unmarshal custom configuration (e.g., SparkJob proto)
	sparkJob := plugins.SparkJob{}
	err = utils.UnmarshalStruct(taskTemplate.GetCustom(), &sparkJob)
	
	// Construct the CRD using helper functions
	sparkConfig := getSparkConfig(taskCtx, &sparkJob)
	driverSpec, err := createDriverSpec(ctx, taskCtx, sparkConfig, &sparkJob)
	executorSpec, err := createExecutorSpec(ctx, taskCtx, sparkConfig, &sparkJob)
	
	app := createSparkApplication(&sparkJob, sparkConfig, driverSpec, executorSpec)
	return app, nil
}

3. Map Resource Status to Flyte Phases​

The GetTaskPhase method monitors the Kubernetes resource and maps its status to a Flyte PhaseInfo. This method is called repeatedly by FlytePropeller until a terminal phase is reached.

// From flyteplugins/go/tasks/plugins/k8s/spark/spark.go
func (sparkResourceHandler) GetTaskPhase(ctx context.Context, pluginContext k8s.PluginContext, resource client.Object) (pluginsCore.PhaseInfo, error) {
	app := resource.(*sparkOp.SparkApplication)
	
	occurredAt := time.Now()
	switch app.Status.AppState.State {
	case sparkOp.NewState:
		return pluginsCore.PhaseInfoQueued(occurredAt, pluginsCore.DefaultPhaseVersion, "job queued"), nil
	case sparkOp.SubmittedState, sparkOp.PendingSubmissionState:
		return pluginsCore.PhaseInfoInitializing(occurredAt, pluginsCore.DefaultPhaseVersion, "job submitted", nil), nil
	case sparkOp.FailedState:
		reason := fmt.Sprintf("Spark Job Failed with Error: %s", app.Status.AppState.ErrorMessage)
		return pluginsCore.PhaseInfoRetryableFailure(errors.DownstreamSystemError, reason, nil), nil
	case sparkOp.CompletedState:
		return pluginsCore.PhaseInfoSuccess(nil), nil
	default:
		return pluginsCore.PhaseInfoRunning(pluginsCore.DefaultPhaseVersion, nil), nil
	}
}

Using flytek8s Utilities​

The flytek8s package provides utilities to simplify building Kubernetes resources from Flyte task definitions. The most critical utility is ToK8sPodSpec, which converts Flyte task metadata, container information, and resource requirements into a standard Kubernetes PodSpec.

// Example usage in flyteplugins/go/tasks/plugins/k8s/spark/spark.go
podSpec, _, primaryContainerName, err := flytek8s.ToK8sPodSpec(ctx, taskCtx)
if err != nil {
    return nil, err
}

This utility automatically handles:

  • Injecting Flyte-specific environment variables.
  • Setting up resource requests and limits (CPU, Memory, GPU).
  • Configuring storage mounts and sidecars (like Co-Pilot).
  • Applying default annotations and labels from the global K8sPluginConfig.

Registering the Plugin​

Plugins must be registered in the global PluginRegistry to be discoverable. This is typically done in an init() function within the plugin package.

// From flyteplugins/go/tasks/plugins/k8s/spark/spark.go
func init() {
	// Register the CRD scheme with the controller-runtime
	if err := sparkOp.AddToScheme(scheme.Scheme); err != nil {
		panic(err)
	}

	pluginmachinery.PluginRegistry().RegisterK8sPlugin(
		k8s.PluginEntry{
			ID:                  "spark",
			RegisteredTaskTypes: []pluginsCore.TaskType{"spark"},
			ResourceToWatch:     &sparkOp.SparkApplication{},
			Plugin:              sparkResourceHandler{},
			IsDefault:           false,
		})
}

Advanced Configuration​

Resource Deletion and Finalizers​

By default, Flyte K8s plugins delete the managed resource upon reaching a terminal state to free up cluster resources. This behavior is controlled by PluginProperties and the global K8sPluginConfig.

  • DisableDeleteResourceOnFinalize: Set this to true in GetProperties() if you need the Kubernetes resource to persist (e.g., for log access) after the task completes.
  • GeneratedNameMaxLength: Kubernetes resources often have strict name length limits. You can override the default (50) if your operator requires shorter names (e.g., Ray uses 47).

GPU and Accelerator Support​

Flyte provides a unified way to request accelerators across different cloud providers and hardware types. The K8sPluginConfig (defined in flyteplugins/go/tasks/pluginmachinery/flytek8s/config/config.go) allows platform operators to map Flyte accelerator types to Kubernetes node labels and resource names.

// Default configuration for NVIDIA GPUs
GpuDeviceNodeLabel:        "k8s.amazonaws.com/accelerator",
GpuResourceName:           "nvidia.com/gpu",

When a task requests a GPU, flytek8s.ToK8sPodSpec uses these configurations to add the appropriate nodeSelector and resources to the generated PodSpec.