Skip to main content

Databricks and Spark Distributed Compute

Choose the execution path​

Use Spark when the task should run as a distributed PySpark application through Flyte's native Spark-on-Kubernetes integration. Use Databricks when Flyte should submit the application to a Databricks Jobs API run and monitor that remote run. In both cases, attach the configuration to a flyte.TaskEnvironment; the task/plugin registry selects PysparkFunctionTask for Spark and DatabricksFunctionTask for Databricks.

Native Spark on Kubernetes​

The native path uses Spark from flyteplugins.spark.task. The example builds a Spark image, configures executor sizing and Kubernetes-related Spark properties, and attaches the configuration to the environment:

import flyte
from flyteplugins.spark.task import Spark

image = (
    flyte.Image.from_base("apache/spark-py:v3.4.0")
    .clone(name="spark", python_version=(3, 10), registry="ghcr.io/flyteorg", extendable=True)
    .with_pip_packages("flyteplugins-spark>=2.0.0")
)

spark_conf = Spark(
    spark_conf={
        "spark.driver.memory": "3000M",
        "spark.executor.memory": "1000M",
        "spark.executor.cores": "1",
        "spark.executor.instances": "2",
        "spark.driver.cores": "1",
        "spark.kubernetes.file.upload.path": "/opt/spark/work-dir",
        "spark.jars": "https://storage.googleapis.com/hadoop-lib/gcs/gcs-connector-hadoop3-latest.jar,https://repo1.maven.org/maven2/org/apache/hadoop/hadoop-aws/3.2.2/hadoop-aws-3.2.2.jar,https://repo1.maven.org/maven2/com/amazonaws/aws-java-sdk-bundle/1.12.262/aws-java-sdk-bundle-1.12.262.jar",
    },
)

spark_env = flyte.TaskEnvironment(
    name="spark_env",
    resources=flyte.Resources(),
    plugin_config=spark_conf,
    image=image,
)

Spark also accepts hadoop_conf, executor_path, applications_path, driver_pod, and executor_pod. When spark_conf or hadoop_conf is omitted, Spark.__post_init__() changes each value to an empty dictionary.

Define the task as an async function and retrieve the session injected by PysparkFunctionTask.pre():

from operator import add

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

    return await get_pi(count, partitions)

At execution time, pre() creates or reuses a PySpark session with the application name FlyteSpark and returns it as {"spark_session": sess}. If the execution is in-cluster and a Flyte code bundle is available, it creates a temporary flyte_code_bundle.zip archive from flyte.ctx().code_bundle.destination and registers that archive with sparkContext.addPyFile().

custom_config() serializes a SparkJob protobuf. It places the configured Spark and Hadoop dictionaries in sparkConf and hadoopConf, uses applications_path or local:// plus SerializationContext.get_entrypoint_path() for mainApplicationFile, and uses executor_path or SerializationContext.interpreter_path for executorPath. It sets the application type to Python and includes converted driver and executor pod templates when supplied.

The task defaults to task_type="spark" and debuggable=True. Its post() hook stops a Spark session only when the current action is named a0, the debug-mode parent action; it does not unconditionally stop a session after every invocation.

Databricks runtime​

Use Databricks from flyteplugins.databricks.task when the same PySpark task shape should be submitted to Databricks:

import flyte
from flyteplugins.databricks.task import Databricks

image = (
    flyte.Image.from_base("databricksruntime/standard:16.4-LTS")
    .clone(name="spark", registry="ghcr.io/flyteorg", extendable=True)
    .with_env_vars({"UV_PYTHON": "/databricks/python3/bin/python"})
    .with_pip_packages("flyteplugins-databricks")
    .with_env_vars({"AWS_REGION": "us-west-1", "LOG_LEVEL": 10})
)

databricks_conf = Databricks(
    spark_conf={
        "spark.driver.memory": "2000M",
        "spark.executor.memory": "1000M",
        "spark.executor.cores": "1",
        "spark.executor.instances": "2",
        "spark.driver.cores": "1",
        "spark.kubernetes.file.upload.path": "/opt/spark/work-dir",
        "spark.jars": "https://storage.googleapis.com/hadoop-lib/gcs/gcs-connector-hadoop3-latest.jar,https://repo1.maven.org/maven2/org/apache/hadoop/hadoop-aws/3.2.2/hadoop-aws-3.2.2.jar,https://repo1.maven.org/maven2/com/amazonaws/aws-java-sdk-bundle/1.12.262/aws-java-sdk-bundle-1.12.262.jar",
    },
    executor_path="/databricks/python3/bin/python",
    databricks_conf={
        "run_name": "flyte databricks plugin",
        "new_cluster": {
            "spark_version": "13.3.x-scala2.12",
            "autoscale": {"min_workers": 1, "max_workers": 1},
            "node_type_id": "m6i.large",
            "num_workers": 1,
        },
        "timeout_seconds": 3600,
        "max_retries": 1,
    },
    databricks_instance="dbc-6f73e2e6-19e4.cloud.databricks.com",
    databricks_token="DATABRICKS_TOKEN",
)

databricks_env = flyte.TaskEnvironment(
    name="databricks_env",
    resources=flyte.Resources(),
    plugin_config=databricks_conf,
    image=image,
)

Databricks subclasses Spark, so its Spark dictionaries, Python and application paths, and optional pod templates have the same meaning. It adds databricks_conf, a Databricks Jobs API configuration; databricks_instance, the workspace hostname; and databricks_token, the name of the Flyte secret containing the API token.

DatabricksFunctionTask combines AsyncConnectorExecutorMixin with PysparkFunctionTask and changes the task type to databricks. Its serialized custom configuration contains the Python SparkJob fields plus databricksConf, databricksInstance, and secrets.databricks_token when those values are set. Protobuf conversion omits unset fields.

The task body uses the same injected session as native Spark:

from operator import add

@databricks_env.task
async def hello_databricks_nested() -> float:
    partitions = 3
    n = 1 * partitions
    spark = flyte.ctx().data["spark_session"]
    count = spark.sparkContext.parallelize(range(1, n + 1), partitions).map(f).reduce(add)
    res = get_pi(count, partitions)
    print(f"result: {res}")
    return res

The connector is registered when the Databricks package is imported. User task code does not instantiate DatabricksConnector; the asynchronous connector executor and global connector registry perform submission, polling, and cancellation.

Configure Spark workloads​

Set application and executor paths explicitly when the runtime needs them​

For native Spark, omitted paths come from the serialization context. For Databricks, the example overrides the executor path with the Databricks Python executable:

databricks_conf = Databricks(
    executor_path="/databricks/python3/bin/python",
    databricks_conf={"existing_cluster_id": "1113-204018-tb9vr2fm"},
    databricks_instance="dbc-6f73e2e6-19e4.cloud.databricks.com",
    databricks_token="DATABRICKS_TOKEN",
)

The existing_cluster_id alternative is shown in the Databricks example as a commented option. A Databricks configuration still needs a valid cluster selector even though Databricks() itself can be constructed with all optional fields unset.

Use pod templates on native Spark​

driver_pod and executor_pod accept Flyte PodTemplate values. PysparkFunctionTask.custom_config() calls to_k8s_pod() for each supplied template and includes the resulting fields in the serialized SparkJob. The same inherited fields are serialized by DatabricksFunctionTask.

Override Spark settings for a nested task​

The Spark example makes a copy of the configuration before changing executor count, then passes it through override(plugin_config=...):

from copy import deepcopy

@task_env.task
async def spark_overrider(executor_instances: int = 3, partitions: int = 4) -> float:
    updated_spark_conf = deepcopy(spark_conf)
    updated_spark_conf.spark_conf["spark.executor.instances"] = str(executor_instances)
    return await hello_spark_nested.override(plugin_config=updated_spark_conf)(partitions=partitions)

Process DataFrames​

Spark tasks can accept a PySpark DataFrame using the annotated transformer pattern, or obtain the injected SparkSession and read a remote file:

import collections
from typing import Annotated, OrderedDict, Type, cast

import pyspark
from pyspark.sql import SparkSession

columns = collections.OrderedDict(name=str, age=int)

@spark_env.task
async def sum_of_all_ages(sd: Annotated[pyspark.sql.DataFrame, columns]) -> int:
    total_age = sd.groupBy().sum("age").collect()[0][0]
    print("Total age sum:", total_age)
    return total_age

@spark_env.task
async def dataframe_transformer() -> int:
    spark = cast(SparkSession, flyte.ctx().data["spark_session"])
    file = await create_new_remote_file("age,name\n10,alice\n20,bob")
    spark_df = spark.read.csv(path=file.path, header=True, inferSchema=True)
    return await sum_of_all_ages(spark_df)

The complete example defines create_new_remote_file with File.new_remote() and returns a Flyte File; the task then reads that remote path and passes the resulting DataFrame to sum_of_all_ages.

Understand Databricks submission and monitoring​

Before submission, _get_databricks_job_spec() reads databricksConf from the serialized task custom configuration. It raises ValueError when that field is absent. The configuration must contain existing_cluster_id or new_cluster; otherwise it raises ValueError with Either existing_cluster_id or new_cluster must be specified.

For a new_cluster, the connector fills missing values from the Flyte task template: docker_image defaults to {"url": container.image}, spark_conf defaults to serialized sparkConf, and spark_env_vars is populated from the container environment. Existing environment values are merged with the container variables.

The connector then overwrites spark_python_task with flyteplugins/databricks/entrypoint.py, passes container arguments as parameters, and sets source to GIT. It also overwrites git_source with the pinned flyteorg/flytetools repository and commit 194364210c47c49ce66c419e8fb68d6f9c06fd7e.

DatabricksConnector.create() sends JSON to https://<databricks_instance>/api/2.1/jobs/runs/submit. The instance comes first from serialized databricksInstance, then from FLYTE_DATABRICKS_INSTANCE; the connector prepends https://. Requests use a Bearer authorization header and content-type: application/json. A successful response produces DatabricksJobMetadata containing the workspace hostname and returned run ID.

Polling calls https://<databricks_instance>/api/2.1/jobs/runs/get?run_id=<run_id>. Non-terminated life_cycle_state values map directly to Flyte phases; only exact TERMINATED uses result_state. The state_message becomes the resource message. A TaskLog named Databricks Console links to https://<databricks_instance>/#job/<job_id>/run/<run_id>.

Cancellation posts {"run_id": run_id} to https://<databricks_instance>/api/2.1/jobs/runs/cancel. All connector HTTP operations require status 200 exactly; non-200 responses raise RuntimeError.

Troubleshooting checklist​

  • Set databricks_conf with new_cluster or existing_cluster_id; Databricks() alone is not submit-ready.
  • Use a workspace hostname without https://, or set FLYTE_DATABRICKS_INSTANCE for connector fallback.
  • Set databricks_token to a Flyte secret name such as DATABRICKS_TOKEN, not a literal token.
  • Use a compatible Python/PySpark runtime. The native example uses apache/spark-py:v3.4.0 and flyteplugins-spark; the Databricks example uses databricksruntime/standard:16.4-LTS, /databricks/python3/bin/python, and flyteplugins-databricks.
  • Native code-bundle distribution occurs only when flyte.ctx().is_in_cluster() is true and flyte.ctx().code_bundle is non-null.
  • The connector consults result_state only for TERMINATED; missing or unexpected state data can leave the phase as TaskExecution.UNDEFINED.
  • Databricks integration requires a Union/Flyte backend with the Databricks connector configured; the repository integration test is skipped otherwise.