Repository: rapidsai/dask-cuda
Issue category: Performance Issue

Issue:
Spill quickly from device to host
I'm curious how fast we are at moving data back and forth between device and host memory.  I suspect that currently we end up serializing with standard Dask serialization, which may not be aware of device memory.  We should take a look at using the `"cuda"` serialization family within Dask, that was added as part of the UCX work.   


```python
from distributed.protocol import serialize, deserialize
x = cupy.arange(100)
header, frames = serialize(x, serializers=("cuda", "dask", "pickle"))
y = deserialize(header, frames, deserializers=("cuda", "dask", "pickle",    "error"))
```

(more here https://github.com/dask/distributed/tree/master/distributed/protocol/tests)

And then we should probably do a bit of profiling.  I would focus this on just the `DeviceHostDisk` class, rather than engaging the full dask.distributed cluster machinery.

This is relevant for a couple of applications that people care about today

1.  Larger than memory image processing (a bit)
2.  Out-of-core sorts and joins (cc @randerzander @kkraus14) (more relevant)

@pentschev , you seem like the obvious person to work on this if you have some time.  I think it'd be useful to know how fast we are today, then look at serialization with the new "cuda" Dask serializer, and then see how fast we can get things relative to theoretical hardware speeds.


Hints:
Yes, I haven't done any profiling, and I'm sure this is something that could be improved. Thanks for pointing me to the CUDA serialization, I wasn't aware of that specialization.

If the is not of utmost urgency, I can take a look at this next week, this week I'm mostly busy with EuroPython presentation and at the conference.
Alrighty, I think that we can probably wait until next week (although @randerzander please jump in if not).  Enjoy the conference !
I did some quick benchmarking:

```python
In [1]: import cupy

In [2]: from distributed.protocol import serialize, deserialize

In [3]: x = cupy.arange(1.25e8)

In [4]: %%timeit
   ...: y = serialize(x)
   ...:
   ...:
1.73 s ± 17.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [5]: x.nbytes / 1.73 / 1e6
Out[5]: 578.0346820809249

In [6]: %%timeit
   ...: y = serialize(serialize(x, serializers=("cuda", "dask", "pickle")))
   ...:
   ...:
1.77 s ± 2.19 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [7]: x.nbytes / 1.77 / 1e6
Out[7]: 564.9717514124293
```

The way we serialize nowadays (in [4]) does roughly 576MB/s, whereas the suggested here does 565MB/s. The new way is slightly slower, and both are _very_ slow. Copying data back to CPU in form of a NumPy array is faster, about 2GB/s:

```python
In [8]: %%timeit
   ...: y = x.get()
   ...:
   ...:
494 ms ± 639 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [9]: x.nbytes / .494 / 1e6
Out[9]: 2024.2914979757086
```

From my experience, we should be able to get about 6GB/s for copying data from GPU to CPU, or 12GB/s when using pinned memory (which I'm not sure we can do with CuPy under all circumstances).

I suspect serialization adds a considerable overhead, but 2GB/s feels still too slow, Numba seems to give us similar performance:

```python
In [1]: import numpy as np

In [2]: from numba import cuda

In [3]: x = np.arange(1.25e8)

In [4]: d_x = cuda.to_device(x)

In [5]: %%timeit
   ...: d_x.copy_to_host()
   ...:
   ...:
514 ms ± 1.41 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [6]: 1e9 / .514 / 1e6
Out[6]: 1945.5252918287938
```

I'm not yet sure how we can improve any of those numbers, but I'll investigate further.
Serializing a NumPy array takes 1.11 s:

```python
In [1]: import numpy as np

In [2]: from distributed.protocol import serialize, deserialize

In [3]: x = np.arange(1.25e8)

In [4]: %%timeit
   ...: y = serialize(x)
   ...:
   ...:
9.26 µs ± 140 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [5]: %%timeit
   ...: y = serialize(serialize(x))
   ...:
   ...:
1.11 s ± 448 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
```

If we split the CuPy copy (from my previous comment, 494 ms) and serialization (1.11 s), we get 1.604 s, almost the same as we do for serializing a CuPy array (1.73 s in my previous comment). Therefore, pickling dominates the serialization. Are there faster ways for pickling data in Python? Since the arrays are basically just a bunch of bytes anyway (assuming we don't need to transpose the array, etc.), I'm a bit surprised that it's so slow.
Should this

```python
y = serialize(serialize(x))
```
Instead be this?

```python
y = deserialize(*serialize(x, serializers=...), serializers=...)
```
 
I would expect to call serialize to get header and frames, and then move those into host memory and write them to disk

```python
header, frames = serialize(x, serializers=["cuda", "dask", "pickle", "msgpack"])
frames = [ensure_on_host(frame) for frame in frames]
```


> `header, frames = serialize(x, serializers=["cuda", "dask", "pickle", "msgpack"])`

This basically just creates a header and wraps the array `x` in a list, but doesn't really serialize anything. The second `serialize` call actually pickles the array, and that's what I wanted to time. In other words, calling `serialize` once, doesn't really do anything, thus it's really fast.
I see.  I suspect that things will be a bit faster if we move each device-side bytes object to a host-side-bytes object independently, rather than bundle them up into a single bytes-like object with serialize/pickle.
The things is that both copying data back to host and pickling are slow. Pickling still dominates, so that should probably be the first one to improve, but the outlook of copying memory back from device is not so good either.
Right, mostly I'm trying to separate overhead from using pickle relative to using whatever numba to/from_device function there is, and also trying to remove the overhead of concatenating all of the results together when we pickle the list.  

Then, once we have a nice number that uses pure numba operations (which it looks like you have above already) we can ask the Numba devs.  Given that this is already the case, maybe we should raise an issue there and ask for help?  I suspect that they have run into this before and can help.
> Right, mostly I'm trying to separate overhead from using pickle relative to using whatever numba to/from_device function there is, and also trying to remove the overhead of concatenating all of the results together when we pickle the list.

The list is a single element, unlike the NumPy implementation, the CuPy one is not creating views of the object to break them down into various blocks. Doing that may help pickling (if that can be done in parallel), but copying will likely perform poorly.

> Then, once we have a nice number that uses pure numba operations (which it looks like you have above already) we can ask the Numba devs. Given that this is already the case, maybe we should raise an issue there and ask for help? I suspect that they have run into this before and can help.

I can do that, but I will first write a small C++ sample to illustrate that and make sure I have numbers for the same machine and compare apples to apples.
@pentschev can I ask you to quickly try an implementation that looks like the following:

```python
def _serialize_if_device(obj):
    """ Serialize an object if it's a device object """
    assert is_device_object(obj)  # this used to be a conditional, under what situations would this not be a device object?
    header, frames = cuda_serialize(obj)
    frames = [numba.DeviceArray(frame).copy_to_host() for frame in frames]
    return header, frames

def _deserialize_if_device(obj):
    header, frames = obj
    frames = [numba.cuda.to_device(frame) for frame in frames]
    return cuda_deserialize(header, frames)
```
    
Previously these functions would check to see if the object was device-like.  Why this condition?  Was this for things like lists and dicts?  If so, lets just implement `cuda_serialize` operations for lists and dicts that concatenate the headers and frames of their inputs as necessary.

Patch:
diff --git a/dask_cuda/device_host_file.py b/dask_cuda/device_host_file.py
index 02c826f7..ac0eff00 100644
--- a/dask_cuda/device_host_file.py
+++ b/dask_cuda/device_host_file.py
@@ -1,28 +1,94 @@
 from zict import Buffer, File, Func
 from zict.common import ZictBase
-from distributed.protocol import deserialize_bytes, serialize_bytes
+from distributed.protocol import (
+    deserialize,
+    deserialize_bytes,
+    serialize,
+    serialize_bytelist,
+    dask_serialize,
+    dask_deserialize,
+)
+from dask.sizeof import sizeof
+from distributed.utils import nbytes
 from distributed.worker import weight
 
-from functools import partial
+from numba import cuda
 import os
 
 from .is_device_object import is_device_object
 
 
-def _serialize_if_device(obj):
-    """ Serialize an object if it's a device object """
-    if is_device_object(obj):
-        return serialize_bytes(obj, on_error="raise")
-    else:
-        return obj
+# Register sizeof for Numba DeviceNDArray while Dask doesn't add it
+if not hasattr(sizeof, "register_numba"):
 
+    @sizeof.register_lazy("numba")
+    def register_numba():
+        import numba
 
-def _deserialize_if_device(obj):
-    """ Deserialize an object if it's an instance of bytes """
-    if isinstance(obj, bytes):
-        return deserialize_bytes(obj)
-    else:
-        return obj
+        @sizeof.register(numba.cuda.cudadrv.devicearray.DeviceNDArray)
+        def sizeof_numba_devicearray(x):
+            return int(x.nbytes)
+
+
+class DeviceSerialized:
+    """ Store device object on the host
+
+    This stores a device-side object as
+
+    1.  A msgpack encodable header
+    2.  A list of objects that are returned by calling
+        `numba.cuda.as_cuda_array(f).copy_to_host()`
+        which are typically NumPy arrays
+    """
+
+    def __init__(self, header, parts, is_cuda):
+        self.header = header
+        self.parts = parts
+        self.is_cuda = is_cuda
+
+    def __sizeof__(self):
+        return sum(map(nbytes, self.parts))
+
+
+@dask_serialize.register(DeviceSerialized)
+def _(obj):
+    headers = []
+    all_frames = []
+    for part in obj.parts:
+        header, frames = serialize(part)
+        header["frame-start-stop"] = [len(all_frames), len(all_frames) + len(frames)]
+        headers.append(header)
+        all_frames.extend(frames)
+
+    header = {"sub-headers": headers, "is-cuda": obj.is_cuda, "main-header": obj.header}
+
+    return header, all_frames
+
+
+@dask_deserialize.register(DeviceSerialized)
+def _(header, frames):
+    parts = []
+    for sub_header in header["sub-headers"]:
+        start, stop = sub_header.pop("frame-start-stop")
+        part = deserialize(sub_header, frames[start:stop])
+        parts.append(part)
+
+    return DeviceSerialized(header["main-header"], parts, header["is-cuda"])
+
+
+def device_to_host(obj: object) -> DeviceSerialized:
+    header, frames = serialize(obj, serializers=["cuda"])
+    is_cuda = [hasattr(f, "__cuda_array_interface__") for f in frames]
+    frames = [
+        cuda.as_cuda_array(f).copy_to_host() if ic else f
+        for ic, f in zip(is_cuda, frames)
+    ]
+    return DeviceSerialized(header, frames, is_cuda)
+
+
+def host_to_device(s: DeviceSerialized) -> object:
+    frames = [cuda.to_device(f) if ic else f for ic, f in zip(s.is_cuda, s.parts)]
+    return deserialize(s.header, frames)
 
 
 class DeviceHostFile(ZictBase):
@@ -55,17 +121,14 @@ def __init__(
         path = os.path.join(local_directory, "storage")
 
         self.host_func = dict()
-        self.disk_func = Func(
-            partial(serialize_bytes, on_error="raise"), deserialize_bytes, File(path)
-        )
+        self.disk_func = Func(serialize_bytelist, deserialize_bytes, File(path))
         self.host_buffer = Buffer(
             self.host_func, self.disk_func, memory_limit, weight=weight
         )
 
+        self.device_keys = set()
         self.device_func = dict()
-        self.device_host_func = Func(
-            _serialize_if_device, _deserialize_if_device, self.host_buffer
-        )
+        self.device_host_func = Func(device_to_host, host_to_device, self.host_buffer)
         self.device_buffer = Buffer(
             self.device_func, self.device_host_func, device_memory_limit, weight=weight
         )
@@ -79,20 +142,18 @@ def __init__(
 
     def __setitem__(self, key, value):
         if is_device_object(value):
+            self.device_keys.add(key)
             self.device_buffer[key] = value
         else:
             self.host_buffer[key] = value
 
     def __getitem__(self, key):
-        if key in self.host_buffer:
-            obj = self.host_buffer[key]
-            del self.host_buffer[key]
-            self.device_buffer[key] = _deserialize_if_device(obj)
-
-        if key in self.device_buffer:
+        if key in self.device_keys:
             return self.device_buffer[key]
+        elif key in self.host_buffer:
+            return self.host_buffer[key]
         else:
-            raise KeyError
+            raise KeyError(key)
 
     def __len__(self):
         return len(self.device_buffer)
@@ -100,5 +161,6 @@ def __len__(self):
     def __iter__(self):
         return iter(self.device_buffer)
 
-    def __delitem__(self, i):
-        del self.device_buffer[i]
+    def __delitem__(self, key):
+        self.device_keys.discard(key)
+        del self.device_buffer[key]


Test patch:
diff --git a/dask_cuda/tests/test_device_host_file.py b/dask_cuda/tests/test_device_host_file.py
index d53cbc3c..4562a00e 100644
--- a/dask_cuda/tests/test_device_host_file.py
+++ b/dask_cuda/tests/test_device_host_file.py
@@ -1,10 +1,11 @@
 import numpy as np
 import cupy
-from dask_cuda.device_host_file import DeviceHostFile
+from dask_cuda.device_host_file import DeviceHostFile, host_to_device, device_to_host
+from distributed.protocol import deserialize_bytes, serialize_bytelist
 from random import randint
+import dask.array as da
 
 import pytest
-from cupy.testing import assert_array_equal
 
 
 @pytest.mark.parametrize("num_host_arrays", [1, 10, 100])
@@ -33,14 +34,15 @@ def test_device_host_file_short(
     full = host + device
     random.shuffle(full)
 
-    for i in full:
-        dhf[i[0]] = i[1]
+    for k, v in full:
+        dhf[k] = v
 
     random.shuffle(full)
 
-    for i in full:
-        assert_array_equal(i[1], dhf[i[0]])
-        del dhf[i[0]]
+    for k, original in full:
+        acquired = dhf[k]
+        da.assert_eq(original, acquired)
+        del dhf[k]
 
     assert set(dhf.device.keys()) == set()
     assert set(dhf.host.keys()) == set()
@@ -94,19 +96,48 @@ def test_device_host_file_step_by_step(tmp_path):
     assert set(dhf.host.keys()) == set(["a2", "b2"])
     assert set(dhf.disk.keys()) == set(["a1", "b1"])
 
-    assert_array_equal(dhf["a1"], a)
+    da.assert_eq(dhf["a1"], a)
     del dhf["a1"]
-    assert_array_equal(dhf["a2"], a)
+    da.assert_eq(dhf["a2"], a)
     del dhf["a2"]
-    assert_array_equal(dhf["b1"], b)
+    da.assert_eq(dhf["b1"], b)
     del dhf["b1"]
-    assert_array_equal(dhf["b2"], b)
+    da.assert_eq(dhf["b2"], b)
     del dhf["b2"]
-    assert_array_equal(dhf["b3"], b)
+    da.assert_eq(dhf["b3"], b)
     del dhf["b3"]
-    assert_array_equal(dhf["b4"], b)
+    da.assert_eq(dhf["b4"], b)
     del dhf["b4"]
 
     assert set(dhf.device.keys()) == set()
     assert set(dhf.host.keys()) == set()
     assert set(dhf.disk.keys()) == set()
+
+
+def test_serialize_cupy():
+    x = cupy.arange(10)
+
+    obj = device_to_host(x)
+    assert all(isinstance(part, (bytes, memoryview, np.ndarray)) for part in obj.parts)
+    btslst = serialize_bytelist(obj)
+
+    bts = deserialize_bytes(b"".join(btslst))
+    y = host_to_device(bts)
+
+    da.assert_eq(x, y)
+
+
+def test_serialize_cudf():
+    cudf = pytest.importorskip("cudf")
+    dd = pytest.importorskip("dask.dataframe")
+
+    df = cudf.DataFrame({"x": [1, 2, 3], "y": [4.0, 5.0, 6.0]})
+
+    obj = device_to_host(df)
+    assert all(isinstance(part, (bytes, memoryview, np.ndarray)) for part in obj.parts)
+    btslst = serialize_bytelist(obj)
+
+    bts = deserialize_bytes(b"".join(btslst))
+    df2 = host_to_device(bts)
+
+    dd.assert_eq(df, df2)
diff --git a/dask_cuda/tests/test_spill.py b/dask_cuda/tests/test_spill.py
index 46315297..9ff2e0ff 100644
--- a/dask_cuda/tests/test_spill.py
+++ b/dask_cuda/tests/test_spill.py
@@ -4,7 +4,7 @@
 from time import sleep
 
 from distributed.metrics import time
-from distributed.utils_test import gen_cluster, loop, gen_test
+from distributed.utils_test import gen_cluster, loop, gen_test  # noqa: F401
 from distributed.worker import Worker
 from distributed import Client, get_worker, wait
 from dask_cuda import LocalCUDACluster
@@ -71,15 +71,17 @@ def delayed_worker_assert(total_size, device_chunk_overhead, serialized_chunk_ov
         {
             "device_memory_limit": 1e9,
             "memory_limit": 4e9,
-            "host_target": 0.6,
-            "host_spill": 0.7,
+            "host_target": 0.0,
+            "host_spill": 0.0,
+            "host_pause": None,
             "spills_to_disk": False,
         },
         {
             "device_memory_limit": 1e9,
-            "memory_limit": 2e9,
-            "host_target": 0.6,
-            "host_spill": 0.6,
+            "memory_limit": 1e9,
+            "host_target": 0.0,
+            "host_spill": 0.0,
+            "host_pause": None,
             "spills_to_disk": True,
         },
     ],
@@ -89,7 +91,7 @@ def test_cupy_device_spill(params):
         client=True,
         nthreads=[("127.0.0.1", 1)],
         Worker=Worker,
-        timeout=300,
+        timeout=60,
         worker_kwargs={
             "memory_limit": params["memory_limit"],
             "data": DeviceHostFile(
@@ -100,25 +102,25 @@ def test_cupy_device_spill(params):
         config={
             "distributed.worker.memory.target": params["host_target"],
             "distributed.worker.memory.spill": params["host_spill"],
+            "distributed.worker.memory.pause": params["host_pause"],
         },
     )
     def test_device_spill(client, scheduler, worker):
         rs = da.random.RandomState(RandomState=cupy.random.RandomState)
         x = rs.random(int(250e6), chunks=10e6)
-        yield wait(x)
 
         xx = x.persist()
         yield wait(xx)
 
         # Allow up to 1024 bytes overhead per chunk serialized
-        yield client.run(worker_assert, x.nbytes, 0, 1024)
+        yield client.run(worker_assert, x.nbytes, 1024, 1024)
 
         y = client.compute(x.sum())
         res = yield y
 
         assert (abs(res / x.size) - 0.5) < 1e-3
 
-        yield client.run(worker_assert, x.nbytes, 0, 1024)
+        yield client.run(worker_assert, x.nbytes, 1024, 1024)
         host_chunks = yield client.run(lambda: len(get_worker().data.host))
         disk_chunks = yield client.run(lambda: len(get_worker().data.disk))
         for hc, dc in zip(host_chunks.values(), disk_chunks.values()):
@@ -137,60 +139,64 @@ def test_device_spill(client, scheduler, worker):
         {
             "device_memory_limit": 1e9,
             "memory_limit": 4e9,
-            "host_target": 0.6,
-            "host_spill": 0.7,
+            "host_target": 0.0,
+            "host_spill": 0.0,
+            "host_pause": None,
             "spills_to_disk": False,
         },
         {
             "device_memory_limit": 1e9,
-            "memory_limit": 2e9,
-            "host_target": 0.6,
-            "host_spill": 0.6,
+            "memory_limit": 1e9,
+            "host_target": 0.0,
+            "host_spill": 0.0,
+            "host_pause": None,
             "spills_to_disk": True,
         },
     ],
 )
 @pytest.mark.asyncio
 async def test_cupy_cluster_device_spill(loop, params):
-    async with LocalCUDACluster(
-        1,
-        scheduler_port=0,
-        processes=True,
-        silence_logs=False,
-        dashboard_address=None,
-        asynchronous=True,
-        death_timeout=300,
-        device_memory_limit=params["device_memory_limit"],
-        memory_limit=params["memory_limit"],
-        memory_target_fraction=params["host_target"],
-        memory_spill_fraction=params["host_spill"],
-    ) as cluster:
-        async with Client(cluster, asynchronous=True) as client:
-
-            rs = da.random.RandomState(RandomState=cupy.random.RandomState)
-            x = rs.random(int(250e6), chunks=10e6)
-            await wait(x)
-
-            xx = x.persist()
-            await wait(xx)
-
-            # Allow up to 1024 bytes overhead per chunk serialized
-            await client.run(worker_assert, x.nbytes, 0, 1024)
-
-            y = client.compute(x.sum())
-            res = await y
-
-            assert (abs(res / x.size) - 0.5) < 1e-3
-
-            await client.run(worker_assert, x.nbytes, 0, 1024)
-            host_chunks = await client.run(lambda: len(get_worker().data.host))
-            disk_chunks = await client.run(lambda: len(get_worker().data.disk))
-            for hc, dc in zip(host_chunks.values(), disk_chunks.values()):
-                if params["spills_to_disk"]:
-                    assert dc > 0
-                else:
-                    assert hc > 0
-                    assert dc == 0
+    with dask.config.set({"distributed.worker.memory.terminate": False}):
+        async with LocalCUDACluster(
+            1,
+            scheduler_port=0,
+            processes=True,
+            silence_logs=False,
+            dashboard_address=None,
+            asynchronous=True,
+            death_timeout=60,
+            device_memory_limit=params["device_memory_limit"],
+            memory_limit=params["memory_limit"],
+            memory_target_fraction=params["host_target"],
+            memory_spill_fraction=params["host_spill"],
+            memory_pause_fraction=params["host_pause"],
+        ) as cluster:
+            async with Client(cluster, asynchronous=True) as client:
+
+                rs = da.random.RandomState(RandomState=cupy.random.RandomState)
+                x = rs.random(int(250e6), chunks=10e6)
+                await wait(x)
+
+                xx = x.persist()
+                await wait(xx)
+
+                # Allow up to 1024 bytes overhead per chunk serialized
+                await client.run(worker_assert, x.nbytes, 1024, 1024)
+
+                y = client.compute(x.sum())
+                res = await y
+
+                assert (abs(res / x.size) - 0.5) < 1e-3
+
+                await client.run(worker_assert, x.nbytes, 1024, 1024)
+                host_chunks = await client.run(lambda: len(get_worker().data.host))
+                disk_chunks = await client.run(lambda: len(get_worker().data.disk))
+                for hc, dc in zip(host_chunks.values(), disk_chunks.values()):
+                    if params["spills_to_disk"]:
+                        assert dc > 0
+                    else:
+                        assert hc > 0
+                        assert dc == 0
 
 
 @pytest.mark.parametrize(
@@ -199,15 +205,17 @@ async def test_cupy_cluster_device_spill(loop, params):
         {
             "device_memory_limit": 1e9,
             "memory_limit": 4e9,
-            "host_target": 0.6,
-            "host_spill": 0.7,
+            "host_target": 0.0,
+            "host_spill": 0.0,
+            "host_pause": None,
             "spills_to_disk": False,
         },
         {
             "device_memory_limit": 1e9,
-            "memory_limit": 2e9,
-            "host_target": 0.6,
-            "host_spill": 0.6,
+            "memory_limit": 1e9,
+            "host_target": 0.0,
+            "host_spill": 0.0,
+            "host_pause": None,
             "spills_to_disk": True,
         },
     ],
@@ -217,7 +225,7 @@ def test_cudf_device_spill(params):
         client=True,
         ncores=[("127.0.0.1", 1)],
         Worker=Worker,
-        timeout=300,
+        timeout=60,
         worker_kwargs={
             "memory_limit": params["memory_limit"],
             "data": DeviceHostFile(
@@ -229,13 +237,19 @@ def test_cudf_device_spill(params):
             "distributed.comm.timeouts.connect": "20s",
             "distributed.worker.memory.target": params["host_target"],
             "distributed.worker.memory.spill": params["host_spill"],
+            "distributed.worker.memory.pause": params["host_pause"],
         },
     )
     def test_device_spill(client, scheduler, worker):
 
-        cdf = dask.datasets.timeseries(
-            dtypes={"x": int, "y": float}, freq="30ms"
-        ).map_partitions(cudf.from_pandas)
+        # There's a known issue with datetime64:
+        # https://github.com/numpy/numpy/issues/4983#issuecomment-441332940
+        # The same error above happens when spilling datetime64 to disk
+        cdf = (
+            dask.datasets.timeseries(dtypes={"x": int, "y": float}, freq="20ms")
+            .reset_index(drop=True)
+            .map_partitions(cudf.from_pandas)
+        )
 
         sizes = yield client.compute(cdf.map_partitions(lambda df: df.__sizeof__()))
         sizes = sizes.tolist()
@@ -271,59 +285,70 @@ def test_device_spill(client, scheduler, worker):
         {
             "device_memory_limit": 1e9,
             "memory_limit": 4e9,
-            "host_target": 0.6,
-            "host_spill": 0.7,
+            "host_target": 0.0,
+            "host_spill": 0.0,
+            "host_pause": None,
             "spills_to_disk": False,
         },
         {
             "device_memory_limit": 1e9,
-            "memory_limit": 2e9,
-            "host_target": 0.6,
-            "host_spill": 0.6,
+            "memory_limit": 1e9,
+            "host_target": 0.0,
+            "host_spill": 0.0,
+            "host_pause": None,
             "spills_to_disk": True,
         },
     ],
 )
 @pytest.mark.asyncio
 async def test_cudf_cluster_device_spill(loop, params):
-    async with LocalCUDACluster(
-        1,
-        device_memory_limit=params["device_memory_limit"],
-        memory_limit=params["memory_limit"],
-        memory_target_fraction=params["host_target"],
-        memory_spill_fraction=params["host_spill"],
-        death_timeout=300,
-        asynchronous=True,
-    ) as cluster:
-        async with Client(cluster, asynchronous=True) as client:
-
-            cdf = dask.datasets.timeseries(
-                dtypes={"x": int, "y": float}, freq="30ms"
-            ).map_partitions(cudf.from_pandas)
-
-            sizes = await client.compute(cdf.map_partitions(lambda df: df.__sizeof__()))
-            sizes = sizes.tolist()
-            nbytes = sum(sizes)
-            part_index_nbytes = (
-                await client.compute(cdf.partitions[0].index)
-            ).__sizeof__()
-
-            cdf2 = cdf.persist()
-            await wait(cdf2)
-
-            del cdf
-
-            await client.run(worker_assert, nbytes, 32, 2048 + part_index_nbytes)
-
-            host_chunks = await client.run(lambda: len(get_worker().data.host))
-            disk_chunks = await client.run(lambda: len(get_worker().data.disk))
-            for hc, dc in zip(host_chunks.values(), disk_chunks.values()):
-                if params["spills_to_disk"]:
-                    assert dc > 0
-                else:
-                    assert hc > 0
-                    assert dc == 0
-
-            del cdf2
-
-            await client.run(delayed_worker_assert, 0, 0, 0)
+    with dask.config.set({"distributed.worker.memory.terminate": False}):
+        async with LocalCUDACluster(
+            1,
+            device_memory_limit=params["device_memory_limit"],
+            memory_limit=params["memory_limit"],
+            memory_target_fraction=params["host_target"],
+            memory_spill_fraction=params["host_spill"],
+            memory_pause_fraction=params["host_pause"],
+            death_timeout=60,
+            asynchronous=True,
+        ) as cluster:
+            async with Client(cluster, asynchronous=True) as client:
+
+                # There's a known issue with datetime64:
+                # https://github.com/numpy/numpy/issues/4983#issuecomment-441332940
+                # The same error above happens when spilling datetime64 to disk
+                cdf = (
+                    dask.datasets.timeseries(dtypes={"x": int, "y": float}, freq="20ms")
+                    .reset_index(drop=True)
+                    .map_partitions(cudf.from_pandas)
+                )
+
+                sizes = await client.compute(
+                    cdf.map_partitions(lambda df: df.__sizeof__())
+                )
+                sizes = sizes.tolist()
+                nbytes = sum(sizes)
+                part_index_nbytes = (
+                    await client.compute(cdf.partitions[0].index)
+                ).__sizeof__()
+
+                cdf2 = cdf.persist()
+                await wait(cdf2)
+
+                del cdf
+
+                await client.run(worker_assert, nbytes, 32, 2048 + part_index_nbytes)
+
+                host_chunks = await client.run(lambda: len(get_worker().data.host))
+                disk_chunks = await client.run(lambda: len(get_worker().data.disk))
+                for hc, dc in zip(host_chunks.values(), disk_chunks.values()):
+                    if params["spills_to_disk"]:
+                        assert dc > 0
+                    else:
+                        assert hc > 0
+                        assert dc == 0
+
+                del cdf2
+
+                await client.run(delayed_worker_assert, 0, 0, 0)