[BUG] DASK_JIT_UNSPILL=True causes failures

[ChrisJar]

Some operations that use dask-cuda with DASK_JIT_UNSPILL=True fail

For example:

import cudf
import numpy as np

import dask_cudf
import dask_cuda
from distributed import Client

cluster = dask_cuda.LocalCUDACluster(jit_unspill=True)
client = Client(cluster)

df1 = cudf.DataFrame({"id":np.random.randint(0,100,100000000), "val":np.random.randint(0,100000,100000000)})
df2 = cudf.DataFrame({"id":np.random.randint(0,100,100000000), "val":np.random.randint(0,100000,100000000)})

ddf1 = dask_cudf.from_cudf(df1, npartitions=8)
ddf2 = dask_cudf.from_cudf(df2, npartitions=8)

ddf1.merge(ddf2, on="id", how="inner").compute()

returns

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/nfs/cjarrett/anaconda3/envs/rapids-gpu-bdb/lib/python3.8/site-packages/dask/base.py", line 285, in compute
    (result,) = compute(self, traverse=False, **kwargs)
  File "/home/nfs/cjarrett/anaconda3/envs/rapids-gpu-bdb/lib/python3.8/site-packages/dask/base.py", line 567, in compute
    results = schedule(dsk, keys, **kwargs)
  File "/home/nfs/cjarrett/anaconda3/envs/rapids-gpu-bdb/lib/python3.8/site-packages/distributed/client.py", line 2705, in get
    results = self.gather(packed, asynchronous=asynchronous, direct=direct)
  File "/home/nfs/cjarrett/anaconda3/envs/rapids-gpu-bdb/lib/python3.8/site-packages/distributed/client.py", line 2014, in gather
    return self.sync(
  File "/home/nfs/cjarrett/anaconda3/envs/rapids-gpu-bdb/lib/python3.8/site-packages/distributed/client.py", line 855, in sync
    return sync(
  File "/home/nfs/cjarrett/anaconda3/envs/rapids-gpu-bdb/lib/python3.8/site-packages/distributed/utils.py", line 338, in sync
    raise exc.with_traceback(tb)
  File "/home/nfs/cjarrett/anaconda3/envs/rapids-gpu-bdb/lib/python3.8/site-packages/distributed/utils.py", line 321, in f
    result[0] = yield future
  File "/home/nfs/cjarrett/anaconda3/envs/rapids-gpu-bdb/lib/python3.8/site-packages/tornado/gen.py", line 762, in run
    value = future.result()
  File "/home/nfs/cjarrett/anaconda3/envs/rapids-gpu-bdb/lib/python3.8/site-packages/distributed/client.py", line 1879, in _gather
    raise exception.with_traceback(traceback)
  File "/home/nfs/cjarrett/anaconda3/envs/rapids-gpu-bdb/lib/python3.8/site-packages/dask_cuda/proxify_device_objects.py", line 158, in wrapper
    ret = func(*args, **kwargs)
  File "/home/nfs/cjarrett/anaconda3/envs/rapids-gpu-bdb/lib/python3.8/site-packages/dask/dataframe/core.py", line 107, in _concat
    else methods.concat(args2, uniform=True, ignore_index=ignore_index)
  File "/home/nfs/cjarrett/anaconda3/envs/rapids-gpu-bdb/lib/python3.8/site-packages/dask/dataframe/dispatch.py", line 61, in concat
    return func(
  File "/home/nfs/cjarrett/anaconda3/envs/rapids-gpu-bdb/lib/python3.8/site-packages/dask_cuda/proxy_object.py", line 730, in wrapper
    args = [unproxy(d) for d in args]
  File "/home/nfs/cjarrett/anaconda3/envs/rapids-gpu-bdb/lib/python3.8/site-packages/dask_cuda/proxy_object.py", line 730, in <listcomp>
    args = [unproxy(d) for d in args]
  File "/home/nfs/cjarrett/anaconda3/envs/rapids-gpu-bdb/lib/python3.8/site-packages/dask_cuda/proxy_object.py", line 111, in unproxy
    return type(obj)(unproxy(o) for o in obj)
  File "/home/nfs/cjarrett/anaconda3/envs/rapids-gpu-bdb/lib/python3.8/site-packages/dask_cuda/proxy_object.py", line 111, in <genexpr>
    return type(obj)(unproxy(o) for o in obj)
  File "/home/nfs/cjarrett/anaconda3/envs/rapids-gpu-bdb/lib/python3.8/site-packages/dask_cuda/proxy_object.py", line 108, in unproxy
    obj = obj._obj_pxy_deserialize()
  File "/home/nfs/cjarrett/anaconda3/envs/rapids-gpu-bdb/lib/python3.8/site-packages/dask_cuda/proxy_object.py", line 344, in _obj_pxy_deserialize
    hostfile.proxies_tally.unspill_proxy(external)
  File "/home/nfs/cjarrett/anaconda3/envs/rapids-gpu-bdb/lib/python3.8/site-packages/dask_cuda/proxify_host_file.py", line 108, in unspill_proxy
    self.unspilled_proxies.add(proxy)
  File "/home/nfs/cjarrett/anaconda3/envs/rapids-gpu-bdb/lib/python3.8/site-packages/dask_cuda/proxify_host_file.py", line 34, in add
    for dev_mem in proxy._obj_pxy_get_device_memory_objects():
  File "/home/nfs/cjarrett/anaconda3/envs/rapids-gpu-bdb/lib/python3.8/site-packages/dask_cuda/proxy_object.py", line 124, in wrapper2
    ret = func(self)
  File "/home/nfs/cjarrett/anaconda3/envs/rapids-gpu-bdb/lib/python3.8/site-packages/dask_cuda/proxy_object.py", line 372, in _obj_pxy_get_device_memory_objects
    return get_device_memory_objects(self._obj_pxy["obj"])
  File "/home/nfs/cjarrett/anaconda3/envs/rapids-gpu-bdb/lib/python3.8/site-packages/dask_cuda/get_device_memory_objects.py", line 27, in get_device_memory_objects
    return set(dispatch(obj))
TypeError: unhashable type: 'Int64Index'

Environment:

cudf                      21.08.00a210622 cuda_11.0_py38_gaa67b138df_201    rapidsai-nightly
dask                      2021.6.1           pyhd8ed1ab_0    conda-forge
dask-core                 2021.6.1           pyhd8ed1ab_0    conda-forge
dask-cuda                 21.08.00a210622         py38_27    rapidsai-nightly
dask-cudf                 21.08.00a210622 py38_gaa67b138df_201    rapidsai-nightly
distributed               2021.6.1         py38h578d9bd_0    conda-forge

[jakirkham]

Did you figure it out or should this still be open?

[madsbk]

#658 should fix the issue