Repository: dask/distributed

Issue:
Allow specifying worker resources on a per thread basis
Currently specifying worker resources happens on a per-process basis: https://distributed.dask.org/en/latest/resources.html#resources-are-applied-separately-to-each-worker-process

We have use cases where we'd ideally like to specify resources on a per thread basis. I.E. I'd want a worker with say 5 threads, 1 of which has a GPU resource, the other 4 of which have a CPU resource. We could then use those CPU resources to stage data for the GPU resourced thread to process. Ideally these would be in the same process so that the memory doesn't need to be transported across processes.


Discussion:
I guess, you need the same thread to run the GPU tasks, right? Otherwise setting something like `--resources "GPU=1, CPU=4"` should be sufficient. 

I see two approaches: 
 * Run a non-nannied worker per thread, which would make it possible to specify resources per thread. However, if the cluster includes multiple machines, we cannot use the `inproc` protocol between workers since workers doesn’t support the use of multiple protocols such as `tcp` or `ucx`for inter-node communication and `inproc` for local communication. This would be another motivation to implement support of multiple protocols for inter-worker communication. 
 * Implement resource support into [`ThreadPoolExecutor`](https://github.com/dask/distributed/blob/6f95c6323454825f1a64b554ffc7267faeb27de5/distributed/threadpoolexecutor.py#L68) that makes is possible to specify thread resources. This approach is properly the easiest of the two but also the least general. 
Awhile back @crusaderky was contemplating how dask could operate in hybrid cpu-gpu environment in https://github.com/dask/distributed/issues/1226

I also think this issue is better suited for dask/distributed so I'll move this to the distributed GH repo
I think that it might be fun to consider Dask workers that have multiple concurrent.future.Executor types.  This would solve a few problems that we see today:

1. People who want to run tasks in separate processes for better isolation
2. Situations where we have async tasks like IO (cc @martindurant )
3. This situation, where we want to associate and constrain certain resources 

Anyway, I think that thinking about multiple executors might be a possibly more general modification that would give you all what you want, and also engage a few other groups at the same time.  
I would add "actors" to this list, which is the only case that already has its own executor, but probably shouldn't.

For IO, fsspec will no longer (on current main) create a separate thread for the event loop, but will run a loop with `run_until_complete` in the thread where it's called. I'm not sure if that affects the discussion - now the workload will be mixed between IO waiting and any other computing.
> I'd want a worker with say 5 threads, 1 of which has a GPU resource, the other 4 of which have a CPU resource.

Is there a particular need to pin these resources to a specific thread? As of now, if you specify ``--nthreads 5 --resources GPU 1 CPU 4`` you should get exactly what you described but the jobs will be run by arbitrary threads.
> Is there a particular need to pin these resources to a specific thread? As of now, if you specify `--nthreads 5 --resources GPU 1 CPU 4` you should get exactly what you described but the jobs will be run by arbitrary threads.

Sorry yes, I should have clarified that. The use cases I have in mind are related to using the CPU threads to stage data for a GPU thread to asynchronously execute work. If the GPU thread changes, then we'd need to enforce some type of synchronization behavior between the threads for the GPU to do things safely and that gets pretty gross pretty quick.
@kkraus14 could you make an example? Isn't your issue solved by using either the legacy default stream or a process-wide explicit stream? On the other hand, if you're already using multiple streams to allow for pipelining between kernel and RAM transfers, aren't they supposedly thread-safe?

Ref. https://docs.nvidia.com/cuda/cuda-driver-api/stream-sync-behavior.html

@crusaderky Would be hard to write up a code example but let me write up a pseudo-example that I had in mind.

Say I'm using `per-thread default stream` (https://docs.nvidia.com/cuda/cuda-runtime-api/stream-sync-behavior.html) and my library / tasks don't explicitly synchronize the stream. The task would return effectively as soon as the kernel launch is done as opposed to when the work is actually done. Say I then have some type of task graph, we can assume it's a linear chain of tasks for now for simplicity sake, `data (CPU) --> A --> B --> C --> D --> E`, if each task can get scheduled on an arbitrary thread, then task `B` could have a kernel launch on a different thread which could attempt to read the result of task `A` before it's safe to do so. If I'm only doing a single task on the GPU at a time then there's no real benefit from using the `per-thread default stream` vs the `legacy default stream`, but we could imagine a world where I have say a `GPU0:0` resource and a `GPU0:1` resource and want those resources assigned to specific threads in order to handle multiple of the linear task graphs concurrently on a single GPU.
@kkraus14 I understand the use case now, but it feels to me like yours may be the wrong solution to the problem.
Namely, a design where you manually name 2+ swimlanes to allow pipelining in the GPU falls apart when you have to deal with it inside a wrapper - namely, if you want to pipeline different chunks of the same dask.array.

Maybe a better design would be a NEP-18 compliant wrapper around cupy, that also rewraps the outputs (see pint.Quantity) and that keeps track of the streams?
The challenge we face is that our users are often times using multiple libraries which can't share explicitly defined streams (larger problem we're working on solving), where the only way to get them to all target the same stream so we can enqueue kernels properly is to use either the legacy default stream or per-thread default stream.

We aren't explicitly talking about Arrays or DataFrames here, where it could be dealing with more general Dask collections, i.e. I use my CPU threads to do IO and stage data in pinned host memory in some type of wrapper objects that exposes a protocol like `__array_interface__` or python buffer protocol, then have my GPU threads running tasks that could use CuPy, Numba, RAPIDS, PyTorch, Tensorflow, etc. where they can all share data zero-copy between one another via `__cuda_array_interface__`. There's no requirement of them to synchronize before returning their object, so even though Dask thinks a task has finished, the kernel / memcpy / etc. could still be running under the hood, meanwhile Dask could get going on the next Python task which could enqueue more work on the stream.
> multiple libraries which can't share explicitly defined streams

It shouldn't be a problem as long as each library can build its own flavour of Stream python wrapper around an external int*?
> It shouldn't be a problem as long as each library can build its own flavour of Stream python wrapper around an external int*?

Someone recently explored this and it got pretty gross pretty quickly between some of the more "cooperative" libraries: https://gist.github.com/leofang/4a043e5d94b4702d04fde2b9e7dcebbd

When you add in a DL framework things get much messier. But even just with the above I don't think end users should be expected to do this type of ninja coding, and solving the underlying interoperability challenges are going to take a long time. Even if a user is only using a single library doing explicit stream management is non-trivial versus using per-thread default stream in many situations.

Resolving patch:
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -7366,6 +7366,8 @@ def _task_to_msg(state: SchedulerState, ts: TaskState, duration: double = -1) ->
     else:
         msg["task"] = task
 
+    if ts._annotations:
+        msg["annotations"] = ts._annotations
     return msg
 
 
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -1,5 +1,6 @@
 import asyncio
 import bisect
+import concurrent.futures
 import errno
 import heapq
 import logging
@@ -16,7 +17,7 @@
 from functools import partial
 from inspect import isawaitable
 from pickle import PicklingError
-from typing import Iterable
+from typing import Dict, Iterable
 
 from tlz import first, keymap, merge, pluck  # noqa: F401
 from tornado import gen
@@ -28,7 +29,7 @@
 from dask.system import CPU_COUNT
 from dask.utils import format_bytes, funcname
 
-from . import comm, preloading, profile, system
+from . import comm, preloading, profile, system, utils
 from .batched import BatchedSend
 from .comm import connect, get_address_host
 from .comm.addressing import address_from_user_args
@@ -154,6 +155,8 @@ class TaskState:
         serializable (e.g. int, string, list, dict).
     * **nbytes**: ``int``
         The size of a particular piece of data
+    * **annotations**: ``dict``
+        Task annotations
 
     Parameters
     ----------
@@ -188,6 +191,7 @@ def __init__(self, key, runspec=None):
         self.stop_time = None
         self.metadata = {}
         self.nbytes = None
+        self.annotations = None
 
     def __repr__(self):
         return "<Task %r %s>" % (self.key, self.state)
@@ -227,11 +231,9 @@ class Worker(ServerNode):
 
     * **nthreads:** ``int``:
         Number of nthreads used by this worker process
-    * **executor:** ``concurrent.futures.ThreadPoolExecutor``:
-        Executor used to perform computation
-        This can also be the string "offload" in which case this uses the same
-        thread pool used for offloading communications.  This results in the
-        same thread being used for deserialization and computation.
+    * **executors:** ``Dict[str, concurrent.futures.Executor]``:
+        Executors used to perform computation. Always contains the default
+        executor.
     * **local_directory:** ``path``:
         Path on local machine to store temporary files
     * **scheduler:** ``rpc``:
@@ -328,7 +330,15 @@ class Worker(ServerNode):
         Fraction of memory at which we start spilling to disk
     memory_pause_fraction: float
         Fraction of memory at which we stop running new tasks
-    executor: concurrent.futures.Executor
+    executor: concurrent.futures.Executor, dict[str, concurrent.futures.Executor], str
+        The executor(s) to use. Depending on the type, it has the following meanings:
+            - Executor instance: The default executor.
+            - Dict[str, Executor]: mapping names to Executor instances. If the
+              "default" key isn't in the dict, a "default" executor will be created
+              using ``ThreadPoolExecutor(nthreads)``.
+            - Str: The string "offload", which refer to the same thread pool used for
+              offloading communications. This results in the same thread being used
+              for deserialization and computation.
     resources: dict
         Resources that this worker has like ``{'GPU': 2}``
     nanny: str
@@ -630,14 +640,25 @@ def __init__(
         self.actors = {}
         self.loop = loop or IOLoop.current()
         self.reconnect = reconnect
+
+        # Common executors always available
+        self.executors: Dict[str, concurrent.futures.Executor] = {
+            "offload": utils._offload_executor,
+            "actor": ThreadPoolExecutor(1, thread_name_prefix="Dask-Actor-Threads"),
+        }
+
+        # Find the default executor
         if executor == "offload":
-            from distributed.utils import _offload_executor as executor
-        self.executor = executor or ThreadPoolExecutor(
-            self.nthreads, thread_name_prefix="Dask-Worker-Threads'"
-        )
-        self.actor_executor = ThreadPoolExecutor(
-            1, thread_name_prefix="Dask-Actor-Threads"
-        )
+            self.executors["default"] = self.executors["offload"]
+        elif isinstance(executor, dict):
+            self.executors.update(executor)
+        elif executor is not None:
+            self.executors["default"] = executor
+        if "default" not in self.executors:
+            self.executors["default"] = ThreadPoolExecutor(
+                self.nthreads, thread_name_prefix="Dask-Default-Threads"
+            )
+
         self.batched_stream = BatchedSend(interval="2ms", loop=self.loop)
         self.name = name
         self.scheduler_delay = 0
@@ -812,6 +833,10 @@ def local_dir(self):
         )
         return self.local_directory
 
+    @property
+    def executor(self):
+        return self.executors["default"]
+
     async def get_metrics(self):
         out = dict(
             executing=self.executing_count,
@@ -1284,13 +1309,14 @@ async def close(
                 with suppress(TimeoutError):
                     await self.batched_stream.close(timedelta(seconds=timeout))
 
-            self.actor_executor._work_queue.queue.clear()
-            if isinstance(self.executor, ThreadPoolExecutor):
-                self.executor._work_queue.queue.clear()
-                self.executor.shutdown(wait=executor_wait, timeout=timeout)
-            else:
-                self.executor.shutdown(wait=False)
-            self.actor_executor.shutdown(wait=executor_wait, timeout=timeout)
+            for executor in self.executors.values():
+                if executor is utils._offload_executor:
+                    continue  # Never shutdown the offload executor
+                if isinstance(executor, ThreadPoolExecutor):
+                    executor._work_queue.queue.clear()
+                    executor.shutdown(wait=executor_wait, timeout=timeout)
+                else:
+                    executor.shutdown(wait=executor_wait)
 
             self.stop()
             await self.rpc.close()
@@ -1503,6 +1529,7 @@ def add_task(
         duration=None,
         resource_restrictions=None,
         actor=False,
+        annotations=None,
         **kwargs2,
     ):
         try:
@@ -1549,6 +1576,7 @@ def add_task(
             ts.duration = duration
             if resource_restrictions:
                 ts.resource_restrictions = resource_restrictions
+            ts.annotations = annotations
 
             who_has = who_has or {}
 
@@ -2564,7 +2592,7 @@ def executor_submit(self, key, function, args=(), kwargs=None, executor=None):
         callbacks to ensure things run smoothly.  This can get tricky, so we
         pull it off into an separate method.
         """
-        executor = executor or self.executor
+        executor = executor or self.executors["default"]
         job_counter[0] += 1
         # logger.info("%s:%d Starts job %d, %s", self.ip, self.port, i, key)
         kwargs = kwargs or {}
@@ -2660,7 +2688,7 @@ async def actor_execute(
                     self.active_threads,
                     self.active_threads_lock,
                 ),
-                executor=self.actor_executor,
+                executor=self.executors["actor"],
             )
         else:
             result = func(*args, **kwargs)
@@ -2794,8 +2822,17 @@ async def execute(self, key, report=False):
                 if self.digests is not None:
                     self.digests["disk-load-duration"].add(stop - start)
 
+            if ts.annotations is not None and "executor" in ts.annotations:
+                executor = ts.annotations["executor"]
+            else:
+                executor = "default"
+            assert executor in self.executors
+
             logger.debug(
-                "Execute key: %s worker: %s", ts.key, self.address
+                "Execute key: %s worker: %s, executor: %s",
+                ts.key,
+                self.address,
+                executor,
             )  # TODO: comment out?
             assert key == ts.key
             try:
@@ -2812,6 +2849,7 @@ async def execute(self, key, report=False):
                         self.active_threads_lock,
                         self.scheduler_delay,
                     ),
+                    executor=self.executors[executor],
                 )
             except RuntimeError as e:
                 executor_error = e


Test patch:
diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py
--- a/distributed/tests/test_worker.py
+++ b/distributed/tests/test_worker.py
@@ -1997,3 +1997,27 @@ def get_worker_client_id():
 
         with pytest.raises(ValueError):
             default_client()
+
+
+@pytest.mark.asyncio
+async def test_multiple_executors(cleanup):
+    def get_thread_name():
+        return threading.current_thread().name
+
+    async with Scheduler() as s:
+        async with Worker(
+            s.address,
+            nthreads=2,
+            executor={
+                "GPU": ThreadPoolExecutor(1, thread_name_prefix="Dask-GPU-Threads")
+            },
+        ) as w:
+            async with Client(s.address, asynchronous=True) as c:
+                futures = []
+                with dask.annotate(executor="default"):
+                    futures.append(c.submit(get_thread_name, pure=False))
+                with dask.annotate(executor="GPU"):
+                    futures.append(c.submit(get_thread_name, pure=False))
+                default_result, gpu_result = await c.gather(futures)
+                assert "Dask-Default-Threads" in default_result
+                assert "Dask-GPU-Threads" in gpu_result