Repository: dask/dask

Issue:
Dask.bag with partition_size>1 has potential memory issues
## Abstract
Linear fused parts in dask.bag clean intermediate results too late, leading to potential memory issues.

## Technical Demonstration

### Code

```python
import weakref

import dask.bag as db
import dask.config

import objgraph

# to simplify debugging
dask.config.set(scheduler="synchronous")

# global weakref so we can feed objects into objgraph for analysis
memory = weakref.WeakValueDictionary()


class C:
    def __init__(self, i):
        print(f"init: {i}")
        self.i = i

    def __del__(self):
        print(f"del: {self.i}")


def f1(i):
    print(f"f1: {i}")
    obj = C(i)
    memory[i] = obj
    return obj


def f2(obj):
    print(f"f2: {obj.i}")
    return obj.i


def f3(i):
    print(f"f3: {i}")
    if i in memory:
        objgraph.show_backrefs(memory[i], filename=f'refs{i}.png', shortnames=False, max_depth=20)
    return i


bag = db.from_sequence(range(4), partition_size=2).map(f1).map(f2).map(f3).sum()
bag.visualize(optimize_graph=False, filename="graph.png")
bag.visualize(optimize_graph=True, filename="graph_opt.png")
bag.compute()
```

### Compute Graph
The compute graph *without fusing* looks like this:
![graph](https://user-images.githubusercontent.com/1529400/62218909-a62a1280-b3ad-11e9-81e7-eb70614a9c93.png)

The optimized graph (with linear fusing) looks like this:
![graph_opt](https://user-images.githubusercontent.com/1529400/62218937-b215d480-b3ad-11e9-9102-473fc693e0a8.png)

Notice that the chain `f1->f2->f3` is fused into a single task (along with the per-partition aggregation).

### Actual Output
```
f1: 0
init: 0
f2: 0
f3: 0
f1: 1
init: 1
f2: 1
f3: 1
del: 0
del: 1
f1: 2
init: 2
f2: 2
f3: 2
f1: 3
init: 3
f2: 3
f3: 3
del: 2
del: 3
```

Notice that the payload objects that are created in `f1` and destroyed in `f2` are kept until the whole partition is finished.

### Expected Output
```
f1: 0
init: 0
f2: 0
del: 0
f3: 0
f1: 1
init: 1
f2: 1
del: 1
f3: 1
f1: 2
init: 2
f2: 2
del: 2
f3: 2
f1: 3
init: 3
f2: 3
del: 3
f3: 3
```

## Reference Graphs
`i=0`:
![refs0](https://user-images.githubusercontent.com/1529400/62219468-ad9deb80-b3ae-11e9-9bde-20a86d6b7ecb.png)

`i=1`:
![refs1](https://user-images.githubusercontent.com/1529400/62219478-b2fb3600-b3ae-11e9-87d0-923ff0949efc.png)
Dask.bag with partition_size>1 has potential memory issues
## Abstract
Linear fused parts in dask.bag clean intermediate results too late, leading to potential memory issues.

## Technical Demonstration

### Code

```python
import weakref

import dask.bag as db
import dask.config

import objgraph

# to simplify debugging
dask.config.set(scheduler="synchronous")

# global weakref so we can feed objects into objgraph for analysis
memory = weakref.WeakValueDictionary()


class C:
    def __init__(self, i):
        print(f"init: {i}")
        self.i = i

    def __del__(self):
        print(f"del: {self.i}")


def f1(i):
    print(f"f1: {i}")
    obj = C(i)
    memory[i] = obj
    return obj


def f2(obj):
    print(f"f2: {obj.i}")
    return obj.i


def f3(i):
    print(f"f3: {i}")
    if i in memory:
        objgraph.show_backrefs(memory[i], filename=f'refs{i}.png', shortnames=False, max_depth=20)
    return i


bag = db.from_sequence(range(4), partition_size=2).map(f1).map(f2).map(f3).sum()
bag.visualize(optimize_graph=False, filename="graph.png")
bag.visualize(optimize_graph=True, filename="graph_opt.png")
bag.compute()
```

### Compute Graph
The compute graph *without fusing* looks like this:
![graph](https://user-images.githubusercontent.com/1529400/62218909-a62a1280-b3ad-11e9-81e7-eb70614a9c93.png)

The optimized graph (with linear fusing) looks like this:
![graph_opt](https://user-images.githubusercontent.com/1529400/62218937-b215d480-b3ad-11e9-9102-473fc693e0a8.png)

Notice that the chain `f1->f2->f3` is fused into a single task (along with the per-partition aggregation).

### Actual Output
```
f1: 0
init: 0
f2: 0
f3: 0
f1: 1
init: 1
f2: 1
f3: 1
del: 0
del: 1
f1: 2
init: 2
f2: 2
f3: 2
f1: 3
init: 3
f2: 3
f3: 3
del: 2
del: 3
```

Notice that the payload objects that are created in `f1` and destroyed in `f2` are kept until the whole partition is finished.

### Expected Output
```
f1: 0
init: 0
f2: 0
del: 0
f3: 0
f1: 1
init: 1
f2: 1
del: 1
f3: 1
f1: 2
init: 2
f2: 2
del: 2
f3: 2
f1: 3
init: 3
f2: 3
del: 3
f3: 3
```

## Reference Graphs
`i=0`:
![refs0](https://user-images.githubusercontent.com/1529400/62219468-ad9deb80-b3ae-11e9-9bde-20a86d6b7ecb.png)

`i=1`:
![refs1](https://user-images.githubusercontent.com/1529400/62219478-b2fb3600-b3ae-11e9-87d0-923ff0949efc.png)


Discussion:
FYI: I have a WIP solution that for sure needs some discussion. Will publish a WIP-PR later.
Thank you for the detailed analysis here @crepererum .  I'm glad to know that someone is thinking about this.  I also hadn't seen the `objgraph.show_backrefs` operation before.  That's handy.  I've just used endless `gc.get_referrers` calls before, which is far less helpful ... 

After glancing through the proposed solution I have two initial reactions:

1.  "I can see how this could work" 
2.  "This scares the heck out of me"

It scares me mostly because it's complex enough that I expect that it will eventually interact with something in a bad way, and that our average maintainer won't be able to understand it well enough to quickly figure out what's going on.

> leading to potential memory issues.

Has this issue affected you in a significant way before?  (I assume so, otherwise you probably wouldn't be investigating this).  I wouldn't mind zooming out a bit first to see if there is some other way to resolve the underlying problem that might involve a less technically sophisticated solution.
Also, is this specific to how Dask handles things, or is this possibly a lower level Python problem?
> "This scares the heck out of me"

I can very much understand this. I think I could tackle this with proper documentation. Honestly, I had quite a hard time to understand the iterator chaining in bags and I think documenting both (the iterator concept and the `_Ref` class) could help maintainers as well as developers. So the WIP PR currently only contains very basic documentation.

> Has this issue affected you in a significant way before?

Yes. This is an issue in one of our ETL pipelines. Part of data that we pass through the bags are dataframes as well as models. We have loads of items and expected that using the partition-size, we could reduce the scheduler overhead and the fusing even suggested that this would work properly. Turns out that keeping the intermediate results for too long just blows up our distributed workers.

> Also, is this specific to how Dask handles things, or is this possibly a lower level Python problem?

So the very conservative refcounting of Python is not a Dask problem, but the iterator chains definitely make this very prominent. The issue is basically that all abstraction and their interactions in Python (iterators, function calls, yields) increase the refcount. Stepping back a little bit, here is a more basic example that does not require fancy iterators:

```python
class C:
    def __init__(self, i):
        print(f"init: {i}")
        self.i = i

    def __del__(self):
        print(f"init: {self.i}")


def f1(o):
    print("f1: start")
    o = C(o.i + 1)
    o = C(o.i + 1)
    print("f1: end")
    return o


def main():
    print("main: start")
    f1(C(1))
    print("main: end")


if __name__ == "__main__":
    main()
```

```
main: start
init: 1
f1: start
init: 2
init: 3
init: 2
f1: end
init: 1
init: 3
main: end
```

As you can see, `C(1)` is created in `main` but passing things via an argument increases the refcount and the count is only decreased again if the object is passed back to the parent scope. This late refcount decreasing is also the underlying issue with the yield statements.

In general, libs can workaround these issue using multiple measurements. One is the ref-wrapper in the proposed PR. Another involves sneaky cython tricks. There might be more that I am not aware of.
We might also consider fusing multiple maps into a single map that runs one compound function.

To do this we would reuse the HighLevelGraph and Blockwise functionality that we've adopted for dask array and dataframe.  This is also some scary machinery, but at least it's machinery that we already have to deal with.
> We might also consider fusing multiple maps into a single map that runs one compound function.

I'm not too much concerned about the trivial case where we have linear chains like in the above example. I think the fusing gets really tricky if one of the parameters to the map functions is also a bag or when the user zips to bags.

Thinking about the python refcounting, I think one of the issues here is the edge case around `yield`, where the refcount is only decremened when the control flow returns to the generator (so `yield` seems to behave more like a function call and is therefore kinda reversed to `return`). I would like to try to rewrite `map_chunk` into an iterator where `__iter__` can just return the mapped result, which (to my understanding) would not increase the refcount. That's for sure is slightly more code than the master implementation, but might be really elegant and could also work for zipped bags. 
FYI: I have a WIP solution that for sure needs some discussion. Will publish a WIP-PR later.
Thank you for the detailed analysis here @crepererum .  I'm glad to know that someone is thinking about this.  I also hadn't seen the `objgraph.show_backrefs` operation before.  That's handy.  I've just used endless `gc.get_referrers` calls before, which is far less helpful ... 

After glancing through the proposed solution I have two initial reactions:

1.  "I can see how this could work" 
2.  "This scares the heck out of me"

It scares me mostly because it's complex enough that I expect that it will eventually interact with something in a bad way, and that our average maintainer won't be able to understand it well enough to quickly figure out what's going on.

> leading to potential memory issues.

Has this issue affected you in a significant way before?  (I assume so, otherwise you probably wouldn't be investigating this).  I wouldn't mind zooming out a bit first to see if there is some other way to resolve the underlying problem that might involve a less technically sophisticated solution.
Also, is this specific to how Dask handles things, or is this possibly a lower level Python problem?
> "This scares the heck out of me"

I can very much understand this. I think I could tackle this with proper documentation. Honestly, I had quite a hard time to understand the iterator chaining in bags and I think documenting both (the iterator concept and the `_Ref` class) could help maintainers as well as developers. So the WIP PR currently only contains very basic documentation.

> Has this issue affected you in a significant way before?

Yes. This is an issue in one of our ETL pipelines. Part of data that we pass through the bags are dataframes as well as models. We have loads of items and expected that using the partition-size, we could reduce the scheduler overhead and the fusing even suggested that this would work properly. Turns out that keeping the intermediate results for too long just blows up our distributed workers.

> Also, is this specific to how Dask handles things, or is this possibly a lower level Python problem?

So the very conservative refcounting of Python is not a Dask problem, but the iterator chains definitely make this very prominent. The issue is basically that all abstraction and their interactions in Python (iterators, function calls, yields) increase the refcount. Stepping back a little bit, here is a more basic example that does not require fancy iterators:

```python
class C:
    def __init__(self, i):
        print(f"init: {i}")
        self.i = i

    def __del__(self):
        print(f"init: {self.i}")


def f1(o):
    print("f1: start")
    o = C(o.i + 1)
    o = C(o.i + 1)
    print("f1: end")
    return o


def main():
    print("main: start")
    f1(C(1))
    print("main: end")


if __name__ == "__main__":
    main()
```

```
main: start
init: 1
f1: start
init: 2
init: 3
init: 2
f1: end
init: 1
init: 3
main: end
```

As you can see, `C(1)` is created in `main` but passing things via an argument increases the refcount and the count is only decreased again if the object is passed back to the parent scope. This late refcount decreasing is also the underlying issue with the yield statements.

In general, libs can workaround these issue using multiple measurements. One is the ref-wrapper in the proposed PR. Another involves sneaky cython tricks. There might be more that I am not aware of.
We might also consider fusing multiple maps into a single map that runs one compound function.

To do this we would reuse the HighLevelGraph and Blockwise functionality that we've adopted for dask array and dataframe.  This is also some scary machinery, but at least it's machinery that we already have to deal with.
> We might also consider fusing multiple maps into a single map that runs one compound function.

I'm not too much concerned about the trivial case where we have linear chains like in the above example. I think the fusing gets really tricky if one of the parameters to the map functions is also a bag or when the user zips to bags.

Thinking about the python refcounting, I think one of the issues here is the edge case around `yield`, where the refcount is only decremened when the control flow returns to the generator (so `yield` seems to behave more like a function call and is therefore kinda reversed to `return`). I would like to try to rewrite `map_chunk` into an iterator where `__iter__` can just return the mapped result, which (to my understanding) would not increase the refcount. That's for sure is slightly more code than the master implementation, but might be really elegant and could also work for zipped bags. 

Resolving patch:
diff --git a/dask/bag/core.py b/dask/bag/core.py
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -1993,43 +1993,63 @@ def bag_zip(*bags):
     return Bag(graph, name, npartitions)
 
 
-def map_chunk(f, args, bag_kwargs, kwargs):
+def map_chunk(f, iters, iter_kwarg_keys=None, kwargs=None):
+    """Map ``f`` across one or more iterables, maybe with keyword arguments.
+
+    Low-level function used in ``bag_map``, not user facing.
+
+    Arguments
+    ---------
+    f : callable
+    iters : List[Iterable]
+    iter_kwarg_keys : List[str] or None
+        Keyword names to use for pair with the tail end of ``iters``, allowing
+        keyword arguments to be passed in from iterators.
+    kwargs : dict or None
+        Additional constant keyword arguments to use on every call to ``f``.
+    """
     if kwargs:
         f = partial(f, **kwargs)
+    iters = [iter(a) for a in iters]
+    return _MapChunk(f, iters, kwarg_keys=iter_kwarg_keys)
 
-    args = [iter(a) for a in args]
-    iters = list(args)
-    if bag_kwargs:
-        keys = list(bag_kwargs)
-        kw_val_iters = [iter(v) for v in bag_kwargs.values()]
-        iters.extend(kw_val_iters)
-        kw_iter = (dict(zip(keys, k)) for k in zip(*kw_val_iters))
-        if args:
-            for a, k in zip(zip(*args), kw_iter):
-                yield f(*a, **k)
-        else:
-            for k in kw_iter:
-                yield f(**k)
-    else:
-        for a in zip(*args):
-            yield f(*a)
-
-    # Check that all iterators are fully exhausted
-    if len(iters) > 1:
-        for i in iters:
-            if isinstance(i, itertools.repeat):
-                continue
-            try:
-                next(i)
-            except StopIteration:
-                pass
-            else:
-                msg = (
-                    "map called with multiple bags that aren't identically "
-                    "partitioned. Please ensure that all bag arguments "
-                    "have the same partition lengths"
-                )
-                raise ValueError(msg)
+
+class _MapChunk(Iterator):
+    def __init__(self, f, iters, kwarg_keys=None):
+        self.f = f
+        self.iters = iters
+        self.kwarg_keys = kwarg_keys or ()
+        self.nkws = len(self.kwarg_keys)
+
+    def __next__(self):
+        try:
+            vals = [next(i) for i in self.iters]
+        except StopIteration:
+            self.check_all_iterators_consumed()
+            raise
+
+        if self.nkws:
+            args = vals[: -self.nkws]
+            kwargs = dict(zip(self.kwarg_keys, vals[-self.nkws :]))
+            return self.f(*args, **kwargs)
+        return self.f(*vals)
+
+    def check_all_iterators_consumed(self):
+        if len(self.iters) > 1:
+            for i in self.iters:
+                if isinstance(i, itertools.repeat):
+                    continue
+                try:
+                    next(i)
+                except StopIteration:
+                    pass
+                else:
+                    msg = (
+                        "map called with multiple bags that aren't identically "
+                        "partitioned. Please ensure that all bag arguments "
+                        "have the same partition lengths"
+                    )
+                    raise ValueError(msg)
 
 
 def starmap_chunk(f, x, kwargs):
@@ -2151,21 +2171,21 @@ def bag_map(func, *args, **kwargs):
         raise ValueError("All bags must have the same number of partitions.")
     npartitions = npartitions.pop()
 
-    def build_args(n):
-        return [(a.name, n) if isinstance(a, Bag) else a for a in args2]
+    def build_iters(n):
+        args = [(a.name, n) if isinstance(a, Bag) else a for a in args2]
+        if bag_kwargs:
+            args.extend((b.name, n) for b in bag_kwargs.values())
+        return args
 
-    def build_bag_kwargs(n):
-        if not bag_kwargs:
-            return None
-        return (
-            dict,
-            (zip, list(bag_kwargs), [(b.name, n) for b in bag_kwargs.values()]),
-        )
+    if bag_kwargs:
+        iter_kwarg_keys = list(bag_kwargs)
+    else:
+        iter_kwarg_keys = None
 
     dsk = {
         (name, n): (
             reify,
-            (map_chunk, func, build_args(n), build_bag_kwargs(n), other_kwargs),
+            (map_chunk, func, build_iters(n), iter_kwarg_keys, other_kwargs),
         )
         for n in range(npartitions)
     }


Test patch:
diff --git a/dask/bag/tests/test_bag.py b/dask/bag/tests/test_bag.py
--- a/dask/bag/tests/test_bag.py
+++ b/dask/bag/tests/test_bag.py
@@ -1,11 +1,13 @@
 # coding=utf-8
 from __future__ import absolute_import, division, print_function
 
+import gc
 from itertools import repeat
 import multiprocessing
 import math
 import os
 import random
+import weakref
 
 import partd
 import pytest
@@ -1419,3 +1421,55 @@ def test_map_keynames():
     assert set(b.map(inc).__dask_graph__()) != set(
         b.map_partitions(inc).__dask_graph__()
     )
+
+
+def test_map_releases_element_references_as_soon_as_possible():
+    # Ensure that Bag.map doesn't keep *element* references longer than
+    # necessary. Previous map implementations used ``yield``, which would keep
+    # a reference to the yielded element until the yielded method resumed (this
+    # is just how generator functions work in CPython).
+    #
+    # See https://github.com/dask/dask/issues/5189
+    #
+    # We test 2 variant of potential extra references here:
+    # 1. Within an element of a partition:
+    #    At the time of the second `f_create` for each element, the `C` from
+    #    the first `f_create` should be dropped.
+    # 2. Within a partition:
+    #    When the second item within a partition is processed, `C` from the
+    #    first item should already be dropped.
+    class C:
+        def __init__(self, i):
+            self.i = i
+
+    # keep a weakref to all existing instances of `C`
+    in_memory = weakref.WeakSet()
+
+    def f_create(i):
+        # check that there are no instances of `C` left
+        assert len(in_memory) == 0
+
+        # create new instance
+        o = C(i)
+        in_memory.add(o)
+
+        return o
+
+    def f_drop(o):
+        # o reference dropped on return, should collect
+        return o.i + 100
+
+    b = (
+        db.from_sequence(range(2), npartitions=1)
+        .map(f_create)
+        .map(f_drop)
+        .map(f_create)
+        .map(f_drop)
+        .sum()
+    )
+    try:
+        # Disable gc to ensure refcycles don't matter here
+        gc.disable()
+        b.compute(scheduler="sync")
+    finally:
+        gc.enable()