Repository: rapidsai/dask-cuda
Issue category: Performance 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.


Functions to edit:
dask_cuda/device_host_file.py:_serialize_if_device
dask_cuda/device_host_file.py:_deserialize_if_device
dask_cuda/device_host_file.py:DeviceHostFile.__init__
dask_cuda/device_host_file.py:DeviceHostFile.__getitem__
dask_cuda/device_host_file.py:DeviceHostFile.__setitem__
dask_cuda/device_host_file.py:DeviceHostFile.__delitem__
Functions added: dask_cuda/device_host_file.py:register_numba
dask_cuda/device_host_file.py:DeviceSerialized
dask_cuda/device_host_file.py:DeviceSerialized.__init__
dask_cuda/device_host_file.py:DeviceSerialized.__sizeof__
dask_cuda/device_host_file.py:_
dask_cuda/device_host_file.py:device_to_host
dask_cuda/device_host_file.py:host_to_device

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)