How to apply large python model to pyspark-dataframe?

Question

I have:

• Large dataframe (parquet format, 100.000.000 rows, 4.5TB size) that contains some data (features)
• Several huge ML models (each one takes 5-15GB of RAM)
• Spark cluster (AWS EMR), typical node configuration is 8 CPU, 32 RAM, can be changed if needed.

I want to apply them using PySpark, but I always get some wired errors like:

• OOM
• Random timeouts (nodes doesn't return any result) -> node killed by YARN manager

I typically used code like

def apply_model(partition):
    model = load(...)  # load model only when apply this function to avoid serialization issue
    for row in partition:
        yield model.infer(row)

or

def apply_model(partition):
    model = load(...)  # load model only when apply this function to 
    yield from model.infer(partition)

and apply that using

df.select(...).rdd.mapPartitions(apply_model)

I can't broadcast model, by serialization reasons.

The question - how to apply the big python/any-non-jvm-based model to spark dataframe & avoid spark exceptions?

Answer 1

Here are some additional suggestions that could help improving the performance of your job:

• The first change I would do is to reduce the partition size. If I understood correctly at the moment you have input data of 4.5TB. That means if you have 1000 partitions then you will end up sending 4,5GB per partition on each executor! This size is considered quite large, instead I would try to keep the partition size between 250-500MB. Roughly in your case that would mean ~10000 (4.5TB / 500MB) partitions.

• Increase parallelism by adding more executors. That would increase the level of data locality and consequently reduce the execution time. Ideally you should have 5cores per executor and two executors (if possible) for each cluster node. The max cores per executor should not be higher than 5 since that would cause I/O bottlenecks (when/if disk storage is used).

• As for the memory the suggestions from @rluta I think are more than sufficient. In general too large values for executor's memory would have a negative effect on Java GC time therefore an upper limit of 10GB should be the ideal value for spark.executor.memory.

Answer 2

When you apply a python function on a partition with mapPartitions, Spark will evaluate it in a companion python process for each JVM based executor.

Usually, the memory used by the python process is small and well within the YARN memoryOverhead setting used on EMR. In your specific case, this assumption does not hold since the python process has to hold your large model in memory so you need to adapt your configuration.

If you stick to 8 CPU/32G RAM per executor host, you can try the following as a base config:

spark.executor.cores=6
spark.executor.memory=8G
spark.executor.pyspark.memory=20G

Note that setting spark.executor.pyspark.memory will impose a hard limit on the memory usage of your python process (it's unbounded by default) so you may need to experiment to find the constraint that will fit your process.

An alternative configuration would be to keep the pyspark process memory unbounded but reserve enough YARN memoryOverhead to fit it in, for example:

spark.executor.cores=6
spark.executor.memory=8G
spark.executor.memoryOverhead=22G

Answer 3

Do your partitions fit into memory of a single executor? You could try to increase the number of partitions, and see whether the situation improves:

df.select(...).repartition(1000).rdd.mapPartitions(apply_model)

Validate that this is an improvement by looking into Spark UI's metrics, such as:

• Input Size / Records
• Shuffle Write Size / Records
• GC Time

Compare Median, 75th percentile and Max values to see whether your data is not skewed.

Answer 4

If I understood your question correctly, then you already have model built and you need to apply it on 4.5TB of data. If so, then you can process it like stream. You divide the file into manageable size and then provide the directory for spark stream to read and process.

5/23: As per my understanding in mapPartition, you read the data entirely and then split it and thats why there is a chance of memory overflow. Where as in stream, you would create a mini batch of the input data and then process that mini batch at a time. Reference: https://spark.apache.org/docs/2.2.0/structured-streaming-programming-guide.html#basic-concepts

If you can break the big data file into smaller one then as indicated in the below link, you would be able to work process them one at a time or bunch of them at a time till the big data is entirely consumed. reference: https://spark.apache.org/docs/2.2.0/structured-streaming-programming-guide.html#input-sources