patch_spark_memory.py

#!/opt/conda/default/bin/python3
import json
import os
import subprocess as sp
import sys
import errno
from subprocess import check_output

assert sys.version_info > (3, 0), sys.version_info


def safe_call(*args, **kwargs):
    try:
        sp.check_output(args, stderr=sp.STDOUT, **kwargs)
    except sp.CalledProcessError as e:
        print(e.output.decode())
        raise e


def get_metadata(key):
    return check_output(['/usr/share/google/get_metadata_value', 'attributes/{}'.format(key)]).decode()


# get role of machine (master or worker)
role = get_metadata('dataproc-role')

if role == 'Master':
    conf_to_set = [
        'spark.driver.memory=85g'
    ]

    print('setting spark-defaults.conf')

    with open('/etc/spark/conf/spark-defaults.conf', 'a') as out:
        out.write('\n')
        for c in conf_to_set:
            out.write(c)
            out.write('\n')