Redesign discussion: Launch Tasks from Tasks
- Dominant language
- Python
- Stars
- 1.7k
- Forks
- 778
- Avg merge
- 2h 50m
- Merged PRs (30d)
- 3
Description
Dask performs best when the whole graph of tasks is defined ahead of time from the client and submitted at once. This is however not always possible.
# Use case and current best practice
A typical example is a top-down discovery + bottom-up aggregation of a tree where the discovery of the parent-child relationships is an operation too expensive to be performed on the client.
Use case in pure Python:
```python
def get_children(node):
"""Expensive operation that discovers the direct children of a node
"""
...
def aggregate(node, children_outputs):
"""Expensive operation that calculates the value of a node based on its own
properties plus the output of the aggregate function for each of its children
"""
...
def crawl(node):
children = get_children(node)
children_outputs = [crawl(child) for child in children] # top-down recursion
return aggregate(node, children_outputs) # bottom-up aggregation
out = crawl(root)
```
The first way to solve this problem today with Dask is to have the client invoke ``client.submit(get_children, node)`` for every node and wait for results. This can be very network and CPU intensive for the client.
The second approach is to use secede/rejoin, as described in https://distributed.dask.org/en/latest/task-launch.html:
```python
import distributed
def crawl(node):
children = get_children(node)
client = distributed.get_client()
children_futures = client.map(crawl, children)
distributed.secede()
children_outputs = client.gather(children_futures)
distributed.rejoin()
return aggregate(node, children_outputs)
client = distributed.Client(...)
out = client.submit(crawl, root).result()
```
The above is problematic, because:
1. there will be an uncontrolled increase in unmanaged memory on the workers, caused by all the local variables of the seceded task plus the accessory data needed to track the seceded tasks
2. each seceded task adds a thread to the worker, which adds a burden to the Linux kernel. If the graph has enough nodes, the cluster will eventually die as the workers hit their ulimit.
3. Because of point 2, it works very poorly for very large quantities of individually small nodes. You may redesign the algorithm to have a single crawl function go through a fixed-size cluster of contiguous nodes; such a change however is algorithmically complex.
and last but not least,
4. If *any* worker dies during the computation, you have to restart from scratch.
A situationally slightly better variant is as follows:
```python
import distributed
def crawl(node):
children = get_children(node)
client = distributed.get_client()
children_futures = client.map(crawl, children)
out_future = client.submit(aggregate, node, children_futures)
distributed.secede()
return out_future.result()
client = distributed.Client(...)
out = client.submit(crawl, root).result()
```
The difference is subtle - as the subgraph of each child gets resolved, its (potentially large) output does not get stored in the stack of client.aggregate (which is unmanaged memory), but it goes into the managed memory instead with all the benefits of the case. On the flip side, the scheduler is now burdened with two futures per node instead of one. Regardless, all of the problems listed above remain.
# Proposed redesign
I would like to suggest deprecating secede()/rejoin().
In its place, I would like to introduce the following rule:
**If a task returns a Future, then the scheduler will wait for it and return its result instead. This may be nested (the result of the Future may itself be a Future).**
The use case code becomes as follows:
```python
import distributed
def crawl(node):
children = get_children(node)
client = distributed.get_client()
children_futures = client.map(crawl, children)
return client.submit(aggregate, node, children_futures)
client = distributed.Client(...)
out = client.submit(crawl, root).result()
```
No extra threads are ever created. Everything is managed by the scheduler - as it should. The network and CPU load on the user's client (e.g. a jupyter notebook) remain trivial.
Nested resolution of futures aside, the above code currently does not work because, after you return a future, as soon as the future is serialised and removed from Worker.data the future destructor kicks in, which in turn releases the refcount on the scheduler, so by the time the future is rebuilt on the opposite side the data it references may have been lost. Same if the future is spilled to disk.
# Challenges
### Publish/unpublish
It is currently possible to work around the scheduler forgetting the future upon return by publishing it temporarily. This is generally a bad idea because, short of implementing a user-defined garbage collector, you may end up with cluster-wide memory leaks of managed memory (datasets that are published and then forgotten, because the task that was supposed to unpublish them crashed or never started). Nonetheless, automatically resolving returned futures will break this pattern.
#### Workaround
Users can still use this hack but return the name of the temporary dataset instead of the Future.
# Additions and nice-to-haves
### Client-side tracking
It would be nice to see ``distributed.diagnostics.progressbar.progress`` display the increasing tasks in real time. This is not something that's happening with the current secede/rejoin design either.
### Collections
Returned dask collections could be treated specially like Futures. For example, the below would halve the number of worker->scheduler comms and (personal preference) would also look nicer:
```python
import distributed
from dask import delayed
def get_children(node):
...
@delayed
def aggregate(node, children_outputs):
...
@delayed
def crawl(node):
children = get_children(node)
children_delayeds = [crawl(child) for child in children]
return aggregate(node, children_delayeds)
client = distributed.Client(...)
out = crawl(root).compute()
```
Under the hood, all it's happening is a two-liner that converts the collection into a future to revert to the base use case:
```python
if is_dask_collection(retvalue):
retvalue = get_client().compute(retvalue)
```
The same should be implemented in dask/dask, so that it works on the threading/multiprocessing schedulers too.
Contributor guide
Assessment
This issue has not been assessed yet.