Nexus dependency node causes huge memory spike on one worker
- Dominant language
- Python
- Stars
- 1.7k
- Forks
- 778
- Avg merge
- 2h 50m
- Merged PRs (30d)
- 3
Description
distributed 2.15.1 git tip
CPython 3.8 x64
Linux
# Use case
I'm using dask.distributed (1 thread per worker) to perform a leaf-to-root aggregation of a huge tree structure of 500k nodes.
Each node in my graph is defined as follows:
- callable
- reference to a global state data node, which is 500kb large and is the same for all 500k nodes
- a substantial number of inlined arguments, worth 1 kb per node
- list of references to children nodes (empty list for leaf nodes)
Example:
```python
{
"node1": (somefunc, "state", <1kb stuff>, ["node2", "node3"]),
"node2": (somefunc, "state", <1kb stuff>, ["node4", "node5"]),
"node3": (somefunc, "state", <1kb stuff>, []),
"node4": (somefunc, "state", <1kb stuff>, []),
"node5": (somefunc, "state", <1kb stuff>, []),
"state": <500 kb stuff>
}
```
# Issue
The existence of the "state" node, that has an edge connecting it directly to all other 500k nodes, causes a huge memory spike - hundreds of megabytes - always exactly on the last worker (as seen in the GUI). I see very briefly a spike of 1.3GB as soon as the computation starts, which goes down after a second or two to 800MB.
After the graph computation is terminated and the client is disconnected, the memory remains allocated. Running pympler on the workers tells me that the vast majority of RAM is not accounted by Python objects, which makes me suspect that the memory may not be deallocated due to fragmentation.
If I remove the global node, the problem disappears.
Scattering the node to a future and then inlining the future does not help.
The size of the spike is crudely proportional to the number of nodes * size of each node (the "1kb stuff" above). The size of the global state does not matter.
# Impact
The worker exhibiting the spike risks being killed off by either the nanny or Kubernetes/other.
A lot of network bandwidth is wasted in the transfer.
# Expected behaviour
No one worker should ever be sent the whole graph.
# POC
In this POC I'm generating 100k nodes with 10kb payload each, for a total dummy payload of ~980 MB. You can play around with the number of nodes and the size of the payload to see how the RAM usage changes.
```python
from os import urandom
from time import sleep
import distributed
from dask import delayed
from dask.delayed import Delayed
from dask.optimization import cull, inline
from dask.utils import ensure_dict
cluster = distributed.LocalCluster(8, threads_per_worker=1)
client = distributed.Client(cluster)
state = delayed(None)
# state = None # Uncomment to make the issue disappear
@delayed(pure=False)
def f(state, payload, children):
sleep(.01)
def build_graph(n):
if n == 1:
children = []
else:
children = [build_graph(n // 2 + n % 2), build_graph(n // 2)]
return f(state, urandom(10240), children)
root = build_graph(50000) # half of the total nodes
future = root.persist()
future.result()
```
### RAM measures:
Total payload size: (100k nodes * 10 kiB) = 980 MiB
| Stage | Scheduler | Worker 1-7 | Worker 8 |
| - | - | - | - |
| Idle, empty cluster | 65 MB | 92 MB | 92 MB |
| Brief spike immediately after persist() | 1700 MB | 240~360 MB | 1300 MB |
| During computation | 1700 MB | 240~360 MB | 820 MB |
| Idle, holding future to the result | 1700 MB | 240~360 MB | 820 MB |
| After closing the client | 1700 MB | 240~360 MB | 820 MB |
### Pympler analysis
After releasing all keys:
```python
import gc
import psutil
from pympler import muppy, summary
def memory_summary():
"""Return:
- Total RAM used by Python objects, in MB
- Total process RAM, in MB
"""
gc.collect()
total_ram = psutil.Process().memory_info().rss
all_objects = muppy.get_objects(include_frames=True)
s = summary.summarize(all_objects)
python_objects = sum(row[2] for row in s)
return python_objects / 2**20, total_ram / 2**20
client.run(memory_summary)
{'tcp://127.0.0.1:33669': (102.12534713745117, 824.49609375),
'tcp://127.0.0.1:33735': (79.34993076324463, 354.8828125),
'tcp://127.0.0.1:35381': (80.06385040283203, 339.31640625),
'tcp://127.0.0.1:37751': (79.50504875183105, 238.703125),
'tcp://127.0.0.1:38299': (80.73084831237793, 315.42578125),
'tcp://127.0.0.1:41453': (80.35994625091553, 306.91015625),
'tcp://127.0.0.1:44371': (80.13745212554932, 296.57421875),
'tcp://127.0.0.1:46077': (80.209303855896, 240.33203125)}
client.run_on_scheduler(memory_summary)
(275.25281620025635, 1716.72265625)
```
Contributor guide
Assessment
This issue has not been assessed yet.