ManualEvictProto (missing unit tests+docs)

crusaderky · crusaderky · commit a8ac78df960a · 2022-03-09T22:55:43.000Z
diff --git a/distributed/spill.py b/distributed/spill.py
@@ -2,10 +2,10 @@
 
 import logging
 import time
-from collections.abc import Mapping
+from collections.abc import Mapping, Sized
 from contextlib import contextmanager
 from functools import partial
-from typing import Any, Literal, NamedTuple
+from typing import Any, Literal, NamedTuple, Protocol
 
 from packaging.version import parse as parse_version
 
@@ -33,6 +33,36 @@ def __sub__(self, other: SpilledSize) -> SpilledSize:  # type: ignore
         return SpilledSize(self.memory - other.memory, self.disk - other.disk)
 
 
+class ManualEvictProto(Protocol):
+    """Duck-type API that a third-party alternative to SpillBuffer must respect (in
+    addition to MutableMapping) if it wishes to support spilling when the
+    ``distributed.worker.memory.spill`` threshold is surpassed.
+
+    This is public API. At the moment of writing, Dask-CUDA implements this protocol in
+    the ProxifyHostFile class.
+    """
+
+    @property
+    def fast(self) -> Sized | bool:
+        """Access to fast memory. This is normally a MutableMapping, but for the purpose
+        of the manual eviction API it is just tested for emptiness to know if there is
+        anything to evict.
+        """
+        ...  # pragma: nocover
+
+    def evict(self) -> int:
+        """Manually evict a key/value pair from fast to slow memory.
+        Return size of the evicted value in fast memory.
+
+        If the eviction failed for whatever reason, return -1. This method must
+        guarantee that the key/value pair that caused the issue has been retained in
+        fast memory and that the problem has been logged internally.
+
+        This method never raises.
+        """
+        ...  # pragma: nocover
+
+
 class SpillBuffer(zict.Buffer):
     """MutableMapping that automatically spills out dask key/value pairs to disk when
     the total size of the stored data exceeds the target. If max_spill is provided the
@@ -163,11 +193,14 @@ def __setitem__(self, key: str, value: Any) -> None:
             assert key not in self.slow
 
     def evict(self) -> int:
-        """Manually evict the oldest key/value pair, even if target has not been reached.
-        Returns sizeof(value).
+        """Implementation of :meth:`ManualEvictProto.evict`.
+
+        Manually evict the oldest key/value pair, even if target has not been
+        reached. Returns sizeof(value).
         If the eviction failed (value failed to pickle, disk full, or max_spill
         exceeded), return -1; the key/value pair that caused the issue will remain in
-        fast. This method never raises.
+        fast. The exception has been logged internally.
+        This method never raises.
         """
         try:
             with self.handle_errors(None):
diff --git a/distributed/worker_memory.py b/distributed/worker_memory.py
@@ -28,7 +28,7 @@
 from collections.abc import Callable, MutableMapping
 from contextlib import suppress
 from functools import partial
-from typing import TYPE_CHECKING, Any, Container, Literal
+from typing import TYPE_CHECKING, Any, Container, Literal, cast
 
 import psutil
 from tornado.ioloop import PeriodicCallback
@@ -39,7 +39,7 @@
 
 from . import system
 from .core import Status
-from .spill import SpillBuffer
+from .spill import ManualEvictProto, SpillBuffer
 from .utils import log_errors
 from .utils_perf import ThrottledGC
 
@@ -208,8 +208,12 @@ def _maybe_pause_or_unpause(self, worker: Worker, memory: int) -> None:
     async def _maybe_spill(self, worker: Worker, memory: int) -> None:
         if self.memory_spill_fraction is False:
             return
-        if not isinstance(self.data, SpillBuffer):
+
+        # SpillBuffer or a duct-type compatible MutableMapping which offers the
+        # fast property and evict() methods. Dask-CUDA uses this.
+        if not hasattr(self.data, "fast") or not hasattr(self.data, "evict"):
             return
+        data = cast(ManualEvictProto, self.data)
 
         assert self.memory_limit
         frac = memory / self.memory_limit
@@ -231,7 +235,7 @@ async def _maybe_spill(self, worker: Worker, memory: int) -> None:
         count = 0
         need = memory - target
         while memory > target:
-            if not self.data.fast:
+            if not data.fast:
                 logger.warning(
                     "Unmanaged memory use is high. This may indicate a memory leak "
                     "or the memory may not be released to the OS; see "
@@ -242,7 +246,8 @@ async def _maybe_spill(self, worker: Worker, memory: int) -> None:
                     format_bytes(self.memory_limit),
                 )
                 break
-            weight = self.data.evict()
+
+            weight = data.evict()
             if weight == -1:
                 # Failed to evict:
                 # disk full, spill size limit exceeded, or pickle error
