Repository: dask/distributed

Issue:
Allowing custom worker `data` mappings access to the worker's `local_directory`
In dask-cuda we implement a spilling mapping very similar to `SpillBuffer` except that it has two LRU caches to handle spilling from device to host memory and host memory to disk. This all works fine _except_ that the on-disk `storage` directory should live under `worker.local_directory` as is done for the builtin `SpillBuffer`

https://github.com/dask/distributed/blob/bec621e40129f97b9130c3e359ab18fb9a08fabf/distributed/worker_memory.py#L125-L129

Consequently, we incorrectly put stuff in `dask-worker-space/storage` rather than `dask-worker-space/worker-XXX/storage` which can cause problems for repeated runs since this is not automatically cleaned up on successful exit by the usual distributed processes.

It would be useful if custom constructors could (optionally?) receive the `worker.local_directory`. 

Minimally something like this:
```patch
diff --git a/distributed/worker_memory.py b/distributed/worker_memory.py
index a5516f32..90984ba5 100644
--- a/distributed/worker_memory.py
+++ b/distributed/worker_memory.py
@@ -42,7 +42,7 @@ from distributed.compatibility import WINDOWS
 from distributed.core import Status
 from distributed.metrics import monotonic
 from distributed.spill import ManualEvictProto, SpillBuffer
-from distributed.utils import log_errors
+from distributed.utils import has_arg, log_errors
 from distributed.utils_perf import ThrottledGC
 
 if TYPE_CHECKING:
@@ -75,6 +75,7 @@ class WorkerMemoryManager:
         data: (
             MutableMapping[str, Any]  # pre-initialised
             | Callable[[], MutableMapping[str, Any]]  # constructor
+            | Callable[[str], MutableMapping[str, Any]]  # constructor
             | tuple[
                 Callable[..., MutableMapping[str, Any]], dict[str, Any]
             ]  # (constructor, kwargs to constructor)
@@ -109,9 +110,21 @@ class WorkerMemoryManager:
         if isinstance(data, MutableMapping):
             self.data = data
         elif callable(data):
-            self.data = data()
+            if has_arg(data, "worker_local_directory"):
+                func = cast(Callable[[str], MutableMapping[str, Any]], data)
+                self.data = func(worker.local_directory)
+            else:
+                func = cast(Callable[[], MutableMapping[str, Any]], data)
+                self.data = func()
         elif isinstance(data, tuple):
-            self.data = data[0](**data[1])
+            func, kwargs = data
+            if not callable(func):
+                raise ValueError("Expecting a callable")
+            if has_arg(func, "worker_local_directory"):
+                args = (worker.local_directory,)
+            else:
+                args = ()
+            self.data = func(*args, **kwargs)
         elif self.memory_limit and (
             self.memory_target_fraction or self.memory_spill_fraction
         ):
```

Sound reasonable? We could work around fully on our end but this would, I think, require less code overall (and we'd avoid using two workspace directories per worker).



Discussion:
No strong preferences on my side. I'm not even sure if there are besides dask-cuda other users for this interface.

I would even be fine if the constructor gets a reference to the worker as long as devs stick to only "public" attributes
@crusaderky you might have an opinion?

Resolving patch:
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -303,7 +303,12 @@ class Worker(BaseWorker, ServerNode):
     scheduler_file: str, optional
     host: str, optional
     data: MutableMapping, type, None
-        The object to use for storage, builds a disk-backed LRU dict by default
+        The object to use for storage, builds a disk-backed LRU dict by default.
+
+        If a callable to construct the storage object is provided, it
+        will receive the worker's attr:``local_directory`` as an
+        argument if the calling signature has an argument named
+        ``worker_local_directory``.
     nthreads: int, optional
     local_directory: str, optional
         Directory where we place local resources
@@ -484,6 +489,8 @@ def __init__(
         data: (
             MutableMapping[str, Any]  # pre-initialised
             | Callable[[], MutableMapping[str, Any]]  # constructor
+            # constructor receiving self.local_directory
+            | Callable[[str], MutableMapping[str, Any]]
             | tuple[
                 Callable[..., MutableMapping[str, Any]], dict[str, Any]
             ]  # (constructor, kwargs to constructor)
diff --git a/distributed/worker_memory.py b/distributed/worker_memory.py
--- a/distributed/worker_memory.py
+++ b/distributed/worker_memory.py
@@ -42,7 +42,7 @@
 from distributed.core import Status
 from distributed.metrics import monotonic
 from distributed.spill import ManualEvictProto, SpillBuffer
-from distributed.utils import log_errors
+from distributed.utils import has_arg, log_errors
 from distributed.utils_perf import ThrottledGC
 
 if TYPE_CHECKING:
@@ -54,6 +54,24 @@
 
 
 class WorkerMemoryManager:
+    """Management of worker memory usage
+
+    Parameters
+    ----------
+    worker
+        Worker to manage
+
+    For meaning of the remaining parameters, see the matching
+    parameter names in :class:`~.distributed.worker.Worker`.
+
+    Notes
+    -----
+
+    If data is a callable and has the argument ``worker_local_directory`` in its
+    signature, it will be filled with the worker's attr:``local_directory``.
+
+    """
+
     data: MutableMapping[str, object]  # {task key: task payload}
     memory_limit: int | None
     memory_target_fraction: float | Literal[False]
@@ -75,6 +93,8 @@ def __init__(
         data: (
             MutableMapping[str, Any]  # pre-initialised
             | Callable[[], MutableMapping[str, Any]]  # constructor
+            # constructor, passed worker.local_directory
+            | Callable[[str], MutableMapping[str, Any]]
             | tuple[
                 Callable[..., MutableMapping[str, Any]], dict[str, Any]
             ]  # (constructor, kwargs to constructor)
@@ -109,9 +129,22 @@ def __init__(
         if isinstance(data, MutableMapping):
             self.data = data
         elif callable(data):
-            self.data = data()
+            if has_arg(data, "worker_local_directory"):
+                data = cast("Callable[[str], MutableMapping[str, Any]]", data)
+                self.data = data(worker.local_directory)
+            else:
+                data = cast("Callable[[], MutableMapping[str, Any]]", data)
+                self.data = data()
         elif isinstance(data, tuple):
-            self.data = data[0](**data[1])
+            func, kwargs = data
+            if not callable(func):
+                raise ValueError("Expecting a callable")
+            if has_arg(func, "worker_local_directory"):
+                self.data = func(
+                    worker_local_directory=worker.local_directory, **kwargs
+                )
+            else:
+                self.data = func(**kwargs)
         elif self.memory_limit and (
             self.memory_target_fraction or self.memory_spill_fraction
         ):


Test patch:
diff --git a/distributed/tests/test_worker_memory.py b/distributed/tests/test_worker_memory.py
--- a/distributed/tests/test_worker_memory.py
+++ b/distributed/tests/test_worker_memory.py
@@ -97,6 +97,46 @@ async def test_dict_data_if_no_spill_to_disk(s, w):
     assert type(w.data) is dict
 
 
+class WorkerData(dict):
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.kwargs = kwargs
+
+
+class WorkerDataLocalDirectory(dict):
+    def __init__(self, worker_local_directory, **kwargs):
+        super().__init__()
+        self.local_directory = worker_local_directory
+        self.kwargs = kwargs
+
+
+@gen_cluster(
+    nthreads=[("", 1)], Worker=Worker, worker_kwargs={"data": WorkerDataLocalDirectory}
+)
+async def test_worker_data_callable_local_directory(s, w):
+    assert type(w.memory_manager.data) is WorkerDataLocalDirectory
+    assert w.memory_manager.data.local_directory == w.local_directory
+
+
+@gen_cluster(
+    nthreads=[("", 1)],
+    Worker=Worker,
+    worker_kwargs={"data": (WorkerDataLocalDirectory, {"a": "b"})},
+)
+async def test_worker_data_callable_local_directory_kwargs(s, w):
+    assert type(w.memory_manager.data) is WorkerDataLocalDirectory
+    assert w.memory_manager.data.local_directory == w.local_directory
+    assert w.memory_manager.data.kwargs == {"a": "b"}
+
+
+@gen_cluster(
+    nthreads=[("", 1)], Worker=Worker, worker_kwargs={"data": (WorkerData, {"a": "b"})}
+)
+async def test_worker_data_callable_kwargs(s, w):
+    assert type(w.memory_manager.data) is WorkerData
+    assert w.memory_manager.data.kwargs == {"a": "b"}
+
+
 class CustomError(Exception):
     pass