Multiple worker executors

[madsbk]

This is the first iteration of implementing multiple worker executors as discussed in #4655

• Closes Multiple ThreadPoolExecutors #4655
• Closes [FEATURE] Transmit task annotations to workers #4560
• Closes Allow specifying worker resources on a per thread basis #4598
• Tests added / passed
• Passes black distributed / flake8 distributed / isort distributed

The idea is to send task annotations to the workers and let workers have a dict of executors name -> ThreadPoolExecutor. The default value of this dict is:

        self.executors = {
            "offload": utils._offload_executor,
            "actor": ThreadPoolExecutor(1, thread_name_prefix="Dask-Actor-Threads"),
            "default": ThreadPoolExecutor(self.nthreads, thread_name_prefix="Dask-Default-Threads'"),
        }

If nothing is specified through the executor annotation, the worker will use the "default" executor.

NB: this doesn't control the number of concurrent tasks running on a worker. For now, use the already available resource management.

Example of running CPU tasks in the default ThreadPoolExecutor and GPU tasks in an dedicated ThreadPoolExecutor:

import asyncio
import threading
import dask
from dask.distributed import Client, Scheduler, Worker
from distributed.threadpoolexecutor import ThreadPoolExecutor

def get_thread_name(prefix):
    return prefix + threading.current_thread().name

async def main():
    async with Scheduler() as s:
        async with Worker(
            s.address,
            nthreads=5,
            executor={
                "GPU": ThreadPoolExecutor(1, thread_name_prefix="Dask-GPU-Threads")
            },
            resources={"GPU": 1, "CPU": 4},
        ) as w:
            async with Client(s.address, asynchronous=True) as c:
                with dask.annotate(resources={"CPU": 1}, executor="default"):
                    print(await c.submit(get_thread_name, "CPU-"))
                with dask.annotate(resources={"GPU": 1}, executor="GPU"):
                    print(await c.submit(get_thread_name, "GPU-"))

if __name__ == "__main__":
    asyncio.get_event_loop().run_until_complete(main())

Output

CPU-Dask-Default-Threads'-29802-2
GPU-Dask-GPU-Threads-29802-3

[sjperkins]

I'm excited by the potential introduced by this PR :-)

[mrocklin]

In general this looks good to me. I like the clean generalization.

If you wanted to maintain backwards compatibility, and also maybe slightly reduce the cost of updating tests (which seems minimal anyway) then we might consider adding the following property

class Worker:
    @property
    def executor(self):
        return self.executors["default"]

[sjperkins]

I suspect this also closes #4560

[madsbk]

In general this looks good to me. I like the clean generalization.

If you wanted to maintain backwards compatibility, and also maybe slightly reduce the cost of updating tests (which seems minimal anyway) then we might consider adding the following property

class Worker:
    @property
    def executor(self):
        return self.executors["default"]

I did this initially but decide against it to limit the API maintenance. But of cause, if people are using Worker.executor we should add it as a property. Do you know of anyone else are using it?

[mrocklin]

Do you know of anyone else are using it?

I know of groups that I wouldn't be surprised to learn that they used it. I think that providing the property is a good idea for backwards compatibility. I don't think that we lose anything from it.

[madsbk]

I know of groups that I wouldn't be surprised to learn that they used it. I think that providing the property is a good idea for backwards compatibility. I don't think that we lose anything from it.

Alright, added the property and a test

[mrocklin]

This seems great to me. +1

I don't know enough about the current state of testing to know if the current failure is relevant, but it looks familiar to me, so I suspect that it's fine.

For future work I think that we might want to start thinking about how to expose this through configuration options. Maybe we optionally accept strings instead of Executor objects, and if it is a string we use distributed.utils.import_term. Then users can use things like dask-spec to create a Worker with custom executor objects from the command line.

[mrocklin]

This seems great to me. +1

I don't know enough about the current state of testing to know if the current failure is relevant, but it looks familiar to me, so I suspect that it's fine.

For future work I think that we might want to start thinking about how to expose this through configuration options. Maybe we optionally accept strings instead of Executor objects, and if it is a string we use distributed.utils.import_term. Then users can use things like dask-spec to create a Worker with custom executor objects from the command line.

[madsbk]

@jrbourbeau do you recognize the CI errors? I cannot reproduce them locally :/

[quasiben]

test_thread_time and test_AllProgress_lost_key are both known to be flaky tests: xref #4790

[madsbk]

test_thread_time and test_AllProgress_lost_key are both known to be flaky tests: xref #4790

In that case, this PR is ready for review :)

[quasiben]

I kicked the CI to rerun the tests

[jrbourbeau]

Thanks for your work on this @madsbk! This looks really nice, the diff is much smaller than I was expecting : )

[jrbourbeau]

I'm curious why we're doing this step manually (I realize this predates this PR -- this isn't meant to be a blocker for merging the changes here)

[madsbk]

It is because even with shutdown(wait=False), the ThreadPoolExecutor will still run pending tasks. By calling executor._work_queue.queue.clear(), we cancel them before shutting down.

[jrbourbeau]

Ah, I see. Thanks for clarifying! FWIW it looks like a cancel_futures= keyword was added to Executor.shutdown in Python 3.9. Just pointing this out as something we can use in the future once 3.9 is our minimum Python version

[jrbourbeau]

Should we also pass wait=executor_wait and timeout=timeout keywords here like we do above? For example, I might have a custom ProcessPoolExecutor I'm using and want to have wait=True.

[madsbk]

Added wait=executor_wait but timeout=timeout isn't a valid concurrent.futures.Executor argument. It is introduced in ThreadPoolExecutor.

[jrbourbeau]

Thanks @madsbk! This is in

(the CI failures here are known flaky tests)