Dagster PySpark not running on EMR

Question

I am trying to build an pipeline in Dagster which does the following:

1. Launch an EMR cluster using the EmrJobRunner class, by using its run_job_flow function.

2. Add one or more steps to that cluster to process data in PySpark by using the emr_pyspark_step_launcher resource.

3. Shut down the cluster once all steps are finished.

I followed this tutorial first, which assumes that you have an EMR cluster running and you hard code the EMR cluster ID as part of the Job specification. This way worked, as I could see my steps being run on EMR. However, when I try to automate the process I noticed that PySpark was running locally and not on EMR. I tried to wrap the emr_pyspark_step_launcher as a Resource which sets the cluster ID as part of the pipeline. The cluster ID can be obtained by using a function in the EmrJobRunner class which returns a cluster ID when providing a cluster name. I am trying to dynamically add the cluster ID during the job after launching the cluster but this isn't working as expected.

This is my code, any help would be appreciated.

from pathlib import Path
from dagster_aws.emr import emr_pyspark_step_launcher
from dagster_aws.emr.emr import EmrJobRunner
from dagster_aws.s3 import s3_resource
from dagster_pyspark import pyspark_resource
from pyspark.sql import DataFrame
from transformations import execute_transformation
from dagster import IOManager, graph, io_manager, op, resource, In, Nothing, Out
from utils.configs import get_emr_cluster_config
import logging


class ParquetIOManager(IOManager):
    def _get_path(self, context):
        return "/".join(
            [
                context.resource_config["path_prefix"],
                context.run_id,
                context.step_key,
                context.name,
            ]
        )

    def handle_output(self, context, obj):
        if isinstance(obj, DataFrame):
            obj.write.parquet(self._get_path(context))
        # return obj

    def load_input(self, context):
        spark = context.resources.pyspark.spark_session
        return spark.read.parquet(self._get_path(context.upstream_output))


@io_manager(required_resource_keys={"pyspark"}, config_schema={"path_prefix": str})
def parquet_io_manager():
    return ParquetIOManager()


@resource
def emr_job_runner(init_context):
    return EmrJobRunner(region="eu-central-1")


@resource(
    config_schema={"cluster_name": str}, required_resource_keys={"emr_job_runner"}
)
def my_pyspark_step_launcher(init_context):
    cluster_id = init_context.resources.emr_job_runner.cluster_id_from_name(
        cluster_name=init_context.resource_config["cluster_name"]
    )
    init_context.log.info(f"CLUSTER ID during resource initilization: {cluster_id}")

    return emr_pyspark_step_launcher.configured(
        {
            "cluster_id": cluster_id,
            "local_pipeline_package_path": str(Path(__file__).parent.parent),
            "deploy_local_pipeline_package": True,
            "region_name": "eu-central-1",
            "staging_bucket": "EMR_STAGING_BUCKET",
            "wait_for_logs": True,
        }
    )
    

def launch_cluster(emr: EmrJobRunner, log: logging.Logger, emr_config: dict) -> None:
    emr_config = get_emr_cluster_config(
        release_label=emr_config["emr_release_label"],
        cluster_name=emr_config["cluster_name"],
        master_node_instance_type=emr_config["master_node_instance_type"],
        worker_node_instance_type=emr_config["worker_node_instance_type"],
        worker_node_instance_count=emr_config["worker_node_instance_count"],
        ec2_subnet_id=emr_config["ec2_subnet_id"],
        bid_price=emr_config["worker_node_spot_bid_price"],
    )

    return emr.run_job_flow(log=log, cluster_config=emr_config)


@op(
    config_schema={
        "emr_release_label": str,
        "cluster_name": str,
        "master_node_instance_type": str,
        "worker_node_instance_type": str,
        "worker_node_instance_count": int,
        "ec2_subnet_id": str,
        "worker_node_spot_bid_price": str,
    },
    required_resource_keys={"emr_job_runner"},
    out=Out(Nothing),
)
def launch_emr_cluster(context) -> None:
    op_config = context.op_config

    cluster_id = launch_cluster(
        emr=context.resources.emr_job_runner, log=context.log, emr_config=op_config
    )

    context.log.info(f"CLUSTER ID: {cluster_id}")


@op(
    ins={"start": In(Nothing)},
    required_resource_keys={"pyspark", "pyspark_step_launcher"},
)
def get_dataframe(context) -> DataFrame:
    return execute_transformation(spark_session=context.resources.pyspark.spark_session)


@graph
def make_and_filter_data():
    get_dataframe(launch_emr_cluster())


run_data_emr = make_and_filter_data.to_job(
    name="prod",
    resource_defs={
        "pyspark_step_launcher": my_pyspark_step_launcher,
        "pyspark": pyspark_resource,
        "s3": s3_resource.configured({"region_name": "eu-central-1"}),
        "io_manager": parquet_io_manager.configured(
            {"path_prefix": "s3://EMR_STEP_OUTPUT"}
        ),
        "emr_job_runner": emr_job_runner,
    },
)

Answer 1

This is a bit tricky because resources are initialized for every op, and are initialized prior to the op starting. So you can't modify / pass values from an op to a resource initialization. In this particular case, I think what you might want to try is doing the cluster initialization in the step launcher resource definition - that way you have access to the cluster ID at the time that the step launcher is being initialized. Something like this:

@resource
def emr_job_runner(init_context):
    return EmrJobRunner(region="eu-central-1")
    

def launch_cluster(emr: EmrJobRunner, log: logging.Logger, emr_config: dict) -> None:
    emr_config = get_emr_cluster_config(
        release_label=emr_config["emr_release_label"],
        cluster_name=emr_config["cluster_name"],
        master_node_instance_type=emr_config["master_node_instance_type"],
        worker_node_instance_type=emr_config["worker_node_instance_type"],
        worker_node_instance_count=emr_config["worker_node_instance_count"],
        ec2_subnet_id=emr_config["ec2_subnet_id"],
        bid_price=emr_config["worker_node_spot_bid_price"],
    )

    return emr.run_job_flow(log=log, cluster_config=emr_config)

@resource(config_schema={
        "emr_release_label": str,
        "cluster_name": str,
        "master_node_instance_type": str,
        "worker_node_instance_type": str,
        "worker_node_instance_count": int,
        "ec2_subnet_id": str,
        "worker_node_spot_bid_price": str,
    },
    required_resource_keys={"emr_job_runner"})
def cluster_launcher(init_context):
    config = init_context.resource_config

    # TODO: handle if cluster already exists, as this resource will be initialized for each op / resource
    # that requires it
    cluster_id = launch_cluster(
        emr=context.resources.emr_job_runner, log=context.log, emr_config=config
    )

    context.log.info(f"CLUSTER ID: {cluster_id}")
    return cluster_id


@resource(
    config_schema={"cluster_name": str}, required_resource_keys={"emr_job_runner", "cluster_launcher"}
)
def my_pyspark_step_launcher(init_context):
    cluster_id = init_context.resources.cluster_launcher
    init_context.log.info(f"CLUSTER ID during resource initilization: {cluster_id}")

    return emr_pyspark_step_launcher.configured(
        {
            "cluster_id": cluster_id,
            "local_pipeline_package_path": str(Path(__file__).parent.parent),
            "deploy_local_pipeline_package": True,
            "region_name": "eu-central-1",
            "staging_bucket": "EMR_STAGING_BUCKET",
            "wait_for_logs": True,
        }
    )