Client-side memory leak in map_partitions + persist
- Dominant language
- Python
- Stars
- 1.7k
- Forks
- 778
- Avg merge
- 2h 50m
- Merged PRs (30d)
- 3
Description
I'm fighting some out-of-memory issues on client side in an operation, which acquires (and probably leaks) much more client memory than expected. The problem occurs in a combination of `map_partitions` and `persist`. Complete reproducing example:
```python
from __future__ import division, print_function
import pandas as pd
import numpy as np
import os
import gc
import psutil
import dask.dataframe as dd
from dask import delayed
from dask.distributed import Client
LAST_MEM_USAGE = 0
def log_memory(label):
global LAST_MEM_USAGE
gc.collect()
process = psutil.Process(os.getpid())
mem_usage = process.memory_info().rss / float(2 ** 20)
mem_usage_diff = mem_usage - LAST_MEM_USAGE
print(
"[MemUsage] {:8.1f} MB [{:+8.1f} MB] @ {}".format(
mem_usage, mem_usage_diff, label
)
)
LAST_MEM_USAGE = mem_usage
def run_reproducing_example(dask_endpoint):
client = Client(dask_endpoint)
log_memory("Start")
vector = dd.from_delayed([
delayed(lambda: pd.Series([1]*10000000)) for _ in xrange(5)
], meta=pd.Series(0))
log_memory("Initializing dask.Series")
matrices = [
np.random.uniform(-1, +1, size=(1000, 1000)) for _ in range(50)
]
log_memory("Initializing matrices")
for partition_id, matrix in enumerate(matrices):
def identity_mapper(vector, matrix):
return vector
vector = dd.map_partitions(
identity_mapper, vector, matrix
).persist()
log_memory("After iteration: {}".format(partition_id + 1))
log_memory("End")
# run example on remote cluster to isolate client memory behavior
run_reproducing_example("")
```
Output:
```
[MemUsage] 72.3 MB [ +72.3 MB] @ Start
[MemUsage] 72.4 MB [ +0.1 MB] @ Initializing dask.Series
[MemUsage] 454.1 MB [ +381.7 MB] @ Initializing matrices
[MemUsage] 516.4 MB [ +62.4 MB] @ After iteration: 1
[MemUsage] 585.9 MB [ +69.4 MB] @ After iteration: 2
[MemUsage] 624.0 MB [ +38.2 MB] @ After iteration: 3
[MemUsage] 662.2 MB [ +38.2 MB] @ After iteration: 4
[MemUsage] 700.3 MB [ +38.2 MB] @ After iteration: 5
[MemUsage] 738.5 MB [ +38.1 MB] @ After iteration: 6
[MemUsage] 776.6 MB [ +38.2 MB] @ After iteration: 7
[MemUsage] 814.8 MB [ +38.2 MB] @ After iteration: 8
[MemUsage] 837.6 MB [ +22.8 MB] @ After iteration: 9
[MemUsage] 875.7 MB [ +38.2 MB] @ After iteration: 10
[MemUsage] 913.9 MB [ +38.2 MB] @ After iteration: 11
[MemUsage] 943.9 MB [ +30.1 MB] @ After iteration: 12
[MemUsage] 982.1 MB [ +38.2 MB] @ After iteration: 13
[MemUsage] 1015.0 MB [ +32.9 MB] @ After iteration: 14
[MemUsage] 1028.8 MB [ +13.8 MB] @ After iteration: 15
[MemUsage] 1098.3 MB [ +69.6 MB] @ After iteration: 16
[MemUsage] 1136.5 MB [ +38.2 MB] @ After iteration: 17
[MemUsage] 1174.6 MB [ +38.2 MB] @ After iteration: 18
[MemUsage] 1212.8 MB [ +38.2 MB] @ After iteration: 19
[MemUsage] 1250.9 MB [ +38.1 MB] @ After iteration: 20
[MemUsage] 1289.1 MB [ +38.2 MB] @ After iteration: 21
[MemUsage] 1307.3 MB [ +18.2 MB] @ After iteration: 22
[MemUsage] 1345.5 MB [ +38.2 MB] @ After iteration: 23
[MemUsage] 1383.6 MB [ +38.2 MB] @ After iteration: 24
[MemUsage] 1421.8 MB [ +38.2 MB] @ After iteration: 25
[MemUsage] 1450.2 MB [ +28.5 MB] @ After iteration: 26
[MemUsage] 1483.4 MB [ +33.2 MB] @ After iteration: 27
[MemUsage] 1519.2 MB [ +35.7 MB] @ After iteration: 28
[MemUsage] 1578.0 MB [ +58.8 MB] @ After iteration: 29
[MemUsage] 1623.9 MB [ +45.9 MB] @ After iteration: 30
[MemUsage] 1662.1 MB [ +38.2 MB] @ After iteration: 31
[MemUsage] 1700.3 MB [ +38.2 MB] @ After iteration: 32
[MemUsage] 1738.4 MB [ +38.1 MB] @ After iteration: 33
[MemUsage] 1776.5 MB [ +38.2 MB] @ After iteration: 34
[MemUsage] 1799.3 MB [ +22.8 MB] @ After iteration: 35
[MemUsage] 1837.5 MB [ +38.2 MB] @ After iteration: 36
[MemUsage] 1875.7 MB [ +38.2 MB] @ After iteration: 37
[MemUsage] 1906.2 MB [ +30.5 MB] @ After iteration: 38
[MemUsage] 1944.3 MB [ +38.2 MB] @ After iteration: 39
[MemUsage] 1976.8 MB [ +32.4 MB] @ After iteration: 40
[MemUsage] 2020.8 MB [ +44.0 MB] @ After iteration: 41
[MemUsage] 2081.8 MB [ +61.1 MB] @ After iteration: 42
[MemUsage] 2120.0 MB [ +38.2 MB] @ After iteration: 43
[MemUsage] 2158.2 MB [ +38.2 MB] @ After iteration: 44
[MemUsage] 2196.4 MB [ +38.2 MB] @ After iteration: 45
[MemUsage] 2219.2 MB [ +22.9 MB] @ After iteration: 46
[MemUsage] 2257.4 MB [ +38.2 MB] @ After iteration: 47
[MemUsage] 2295.6 MB [ +38.2 MB] @ After iteration: 48
[MemUsage] 2326.0 MB [ +30.5 MB] @ After iteration: 49
[MemUsage] 2359.9 MB [ +33.9 MB] @ After iteration: 50
```
A guess would be that this combination of `map_partitions` + `persist` makes copies of the `matrix` parameters and there are unreleased references to these copies. Each iteration even seems to hold on to several copies, because each individual matrix is just ~7.6 MB in size, but the amount of leaked memory is ~5 times larger. In general I would not expect the client to keep references to these parameter matrices at all (but I assume that they are stored in the scheduler). Using a single `persist` after the loop without persisting in the loop seems to leak memory in a similar way.
Versions:
- `dask==0.15.2`
- `distributed==1.18.3`
- `tornado==4.5.2`
Contributor guide
Assessment
This issue has not been assessed yet.