Skip to main content

Plugins and Integrations

Two extension paths​

flyte-sdk extends execution in two ways. Distributed-runtime plugins are selected through TaskPluginRegistry: Spark, Dask, Ray, Databricks, and PyTorch provide task implementations and serialize operator-specific configuration. External-service integrations use ConnectorRegistry: a task serializes provider configuration, and AsyncConnectorExecutorMixin submits and polls the external job.

TaskTemplate.custom_config() -> task custom configuration
        -> ConnectorRegistry[(task_type, task_type_version)]
        -> create() -> ResourceMeta -> get() ... -> Resource

ConnectorRegistry.register() indexes an instance by task type and version. Duplicate keys raise ValueError unless override=True; get_connector() raises FlyteConnectorNotFound for an unknown key. Provider connectors register at import time, so the relevant plugin must be installed and imported in the executing runtime.

Build a connector​

Implement AsyncConnector, a ResourceMeta dataclass, and a task inheriting AsyncConnectorExecutorMixin and TaskTemplate. The custom connector example in examples/apps/custom_connector_app/my_connector/connector.py defines BatchJobMetadata(job_id, created_at), returns it from create(), reports RUNNING until five seconds have elapsed, then returns a successful Resource with outputs={"result": ...}, and calls ConnectorRegistry.register(BatchJobConnector()).

ResourceMeta.encode() JSON-serializes dataclasses.asdict(self) to UTF-8, and decode() reconstructs the declared dataclass. Resource carries a Flyte phase, message, log links, native outputs, and custom information. AsyncConnector requires asynchronous create(task_template, output_prefix, inputs, ...), get(resource_meta, ...), and delete(resource_meta, ...); deletion is expected to be idempotent. The base metrics and log methods currently raise `NotImplementedError.

AsyncConnectorExecutorMixin.execute() resolves the connector, prepares code bundles and inputs when needed, resolves secret references through environment variables, calls create(), polls get() every three seconds, records log links, and converts outputs. A terminal failure becomes RuntimeError. Secret entries are environment-variable names, not values; a missing variable raises ValueError. The companion BatchJobTask uses task type batch_job, builds a NativeInterface, and serializes its timeout_seconds configuration.

SQL connectors​

BigQuery, Snowflake, and Databricks use the same lifecycle but provider-specific configuration, authentication, links, and dataframe URIs.

BigQuery​

Use BigQueryTask with BigQueryConfig. An optional output_dataframe_type creates a results output, and typed inputs become query parameters:

from flyteplugins.bigquery.task import BigQueryConfig, BigQueryTask

import flyte
from flyte.io._dataframe.dataframe import DataFrame

bigquery_task = BigQueryTask(
    name="bigquery",
    inputs={"version": int},
    output_dataframe_type=DataFrame,
    plugin_config=BigQueryConfig(ProjectID="dogfood-gcp-dataplane"),
    query_template="SELECT * from dataset.flyte_table3;",
)

flyte.TaskEnvironment.from_task("bigquery_env", bigquery_task)

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.with_runcontext(mode="local").run(bigquery_task, 123)
    print(run.url)

BigQueryTask normalizes SQL whitespace and emits ANSI tasks_pb2.Sql. Its custom config contains project, location, domain, optional QueryJobConfig.to_api_repr()["query"], and a credential secret reference. The resolved credential must be JSON service-account content, parsed with service_account.Credentials.from_service_account_info.

BigQueryConnector.create() derives input Python types from the Flyte interface, maps them to BigQuery types, and submits task_template.sql.statement. get() maps job state to a Flyte phase and adds a BigQuery Console link. When a destination exists, successful output is DataFrame(uri="bq://project:dataset.table"); a successful query without a destination produces no dataframe result. delete() calls cancel_job().

Snowflake​

SnowflakeConfig contains account, database, schema, warehouse, and user, plus optional connection_kwargs. Snowflake supports named %(name)s parameters and batch=True:

from flyteplugins.snowflake.task import Snowflake, SnowflakeConfig

sf_config = SnowflakeConfig(
    user="KEVIN",
    account="PWGJLTH-XKB21544",
    database="FLYTE",
    schema="PUBLIC",
    warehouse="COMPUTE_WH",
)

snowflake_insert_task = Snowflake(
    name="snowflake",
    inputs={"id": list[int], "name": list[str], "age": list[int]},
    plugin_config=sf_config,
    query_template=(
        "INSERT INTO FLYTE.PUBLIC.TEST (ID, NAME, AGE) "
        "VALUES (%(id)s, %(name)s, %(age)s);"
    ),
    snowflake_private_key="snowflake",
    batch=True,
)

_expand_batch_query() expands one VALUES (...) clause. Every list must be non-empty and all lists must have equal lengths. The task serializes connection kwargs, batch, and key-pair secret references; _to_env_var() optionally adds secret_group, replaces hyphens with underscores, and uppercases the name.

SnowflakeConnector.create() validates required metadata and submits execute_async() in an executor. get() polls status, adds a dashboard link, and returns snowflake://user/account/warehouse/database/schema/query_id as the dataframe URI when output exists. delete() executes SYSTEM$CANCEL_QUERY and closes the connection. Key-pair decoding uses SNOWFLAKE_PRIVATE_KEY and optionally SNOWFLAKE_PRIVATE_KEY_PASSPHRASE.

Databricks​

Databricks extends Spark with databricks_conf, databricks_instance, and databricks_token. databricks_conf must contain either existing_cluster_id or new_cluster. DatabricksFunctionTask serializes Spark and Databricks settings into SparkJob and emits the token as a secret reference.

DatabricksConnector.create() posts to /api/2.1/jobs/runs/submit with a bearer token. The instance comes from task configuration, then FLYTE_DATABRICKS_INSTANCE; neither raises ValueError. The job-specification helper injects a pinned GitHub spark_python_task entrypoint. get() polls /runs/get, uses result_state once available after the lifecycle state, and returns a Databricks console link. Non-200 calls raise RuntimeError; delete() posts to /runs/cancel.

Distributed computing​

RuntimeConfiguration/taskRuntime objectSerialized deployment
SparkSpark, PysparkFunctionTaskflyte.ctx().data["spark_session"]SparkJob, Spark/Hadoop settings, paths, pod templates
DaskDask, DaskTaskdistributed.ClientDaskJob, scheduler and worker groups
RayRayJobConfig, RayFunctionTaskRay APIs after ray.init()KubeRay head/worker specs and runtime environment
PyTorchElastic, TorchFunctionTaskElastic worker processesKubeflow DistributedPyTorchTrainingTask

Spark, Dask, and Ray​

PysparkFunctionTask.pre() creates or reuses a session named FlyteSpark. In cluster mode it zips the code bundle and calls sparkContext.addPyFile. Spark.__post_init__() initializes missing Spark/Hadoop dictionaries, while post() stops the session only for debug action a0.

spark_conf = Spark(
    spark_conf={
        "spark.driver.memory": "3000M",
        "spark.executor.memory": "1000M",
        "spark.executor.cores": "1",
        "spark.executor.instances": "2",
    },
)

spark_env = flyte.TaskEnvironment(
    name="spark_env",
    resources=flyte.Resources(cpu=(1, 2), memory=("3000Mi", "5000Mi")),
    plugin_config=spark_conf,
    image=image,
    depends_on=[task_env],
)

@spark_env.task
async def hello_spark_nested(partitions: int = 3) -> float:
    spark = flyte.ctx().data["spark_session"]
    count = spark.sparkContext.parallelize(range(1, partitions + 1), partitions).map(f).reduce(add)
    return await get_pi(count, partitions)

Dask scheduler and worker images must contain dask[distributed] and compatible Python environments. DaskTask.pre() registers DownloadCodeBundleWorkerPlugin and DownloadCodeBundleSchedulerPlugin with a live client; their setup methods put "." on sys.path and download the bundle. Ray's pre() calls ray.init(address=...) when needed, and in cluster mode supplies the working directory and excludes local archive/history paths. Ray serialization emits both deprecated base64-JSON runtime_env and runtime_env_yaml.

ray_config = RayJobConfig(
    head_node_config=HeadNodeConfig(ray_start_params={"log-color": "True"}),
    worker_node_config=[WorkerNodeConfig(group_name="ray-group", replicas=2)],
    runtime_env={"pip": ["numpy", "pandas"]},
    enable_autoscaling=False,
    shutdown_after_job_finishes=True,
    ttl_seconds_after_finished=300,
)

ray_env = flyte.TaskEnvironment(
    name="ray_env",
    plugin_config=ray_config,
    image=image,
    resources=flyte.Resources(cpu=(3, 4), memory=("3000Mi", "5000Mi")),
    depends_on=[task_env],
)

PyTorch Elastic​

Elastic controls nnodes (integer or range such as "2:4"), processes per node, rendezvous, monitoring, restart count, and NCCL timeouts. RunPolicy adds cleanup, TTL, deadline, and backoff settings.

torch_env = flyte.TaskEnvironment(
    name="torch_env",
    resources=flyte.Resources(cpu=(1, 2), memory=("1Gi", "2Gi"), gpu="T4:1"),
    plugin_config=Elastic(
        nproc_per_node=1,
        # if you want to do local testing set nnodes=1
        nnodes=2,
    ),
    image=image,
)

TorchFunctionTask uses task type python-task for nnodes == 1 and pytorch otherwise. Local execution calls the function directly. Remote execution invokes elastic_launch on the main thread, returns rank 0's value, and converts ChildFailedError to RuntimeError; the main thread is required for signal-handler installation. pre() conditionally sets PYTHONUNBUFFERED=1, OMP_NUM_THREADS=1 for multi-process tasks, and the configured NCCL variables, preserving existing values. A zombie watchdog handles the all-workers-dead NCCL deadlock and can force process exit.

Type integrations​

OmegaConf​

OmegaConf is a TypeEngine extension. Importing the package registers DictConfigTransformer and ListConfigTransformer. OmegaConfTransformerBase advertises a STRUCT literal with MessagePack metadata, serializes with serialize_omegaconf() and msgpack, and accepts only MESSAGEPACK binary tags. Wrong tags or container types raise TypeError.

Pandera dataframes​

Pandera combines schema validation with Flyte's DataFrameTransformerEngine. Use typed annotations and Annotated for validation behavior:

class EmployeeSchema(pa.DataFrameModel):
    employee_id: int = pa.Field(ge=0)
    name: str

@env.task(report=True)
async def build_valid_employees() -> pt.DataFrame[EmployeeSchema]:
    return pd.DataFrame({
        "employee_id": [1, 2, 3],
        "name": ["Ada", "Grace", "Barbara"],
    })

@env.task(report=True)
async def pass_through_with_error_warn(
    df: Annotated[pt.DataFrame[EmployeeSchema], ValidationConfig(on_error="warn")],
) -> Annotated[pt.DataFrame[EmployeeSchema], ValidationConfig(on_error="warn")]:
    del df["name"]
    return df

PanderaDataFrameTransformer validates with lazy=True on encode and decode, delegates storage to DataFrameTransformerEngine, and emits HTML through flyte.report. on_error="warn" logs and returns original data; "raise" propagates failure. Output URI memoization and TaskContext.in_driver_literal_conversion suppress duplicate reports.

The pandas transformer resolves pandera.typing.pandas.DataFrame to pandas. Polars supports DataFrame and LazyFrame. PySpark SQL supports pandera.typing.pyspark_sql.DataFrame, reads validated.pandera.errors, and raises RuntimeError for validation errors; Spark parquet dataframe handlers are also required. Registration is conditional and warns when optional extras are absent. Import Pandera typing modules before registration to avoid duplicate class objects and pickle fallback.

MLflow​

Put @mlflow_run outside @env.task. auto reuses a propagated parent, new starts an independent run, and nested starts a child with mlflow.parentRunId. mlflow_config() overrides tracking, experiment/run identity, tags, autologging, links, and start_run kwargs. Experiment name/ID and run name/ID are mutually exclusive.

@mlflow_run
@env.task(links=(Mlflow()))
async def child_task_inherits_mode(x: int) -> str:
    run = get_mlflow_run()
    mlflow.log_metric("child_metric", x * 2)
    return run.info.run_id

@mlflow_run(run_mode="new")
@env.task(links=(Mlflow()))
async def child_task_new_run(x: int) -> str:
    run = get_mlflow_run()
    mlflow.log_metric("child_metric", x * 2)
    return run.info.run_id

Mlflow.get_link() prefers an explicit URL, suppresses inherited links for run_mode="new", and otherwise returns _mlflow_link. Decorating a trace with @mlflow_run raises; traces should call get_mlflow_run().

Weights & Biases​

W&B uses the outermost @wandb_init/@wandb_sweep decorators. wandb_config() overrides child settings and propagates IDs. Wandb.get_link() resolves project, entity, host, and ID from decorator fields first and context second. Non-distributed auto uses a parent ID when available; shared requires a parent; new creates an independent target. Distributed links additionally use rank_scope="global" or "worker".

@wandb_init
@env.task
async def parent_task() -> str:
    run = get_wandb_run()
    run.log({"parent_metric": 100})

    with wandb_config(name="child-run", tags=["child-task"]):
        result1 = await child_task_new_run(5)

    result2 = await child_task_inherits_mode(10)

    with wandb_config(run_mode="new"):
        result3 = await child_task_inherits_mode(25)

    return run.id

Without project/entity, W&B link generation returns only the host. get_wandb_run() returns None outside a decorated task or on non-logging ranks, so callers must guard before .log(). WandbSweep links to a propagated sweep or the project sweep list.

Hydra launcher​

Hydra selects Flyte through configuration rather than a Flyte task decorator. FlyteLauncherConfig registers hydra/launcher=flyte; defaults are mode="remote", wait=True, and wait_max_workers=32.

python train.py hydra/launcher=flyte hydra.launcher.mode=remote

python train.py --multirun \
    hydra/launcher=flyte hydra.launcher.mode=remote \
    optimizer.lr=0.001,0.01,0.1 training.epochs=10,20

python train.py --multirun \
    hydra/launcher=flyte hydra.launcher.wait=false \
    optimizer.lr=0.001,0.01,0.1

Local mode runs sequentially in-process. Remote mode submits jobs, prints URLs immediately, and optionally waits concurrently; wait_max_workers=None removes the cap. Results are wrapped in FlyteHydraRunResult, preserving URL/run delegation and float/iteration conversion.

Configuration and troubleshooting​

IntegrationRequired setupCommon observed failure
ConnectorsImport provider package; expose referenced secrets as environment variablesUnknown type/version raises FlyteConnectorNotFound; missing secret raises ValueError
BigQueryProjectID, optional location, BigQuery and Google-auth extrasUnsupported input misses the type map; no destination means no results
DatabricksdatabricksConf with existing_cluster_id or new_cluster, token, and instance or FLYTE_DATABRICKS_INSTANCEMissing instance raises ValueError; non-200 calls raise RuntimeError
SnowflakeRequired metadata and optional key/passphrase or connection_kwargsBatch lists must be non-empty, equal length, and match one named VALUES clause
Spark/Dask/RayCompatible runtime/operator packages and images; Dask images need dask[distributed]Code-bundle setup depends on cluster mode; Spark teardown is debug-only; Dask plugins need a live client
PyTorchElastic, rendezvous settings, and suitable runtime/GPU imageSingle-node uses python-task; deterministic OOM may repeat until restart limit; NCCL watchdog can exit
OmegaConf/PanderaRegister transformers with optional extras; import Pandera typing modules firstNon-MessagePack raises TypeError; duplicate Pandera classes can cause pickle fallback; PySpark needs Spark handlers
MLflow/W&BOutermost decorators, provider credentials, and project/entity configurationActive runs may be absent; W&B without project/entity returns only host; traces use accessors
Hydrahydra/launcher=flyte, mode/wait settingsRemote sweeps submit before waiting; wait=false intentionally does not wait