Github user davies commented on a diff in the pull request:
https://github.com/apache/spark/pull/1460#discussion_r15300037
--- Diff: python/pyspark/shuffle.py ---
@@ -0,0 +1,416 @@
+#
+# 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.
+#
+
+import os
+import sys
+import platform
+import shutil
+import warnings
+
+from pyspark.serializers import BatchedSerializer, PickleSerializer
+
+try:
+ import psutil
+
+ def get_used_memory():
+ """ return the used memory in MB """
+ self = psutil.Process(os.getpid())
+ return self.memory_info().rss >> 20
+
+except ImportError:
+
+ def get_used_memory():
+ """ return the used memory in MB """
+ if platform.system() == 'Linux':
+ for line in open('/proc/self/status'):
+ if line.startswith('VmRSS:'):
+ return int(line.split()[1]) >> 10
+ else:
+ warnings.warn("please install psutil to have better "
+ "support with spilling")
+ if platform.system() == "Darwin":
+ import resource
+ rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+ return rss >> 20
+ # TODO: support windows
+ return 0
+
+
+class Aggregator(object):
+
+ """
+ Aggregator has tree functions to merge values into combiner.
+
+ createCombiner: (value) -> combiner
+ mergeValue: (combine, value) -> combiner
+ mergeCombiners: (combiner, combiner) -> combiner
+ """
+
+ def __init__(self, createCombiner, mergeValue, mergeCombiners):
+ self.createCombiner = createCombiner
+ self.mergeValue = mergeValue
+ self.mergeCombiners = mergeCombiners
+
+
+class SimpleAggregator(Aggregator):
+
+ """
+ SimpleAggregator is useful for the cases that combiners have
+ same type with values
+ """
+
+ def __init__(self, combiner):
+ Aggregator.__init__(self, lambda x: x, combiner, combiner)
+
+
+class Merger(object):
+
+ """
+ merge shuffled data together by aggregator
+ """
+
+ def __init__(self, aggregator):
+ self.agg = aggregator
+
+ def mergeValues(self, iterator):
+ """ combine the items by creator and combiner """
+ raise NotImplementedError
+
+ def mergeCombiners(self, iterator):
+ """ merge the combined items by mergeCombiner """
+ raise NotImplementedError
+
+ def iteritems(self):
+ """ return the merged items ad iterator """
+ raise NotImplementedError
+
+
+class InMemoryMerger(Merger):
+
+ """
+ In memory merger based on in-memory dict.
+ """
+
+ def __init__(self, aggregator):
+ Merger.__init__(self, aggregator)
+ self.data = {}
+
+ def mergeValues(self, iterator):
+ """ combine the items by creator and combiner """
+ # speed up attributes lookup
+ d, creator = self.data, self.agg.createCombiner
+ comb = self.agg.mergeValue
+ for k, v in iterator:
+ d[k] = comb(d[k], v) if k in d else creator(v)
+
+ def mergeCombiners(self, iterator):
+ """ merge the combined items by mergeCombiner """
+ # speed up attributes lookup
+ d, comb = self.data, self.agg.mergeCombiners
+ for k, v in iterator:
+ d[k] = comb(d[k], v) if k in d else v
+
+ def iteritems(self):
+ """ return the merged items ad iterator """
+ return self.data.iteritems()
+
+
+class ExternalMerger(Merger):
+
+ """
+ External merger will dump the aggregated data into disks when
+ memory usage goes above the limit, then merge them together.
+
+ This class works as follows:
+
+ - It repeatedly combine the items and save them in one dict in
+ memory.
+
+ - When the used memory goes above memory limit, it will split
+ the combined data into partitions by hash code, dump them
+ into disk, one file per partition.
+
+ - Then it goes through the rest of the iterator, combine items
+ into different dict by hash. Until the used memory goes over
+ memory limit, it dump all the dicts into disks, one file per
+ dict. Repeat this again until combine all the items.
+
+ - Before return any items, it will load each partition and
+ combine them seperately. Yield them before loading next
+ partition.
+
+ - During loading a partition, if the memory goes over limit,
+ it will partition the loaded data and dump them into disks
+ and load them partition by partition again.
+
+ `data` and `pdata` are used to hold the merged items in memory.
+ At first, all the data are merged into `data`. Once the used
+ memory goes over limit, the items in `data` are dumped indo
+ disks, `data` will be cleared, all rest of items will be merged
+ into `pdata` and then dumped into disks. Before returning, all
+ the items in `pdata` will be dumped into disks.
+
+ Finally, if any items were spilled into disks, each partition
+ will be merged into `data` and be yielded, then cleared.
+
+ >>> agg = SimpleAggregator(lambda x, y: x + y)
+ >>> merger = ExternalMerger(agg, 10)
+ >>> N = 10000
+ >>> merger.mergeValues(zip(xrange(N), xrange(N)) * 10)
+ >>> assert merger.spills > 0
+ >>> sum(v for k,v in merger.iteritems())
+ 499950000
+
+ >>> merger = ExternalMerger(agg, 10)
+ >>> merger.mergeCombiners(zip(xrange(N), xrange(N)) * 10)
+ >>> assert merger.spills > 0
+ >>> sum(v for k,v in merger.iteritems())
+ 499950000
+ """
+
+ def __init__(self, aggregator, memory_limit=512, serializer=None,
+ localdirs=None, scale=1, partitions=64, batch=10000):
+ Merger.__init__(self, aggregator)
+ self.memory_limit = memory_limit
+ # default serializer is only used for tests
+ self.serializer = serializer or \
+ BatchedSerializer(PickleSerializer(), 1024)
+ self.localdirs = localdirs or self._get_dirs()
+ # number of partitions when spill data into disks
+ self.partitions = partitions
+ # check the memory after # of items merged
+ self.batch = batch
+ # scale is used to scale down the hash of key for recursive hash
map,
+ self.scale = scale
+ # unpartitioned merged data
+ self.data = {}
+ # partitioned merged data
+ self.pdata = []
+ # number of chunks dumped into disks
+ self.spills = 0
+
+ def _get_dirs(self):
+ """ get all the directories """
+ path = os.environ.get("SPARK_LOCAL_DIR", "/tmp/spark")
+ dirs = path.split(",")
+ return [os.path.join(d, "python", str(os.getpid()), str(id(self)))
+ for d in dirs]
+
+ def _get_spill_dir(self, n):
+ """ choose one directory for spill by number n """
+ return os.path.join(self.localdirs[n % len(self.localdirs)],
str(n))
+
+ def next_limit(self):
+ """
+ return the next memory limit. If the memory is not released
+ after spilling, it will dump the data only when the used memory
+ starts to increase.
+ """
+ return max(self.memory_limit, get_used_memory() * 1.05)
+
+ def mergeValues(self, iterator):
+ """ combine the items by creator and combiner """
+ iterator = iter(iterator)
+ # speedup attribute lookup
+ creator, comb = self.agg.createCombiner, self.agg.mergeValue
+ d, c, batch = self.data, 0, self.batch
+
+ for k, v in iterator:
+ d[k] = comb(d[k], v) if k in d else creator(v)
+
+ c += 1
+ if c % batch == 0 and get_used_memory() > self.memory_limit:
+ self._first_spill()
+ self._partitioned_mergeValues(iterator, self.next_limit())
+ break
+
+ def _partition(self, key):
+ """ return the partition for key """
+ return (hash(key) / self.scale) % self.partitions
+
+ def _partitioned_mergeValues(self, iterator, limit=0):
+ """ partition the items by key, then combine them """
+ # speedup attribute lookup
+ creator, comb = self.agg.createCombiner, self.agg.mergeValue
+ c, pdata, hfun, batch = 0, self.pdata, self._partition, self.batch
+
+ for k, v in iterator:
+ d = pdata[hfun(k)]
+ d[k] = comb(d[k], v) if k in d else creator(v)
+ if not limit:
+ continue
+
+ c += 1
+ if c % batch == 0 and get_used_memory() > limit:
+ self._spill()
+ limit = self.next_limit()
+
+ def mergeCombiners(self, iterator, check=True):
+ """ merge (K,V) pair by mergeCombiner """
+ iterator = iter(iterator)
+ # speedup attribute lookup
+ d, comb, batch = self.data, self.agg.mergeCombiners, self.batch
+ c = 0
+ for k, v in iterator:
+ d[k] = comb(d[k], v) if k in d else v
+ if not check:
+ continue
+
+ c += 1
+ if c % batch == 0 and get_used_memory() > self.memory_limit:
+ self._first_spill()
+ self._partitioned_mergeCombiners(iterator,
self.next_limit())
+ break
+
+ def _partitioned_mergeCombiners(self, iterator, limit=0):
+ """ partition the items by key, then merge them """
+ comb, pdata = self.agg.mergeCombiners, self.pdata
+ c, hfun = 0, self._partition
+ for k, v in iterator:
+ d = pdata[hfun(k)]
+ d[k] = comb(d[k], v) if k in d else v
+ if not limit:
+ continue
+
+ c += 1
+ if c % self.batch == 0 and get_used_memory() > limit:
+ self._spill()
+ limit = self.next_limit()
+
+ def _first_spill(self):
+ """
+ dump all the data into disks partition by partition.
+
+ The data has not been partitioned, it will iterator the
+ dataset once, write them into different files, has no
+ additional memory. It only called when the memory goes
+ above limit at the first time.
+ """
+ path = self._get_spill_dir(self.spills)
+ if not os.path.exists(path):
+ os.makedirs(path)
+ # open all the files for writing
+ streams = [open(os.path.join(path, str(i)), 'w')
+ for i in range(self.partitions)]
+
+ for k, v in self.data.iteritems():
+ h = self._partition(k)
+ # put one item in batch, make it compatitable with load_stream
+ # it will increase the memory if dump them in batch
+ self.serializer.dump_stream([(k, v)], streams[h])
+ for s in streams:
+ s.close()
+ self.data.clear()
+ self.pdata = [{} for i in range(self.partitions)]
+ self.spills += 1
+
+ def _spill(self):
+ """
+ dump already partitioned data into disks.
+
+ It will dump the data in batch for better performance.
+ """
+ path = self._get_spill_dir(self.spills)
+ if not os.path.exists(path):
+ os.makedirs(path)
+
+ for i in range(self.partitions):
+ p = os.path.join(path, str(i))
+ with open(p, "w") as f:
+ # dump items in batch
+ self.serializer.dump_stream(self.pdata[i].iteritems(), f)
+ self.pdata[i].clear()
+ self.spills += 1
+
+ def iteritems(self):
+ """ return all merged items as iterator """
+ if not self.pdata and not self.spills:
+ return self.data.iteritems()
+ return self._external_items()
+
+ def _external_items(self):
+ """ return all partitioned items as iterator """
+ assert not self.data
+ if any(self.pdata):
+ self._spill()
+ hard_limit = self.next_limit()
+
+ try:
+ for i in range(self.partitions):
+ self.data = {}
+ for j in range(self.spills):
+ path = self._get_spill_dir(j)
+ p = os.path.join(path, str(i))
+ # do not check memory during merging
+
self.mergeCombiners(self.serializer.load_stream(open(p)),
+ False)
+
+ if get_used_memory() > hard_limit and j < self.spills
- 1:
+ self.data.clear() # will read from disk again
+ for v in self._recursive_merged_items(i):
+ yield v
+ return
+
+ for v in self.data.iteritems():
+ yield v
+ self.data.clear()
+ finally:
+ self._cleanup()
+
+ def _cleanup(self):
+ """ clean up all the files in disks """
+ for d in self.localdirs:
+ shutil.rmtree(d, True)
--- End diff --
It will remove the older files for each recursive spill. Also, I will
remove the files after merge each partition.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at [email protected] or file a JIRA ticket
with INFRA.
---
