<commit_before>from __future__ import absolute_import
import marshal
import time
import six.moves.cPickle
import os
import os.path

import dpark.conf
from dpark.env import env
from dpark.util import compress, get_logger, ERROR_TASK_OOM
from dpark.serialize import marshalable, load_func, dump_func, dumps, loads
from dpark.shuffle import LocalFileShuffle, get_serializer, Merger, pack_header
from six.moves import range

logger = get_logger(__name__)

class Task:
    def __init__(self):
        self.id = Task.newId()

    nextId = 0
    @classmethod
    def newId(cls):
        cls.nextId += 1
        return cls.nextId

    def run(self, task_id):
        raise NotImplementedError

    def preferredLocations(self):
        raise NotImplementedError


class DAGTask(Task):
    def __init__(self, stageId):
        Task.__init__(self)
        self.stageId = stageId
        self.mem = 0

    def __repr__(self):
        return '<task %d:%d>'%(self.stageId, self.id)

    def run(self, task_id):
        try:
            if self.mem != 0:
                env.meminfo.start(task_id, int(self.mem))
                if dpark.conf.MULTI_SEGMENT_DUMP:
                    env.meminfo.check = False
            return self._run(task_id)
        except Exception as e:
            logger.exception("Task.run error")
        except KeyboardInterrupt as e:
            if self.mem != 0 and env.meminfo.oom:
                os._exit(ERROR_TASK_OOM)
            else:
                raise e
        finally:
            if self.mem != 0:
                env.meminfo.check = True
                env.meminfo.stop()


class ResultTask(DAGTask):
    def __init__(self, stageId, rdd, func, partition, locs, outputId):
        DAGTask.__init__(self, stageId)
        self.rdd = rdd
        self.func = func
        self.partition = partition
        self.split = rdd.splits[partition]
        self.locs = locs
        self.outputId = outputId

    def _run(self, task_id):
        logger.debug("run task %s: %s", task_id, self)
        return self.func(self.rdd.iterator(self.split))

    def preferredLocations(self):
        return self.locs

    def __repr__(self):
        partition = getattr(self, 'partition', None)
        rdd = getattr(self, 'rdd', None)
        return "<ResultTask(%s) of %s" % (partition, rdd)

    def __getstate__(self):
        d = dict(self.__dict__)
        del d['func']
        del d['rdd']
        del d['split']
        return d, dumps(self.rdd), dump_func(self.func), dumps(self.split)

    def __setstate__(self, state):
        d, rdd, func, split = state
        self.__dict__.update(d)
        self.rdd = loads(rdd)
        self.func = load_func(func)
        self.split = loads(split)


class ShuffleMapTask(DAGTask):
    def __init__(self, stageId, rdd, dep, partition, locs):
        DAGTask.__init__(self, stageId)
        self.rdd = rdd
        self.shuffleId = dep.shuffleId
        self.aggregator = dep.aggregator
        self.partitioner = dep.partitioner
        self.shuffle_config = dep.shuffle_config
        self.partition = partition
        self.split = rdd.splits[partition]
        self.locs = locs

    def __repr__(self):
        shuffleId = getattr(self, 'shuffleId', None)
        partition = getattr(self, 'partition', None)
        rdd = getattr(self, 'rdd', None)
        return '<ShuffleTask(%s, %s) of %s>' % (shuffleId, partition, rdd)

    def __getstate__(self):
        d = dict(self.__dict__)
        del d['rdd']
        del d['split']
        return d, dumps(self.rdd), dumps(self.split)

    def __setstate__(self, state):
        d, rdd, split = state
        self.__dict__.update(d)
        self.rdd = loads(rdd)
        self.split = loads(split)

    def preferredLocations(self):
        return self.locs

    def _run(self, task_id):
        mem_limit = env.meminfo.mem_limit_soft
        t0 = time.time()
        logger.debug("run task with shuffle_flag %r" % (self.shuffle_config, ))
        rdd = self.rdd
        meminfo = env.meminfo
        n = self.partitioner.numPartitions
        get_partition = self.partitioner.getPartition
        merge_value = self.aggregator.mergeValue
        create_combiner = self.aggregator.createCombiner
        dumper_cls = SortMergeBucketDumper if self.shuffle_config.is_sort_merge else BucketDumper
        dumper = dumper_cls(self.shuffleId, self.partition, n, self.shuffle_config)
        buckets = [{} for _ in range(n)]

        last_i = 0
        for i, item in enumerate(rdd.iterator(self.split)):
            try:
                k, v = item
                bucket = buckets[get_partition(k)]
                r = bucket.get(k, None)
                if r is not None:
                    bucket[k] = merge_value(r, v)
                else:
                    bucket[k] = create_combiner(v)<commit_msg>Reserve enough mem for dumps for map dumps<commit_after>from __future__ import absolute_import
import marshal
import time
import six.moves.cPickle
import os
import os.path

import dpark.conf
from dpark.env import env
from dpark.util import compress, get_logger, ERROR_TASK_OOM
from dpark.serialize import marshalable, load_func, dump_func, dumps, loads
from dpark.shuffle import LocalFileShuffle, get_serializer, Merger, pack_header
from six.moves import range

logger = get_logger(__name__)

class Task:
    def __init__(self):
        self.id = Task.newId()

    nextId = 0
    @classmethod
    def newId(cls):
        cls.nextId += 1
        return cls.nextId

    def run(self, task_id):
        raise NotImplementedError

    def preferredLocations(self):
        raise NotImplementedError


class DAGTask(Task):
    def __init__(self, stageId):
        Task.__init__(self)
        self.stageId = stageId
        self.mem = 0

    def __repr__(self):
        return '<task %d:%d>'%(self.stageId, self.id)

    def run(self, task_id):
        try:
            if self.mem != 0:
                env.meminfo.start(task_id, int(self.mem))
                if dpark.conf.MULTI_SEGMENT_DUMP:
                    env.meminfo.check = False
            return self._run(task_id)
        except Exception as e:
            logger.exception("Task.run error")
        except KeyboardInterrupt as e:
            if self.mem != 0 and env.meminfo.oom:
                os._exit(ERROR_TASK_OOM)
            else:
                raise e
        finally:
            if self.mem != 0:
                env.meminfo.check = True
                env.meminfo.stop()


class ResultTask(DAGTask):
    def __init__(self, stageId, rdd, func, partition, locs, outputId):
        DAGTask.__init__(self, stageId)
        self.rdd = rdd
        self.func = func
        self.partition = partition
        self.split = rdd.splits[partition]
        self.locs = locs
        self.outputId = outputId

    def _run(self, task_id):
        logger.debug("run task %s: %s", task_id, self)
        return self.func(self.rdd.iterator(self.split))

    def preferredLocations(self):
        return self.locs

    def __repr__(self):
        partition = getattr(self, 'partition', None)
        rdd = getattr(self, 'rdd', None)
        return "<ResultTask(%s) of %s" % (partition, rdd)

    def __getstate__(self):
        d = dict(self.__dict__)
        del d['func']
        del d['rdd']
        del d['split']
        return d, dumps(self.rdd), dump_func(self.func), dumps(self.split)

    def __setstate__(self, state):
        d, rdd, func, split = state
        self.__dict__.update(d)
        self.rdd = loads(rdd)
        self.func = load_func(func)
        self.split = loads(split)


class ShuffleMapTask(DAGTask):
    def __init__(self, stageId, rdd, dep, partition, locs):
        DAGTask.__init__(self, stageId)
        self.rdd = rdd
        self.shuffleId = dep.shuffleId
        self.aggregator = dep.aggregator
        self.partitioner = dep.partitioner
        self.shuffle_config = dep.shuffle_config
        self.partition = partition
        self.split = rdd.splits[partition]
        self.locs = locs

    def __repr__(self):
        shuffleId = getattr(self, 'shuffleId', None)
        partition = getattr(self, 'partition', None)
        rdd = getattr(self, 'rdd', None)
        return '<ShuffleTask(%s, %s) of %s>' % (shuffleId, partition, rdd)

    def __getstate__(self):
        d = dict(self.__dict__)
        del d['rdd']
        del d['split']
        return d, dumps(self.rdd), dumps(self.split)

    def __setstate__(self, state):
        d, rdd, split = state
        self.__dict__.update(d)
        self.rdd = loads(rdd)
        self.split = loads(split)

    def preferredLocations(self):
        return self.locs

    def _run(self, task_id):
        mem_limit = env.meminfo.mem_limit_soft
        t0 = time.time()
        logger.debug("run task with shuffle_flag %r" % (self.shuffle_config, ))
        rdd = self.rdd
        meminfo = env.meminfo
        n = self.partitioner.numPartitions
        get_partition = self.partitioner.getPartition
        merge_value = self.aggregator.mergeValue
        create_combiner = self.aggregator.createCombiner
        dumper_cls = SortMergeBucketDumper if self.shuffle_config.is_sort_merge else BucketDumper
        dumper = dumper_cls(self.shuffleId, self.partition, n, self.shuffle_config)
        buckets = [{} for _ in range(n)]
        env.meminfo.ratio = min(float(n) / (n+1), env.meminfo.ratio)

        last_i = 0
        for i, item in enumerate(rdd.iterator(self.split)):
            try:
                k, v = item
                bucket = buckets[get_partition(k)]
                r = bucket.get(k, None)
                if r is not None:
                    bucket[k] = merge_value(r, v)
                else:
                    bucket[k] = create_combiner(v)
{'fraction_non_alphanumeric': 0.10213727907932593, 'fraction_numerical': 0.0017468146321413893, 'mean_word_length': 2.154943273905997, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 7, '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': '599027', 'n_tokens_mistral': 3129, 'n_tokens_neox': 3041, 'n_words': 724}