<commit_before>from __future__ import absolute_import
from __future__ import print_function
import os
import os.path
import sys
import time
import random
from six.moves import urllib
import marshal
import struct
import time
import six.moves.cPickle as pickle
import gzip
import zlib
import six.moves.queue
import heapq
import uuid
import gc
import six
import itertools
from operator import itemgetter
from six.moves import range, zip, reduce

try:
    import cStringIO as StringIO
except ImportError:
    from six import BytesIO as StringIO

from dpark.util import decompress, spawn, mkdir_p, atomic_file, get_logger
from dpark.env import env
from dpark.tracker import GetValueMessage, SetValueMessage
from dpark.heaponkey import HeapOnKey
from dpark.dependency import AggregatorBase
from dpark.nested_groupby import GroupByNestedIter, cogroup_no_dup


MAX_SHUFFLE_MEMORY = 2000  # 2 GB

logger = get_logger(__name__)


class LocalFileShuffle:
    @classmethod
    def getOutputFile(cls, shuffle_id, input_id, output_id, datasize=0):
        shuffleDir = env.get('WORKDIR')
        path = os.path.join(shuffleDir[0], str(shuffle_id), str(input_id))
        mkdir_p(path)
        p = os.path.join(path, str(output_id))
        if datasize > 0 and len(shuffleDir) > 1:
            # datasize > 0 means its writing
            st = os.statvfs(path)
            free = st.f_bfree * st.f_bsize
            ratio = st.f_bfree * 1.0 / st.f_blocks
            if free < max(datasize, 1 << 30) or ratio < 0.66:
                d2 = os.path.join(
                    random.choice(shuffleDir[1:]),
                    str(shuffle_id), str(input_id))
                mkdir_p(d2)
                p2 = os.path.join(d2, str(output_id))
                if os.path.exists(p):
                    os.remove(p)
                os.symlink(p2, p)
                if os.path.islink(p2):
                    os.unlink(p2)  # p == p2
                return p2
        return p

    @classmethod
    def getServerUri(cls):
        return env.get('SERVER_URI')


class BadShuffleStreamException(Exception):
    pass


def write_buf(stream, buf, is_mashal):
    buf = zlib.compress(buf, 1)
    size = len(buf)
    stream.write(struct.pack("!I?", size, is_mashal))
    stream.write(buf)


class AutoBatchedSerializer(object):
    """
    Choose the size of batch automatically based on the size of object
    """

    size_loaded = 0

    def __init__(self, best_size=1 << 16):
        self.best_size = best_size
        self.max_num = 0
        self.max_size = 0
        self.use_marshal = True

    def load_stream(self, stream):
        while True:
            head = stream.read(5)
            if not head:
                return
            length, is_mashal = struct.unpack("!I?", head)
            buf = stream.read(length)
            if len(buf) < length:
                raise BadShuffleStreamException("@%d" % stream.tell())
            buf = zlib.decompress(buf)
            AutoBatchedSerializer.size_loaded += len(buf)
            if is_mashal:
                vs = marshal.loads(buf)
            else:
                vs = pickle.loads(buf)
            for v in vs:
                yield v

    def dump_stream(self, iterator, stream):
        self._dump_stream(iter(iterator), stream)
        logger.debug("max batch num = %d, max batch size = %d", self.max_num, self.max_size)

    def _dump_stream(self, iterator, stream):
        batch_num = 1

        while True:
            vs = list(itertools.islice(iterator, batch_num))
            if not vs:
                break
            batch_num = self._dump_batch(stream, vs, batch_num)

    def _dump_batch(self, stream, vs, batch_num):
        if self.use_marshal:
            try:
                buf = marshal.dump(vs)
            except:
                buf = pickle.dumps(vs, -1)
                self.use_marshal = False
        else:
            buf = pickle.dumps(vs, -1)

        mem_size = len(buf)
        write_buf(stream, buf, self.use_marshal)

        if mem_size < self.best_size:
            batch_num *= 2
            if batch_num > self.max_num:<commit_msg>Fix marshal<commit_after>from __future__ import absolute_import
from __future__ import print_function
import os
import os.path
import sys
import time
import random
from six.moves import urllib
import marshal
import struct
import time
import six.moves.cPickle as pickle
import gzip
import zlib
import six.moves.queue
import heapq
import uuid
import gc
import six
import itertools
from operator import itemgetter
from six.moves import range, zip, reduce

try:
    import cStringIO as StringIO
except ImportError:
    from six import BytesIO as StringIO

from dpark.util import decompress, spawn, mkdir_p, atomic_file, get_logger
from dpark.env import env
from dpark.tracker import GetValueMessage, SetValueMessage
from dpark.heaponkey import HeapOnKey
from dpark.dependency import AggregatorBase
from dpark.nested_groupby import GroupByNestedIter, cogroup_no_dup


MAX_SHUFFLE_MEMORY = 2000  # 2 GB

logger = get_logger(__name__)


class LocalFileShuffle:
    @classmethod
    def getOutputFile(cls, shuffle_id, input_id, output_id, datasize=0):
        shuffleDir = env.get('WORKDIR')
        path = os.path.join(shuffleDir[0], str(shuffle_id), str(input_id))
        mkdir_p(path)
        p = os.path.join(path, str(output_id))
        if datasize > 0 and len(shuffleDir) > 1:
            # datasize > 0 means its writing
            st = os.statvfs(path)
            free = st.f_bfree * st.f_bsize
            ratio = st.f_bfree * 1.0 / st.f_blocks
            if free < max(datasize, 1 << 30) or ratio < 0.66:
                d2 = os.path.join(
                    random.choice(shuffleDir[1:]),
                    str(shuffle_id), str(input_id))
                mkdir_p(d2)
                p2 = os.path.join(d2, str(output_id))
                if os.path.exists(p):
                    os.remove(p)
                os.symlink(p2, p)
                if os.path.islink(p2):
                    os.unlink(p2)  # p == p2
                return p2
        return p

    @classmethod
    def getServerUri(cls):
        return env.get('SERVER_URI')


class BadShuffleStreamException(Exception):
    pass


def write_buf(stream, buf, is_mashal):
    buf = zlib.compress(buf, 1)
    size = len(buf)
    stream.write(struct.pack("!I?", size, is_mashal))
    stream.write(buf)


class AutoBatchedSerializer(object):
    """
    Choose the size of batch automatically based on the size of object
    """

    size_loaded = 0

    def __init__(self, best_size=1 << 16):
        self.best_size = best_size
        self.max_num = 0
        self.max_size = 0
        self.use_marshal = True

    def load_stream(self, stream):
        while True:
            head = stream.read(5)
            if not head:
                return
            length, is_marshal = struct.unpack("!I?", head)
            buf = stream.read(length)
            if len(buf) < length:
                raise BadShuffleStreamException("@%d" % stream.tell())
            buf = zlib.decompress(buf)
            AutoBatchedSerializer.size_loaded += len(buf)
            if is_marshal:
                vs = marshal.loads(buf)
            else:
                vs = pickle.loads(buf)
            for v in vs:
                yield v

    def dump_stream(self, iterator, stream):
        self._dump_stream(iter(iterator), stream)
        logger.debug("max batch num = %d, max batch size = %d", self.max_num, self.max_size)

    def _dump_stream(self, iterator, stream):
        batch_num = 1

        while True:
            vs = list(itertools.islice(iterator, batch_num))
            if not vs:
                break
            batch_num = self._dump_batch(stream, vs, batch_num)

    def _dump_batch(self, stream, vs, batch_num):
        if self.use_marshal:
            try:
                buf = marshal.dumps(vs)
            except:
                buf = pickle.dumps(vs, -1)
                self.use_marshal = False
        else:
            buf = pickle.dumps(vs, -1)

        mem_size = len(buf)
        write_buf(stream, buf, self.use_marshal)

        if mem_size < self.best_size:
            batch_num *= 2
            if batch_num > self.max_num:
{'fraction_non_alphanumeric': 0.08842463257621766, 'fraction_numerical': 0.00971699259079315, 'mean_word_length': 2.060966542750929, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '299695', 'n_tokens_mistral': 2674, 'n_tokens_neox': 2535, 'n_words': 668}