[SPARK-1777] Prevent OOMs from single partitions #1165
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -17,9 +17,9 @@ | |
|
|
||
| package org.apache.spark | ||
|
|
||
| import scala.collection.mutable | ||
| import scala.collection.mutable.ArrayBuffer | ||
|
|
||
| import org.apache.spark.rdd.RDD | ||
| import org.apache.spark.storage._ | ||
|
|
||
|
|
@@ -30,7 +30,7 @@ import org.apache.spark.storage._ | |
| private[spark] class CacheManager(blockManager: BlockManager) extends Logging { | ||
|
|
||
| /** Keys of RDD partitions that are being computed/loaded. */ | ||
| private val loading = new mutable.HashSet[RDDBlockId] | ||
|
|
||
| /** Gets or computes an RDD partition. Used by RDD.iterator() when an RDD is cached. */ | ||
| def getOrCompute[T]( | ||
|
|
@@ -118,36 +118,66 @@ private[spark] class CacheManager(blockManager: BlockManager) extends Logging { | |
| } | ||
|
|
||
| /** | ||
| * Cache the values of a partition, keeping track of any updates in the storage statuses of | ||
| * other blocks along the way. | ||
| * | ||
| * The effective storage level refers to the level that actually specifies BlockManager put | ||
| * behavior, not the level originally specified by the user. This is mainly for forcing a | ||
| * MEMORY_AND_DISK partition to disk if there is not enough room to unroll the partition, | ||
| * while preserving the the original semantics of the RDD as specified by the application. | ||
| */ | ||
| private def putInBlockManager[T]( | ||
| key: BlockId, | ||
| values: Iterator[T], | ||
| level: StorageLevel, | ||
| updatedBlocks: ArrayBuffer[(BlockId, BlockStatus)], | ||
| effectiveStorageLevel: Option[StorageLevel] = None): Iterator[T] = { | ||
|
There was a problem hiding this comment. This should be documented in the javadoc here - it's not at all obvious what it means. There was a problem hiding this comment. @andrewor14 , since the Take another case for example, if original Level is MEM_ONLY, and the user force it to cache on disk via effectiveStorageLevel, then the BlockManager will never know the block is on disk, and the block will be a zombie. Of course , such case would not likely happen. |
||
|
|
||
| val putLevel = effectiveStorageLevel.getOrElse(level) | ||
| if (!putLevel.useMemory) { | ||
| /* | ||
| * This RDD is not to be cached in memory, so we can just pass the computed values as an | ||
| * iterator directly to the BlockManager rather than first fully unrolling it in memory. | ||
| */ | ||
| updatedBlocks ++= | ||
| blockManager.putIterator(key, values, level, tellMaster = true, effectiveStorageLevel) | ||
| blockManager.get(key) match { | ||
| case Some(v) => v.data.asInstanceOf[Iterator[T]] | ||
| case None => | ||
| logInfo(s"Failure to store $key") | ||
| throw new BlockException(key, s"Block manager failed to return cached value for $key!") | ||
| } | ||
| } else { | ||
| /* | ||
| * This RDD is to be cached in memory. In this case we cannot pass the computed values | ||
| * to the BlockManager as an iterator and expect to read it back later. This is because | ||
| * we may end up dropping a partition from memory store before getting it back. | ||
| * | ||
| * In addition, we must be careful to not unroll the entire partition in memory at once. | ||
| * Otherwise, we may cause an OOM exception if the JVM does not have enough space for this | ||
| * single partition. Instead, we unroll the values cautiously, potentially aborting and | ||
| * dropping the partition to disk if applicable. | ||
| */ | ||
| blockManager.memoryStore.unrollSafely(key, values, updatedBlocks) match { | ||
| case Left(arr) => | ||
| // We have successfully unrolled the entire partition, so cache it in memory | ||
| updatedBlocks ++= | ||
| blockManager.putArray(key, arr, level, tellMaster = true, effectiveStorageLevel) | ||
| arr.iterator.asInstanceOf[Iterator[T]] | ||
| case Right(it) => | ||
| // There is not enough space to cache this partition in memory | ||
| logWarning(s"Not enough space to cache partition $key in memory! " + | ||
| s"Free memory is ${blockManager.memoryStore.freeMemory} bytes.") | ||
| val returnValues = it.asInstanceOf[Iterator[T]] | ||
| if (putLevel.useDisk) { | ||
| logWarning(s"Persisting partition $key to disk instead.") | ||
| val diskOnlyLevel = StorageLevel(useDisk = true, useMemory = false, | ||
| useOffHeap = false, deserialized = false, putLevel.replication) | ||
| putInBlockManager[T](key, returnValues, level, updatedBlocks, Some(diskOnlyLevel)) | ||
| } else { | ||
| returnValues | ||
| } | ||
|
There was a problem hiding this comment. In case of exceptions/etc in this method, we end up with every increasing cacheMemoryMap. |
||
| } | ||
|
There was a problem hiding this comment. I wonder if we could move this logic to BlockManager later. Probably can be part of another PR, but it's pretty complicated to have this here and then similar logic in there when it's doing a get (not to mention in its put methods). There was a problem hiding this comment. Yes. The existing BM interface does not return the values you just put (for good reasons), but once we add it, this entire method can be simplified to a wrapper, and there won't be duplicate logic between here and the |
||
| } | ||
| } | ||
|
|
||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
As we have
putIteratorinterface now, could we just let theblockManagerto handle this logic? CallingunrollSafelyhere looks a little coupling, maybe it's better to markunrollSafelyas private and hide it outside memory store.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The logic here is first introduced at 40566e1.
It was aimed to avoid unroll the iterator into an ArrayBuffer if block is
DISK_ONLY. But with the appearance ofputIteratorinterface inblockManager, this logic can be easily done inblockManagerand we can clean the code here. Did I miss something here?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Ah, this part is actually really tricky. If you just used
putIteratorhere, the result would be incorrect and the reason behind that is quite subtle. Here ingetOrComputewe need to return the actual iterator in addition to storing it inBlockManager, so if we just useputIteratorwithMEMORY_ONLYlevel, then other threads might drop our block before we get to read it back, in which case we will have nothing to return because our original iterator was already exhausted. We have no other choice but to attempt to unroll it here first.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Ah, got it. So it's a follow of the origin implementation:
Then we can ensure the returned iterator will always have data for user.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes, except the original implementation OOM's easily.