dask / dask/distributed

Detect GPU tasks by inspecting inputs/outputs

Open
#4,656 4 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
1.7k
Forks
778
Avg merge
2h 50m
Merged PRs (30d)
3

Description

It would be useful to automagically detect which tasks engaged the GPU. This would allow us to more easily use both the CPU and GPU in mixed workloads, and require less configuration by the user. Unfortunately automatically detecting GPU tasks is hard.

There are a few approaches to this:

1. Years ago I tried achieving this by looking at the serialized form of the task, and looking for text like `b"cudf"` or `b"torch"`. This was surprisingly effective, but also cludgy as heck.
2. Libraries like cudf could annotate layers, this may help less with PyTorch and delayed/futures though
3. Users can handle this themselves with annotations and resource restrictions
3. **New idea!** we could learn this by looking at the inputs and outputs of a function for common protocols like `__cuda_array_interface__` and send that information back to the scheduler

So, the new idea again would be that whenever a task created a result that engaged the `__cuda_array_interface__` protocol we would include that information as we sent it up to the scheduler. Probably this requires a new attribute like `cuda_nbytes` on the `TaskState` (which I'm personally fine with). The scheduler would watch for this signal, and if it occurred it would probably flip a `cuda` flag on the `TaskPrefix`, which would then trigger a signal that got sent down to all of the workers, and maybe pushed that task to run in a different ThreadPoolExecutor (see #4655 )

This would mis-allocate the first few tasks to the CPU Executor, but mostly it would do the right thing, and wouldn't require any intervention from the user

cc @dask/gpu

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.