Repository: dask/distributed

Issue:
Race condition between Client.normalize_collection and Future.release
It seems there's a race condition between futures being cleaned up, and futures being reused by a collection:


```python
from __future__ import absolute_import, division

from dask import array as da
from dask.distributed import Client, wait


if __name__ == '__main__':
    with Client() as client:
        # Repeat until failure
        while True:
            x = da.arange(2**20, chunks=2**10)
            y = x.persist()
            wait(y)
            # Start cleaning up futures...
            del y
            # Find existing futures and reuse them, *race condition here*
            normalized = client.normalize_collection(x)
            normalized.compute()
```

Result:

```
$ python test.py
Traceback (most recent call last):
  File "test.py", line 18, in <module>
    normalized.compute()
  File "/Users/jcrist/Code/dask/dask/base.py", line 156, in compute
    (result,) = compute(self, traverse=False, **kwargs)
  File "/Users/jcrist/Code/dask/dask/base.py", line 398, in compute
    return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
  File "/Users/jcrist/Code/dask/dask/base.py", line 398, in <listcomp>
    return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
  File "/Users/jcrist/Code/dask/dask/array/core.py", line 904, in finalize
    return concatenate3(results)
  File "/Users/jcrist/Code/dask/dask/array/core.py", line 3637, in concatenate3
    result[idx] = arr
ValueError: invalid literal for int() with base 10: "('arange-96617f71e88551ee72155f5966b630aa', 45)"
```

It looks like futures are found to still exist in `normalize_collection`, but are fully released before they can be reused by the `.compute()` call (resulting in keys being treated as raw strings instead).


Discussion:
Hrm, indeed I can see how that would be a problem.  Thank you for tracking this down.  

This is slightly tricky to resolve because garbage collection always happens on another thread, so it's hard to guarantee consistency.  

It looks like `normalize_collection` calls `_optimize_insert_futures` which does create a new `Future` object that correctly calls `Client._inc_ref`.  This all happens immediately (there are no yields or `loop.add_callbacks` anywhere) so my guess is that if `normalize_collection` happens before the `gc.collect` then things are fine.  

In the other direction, it looks like `Future.__del__` calls `Future.release` calls `loop.add_callback(Client._dec_ref, ...)` which calls `Client._release_key` which sends the `client-releases-keys` message to the scheduler asynchronously using the `Client._send_to_scheduler` method.   There are several gaps here where things could go badly, including the gap while this information is being transmitted across the wire to the scheduler.

(sorry for typing out my thought process above, I do recommend tracing through those just for personal education)

After thinking about this for 30 minutes I can't yet find a situation where we're going to tell the scheduler to clear out the task, but also think that the future is valid.  To track this down I would probably try getting more telemetry about exactly what happens to that task and when.  This might be done through the structured logs with `scheduler.story(key)`, though this only gives information on the scheduler's side, or instead one could just insert a bunch of print statements.  Hopefully after getting more information it becomes more obvious about what sequence of events is triggering the bad outcome.
The following naive patch fixes the issue, but may not be the best method. It's not clear to me what the difference between `_lock` and `_refcount_lock` are.

```diff
diff --git a/distributed/client.py b/distributed/client.py
index 04e0e4c..251aae9 100644
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -544,7 +544,7 @@ class Client(Node):
         self.scheduler = None
         self._scheduler_identity = {}
         self._lock = threading.Lock()
-        self._refcount_lock = threading.Lock()
+        self._refcount_lock = threading.RLock()
         self.datasets = Datasets(self)
         self._serializers = serializers
         if deserializers is None:
@@ -2371,14 +2371,15 @@ class Client(Node):
         Client.persist: trigger computation of collection's tasks
         """
         with self._lock:
-            dsk = self._optimize_insert_futures(
-                    collection.__dask_graph__(),
-                    collection.__dask_keys__())
+            with self._refcount_lock:
+                dsk = self._optimize_insert_futures(
+                        collection.__dask_graph__(),
+                        collection.__dask_keys__())
 
-            if dsk is collection.__dask_graph__():
-                return collection
-            else:
-                return redict_collection(collection, dsk)
+                if dsk is collection.__dask_graph__():
+                    return collection
+                else:
+                    return redict_collection(collection, dsk)
 
     def compute(self, collections, sync=False, optimize_graph=True,
                 workers=None, allow_other_workers=False, resources=None,
```
I recommend putting it into a PR and seeing if it passed on CI

On Fri, Nov 30, 2018 at 8:11 AM Jim Crist <notifications@github.com> wrote:

> The following naive patch fixes the issue, but may not be the best method.
> It's not clear to me what the difference between _lock and _refcount_lock
> are.
>
> diff --git a/distributed/client.py b/distributed/client.py
> index 04e0e4c..251aae9 100644--- a/distributed/client.py+++ b/distributed/client.py@@ -544,7 +544,7 @@ class Client(Node):
>          self.scheduler = None
>          self._scheduler_identity = {}
>          self._lock = threading.Lock()-        self._refcount_lock = threading.Lock()+        self._refcount_lock = threading.RLock()
>          self.datasets = Datasets(self)
>          self._serializers = serializers
>          if deserializers is None:@@ -2371,14 +2371,15 @@ class Client(Node):
>          Client.persist: trigger computation of collection's tasks
>          """
>          with self._lock:-            dsk = self._optimize_insert_futures(-                    collection.__dask_graph__(),-                    collection.__dask_keys__())+            with self._refcount_lock:+                dsk = self._optimize_insert_futures(+                        collection.__dask_graph__(),+                        collection.__dask_keys__())
>  -            if dsk is collection.__dask_graph__():-                return collection-            else:-                return redict_collection(collection, dsk)+                if dsk is collection.__dask_graph__():+                    return collection+                else:+                    return redict_collection(collection, dsk)
>
>      def compute(self, collections, sync=False, optimize_graph=True,
>                  workers=None, allow_other_workers=False, resources=None,
>
> —
> You are receiving this because you commented.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/distributed/issues/2385#issuecomment-443253298>,
> or mute the thread
> <https://github.com/notifications/unsubscribe-auth/AASszDO68NwYo53nYrygHRbUTSBQ68OPks5u0ViegaJpZM4Y73pw>
> .
>

I don't fully remember the intention behind the lock and refcount_lock.  I
also wasn't aware of RLock (seems handy).  It's entirely possible that we
might be able to merge them safely (recall that most of this codebase has
one author, and so there are likely many opportunities for improvement).

On Fri, Nov 30, 2018 at 8:13 AM Matthew Rocklin <mrocklin@gmail.com> wrote:

> I recommend putting it into a PR and seeing if it passed on CI
>
> On Fri, Nov 30, 2018 at 8:11 AM Jim Crist <notifications@github.com>
> wrote:
>
>> The following naive patch fixes the issue, but may not be the best
>> method. It's not clear to me what the difference between _lock and
>> _refcount_lock are.
>>
>> diff --git a/distributed/client.py b/distributed/client.py
>> index 04e0e4c..251aae9 100644--- a/distributed/client.py+++ b/distributed/client.py@@ -544,7 +544,7 @@ class Client(Node):
>>          self.scheduler = None
>>          self._scheduler_identity = {}
>>          self._lock = threading.Lock()-        self._refcount_lock = threading.Lock()+        self._refcount_lock = threading.RLock()
>>          self.datasets = Datasets(self)
>>          self._serializers = serializers
>>          if deserializers is None:@@ -2371,14 +2371,15 @@ class Client(Node):
>>          Client.persist: trigger computation of collection's tasks
>>          """
>>          with self._lock:-            dsk = self._optimize_insert_futures(-                    collection.__dask_graph__(),-                    collection.__dask_keys__())+            with self._refcount_lock:+                dsk = self._optimize_insert_futures(+                        collection.__dask_graph__(),+                        collection.__dask_keys__())
>>  -            if dsk is collection.__dask_graph__():-                return collection-            else:-                return redict_collection(collection, dsk)+                if dsk is collection.__dask_graph__():+                    return collection+                else:+                    return redict_collection(collection, dsk)
>>
>>      def compute(self, collections, sync=False, optimize_graph=True,
>>                  workers=None, allow_other_workers=False, resources=None,
>>
>> —
>> You are receiving this because you commented.
>> Reply to this email directly, view it on GitHub
>> <https://github.com/dask/distributed/issues/2385#issuecomment-443253298>,
>> or mute the thread
>> <https://github.com/notifications/unsubscribe-auth/AASszDO68NwYo53nYrygHRbUTSBQ68OPks5u0ViegaJpZM4Y73pw>
>> .
>>
>


Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -543,8 +543,11 @@ def __init__(self, address=None, loop=None, timeout=no_default,
         self.cluster = None
         self.scheduler = None
         self._scheduler_identity = {}
-        self._lock = threading.Lock()
-        self._refcount_lock = threading.Lock()
+        # A reentrant-lock on the refcounts for futures associated with this
+        # client. Should be held by individual operations modifying refcounts,
+        # or any bulk operation that needs to ensure the set of futures doesn't
+        # change during operation.
+        self._refcount_lock = threading.RLock()
         self.datasets = Datasets(self)
         self._serializers = serializers
         if deserializers is None:
@@ -1249,7 +1252,7 @@ def submit(self, func, *args, **kwargs):
 
         skey = tokey(key)
 
-        with self._lock:
+        with self._refcount_lock:
             if skey in self.futures:
                 return Future(key, self, inform=False)
 
@@ -2184,7 +2187,7 @@ def _graph_to_futures(self, dsk, keys, restrictions=None,
                           loose_restrictions=None, priority=None,
                           user_priority=0, resources=None, retries=None,
                           fifo_timeout=0, actors=None):
-        with self._lock:
+        with self._refcount_lock:
             if resources:
                 resources = self._expand_resources(resources,
                                                    all_keys=itertools.chain(dsk, keys))
@@ -2334,13 +2337,14 @@ def _optimize_insert_futures(self, dsk, keys):
         This returns the same graph if unchanged but a new graph if any changes
         were necessary.
         """
-        changed = False
-        for key in list(dsk):
-            if tokey(key) in self.futures:
-                if not changed:
-                    changed = True
-                    dsk = dict(dsk)
-                dsk[key] = Future(key, self, inform=False)
+        with self._refcount_lock:
+            changed = False
+            for key in list(dsk):
+                if tokey(key) in self.futures:
+                    if not changed:
+                        changed = True
+                        dsk = dict(dsk)
+                    dsk[key] = Future(key, self, inform=False)
 
         if changed:
             dsk, _ = dask.optimization.cull(dsk, keys)
@@ -2370,15 +2374,13 @@ def normalize_collection(self, collection):
         --------
         Client.persist: trigger computation of collection's tasks
         """
-        with self._lock:
-            dsk = self._optimize_insert_futures(
-                    collection.__dask_graph__(),
-                    collection.__dask_keys__())
+        dsk_orig = collection.__dask_graph__()
+        dsk = self._optimize_insert_futures(dsk_orig, collection.__dask_keys__())
 
-            if dsk is collection.__dask_graph__():
-                return collection
-            else:
-                return redict_collection(collection, dsk)
+        if dsk is dsk_orig:
+            return collection
+        else:
+            return redict_collection(collection, dsk)
 
     def compute(self, collections, sync=False, optimize_graph=True,
                 workers=None, allow_other_workers=False, resources=None,


Test patch:
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
--- a/distributed/tests/test_client.py
+++ b/distributed/tests/test_client.py
@@ -4118,6 +4118,24 @@ def test_normalize_collection_dask_array(c, s, a, b):
     assert result1 == result2
 
 
+@slow
+def test_normalize_collection_with_released_futures(c):
+    da = pytest.importorskip('dask.array')
+
+    x = da.arange(2**20, chunks=2**10)
+    y = x.persist()
+    wait(y)
+    sol = y.sum().compute()
+    # Start releasing futures
+    del y
+    # Try to reuse futures. Previously this was a race condition,
+    # and the call to `.compute()` would error out due to missing
+    # futures on the scheduler at compute time.
+    normalized = c.normalize_collection(x)
+    res = normalized.sum().compute()
+    assert res == sol
+
+
 @gen_cluster(client=True)
 def test_auto_normalize_collection(c, s, a, b):
     da = pytest.importorskip('dask.array')