From 4354db5990750a40c97787044e49f27ea568af5f Mon Sep 17 00:00:00 2001
From: "Mads R. B. Kristensen" <madsbk@gmail.com>
Date: Tue, 22 Feb 2022 09:03:54 +0100
Subject: [PATCH] Implement ProxifyHostFile.evict()

Closes #861
---
 dask_cuda/proxify_host_file.py | 35 ++++++++++++++++++++++++++--------
 1 file changed, 27 insertions(+), 8 deletions(-)

diff --git a/dask_cuda/proxify_host_file.py b/dask_cuda/proxify_host_file.py
index b3e227e1..9087299a 100644
--- a/dask_cuda/proxify_host_file.py
+++ b/dask_cuda/proxify_host_file.py
@@ -583,25 +583,44 @@ def oom(nbytes: int) -> bool:
                 mr = rmm.mr.FailureCallbackResourceAdaptor(current_mr, oom)
                 rmm.mr.set_current_device_resource(mr)
 
+    def evict(self) -> int:
+        """Manually evict 1% of host limit.
+
+        Dask use this to trigger CPU-to-Disk spilling. We don't know how much
+        we need to spill but Dask will call `evict()` repeatedly until enough
+        is spilled. We ask for 1% each time.
+
+        Return
+        ------
+        nbytes: int
+            Number of bytes spilled or -1 if nothing to spill.
+        """
+
+        ret = self.manager.evict(
+            nbytes=int(self.manager._host_memory_limit * 0.01),
+            proxies_access=self.manager._host.buffer_info,
+            serializer=ProxifyHostFile.serialize_proxy_to_disk_inplace,
+        )
+        gc.collect()
+        return ret if ret > 0 else -1
+
     @property
     def fast(self):
-        """Dask use this to trigger CPU-to-Disk spilling"""
+        """Alternative access to `.evict()` used by Dask
+
+        Dask expect `.fast.evict()` to be availabe for manually triggering
+        of CPU-to-Disk spilling.
+        """
         if len(self.manager._host) == 0:
             return False  # We have nothing in host memory to spill
 
         class EvictDummy:
             @staticmethod
             def evict():
-                # We don't know how much we need to spill but Dask will call evict()
-                # repeatedly until enough is spilled. We ask for 1% each time.
                 ret = (
                     None,
                     None,
-                    self.manager.evict(
-                        nbytes=int(self.manager._host_memory_limit * 0.01),
-                        proxies_access=self.manager._host.buffer_info,
-                        serializer=ProxifyHostFile.serialize_proxy_to_disk_inplace,
-                    ),
+                    self.evict(),
                 )
                 gc.collect()
                 return ret