[BUG] Failing to find dependencies in tuple of futures

[rjzamora]

What happened: After #4139 there seem to be a number of broken tests in custreamz CI. 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, but fails on master:

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

[jrbourbeau]

Thanks for the nice reproducer @rjzamora! I haven't had a chance to look into the details, but #4139 definitely introduced this regression

[rjzamora]

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.

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 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, which resolves the issue, but requires changes in both dask and distributed.