Automatically assign resources based on serialized form of functions

[mrocklin]

This commit adds functionality to the scheduler to automatically assign
resources to tasks based on terms found in their serialized forms.

This would be useful, for example, to direct all tasks with cupy, cudf,
torch, or tensorflow module names in their serialized form, to be
automatically tagged with a GPU resource tag. This would simplify
deploying GPU systems considerably, and would also allow for nicer mixed
CPU/GPU workloads.

This should maybe be coupled, in the future, with code that
automatically assigns resources based on the presence of GPUs.

In practice, I'm thinking that, at least for GPU computing, we might do something like the following:

resources_auto={
    "GPU": {
        "names": ["cupy", "cuda", "cudf", "cuml", "pytorch", "tensorflow", "keras"], 
        "value": 1}
    }
}

This isn't perfect, but I'll bet that it does a decent job in practice. This is free in the common case that nothing is registered. I'm inclined to try this for a while and possibly remove it without warning in the future.

cc @pentschev

[pentschev]

I'm not familiar with the scheduler code, thus I'm having some difficulty understanding what are the effects that this feature adds in practice, could you clarify that @mrocklin?

[mrocklin]

We can optionally make it so that every task whose serialized form contains the terms cupy, cudf, or similar would be automatically tagged as a GPU task. We would then automatically register workers with GPUs as having that tag, and would restrict execution so that only one GPU task could run at a time.

This would help in two ways:

1. We could have a mix of CPU and GPU workers
2. The GPU enabled workers could have a full thread pool, and could run a bunch of CPU stuff on the side while only running one GPU task at a time.

This uses a rarely used feature documented here: https://distributed.dask.org/en/latest/resources.html

[TomAugspurger]

We would then automatically register workers with GPUs as having that tag

Whose "we" in this case? Do you envision the default resources.auto value ever changing in distributed to include libraries by default?

[mrocklin]

Whose "we" in this case? Do you envision the default resources.auto value ever changing in distributed to include libraries by default?

No. I do not think that this should be the default (it adds a couple microseconds per task) (which is a small, but not-insignificant time).

But it would make it so that this behavior would just be a configuration change. We can easily ship around config file snippets.

I do think that it might make sense to automatically register a GPU=1 resource if we detect a GPU though. This can be done quickly at no cost, and is fairly separable in terms of code maintenance.

[pentschev]

Thanks for clarifying @mrocklin. I can see how this is helpful now and I agree this is something nice to have, in fact, we had already users asking for ways to do this in rapidsai/dask-cuda#108.

My only comment is that we could have some short documentation here to let users know this is supported and make it easily searchable. Apart from that, PR is +1 from me.

Thanks for working on it @mrocklin.

[pentschev]

Actually, I think the documentation is mostly covered in https://distributed.dask.org/en/latest/resources.html already. So perhaps there's nothing that we really need to do here, and only add a comment later on if we add something by default, as per your sample in #3036 (comment).

[pentschev]

Just in case people are still waiting for more feedback, from my side, this is +1 for merging.

[jrbourbeau]

This generally looks good to me, thanks @mrocklin. Could you add a section to https://distributed.dask.org/en/latest/resources.html about this?

[mrocklin]

I would like to wait on merging this until I have some time to play with it in the wild if that's ok.

cc @jacobtomlinson @rjzamora one of you might also find this interesting. The goal here would be to let dask-workers have many threads, but we would identify which tasks were GPU tasks and only run one of them at a time. This would allow us to use the full CPU for things on the CPU and use the GPU efficiently with only one thread.