Skip to content
Snippets Groups Projects
  • Davies Liu's avatar
    4a377aff
    [SPARK-3721] [PySpark] broadcast objects larger than 2G · 4a377aff
    Davies Liu authored
    This patch will bring support for broadcasting objects larger than 2G.
    
    pickle, zlib, FrameSerializer and Array[Byte] all can not support objects larger than 2G, so this patch introduce LargeObjectSerializer to serialize broadcast objects, the object will be serialized and compressed into small chunks, it also change the type of Broadcast[Array[Byte]]] into Broadcast[Array[Array[Byte]]]].
    
    Testing for support broadcast objects larger than 2G is slow and memory hungry, so this is tested manually, could be added into SparkPerf.
    
    Author: Davies Liu <davies@databricks.com>
    Author: Davies Liu <davies.liu@gmail.com>
    
    Closes #2659 from davies/huge and squashes the following commits:
    
    7b57a14 [Davies Liu] add more tests for broadcast
    28acff9 [Davies Liu] Merge branch 'master' of github.com:apache/spark into huge
    a2f6a02 [Davies Liu] bug fix
    4820613 [Davies Liu] Merge branch 'master' of github.com:apache/spark into huge
    5875c73 [Davies Liu] address comments
    10a349b [Davies Liu] address comments
    0c33016 [Davies Liu] Merge branch 'master' of github.com:apache/spark into huge
    6182c8f [Davies Liu] Merge branch 'master' into huge
    d94b68f [Davies Liu] Merge branch 'master' of github.com:apache/spark into huge
    2514848 [Davies Liu] address comments
    fda395b [Davies Liu] Merge branch 'master' of github.com:apache/spark into huge
    1c2d928 [Davies Liu] fix scala style
    091b107 [Davies Liu] broadcast objects larger than 2G
    4a377aff
    History
    [SPARK-3721] [PySpark] broadcast objects larger than 2G
    Davies Liu authored
    This patch will bring support for broadcasting objects larger than 2G.
    
    pickle, zlib, FrameSerializer and Array[Byte] all can not support objects larger than 2G, so this patch introduce LargeObjectSerializer to serialize broadcast objects, the object will be serialized and compressed into small chunks, it also change the type of Broadcast[Array[Byte]]] into Broadcast[Array[Array[Byte]]]].
    
    Testing for support broadcast objects larger than 2G is slow and memory hungry, so this is tested manually, could be added into SparkPerf.
    
    Author: Davies Liu <davies@databricks.com>
    Author: Davies Liu <davies.liu@gmail.com>
    
    Closes #2659 from davies/huge and squashes the following commits:
    
    7b57a14 [Davies Liu] add more tests for broadcast
    28acff9 [Davies Liu] Merge branch 'master' of github.com:apache/spark into huge
    a2f6a02 [Davies Liu] bug fix
    4820613 [Davies Liu] Merge branch 'master' of github.com:apache/spark into huge
    5875c73 [Davies Liu] address comments
    10a349b [Davies Liu] address comments
    0c33016 [Davies Liu] Merge branch 'master' of github.com:apache/spark into huge
    6182c8f [Davies Liu] Merge branch 'master' into huge
    d94b68f [Davies Liu] Merge branch 'master' of github.com:apache/spark into huge
    2514848 [Davies Liu] address comments
    fda395b [Davies Liu] Merge branch 'master' of github.com:apache/spark into huge
    1c2d928 [Davies Liu] fix scala style
    091b107 [Davies Liu] broadcast objects larger than 2G
worker.py 5.38 KiB
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""
Worker that receives input from Piped RDD.
"""
import os
import sys
import time
import socket
import traceback
import cProfile
import pstats

from pyspark.accumulators import _accumulatorRegistry
from pyspark.broadcast import Broadcast, _broadcastRegistry
from pyspark.files import SparkFiles
from pyspark.serializers import write_with_length, write_int, read_long, \
    write_long, read_int, SpecialLengths, UTF8Deserializer, PickleSerializer, \
    SizeLimitedStream, LargeObjectSerializer
from pyspark import shuffle

pickleSer = PickleSerializer()
utf8_deserializer = UTF8Deserializer()


def report_times(outfile, boot, init, finish):
    write_int(SpecialLengths.TIMING_DATA, outfile)
    write_long(1000 * boot, outfile)
    write_long(1000 * init, outfile)
    write_long(1000 * finish, outfile)


def add_path(path):
    # worker can be used, so donot add path multiple times
    if path not in sys.path:
        # overwrite system packages
        sys.path.insert(1, path)


def main(infile, outfile):
    try:
        boot_time = time.time()
        split_index = read_int(infile)
        if split_index == -1:  # for unit tests
            exit(-1)

        # initialize global state
        shuffle.MemoryBytesSpilled = 0
        shuffle.DiskBytesSpilled = 0
        _accumulatorRegistry.clear()

        # fetch name of workdir
        spark_files_dir = utf8_deserializer.loads(infile)
        SparkFiles._root_directory = spark_files_dir
        SparkFiles._is_running_on_worker = True

        # fetch names of includes (*.zip and *.egg files) and construct PYTHONPATH
        add_path(spark_files_dir)  # *.py files that were added will be copied here
        num_python_includes = read_int(infile)
        for _ in range(num_python_includes):
            filename = utf8_deserializer.loads(infile)
            add_path(os.path.join(spark_files_dir, filename))

        # fetch names and values of broadcast variables
        num_broadcast_variables = read_int(infile)
        bser = LargeObjectSerializer()
        for _ in range(num_broadcast_variables):
            bid = read_long(infile)
            if bid >= 0:
                size = read_long(infile)
                s = SizeLimitedStream(infile, size)
                value = list((bser.load_stream(s)))[0]  # read out all the bytes
                _broadcastRegistry[bid] = Broadcast(bid, value)
            else:
                bid = - bid - 1
                _broadcastRegistry.pop(bid)

        _accumulatorRegistry.clear()
        command = pickleSer._read_with_length(infile)
        if isinstance(command, Broadcast):
            command = pickleSer.loads(command.value)
        (func, stats, deserializer, serializer) = command
        init_time = time.time()

        def process():
            iterator = deserializer.load_stream(infile)
            serializer.dump_stream(func(split_index, iterator), outfile)

        if stats:
            p = cProfile.Profile()
            p.runcall(process)
            st = pstats.Stats(p)
            st.stream = None  # make it picklable
            stats.add(st.strip_dirs())
        else:
            process()
    except Exception:
        try:
            write_int(SpecialLengths.PYTHON_EXCEPTION_THROWN, outfile)
            write_with_length(traceback.format_exc(), outfile)
        except IOError:
            # JVM close the socket
            pass
        except Exception:
            # Write the error to stderr if it happened while serializing
            print >> sys.stderr, "PySpark worker failed with exception:"
            print >> sys.stderr, traceback.format_exc()
        exit(-1)
    finish_time = time.time()
    report_times(outfile, boot_time, init_time, finish_time)
    write_long(shuffle.MemoryBytesSpilled, outfile)
    write_long(shuffle.DiskBytesSpilled, outfile)

    # Mark the beginning of the accumulators section of the output
    write_int(SpecialLengths.END_OF_DATA_SECTION, outfile)
    write_int(len(_accumulatorRegistry), outfile)
    for (aid, accum) in _accumulatorRegistry.items():
        pickleSer._write_with_length((aid, accum._value), outfile)

    # check end of stream
    if read_int(infile) == SpecialLengths.END_OF_STREAM:
        write_int(SpecialLengths.END_OF_STREAM, outfile)
    else:
        # write a different value to tell JVM to not reuse this worker
        write_int(SpecialLengths.END_OF_DATA_SECTION, outfile)
        exit(-1)


if __name__ == '__main__':
    # Read a local port to connect to from stdin
    java_port = int(sys.stdin.readline())
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.connect(("127.0.0.1", java_port))
    sock_file = sock.makefile("a+", 65536)
    main(sock_file, sock_file)