We run a lot of Pyspark jobs on EMR. The pipeline executed is the same, but the inputs can wildly change the peak memory utilization, and that utilization is growing over time. I would like to automatically write out the peak memory utilization of each step submitted to the EMR cluster. If it matters, we are using cluster mode with yarn as the cluster manager. We are also submitting these jobs as Docker containers.

I have tried the following two approaches, but I am very open to the idea that I’m wildly overthinking this… it doesn’t feel like it should be that hard to just do:

spark = init_spark()

do_some_stuff()

spark.metrics.peakMemoryUtilization()

So if I’m missing something please educate me.

Attempt 1 – Custom Class extending JVM’s Spark Listener

Fiddling with some suggestions from Github CoPilot (I know) I was pointed towards using SparkListeners for this task, however I have not had any success with this approach. The suggested approach appears to be to create a custom class which uses the JVM’s access to the tasks.

from pyspark.sql import SparkSession
from pyspark import SparkConf
from py4j.java_gateway import java_import

class MemoryUsageListener:
    def __init__(self):
        self.peak_memory_usage = 0

    def on_task_end(self, task_end):
        metrics = task_end.taskMetrics()
        memory_used = metrics.peakExecutionMemory()
        if memory_used > self.peak_memory_usage:
            self.peak_memory_usage = memory_used

    # Getter function for adding this to the SparkContext's Listeners
    def get_listener(self, sc):
        gw = sc._gateway
        java_import(gw.jvm, "org.apache.spark.scheduler.SparkListener")
        java_import(gw.jvm, "org.apache.spark.scheduler.SparkListenerTaskEnd")

        class JavaListener(gw.jvm.SparkListener):
            def __init__(self, parent):
                super(gw.jvm.SparkListener, self).__init__()
                self.parent = parent

            def onTaskEnd(self, taskEnd):
                self.parent.on_task_end(taskEnd)

        return JavaListener(self)

    def get_peak_memory_usage(self):
        return self.peak_memory_usage

With the setup being something like:

conf = SparkConf().setAppName("MemoryUsageTracker")
spark = SparkSession.builder.config(conf=conf).getOrCreate()
sc = spark.sparkContext

# Create and register the listener
memory_listener = MemoryUsageListener()
java_listener = memory_listener.get_listener(sc)
sc._jsc.sc().addSparkListener(java_listener)

... do some spark stuff here (run the actual job) ...

peak_memory_usage = memory_listener.get_peak_memory_usage()
write_to_file(peak_memory_usage)

I played with the above suggestion extensively but struggled with the following: JavaClass.__init__() "takes 3 positional arguments but 4 were given" when the class is initialized, specifically it hates class JavaListener(gw.jvm.SparkListener). I’ve tried moving the super declaration around, moving the Java class into a completely separate function, no cigar. It seems intuitive, but clearly I don’t understand how this OOP approach works unfortunately.

Attempt 2 – API Calls to Yarn

Okay, so the intuitive solution of just asking Spark what it’s doing didn’t work. Another suggestion I found was to use the requests library and hit the api endpoint to get details on the executors.

import requests
from pyspark.sql import SparkSession
from pyspark import SparkConf

def get_executor_memory_metrics(app_id, spark_history_server_url):
    response = requests.get(f"{spark_history_server_url}/api/v1/applications/{app_id}/executors")
    if response.status_code != 200:
        raise Exception(f"Failed to fetch executor metrics, status code: {response.status_code}")
    executors = response.json()
    return executors

def calculate_peak_memory_usage(executors):
    # i'm not certain this is the correct dict access, it's what copilot generated, but it's not my current issue
    peak_memory_usage = max(executor['peakMemoryMetrics']['JVMHeapMemory'] for executor in executors if 'peakMemoryMetrics' in executor)
    return peak_memory_usage

sc = spark.sparkContext
app_id = sc.applicationId

master_url = sc.master
# I'm not convinced this logic is actually sound, but I don't have access to modify the cluster / yarn settings. 
yarn_resource_manager_host = socket.getfqdn()
yarn_port = 18080

# At the end of the job
executors = get_executor_memory_metrics(app_id, f"{yarn_resource_manager_host}:{yarn_port}")
peak_memory_usage = calculate_peak_memory_usage(executors)

And obviously the above code would be packaged in a thread to loop and continuously monitor the peak usage — however, before moving to that piece, I’m unable to get the request portion actually working with a persistent “Connection Refused” error. I’ve confirmed it’s receiving the cluster’s IP and application id correctly and imputing the URL, though I can’t share them here — I’m not certain the port number is right, but I’m not sure if there’s a better way to access it at runtime.