Effortless EMR: A Guide to Seamlessly Running PySpark Code

2.1 Java packages

Adding Java packages is quite straightforward:

1. Download the package as a JAR file, for example, spark-3.3-bigquery-0.31.1.jar
2. Add this file somewhere in S3 where you have reading access, for example, s3://your-bucket/emr/jars/
3. Set up the cluster to include that package by adding the following configuration:

--jars s3://your-bucket/emr/jars/spark-3.3-bigquery-0.31.1.jar

With that, you should be able to use that package.

2.2. Python packages

The simplest way of adding a Python package is by Setting up a bootstrap action (as recomended in QA: EMR install python libraries)

The main problem with that is that you’ll be downloading the package at runtime. You could download the package (an egg or a wheel) to S3 and install it from there, but it adds complexity.

The way to overcome those problems is to create a docker image.

3.1. Add external packages

The overall idea is to follow EMR serverless samples. In this example I’m using Poetry python package manager since I think it’s better than regular pip.

The steps I follow are:

1. Get pip
2. Install poetry
3. Copy the needed files in order to install the poetry dependencies
4. Create the virtual environment
5. Create a tar.gz with the virtual environment
6. Export the tar.gz

Dockerfile

FROM --platform=linux/amd64 amazonlinux:2023 AS base

RUN dnf install python3-pip -y

# Update and install python packages
RUN pip install poetry --no-cache-dir

# Copy needed files to install the project
COPY pyproject.toml .
COPY README.md .
COPY my_package my_package

# Set up poetry and install dependencies
RUN poetry config virtualenvs.in-project true && \
    poetry install

# Create a '.tar.gz' with the virtual environment
RUN mkdir /output && \
    poetry run venv-pack -o /output/my_package-latest.tar.gz

# Export the '.tar.gz'
FROM scratch AS export
COPY --from=base /output/*.tar.gz /

To create the tar.gz simply run:

docker build --output . .

Once the tar.gz is created, you will need to upload to S3. Then you can set up an EMR cluster to use that environment with:

--conf spark.emr-serverless.driverEnv.PYSPARK_DRIVER_PYTHON=./environment/bin/python
--conf spark.emr-serverless.driverEnv.PYSPARK_PYTHON=./environment/bin/python
--conf spark.emr-serverless.executorEnv.PYSPARK_PYTHON=./environment/bin/python
--conf spark.archives=s3:/your-bucket/venv/my_package-latest.tar.gz#environment

3.2. Create a ‘Package’ for Adding Code

Now that we know how to add external packages, we can reuse it to also be able to import our own code. As stated in How to run a Python project (package) on AWS EMR serverless?, the idea is to create a Python package and then install it.

In my case, it would be done with Docker by:

# Install 'my_package' as a package
RUN poetry build && \
    poetry run pip install dist/*.whl

Another thing I added is the ability to pass the version as an argument. The final Dockerfile code would be:

Dockerfile

FROM --platform=linux/amd64 amazonlinux:2023 AS base

ARG PACKAGE_VERSION

# Should match CI config
ARG POETRY_VERSION=1.6.1 

RUN dnf install python3-pip -y

# Update and install python packages
RUN pip install \
  "poetry==${POETRY_VERSION}" \
  boto3 \
  loguru \
  toml \
  --no-cache-dir

# Copy needed files to install the project
COPY pyproject.toml .
COPY README.md .
COPY my_package my_package

# Set up poetry and install dependencies
RUN poetry config virtualenvs.in-project true && \
    poetry install

# Install 'my_package' as a package
RUN poetry build && \
    poetry run pip install dist/*.whl

# Create a '.tar.gz' with the virtual environment
RUN mkdir /output && \
    poetry run venv-pack -o /output/my_package-${PACKAGE_VERSION}.tar.gz

# Export the '.tar.gz'
FROM scratch AS export
COPY --from=base /output/*.tar.gz /

And it is compiled with:

# Set a version that is unique and will not collide with other people
docker build --output . . --build-arg PACKAGE_VERSION=2.1.7 # Any version you want

3.3. Export the environment with CD

Here, the idea is to have a GitHub action that uploads the virtual environment whenever there is a commit to main.

The general steps are:

1. Fetch the code
2. Create the tar.gz as seen in the previous section
3. Upload it to S3

One thing that I find very useful is to always export the real version, such as 2.1.7, and then also create one called latest. With this, I can easily point to the latest.tar.gz in production and/or point to a specific version if I need to.

You might want to read Managing package versions with Poetry for more information about how to manage versions with Poetry.

5.1. Catching Generic Parameters with click

The idea here is to define the entrypoint as a parameter in click and then catch the rest of the parameters as explained in Add unspecified options to cli command using python-click.

The basic code for that would be:

entrypoint.py

import click
from loguru import logger

@click.command(
    name="entrypoint",
    context_settings=dict(
        ignore_unknown_options=True,
        allow_extra_args=True,
    ),
)
@click.pass_context
@click.option(
    "--entrypoint",
    "-e",
    required=True,
    type=str,
    help="File used as an entrypoint. Must have a 'main' function",
)
def main(ctx, entrypoint):
    logger.debug("Extracting parameters from context")
    # https://stackoverflow.com/a/32946412/3488853
    kwargs = {ctx.args[i][2:]: ctx.args[i + 1] for i in range(0, len(ctx.args), 2)}
    kwargs.update({"entrypoint": entrypoint})

    # Replace dashes for underscores in param names to ensure compatibility
    kwargs = {k.replace("-", "_"): v for k, v in kwargs.items()}

    logger.info(f"Importing '{entrypoint}'")
    module = importlib.import_module(entrypoint)
    module.main(**kwargs)


if __name__ == "__main__":
    print("Starting EMR job")
    main()

Then you can define the parameters in each file as usual:

my_package/salesforce/main.py

# Other code

def main(name, num):
    # This is what will be called
    print(f"Hello {name} with {num=}")

5.2. Define Script Parameters with pydantic

One way to improve the previous code is to define the parameters with Pydantic. This way you don’t need to handle the castings and can use the powerful validators from pydantic.

You would do that with:

my_package/salesforce/main.py

from pydantic import BaseModel

class MyParams(BaseModel):
    entrypoint: str # This is needed
    name: str
    num: int

def main(**kwargs):
    args = MyParams(**kwargs)
    print(f"Hello {args.name} with {args.num=}")

6.1. Jobs as Flows

The idea here is to create a flow at the entrypoint so that you can see each run. The idea is to create a flow dynamically with flow(name="flow_name")(function)(*args, **kwargs).

So the modified lines look like as follows:

entrypoint.py

from prefect import flow

def main(ctx, entrypoint):
    logger.info(f"Importing '{entrypoint}'")
    module = importlib.import_module(entrypoint)
    flow(name=entrypoint)(module.main)(**kwargs)

6.2. Adding Tasks and/or Subflows

You can follow the same approach to add different tasks and/or subflows. For example, let’s imagine you are extracting multiple tables from somewhere. Here you could create a Prefect task for each table. This way you will have a much better observability system. It would be done with something like:

from prefect import task

for table in tables:
    name = f"{params.entrypoint}.{table.name}"
    task(name=name)(process_one_table)(name=table.name, **any_other_arguments)

7.1. Entrypoint

This takes care of:

• Calling the main function of other files
• Catching parameters as a dict
• Adding prefect.tags

entrypoint.py

import importlib
import sys

import click
from loguru import logger
from prefect import tags


PACKAGE_NAME = "my_package"

logger.configure(handlers=[{"sink": sys.stdout, "level": "INFO"}])
logger.enable(PACKAGE_NAME)


def get_run_tags(entrypoint):
    tags = [
        "type:emr",
        "entrypoint:" + entrypoint.replace("my_package.", ""),
        # Feel free to add any other tags
    ]
    logger.info(f"Running with {tags=}")
    return tags


@click.command(
    name="entrypoint",
    context_settings=dict(
        ignore_unknown_options=True,
        allow_extra_args=True,
    ),
)
@click.pass_context
@click.option(
    "--entrypoint",
    "-e",
    required=True,
    type=str,
    help="File used as an entrypoint. Must have a 'main' function",
)
def main(ctx, entrypoint):
    if not entrypoint.startswith("my_package."):
        entrypoint = f"my_package.{entrypoint}"

    logger.debug("Extracting parameters from context")
    # https://stackoverflow.com/a/32946412/3488853
    kwargs = {ctx.args[i][2:]: ctx.args[i + 1] for i in range(0, len(ctx.args), 2)}
    kwargs.update({"entrypoint": entrypoint})

    # Replace dashes for underscores in param names to ensure compatibility
    kwargs = {k.replace("-", "_"): v for k, v in kwargs.items()}

    logger.info(f"Importing '{entrypoint}'")
    module = importlib.import_module(entrypoint)

    logger.info(f"Running '{entrypoint}.main()' with kwargs='{kwargs}'")
    with tags(*get_run_tags(entrypoint)):
        module.main(**kwargs)


if __name__ == "__main__":
    print("Starting EMR job")
    main()

7.2. Arguments code

Here we define:

• BaseParams to be extended in each module
• A decorator that catches the Params and creates a prefect.flow

my_package/common/parameters.py

from time import time

from pydantic import BaseModel
from prefect import flow, get_run_logger

class BaseParams(BaseModel):
    """
    Parent class for the Params meant to be extended.
    It adds 'entrypoint' as a forced parameter.
    """

    entrypoint: str

def catch_params(ParamsClass):
    """
    This is used to easily parse params. It takes a class that inherits from 'BaseParams'.

    For example:
        class Params(params.BaseParams):
            schema_out: str = "nt_bronze__big_query"
            files_per_partition: str = 1


        @params.catch_params(Params)
        def main(params):
            do_stuff
    """

    def decorator(func):
        def inner(**kwargs):
            params = ParamsClass(**kwargs)
            entrypoint = params.entrypoint

            @flow(name=entrypoint)
            def prefect_flow(params):
                t0 = time()
                logger = get_run_logger()

                logger.info(f"Running with {params=}")

                # Make sure to log the outcome of the process
                try:
                    func(params)

                except Exception as ex:
                    logger.exception(
                        f"Job '{entrypoint}.main()' failed after {(time() - t0)/ 60:.2f} minutes"
                    )
                    raise ex
                else:
                    logger.info(
                        f"Job '{entrypoint}.main()' finished successfully "
                        f"after {(time() - t0)/ 60:.2f} minutes"
                    )

            return prefect_flow(params)

        return inner

    return decorator

7.3. Sample module

Here is an example of how to define a job:

my_package/salesforce/main.py

from datetime import datetime

from my_package.common import BaseParams, catch_params

# Here we define all params
class Params(BaseParams):
    n_jobs: int = 10
    start_time: datetime = None

@catch_params(Params)
def main(params):
    print(f"Starting {params.n_jobs} jobs with start_time={params.start_time}")