Repository: dask/distributed

Issue:
[BUG] Failing to find dependencies in tuple of futures
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->


**What happened**:  After #4139  there seem to be a number of broken tests in [custreamz CI](https://gpuci.gpuopenanalytics.com/view/RAPIDS%20-%20GPU%20Matrix/job/rapidsai/job/gpuci/job/cudf/job/branches/job/cudf-gpu-matrix-branch-0.16/CUDA=10.1,LINUX_VERSION=ubuntu16.04,PYTHON_VERSION=3.7/49/testReport/custreamz.tests/test_dataframes/test_getitem_dask_/).  The failures seem to be related to the way streamz/custreamz is using futures.  More specifically, the task graph contains tuples with futures objects that the new `get_all_dependencies` method is failing to search for dependencies.

**What you expected to happen**:
I would expect distributed to handle a tuple of futures as an argument.

**Minimal Complete Verifiable Example**:

For example, the following seems to work with [3e5b506](https://github.com/dask/distributed/commit/3e5b506ed94dea47c2ce565fca91709d964a8716), but fails on master:

```python
import pandas as pd
from distributed import LocalCluster, Client

def _create(size):
    return pd.DataFrame({"a": range(size)})

client = Client(LocalCluster(n_workers=1))

x = client.submit(_create, 5)
df = client.submit(pd.concat, (x,))  # Note the tuple argument
df.result()
```
**Traceback**:
```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-1-0bc1dd2ca416> in <module>
      9 x = client.submit(_create, 5)
     10 df = client.submit(pd.concat, (x,))
---> 11 df.result()

~/workspace/cudf-0.16/distributed/distributed/client.py in result(self, timeout)
    224         if self.status == "error":
    225             typ, exc, tb = result
--> 226             raise exc.with_traceback(tb)
    227         elif self.status == "cancelled":
    228             raise result

/datasets/rzamora/miniconda3/envs/cudf_16/lib/python3.7/site-packages/pandas/core/reshape/concat.py in concat()
    282         verify_integrity=verify_integrity,
    283         copy=copy,
--> 284         sort=sort,
    285     )
    286 

/datasets/rzamora/miniconda3/envs/cudf_16/lib/python3.7/site-packages/pandas/core/reshape/concat.py in __init__()
    357                     "only Series and DataFrame objs are valid"
    358                 )
--> 359                 raise TypeError(msg)
    360 
    361             # consolidate

TypeError: cannot concatenate object of type '<class 'str'>'; only Series and DataFrame objs are valid
```
Note that the code works fine when a `list` of futures is used in place of a `tuple`: `df = client.submit(pd.concat, [x])`

cc @madsbk 


Discussion:
Thanks for the nice reproducer @rjzamora! I haven't had a chance to look into the details, but https://github.com/dask/distributed/pull/4139 definitely introduced this regression
Since we are no longer operating on a raw dict in `_graph_to_futures`, we are now relying on `get_all_dependencies` to gather the required dependency structure.  For `BasicLayer`, this is ultimately relying on `keys_in_tasks`, which only searches for keys in tuples [if the first element is callable](https://github.com/dask/dask/blob/46ef300d4174983e769702a38a51c6884a8765b5/dask/core.py#L178).

It probably doesn't make sense to modify `keys_in_tasks`, but we do need a way to gather dependencies of the HLG.  One possible solution is to add a `key=bool` option to [`map_tasks`](https://github.com/dask/dask/blob/46ef300d4174983e769702a38a51c6884a8765b5/dask/highlevelgraph.py#L97) to indicate that the supplied function will accept both a key and task (rather than just a task).  I have confirmed that this allows us to compile a list of dependencies in [unpack_remote_data](https://github.com/dask/distributed/blob/dd28f7af1446e71091e079f7fdceacf489899911/distributed/client.py#L2605), which resolves the issue, but requires changes in both dask and distributed.

Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -2607,7 +2607,21 @@ def unpack_remote_data(tasks):
                 unpacked_futures.update(futures)
                 return tasks
 
-            dsk = dsk.map_tasks(unpack_remote_data)
+            # HACK: Need to loop through each layer manually to collect
+            # futures-related dependencies (for now).  This should be removed
+            # once the scheduler doesn't need these dependencies.
+            _dsk = {}
+            future_deps = defaultdict(set)
+            for key, layer in dsk.layers.items():
+                _dsk[key] = layer.map_tasks(unpack_remote_data)
+                # Here is the "manual" loop where the full graph
+                # needs to be materialized.
+                if unpacked_futures:
+                    for k, v in dict(layer).items():
+                        futures = unpack_remotedata(v)[1]
+                        future_deps[k].update({f.key for f in futures})
+            if _dsk:
+                dsk = HighLevelGraph(_dsk, dsk.dependencies, dsk.key_dependencies)
 
             for future in unpacked_futures:
                 if future.client is not self:
@@ -2618,12 +2632,14 @@ def unpack_remote_data(tasks):
 
             # HACK: currently when submitting work to the scheduler, the client need to
             # send all key dependencies along with the task graph. Since `dsk` doesn't
-            # know about the unpacked futures, we add them manually to `dsk` before
-            # calculating dependencies. This hack shouldn't be necessary when the
-            # scheduler accepts high level graphs.
+            # know about the unpacked futures, we add futures-related dependencies from
+            # `future_deps`. This hack shouldn't be necessary when the scheduler accepts
+            # high level graphs.
             dsk.keyset()
-            dsk._keys.update({f.key for f in unpacked_futures})
             dependencies = dsk.get_all_dependencies()
+            if future_deps:
+                for key in dependencies.keys():
+                    dependencies[key] |= future_deps.get(key, set())
 
             if priority is None:
                 # Removing all unpacked futures before calling order()


Test patch:
diff --git a/distributed/tests/test_collections.py b/distributed/tests/test_collections.py
--- a/distributed/tests/test_collections.py
+++ b/distributed/tests/test_collections.py
@@ -210,3 +210,19 @@ async def test_delayed_none(c, s, w):
     [xx, yy] = c.compute([x, y])
     assert await xx is None
     assert await yy == 123
+
+
+@pytest.mark.parametrize("typ", [tuple, list])
+def test_tuple_futures_arg(client, typ):
+    x = client.submit(
+        make_time_dataframe,
+    )
+    df2 = client.submit(
+        pd.concat,
+        typ(
+            [
+                x,
+            ]
+        ),
+    )
+    dd.assert_eq(df2.result().iloc[:0], make_time_dataframe().iloc[:0])