Web API and Async Plugins

Flyte provides a specialized framework for developing plugins that offload task execution to external web services (e.g., REST, gRPC). This framework handles common concerns like rate limiting, caching, and state management, allowing you to focus on the interaction with the remote service.

Registering a Web API Plugin

To register a Web API plugin, you define a PluginEntry and use the PluginRegistry to make it available to Flyte. The PluginLoader is a lazy-loading function that initializes your plugin instance.

import (
    "context"
    "encoding/gob"
    "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery"
    "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/webapi"
    pluginsCore "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/core"
)

func RegisterMyPlugin() {
    // Custom state types must be registered with gob for persistence
    gob.Register(MyResourceMeta{})
    gob.Register(MyResource{})

    pluginmachinery.PluginRegistry().RegisterRemotePlugin(webapi.PluginEntry{
        ID:                 "my_remote_task",
        SupportedTaskTypes: []pluginsCore.TaskType{"my_remote_task"},
        PluginLoader: func(ctx context.Context, iCtx webapi.PluginSetupContext) (webapi.AsyncPlugin, error) {
            return &MyPlugin{
                metricScope: iCtx.MetricsScope(),
            }, nil
        },
    })
}

Implementing the AsyncPlugin Interface

The AsyncPlugin interface is designed for long-running remote tasks. It separates the concerns of launching a task, fetching its state from the network, and mapping that state to Flyte phases.

1. Resource Requirements

Use ResourceRequirements to define the resource namespace for token allocation. This helps Flyte respect quotas on the remote service.

func (p *MyPlugin) ResourceRequirements(ctx context.Context, tCtx webapi.TaskExecutionContextReader) (
    namespace pluginsCore.ResourceNamespace, constraints pluginsCore.ResourceConstraintsSpec, err error) {
    // Use a namespace corresponding to the remote service or project
    return "my-remote-service", pluginsCore.ResourceConstraintsSpec{}, nil
}

2. Creating the Remote Resource

The Create method is called to launch the task. It must be idempotent because Flyte may call it multiple times (e.g., after a restart). It returns a ResourceMeta object which is persisted by Flyte.

func (p *MyPlugin) Create(ctx context.Context, tCtx webapi.TaskExecutionContextReader) (
    webapi.ResourceMeta, webapi.Resource, error) {
    
    // Prepare the request using task inputs
    inputReader := tCtx.InputReader()
    // ... logic to call remote API ...

    resp, err := p.client.CreateTask(ctx, request)
    if err != nil {
        return nil, nil, err
    }

    // Return metadata needed to track this task later
    return MyResourceMeta{
        RemoteID: resp.ID,
    }, nil, nil
}

3. Fetching and Mapping Status

Flyte uses an internal auto-refresh cache to periodically sync the state of active tasks.

• Get: Performs the actual network call to fetch the latest state.
• Status: Maps the fetched state to a Flyte PhaseInfo. This method should be efficient and avoid network calls.

func (p *MyPlugin) Get(ctx context.Context, tCtx webapi.GetContext) (webapi.Resource, error) {
    meta := tCtx.ResourceMeta().(MyResourceMeta)
    // Perform network call to get status
    resp, err := p.client.GetTaskStatus(ctx, meta.RemoteID)
    if err != nil {
        return nil, err
    }
    return MyResource{State: resp.State, Message: resp.Message}, nil
}

func (p *MyPlugin) Status(ctx context.Context, tCtx webapi.StatusContext) (pluginsCore.PhaseInfo, error) {
    // Retrieve the resource returned by Get()
    resource := tCtx.Resource().(MyResource)
    
    switch resource.State {
    case "SUCCEEDED":
        return pluginsCore.PhaseInfoSuccess(nil), nil
    case "FAILED":
        return pluginsCore.PhaseInfoFailure("RemoteError", resource.Message, nil), nil
    case "RUNNING":
        return pluginsCore.PhaseInfoRunning(1, nil), nil
    default:
        return pluginsCore.PhaseInfoQueued(1, "waiting"), nil
    }
}

4. Deleting the Resource

The Delete method is called when a task is aborted or finalized. It should clean up resources on the remote service.

func (p *MyPlugin) Delete(ctx context.Context, tCtx webapi.DeleteContext) error {
    if tCtx.ResourceMeta() == nil {
        return nil
    }
    meta := tCtx.ResourceMeta().(MyResourceMeta)
    return p.client.AbortTask(ctx, meta.RemoteID)
}

Configuring the Plugin

You can control the behavior of the Web API machinery using PluginConfig. This includes rate limiting to protect the remote service and caching parameters for the auto-refresh cycle.

func (p *MyPlugin) GetConfig() webapi.PluginConfig {
    return webapi.PluginConfig{
        ReadRateLimiter: webapi.RateLimiterConfig{
            QPS:   30,
            Burst: 300,
        },
        WriteRateLimiter: webapi.RateLimiterConfig{
            QPS:   20,
            Burst: 200,
        },
        Caching: webapi.CachingConfig{
            Size:              100000,
            ResyncInterval:    config.Duration{Duration: 30 * time.Second},
            Workers:           10,
            MaxSystemFailures: 5,
        },
    }
}

Best Practices and Gotchas

• Idempotency: Both Create and Delete must be idempotent. Use the task execution ID from tCtx.TaskExecutionMetadata().GetID() as a client token if the remote service supports it.
• State Persistence: Any struct returned as ResourceMeta in Create must be registered with gob.Register() in your plugin's initialization (e.g., in RegisterMyPlugin). This object is serialized and stored in the Flyte state.
• Error Handling: Distinguish between system errors (e.g., network timeout) and user errors. System errors returned from Create or Get will trigger retries based on the plugin configuration.
• Efficiency: Keep the Status method logic-only. The Get method is where network latency is expected and handled by background workers in the autorefreshcache.
• ResourceMeta vs Resource: ResourceMeta is the minimal set of keys (like a job ID) needed to identify the task. Resource is the full state object fetched from the remote service. Only ResourceMeta is persisted; Resource is kept in an in-memory cache.