How to avoid out-of-memory error when creating large arrays from PySpark and Parquet for data analysis?

Viewed 625

We have a somewhat unusual requirement that we're experimenting with. We have a parquet table with ~320 GB of frequency pixel values in arrays - the data is from a 3-dimensional spectral line image cube (A FITs source file).

+--------------------+------------+-------+
|            col_name|   data_type|comment|
+--------------------+------------+-------+
|           spi_index|         int|   null|
|           spi_image|array<float>|   null|
|        spi_filename|      string|   null|
|            spi_band|         int|   null|
|# Partition Infor...|            |       |
|          # col_name|   data_type|comment|
|        spi_filename|      string|   null|
|            spi_band|         int|   null|
+--------------------+------------+-------+

There is a legacy source finding algorithm written in C that is used to analyse the original data source files, but it's restrained to the physical memory available on one machine. We're running a CDH Express cluster, 3 management nodes (8 cores, 32 GB ram) 10 worker nodes (16 cores, 64 GB ram) and a gateway node running Jupyterhub (8 cores, 32 GB of ram). We've modified the original C program into a shared object, and distributed it across the cluster. We've incorporated the C shared object into a partitioner class so we can run multiple partitions in multiple executors across the cluster. We have this up and running using pyspark.

The problem that we're experiencing is that we optimally need an input pixel array of a minimum of about 15GB and we seem to be hitting a wall at creating arrays of around 7.3 GB, and we're unsure of why this is.

YARN maximum allocation setting.

yarn.scheduler.maximum-allocation-mb = 40GB

Spark configuration settings

--executor-memory 18g \
--num-executors 29 \
--driver-memory 18g \
--executor-cores 5 \
--conf spark.executor.memoryOverhead='22g' \
--conf spark.driver.memoryOverhead='22g' \
--conf spark.driver.maxResultSize='24g' \
--conf "spark.executor.extraJavaOptions-XX:MaxDirectMemorySize=20G -XX:+UseCompressedOops" \

Summary of partitioner class

class Partitioner:
def __init__(self):
    self.callPerDriverSetup
    
def callPerDriverSetup(self):
    pass

def callPerPartitionSetup(self):
    sys.path.append('sofia')
    #import example
    
    import sofia
    import myLib
    
    import faulthandler
    import time as tm
    from time import time, clock, process_time, sleep 

     
    self.sofia=sofia
    self.myLib=myLib
    #self.example=example
    self.parameterFile=SparkFiles.get('sofia.par')
    
def doProcess(self, element):
    ###  here's the call to the C library for each row of the dataframe partition
    ### In here we have to transform the flattened array data to the format SoFiA
    ### requires, as well as the
    
    ra=np.array(element.right_ascension, dtype='<f4')
    dec=np.array(element.declination, dtype='<f4')
    frequencies=np.array(element.frequency, dtype='<f4')
    Pixels=np.array(element.pixels, dtype='<f4')
    
    
    dataPtr= Pixels.ravel()
            
    #
    # create the dictionary of the original header
    #
    hdrKey = np.array(element.keys, dtype='U')
    hdrValue = np.array(element.values, dtype='U')
    hdrDict = dict(zip(hdrKey,hdrValue))
    
    newHdr=self.myLib.CreateHeader(hdrDict)
    
    # Get the crpix adjustment values for the new header
    crpix1,crpix2,crpix3,crpix4=self.myLib.GetHeaderUpdates(newHdr,\
                                      element.raHeaderIndex, \
                                      element.decHeaderIndex, \
                                      element.freqHeaderIndex)
    # Get the new axis values
    naxis1 = len(ra)
    naxis2 = len(dec)
    naxis4 = len(frequencies)
    
    newHdr['CRPIX1']=float(crpix1)
    newHdr['CRPIX2']=float(crpix2)
    newHdr['CRPIX3']=float(crpix3)
    newHdr['CRPIX4']=float(crpix4)
    
    newHdr['NAXIS1']=float(naxis1)
    newHdr['NAXIS2']=float(naxis2)
    newHdr['NAXIS4']=float(naxis4)
    
    newHdr.pop("END",None)
    
    hdrstr,hdrsize = self.myLib.dict2FITSstr(newHdr)
    
    path_to_par = self.parameterFile
    parsize = len(path_to_par)
    # pass off to sofia C library
    try:
        # This is the call to the shared object
        ret = self.sofia.sofia_mainline(dataPtr,hdrstr,hdrsize,path_to_par,parsize)
        returnCode= ret[0]
        sofiaMsg="Call to sofia has worked!"
    except RuntimeError:
        print("Caught general exception\n")
        sofiaMsg="Caught general exception"
        returnCode=1
        #sys.exit()
    except StopIteration:
        print("Caught Null pointer\n")
        sofiaMsg="Caught Null pointer"
        returnCode=2
        #sys.exit()
    except MemoryError:
        print("Caught ALLOC error\n")
        sofiaMsg="Caught ALLOC error"
        returnCode=3
        #sys.exit()
    except IndexError:
        print("Caught index range error\n")
        sofiaMsg="Caught index range error"
        returnCode=4
        #ys.exit()
    except IOError:
        print("Caught file error\n")
        sofiaMsg="Caught file error"
        returnCode=5
        #sys.exit()
    except OverflowError:
        print("Caught integer overflow error\n")
        sofiaMsg="Caught integer overflow error"
        returnCode=6
        #sys.exit()
    except TypeError:
        print("Caught user input error\n")
        sofiaMsg="Caught user input error"
        returnCode=7
        #sys.exit()
    except SystemExit:
        print("Caught no sources error\n")
        sofiaMsg="Caught no sources error"
        returnCode=8
        #sys.exit()


    if returnCode==0:
        catalogXML=ret[5].tobytes()
        pass
    else:
        catalogXML=""

    DELIMITER=chr(255)                      #"{}{}"
    msg="{}{}{}{}{}{}{}{}{}{}{}{}{}{}{}{}{}"\
    .format( str(returnCode), DELIMITER, \
            str(element.raHeaderIndex), DELIMITER, \
            str(element.decHeaderIndex), DELIMITER, \
            str(len(ra)), DELIMITER, \
            str(len(dec)) , DELIMITER, \
            str(len(frequencies)), DELIMITER, \
            str(element.freqHeaderIndex) , DELIMITER, \
            str(catalogXML), DELIMITER, \
            str(dataPtr.nbytes)
           )
    
    return msg

def processPartition(self, partition):
    self.callPerPartitionSetup()
    for element in partition:
        yield self.doProcess(element)

We create a dataframe where each row of the dataframe represents a 3-dimensional array of data and containing columns of array data for the 3 position dimensions, and the array containing the pixel values (the call in doProcess for ra, dec, frequencies and pixels). The dataframe row also contains coordinate system information from the original source file, which we use to build a new set of header information, which is passed to an instantiated instance of the Partitioner class via a df.rdd.mapPartitions call.

p=Partitioner()
try:           
    ...
    ...
    ...
    ...
    # Creates the positional array dataframes, 
    # and the single rows representing the 3d images in finalSubCubeDF
    # finalSubCube
    #
    ...
    ...
    ...                  
    finalSubcubeDF=finalSubcubeDF\
    .join(broadcast(raRangeDF), finalSubcubeDF.bins == raRangeDF.grp, "inner")\
    .join(broadcast(decRangeDF), finalSubcubeDF.bins == decRangeDF.grp, "inner")\
    .join(broadcast(freqRangeDF), finalSubcubeDF.bins == freqRangeDF.grp, "inner")\
    .join(broadcast(hdrDF), finalSubcubeDF.bins == hdrDF.grp, "inner")\
    .select("bins","right_ascension","declination","frequency","raHeaderIndex",\
            "decHeaderIndex","freqHeaderIndex","pixels","keys","values")

    # repartition on the bins column to distribute the processing 
    finalSubcubeDF=finalSubcubeDF.repartition("bins")
    finalSubcubeDF.persist(StorageLevel.MEMORY_AND_DISK_SER)
    
    ...
    # Calls the partitioner class which containes the C shared object calls, as above
    ...
    
    rddout=finalSubcubeDF.rdd.mapPartitions(p.processPartition)
    DELIMITER=chr(255)
    rddout= rddout.map(lambda x:x.split(DELIMITER))
    ...
    ...
    # Write the results (which activates the computation) including the catalogue XML file
                
    rddout.saveAsTextFile("hdfs:///<file path>")
except Exception as e:
    ...
    ...
    ...     

Error message captured from the calling process

2021-02-16 11:03:42,830 INFO ERROR! ProcessThread FAILURE...An error occurred while calling o1989.collectToPython.
: org.apache.spark.SparkException: Job aborted due to stage failure: Task 5 in stage 47.1 failed 4 times, most recent failure: Lost task 5.3 in stage 47.1 (TID 5875, hercules-1-2.nimbus.pawsey.org.au, executor 12): ExecutorLostFailure (executor 12 exited caused by one of the running tasks) Reason: Container marked as failed: container_1612679558527_0181_01_000014 on host: hercules-1-2.nimbus.pawsey.org.au. Exit status: 143. Diagnostics: Container killed on request. Exit code is 143
Container exited with a non-zero exit code 143
Killed by external signal
2021-02-16 11:03:42,830 INFO ERROR! ProcessThread FAILURE...An error occurred while calling o1989.collectToPython.
: org.apache.spark.SparkException: Job aborted due to stage failure: Task 5 in stage 47.1 failed 4 times, most recent failure: Lost task 5.3 in stage 47.1 (TID 5875, hercules-1-2.nimbus.pawsey.org.au, executor 12): ExecutorLostFailure (executor 12 exited caused by one of the running tasks) Reason: Container marked as failed: container_1612679558527_0181_01_000014 on host: hercules-1-2.nimbus.pawsey.org.au. Exit status: 143. Diagnostics: Container killed on request. Exit code is 143
Container exited with a non-zero exit code 143
Killed by external signal

And the error message from the YARN Container log

LogType:stdout
Log Upload Time:Tue Feb 16 11:19:03 +0000 2021
LogLength:124
Log Contents:
#
# java.lang.OutOfMemoryError: Java heap space
# -XX:OnOutOfMemoryError="kill %p"
#   Executing /bin/sh -c "kill 20417"...

So the problem is obviously memory related, but not entirely sure why that's the case given the executor and driver process memory settings are set quite high? At the moment we're just grasping at straws.

We're aware that collecting distributed data into an array under normal circumstances isn't recommended; however, it seems that being able to run multiple C shared objects in parallel across the cluster could be more efficient than running 30-40 GB extractions serially on a single machine.

Thanks in advance for any thoughts and assistance.

0 Answers
Related