[SPARK-21319][SQL] Fix memory leak in UnsafeExternalRowSorter.RowComparator #18543
Conversation
UnsafeExternalRowSorter.RowComparator contains references to the objects backing the last arrays sorted. This causes memory leaks, since those objects become deallocated but are still live. We make sure to unset them.
|
Seems reasonable, though I don't know this code well. There is a similar pattern in KVComparator. I wonder if it's even more efficient to have a |
j-baker
force-pushed
the
jbaker/fix_memory_leak
branch
from
cdf869a to
5e7a935
Compare
|
ok to test |
|
Test build #79247 has finished for PR 18543 at commit
|
|
Can confirm that my job which was exhibiting OOMs no longer OOMs after this change. |
| return ordering.compare(row1, row2); | ||
| int comparison = ordering.compare(row1, row2); | ||
| row1.pointTo(null, 0L, -1); | ||
| row2.pointTo(null, 0L, -1); |
There was a problem hiding this comment.
can we avoid to do it per comparison? is there any places we can do a cleanup at the end?
There was a problem hiding this comment.
Good idea @cloud-fan . Looks like RowComparator and KVComparator could be cleaned up in UnsafeExternalSorter.cleanupResources and UnsafeInMemorySorter.free. @davies I know this is from a long while back, but does that make sense? Seems like reasonable places to simply 'flush' the references, and won't hurt anything.
I am not 100% sure there's not another path to take care of, or if this frees the ref soon enough to avoid the problem. Thoughts @j-baker ?
There was a problem hiding this comment.
@cloud-fan @srowen It is good idea to do this cleanup only once at the end. Now, I am curious how to implement this cleanup.
While @srowen proposed to use nsafeExternalSorter.cleanupResources and UnsafeInMemorySorter.free that will be called when a task is finished, to do cleanup here does not seem to work in this case. This is because this issue occurs before completing a task since UnsafeExternalSorter instance is registered into the task taskContext at here. This cleanup approach will not be performed before an OOM occurs during execution of the task.
IIUC, the end of sort is here. This line calls this sort method. Either to do the cleanup at the first part or to do the cleanup after checking type of a given comparator at the second part could work.
What do you think?
There was a problem hiding this comment.
So I feel like if the thing ends up in memory, this is correct - but otherwise the comparator is used in the UnsafeSorterSpillMerger.
Since we're handing back an iterator, am I right in thinking that without some periodic cleanup task you always stand a risk from this kind of leak unless you clear after each comparison or have some kind of async cleanup task?
There was a problem hiding this comment.
I suppose that if you assume that you can only use these sorters once, then you can probably null out the reference to the comparator in the UnsafeExternalSorter once you've constructed the UnsafeSorterSpillMerger using it, and that would also solve the problem I've been seeing (the callback keeping a strong reference).
It'd still feel weird that you have a comparator that could potentially be responsible for the liveness of hundreds of megabytes of memory, though.
There was a problem hiding this comment.
maybe it's worth having a clone method on the comparator and making sure we clone the comparator before passing it to anything?
There was a problem hiding this comment.
The merger reads data from spill files lazily, so when the merging finishes, it's end of the task.
There was a problem hiding this comment.
Then why would we be seeing the OOMs? If at the end of the task the taskcompletionlistener fires and is removed, then the whole comparator becomes unreachable and we have no problem here.
My job looks something like:
dataset.sortWithinPartitions().coalescePartitions() - would not we potentially finish doing some merging before reaching the end of the task?
There was a problem hiding this comment.
i see, some partition may finish merging but some do not, and the merger which is finished is not referred, but it changed the comparator and make it keep the input rows it compared last time.
|
Hi @j-baker , do you have any code snippet to reproduce this memory leak? |
|
This program works well with this PR while it causes an infinite loop in org.apache.spark.memory.TaskMemoryManager.allocatePage without this PR due to OOM. |
Instead of holding a single comparator, keep a factory around which produces objects which may be more tightly scoped. No explicit free work is provided - instead long lived object should not keep around a comparator.
j-baker
force-pushed
the
jbaker/fix_memory_leak
branch
from
1ee2b59 to
51ec9a0
Compare
| final MemoryConsumer consumer, | ||
| final TaskMemoryManager memoryManager, | ||
| final RecordComparator recordComparator, | ||
| final RecordComparator.Factory recordComparatorFactory, |
There was a problem hiding this comment.
Why do we need to change UnsafeInMemorySorter. The TaskContext refers UnsafeExternalSorter, so we only need the comparator factory in UnsafeExternalSorter.
|
BTW, since the GC root is |
|
Test build #79756 has finished for PR 18543 at commit
|
|
Test build #79757 has finished for PR 18543 at commit
|
|
ping |
6 participants
What changes were proposed in this pull request?
UnsafeExternalRowSorter.RowComparator contains references to the objects
backing the last arrays sorted. This causes memory leaks, since those
objects become deallocated but are still live.
We make sure to unset them.
How was this patch tested?
Not sure how to explicitly test this. Would appreciate guidance on what kind of testing might be necessary here. Evidence for the bug is provided in the JIRA ticket here.