Detect GPU tasks by inspecting inputs/outputs
- 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
Assessment
This issue has not been assessed yet.