Cannot use SparkSession in resource_manager, when using LocalDaskExecutor with "processes" scheduler

[Sinha-Ujjawal]

Hi
Cannot use SparkSession in resource_manager, when using LocalDaskExecutor with "processes" scheduler

from prefect import task, Flow, Parameter
from prefect.engine.executors import LocalDaskExecutor

from prefect import resource_manager
import pyspark

@resource_manager
class SparkContext:
    def __init__(self, app_name: str, master: str = "local[*]"):
        self._app_name = app_name
        self._master = master
    
    def setup(self):
        return (
            pyspark.sql.SparkSession
            .builder
            .appName(self._app_name)
            .master(self._master)
            .getOrCreate()
        )
    
    def cleanup(self, spark: pyspark.sql.session.SparkSession):
        spark.stop()

import pyspark

@task
def task_1(spark: pyspark.sql.session.SparkSession, output_file: str):
    (
        spark
        .read
        .option("header", "true")
        .csv("./abc.csv")
        .limit(10)
        .toPandas()
        .to_csv(output_file)
    )

with Flow("LocalDask+Spark") as flow:
    spark_app_name = Parameter(name="spark_app_name", required=True)
    spark_master = Parameter(name="spark_master", default="local[*]")
    
    with SparkContext(spark_app_name, spark_master) as spark:
        t1 = task_1(spark, "local_dask__spark.csv")
        t1.name = "t1"
        
        t2 = task_1(spark, "local_dask__spark2.csv")
        t2.name = "t2"

state = flow.run(
    executor=LocalDaskExecutor(scheduler="processes"),
    spark_app_name="app-demo",
)

Error-

Unexpected error: Exception('It appears that you are attempting to reference SparkContext from a broadcast variable, action, or transformation. SparkContext can only be used on the driver, not in code that it run on workers. For more information, see SPARK-5063.')

When I used scheduler="threads", it worked! Can anyone help me on that?

Ragards,
Ujjawal

[cicdw]

Hi @Sinha-Ujjawal and sincere apologies for the delayed response! That error arises out of pickling pyspark's context object (you can find the error message within this source code: https://people.eecs.berkeley.edu/~jegonzal/pyspark/_modules/pyspark/context.html) which is only necessary when sharing the spark objects across processes, and is not necessary when sharing the objects amongst threads. I recommending using threads when using both Spark and Dask together - this is still a pattern we are learning about from our community.

Also, completely unrelated: note that you can change the name of your tasks in one line as follows:

with Flow("LocalDask+Spark") as flow:
    spark_app_name = Parameter(name="spark_app_name", required=True)
    spark_master = Parameter(name="spark_master", default="local[*]")
    
    with SparkContext(spark_app_name, spark_master) as spark:
        t1 = task_1(spark, "local_dask__spark.csv", task_args={"name": "t1"})
        
        t2 = task_1(spark, "local_dask__spark2.csv", task_args={"name": "t2"})

[Sinha-Ujjawal]

Hey thx for the reply. Before this I always had to change the task name by - task.name = "name". I didn't know about task_args. This is useful 😁